diff --git a/.dev_scripts/ci_container_test.sh b/.dev_scripts/ci_container_test.sh index 722e977a4..04cd9d4ac 100644 --- a/.dev_scripts/ci_container_test.sh +++ b/.dev_scripts/ci_container_test.sh @@ -26,7 +26,7 @@ if [ "$MODELSCOPE_SDK_DEBUG" == "True" ]; then pip uninstall autoawq -y pip uninstall lmdeploy -y pip uninstall tensorflow -y - pip install ray==2.48 + pip install ray==2.55.1 pip install optimum # test with install diff --git a/.gitignore b/.gitignore index 6ba7b5381..8c6662abd 100644 --- a/.gitignore +++ b/.gitignore @@ -169,3 +169,8 @@ tests/server/contract/client_api_baseline.json # Redis dump file produced by a local redis-server (test infra), never source. *.rdb + +# Local ST dependency reproducers and acceptance notes. +/cookbook/client/swift_model/README-st.md +/cookbook/client/swift_model/reproduce_st_dependency.py +/cookbook/client/swift_model/requirements-st.txt diff --git a/cookbook/client/async_rl/server_config.yaml b/cookbook/client/async_rl/server_config.yaml index 6e9c46441..fe86dd552 100644 --- a/cookbook/client/async_rl/server_config.yaml +++ b/cookbook/client/async_rl/server_config.yaml @@ -64,6 +64,10 @@ applications: args: backend: transformers model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} + min_replicas: 1 + max_replicas: 1 + adapter_config: + adapter_timeout: 30 max_loras: 8 max_length: 10240 data_plane_url: http://127.0.0.1:8000/api/v1/data-plane @@ -78,8 +82,6 @@ applications: queue_config: rps_limit: 100 tps_limit: 100000 - adapter_config: - adapter_timeout: 30 deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/server/megatron/server_config.yaml b/cookbook/client/server/megatron/server_config.yaml index f10de4754..ebae1bff0 100644 --- a/cookbook/client/server/megatron/server_config.yaml +++ b/cookbook/client/server/megatron/server_config.yaml @@ -1,4 +1,8 @@ -# Twinkle Server Configuration - Megatron Backend +# Twinkle Server Configuration - Megatron LoRA Backend +# Candidate for 8 x L20D (~269 GiB each): 2-4 single-GPU sampler instances +# plus up to 2 two-GPU LoRA training instances (3 handles each, 6 total). +# Validate 27B/32K peak memory, throughput and optimizer resume on L20D +# before production rollout; local smaller-model tests do not prove those limits. # proxy_location: determines where the HTTP proxy runs. # "EveryNode" means each Ray node runs its own proxy (good for multi-node). @@ -26,6 +30,13 @@ persistence: mode: redis redis_url: redis://localhost:6380 +# Independent allocation, startup and confirmed-shutdown budgets. +resources: + allocation_timeout_s: 3 + startup_timeout_s: 1800 + shutdown_timeout_s: 180 + sweep_interval_s: 5 + # Applications: each entry defines a service component deployed on the server applications: @@ -36,7 +47,7 @@ applications: import_path: server # Python module to import args: server_config: - per_token_model_limit: 3 # Maximum number of models (adapters) per token (server-globally enforced) + per_token_model_limit: 3 # Global per-token limit; shared-token clients share these 3 slots supported_models: - Qwen/Qwen3.8-27B @@ -56,30 +67,33 @@ applications: # 3. Sampler Service - Runs inference / sampling using vLLM engine # Used for generating text from the model (e.g., evaluating LoRA results). - # Config: TP=2 x DP=2 on 4 GPUs, ~27GB weights/GPU, ~37GB for KV cache + LoRA + # Each instance has DP=1 and one GPU. Two are preloaded; scale to four. + # Memory usage includes each engine's preallocated KV cache, not just weights. - name: sampler-Qwen3.8-27B route_prefix: /api/v1/sampler/Qwen/Qwen3.8-27B import_path: sampler args: model_id: "ms://Qwen/Qwen3.8-27B" # ModelScope model identifier - nproc_per_node: 4 # Number of GPU processes per node + nproc_per_node: 1 # One GPU per independently allocated instance + min_replicas: 2 + max_replicas: 4 + target_ongoing_requests: 16 + idle_timeout_s: 600 sampler_type: vllm # Inference engine: vllm engine_args: # vLLM engine-specific settings max_model_len: 32768 # Maximum sequence length the engine supports - gpu_memory_utilization: 0.75 # 80% utilization, ~64GB/GPU, leaves buffer for safety + gpu_memory_utilization: 0.75 # 75% engine memory budget; does not free Ray GPU reservations enable_lora: true # Allow loading LoRA adapters during inference max_loras: 5 # Max allowed loras working on vLLM at the same time - max_lora_rank: 32 # Support up to rank 64 LoRA adapters + max_lora_rank: 32 # Maximum supported LoRA rank enable_tower_connector_lora: true device_group: # Logical device group for the sampler name: sampler - gpus_per_worker: 2 - ranks: 4 # GPU rank indices to use + ranks: 1 # Per-instance device count, not fixed physical GPU IDs device_type: cuda device_mesh: device_type: cuda - dp_size: 2 - tp_size: 2 # 2 TP replicas for multi-tenant throughput + dp_size: 1 # Pool instance count is controlled by args.min/max_replicas queue_config: rps_limit: 20 # Max requests per second tps_limit: 131072 # Max tokens per second @@ -97,8 +111,9 @@ applications: TWINKLE_TRUST_REMOTE_CODE: "0" TWINKLE_LONG_POLL_TIMEOUT: "120" - # 2. Model Service - Hosts the base model for training. - # Config: PP=2 x DP=2 on 4 GPUs, ~27GB weights/GPU, comfortable for LoRA training + # 2. Model Service - LoRA-only, independently queued two-GPU instances. + # Megatron instances require at least two devices; PP=2 shards each model. + # min=0 enables cold startup; max=2 limits total training devices to four. - name: models-Qwen3.8-27B route_prefix: /api/v1/model/Qwen/Qwen3.8-27B import_path: model @@ -106,24 +121,28 @@ applications: backend: megatron # Use Megatron-LM backend model_id: "ms://Qwen/Qwen3.8-27B" # ModelScope model identifier max_length: 32768 # model max length - max_loras: 3 # model max loras - nproc_per_node: 4 # Number of GPU processes per node + train_mode: lora + adapter_config: + adapter_timeout: 120 # Seconds without a session heartbeat before release + adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) + max_loras: 3 # Handles per instance, not simultaneous backward calls + min_replicas: 0 # Do not reserve training GPUs at startup + max_replicas: 2 # Up to 2 independent two-GPU instances + idle_timeout_s: 600 # Business idle; heartbeats/polling do not renew it + nproc_per_node: 2 # Two workers per training instance device_group: name: model - ranks: 4 + ranks: 2 device_type: cuda device_mesh: device_type: cuda - dp_size: 2 + dp_size: 1 pp_size: 2 queue_config: rps_limit: 20 # Max requests per second tps_limit: 131072 # Max tokens per second max_input_tokens: 32768 - adapter_config: - adapter_timeout: 120 # Seconds before idle adapter unload - adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/server/megatron/server_config_4b.yaml b/cookbook/client/server/megatron/server_config_4b.yaml index 5b6cd9feb..a0de5caec 100644 --- a/cookbook/client/server/megatron/server_config_4b.yaml +++ b/cookbook/client/server/megatron/server_config_4b.yaml @@ -39,6 +39,8 @@ applications: args: backend: megatron model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier + min_replicas: 1 + max_replicas: 1 max_length: 10240 nproc_per_node: 2 # Number of GPU processes per node device_group: @@ -53,7 +55,7 @@ applications: tps_limit: 100000 # Max tokens per second for a single user max_input_tokens: 60000 # Maximum input tokens per request adapter_config: - adapter_timeout: 30 # Seconds before idle adapter unload + adapter_timeout: 30 # Seconds without a session heartbeat before release adapter_max_lifetime: 36000 # Maximum lifetime of an adapter in seconds (e.g., 10 hours) max_loras: 5 # Maximum number of LoRA adapters per model deployments: diff --git a/cookbook/client/server/transformer/server_config.yaml b/cookbook/client/server/transformer/server_config.yaml index dd6cec311..e9a8bb5e7 100644 --- a/cookbook/client/server/transformer/server_config.yaml +++ b/cookbook/client/server/transformer/server_config.yaml @@ -55,6 +55,8 @@ applications: args: backend: transformers # Model backend: transformers | megatron model_id: "ms://Qwen/Qwen3.5-4B" # ModelScope model identifier + min_replicas: 1 + max_replicas: 1 max_length: 10240 nproc_per_node: 1 # Number of GPU processes per node device_group: @@ -68,7 +70,7 @@ applications: rps_limit: 100 # Max requests per second tps_limit: 100000 # Max tokens per second adapter_config: - adapter_timeout: 30 # Seconds before idle adapter unload + adapter_timeout: 30 # Seconds without a session heartbeat before release deployments: - name: ModelManagement autoscaling_config: diff --git a/cookbook/client/server/transformer/server_config_seq_cls.yaml b/cookbook/client/server/transformer/server_config_seq_cls.yaml new file mode 100644 index 000000000..3db2cd9fd --- /dev/null +++ b/cookbook/client/server/transformer/server_config_seq_cls.yaml @@ -0,0 +1,77 @@ +# Twinkle Server Configuration - Reranker / Sequence Classification (Transformers backend) +# +# Serves the two bases written by cookbook/client/twinkle/prepare_seq_cls_model.py. +# Launch from the repository root so the relative model paths resolve (one GPU per model): +# ray start --head --num-gpus=2 +# python -m twinkle.server launch --config cookbook/client/server/transformer/server_config_seq_cls.yaml +# Then run cookbook/client/twinkle/reranker.py and cookbook/client/twinkle/seq_cls.py. + +http_options: + host: 0.0.0.0 + port: 8000 + +persistence: + mode: memory + +applications: + + # 1. GatewayServer - The central API server + - name: server + route_prefix: /api/v1 + import_path: server + deployments: + - name: GatewayServer + ray_actor_options: + num_cpus: 0.1 + + # 2. Reranker - one relevance score per (query, document) sequence + - name: models-Qwen3.5-4B-reranker + route_prefix: /api/v1/model/Qwen3.5-4B-reranker # clients call client.model('Qwen3.5-4B-reranker') + import_path: model + args: + backend: transformers + model_id: output/Qwen3.5-4B-reranker # prepared with --num-labels 1 + task_type: seq_cls # load the ForSequenceClassification head + num_labels: 1 # must match the prepared head + target_modules: [q_proj, v_proj, score] # LoRA slots, including the score head + min_replicas: 1 + max_replicas: 1 + max_length: 2048 + nproc_per_node: 1 + device_group: + name: model + ranks: 1 + device_type: cuda + device_mesh: + device_type: cuda + dp_size: 1 + deployments: + - name: ModelManagement + ray_actor_options: + num_cpus: 0.1 + + # 3. Sequence classification - three classes + - name: models-Qwen3.5-4B-seq-cls + route_prefix: /api/v1/model/Qwen3.5-4B-seq-cls # clients call client.model('Qwen3.5-4B-seq-cls') + import_path: model + args: + backend: transformers + model_id: output/Qwen3.5-4B-seq-cls # prepared with --num-labels 3 + task_type: seq_cls + num_labels: 3 + target_modules: [q_proj, v_proj, score] + min_replicas: 1 + max_replicas: 1 + max_length: 2048 + nproc_per_node: 1 + device_group: + name: model + ranks: 1 + device_type: cuda + device_mesh: + device_type: cuda + dp_size: 1 + deployments: + - name: ModelManagement + ray_actor_options: + num_cpus: 0.1 diff --git a/cookbook/client/swift_model/sentence_transformer_client.py b/cookbook/client/swift_model/sentence_transformer_client.py new file mode 100644 index 000000000..4ece80392 --- /dev/null +++ b/cookbook/client/swift_model/sentence_transformer_client.py @@ -0,0 +1,86 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""ST LoRA over HTTP: train, read embeddings, save, recreate, resume and release.""" + +import argparse + +import json +import numpy as np +from peft import LoraConfig +from transformers import AutoTokenizer + +from twinkle import init_twinkle_client + +MODEL = "sentence-transformers/all-MiniLM-L6-v2" +TEXTS = [ + "A cat sits on a mat.", + "There is a cat on the mat.", + "A plane is flying.", + "The dog runs in a park.", + "A dog is running outdoors.", + "The stock market closed.", +] + + +def encoded_batch(tokenizer): + rows = [dict(tokenizer(text, truncation=True, max_length=128)) for text in TEXTS] + for index, row in enumerate(rows): + row["labels"] = [int(index % 3 == 0)] + return rows + + +def configure(model, *, lr=1e-3, max_length=128): + model.set_template("Template", max_length=max_length) + model.set_processor("InputProcessor", padding_side="right") + model.set_loss("InfonceLoss", temperature=0.1) + model.set_optimizer("AdamW", lr=lr) + + +def vectors(model, batch): + return np.asarray(model.forward_only(inputs=batch).result["embeddings"], dtype=float) + + +def step(model, batch): + model.forward_backward(inputs=batch, task="embedding") + model.clip_grad_and_step() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:8000") + parser.add_argument("--model", default=MODEL, help="Public model name in the server YAML") + parser.add_argument("--tokenizer", default=MODEL, help="Original ST model or local cache directory") + parser.add_argument("--api-key", default="st-example") + args = parser.parse_args() + batch = encoded_batch(AutoTokenizer.from_pretrained(args.tokenizer)) + client = init_twinkle_client(base_url=args.base_url, api_key=args.api_key) + active = None + try: + active = client.model(args.model) + active.add_adapter_to_model("first", LoraConfig(r=4, target_modules=["query", "value"])) + configure(active) + before = vectors(active, batch) + for _ in range(2): + step(active, batch) + trained = vectors(active, batch) + delta = float(np.max(np.abs(trained - before))) + assert delta > 1e-6, "Training did not change embeddings" + saved = active.save("st-example", save_optimizer=True, consumed_train_samples=2 * len(batch)) + active.remove_adapter() + active = client.model(args.model) + active.add_adapter_to_model("restored", LoraConfig(r=4, target_modules=["query", "value"])) + configure(active) + progress = active.resume_from_checkpoint(saved.twinkle_path) + np.testing.assert_allclose(vectors(active, batch), trained, rtol=0, atol=2e-5) + assert progress["consumed_train_samples"] == 2 * len(batch) + step(active, batch) + print(json.dumps({"checkpoint": saved.twinkle_path, "progress": progress, "embedding_delta": delta})) # noqa: T201 + finally: + try: + if active is not None: + active.remove_adapter() + finally: + client.close() + + +if __name__ == "__main__": + main() diff --git a/cookbook/client/swift_model/sentence_transformer_config.yaml b/cookbook/client/swift_model/sentence_transformer_config.yaml index ad5df8de5..efc28a2e2 100644 --- a/cookbook/client/swift_model/sentence_transformer_config.yaml +++ b/cookbook/client/swift_model/sentence_transformer_config.yaml @@ -1,98 +1,61 @@ -# Twinkle Server Configuration - serving a model family whose runtime object is -# not a plain AutoModel: a sentence-transformers encoder (BAAI/bge-m5) served -# for multi-tenant LoRA fine-tuning through the transformers backend. -# -# The `model_cls` field on the model app is the whole point of this example. -# `model_loader` (see the gemma-4 example next door) lets a factory take over -# config / processor / model construction while twinkle's own model class keeps -# running the show; `model_cls` goes further and names the model class itself -- -# here swift's `MultiLoraSentenceTransformerModel`, which wraps the -# sentence-transformers pipeline (the ST pipeline is an nn.Sequential, not an -# AutoModel, so the generic path cannot build it) and owns the multi-tenant LoRA -# machinery on its HF backbone. -# -# twinkle never hardcodes a swift symbol: it imports whatever spec this file -# names *inside the model actor*, wraps the class in its own tinker/twinkle -# compat mixin, and serves it exactly like a built-in backend. The spec crosses -# the Ray boundary as plain data, so it is gated behind TWINKLE_TRUST_REMOTE_CODE -# (set below) exactly like any other remote-code import. -# -# Constraints (enforced by the config schema, not silently ignored): -# - `model_cls` requires backend=transformers, train_mode=lora and no hybrid. -# - the named class must subclass twinkle's TrainableModel. -# -# Requirements: -# - ms-swift importable in the server env (see the gemma-4 example). -# - sentence-transformers installed in the model actor env. - +# ST multi-tenant LoRA over the on-demand training pool. +# The original ST pipeline remains the base model; checkpoints contain tenant adapters. proxy_location: EveryNode - http_options: - host: 0.0.0.0 + host: 127.0.0.1 port: 8000 - telemetry: enabled: false - otlp_endpoint: http://localhost:4317 - persistence: - mode: file - file_path: /tmp/twinkle_state.json - + mode: memory +resources: + startup_timeout_s: 600 + shutdown_timeout_s: 120 applications: - - # 1. Gateway / API server. - - name: server - route_prefix: /api/v1 - import_path: server - args: - supported_models: - - BAAI/bge-m5 - deployments: - - name: TinkerCompatServer - max_ongoing_requests: 50 - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 128 - ray_actor_options: - num_cpus: 0.1 - - # 2. Training model. The sentence-transformers pipeline runs on one GPU of the - # `model` device group; LoRA slots go on its HF backbone. - - name: models-bge-m5 - route_prefix: /api/v1/model/BAAI/bge-m5 - import_path: model - args: - backend: transformers - model_id: "ms://BAAI/bge-m5" - train_mode: lora - # The swift class that builds and runs the ST pipeline. Remove this line - # and the generic AutoModel path takes over -- which cannot build a - # sentence-transformers pipeline. - model_cls: swift.dev.model.multi_lora_sentence_transformer_model:MultiLoraSentenceTransformerModel - max_loras: 4 - nproc_per_node: 1 - device_group: - name: model - ranks: 1 - device_type: cuda - device_mesh: - device_type: cuda - dp_size: 1 - queue_config: - rps_limit: 100 - tps_limit: 100000 - adapter_config: - adapter_timeout: 30 - deployments: - - name: ModelManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 16 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_TRUST_REMOTE_CODE: "1" +- name: server + route_prefix: /api/v1 + import_path: server + args: + supported_models: + - sentence-transformers/all-MiniLM-L6-v2 + deployments: + - name: GatewayServer + ray_actor_options: + num_cpus: 0.1 +- name: models-minilm + route_prefix: /api/v1/model/sentence-transformers/all-MiniLM-L6-v2 + import_path: model + args: + backend: transformers + model_id: hf://sentence-transformers/all-MiniLM-L6-v2 + train_mode: lora + model_cls: swift.dev.model.multi_lora_sentence_transformer_model:MultiLoraSentenceTransformerModel + gradient_checkpointing: false + mixed_precision: 'no' + min_replicas: 0 + max_replicas: 1 + idle_timeout_s: 60 + max_loras: 2 + max_r: 8 + max_length: 256 + target_modules: + - query + - value + nproc_per_node: 1 + device_group: + name: model + ranks: 1 + device_type: cuda + device_mesh: + device_type: cuda + dp_size: 1 + env_vars: + TWINKLE_TRUST_REMOTE_CODE: '1' + queue_config: + execution_timeout: 300 + adapter_config: + adapter_timeout: 300 + deployments: + - name: ModelManagement + ray_actor_options: + num_cpus: 0.1 diff --git a/cookbook/client/tinker/upload_to_hub.py b/cookbook/client/tinker/upload_to_hub.py index d35822144..96491aff7 100644 --- a/cookbook/client/tinker/upload_to_hub.py +++ b/cookbook/client/tinker/upload_to_hub.py @@ -9,8 +9,6 @@ # Task_Envelope with a request_id immediately, so the HTTP call never times out. # 2. The client's future layer long-polls /twinkle/retrieve_future and blocks # until the upload reaches a terminal state, raising on failure. -# (`upload_to_hub` keeps its `poll_interval` / `async_upload` arguments for -# signature compatibility; both are deprecated and have no effect.) # # Prerequisites: # - Server must be running (see server.py / server_config.yaml) diff --git a/cookbook/client/twinkle/embedding.py b/cookbook/client/twinkle/embedding.py index 1e8eb36dc..a7320f369 100644 --- a/cookbook/client/twinkle/embedding.py +++ b/cookbook/client/twinkle/embedding.py @@ -11,16 +11,16 @@ # -> add_metric('EmbeddingMetric', is_training=True) # -> loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step # -> calculate_metric(is_training=True) +# -> save(...) # # Data format (anchor/positive contrastive pairs): # inputs: [anchor_0, positive_0, anchor_1, positive_1, ...] # labels: [ 1, 0, 1, 0, ...] # (labels==1 marks a new group's anchor; labels==0 marks its positive) # -# The server must be running first (see server.py / server_config.yaml). Only the -# model service is required (no sampler). Works against both the transformers and -# megatron backends: each single-sequence feature carries ``position_ids`` so the -# megatron mrope model gets valid positions (transformers derives them internally). +# The server must be running first, e.g. with +# cookbook/client/server/transformer/server_config.yaml. Only the model service is +# required (no sampler), and the base model is used as is (no extra head). import dotenv import os @@ -88,7 +88,7 @@ def train(): # Step 2: build the client model with a fresh LoRA adapter. model = client.model(MODEL_ID) model.add_adapter_to_model(ADAPTER_NAME, LoraConfig(target_modules='all-linear')) - model.set_template('Qwen3_5Template', model_id=MODEL_ID) + model.set_template('Qwen3_5Template') # Step 3: configure the embedding-training pipeline (ORDER MATTERS). # Pass class names as strings; the server resolves them to core-lib classes. @@ -107,15 +107,20 @@ def train(): # Step 5: training loop. task='embedding' selects the embedding pooling + # InfoNCE path on the server (TransformersEmbeddingPatch/MegatronEmbeddingPatch # is applied and rolled back automatically inside the forward). + # On a megatron server, also pass micro_batch_size=: its + # default micro batch of 2 would leave a single pair (no negatives) per step. for step in range(MAX_STEPS): model.forward_backward(inputs=minibatch, task='embedding') model.clip_grad_and_step(max_grad_norm=1.0) metric = model.calculate_metric(is_training=True) - metric = getattr(metric, 'result', metric) - logger.info(f'Step {step}: {metric}') - - logger.info('Embedding training finished. Healthy training shows pos_sim rising ' - 'and neg_sim stable/decreasing.') + logger.info(f'Step {step}: {metric.result}') + + # Step 6: save the LoRA adapter (and optimizer state, so training can be + # resumed later with model.resume_from_checkpoint(saved.twinkle_path)). + saved = model.save('embedding-final', save_optimizer=True) + logger.info(f'Saved checkpoint: {saved.twinkle_path}') + logger.info('Embedding training finished. Healthy training shows loss falling, ' + 'pos_sim rising and neg_sim stable/decreasing.') if __name__ == '__main__': diff --git a/cookbook/client/twinkle/prepare_seq_cls_model.py b/cookbook/client/twinkle/prepare_seq_cls_model.py new file mode 100644 index 000000000..7844c21bf --- /dev/null +++ b/cookbook/client/twinkle/prepare_seq_cls_model.py @@ -0,0 +1,44 @@ +# Twinkle Client - Prepare a base model with a classification head +# +# Reranker and sequence-classification training put LoRA on a FIXED score head. +# The head must be saved with the base model, so every server worker (including +# one restarted later to resume a checkpoint) loads exactly the same head. +# This script adds a randomly initialized ``num_labels``-wide head to the text +# backbone of a causal LM and saves the result, together with the tokenizer. +# +# Run it once on the server machine, from the repository root: +# python cookbook/client/twinkle/prepare_seq_cls_model.py --num-labels 1 --output-dir output/Qwen3.5-4B-reranker +# python cookbook/client/twinkle/prepare_seq_cls_model.py --num-labels 3 --output-dir output/Qwen3.5-4B-seq-cls +# +# The output directories are the model_id values used in +# cookbook/client/server/transformer/server_config_seq_cls.yaml. + +import argparse +import torch +from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer + +from twinkle.hub import HubOperation + +parser = argparse.ArgumentParser() +parser.add_argument('--model', default='ms://Qwen/Qwen3.5-4B') +parser.add_argument('--num-labels', type=int, required=True, help='1 for a reranker score, N for N classes') +parser.add_argument('--output-dir', required=True) +args = parser.parse_args() + +model_dir = HubOperation.download_model(args.model) +tokenizer = AutoTokenizer.from_pretrained(model_dir) + +# Qwen3.5 is a vision-language checkpoint; the classifier is built on its text part only. +config = AutoConfig.from_pretrained(model_dir) +text_config = getattr(config, 'text_config', config) +text_config.num_labels = args.num_labels +# The classifier scores the last non-padding token; batched inputs need the pad id to find it. +text_config.pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + +torch.manual_seed(42) +# The load report lists the vision weights as UNEXPECTED and score.weight as MISSING: +# both are expected, the vision tower is dropped and the new head is initialized here. +model = AutoModelForSequenceClassification.from_pretrained(model_dir, config=text_config, torch_dtype='auto') +model.save_pretrained(args.output_dir) +tokenizer.save_pretrained(args.output_dir) +print(f'Saved {type(model).__name__} with {args.num_labels} label(s) to {args.output_dir}') # noqa: T201 diff --git a/cookbook/client/twinkle/reranker.py b/cookbook/client/twinkle/reranker.py new file mode 100644 index 000000000..d7fa7098d --- /dev/null +++ b/cookbook/client/twinkle/reranker.py @@ -0,0 +1,100 @@ +# Twinkle Client - Pointwise Reranker Training Example +# +# Trains a reranker that gives every (query, document) pair one relevance score: +# +# add_adapter_to_model(LoRA on the backbone + the 'score' head) +# -> set_processor('InputProcessor') +# -> set_loss('PointwiseRerankerLoss') # binary cross entropy on the score +# -> set_optimizer('Adam', ...) +# -> loop: forward_backward(inputs=batch, task='reranker') + clip_grad_and_step +# -> save(...) +# +# Data format: query and document are joined into ONE text per sample, and +# ``labels`` = [1] for a relevant document or [0] for an irrelevant one (int only). +# +# Before running (from the repository root): +# 1. python cookbook/client/twinkle/prepare_seq_cls_model.py --num-labels 1 --output-dir output/Qwen3.5-4B-reranker +# 2. ray start --head --num-gpus=2 +# 3. python -m twinkle.server launch --config cookbook/client/server/transformer/server_config_seq_cls.yaml + +import dotenv +import os +from peft import LoraConfig + +from twinkle import get_logger +from twinkle import init_twinkle_client +from twinkle.template import Qwen3_5Template + +dotenv.load_dotenv('.env') +logger = get_logger() + +# ========== Configuration ========== +MODEL = 'Qwen3.5-4B-reranker' # the model route in server_config_seq_cls.yaml +TOKENIZER_ID = 'ms://Qwen/Qwen3.5-4B' # the prepared base keeps the original tokenizer +LEARNING_RATE = 1e-4 +MAX_STEPS = 10 + +# Each query has one relevant and one irrelevant document. +SAMPLES = [ + ('What is the capital of France?', 'Paris is the capital city of France.', 1), + ('What is the capital of France?', 'Bananas are rich in potassium.', 0), + ('How do plants make their food?', 'Plants use photosynthesis to produce food from sunlight.', 1), + ('How do plants make their food?', 'The train arrived at the station on time.', 0), + ('Who wrote the play Hamlet?', 'William Shakespeare wrote the tragedy Hamlet.', 1), + ('Who wrote the play Hamlet?', 'Water boils at 100 degrees Celsius at sea level.', 0), + ('What language is spoken in Brazil?', 'The main language spoken in Brazil is Portuguese.', 1), + ('What language is spoken in Brazil?', 'A dog barked at the mailman.', 0), +] + + +def build_batch(tokenizer): + """One feature per (query, document) pair; ``labels`` is its 0/1 relevance.""" + batch = [] + for query, document, relevance in SAMPLES: + input_ids = tokenizer.encode(f'Query: {query}\nDocument: {document}') + batch.append({ + 'input_ids': input_ids, + 'attention_mask': [1] * len(input_ids), + 'labels': [relevance], + }) + return batch + + +def train(): + # Step 1: connect to the running Twinkle server. + client = init_twinkle_client( + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), + ) + + # Step 2: create a LoRA adapter. 'score' is the one-output relevance head: give it + # its own LoRA too, otherwise the head stays at its random initialization. + model = client.model(MODEL) + model.add_adapter_to_model('reranker', LoraConfig(target_modules=['q_proj', 'v_proj', 'score'])) + model.set_template('Template') + + # Step 3: configure processor, loss and optimizer. + model.set_processor('InputProcessor') + model.set_loss('PointwiseRerankerLoss') + model.set_optimizer('Adam', lr=LEARNING_RATE) + + # Step 4: tokenize the pairs locally. + tokenizer = Qwen3_5Template(model_id=TOKENIZER_ID).tokenizer + batch = build_batch(tokenizer) + + # Step 5: training loop. task='reranker' returns one score per pair. + for step in range(MAX_STEPS): + output = model.forward_backward(inputs=batch, task='reranker') + model.clip_grad_and_step(max_grad_norm=1.0) + logger.info(f'Step {step}: loss={output.result["loss"]:.4f}') + + # Step 6: relevant documents should now score higher, then save the adapter. + scores = model.forward_only(inputs=batch, task='reranker', return_logits=True).result['logits'] + for (query, document, relevance), score in zip(SAMPLES, scores): + logger.info(f'relevance={relevance} score={float(score[0]):+.3f} | {query} -> {document}') + saved = model.save('reranker-final', save_optimizer=True) + logger.info(f'Saved checkpoint: {saved.twinkle_path}') + + +if __name__ == '__main__': + train() diff --git a/cookbook/client/twinkle/seq_cls.py b/cookbook/client/twinkle/seq_cls.py new file mode 100644 index 000000000..c01f17444 --- /dev/null +++ b/cookbook/client/twinkle/seq_cls.py @@ -0,0 +1,106 @@ +# Twinkle Client - Sequence Classification Training Example +# +# Trains a three-class text classifier (animal / fruit / vehicle) over HTTP: +# +# add_adapter_to_model(LoRA on the backbone + the 'score' head) +# -> set_processor('InputProcessor') +# -> set_loss('SeqClsLoss', problem_type='single_label_classification', num_labels=3) +# -> set_optimizer('Adam', ...) +# -> loop: forward_backward(inputs=batch, task='seq_cls') + clip_grad_and_step +# -> save(...) +# +# Data format: one text per sample, ``labels`` = [class_id] with an int class_id +# in [0, num_labels). +# +# Before running (from the repository root): +# 1. python cookbook/client/twinkle/prepare_seq_cls_model.py --num-labels 3 --output-dir output/Qwen3.5-4B-seq-cls +# 2. ray start --head --num-gpus=2 +# 3. python -m twinkle.server launch --config cookbook/client/server/transformer/server_config_seq_cls.yaml + +import dotenv +import numpy as np +import os +from peft import LoraConfig + +from twinkle import get_logger +from twinkle import init_twinkle_client +from twinkle.template import Qwen3_5Template + +dotenv.load_dotenv('.env') +logger = get_logger() + +# ========== Configuration ========== +MODEL = 'Qwen3.5-4B-seq-cls' # the model route in server_config_seq_cls.yaml +TOKENIZER_ID = 'ms://Qwen/Qwen3.5-4B' # the prepared base keeps the original tokenizer +NUM_LABELS = 3 # must match num_labels in the server config +LEARNING_RATE = 1e-4 +MAX_STEPS = 10 + +LABELS = ['animal', 'fruit', 'vehicle'] +SAMPLES = [ + ('The cat is sleeping on the sofa.', 'animal'), + ('A dog barked at the mailman.', 'animal'), + ('The horse galloped across the field.', 'animal'), + ('An eagle flew over the mountains.', 'animal'), + ('I ate a crisp red apple.', 'fruit'), + ('Bananas are rich in potassium.', 'fruit'), + ('She squeezed fresh oranges for juice.', 'fruit'), + ('The grapes on the vine are ripe.', 'fruit'), + ('He drove his car to work.', 'vehicle'), + ('The train arrived at the station on time.', 'vehicle'), + ('A bus stopped at the corner.', 'vehicle'), + ('The airplane took off from the runway.', 'vehicle'), +] + + +def build_batch(tokenizer): + """One feature per text; ``labels`` holds the class index of that text.""" + batch = [] + for text, label in SAMPLES: + input_ids = tokenizer.encode(text) + batch.append({ + 'input_ids': input_ids, + 'attention_mask': [1] * len(input_ids), + 'labels': [LABELS.index(label)], + }) + return batch + + +def train(): + # Step 1: connect to the running Twinkle server. + client = init_twinkle_client( + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), + ) + + # Step 2: create a LoRA adapter. 'score' is the classification head: give it its + # own LoRA too, otherwise the head stays at its random initialization. + model = client.model(MODEL) + model.add_adapter_to_model('seq_cls', LoraConfig(target_modules=['q_proj', 'v_proj', 'score'])) + model.set_template('Template') + + # Step 3: configure processor, loss and optimizer. + model.set_processor('InputProcessor') + model.set_loss('SeqClsLoss', problem_type='single_label_classification', num_labels=NUM_LABELS) + model.set_optimizer('Adam', lr=LEARNING_RATE) + + # Step 4: tokenize the samples locally. + tokenizer = Qwen3_5Template(model_id=TOKENIZER_ID).tokenizer + batch = build_batch(tokenizer) + + # Step 5: training loop. task='seq_cls' returns one [num_labels] score row per text. + for step in range(MAX_STEPS): + output = model.forward_backward(inputs=batch, task='seq_cls') + model.clip_grad_and_step(max_grad_norm=1.0) + logger.info(f'Step {step}: loss={output.result["loss"]:.4f}') + + # Step 6: check the predictions, then save the adapter and optimizer state. + logits = model.forward_only(inputs=batch, task='seq_cls', return_logits=True).result['logits'] + for (text, label), row in zip(SAMPLES, logits): + logger.info(f'{text!r}: predicted={LABELS[int(np.argmax(row))]}, expected={label}') + saved = model.save('seq-cls-final', save_optimizer=True) + logger.info(f'Saved checkpoint: {saved.twinkle_path}') + + +if __name__ == '__main__': + train() diff --git a/cookbook/client/twinkle/upload_to_hub.py b/cookbook/client/twinkle/upload_to_hub.py index 7a2bc2fd2..e8e6921cb 100644 --- a/cookbook/client/twinkle/upload_to_hub.py +++ b/cookbook/client/twinkle/upload_to_hub.py @@ -9,8 +9,6 @@ # Task_Envelope with a request_id immediately, so the HTTP call never times out. # 2. The client's future layer long-polls /twinkle/retrieve_future and blocks # until the upload reaches a terminal state, raising on failure. -# (`upload_to_hub` keeps its `poll_interval` / `async_upload` arguments for -# signature compatibility; both are deprecated and have no effect.) # # Prerequisites: # - Server must be running (see server.py / server_config.yaml) diff --git a/docs/source_en/Usage Guide/Embedding-Training.md b/docs/source_en/Usage Guide/Embedding-Training.md index 4198f6991..7d27f40f5 100644 --- a/docs/source_en/Usage Guide/Embedding-Training.md +++ b/docs/source_en/Usage Guide/Embedding-Training.md @@ -117,6 +117,7 @@ Client-side embedding training follows the same call order as the bare library: ``` set_processor('InputProcessor') -> set_loss('InfonceLoss', ...) + -> set_optimizer('Adam', ...) -> add_metric('EmbeddingMetric', is_training=True) -> loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step(...) -> calculate_metric(is_training=True) @@ -129,13 +130,12 @@ The difference from the bare library is that the client passes **class-name stri ```python from peft import LoraConfig from twinkle import init_twinkle_client -from twinkle_client.model import MultiLoraTransformersModel # --- Connect to the running Twinkle server --- -init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') +client = init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') # --- Build the client model with a LoRA adapter for embedding --- -model = MultiLoraTransformersModel(model_id='ms://Qwen/Qwen3.5-4B') +model = client.model('ms://Qwen/Qwen3.5-4B') model.add_adapter_to_model('emb_adapter', LoraConfig(target_modules='all-linear')) model.set_template('Qwen3_5Template') @@ -143,6 +143,7 @@ model.set_template('Qwen3_5Template') # NOTE: pass class names as strings; the server resolves them to core-lib classes. model.set_processor('InputProcessor') model.set_loss('InfonceLoss', temperature=0.07, use_batch=True, hard_negatives=None) +model.set_optimizer('Adam', lr=1e-4) model.add_metric('EmbeddingMetric', is_training=True) # --- Training loop --- @@ -158,7 +159,7 @@ metric = model.calculate_metric(is_training=True) **No manual apply_patch needed.** As long as you pass `task='embedding'` to `forward_backward` (or `forward_only`), the server automatically switches to embedding mode for that single `forward` call and rolls back afterwards — you do **not** need to call `apply_patch(...)` explicitly. So after a `forward_backward(task='embedding')`, calling `forward_only` without `task` on the same adapter still returns normal vocab-dimension `logits`, unaffected. -The only three explicit setup steps on the client side are: `set_processor('InputProcessor')`, `set_loss('InfonceLoss', ...)`, `add_metric('EmbeddingMetric', is_training=True)`. Parameter meanings and recommended values are in "Key Parameters" above; returned metrics are in "Monitoring" below. +After creating the adapter and configuring its template, set the processor, loss and optimizer explicitly. Metrics are optional. See `cookbook/client/twinkle/embedding.py` for a standalone example including the final save; parameter meanings and recommended values are in "Key Parameters" above. --- diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index dd0064aba..ecb5f1730 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -233,9 +233,6 @@ applications: rps_limit: 100 # Max requests per second tps_limit: 10000 # Max tokens per second per user max_input_tokens: 10000 # Maximum input tokens per request - adapter_config: - adapter_timeout: 30 # Idle adapter timeout unload time (seconds) - adapter_max_lifetime: 36000 # Maximum adapter lifetime (seconds) max_loras: 1 # Maximum number of LoRA adapters per model deployments: - name: ModelManagement @@ -328,9 +325,6 @@ The difference from the Megatron backend is only in the `backend` parameter of t device_mesh: device_type: cuda dp_size: 2 - adapter_config: - adapter_timeout: 1800 # Idle adapter timeout unload time (seconds) - adapter_max_lifetime: 36000 deployments: - name: ModelManagement autoscaling_config: @@ -479,3 +473,58 @@ raises. During that window newly arriving tasks fail fast with a `server`/503 er Each persisted future stores its immutable `absolute_deadline` when it is created. Cleanup therefore reaches the same decision regardless of which deployment process holds the cleanup lease. + +## On-demand training and device reclamation + +Training services can load models when a request arrives and release devices when they are no longer needed. Keep the existing `applications` format and configure the `args` of an `import_path: model` entry: + +```yaml +- name: training-qwen + route_prefix: /api/v1/model/Qwen/Qwen3.5-4B + import_path: model + args: + model_id: ms://Qwen/Qwen3.5-4B + backend: transformers + train_mode: lora + nproc_per_node: 2 + device_group: {name: model, ranks: 2, device_type: cuda} + device_mesh: {device_type: cuda, dp_size: 2} + max_loras: 3 + min_replicas: 0 + max_replicas: 2 + idle_timeout_s: 600 +``` + +| Setting | How to use it | +|---------|---------------| +| `device_group`, `device_mesh` | Set the devices and parallel layout for each training instance. Adjust them to fit your model and available devices. Megatron requires at least two devices per instance. | +| `min_replicas` | Training instances to keep loaded at startup. With `0`, the first training request waits for model loading. | +| `max_replicas` | Maximum number of training instances. | +| `max_loras` | Maximum LoRA trainings per instance. Used only for LoRA training. | +| `idle_timeout_s` | Seconds of training inactivity before its handle is released. Heartbeats and result polling do not extend this time. Queued or running training operations are protected. | + +This example allows two instances, each holding three LoRA trainings. Clients are also subject to `server_config.per_token_model_limit`; clients using the same token share that token's quota. + +For full-parameter training, set `train_mode: full` and remove `max_loras`. Each full training uses its own instance. To offer LoRA and full training for the same weights, configure two model applications with different `name` and `route_prefix` values and select the mode in each. Clients choose the model alias from the corresponding route. See [Hybrid LoRA Training](../Hybrid-LoRA-Training.md) for Hybrid LoRA configuration. + +Release models when training finishes; Twinkle clients can call `remove_adapter()`. After the last training releases an instance, empty instances above `min_replicas` shut down. Save checkpoints regularly. If a handle expires or the service restarts, create a new training and restore from a checkpoint. Upgrade the Twinkle SDK together with the server. + +Complete configurations are available in `cookbook/client/server/transformer/server_config.yaml` and `cookbook/client/server/megatron/server_config.yaml`. + +### Automatic sampling scaling + +Set instance counts and the scaling target in the `args` of your `import_path: sampler` entry. For example: + +```yaml +nproc_per_node: 1 +device_group: {name: sampler, ranks: 1, device_type: cuda} +device_mesh: {device_type: cuda, dp_size: 1} +min_replicas: 2 +max_replicas: 4 +target_ongoing_requests: 16 +idle_timeout_s: 600 +``` + +This configuration loads two sampling instances at startup and scales up to four as requests increase. Extra instances shut down after 600 seconds of continuous inactivity, leaving two loaded. `target_ongoing_requests` is the target number of in-flight requests per instance; a smaller value triggers earlier scaling. `device_group` and `device_mesh` determine the devices used by each instance. Allow enough devices for both training and sampling. + +Clients keep using the existing sampling API without selecting an instance. LoRA checkpoints load as requests need them. After a service restart, reapply custom templates and other runtime settings. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Embedding\350\256\255\347\273\203.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Embedding\350\256\255\347\273\203.md" index b3b5d3084..c2321a96c 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Embedding\350\256\255\347\273\203.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Embedding\350\256\255\347\273\203.md" @@ -117,6 +117,7 @@ client 化的 Embedding 训练遵循与裸库一致的调用顺序： ``` set_processor('InputProcessor') -> set_loss('InfonceLoss', ...) + -> set_optimizer('Adam', ...) -> add_metric('EmbeddingMetric', is_training=True) -> loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step(...) -> calculate_metric(is_training=True) @@ -129,13 +130,12 @@ set_processor('InputProcessor') ```python from peft import LoraConfig from twinkle import init_twinkle_client -from twinkle_client.model import MultiLoraTransformersModel # --- Connect to the running Twinkle server --- -init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') +client = init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') # --- Build the client model with a LoRA adapter for embedding --- -model = MultiLoraTransformersModel(model_id='ms://Qwen/Qwen3.5-4B') +model = client.model('ms://Qwen/Qwen3.5-4B') model.add_adapter_to_model('emb_adapter', LoraConfig(target_modules='all-linear')) model.set_template('Qwen3_5Template') @@ -143,6 +143,7 @@ model.set_template('Qwen3_5Template') # NOTE: pass class names as strings; the server resolves them to core-lib classes. model.set_processor('InputProcessor') model.set_loss('InfonceLoss', temperature=0.07, use_batch=True, hard_negatives=None) +model.set_optimizer('Adam', lr=1e-4) model.add_metric('EmbeddingMetric', is_training=True) # --- Training loop --- @@ -158,7 +159,7 @@ metric = model.calculate_metric(is_training=True) **无需手动 apply_patch。** 只要在 `forward_backward`（或 `forward_only`）中传入 `task='embedding'`，服务端就会在这一次 `forward` 期间自动切换到 embedding 模式、并在结束后自动回滚——你**不需要**显式调用 `apply_patch(...)`。因此同一个 adapter 在 `forward_backward(task='embedding')` 之后，再调用不带 `task` 的 `forward_only` 仍会返回正常的词表维度 `logits`，互不影响。 -client 侧真正需要显式配置的只有三步：`set_processor('InputProcessor')`、`set_loss('InfonceLoss', ...)`、`add_metric('EmbeddingMetric', is_training=True)`；各参数含义与取值建议见上文「关键参数」，返回指标见下文「监控指标」。 +创建 adapter 并配置模板后，需要显式设置 processor、loss 和 optimizer；metric 按需添加。可直接运行的完整示例见 `cookbook/client/twinkle/embedding.py`，包含训练后的保存。各参数含义与取值建议见上文「关键参数」。 --- diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 81fcacfe5..7b51eb47e 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -233,9 +233,6 @@ applications: rps_limit: 100 # 最大请求速率（每秒） tps_limit: 10000 # 单用户最大 token 速率（每秒） max_input_tokens: 10000 # 每次请求最大输入 token 数 - adapter_config: - adapter_timeout: 30 # 空闲适配器超时卸载时间（秒） - adapter_max_lifetime: 36000 # 适配器最大生命周期（秒） max_loras: 1 # 每个模型最多加载的 LoRA 数量 deployments: - name: ModelManagement @@ -328,9 +325,6 @@ Transformers 后端与 Megatron 后端的区别仅在 Model 服务的 `backend` device_mesh: device_type: cuda dp_size: 2 - adapter_config: - adapter_timeout: 1800 # 空闲适配器超时卸载时间（秒） - adapter_max_lifetime: 36000 deployments: - name: ModelManagement autoscaling_config: @@ -466,3 +460,65 @@ twinkle-server check-config -c server_config.yaml 任务记录进入终态后，per-replica 准入闸门额外保持关闭的最长时长为 `max(0, Collect_Width × R − T)`：此时记录已是终态，但泄漏的 executor 线程可能仍持有闸门，直到其 `ray.get` 返回或抛出。在该窗口内新到达的任务会以 `server`/503 错误快速失败。 每条持久化 future 在创建时写入不可变的 `absolute_deadline`，因此无论哪个 deployment 进程持有 cleanup lease，清理结果都由任务自身契约决定。 + +## 按需训练与设备回收 + +训练服务可以在收到请求时再加载模型，用完后自动释放设备。继续使用原来的 `applications` 配置即可，在 `import_path: model` 的 `args` 中设置： + +```yaml +- name: training-qwen + route_prefix: /api/v1/model/Qwen/Qwen3.5-4B + import_path: model + args: + model_id: ms://Qwen/Qwen3.5-4B + backend: transformers + train_mode: lora + nproc_per_node: 2 + device_group: {name: model, ranks: 2, device_type: cuda} + device_mesh: {device_type: cuda, dp_size: 2} + max_loras: 3 + min_replicas: 0 + max_replicas: 2 + idle_timeout_s: 600 + adapter_config: + adapter_timeout: 30 + adapter_max_lifetime: 36000 +``` + +| 配置项 | 如何使用 | +|--------|----------| +| `device_group`、`device_mesh` | 配置每个训练实例使用的设备和并行方式，请按模型大小与可用设备调整。Megatron 每实例至少需要 2 张设备。 | +| `min_replicas` | 启动时保留的训练实例数。设为 `0` 时，首次训练需要等待模型加载。 | +| `max_replicas` | 最多启动多少个训练实例。 | +| `max_loras` | 每个实例最多同时容纳多少个 LoRA 训练，仅用于 LoRA 模式。 | +| `idle_timeout_s` | 训练闲置多久后自动释放句柄，单位为秒。心跳和结果轮询不会延长这个时间，排队中或正在执行的训练不会因此中断。 | +| `adapter_config.adapter_timeout` | 客户端多久没有发送会话心跳后释放训练句柄，单位为秒，默认 `1800`。SDK 会自动发送心跳；设为 `null` 可关闭这项检查，但会话删除后句柄仍会失效。 | +| `adapter_config.adapter_max_lifetime` | 训练句柄从创建起最长保留多久，单位为秒，默认 `36000`（10 小时）。设为 `null` 可取消这个上限。 | + +这三个时间分别控制业务闲置、心跳中断和最长寿命，互不替代。已经排队或正在执行的任务会先完成，再回收句柄。长时间训练请定期保存 checkpoint；句柄过期后需要重新创建，并显式恢复 checkpoint。 + +上例最多提供 2 个实例，每个实例容纳 3 个 LoRA 训练。客户端同时还受 `server_config.per_token_model_limit` 限制；多个客户端使用同一个 token 时，共享该 token 的配额。 + +使用全参数训练时，将 `train_mode` 改为 `full`，并删除 `max_loras`。每个全参数训练单独使用一个实例。同一权重需要同时提供 LoRA 和全参数训练时，配置两个 model application，分别设置 `train_mode`，并使用不同的 `name` 和 `route_prefix`。客户端通过对应路由中的模型别名选择服务。Hybrid LoRA 的配置见 [Hybrid LoRA 训练](../Hybrid%20LoRA训练.md)。 + +训练结束后请主动释放模型，Twinkle 客户端可调用 `remove_adapter()`。最后一个训练释放后，超出 `min_replicas` 的空闲实例会关闭。训练期间应定期保存 checkpoint；句柄过期或服务重启后，需要重新创建训练并从 checkpoint 恢复。升级服务端时，也请配套升级 Twinkle SDK。 + +完整配置可参考仓库中的 `cookbook/client/server/transformer/server_config.yaml` 和 `cookbook/client/server/megatron/server_config.yaml`。 + +### 采样服务自动扩缩容 + +在 `import_path: sampler` 的 `args` 中设置实例数量和扩容目标。例如： + +```yaml +nproc_per_node: 1 +device_group: {name: sampler, ranks: 1, device_type: cuda} +device_mesh: {device_type: cuda, dp_size: 1} +min_replicas: 2 +max_replicas: 4 +target_ongoing_requests: 16 +idle_timeout_s: 600 +``` + +这份配置启动时加载 2 个采样实例，请求增多时最多扩到 4 个；额外实例连续空闲 600 秒后关闭，保留 2 个。`target_ongoing_requests` 是每个实例的目标在途请求数，调小后会更早扩容。每个实例的设备用量由 `device_group` 和 `device_mesh` 决定，请为训练和采样预留足够的设备。 + +客户端继续使用原来的采样接口，无需指定实例。采样所需的 LoRA checkpoint 按请求加载；服务重启后，请重新设置自定义 template 等运行时配置。 diff --git a/poetry.lock b/poetry.lock index 773c9b14d..ae07f1af0 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2261,43 +2261,39 @@ files = [ [[package]] name = "huggingface-hub" -version = "0.36.2" +version = "1.29.0" description = "Client library to download and publish models, datasets and other repos on the huggingface.co hub" optional = false -python-versions = ">=3.8.0" +python-versions = ">=3.10.0" groups = ["main"] files = [ - {file = "huggingface_hub-0.36.2-py3-none-any.whl", hash = "sha256:48f0c8eac16145dfce371e9d2d7772854a4f591bcb56c9cf548accf531d54270"}, - {file = "huggingface_hub-0.36.2.tar.gz", hash = "sha256:1934304d2fb224f8afa3b87007d58501acfda9215b334eed53072dd5e815ff7a"}, + {file = "huggingface_hub-1.29.0-py3-none-any.whl", hash = "sha256:b00f7782afc14db4bc6572763810a635bdfbab8623d957bfb553bd18e03852cd"}, + {file = "huggingface_hub-1.29.0.tar.gz", hash = "sha256:6ebb385a581435325cf6d5c5b233d5d4bc91175834d99fd65dae14379b36e9ad"}, ] [package.dependencies] -filelock = "*" +click = ">=8.4.2,<9.0.0" +filelock = ">=3.10.0" fsspec = ">=2023.5.0" -hf-xet = {version = ">=1.1.3,<2.0.0", markers = "platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\""} +hf-xet = {version = ">=1.5.2,<2.0.0", markers = "platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\""} +httpx = ">=0.23.0,<1" packaging = ">=20.9" pyyaml = ">=5.1" -requests = "*" tqdm = ">=4.42.1" -typing-extensions = ">=3.7.4.3" +typing-extensions = ">=4.1.0" [package.extras] -all = ["InquirerPy (==0.3.4)", "Jinja2", "Pillow", "aiohttp", "authlib (>=1.3.2)", "fastapi", "fastapi", "gradio (>=4.0.0)", "httpx", "itsdangerous", "jedi", "libcst (>=1.4.0)", "mypy (==1.15.0) ; python_version >= \"3.9\"", "mypy (>=1.14.1,<1.15.0) ; python_version == \"3.8\"", "numpy", "pytest (>=8.1.1,<8.2.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (<16.0)", "pytest-vcr", "pytest-xdist", "ruff (>=0.9.0)", "soundfile", "ty", "types-PyYAML", "types-requests", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)", "urllib3 (<2.0)"] -cli = ["InquirerPy (==0.3.4)"] -dev = ["InquirerPy (==0.3.4)", "Jinja2", "Pillow", "aiohttp", "authlib (>=1.3.2)", "fastapi", "fastapi", "gradio (>=4.0.0)", "httpx", "itsdangerous", "jedi", "libcst (>=1.4.0)", "mypy (==1.15.0) ; python_version >= \"3.9\"", "mypy (>=1.14.1,<1.15.0) ; python_version == \"3.8\"", "numpy", "pytest (>=8.1.1,<8.2.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (<16.0)", "pytest-vcr", "pytest-xdist", "ruff (>=0.9.0)", "soundfile", "ty", "types-PyYAML", "types-requests", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)", "urllib3 (<2.0)"] +all = ["Jinja2", "Pillow", "authlib (>=1.3.2)", "duckdb", "fastapi", "fastapi", "httpx", "itsdangerous", "jedi", "libcst (>=1.4.0)", "mypy (==1.15.0)", "numpy", "pytest (>=8.4.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (>=16.2)", "pytest-vcr", "pytest-xdist", "ruff (>=0.9.0)", "soundfile", "ty", "types-PyYAML", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)", "urllib3 (<2.0)"] +dev = ["Jinja2", "Pillow", "authlib (>=1.3.2)", "duckdb", "fastapi", "fastapi", "httpx", "itsdangerous", "jedi", "libcst (>=1.4.0)", "mypy (==1.15.0)", "numpy", "pytest (>=8.4.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (>=16.2)", "pytest-vcr", "pytest-xdist", "ruff (>=0.9.0)", "soundfile", "ty", "types-PyYAML", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)", "urllib3 (<2.0)"] fastai = ["fastai (>=2.4)", "fastcore (>=1.3.27)", "toml"] -hf-transfer = ["hf_transfer (>=0.1.4)"] -hf-xet = ["hf-xet (>=1.1.2,<2.0.0)"] -inference = ["aiohttp"] -mcp = ["aiohttp", "mcp (>=1.8.0)", "typer"] +gradio = ["gradio (>=5.0.0)", "requests"] +hf-xet = ["hf-xet (>=1.5.2,<2.0.0)"] +mcp = ["mcp (>=1.8.0,<2.0.0)"] oauth = ["authlib (>=1.3.2)", "fastapi", "httpx", "itsdangerous"] -quality = ["libcst (>=1.4.0)", "mypy (==1.15.0) ; python_version >= \"3.9\"", "mypy (>=1.14.1,<1.15.0) ; python_version == \"3.8\"", "ruff (>=0.9.0)", "ty"] -tensorflow = ["graphviz", "pydot", "tensorflow"] -tensorflow-testing = ["keras (<3.0)", "tensorflow"] -testing = ["InquirerPy (==0.3.4)", "Jinja2", "Pillow", "aiohttp", "authlib (>=1.3.2)", "fastapi", "fastapi", "gradio (>=4.0.0)", "httpx", "itsdangerous", "jedi", "numpy", "pytest (>=8.1.1,<8.2.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (<16.0)", "pytest-vcr", "pytest-xdist", "soundfile", "urllib3 (<2.0)"] +quality = ["libcst (>=1.4.0)", "mypy (==1.15.0)", "ruff (>=0.9.0)", "ty"] +testing = ["Jinja2", "Pillow", "authlib (>=1.3.2)", "duckdb", "fastapi", "fastapi", "httpx", "itsdangerous", "jedi", "numpy", "pytest (>=8.4.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (>=16.2)", "pytest-vcr", "pytest-xdist", "soundfile", "urllib3 (<2.0)"] torch = ["safetensors[torch]", "torch"] -typing = ["types-PyYAML", "types-requests", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)"] - +typing = ["types-PyYAML", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)"] [[package]] name = "hyperframe" version = "6.1.0" @@ -8027,79 +8023,58 @@ transformer_engine_cu12 = "2.12.0" [[package]] name = "transformers" -version = "4.57.6" -description = "State-of-the-art Machine Learning for JAX, PyTorch and TensorFlow" +version = "5.8.1" +description = "Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training." optional = false -python-versions = ">=3.9.0" +python-versions = ">=3.10.0" groups = ["main"] files = [ - {file = "transformers-4.57.6-py3-none-any.whl", hash = "sha256:4c9e9de11333ddfe5114bc872c9f370509198acf0b87a832a0ab9458e2bd0550"}, - {file = "transformers-4.57.6.tar.gz", hash = "sha256:55e44126ece9dc0a291521b7e5492b572e6ef2766338a610b9ab5afbb70689d3"}, + {file = "transformers-5.8.1-py3-none-any.whl", hash = "sha256:5340fb95962162cdfdae5cc91d7f8fedd92ed75216c1154c5e1f590fcf56dd0e"}, + {file = "transformers-5.8.1.tar.gz", hash = "sha256:4dd5b6de4105725104d84fd6abd74b305f4debfc251b38c648ee5dd087cf543b"}, ] [package.dependencies] -filelock = "*" -huggingface-hub = ">=0.34.0,<1.0" +huggingface-hub = ">=1.5.0,<2.0" numpy = ">=1.17" packaging = ">=20.0" pyyaml = ">=5.1" -regex = "!=2019.12.17" -requests = "*" +regex = ">=2025.10.22" safetensors = ">=0.4.3" tokenizers = ">=0.22.0,<=0.23.0" tqdm = ">=4.27" +typer = "*" [package.extras] -accelerate = ["accelerate (>=0.26.0)"] -all = ["Pillow (>=10.0.1,<=15.0)", "Pillow (>=10.0.1,<=15.0)", "accelerate (>=0.26.0)", "accelerate (>=0.26.0)", "av", "codecarbon (>=2.8.1)", "flax (>=0.4.1,<=0.7.0)", "jax (>=0.4.1,<=0.4.13)", "jaxlib (>=0.4.1,<=0.4.13)", "jinja2 (>=3.1.0)", "kenlm", "keras-nlp (>=0.3.1,<0.14.0)", "kernels (>=0.6.1,<=0.9)", "librosa", "mistral-common[opencv] (>=1.6.3)", "num2words", "onnxconverter-common", "optax (>=0.0.8,<=0.1.4)", "optuna", "phonemizer", "protobuf", "pyctcdecode (>=0.4.0)", "ray[tune] (>=2.7.0)", "scipy (<1.13.0)", "sentencepiece (>=0.1.91,!=0.1.92)", "tensorflow (>2.9,<2.16)", "tensorflow-text (<2.16)", "tf2onnx", "timm (!=1.0.18,<=1.0.19)", "tokenizers (>=0.22.0,<=0.23.0)", "torch (>=2.2)", "torchaudio", "torchvision"] -audio = ["kenlm", "librosa", "phonemizer", "pyctcdecode (>=0.4.0)"] +accelerate = ["accelerate (>=1.1.0)"] +all = ["Pillow (>=10.0.1,<=15.0)", "accelerate (>=1.1.0)", "av", "blobfile", "jinja2 (>=3.1.0)", "jmespath (>=1.0.1)", "kernels (>=0.12.0,<0.13)", "librosa", "mistral-common[image] (>=1.10.0)", "num2words", "phonemizer", "protobuf", "pyctcdecode (>=0.4.0)", "sentencepiece (>=0.1.91,!=0.1.92)", "tiktoken", "timm (>=1.0.23)", "torch (>=2.4)", "torchaudio", "torchvision"] +audio = ["librosa", "phonemizer", "pyctcdecode (>=0.4.0)", "torchaudio"] benchmark = ["optimum-benchmark (>=0.3.0)"] -chat-template = ["jinja2 (>=3.1.0)"] +chat-template = ["jinja2 (>=3.1.0)", "jmespath (>=1.0.1)"] codecarbon = ["codecarbon (>=2.8.1)"] -deepspeed = ["accelerate (>=0.26.0)", "deepspeed (>=0.9.3)"] -deepspeed-testing = ["GitPython (<3.1.19)", "accelerate (>=0.26.0)", "accelerate (>=0.26.0)", "beautifulsoup4", "cookiecutter (==1.7.3)", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "deepspeed (>=0.9.3)", "dill (<0.3.5)", "evaluate (>=0.2.0)", "faiss-cpu", "fastapi", "libcst", "mistral-common[opencv] (>=1.6.3)", "nltk (<=3.8.1)", "openai (>=1.98.0)", "optuna", "parameterized (>=0.9)", "protobuf", "psutil", "pydantic (>=2)", "pydantic (>=2)", "pytest (>=7.2.0)", "pytest-asyncio", "pytest-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.13.1)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "sentencepiece (>=0.1.91,!=0.1.92)", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "tensorboard", "timeout-decorator", "torch (>=2.2)", "uvicorn"] -dev = ["GitPython (<3.1.19)", "GitPython (<3.1.19)", "Pillow (>=10.0.1,<=15.0)", "Pillow (>=10.0.1,<=15.0)", "accelerate (>=0.26.0)", "accelerate (>=0.26.0)", "accelerate (>=0.26.0)", "av", "beautifulsoup4", "codecarbon (>=2.8.1)", "cookiecutter (==1.7.3)", "cookiecutter (==1.7.3)", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "dill (<0.3.5)", "evaluate (>=0.2.0)", "faiss-cpu", "fastapi", "flax (>=0.4.1,<=0.7.0)", "fugashi (>=1.0)", "ipadic (>=1.0.0,<2.0)", "jax (>=0.4.1,<=0.4.13)", "jaxlib (>=0.4.1,<=0.4.13)", "jinja2 (>=3.1.0)", "kenlm", "keras-nlp (>=0.3.1,<0.14.0)", "kernels (>=0.6.1,<=0.9)", "libcst", "libcst", "librosa", "mistral-common[opencv] (>=1.6.3)", "mistral-common[opencv] (>=1.6.3)", "nltk (<=3.8.1)", "num2words", "onnxconverter-common", "openai (>=1.98.0)", "optax (>=0.0.8,<=0.1.4)", "optuna", "pandas (<2.3.0)", "parameterized (>=0.9)", "phonemizer", "protobuf", "psutil", "pyctcdecode (>=0.4.0)", "pydantic (>=2)", "pydantic (>=2)", "pytest (>=7.2.0)", "pytest-asyncio", "pytest-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "ray[tune] (>=2.7.0)", "rhoknp (>=1.1.0,<1.3.1)", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.13.1)", "ruff (==0.13.1)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "scikit-learn", "scipy (<1.13.0)", "sentencepiece (>=0.1.91,!=0.1.92)", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "sudachidict_core (>=20220729)", "sudachipy (>=0.6.6)", "tensorboard", "tensorflow (>2.9,<2.16)", "tensorflow-text (<2.16)", "tf2onnx", "timeout-decorator", "timm (!=1.0.18,<=1.0.19)", "tokenizers (>=0.22.0,<=0.23.0)", "torch (>=2.2)", "torch (>=2.2)", "torchaudio", "torchvision", "unidic (>=1.0.2)", "unidic_lite (>=1.0.7)", "urllib3 (<2.0.0)", "uvicorn"] -dev-tensorflow = ["GitPython (<3.1.19)", "GitPython (<3.1.19)", "Pillow (>=10.0.1,<=15.0)", "accelerate (>=0.26.0)", "beautifulsoup4", "cookiecutter (==1.7.3)", "cookiecutter (==1.7.3)", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "dill (<0.3.5)", "evaluate (>=0.2.0)", "faiss-cpu", "fastapi", "kenlm", "keras-nlp (>=0.3.1,<0.14.0)", "libcst", "libcst", "librosa", "mistral-common[opencv] (>=1.6.3)", "nltk (<=3.8.1)", "onnxconverter-common", "onnxconverter-common", "onnxruntime (>=1.4.0)", "onnxruntime-tools (>=1.4.2)", "openai (>=1.98.0)", "pandas (<2.3.0)", "parameterized (>=0.9)", "phonemizer", "protobuf", "psutil", "pyctcdecode (>=0.4.0)", "pydantic (>=2)", "pydantic (>=2)", "pytest (>=7.2.0)", "pytest-asyncio", "pytest-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.13.1)", "ruff (==0.13.1)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "scikit-learn", "sentencepiece (>=0.1.91,!=0.1.92)", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "tensorboard", "tensorflow (>2.9,<2.16)", "tensorflow-text (<2.16)", "tf2onnx", "tf2onnx", "timeout-decorator", "tokenizers (>=0.22.0,<=0.23.0)", "torch (>=2.2)", "urllib3 (<2.0.0)", "uvicorn"] -dev-torch = ["GitPython (<3.1.19)", "GitPython (<3.1.19)", "Pillow (>=10.0.1,<=15.0)", "Pillow (>=10.0.1,<=15.0)", "accelerate (>=0.26.0)", "accelerate (>=0.26.0)", "beautifulsoup4", "codecarbon (>=2.8.1)", "cookiecutter (==1.7.3)", "cookiecutter (==1.7.3)", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "dill (<0.3.5)", "evaluate (>=0.2.0)", "faiss-cpu", "fastapi", "fugashi (>=1.0)", "ipadic (>=1.0.0,<2.0)", "kenlm", "kernels (>=0.6.1,<=0.9)", "libcst", "libcst", "librosa", "mistral-common[opencv] (>=1.6.3)", "nltk (<=3.8.1)", "num2words", "onnxruntime (>=1.4.0)", "onnxruntime-tools (>=1.4.2)", "openai (>=1.98.0)", "optuna", "pandas (<2.3.0)", "parameterized (>=0.9)", "phonemizer", "protobuf", "psutil", "pyctcdecode (>=0.4.0)", "pydantic (>=2)", "pydantic (>=2)", "pytest (>=7.2.0)", "pytest-asyncio", "pytest-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "ray[tune] (>=2.7.0)", "rhoknp (>=1.1.0,<1.3.1)", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.13.1)", "ruff (==0.13.1)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "scikit-learn", "sentencepiece (>=0.1.91,!=0.1.92)", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "sudachidict_core (>=20220729)", "sudachipy (>=0.6.6)", "tensorboard", "timeout-decorator", "timm (!=1.0.18,<=1.0.19)", "tokenizers (>=0.22.0,<=0.23.0)", "torch (>=2.2)", "torch (>=2.2)", "torchaudio", "torchvision", "unidic (>=1.0.2)", "unidic_lite (>=1.0.7)", "urllib3 (<2.0.0)", "uvicorn"] -flax = ["flax (>=0.4.1,<=0.7.0)", "jax (>=0.4.1,<=0.4.13)", "jaxlib (>=0.4.1,<=0.4.13)", "optax (>=0.0.8,<=0.1.4)", "scipy (<1.13.0)"] -flax-speech = ["kenlm", "librosa", "phonemizer", "pyctcdecode (>=0.4.0)"] -ftfy = ["ftfy"] -hf-xet = ["hf_xet"] -hub-kernels = ["kernels (>=0.6.1,<=0.9)"] -integrations = ["kernels (>=0.6.1,<=0.9)", "optuna", "ray[tune] (>=2.7.0)"] +deepspeed = ["accelerate (>=1.1.0)", "deepspeed (>=0.9.3)"] +deepspeed-testing = ["GitPython (<3.1.19)", "accelerate (>=1.1.0)", "accelerate (>=1.1.0)", "beautifulsoup4", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "deepspeed (>=0.9.3)", "dill (<0.3.5)", "evaluate (>=0.4.6)", "faiss-cpu", "fastapi", "filelock", "hf-doc-builder", "libcst", "mistral-common[image] (>=1.10.0)", "nltk (<=3.8.1)", "openai (>=1.98.0)", "optuna", "parameterized (>=0.9)", "protobuf", "protobuf", "psutil", "pydantic (>=2)", "pytest (>=7.2.0,<9.0.0)", "pytest-asyncio (>=1.2.0)", "pytest-env", "pytest-order", "pytest-random-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rich", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.14.10)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "sentencepiece (>=0.1.91,!=0.1.92)", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "tensorboard", "timeout-decorator", "tomli", "torch (>=2.4)", "transformers-mlinter (==0.1.1)", "ty (==0.0.20)", "urllib3 (<2.0.0)", "uvicorn"] +dev = ["GitPython (<3.1.19)", "Pillow (>=10.0.1,<=15.0)", "accelerate (>=1.1.0)", "accelerate (>=1.1.0)", "av", "beautifulsoup4", "blobfile", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "dill (<0.3.5)", "evaluate (>=0.4.6)", "faiss-cpu", "fastapi", "filelock", "fugashi (>=1.0)", "hf-doc-builder", "ipadic (>=1.0.0,<2.0)", "jinja2 (>=3.1.0)", "jmespath (>=1.0.1)", "kernels (>=0.12.0,<0.13)", "libcst", "librosa", "mistral-common[image] (>=1.10.0)", "mistral-common[image] (>=1.10.0)", "nltk (<=3.8.1)", "num2words", "openai (>=1.98.0)", "parameterized (>=0.9)", "phonemizer", "protobuf", "protobuf", "psutil", "pyctcdecode (>=0.4.0)", "pydantic (>=2)", "pytest (>=7.2.0,<9.0.0)", "pytest-asyncio (>=1.2.0)", "pytest-env", "pytest-order", "pytest-random-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rhoknp (>=1.1.0,<1.3.1)", "rich", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.14.10)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "scikit-learn", "sentencepiece (>=0.1.91,!=0.1.92)", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "sudachidict_core (>=20220729)", "sudachipy (>=0.6.6)", "tensorboard", "tiktoken", "timeout-decorator", "timm (>=1.0.23)", "tomli", "torch (>=2.4)", "torch (>=2.4)", "torchaudio", "torchvision", "transformers-mlinter (==0.1.1)", "ty (==0.0.20)", "unidic (>=1.0.2)", "unidic_lite (>=1.0.7)", "urllib3 (<2.0.0)", "uvicorn"] +docs = ["hf-doc-builder"] +integrations = ["codecarbon (>=2.8.1)", "kernels (>=0.12.0,<0.13)", "optuna", "ray[tune] (>=2.7.0)"] ja = ["fugashi (>=1.0)", "ipadic (>=1.0.0,<2.0)", "rhoknp (>=1.1.0,<1.3.1)", "sudachidict_core (>=20220729)", "sudachipy (>=0.6.6)", "unidic (>=1.0.2)", "unidic_lite (>=1.0.7)"] -mistral-common = ["mistral-common[opencv] (>=1.6.3)"] -modelcreation = ["cookiecutter (==1.7.3)"] -natten = ["natten (>=0.14.6,<0.15.0)"] +kernels = ["kernels (>=0.12.0,<0.13)"] +mistral-common = ["mistral-common[image] (>=1.10.0)"] num2words = ["num2words"] -onnx = ["onnxconverter-common", "onnxruntime (>=1.4.0)", "onnxruntime-tools (>=1.4.2)", "tf2onnx"] -onnxruntime = ["onnxruntime (>=1.4.0)", "onnxruntime-tools (>=1.4.2)"] open-telemetry = ["opentelemetry-api", "opentelemetry-exporter-otlp", "opentelemetry-sdk"] optuna = ["optuna"] -quality = ["GitPython (<3.1.19)", "datasets (>=2.15.0)", "libcst", "pandas (<2.3.0)", "rich", "ruff (==0.13.1)", "urllib3 (<2.0.0)"] +quality = ["GitPython (<3.1.19)", "datasets (>=2.15.0)", "libcst", "rich", "ruff (==0.14.10)", "tomli", "transformers-mlinter (==0.1.1)", "ty (==0.0.20)", "urllib3 (<2.0.0)"] ray = ["ray[tune] (>=2.7.0)"] retrieval = ["datasets (>=2.15.0)", "faiss-cpu"] -ruff = ["ruff (==0.13.1)"] sagemaker = ["sagemaker (>=2.31.0)"] sentencepiece = ["protobuf", "sentencepiece (>=0.1.91,!=0.1.92)"] -serving = ["accelerate (>=0.26.0)", "fastapi", "openai (>=1.98.0)", "pydantic (>=2)", "starlette", "torch (>=2.2)", "uvicorn"] -sigopt = ["sigopt"] +serving = ["accelerate (>=1.1.0)", "fastapi", "openai (>=1.98.0)", "pydantic (>=2)", "rich", "starlette", "torch (>=2.4)", "uvicorn"] sklearn = ["scikit-learn"] -speech = ["kenlm", "librosa", "phonemizer", "pyctcdecode (>=0.4.0)", "torchaudio"] -testing = ["GitPython (<3.1.19)", "accelerate (>=0.26.0)", "beautifulsoup4", "cookiecutter (==1.7.3)", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "dill (<0.3.5)", "evaluate (>=0.2.0)", "faiss-cpu", "fastapi", "libcst", "mistral-common[opencv] (>=1.6.3)", "nltk (<=3.8.1)", "openai (>=1.98.0)", "parameterized (>=0.9)", "psutil", "pydantic (>=2)", "pydantic (>=2)", "pytest (>=7.2.0)", "pytest-asyncio", "pytest-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.13.1)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "tensorboard", "timeout-decorator", "torch (>=2.2)", "uvicorn"] -tf = ["keras-nlp (>=0.3.1,<0.14.0)", "onnxconverter-common", "tensorflow (>2.9,<2.16)", "tensorflow-text (<2.16)", "tf2onnx"] -tf-cpu = ["keras (>2.9,<2.16)", "keras-nlp (>=0.3.1,<0.14.0)", "onnxconverter-common", "tensorflow-cpu (>2.9,<2.16)", "tensorflow-probability (<0.24)", "tensorflow-text (<2.16)", "tf2onnx"] -tf-speech = ["kenlm", "librosa", "phonemizer", "pyctcdecode (>=0.4.0)"] +testing = ["GitPython (<3.1.19)", "accelerate (>=1.1.0)", "beautifulsoup4", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "dill (<0.3.5)", "evaluate (>=0.4.6)", "faiss-cpu", "fastapi", "filelock", "hf-doc-builder", "libcst", "mistral-common[image] (>=1.10.0)", "nltk (<=3.8.1)", "openai (>=1.98.0)", "parameterized (>=0.9)", "protobuf", "psutil", "pydantic (>=2)", "pytest (>=7.2.0,<9.0.0)", "pytest-asyncio (>=1.2.0)", "pytest-env", "pytest-order", "pytest-random-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rich", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.14.10)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "tensorboard", "timeout-decorator", "tomli", "torch (>=2.4)", "transformers-mlinter (==0.1.1)", "ty (==0.0.20)", "urllib3 (<2.0.0)", "uvicorn"] tiktoken = ["blobfile", "tiktoken"] -timm = ["timm (!=1.0.18,<=1.0.19)"] -tokenizers = ["tokenizers (>=0.22.0,<=0.23.0)"] -torch = ["accelerate (>=0.26.0)", "torch (>=2.2)"] -torch-speech = ["kenlm", "librosa", "phonemizer", "pyctcdecode (>=0.4.0)", "torchaudio"] -torch-vision = ["Pillow (>=10.0.1,<=15.0)", "torchvision"] -torchhub = ["filelock", "huggingface-hub (>=0.34.0,<1.0)", "importlib_metadata", "numpy (>=1.17)", "packaging (>=20.0)", "protobuf", "regex (!=2019.12.17)", "requests", "sentencepiece (>=0.1.91,!=0.1.92)", "tokenizers (>=0.22.0,<=0.23.0)", "torch (>=2.2)", "tqdm (>=4.27)"] +timm = ["timm (>=1.0.23)"] +torch = ["accelerate (>=1.1.0)", "torch (>=2.4)"] video = ["av"] -vision = ["Pillow (>=10.0.1,<=15.0)"] - +vision = ["Pillow (>=10.0.1,<=15.0)", "torchvision"] [[package]] name = "triton" version = "3.5.1" @@ -8134,10 +8109,9 @@ tutorials = ["matplotlib", "pandas", "tabulate"] name = "typer" version = "0.24.1" description = "Typer, build great CLIs. Easy to code. Based on Python type hints." -optional = true +optional = false python-versions = ">=3.10" groups = ["main"] -markers = "extra == \"vllm\"" files = [ {file = "typer-0.24.1-py3-none-any.whl", hash = "sha256:112c1f0ce578bfb4cab9ffdabc68f031416ebcc216536611ba21f04e9aa84c9e"}, {file = "typer-0.24.1.tar.gz", hash = "sha256:e39b4732d65fbdcde189ae76cf7cd48aeae72919dea1fdfc16593be016256b45"}, @@ -9076,4 +9050,4 @@ vllm = ["vllm"] [metadata] lock-version = "2.1" python-versions = ">=3.11,<3.13" -content-hash = "6bc839d412edaa773717488aaf25cb1e5b663f305e3f9945d9f92fe1524160d2" +content-hash = "dfacac6dc1845f345e1104e0d805b8ccc32542eb991c3d79dde6738a50c319ab" diff --git a/pyproject.toml b/pyproject.toml index 5ec67ed3c..82babced7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ version = "0.4.0.dev0" description = "Training API for large language models with efficient data handling and advanced optimization techniques." readme = "README.md" authors = [{ name = "ModelScope", email = "contact@modelscope.cn" }] -requires-python = ">=3.10,<=3.13" +requires-python = ">=3.11,<=3.13" dependencies = [ "numpy>=2.0.0,<2.3.0", "datasets", @@ -13,7 +13,7 @@ dependencies = [ "modelscope[framework]>=1.34.0", "safetensors", "peft>=0.11.0,<=0.19.0", - "transformers", + "transformers>=5.8,<6.0", "typer>=0.9.0", "pyzmq", "accelerate", @@ -30,11 +30,11 @@ megatron = ["megatron-core>=0.12.0", "transformer-engine[pytorch]", "mcore_bridg data = ["py-data-juicer"] rl = [ "vllm>=0.11", - "ray[serve]" + "ray[serve]>=2.55" ] async-rl = [ "vllm>=0.11", - "ray[serve]", + "ray[serve]>=2.55", "TransferQueue>=0.1.9.dev0", "math-verify==0.8.0", "swanlab>=0.8", @@ -124,5 +124,8 @@ extra-standard-library = ["setuptools"] no-lines-before = ["standard-library", "local-folder"] default-section = "third-party" +[tool.ruff.lint.flake8-bugbear] +extend-immutable-calls = ["fastapi.Depends"] + [tool.ruff.format] quote-style = "double" diff --git a/src/twinkle/hub/__init__.py b/src/twinkle/hub/__init__.py index 4180d5f63..933a2a5c9 100644 --- a/src/twinkle/hub/__init__.py +++ b/src/twinkle/hub/__init__.py @@ -1,3 +1,13 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .hub import HFHub, HubOperation, MSHub -from .model_alias import MODEL_ID_ALIASES_ENV, build_model_alias_map, load_model_alias_map, resolve_model_id_alias +import sys + +from twinkle._lazy_module import _LazyModule + +sys.modules[__name__] = _LazyModule( + __name__, + __file__, { + 'hub': ['HFHub', 'HubOperation', 'MSHub'], + 'model_alias': + ['MODEL_ID_ALIASES_ENV', 'build_model_alias_map', 'load_model_alias_map', 'resolve_model_id_alias'], + }, + module_spec=__spec__) diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index 9b386f517..ed3e45c54 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -174,13 +174,20 @@ def get_config(self, key: str): def clear(self): self.config.clear() + registry_name = os.environ.get('TWINKLE_RAY_REGISTRY_NAME', 'config_registry') + registry_namespace = os.environ.get('TWINKLE_RAY_NAMESPACE') try: - RayHelper._registry = ray.get_actor('config_registry') + RayHelper._registry = ray.get_actor(registry_name, namespace=registry_namespace) except ValueError: try: - RayHelper._registry = WorkerRegistry.options(name='config_registry', lifetime='detached').remote() + RayHelper._registry = WorkerRegistry.options( + name=registry_name, + namespace=registry_namespace, + lifetime='detached', + # Preserve Ray's default CPU policy outside managed instances. + num_cpus=0.01 if registry_namespace else None).remote() except ValueError: - RayHelper._registry = ray.get_actor('config_registry') + RayHelper._registry = ray.get_actor(registry_name, namespace=registry_namespace) assert RayHelper._registry is not None @staticmethod @@ -444,7 +451,7 @@ def create_workers(worker_cls: Type[T], 'scheduling_strategy': PlacementGroupSchedulingStrategy(placement_group=deploy_pg['placement_group']), 'name': worker_name, - 'namespace': 'default', + 'namespace': os.environ.get('TWINKLE_RAY_NAMESPACE', 'default'), 'runtime_env': runtime_env, 'num_cpus': 0.01, } @@ -468,6 +475,8 @@ def create_workers(worker_cls: Type[T], if max_concurrency is not None: worker_options['max_concurrency'] = max_concurrency + if os.environ.get('TWINKLE_RAY_NAMESPACE'): + worker_options.update(max_restarts=0, max_task_retries=0) try: worker = worker_cls.options(**worker_options).remote(*args, **kwargs) except BaseException: @@ -510,7 +519,7 @@ def create_workers(worker_cls: Type[T], 'scheduling_strategy': PlacementGroupSchedulingStrategy(placement_group=deploy_pg['placement_group']), 'name': worker_name, - 'namespace': 'default', + 'namespace': os.environ.get('TWINKLE_RAY_NAMESPACE', 'default'), 'runtime_env': runtime_env, 'num_cpus': 0.01, } @@ -518,6 +527,8 @@ def create_workers(worker_cls: Type[T], if max_concurrency is not None: worker_options['max_concurrency'] = max_concurrency + if os.environ.get('TWINKLE_RAY_NAMESPACE'): + worker_options.update(max_restarts=0, max_task_retries=0) try: worker = worker_cls.options(**worker_options).remote(*args, **kwargs) except BaseException: diff --git a/src/twinkle/infra/_ray/resource_manager.py b/src/twinkle/infra/_ray/resource_manager.py index 27656839e..ba9ec361a 100644 --- a/src/twinkle/infra/_ray/resource_manager.py +++ b/src/twinkle/infra/_ray/resource_manager.py @@ -20,6 +20,39 @@ def _gpu_placement_group_cpus(node_cpu: int, nproc_per_node: int) -> int: class ResourceManager: + @classmethod + def from_allocation(cls, group: DeviceGroup, placement_group, devices: list[str]): + """Use an allocation owned by a lifecycle controller, without reserving again. + + Ranks are local to this allocation; devices are the actual Ray-selected IDs. + The controller owns removal, so this view's destroy method is a no-op. + """ + import ray + manager = cls.__new__(cls) + width = group.gpus_per_worker + if len(devices) % width: + raise ValueError('Allocation device count must be divisible by gpus_per_worker') + chunks = [devices[index:index + width] for index in range(0, len(devices), width)] + ranks = list(range(len(chunks))) + group.ranks = ranks + manager.group_configs = [group] + manager.placement_groups = [placement_group] + manager.cpu_placement_groups = [] + manager._external_allocation = True + manager.device_groups = { + group.name: [{ + 'gpu_rank': chunk, + 'placement_group': placement_group, + 'ray_address': str(ray.get_runtime_context().gcs_address) + } for chunk in chunks] + } + if (group.device_type or '').upper() == 'CPU': + manager.device_groups[group.name] = [{ + 'placement_group': placement_group, + 'ray_address': str(ray.get_runtime_context().gcs_address) + } for _ in devices] + return manager + def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[DeviceGroup]): # CPU placement group default strategy: # - Old approach: use node_cpu//4 as CPU bundle per node, even if only 1~2 CPU processes are needed, @@ -319,6 +352,8 @@ def get_group(self, group: str): return self.device_groups[group] def destroy_placement_group(self): + if getattr(self, '_external_allocation', False): + return import ray for pg in self.placement_groups: ray.util.remove_placement_group(pg) diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py index cf2c3b282..b27c926ae 100644 --- a/src/twinkle/loss/base.py +++ b/src/twinkle/loss/base.py @@ -12,6 +12,14 @@ class Loss: def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput: ... + def validate_inputs(self, inputs, **kwargs) -> None: + """Reject a batch before any rank starts its forward pass. + + Loss implementations whose alignment rules can fail per sample override + this hook. The default accepts every batch. + """ + return None + def micro_batch_scale(self, inputs: list[InputFeature], indices: list[int]) -> float: if len(indices) == len(inputs): return 1.0 diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index d0459fa4c..07ca9ab5e 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -232,10 +232,10 @@ def _pad_and_align_to_batch( * Response-only form (e.g. ``old_logps`` from vLLM): length equals the number of trainable positions in ``mask[i]`` and is scattered directly onto those positions. - * Full-sequence form (e.g. ``ref_logps`` from a ref-model forward, - right-padded to ``mask.shape[1]``): length ``>= mask.shape[1]``; - we slice to ``seq_len`` and index by ``mask[i]`` to extract the - trainable positions, then scatter. + * Full-sequence form (e.g. ``ref_logps`` from a ref-model forward): + every trainable mask position must exist in the per-sample row. + Rows need not include the batch's right-padding; we index by + ``mask[i]`` to extract the trainable positions, then scatter. Any other length is a real bug and triggers a hard assert — never silently truncate, since that misaligns IS ratios to the wrong tokens. @@ -276,25 +276,96 @@ def _pad_and_align_to_batch( pos = mask[i].nonzero(as_tuple=True)[0] n_pos = len(pos) n_sample = sample.numel() + last_pos = int(pos[-1].item()) if n_pos else -1 + if not self._alignment_supported(n_pos, n_sample, last_pos): + self._raise_alignment_mismatch(i, n_pos, n_sample, seq_len) if n_sample == 1: result[i, pos] = sample.item() elif n_sample == n_pos: # Response-only form (e.g. old_logps from vLLM). result[i, pos] = sample - elif n_pos == 0 or (n_sample > 0 and pos[-1].item() < n_sample): + else: # Variable-length full-sequence form. The processor right-pads the # batch, but per-sample RL fields from Tinker remain unpadded. They # are valid when every selected mask position exists in this row. result[i, pos] = sample[pos] - else: - raise AssertionError(f'data/mask length mismatch at sample {i}: ' - f'n_pos={n_pos}, n_sample={n_sample}, seq_len={seq_len} ' - '(expected n_sample == n_pos for response-only form, or all masked positions ' - 'to exist in the per-sample full-sequence form)') return result + @staticmethod + def _alignment_supported(n_pos: int, n_sample: int, last_pos: int) -> bool: + """Whether one ragged row fits the response-only or full-sequence form.""" + return (n_sample == 1 or n_sample == n_pos or n_pos == 0 or (n_sample > 0 and last_pos < n_sample)) + + @staticmethod + def _raise_alignment_mismatch(sample_index: int, n_pos: int, n_sample: int, seq_len: int): + raise AssertionError(f'data/mask length mismatch at sample {sample_index}: ' + f'n_pos={n_pos}, n_sample={n_sample}, seq_len={seq_len} ' + '(expected n_sample == n_pos for response-only form, or all masked positions ' + 'to exist in the per-sample full-sequence form)') + + @staticmethod + def _alignment_rows(data, batch_size: int): + """Return one row per sample, or None when the shape is not attributable.""" + import numpy as np + import torch + + if torch.is_tensor(data): + if data.dim() == 1 and data.shape[0] == batch_size: + return [data[index].reshape(1) for index in range(batch_size)] + if data.dim() >= 2 and data.shape[0] == batch_size: + return list(data) + return None + if isinstance(data, np.ndarray): + return GRPOLoss._alignment_rows(torch.as_tensor(data), batch_size) + if isinstance(data, (list, tuple)) and len(data) == batch_size: + return list(data) + return None + + def validate_inputs(self, inputs, *, old_logps=None, ref_logps=None, advantages=None, **kwargs) -> None: + """Reject ragged per-sample lengths before any rank starts its forward. + + ``_pad_and_align_to_batch`` already rejects these shapes, but it runs inside + the loss on each rank independently. Raising there lets one rank return + while the others enter the DDP backward collective, which deadlocks the + replica. Running the same rule on the full, unsliced batch first makes + every rank fail together before any collective is entered. + """ + import torch + + if old_logps is None and ref_logps is None and advantages is None: + return + samples = inputs if isinstance(inputs, list) else [inputs] + if not samples: + return + for data in (old_logps, ref_logps, advantages): + if data is None: + continue + rows = self._alignment_rows(data, len(samples)) + if rows is None: + # The shape does not map one row per sample; the in-loss check + # remains the only applicable validation for this batch. + continue + for index, (sample, row) in enumerate(zip(samples, rows)): + if not isinstance(sample, dict): + continue + labels = sample.get('labels') + if labels is None: + continue + if not torch.is_tensor(labels): + labels = torch.as_tensor(labels) + if labels.dim() == 1: + labels = labels.unsqueeze(0) + mask = self._resolve_loss_mask(sample, labels)[0] + pos = mask.nonzero(as_tuple=True)[0] + row_tensor = row if torch.is_tensor(row) else torch.as_tensor(row) + n_sample = row_tensor.numel() + n_pos = len(pos) + last_pos = int(pos[-1].item()) if n_pos else -1 + if not self._alignment_supported(n_pos, n_sample, last_pos): + self._raise_alignment_mismatch(index, n_pos, n_sample, labels.shape[-1]) + def _resolve_loss_mask(self, inputs: Dict, labels: 'torch.Tensor') -> 'torch.Tensor': """Positions this loss may score: trainable *and* log-prob-bearing. diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 2f487aec7..398597e63 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -473,7 +473,7 @@ def forward_backward(self, # a KTO z_KL / adapter-disabled reference batch carries no label/ref_logps, which the loss # would reject). eval/training omit it; megatron's calculate_loss is a raise, so eval_loss # depends on forward_only computing it here. - compute_loss = kwargs.pop('compute_loss', True) + compute_loss = kwargs.pop('compute_loss', not (forward_only and disable_lora)) return_logits = kwargs.pop('return_logits', False) task = kwargs.pop('task', 'causal_lm') # MoE routing replay (R2/R3): the driver passes the MODE scalar, resolved below to this @@ -519,11 +519,17 @@ def forward_backward(self, # Compatible with DPO micro_batch_size = min(2, len(inputs)) unwrapped_model = self.strategy.unwrap_model(self.model)[0] + # No config in the Megatron stack declares ``attention_mask_type``, so this + # lookup always fell through to None and the processor could never tell that + # the task is causal. Default it for decoder-only causal_lm. + attention_mask_type = getattr(unwrapped_model.config, 'attention_mask_type', None) + if attention_mask_type is None and task == 'causal_lm': + attention_mask_type = 'causal' inputs = processor( inputs, micro_batch_size=micro_batch_size, variable_seq_lengths=self.variable_seq_lengths, - attention_mask_type=getattr(unwrapped_model.config, 'attention_mask_type', None), + attention_mask_type=attention_mask_type, ) # Get parallelism settings for sequence padding and splitting @@ -594,17 +600,18 @@ def post_loss_function(output_tensor, inputs, logps, unpacked_logits=None, entro if entropies is not None: outputs['entropies'] = entropies if not compute_loss: - # Empty-loss probe: skip the loss but still route a zero loss through reduce_loss so - # the logps reach its returned dict unchanged (the transformers forward_only contract). + # Empty-loss probes still use the common task-output reporting path. zero = torch.zeros((), device=Platform.get_local_device(), dtype=torch.float32) - return self.strategy.reduce_loss(zero, zero.clone(), output_tensor, logps) - result = loss_instance(inputs, outputs, **current_kwargs) + result = {'loss': zero, 'num_tokens': zero.clone()} + else: + result = loss_instance(inputs, outputs, **current_kwargs) + reported_logits = outputs.get('logits') if return_logits else None if unpacked_logits is not None: outputs.pop('logits', None) del unpacked_logits losses = result['loss'] counts = result['num_tokens'] - if not counts: + if compute_loss and not counts: # Later will gather this value, so it becomes: # 1. SUM loss: gather_sum(local_num_tokens) = global_num_tokens # 2. PER TOKEN MEAN loss: gather_sum(1 * gradient_accumulation_steps ) @@ -614,7 +621,9 @@ def post_loss_function(output_tensor, inputs, logps, unpacked_logits=None, entro # 2. PER TOKEN MEAN loss: (gather_sum(per_token_grad * gradient_accumulation_steps)) # / (gradient_accumulation_steps * world_size ) = avg_per_token_grad counts = torch.tensor(1, device=losses.device) - reduced = self.strategy.reduce_loss(losses, counts, output_tensor, logps) + reduced = self.strategy.reduce_loss(losses, counts, reported_logits, logps) + if embeddings is not None: + reduced[2]['embeddings'] = embeddings.detach() if result.get('channel_loss') is not None: reduced[2]['channel_loss'] = result['channel_loss'] return reduced @@ -811,6 +820,7 @@ def forward_step_func(data_iterator, model): loss = torch.tensor(0.0).to(Platform.get_local_device()) logits = [] logps = [] + embedding_batches = [] channel_loss = {} count = 0 n_loss_entries = 0 @@ -820,10 +830,12 @@ def forward_step_func(data_iterator, model): if 'loss' in loss_dict: loss += loss_dict['loss'] n_loss_entries += 1 - if 'logits' in loss_dict: + if loss_dict.get('logits') is not None: logits.append(loss_dict['logits']) - if 'logps' in loss_dict: + if loss_dict.get('logps') is not None: logps.append(loss_dict['logps']) + if loss_dict.get('embeddings') is not None: + embedding_batches.append(loss_dict['embeddings']) if 'num_tokens' in loss_dict: count += loss_dict['num_tokens'] for channel, stats in loss_dict.get('channel_loss', {}).items(): @@ -868,7 +880,7 @@ def forward_step_func(data_iterator, model): logps = torch.cat(logps, dim=0) elif not logps: logps = None - if logits and not self.variable_seq_lengths: + if logits and (not self.variable_seq_lengths or task in ('seq_cls', 'reranker', 'generative_reranker')): logits = torch.cat(logits, dim=0) if transactional: dp_world_size = self.device_mesh._get_dp_fsdp_world_size() @@ -880,10 +892,12 @@ def forward_step_func(data_iterator, model): loss = loss.detach().float().item() if isinstance(num_tokens, torch.Tensor): num_tokens = num_tokens.detach().float().item() - if not return_logits: + if not return_logits or (isinstance(logits, list) and not logits): logits = None inputs = processor.unpack_inputs(inputs, task=task) model_output = ModelOutput(logits=logits, loss=loss, logps=logps, num_tokens=num_tokens) + if embedding_batches: + model_output['embeddings'] = torch.cat(embedding_batches, dim=0) if channel_loss: model_output['channel_loss'] = channel_loss # R2 RECORD: hand the captured whole-model routing back to the driver (symmetric with @@ -1409,6 +1423,9 @@ def reload_initial_weights(self, **kwargs): for _model in self.strategy.unwrap_model(self.model): for param in _model.parameters(): param.grad = None + if kwargs.get('reset_training_state', False): + self.optimizer_group = {_default_adapter_name: self._construct_default_optimizer_group()} + self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name if dist.is_initialized(): dist.barrier() diff --git a/src/twinkle/model/megatron/multi_lora_megatron.py b/src/twinkle/model/megatron/multi_lora_megatron.py index 90e29eb26..19c3db52c 100644 --- a/src/twinkle/model/megatron/multi_lora_megatron.py +++ b/src/twinkle/model/megatron/multi_lora_megatron.py @@ -8,6 +8,7 @@ import torch.distributed as dist import torch.nn as nn from contextlib import contextmanager +from copy import deepcopy from functools import partial from peft import LoraConfig from torch.optim import Optimizer @@ -35,6 +36,36 @@ logger = get_logger() +@contextmanager +def _seq_cls_head_lora_load_gate(bridge): + """Let mcore-bridge restore a LoRA-only sequence classification head. + + mcore-bridge 1.6.4 checks for ``score.weight`` before entering its head loader, + but a PEFT checkpoint contains only ``score.lora_A/B.weight``. Provide the key + as a load-time gate; the LoRA branch never reads the corresponding base tensor. + """ + required = ('_convert_post_process', '_peft_format', 'hf_score_key', 'config') + if any(not hasattr(bridge, name) for name in required): + raise RuntimeError('seq_cls head LoRA restore requires the mcore-bridge 1.6.4 loader API') + original = bridge._convert_post_process + + def convert_post_process(mg_model, hf_state_dict, hf_prefix, to_mcore): + if to_mcore and bridge._peft_format and bridge.config.task_type == 'seq_cls': + head_key = bridge.hf_score_key + module_key, _, param_key = head_key.rpartition('.') + lora_key = f'{module_key}.lora_A.{param_key}' + if head_key not in hf_state_dict and lora_key in hf_state_dict: + hf_state_dict = dict(hf_state_dict) + hf_state_dict[head_key] = hf_state_dict[lora_key] + return original(mg_model, hf_state_dict, hf_prefix, to_mcore) + + bridge._convert_post_process = convert_post_process + try: + yield + finally: + bridge._convert_post_process = original + + @remote_class(execute='all') class MultiLoraMegatronModel(MegatronModel): @@ -377,7 +408,7 @@ def load(self, name: str, output_dir: Optional[str] = None, **kwargs): else: checkpoint_dir = os.path.join(output_dir, name) bridge = self.strategy.bridge - with self.multi_adapter.save_context(adapter_name) as real_adapter_name: + with self.multi_adapter.save_context(adapter_name) as real_adapter_name, _seq_cls_head_lora_load_gate(bridge): model = self.strategy.unwrap_model(self.model) bridge.load_weights( model, @@ -454,8 +485,12 @@ def add_adapter_to_model( 'because loading config from modelhub may causing unexpected behavior') if not isinstance(config_or_dir, LoraConfig): raise ValueError('config_or_dir must be a LoraConfig instance') + config_or_dir = deepcopy(config_or_dir) + if config_or_dir.modules_to_save: + raise ValueError('modules_to_save is not supported for a multi-tenant LoRA adapter.') + self.multi_adapter.validate_tenant_target_modules(config_or_dir.target_modules, + getattr(config_or_dir, 'target_parameters', None)) # Limit the max peft version in pyproject.toml, in case any newer version opens some untested module grad. - config_or_dir.modules_to_save = None config_or_dir.bias = 'none' config_or_dir.init_lora_weights = False config_or_dir.trainable_token_indices = None diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 0aa8e62fe..4b583a8e5 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -12,7 +12,7 @@ from twinkle import torch_util from twinkle.data_format import InputFeature from twinkle.utils import get_logger -from .multi_lora_target_parameters import TargetParameterLoraManager +from .multi_lora_target_parameters import TargetParameterLoraManager, TargetParameterLoraWrapper logger = get_logger() @@ -124,6 +124,54 @@ def reset_adapter_status(self): else: self.module.set_adapter('lora_0') + def _named_slot_parameters(self, module): + """Identify owned adapter parameters through their actual layer containers.""" + slots = {tenant.adapter_name for tenant in self.loras} + for layer_name, layer in module.named_modules(): + if isinstance(layer, LoraLayer): + containers = layer.adapter_layer_names + elif isinstance(layer, TargetParameterLoraWrapper): + containers = ('lora_A', 'lora_B') + else: + continue + for container_name in containers: + container = getattr(layer, container_name) + for slot in slots & container.keys(): + value = container[slot] + prefix = '.'.join(part for part in (layer_name, container_name, slot) if part) + if isinstance(value, torch.nn.Parameter): + yield prefix, value + else: + for name, parameter in value.named_parameters(): + yield prefix + '.' + name, parameter + + def enable_slot_gradients(self): + """Keep owned adapter slots visible to reducers without unfreezing the base.""" + modules = self.module if isinstance(self.module, list) else [self.module] + for module in modules: + for _, parameter in self._named_slot_parameters(module): + parameter.requires_grad_(True) + + @contextmanager + def trainable_slots(self): + """Expose all slots while wrapping DDP/FSDP, then restore the active mask. + + Reducers discover parameters once during wrapping. The mask is restored + by adapter path because FSDP can replace the Parameter objects themselves. + """ + modules = self.module if isinstance(self.module, list) else [self.module] + masks = [{ + name: parameter.requires_grad + for name, parameter in self._named_slot_parameters(module) + } for module in modules] + try: + self.enable_slot_gradients() + yield + finally: + for module, mask in zip(modules, masks): + for name, parameter in self._named_slot_parameters(module): + parameter.requires_grad_(mask[name]) + def activate_adapter(self, tenant_adapter_name: str, call_enable=False): if not self.has_lora(tenant_adapter_name): raise ValueError(f'Adapter {tenant_adapter_name} does not exist') @@ -283,7 +331,8 @@ def validate_tenant_target_modules(self, target_modules, target_parameters=None) """Ensure every requested target resolves inside the preallocated LoRA layers.""" if not target_modules and target_parameters: return - layers = [name for name, layer in self.module.named_modules() if isinstance(layer, LoraLayer)] + modules = self.module if isinstance(self.module, list) else [self.module] + layers = [name for module in modules for name, layer in module.named_modules() if isinstance(layer, LoraLayer)] if target_modules == 'all-linear' or (isinstance(target_modules, (list, set)) and 'all-linear' in target_modules): return @@ -559,18 +608,12 @@ def _patch_megatron(_module): # PEFT's add_adapter calls set_adapter(active_adapters) which only keeps the # first adapter's requires_grad=True. We need ALL LoRA params to be trainable # so that MegatronDDP registers them all in its gradient buffers (main_grad). - def _enable_all_lora_grad(_module): - for name, param in _module.named_parameters(): - if 'lora_' in name and not param.requires_grad: - param.requires_grad_(True) - - if isinstance(module, list): - for _m in module: - _enable_all_lora_grad(_m) - else: - _enable_all_lora_grad(module) - self.module = module + modules = module if isinstance(module, list) else [module] + for chunk in modules: + for _, parameter in self._named_slot_parameters(chunk): + parameter.requires_grad_(True) + if not isinstance(module, list): self.lora_layer_names = [name for name, layer in module.named_modules() if isinstance(layer, Linear)] return module diff --git a/src/twinkle/model/optimizer_group.py b/src/twinkle/model/optimizer_group.py index cced80a67..566bc03ca 100644 --- a/src/twinkle/model/optimizer_group.py +++ b/src/twinkle/model/optimizer_group.py @@ -135,36 +135,31 @@ def _get_lr(self): """Get learning rates from optimizer. Override in subclass.""" return [] + def _metric_kwargs(self, status: TrainStatus, metric_step: Optional[int] = None) -> Dict[str, Any]: + """Use actual optimizer state while retaining the selected loss's reduction.""" + forward_kwargs = dict(status.forward_kwargs) + loss_key = forward_kwargs.pop('_metric_loss_key', None) + return { + **forward_kwargs, + 'lr': self._get_lr(), + 'step': self.cur_step - 1 if metric_step is None else metric_step, + 'gradient_accumulation_steps': self.gradient_accumulation_steps, + 'grad_norm': self._last_grad_norm, + 'loss_reduction': status.loss_reduction or getattr(self.get_loss(loss_key), 'reduction', 'mean'), + } + def accumulate_metrics(self, is_training, include_default: bool = True, metric_step: Optional[int] = None): - """Accumulate metrics for train/eval status. Override in subclass if needed.""" + """Accumulate metrics through one argument merge and keyed-loss dispatch.""" status = self.train_status if is_training else self.eval_status - if metric_step is None: - metric_step = self.cur_step - 1 - if len(status.metrics) > 0 and status.inputs is not None and status.outputs is not None: - # forward_kwargs is the raw kwargs the forward was called with. A caller that drives - # gradient accumulation threads gradient_accumulation_steps through it (the Megatron - # forward_backward pops it for do_grad_sync, but the dict stored on train_status still - # carries it). It must not be forwarded to the metric verbatim: this method passes the - # authoritative self.gradient_accumulation_steps explicitly, so a second copy inside - # **forward_kwargs raises "got multiple values for keyword argument". Pop it here -- the - # TransformersOptimizerGroup override does the same -- so both backends accumulate metrics - # identically (basic principle: backend equivalence, no per-backend special casing). - forward_kwargs = dict(status.forward_kwargs) - forward_kwargs.pop('gradient_accumulation_steps', None) - loss_key = forward_kwargs.pop('_metric_loss_key', None) - for metric in status.metrics: - metric_loss_key = getattr(metric, '_twinkle_loss_key', None) - if metric_loss_key != loss_key or (metric_loss_key is None and not include_default): - continue - metric.accumulate( - status.inputs, - status.outputs, - lr=self._get_lr(), - step=metric_step, - gradient_accumulation_steps=self.gradient_accumulation_steps, - grad_norm=self._last_grad_norm, - loss_reduction=status.loss_reduction or getattr(self.get_loss(loss_key), 'reduction', 'mean'), - **forward_kwargs) + if not status.metrics or status.inputs is None or status.outputs is None: + return + loss_key = status.forward_kwargs.get('_metric_loss_key') + metric_kwargs = self._metric_kwargs(status, metric_step) + for metric in status.metrics: + metric_loss_key = getattr(metric, '_twinkle_loss_key', None) + if metric_loss_key != loss_key or (metric_loss_key is None and not include_default): + continue + metric.accumulate(status.inputs, status.outputs, **metric_kwargs) def calculate_metrics(self, is_training): """Calculate and return metrics.""" diff --git a/src/twinkle/model/transformers/hybrid/fft_slots.py b/src/twinkle/model/transformers/hybrid/fft_slots.py index 39cde9579..38d675d2b 100644 --- a/src/twinkle/model/transformers/hybrid/fft_slots.py +++ b/src/twinkle/model/transformers/hybrid/fft_slots.py @@ -120,6 +120,10 @@ def activate_fft_slot(self, adapter_name: str) -> None: fft_adapter_name = self._fft_adapter_name(adapter_name) if self.is_hybrid(adapter_name) else None selected_adapters = [fft_adapter_name] if fft_adapter_name is not None else [] for wrapper in self._iter_fft_wrappers(): + # Selecting a regular LoRA can disable auxiliary wrappers. Restoring + # an FFT slot name alone does not restore its execution state. + if fft_adapter_name is not None and wrapper.disable_adapters: + wrapper.enable_adapters(True) if wrapper.active_adapters != selected_adapters: wrapper.set_adapter(selected_adapters) self._enable_preallocated_fft_grad(wrapper) diff --git a/src/twinkle/model/transformers/hybrid/model.py b/src/twinkle/model/transformers/hybrid/model.py index 26d25f674..6840889f5 100644 --- a/src/twinkle/model/transformers/hybrid/model.py +++ b/src/twinkle/model/transformers/hybrid/model.py @@ -64,27 +64,17 @@ def _adapter_context(self, adapter_name: str, disable_lora: bool = False): return with super()._adapter_context(adapter_name) as slot_name: - self._restore_preallocated_lora_grad() self.fft_slots.activate_fft_slot(adapter_name) yield slot_name @contextmanager def _disabled_adapter_context(self, adapter_name: str): - try: - with super()._adapter_context(adapter_name, disable_lora=True) as slot_name: - self._restore_preallocated_lora_grad() - self.fft_slots.deactivate_fft_slots() - try: - yield slot_name - finally: - self.fft_slots.activate_fft_slot(adapter_name) - finally: - self._restore_preallocated_lora_grad() - - def _restore_preallocated_lora_grad(self) -> None: - for name, parameter in self.multi_adapter.module.named_parameters(): - if 'lora_' in name and not parameter.requires_grad: - parameter.requires_grad_(True) + with super()._adapter_context(adapter_name, disable_lora=True) as slot_name: + self.fft_slots.deactivate_fft_slots() + try: + yield slot_name + finally: + self.fft_slots.activate_fft_slot(adapter_name) @remote_function() def add_adapter_to_model(self, adapter_name: str, config_or_dir: Union[PeftConfig, str], **kwargs): diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index 3c8e57dea..48ccd720e 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -43,6 +43,8 @@ def __init__( max_length: int = 8192, target_modules: Union[List[str], str] = 'all-linear', model_loader: Optional[ModelLoaderProtocol] = None, + gradient_checkpointing: bool = True, + gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None, acc_strategy: str = 'token', model_arch: Optional[Dict[str, Any]] = None, **kwargs): @@ -57,7 +59,12 @@ def __init__( self.device_mesh = device_mesh self.mixed_precision = mixed_precision self._fsdp_config = dict(fsdp_config or {}) - self._ddp_config = ddp_config or {} + if self._fsdp_config.get('reshard_after_forward', True) is not True: + raise ValueError('Multi-LoRA FSDP requires reshard_after_forward=True for safe adapter switching.') + self._ddp_config = dict(ddp_config or {}) + if self._ddp_config.get('find_unused_parameters') is False or self._ddp_config.get('static_graph', False): + raise ValueError('Multi-LoRA DDP requires find_unused_parameters=True and static_graph=False.') + self._ddp_config.setdefault('find_unused_parameters', True) self._memory_efficient_init = memory_efficient_init self._router_replay_enabled = bool(kwargs.pop('enable_router_replay', False)) self._router_replay_applied = False @@ -118,7 +125,9 @@ def __init__( # Initialize expert parallel attributes (required by set_optimizer in TransformersModel) self.optimizer_group: Dict[str, OptimizerGroup] = {} self.multi_adapter = MultiLora(max_loras=max_loras, max_r=max_r, max_length=max_length) - self.model.gradient_checkpointing_enable() + if gradient_checkpointing: + self.model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={'use_reentrant': False, **(gradient_checkpointing_kwargs or {})}) self.model = self.multi_adapter.patch(self.model, target_modules=target_modules, lora_config=self.lora_config) # PEFT creates LoRA parameters in FP32 even when the base model uses a # lower-precision dtype. Align the preallocated slots before FSDP2 @@ -153,6 +162,17 @@ def register_mm_forward_hook(self, optimizer_group: OptimizerGroup): def unregister_mm_forward_hook(self, optimizer_group: OptimizerGroup): pass + def _lazy_wrap_model(self): + if self._model_wrapped: + return + # PEFT freezes inactive slots. DDP must nevertheless register all slots, + # since a later tenant may activate any of them on this resident model. + with self.multi_adapter.trainable_slots(): + super()._lazy_wrap_model() + # Reset must restore the synchronized (and possibly sharded) basis, + # not the rank-local initialization captured before model wrapping. + self.multi_adapter.save_initial_weights() + def _maybe_apply_expert_parallel(self): if self._memory_efficient_init: raise NotImplementedError('Expert parallel is not supported with memory_efficient_init') @@ -172,9 +192,14 @@ def _ensure_target_parameter_lora_installed(self, config: LoraConfig) -> None: @contextmanager def _adapter_context(self, adapter_name: str, disable_lora: bool = False): - """Activate one LoRA tenant for a model operation.""" - with self.multi_adapter.adapter(adapter_name, disable_lora=disable_lora) as slot_name: - yield slot_name + """Activate a tenant while keeping FSDP2's trainable parameter set stable.""" + try: + with self.multi_adapter.adapter(adapter_name, disable_lora=disable_lora) as slot_name: + self.multi_adapter.enable_slot_gradients() + yield slot_name + finally: + # Leaving a disable_lora context re-enables adapter layers, which flips the flags again. + self.multi_adapter.enable_slot_gradients() @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) def forward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs): @@ -344,6 +369,7 @@ def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): adapter_name = kwargs.get('adapter_name') self._check_adapter_valid(adapter_name) + self._lazy_wrap_model() with self._adapter_context(adapter_name), self.multi_adapter.save_context(adapter_name): checkpoint_dir = super().save(name, output_dir, interval, **kwargs) if dist.is_initialized(): @@ -354,6 +380,9 @@ def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): def load(self, name: str, output_dir: Optional[str] = None, **kwargs): adapter_name = kwargs.get('adapter_name') self._check_adapter_valid(adapter_name) + # Establish the synchronized reset basis before checkpoint weights can + # replace it, including when load is the first model operation. + self._lazy_wrap_model() with self._adapter_context(adapter_name), self.multi_adapter.save_context(adapter_name): load_optimizer = kwargs.get('load_optimizer', False) if output_dir is None: diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 43041cb11..2884e61a4 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -233,26 +233,7 @@ def _get_lr(self): def accumulate_metrics(self, is_training, include_default: bool = True, metric_step: Optional[int] = None): self._ensure_dp_group() - status = self.train_status if is_training else self.eval_status - if metric_step is None: - metric_step = self.cur_step - 1 - if len(status.metrics) > 0 and status.inputs is not None and status.outputs is not None: - forward_kwargs = copy(status.forward_kwargs) - forward_kwargs.pop('gradient_accumulation_steps', None) - loss_key = forward_kwargs.pop('_metric_loss_key', None) - for metric in status.metrics: - metric_loss_key = getattr(metric, '_twinkle_loss_key', None) - if metric_loss_key != loss_key or (metric_loss_key is None and not include_default): - continue - metric.accumulate( - status.inputs, - status.outputs, - lr=self._get_lr(), - step=metric_step, - gradient_accumulation_steps=self.gradient_accumulation_steps, - grad_norm=self._last_grad_norm, - loss_reduction=status.loss_reduction or getattr(self.get_loss(loss_key), 'reduction', 'mean'), - **forward_kwargs) + super().accumulate_metrics(is_training, include_default=include_default, metric_step=metric_step) _default_adapter_name = '' @@ -1722,7 +1703,10 @@ def get_parameter_names(model, forbidden_layer_types, forbidden_layer_names=None return result forbidden_name_patterns = [r'bias', r'layernorm', r'rmsnorm', r'(?:^|\.)norm(?:$|\.)', r'_norm(?:$|\.)'] - decay_parameters = get_parameter_names(self.model, [torch.nn.LayerNorm], forbidden_name_patterns) + # Match names from the same unwrapped model used by _get_trainable_parameters. + # DDP's module prefix otherwise changes optimizer groups for later tenants. + decay_parameters = get_parameter_names( + self.strategy.unwrap_model(self.model), [torch.nn.LayerNorm], forbidden_name_patterns) params = self._get_trainable_parameters(adapter_name) decay_param_names = [n for n, p in params.items() if (n in decay_parameters and p.requires_grad)] no_decay_param_names = [n for n, p in params.items() if (n not in decay_parameters and p.requires_grad)] @@ -2260,6 +2244,15 @@ def reload_initial_weights(self, **kwargs): (exclusive) model, so the next tenant starts from clean pretrained weights instead of the previous tenant's trained weights. """ + if kwargs.get('reset_training_state', False): + for optimizer_group in self.optimizer_group.values(): + self.unregister_mm_forward_hook(optimizer_group) + self.optimizer_group = {_default_adapter_name: self._construct_default_optimizer_group()} + self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name + self.active_group = _default_adapter_name + if self._model_wrapped: + self.register_mm_forward_hook(self.optimizer_group[_default_adapter_name]) + if not self.model_id: logger.warning('reload_initial_weights skipped: model_id is not set (blank model).') return diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index cbea6cd3b..2d0413752 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -550,13 +550,13 @@ def prepare_gdn_padding_free_args(self, inputs: list[InputFeature], **kwargs) -> _inp["max_length_k"] = int(packed_seq_params.max_seqlen_kv) return inputs + @staticmethod + def _backend_builds_causal_mask(attention_mask_type: str | None) -> bool: + return attention_mask_type == "causal" and Platform.device_prefix() == "npu" + def drop_causal_4d_mask(self, inputs: list[InputFeature], **kwargs) -> list[InputFeature]: - """On NPU, drop the generic 4D dense mask so MindSpeed can build - its own compressed causal mask for FlashAttention.""" - if Platform.device_prefix() != "npu": - return inputs - attention_mask_type = kwargs.get("attention_mask_type") - if attention_mask_type != "causal": + """Drop a dense mask when the accelerator synthesizes its causal mask.""" + if not self._backend_builds_causal_mask(kwargs.get("attention_mask_type")): return inputs for _inp in inputs: attention_mask = _inp.get("attention_mask") @@ -862,7 +862,9 @@ def _fill_optional_sequence_fields(self, batch: list[InputFeature]) -> None: #: these fields and must change in the same commit. _PER_AUDIO_BATCHED_FIELDS = ("input_features", "feature_attention_mask") - def _collate_macro_batch(self, inputs: list[InputFeature]) -> InputFeature: + def _collate_macro_batch( + self, inputs: list[InputFeature], *, omit_megatron_attention_mask: bool = False + ) -> InputFeature: # Two mutually-exclusive audio contracts share the `input_features` field name and are told # apart by the mask field that rides alongside it -- the producer (the model's template) # emits one pair or the other, so the field name IS the contract marker: @@ -931,7 +933,7 @@ def is_mm_position_ids(position_ids): for key in text_keys: values = [item[key] for item in text_inputs] if self.framework == "megatron" and key == "attention_mask": - result[key] = self._create_4d_attention_mask(values) + result[key] = None if omit_megatron_attention_mask else self._create_4d_attention_mask(values) elif key == "position_ids" and is_mm_position_ids(values[0]): result[key] = InputProcessor._pad_sequence(values, self.padding_map[key], self.padding_side) num_axes = values[0].shape[0] @@ -980,9 +982,12 @@ def collate_fn( ) -> list[InputFeature]: if len(inputs) == 1 and self.framework != "megatron": return inputs + omit_megatron_attention_mask = ( + self.framework == "megatron" and self._backend_builds_causal_mask(kwargs.get("attention_mask_type")) + ) if micro_batch_size is None: # normal collate - outputs = self._collate_macro_batch(inputs) + outputs = self._collate_macro_batch(inputs, omit_megatron_attention_mask=omit_megatron_attention_mask) for key in outputs: if key in self.VLM_CONCAT_FIELDS: outputs[key] = torch.cat(outputs[key], dim=0) @@ -993,7 +998,9 @@ def collate_fn( assert len(inputs) >= micro_batch_size outputs = [] for i in range(0, len(inputs), micro_batch_size): - _output = self._collate_macro_batch(inputs[i : i + micro_batch_size]) + _output = self._collate_macro_batch( + inputs[i : i + micro_batch_size], omit_megatron_attention_mask=omit_megatron_attention_mask + ) for key in _output: if key in self.VLM_CONCAT_FIELDS: _output[key] = torch.cat(_output[key], dim=0) @@ -1001,14 +1008,16 @@ def collate_fn( return outputs else: # each macro batch shares the same length - res = self._collate_macro_batch(inputs) + res = self._collate_macro_batch(inputs, omit_megatron_attention_mask=omit_megatron_attention_mask) keys = list(res.keys()) outputs = [] for i in range(0, len(inputs), micro_batch_size): end = i + micro_batch_size output = {} for key in keys: - if key == "position_ids" and res[key].dim() > 2: + if res[key] is None: + output[key] = None + elif key == "position_ids" and res[key].dim() > 2: output[key] = res[key][:, i:end, :] elif key in self.VLM_CONCAT_FIELDS: output[key] = torch.cat(res[key][i:end], dim=0) diff --git a/src/twinkle/protocol/types/base.py b/src/twinkle/protocol/types/base.py index d428afc81..a3ef46cd7 100644 --- a/src/twinkle/protocol/types/base.py +++ b/src/twinkle/protocol/types/base.py @@ -22,7 +22,7 @@ ``status_code`` -- an execution-time failure is delivered with HTTP 200, so the value is systematically unequal to the response status code. 4. A closed value set on a wire field is declared as ``Literal`` / enum, never a - bare ``str`` (see ``QueueStateLiteral`` in ``errors.py``). + bare ``str`` (see ``QueueState`` in ``errors.py``). Field roles ----------- diff --git a/src/twinkle/protocol/types/errors.py b/src/twinkle/protocol/types/errors.py index b896e4eed..5d3e5e32b 100644 --- a/src/twinkle/protocol/types/errors.py +++ b/src/twinkle/protocol/types/errors.py @@ -20,10 +20,14 @@ class can be lined up against its tinker counterpart): from .base import ResponseModel -# Closed value set, kept in sync with the server-side ``QueueState`` enum values -# (a consistency test asserts the two sets are equal). Wire fields carrying a queue -# state declare this alias, never a bare ``str`` (naming ruling 4). -QueueStateLiteral = Literal['active', 'paused_rate_limit', 'paused_capacity', 'unknown'] + +class QueueState(StrEnum): + """Scheduling hints; these are independent of a task's lifecycle status.""" + + ACTIVE = 'active' + PAUSED_RATE_LIMIT = 'paused_rate_limit' + PAUSED_CAPACITY = 'paused_capacity' + UNKNOWN = 'unknown' class ErrorCategory(StrEnum): diff --git a/src/twinkle/protocol/types/lifecycle.py b/src/twinkle/protocol/types/lifecycle.py index 32555e24b..641c47fbb 100644 --- a/src/twinkle/protocol/types/lifecycle.py +++ b/src/twinkle/protocol/types/lifecycle.py @@ -7,18 +7,25 @@ """ from __future__ import annotations -from typing import Any, Literal, Optional +from enum import StrEnum +from typing import Any, Optional from .base import ResponseModel, StrictRequest -from .errors import ErrorPayload, QueueStateLiteral +from .errors import ErrorPayload, QueueState -# The lifecycle status of a queued task. Kept in sync with the server-side -# ``TaskStatus`` enum values (a consistency test asserts the two sets are equal). -TaskStatus = Literal['pending', 'queued', 'running', 'completed', 'failed', 'cancelled'] -# The two states past which a task never changes again. ``frozenset`` so a caller -# cannot mutate the shared set. -TERMINAL_STATUSES: frozenset[str] = frozenset({'completed', 'failed', 'cancelled'}) +class TaskStatus(StrEnum): + """Canonical task states shared by the wire envelope and server execution.""" + + PENDING = 'pending' + QUEUED = 'queued' + RUNNING = 'running' + COMPLETED = 'completed' + FAILED = 'failed' + CANCELLED = 'cancelled' + + +TERMINAL_STATUSES = frozenset({TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, TaskStatus.CANCELLED.value}) class RetrieveFutureRequest(StrictRequest): @@ -79,5 +86,5 @@ class TaskEnvelope(ResponseModel): status: TaskStatus result: Any | None = None # set iff status == 'completed' error: ErrorPayload | None = None # set iff status == 'failed' - queue_state: QueueStateLiteral | None = None + queue_state: QueueState | None = None queue_state_reason: str | None = None diff --git a/src/twinkle/protocol/types/model.py b/src/twinkle/protocol/types/model.py index f2e6bd976..51e5c9ce9 100644 --- a/src/twinkle/protocol/types/model.py +++ b/src/twinkle/protocol/types/model.py @@ -28,7 +28,9 @@ class CreateRequest(StrictRequest): - """Body of ``POST /twinkle/create``: a session-establishing no-op.""" + """Body of ``POST /twinkle/create``: CPU-only capability handshake.""" + + training_handles: bool = False # --------------------------------------------------------------------------- # @@ -36,7 +38,14 @@ class CreateRequest(StrictRequest): # --------------------------------------------------------------------------- # -class AdapterRequest(StrictRequest): +class TrainingRequest(StrictRequest): + """A training handle is control metadata, never a backend kwarg.""" + + training_id: str = Field(min_length=1) + adapter_name: str + + +class AdapterRequest(TrainingRequest): """The shared shape of an adapter-scoped operation. ``seq_id`` is an idempotency key, not a backend parameter: the submit shell @@ -46,7 +55,6 @@ class AdapterRequest(StrictRequest): which would silently disable that dedup. """ - adapter_name: str seq_id: int | None = None gradient_accumulation_steps: int | None = backend_kwarg(default=None, ge=1) @@ -83,10 +91,9 @@ class ClipGradAndStepRequest(ClipGradNormRequest): optim_params: dict[str, JsonValue] | None = backend_kwarg(default=None) -class CalculateMetricRequest(StrictRequest): +class CalculateMetricRequest(TrainingRequest): """Body of ``POST /twinkle/calculate_metric``.""" - adapter_name: str is_training: bool = True @@ -105,7 +112,7 @@ class CalculateMetricRequest(StrictRequest): # --------------------------------------------------------------------------- # -class _InlineForwardBase(StrictRequest): +class _InlineForwardBase(TrainingRequest): """Fields common to the three inline forward endpoints.""" inputs: WireInputBatch @@ -127,7 +134,6 @@ class _InlineForwardBase(StrictRequest): class ForwardRequest(_InlineForwardBase): """Body of ``POST /twinkle/forward``: keeps the graph, mutates no gradients.""" - adapter_name: str disable_lora: bool | None = backend_kwarg(default=None) @@ -139,14 +145,12 @@ class ForwardOnlyRequest(_InlineForwardBase): nothing is mutated to be idempotent about. """ - adapter_name: str disable_lora: bool | None = backend_kwarg(default=None) class ForwardBackwardTaskRequest(_InlineForwardBase): """Body of ``POST /twinkle/forward_backward``: accumulates gradients.""" - adapter_name: str seq_id: int | None = None sync_gradients: bool | None = backend_kwarg(default=None) loss_scale: float | None = backend_kwarg(default=None) @@ -161,7 +165,7 @@ class ForwardBackwardTaskRequest(_InlineForwardBase): # --------------------------------------------------------------------------- # -class DataPlaneForwardRequest(StrictRequest): +class DataPlaneForwardRequest(TrainingRequest): """Body of the ``*_from_data_plane`` forward endpoints.""" input_refs: list[DataRef] = Field(min_length=1) @@ -169,7 +173,6 @@ class DataPlaneForwardRequest(StrictRequest): # Values are *field paths*, not parameter values, so this is not a passthrough # region: nothing in it is forwarded verbatim. kwarg_fields: dict[str, str] = Field(default_factory=dict) - adapter_name: str seq_id: int | None = None task: str | None = backend_kwarg(default=None) temperature: float | None = backend_kwarg(default=None, gt=0) @@ -180,7 +183,7 @@ class DataPlaneForwardRequest(StrictRequest): loss_kwargs: dict[str, JsonValue] = passthrough() -class DataPlaneForwardOnlyRequest(StrictRequest): +class DataPlaneForwardOnlyRequest(TrainingRequest): """Body of ``POST /twinkle/forward_only_from_data_plane``. This endpoint is read-only, so it deliberately has no ``seq_id`` idempotency @@ -191,7 +194,6 @@ class DataPlaneForwardOnlyRequest(StrictRequest): input_refs: list[DataRef] = Field(min_length=1) input_field: str | None = None kwarg_fields: dict[str, str] = Field(default_factory=dict) - adapter_name: str task: str | None = backend_kwarg(default=None) temperature: float | None = backend_kwarg(default=None, gt=0) return_logits: bool | None = backend_kwarg(default=None) @@ -222,25 +224,22 @@ def validate_output(self) -> DataPlaneForwardOnlyRequest: # --------------------------------------------------------------------------- # -class SetLossRequest(StrictRequest): +class SetLossRequest(TrainingRequest): loss_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class SetOptimizerRequest(StrictRequest): +class SetOptimizerRequest(TrainingRequest): optimizer_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class SetLrSchedulerRequest(StrictRequest): +class SetLrSchedulerRequest(TrainingRequest): scheduler_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class SetTemplateRequest(StrictRequest): +class SetTemplateRequest(TrainingRequest): """Body of ``POST /twinkle/set_template``. No top-level ``model_id``: the backend always overrides it with its own @@ -250,26 +249,22 @@ class SetTemplateRequest(StrictRequest): """ template_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class SetProcessorRequest(StrictRequest): +class SetProcessorRequest(TrainingRequest): processor_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() -class AddMetricRequest(StrictRequest): +class AddMetricRequest(TrainingRequest): metric_cls: str - adapter_name: str is_training: bool | None = None init_kwargs: dict[str, JsonValue] = passthrough() -class ApplyPatchRequest(StrictRequest): +class ApplyPatchRequest(TrainingRequest): patch_cls: str - adapter_name: str init_kwargs: dict[str, JsonValue] = passthrough() @@ -278,8 +273,7 @@ class ApplyPatchRequest(StrictRequest): # --------------------------------------------------------------------------- # -class SaveRequest(StrictRequest): - adapter_name: str +class SaveRequest(TrainingRequest): name: str | None = None save_optimizer: bool = False is_sampler: bool = False # If True, delete existing sampler weights before saving @@ -287,8 +281,7 @@ class SaveRequest(StrictRequest): merge_lora: bool | None = backend_only('megatron', default=None) -class LoadRequest(StrictRequest): - adapter_name: str +class LoadRequest(TrainingRequest): name: str load_optimizer: bool = False no_load_optim: bool | None = backend_only('megatron', default=None) @@ -296,7 +289,7 @@ class LoadRequest(StrictRequest): strict: bool | None = backend_only('transformers', default=None) -class ResumeFromCheckpointRequest(StrictRequest): +class ResumeFromCheckpointRequest(TrainingRequest): """Body of ``POST /twinkle/resume_from_checkpoint``.""" name: str @@ -305,6 +298,7 @@ class ResumeFromCheckpointRequest(StrictRequest): class AddAdapterRequest(StrictRequest): + client_create_id: str | None = None adapter_name: str # ``config`` is None for full-parameter training (no LoRA adapter) and a # serialized LoraConfig string for LoRA training. @@ -418,6 +412,12 @@ class TrainingProgressResponse(ResponseModel): # --- Other responses --- +class TrainingHandleResponse(ResponseModel): + training_id: str + adapter_name: str + + class CreateResponse(ResponseModel): """Response for /create endpoint.""" status: str = 'ok' + training_handles: bool = False diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 3ebdce639..c3b7b49be 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -192,17 +192,24 @@ async def _drain(): except Exception as exc: q.put(exc) finally: - q.put(_SENTINEL) + try: + await async_gen_coro.aclose() + finally: + q.put(_SENTINEL) - asyncio.run_coroutine_threadsafe(_drain(), self._async_loop) - - while True: - item = q.get() - if item is _SENTINEL: - break - if isinstance(item, Exception): - raise item - yield item + generation = asyncio.run_coroutine_threadsafe(_drain(), self._async_loop) + try: + while True: + item = q.get() + if item is _SENTINEL: + break + if isinstance(item, Exception): + raise item + yield item + finally: + # Queue termination or closing the Python iterator is not evidence + # that the engine-loop task has exited. Join before releasing weights. + generation.result() async def _create_engine_async(self, engine_cls, model_id, engine_kwargs): """Create engine in async context to ensure output_handler starts correctly.""" @@ -781,7 +788,7 @@ def sample_stream( def sample_stream_to_queue(self, queue, inputs, sampling_params=None, adapter_name='', adapter_path=None): """Push streaming deltas to a cross-process Ray queue.""" - from twinkle.server.sampler.backends import stream_to_queue + from twinkle.server.sampler.backends.streaming import stream_to_queue stream_to_queue(self, queue, inputs, sampling_params, adapter_name, adapter_path) @remote_function(dispatch='all', collect='first') diff --git a/src/twinkle/server/__init__.py b/src/twinkle/server/__init__.py index 5bdb1f97c..264552bfc 100644 --- a/src/twinkle/server/__init__.py +++ b/src/twinkle/server/__init__.py @@ -1,7 +1,8 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .launcher import ServerLauncher, launch_server +import sys -__all__ = [ - 'ServerLauncher', - 'launch_server', -] +from twinkle._lazy_module import _LazyModule + +sys.modules[__name__] = _LazyModule( + __name__, __file__, {'launcher.server_launcher': ['ServerLauncher', 'launch_server']}, module_spec=__spec__) +sys.modules[__name__].__all__ = ['ServerLauncher', 'launch_server'] diff --git a/src/twinkle/server/checkpoint/checkpoint_manager.py b/src/twinkle/server/checkpoint/checkpoint_manager.py index 36cd30da4..322fea2ec 100644 --- a/src/twinkle/server/checkpoint/checkpoint_manager.py +++ b/src/twinkle/server/checkpoint/checkpoint_manager.py @@ -27,7 +27,7 @@ from twinkle import get_logger from twinkle.hub import HubOperation from twinkle.protocol.types import ResolvedLoadPath -from twinkle.server.task_queue.types import UserTaskError +from twinkle.server.exceptions import RequestRejectedError from twinkle.utils.parallel import FileLock, acquire_lock, release_lock from .paths import CHECKPOINT_INFO_FILENAME, ensure_inside, validate_user_path from .training_run_manager import BaseFileManager, BaseTrainingRunManager @@ -371,7 +371,7 @@ async def save_sampler(self, model_id: str, name: str | None, """ explicit_name = self.get_ckpt_name(name) if name else None if explicit_name == 'latest': - raise UserTaskError('Checkpoint name "latest" is reserved for unnamed sampler saves.') + raise RequestRejectedError('Checkpoint name "latest" is reserved for unnamed sampler saves.') version_id = uuid.uuid4().hex checkpoint_name = explicit_name or 'latest' uri_name = explicit_name or f'live_{version_id}' @@ -457,12 +457,24 @@ def delete(self, model_id: str, checkpoint_id: str) -> bool: """Remove a checkpoint and refresh only its kind's last-checkpoint pointer. Sampler links own versions under .versions; legacy directories remain - supported. Missing checkpoints or traversal attempts return False. + supported. Missing checkpoints, traversal attempts, and paths that are + not a real checkpoint all return False. """ # Basic safety check to prevent directory traversal if '..' in checkpoint_id: return False + # Delete only a checkpoint that the manager itself can resolve. This + # rejects kind roots (``weights``), the run root (``.``), aliases whose + # metadata names another checkpoint (``sampler_weights/latest``), and + # any directory that has no checkpoint metadata. + try: + checkpoint = self.get(model_id, checkpoint_id) + except ValueError: + return False + if checkpoint is None or checkpoint.checkpoint_id != checkpoint_id: + return False + ckpt_dir = self.get_ckpt_dir(model_id, checkpoint_id) if not ckpt_dir.exists(): diff --git a/src/twinkle/server/config/__init__.py b/src/twinkle/server/config/__init__.py index 3e3fe2f71..f6cb0ca2e 100644 --- a/src/twinkle/server/config/__init__.py +++ b/src/twinkle/server/config/__init__.py @@ -2,10 +2,12 @@ """Server configuration package — aggregate root and per-deployment specs.""" from .application_spec import (ApplicationSpec, DataPlaneArgs, HttpOptions, ModelArgs, ProcessorArgs, SamplerArgs, - ServerArgs, SpectralHybridArgs) + ServerArgs) from .persistence import PersistenceConfig from .server_config import ServerConfig +from .task_queue import TaskQueueConfig from .telemetry import TelemetryConfig +from .training import SpectralHybridArgs __all__ = [ 'ApplicationSpec', @@ -18,5 +20,6 @@ 'ServerArgs', 'SpectralHybridArgs', 'ServerConfig', + 'TaskQueueConfig', 'TelemetryConfig', ] diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index 545dff46e..f80132650 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -11,12 +11,12 @@ """ from __future__ import annotations -import math import os -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any, Literal -from twinkle.server.task_queue.config import TaskQueueConfig +from .task_queue import TaskQueueConfig +from .training import ModelArgs # Env var keys the launcher sets from the gateway ``server_config`` so that any # Ray worker (model / sampler / processor), not just the gateway, applies the @@ -66,69 +66,6 @@ def _ssl_pair_set_together(self) -> 'HttpOptions': # ---------- per-deployment args schemas ------------------------------------ # -class SpectralHybridArgs(_ArgsBase): - """Strict server-owned Spectral Hybrid allocation and optimizer defaults.""" - - allocation_path: str = Field(min_length=1) - default_lr_lora: float = Field(default=2.5e-5, gt=0) - default_lr_fft: float = Field(default=1.0e-6, gt=0) - - @field_validator('default_lr_lora', 'default_lr_fft') - @classmethod - def _finite_learning_rate(cls, value: float) -> float: - if not math.isfinite(value): - raise ValueError('learning rate must be finite') - return value - - -class ModelArgs(_ArgsBase): - """Args for the ``model`` deployment. - - The ``backend`` field selects the model implementation (transformers / - megatron / mock) — see ``model/app.py`` for the dispatch site that picks - the runtime class based on this value. - """ - - model_id: str - nproc_per_node: int = 1 - device_group: dict[str, Any] - device_mesh: dict[str, Any] - backend: Literal['mock', 'transformers', 'megatron'] - train_mode: Literal['lora', 'hybrid', 'full'] = 'lora' - strategy: Literal['accelerate', 'native_fsdp'] | None = None - adapter_config: dict[str, Any] | None = None - queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) - max_loras: int = 5 - max_r: int = Field(default=32, gt=0) - max_length: int | None = None - target_modules: str | list[str] = 'all-linear' - hybrid: SpectralHybridArgs | None = None - data_plane_url: str | None = None - model_loader: str | None = None - #: ``module:attr`` import spec naming the model class to serve, e.g. - #: ``swift.dev.model.multi_lora_sentence_transformer_model:MultiLoraSentenceTransformerModel``. - #: twinkle resolves it inside the replica and wraps the class in its own tinker/twinkle compat - #: layer, so a family whose runtime object is not a plain AutoModel (a sentence-transformers - #: pipeline, say) can be served by the same deployment. Only meaningful for a LoRA-mode - #: transformers deployment; every other combination is rejected below instead of being ignored. - model_cls: str | None = None - - @model_validator(mode='after') - def _validate_spectral_backend(self): - if self.strategy is not None and self.backend != 'transformers': - raise ValueError('strategy is only supported by the transformers backend') - if self.hybrid is not None: - if self.backend != 'transformers': - raise ValueError('hybrid is only supported by the transformers backend') - if self.train_mode == 'full': - raise ValueError('hybrid is not supported with train_mode="full"') - if self.model_cls is not None: - if self.backend != 'transformers' or self.train_mode != 'lora' or self.hybrid is not None: - raise ValueError('model_cls is only supported by a transformers backend in train_mode="lora" ' - 'without hybrid') - return self - - class SamplerArgs(_ArgsBase): """Args for the ``sampler`` deployment. @@ -150,6 +87,27 @@ class SamplerArgs(_ArgsBase): queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) data_plane_url: str | None = None template: dict[str, Any] | None = None + min_replicas: int = Field(default=1, ge=0, strict=True) + max_replicas: int = Field(default=1, gt=0, strict=True) + idle_timeout_s: float = Field(default=600, gt=0, allow_inf_nan=False) + target_ongoing_requests: int = Field(default=16, gt=0, strict=True) + + @model_validator(mode='after') + def _validate_pool(self): + if self.min_replicas > self.max_replicas: + raise ValueError('sampler min_replicas must not exceed max_replicas') + ranks = self.device_group.get('ranks', self.nproc_per_node) + width = self.device_group.get('gpus_per_worker', 1) + if type(ranks) is not int or ranks <= 0 or type(width) is not int or width <= 0 or ranks % width: + raise ValueError('sampler device_group requires a positive device count divisible by gpus_per_worker') + if self.sampler_type != 'mock': + from twinkle.utils.device_mesh import DeviceMesh + if DeviceMesh.from_sizes(**self.device_mesh).world_size != ranks: + raise ValueError('sampler device_mesh world_size must equal its device count') + tp = (self.engine_args or {}).get('tensor_parallel_size', width) + if tp != width: + raise ValueError('sampler tensor_parallel_size must equal gpus_per_worker') + return self class ServerStateArgs(_ArgsBase): @@ -283,7 +241,7 @@ class ApplicationSpec(BaseModel): # ``import_path`` and defaults a missing block to ``{}``), so the field is # required here — the validator runs first and fills it. args: ServerArgs | ModelArgs | SamplerArgs | ProcessorArgs | DataPlaneArgs - deployments: list[dict[str, Any]] = Field(default_factory=list) + deployments: list[dict[str, Any]] = Field(default_factory=list, max_length=1) @model_validator(mode='before') @classmethod @@ -309,23 +267,3 @@ def _coerce_args_to_schema(cls, data: Any) -> Any: # ``schema.model_validate`` rejects a non-dict itself with a clean # error, so no separate non-dict guard is needed here. return {**data, 'args': schema.model_validate(raw_args)} - - @model_validator(mode='after') - def _validate_full_mode_single_replica(self) -> ApplicationSpec: - """``train_mode: full`` requires exactly one replica. - - The exclusive-tenant lock lives in per-replica memory - (``ModelManagement._resource_records``), so more than one replica would - silently allow one full-parameter tenant per replica, each rewriting - its own copy of the base weights. Reject that at config-load time. - """ - if not (isinstance(self.args, ModelArgs) and self.args.train_mode == 'full'): - return self - for dep in self.deployments: - num_replicas = dep.get('num_replicas') - max_replicas = (dep.get('autoscaling_config') or {}).get('max_replicas') - if (num_replicas or 1) > 1 or (max_replicas or 1) > 1: - raise ValueError(f"Application '{self.name}': train_mode='full' is an exclusive single-tenant " - 'mode and requires a single replica; set num_replicas/autoscaling_config.' - 'max_replicas to 1.') - return self diff --git a/src/twinkle/server/config/resources.py b/src/twinkle/server/config/resources.py new file mode 100644 index 000000000..da022d986 --- /dev/null +++ b/src/twinkle/server/config/resources.py @@ -0,0 +1,59 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared instance startup, allocation, shutdown and monitoring budgets.""" + +import os + +from pydantic import BaseModel, ConfigDict, Field + + +class ResourceOptions(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + allocation_timeout_s: float = Field(default=3, gt=0, allow_inf_nan=False) + startup_timeout_s: float = Field(default=1800, gt=0, allow_inf_nan=False) + shutdown_timeout_s: float = Field(default=60, gt=0, allow_inf_nan=False) + sweep_interval_s: float = Field(default=5, gt=0, allow_inf_nan=False) + + +SHARED_WORKER_ENV_KEYS: tuple[str, ...] = ( + "TWINKLE_STATE_RAY_NAMESPACE", + "TWINKLE_TELEMETRY_ENABLED", + "TWINKLE_TELEMETRY_DEBUG", + "TWINKLE_TELEMETRY_SERVICE", + "TWINKLE_TELEMETRY_ENDPOINT", + "TWINKLE_TELEMETRY_INTERVAL", + "TWINKLE_MODEL_ID_ALIASES", +) + + +def build_shared_env_vars() -> dict[str, str]: + return {key: os.environ[key] for key in SHARED_WORKER_ENV_KEYS if key in os.environ} + + +def build_persistence_env_vars() -> dict[str, str]: + from twinkle.server.config.persistence import PERSISTENCE_ENV_KEYS + + return {key: os.environ[key] for key in PERSISTENCE_ENV_KEYS if key in os.environ} + + +def build_server_state_env_vars() -> dict[str, str]: + from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS + + return {key: os.environ[key] for key in SERVER_STATE_ENV_KEYS if key in os.environ} + + +def build_propagated_env_vars() -> dict[str, str]: + merged: dict[str, str] = {} + merged.update(build_shared_env_vars()) + merged.update(build_persistence_env_vars()) + merged.update(build_server_state_env_vars()) + return merged + + +def build_worker_env_vars() -> dict[str, str]: + """Policy plus the explicit process settings shared by launcher and resource owners.""" + env = build_propagated_env_vars() + for key in ("PYTHONPATH", "TWINKLE_DEFAULT_SAVE_DIR", "TMPDIR", "OMP_NUM_THREADS"): + if key in os.environ: + env.setdefault(key, os.environ[key]) + return env diff --git a/src/twinkle/server/config/server_config.py b/src/twinkle/server/config/server_config.py index 82bac4d44..7de25f094 100644 --- a/src/twinkle/server/config/server_config.py +++ b/src/twinkle/server/config/server_config.py @@ -19,10 +19,12 @@ from typing import Any from twinkle.server.exceptions import ConfigParseError -from twinkle.server.task_queue.config import TaskQueueConfig from .application_spec import ApplicationSpec, HttpOptions from .persistence import PersistenceConfig +from .resources import ResourceOptions +from .task_queue import TaskQueueConfig from .telemetry import TelemetryConfig +from .training import TrainingConfig, TrainingModelConfig class ServerConfig(BaseModel): @@ -36,6 +38,7 @@ class ServerConfig(BaseModel): telemetry: TelemetryConfig = Field(default_factory=TelemetryConfig) persistence: PersistenceConfig = Field(default_factory=PersistenceConfig) task_queue: TaskQueueConfig = Field(default_factory=TaskQueueConfig) + resources: ResourceOptions = Field(default_factory=ResourceOptions) applications: list[ApplicationSpec] = Field(default_factory=list) # ---- loading ---------------------------------------------------------- # @@ -69,10 +72,30 @@ def from_yaml(cls, path: str | Path) -> ServerConfig: @model_validator(mode='after') def _validate_cross_field(self) -> ServerConfig: + if sum(app.import_path == 'server' for app in self.applications) > 1: + raise ValueError('configure one gateway per server') if self.persistence.mode == 'redis' and not self.persistence.redis_url: raise ValueError("persistence.redis_url is required when persistence.mode == 'redis'", ) + if self.training_config.models: + static = [app for app in self.applications if app.import_path != 'model'] + names = {app.name for app in static} + for model in self.training_config.models: + if model.name in names: + raise ValueError(f'training model name conflicts with application: {model.name}') + path = model.route_prefix + if any(app.route_prefix.rstrip('/') == path for app in static): + raise ValueError(f'training route conflicts with application: {path}') return self + @property + def training_config(self) -> TrainingConfig: + """Normalize model application args into the internal runtime catalog.""" + return TrainingConfig( + **self.resources.model_dump(), + models=[ + TrainingModelConfig.from_application(app) for app in self.applications if app.import_path == 'model' + ]) + # ---- round-trip / serialization -------------------------------------- # def to_yaml_dict(self) -> dict[str, Any]: diff --git a/src/twinkle/server/task_queue/config.py b/src/twinkle/server/config/task_queue.py similarity index 78% rename from src/twinkle/server/task_queue/config.py rename to src/twinkle/server/config/task_queue.py index 76f1813e8..4dede86b6 100644 --- a/src/twinkle/server/task_queue/config.py +++ b/src/twinkle/server/config/task_queue.py @@ -1,12 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -""" -Task queue configuration. - -Provides TaskQueueConfig (Pydantic) for controlling rate limits, timeouts, -and queue behavior. Constraints are validated at construction time so an -invalid YAML/dict value is rejected before the deployment reaches a ready -state. -""" +"""Task queue configuration.""" from __future__ import annotations from pydantic import BaseModel, ConfigDict, Field @@ -55,20 +48,13 @@ class TaskQueueConfig(BaseModel): @property def effective_execution_timeout(self) -> float: - """The single source of the execution time bound. - - ``0`` is not rejected (that would fail existing deployments); it is read - as "no configured limit" and replaced by a finite fallback so the bound - is always positive. This value feeds both ``_ray_get_timeout`` and the - ComputeWorker's ``asyncio.wait_for`` -- there is no second, independently - configurable timeout. - """ + """Return the single execution time bound used by all queue paths.""" if self.execution_timeout > 0: return self.execution_timeout return _ZERO_EXECUTION_TIMEOUT_FALLBACK def absolute_future_ttl(self, collect_width: int) -> float: - """Conservative lifetime for a non-terminal future record.""" + """Return a conservative lifetime for a non-terminal future record.""" ray_timeout = max(self.effective_execution_timeout, _MAX_DECLARED_BACKEND_TIMEOUT) resource_bound = max(1, collect_width) * ray_timeout return _ABSOLUTE_TTL_MULTIPLIER * (self.queue_timeout + resource_bound) diff --git a/src/twinkle/server/config/training.py b/src/twinkle/server/config/training.py new file mode 100644 index 000000000..78723fe0b --- /dev/null +++ b/src/twinkle/server/config/training.py @@ -0,0 +1,289 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Training YAML arguments and the controller's validated runtime configuration.""" + +from __future__ import annotations +import math +from typing import TYPE_CHECKING, Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny, ValidationInfo, field_validator, model_validator + +from twinkle.hub.model_alias import extract_model_alias +from twinkle.server.exceptions import RequestRejectedError +from .resources import ResourceOptions +from .task_queue import TaskQueueConfig + +if TYPE_CHECKING: + from .application_spec import ApplicationSpec + + +class SpectralHybridArgs(BaseModel): + """Strict server-owned Spectral Hybrid allocation and optimizer defaults.""" + + model_config = ConfigDict(extra="forbid") + + allocation_path: str = Field(min_length=1) + default_lr_lora: float = Field(default=2.5e-5, gt=0) + default_lr_fft: float = Field(default=1.0e-6, gt=0) + + @field_validator("default_lr_lora", "default_lr_fft") + @classmethod + def _finite_learning_rate(cls, value: float) -> float: + if not math.isfinite(value): + raise ValueError("learning rate must be finite") + return value + + +class AdapterConfig(BaseModel): + """Session liveness and absolute lifetime of a training handle, in seconds.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + adapter_timeout: float | None = Field(default=1800, gt=0, allow_inf_nan=False) + adapter_max_lifetime: float | None = Field(default=36000, gt=0, allow_inf_nan=False) + + +class TrainingArguments(BaseModel): + """Training options shared by model YAML and resolved instance configuration.""" + + model_config = ConfigDict(extra="forbid") + + min_replicas: int = Field(default=0, ge=0, strict=True) + max_replicas: int | None = Field(default=None, gt=0, strict=True) + idle_timeout_s: float = Field(default=1800, gt=0, allow_inf_nan=False) + adapter_config: AdapterConfig = Field(default_factory=AdapterConfig) + strategy: Literal["accelerate", "native_fsdp"] | None = None + queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) + max_length: int | None = Field(default=None, gt=0, strict=True) + model_cls: str | None = None + model_loader: str | None = None + gradient_checkpointing: bool | None = Field(default=None, strict=True) + mixed_precision: Literal["no", "fp8", "fp16", "bf16"] | None = None + task_type: Literal["causal_lm", "seq_cls"] | None = None + num_labels: int | None = Field(default=None, gt=0, strict=True) + language_model_only: bool | None = Field(default=None, strict=True) + + @model_validator(mode="after") + def _validate_head(self): + if self.task_type == "seq_cls" and self.num_labels is None: + raise ValueError('task_type="seq_cls" requires num_labels') + if self.num_labels is not None and self.task_type != "seq_cls": + raise ValueError('num_labels requires task_type="seq_cls"') + return self + + @model_validator(mode="after") + def _validate_replicas(self): + if self.max_replicas is not None and self.min_replicas > self.max_replicas: + raise ValueError("min_replicas must not exceed max_replicas") + return self + + +class LoraArguments(BaseModel): + """LoRA parameters shared by YAML and LoRA instance configuration.""" + + max_loras: int = Field(default=5, gt=0, strict=True) + max_r: int = Field(default=32, gt=0, strict=True) + target_modules: str | list[str] = "all-linear" + hybrid: SpectralHybridArgs | None = None + + +class ModelArgs(TrainingArguments, LoraArguments): + """Args for the ``model`` deployment. + + The ``backend`` field selects the model implementation (transformers / + megatron / mock) — see ``model/app.py`` for the dispatch site that picks + the runtime class based on this value. + """ + + model_id: str + nproc_per_node: int = 1 + device_group: dict[str, Any] + device_mesh: dict[str, Any] + backend: Literal["mock", "transformers", "megatron"] + env_vars: dict[str, str] = Field(default_factory=dict) + train_mode: Literal["lora", "hybrid", "full"] = "lora" + data_plane_url: str | None = None + + @model_validator(mode="after") + def _validate_spectral_backend(self): + if self.strategy is not None and self.backend != "transformers": + raise ValueError("strategy is only supported by the transformers backend") + if self.backend != "transformers" and ( + self.gradient_checkpointing is not None or self.mixed_precision is not None + ): + raise ValueError("gradient_checkpointing and mixed_precision require the transformers backend") + if self.hybrid is not None: + if self.backend != "transformers": + raise ValueError("hybrid is only supported by the transformers backend") + if self.train_mode == "full": + raise ValueError('hybrid is not supported with train_mode="full"') + if self.model_cls is not None: + if self.backend != "transformers" or self.train_mode != "lora" or self.hybrid is not None: + raise ValueError("model_cls requires a transformers LoRA deployment without hybrid") + if self.model_loader is not None and self.backend != "transformers": + raise ValueError("model_loader requires the transformers backend") + if self.language_model_only is not None and self.backend != "megatron": + raise ValueError("language_model_only requires the megatron backend") + return self + + +class TrainingMesh(BaseModel): + """Sizes accepted by DeviceMesh.from_sizes, including overlapping EP axes.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + device_type: str | None = None + dp_size: int = Field(default=1, gt=0, strict=True) + fsdp_size: int | None = Field(default=None, gt=0, strict=True) + tp_size: int | None = Field(default=None, gt=0, strict=True) + pp_size: int | None = Field(default=None, gt=0, strict=True) + cp_size: int | None = Field(default=None, gt=0, strict=True) + ep_size: int | None = Field(default=None, gt=0, strict=True) + etp_size: int = Field(default=1, gt=0, strict=True) + ep_fsdp_size: int | None = Field(default=None, gt=0, strict=True) + vpp_size: int | None = Field(default=None, gt=0, strict=True) + ulysses_size: int | None = Field(default=None, gt=0, strict=True) + sequence_parallel: bool = False + + def world_size(self) -> int: + from twinkle.utils.device_mesh import DeviceMesh + + return DeviceMesh.from_sizes(**self.model_dump(exclude_none=True)).world_size + + +class TrainingModeConfig(TrainingArguments): + model_config = ConfigDict(extra="forbid", frozen=True) + + num_devices: int = Field(gt=0, strict=True) + device_mesh: TrainingMesh + + @model_validator(mode="after") + def _validate_resources(self): + if self.device_mesh.world_size() != self.num_devices: + raise ValueError("device_mesh world_size must equal num_devices") + return self + + +class LoraTrainingConfig(TrainingModeConfig, LoraArguments): + pass + + +class TrainingModelConfig(BaseModel): + """One application exposes one training mode and one instance specification.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + name: str = Field(min_length=1) + model_id: str = Field(min_length=1) + backend: Literal["mock", "transformers", "megatron"] + device_type: Literal["cuda", "npu", "mps", "cpu"] + mode: Literal["lora", "full"] + config: SerializeAsAny[TrainingModeConfig] + route_prefix: str | None = None + data_plane_url: str | None = None + env_vars: dict[str, str] = Field(default_factory=dict) + deploy_options: dict = Field(default_factory=dict) + + @field_validator("config", mode="before") + @classmethod + def _instance_config(cls, value, info: ValidationInfo): + config_type = LoraTrainingConfig if info.data.get("mode") == "lora" else TrainingModeConfig + data = value.model_dump() if isinstance(value, TrainingModeConfig) else value + return config_type.model_validate(data) + + @model_validator(mode="after") + def _validate_backend(self): + if ( + self.name.startswith("/") + or any(part in ("", ".", "..") for part in self.name.split("/")) + or any(char.isspace() or char in "?#%\\" for char in self.name) + ): + raise ValueError("name must be a nonempty relative model path without dot segments") + if self.device_type == "cpu" and self.backend != "mock": + raise ValueError("cpu is supported only by the mock training backend") + if self.backend == "megatron" and self.config.num_devices < 2: + raise ValueError("megatron training instances require at least 2 devices") + if self.config.device_mesh.device_type not in (None, self.device_type): + raise ValueError("device_mesh.device_type must match device_type") + if self.config.strategy is not None and self.backend != "transformers": + raise ValueError("strategy requires the transformers backend") + if self.backend != "transformers" and ( + self.config.gradient_checkpointing is not None or self.config.mixed_precision is not None + ): + raise ValueError("gradient_checkpointing and mixed_precision require the transformers backend") + if isinstance(self.config, LoraTrainingConfig) and self.config.hybrid and self.backend != "transformers": + raise ValueError("hybrid requires the transformers backend") + if self.config.model_cls is not None: + if self.backend != "transformers" or self.mode != "lora" or self.config.hybrid is not None: + raise ValueError("model_cls requires a transformers LoRA deployment without hybrid") + if self.config.model_loader is not None and self.backend != "transformers": + raise ValueError("model_loader requires the transformers backend") + return self + + @property + def key(self) -> tuple[str, str]: + return self.name, self.mode + + @property + def capacity(self) -> int: + return self.config.max_loras if isinstance(self.config, LoraTrainingConfig) else 1 + + @classmethod + def from_application(cls, app: ApplicationSpec) -> TrainingModelConfig: + args = app.args + if not isinstance(args, ModelArgs): + raise TypeError("Training configuration requires model application arguments") + if len(app.deployments) > 1: + raise ValueError(f"{app.name}: a model application must have one deployment configuration") + deployment = app.deployments[0] if app.deployments else {} + device_type = str( + args.device_group.get( + "device_type", args.device_mesh.get("device_type", "cpu" if args.backend == "mock" else "cuda") + ) + ).lower() + mode = "full" if args.train_mode == "full" else "lora" + fields = set(TrainingArguments.model_fields) + if mode == "lora": + fields.update(LoraArguments.model_fields) + config = {key: getattr(args, key) for key in fields} + config.update( + num_devices=args.device_group.get("ranks", args.nproc_per_node), + device_mesh={**args.device_mesh, "device_type": device_type}, + ) + env_vars = { + **(deployment.get("ray_actor_options") or {}).get("runtime_env", {}).get("env_vars", {}), + **args.env_vars, + } + return cls( + name=extract_model_alias(app.route_prefix) or app.name, + model_id=args.model_id, + backend=args.backend, + device_type=device_type, + mode=mode, + config=config, + route_prefix=app.route_prefix.rstrip("/") or "/", + data_plane_url=args.data_plane_url, + env_vars=env_vars, + deploy_options={key: value for key, value in deployment.items() if key != "name"}, + ) + + +class TrainingConfig(ResourceOptions): + """Resolved catalog; public model entries stay in applications.""" + + models: tuple[TrainingModelConfig, ...] = () + + @model_validator(mode="after") + def _unique_names(self): + names = [model.name for model in self.models] + if len(names) != len(set(names)): + raise ValueError("training model names must be unique") + return self + + def get(self, name: str, mode: Literal["lora", "full"]) -> TrainingModelConfig: + """Resolve a declared specification without allocating resources.""" + model = next((model for model in self.models if model.name == name), None) + if model is None: + raise RequestRejectedError(f"Unknown training model: {name}") + if model.mode != mode: + raise RequestRejectedError(f"Training model {name} does not provide {mode} mode") + return model diff --git a/src/twinkle/server/deployment.py b/src/twinkle/server/deployment.py index 45ca43fbc..ca38798f8 100644 --- a/src/twinkle/server/deployment.py +++ b/src/twinkle/server/deployment.py @@ -329,4 +329,5 @@ async def _ensure_state_cleanup_started(self) -> None: await self.state.start_cleanup_task() except Exception as e: logger.warning(f'Failed to start ServerState cleanup task: {e}') - self._state_cleanup_started = True + else: + self._state_cleanup_started = True diff --git a/src/twinkle/server/exceptions.py b/src/twinkle/server/exceptions.py index b8985879d..5fc3cc044 100644 --- a/src/twinkle/server/exceptions.py +++ b/src/twinkle/server/exceptions.py @@ -9,6 +9,8 @@ from __future__ import annotations +from typing import Literal + from twinkle.protocol.types.errors import ErrorCategory @@ -176,3 +178,25 @@ def __init__(self, current_holder: str) -> None: super().__init__('This deployment runs in full-parameter (exclusive) mode and is already ' f'held by another training task ({current_holder}). Only one full-parameter ' 'training task is allowed at a time; retry after it is released.') + + +class TrainingExpiredError(ResourceNotFoundError): + """The training handle expired and must be recreated before explicit resume.""" + + +class ProtocolMismatchError(RequestRejectedError): + """The requested client protocol does not match the training handle.""" + + +class ResourceReleaseError(TwinkleServerError, RuntimeError): + """Owned resource removal failed at a known phase; capacity remains reserved.""" + + error_code = 503 + + def __init__(self, message: str, *, phase: Literal['worker_death', 'pg_removal']): + super().__init__(message) + self.phase = phase + + +class BackendBusyError(ResourceExhaustedError, RuntimeError): + """A timed-out backend still owns execution/weight resources on this replica.""" diff --git a/src/twinkle/server/gateway/app.py b/src/twinkle/server/gateway/app.py index 8502e5a93..415352bfe 100644 --- a/src/twinkle/server/gateway/app.py +++ b/src/twinkle/server/gateway/app.py @@ -13,10 +13,10 @@ import twinkle.protocol.types as types from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app +from twinkle.server.lifecycle.proxy import ServiceProxy from twinkle.server.state import get_server_state from twinkle.utils.logger import get_logger from .openai_handlers import _register_api_key_gate, _register_openai_routes -from .proxy import ServiceProxy from .tinker_handlers import _register_gateway_tinker_routes from .twinkle_handlers import _register_gateway_twinkle_routes @@ -41,6 +41,8 @@ def __init__(self, self.http_options = http_options or {} self.proxy = ServiceProxy(http_options=http_options, route_prefix=self.route_prefix) self.supported_models = self._normalize_models(supported_models) + self.training_models = self._normalize_models(kwargs.get('training_models', supported_models)) + self._training_model_names = frozenset(m.model_name for m in self.training_models) self._supported_model_names = frozenset(m.model_name for m in self.supported_models) # OpenAI-compat surface knobs: ``api_key`` gates every inference route with a Bearer check # (None = unauthenticated), ``owned_by`` is the ``owned_by`` field of ``/models`` entries. @@ -93,11 +95,11 @@ def _normalize_models(self, supported_models): return normalized def _validate_base_model(self, base_model: str) -> None: - if base_model not in self._supported_model_names: + if base_model not in self._training_model_names: raise HTTPException( status_code=400, detail=f"Base model '{base_model}' is not supported. " - f"Supported models: {', '.join(sorted(self._supported_model_names))}") + f"Supported models: {', '.join(sorted(self._training_model_names))}") async def _get_base_model(self, model_id: str) -> str: metadata = await self.state.get_model_metadata(model_id) diff --git a/src/twinkle/server/gateway/openai_handlers.py b/src/twinkle/server/gateway/openai_handlers.py index 5495cb0c7..5a5d63733 100644 --- a/src/twinkle/server/gateway/openai_handlers.py +++ b/src/twinkle/server/gateway/openai_handlers.py @@ -8,7 +8,9 @@ """ from __future__ import annotations +import asyncio import json +import os import secrets import time import uuid @@ -19,16 +21,16 @@ from twinkle.protocol.headers import H_AUTH, H_AUTH_TWINKLE, build_routing_headers from twinkle.server.lifecycle.envelope import envelope_from_record -from twinkle.server.lifecycle.poll_config import long_poll_window if TYPE_CHECKING: from .app import GatewayServer import httpx +import twinkle.protocol.types as types + from twinkle.server.utils import get_template_for_model from twinkle.utils.logger import get_logger -from . import routes from .openai_bridge import (make_error, translate_chat_request, translate_completion_request, translate_completion_response, translate_completion_stream_chunk, translate_embedding_request, translate_embedding_response, @@ -39,6 +41,59 @@ logger = get_logger() +_OPENAI_SAMPLE_TIMEOUT_ENV = 'TWINKLE_OPENAI_SAMPLE_TIMEOUT' +_DEFAULT_OPENAI_SAMPLE_TIMEOUT = 1800.0 + + +class SamplerWaitTimeoutError(TimeoutError): + """The sampler task did not reach a terminal state inside the gateway's budget.""" + + +def _sample_wait_timeout() -> float: + return float(os.environ.get(_OPENAI_SAMPLE_TIMEOUT_ENV, str(_DEFAULT_OPENAI_SAMPLE_TIMEOUT))) + + +async def _resolve_sampler_envelope(gateway: GatewayServer, sampler_data: dict[str, Any], + request: Request) -> types.TaskEnvelope | None: + """Resolve the sampler's TaskEnvelope to a terminal state. + + ``/twinkle/sample`` always answers with a TaskEnvelope; a slow sample is + still ``queued``/``running`` when the proxy returns. Polling the shared + state backend is the only way to obtain its result without changing the + upstream wire contract. Returns ``None`` when the future record was lost. + + ``TWINKLE_OPENAI_SAMPLE_TIMEOUT`` bounds the entire terminal-state wait, + including state reads. Client disconnection is checked between long polls. + """ + envelope = types.TaskEnvelope.model_validate(sampler_data) + if envelope.status in types.TERMINAL_STATUSES: + return envelope + + async def wait_for_terminal(envelope: types.TaskEnvelope) -> types.TaskEnvelope | None: + while envelope.status not in types.TERMINAL_STATUSES: + if await request.is_disconnected(): + raise SamplerWaitTimeoutError('Client disconnected while waiting for the sampling task.') + # ``poll_future`` waits out its whole window before reporting a missing + # record, so check once here: a lost record is a terminal answer, not a + # reason to keep the coroutine alive. + if await gateway.state.get_future(envelope.request_id) is None: + return None + outcome = await poll_future(gateway.state, envelope.request_id) + if outcome.record is None: + return None + envelope = envelope_from_record(envelope.request_id, outcome.record) + return envelope + + timeout = _sample_wait_timeout() + try: + return await asyncio.wait_for(wait_for_terminal(envelope), timeout=timeout) + except SamplerWaitTimeoutError: + raise + except asyncio.TimeoutError: + raise SamplerWaitTimeoutError( + f'Sampling task {envelope.request_id} did not finish within {timeout:g}s.') from None + + def _register_openai_routes(app: FastAPI, self_fn: Callable[[], GatewayServer]) -> None: """Register OpenAI-compatible routes on the gateway FastAPI app.""" @@ -293,30 +348,6 @@ def _deny_without_key(request: Request, api_key: str | None) -> JSONResponse | N return None -async def _await_terminal_envelope(gateway: GatewayServer, request_id: str) -> dict[str, Any]: - """Long-poll one queued task's future record to a terminal state (or the window's end). - - Uses the same future store, the same ``poll_future`` use case, and the same Long_Poll_Window - that the twinkle-native ``/twinkle/retrieve_future`` endpoint reads, so the OpenAI surface and - the native surface wait identically. - """ - outcome = await poll_future(gateway.state, request_id) - if outcome.record is None: - return {'request_id': request_id, 'status': 'pending'} - return envelope_from_record(request_id, outcome.record).model_dump() - - -def _task_failed(envelope: dict[str, Any]) -> JSONResponse: - """Map a failed/cancelled TaskEnvelope's ErrorPayload to an OpenAI-shaped error response.""" - error = envelope.get('error') or {} - message = error.get('message') or envelope.get('status') or 'task failed' - error_code = error.get('error_code') - return JSONResponse( - status_code=error_code if isinstance(error_code, int) else 500, - content=make_error(message=f'Task failed: {message}', error_type='server_error'), - ) - - async def _proxy_sample( gateway: GatewayServer, request: Request, @@ -365,24 +396,31 @@ async def _proxy_sample( ), ) - envelope = json.loads(response.body) - if envelope.get('status') != 'completed': - envelope = await _await_terminal_envelope(gateway, envelope.get('request_id', '')) - if envelope.get('status') == 'failed' or envelope.get('status') == 'cancelled': - return _task_failed(envelope) - if envelope.get('status') != 'completed': - return JSONResponse( - status_code=504, - content=make_error( - message=(f'Task {envelope.get("request_id")} is still running after the ' - f'{long_poll_window():.0f}s long-poll window; poll /twinkle/retrieve_future ' - 'with its request_id'), - error_type='server_error', - code='task_timeout', - ), - ) - - result = translate(envelope.get('result'), model) + sampler_data = json.loads(response.body) + try: + envelope = await _resolve_sampler_envelope(gateway, sampler_data, request) + except SamplerWaitTimeoutError as exc: + return JSONResponse(status_code=504, content=make_error(message=str(exc), error_type='timeout')) + if envelope is None: + return JSONResponse( + status_code=500, + content=make_error( + message=f"Sampling task {sampler_data.get('request_id')} was lost by the server", + error_type='server_error')) + if envelope.status != 'completed': + error = envelope.error + return JSONResponse( + status_code=error.error_code if error is not None else 500, + content=make_error( + message=error.error if error is not None else f'Sampling task {envelope.status}', + error_type=('invalid_request_error' + if error is not None and error.category.value == 'user' else 'server_error'))) + payload = envelope.result or {} + expected_key = 'data' if endpoint == 'twinkle/encode' else 'samples' + if not payload.get(expected_key): + return JSONResponse(status_code=502, content=make_error( + message=f'Sampler completed without returning any {expected_key}', error_type='server_error')) + result = translate(payload, model) resp_headers = {} replica_id = response.headers.get('x-twinkle-replica-id') if replica_id: @@ -479,14 +517,13 @@ async def _ensure_template( template_cls = get_template_for_model(base_model) set_template_body = json.dumps({ 'template_cls': template_cls, - 'model_id': base_model, 'adapter_name': '', }).encode() try: resp = await gateway.proxy.proxy_request( request, - endpoint=routes.TWINKLE_SET_TEMPLATE, + endpoint='twinkle/set_template', base_model=base_model, service_type='sampler', body_override=set_template_body, diff --git a/src/twinkle/server/gateway/proxy.py b/src/twinkle/server/gateway/proxy.py deleted file mode 100644 index 181749e90..000000000 --- a/src/twinkle/server/gateway/proxy.py +++ /dev/null @@ -1,242 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -""" -Proxy utilities for forwarding requests to internal model/sampler services. - -Moved from tinker/proxy.py. Updated proxy_to_model and proxy_to_sampler -to prepend the 'tinker/' prefix to endpoints so they route to /tinker/* paths -on the unified model/sampler deployments. -""" -from __future__ import annotations - -import httpx -from fastapi import Request, Response -from fastapi.responses import JSONResponse -from typing import Any - -from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID, H_REQUEST_ID_LEGACY -from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload -from twinkle.server.telemetry.tracing import inject_context -from twinkle.utils.logger import get_logger -from . import routes - -logger = get_logger() - - -class ServiceProxy: - """HTTP proxy for routing requests to internal model and sampler services. - - Handles: - 1. URL construction using localhost to avoid external routing loops - 2. Header forwarding with appropriate cleanup - 3. Debug logging for troubleshooting - 4. Error handling and response forwarding - - Tinker endpoints are routed to /tinker/ on the unified deployments. - """ - - def __init__( - self, - http_options: dict[str, Any] | None = None, - route_prefix: str = '/api/v1', - ): - self.http_options = http_options or {} - self.route_prefix = route_prefix - # Disable proxy env vars to avoid external routing. Verification is off iff TLS is on: - # every outbound target is this deployment's own listener (localhost), whose certificate is - # whatever ssl_certfile configured -- typically self-signed and in no CA store, so default - # verification would fail every internal call with CERTIFICATE_VERIFY_FAILED. This is a - # self-connection, not a trust decision about a remote peer. - self.client = httpx.AsyncClient( - timeout=None, - trust_env=False, - verify=not self.http_options.get('ssl_certfile'), - ) - - async def close(self) -> None: - """Close the underlying httpx.AsyncClient to release connections.""" - await self.client.aclose() - - def _build_target_url(self, service_type: str, base_model: str, endpoint: str) -> str: - """Build the target URL for internal service routing. - - Args: - service_type: Either 'model' or 'sampler' - base_model: The base model name for routing - endpoint: The target endpoint name (already includes tinker/ or twinkle/ prefix) - - Returns: - Complete target URL for the internal service - """ - host = self.http_options.get('host', 'localhost') - port = self.http_options.get('port', 8000) - - if host == '0.0.0.0': - host = 'localhost' - - # The listener these URLs target serves HTTPS when the deployment set an ssl_certfile - # (HttpOptions forwards the pair to Ray Serve verbatim), so the outbound scheme has to - # follow it -- a fixed ``http://`` would send plaintext to a TLS port and fail every proxy. - scheme = 'https' if self.http_options.get('ssl_certfile') else 'http' - base_url = f'{scheme}://{host}:{port}' - return f'{base_url}{routes.target_url(self.route_prefix, service_type, base_model, endpoint)}' - - def _prepare_headers(self, request_headers) -> dict[str, str]: - """Prepare headers for proxying by removing problematic headers.""" - logger.debug('prepare_headers request_headers=%s', request_headers) - headers = dict(request_headers) - headers.pop('host', None) - headers.pop('content-length', None) - request_id = request_headers.get(H_REQUEST_ID) or request_headers.get(H_REQUEST_ID_LEGACY) - if request_id is not None and not request_headers.get(H_MULTIPLEX_LEGACY): - headers[H_MULTIPLEX_LEGACY] = request_id - if request_id is not None: - headers.setdefault(H_REQUEST_ID, request_id) - headers.setdefault(H_MULTIPLEX, request_id) - return headers - - async def proxy_request( - self, - request: Request, - endpoint: str, - base_model: str, - service_type: str, - *, - body_override: bytes | None = None, - extra_headers: dict[str, str] | None = None, - ) -> Response: - """Generic proxy method to forward requests to model or sampler services. - - Args: - request: The incoming FastAPI request - endpoint: The target endpoint path (e.g., 'tinker/create_model') - base_model: The base model name for routing - service_type: Either 'model' or 'sampler' - body_override: If set, use this as the request body instead of reading from request - extra_headers: Additional headers to merge into the forwarded request - - Returns: - Proxied response from the target service - """ - body_bytes = body_override if body_override is not None else await request.body() - target_url = self._build_target_url(service_type, base_model, endpoint) - headers = self._prepare_headers(request.headers) - if extra_headers: - headers.update(extra_headers) - - # Inject current trace context into outgoing headers for distributed tracing. - # When telemetry is not initialized, this is a noop. - inject_context(headers) - - try: - logger.debug( - 'proxy_request service=%s endpoint=%s target_url=%s request_id=%s', - service_type, - endpoint, - target_url, - headers.get(H_MULTIPLEX_LEGACY), - ) - - response = await self.client.request( - method=request.method, - url=target_url, - content=body_bytes, - headers=headers, - params=request.query_params, - ) - - logger.debug( - 'proxy_response status=%s body_preview=%s', - response.status_code, - response.text[:200], - ) - - return Response( - content=response.content, - status_code=response.status_code, - headers=dict(response.headers), - media_type=response.headers.get('content-type'), - ) - except Exception as e: - logger.error('Proxy error: %s', str(e), exc_info=True) - # The gateway could not reach the upstream deployment. Return the - # unified ``ErrorPayload`` (502/Server) instead of a plain-text body - # so every gateway failure has the same wire shape. Upstream error - # responses are passed through unchanged above, preserving their own - # ErrorPayload body. - request_id = request.headers.get(H_REQUEST_ID) or request.headers.get(H_REQUEST_ID_LEGACY) or '' - payload = ErrorPayload( - error=f'Proxy Error: {str(e)}', - category=ErrorCategory.Server, - error_code=502, - request_id=request_id, - ) - return JSONResponse(status_code=502, content=payload.model_dump(mode='json', exclude_none=True)) - - async def proxy_request_stream( - self, - request: Request, - endpoint: str, - base_model: str, - service_type: str, - *, - body_override: bytes | None = None, - extra_headers: dict[str, str] | None = None, - ): - """Proxy with httpx streaming — returns an async generator of response lines. - - Yields: - str lines from the upstream response (newline-delimited JSON). - - Raises: - httpx.HTTPStatusError: If the upstream returns a non-2xx status. - """ - body_bytes = body_override if body_override is not None else await request.body() - target_url = self._build_target_url(service_type, base_model, endpoint) - headers = self._prepare_headers(request.headers) - if extra_headers: - headers.update(extra_headers) - inject_context(headers) - - logger.debug( - 'proxy_request_stream service=%s endpoint=%s target_url=%s', - service_type, - endpoint, - target_url, - ) - - async with self.client.stream( - method='POST', - url=target_url, - content=body_bytes, - headers=headers, - ) as response: - if response.status_code >= 400: - body = await response.aread() - raise httpx.HTTPStatusError( - f'Upstream returned {response.status_code}', - request=response.request, - response=httpx.Response(status_code=response.status_code, content=body), - ) - async for line in response.aiter_lines(): - if line: - yield line - - async def proxy_to_model(self, request: Request, endpoint: str, base_model: str) -> Response: - """Proxy request to model's tinker endpoint (/tinker/). - - Args: - request: The incoming FastAPI request - endpoint: The tinker endpoint name (e.g., 'create_model', 'forward') - base_model: The base model name for routing - """ - return await self.proxy_request(request, routes.tinker_endpoint(endpoint), base_model, 'model') - - async def proxy_to_sampler(self, request: Request, endpoint: str, base_model: str) -> Response: - """Proxy request to sampler's tinker endpoint (/tinker/). - - Args: - request: The incoming FastAPI request - endpoint: The tinker endpoint name (e.g., 'asample') - base_model: The base model name for routing - """ - return await self.proxy_request(request, routes.tinker_endpoint(endpoint), base_model, 'sampler') diff --git a/src/twinkle/server/gateway/routes.py b/src/twinkle/server/gateway/routes.py deleted file mode 100644 index 0cececc05..000000000 --- a/src/twinkle/server/gateway/routes.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Route-path constants shared by the gateway's proxy and its OpenAI bridge. - -Single source for the ``route_prefix`` convention these three places must agree on; they -previously agreed by having the same string typed out in three files. Placed on -the consumer side (gateway) rather than in ``launcher`` -- launcher *produces* the -``route_prefix``, gateway consumes it -- to avoid a new ``gateway -> launcher`` edge. -""" -from __future__ import annotations - -# Downstream endpoint paths the gateway proxies to (relative to the service's route prefix). -TWINKLE_SAMPLE = 'twinkle/sample' -TWINKLE_SAMPLE_STREAM = 'twinkle/sample_stream' -TWINKLE_SET_TEMPLATE = 'twinkle/set_template' - -TINKER_PREFIX = 'tinker' - - -def tinker_endpoint(endpoint: str) -> str: - """``tinker/`` -- the shape ``proxy_to_model`` / ``proxy_to_sampler`` build.""" - return f'{TINKER_PREFIX}/{endpoint}' - - -def target_url(route_prefix: str, service_type: str, base_model: str, endpoint: str) -> str: - """The one definition of ``{route_prefix}/{service_type}/{base_model}/{endpoint}``.""" - return f'{route_prefix.rstrip("/")}/{service_type}/{base_model}/{endpoint}' diff --git a/src/twinkle/server/gateway/tinker_handlers.py b/src/twinkle/server/gateway/tinker_handlers.py index 2161e712e..7cb4ed5ec 100644 --- a/src/twinkle/server/gateway/tinker_handlers.py +++ b/src/twinkle/server/gateway/tinker_handlers.py @@ -17,9 +17,8 @@ from twinkle.hub import HubOperation from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager +from twinkle.server.lifecycle.envelope import error_payload_from_failure from twinkle.server.middleware.auth import get_token_from_request -from twinkle.server.state.models import FutureFailureRecord -from twinkle.server.task_errors import trim_traceback from twinkle.utils.logger import get_logger from .use_cases import delete_checkpoint from .use_cases import get_training_run as get_training_run_use_case @@ -27,42 +26,10 @@ logger = get_logger() -# Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). -_TINKER_FAILURE_WIRE: dict[str, tuple[int, str]] = { - 'invalid_request': (400, 'user'), - 'request_rejected': (400, 'user'), - 'resource_not_found': (404, 'user'), - 'full_mode_busy': (409, 'user'), - 'input_tokens_exceeded': (422, 'user'), - 'batch_size_invalid': (422, 'user'), - 'rate_limit_exceeded': (429, 'user'), - 'resource_quota_exceeded': (429, 'user'), - 'cancelled': (499, 'user'), - 'endpoint_unavailable': (501, 'server'), - 'state_contention': (503, 'server'), - 'backend_gate_unavailable': (503, 'server'), - 'orphaned_replica': (503, 'server'), - 'execution_timeout': (504, 'server'), - 'deadline_exceeded': (500, 'server'), - 'internal_error': (500, 'server'), -} - def _tinker_error_from_failure(stored: Any, *, request_id: str) -> dict[str, Any]: - """Map domain failure state to the existing Tinker-compatible error body.""" - failure = FutureFailureRecord.model_validate(stored) - error_code, category = _TINKER_FAILURE_WIRE.get(failure.reason_code, (500, 'server')) - payload: dict[str, Any] = { - 'error': failure.message[:1024], - 'category': category, - 'error_code': error_code, - 'request_id': request_id, - } - if failure.details is not None: - payload['details'] = failure.details - if failure.diagnostic and category == 'server': - payload['traceback'] = trim_traceback(failure.diagnostic) - return payload + """Serialize the shared failure presentation with Tinker's field names.""" + return error_payload_from_failure(stored, request_id=request_id).model_dump(mode='json', exclude_none=True) def _register_gateway_tinker_routes(app: FastAPI, self_fn: Callable[[], GatewayServer]) -> None: @@ -83,7 +50,7 @@ async def get_server_capabilities( self: GatewayServer = Depends(self_fn), ) -> types.GetServerCapabilitiesResponse: # Convert twinkle.protocol.types.SupportedModel to tinker.types.SupportedModel - tinker_supported_models = [types.SupportedModel(model_name=m.model_name) for m in self.supported_models] + tinker_supported_models = [types.SupportedModel(model_name=m.model_name) for m in self.training_models] return types.GetServerCapabilitiesResponse(supported_models=tinker_supported_models) @app.post('/telemetry') @@ -121,6 +88,7 @@ async def retrieve_future(request: Request, self: GatewayServer = Depends(self_fn)) -> Any: """Retrieve the result of an async task with long polling.""" request_id = body.request_id + await self.state.assert_future_owner(request_id, get_token_from_request(request)) outcome = await poll_future(self.state, request_id) record = outcome.record if outcome.timed_out: diff --git a/src/twinkle/server/gateway/twinkle_handlers.py b/src/twinkle/server/gateway/twinkle_handlers.py index 3df44ffc8..91962c8d6 100644 --- a/src/twinkle/server/gateway/twinkle_handlers.py +++ b/src/twinkle/server/gateway/twinkle_handlers.py @@ -57,11 +57,14 @@ async def healthz_deep( results = {} all_healthy = True - for model in self.supported_models: + for model in self.training_models: model_name = model.model_name try: resp = await self.proxy.proxy_request(request, 'healthz', model_name, 'model') healthy = (resp.status_code == 200) + if healthy: + import json + healthy = json.loads(resp.body).get('healthy', True) if not healthy: all_healthy = False results[model_name] = { @@ -86,7 +89,7 @@ async def get_server_capabilities( self: GatewayServer = Depends(self_fn), ) -> types.GetServerCapabilitiesResponse: return types.GetServerCapabilitiesResponse( - supported_models=self.supported_models, + supported_models=self.training_models, protocol_version=1, features=types.ClientFeatures( task_envelope=True, @@ -139,6 +142,7 @@ async def retrieve_future( finish. See ``poll_config`` for the numbers. """ request_id = body.request_id + await self.state.assert_future_owner(request_id, get_token_from_request(request)) outcome = await poll_future(self.state, request_id) if outcome.record is None: raise ResourceNotFoundError(f'request_id {request_id} not found or expired') @@ -156,6 +160,7 @@ async def cancel_future( running or already-terminal task is reported but never interrupted, so cancel can never corrupt in-flight GPU/optimizer state. """ + await self.state.assert_future_owner(body.request_id, get_token_from_request(request)) result = await self.state.cancel_future(body.request_id) return types.CancelResponse(**result) diff --git a/src/twinkle/server/gateway/use_cases.py b/src/twinkle/server/gateway/use_cases.py index 26691014b..f3c5df5dd 100644 --- a/src/twinkle/server/gateway/use_cases.py +++ b/src/twinkle/server/gateway/use_cases.py @@ -18,6 +18,8 @@ from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager from twinkle.server.lifecycle.poll_config import long_poll_window, retrieve_poll_interval +# Deliberately local: this layer works on canonical records and must not import the +# protocol models (see test_gateway_services_do_not_import_protocol_models). _TERMINAL_STATUSES = frozenset({'completed', 'failed', 'cancelled'}) diff --git a/src/twinkle/server/launcher/__init__.py b/src/twinkle/server/launcher/__init__.py index 4a2514ceb..9a5b0c1ff 100644 --- a/src/twinkle/server/launcher/__init__.py +++ b/src/twinkle/server/launcher/__init__.py @@ -8,7 +8,6 @@ - ``server_launcher`` — the ``ServerLauncher`` class and the ``launch_server`` entry point (the real launcher logic); - ``builder_registry`` — ``import_path`` → deployment-builder resolution; -- ``env_propagation`` — telemetry / persistence env-var collection for workers. This ``__init__`` is a thin aggregator: it only re-exports the public surface so the dotted paths ``twinkle.server.launcher.ServerLauncher`` / ``launch_server`` diff --git a/src/twinkle/server/launcher/env_propagation.py b/src/twinkle/server/launcher/env_propagation.py deleted file mode 100644 index c478444d9..000000000 --- a/src/twinkle/server/launcher/env_propagation.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Collection of telemetry / persistence env vars for propagation to Ray workers. - -These variables are read inside each Ray Serve worker process — telemetry by -``ensure_telemetry_initialized()`` and persistence -by ``PersistenceConfig.from_env()`` — so the chosen backend / telemetry config -is independent of deployment startup order. -""" -from __future__ import annotations - -import os - -# Telemetry env var keys that need to be propagated to Ray worker processes. -TELEMETRY_ENV_KEYS: tuple[str, ...] = ( - 'TWINKLE_TELEMETRY_ENABLED', - 'TWINKLE_TELEMETRY_DEBUG', - 'TWINKLE_TELEMETRY_SERVICE', - 'TWINKLE_TELEMETRY_ENDPOINT', - 'TWINKLE_TELEMETRY_INTERVAL', - 'TWINKLE_MODEL_ID_ALIASES', -) - - -def build_telemetry_env_vars() -> dict[str, str]: - """Collect telemetry env vars from ``os.environ`` for worker propagation.""" - return {k: os.environ[k] for k in TELEMETRY_ENV_KEYS if k in os.environ} - - -def build_persistence_env_vars() -> dict[str, str]: - """Collect persistence env vars from ``os.environ`` for worker propagation.""" - from twinkle.server.config.persistence import PERSISTENCE_ENV_KEYS - return {k: os.environ[k] for k in PERSISTENCE_ENV_KEYS if k in os.environ} - - -def build_server_state_env_vars() -> dict[str, str]: - """Collect ServerState-policy env vars from ``os.environ`` for worker propagation. - - Read inside each worker by ``ServerStateArgs.from_env()`` (via - ``get_server_state``) so the configured quota / expiry / metrics interval is - applied everywhere, not only in the gateway that first built the state. - """ - from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS - return {k: os.environ[k] for k in SERVER_STATE_ENV_KEYS if k in os.environ} - - -def build_propagated_env_vars() -> dict[str, str]: - """Aggregate all env vars that must reach Ray worker processes.""" - merged: dict[str, str] = {} - merged.update(build_telemetry_env_vars()) - merged.update(build_persistence_env_vars()) - merged.update(build_server_state_env_vars()) - return merged diff --git a/src/twinkle/server/launcher/server_launcher.py b/src/twinkle/server/launcher/server_launcher.py index ca3573abe..e53a4f26b 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -1,11 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""``ServerLauncher`` and the ``launch_server`` entry point. - -The real launcher module. It lives here (rather than in the package -``__init__.py``) so the package's ``__init__`` stays a thin aggregator; the -public dotted paths ``twinkle.server.launcher.ServerLauncher`` / -``launch_server`` are preserved by the ``__init__`` re-export. -""" +"""Initialize server configuration, deploy applications, and own their shutdown.""" from __future__ import annotations import json @@ -13,16 +7,17 @@ import signal import threading from collections.abc import Callable +from copy import deepcopy from pathlib import Path from typing import Any from twinkle import get_logger -from twinkle.hub.model_alias import MODEL_ID_ALIASES_ENV, build_model_alias_map +from twinkle.hub.model_alias import MODEL_ID_ALIASES_ENV, build_model_alias_map, extract_model_alias from twinkle.patch.ray_serve import apply_ray_serve_patches from twinkle.server.config import ServerConfig from twinkle.server.config.application_spec import ApplicationSpec +from twinkle.server.config.resources import build_propagated_env_vars, build_worker_env_vars from .builder_registry import get_builders, resolve_builder -from .env_propagation import build_propagated_env_vars logger = get_logger() @@ -87,55 +82,57 @@ def ensure_ray_cluster(*, namespace: str, address: str | None = None) -> None: class ServerLauncher: - """ - Unified server launcher. + """Own startup and cleanup; business scheduling stays in the controllers.""" - This class handles Ray/Serve initialization and application deployment. - - Attributes: - config: The server configuration dictionary - ray_namespace: The Ray namespace for the cluster - """ - - def __init__( - self, - config: ServerConfig, - ray_namespace: str | None = None, - ): - """ - Initialize the server launcher. - - Args: - config: A validated :class:`ServerConfig` instance. Raw dicts are - rejected — operators must build ``ServerConfig`` via - ``ServerConfig.from_yaml`` or its constructor so cross-field - validation runs before the launcher consumes anything. - ray_namespace: Ray namespace (default: 'twinkle_cluster') - """ + def __init__(self, config: ServerConfig, ray_namespace: str | None = None): if not isinstance(config, ServerConfig): raise TypeError('ServerLauncher requires a typed ServerConfig instance; ' f'got {type(config).__name__}. Build one with ' 'ServerConfig.from_yaml(path) or ServerConfig(...).') - self.config: ServerConfig = config - self.ray_namespace = ray_namespace + self.config = config + self.ray_namespace = ray_namespace or config.ray_namespace or 'twinkle_cluster' self._builders: dict[str, Callable] = {} self._ray_initialized = False self._serve_started = False + self._training_controller = None + self._sampler_controllers = [] def _get_builders(self) -> dict[str, Callable]: - """Get (and cache) the builder functions for all app types.""" if not self._builders: self._builders = get_builders() return self._builders def _resolve_builder(self, import_path: str) -> Callable: - """Resolve an import_path to a builder function. - - Raises: - ValueError: If the import_path cannot be resolved. - """ return resolve_builder(import_path, self._get_builders()) + def _configure_environment(self) -> None: + if any(app.import_path in ('model', 'sampler') for app in self.config.applications): + os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = self.ray_namespace + os.environ.update(self.config.persistence.to_env_vars()) + gateways = [app for app in self.config.applications if app.import_path == 'server'] + if gateways: + os.environ.update(gateways[0].args.server_config.to_env_vars()) + aliases = build_model_alias_map(self.config.applications) + if aliases: + os.environ[MODEL_ID_ALIASES_ENV] = json.dumps(aliases, ensure_ascii=False) + else: + os.environ.pop(MODEL_ID_ALIASES_ENV, None) + self._configure_telemetry() + + def _configure_telemetry(self) -> None: + telemetry = self.config.telemetry + os.environ['TWINKLE_TELEMETRY_ENABLED'] = str(int(telemetry.enabled)) + if not telemetry.enabled: + return + from twinkle.server.telemetry import init_telemetry + init_telemetry(telemetry) + os.environ.update({ + 'TWINKLE_TELEMETRY_DEBUG': str(int(telemetry.debug)), + 'TWINKLE_TELEMETRY_SERVICE': telemetry.service_name, + 'TWINKLE_TELEMETRY_ENDPOINT': telemetry.otlp_endpoint, + 'TWINKLE_TELEMETRY_INTERVAL': str(telemetry.export_interval_ms), + }) + def _init_ray(self) -> None: """Initialize Ray with the resolved namespace (attach-or-local via :func:`ensure_ray_cluster`).""" if self._ray_initialized: @@ -155,7 +152,6 @@ def _start_serve(self) -> None: """ if self._serve_started: return - from ray import serve try: @@ -179,179 +175,172 @@ def _start_serve(self) -> None: if self.config.proxy_location: serve_kwargs['proxy_location'] = self.config.proxy_location serve.start(**serve_kwargs) - logger.info(f'Ray Serve started with http_options={http_options}, ' - f'proxy_location={self.config.proxy_location!r}') - self._serve_started = True - def _deploy_application(self, app_spec: ApplicationSpec) -> None: - """Deploy a single application. - - Args: - app_spec: Validated :class:`ApplicationSpec` from the typed config. - """ + def _deployment_options(self, options: dict, env: dict | None = None) -> dict: + options = deepcopy(options) + options.pop('name', None) + actor = options.setdefault('ray_actor_options', {}) + runtime_env = actor.setdefault('runtime_env', {}) + runtime_env['env_vars'] = {**build_worker_env_vars(), **runtime_env.get('env_vars', {}), **(env or {})} + return options + + def _gateway_arguments(self) -> dict[str, Any]: + args = {'http_options': self.config.http_options.model_dump()} + models = self.config.training_config.models + if models: + args['training_models'] = [model.name for model in models] + args['supported_models'] = list( + dict.fromkeys( + alias for app in self.config.applications + if app.import_path == 'sampler' and (alias := extract_model_alias(app.route_prefix, 'sampler')))) + return args + + def _deploy_sampler(self, app: ApplicationSpec) -> None: + import ray from ray import serve - name = app_spec.name - route_prefix = app_spec.route_prefix - import_path = app_spec.import_path - # Shallow-dump the typed args to a kwargs dict WITHOUT recursing into - # nested models: ``dict(model)`` yields top-level (field, value) pairs - # with nested models left as instances, so ``queue_config`` stays a - # typed ``TaskQueueConfig`` and is not re-serialized to a dict only to be - # revived via ``from_dict`` inside the builder (the prior triple - # validation). ``None`` values are dropped to match ``exclude_none``. - args = {k: v for k, v in dict(app_spec.args).items() if v is not None} - deployments = list(app_spec.deployments or []) - - logger.info(f'Starting {name} at {route_prefix}...') - - builder = self._resolve_builder(import_path) - - deploy_options = {} - if deployments: - if len(deployments) > 1: - logger.warning(f'Application "{name}" has {len(deployments)} deployments configured, ' - f'but only the first deployment will be used.') - deploy_config = deployments[0] - if isinstance(deploy_config, dict): - deploy_options = {k: v for k, v in deploy_config.items() if k != 'name'} - - # Inject telemetry + persistence env vars into the deployment's - # runtime_env so that Ray Serve replicas (worker processes) can - # initialize telemetry and resolve the configured persistence backend - # regardless of deployment startup order. - # User-specified env_vars take precedence over our defaults. - propagated_env_vars = build_propagated_env_vars() - if propagated_env_vars: - ray_actor_options = dict(deploy_options.get('ray_actor_options') or {}) - runtime_env = dict(ray_actor_options.get('runtime_env') or {}) - env_vars = dict(runtime_env.get('env_vars') or {}) - for k, v in propagated_env_vars.items(): - env_vars.setdefault(k, v) - runtime_env['env_vars'] = env_vars - ray_actor_options['runtime_env'] = runtime_env - deploy_options['ray_actor_options'] = ray_actor_options - - # Pass http_options to server apps for internal proxy routing - if import_path == 'server': - args.setdefault('http_options', self.config.http_options.model_dump()) - - app = builder(deploy_options=deploy_options, **args) - - serve.run(app, name=name, route_prefix=route_prefix) - logger.info(f'Deployed {name} at {route_prefix}') - - def launch(self, blocking: bool = True) -> None: - """Launch the server with all configured applications. - - Blocks the calling thread to keep the server running. Installs signal - handlers for SIGINT/SIGTERM so that ``serve.shutdown()`` is called on - termination instead of leaving orphaned deployments. With - ``blocking=False`` the call returns once every application was - submitted -- detached deployments own their lifetime on the Ray - cluster, not in this driver. - """ - # Apply Ray Serve patches before initializing Ray - apply_ray_serve_patches() - - # Initialize telemetry if configured - telemetry = self.config.telemetry - if telemetry.enabled: - from twinkle.server.telemetry import init_telemetry - init_telemetry(telemetry) - # Export config to env vars for Ray worker processes - os.environ['TWINKLE_TELEMETRY_ENABLED'] = '1' - os.environ['TWINKLE_TELEMETRY_DEBUG'] = '1' if telemetry.debug else '0' - os.environ['TWINKLE_TELEMETRY_SERVICE'] = telemetry.service_name - os.environ['TWINKLE_TELEMETRY_ENDPOINT'] = telemetry.otlp_endpoint - os.environ['TWINKLE_TELEMETRY_INTERVAL'] = str(telemetry.export_interval_ms) - - # Export top-level persistence to env vars so any worker - # (not just Gateway) can build the same backend on first call to - # get_server_state(). - persistence = self.config.persistence - for k, v in persistence.to_env_vars().items(): - os.environ[k] = v - logger.info(f'Persistence backend configured: mode={persistence.mode}') - - # Export the gateway ``server`` application's ServerState policy (quota / - # expiry / cleanup / metrics interval) to env vars for the same reason: - # so every worker's first ``get_server_state()`` applies the configured - # values instead of the hardcoded defaults. Without this the model worker - # that enforces ``per_token_model_limit`` runs on the default (30), - # silently ignoring the YAML value. - server_specs = [a for a in self.config.applications if a.import_path == 'server'] - if len(server_specs) > 1: - logger.warning(f'{len(server_specs)} "server" applications declared; using the first ' - 'for ServerState policy env propagation.') - if server_specs: - server_state_env = server_specs[0].args.server_config.to_env_vars() - for k, v in server_state_env.items(): - os.environ[k] = v - if server_state_env: - logger.info(f'ServerState policy exported to worker env: {server_state_env}') - - model_alias_map = build_model_alias_map(self.config.applications) - if model_alias_map: - os.environ[MODEL_ID_ALIASES_ENV] = json.dumps(model_alias_map, ensure_ascii=False) - else: - os.environ.pop(MODEL_ID_ALIASES_ENV, None) - - self._init_ray() - self._start_serve() + from twinkle.server.sampler.pool.actor import create_sampler_controller + from twinkle.server.sampler.pool.router import build_sampler_router + controller, options, pool_id = create_sampler_controller(app, self.ray_namespace, + self.config.resources.model_dump(), + build_worker_env_vars()) + self._sampler_controllers.append(controller) + ray.get(controller.initialize.remote()) + serve.run( + build_sampler_router(self.ray_namespace, pool_id, options), name=app.name, route_prefix=app.route_prefix) - applications = self.config.applications - if not applications: - logger.warning('No applications configured') + def _deploy_application(self, app_spec: ApplicationSpec) -> None: + """Model declarations go to the training controller; other apps deploy here.""" + if app_spec.import_path == 'model': return - - for app_spec in applications: - self._deploy_application(app_spec) - - host = self.config.http_options.host - port = self.config.http_options.port - - logger.info('All applications started!') - for app_spec in applications: - logger.info(' Endpoint: http://%s:%s%s', host, port, app_spec.route_prefix) - - if not blocking: + if app_spec.import_path == 'sampler': + self._deploy_sampler(app_spec) return + from ray import serve + args = {key: value for key, value in dict(app_spec.args).items() if value is not None} + if app_spec.import_path == 'server': + defaults = self._gateway_arguments() + for key, value in defaults.items(): + args.setdefault(key, value) + deployment = app_spec.deployments[0] if app_spec.deployments else {} + options = self._deployment_options(deployment) + builder = self._resolve_builder(app_spec.import_path) + serve.run(builder(deploy_options=options, **args), name=app_spec.name, route_prefix=app_spec.route_prefix) + + def _training_router_options(self, model) -> dict: + options = self._deployment_options(model.deploy_options, model.env_vars) + if 'autoscaling_config' not in options: + options.setdefault('num_replicas', 1) + options['ray_actor_options'].setdefault('num_cpus', 0.1) + return options + + def _deploy_training(self) -> None: + import ray + from ray import serve - # Graceful shutdown via signal handling - shutdown_event = threading.Event() - - def _handle_signal(signum, frame): - sig_name = signal.Signals(signum).name - logger.info(f'Received {sig_name}, shutting down gracefully...') - shutdown_event.set() + from twinkle.server.model.pool.actor import TrainingControllerActor, controller_name + from twinkle.server.model.pool.router import build_training_router + config = self.config.training_config + controller = TrainingControllerActor.options( + name=controller_name(self.ray_namespace), + namespace=self.ray_namespace, + runtime_env={ + 'env_vars': build_worker_env_vars() + }).remote(config.model_dump(), self.ray_namespace) + self._training_controller = controller + try: + ray.get(controller.initialize.remote()) + for model in config.models: + app = build_training_router(config, model, self.ray_namespace, self.config.http_options.model_dump(), + self._training_router_options(model)) + serve.run( + app, + name=f'training-router-{self.ray_namespace}-{model.name}', + route_prefix=model.route_prefix or f'/api/v1/model/{model.name}') + except Exception: + self._run_cleanup([('training controller rollback', self._stop_training_controller)], raise_errors=False) + raise + + def _stop_controller(self, controller) -> None: + import ray + + # Keep an unconfirmed owner available for cleanup retry; a timeout alone + # does not certify its workers or devices were released. + ray.get(controller.shutdown.remote(), timeout=self.config.resources.shutdown_timeout_s) + ray.kill(controller, no_restart=True) + + def _stop_training_controller(self) -> None: + if self._training_controller is not None: + self._stop_controller(self._training_controller) + self._training_controller = None + + def _stop_sampler_controller(self, controller) -> None: + self._stop_controller(controller) + self._sampler_controllers.remove(controller) + + @staticmethod + def _run_cleanup(actions: list[tuple[str, Callable]], *, raise_errors: bool) -> None: + """Attempt every independent cleanup; rollback preserves the startup error.""" + failures = [] + for name, action in actions: + try: + action() + except Exception as exc: + failures.append((name, exc)) + logger.exception('%s failed', name) + if failures and raise_errors: + names = ', '.join(name for name, _ in failures) + raise RuntimeError(f'Server cleanup is incomplete: {names}') from failures[0][1] + + def _shutdown(self, *, raise_errors: bool = True) -> None: + from functools import partial + from ray import serve - signal.signal(signal.SIGINT, _handle_signal) - signal.signal(signal.SIGTERM, _handle_signal) + from twinkle.server.telemetry import flush_telemetry_safely + actions = [('training controller shutdown', self._stop_training_controller)] + actions.extend(('sampler controller shutdown', partial(self._stop_sampler_controller, controller)) + for controller in tuple(self._sampler_controllers)) + if self._serve_started: + actions.append(('Ray Serve shutdown', serve.shutdown)) + actions.append(('telemetry flush', flush_telemetry_safely)) + self._run_cleanup(actions, raise_errors=raise_errors) + self._serve_started = False - # Block until a termination signal is received - shutdown_event.wait() + @staticmethod + def _wait_for_shutdown() -> None: + stopped = threading.Event() - self._shutdown() + def signal_shutdown(signum, frame): + logger.info('Received %s', signal.Signals(signum).name) + stopped.set() - def _shutdown(self) -> None: - """Tear down Serve and flush telemetry on graceful termination. + signal.signal(signal.SIGINT, signal_shutdown) + signal.signal(signal.SIGTERM, signal_shutdown) + stopped.wait() - Telemetry flush runs AFTER ``serve.shutdown()`` and is wrapped so a - telemetry-shutdown failure cannot mask the user-facing shutdown path. - """ - from ray import serve + def launch(self, blocking: bool = True) -> None: + """Launch applications and their controllers, optionally waiting for shutdown.""" + if not self.config.applications: + logger.warning('No applications configured') + return + apply_ray_serve_patches() + self._configure_environment() + self._init_ray() try: - serve.shutdown() - logger.info('Ray Serve shut down successfully') + self._start_serve() + for app in self.config.applications: + self._deploy_application(app) + if self.config.training_config.models: + self._deploy_training() + logger.info('All applications started') + if not blocking: + return + self._wait_for_shutdown() except Exception: - logger.warning('Error during Ray Serve shutdown', exc_info=True) - - # Flush buffered OTLP batches (traces / metrics / logs) on the driver - # after Serve is down. Errors are swallowed by the helper. - from twinkle.server.telemetry import flush_telemetry_safely - flush_telemetry_safely() + self._shutdown(raise_errors=False) + raise + self._shutdown() def _warn_public_unauthenticated(config: ServerConfig) -> None: diff --git a/src/twinkle/server/lifecycle/envelope.py b/src/twinkle/server/lifecycle/envelope.py index 25343c9b0..dc193aadf 100644 --- a/src/twinkle/server/lifecycle/envelope.py +++ b/src/twinkle/server/lifecycle/envelope.py @@ -13,41 +13,18 @@ from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload from twinkle.protocol.types.lifecycle import TaskEnvelope from twinkle.server.state.models import FutureFailureRecord -from twinkle.server.task_errors import trim_traceback - -# Keys must equal ``state.models.FAILURE_REASON_CODES`` (guarded by test_envelope). -_FAILURE_WIRE: dict[str, tuple[int, ErrorCategory]] = { - 'invalid_request': (400, ErrorCategory.User), - 'request_rejected': (400, ErrorCategory.User), - 'resource_not_found': (404, ErrorCategory.User), - 'full_mode_busy': (409, ErrorCategory.User), - 'input_tokens_exceeded': (422, ErrorCategory.User), - 'batch_size_invalid': (422, ErrorCategory.User), - 'rate_limit_exceeded': (429, ErrorCategory.User), - 'resource_quota_exceeded': (429, ErrorCategory.User), - 'cancelled': (499, ErrorCategory.User), - 'endpoint_unavailable': (501, ErrorCategory.Server), - 'state_contention': (503, ErrorCategory.Server), - 'backend_gate_unavailable': (503, ErrorCategory.Server), - 'orphaned_replica': (503, ErrorCategory.Server), - 'execution_timeout': (504, ErrorCategory.Server), - 'deadline_exceeded': (500, ErrorCategory.Server), - 'internal_error': (500, ErrorCategory.Server), -} +from twinkle.server.task_errors import FailureResponse, trim_traceback def error_payload_from_failure(stored: Any, *, request_id: str) -> ErrorPayload: """Map one protocol-independent failure record to Twinkle's wire model.""" failure = FutureFailureRecord.model_validate(stored) - error_code, category = _FAILURE_WIRE.get( - failure.reason_code, - (500, ErrorCategory.Server), - ) - diagnostic = failure.diagnostic if category is ErrorCategory.Server else None + response = FailureResponse.for_reason(failure.reason_code) + diagnostic = failure.diagnostic if response.category is ErrorCategory.Server else None return ErrorPayload( error=failure.message[:1024], - category=category, - error_code=error_code, + category=response.category, + error_code=response.error_code, request_id=request_id, traceback=trim_traceback(diagnostic) if diagnostic else None, details=failure.details, diff --git a/src/twinkle/server/lifecycle/protocols.py b/src/twinkle/server/lifecycle/protocols.py index c05ea12bb..3bee91ade 100644 --- a/src/twinkle/server/lifecycle/protocols.py +++ b/src/twinkle/server/lifecycle/protocols.py @@ -19,9 +19,9 @@ from fastapi import Request from typing import Any, Protocol, runtime_checkable +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.lifecycle.envelope import TaskEnvelope from twinkle.server.state import ServerState -from twinkle.server.task_queue.config import TaskQueueConfig @runtime_checkable diff --git a/src/twinkle/server/lifecycle/proxy.py b/src/twinkle/server/lifecycle/proxy.py new file mode 100644 index 000000000..24043b616 --- /dev/null +++ b/src/twinkle/server/lifecycle/proxy.py @@ -0,0 +1,137 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared HTTP forwarding between services; independent of gateway and training controller.""" + +from __future__ import annotations +from typing import Any, TypedDict + +import httpx +from fastapi import Request, Response +from fastapi.responses import JSONResponse + +from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID, H_REQUEST_ID_LEGACY +from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.telemetry.tracing import inject_context +from twinkle.utils.logger import get_logger + +logger = get_logger() + + +class ForwardedRequest(TypedDict): + """HTTP request fields carried intact across an internal Ray RPC.""" + + endpoint: str + body: bytes + method: str + headers: list[tuple[bytes, bytes]] + query_string: bytes + + +class ForwardedReply(TypedDict): + status: int + headers: list[tuple[bytes, bytes]] + body: bytes + + +class ServiceProxy: + def __init__(self, http_options: dict[str, Any] | None = None, route_prefix: str = "/api/v1"): + self.http_options = http_options or {} + self.route_prefix = route_prefix + queue = TaskQueueConfig() + budget = queue.queue_timeout + queue.effective_execution_timeout + self.client = httpx.AsyncClient( + timeout=httpx.Timeout(budget, connect=10, pool=10), + trust_env=False, + verify=not self.http_options.get("ssl_certfile"), + ) + + async def close(self) -> None: + await self.client.aclose() + + def _build_target_url(self, service_type: str, base_model: str, endpoint: str) -> str: + host = self.http_options.get("host", "localhost") + port = self.http_options.get("port", 8000) + if host == "0.0.0.0": + host = "localhost" + scheme = "https" if self.http_options.get("ssl_certfile") else "http" + return f"{scheme}://{host}:{port}{self.route_prefix.rstrip('/')}/{service_type}/{base_model}/{endpoint}" + + def _prepare_headers(self, request_headers) -> dict[str, str]: + headers = dict(request_headers) + headers.pop("host", None) + headers.pop("content-length", None) + request_id = request_headers.get(H_REQUEST_ID) or request_headers.get(H_REQUEST_ID_LEGACY) + if request_id is not None and not request_headers.get(H_MULTIPLEX_LEGACY): + headers[H_MULTIPLEX_LEGACY] = request_id + if request_id is not None: + headers.setdefault(H_REQUEST_ID, request_id) + headers.setdefault(H_MULTIPLEX, request_id) + return headers + + async def proxy_request( + self, + request: Request, + endpoint: str, + base_model: str, + service_type: str, + *, + body_override: bytes | None = None, + extra_headers: dict[str, str] | None = None, + ) -> Response: + body_bytes = body_override if body_override is not None else await request.body() + target_url = self._build_target_url(service_type, base_model, endpoint) + headers = self._prepare_headers(request.headers) + if extra_headers: + headers.update(extra_headers) + inject_context(headers) + try: + response = await self.client.request( + method=request.method, url=target_url, content=body_bytes, headers=headers, params=request.query_params + ) + return Response( + content=response.content, + status_code=response.status_code, + headers=dict(response.headers), + media_type=response.headers.get("content-type"), + ) + except Exception as exc: + logger.error("Proxy error: %s", str(exc), exc_info=True) + request_id = request.headers.get(H_REQUEST_ID) or request.headers.get(H_REQUEST_ID_LEGACY) or "" + payload = ErrorPayload( + error=f"Proxy Error: {exc}", category=ErrorCategory.Server, error_code=502, request_id=request_id + ) + return JSONResponse(status_code=502, content=payload.model_dump(mode="json", exclude_none=True)) + + async def proxy_request_stream( + self, + request: Request, + endpoint: str, + base_model: str, + service_type: str, + *, + body_override: bytes | None = None, + extra_headers: dict[str, str] | None = None, + ): + body_bytes = body_override if body_override is not None else await request.body() + target_url = self._build_target_url(service_type, base_model, endpoint) + headers = self._prepare_headers(request.headers) + if extra_headers: + headers.update(extra_headers) + inject_context(headers) + async with self.client.stream(method="POST", url=target_url, content=body_bytes, headers=headers) as response: + if response.status_code >= 400: + body = await response.aread() + raise httpx.HTTPStatusError( + f"Upstream returned {response.status_code}", + request=response.request, + response=httpx.Response(status_code=response.status_code, content=body), + ) + async for line in response.aiter_lines(): + if line: + yield line + + async def proxy_to_model(self, request: Request, endpoint: str, base_model: str) -> Response: + return await self.proxy_request(request, f"tinker/{endpoint}", base_model, "model") + + async def proxy_to_sampler(self, request: Request, endpoint: str, base_model: str) -> Response: + return await self.proxy_request(request, f"tinker/{endpoint}", base_model, "sampler") diff --git a/src/twinkle/server/lifecycle/submit.py b/src/twinkle/server/lifecycle/submit.py index b1b409a47..5c432281b 100644 --- a/src/twinkle/server/lifecycle/submit.py +++ b/src/twinkle/server/lifecycle/submit.py @@ -16,6 +16,8 @@ from twinkle.protocol.types.base import FieldRole, fields_with_role from twinkle.protocol.types.data import export_batch from twinkle.protocol.types.lifecycle import TaskEnvelope +from twinkle.protocol.types.model import (AdapterRequest, DataPlaneForwardRequest, ForwardBackwardTaskRequest, + TrainingRequest) from twinkle.server.middleware.auth import get_session_id_from_request from twinkle.server.validation import assert_request_supported @@ -85,11 +87,11 @@ def backend_kwargs(body: Any) -> dict[str, Any]: model_cls = type(body) kwargs: dict[str, Any] = {} for name in fields_with_role(model_cls, FieldRole.BackendKwarg): - value = getattr(body, name, None) + value = getattr(body, name) if value is not None: kwargs[name] = value for name in fields_with_role(model_cls, FieldRole.Passthrough): - region = getattr(body, name, None) or {} + region = getattr(body, name) or {} overlap = set(region) & set(kwargs) if overlap: raise ValueError(f'{name} collides with declared backend parameters: {", ".join(sorted(overlap))}') @@ -123,7 +125,7 @@ def resolve_twinkle_adapter_name(request: Request, adapter_name: str | None) -> async def run_submit( self: QueuedDeployment, request: Request, - body: Any, + body: TrainingRequest, *, task_type: str, backend_call: Callable[..., Coroutine], @@ -152,13 +154,10 @@ async def run_submit( signature-patching wrapper once deepened the route graph past CPython's C-stack recursion guard during ``serve.ingress``). - Four queued endpoints deliberately do NOT route through here and call - ``submit_and_peek`` / ``submit_background_and_peek`` directly: model - ``add_adapter_to_model`` (creates the adapter, so the resource assertion cannot - apply and it owns the train_mode/full-mode checks), model ``upload_to_hub`` (pure - I/O, background task), and sampler ``sample`` / ``sample_to_data_plane`` (no adapter - semantics). Zero-write and admission still hold for them because both live inside - ``schedule_task`` -> ``_perform_preflight_checks``, not in this shell. + Upload uses the background I/O path; sampling uses concurrent admission. + Their scheduling entrypoints apply the same preflight and Future ownership rules. + Training creation/release belongs to the public Router and Controller, rather + than this compute-operation shell. ``backend_call(self, body, adapter_name, token)`` runs the endpoint-specific call and returns the JSON-safe task result. ``metrics(self, body)`` supplies scheduling @@ -168,7 +167,7 @@ async def run_submit( not implemented by every backend. """ token = await self._on_request_start(request) - adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) + adapter_name = body.training_id # ---- Preflight: decidable from the body plus this deployment's backend, so it # runs before any state write and before the enqueue. ---- @@ -192,29 +191,40 @@ async def _task(): # forward_backward would be dropped as a duplicate AND handed the first adapter's # loss. That is a silent wrong-result bug, i.e. the exact failure this dedup # exists to prevent, one level up. ---- - request_id = f'req_{uuid.uuid4().hex}' - seq_id = getattr(body, 'seq_id', None) - dedup_key = None + sequenced = isinstance(body, (AdapterRequest, ForwardBackwardTaskRequest, DataPlaneForwardRequest)) + return await submit_training_task( + self, + _task, + training_id=body.training_id, + stream=f'native::{get_session_id_from_request(request)}::{task_type}', + seq_id=body.seq_id if sequenced else None, + peek=True, + token=token, + task_type=task_type, + **schedule_kwargs) + + +async def submit_training_task(self, + factory, + *, + training_id: str, + stream: str, + seq_id: int | None, + peek: bool = False, + **kwargs): + """Submit with explicit request identity; retain claims once any future was written.""" + request_id = kwargs.pop('request_id', None) or f'req_{uuid.uuid4().hex}' if seq_id is not None: - session_id = get_session_id_from_request(request) or request.state.request_id - dedup_key = f'seq::{session_id}::{adapter_name or "-"}::{seq_id}' - ttl = int(self.task_queue_config.effective_execution_timeout) + 60 - prior_request_id = await self.state.claim_seq(dedup_key, request_id, ttl) - if prior_request_id is not None: - return await self._peek_terminal(prior_request_id, fallback_status='pending') - - # ---- Decision_Boundary: preflight (in schedule_task) then peek ---- + prior = await self.state.claim_training_seq(training_id, stream, seq_id, request_id, self.sequence_retention_s) + if prior is not None: + if peek: + return await self._peek_terminal(prior, fallback_status='pending') + return {'request_id': prior, 'model_id': training_id} + submit = self.submit_and_peek if peek else self.schedule_task try: - return await self.submit_and_peek( - _task, model_id=adapter_name, token=token, task_type=task_type, request_id=request_id, **schedule_kwargs) + return await submit(factory, model_id=training_id, request_id=request_id, **kwargs) except Exception: - # Release the seq claim only when the task never made it onto the queue -- - # decided by whether a future record exists, NOT by the exception type. A - # preflight rejection raises before any record is written, so releasing lets - # a retry re-enqueue. But a failure *after* enqueue (e.g. a transient state - # error inside the peek) leaves a live task that will still run; releasing - # there would let a retry enqueue a duplicate -> the exact double-apply this - # dedup prevents. When unsure (record exists), keep the claim. - if dedup_key is not None and await self.state.get_future(request_id) is None: - await self.state.release_seq(dedup_key) + # A failed response after enqueue must not allow the same update to run twice. + if seq_id is not None and await self.state.get_future(request_id) is None: + await self.state.release_training_seq(training_id, stream, seq_id, request_id) raise diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index bbcc918a5..908dddbd8 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -10,25 +10,22 @@ import asyncio import importlib -from fastapi import FastAPI, HTTPException, Request -from ray import serve -from ray.serve.config import RequestRouterConfig +from fastapi import FastAPI, Request from typing import Any from twinkle import DeviceGroup, remote_class from twinkle.model.base import TrainableModel from twinkle.server.config.backend_dispatch import BackendSelector -from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app -from twinkle.server.exceptions import FullModeBusyError -from twinkle.server.middleware.auth import get_token_from_request -from twinkle.server.model.routing import StickyLoraRequestRouter +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.deployment import bind_deployment, build_deployment_app from twinkle.server.runtime import init_twinkle_runtime -from twinkle.server.session_resource import AdapterManagerMixin from twinkle.server.state import ServerState, get_server_state -from twinkle.server.task_queue import TaskQueueConfig, TaskQueueMixin +from twinkle.server.state.training_records import TrainingRuntimeContext +from twinkle.server.task_queue import TaskQueueMixin from twinkle.server.utils import wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger from twinkle.utils.unsafe import trust_remote_code +from .runtime import TrainingRuntime from .tinker_handlers import _register_model_tinker_routes from .twinkle_handlers import _register_model_twinkle_routes @@ -42,6 +39,42 @@ _LORA_ONLY_CTOR_KWARGS = ('max_loras', 'max_r', 'max_length', 'target_modules', 'lora_config') +def _validate_model_task(kw: dict[str, Any]) -> None: + """Check the persisted head once, before choosing a CS backend.""" + task_type = kw.get('task_type') + if task_type is None: + return + from transformers import AutoConfig + from twinkle.hub import HubOperation + + path = HubOperation.download_model(kw['model_id']) + loader = kw.get('model_loader') + config = (loader.process_config(loader.build_config(path)) if loader is not None else + AutoConfig.from_pretrained(path, trust_remote_code=trust_remote_code())) + architectures = getattr(config, 'architectures', None) or [] + classification = any(name.endswith('ForSequenceClassification') for name in architectures) + num_labels = kw.get('num_labels') + if task_type == 'seq_cls': + if not classification: + raise ValueError('seq_cls requires a persisted ForSequenceClassification checkpoint') + if num_labels is None or getattr(config, 'num_labels', None) != num_labels: + raise ValueError('num_labels does not match the persisted classification head') + from pathlib import Path + from safetensors import safe_open + head_found = False + for file in Path(path).glob('model*.safetensors'): + with safe_open(str(file), framework='pt', device='cpu') as weights: + for key in weights.keys(): + if key.endswith(('score.weight', 'classifier.weight')): + if weights.get_slice(key).get_shape()[0] != num_labels: + raise ValueError('Persisted head weight shape does not match num_labels') + head_found = True + if not head_found: + raise ValueError('seq_cls requires persisted score/classifier weights in safetensors format') + elif classification: + raise ValueError('causal_lm cannot use a classification checkpoint') + + def _make_mock_model(kw: dict[str, Any]) -> Any: from .backends.mock_model import TwinkleCompatMockModel @@ -164,14 +197,14 @@ def _resolve_model_cls(spec: str) -> type: ) -class ModelManagement(LazyCleanupMixin, TaskQueueMixin, AdapterManagerMixin): +class ModelManagement(TrainingRuntime, TaskQueueMixin): """Unified model management service. Handles: - Base model and multiple LoRA adapters (multi-user) - Tinker training operations via /tinker/* endpoints (async/polling) - Twinkle training operations via /twinkle/* endpoints (synchronous) - - Adapter lifecycle via AdapterManagerMixin + - Worker-local gradient state is keyed by the controller-issued training handle. - Per-user rate limiting via TaskQueueMixin """ @@ -181,19 +214,21 @@ async def __init__(self, device_group: dict[str, Any], device_mesh: dict[str, Any], backend: str, - adapter_config: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, data_plane_url: str | None = None, **kwargs): self.backend = backend self.device_group = DeviceGroup(**device_group) + self.training_context = TrainingRuntimeContext.model_validate(kwargs.pop('training_context')) + await self.initialize_training_runtime(self.training_context, kwargs.pop('placement_group'), + kwargs.pop('allocation_devices')) self.device_mesh = init_twinkle_runtime( is_mock=(backend == 'mock'), nproc_per_node=nproc_per_node, device_group=self.device_group, device_mesh_dict=device_mesh, ) - self.replica_id = serve.get_replica_context().replica_id.unique_id + self.replica_id = self.training_context.instance_id self.max_loras = kwargs.get('max_loras', 5) self.base_model = model_id # ``train_mode`` selects LoRA multi-tenant (default) or exclusive @@ -217,7 +252,10 @@ async def __init__(self, ctor_kwargs['model_loader'] = _resolve_model_loader(model_loader_spec, model_id) if self.device_mesh is not None: ctor_kwargs['device_mesh'] = self.device_mesh + if backend != 'mock': + await asyncio.to_thread(_validate_model_task, ctor_kwargs) self.model = MODEL_SELECTOR.construct(backend, ctor_kwargs) + await self.wait_for_training_workers() self.state: ServerState = get_server_state() from twinkle.server.data_plane import DataPlaneProxy @@ -235,9 +273,8 @@ async def __init__(self, collect_width=len(actors) if actors else 1, ) self.model._ray_get_timeout = self.task_queue_config.effective_execution_timeout - self._init_adapter_manager(**(adapter_config or {})) + self._training_states = {} await self._register_replica_on_startup() - # Note: countdown task is started lazily in _ensure_sticky() @property def data_world_size(self) -> int: @@ -252,34 +289,15 @@ def data_world_size(self) -> int: async def _register_replica_on_startup(self) -> None: """Register this replica's capacity before Ray Serve marks it ready.""" if not self._replica_registered: - await self.state.register_replica(self.replica_id, self.max_loras) + await self.state.register_replica(self.replica_id, 0 if self.is_full_mode else self.max_loras) await self.state.touch_replica_last_seen(self.replica_id) self._replica_registered = True - @serve.multiplexed(max_num_models_per_replica=5) - async def _sticky_entry(self, sticky_key: str): - return sticky_key - - async def _ensure_sticky(self): - sticky_key = serve.get_multiplexed_model_id() - await self._sticky_entry(sticky_key) - # Lazy-start countdown task on first request (requires running event loop) - self._ensure_countdown_started() - async def _on_request_start(self, request: Request) -> str: - await self._ensure_sticky() - await self.state.touch_replica_last_seen(self.replica_id) - await self._ensure_state_cleanup_started() - if self._model_unhealthy: - raise HTTPException(status_code=503, detail='Model actors are unavailable') - return get_token_from_request(request) + return await self.training_request_start(request) async def shutdown(self) -> None: """Explicit async cleanup — called via FastAPI shutdown event.""" - try: - await self.state.unregister_replica(self.replica_id) - except Exception: - pass await self.shutdown_task_queue() await self.data_plane.close() @@ -287,8 +305,7 @@ async def _run_model_health_probe(self) -> dict: try: result = await self.call_backend(self.model.ping, admit=False) if result is True: - self._model_unhealthy = False - return {'healthy': True, 'detail': 'model actors alive'} + return {'healthy': not self._model_unhealthy, 'detail': 'model actors alive'} self._model_unhealthy = True return {'healthy': False, 'detail': f'unexpected ping result: {result}'} except Exception as e: @@ -303,28 +320,11 @@ async def check_model_health(self) -> dict: return await asyncio.shield(self._health_probe_task) def mark_unhealthy(self) -> None: - """Flag the deployment unhealthy; /healthz returns 503 until a probe recovers it.""" + """Flag this training runtime unhealthy until the controller replaces it.""" self._model_unhealthy = True async def _probe_after_timeout(self) -> None: - """Fired by ComputeWorker on a backend timeout: probe and log liveness.""" - result = await self.check_model_health() - logger.warning('[Model] post-timeout liveness probe: %s', result) - - async def _cleanup_adapter(self, adapter_name: str) -> None: - if self.get_resource_info(adapter_name): - if self.train_mode == 'full': - # No PEFT adapter to remove; restore clean base weights so the - # next tenant does not inherit this tenant's trained weights. - # Takes the Admission_Gate: this path is driven by the background - # countdown and never enters Task_Queue, so the gate is what keeps - # it from colliding with an in-flight training call. - await self.call_backend(self.model.reload_initial_weights) - else: - await self.call_backend(self.model.remove_adapter, adapter_name) - self.clear_resource_state(adapter_name) - self.unregister_resource(adapter_name) - await self.state.unload_model(adapter_name) + self._model_unhealthy = True @property def is_full_mode(self) -> bool: @@ -341,26 +341,6 @@ def resolve_model_adapter_name(self, adapter_name: str | None) -> str | None: return '' return adapter_name - def assert_full_mode_available(self, adapter_name: str | None = None) -> None: - """In full mode, ensure no other tenant currently holds this deployment. - - Call before registering any state so a rejected request leaves nothing - behind. ``adapter_name`` (when given) is excluded from the check so a - tenant re-issuing against its own resource is not blocked. - - Raises: - FullModeBusyError: another (non-expiring) tenant already holds it. - """ - if not self.is_full_mode: - return - holder = self.find_active_resource(exclude=adapter_name) - if holder is not None: - raise FullModeBusyError(holder) - - async def _on_adapter_expired(self, adapter_name: str) -> None: - self.fail_pending_tasks_for_model(adapter_name, reason='Adapter expired') - await self._cleanup_adapter(adapter_name) - def build_model_app(model_id: str, nproc_per_node: int, @@ -368,7 +348,6 @@ def build_model_app(model_id: str, device_mesh: dict[str, Any], deploy_options: dict[str, Any], backend: str, - adapter_config: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, data_plane_url: str | None = None, **kwargs): @@ -385,7 +364,6 @@ def build_model_app(model_id: str, backend: Model backend selector — ``mock`` | ``transformers`` | ``megatron``. Validated up front; bad values raise :class:`ConfigError` before any side effect. - adapter_config: Adapter lifecycle config (timeout, per-token limits) queue_config: Task queue configuration (rate limiting, etc.) **kwargs: Additional model initialization arguments @@ -398,8 +376,7 @@ def build_model_app(model_id: str, # Build the FastAPI app + middleware stack + routes via the shared scaffold, # then bind the Ray Serve deployment. The Model passes its ``shutdown()`` - # teardown via ``on_shutdown`` and its sticky-LoRA router via - # ``request_router_config``. + # teardown via ``on_shutdown``; the training controller owns handle routing. def register_routes(app: FastAPI, get_self: Any) -> None: _register_model_tinker_routes(app, get_self) _register_model_twinkle_routes(app, get_self) @@ -414,9 +391,7 @@ async def _on_shutdown(servable: Any) -> None: ModelManagement, deploy_options, deployment_name='ModelManagement', - request_router_config=RequestRouterConfig(request_router_class=StickyLoraRequestRouter), - bind_args=(model_id, nproc_per_node, device_group, device_mesh, backend, adapter_config, queue_config, - data_plane_url), + bind_args=(model_id, nproc_per_node, device_group, device_mesh, backend, queue_config, data_plane_url), bind_kwargs=kwargs, ) diff --git a/src/twinkle/server/model/backends/common.py b/src/twinkle/server/model/backends/common.py index 2ca5f8a74..0c62f3090 100644 --- a/src/twinkle/server/model/backends/common.py +++ b/src/twinkle/server/model/backends/common.py @@ -10,7 +10,8 @@ from tinker import types from typing import Any, List -from twinkle import DeviceMesh +from twinkle import DeviceMesh, remote_function +from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss from twinkle.server.validation.tinker_loss import validate_tinker_loss from twinkle.template import Template @@ -173,6 +174,35 @@ def _tinker_setup_loss(self, loss_fn: str, inputs, adapter_name: str, kwargs: di grpo_beta = kwargs.pop('beta', 0.0) self.set_loss('GRPOLoss', adapter_name=adapter_name, epsilon=epsilon, beta=grpo_beta) + @remote_function(dispatch='all', collect='none') + def preflight_loss_inputs(self, *, inputs, adapter_name: str = None, loss_fn: str = None, **kwargs) -> None: + """Validate the full, unsliced batch on every rank before any forward. + + Each backend's ``forward_backward`` is dispatched with ``slice_dp``, so a + rank that rejects its slice inside the loss has already forwarded while + its peers enter the backward collective, which deadlocks the replica + (BUG-05). Running the loss's own rule on the full batch first, on every + rank, makes them fail together before any collective is entered. + + ``loss_fn`` is the tinker route's loss selector: tinker batches carry + their RL features inside the datums, so the loss has to be set up and + the datums converted before the same rule applies. Twinkle-native + batches already pass ``old_logps``/``ref_logps``/``advantages`` as + keyword arguments and leave ``loss_fn`` unset. + """ + adapter_name = adapter_name or self._get_default_group() + if loss_fn is not None: + loss_values = extract_rl_features_for_loss(inputs) + loss_values.pop('ref_outputs', None) + self._tinker_setup_loss(loss_fn, inputs, adapter_name, kwargs) + inputs = datum_to_input_feature(inputs, self.get_template(adapter_name)) + loss_values.update(kwargs) + else: + loss_values = kwargs + loss = getattr(self.optimizer_group.get(adapter_name), 'loss_instance', None) + if loss is not None: + loss.validate_inputs(inputs, **loss_values) + def _ensure_dpo_metric(self, adapter_name: str, beta: float): """Add DPOMetric for the adapter if not already present. diff --git a/src/twinkle/server/model/backends/mock_model.py b/src/twinkle/server/model/backends/mock_model.py index 22abfb64b..55309fe14 100644 --- a/src/twinkle/server/model/backends/mock_model.py +++ b/src/twinkle/server/model/backends/mock_model.py @@ -181,7 +181,7 @@ def save(self, name: str = 'latest', output_dir: str | None = None, **kwargs: An The real backend creates ``//`` on disk and returns its path. Tinker /asample later does ``os.path.exists(...)`` on a symlink - pointing here (see ``sampler/tinker_handlers.py:72``), so the directory + pointing here (see ``sampler/tinker_handlers.py``), so the directory has to really exist for the mock path to clear. """ target = os.path.join(output_dir or '/tmp/twinkle_mock_ckpt', name) diff --git a/src/twinkle/server/model/pool/__init__.py b/src/twinkle/server/model/pool/__init__.py new file mode 100644 index 000000000..720c35471 --- /dev/null +++ b/src/twinkle/server/model/pool/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Training instance pool: routing, lifecycle decisions and Ray composition.""" diff --git a/src/twinkle/server/model/pool/actor.py b/src/twinkle/server/model/pool/actor.py new file mode 100644 index 000000000..b2d317613 --- /dev/null +++ b/src/twinkle/server/model/pool/actor.py @@ -0,0 +1,184 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""One named controller actor per server namespace.""" + +from __future__ import annotations +import asyncio +import time + +import ray + +from twinkle.server.config.training import TrainingConfig +from twinkle.server.state import get_server_state +from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.state.training_records import CreateTraining +from twinkle.server.telemetry.metrics import get_training_metrics +from twinkle.server.telemetry.worker_init import ensure_telemetry_initialized +from twinkle.utils.logger import get_logger +from .controller import TrainingController +from .driver import RayInstanceDriver + +logger = get_logger() + + +def controller_name(namespace): + return f"training-controller-{namespace}" + + +@ray.remote(num_cpus=0.1, max_restarts=-1, max_task_retries=0) +class TrainingControllerActor: + def __init__(self, config, namespace): + ensure_telemetry_initialized() + self.config = TrainingConfig.model_validate(config) + self.namespace = namespace + self.store = get_server_state().create_training_store(namespace) + self.driver = RayInstanceDriver(self.config, self.store, namespace, controller_name(namespace)) + self.core = TrainingController(self.config, self.store, self.driver) + self._metrics = get_training_metrics() + self._monitors = [] + self._initialization = asyncio.get_event_loop().create_task(self._initialize()) + + async def initialize(self): + await asyncio.shield(self._initialization) + + async def _initialize(self): + try: + await self.core.initialize() + except Exception: + logger.exception("Training controller initialization failed; restart is deferred") + # First-time startup is rolled back by the launcher; auto-restart waits before retrying. + self._restart_task = asyncio.create_task(self._restart_after_failure()) + raise + self._monitors = [asyncio.create_task(self._heartbeat()), asyncio.create_task(self._sweep())] + + async def _restart_after_failure(self): + await asyncio.sleep(max(5.0, self.config.sweep_interval_s)) + ray.kill(ray.get_actor(controller_name(self.namespace), namespace=self.namespace), no_restart=False) + + async def _heartbeat(self): + while True: + try: + active = await self.store.backend.get(self.store.key("controller", "active")) + if active is None: + await self._recover_lost_state() + return + await self.store.models.touch_replica_last_seen(self.core.epoch) + except Exception: + logger.exception("Training state heartbeat failed; admission remains fenced by the state backend") + await asyncio.sleep(5) + continue + await asyncio.gather( + *( + self._check_instance_health(instance) + for instance in list(self.core.instances.values()) + if instance.state in ("READY", "RESETTING") + ) + ) + await asyncio.sleep(5) + + async def _check_instance_health(self, instance): + async def probe(): + identity = await self.driver.health(instance) + if not identity["unhealthy"]: + await self.store.models.touch_replica_last_seen(instance.instance_id) + return identity + + healthy = await self.core.lifecycle.health(instance.instance_id, probe) + if healthy is False: + self._fail_unhealthy(instance, "Training runtime became unavailable or reported a failed worker") + + def _fail_unhealthy(self, instance, message): + self.core.lifecycle.spawn_once( + ("health", instance.instance_id), + lambda: self.core.fail_instance( + instance.instance_id, + FutureFailureRecord(reason_code="orphaned_replica", attribution="server", message=message), + ), + ) + + async def _sweep(self): + while True: + started, status = time.monotonic(), "success" + try: + await self.core.sweep() + except Exception: + status = "error" + logger.exception("Training lifecycle sweep failed; resources remain owned") + finally: + self._metrics.sweep_total.inc(tags={"status": status}) + self._metrics.sweep_seconds.observe(time.monotonic() - started, tags={"status": status}) + await asyncio.sleep(self.config.sweep_interval_s) + + async def _recover_lost_state(self): + # Reconstruct enough ownership for the next epoch to fail pending + # futures and terminate compute. Sessions were lost, so old handles + # cannot be reused, even while these rows are being restored. + failure = FutureFailureRecord( + reason_code="orphaned_replica", + attribution="server", + message="Training state was lost; create a new session and training handle.", + ) + for instance in self.core.instances.values(): + if instance.state != "STOPPED": + await self.store.save_instance(instance) + for training in self.core.trainings.values(): + if training.state in ("PENDING", "ACTIVE", "CLOSING"): + await self.store.save_training(training) + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.future( + request_id, "failed", training.training_id, failure=failure, owner_token=training.token + ) + ray.kill(ray.get_actor(controller_name(self.namespace), namespace=self.namespace), no_restart=False) + + async def register_runtime(self, instance_id, epoch, incarnation): + await self.core.register_runtime(instance_id, epoch, incarnation) + + async def create(self, request): + await self.initialize() + return await self.core.create(CreateTraining.model_validate(request)) + + async def route(self, *args): + await self.initialize() + return await self.core.route(*args) + + async def admit(self, *args, **kwargs): + await self.initialize() + return await self.core.admit(*args, **kwargs) + + async def finish_operation(self, *args, **kwargs): + return await self.core.finish_operation(*args, **kwargs) + + async def release(self, *args): + return await self.core.release(*args) + + async def fail_instance(self, *args): + return await self.core.fail_instance(*args) + + async def mark_dirty(self, instance_id, epoch): + await self.core.mark_dirty(instance_id, epoch) + + async def diagnostics(self, base_model=None): + await self.initialize() + instances = [ + instance.model_dump() + for instance in self.core.instances.values() + if base_model is None or instance.base_model == base_model + ] + return { + "ready": self.core._ready and self.core.residents_ready, + "epoch": self.core.epoch, + "instances": instances, + } + + async def shutdown(self): + if not self._initialization.done(): + self._initialization.cancel() + await asyncio.gather(self._initialization, return_exceptions=True) + restart = getattr(self, "_restart_task", None) + if restart is not None: + restart.cancel() + await asyncio.gather(restart, return_exceptions=True) + for task in self._monitors: + task.cancel() + await asyncio.gather(*self._monitors, return_exceptions=True) + await self.core.shutdown() + await self.driver.runtime.wait_for_cleanup() diff --git a/src/twinkle/server/model/pool/controller.py b/src/twinkle/server/model/pool/controller.py new file mode 100644 index 000000000..c92292624 --- /dev/null +++ b/src/twinkle/server/model/pool/controller.py @@ -0,0 +1,875 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Single-writer training lifecycle; slow runtime work stays outside the lock.""" + +from __future__ import annotations +import asyncio +import logging +import time +import traceback +import uuid +from collections.abc import Callable +from datetime import datetime +from typing import Any, Protocol + +from twinkle.server.config.training import TrainingConfig, TrainingModelConfig +from twinkle.server.exceptions import ( + RequestRejectedError, + ResourceExhaustedError, + ResourceNotFoundError, + ResourceQuotaExceededError, + TrainingExpiredError, +) +from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.resources.instances import InstanceLifecycle +from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.state.training_records import ( + LIVE_INSTANCES, + TERMINAL_TRAINING, + CreateTraining, + InstanceRecord, + RuntimeOwner, + TrainingRecord, + TrainingRoute, +) +from twinkle.server.state.training_store import TrainingStore +from twinkle.server.telemetry.metrics import TrainingInstanceMetrics, TrainingInstanceSnapshot, get_training_metrics + +logger = logging.getLogger(__name__) + + +class InstanceDriver(Protocol): + async def allocate(self, instance: InstanceRecord, spec: TrainingModelConfig) -> dict[str, Any]: ... + + async def deploy(self, instance: InstanceRecord, spec: TrainingModelConfig) -> None: ... + + async def create_training( + self, instance: InstanceRecord, training: TrainingRecord, request: CreateTraining + ) -> None: ... + + async def release_training(self, instance: InstanceRecord, training: TrainingRecord) -> None: ... + + async def shutdown(self, instance: InstanceRecord) -> None: ... + + async def verify_operation(self, instance: InstanceRecord, request_id: str) -> dict[str, Any]: ... + + async def recover_orphans(self) -> None: ... + + +def failure_for(exc: Exception) -> FutureFailureRecord: + if isinstance(exc, CapacityUnavailableError): + reason, attribution = "capacity_unavailable", "user" + elif isinstance(exc, ResourceQuotaExceededError): + reason, attribution = "resource_quota_exceeded", "user" + elif isinstance(exc, TimeoutError): + reason, attribution = "startup_timeout", "server" + elif isinstance(exc, RequestRejectedError): + reason, attribution = "invalid_request", "user" + else: + reason, attribution = "startup_failed", "server" + return FutureFailureRecord( + reason_code=reason, + message=str(exc)[:1024], + attribution=attribution, + diagnostic=traceback.format_exc() if attribution == "server" else None, + ) + + +class TrainingController: + def __init__( + self, + config: TrainingConfig, + store: TrainingStore, + driver: InstanceDriver, + *, + clock: Callable[[], float] = time.monotonic, + ): + self.config = config + self.store = store + self.driver = driver + self.clock = clock + self.epoch = "controller_" + uuid.uuid4().hex + self.instances: dict[str, InstanceRecord] = {} + self.trainings: dict[str, TrainingRecord] = {} + self.lifecycle = InstanceLifecycle("Training") + self._cleaning: set[str] = set() + self._reconciling: set[str] = set() + self._failed_instances: dict[str, FutureFailureRecord] = {} + self._lock = asyncio.Lock() + self._ready = False + self._metrics = get_training_metrics() + self._instance_metrics = TrainingInstanceMetrics(self._metrics) + + def _instance_metric_tags(self, instance: InstanceRecord) -> dict[str, str]: + spec = self.config.get(instance.base_model, instance.mode) + return { + "backend": spec.backend, + "mode": instance.mode, + "state": instance.state, + "resident": str(instance.resident).lower(), + } + + def _report_instance_metrics(self) -> None: + snapshots = [] + for instance in self.instances.values(): + if instance.state == "STOPPED": + continue + model = self.config.get(instance.base_model, instance.mode) + snapshots.append( + TrainingInstanceSnapshot( + **self._instance_metric_tags(instance), + device_type=model.device_type, + device_count=len(instance.devices), + placement_group_count=int(bool(instance.application_name)), + ) + ) + self._instance_metrics.report(snapshots) + + async def initialize(self) -> None: + """Discard old state only after owned compute has been confirmed dead.""" + old_instances = await self.store.instances() + old_trainings = await self.store.trainings() + old_controller = await self.store.backend.get(self.store.key("controller", "active")) + await self.store.begin_epoch(self.epoch) + failure = FutureFailureRecord( + reason_code="orphaned_replica", + message="Training controller restarted; create a new training handle.", + attribution="server", + ) + for training in old_trainings: + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.future(request_id, "failed", training.training_id, failure=failure) + await self.driver.recover_orphans() + for instance in old_instances: + if instance.state != "STOPPED": + await asyncio.wait_for(self.driver.shutdown(instance), self.config.shutdown_timeout_s) + instance.state = "STOPPED" + instance.stopped_at = time.time() + await self.store.save_instance(instance) + for training in old_trainings: + training.state = "FAILED" + training.terminal_reason = "orphaned_replica" + training.terminal_at = time.time() + await self.store.retire(training) + await self.store.models.rebuild_indexes() + if old_controller: + await self.store.models.unregister_replica(old_controller["epoch"]) + await self.store.models.register_replica(self.epoch, 0) + await self.store.models.touch_replica_last_seen(self.epoch) + try: + for spec in self.config.models: + for _ in range(spec.config.min_replicas): + async with self._lock: + instance = await self._new_instance(spec, resident=True) + await asyncio.shield(self.lifecycle.startups[instance.instance_id]) + except Exception: + await self.shutdown() + raise + self._ready = True + + @property + def residents_ready(self) -> bool: + return all(self.residents_ready_for(spec) for spec in self.config.models) + + def residents_ready_for(self, spec: TrainingModelConfig) -> bool: + return ( + sum( + instance.resident + and instance.instance_id not in self._failed_instances + and (instance.base_model, instance.mode) == spec.key + and instance.state in ("READY", "RESETTING") + for instance in self.instances.values() + ) + >= spec.config.min_replicas + ) + + async def register_runtime(self, instance_id: str, epoch: str, incarnation: str) -> None: + async with self._lock: + await self._assert_epoch() + instance = self.instances.get(instance_id) + if ( + epoch != self.epoch + or instance is None + or instance.incarnation is not None + or instance.state not in ("ALLOCATING", "STARTING") + ): + raise ResourceNotFoundError("Runtime incarnation is no longer allowed") + instance.incarnation = incarnation + instance.state = "STARTING" + await self.store.save_instance(instance) + self._report_instance_metrics() + + async def mark_dirty(self, instance_id: str, epoch: str) -> None: + async with self._lock: + await self._assert_epoch() + instance = self.instances.get(instance_id) + if epoch != self.epoch or instance is None: + raise ResourceNotFoundError("Unknown runtime") + instance.dirty = True + await self.store.save_instance(instance) + + async def _assert_epoch(self) -> None: + active = await self.store.backend.get(self.store.key("controller", "active")) + if not active or active["epoch"] != self.epoch: + raise ResourceNotFoundError("Training controller epoch is no longer active") + + async def _new_instance(self, spec: TrainingModelConfig, *, resident: bool = False) -> InstanceRecord: + matching = [item for item in self.instances.values() if (item.base_model, item.mode) == spec.key] + if not self.lifecycle.headroom(matching, spec.config.max_replicas): + raise CapacityUnavailableError("Training replica limit reached") + instance = InstanceRecord( + instance_id="instance_" + uuid.uuid4().hex, + controller_epoch=self.epoch, + base_model=spec.name, + mode=spec.mode, + resident=resident, + startup_deadline=self.clock() + self.config.startup_timeout_s, + ) + await self.store.save_instance(instance) + self.instances[instance.instance_id] = instance + self._report_instance_metrics() + self.lifecycle.start(instance.instance_id, lambda: self._start_instance(instance)) + return instance + + async def _start_instance(self, instance: InstanceRecord) -> None: + spec = self.config.get(instance.base_model, instance.mode) + started, status = time.monotonic(), "error" + tags = {"backend": spec.backend, "mode": spec.mode} + try: + + async def startup(): + details = await self.driver.allocate(instance, spec) + async with self._lock: + await self._assert_epoch() + if instance.state != "ALLOCATING": + raise RuntimeError("Instance became unavailable during allocation") + for name in details: + if name not in ("application_name", "placement_group_id", "node_id", "devices"): + raise ValueError(f"Unknown runtime allocation field: {name}") + allocated = instance.model_copy(deep=True, update=details) + await self.store.save_instance(allocated) + for name, value in details.items(): + setattr(instance, name, value) + self._report_instance_metrics() + # Ownership is persisted before any worker is constructed; the driver does not write records. + await self.driver.deploy(instance.model_copy(deep=True), spec) + + await asyncio.wait_for(startup(), self.config.startup_timeout_s) + async with self._lock: + await self._assert_epoch() + if self.clock() >= instance.startup_deadline: + raise TimeoutError("Training runtime became ready after its startup deadline") + if instance.state not in ("ALLOCATING", "STARTING"): + raise RuntimeError("Instance became unavailable during startup") + ready = instance.model_copy(update={"state": "READY"}) + await self.store.save_instance(ready) + instance.state = "READY" + self._report_instance_metrics() + status = "success" + except Exception as exc: + await self.fail_instance(instance.instance_id, failure_for(exc)) + raise + finally: + self._metrics.startup_seconds.observe(time.monotonic() - started, tags={**tags, "status": status}) + + def _select(self, spec: TrainingModelConfig) -> InstanceRecord | None: + candidates = [ + instance + for instance in self.instances.values() + if instance.instance_id not in self._failed_instances + and (instance.base_model, instance.mode) == spec.key + and instance.state in ("READY", "ALLOCATING", "STARTING") + and len(instance.training_ids) < spec.capacity + ] + return min( + candidates, + key=lambda instance: (instance.state != "READY", -len(instance.training_ids), instance.instance_id), + default=None, + ) + + async def create(self, request: CreateTraining) -> dict[str, str]: + # Catalog/mode errors are rejected before a future or any resource exists. + spec = self.config.get(request.base_model, request.mode) + identity, fingerprint = self.store.create_identity(request) + async with self._lock: + await self._assert_epoch() + if not self._ready: + raise ResourceExhaustedError("Training service is not ready") + previous = await self.store.get_create(identity) + if previous is not None: + if previous["fingerprint"] != fingerprint: + raise RequestRejectedError("Create identity was reused with different parameters") + return {"request_id": previous["request_id"]} + request_id = "req_" + uuid.uuid4().hex + training_id = "training_" + uuid.uuid4().hex + await self.store.future( + request_id, + "pending", + training_id, + epoch=self.epoch, + deadline=time.time() + 2 * self.config.startup_timeout_s, + owner_token=request.token, + ) + await self.store.save_create(identity, fingerprint, request_id, request.session_id) + training = None + try: + # Reserve quota before starting an instance, but leave no usable route. + training = TrainingRecord( + training_id=training_id, + token=request.token, + session_id=request.session_id, + base_model=request.base_model, + mode=request.mode, + controller_epoch=self.epoch, + training_namespace=self.store.prefix, + instance_id="", + adapter_label=request.adapter_label, + create_request_id=request_id, + creation_deadline=self.clock() + self.config.startup_timeout_s, + last_activity_at=self.clock(), + lora_config=request.payload.get("lora_config"), + user_metadata=request.payload.get("user_metadata") or {}, + ) + await self.store.reserve(training) + instance = self._select(spec) + if instance is None: + if not self.residents_ready_for(spec): + raise CapacityUnavailableError( + "Resident training capacity is recovering; retry a new create later." + ) + matching = [item for item in self.instances.values() if (item.base_model, item.mode) == spec.key] + if not self.lifecycle.headroom(matching, spec.config.max_replicas): + raise CapacityUnavailableError("Training replica limit reached; retry a new create later.") + instance = await self._new_instance(spec) + training.instance_id = instance.instance_id + training.replica_id = instance.instance_id + instance.training_ids.append(training_id) + self.trainings[training_id] = training + await self.store.save_training(training) + await self.store.save_instance(instance) + self.lifecycle.spawn(self._create_user(training, request)) + except Exception as exc: + if training is not None: + self.trainings.pop(training_id, None) + if training.instance_id in self.instances: + instance = self.instances[training.instance_id] + if training_id in instance.training_ids: + instance.training_ids.remove(training_id) + await self.store.save_instance(instance) + if not instance.training_ids and not instance.resident: + self.lifecycle.spawn(self.fail_instance(instance.instance_id, failure_for(exc))) + training.state = "FAILED" + training.terminal_reason = failure_for(exc).reason_code + await self.store.retire(training) + await self.store.future(request_id, "failed", training_id, failure=failure_for(exc)) + return {"request_id": request_id} + + async def _create_user(self, training: TrainingRecord, request: CreateTraining) -> None: + instance = self.instances[training.instance_id] + try: + await asyncio.shield(self.lifecycle.startups[instance.instance_id]) + async with self._lock: + future = await self.store.futures.get(training.create_request_id) + if future and future.status == "cancelled": + training.state = "CLOSING" + training.terminal_reason = "cancelled" + if training.state != "PENDING" or instance.state != "READY": + raise RequestRejectedError("Training creation is no longer pending") + training.creation_deadline = self.clock() + self.config.startup_timeout_s + await self.store.save_training(training) + if not await self.store.futures.try_start(training.create_request_id): + future = await self.store.futures.get(training.create_request_id) + if future and future.status == "cancelled": + training.terminal_reason = "cancelled" + raise RequestRejectedError("Training creation did not acquire execution") + remaining = training.creation_deadline - self.clock() + await asyncio.wait_for(self.driver.create_training(instance, training, request), max(0, remaining)) + # Driver includes training-run metadata commit, before public success. + async with self._lock: + await self._assert_epoch() + if self.clock() >= training.creation_deadline: + raise TimeoutError("Training initialization completed after its creation deadline") + if training.state != "PENDING" or instance.state != "READY": + raise RuntimeError("Training became unavailable during initialization") + training.state = "ACTIVE" + training.last_activity_at = self.clock() + await self.store.save_training(training) + result = ( + {"model_id": training.training_id} + if request.client_type == "tinker" + else {"training_id": training.training_id, "adapter_name": training.adapter_label} + ) + await self.store.future(training.create_request_id, "completed", training.training_id, result=result) + except Exception as exc: + if training.state not in TERMINAL_TRAINING: + await self.store.future( + training.create_request_id, "failed", training.training_id, failure=failure_for(exc) + ) + async with self._lock: + training.state = "CLOSING" + training.terminal_reason = ( + "cancelled" if training.terminal_reason == "cancelled" else failure_for(exc).reason_code + ) + await self.store.save_training(training) + await self._cleanup_user(training) + + async def _expired(self, training: TrainingRecord) -> bool: + heartbeat = await self.store.sessions.get_last_heartbeat(training.session_id) + return self._expired_with_heartbeat(training, heartbeat) + + def _session_expired(self, training: TrainingRecord, heartbeat: float | None) -> bool: + timeout = self.config.get(training.base_model, training.mode).config.adapter_config.adapter_timeout + return heartbeat is None or (timeout is not None and time.time() - heartbeat >= timeout) + + def _lifetime_expired(self, training: TrainingRecord) -> bool: + lifetime = self.config.get(training.base_model, training.mode).config.adapter_config.adapter_max_lifetime + if lifetime is None: + return False + created_at = datetime.fromisoformat(training.created_at).timestamp() + return time.time() - created_at >= lifetime + + def _expired_with_heartbeat(self, training: TrainingRecord, heartbeat: float | None) -> bool: + spec = self.config.get(training.base_model, training.mode) + if self._session_expired(training, heartbeat) or self._lifetime_expired(training): + return True + return self.clock() - training.last_activity_at >= spec.config.idle_timeout_s + + async def _authorized(self, training_id: str, token: str, session_id: str | None) -> TrainingRecord: + await self._assert_epoch() + training = self.trainings.get(training_id) + if ( + training is None + or training.token != token + or (session_id is not None and session_id != training.session_id) + ): + raise ResourceNotFoundError("Unknown training handle") + if training.state == "EXPIRED": + raise TrainingExpiredError("Training handle expired; create a new handle and explicitly load a checkpoint.") + if training.state != "ACTIVE" or training.instance_id in self._failed_instances: + raise ResourceNotFoundError("Training handle is no longer active") + heartbeat = await self.store.sessions.get_last_heartbeat(training.session_id) + if self._session_expired(training, heartbeat): + raise TrainingExpiredError("Training session expired; create a new training handle") + return training + + async def admit( + self, training_id: str, token: str, session_id: str | None, request_id: str, *, renew: bool = True + ) -> RuntimeOwner: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + expired = self._lifetime_expired(training) or ( + not training.pending_operation_ids and await self._expired(training) + ) + if request_id not in training.pending_operation_ids and expired: + training.state = "CLOSING" + training.terminal_reason = "expired" + await self.store.save_training(training) + self.lifecycle.spawn(self._cleanup_user(training)) + raise TrainingExpiredError( + "Training handle expired; create a new handle and explicitly load a checkpoint." + ) + instance = self.instances[training.instance_id] + if instance.state != "READY": + raise ResourceNotFoundError("Training runtime is unavailable") + if request_id not in training.pending_operation_ids: + training.pending_operation_ids[request_id] = {"accepted_at": self.clock(), "renew": renew} + await self.store.save_training(training) + return {"instance_id": instance.instance_id, "epoch": self.epoch} + + async def finish_operation( + self, training_id: str, request_id: str, epoch: str, *, executed: bool, finished_at: float | None = None + ) -> None: + async with self._lock: + try: + await self._assert_epoch() + except ResourceNotFoundError: + return + training = self.trainings.get(training_id) + if epoch != self.epoch or training is None: + return + ticket = training.pending_operation_ids.get(request_id) + if ticket is None: + return + updated = training.model_copy(deep=True) + updated.pending_operation_ids.pop(request_id) + if executed and ticket["renew"]: + age = max(0, time.time() - finished_at) if finished_at is not None else 0 + updated.last_activity_at = max(training.last_activity_at, self.clock() - age) + # A failed completion write must leave the original ticket retryable. + await self.store.save_training(updated) + training.pending_operation_ids = updated.pending_operation_ids + training.last_activity_at = updated.last_activity_at + cleanup = training.state == "CLOSING" and not training.pending_operation_ids + if cleanup: + self.lifecycle.spawn(self._cleanup_user(training)) + + async def route(self, training_id: str, token: str, session_id: str | None = None) -> TrainingRoute: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + return {"instance_id": training.instance_id, "base_model": training.base_model, "epoch": self.epoch} + + async def release(self, training_id: str, token: str, session_id: str | None = None) -> None: + async with self._lock: + training = await self._authorized(training_id, token, session_id) + training.state = "CLOSING" + training.terminal_reason = "released" + await self.store.save_training(training) + cleanup = not training.pending_operation_ids + if cleanup: + await self._cleanup_user(training) + config = self.config.get(training.base_model, training.mode).config.queue_config + deadline = asyncio.get_running_loop().time() + config.absolute_future_ttl(1) + self.config.shutdown_timeout_s + while training.state == "CLOSING" or self.instances[training.instance_id].state in ("DRAINING", "QUARANTINED"): + if self.instances[training.instance_id].state == "QUARANTINED": + raise RuntimeError("Runtime cleanup is quarantined; device release is not confirmed") + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError("Training cleanup is still pending; its slot remains reserved") + await asyncio.sleep(config.inline_fast_path_timeout) + + async def _cleanup_user(self, training: TrainingRecord) -> None: + instance = self.instances[training.instance_id] + claimed = False + resetting = False + closing_full = False + try: + # Each training must have exactly one cleanup, including concurrent finish callbacks. + async with self._lock: + if ( + instance.state not in ("READY", "RESETTING") + or training.state != "CLOSING" + or training.pending_operation_ids + or training.training_id in self._cleaning + ): + return + self._cleaning.add(training.training_id) + claimed = True + resetting = instance.resident and instance.mode == "full" and not instance.dirty + closing_full = instance.mode == "full" and (not instance.resident or instance.dirty) + if resetting: + updated_instance = instance.model_copy(update={"state": "RESETTING"}) + await self.store.save_instance(updated_instance) + instance.state = "RESETTING" + if closing_full: + updated_instance = instance.model_copy(update={"state": "DRAINING"}) + await self.store.save_instance(updated_instance) + instance.state = "DRAINING" + if resetting or closing_full: + self._report_instance_metrics() + if closing_full: + await self._close_instance(instance) + return + try: + spec = self.config.get(training.base_model, training.mode) + await asyncio.wait_for( + self.driver.release_training(instance, training), + spec.config.queue_config.absolute_future_ttl(spec.config.num_devices), + ) + except Exception as exc: + # A timed-out create/cleanup RPC does not prove that shared workers failed. + # Keep this slot reserved until serial cleanup can be confirmed. + if instance.mode == "lora" and any( + self.trainings[identity].state in ("ACTIVE", "PENDING") + or self.trainings[identity].pending_operation_ids + for identity in instance.training_ids + if identity != training.training_id + ): + logger.warning("Training cleanup remains pending for %s: %s", training.training_id, exc) + else: + await self.fail_instance( + instance.instance_id, + FutureFailureRecord( + reason_code="internal_error", + attribution="server", + message=f"Training cleanup failed: {exc}"[:1024], + diagnostic=traceback.format_exc(), + ), + ) + return + async with self._lock: + # A health failure may have taken ownership while release/reset + # was in flight. Its late ACK cannot retire quota or revive READY. + if ( + instance.instance_id in self._failed_instances + or instance.state not in ("READY", "RESETTING") + or training.state != "CLOSING" + ): + return + terminal = training.model_copy( + deep=True, + update={ + "state": {"expired": "EXPIRED", "released": "RELEASED", "cancelled": "CANCELLED"}.get( + training.terminal_reason, "FAILED" + ), + "terminal_at": time.time(), + }, + ) + await self.store.retire(terminal) + remaining = [identity for identity in instance.training_ids if identity != training.training_id] + close = not remaining and (not instance.resident or instance.dirty) + updated_instance = instance.model_copy( + deep=True, + update={ + "training_ids": remaining, + "state": "DRAINING" if close else ("READY" if resetting else instance.state), + }, + ) + await self.store.save_instance(updated_instance) + training.state = terminal.state + training.terminal_at = terminal.terminal_at + instance.training_ids = remaining + instance.state = updated_instance.state + self._report_instance_metrics() + if close: + await self._close_instance(instance) + finally: + if claimed: + self._cleaning.discard(training.training_id) + + async def _close_instance(self, instance: InstanceRecord) -> None: + started = time.monotonic() + tags = self._instance_metric_tags(instance) + timing_tags = {"backend": tags["backend"], "mode": tags["mode"]} + + async def quarantined(exc): + async with self._lock: + entered_quarantine = instance.state != "QUARANTINED" + quarantined = instance.model_copy(update={"state": "QUARANTINED", "failure": str(exc)}) + await self.store.save_instance(quarantined) + instance.state = "QUARANTINED" + instance.failure = str(exc) + self._report_instance_metrics() + tags = self._instance_metric_tags(instance) + if entered_quarantine: + self._metrics.quarantined_total.inc( + tags={ + "backend": tags["backend"], + "mode": tags["mode"], + "reason": TrainingInstanceMetrics.quarantine_reason(exc), + } + ) + self._metrics.shutdown_seconds.observe(time.monotonic() - started, tags={**timing_tags, "status": "error"}) + + async def stopped(): + self._metrics.shutdown_seconds.observe( + time.monotonic() - started, tags={**timing_tags, "status": "success"} + ) + await self._retire_instance(instance) + + await self.lifecycle.close( + instance.instance_id, + lambda: self.driver.shutdown(instance), + self.config.shutdown_timeout_s, + stopped, + quarantined, + ) + + async def _retire_instance(self, instance): + async with self._lock: + terminal_records = [] + terminal_at = time.time() + failure = self._failed_instances.get(instance.instance_id) + for training_id in list(instance.training_ids): + training = self.trainings[training_id] + if failure is not None: + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.future(request_id, "failed", training_id, failure=failure) + terminal = training.model_copy( + deep=True, + update={ + "pending_operation_ids": {}, + "state": {"expired": "EXPIRED", "released": "RELEASED", "cancelled": "CANCELLED"}.get( + training.terminal_reason, "FAILED" + ), + "terminal_at": terminal_at, + }, + ) + await self.store.retire(terminal) + terminal_records.append((training, terminal)) + stopped = instance.model_copy( + deep=True, update={"state": "STOPPED", "stopped_at": terminal_at, "training_ids": []} + ) + await self.store.save_instance(stopped) + for training, terminal in terminal_records: + training.pending_operation_ids.clear() + training.state = terminal.state + training.terminal_at = terminal.terminal_at + instance.state = "STOPPED" + instance.stopped_at = terminal_at + instance.training_ids.clear() + self._failed_instances.pop(instance.instance_id, None) + self._report_instance_metrics() + + async def fail_instance(self, instance_id: str, failure: FutureFailureRecord) -> None: + async with self._lock: + instance = self.instances[instance_id] + if instance.state == "STOPPED": + self._failed_instances.pop(instance_id, None) + return + # Record the retry intent without promoting unpersisted state; block new routing. + self._failed_instances[instance_id] = failure + if instance.state not in ("DRAINING", "QUARANTINED"): + draining = instance.model_copy(update={"state": "DRAINING", "failure": failure.message}) + await self.store.save_instance(draining) + instance.state = "DRAINING" + instance.failure = failure.message + self._report_instance_metrics() + try: + async with self._lock: + for training_id in instance.training_ids: + training = self.trainings[training_id] + closing = training.model_copy( + deep=True, update={"state": "CLOSING", "terminal_reason": failure.reason_code} + ) + await self.store.save_training(closing) + training.state = "CLOSING" + training.terminal_reason = failure.reason_code + for request_id in (training.create_request_id, *training.pending_operation_ids): + await self.store.future(request_id, "failed", training_id, failure=failure) + await self.lifecycle.cancel_start(instance_id) + finally: + # Once DRAINING is durable, shutdown must be attempted even if a + # later training/future write fails. Sweep will retry if this fails. + await self._close_instance(instance) + + async def sweep(self) -> None: + # Batch I/O outside the decision lock; re-verify state after re-acquiring, ignore handles added since snapshot. + session_ids = list( + { + training.session_id + for training in self.trainings.values() + if training.state == "ACTIVE" and not training.pending_operation_ids + } + ) + heartbeats = await self.store.sessions.get_last_heartbeats(session_ids) + async with self._lock: + await self._assert_epoch() + for training in self.trainings.values(): + if ( + training.state != "ACTIVE" + or training.pending_operation_ids + or training.session_id not in heartbeats + ): + continue + if self._expired_with_heartbeat(training, heartbeats[training.session_id]): + closing_record = training.model_copy(update={"state": "CLOSING", "terminal_reason": "expired"}) + await self.store.save_training(closing_record) + training.state = "CLOSING" + training.terminal_reason = "expired" + uncertain = [ + (training, request_id) + for training in self.trainings.values() + for request_id, ticket in training.pending_operation_ids.items() + if request_id not in self._reconciling + and self.clock() - ticket["accepted_at"] >= self.config.sweep_interval_s + ] + self._reconciling.update(request_id for _, request_id in uncertain) + closing = [ + training + for training in self.trainings.values() + if training.state == "CLOSING" + and not training.pending_operation_ids + and training.training_id not in self._cleaning + ] + for training in closing: + self.lifecycle.spawn(self._cleanup_user(training)) + for training, request_id in uncertain: + self.lifecycle.spawn(self._reconcile_operation(training, request_id)) + for instance in list(self.instances.values()): + failure = self._failed_instances.get(instance.instance_id) + if failure is not None: + self.lifecycle.spawn(self.fail_instance(instance.instance_id, failure)) + elif instance.state in ("DRAINING", "QUARANTINED"): + self.lifecycle.spawn(self._close_instance(instance)) + # Rebuilding one specification must not stop expiry or other models. + for spec in self.config.models: + async with self._lock: + matching = [item for item in self.instances.values() if (item.base_model, item.mode) == spec.key] + missing = self.lifecycle.missing( + matching, + spec.config.min_replicas, + spec.config.max_replicas, + lambda instance: instance.resident + and instance.instance_id not in self._failed_instances + and instance.state in LIVE_INSTANCES, + ) + for _ in range(missing): + await self._new_instance(spec, resident=True) + await self.store.cleanup_receipts() + # Terminal diagnostics expire like their backend tombstones. Repeated + # model creation must not grow controller memory for the server lifetime. + async with self._lock: + cutoff = time.time() - self.store.expiration_timeout + for identity, training in list(self.trainings.items()): + if training.state in TERMINAL_TRAINING and training.terminal_at is not None: + if training.terminal_at <= cutoff: + self.trainings.pop(identity) + for identity, instance in list(self.instances.items()): + if ( + instance.state == "STOPPED" + and instance.stopped_at is not None + and instance.stopped_at <= cutoff + and not instance.training_ids + ): + self.instances.pop(identity) + self.lifecycle.startups.pop(identity, None) + self._report_instance_metrics() + + async def _reconcile_operation(self, training: TrainingRecord, request_id: str) -> None: + try: + instance = self.instances[training.instance_id] + if instance.state != "READY": + return + result = await asyncio.wait_for( + self.driver.verify_operation(instance, request_id), self.config.shutdown_timeout_s + ) + if result["state"] == "finished": + await self.finish_operation( + training.training_id, + request_id, + self.epoch, + executed=result["executed"], + finished_at=result.get("finished_at"), + ) + elif result["state"] == "missing": + await self.store.future( + request_id, + "failed", + training.training_id, + failure=FutureFailureRecord( + reason_code="internal_error", + attribution="server", + message="Operation was not committed to the runtime queue", + ), + ) + await self.finish_operation(training.training_id, request_id, self.epoch, executed=False) + # An unsafe/unreachable runtime is handled by the worker health monitor. + except Exception: + logger.exception("Operation reconciliation deferred for %s", request_id) + finally: + self._reconciling.discard(request_id) + + async def wait_for_jobs(self) -> None: + await self.lifecycle.join() + + async def shutdown(self) -> None: + self._ready = False + outcomes = await asyncio.gather( + *( + self.fail_instance( + instance.instance_id, + FutureFailureRecord( + reason_code="orphaned_replica", + message="Training service is shutting down", + attribution="server", + ), + ) + for instance in list(self.instances.values()) + ), + return_exceptions=True, + ) + await self.wait_for_jobs() + errors = [outcome for outcome in outcomes if isinstance(outcome, BaseException)] + for error in errors: + logger.error("Training instance shutdown failed", exc_info=(type(error), error, error.__traceback__)) + if errors: + raise errors[0] + self.lifecycle.require_released(self.instances.values()) diff --git a/src/twinkle/server/model/pool/driver.py b/src/twinkle/server/model/pool/driver.py new file mode 100644 index 000000000..f3140b32d --- /dev/null +++ b/src/twinkle/server/model/pool/driver.py @@ -0,0 +1,150 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Ray/Serve effects behind the controller's small runtime interface.""" + +from __future__ import annotations +import asyncio +import re +import time + +from twinkle.server.resources.runtime import RayRuntime +from twinkle.server.state.training_records import InstanceRecord, TrainingRuntimeContext +from twinkle.server.telemetry.metrics import get_training_metrics +from twinkle.utils.logger import get_logger + +logger = get_logger() + + +class RayInstanceDriver: + def __init__(self, config, store, namespace, controller_name): + self.config, self.store = config, store + self.namespace, self.controller_name = namespace, controller_name + self._metrics = get_training_metrics() + self.runtime = RayRuntime( + namespace, allocation_timeout=config.allocation_timeout_s, shutdown_timeout=config.shutdown_timeout_s + ) + + def application_name(self, instance): + return self.runtime.pool.name(instance.controller_epoch, instance.instance_id) + + async def allocate(self, instance, spec): + started, status = time.monotonic(), "error" + tags = {"backend": spec.backend, "mode": spec.mode, "device_type": spec.device_type} + try: + allocation = await self.runtime.pool.allocate( + spec.config.num_devices, spec.device_type, instance.controller_epoch, instance.instance_id + ) + status = "success" + return { + "application_name": allocation.name, + "placement_group_id": allocation.placement_group_id, + "node_id": allocation.node_id, + "devices": list(allocation.devices), + } + finally: + self._metrics.allocation_seconds.observe(time.monotonic() - started, tags={**tags, "status": status}) + + async def deploy(self, instance, spec): + from ray.util.placement_group import get_placement_group + + from twinkle.server.model.app import build_model_app + + name = self.application_name(instance) + context = TrainingRuntimeContext( + instance_id=instance.instance_id, + epoch=instance.controller_epoch, + namespace=self.namespace, + controller_name=self.controller_name, + registry_name=name + "-registry", + ) + kwargs = spec.config.model_dump( + exclude={ + "num_devices", + "device_mesh", + "min_replicas", + "max_replicas", + "idle_timeout_s", + "adapter_config", + "queue_config", + }, + exclude_none=True, + ) + if "hybrid" in kwargs: + kwargs["hybrid"] = spec.config.hybrid + kwargs["train_mode"] = "hybrid" if kwargs.get("hybrid") else spec.mode + options = self.runtime.deployment_options(spec.deploy_options, name, spec.env_vars) + app = build_model_app( + model_id=spec.model_id, + backend=spec.backend, + nproc_per_node=spec.config.num_devices, + device_group={"name": "model", "ranks": spec.config.num_devices, "device_type": spec.device_type}, + device_mesh={**spec.config.device_mesh.model_dump(exclude_none=True), "device_type": spec.device_type}, + deploy_options=options, + queue_config=spec.config.queue_config, + data_plane_url=spec.data_plane_url, + training_context=context, + placement_group=await asyncio.to_thread(get_placement_group, name), + allocation_devices=list(instance.devices), + **kwargs, + ) + await self.runtime.deploy( + instance.instance_id, + name, + app, + f"/api/v1/training-runtime/{instance.instance_id}", + lambda: self.shutdown(instance), + ) + + async def handle(self, instance): + return await self.runtime.handle(instance.application_name or self.application_name(instance)) + + async def create_training(self, instance, training, request): + handle = await self.handle(instance) + await handle.initialize_training.remote(training.model_dump(), request.model_dump()) + + async def release_training(self, instance, training): + handle = await self.handle(instance) + await handle.release_training.remote(training.model_dump(), instance.resident and not instance.dirty) + + async def verify_operation(self, instance, request_id): + handle = await self.handle(instance) + return await handle.verify_operation.remote(request_id) + + async def shutdown(self, instance: InstanceRecord): + name = instance.application_name or self.application_name(instance) + try: + spec = self.config.get(instance.base_model, instance.mode) + except Exception: + tags = {"backend": "unknown", "mode": instance.mode, "device_type": "unknown"} + else: + tags = {"backend": spec.backend, "mode": spec.mode, "device_type": spec.device_type} + + def observe_removal(duration, status): + self._metrics.removal_seconds.observe(duration, tags={**tags, "status": status}) + + await self.runtime.shutdown(instance.instance_id, name, observe_removal=observe_removal) + await self.store.models.unregister_replica_capacity(instance.instance_id) + + async def health(self, instance): + async with asyncio.timeout(5): + handle = await self.handle(instance) + return await handle.runtime_identity.remote() + + async def recover_orphans(self): + """Ownership names survive loss of the memory state actor.""" + pattern = re.compile( + r"resources-" + re.escape(self.namespace) + r"-(controller_[0-9a-f]{32})-(instance_[0-9a-f]{32})$" + ) + + def decode(name): + match = pattern.fullmatch(name) + if match: + return InstanceRecord( + controller_epoch=match[1], + instance_id=match[2], + application_name=name, + base_model="", + mode="full", + startup_deadline=0, + ) + + await self.runtime.recover_orphans(decode, self.shutdown) diff --git a/src/twinkle/server/model/pool/router.py b/src/twinkle/server/model/pool/router.py new file mode 100644 index 000000000..ca5b77734 --- /dev/null +++ b/src/twinkle/server/model/pool/router.py @@ -0,0 +1,265 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Stable CPU endpoint for one catalog model; requests route by training ID.""" + +from __future__ import annotations +import asyncio +import uuid +from types import SimpleNamespace + +import json +import ray +from fastapi import Depends, FastAPI, Request +from fastapi.exceptions import RequestValidationError +from pydantic import ValidationError +from tinker import types as tinker + +from twinkle.hub import HubOperation +from twinkle.protocol.types import model as native +from twinkle.server.checkpoint import _resolve_client_save_dir, create_checkpoint_manager, validate_user_path +from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app +from twinkle.server.exceptions import ProtocolMismatchError, RequestRejectedError, ResourceNotFoundError +from twinkle.server.lifecycle.envelope import envelope_from_record +from twinkle.server.lifecycle.proxy import ServiceProxy +from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request +from twinkle.server.state import get_server_state +from twinkle.server.state.training_records import CreateTraining +from twinkle.server.task_queue import TaskQueueMixin +from twinkle.server.validation import assert_request_supported +from .actor import controller_name + +NATIVE_REQUESTS = { + "forward": native.ForwardRequest, + "forward_only": native.ForwardOnlyRequest, + "forward_backward": native.ForwardBackwardTaskRequest, + "calculate_loss": native.AdapterRequest, + "backward": native.AdapterRequest, + "forward_from_data_plane": native.DataPlaneForwardRequest, + "forward_backward_from_data_plane": native.DataPlaneForwardRequest, + "forward_only_from_data_plane": native.DataPlaneForwardOnlyRequest, + "step": native.StepRequest, + "zero_grad": native.AdapterRequest, + "lr_step": native.LrStepRequest, + "clip_grad_norm": native.ClipGradNormRequest, + "clip_grad_and_step": native.ClipGradAndStepRequest, + "get_train_configs": native.AdapterRequest, + "set_loss": native.SetLossRequest, + "set_optimizer": native.SetOptimizerRequest, + "set_lr_scheduler": native.SetLrSchedulerRequest, + "set_template": native.SetTemplateRequest, + "set_processor": native.SetProcessorRequest, + "add_metric": native.AddMetricRequest, + "apply_patch": native.ApplyPatchRequest, + "calculate_metric": native.CalculateMetricRequest, + "save": native.SaveRequest, + "load": native.LoadRequest, + "resume_from_checkpoint": native.ResumeFromCheckpointRequest, + "remove_adapter": native.AdapterRequest, + "add_adapter_to_model": native.AddAdapterRequest, + "upload_to_hub": native.UploadToHubRequest, + "create": native.CreateRequest, +} +TINKER_REQUESTS = { + "create_model": tinker.CreateModelRequest, + "get_info": tinker.GetInfoRequest, + "unload_model": tinker.UnloadModelRequest, + "forward": tinker.ForwardRequest, + "forward_backward": tinker.ForwardBackwardRequest, + "optim_step": tinker.OptimStepRequest, + "save_weights": tinker.SaveWeightsRequest, + "load_weights": tinker.LoadWeightsRequest, + "save_weights_for_sampler": tinker.SaveWeightsForSamplerRequest, +} + + +class TrainingRouter(TaskQueueMixin, LazyCleanupMixin): + async def __init__(self, config, model, namespace, http_options): + self.config = config + self.model = model + self.backend = model.backend + self.state = get_server_state() + self.store = self.state.create_training_store(namespace) + self.controller = ray.get_actor(controller_name(namespace), namespace=namespace) + self.proxy = ServiceProxy(http_options) + self.replica_id = "training-router_" + uuid.uuid4().hex + self._init_task_queue(None, deployment_name="TrainingRouter") + await self.state.register_replica(self.replica_id, 0) + self._heartbeat = asyncio.create_task(self._keep_alive()) + + async def _keep_alive(self): + while True: + await self.state.touch_replica_last_seen(self.replica_id) + await asyncio.sleep(5) + + async def shutdown(self): + self._heartbeat.cancel() + await asyncio.gather(self._heartbeat, return_exceptions=True) + await self.shutdown_task_queue() + await self.proxy.close() + await self.state.unregister_replica(self.replica_id) + + async def _check_session(self, session): + if not session or await self.state.get_session_last_heartbeat(session) is None: + raise ResourceNotFoundError("Unknown or expired session") + + async def native_request(self, request, endpoint): + schema = NATIVE_REQUESTS.get(endpoint) + if schema is None: + raise ResourceNotFoundError("Unknown training endpoint") + try: + body = schema.model_validate(await request.json()) + except ValidationError as exc: + raise RequestValidationError(exc.errors()) from exc + token = get_token_from_request(request) + if endpoint == "create": + if not body.training_handles: + raise ProtocolMismatchError("Upgrade Twinkle SDK: training_handles handshake is required") + return native.CreateResponse(training_handles=True) + if endpoint == "add_adapter_to_model": + if not body.client_create_id: + raise ProtocolMismatchError("Upgrade Twinkle SDK: client_create_id is required") + if not body.adapter_name or not validate_user_path(token, body.adapter_name): + raise RequestRejectedError("Invalid adapter_name") + try: + config = json.loads(body.config) if body.config is not None else None + except (ValueError, TypeError) as exc: + raise RequestRejectedError("Expected a serialized LoRA config") from exc + mode = "full" if config is None else "lora" + spec = self.config.get(self.model.name, mode) + if config is not None: + rank = config.get("r") if isinstance(config, dict) else None + if ( + not isinstance(config, dict) + or config.get("_TWINKLE_TYPE_") != "LoraConfig" + or type(rank) is not int + or not 0 < rank <= spec.config.max_r + ): + raise RequestRejectedError("LoRA config/rank is not supported by this training specification") + payload = body.model_dump(mode="json") + if body.save_dir: + try: + payload["save_dir"] = _resolve_client_save_dir(body.save_dir).as_posix() + except ValueError as exc: + raise RequestRejectedError(str(exc)) from exc + session = get_session_id_from_request(request) + await self._check_session(session) + ref = await self.controller.create.remote( + CreateTraining( + base_model=self.model.name, + mode=mode, + token=token, + session_id=session, + create_key=body.client_create_id, + client_type="twinkle", + adapter_label=body.adapter_name, + payload=payload, + ).model_dump() + ) + return envelope_from_record(ref["request_id"], await self.state.get_future(ref["request_id"])) + if endpoint == "upload_to_hub": + + async def upload(): + checkpoint = body.checkpoint_dir + if checkpoint.startswith("twinkle://"): + checkpoint = create_checkpoint_manager(token, client_type="twinkle").resolve_checkpoint_dir( + checkpoint + ) + await self.call_backend( + HubOperation.push_to_hub, + repo_id=body.hub_model_id, + folder_path=checkpoint, + token=body.hub_token or token, + private=True, + ) + + ref = await self.schedule_background_task(upload, task_type=endpoint, token=token) + return envelope_from_record(ref["request_id"], await self.state.get_future(ref["request_id"])) + session = get_session_id_from_request(request) + await self._check_session(session) + if endpoint == "remove_adapter": + return await self._release(body.training_id, token, session, "twinkle") + capability = endpoint if endpoint in ("forward", "backward", "calculate_loss") else None + assert_request_supported(SimpleNamespace(backend=self.backend), body, capability=capability) + route = await self.controller.route.remote(body.training_id, token, session) + if route["base_model"] != self.model.name: + raise ResourceNotFoundError("Training handle belongs to another model endpoint") + return await self.proxy.proxy_request(request, f"twinkle/{endpoint}", route["instance_id"], "training-runtime") + + async def tinker_request(self, request, endpoint): + schema = TINKER_REQUESTS.get(endpoint) + if schema is None: + raise ResourceNotFoundError("Unknown training endpoint") + try: + body = schema.model_validate(await request.json()) + except ValidationError as exc: + raise RequestValidationError(exc.errors()) from exc + token = get_token_from_request(request) + if endpoint == "create_model": + if body.base_model != self.model.name: + raise RequestRejectedError("Model does not match the catalog endpoint") + mode = "lora" if body.lora_config is not None else "full" + spec = self.config.get(self.model.name, mode) + if body.lora_config is not None and not 0 < body.lora_config.rank <= spec.config.max_r: + raise RequestRejectedError("LoRA rank exceeds the training specification") + await self._check_session(body.session_id) + return await self.controller.create.remote( + CreateTraining( + base_model=self.model.name, + mode=mode, + token=token, + session_id=body.session_id, + create_key=str(body.model_seq_id), + client_type="tinker", + payload=body.model_dump(mode="json"), + ).model_dump() + ) + if endpoint == "unload_model": + return await self._release(body.model_id, token, None, "tinker") + route = await self.controller.route.remote(body.model_id, token) + if route["base_model"] != self.model.name: + raise ResourceNotFoundError("Training handle belongs to another model endpoint") + return await self.proxy.proxy_request(request, f"tinker/{endpoint}", route["instance_id"], "training-runtime") + + async def _release(self, training_id, token, session, client_type): + await self.controller.route.remote(training_id, token, session) + + async def release(): + await self.controller.release.remote(training_id, token, session) + return {"model_id": training_id} if client_type == "tinker" else None + + ref = await self.schedule_background_task(release, task_type="release_training", token=token) + if client_type == "tinker": + return {"request_id": ref["request_id"]} + return envelope_from_record(ref["request_id"], await self.state.get_future(ref["request_id"])) + + +def build_training_router(config, model, namespace, http_options, deploy_options): + def routes(app: FastAPI, get_self): + @app.get("/healthz") + async def health(self=Depends(get_self)): + diagnostic = await self.controller.diagnostics.remote(self.model.name) + active = [item for item in diagnostic["instances"] if item["state"] != "STOPPED"] + return { + "healthy": diagnostic["ready"] and not any(item["state"] == "QUARANTINED" for item in active), + "status": "unloaded" if not active else "loaded", + "instances": active, + } + + @app.post("/twinkle/{endpoint}") + async def native_route(request: Request, endpoint: str, self=Depends(get_self)): + return await self.native_request(request, endpoint) + + @app.post("/tinker/{endpoint}") + async def tinker_route(request: Request, endpoint: str, self=Depends(get_self)): + return await self.tinker_request(request, endpoint) + + async def shutdown(servable): + await servable.shutdown() + + app = build_deployment_app("TrainingRouter", routes, on_shutdown=shutdown) + return bind_deployment( + app, + TrainingRouter, + deploy_options, + deployment_name="TrainingRouter", + bind_args=(config, model, namespace, http_options), + ) diff --git a/src/twinkle/server/model/routing.py b/src/twinkle/server/model/routing.py deleted file mode 100644 index 02b7669a3..000000000 --- a/src/twinkle/server/model/routing.py +++ /dev/null @@ -1,75 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from ray.serve.request_router import FIFOMixin, MultiplexMixin, PendingRequest, ReplicaID, RequestRouter, RunningReplica - -from twinkle.server.state import ServerState, get_server_state -from twinkle.utils.logger import get_logger - -logger = get_logger() - - -class StickyLoraRequestRouter(FIFOMixin, MultiplexMixin, RequestRouter): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - self.state: ServerState = get_server_state() - - async def choose_replicas( - self, - candidate_replicas: list[RunningReplica], - pending_request: PendingRequest | None = None, - ) -> list[list[RunningReplica]]: - """ - This method chooses the best replica for the request based on - multiplexed and available lora count. The algorithm - works as follows: - - 1. Populate top_ranked_replicas based on available replicas based on - multiplex_id (only one replica is chosen) - 2. Populate and override top_ranked_replicas info based on available lora - slots of the replica. - """ - - # Take the best set of replicas for the multiplexed model - if (pending_request is not None and pending_request.metadata.multiplexed_model_id): - ranked_replicas_multiplex: list[RunningReplica] = (self.rank_replicas_via_multiplex( - replicas=candidate_replicas, - multiplexed_model_id=pending_request.metadata.multiplexed_model_id, - ))[0] - - # If found any replica, return it - if ranked_replicas_multiplex: - logger.debug('[Router] Found replica for multiplexed model !!!') - return [ranked_replicas_multiplex] - - # Dictionary to hold the top-ranked replicas - top_ranked_replicas: dict[ReplicaID, RunningReplica] = {} - - # Filter out replicas that are not available (queue length exceed max ongoing request) - ranked_replicas_locality = self.select_available_replicas(candidates=candidate_replicas) - - for replica in ranked_replicas_locality: - top_ranked_replicas[replica.replica_id] = replica - - # Filter out replicas that exceed max lora count (query from server state) - candidate_ids = [r.replica_id.unique_id for r in top_ranked_replicas.values()] - available_ids = set(await self.state.get_available_replica_ids(candidate_ids)) - if available_ids: - top_ranked_replicas = { - rid: r - for rid, r in top_ranked_replicas.items() if r.replica_id.unique_id in available_ids - } - - if not top_ranked_replicas: - # No replica has remaining LoRA capacity – fall back to all candidates - logger.debug('[Router] No replica has remaining LoRA capacity') - return [candidate_replicas] - - logger.debug('[Router] StickyLoraRequestRouter choosing replica for request') - - # Take the replica with minimum throughput. - min_throughput_replicas = min( - [replica for replica in top_ranked_replicas.values()], - key=lambda r: r.routing_stats.get('throughput', 0), - ) - return [[min_throughput_replicas]] diff --git a/src/twinkle/server/model/runtime.py b/src/twinkle/server/model/runtime.py new file mode 100644 index 000000000..2f4916628 --- /dev/null +++ b/src/twinkle/server/model/runtime.py @@ -0,0 +1,221 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Managed training operations on the existing model queue and backend gate.""" + +from __future__ import annotations +import asyncio +import os +import uuid +from dataclasses import dataclass + +from twinkle.protocol.serialize import deserialize_object +from twinkle.protocol.types import model as native +from twinkle.protocol.types.training import CreateModelRequest +from twinkle.protocol.types.training import LoraConfig as RunLoraConfig +from twinkle.server.checkpoint import create_training_run_manager +from twinkle.server.exceptions import ResourceNotFoundError +from twinkle.server.lifecycle.submit import backend_kwargs +from twinkle.server.middleware.auth import get_session_id_from_request, get_token_from_request +from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.state.training_records import CreateTraining, TrainingRecord, TrainingRuntimeContext +from twinkle.server.utils import get_template_for_model +from twinkle.utils.logger import get_logger + +logger = get_logger() + + +@dataclass +class TrainingExecutionState: + """Worker-local state; ownership and expiration belong to the controller.""" + + grad_ready: bool = False + + +class TrainingRuntime: + """Execute controller-owned training handles on one allocated backend.""" + + training_context: TrainingRuntimeContext + + _training_states: dict[str, TrainingExecutionState] + + def get_training_state(self, training_id: str) -> TrainingExecutionState: + state = self._training_states.get(training_id) + if state is None: + raise ResourceNotFoundError("Training is not registered on this runtime") + return state + + def assert_resource_exists(self, training_id: str) -> None: + self.get_training_state(training_id) + + async def wait_for_training_workers(self): + actors = getattr(self.model, "_actors", None) or [] + await asyncio.gather(*(actor.__ray_ready__.remote() for actor in actors)) + + async def initialize_training_runtime(self, context, placement_group, devices): + import ray + + from twinkle.infra._ray.ray_helper import RayHelper + from twinkle.infra._ray.resource_manager import ResourceManager + + self._training_controller = ray.get_actor(context.controller_name, namespace=context.namespace) + await self._training_controller.register_runtime.remote(context.instance_id, context.epoch, uuid.uuid4().hex) + os.environ["TWINKLE_RAY_REGISTRY_NAME"] = context.registry_name + os.environ["TWINKLE_RAY_NAMESPACE"] = context.namespace + RayHelper.resource_manager = ResourceManager.from_allocation(self.device_group, placement_group, devices) + + async def training_request_start(self, request): + token = get_token_from_request(request) + if self._model_unhealthy: + raise ResourceNotFoundError("Training runtime is unavailable") + body = await request.json() + training_id = body.get("training_id") or body.get("model_id") + session = get_session_id_from_request(request) if "/twinkle/" in request.url.path else None + if "/twinkle/" in request.url.path and not session: + raise ResourceNotFoundError("Training requests require the creating session") + route = await self._training_controller.route.remote(training_id, token, session) + if route["instance_id"] != self.training_context.instance_id or route["epoch"] != self.training_context.epoch: + raise ResourceNotFoundError("Runtime does not own this training handle") + return token + + async def _prepare_queued_task(self, task): + if task.model_id is None: + return + context = self.training_context + await self._training_controller.admit.remote( + task.model_id, + task.token, + None, + task.request_id, + renew=task.task_type not in ("get_info", "get_train_configs"), + ) + + async def finished(executed, unsafe): + if unsafe: + self._model_unhealthy = True + self._training_controller.fail_instance.remote( + context.instance_id, + FutureFailureRecord( + reason_code="execution_timeout", + attribution="server", + message="Backend execution timed out; runtime is being terminated.", + ), + ) + else: + try: + outcome = self._compute_worker.get_task_outcome(task.request_id) + await self._training_controller.finish_operation.remote( + task.model_id, + task.request_id, + context.epoch, + executed=executed, + finished_at=outcome["finished_at"], + ) + return True + except Exception: + logger.exception("Training completion notification failed for %s", task.request_id) + + task.on_finished = finished + if task.task_type == "apply_patch": + await self._training_controller.mark_dirty.remote(context.instance_id, context.epoch) + + async def _run_lifecycle_operation(self, factory, task_type): + ref = await self.schedule_task(factory, task_type=task_type) + deadline = asyncio.get_running_loop().time() + self.task_queue_config.absolute_future_ttl(1) + while True: + future = await self.state.get_future(ref["request_id"]) + if future and future["status"] in ("failed", "cancelled"): + raise RuntimeError(future["failure"]["message"]) + if future and future["status"] == "completed": + return future["result"] + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError("Training lifecycle operation has not finished") + await asyncio.sleep(self.task_queue_config.inline_fast_path_timeout) + + async def initialize_training(self, training_data, request_data): + training = TrainingRecord.model_validate(training_data) + request = CreateTraining.model_validate(request_data) + + async def initialize(): + adapter = training.training_id + backend_adapter = self.resolve_model_adapter_name(adapter) + self._training_states[adapter] = TrainingExecutionState() + if request.client_type == "twinkle": + body = native.AddAdapterRequest.model_validate(request.payload) + # Dynamic imports (e.g. PEFT) and deserialization must not block the health-probe event loop. + config = await asyncio.to_thread(deserialize_object, body.config) + if not self.is_full_mode: + await self.call_backend(self.model.add_adapter_to_model, adapter, config, **backend_kwargs(body)) + lora = ( + None + if config is None + else RunLoraConfig(rank=config.r, train_unembed=False, train_mlp=True, train_attn=True) + ) + run_config = CreateModelRequest( + base_model=training.base_model, + lora_config=lora, + save_dir=body.save_dir, + user_metadata={"adapter_name": training.adapter_label}, + ) + else: + from tinker import types + + run_config = types.CreateModelRequest.model_validate(request.payload) + if not self.is_full_mode: + + def build_lora_config(): + from peft import LoraConfig + + return LoraConfig(r=run_config.lora_config.rank, target_modules="all-linear") + + config = await asyncio.to_thread(build_lora_config) + await self.call_backend(self.model.add_adapter_to_model, adapter_name=adapter, config_or_dir=config) + await self.call_backend( + self.model.set_template, + get_template_for_model(self.base_model), + adapter_name=backend_adapter, + model_id=self.base_model, + ) + await self.call_backend(self.model.set_processor, "InputProcessor", adapter_name=backend_adapter) + await self.call_backend(self.model.set_optimizer, "Adam", adapter_name=backend_adapter) + + def save_training_metadata(): + create_training_run_manager(training.token, client_type=request.client_type).save(adapter, run_config) + + await asyncio.to_thread(save_training_metadata) + + await self._run_lifecycle_operation(initialize, "initialize_training") + + async def release_training(self, training_data, reset_full): + training = TrainingRecord.model_validate(training_data) + + async def release(): + adapter = training.training_id + if adapter not in self._training_states: + return + if self.is_full_mode: + if reset_full: + await self.call_backend(self.model.reload_initial_weights, reset_training_state=True) + else: + await self.call_backend(self.model.remove_adapter, adapter) + self._training_states.pop(adapter) + + await self._run_lifecycle_operation(release, "release_training") + + async def runtime_identity(self): + import ray + + actors = getattr(self.model, "_actors", None) or [] + states = await asyncio.to_thread( + lambda: [ray._private.state.actors(actor._actor_id.hex()).get("State") for actor in actors] + ) + return { + "pid": os.getpid(), + "instance_id": self.training_context.instance_id, + "unhealthy": self._model_unhealthy or any(state != "ALIVE" for state in states), + } + + async def verify_operation(self, request_id): + if self._model_unhealthy: + return {"state": "unsafe"} + if self._compute_worker.has_task(request_id): + return {"state": "pending"} + return self._compute_worker.get_task_outcome(request_id) or {"state": "missing"} diff --git a/src/twinkle/server/model/tinker_handlers.py b/src/twinkle/server/model/tinker_handlers.py index b455e63d2..b1fd47f69 100644 --- a/src/twinkle/server/model/tinker_handlers.py +++ b/src/twinkle/server/model/tinker_handlers.py @@ -9,7 +9,6 @@ """ from __future__ import annotations -import traceback from collections.abc import Callable from fastapi import Depends, FastAPI, Request from functools import partial @@ -20,9 +19,8 @@ from .app import ModelManagement from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager -from twinkle.server.exceptions import FullModeBusyError -from twinkle.server.task_queue.types import UserTaskError -from twinkle.server.utils import get_template_for_model +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.lifecycle.submit import submit_training_task from twinkle.server.validation.tinker_loss import validate_tinker_loss from twinkle.utils.logger import get_logger @@ -42,60 +40,7 @@ async def create_model( body: types.CreateModelRequest, self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _create_adapter(): - _model_id = None - try: - # Validate lora_config against the deployment's train_mode up front. - if self.is_full_mode and body.lora_config: - raise UserTaskError('This deployment runs in full-parameter (exclusive) mode; do not pass ' - 'lora_config. Use create_full_training_client (or omit lora_config).') - if (not self.is_full_mode) and (not body.lora_config): - raise UserTaskError('This deployment runs in LoRA mode; lora_config is required. ' - 'Use create_lora_training_client.') - # Exclusive full-parameter training: reject early (before touching - # state) if another tenant already holds the deployment. - if self.is_full_mode: - self.assert_full_mode_available() - - _model_id = await self.state.register_model( - body.model_dump(), token=token, replica_id=self.replica_id, session_id=body.session_id) - adapter_name = self.get_adapter_name(adapter_name=_model_id) - model_adapter = self.resolve_model_adapter_name(adapter_name) - # Select template based on model type - template = get_template_for_model(self.base_model) - if self.is_full_mode: - self.register_resource(adapter_name, token, session_id=body.session_id) - await self.call_backend( - self.model.set_template, template, adapter_name=model_adapter, model_id=self.base_model) - await self.call_backend(self.model.set_processor, 'InputProcessor', adapter_name=model_adapter) - await self.call_backend(self.model.set_optimizer, 'Adam', adapter_name=model_adapter) - self.set_resource_state(adapter_name, 'grad_ready', False) - else: - from peft import LoraConfig - lora_cfg = LoraConfig(r=body.lora_config.rank, target_modules='all-linear') - self.register_resource(adapter_name, token, session_id=body.session_id) - await self.call_backend( - self.model.add_adapter_to_model, adapter_name=adapter_name, config_or_dir=lora_cfg) - await self.call_backend( - self.model.set_template, template, adapter_name=adapter_name, model_id=self.base_model) - await self.call_backend(self.model.set_processor, 'InputProcessor', adapter_name=adapter_name) - await self.call_backend(self.model.set_optimizer, 'Adam', adapter_name=adapter_name) - self.set_resource_state(adapter_name, 'grad_ready', False) - training_run_manager = create_training_run_manager(token, client_type='tinker') - training_run_manager.save(_model_id, body) - return types.CreateModelResponse(model_id=_model_id) - except FullModeBusyError as e: - raise UserTaskError(str(e)) from e - except Exception: - if _model_id: - adapter_name = self.get_adapter_name(adapter_name=_model_id) - await self._cleanup_adapter(adapter_name) - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_create_adapter, token=token, task_type='create_model') + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') @app.post('/tinker/get_info') async def get_info(request: Request, body: types.GetInfoRequest, @@ -123,14 +68,7 @@ async def unload_model( body: types.UnloadModelRequest, self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) - - async def _do_unload(): - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - await self._cleanup_adapter(adapter_name) - return types.UnloadModelResponse(model_id=body.model_id) - - return await self.schedule_task(_do_unload, model_id=body.model_id, token=token, task_type='unload_model') + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') @app.post('/tinker/forward') async def forward(request: Request, body: types.ForwardRequest, @@ -138,29 +76,28 @@ async def forward(request: Request, body: types.ForwardRequest, token = await self._on_request_start(request) async def _do_forward(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - model_adapter = self.resolve_model_adapter_name(adapter_name) - datum_list = body.forward_input.data - loss_fn_config = body.forward_input.loss_fn_config or {} - output, loss = await self.call_backend( - self.model.tinker_forward_only, inputs=datum_list, adapter_name=model_adapter, **loss_fn_config) - return types.ForwardBackwardOutput( - loss_fn_output_type='CrossEntropyLossReturn', - loss_fn_outputs=output, - metrics={'loss:mean': loss}, - ) - except Exception: - logger.error(traceback.format_exc()) - raise + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + model_adapter = self.resolve_model_adapter_name(adapter_name) + datum_list = body.forward_input.data + loss_fn_config = body.forward_input.loss_fn_config or {} + output, loss = await self.call_backend( + self.model.tinker_forward_only, inputs=datum_list, adapter_name=model_adapter, **loss_fn_config) + return types.ForwardBackwardOutput( + loss_fn_output_type='CrossEntropyLossReturn', + loss_fn_outputs=output, + metrics={'loss:mean': loss}, + ) datum_list = body.forward_input.data input_tokens = sum(len(d.model_input.to_ints()) for d in datum_list) batch_size = len(datum_list) - return await self.schedule_task( + return await submit_training_task( + self, _do_forward, - model_id=body.model_id, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::forward', token=token, input_tokens=input_tokens, batch_size=batch_size, @@ -177,39 +114,46 @@ async def forward_backward( token = await self._on_request_start(request) async def _do_forward_backward(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - model_adapter = self.resolve_model_adapter_name(adapter_name) - datum_list = body.forward_backward_input.data - loss_fn = body.forward_backward_input.loss_fn - loss_fn_config = body.forward_backward_input.loss_fn_config or {} - output, loss = await self.call_backend( - self.model.tinker_forward_backward, - inputs=datum_list, - adapter_name=model_adapter, - loss_fn=loss_fn, - **loss_fn_config) - output_type = ('ImportanceSamplingLossReturn' - if loss_fn == 'importance_sampling' else 'CrossEntropyLossReturn') - self.set_resource_state(adapter_name, 'grad_ready', True) - return types.ForwardBackwardOutput( - loss_fn_output_type=output_type, - loss_fn_outputs=output, - metrics={'loss:mean': loss}, - ) - except Exception: - logger.error(traceback.format_exc()) - raise + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + model_adapter = self.resolve_model_adapter_name(adapter_name) + datum_list = body.forward_backward_input.data + loss_fn = body.forward_backward_input.loss_fn + loss_fn_config = body.forward_backward_input.loss_fn_config or {} + preflight = getattr(self.model, 'preflight_loss_inputs', None) + if preflight is not None and loss_fn == 'importance_sampling' and not validate_tinker_loss( + loss_fn, datum_list, data_world_size=self.data_world_size): + # GRPO datums carry ragged per-sample targets; the loss check inside + # the sliced backend would deadlock the peer ranks (BUG-05), so the + # full batch is validated on every rank first. + await self.call_backend( + preflight, inputs=datum_list, adapter_name=model_adapter, loss_fn=loss_fn, **loss_fn_config) + output, loss = await self.call_backend( + self.model.tinker_forward_backward, + inputs=datum_list, + adapter_name=model_adapter, + loss_fn=loss_fn, + **loss_fn_config) + output_type = ('ImportanceSamplingLossReturn' + if loss_fn == 'importance_sampling' else 'CrossEntropyLossReturn') + self.get_training_state(adapter_name).grad_ready = True + return types.ForwardBackwardOutput( + loss_fn_output_type=output_type, + loss_fn_outputs=output, + metrics={'loss:mean': loss}, + ) datum_list = body.forward_backward_input.data input_tokens = sum(len(d.model_input.to_ints()) for d in datum_list) batch_size = len(datum_list) loss_fn = body.forward_backward_input.loss_fn is_dpo = validate_tinker_loss(loss_fn, datum_list, data_world_size=self.data_world_size) - return await self.schedule_task( + return await submit_training_task( + self, _do_forward_backward, - model_id=body.model_id, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::forward_backward', token=token, input_tokens=input_tokens, batch_size=batch_size, @@ -227,24 +171,26 @@ async def optim_step( token = await self._on_request_start(request) async def _do_optim(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - model_adapter = self.resolve_model_adapter_name(adapter_name) - if not self.get_resource_state(adapter_name, 'grad_ready', False): - raise RuntimeError(f'No accumulated gradients for adapter={adapter_name}; ' - 'call forward_backward before optim_step') - await self.call_backend( - self.model.tinker_step, adam_params=body.adam_params, adapter_name=model_adapter) - self.set_resource_state(adapter_name, 'grad_ready', False) - metrics = await self.call_backend( - self.model.tinker_calculate_metric, is_training=True, adapter_name=model_adapter) - return types.OptimStepResponse(metrics=metrics) - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_do_optim, model_id=body.model_id, token=token, task_type='optim_step') + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + model_adapter = self.resolve_model_adapter_name(adapter_name) + if not self.get_training_state(adapter_name).grad_ready: + raise RuntimeError(f'No accumulated gradients for adapter={adapter_name}; ' + 'call forward_backward before optim_step') + await self.call_backend(self.model.tinker_step, adam_params=body.adam_params, adapter_name=model_adapter) + self.get_training_state(adapter_name).grad_ready = False + metrics = await self.call_backend( + self.model.tinker_calculate_metric, is_training=True, adapter_name=model_adapter) + return types.OptimStepResponse(metrics=metrics) + + return await submit_training_task( + self, + _do_optim, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::optim_step', + token=token, + task_type='optim_step') @app.post('/tinker/save_weights') async def save_weights( @@ -255,22 +201,25 @@ async def save_weights( token = await self._on_request_start(request) async def _do_save(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') - save_weights = partial( - self.call_backend, - self.model.save, - adapter_name=self.resolve_model_adapter_name(adapter_name), - save_optimizer=True) - tinker_path, _ = await checkpoint_manager.save_training(body.model_id, body.path, save_weights) - return types.SaveWeightsResponse(path=tinker_path, type='save_weights') - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_do_save, model_id=body.model_id, token=token, task_type='save_weights') + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + save_weights = partial( + self.call_backend, + self.model.save, + adapter_name=self.resolve_model_adapter_name(adapter_name), + save_optimizer=True) + tinker_path, _ = await checkpoint_manager.save_training(body.model_id, body.path, save_weights) + return types.SaveWeightsResponse(path=tinker_path, type='save_weights') + + return await submit_training_task( + self, + _do_save, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::save_weights', + token=token, + task_type='save_weights') @app.post('/tinker/save_weights_for_sampler') async def save_weights_for_sampler( @@ -281,39 +230,40 @@ async def save_weights_for_sampler( token = await self._on_request_start(request) async def _do_save_for_sampler(): - try: - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') - - save_weights = partial( - self.call_backend, - self.model.save, - adapter_name=self.resolve_model_adapter_name(adapter_name), - save_optimizer=False) - - tinker_path, _ = await checkpoint_manager.save_sampler(body.model_id, body.path, save_weights) - payload = body.model_dump() - payload['model_path'] = tinker_path - metadata = await self.state.get_model_metadata(body.model_id) or {} - if metadata.get('base_model'): - payload['base_model'] = metadata['base_model'] - sampling_session_id = await self.state.create_sampling_session(payload) - # Tinker SDK distinguishes two modes by whether sampling_session_seq_id is set: - # 1. save_weights_for_sampler(name): expects path != None - # 2. save_weights_and_get_sampling_client(): expects path == None, sampling_session_id != None - if body.sampling_session_seq_id is not None: - return types.SaveWeightsForSamplerResponseInternal( - path=None, sampling_session_id=sampling_session_id) - else: - return types.SaveWeightsForSamplerResponseInternal( - path=tinker_path, sampling_session_id=sampling_session_id) - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task( - _do_save_for_sampler, model_id=body.model_id, token=token, task_type='save_weights_for_sampler') + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + + save_weights = partial( + self.call_backend, + self.model.save, + adapter_name=self.resolve_model_adapter_name(adapter_name), + save_optimizer=False) + + tinker_path, _ = await checkpoint_manager.save_sampler(body.model_id, body.path, save_weights) + payload = body.model_dump() + payload['model_path'] = tinker_path + metadata = await self.state.get_model_metadata(body.model_id) or {} + if metadata.get('base_model'): + payload['base_model'] = metadata['base_model'] + sampling_session_id = await self.state.create_sampling_session(payload) + # Tinker SDK distinguishes two modes by whether sampling_session_seq_id is set: + # 1. save_weights_for_sampler(name): expects path != None + # 2. save_weights_and_get_sampling_client(): expects path == None, sampling_session_id != None + if body.sampling_session_seq_id is not None: + return types.SaveWeightsForSamplerResponseInternal(path=None, sampling_session_id=sampling_session_id) + else: + return types.SaveWeightsForSamplerResponseInternal( + path=tinker_path, sampling_session_id=sampling_session_id) + + return await submit_training_task( + self, + _do_save_for_sampler, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::save_weights_for_sampler', + token=token, + task_type='save_weights_for_sampler') @app.post('/tinker/load_weights') async def load_weights( @@ -324,25 +274,28 @@ async def load_weights( token = await self._on_request_start(request) async def _do_load(): - try: - assert self.model is not None, 'Model not loaded, please load model first' - adapter_name = self.get_adapter_name(adapter_name=body.model_id) - self.assert_resource_exists(adapter_name) - # Path resolution (token isolation, twinkle-vs-external path shapes) is a - # server storage policy and belongs in the handler, not in the GPU-side - # backend actor. The backend receives an already-resolved location. - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') - resolved = checkpoint_manager.resolve_load_path(body.path) - await self.call_backend( - self.model.tinker_load, - checkpoint_name=resolved.checkpoint_name, - output_dir=resolved.checkpoint_dir if resolved.is_twinkle_path else None, - load_optimizer=body.optimizer, - adapter_name=self.resolve_model_adapter_name(adapter_name)) - self.set_resource_state(adapter_name, 'grad_ready', False) - return types.LoadWeightsResponse(path=body.path, type='load_weights') - except Exception: - logger.error(traceback.format_exc()) - raise - - return await self.schedule_task(_do_load, model_id=body.model_id, token=token, task_type='load_weights') + assert self.model is not None, 'Model not loaded, please load model first' + adapter_name = body.model_id + self.assert_resource_exists(adapter_name) + # Path resolution (token isolation, twinkle-vs-external path shapes) is a + # server storage policy and belongs in the handler, not in the GPU-side + # backend actor. The backend receives an already-resolved location. + checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + resolved = checkpoint_manager.resolve_load_path(body.path) + await self.call_backend( + self.model.tinker_load, + checkpoint_name=resolved.checkpoint_name, + output_dir=resolved.checkpoint_dir if resolved.is_twinkle_path else None, + load_optimizer=body.optimizer, + adapter_name=self.resolve_model_adapter_name(adapter_name)) + self.get_training_state(adapter_name).grad_ready = False + return types.LoadWeightsResponse(path=body.path, type='load_weights') + + return await submit_training_task( + self, + _do_load, + training_id=body.model_id, + seq_id=body.seq_id, + stream='tinker::load_weights', + token=token, + task_type='load_weights') diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index 292960b0f..e2d2a707f 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -22,12 +22,9 @@ import twinkle.protocol.types as types from twinkle.protocol.serialize import deserialize_object -from twinkle.server.checkpoint import (_resolve_client_save_dir, create_checkpoint_manager, create_training_run_manager, - validate_user_path) -from twinkle.server.exceptions import RequestRejectedError, TrainModeMismatchError -from twinkle.server.lifecycle.submit import (backend_kwargs, input_metrics, resolve_twinkle_adapter_name, run_submit, - to_backend_inputs) -from twinkle.server.middleware.auth import get_session_id_from_request +from twinkle.server.checkpoint import create_checkpoint_manager +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.lifecycle.submit import backend_kwargs, input_metrics, run_submit, to_backend_inputs from twinkle.server.model.data_plane_inputs import (data_plane_request_shape, merge_forward_kwargs, resolve_data_plane_model_inputs, select_output_rows) from twinkle.server.validation import BackendCapability @@ -74,8 +71,7 @@ async def model_healthz( @app.post('/twinkle/create', response_model=types.CreateResponse) async def create(request: Request, body: types.CreateRequest, self: ModelManagement = Depends(self_fn)) -> types.CreateResponse: - await self._on_request_start(request) - return types.CreateResponse() + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') # ------------------------------------------------------------------ # # Inline data / forward family @@ -137,11 +133,19 @@ def first_element(data): for key in inputs: if isinstance(inputs[key], list) and isinstance(first_element(inputs[key]), (int, float)): inputs[key] = torch.tensor(inputs[key]) + kwargs = backend_kwargs(body) + resolved_adapter_name = self.resolve_model_adapter_name(adapter_name) + preflight = getattr(self.model, 'preflight_loss_inputs', None) + if preflight is not None and any(key in kwargs for key in ('old_logps', 'ref_logps', 'advantages')): + # Full, unsliced batch on every rank: a loss alignment error must + # fail all ranks before any of them enters its backward collective. + await self.call_backend(preflight, inputs=all_inputs, adapter_name=resolved_adapter_name, **kwargs) + elif preflight is None: + logger.warning( + 'Backend %s has no preflight_loss_inputs; a loss alignment error can still deadlock the DP peers.', + type(self.model).__name__) ret = await self.call_backend( - self.model.forward_backward, - inputs=all_inputs, - adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) + self.model.forward_backward, inputs=all_inputs, adapter_name=resolved_adapter_name, **kwargs) return {'result': ret} return await run_submit( @@ -241,11 +245,17 @@ async def forward_backward_from_data_plane( async def _call(self, body, adapter_name, token): raw_inputs, field_kwargs = await resolve_data_plane_model_inputs(body, self.data_plane) kwargs = merge_forward_kwargs(backend_kwargs(body), field_kwargs) + inputs = to_backend_inputs(raw_inputs) + resolved_adapter_name = self.resolve_model_adapter_name(adapter_name) + preflight = getattr(self.model, 'preflight_loss_inputs', None) + if preflight is not None and any(key in kwargs for key in ('old_logps', 'ref_logps', 'advantages')): + await self.call_backend(preflight, inputs=inputs, adapter_name=resolved_adapter_name, **kwargs) + elif preflight is None: + logger.warning( + 'Backend %s has no preflight_loss_inputs; a loss alignment error can still deadlock the DP peers.', + type(self.model).__name__) ret = await self.call_backend( - self.model.forward_backward, - inputs=to_backend_inputs(raw_inputs), - adapter_name=self.resolve_model_adapter_name(adapter_name), - **kwargs) + self.model.forward_backward, inputs=inputs, adapter_name=resolved_adapter_name, **kwargs) return {'result': ret} return await run_submit( @@ -514,14 +524,7 @@ async def _call(self, body, adapter_name, token): @app.post('/twinkle/remove_adapter', response_model=types.TaskEnvelope) async def remove_adapter( request: Request, body: types.AdapterRequest, self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - """Release a drained tenant's in-memory training adapter.""" - - async def _call(self, body, adapter_name, token): - await self._cleanup_adapter(adapter_name) - return {'status': 'ok'} - - return await run_submit( - self, request, body, task_type='remove_adapter', backend_call=_call, assert_resource=False) + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') @app.post('/twinkle/add_adapter_to_model', response_model=types.TaskEnvelope) async def add_adapter_to_model( @@ -529,75 +532,7 @@ async def add_adapter_to_model( body: types.AddAdapterRequest, self: ModelManagement = Depends(self_fn), ) -> types.TaskEnvelope: - # This endpoint creates the adapter, so it cannot use the standard resource - # assertion. The Decision_Boundary left checks (train_mode 400 / full-mode - # 409) run here, before any state write, raising RequestRejectedError - # subclasses (zero future writes). - # - # Raised, not asserted: a missing adapter_name is decidable from the request body - # alone, so it owes the caller a real 400. A bare `assert` would surface as a 500 - # ('the server broke') and would vanish entirely under `python -O`, letting an - # empty adapter_name through to the backend. - if not body.adapter_name: - raise RequestRejectedError('`adapter_name` is required and must be non-empty.') - token = await self._on_request_start(request) - if not validate_user_path(token, body.adapter_name): - raise RequestRejectedError(f'Invalid adapter_name: {body.adapter_name}') - adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) - session_id = get_session_id_from_request(request) - try: - resolved_save_dir = _resolve_client_save_dir(body.save_dir).as_posix() if body.save_dir else None - except ValueError as exc: - raise RequestRejectedError(str(exc)) from exc - - config = deserialize_object(body.config) - - # ---- Decision_Boundary left: validate against the deployment's train_mode ---- - if self.is_full_mode and config is not None: - raise TrainModeMismatchError('This deployment runs in full-parameter (exclusive) mode; pass ' - 'config=None (do not send a LoraConfig).') - if (not self.is_full_mode) and config is None: - raise TrainModeMismatchError('This deployment runs in LoRA mode; a LoraConfig is required.') - if self.is_full_mode: - # Raises FullModeBusyError (409) if another tenant holds the exclusive deployment. - self.assert_full_mode_available(adapter_name) - - async def _task(): - from peft import LoraConfig - extra_kwargs = backend_kwargs(body) - training_run_manager = create_training_run_manager(token, client_type='twinkle') - lora_config = None - if isinstance(config, LoraConfig): - lora_config = types.LoraConfig(rank=config.r, train_unembed=False, train_mlp=True, train_attn=True) - run_config = types.CreateModelRequest( - base_model=self.base_model, - lora_config=lora_config, - save_dir=resolved_save_dir, - user_metadata={'adapter_name': body.adapter_name}) - await self.state.register_model( - run_config.model_dump(), - token=token, - model_id=adapter_name, - replica_id=self.replica_id, - session_id=session_id, - ) - try: - self.register_resource(adapter_name, token, session_id) - if self.is_full_mode: - # No PEFT adapter to add; the default optimizer group is used. - self.set_resource_state(adapter_name, 'grad_ready', False) - else: - await self.call_backend(self.model.add_adapter_to_model, adapter_name, config, **extra_kwargs) - # Inside the same guard: a failed run-metadata save must roll the registration - # back too, or the adapter is served while list/get report no such run. - training_run_manager.save(adapter_name, run_config) - except Exception: - self.unregister_resource(adapter_name) - await self.state.unload_model(adapter_name) - raise - return {'status': 'ok', 'adapter_name': adapter_name} - - return await self.submit_and_peek(_task, model_id=adapter_name, token=token, task_type='add_adapter_to_model') + raise RequestRejectedError('Training lifecycle operations must use the catalog endpoint') # ------------------------------------------------------------------ # # Hub upload (pure I/O -> background task; state-tracked via Retrieve_Endpoint) diff --git a/src/twinkle/server/resources/__init__.py b/src/twinkle/server/resources/__init__.py new file mode 100644 index 000000000..5fe876191 --- /dev/null +++ b/src/twinkle/server/resources/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared device-resource allocation for training, sampling, and future workloads.""" diff --git a/src/twinkle/server/resources/allocation.py b/src/twinkle/server/resources/allocation.py new file mode 100644 index 000000000..007f9cf46 --- /dev/null +++ b/src/twinkle/server/resources/allocation.py @@ -0,0 +1,144 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Single-node device allocation; only the resource owner removes placement groups.""" + +from __future__ import annotations +import asyncio +from dataclasses import dataclass + +from twinkle import Platform +from twinkle.infra._ray.resource_manager import _gpu_placement_group_cpus +from twinkle.server.exceptions import RequestRejectedError, ResourceReleaseError + + +class CapacityUnavailableError(RequestRejectedError): + error_code = 409 + + +@dataclass(frozen=True) +class DeviceAllocation: + name: str + placement_group_id: str + node_id: str + devices: tuple[str, ...] + + +def _probe_allocation(resource: str, device_type: str): + import os + + import ray + + context = ray.get_runtime_context() + if resource == "CPU": + devices = [] + else: + env = os.environ.get(Platform.get_platform(device_type).visible_device_env()) + devices = env.split(",") if env else context.get_accelerator_ids().get(resource, []) + return context.get_node_id(), [str(device) for device in devices] + + +class DevicePool: + def __init__(self, namespace: str, *, allocation_timeout: float, shutdown_timeout: float): + self.namespace = namespace + self.allocation_timeout = allocation_timeout + self.shutdown_timeout = shutdown_timeout + self._creating = {} + + def name(self, epoch: str, instance_id: str) -> str: + return f"resources-{self.namespace}-{epoch}-{instance_id}" + + @staticmethod + def _create_group(resource, count, name): + import ray + from ray.util.placement_group import placement_group + + candidates = [node for node in ray.nodes() if node.get("Alive") and node["Resources"].get(resource, 0) >= count] + if not candidates: + raise CapacityUnavailableError(f"No single-node allocation available for {count} {resource} devices") + # The aggregate free-resource snapshot can lag a confirmed group removal. + # pg.ready() is the allocation witness and is bounded by allocate(). + cpus = min(_gpu_placement_group_cpus(int(node["Resources"].get("CPU", 0)), count) for node in candidates) + bundle = {resource: count, "CPU": count if resource == "CPU" else cpus} + return placement_group([bundle], strategy="STRICT_PACK", name=name) + + async def allocate(self, num_devices: int, device_type: str, epoch: str, instance_id: str) -> DeviceAllocation: + import ray + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + + if type(num_devices) is not int or num_devices <= 0: + raise ValueError("num_devices must be a positive integer") + resource = "CPU" if device_type == "cpu" else Platform.get_platform(device_type).__name__ + count = num_devices + name = self.name(epoch, instance_id) + creating = asyncio.create_task(asyncio.to_thread(self._create_group, resource, count, name)) + self._creating[name] = creating + try: + async with asyncio.timeout(self.allocation_timeout): + pg = await asyncio.shield(creating) + await pg.ready() + options = {"num_cpus": 0.01, "scheduling_strategy": PlacementGroupSchedulingStrategy(placement_group=pg)} + if resource == "GPU": + options["num_gpus"] = count + elif resource != "CPU": + options["resources"] = {resource: count} + node, devices = await asyncio.wait_for( + ray.remote(_probe_allocation).options(**options).remote(resource, device_type), self.shutdown_timeout + ) + if resource == "CPU": + devices = [str(index) for index in range(count)] + if len(devices) != count or len(set(devices)) != count: + raise RuntimeError("Ray allocation did not return the requested unique device IDs") + return DeviceAllocation(name, pg.id.hex(), node, tuple(devices)) + except BaseException as exc: + # Cancellation cannot stop a synchronous GCS create. Keep its task + # owned until removal has observed the late result; a close timeout + # leaves the enclosing instance quarantined for a later retry. + await self.remove(name) + if isinstance(exc, TimeoutError): + raise CapacityUnavailableError( + "Single-node allocation could not be confirmed; retry a new create later." + ) from exc + raise + finally: + if creating.done() and self._creating.get(name) is creating: + self._creating.pop(name) + + async def _finish_creation(self, name): + creating = self._creating.get(name) + if creating is None: + return + # A failed create has no result to free. A pending create must stay tracked + # across cancellation, so do not remove it until this await has completed. + try: + await asyncio.shield(creating) + except Exception: + pass # A failed create has no allocation result; cancellation remains unsafe. + if self._creating.get(name) is creating: + self._creating.pop(name) + + @staticmethod + async def _find_group(name): + from ray.util.placement_group import get_placement_group + + try: + return await asyncio.to_thread(get_placement_group, name) + except ValueError: + return None # A missing named group is already released. + + async def _remove(self, name): + from ray.util.placement_group import placement_group_table, remove_placement_group + + await self._finish_creation(name) + pg = await self._find_group(name) + if pg is None: + return + await asyncio.to_thread(remove_placement_group, pg) + while (await asyncio.to_thread(placement_group_table, pg)).get("state") != "REMOVED": + await asyncio.sleep(0.05) + + async def remove(self, name: str) -> None: + try: + await asyncio.wait_for(self._remove(name), self.shutdown_timeout) + except TimeoutError as exc: + raise TimeoutError(f"Placement group removal was not confirmed: {name}") from exc + except Exception as exc: + raise ResourceReleaseError(f"Placement group removal failed: {name}", phase="pg_removal") from exc diff --git a/src/twinkle/server/resources/instances.py b/src/twinkle/server/resources/instances.py new file mode 100644 index 000000000..f732b9c82 --- /dev/null +++ b/src/twinkle/server/resources/instances.py @@ -0,0 +1,123 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared instance supervision and bounds, with feature-owned records and policies.""" + +from __future__ import annotations +import asyncio +import logging +import time + +logger = logging.getLogger(__name__) + + +class InstanceLifecycle: + """Own tasks, not a second copy of instance records or persisted status.""" + + def __init__(self, label, *, clock=time.monotonic): + self.label = label + self.clock = clock + self.unreachable_since = {} + self.jobs = set() + self.startups = {} + self.closing = {} + self.named_jobs = {} + + def spawn(self, coroutine): + task = asyncio.create_task(coroutine) + self.jobs.add(task) + + def done(completed): + self.jobs.discard(completed) + if not completed.cancelled(): + error = completed.exception() + if error is not None: + logger.error( + "%s background job failed", self.label, exc_info=(type(error), error, error.__traceback__) + ) + + task.add_done_callback(done) + return task + + def spawn_once(self, key, factory): + task = self.named_jobs.get(key) + if task is None: + task = self.spawn(factory()) + self.named_jobs[key] = task + task.add_done_callback(lambda completed: self.named_jobs.pop(key, None)) + return task + + def start(self, identity, factory): + if identity not in self.startups: + self.startups[identity] = self.spawn(factory()) + return self.startups[identity] + + async def cancel_start(self, identity): + task = self.startups.get(identity) + if task is not None and task is not asyncio.current_task() and not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + async def close(self, identity, stop, timeout, stopped, quarantined): + task = self.closing.get(identity) + if task is None: + + async def close_once(): + try: + await asyncio.wait_for(stop(), timeout) + except Exception as error: + logger.warning( + "%s instance %s shutdown is unconfirmed", + self.label, + identity, + exc_info=(type(error), error, error.__traceback__), + ) + await quarantined(error) + else: + # Only the feature can retire its users/quota and commit STOPPED. + await stopped() + self.unreachable_since.pop(identity, None) + + task = self.spawn(close_once()) + self.closing[identity] = task + + def clear(completed): + if self.closing.get(identity) is completed: + self.closing.pop(identity, None) + + task.add_done_callback(clear) + await asyncio.shield(task) + + async def health(self, identity, probe, *, unreachable_timeout=30): + """Distinguish explicit backend failure from a transient control-plane timeout.""" + try: + result = await probe() + except Exception: + since = self.unreachable_since.setdefault(identity, self.clock()) + return False if self.clock() - since >= unreachable_timeout else None + self.unreachable_since.pop(identity, None) + return not result["unhealthy"] + + @staticmethod + def owned_count(instances): + return sum(instance.state != "STOPPED" for instance in instances) + + def require_released(self, instances): + """A shutdown ACK must certify release before its owner can be removed.""" + owned = self.owned_count(instances) + if owned: + raise RuntimeError(f"{self.label} shutdown is unconfirmed: {owned} instances remain owned") + + @classmethod + def headroom(cls, instances, maximum): + return float("inf") if maximum is None else max(0, maximum - cls.owned_count(instances)) + + @classmethod + def missing(cls, instances, minimum, maximum, eligible): + instances = tuple(instances) + needed = max(0, minimum - sum(eligible(instance) for instance in instances)) + return min(needed, cls.headroom(instances, maximum)) + + async def join(self): + while self.jobs: + tasks = tuple(self.jobs) + await asyncio.gather(*tasks, return_exceptions=True) + self.jobs.difference_update(task for task in tasks if task.done()) diff --git a/src/twinkle/server/resources/runtime.py b/src/twinkle/server/resources/runtime.py new file mode 100644 index 000000000..3a1c3cd20 --- /dev/null +++ b/src/twinkle/server/resources/runtime.py @@ -0,0 +1,158 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Ray/Serve resource effects shared by training and sampling owners.""" + +from __future__ import annotations +import asyncio +import logging +import time +from copy import deepcopy + +from twinkle.server.config.resources import build_worker_env_vars +from twinkle.server.exceptions import ResourceReleaseError +from .allocation import DevicePool +from .instances import InstanceLifecycle + +logger = logging.getLogger(__name__) + + +class RayRuntime: + """Own deployment threads and certify worker/PG death before releasing capacity.""" + + def __init__(self, namespace, *, allocation_timeout, shutdown_timeout): + self.namespace = namespace + self.pool = DevicePool(namespace, allocation_timeout=allocation_timeout, shutdown_timeout=shutdown_timeout) + self.deployments = {} + self.lifecycle = InstanceLifecycle("Ray deployment cleanup") + + def deployment_options(self, options, name, env_vars=None): + options = deepcopy(options) + options.pop("autoscaling_config", None) + options["num_replicas"] = 1 + actor = options.setdefault("ray_actor_options", {}) + actor.setdefault("num_cpus", 0.1) + runtime_env = actor.setdefault("runtime_env", {}) + env = build_worker_env_vars() + env.update(runtime_env.get("env_vars", {})) + env.update(env_vars or {}) + env.update(TWINKLE_RAY_NAMESPACE=self.namespace, TWINKLE_RAY_REGISTRY_NAME=name + "-registry") + runtime_env["env_vars"] = env + return options + + async def deploy(self, identity, name, app, route_prefix, cleanup): + from ray import serve + + deployment = asyncio.create_task(asyncio.to_thread(serve.run, app, name=name, route_prefix=route_prefix)) + self.deployments[identity] = deployment + try: + await asyncio.shield(deployment) + except asyncio.CancelledError: + + async def clean_late(): + await asyncio.gather(deployment, return_exceptions=True) + await cleanup() + + self.lifecycle.spawn(clean_late()) + raise + finally: + if deployment.done(): + self.deployments.pop(identity, None) + + async def wait_for_cleanup(self): + await self.lifecycle.join() + + @staticmethod + async def handle(name): + def lookup(): + from ray import serve + + return serve.get_app_handle(name) + + # get_app_handle waits synchronously for the Serve controller. Keep + # that RPC inside the caller's cancellable asynchronous deadline. + return await asyncio.to_thread(lookup) + + async def shutdown(self, identity, name, *, observe_removal=None): + from ray import serve + from ray.util import list_named_actors + + deployment = self.deployments.get(identity) + if deployment is not None: + await asyncio.gather(asyncio.shield(deployment), return_exceptions=True) + self.deployments.pop(identity, None) + from ray._private.gcs_utils import create_gcs_channel + from ray._private.worker import global_worker + from ray.core.generated.gcs_pb2 import ActorTableData + from ray.core.generated.gcs_service_pb2 import ( + GetActorInfoRequest, + GetNamedActorInfoRequest, + KillActorViaGcsRequest, + ) + from ray.core.generated.gcs_service_pb2_grpc import ActorInfoGcsServiceStub + + # ActorHandle deserialization imports the worker's model class. Cleanup + # needs only GCS identities, so it must not load that training stack. + async with create_gcs_channel(global_worker.gcs_client.address, aio=True) as channel: + actors = ActorInfoGcsServiceStub(channel) + owned_ids = [] + for row in await asyncio.to_thread(list_named_actors, all_namespaces=True): + if row["namespace"] != self.namespace or identity not in row["name"]: + continue + reply = await actors.GetNamedActorInfo( + GetNamedActorInfoRequest(name=row["name"], ray_namespace=self.namespace), + timeout=self.pool.shutdown_timeout, + ) + actor = reply.actor_table_data + if ( + reply.status.code + or not actor.actor_id + or actor.name != row["name"] + or actor.ray_namespace != self.namespace + ): + raise ResourceReleaseError("Owned worker identity could not be confirmed", phase="worker_lookup") + owned_ids.append(actor.actor_id) + reply = await actors.KillActorViaGcs( + KillActorViaGcsRequest(actor_id=actor.actor_id, force_kill=True, no_restart=True), + timeout=self.pool.shutdown_timeout, + ) + if reply.status.code: + raise ResourceReleaseError("Owned worker termination failed", phase="worker_death") + await asyncio.to_thread(serve.delete, name) + for actor_id in owned_ids: + try: + async with asyncio.timeout(5): + while True: + reply = await actors.GetActorInfo(GetActorInfoRequest(actor_id=actor_id), timeout=5) + if reply.status.code or reply.actor_table_data.actor_id != actor_id: + raise ResourceReleaseError( + "Owned worker death could not be confirmed", phase="worker_death" + ) + if reply.actor_table_data.state == ActorTableData.DEAD: + break + await asyncio.sleep(0.05) + except Exception as exc: + raise ResourceReleaseError( + "Owned worker death could not be confirmed", phase="worker_death" + ) from exc + + started, status = time.monotonic(), "error" + try: + await self.pool.remove(name) + status = "success" + finally: + if observe_removal is not None: + observe_removal(time.monotonic() - started, status) + + async def recover_orphans(self, decode, shutdown): + """Decode only this owner's names; never sweep another service's resources.""" + from ray import serve + from ray.util.placement_group import placement_group_table + + applications, groups = await asyncio.gather( + asyncio.to_thread(serve.status), asyncio.to_thread(placement_group_table) + ) + names = set(applications.applications) + names.update(row.get("name", "") for row in groups.values() if row.get("state") != "REMOVED") + for name in sorted(names): + instance = decode(name) + if instance is not None: + await asyncio.wait_for(shutdown(instance), self.pool.shutdown_timeout) diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampler/app.py index b7b33f78a..650254c2b 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -7,21 +7,26 @@ """ from __future__ import annotations +import asyncio +import os from fastapi import FastAPI, Request from ray import serve from typing import Any from twinkle import DeviceGroup from twinkle.server.config.backend_dispatch import BackendSelector +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.deployment import LazyCleanupMixin, bind_deployment, build_deployment_app +from twinkle.server.exceptions import RequestRejectedError from twinkle.server.middleware.auth import get_token_from_request from twinkle.server.runtime import init_twinkle_runtime from twinkle.server.state import ServerState, get_server_state -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.state.models import RuntimeContext from twinkle.server.task_queue.mixin import TaskQueueMixin from twinkle.server.utils import get_template_for_model, wrap_builder_with_device_group_env from twinkle.utils.logger import get_logger from .concurrency import SamplerConcurrency +from .runtime import SamplerRuntime from .tinker_handlers import _register_tinker_sampler_routes from .twinkle_handlers import _register_twinkle_sampler_routes @@ -112,7 +117,18 @@ async def __init__(self, data_plane_url: str | None = None, template: dict[str, Any] | None = None, **kwargs): + self.request_runtime = SamplerRuntime(self) self.device_group = DeviceGroup(**device_group) + context = kwargs.pop('sampling_context', None) + if context is not None: + context = RuntimeContext.model_validate(context) + from twinkle.infra._ray.ray_helper import RayHelper + from twinkle.infra._ray.resource_manager import ResourceManager + os.environ['TWINKLE_RAY_NAMESPACE'] = context.namespace + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = context.registry_name + RayHelper.resource_manager = ResourceManager.from_allocation(self.device_group, + kwargs.pop('placement_group'), + kwargs.pop('allocation_devices')) self.device_mesh = init_twinkle_runtime( is_mock=(sampler_type == 'mock'), nproc_per_node=nproc_per_node, @@ -121,8 +137,10 @@ async def __init__(self, ) self.sampler_type = sampler_type self.model_id = model_id - replica_context = serve.get_replica_context() - self.replica_id = replica_context.replica_id.unique_id + if context is not None: + self.replica_id = context.instance_id + else: + self.replica_id = serve.get_replica_context().replica_id.unique_id sampler_kwargs: dict[str, Any] = { 'model_id': model_id, @@ -141,6 +159,8 @@ async def __init__(self, else: sampler_kwargs.update(kwargs) self.sampler = _construct_sampler_backend(sampler_type, sampler_kwargs, data_plane_url) + actors = getattr(self.sampler, '_actors', ()) + await asyncio.gather(*(actor.__ray_ready__.remote() for actor in actors)) self.state: ServerState = get_server_state() from twinkle.server.data_plane import DataPlaneProxy @@ -187,21 +207,28 @@ async def shutdown(self) -> None: finally: await self.shutdown_task_queue() await self.sampling_guard.wait_closed() + await self.request_runtime.shutdown() await self.data_plane.close() - async def check_health(self) -> None: - self.sampling_guard.check_available() + async def invoke(self, ticket, payload, stream_queue=None): + return await self.request_runtime.invoke(ticket, payload, stream_queue) + + async def wait_operation(self, ticket): + await self.request_runtime.wait_operation(ticket) + + async def install_base_weights(self, path): + await self.request_runtime.install_base_weights(path) + + async def pending_futures(self): + return await self.request_runtime.pending_futures() - @serve.multiplexed(max_num_models_per_replica=5) - async def _sticky_entry(self, sticky_key: str): - return sticky_key + async def runtime_identity(self): + return await self.request_runtime.runtime_identity() - async def _ensure_sticky(self): - sticky_key = serve.get_multiplexed_model_id() - await self._sticky_entry(sticky_key) + async def check_health(self) -> None: + self.sampling_guard.check_available() async def _on_request_start(self, request: Request) -> str: - await self._ensure_sticky() await self.state.touch_replica_last_seen(self.replica_id) await self._ensure_state_cleanup_started() return get_token_from_request(request) @@ -250,9 +277,20 @@ def build_sampler_app(model_id: str, # Build the FastAPI app + middleware stack + routes via the shared scaffold, # then bind the Ray Serve deployment. The Sampler passes its FastAPI # title/description/version through ``fastapi_kwargs``. + managed = kwargs.get('sampling_context') is not None + def register_routes(app: FastAPI, get_self: Any) -> None: - _register_tinker_sampler_routes(app, get_self) - _register_twinkle_sampler_routes(app, get_self) + if managed: + + def managed_self(): + # Internal compute URLs must not bypass the controller's tickets. + raise RequestRejectedError('Sampling requests must use the stable sampler endpoint') + + dependency = managed_self + else: + dependency = get_self + _register_tinker_sampler_routes(app, dependency) + _register_twinkle_sampler_routes(app, dependency) app = build_deployment_app( 'Sampler', diff --git a/src/twinkle/server/sampler/backends/__init__.py b/src/twinkle/server/sampler/backends/__init__.py index c2caf8e9e..a734d4690 100644 --- a/src/twinkle/server/sampler/backends/__init__.py +++ b/src/twinkle/server/sampler/backends/__init__.py @@ -1,11 +1,2 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Sampler backend implementations. - -The cross-process streaming bridge lives in ``streaming`` and is re-exported here -so a sibling backend module can import it from ``.streaming`` without importing -this package root, while external consumers keep using -``from twinkle.server.sampler.backends import stream_to_queue``. -""" -from .streaming import STREAM_SENTINEL, stream_to_queue - -__all__ = ['STREAM_SENTINEL', 'stream_to_queue'] +"""Sampling backend implementations; helpers are imported from their defining modules.""" diff --git a/src/twinkle/server/sampler/backends/streaming.py b/src/twinkle/server/sampler/backends/streaming.py index dc61ff7e7..5d07e1b59 100644 --- a/src/twinkle/server/sampler/backends/streaming.py +++ b/src/twinkle/server/sampler/backends/streaming.py @@ -1,26 +1,44 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Cross-process streaming bridge for sampler backends. +"""Lightweight queue bridge shared by real and mock sampling backends. -Lives in its own module (rather than the ``backends`` package ``__init__``) so -that a backend implementation (e.g. ``mock_sampler``) can import it without -importing its own package root — which would be a package-initialisation-order -dependency. ``backends/__init__`` re-exports these names, so external consumers -(``from twinkle.server.sampler.backends import stream_to_queue``) are unchanged. +Kept independent of the HTTP runtime so backend workers do not import FastAPI +or depend on the mock backend. Both consumers import this defining module. """ from __future__ import annotations STREAM_SENTINEL = '__STREAM_END__' +class StreamPublisher: + """A disconnected consumer cannot end the backend generation lifetime.""" + + def __init__(self, queue): + self.queue = queue + self.connected = True + + def put(self, item): + if not self.connected: + return + try: + self.queue.put(item) + except Exception: + # Ray queue transport can disappear after the HTTP reader closes. + self.connected = False + + def stream_to_queue(sampler, queue, inputs, sampling_params=None, adapter_name='', adapter_path=None): - """Push streaming deltas from *sampler* to a cross-process Ray queue. - - Works with any object that exposes a ``sample_stream`` iterator. - """ - try: - for delta, reason in sampler.sample_stream(inputs, sampling_params, adapter_name, adapter_path): - queue.put((delta, reason)) - except Exception as e: - queue.put(e) - finally: - queue.put(STREAM_SENTINEL) + """Drain and close the generator even when its HTTP consumer disconnects.""" + from contextlib import ExitStack + + publisher = StreamPublisher(queue) + generation = sampler.sample_stream(inputs, sampling_params, adapter_name, adapter_path) + with ExitStack() as cleanup: + cleanup.callback(publisher.put, STREAM_SENTINEL) + close = getattr(generation, 'close', None) + if close is not None: + cleanup.callback(close) + try: + for delta, reason in generation: + publisher.put((delta, reason)) + except Exception as error: + publisher.put(error) diff --git a/src/twinkle/server/sampler/concurrency.py b/src/twinkle/server/sampler/concurrency.py index 71d7e7c68..0acb32232 100644 --- a/src/twinkle/server/sampler/concurrency.py +++ b/src/twinkle/server/sampler/concurrency.py @@ -4,14 +4,14 @@ import asyncio -from twinkle.server.task_queue.types import BackendBusyError +from twinkle.server.exceptions import BackendBusyError class SamplerConcurrency: """Readers share stable weights; mutations wait for every reader to finish. Jobs retain their lock if an HTTP waiter leaves. Admission and execution - timeouts belong to ComputeWorker, not this lock. + timeouts belong to the caller; the pool may also bound lock admission. """ def __init__(self): @@ -35,13 +35,17 @@ def invalidate_weights(self, error: Exception): def close(self): self._closed = True - async def _execute(self, factory, exclusive: bool): + async def _execute(self, factory, exclusive: bool, queue_timeout): async with self._condition: if exclusive: self._waiting_writers += 1 try: - await self._condition.wait_for(lambda: not self._writer and - (self._readers == 0 if exclusive else self._waiting_writers == 0)) + + def available(): + return not self._writer and (self._readers == 0 if exclusive else self._waiting_writers == 0) + + if not available(): + await asyncio.wait_for(self._condition.wait_for(available), queue_timeout) self.check_available() if exclusive: self._writer = True @@ -61,7 +65,7 @@ async def _execute(self, factory, exclusive: bool): self._readers -= 1 self._condition.notify_all() - async def run(self, factory, *, exclusive: bool = False): + async def run(self, factory, *, exclusive: bool = False, queue_timeout: float | None = None): self.check_available() started = False @@ -70,7 +74,7 @@ async def call(): started = True return await factory() - job = asyncio.create_task(self._execute(call, exclusive)) + job = asyncio.create_task(self._execute(call, exclusive, queue_timeout)) self._jobs.add(job) def finished(task): diff --git a/src/twinkle/server/sampler/pool/__init__.py b/src/twinkle/server/sampler/pool/__init__.py new file mode 100644 index 000000000..78ee5693d --- /dev/null +++ b/src/twinkle/server/sampler/pool/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sampling instance pool: routing, lifecycle decisions and Ray composition.""" diff --git a/src/twinkle/server/sampler/pool/actor.py b/src/twinkle/server/sampler/pool/actor.py new file mode 100644 index 000000000..3aa7027de --- /dev/null +++ b/src/twinkle/server/sampler/pool/actor.py @@ -0,0 +1,53 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Ray composition for the sampling controller and its owned driver.""" + +import asyncio +from copy import deepcopy + +from twinkle.server.config.application_spec import SamplerArgs +from .controller import SamplerController + + +def create_controller_actor(): + import ray + + from .driver import SamplerDriver + + @ray.remote(num_cpus=0.1, max_restarts=-1, max_task_retries=0) + class SamplerControllerActor(SamplerController): + def __init__(self, values, namespace, pool_id, options, budgets): + args = SamplerArgs.model_validate(values) + super().__init__(args, SamplerDriver(args, namespace, pool_id, options, budgets), budgets) + + async def initialize(self): + await super().initialize() + if not self._stopping and self._monitor_task is None: + self._monitor_task = asyncio.create_task(self._monitor()) + + async def shutdown(self): + await super().shutdown() + await self.driver.runtime.wait_for_cleanup() + + return SamplerControllerActor + + +def create_sampler_controller(spec, namespace, budgets, env): + """Driver-only composition. Each compute instance owns its external allocation.""" + from .controller import sampler_controller_name, sampling_pool_id + + options = deepcopy(spec.deployments[0]) if spec.deployments else {} + options.pop("name", None) + options.pop("autoscaling_config", None) + options["num_replicas"] = 1 + actor_options = options.setdefault("ray_actor_options", {}) + actor_options.setdefault("num_cpus", 0.1) + runtime_env = actor_options.setdefault("runtime_env", {}) + runtime_env["env_vars"] = {**env, **runtime_env.get("env_vars", {}), "TWINKLE_STATE_RAY_NAMESPACE": namespace} + pool_id = sampling_pool_id(spec.route_prefix) + name = sampler_controller_name(namespace, pool_id) + controller = ( + create_controller_actor() + .options(name=name, namespace=namespace, runtime_env=runtime_env) + .remote(spec.args.model_dump(), namespace, pool_id, options, budgets) + ) + return controller, options, pool_id diff --git a/src/twinkle/server/sampler/pool/controller.py b/src/twinkle/server/sampler/pool/controller.py new file mode 100644 index 000000000..9e8a9c95a --- /dev/null +++ b/src/twinkle/server/sampler/pool/controller.py @@ -0,0 +1,379 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sampling pool decisions; generation ownership outlives HTTP and future timeouts.""" + +from __future__ import annotations +import asyncio +import hashlib +import logging +import math +import time +import uuid +from copy import deepcopy +from dataclasses import asdict, dataclass, field +from typing import TYPE_CHECKING, Literal + +import json + +if TYPE_CHECKING: + from twinkle.server.lifecycle.proxy import ForwardedReply, ForwardedRequest + +from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.config.resources import ResourceOptions +from twinkle.server.exceptions import ResourceExhaustedError +from twinkle.server.resources.instances import InstanceLifecycle +from twinkle.server.sampler.concurrency import SamplerConcurrency + +logger = logging.getLogger(__name__) +_CONTROL_ENDPOINTS = frozenset( + {"twinkle/set_template", "twinkle/add_adapter_to_sampler", "twinkle/apply_patch", "twinkle/unload_adapter_paths"} +) + + +def sampler_controller_name(namespace, pool_id): + return f"sampler-controller-{namespace}-{pool_id}" + + +def sampling_pool_id(route): + return hashlib.sha256(route.encode()).hexdigest()[:16] + + +@dataclass +class SamplingInstance: + identity: str + epoch: str + state: Literal["STARTING", "READY", "DRAINING", "QUARANTINED", "STOPPED"] = "STARTING" + application: str | None = None + devices: list[str] = field(default_factory=list) + requests: set[str] = field(default_factory=set) + last_finished: float = 0 + failure: str | None = None + + +class SamplerController: + def __init__(self, args: SamplerArgs, driver, budgets: dict, *, clock=time.monotonic): + self.args, self.driver, self.clock = args, driver, clock + self.budgets = ResourceOptions.model_validate(budgets) + self.epoch = uuid.uuid4().hex + self.instances = {} + self.lifecycle = InstanceLifecycle("Sampling") + self._actions = [] + self._base_path = None + self._lock = asyncio.Lock() + self._config_lock = asyncio.Lock() + self._guard = SamplerConcurrency() + self._ready = False + self._stopping = False + self._initialization = None + self._monitor_task = None + + def _owned_count(self): + return self.lifecycle.owned_count(self.instances.values()) + + def _ready_instances(self): + return [instance for instance in self.instances.values() if instance.state == "READY"] + + def _new_instance(self): + if not self.lifecycle.headroom(self.instances.values(), self.args.max_replicas): + raise ResourceExhaustedError("Sampling replica limit reached") + instance = SamplingInstance("sampler_" + uuid.uuid4().hex, self.epoch, last_finished=self.clock()) + self.instances[instance.identity] = instance + self.lifecycle.start(instance.identity, lambda: self._start(instance)) + return instance + + async def initialize(self): + if self._initialization is None: + self._initialization = asyncio.create_task(self._initialize()) + await asyncio.shield(self._initialization) + + async def _initialize(self): + await self.driver.recover_orphans() + try: + for _ in range(self.args.min_replicas): + instance = self._new_instance() + await asyncio.shield(self.lifecycle.startups[instance.identity]) + self._ready = True + except Exception: + self._stopping = True + await asyncio.gather(*(self._close(instance) for instance in list(self.instances.values()))) + raise + + async def _start(self, instance): + async def startup(): + await self.driver.start(instance) + async with self._config_lock: + if self._stopping or instance.state != "STARTING": + raise ResourceExhaustedError("Sampling startup no longer owns this instance") + for action in self._actions: + if "base_path" in action: + await self.driver.load_weights(instance, action["base_path"]) + else: + reply = await self._invoke_finished(instance, action) + if reply["status"] >= 400: + raise RuntimeError("Sampler control replay failed") + instance.state = "READY" + instance.last_finished = self.clock() + + try: + # Replaying control state is part of readiness, inside the same budget. + await asyncio.wait_for(startup(), self.budgets.startup_timeout_s) + except BaseException: + await self._close(instance) + raise + + async def _select(self): + while True: + async with self._lock: + if self._stopping or not self._ready: + raise ResourceExhaustedError("Sampling pool is not ready") + ready = self._ready_instances() + active = sum(len(instance.requests) for instance in self.instances.values()) + desired = min( + self.args.max_replicas, + max(self.args.min_replicas, math.ceil((active + 1) / self.args.target_ongoing_requests)), + ) + if self._owned_count() < desired: + self._new_instance() + if ready: + instance = min(ready, key=lambda item: (len(item.requests), item.identity)) + ticket = uuid.uuid4().hex + instance.requests.add(ticket) + return instance, ticket + pending = [task for task in self.lifecycle.startups.values() if not task.done()] + if not pending: + raise ResourceExhaustedError("No available sampler instance; device allocation was not confirmed") + await asyncio.wait_for(asyncio.shield(pending[0]), self.budgets.startup_timeout_s) + + async def _invoke_finished(self, instance, payload): + ticket = uuid.uuid4().hex + reply = await self.driver.invoke(instance, ticket, payload) + await self.driver.wait_operation(instance, ticket) + return reply + + @staticmethod + def _action_key(action): + if "base_path" in action: + return ("base_path",) + endpoint = action["endpoint"] + if endpoint == "twinkle/set_template": + return (endpoint,) + if endpoint == "twinkle/add_adapter_to_sampler": + # Adapter names are scoped by the caller's token in the handler. + headers = {key.lower(): value for key, value in action["headers"]} + owner = headers.get(b"twinkle-authorization") or headers.get(b"authorization") + return endpoint, owner, json.loads(action["body"])["adapter_name"] + return None + + def _remember(self, action): + """Compact overwritten state only within the segment after the last opaque patch.""" + key = self._action_key(action) + if key is not None: + for index in range(len(self._actions) - 1, -1, -1): + previous = self._actions[index] + if previous.get("endpoint") == "twinkle/apply_patch": + break + if self._action_key(previous) == key: + del self._actions[index] + self._actions.append(deepcopy(action)) + + async def _release(self, instance, ticket): + async with self._lock: + # Unknown completion retains ownership until shutdown is confirmed. + if instance.state not in ("DRAINING", "QUARANTINED"): + instance.requests.discard(ticket) + instance.last_finished = self.clock() + + async def _fence(self, replicas): + async with self._lock: + for replica in replicas: + if replica.state != "STOPPED": + replica.state = "DRAINING" + await asyncio.gather(*(self._close(replica) for replica in replicas)) + + async def _mutate(self, payload=None, full_path=None): + # Lock order: config -> state. Never wait for startup while holding config, + # because startup also replays the committed configuration under that lock. + seed, seed_ticket = await self._select() + replicas = [] + ticket = uuid.uuid4().hex + try: + async with self._config_lock: + async with self._lock: + if self._stopping: + raise ResourceExhaustedError("Sampling pool is stopping") + replicas = self._ready_instances() + if not replicas: + raise ResourceExhaustedError("No ready sampler for control update") + for replica in replicas: + replica.requests.add(ticket) + + async def broadcast(): + response = None + for replica in replicas: + if full_path is not None: + await self.driver.load_weights(replica, full_path) + else: + response = await self._invoke_finished(replica, payload) + if response["status"] >= 400: + return response + return response + + try: + response = await asyncio.wait_for(broadcast(), self.args.queue_config.effective_execution_timeout) + if response is not None and response["status"] >= 400: + await self._fence(replicas) + return response + # A health probe or shutdown can invalidate a participant while + # its RPC is pending. Such an ACK must not commit pool state. + if any(replica.state != "READY" for replica in replicas): + raise ResourceExhaustedError("Sampler changed state during control update") + if full_path is not None: + self._base_path = full_path + self._remember({"base_path": full_path}) + elif payload["endpoint"] != "twinkle/unload_adapter_paths": + self._remember(payload) + return response + except BaseException: + await self._fence(replicas) + raise + finally: + for replica in replicas: + await self._release(replica, ticket) + await self._release(seed, seed_ticket) + + async def request(self, payload: ForwardedRequest, stream_queue=None, full_path=None) -> ForwardedReply: + await self.initialize() + endpoint = payload["endpoint"] + if endpoint == "twinkle/create": + return {"status": 200, "headers": [(b"content-type", b"application/json")], "body": b'{"status":"ok"}'} + control = endpoint in _CONTROL_ENDPOINTS + reply_ready = asyncio.get_running_loop().create_future() + + async def operation(): + instance = None + ticket = None + try: + if control: + reply_ready.set_result(await self._mutate(payload=payload)) + return + if full_path is not None and full_path != self._base_path: + await self._mutate(full_path=full_path) + instance, ticket = await self._select() + response = await self.driver.invoke(instance, ticket, payload, stream_queue) + if not reply_ready.done(): + reply_ready.set_result(response) + # A pending future reply is not evidence that the generation stopped. + await asyncio.wait_for( + self.driver.wait_operation(instance, ticket), self.args.queue_config.absolute_future_ttl(1) + ) + except BaseException as error: + if not reply_ready.done(): + reply_ready.set_exception(error) + if instance is not None: + await self._close(instance) + raise + finally: + if instance is not None and ticket is not None: + await self._release(instance, ticket) + + async def guarded(): + try: + await self._guard.run( + operation, + exclusive=control or full_path is not None, + queue_timeout=self.args.queue_config.queue_timeout, + ) + except BaseException as error: + if not reply_ready.done(): + reply_ready.set_exception(error) + raise + + self.lifecycle.spawn(guarded()) + return await asyncio.shield(reply_ready) + + async def _close(self, instance): + if instance.state == "STOPPED": + return + instance.state = "DRAINING" + + async def quarantined(error): + instance.state = "QUARANTINED" + instance.failure = str(error) + + async def stopped(): + instance.state = "STOPPED" + instance.requests.clear() + + await self.lifecycle.close( + instance.identity, + lambda: self.driver.shutdown(instance), + self.budgets.shutdown_timeout_s, + stopped, + quarantined, + ) + + async def sweep(self): + async with self._lock: + ready = sorted(self._ready_instances(), key=lambda instance: instance.last_finished) + removable = max(0, len(ready) - self.args.min_replicas) + victims = [ + instance + for instance in ready + if not instance.requests and self.clock() - instance.last_finished >= self.args.idle_timeout_s + ][:removable] + for instance in victims: + instance.state = "DRAINING" + missing = self.lifecycle.missing( + self.instances.values(), + self.args.min_replicas, + self.args.max_replicas, + lambda instance: instance.state in ("READY", "STARTING"), + ) + if not self._stopping: + for _ in range(missing): + self._new_instance() + for instance in list(self.instances.values()): + if instance.state in ("DRAINING", "QUARANTINED"): + self.lifecycle.spawn(self._close(instance)) + elif instance.state == "STOPPED" and not instance.requests: + self.instances.pop(instance.identity, None) + self.lifecycle.startups.pop(instance.identity, None) + + async def _monitor(self): + async def health(instance): + healthy = await self.lifecycle.health(instance.identity, lambda: self.driver.health(instance)) + if healthy is False: + await self._close(instance) + + while not self._stopping: + try: + for instance in self._ready_instances(): + self.lifecycle.spawn_once(("health", instance.identity), lambda item=instance: health(item)) + await self.sweep() + except Exception: + logger.exception("Sampler monitor failed; ownership remains reserved") + await asyncio.sleep(self.budgets.sweep_interval_s) + + async def diagnostics(self): + await self.initialize() + return { + "ready": self._ready and not self._stopping and len(self._ready_instances()) >= self.args.min_replicas, + "epoch": self.epoch, + "instances": [ + {**asdict(instance), "requests": len(instance.requests)} for instance in self.instances.values() + ], + } + + async def shutdown(self): + self._stopping = True + self._ready = False + if self._monitor_task is not None: + self._monitor_task.cancel() + await asyncio.gather(self._monitor_task, return_exceptions=True) + for task in list(self.lifecycle.startups.values()): + if not task.done(): + task.cancel() + await asyncio.gather(*tuple(self.lifecycle.startups.values()), return_exceptions=True) + await asyncio.gather(*(self._close(instance) for instance in list(self.instances.values()))) + self._guard.close() + await self.lifecycle.join() + self.lifecycle.require_released(self.instances.values()) diff --git a/src/twinkle/server/sampler/pool/driver.py b/src/twinkle/server/sampler/pool/driver.py new file mode 100644 index 000000000..e71393ac7 --- /dev/null +++ b/src/twinkle/server/sampler/pool/driver.py @@ -0,0 +1,121 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Owned Ray/Serve effects for sampling instances; no training lifecycle dependency.""" + +from __future__ import annotations +import asyncio + +from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.config.resources import ResourceOptions +from twinkle.server.resources.runtime import RayRuntime +from twinkle.server.state import get_server_state +from twinkle.server.state.models import FutureFailureRecord, RuntimeContext + + +class SamplerDriver: + def __init__(self, args: SamplerArgs, namespace: str, pool_id: str, options: dict, budgets: dict): + self.args, self.namespace, self.pool_id = args, namespace, pool_id + self.options = options + self.budgets = ResourceOptions.model_validate(budgets) + self.runtime = RayRuntime( + namespace, + allocation_timeout=self.budgets.allocation_timeout_s, + shutdown_timeout=self.budgets.shutdown_timeout_s, + ) + + @property + def prefix(self): + return f"resources-{self.namespace}-sampling_{self.pool_id}_" + + async def start(self, instance): + from ray.util.placement_group import get_placement_group + + from twinkle.server.sampler.app import build_sampler_app + + group = self.args.device_group + device_type = str(group.get("device_type", "cpu" if self.args.sampler_type == "mock" else "cuda")).lower() + allocation = await self.runtime.pool.allocate( + group.get("ranks", self.args.nproc_per_node), + device_type, + f"sampling_{self.pool_id}_{instance.epoch}", + instance.identity, + ) + instance.application = allocation.name + instance.devices = list(allocation.devices) + options = self.runtime.deployment_options(self.options, allocation.name) + values = self.args.model_dump( + exclude={"min_replicas", "max_replicas", "idle_timeout_s", "target_ongoing_requests", "queue_config"} + ) + values["device_group"].setdefault("name", "sampler") + values["device_group"].setdefault("ranks", self.args.nproc_per_node) + values["device_group"].setdefault("device_type", device_type) + app = build_sampler_app( + **values, + queue_config=self.args.queue_config, + deploy_options=options, + sampling_context=RuntimeContext( + namespace=self.namespace, registry_name=allocation.name + "-registry", instance_id=instance.identity + ), + placement_group=await asyncio.to_thread(get_placement_group, allocation.name), + allocation_devices=list(allocation.devices), + ) + await self.runtime.deploy( + instance.identity, + allocation.name, + app, + f"/api/v1/sampling-runtime/{instance.identity}", + lambda: self.shutdown(instance), + ) + + async def handle(self, instance): + return await self.runtime.handle(instance.application) + + async def invoke(self, instance, ticket, payload, stream_queue=None): + handle = await self.handle(instance) + return await handle.invoke.remote(ticket, payload, stream_queue) + + async def wait_operation(self, instance, ticket): + handle = await self.handle(instance) + await handle.wait_operation.remote(ticket) + + async def load_weights(self, instance, path): + handle = await self.handle(instance) + await handle.install_base_weights.remote(path) + + async def health(self, instance): + async with asyncio.timeout(5): + handle = await self.handle(instance) + return await handle.runtime_identity.remote() + + async def shutdown(self, instance): + name = instance.application or self.runtime.pool.name( + f"sampling_{self.pool_id}_{instance.epoch}", instance.identity + ) + futures = [] + try: + async with asyncio.timeout(5): + handle = await self.handle(instance) + futures = await handle.pending_futures.remote() + except Exception: + pass + await self.runtime.shutdown(instance.identity, name) + failure = FutureFailureRecord( + reason_code="orphaned_replica", + attribution="server", + message="Sampler instance stopped; retry with the checkpoint path.", + ) + state = get_server_state() + for request_id in futures: + await state.store_future_status(request_id, "failed", None, failure=failure) + + async def recover_orphans(self): + import re + from types import SimpleNamespace + + pattern = re.compile(re.escape(self.prefix) + r"([0-9a-f]{32})-(sampler_[0-9a-f]{32})$") + + def decode(name): + match = pattern.fullmatch(name) + if match: + return SimpleNamespace(identity=match[2], epoch=match[1], application=name) + + await self.runtime.recover_orphans(decode, self.shutdown) diff --git a/src/twinkle/server/sampler/pool/router.py b/src/twinkle/server/sampler/pool/router.py new file mode 100644 index 000000000..fa6e2b9d6 --- /dev/null +++ b/src/twinkle/server/sampler/pool/router.py @@ -0,0 +1,107 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Stable CPU ingress for an owned sampling pool; existing sampler wire is forwarded intact.""" + +from __future__ import annotations +import asyncio + +import json +from fastapi import Depends, Request, Response +from fastapi.responses import StreamingResponse + +from twinkle.server.checkpoint import create_checkpoint_manager +from twinkle.server.deployment import bind_deployment, build_deployment_app +from twinkle.server.lifecycle.proxy import ForwardedRequest +from twinkle.server.middleware.auth import get_token_from_request +from twinkle.server.sampler.weights import is_full_checkpoint +from twinkle.server.state import get_server_state +from .controller import sampler_controller_name + + +class SamplerRouter: + async def __init__(self, namespace, pool_id): + import ray + + self.state = get_server_state() + self.controller = ray.get_actor(sampler_controller_name(namespace, pool_id), namespace=namespace) + + async def forward(self, request, endpoint): + body = await request.body() + payload: ForwardedRequest = { + "endpoint": endpoint, + "body": body, + "method": request.method, + "headers": list(request.headers.raw), + "query_string": request.scope.get("query_string", b""), + } + full_path = None + if endpoint in ("twinkle/sample", "twinkle/sample_stream", "twinkle/sample_to_data_plane", "tinker/asample"): + try: + data = json.loads(body) + except (ValueError, UnicodeDecodeError): + data = {} + if not isinstance(data, dict): + data = {} + # The existing handler remains responsible for its validation/wire errors. + path = data.get("adapter_uri") or data.get("model_path") + if not path and endpoint == "tinker/asample" and data.get("sampling_session_id"): + session = await self.state.get_sampling_session(data["sampling_session_id"]) + path = session.get("model_path") if session else None + if path: + client_type = "tinker" if endpoint.startswith("tinker/") else "twinkle" + _, resolved = create_checkpoint_manager( + get_token_from_request(request), client_type=client_type + ).parse_adapter_uri(path) + if await asyncio.to_thread(is_full_checkpoint, resolved): + full_path = resolved + stream = endpoint == "twinkle/sample_stream" + queue = None + if stream: + from ray.util.queue import Queue + + queue = Queue(maxsize=128) + try: + reply = await self.controller.request.remote(payload, queue, full_path) + except BaseException: + if queue is not None: + queue.shutdown(force=True) + raise + headers = { + key.decode(): value.decode() + for key, value in reply["headers"] + if key.lower() not in (b"content-length", b"transfer-encoding") + } + if not stream: + return Response(reply["body"], status_code=reply["status"], headers=headers) + + async def generate(): + try: + while True: + item = await queue.get_async() + if item is None: + return + yield item + finally: + queue.shutdown(force=True) + + return StreamingResponse(generate(), status_code=reply["status"], headers=headers) + + +def build_sampler_router(namespace, pool_id, deploy_options): + def routes(app, get_self): + @app.get("/healthz") + async def health(self=Depends(get_self)): + diagnostic = await self.controller.diagnostics.remote() + return {"healthy": diagnostic["ready"], **diagnostic} + + @app.post("/twinkle/{endpoint}") + async def native(request: Request, endpoint: str, self=Depends(get_self)): + return await self.forward(request, f"twinkle/{endpoint}") + + @app.post("/tinker/{endpoint}") + async def tinker(request: Request, endpoint: str, self=Depends(get_self)): + return await self.forward(request, f"tinker/{endpoint}") + + app = build_deployment_app("SamplerRouter", routes) + return bind_deployment( + app, SamplerRouter, deploy_options, deployment_name="SamplerRouter", bind_args=(namespace, pool_id) + ) diff --git a/src/twinkle/server/sampler/runtime.py b/src/twinkle/server/sampler/runtime.py new file mode 100644 index 000000000..94976e4ec --- /dev/null +++ b/src/twinkle/server/sampler/runtime.py @@ -0,0 +1,159 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Request ownership inside one allocated sampling replica, independent of HTTP waiters.""" + +from __future__ import annotations +import asyncio +from dataclasses import dataclass + +from twinkle.server.deployment import build_deployment_app +from twinkle.server.exceptions import ResourceNotFoundError +from twinkle.server.lifecycle.proxy import ForwardedReply, ForwardedRequest +from twinkle.server.task_queue.types import ExecutionScope + + +def request_execution_scope(request) -> ExecutionScope | None: + """Internal ASGI metadata, absent for direct standalone endpoint calls.""" + return request.scope.get("state", {}).get("execution_scope") + + +@dataclass +class SamplingOperation: + scope: ExecutionScope + task: asyncio.Task | None = None + + +class SamplerRuntime: + def __init__(self, owner): + self.owner = owner + self._operations: dict[str, SamplingOperation] = {} + self._stream_tasks = set() + self._rpc_app = None + + def start_stream_job(self, factory, *, execution_scope: ExecutionScope | None = None): + job = asyncio.create_task(factory()) + if execution_scope is not None: + execution_scope.track(job) + self._stream_tasks.add(job) + + def done(task): + self._stream_tasks.discard(task) + if not task.cancelled(): + task.exception() + + job.add_done_callback(done) + return job + + async def shutdown(self): + await asyncio.gather(*tuple(self._stream_tasks), return_exceptions=True) + + async def invoke(self, ticket: str, payload: ForwardedRequest, stream_queue=None) -> ForwardedReply: + """Return the HTTP reply early while retaining every actual generation job.""" + if ticket in self._operations: + raise ResourceNotFoundError("Sampling request was already dispatched") + if self._rpc_app is None: + from .tinker_handlers import _register_tinker_sampler_routes + from .twinkle_handlers import _register_twinkle_sampler_routes + + def routes(app, get_self): + _register_tinker_sampler_routes(app, lambda: self.owner) + _register_twinkle_sampler_routes(app, lambda: self.owner) + + self._rpc_app = build_deployment_app("Sampler", routes) + loop = asyncio.get_running_loop() + reply = loop.create_future() + record = SamplingOperation(ExecutionScope()) + self._operations[ticket] = record + + async def dispatch(): + response: ForwardedReply = {"status": 500, "headers": [], "body": b""} + sent = False + + async def receive(): + nonlocal sent + if not sent: + sent = True + return {"type": "http.request", "body": payload["body"], "more_body": False} + await asyncio.Event().wait() + + async def send(message): + if message["type"] == "http.response.start": + response["status"] = message["status"] + response["headers"] = message["headers"] + if stream_queue is not None and not reply.done(): + reply.set_result(dict(response)) + elif message["type"] == "http.response.body": + chunk = message.get("body", b"") + if stream_queue is None: + response["body"] += chunk + elif chunk: + try: + await stream_queue.put_async(chunk) + except Exception: + pass # Keep draining generation after an HTTP consumer disconnects. + + scope = { + "type": "http", + "state": {"execution_scope": record.scope}, + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "scheme": "http", + "method": payload["method"], + "path": "/" + payload["endpoint"], + "raw_path": ("/" + payload["endpoint"]).encode(), + "root_path": "", + "query_string": payload.get("query_string", b""), + "headers": payload["headers"], + "server": ("sampler", 80), + "client": ("controller", 0), + } + try: + await self._rpc_app(scope, receive, send) + if not reply.done(): + reply.set_result(response) + except BaseException as error: + if not reply.done(): + reply.set_exception(error) + raise + finally: + # ComputeWorker supplied completion handles, not a terminal status read. + await record.scope.wait() + if stream_queue is not None: + try: + await stream_queue.put_async(None) + except Exception: + pass + + record.task = asyncio.create_task(dispatch()) + record.task.add_done_callback(lambda task: task.exception() if not task.cancelled() else None) + return await asyncio.shield(reply) + + async def wait_operation(self, ticket): + record = self._operations.get(ticket) + if record is None: + raise ResourceNotFoundError("Unknown sampling operation") + await asyncio.shield(record.task) + self.owner.sampling_guard.check_available() + self._operations.pop(ticket, None) + + async def install_base_weights(self, path): + from .weights import resolve_sampler_weights + + await self.owner.sampling_guard.run( + lambda: resolve_sampler_weights(self.owner, path, full_checkpoint=True), exclusive=True + ) + + async def pending_futures(self): + return list({identity for record in self._operations.values() for identity in record.scope.request_ids}) + + async def runtime_identity(self): + import ray + + actors = getattr(self.owner.sampler, "_actors", ()) + states = await asyncio.to_thread( + lambda: [ray._private.state.actors(actor._actor_id.hex()).get("State") for actor in actors] + ) + return { + "unhealthy": self.owner.sampling_guard._closed + or self.owner.sampling_guard._weights_error is not None + or any(state != "ALIVE" for state in states) + } diff --git a/src/twinkle/server/sampler/tinker_handlers.py b/src/twinkle/server/sampler/tinker_handlers.py index 719ab0621..b6a6543ed 100644 --- a/src/twinkle/server/sampler/tinker_handlers.py +++ b/src/twinkle/server/sampler/tinker_handlers.py @@ -20,9 +20,10 @@ from twinkle.data_format import SamplingParams from twinkle.server.checkpoint import create_checkpoint_manager +from twinkle.server.exceptions import RequestRejectedError from twinkle.server.sampler.weights import sample_with_weights -from twinkle.server.task_queue.types import UserTaskError from twinkle.utils.logger import get_logger +from .runtime import request_execution_scope logger = get_logger() @@ -84,7 +85,7 @@ async def _do_sample(): # Base-model sampling is valid when no model_path was provided. if adapter_uri and not os.path.exists(adapter_uri): - raise UserTaskError(f'Adapter URI {model_path} does not exist. Please check the model_path.') + raise RequestRejectedError(f'Adapter URI {model_path} does not exist. Please check the model_path.') # Convert tinker SamplingParams to twinkle SamplingParams if needed sampling_params = None @@ -153,4 +154,5 @@ async def generate(lora_path): token=token, input_tokens=input_tokens, task_type='sample', + execution_scope=request_execution_scope(request), ) diff --git a/src/twinkle/server/sampler/twinkle_handlers.py b/src/twinkle/server/sampler/twinkle_handlers.py index 7b8154846..648bfad4f 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Mapping from fastapi import Depends, FastAPI, Request from fastapi.responses import StreamingResponse -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple if TYPE_CHECKING: from .app import SamplerManagement @@ -33,6 +33,7 @@ from twinkle.server.telemetry.correlation import MODEL_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.utils.logger import get_logger +from .runtime import request_execution_scope logger = get_logger() @@ -283,7 +284,12 @@ async def generate(adapter_path): # Calculate metrics for queue scheduling. The body is wire-validated, so the # entries are models and ``input_ids`` is absent or a list of ints. input_tokens = sum(len(getattr(entry, 'input_ids', None) or ()) for entry in body.inputs) - return await self.submit_concurrent_and_peek(_task, token=token, input_tokens=input_tokens, task_type='sample') + return await self.submit_concurrent_and_peek( + _task, + token=token, + input_tokens=input_tokens, + task_type='sample', + execution_scope=request_execution_scope(request)) @app.post('/twinkle/sample_to_data_plane', response_model=types.TaskEnvelope) async def sample_to_data_plane( @@ -335,7 +341,10 @@ async def generate(): return ref.model_dump() return await self.submit_concurrent_and_peek( - _generate_and_store, model_id=full_adapter_name or None, task_type='sample_to_data_plane') + _generate_and_store, + model_id=full_adapter_name or None, + task_type='sample_to_data_plane', + execution_scope=request_execution_scope(request)) @app.post('/twinkle/encode', response_model=types.TaskEnvelope) async def encode( @@ -357,11 +366,12 @@ async def _task(): inputs = to_backend_inputs(body.inputs) pooling_params = PoolingParams.from_dict(body.pooling_params) if body.pooling_params else None adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) or '' - responses = await self.call_backend(self.sampler.encode, inputs, pooling_params, - adapter_name=adapter_name) + responses = await self.sampling_guard.run( + lambda: self.call_backend(self.sampler.encode, inputs, pooling_params, adapter_name=adapter_name)) return {'data': [list(response.data) for response in responses]} - return await self.submit_concurrent_and_peek(_task, token=token, input_tokens=0, task_type='encode') + return await self.submit_concurrent_and_peek( + _task, token=token, input_tokens=0, task_type='encode', execution_scope=request_execution_scope(request)) @app.post('/twinkle/unload_adapter_paths') async def unload_adapter_paths( @@ -426,8 +436,11 @@ async def add_adapter_to_sampler( raise RequestRejectedError('`adapter_name` is required and must be non-empty.') full_adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) - from peft import LoraConfig - config = LoraConfig(**body.config) if isinstance(body.config, dict) else body.config + def build_config(): + from peft import LoraConfig + return LoraConfig(**body.config) if isinstance(body.config, dict) else body.config + + config = await asyncio.to_thread(build_config) with traced_operation('sampler.add_adapter_to_sampler', attrs={MODEL_ID: self.model_id}): await self.sampling_guard.run( @@ -463,7 +476,7 @@ async def sample_stream( """ token = await self._on_request_start(request) - adapter_path = None + resolved_uri = None adapter_name = body.adapter_name or '' full_adapter_name = resolve_twinkle_adapter_name(request, adapter_name) or '' @@ -471,8 +484,6 @@ async def sample_stream( from twinkle.server.checkpoint import create_checkpoint_manager checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') _, resolved_uri = checkpoint_manager.parse_adapter_uri(body.adapter_uri) - await self.call_backend(self.sampler.reset_prefix_cache) - adapter_path = await resolve_sampler_weights(self, resolved_uri) # Streaming accepts exactly one input; the shared seam enforces that and # returns a single parsed object. Its ValueError maps to the same 400 this @@ -488,7 +499,7 @@ async def sample_stream( from ray.util.queue import Queue - from .backends import STREAM_SENTINEL + from .backends.streaming import STREAM_SENTINEL request_id = f'req_{uuid.uuid4().hex}' actors = self.sampler._actors @@ -501,14 +512,31 @@ async def _no_actor_generator(): return StreamingResponse(_no_actor_generator(), media_type='application/x-ndjson') q = Queue(maxsize=128) actor = actors[0] - actor.sample_stream_to_queue.remote( - q, - inputs_parsed, - params, - adapter_name=full_adapter_name, - adapter_path=adapter_path, - ) + async def stream_job(): + + async def _stream_generation(adapter_path): + # Retain the actor result, not merely the HTTP queue iterator. + ref = actor.sample_stream_to_queue.remote( + q, inputs_parsed, params, adapter_name=full_adapter_name, adapter_path=adapter_path) + try: + await ref + except BaseException as error: + # An errored bridge cannot certify the inner engine request + # exited. Controller must confirm worker death before reuse. + self.sampling_guard.invalidate_weights(error) + raise + + try: + await sample_with_weights(self, resolved_uri, _stream_generation) + except Exception as error: + try: + await q.put_async(error) + await q.put_async(STREAM_SENTINEL) + except Exception: + pass + + self.request_runtime.start_stream_job(stream_job, execution_scope=request_execution_scope(request)) return StreamingResponse( _stream_queue( q, diff --git a/src/twinkle/server/sampler/weights.py b/src/twinkle/server/sampler/weights.py index f2d1077f7..dac788542 100644 --- a/src/twinkle/server/sampler/weights.py +++ b/src/twinkle/server/sampler/weights.py @@ -7,7 +7,8 @@ from safetensors import SafetensorError from twinkle.exceptions import FullWeightLoadError -from twinkle.server.task_queue.types import BACKEND_TIMEOUT_EXCEPTIONS, UserTaskError +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.task_queue.backend_gate import BACKEND_TIMEOUT_EXCEPTIONS def is_full_checkpoint(resolved_uri: str | None) -> bool: @@ -16,10 +17,10 @@ def is_full_checkpoint(resolved_uri: str | None) -> bool: if os.path.isfile(os.path.join(resolved_uri, 'adapter_config.json')): return False if os.path.isabs(resolved_uri) and not os.path.isdir(resolved_uri): - raise UserTaskError(f'Checkpoint directory does not exist: {resolved_uri}') + raise RequestRejectedError(f'Checkpoint directory does not exist: {resolved_uri}') if os.path.isdir(resolved_uri) and not (os.path.isfile(os.path.join(resolved_uri, 'model.safetensors.index.json')) or glob.glob(os.path.join(resolved_uri, '*.safetensors'))): - raise UserTaskError(f'No .safetensors weights found under {resolved_uri}') + raise RequestRejectedError(f'No .safetensors weights found under {resolved_uri}') return True @@ -38,7 +39,7 @@ async def resolve_sampler_weights(service, resolved_uri: str | None, *, full_che service.sampling_guard.invalidate_weights(error) raise except (ValueError, OSError, KeyError, SafetensorError) as error: - raise UserTaskError(str(error)) from error + raise RequestRejectedError(str(error)) from error return None diff --git a/src/twinkle/server/session_resource/__init__.py b/src/twinkle/server/session_resource/__init__.py index fdb7ebdc1..7384532cb 100644 --- a/src/twinkle/server/session_resource/__init__.py +++ b/src/twinkle/server/session_resource/__init__.py @@ -7,8 +7,7 @@ session-driven expiration of session-bound resources. """ -from .adapter import AdapterManagerMixin from .base import SessionResourceMixin from .processor import ProcessorManagerMixin -__all__ = ['AdapterManagerMixin', 'ProcessorManagerMixin', 'SessionResourceMixin'] +__all__ = ['ProcessorManagerMixin', 'SessionResourceMixin'] diff --git a/src/twinkle/server/session_resource/adapter.py b/src/twinkle/server/session_resource/adapter.py deleted file mode 100644 index 392a05bf1..000000000 --- a/src/twinkle/server/session_resource/adapter.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -""" -Adapter Lifecycle Manager Mixin for Twinkle Server. - -This module provides adapter lifecycle management as a mixin class that can be -inherited directly by services. It tracks adapter activity and provides interfaces -for registration, heartbeat updates, and expiration handling. - -By inheriting this mixin, services can override the _on_adapter_expired() method -to handle expired adapters without using callbacks or polling. -""" -from __future__ import annotations - -from abc import abstractmethod - -from twinkle.utils.logger import get_logger -from .base import SessionResourceMixin - -logger = get_logger() - - -class AdapterManagerMixin(SessionResourceMixin): - """Mixin for adapter lifecycle management with session-based expiration. - - This mixin tracks adapter activity and automatically expires adapters - when their associated session expires. - - Inheriting classes should: - 1. Call _init_adapter_manager() in __init__ - 2. Override _on_adapter_expired() to customize expiration handling - - The inactivity timeout / max lifetime are stored on the base mixin as - ``_resource_timeout`` / ``_resource_max_lifetime`` (set via - ``_init_adapter_manager``). - """ - - # Set resource type for logging - _resource_type = 'Adapter' - - def _init_adapter_manager( - self, - adapter_timeout: float = 1800.0, - adapter_max_lifetime: float = 36000.0, - ) -> None: - """Initialize the adapter manager. - - This should be called in the __init__ of the inheriting class. - - Args: - adapter_timeout: Timeout in seconds used to check whether a session is still alive. - Default is 1800.0 (30 minutes). - adapter_max_lifetime: Maximum lifetime in seconds for an adapter regardless of session - liveness. Adapters older than this are treated as expired. Default is 36000.0 (10 hours). - """ - self._init_resource_manager( - resource_timeout=adapter_timeout, - resource_max_lifetime=adapter_max_lifetime, - ) - - async def _on_resource_expired(self, resource_id: str) -> None: - """Base-class expiry hook; forwards to the domain hook ``_on_adapter_expired``. - - ``_on_adapter_expired`` is the supported extension point: the adapter-domain - name is kept deliberately so subclass authors override a method named for - adapters rather than the generic base-class resource hook. - """ - await self._on_adapter_expired(resource_id) - - @abstractmethod - async def _on_adapter_expired(self, adapter_name: str) -> None: - """Hook method called when an adapter expires.""" - ... - - @staticmethod - def get_adapter_name(adapter_name: str) -> str: - """Get the adapter name for a request. - - This is a passthrough method for consistency with the original API. - - Args: - adapter_name: The adapter name (typically model_id) - - Returns: - The adapter name to use - """ - return adapter_name diff --git a/src/twinkle/server/session_resource/processor.py b/src/twinkle/server/session_resource/processor.py index 69a6192e3..dd2a7ff10 100644 --- a/src/twinkle/server/session_resource/processor.py +++ b/src/twinkle/server/session_resource/processor.py @@ -2,7 +2,7 @@ """ Processor Lifecycle Manager Mixin for Twinkle Server. -Mirrors AdapterManagerMixin but adds a global per-token processor limit. +Adds a global per-token processor limit to session resources. Sessions are tracked via session ID; processors expire when their session expires. """ from __future__ import annotations @@ -18,7 +18,7 @@ class ProcessorManagerMixin(SessionResourceMixin): """Mixin for processor lifecycle management with session-based expiration. - Mirrors AdapterManagerMixin with an additional per-token processor limit. + Manages processor sessions with a global per-token processor limit. Inheriting classes should: 1. Call _init_processor_manager() in __init__ @@ -80,9 +80,7 @@ async def _on_resource_expired(self, resource_id: str) -> None: ``_on_processor_expired`` is the supported extension point: the processor-domain name is kept deliberately so subclass authors override a - method named for processors rather than the generic base-class hook. It is - ``async`` to match the sibling ``AdapterManagerMixin._on_adapter_expired`` - contract, so both resource kinds expose the same extension-point shape. + method named for processors rather than the generic base-class hook. """ await self._on_processor_expired(resource_id) diff --git a/src/twinkle/server/state/__init__.py b/src/twinkle/server/state/__init__.py index 5f9a47549..3c270ea30 100644 --- a/src/twinkle/server/state/__init__.py +++ b/src/twinkle/server/state/__init__.py @@ -1,17 +1,26 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .backend import create_backend -from .base import BaseManager -from .cleanup_coordinator import ResourceCleanupCoordinator -from .count_publisher import ResourceCountPublisher -from .future_manager import FutureManager -from .model_manager import ModelManager -from .models import FutureRecord, ModelRecord, SamplingSessionRecord, SessionRecord -from .replica_registry import ReplicaRegistry -from .sampling_manager import SamplingSessionManager -from .server_state import ServerState, get_server_state, reset_server_state_cache -from .session_manager import SessionManager +import sys -__all__ = [ +from twinkle._lazy_module import _LazyModule + +sys.modules[__name__] = _LazyModule( + __name__, + __file__, { + 'backend': ['create_backend'], + 'base': ['BaseManager'], + 'cleanup_coordinator': ['ResourceCleanupCoordinator'], + 'count_publisher': ['ResourceCountPublisher'], + 'future_manager': ['FutureManager'], + 'model_manager': ['ModelManager'], + 'models': ['FutureRecord', 'ModelRecord', 'SamplingSessionRecord', 'SessionRecord'], + 'replica_registry': ['ReplicaRegistry'], + 'sampling_manager': ['SamplingSessionManager'], + 'server_state': ['ServerState', 'get_server_state', 'reset_server_state_cache'], + 'session_manager': ['SessionManager'], + }, + module_spec=__spec__) + +sys.modules[__name__].__all__ = [ # Pydantic record models 'SessionRecord', 'ModelRecord', diff --git a/src/twinkle/server/state/backend/base.py b/src/twinkle/server/state/backend/base.py index b65f7d2b2..4ed5bbc49 100644 --- a/src/twinkle/server/state/backend/base.py +++ b/src/twinkle/server/state/backend/base.py @@ -74,6 +74,8 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: """Read key, call ``transform(current_value)``, write the result atomically. @@ -87,6 +89,10 @@ async def update_atomic( and may raise :class:`ConcurrencyError` after exhausting internal retries (currently 16). + With ``return_changed=True``, return whether THIS call committed a + write, not the stored value. A lost write acknowledgement raises; it + must never be inferred as success from a later read. + ``transform`` must be picklable when running against a Ray-backed backend — pass module-level functions wrapped with ``functools.partial``, not lambdas or local closures. diff --git a/src/twinkle/server/state/backend/factory.py b/src/twinkle/server/state/backend/factory.py index 842a67dfe..d91a29039 100644 --- a/src/twinkle/server/state/backend/factory.py +++ b/src/twinkle/server/state/backend/factory.py @@ -1,6 +1,8 @@ """Backend factory for creating StateBackend instances based on configuration.""" from __future__ import annotations +import os + from twinkle.server.config.persistence import PersistenceConfig from twinkle.utils import get_logger from .base import StateBackend @@ -29,7 +31,8 @@ def create_backend(config: PersistenceConfig | None = None) -> StateBackend: # Deferred import keeps the module-level import graph light; the Twinkle # server always runs on Ray Serve, so ``ray`` is available here. from .memory_backend import RayActorBackend - return RayActorBackend(key_prefix=config.key_prefix) + return RayActorBackend( + key_prefix=config.key_prefix, namespace=os.environ.get('TWINKLE_STATE_RAY_NAMESPACE')) case 'redis': if not config.redis_url: raise ValueError('redis_url is required for redis persistence mode') diff --git a/src/twinkle/server/state/backend/memory_backend.py b/src/twinkle/server/state/backend/memory_backend.py index 433f8baa4..a92561264 100644 --- a/src/twinkle/server/state/backend/memory_backend.py +++ b/src/twinkle/server/state/backend/memory_backend.py @@ -84,14 +84,16 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: current = None if self._is_expired(key) else self._store[key][0] new_value = transform(current) if new_value is None: - return current + return False if return_changed else current expire_at = (time.time() + ttl) if ttl is not None else None self._store[key] = (new_value, expire_at) - return new_value + return True if return_changed else new_value async def mget(self, keys: list[str]) -> list[Any | None]: results: list[Any | None] = [] @@ -121,21 +123,26 @@ class RayActorBackend(StateBackend): persistence mode literal for this backend is still ``"memory"``.) """ - def __init__(self, key_prefix: str = '') -> None: + def __init__(self, key_prefix: str = '', namespace: str | None = None) -> None: if not ray.is_initialized(): raise RuntimeError('RayActorBackend requires an initialized Ray runtime — call ' 'ray.init() first, switch persistence to "redis", or rely on ' 'the deployment launcher to start Ray.') name = _actor_name(key_prefix) try: - self._actor = ray.get_actor(name) + self._actor = ray.get_actor(name, namespace=namespace) except ValueError: try: - self._actor = _StateActor.options(name=name, lifetime='detached').remote() + self._actor = _StateActor.options( + name=name, + namespace=namespace, + lifetime='detached', + max_restarts=-1 if namespace else 0, + max_task_retries=0).remote() except ValueError: # Lost the create race against another worker — the actor is # now resolvable. - self._actor = ray.get_actor(name) + self._actor = ray.get_actor(name, namespace=namespace) async def set(self, key: str, value: Any, ttl: int | None = None) -> None: await self._actor.set.remote(key, value, ttl) @@ -163,8 +170,10 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: - return await self._actor.update_atomic.remote(key, transform, ttl) + return await self._actor.update_atomic.remote(key, transform, ttl, return_changed=return_changed) async def mget(self, keys: list[str]) -> list[Any | None]: return await self._actor.mget.remote(keys) diff --git a/src/twinkle/server/state/backend/redis_backend.py b/src/twinkle/server/state/backend/redis_backend.py index e6769551a..53b4d78ce 100644 --- a/src/twinkle/server/state/backend/redis_backend.py +++ b/src/twinkle/server/state/backend/redis_backend.py @@ -6,10 +6,13 @@ from collections.abc import Callable from typing import Any +from twinkle.server.telemetry.metrics import get_state_backend_metrics from .base import ConcurrencyError, StateBackend try: import redis.asyncio as aioredis + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError from redis.exceptions import WatchError _REDIS_AVAILABLE = True @@ -38,6 +41,7 @@ def __init__(self, redis_url: str, key_prefix: str = '') -> None: raise ImportError('redis package required. Install with: pip install redis') self._client = aioredis.from_url(redis_url, decode_responses=True) self._prefix = key_prefix + self._metrics = get_state_backend_metrics() def _make_key(self, key: str) -> str: """Add namespace prefix to key.""" @@ -56,10 +60,27 @@ async def set(self, key: str, value: Any, ttl: int | None = None) -> None: else: await self._client.set(real_key, data) + async def _retry_read(self, operation, *args): + """Retry only read commands after a stale pooled connection is discarded. + + Writes/EXEC are deliberately excluded: a lost acknowledgement does not + prove that a mutation failed, so replay could apply it twice. + """ + for attempt in range(3): + try: + result = await operation(*args) + self._metrics.redis_health.set(1) + return result + except (RedisConnectionError, RedisTimeoutError): + self._metrics.redis_health.set(0) + if attempt == 2: + raise + await asyncio.sleep(0.05 * (attempt + 1)) + async def get(self, key: str) -> Any | None: """Retrieve value, return None if not found or expired.""" real_key = self._make_key(key) - raw = await self._client.get(real_key) + raw = await self._retry_read(self._client.get, real_key) if raw is None: return None return json.loads(raw) @@ -72,7 +93,7 @@ async def delete(self, key: str) -> None: async def exists(self, key: str) -> bool: """Check if key exists and is not expired.""" real_key = self._make_key(key) - return bool(await self._client.exists(real_key)) + return bool(await self._retry_read(self._client.exists, real_key)) async def keys(self, pattern: str) -> list[str]: """Return all key names matching the pattern using SCAN. @@ -106,6 +127,8 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: """WATCH/GET/MULTI/SET/EXEC with bounded retries + jittered backoff on WatchError. @@ -116,8 +139,10 @@ async def update_atomic( / heartbeat) or reraise (lease renewal). """ real_key = self._make_key(key) + connection_failures = 0 for attempt in range(_UPDATE_ATOMIC_MAX_RETRIES): async with self._client.pipeline() as pipe: + executing = False try: await pipe.watch(real_key) raw = await pipe.get(real_key) @@ -125,21 +150,40 @@ async def update_atomic( new_value = transform(current) if new_value is None: await pipe.unwatch() - return current + self._metrics.update_atomic_retries.observe(attempt) + return False if return_changed else current pipe.multi() if ttl is not None: pipe.set(real_key, json.dumps(new_value), ex=ttl) else: pipe.set(real_key, json.dumps(new_value)) + executing = True await pipe.execute() - return new_value - except WatchError: + self._metrics.update_atomic_retries.observe(attempt) + return True if return_changed else new_value + except WatchError as exc: + # redis-py also wraps a lost watched connection as WatchError. + # Once EXEC was sent, its outcome may be unknown, not aborted. + transport_failure = isinstance(exc.__context__, + (RedisConnectionError, RedisTimeoutError)) or (str(exc).startswith( + ('A ConnectionError occurred', 'A TimeoutError occurred'))) + if executing and transport_failure: + raise RedisConnectionError('Atomic update outcome could not be confirmed') from exc + self._metrics.update_atomic_conflicts.inc() backoff = min( _UPDATE_ATOMIC_BASE_BACKOFF_SECONDS * (2**attempt), _UPDATE_ATOMIC_MAX_BACKOFF_SECONDS, ) await asyncio.sleep(backoff * random.random()) continue + except (RedisConnectionError, RedisTimeoutError): + self._metrics.redis_health.set(0) + connection_failures += 1 + if executing or connection_failures >= 3: + raise + # Before EXEC there was no mutation, so reconnecting is safe. + await asyncio.sleep(0.05 * connection_failures) + self._metrics.update_atomic_retries.observe(_UPDATE_ATOMIC_MAX_RETRIES) raise ConcurrencyError(f'update_atomic exhausted {_UPDATE_ATOMIC_MAX_RETRIES} retries on key {key!r}') async def mget(self, keys: list[str]) -> list[Any | None]: @@ -147,7 +191,7 @@ async def mget(self, keys: list[str]) -> list[Any | None]: if not keys: return [] real_keys = [self._make_key(k) for k in keys] - raw_values = await self._client.mget(real_keys) + raw_values = await self._retry_read(self._client.mget, real_keys) return [json.loads(v) if v is not None else None for v in raw_values] async def close(self) -> None: @@ -157,6 +201,8 @@ async def close(self) -> None: async def health_check(self) -> bool: """Check if Redis is healthy and available.""" try: - return bool(await self._client.ping()) + healthy = bool(await self._retry_read(self._client.ping)) except Exception: - return False + healthy = False + self._metrics.redis_health.set(1 if healthy else 0) + return healthy diff --git a/src/twinkle/server/state/cleanup_coordinator.py b/src/twinkle/server/state/cleanup_coordinator.py index c6278c248..419846181 100644 --- a/src/twinkle/server/state/cleanup_coordinator.py +++ b/src/twinkle/server/state/cleanup_coordinator.py @@ -71,6 +71,7 @@ def __init__( self._on_become_leader = on_become_leader self._on_lose_leader = on_lose_leader + self._lifecycle_lock = asyncio.Lock() self._cleanup_task: asyncio.Task | None = None self._cleanup_running = False self._leader_id = uuid.uuid4().hex @@ -143,16 +144,14 @@ async def _cleanup_loop(self) -> None: async def _leader_loop(self) -> None: """Acquire and renew the cleanup-leader lease every LEASE_RENEW seconds.""" - await self._try_acquire_or_renew() # Race for leadership at startup while self._leader_running: try: - await asyncio.sleep(LEASE_RENEW) await self._try_acquire_or_renew() except asyncio.CancelledError: break except Exception as e: logger.warning(f'[ServerState Leader] renew error: {e}') - continue + await asyncio.sleep(LEASE_RENEW) async def _try_acquire_or_renew(self) -> None: was_leader = self._is_leader @@ -169,17 +168,8 @@ async def _try_acquire_or_renew(self) -> None: except Exception as e: logger.warning(f'[ServerState Leader] backend error during election: {e}') self._is_leader = False - if was_leader: - # Our renewal failed but our lease value may still be sitting in the - # backend, so a plain ``set_nx`` would keep returning False for up to - # LEASE_TTL and leadership would stall unclaimed. Best-effort delete - # ONLY when we were the leader (never steal a lease another replica - # legitimately holds), swallowing errors so a delete failure cannot - # escape the election loop. The next tick can then re-acquire. - try: - await self._backend.delete(LEADER_KEY) - except Exception: - pass + # A renewal failure does not prove current ownership. Leave the key + # untouched: our lease expires naturally, or a successor owns it. if self._is_leader and not was_leader: logger.info(f'[ServerState] became cleanup leader (id={self._leader_id[:8]})') @@ -199,36 +189,32 @@ async def start(self) -> bool: ``ServerState``) so the implementation and its guard can never drift apart into a double-start. """ - if self._cleanup_running: - return False - # Rebuild in-memory indexes from backend data before the loops start. - await self._managers['models'].rebuild_indexes() - self._cleanup_running = True - self._cleanup_task = asyncio.create_task(self._cleanup_loop()) - self._leader_running = True - self._leader_task = asyncio.create_task(self._leader_loop()) - return True + async with self._lifecycle_lock: + if self._cleanup_running: + return False + await self._managers['models'].rebuild_indexes() + self._cleanup_running = True + self._leader_running = True + self._cleanup_task = asyncio.create_task(self._cleanup_loop()) + self._leader_task = asyncio.create_task(self._leader_loop()) + return True async def stop(self) -> bool: """Stop the background cleanup + leader-election tasks. Returns ``False`` if not running.""" - if not self._cleanup_running: - return False - self._cleanup_running = False - if self._cleanup_task: - self._cleanup_task.cancel() - self._cleanup_task = None - self._leader_running = False - if self._leader_task: - self._leader_task.cancel() - self._leader_task = None - if self._is_leader: - # Release callback registration; the lease itself expires on its own TTL — - # update_atomic can't express "atomic delete", so we accept a short outage - # where the gauge reads 0 between leaders. - if self._on_lose_leader is not None: + async with self._lifecycle_lock: + if not self._cleanup_running: + return False + self._cleanup_running = self._leader_running = False + tasks = [task for task in (self._cleanup_task, self._leader_task) if task is not None] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + self._cleanup_task = self._leader_task = None + was_leader, self._is_leader = self._is_leader, False + # Do not delete an uncertain lease; its original TTL is the fence. + if was_leader and self._on_lose_leader is not None: await self._on_lose_leader() - self._is_leader = False - return True + return True async def get_cleanup_stats(self) -> dict[str, Any]: """Get current cleanup configuration and resource counts.""" diff --git a/src/twinkle/server/state/future_manager.py b/src/twinkle/server/state/future_manager.py index 54bca4c70..4e36d28b4 100644 --- a/src/twinkle/server/state/future_manager.py +++ b/src/twinkle/server/state/future_manager.py @@ -5,6 +5,7 @@ import time from typing import Any +from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES from twinkle.utils.logger import get_logger from .backend.base import StateBackend from .base import BaseManager @@ -12,100 +13,114 @@ logger = get_logger() -# Status sets used by the do-not-regress guard inside the atomic transform. -_TERMINAL_STATUSES = frozenset({'completed', 'failed', 'cancelled'}) -_NON_TERMINAL_STATUSES = frozenset({'pending', 'queued', 'running'}) - - -def _future_record_transform( - existing: dict | None, - *, - request_id: str, - new_status: str, - model_id: str | None, - reason: str | None, - result: Any, - failure: dict[str, Any] | None, - queue_state: str | None, - queue_state_reason: str | None, - replica_id: str | None, - absolute_deadline: float | None, - now: str, -) -> dict | None: - """Atomic transform body for :meth:`FutureManager.store_status`. - - Module-level so it remains picklable when forwarded across the Ray actor - boundary (closures and lambdas cannot be). - - A record already in a terminal state is never overwritten (returns ``None``, - which ``update_atomic`` treats as "keep current value"). A write of a - *different* terminal state is logged; a write of the *same* terminal state is - dropped silently (State_Backend idempotent retries produce these and they - indicate no defect). - """ - existing_status = existing.get('status') if existing is not None else None - if existing_status in _TERMINAL_STATUSES: - if new_status != existing_status: - logger.warning('future %s already terminal as %r; refusing %r', request_id, existing_status, new_status) - return None - - if existing is None: - record = FutureRecord( - status=new_status, - model_id=model_id, - reason=reason, - result=result, - failure=failure, - queue_state=queue_state, - queue_state_reason=queue_state_reason, - replica_id=replica_id, - absolute_deadline=absolute_deadline, - created_at=now, - updated_at=now, - ) - return record.model_dump() - - updated = dict(existing) - updated['status'] = new_status - updated['model_id'] = model_id - updated['updated_at'] = now - # replica_id is set at creation and is deliberately NOT overwritten here. - if reason is not None: - updated['reason'] = reason - if new_status in _TERMINAL_STATUSES: - if new_status in ('failed', 'cancelled'): - updated['result'] = None - updated['failure'] = failure - else: - updated['failure'] = None - updated['result'] = result - elif result is not None: - updated['result'] = result - if queue_state is not None: - updated['queue_state'] = queue_state - if queue_state_reason is not None: - updated['queue_state_reason'] = queue_state_reason - return updated - -_CANCELLABLE_STATUSES = frozenset({'pending', 'queued'}) +class FutureTransitions: + """Pure, picklable atomic transitions; the manager owns I/O and retention.""" + terminal = TERMINAL_STATUSES + cancellable = frozenset({'pending', 'queued'}) + order = {'pending': 0, 'queued': 1, 'running': 2} -def _cancel_if_not_started_transform(existing: dict | None, *, failure: dict, now: str) -> dict | None: - """Atomic transform: cancel iff the task has not started (pending/queued). - - Returns ``None`` (no change) for running/terminal/missing records so a task - already executing is never interrupted -- cancel is best-effort on the queue. - """ - status = existing.get('status') if existing is not None else None - if status not in _CANCELLABLE_STATUSES: - return None - updated = dict(existing) - updated['status'] = 'cancelled' - updated['result'] = None - updated['failure'] = failure - updated['updated_at'] = now - return updated + @staticmethod + def update( + existing: dict | None, + *, + request_id: str, + new_status: str, + model_id: str | None, + reason: str | None, + result: Any, + failure: dict[str, Any] | None, + queue_state: str | None, + queue_state_reason: str | None, + replica_id: str | None, + absolute_deadline: float | None, + now: str, + owner_token: str | None = None, + ) -> dict | None: + """Atomic transform body for :meth:`FutureManager.store_status`. + + A static method remains picklable across the Ray actor boundary. + + A record already in a terminal state is never overwritten (returns ``None``, + which ``update_atomic`` treats as "keep current value"). A write of a + *different* terminal state is logged; a write of the *same* terminal state is + dropped silently (State_Backend idempotent retries produce these and they + indicate no defect). + """ + existing_status = existing.get('status') if existing is not None else None + if existing_status in FutureTransitions.terminal: + if new_status != existing_status: + logger.warning('future %s already terminal as %r; refusing %r', request_id, existing_status, new_status) + return None + if new_status in FutureTransitions.order and existing_status in FutureTransitions.order: + if FutureTransitions.order[new_status] <= FutureTransitions.order[existing_status]: + return None + if new_status == 'cancelled' and existing_status == 'running': + return None + + if existing is None: + record = FutureRecord( + owner_token=owner_token, + status=new_status, + model_id=model_id, + reason=reason, + result=result, + failure=failure, + queue_state=queue_state, + queue_state_reason=queue_state_reason, + replica_id=replica_id, + absolute_deadline=absolute_deadline, + created_at=now, + updated_at=now, + ) + return record.model_dump() + + updated = dict(existing) + updated['status'] = new_status + updated['model_id'] = model_id + updated['updated_at'] = now + # replica_id is set at creation and is deliberately NOT overwritten here. + if reason is not None: + updated['reason'] = reason + if new_status in FutureTransitions.terminal: + if new_status in ('failed', 'cancelled'): + updated['result'] = None + updated['failure'] = failure + else: + updated['failure'] = None + updated['result'] = result + elif result is not None: + updated['result'] = result + if queue_state is not None: + updated['queue_state'] = queue_state + if queue_state_reason is not None: + updated['queue_state_reason'] = queue_state_reason + return updated + + @staticmethod + def start(existing: dict | None, *, now: str) -> dict | None: + """Acquire execution only from pending/queued, using the existing record.""" + if existing is None or existing.get('status') not in FutureTransitions.cancellable: + return None + return {**existing, 'status': 'running', 'queue_state': 'active', 'updated_at': now} + + @staticmethod + def cancel(existing: dict | None, *, failure: dict, now: str) -> dict | None: + """Atomic transform: cancel iff the task has not started (pending/queued). + + Returns ``None`` (no change) for running/terminal/missing records so a task + already executing is never interrupted -- cancel is best-effort on the queue. + """ + status = existing.get('status') if existing is not None else None + if status not in FutureTransitions.cancellable: + return None + updated = dict(existing) + updated['status'] = 'cancelled' + updated['result'] = None + updated['failure'] = failure + updated['updated_at'] = now + return updated class FutureManager(BaseManager[FutureRecord]): @@ -128,6 +143,7 @@ async def store_status( queue_state_reason: str | None = None, replica_id: str | None = None, absolute_deadline: float | None = None, + owner_token: str | None = None, ) -> None: """Create or update a future record with the latest status. @@ -153,7 +169,7 @@ async def store_status( await self._backend.update_atomic( self._make_key(request_id), functools.partial( - _future_record_transform, + FutureTransitions.update, request_id=request_id, new_status=status, model_id=model_id, @@ -164,10 +180,19 @@ async def store_status( queue_state_reason=queue_state_reason, replica_id=replica_id, absolute_deadline=absolute_deadline, + owner_token=owner_token, now=now, ), ) + async def try_start(self, request_id: str) -> bool: + """True only when this attempt atomically acquired execution rights.""" + return await self._backend.update_atomic( + self._make_key(request_id), + functools.partial(FutureTransitions.start, now=_now_iso()), + return_changed=True, + ) + async def cancel_if_pending(self, request_id: str) -> str | None: """Cancel a task iff it has not started; return the resulting status. @@ -183,7 +208,7 @@ async def cancel_if_pending(self, request_id: str) -> str | None: ) result = await self._backend.update_atomic( self._make_key(request_id), - functools.partial(_cancel_if_not_started_transform, failure=failure.model_dump(), now=_now_iso()), + functools.partial(FutureTransitions.cancel, failure=failure.model_dump(), now=_now_iso()), ) return result.get('status') if result else None @@ -219,14 +244,14 @@ async def cleanup_expired( now = time.time() expired_ids: list[str] = [] for request_id, record in all_records.items(): - if record.status in _TERMINAL_STATUSES: + if record.status in FutureTransitions.terminal: timestamp = self._parse_timestamp(record.updated_at or record.created_at) if timestamp < cutoff_time: expired_ids.append(request_id) continue # Non-terminal records are never deleted -- only ever written ``failed``. - # replica_id None (pre-upgrade) => ownership unknown => treated as alive. + # Unowned background tasks have no replica; only explicit owners can be declared dead. replica_id = record.replica_id replica_alive = (replica_id is None or alive_replica_ids is None or replica_id in alive_replica_ids) if not replica_alive: diff --git a/src/twinkle/server/state/model_manager.py b/src/twinkle/server/state/model_manager.py index aaf520fca..39c327e72 100644 --- a/src/twinkle/server/state/model_manager.py +++ b/src/twinkle/server/state/model_manager.py @@ -93,7 +93,8 @@ async def rebuild_indexes(self) -> None: async def get_capacity_info(self) -> dict[str, int]: """Return global LoRA capacity across all registered replicas.""" - replicas = await self._replicas.get_all() + capacities = await self._replicas.get_all() + replicas = {identity: capacity for identity, capacity in capacities.items() if capacity > 0} loaded_per_replica = await self._loaded_per_replica() total_max = sum(replicas.values()) total_used = sum(loaded_per_replica.get(rid, 0) for rid in replicas) @@ -116,6 +117,22 @@ async def unregister_replica(self, replica_id: str) -> None: await self.remove(model_id) await self._replicas.unregister(replica_id) + async def unregister_replica_capacity(self, replica_id: str) -> None: + """Remove capacity and heartbeat only; model records are retired by the owning lifecycle controller.""" + await self._replicas.unregister(replica_id) + + async def save_record(self, model_id: str, record: ModelRecord) -> None: + # Update without re-applying quota; cannot reuse add(). + await self._backend.set(self._make_key(model_id), record.model_dump()) + + async def get_training_records(self, namespace: str): + from .training_records import TrainingRecord + keys = await self._backend.keys('model::*') + return [ + TrainingRecord.model_validate(value) for value in await self._backend.mget(keys) if isinstance(value, dict) + and value.get('controller_epoch') and value.get('training_namespace') == namespace + ] + async def touch_replica_last_seen(self, replica_id: str) -> None: """Refresh a replica's liveness timestamp.""" await self._replicas.touch_last_seen(replica_id) @@ -220,6 +237,9 @@ async def cleanup_expired(self, cutoff_time: float, expired_session_ids: list[st all_records = await self.get_all() expired_ids: list[str] = [] for model_id, record in all_records.items(): + if record.controller_epoch is not None: + # Managed training has one lifecycle writer: TrainingController. + continue if record.session_id and record.session_id in session_set: expired_ids.append(model_id) continue diff --git a/src/twinkle/server/state/models.py b/src/twinkle/server/state/models.py index 0364aa6ce..d84760d38 100644 --- a/src/twinkle/server/state/models.py +++ b/src/twinkle/server/state/models.py @@ -3,7 +3,7 @@ import time from datetime import datetime, timezone -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any, Literal @@ -14,6 +14,16 @@ def _now_iso() -> str: return datetime.now(timezone.utc).isoformat() +class RuntimeContext(BaseModel): + """Allocation identity passed to a worker; never a client request payload.""" + + model_config = ConfigDict(frozen=True, extra='forbid') + + namespace: str + registry_name: str + instance_id: str + + class SessionRecord(BaseModel): """Represents a client session.""" @@ -35,6 +45,9 @@ class ModelRecord(BaseModel): lora_config: Any = None replica_id: str | None = None created_at: str = Field(default_factory=_now_iso) + controller_epoch: str | None = None + training_namespace: str | None = None + instance_id: str | None = None class SamplingSessionRecord(BaseModel): @@ -57,32 +70,10 @@ class FutureFailureRecord(BaseModel): diagnostic: str | None = None -# Canonical set of protocol-independent failure reason codes. The Twinkle and Tinker -# gateways each map this same key set to their own wire vocabularies; keeping the set -# in one place lets a consistency test catch a map that drifts out of coverage. -FAILURE_REASON_CODES: frozenset[str] = frozenset({ - 'invalid_request', - 'request_rejected', - 'resource_not_found', - 'full_mode_busy', - 'input_tokens_exceeded', - 'batch_size_invalid', - 'rate_limit_exceeded', - 'resource_quota_exceeded', - 'cancelled', - 'endpoint_unavailable', - 'state_contention', - 'backend_gate_unavailable', - 'orphaned_replica', - 'execution_timeout', - 'deadline_exceeded', - 'internal_error', -}) - - class FutureRecord(BaseModel): """Represents an async task future / request status.""" + owner_token: str | None = None status: str model_id: str | None = None reason: str | None = None diff --git a/src/twinkle/server/state/server_state.py b/src/twinkle/server/state/server_state.py index dd926aa57..01cf54607 100644 --- a/src/twinkle/server/state/server_state.py +++ b/src/twinkle/server/state/server_state.py @@ -10,7 +10,7 @@ from typing import Any from twinkle.server.config.persistence import PersistenceConfig -from twinkle.server.exceptions import ResourceQuotaExceededError +from twinkle.server.exceptions import ResourceNotFoundError, ResourceQuotaExceededError from twinkle.server.telemetry.correlation import (BASE_MODEL, MODEL_ID, REPLICA_ID, SAMPLING_SESSION_ID, SESSION_ID, TOKEN_ID) from twinkle.server.telemetry.tracing import traced_operation @@ -24,6 +24,7 @@ from .models import FutureFailureRecord, ModelRecord, SamplingSessionRecord, SessionRecord from .sampling_manager import SamplingSessionManager from .session_manager import SessionManager +from .training_store import TrainingSequences, TrainingStore logger = get_logger() @@ -131,6 +132,7 @@ def __init__( self._model_mgr = ModelManager(self._backend, expiration_timeout, per_token_model_limit) self._sampling_mgr = SamplingSessionManager(self._backend, expiration_timeout) self._future_mgr = FutureManager(self._backend, expiration_timeout) + self.training_sequences = TrainingSequences(self._backend) self.expiration_timeout = expiration_timeout self.cleanup_interval = cleanup_interval @@ -173,6 +175,33 @@ async def _on_lose_leader(self) -> None: async def get_capacity_info(self) -> dict[str, int]: return await self._model_mgr.get_capacity_info() + @property + def models(self) -> ModelManager: + return self._model_mgr + + @property + def sessions(self) -> SessionManager: + return self._session_mgr + + @property + def futures(self) -> FutureManager: + return self._future_mgr + + @property + def sampling_sessions(self) -> SamplingSessionManager: + return self._sampling_mgr + + def create_training_store(self, namespace: str) -> TrainingStore: + """Compose a namespaced repository; construction performs no I/O.""" + return TrainingStore( + namespace, + self._backend, + models=self.models, + sessions=self.sessions, + futures=self.futures, + sequences=self.training_sequences, + expiration_timeout=self.expiration_timeout) + # ----- Session Management ----- async def create_session(self, payload: dict[str, Any]) -> str: @@ -213,6 +242,9 @@ async def get_session_last_heartbeat(self, session_id: str) -> float | None: """ return await self._session_mgr.get_last_heartbeat(session_id) + async def get_session_last_heartbeats(self, session_ids: list[str]) -> dict[str, float | None]: + return await self._session_mgr.get_last_heartbeats(session_ids) + # ----- Model Registration ----- async def register_model(self, @@ -429,7 +461,13 @@ async def get_future(self, request_id: str) -> dict[str, Any] | None: record = await self._future_mgr.get(request_id) return record.model_dump() if record is not None else None - async def claim_seq(self, dedup_key: str, request_id: str, ttl: int) -> str | None: + async def assert_future_owner(self, request_id: str, token: str) -> None: + record = await self._future_mgr.get(request_id) + owner = record.owner_token if record is not None else None + if owner is not None and owner != token: + raise ResourceNotFoundError('Unknown future') + + async def claim_seq(self, dedup_key: str, request_id: str, ttl: int | None) -> str | None: """Idempotency claim for a client seq_id. Atomically records ``dedup_key -> request_id`` if unseen and returns ``None`` @@ -446,6 +484,17 @@ async def release_seq(self, dedup_key: str) -> None: preflight rejected it) so a retry can be admitted rather than see a phantom.""" await self._backend.delete(dedup_key) + async def claim_training_seq(self, training_id: str, stream: str, sequence: int, request_id: str, + retention: float) -> str | None: + return await self.training_sequences.claim(training_id, stream, sequence, request_id, retention) + + async def release_training_seq(self, training_id: str, stream: str, sequence: int, request_id: str) -> None: + await self.training_sequences.release(training_id, stream, sequence, request_id) + + async def try_start_future(self, request_id: str) -> bool: + """Acquire task execution atomically, without inferring ownership from reads.""" + return await self._future_mgr.try_start(request_id) + async def cancel_future(self, request_id: str) -> dict[str, Any]: """Best-effort cancel: drop the task iff it has not started running. @@ -470,6 +519,7 @@ async def store_future_status( queue_state_reason: str | None = None, replica_id: str | None = None, absolute_deadline: float | None = None, + owner_token: str | None = None, ) -> None: """Store task status with either a success result or domain failure. @@ -501,6 +551,7 @@ async def store_future_status( queue_state_reason=queue_state_reason, replica_id=replica_id, absolute_deadline=absolute_deadline, + owner_token=owner_token, ) # ----- Resource Cleanup ----- @@ -613,7 +664,7 @@ def get_server_state(cache_key: str = 'twinkle_server_state', # launcher-propagated env (so a non-gateway worker honours the operator's # YAML instead of the hardcoded default), else the module default. from twinkle.server.config.application_spec import ServerStateArgs - env_policy = ServerStateArgs.from_env() + env_policy = ServerStateArgs.from_env() or ServerStateArgs() def _resolve(explicit, env_value, default): if explicit is not None: @@ -622,13 +673,11 @@ def _resolve(explicit, env_value, default): return env_value return default - expiration_timeout = _resolve(expiration_timeout, getattr(env_policy, 'expiration_timeout', None), - _DEFAULT_EXPIRATION_TIMEOUT) - cleanup_interval = _resolve(cleanup_interval, getattr(env_policy, 'cleanup_interval', None), - _DEFAULT_CLEANUP_INTERVAL) - per_token_model_limit = _resolve(per_token_model_limit, getattr(env_policy, 'per_token_model_limit', None), + expiration_timeout = _resolve(expiration_timeout, env_policy.expiration_timeout, _DEFAULT_EXPIRATION_TIMEOUT) + cleanup_interval = _resolve(cleanup_interval, env_policy.cleanup_interval, _DEFAULT_CLEANUP_INTERVAL) + per_token_model_limit = _resolve(per_token_model_limit, env_policy.per_token_model_limit, _DEFAULT_PER_TOKEN_MODEL_LIMIT) - metrics_update_interval = _resolve(metrics_update_interval, getattr(env_policy, 'metrics_update_interval', None), + metrics_update_interval = _resolve(metrics_update_interval, env_policy.metrics_update_interval, _DEFAULT_METRICS_UPDATE_INTERVAL) state = ServerState( diff --git a/src/twinkle/server/state/session_manager.py b/src/twinkle/server/state/session_manager.py index 0e4beeb37..6249ffc9c 100644 --- a/src/twinkle/server/state/session_manager.py +++ b/src/twinkle/server/state/session_manager.py @@ -67,6 +67,13 @@ async def get_last_heartbeat(self, session_id: str) -> float | None: return None return record.last_heartbeat + async def get_last_heartbeats(self, session_ids: list[str]) -> dict[str, float | None]: + records = await self._backend.mget([self._make_key(identity) for identity in session_ids]) + return { + identity: SessionRecord.model_validate(record).last_heartbeat if record is not None else None + for identity, record in zip(session_ids, records, strict=True) + } + # ----- Cleanup ----- def _is_expired(self, record: SessionRecord, cutoff_time: float) -> bool: diff --git a/src/twinkle/server/state/training_records.py b/src/twinkle/server/state/training_records.py new file mode 100644 index 000000000..425955596 --- /dev/null +++ b/src/twinkle/server/state/training_records.py @@ -0,0 +1,81 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Serializable training identities and states shared by controller and model layers.""" + +from __future__ import annotations +from typing import Any, Literal + +from pydantic import BaseModel, Field +from typing_extensions import TypedDict + +from .models import ModelRecord, RuntimeContext + +TrainingState = Literal["PENDING", "ACTIVE", "CLOSING", "RELEASED", "EXPIRED", "FAILED", "CANCELLED"] +InstanceState = Literal["ALLOCATING", "STARTING", "READY", "RESETTING", "DRAINING", "STOPPED", "QUARANTINED"] +TERMINAL_TRAINING = frozenset({"RELEASED", "EXPIRED", "FAILED", "CANCELLED"}) +LIVE_INSTANCES = frozenset({"ALLOCATING", "STARTING", "READY", "RESETTING"}) + + +class TrainingOperation(TypedDict): + """Admission facts retained until actual backend completion is confirmed.""" + + accepted_at: float + renew: bool + + +class RuntimeOwner(TypedDict): + instance_id: str + epoch: str + + +class TrainingRoute(RuntimeOwner): + base_model: str + + +class TrainingRecord(ModelRecord): + training_id: str + controller_epoch: str + instance_id: str + mode: Literal["lora", "full"] + adapter_label: str | None = None + state: TrainingState = "PENDING" + create_request_id: str + creation_deadline: float + last_activity_at: float + pending_operation_ids: dict[str, TrainingOperation] = Field(default_factory=dict) + terminal_reason: str | None = None + terminal_at: float | None = None + + +class InstanceRecord(BaseModel): + instance_id: str + controller_epoch: str + base_model: str + mode: Literal["lora", "full"] + resident: bool = False + state: InstanceState = "ALLOCATING" + training_ids: list[str] = Field(default_factory=list) + startup_deadline: float + application_name: str | None = None + placement_group_id: str | None = None + node_id: str | None = None + devices: list[str] = Field(default_factory=list) + incarnation: str | None = None + dirty: bool = False + failure: str | None = None + stopped_at: float | None = None + + +class CreateTraining(BaseModel): + base_model: str + mode: Literal["lora", "full"] + token: str + session_id: str + create_key: str + client_type: Literal["twinkle", "tinker"] + payload: dict[str, Any] + adapter_label: str | None = None + + +class TrainingRuntimeContext(RuntimeContext): + controller_name: str + epoch: str diff --git a/src/twinkle/server/state/training_store.py b/src/twinkle/server/state/training_store.py new file mode 100644 index 000000000..7e787dbf1 --- /dev/null +++ b/src/twinkle/server/state/training_store.py @@ -0,0 +1,171 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Training persistence and retry receipts, composed from shared state managers.""" + +from __future__ import annotations +import functools +import hashlib +import time +from typing import Any + +import json + +from twinkle.server.exceptions import RequestRejectedError +from .backend.base import StateBackend +from .future_manager import FutureManager +from .model_manager import ModelManager +from .models import FutureFailureRecord +from .session_manager import SessionManager +from .training_records import CreateTraining, InstanceRecord, TrainingRecord + + +class TrainingStore: + def __init__( + self, + namespace: str, + backend: StateBackend, + *, + models: ModelManager, + sessions: SessionManager, + futures: FutureManager, + sequences: TrainingSequences, + expiration_timeout: float, + ): + self.backend = backend + self.models, self.sessions, self.futures, self.sequences = models, sessions, futures, sequences + self.expiration_timeout = expiration_timeout + self.prefix = f"training::{namespace}::" + + def key(self, kind: str, identity: str) -> str: + return f"{self.prefix}{kind}::{identity}" + + @staticmethod + def create_identity(request: CreateTraining) -> tuple[str, str]: + identity = json.dumps([request.client_type, request.token, request.session_id, request.create_key]) + fingerprint = json.dumps(request.model_dump(mode="json"), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(identity.encode()).hexdigest(), hashlib.sha256(fingerprint.encode()).hexdigest() + + async def reserve(self, record: TrainingRecord) -> None: + await self.models.add(record.training_id, record) + + async def save_training(self, record: TrainingRecord) -> None: + await self.models.save_record(record.training_id, record) + + async def retire(self, record: TrainingRecord) -> None: + await self.backend.set( + self.key("terminal", record.training_id), record.model_dump(), ttl=max(1, int(self.expiration_timeout)) + ) + await self.sequences.delete(record.training_id) + # Release the quota row last so a failed earlier step can be idempotently retried. + await self.models.remove(record.training_id) + + async def save_instance(self, record: InstanceRecord) -> None: + ttl = max(1, int(self.expiration_timeout)) if record.state == "STOPPED" else None + await self.backend.set(self.key("instance", record.instance_id), record.model_dump(), ttl=ttl) + + async def instances(self) -> list[InstanceRecord]: + keys = await self.backend.keys(self.key("instance", "*")) + return [InstanceRecord.model_validate(value) for value in await self.backend.mget(keys) if value is not None] + + async def trainings(self) -> list[TrainingRecord]: + return await self.models.get_training_records(self.prefix) + + async def get_create(self, identity: str) -> dict[str, str] | None: + return await self.backend.get(self.key("create", identity)) + + async def save_create(self, identity: str, fingerprint: str, request_id: str, session_id: str) -> None: + await self.backend.set( + self.key("create", identity), + {"fingerprint": fingerprint, "request_id": request_id, "session_id": session_id}, + ) + + async def cleanup_receipts(self) -> None: + for key in await self.backend.keys(self.key("create", "*")): + receipt = await self.backend.get(key) + if receipt and await self.sessions.get_last_heartbeat(receipt["session_id"]) is None: + await self.backend.delete(key) + + async def future( + self, + request_id: str, + status: str, + model_id: str | None, + *, + epoch: str | None = None, + result: Any = None, + failure: FutureFailureRecord | None = None, + deadline: float | None = None, + owner_token: str | None = None, + ) -> None: + await self.futures.store_status( + request_id, + status, + model_id, + result=result, + failure=failure, + replica_id=epoch, + absolute_deadline=deadline, + owner_token=owner_token, + ) + + async def begin_epoch(self, epoch: str) -> None: + await self.backend.set(self.key("controller", "active"), {"epoch": epoch, "started_at": time.time()}) + + +class TrainingSequences: + """Atomic retry receipts and an expired-sequence fence, scoped to a training handle. + + Expiration compacts receipts but retains the sequence floor until that handle is + retired. An old update therefore cannot be applied twice after its receipt ages out. + """ + + def __init__(self, backend: StateBackend): + self.backend = backend + + async def claim( + self, training_id: str, stream: str, sequence: int, request_id: str, retention: float + ) -> str | None: + ledger = await self.backend.update_atomic( + f"seq_ledger::{training_id}::{stream}", + functools.partial( + self.claim_transform, sequence=sequence, request_id=request_id, now=time.time(), retention=retention + ), + ) + receipt = ledger["receipts"].get(str(sequence)) + if receipt is None: + raise RequestRejectedError("Sequence retry window expired; the operation will not be executed again") + return None if receipt["request_id"] == request_id else receipt["request_id"] + + async def release(self, training_id: str, stream: str, sequence: int, request_id: str) -> None: + await self.backend.update_atomic( + f"seq_ledger::{training_id}::{stream}", + functools.partial(self.release_transform, sequence=sequence, request_id=request_id), + ) + + async def delete(self, training_id: str) -> None: + for key in await self.backend.keys(f"seq_ledger::{training_id}::*"): + await self.backend.delete(key) + + @staticmethod + def claim_transform(current, *, sequence, request_id, now, retention): + floor = current.get("floor") if current else None + receipts = {} + for key, receipt in (current or {}).get("receipts", {}).items(): + if receipt["expires_at"] <= now: + floor = int(key) if floor is None else max(floor, int(key)) + else: + receipts[key] = receipt + key = str(sequence) + # Unexpired receipts remain retryable even if another sequence expired first. + if key not in receipts and (floor is None or sequence > floor): + receipts[key] = {"request_id": request_id, "expires_at": now + retention} + return {"floor": floor, "receipts": receipts} + + @staticmethod + def release_transform(current, *, sequence, request_id): + if current is None: + return None + receipts = dict(current["receipts"]) + key = str(sequence) + if receipts.get(key, {}).get("request_id") == request_id: + receipts.pop(key) + return {"floor": current["floor"], "receipts": receipts} diff --git a/src/twinkle/server/task_errors.py b/src/twinkle/server/task_errors.py index cd94d7b70..3c98140eb 100644 --- a/src/twinkle/server/task_errors.py +++ b/src/twinkle/server/task_errors.py @@ -2,10 +2,56 @@ """Helpers for constructing protocol-layer ``ErrorPayload`` values.""" from __future__ import annotations +from dataclasses import dataclass +from functools import lru_cache from typing import Any from twinkle.protocol.types.errors import ErrorCategory, ErrorPayload + +@dataclass(frozen=True) +class FailureResponse: + """How a domain failure is presented inside a response, not its HTTP status.""" + + error_code: int + category: ErrorCategory + + @classmethod + def for_reason(cls, reason: str) -> FailureResponse: + policies = cls._policies() + return policies.get(reason, policies['internal_error']) + + @classmethod + def reasons(cls) -> frozenset[str]: + return frozenset(cls._policies()) + + @classmethod + @lru_cache(maxsize=1) + def _policies(cls) -> dict[str, FailureResponse]: + return { + 'capacity_unavailable': cls(error_code=409, category=ErrorCategory.User), + 'startup_failed': cls(error_code=500, category=ErrorCategory.Server), + 'startup_timeout': cls(error_code=504, category=ErrorCategory.Server), + 'invalid_request': cls(error_code=400, category=ErrorCategory.User), + 'request_rejected': cls(error_code=400, category=ErrorCategory.User), + 'resource_not_found': cls(error_code=404, category=ErrorCategory.User), + 'full_mode_busy': cls(error_code=409, category=ErrorCategory.User), + 'input_tokens_exceeded': cls(error_code=422, category=ErrorCategory.User), + 'batch_size_invalid': cls(error_code=422, category=ErrorCategory.User), + 'rate_limit_exceeded': cls(error_code=429, category=ErrorCategory.User), + 'resource_quota_exceeded': cls(error_code=429, category=ErrorCategory.User), + 'cancelled': cls(error_code=499, category=ErrorCategory.User), + 'endpoint_unavailable': cls(error_code=501, category=ErrorCategory.Server), + 'state_contention': cls(error_code=503, category=ErrorCategory.Server), + 'backend_gate_unavailable': cls(error_code=503, category=ErrorCategory.Server), + 'orphaned_replica': cls(error_code=503, category=ErrorCategory.Server), + 'backend_cancelled': cls(error_code=503, category=ErrorCategory.Server), + 'execution_timeout': cls(error_code=504, category=ErrorCategory.Server), + 'deadline_exceeded': cls(error_code=500, category=ErrorCategory.Server), + 'internal_error': cls(error_code=500, category=ErrorCategory.Server), + } + + _ERROR_MAX = 1024 _TRACEBACK_MAX = 65536 _TRUNCATION_MARKER = '...[traceback truncated, tail kept]...\n' diff --git a/src/twinkle/server/task_queue/__init__.py b/src/twinkle/server/task_queue/__init__.py index bf5391611..7f4805db5 100644 --- a/src/twinkle/server/task_queue/__init__.py +++ b/src/twinkle/server/task_queue/__init__.py @@ -5,22 +5,20 @@ Public exports: - TaskStatus - task lifecycle enum - QueueState - queue state enum for tinker client compatibility -- TaskQueueConfig - queue and rate-limit configuration dataclass - TaskQueueMixin - mixin with schedule_task / schedule_background_task - RateLimiter - sliding-window rate limiter """ -from .config import TaskQueueConfig +from twinkle.protocol.types.errors import QueueState +from twinkle.protocol.types.lifecycle import TaskStatus from .mixin import TaskQueueMixin from .rate_limiter import RateLimiter -from .types import QueuedTask, QueueState, TaskStatus, UserTaskError +from .types import QueuedTask from .worker import ComputeWorker __all__ = [ 'TaskStatus', 'QueueState', - 'UserTaskError', 'QueuedTask', - 'TaskQueueConfig', 'TaskQueueMixin', 'ComputeWorker', 'RateLimiter', diff --git a/src/twinkle/server/task_queue/backend_gate.py b/src/twinkle/server/task_queue/backend_gate.py index 82d17c1e0..2570a1e10 100644 --- a/src/twinkle/server/task_queue/backend_gate.py +++ b/src/twinkle/server/task_queue/backend_gate.py @@ -20,9 +20,13 @@ import functools from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor +from ray.exceptions import GetTimeoutError from typing import Any -from .types import BackendBusyError +from twinkle.server.exceptions import BackendBusyError + +# Backend execution has the same deadline semantics for asyncio and Ray waits. +BACKEND_TIMEOUT_EXCEPTIONS = (TimeoutError, GetTimeoutError) class BackendGate: diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index b06ef8525..78bf0e93d 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -9,12 +9,13 @@ import asyncio import time -import traceback import uuid from collections.abc import Callable, Coroutine from typing import TYPE_CHECKING, Any -from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES, TaskEnvelope +from twinkle.protocol.types.errors import QueueState +from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES, TaskEnvelope, TaskStatus +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.exceptions import BatchSizeError, ConfigError, InputTokensExceededError, RateLimitExceededError from twinkle.server.lifecycle.envelope import envelope_from_record from twinkle.server.lifecycle.poll_config import long_poll_window @@ -22,9 +23,8 @@ from twinkle.server.telemetry.metrics import get_task_metrics from twinkle.utils.logger import get_logger from .backend_gate import BackendGate -from .config import TaskQueueConfig from .rate_limiter import RateLimiter -from .types import QueuedTask, QueueState, TaskStatus +from .types import ExecutionScope, QueuedTask from .worker import ComputeWorker if TYPE_CHECKING: @@ -48,9 +48,8 @@ class TaskQueueMixin: concurrent requests itself, such as vLLM sampling. 3. Background task (schedule_background_task): - asyncio.create_task, runs concurrently with compute queue. Use for I/O - or an engine that provides its own safe concurrency and batching. - Status is still tracked; clients can poll the same status endpoints. + Reuses the supervised concurrent path for pure I/O. Status, capacity + and shutdown are owned by the same worker as concurrent compute. Requirements ------------ @@ -130,6 +129,7 @@ def _init_task_queue( #: Strong references to in-flight fire-and-forget background tasks (see schedule_background_task). self._background_tasks: set = set() self._event_loop: asyncio.AbstractEventLoop | None = None + self._queue_tasks: set[asyncio.Task] = set() @property def task_queue_config(self) -> TaskQueueConfig: @@ -141,6 +141,10 @@ def task_queue_config(self) -> TaskQueueConfig: """ return self._task_queue_config + @property + def sequence_retention_s(self) -> float: + return self.state.expiration_timeout + self._future_absolute_ttl + async def call_backend(self, fn: Callable[..., Any], /, *args: Any, admit: bool = True, **kwargs: Any) -> Any: """Run one backend call outside the event loop. @@ -229,9 +233,33 @@ async def _schedule_task( batch_size_multiple=batch_size_multiple, ) + task = QueuedTask( + request_id=request_id, + coro_factory=coro_factory, + model_id=model_id, + token=token, + input_tokens=input_tokens, + task_type=task_type, + created_at=time.monotonic()) + await self._prepare_queued_task(task) + if self._event_loop is None: self._event_loop = asyncio.get_running_loop() + try: + await self._enqueue_prepared_task(task) + except BaseException: + if task.on_finished is not None and await self.state.get_future(request_id) is None: + await self._compute_worker.finish_task(task, False) + raise + return {'request_id': request_id, 'model_id': model_id} + + async def _prepare_queued_task(self, task: QueuedTask) -> None: + """Optional owner admission, after preflight and before any future write.""" + + async def _enqueue_prepared_task(self, task: QueuedTask) -> None: + request_id, model_id = task.request_id, task.model_id + self._compute_worker.ensure_accepting() await self.state.store_future_status( request_id, TaskStatus.PENDING.value, @@ -239,36 +267,43 @@ async def _schedule_task( queue_state=QueueState.ACTIVE.value, replica_id=getattr(self, 'replica_id', None), absolute_deadline=self._future_deadline(), + owner_token=task.token, ) - queue_key = self._queue_key(model_id=model_id, token=token) + try: + await self._compute_worker.ensure_started() + except BaseException: + try: + await self.state.store_future_status( + request_id, + TaskStatus.FAILED.value, + model_id, + failure=FutureFailureRecord( + reason_code='orphaned_replica', + attribution='server', + message='Replica stopped before the task could be enqueued.')) + finally: + await self._compute_worker.finish_task(task, False) + raise + queue_key = self._queue_key(model_id=model_id, token=task.token) self._compute_worker.ensure_queue_registered(queue_key) - await self._compute_worker.ensure_started() - q = self._compute_worker.get_queue(queue_key) - await q.put( - QueuedTask( - request_id=request_id, - coro_factory=coro_factory, - model_id=model_id, - token=token, - input_tokens=input_tokens, - task_type=task_type, - created_at=time.monotonic(), - )) - await self.state.store_future_status( - request_id, - TaskStatus.QUEUED.value, - model_id, - queue_state=QueueState.ACTIVE.value, - ) - logger.info(f'[TaskQueue] Task {request_id} queued, type={task_type or "unknown"}, ' - f'model_id={model_id}, queue_key={queue_key}, ' - f'queue_depth={q.qsize()}, input_tokens={input_tokens}') - - self._compute_worker.notify_new_task() - - return {'request_id': request_id, 'model_id': model_id} + self._compute_worker.track_task(request_id) + await q.put(task) + try: + await self.state.store_future_status( + request_id, + TaskStatus.QUEUED.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) + logger.info(f'[TaskQueue] Task {request_id} queued, type={task.task_type or "unknown"}, ' + f'model_id={model_id}, queue_key={queue_key}, ' + f'queue_depth={q.qsize()}, input_tokens={task.input_tokens}') + finally: + # Once q.put succeeds the sleeping worker must always be woken, even + # if an acknowledgement or status write fails afterwards. + self._compute_worker.notify_new_task() async def schedule_task( self, @@ -320,6 +355,7 @@ async def schedule_concurrent_task( input_tokens: int = 0, task_type: str | None = None, request_id: str | None = None, + execution_scope: ExecutionScope | None = None, ) -> dict[str, Any]: """Schedule a compute task that runs now instead of in serial order. @@ -344,7 +380,8 @@ async def schedule_concurrent_task( token=token, input_tokens=input_tokens, task_type=task_type, - created_at=time.monotonic()), + created_at=time.monotonic(), + execution_scope=execution_scope), self._queue_key(model_id=model_id, token=token), replica_id=getattr(self, 'replica_id', None), absolute_deadline=self._future_deadline()) @@ -421,73 +458,14 @@ async def schedule_background_task( coro_factory: Callable[[], Coroutine], model_id: str | None = None, task_type: str | None = None, + token: str | None = None, ) -> dict[str, Any]: - """Schedule a fire-and-forget task outside the serial queue. + """Run pure I/O outside serial compute, with the same supervised lifecycle. - The task is launched immediately as an asyncio task. This is suitable - for pure I/O and for an inference engine such as vLLM that performs its - own safe request concurrency and continuous batching. Stateful Model - operations must continue to use :meth:`schedule_task`. - - Status is tracked via state.store_future_status so clients can poll - progress through the same status endpoints as schedule_task(). - - Args: - coro_factory: Zero-argument callable that creates the coroutine. - model_id: Optional model id for result association. - task_type: Label for logging. - - Returns: - {'request_id': str, 'model_id': str | None} + Capacity, timeout, cancellation and shutdown are owned by ComputeWorker; + background jobs must not create a second unsupervised execution path. """ - request_id = f'req_{uuid.uuid4().hex}' - logger.info(f'[TaskQueue] Scheduling background task {request_id}, ' - f'type={task_type or "unknown"}, model_id={model_id}') - - await self.state.store_future_status( - request_id, - TaskStatus.RUNNING.value, - model_id, - queue_state=QueueState.ACTIVE.value, - replica_id=getattr(self, 'replica_id', None), - absolute_deadline=self._future_deadline(), - ) - - async def _run() -> None: - try: - result = await coro_factory() - await self.state.store_future_status( - request_id, - TaskStatus.COMPLETED.value, - model_id, - result=result, - queue_state=QueueState.ACTIVE.value, - ) - logger.info(f'[TaskQueue] Background task {request_id} completed, type={task_type or "unknown"}') - except Exception as exc: - failure = FutureFailureRecord( - reason_code='internal_error', - message=f'{type(exc).__name__}: {exc}'[:1024], - attribution='server', - diagnostic=traceback.format_exc(), - ) - await self.state.store_future_status( - request_id, - TaskStatus.FAILED.value, - model_id, - failure=failure, - queue_state=QueueState.ACTIVE.value, - ) - logger.error(f'[TaskQueue] Background task {request_id} FAILED, type={task_type or "unknown"}:\n' - f'{traceback.format_exc(limit=3)}') - - task = asyncio.create_task(_run()) - # Hold a strong reference: the event loop keeps only a weak one, so an unreferenced task - # can be garbage-collected mid-run -- its state would stay RUNNING forever with no FAILED - # record to explain it (a minutes-long hub upload is the main victim). - self._background_tasks.add(task) - task.add_done_callback(self._background_tasks.discard) - return {'request_id': request_id, 'model_id': model_id} + return await self.schedule_concurrent_task(coro_factory, model_id=model_id, task_type=task_type, token=token) async def _fail_queue_tasks_async(self, queue_key: str, reason: str) -> None: await self._compute_worker.fail_queue_tasks(queue_key, reason) @@ -500,13 +478,23 @@ def fail_pending_tasks_for_model(self, model_id: str, reason: str) -> None: return def _schedule() -> None: - asyncio.create_task(self._fail_queue_tasks_async(queue_key, reason)) + task = asyncio.create_task(self._fail_queue_tasks_async(queue_key, reason)) + self._queue_tasks.add(task) + task.add_done_callback(self._queue_task_done) self._event_loop.call_soon_threadsafe(_schedule) + def _queue_task_done(self, task: asyncio.Task) -> None: + self._queue_tasks.discard(task) + if not task.cancelled(): + error = task.exception() + if error is not None: + logger.error('Queue maintenance failed', exc_info=(type(error), error, error.__traceback__)) + async def shutdown_task_queue(self) -> None: """Gracefully shut down the compute queue and release resources.""" await self._rate_limiter.stop_cleanup_task() await self._compute_worker.stop() + await asyncio.gather(*self._queue_tasks, return_exceptions=True) self._backend_gate.shutdown() logger.debug('[TaskQueue] Task queue shutdown complete') diff --git a/src/twinkle/server/task_queue/types.py b/src/twinkle/server/task_queue/types.py index 936a433a3..abfb23809 100644 --- a/src/twinkle/server/task_queue/types.py +++ b/src/twinkle/server/task_queue/types.py @@ -1,61 +1,32 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -""" -Core type definitions for the task queue system. +"""In-process execution records, never HTTP/Ray wire models. -Provides: -- TaskStatus: Enum for tracking task lifecycle states -- QueueState: Enum for tinker client compatibility (retry behavior hints) -- QueuedTask: Dataclass representing a queued work item +Coroutine factories, asyncio tasks, and callbacks belong here. Public task and +queue state enums live in twinkle.protocol.types and are shared by SDK/server. """ from __future__ import annotations import asyncio from collections.abc import Callable, Coroutine -from dataclasses import dataclass -from enum import Enum - -# Ray_Get_Timeout is classified the same as asyncio.TimeoutError: 504/Server. -try: - from ray.exceptions import GetTimeoutError as _RayGetTimeout - BACKEND_TIMEOUT_EXCEPTIONS: tuple[type[BaseException], ...] = (asyncio.TimeoutError, _RayGetTimeout) -except Exception: # pragma: no cover - ray always present in server runtime - BACKEND_TIMEOUT_EXCEPTIONS = (asyncio.TimeoutError, ) - - -class TaskStatus(Enum): - """Task lifecycle status.""" - PENDING = 'pending' # Task created, waiting to be processed - QUEUED = 'queued' # Task in queue waiting for execution - RUNNING = 'running' # Task currently executing - COMPLETED = 'completed' # Task completed successfully - FAILED = 'failed' # Task failed with error - CANCELLED = 'cancelled' # Task cancelled by the client before it started running - +from dataclasses import dataclass, field -class UserTaskError(ValueError): - """A queued operation rejected because of caller input or usage.""" - - -class BackendBusyError(RuntimeError): - """Raised when the per-replica Admission_Gate is held by a leaked backend call. - - A new backend call arriving while the gate is closed (its holder is a call that - already exceeded ``asyncio.wait_for`` but whose executor thread has not yet - returned) fails fast with this error instead of queueing behind it. The worker - maps it to ``ErrorPayload(category='server', error_code=503)``. - """ +@dataclass +class ExecutionScope: + """Completion witnesses for an operation, independent of its client future.""" + jobs: set[asyncio.Task] = field(default_factory=set) + request_ids: set[str] = field(default_factory=set) -class QueueState(Enum): - """Queue state for tinker client compatibility. + def track(self, job: asyncio.Task, request_id: str | None = None) -> None: + self.jobs.add(job) + if request_id is not None: + self.request_ids.add(request_id) - These states are returned to the tinker client to indicate the current - state of the task queue and help the client adjust its retry behavior. - """ - ACTIVE = 'active' # Queue is actively processing tasks - PAUSED_RATE_LIMIT = 'paused_rate_limit' # Queue paused due to rate limiting - PAUSED_CAPACITY = 'paused_capacity' # Queue paused due to capacity limits - UNKNOWN = 'unknown' # Unknown or unspecified state + async def wait(self): + observed = set() + while pending := self.jobs - observed: + await asyncio.gather(*pending, return_exceptions=True) + observed.update(pending) @dataclass @@ -68,3 +39,5 @@ class QueuedTask: input_tokens: int task_type: str | None created_at: float + on_finished: Callable[[bool, bool], Coroutine] | None = None + execution_scope: ExecutionScope | None = None diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index 93d4e165f..a19cac88e 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -13,20 +13,23 @@ import traceback from collections import deque from dataclasses import replace +from ray.exceptions import RayTaskError from typing import TYPE_CHECKING, Any, Callable, Deque -from twinkle.protocol.types.errors import ErrorCategory -from twinkle.server.exceptions import (BatchSizeError, EndpointUnavailableError, FullModeBusyError, +from twinkle.protocol.types.errors import ErrorCategory, QueueState +from twinkle.protocol.types.lifecycle import TaskStatus +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.exceptions import (BackendBusyError, BatchSizeError, EndpointUnavailableError, FullModeBusyError, InputTokensExceededError, RateLimitExceededError, RequestRejectedError, ResourceExhaustedError, ResourceNotFoundError, ResourceQuotaExceededError, StateBackendError, TrainModeMismatchError, TwinkleServerError) from twinkle.server.state.backend.base import ConcurrencyError from twinkle.server.state.models import FutureFailureRecord +from twinkle.server.task_queue.backend_gate import BACKEND_TIMEOUT_EXCEPTIONS from twinkle.server.telemetry.correlation import MODEL_ID, TOKEN_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.utils.logger import get_logger -from .config import TaskQueueConfig -from .types import BACKEND_TIMEOUT_EXCEPTIONS, BackendBusyError, QueuedTask, QueueState, TaskStatus, UserTaskError +from .types import QueuedTask if TYPE_CHECKING: from twinkle.server.state import ServerState @@ -35,6 +38,27 @@ logger = get_logger() +def _exception_summary(exc: BaseException) -> str: + """Build a single-line summary, unwrapping Ray's transport wrapper. + + A RayTaskError's own ``str`` starts with the Ray call header (actor id, + pid, repr); the actionable error (e.g. the loss assertion) is in ``cause``. + Preferring it keeps the summary useful without enlarging the wire field, + while the full traceback still travels separately. + + Only Ray's wrapper is unwrapped: elsewhere ``raise X from e`` is a + deliberate message and must not be replaced by its inner cause. + """ + cause = None + if isinstance(exc, RayTaskError): + cause = getattr(exc, 'cause', None) + if not isinstance(cause, BaseException): + cause = exc.__cause__ + if isinstance(cause, BaseException): + return f'{type(cause).__name__}: {cause}' + return f'{type(exc).__name__}: {exc}' + + def _reason_code_for_server_error(exc: TwinkleServerError) -> str: """Map typed execution failures to protocol-independent reason codes.""" mappings: tuple[tuple[type[TwinkleServerError], str], ...] = ( @@ -85,6 +109,8 @@ def __init__( self._on_backend_timeout = on_backend_timeout self.task_queues: dict[str, asyncio.Queue] = {} + self._request_ids: set[str] = set() + self._task_outcomes: dict[str, dict[str, Any]] = {} self.queue_order: Deque[str] = deque() self.new_task_event: asyncio.Event = asyncio.Event() @@ -97,14 +123,19 @@ def __init__( self._detached_tasks: dict[asyncio.Task, QueuedTask] = {} self._detached_reserved = 0 + def ensure_accepting(self) -> None: + if self._stopping: + raise ResourceExhaustedError('The replica is shutting down.') + async def ensure_started(self) -> None: - """Ensure the background worker coroutine is running.""" + """Ensure the background worker coroutine is running; never reopen a stopped replica.""" + self.ensure_accepting() if self._started and self._worker_task is not None and not self._worker_task.done(): return async with self._start_lock: + self.ensure_accepting() if self._started and self._worker_task is not None and not self._worker_task.done(): return - self._stopping = False self._worker_task = asyncio.create_task(self._worker_loop()) self._started = True @@ -128,10 +159,11 @@ async def stop(self) -> None: 'Sampler task stopped during replica shutdown.', QueueState.ACTIVE.value, reason_code='orphaned_replica') + for queue_key in list(self.task_queues): + await self.fail_queue_tasks(queue_key, 'Replica stopped before the task could execute.') self._worker_task = None self._started = False - self.task_queues.clear() - self.queue_order.clear() + self._request_ids.clear() def ensure_queue_registered(self, queue_key: str) -> None: """Register a new per-key queue if it does not yet exist.""" @@ -149,6 +181,40 @@ def get_queue(self, queue_key: str) -> asyncio.Queue: """ return self.task_queues[queue_key] + def track_task(self, request_id: str) -> None: + self._request_ids.add(request_id) + + def has_task(self, request_id: str) -> bool: + """True while a task is queued or executing, including a terminal-future race.""" + return request_id in self._request_ids + + def get_task_outcome(self, request_id: str) -> dict[str, Any] | None: + """Read bounded local execution facts, never infer them from a future status.""" + cutoff = time.time() - self._config.absolute_future_ttl(1) + for identity, outcome in list(self._task_outcomes.items()): + if outcome['finished_at'] <= cutoff: + self._task_outcomes.pop(identity) + outcome = self._task_outcomes.get(request_id) + return dict(outcome) if outcome is not None else None + + async def finish_task(self, + task: QueuedTask, + executed: bool, + unsafe: bool = False, + finished_at: float | None = None) -> None: + if task.on_finished is not None: + self.get_task_outcome(task.request_id) + self._task_outcomes[task.request_id] = { + 'state': 'unsafe' if unsafe else 'finished', + 'executed': executed, + 'finished_at': time.time() if finished_at is None else finished_at, + } + self._request_ids.discard(task.request_id) + if task.on_finished is not None: + acknowledged = await task.on_finished(executed, unsafe) + if acknowledged is True: + self._task_outcomes.pop(task.request_id, None) + def total_queued(self) -> int: """Total number of pending tasks across all per-key queues.""" return sum(q.qsize() for q in self.task_queues.values()) @@ -246,6 +312,7 @@ async def fail_queue_tasks(self, queue_key: str, reason: str) -> None: for task in drained: await self._store_task_failed(task, reason, QueueState.UNKNOWN.value, queue_state_reason=reason) q.task_done() + await self.finish_task(task, False) self.task_queues.pop(queue_key, None) try: @@ -272,6 +339,7 @@ async def _fail_timed_out_task(self, task: QueuedTask, waited: float, q: asyncio await self._store_task_failed(task, error, QueueState.PAUSED_CAPACITY.value, queue_state_reason=error) self._record_tasks_total(task.task_type or 'unknown', 'timeout') q.task_done() + await self.finish_task(task, False) async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue) -> None: """Execute a dequeued task, then release its queue slot.""" @@ -292,7 +360,8 @@ async def submit_detached(self, task: QueuedTask, queue_key: str, *, replica_id: task.model_id, queue_state=QueueState.ACTIVE.value, replica_id=replica_id, - absolute_deadline=absolute_deadline) + absolute_deadline=absolute_deadline, + owner_token=task.token) wrote_status = True self.run_detached(task, queue_key) except BaseException: @@ -307,7 +376,8 @@ async def submit_detached(self, task: QueuedTask, queue_key: str, *, replica_id: def reserve_detached(self) -> None: """Reserve before the future write, so capacity rejection leaves no record.""" - if self._stopping or self._detached_reserved >= self._config.concurrent_task_limit: + self.ensure_accepting() + if self._detached_reserved >= self._config.concurrent_task_limit: raise ResourceExhaustedError('Concurrent compute capacity is unavailable on this replica.') self._detached_reserved += 1 @@ -353,10 +423,17 @@ async def run(): running = asyncio.create_task(run()) self._detached_tasks[running] = task + scope = task.execution_scope + if scope is not None: + scope.track(running, task.request_id) def done(completed: asyncio.Task) -> None: self._detached_tasks.pop(completed, None) self.release_detached() + if not completed.cancelled(): + error = completed.exception() + if error is not None: + logger.error('Detached task failed', exc_info=(type(error), error, error.__traceback__)) running.add_done_callback(done) @@ -369,6 +446,8 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: exec_start = time.monotonic() task_status = 'completed' exec_time = 0.0 + executed = False + finished_at = None # One span per queued task execution, tagged `.` # (e.g. `model.forward`, `sampler.sample`). queue_attrs = { @@ -381,28 +460,20 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: handler_attrs = {TOKEN_ID: task.token, MODEL_ID: task.model_id} handler_span_name = f'{self._deployment_name or "deployment"}.{task_type}' try: - await self._state.store_future_status( - task.request_id, - TaskStatus.RUNNING.value, - task.model_id, - queue_state=QueueState.ACTIVE.value, - ) - # A client cancel (or orphan cleanup) can win the race against the - # running write; the terminal record is kept, so do not start work - # nobody will collect. Once running is stored, cancel_if_pending can - # no longer succeed, so this read closes the window. - if await self._is_record_terminal(task.request_id): + if not await self._state.try_start_future(task.request_id): task_status = 'cancelled' - logger.info(f'[ComputeWorker] Task {task.request_id} became terminal before start; skipping.') + logger.info(f'[ComputeWorker] Task {task.request_id} was not acquired; skipping.') return with traced_operation('task_queue.execute', attrs=queue_attrs): logger.debug(f'[ComputeWorker] Task {task.request_id} executing, ' f'type={task_type}, queue_key={queue_key}') with traced_operation(handler_span_name, attrs=handler_attrs): + executed = True coro = task.coro_factory() # effective_execution_timeout is always positive (0 -> finite fallback), # so wait_for is always in effect. result = await asyncio.wait_for(coro, timeout=self._config.effective_execution_timeout) + finished_at = time.time() exec_time = time.monotonic() - exec_start logger.info(f'[ComputeWorker] Task {task.request_id} completed in {exec_time:.2f}s, type={task_type}') await self._state.store_future_status( @@ -438,16 +509,6 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: await self._on_backend_timeout() except Exception: logger.error(f'[ComputeWorker] backend-timeout probe failed:\n{traceback.format_exc(limit=3)}') - except UserTaskError as exc: - task_status = 'failed' - exec_time = time.monotonic() - exec_start - await self._store_task_failed( - task, - f'{type(exc).__name__}: {exc}', - QueueState.UNKNOWN.value, - reason_code='request_rejected', - attribution='user', - ) except BackendBusyError as exc: task_status = 'failed' exec_time = time.monotonic() - exec_start @@ -481,7 +542,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: exec_time = time.monotonic() - exec_start # error is a single-line summary; the full traceback goes only to the # traceback field, never into `error`. - error = f'{type(exc).__name__}: {exc}' + error = _exception_summary(exc) logger.error(f'[ComputeWorker] Task {task.request_id} FAILED after {exec_time:.2f}s, ' f'type={task_type}:\n{traceback.format_exc(limit=3)}') await self._store_task_failed( @@ -492,6 +553,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: traceback_text=traceback.format_exc(), ) finally: + await self.finish_task(task, executed, executed and task_status in ('timeout', 'cancelled'), finished_at) self._record_execution_time(task_type, exec_time) self._record_tasks_total(task_type, task_status) @@ -524,13 +586,6 @@ async def _try_run_one(self) -> bool: await self._fail_timed_out_task(task, queue_wait, q) continue # try the next queue - # A record already in a Terminal_State (e.g. written 'failed' by the - # state-hygiene orphan handling) must not be executed again. - if await self._is_record_terminal(task.request_id): - logger.info(f'[ComputeWorker] Task {task.request_id} already terminal on dequeue; skipping.') - q.task_done() - continue - # Execute the task (serial: stops after the first execution) await self._execute_task(task, queue_key, q) return True @@ -539,11 +594,8 @@ async def _try_run_one(self) -> bool: async def _is_record_terminal(self, request_id: str) -> bool: """True if the future record already holds a Terminal_State.""" - try: - record = await self._state.get_future(request_id) - except Exception: - return False - if not record: + record = await self._state.get_future(request_id) + if record is None: return False return record.get('status') in (TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, TaskStatus.CANCELLED.value) diff --git a/src/twinkle/server/telemetry/metrics.py b/src/twinkle/server/telemetry/metrics.py index a182d57b4..407737cae 100644 --- a/src/twinkle/server/telemetry/metrics.py +++ b/src/twinkle/server/telemetry/metrics.py @@ -17,9 +17,11 @@ from __future__ import annotations +from dataclasses import dataclass from pydantic import BaseModel, ConfigDict from typing import Any +from twinkle.server.exceptions import ResourceReleaseError from .provider import get_meter try: @@ -111,6 +113,33 @@ class TaskMetrics(BaseModel): rate_limiter_active_tokens: _Gauge +class TrainingMetrics(BaseModel): + """Low-cardinality metrics for the managed training lifecycle.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + instances_active: _Gauge + startup_seconds: _Histogram + shutdown_seconds: _Histogram + allocation_seconds: _Histogram + removal_seconds: _Histogram + quarantined_total: _Counter + sweep_total: _Counter + sweep_seconds: _Histogram + devices_allocated: _Gauge + placement_groups_active: _Gauge + + +class StateBackendMetrics(BaseModel): + """Metrics for persistent state backend availability and contention.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + redis_health: _Gauge + update_atomic_conflicts: _Counter + update_atomic_retries: _Histogram + + class _RequestMetrics(BaseModel): """HTTP request metrics container (internal).""" @@ -177,6 +206,46 @@ def __init__(self) -> None: description='Number of tokens currently tracked by the rate limiter', ) + # === Managed training === + self.training_instances_active = meter.create_up_down_counter( + 'twinkle.training.instances.active', description='Active managed training instances by lifecycle state') + self.training_startup_seconds = meter.create_histogram( + 'twinkle.training.instance.startup.duration_seconds', + description='Training runtime startup duration', + unit='s') + self.training_shutdown_seconds = meter.create_histogram( + 'twinkle.training.instance.shutdown.duration_seconds', + description='Training runtime shutdown duration', + unit='s') + self.training_allocation_seconds = meter.create_histogram( + 'twinkle.training.placement_group.allocation.duration_seconds', + description='Placement group allocation duration', + unit='s') + self.training_removal_seconds = meter.create_histogram( + 'twinkle.training.placement_group.removal.duration_seconds', + description='Placement group removal duration', + unit='s') + self.training_quarantined_total = meter.create_counter( + 'twinkle.training.quarantined.total', + description='Transitions to quarantine during cleanup; retries in quarantine are excluded') + self.training_sweep_total = meter.create_counter( + 'twinkle.training.sweep.total', description='Training lifecycle sweep runs by status') + self.training_sweep_seconds = meter.create_histogram( + 'twinkle.training.sweep.duration_seconds', description='Training lifecycle sweep duration', unit='s') + self.training_devices_allocated = meter.create_up_down_counter( + 'twinkle.training.devices.allocated', description='Logical devices owned by managed training instances') + self.training_placement_groups_active = meter.create_up_down_counter( + 'twinkle.training.placement_groups.active', + description='Placement groups owned by managed training instances') + + # === Redis state backend === + self.redis_health = meter.create_up_down_counter( + 'twinkle.state.redis.health', description='Last observed Redis health (1 healthy, 0 unhealthy)') + self.redis_update_atomic_conflicts = meter.create_counter( + 'twinkle.state.redis.update_atomic.conflicts.total', description='Redis WATCH conflicts') + self.redis_update_atomic_retries = meter.create_histogram( + 'twinkle.state.redis.update_atomic.retries', description='Retries used by Redis atomic updates') + # === Resources (ObservableGauge backed by the push cache) === # OTEL holds its own references to the gauges via the meter, so we # only need to keep the cache the callbacks read from. @@ -192,6 +261,8 @@ def __init__(self) -> None: # swaps this singleton to rebind instruments) invalidates them. self._task_metrics: dict[str, TaskMetrics] = {} self._request_metrics: dict[str, _RequestMetrics] = {} + self._training_metrics: TrainingMetrics | None = None + self._state_backend_metrics: StateBackendMetrics | None = None def _make_gauge_callback(self, name: str): """Build the sync OTEL callback that reads ``_resource_cache[name]``.""" @@ -227,6 +298,33 @@ def request_metrics(self, deployment: str) -> _RequestMetrics: ) return cached + def training_metrics(self) -> TrainingMetrics: + """Return process-local managed training metric adapters.""" + if self._training_metrics is None: + self._training_metrics = TrainingMetrics( + instances_active=_Gauge(self.training_instances_active), + startup_seconds=_Histogram(self.training_startup_seconds), + shutdown_seconds=_Histogram(self.training_shutdown_seconds), + allocation_seconds=_Histogram(self.training_allocation_seconds), + removal_seconds=_Histogram(self.training_removal_seconds), + quarantined_total=_Counter(self.training_quarantined_total), + sweep_total=_Counter(self.training_sweep_total), + sweep_seconds=_Histogram(self.training_sweep_seconds), + devices_allocated=_Gauge(self.training_devices_allocated), + placement_groups_active=_Gauge(self.training_placement_groups_active), + ) + return self._training_metrics + + def state_backend_metrics(self) -> StateBackendMetrics: + """Return process-local state backend metric adapters.""" + if self._state_backend_metrics is None: + self._state_backend_metrics = StateBackendMetrics( + redis_health=_Gauge(self.redis_health), + update_atomic_conflicts=_Counter(self.redis_update_atomic_conflicts), + update_atomic_retries=_Histogram(self.redis_update_atomic_retries), + ) + return self._state_backend_metrics + # ----- Push API for the cleanup leader ----- def set_resource_count(self, name: str, value: int) -> None: @@ -268,3 +366,62 @@ def get_task_metrics(deployment: str) -> TaskMetrics: def get_request_metrics(deployment: str) -> _RequestMetrics: """Return the per-deployment HTTP request metric adapters.""" return MetricsRegistry.get().request_metrics(deployment) + + +def get_training_metrics() -> TrainingMetrics: + """Return managed training metric adapters.""" + return MetricsRegistry.get().training_metrics() + + +def get_state_backend_metrics() -> StateBackendMetrics: + """Return persistent state backend metric adapters.""" + return MetricsRegistry.get().state_backend_metrics() + + +@dataclass(frozen=True) +class TrainingInstanceSnapshot: + """Known resource ownership, including draining/quarantined instances.""" + + backend: str + mode: str + state: str + resident: str + device_type: str + device_count: int + placement_group_count: int + + +class TrainingInstanceMetrics: + """Aggregate snapshots and clear retired label sets; never own scheduling state.""" + + def __init__(self, metrics: TrainingMetrics): + self.metrics = metrics + self._instance_keys = set() + self._resource_keys = set() + + def report(self, snapshots): + counts, resources = {}, {} + for snapshot in snapshots: + key = snapshot.backend, snapshot.mode, snapshot.state, snapshot.resident + counts[key] = counts.get(key, 0) + 1 + resource_key = snapshot.backend, snapshot.mode, snapshot.device_type + devices, groups = resources.get(resource_key, (0, 0)) + resources[resource_key] = devices + snapshot.device_count, groups + snapshot.placement_group_count + for key in self._instance_keys | counts.keys(): + tags = dict(zip(('backend', 'mode', 'state', 'resident'), key)) + self.metrics.instances_active.set(counts.get(key, 0), tags=tags) + for key in self._resource_keys | resources.keys(): + tags = dict(zip(('backend', 'mode', 'device_type'), key)) + devices, groups = resources.get(key, (0, 0)) + self.metrics.devices_allocated.set(devices, tags=tags) + self.metrics.placement_groups_active.set(groups, tags=tags) + self._instance_keys = set(counts) + self._resource_keys = set(resources) + + @staticmethod + def quarantine_reason(exc: Exception) -> str: + if isinstance(exc, TimeoutError): + return 'timeout' + if isinstance(exc, ResourceReleaseError): + return exc.phase + return 'other' diff --git a/src/twinkle/utils/parallel.py b/src/twinkle/utils/parallel.py index 26f8c7c15..717654bad 100644 --- a/src/twinkle/utils/parallel.py +++ b/src/twinkle/utils/parallel.py @@ -1,10 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from __future__ import annotations +import getpass import hashlib import inspect import os import re +import tempfile from collections import defaultdict from contextlib import contextmanager from datetime import timedelta @@ -25,8 +27,14 @@ def __getattr__(name: str): return FileLock raise AttributeError(f'module {__name__!r} has no attribute {name!r}') -_LOCK_DIR = '.locks' -os.makedirs(_LOCK_DIR, exist_ok=True) +# Workers may start in read-only or different working directories. Keep their +# locks in one writable runtime directory, without creating it during import. +# TWINKLE_HOME overrides the default per-user temporary directory. +_LOCK_DIR = os.path.abspath( + os.path.join( + os.path.expanduser( + os.environ.get('TWINKLE_HOME') or os.path.join(tempfile.gettempdir(), 'twinkle', getpass.getuser())), + '.locks')) # Coordination store: derived from MASTER_PORT so every rank of a group agrees on it without # talking to each other first. A collision with another job's MASTER_PORT surfaces as a bind @@ -314,6 +322,7 @@ def processing_lock(lock_file: str, sticky: bool = False): # actually runs, and importing it at module scope made ``import twinkle.loss`` require the # whole HF datasets stack even for losses that never touch it. from datasets.utils.filelock import FileLock + os.makedirs(_LOCK_DIR, exist_ok=True) lock: FileLock = FileLock(os.path.join(_LOCK_DIR, f'{lock_name}.lock')) # noqa if acquire_lock(lock, False): diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index bfd00a104..02aa08726 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -61,6 +61,72 @@ def pad_sequence_to_length( return F.pad(tensor, pad_tuple, mode='constant', value=pad_value) +_CHUNKED_SELECTIVE_LOG_SOFTMAX = None + + +def _chunked_selective_log_softmax_function(): + """Lazily build the autograd function used by :func:`selective_log_softmax`. + + ``torch`` is imported lazily throughout this module, so the ``Function`` + subclass cannot live at module scope. It is built once and cached. + """ + global _CHUNKED_SELECTIVE_LOG_SOFTMAX + if _CHUNKED_SELECTIVE_LOG_SOFTMAX is not None: + return _CHUNKED_SELECTIVE_LOG_SOFTMAX + import torch + + class ChunkedSelectiveLogSoftmax(torch.autograd.Function): + """One chunk of ``selective_log_softmax`` that recomputes in backward. + + The wide FP32 upcast is a forward-only temporary: only the original + dtype chunk and its labels are saved, so no full-vocab FP32 copy of the + activation stays resident until the backward pass. + """ + + @staticmethod + def forward(ctx, logits, labels, compute_dtype, return_entropy): + chunk = logits.to(compute_dtype) + logsumexp = torch.logsumexp(chunk, dim=-1) + selected = torch.gather(chunk, dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) + logps = selected - logsumexp + ctx.save_for_backward(logits, labels) + ctx.compute_dtype = compute_dtype + ctx.return_entropy = return_entropy + if not return_entropy: + return logps + probabilities = torch.softmax(chunk, dim=-1) + return logps, logsumexp - (probabilities * chunk).sum(dim=-1) + + @staticmethod + def backward(ctx, *grad_outputs): + logits, labels = ctx.saved_tensors + compute_dtype = ctx.compute_dtype + chunk = logits.to(compute_dtype) + probabilities = torch.softmax(chunk, dim=-1) + grad_logits = torch.zeros_like(chunk) + + grad_logps = grad_outputs[0] + if grad_logps is not None: + grad_logps = grad_logps.to(compute_dtype).unsqueeze(-1) + # d/dz (z[label] - logsumexp(z)) = onehot(label) - softmax(z) + grad_logits.scatter_add_(dim=-1, index=labels.unsqueeze(-1), src=grad_logps) + grad_logits.sub_(probabilities * grad_logps) + + if ctx.return_entropy: + grad_entropy = grad_outputs[1] + if grad_entropy is not None: + grad_entropy = grad_entropy.to(compute_dtype).unsqueeze(-1) + # H(z) = logsumexp(z) - sum(softmax(z) * z) + # dH/dz = softmax(z) * (E_p[z] - z) + expected = (probabilities * chunk).sum(dim=-1, keepdim=True) + grad_logits.add_(probabilities * (grad_entropy * (expected - chunk))) + + return grad_logits.to(logits.dtype), None, None, None + + _CHUNKED_SELECTIVE_LOG_SOFTMAX = ChunkedSelectiveLogSoftmax + return _CHUNKED_SELECTIVE_LOG_SOFTMAX + + def selective_log_softmax(logits, index, return_entropy: bool = False): """ refer: trl/trainer/utils @@ -83,11 +149,13 @@ def selective_log_softmax(logits, index, return_entropy: bool = False): site avoids materializing a second full-vocab tensor later). Returns: - Gathered log probabilities with the same shape as ``index``. + Gathered log probabilities with the same shape as ``index``. The dtype is + float32 (float64 for float64 input) whatever the input dtype: the softmax is + computed wide and the result is not narrowed back, because a peaked row + reaches log-probs whose bfloat16 ULP is 0.5. If ``return_entropy`` is True, returns ``(per_token_logps, per_token_entropy)``. """ import torch - import torch.nn.functional as F try: from megatron.core import parallel_state as mpu @@ -102,40 +170,34 @@ def selective_log_softmax(logits, index, return_entropy: bool = False): except (ImportError, AssertionError, OSError): pass - if logits.dtype in [torch.float32, torch.float64]: - selected_logits = torch.gather(logits, dim=-1, index=index.unsqueeze(-1)).squeeze(-1) - if return_entropy: - # Per-row loop mirrors the logsumexp path below, to keep peak memory bounded. - logsumexp_values = [] - per_token_entropy = [] - for row_logits in logits: - row_lse = torch.logsumexp(row_logits, dim=-1) - logsumexp_values.append(row_lse) - # H = lse - E_p[x] = lse - sum(exp(x - lse) * x) - row_p = torch.exp(row_logits - row_lse.unsqueeze(-1)) - per_token_entropy.append(row_lse - (row_p * row_logits).sum(dim=-1)) - logsumexp_values = torch.stack(logsumexp_values) - per_token_entropy = torch.stack(per_token_entropy) - per_token_logps = selected_logits - logsumexp_values - return per_token_logps, per_token_entropy - # loop to reduce peak mem consumption - logsumexp_values = torch.stack([torch.logsumexp(lg, dim=-1) for lg in logits]) - per_token_logps = selected_logits - logsumexp_values # log_softmax(x_i) = x_i - logsumexp(x) - else: - # logsumexp approach is unstable with bfloat16, fall back to slightly less efficient approach - per_token_logps = [] - per_token_entropy = [] if return_entropy else None - for row_logits, row_labels in zip(logits, index, strict=True): # loop to reduce peak mem consumption - row_logps = F.log_softmax(row_logits, dim=-1) - row_per_token_logps = row_logps.gather(dim=-1, index=row_labels.unsqueeze(-1)).squeeze(-1) - per_token_logps.append(row_per_token_logps) - if return_entropy: - # row_logps is already stable; softmax reuses the same numerics. - row_p = torch.exp(row_logps) - per_token_entropy.append(-(row_p * row_logps).sum(dim=-1)) - per_token_logps = torch.stack(per_token_logps) + vocab_size = logits.shape[-1] + flat_logits = logits.reshape(-1, vocab_size) + flat_index = index.reshape(-1) + compute_dtype = torch.float64 if logits.dtype == torch.float64 else torch.float32 + element_size = torch.empty((), dtype=compute_dtype).element_size() + temporary_bytes = 32 << 20 if return_entropy else 64 << 20 + chunk_rows = max(1, temporary_bytes // (vocab_size * element_size)) + chunk_function = _chunked_selective_log_softmax_function() + + logps = [] + entropies = [] + for start in range(0, flat_logits.shape[0], chunk_rows): + stop = min(start + chunk_rows, flat_logits.shape[0]) + result = chunk_function.apply( + flat_logits[start:stop], + flat_index[start:stop], + compute_dtype, + return_entropy, + ) if return_entropy: - return per_token_logps, torch.stack(per_token_entropy) + chunk_logps, chunk_entropies = result + logps.append(chunk_logps) + entropies.append(chunk_entropies) + else: + logps.append(result) + per_token_logps = torch.cat(logps).reshape(index.shape) + if return_entropy: + return per_token_logps, torch.cat(entropies).reshape(index.shape) return per_token_logps diff --git a/src/twinkle_client/_future.py b/src/twinkle_client/_future.py index b8fc5a3b4..59570f0c4 100644 --- a/src/twinkle_client/_future.py +++ b/src/twinkle_client/_future.py @@ -99,6 +99,7 @@ def _unwrap(env: TaskEnvelope, model_cls) -> Any: request_id=env.request_id, error_code=p.error_code, details=p.details, + traceback=p.traceback, ) if env.status == 'cancelled': p = env.error diff --git a/src/twinkle_client/auto/agent/server_tools.py b/src/twinkle_client/auto/agent/server_tools.py index 09786c274..2324dc547 100644 --- a/src/twinkle_client/auto/agent/server_tools.py +++ b/src/twinkle_client/auto/agent/server_tools.py @@ -368,6 +368,8 @@ def _generate_server_config( sampler_list = samplers or [] if backend not in ('transformers', 'megatron'): raise ValueError(f'Unsupported training backend: {backend!r}') + if backend == 'megatron' and train_gpus < 2: + raise ValueError('megatron training instances require at least 2 devices') # Sanitize model name for use in route/names def _short(mid: str) -> str: @@ -428,6 +430,9 @@ def _short(mid: str) -> str: 'args': { 'backend': backend, 'model_id': f'ms://{model_id}', + 'min_replicas': 0, + 'max_replicas': 1, + 'idle_timeout_s': 600, 'max_length': 500000, # total tokens per forward pass (must match max_input_tokens) 'nproc_per_node': train_gpus, 'device_group': { diff --git a/src/twinkle_client/exceptions.py b/src/twinkle_client/exceptions.py index dd56d6e35..32f59990c 100644 --- a/src/twinkle_client/exceptions.py +++ b/src/twinkle_client/exceptions.py @@ -80,6 +80,7 @@ def __init__( request_id: str, error_code: int | None = None, details: list[dict[str, Any]] | None = None, + traceback: str | None = None, ) -> None: super().__init__(error) self.error = error @@ -87,6 +88,7 @@ def __init__( self.request_id = request_id self.error_code = error_code self.details = details + self.traceback = traceback class TaskCancelledError(Exception): diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index 1f095bc46..b370660a6 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -3,9 +3,10 @@ import itertools import logging import threading +import uuid from collections.abc import Mapping from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from peft import LoraConfig @@ -13,6 +14,7 @@ from twinkle.protocol.types import model as model_types from twinkle.protocol.types.component import DataRef from twinkle_client._request_builder import build_request +from twinkle_client.exceptions import TwinkleHTTPError from twinkle_client.http import ClientTransport from twinkle_client.http.context import capture_transport @@ -33,7 +35,7 @@ class MultiLoraTransformersModel: """Client wrapper for TrainableModel that calls server HTTP endpoints. This client manages adapters and sends training/inference requests to the model server. - The server-side session (managed by TwinkleClient) keeps the model alive. + The session authenticates requests; executed training work renews the idle deadline. Every method builds its endpoint's request model rather than a dict, so a misspelled or wrongly-typed argument fails here -- in the caller's own stack trace, @@ -47,19 +49,20 @@ def __init__( model_id: str, *, transport: ClientTransport | None = None, - **kwargs, ): """Initialize a model wrapper bound to one immutable request identity.""" self._transport = capture_transport(transport) - kwargs.pop('data_plane_url', None) if '://' in model_id: model_id = model_id.split('://')[1] self.model_id = model_id self.server_url = f'{self._transport.context.base_url}/model/{model_id}/twinkle' self.adapter_name = None + self._training_ids: dict[str, str] = {} + self._creating: set[str] = set() + self._training_lock = threading.Lock() # Per-client monotonic sequence for idempotent dedup of stateful training ops: - # the server dedups on (session_id, seq_id) so a retried grad/step call is + # the server dedups on (training_id, session, operation, seq_id), so each update is # applied at most once. Reserved once per call and reused on retry. self._seq_counter = itertools.count(1) self._seq_lock = threading.Lock() @@ -84,13 +87,39 @@ def _ensure_created(self) -> None: with self._create_lock: if self._created: return - self._transport.post(f'{self.server_url}/create') + response = self._transport.post(f'{self.server_url}/create', json_data={'training_handles': True}) + if response.json().get('training_handles') is not True: + raise RuntimeError('Server does not support training handles; upgrade Twinkle server and SDK together.') self._created = True + _SEQUENCED_ENDPOINTS = frozenset({ + 'forward_backward', + 'forward_backward_from_data_plane', + 'backward', + 'step', + 'lr_step', + 'clip_grad_and_step', + }) + + def _training_values(self, values: dict[str, Any]) -> dict[str, Any]: + """Choose the label and its handle together, under one lock.""" + with self._training_lock: + name = values.get('adapter_name', self.adapter_name) + training_id = values.get('training_id') or self._training_ids.get(name) + if training_id is None: + raise RuntimeError('Create this training with add_adapter_to_model on this object before using it.') + return {**values, 'adapter_name': name, 'training_id': training_id} + def _submit(self, endpoint: str, model_cls, response_cls, **values): - """Build, send, and resolve one twinkle-native request.""" - self._ensure_created() + """Validate once, then send with a stable handle and retry identity.""" + if issubclass(model_cls, model_types.TrainingRequest): + values = self._training_values(values) + if endpoint in self._SEQUENCED_ENDPOINTS: + if 'seq_id' in values: + raise TypeError('seq_id is managed by the client') + values['seq_id'] = self._next_seq_id() body = build_request(model_cls, **values) + self._ensure_created() response = self._transport.post_model(f'{self.server_url}/{endpoint}', body) return self._await_task(response, response_cls) @@ -132,22 +161,50 @@ def add_adapter_to_model( save_dir = kwargs.pop('save_dir', None) if save_dir: save_dir = Path(save_dir).expanduser().resolve().as_posix() - self._submit( - 'add_adapter_to_model', - model_types.AddAdapterRequest, - None, - adapter_name=adapter_name, - config=config, - save_dir=save_dir, - **kwargs) - self.adapter_name = adapter_name + with self._training_lock: + if adapter_name in self._creating: + raise RuntimeError('A training with this adapter_name is already being created on this object.') + if adapter_name in self._training_ids: + raise RuntimeError('A training with this adapter_name already exists on this object; release it first.') + self._creating.add(adapter_name) + try: + result = self._submit( + 'add_adapter_to_model', + model_types.AddAdapterRequest, + model_types.TrainingHandleResponse, + client_create_id=uuid.uuid4().hex, + adapter_name=adapter_name, + config=config, + save_dir=save_dir, + **kwargs) + with self._training_lock: + self._training_ids[adapter_name] = result.training_id + self.adapter_name = adapter_name + finally: + with self._training_lock: + self._creating.discard(adapter_name) def remove_adapter(self, adapter_name: str | None = None) -> None: """Release one client-owned adapter from the training component.""" name = adapter_name or self.adapter_name - self._submit('remove_adapter', model_types.AdapterRequest, None, adapter_name=name) - if name == self.adapter_name: - self.adapter_name = None + with self._training_lock: + training_id = self._training_ids.get(name) + try: + self._submit('remove_adapter', model_types.AdapterRequest, None, adapter_name=name, training_id=training_id) + except TwinkleHTTPError as exc: + if exc.status_code == 404: + self._forget_training(name, training_id) + # A missing handle is unusable, but does not certify device release. + raise + self._forget_training(name, training_id) + + def _forget_training(self, name: str, training_id: str) -> None: + """A late response must never erase a newer same-name binding.""" + with self._training_lock: + if self._training_ids.get(name) == training_id: + self._training_ids.pop(name, None) + if name == self.adapter_name: + self.adapter_name = None # ------------------------------------------------------------------ # # Inline forward family @@ -155,23 +212,12 @@ def remove_adapter(self, adapter_name: str | None = None) -> None: def forward(self, inputs: Any, **kwargs) -> model_types.ForwardResponse: """Execute forward pass on inline model inputs.""" - return self._submit( - 'forward', - model_types.ForwardRequest, - model_types.ForwardResponse, - inputs=inputs, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('forward', model_types.ForwardRequest, model_types.ForwardResponse, inputs=inputs, **kwargs) def forward_only(self, inputs: Any, **kwargs) -> model_types.ForwardResponse: """Execute forward pass without gradient computation on inline inputs.""" return self._submit( - 'forward_only', - model_types.ForwardOnlyRequest, - model_types.ForwardResponse, - inputs=inputs, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + 'forward_only', model_types.ForwardOnlyRequest, model_types.ForwardResponse, inputs=inputs, **kwargs) def forward_backward(self, inputs: Any, **kwargs) -> model_types.ForwardBackwardResponse: """Execute combined forward and backward pass on inline inputs.""" @@ -180,37 +226,20 @@ def forward_backward(self, inputs: Any, **kwargs) -> model_types.ForwardBackward model_types.ForwardBackwardTaskRequest, model_types.ForwardBackwardResponse, inputs=inputs, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), **kwargs) def calculate_loss(self, **kwargs) -> model_types.CalculateLossResponse: """Calculate loss from model outputs.""" - return self._submit( - 'calculate_loss', - model_types.AdapterRequest, - model_types.CalculateLossResponse, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('calculate_loss', model_types.AdapterRequest, model_types.CalculateLossResponse, **kwargs) def get_train_configs(self, **kwargs) -> model_types.GetTrainConfigsResponse: """Get training configs.""" - return self._submit( - 'get_train_configs', - model_types.AdapterRequest, - model_types.GetTrainConfigsResponse, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('get_train_configs', model_types.AdapterRequest, model_types.GetTrainConfigsResponse, + **kwargs) def backward(self, **kwargs) -> None: """Execute backward pass.""" - self._submit( - 'backward', - model_types.AdapterRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), - **kwargs) + self._submit('backward', model_types.AdapterRequest, None, **kwargs) # ------------------------------------------------------------------ # # Data-plane forward family @@ -230,7 +259,6 @@ def forward_from_data_plane( model_types.DataPlaneForwardRequest, model_types.ForwardResponse, input_refs=_data_refs(inputs), - adapter_name=kwargs.pop('adapter_name', self.adapter_name), input_field=input_field, kwarg_fields=kwarg_fields or {}, **kwargs) @@ -251,7 +279,6 @@ def forward_only_from_data_plane( model_types.DataPlaneForwardOnlyRequest, model_types.ForwardResponse, input_refs=_data_refs(inputs), - adapter_name=kwargs.pop('adapter_name', self.adapter_name), input_field=input_field, kwarg_fields=kwarg_fields or {}, output_ref=output_ref.model_dump() if output_ref is not None else None, @@ -275,10 +302,8 @@ def forward_backward_from_data_plane( model_types.DataPlaneForwardRequest, model_types.ForwardBackwardResponse, input_refs=_data_refs(inputs), - adapter_name=kwargs.pop('adapter_name', self.adapter_name), input_field=input_field, kwarg_fields=kwarg_fields or {}, - seq_id=self._next_seq_id(), **kwargs) # ------------------------------------------------------------------ # @@ -287,32 +312,15 @@ def forward_backward_from_data_plane( def step(self, **kwargs) -> None: """Execute optimizer step.""" - self._submit( - 'step', - model_types.StepRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), - **kwargs) + self._submit('step', model_types.StepRequest, None, **kwargs) def zero_grad(self, **kwargs) -> None: """Zero out gradients.""" - self._submit( - 'zero_grad', - model_types.AdapterRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('zero_grad', model_types.AdapterRequest, None, **kwargs) def lr_step(self, **kwargs) -> None: """Execute learning rate scheduler step.""" - self._submit( - 'lr_step', - model_types.LrStepRequest, - None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), - **kwargs) + self._submit('lr_step', model_types.LrStepRequest, None, **kwargs) def clip_grad_norm(self, max_grad_norm: float = 1.0, @@ -323,7 +331,6 @@ def clip_grad_norm(self, 'clip_grad_norm', model_types.ClipGradNormRequest, model_types.ClipGradNormResponse, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), max_grad_norm=max_grad_norm, norm_type=norm_type, **kwargs) @@ -334,10 +341,8 @@ def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type: int = 2, **k 'clip_grad_and_step', model_types.ClipGradAndStepRequest, None, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), max_grad_norm=max_grad_norm, norm_type=norm_type, - seq_id=self._next_seq_id(), **kwargs) # ------------------------------------------------------------------ # @@ -346,33 +351,15 @@ def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type: int = 2, **k def set_loss(self, loss_cls: str, **kwargs) -> None: """Set the loss function.""" - self._submit( - 'set_loss', - model_types.SetLossRequest, - None, - loss_cls=loss_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_loss', model_types.SetLossRequest, None, loss_cls=loss_cls, **kwargs) def set_optimizer(self, optimizer_cls: str, **kwargs) -> None: """Set the optimizer.""" - self._submit( - 'set_optimizer', - model_types.SetOptimizerRequest, - None, - optimizer_cls=optimizer_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_optimizer', model_types.SetOptimizerRequest, None, optimizer_cls=optimizer_cls, **kwargs) def set_lr_scheduler(self, scheduler_cls: str, **kwargs) -> None: """Set the learning rate scheduler.""" - self._submit( - 'set_lr_scheduler', - model_types.SetLrSchedulerRequest, - None, - scheduler_cls=scheduler_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_lr_scheduler', model_types.SetLrSchedulerRequest, None, scheduler_cls=scheduler_cls, **kwargs) def set_template(self, template_cls: str, **kwargs) -> None: """Set the template for data processing. @@ -382,44 +369,20 @@ def set_template(self, template_cls: str, **kwargs) -> None: effect. A caller that passes it explicitly still reaches the template constructor through the passthrough region. """ - self._submit( - 'set_template', - model_types.SetTemplateRequest, - None, - template_cls=template_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_template', model_types.SetTemplateRequest, None, template_cls=template_cls, **kwargs) def set_processor(self, processor_cls: str, **kwargs) -> None: """Set the input processor.""" - self._submit( - 'set_processor', - model_types.SetProcessorRequest, - None, - processor_cls=processor_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('set_processor', model_types.SetProcessorRequest, None, processor_cls=processor_cls, **kwargs) def add_metric(self, metric_cls: str, is_training: bool | None = None, **kwargs) -> None: """Add a metric to the model.""" self._submit( - 'add_metric', - model_types.AddMetricRequest, - None, - metric_cls=metric_cls, - is_training=is_training, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + 'add_metric', model_types.AddMetricRequest, None, metric_cls=metric_cls, is_training=is_training, **kwargs) def apply_patch(self, patch_cls: str, **kwargs) -> None: """Apply a patch to the model.""" - self._submit( - 'apply_patch', - model_types.ApplyPatchRequest, - None, - patch_cls=patch_cls, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('apply_patch', model_types.ApplyPatchRequest, None, patch_cls=patch_cls, **kwargs) def calculate_metric(self, is_training: bool = True, **kwargs) -> model_types.CalculateMetricResponse: """Calculate metrics from model outputs.""" @@ -428,7 +391,6 @@ def calculate_metric(self, is_training: bool = True, **kwargs) -> model_types.Ca model_types.CalculateMetricRequest, model_types.CalculateMetricResponse, is_training=is_training, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), **kwargs) # ------------------------------------------------------------------ # @@ -437,23 +399,11 @@ def calculate_metric(self, is_training: bool = True, **kwargs) -> model_types.Ca def save(self, name: str, **kwargs) -> model_types.SaveResponse: """Save model checkpoint.""" - return self._submit( - 'save', - model_types.SaveRequest, - model_types.SaveResponse, - name=name, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + return self._submit('save', model_types.SaveRequest, model_types.SaveResponse, name=name, **kwargs) def load(self, name: str, **kwargs) -> None: """Load model checkpoint.""" - self._submit( - 'load', - model_types.LoadRequest, - None, - name=name, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + self._submit('load', model_types.LoadRequest, None, name=name, **kwargs) def resume_from_checkpoint(self, name: str, *, resume_only_model: bool = False, **kwargs) -> dict[str, Any]: """Resume weights (and optionally optimizer state) from a checkpoint.""" @@ -462,7 +412,6 @@ def resume_from_checkpoint(self, name: str, *, resume_only_model: bool = False, model_types.ResumeFromCheckpointRequest, model_types.TrainingProgressResponse, name=name, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), resume_only_model=resume_only_model, **kwargs) return progress.result @@ -472,8 +421,6 @@ def upload_to_hub( checkpoint_dir: str, hub_model_id: str, hub_token: str | None = None, - async_upload: bool = True, - poll_interval: float = 5.0, ) -> None: """Upload model checkpoint to hub. @@ -484,10 +431,6 @@ def upload_to_hub( checkpoint_dir: The directory path of the checkpoint to upload. hub_model_id: The hub model id. hub_token: The hub token (optional). - async_upload: Deprecated, has no effect. The server always runs the - upload in the background and the client waits via the future layer. - poll_interval: Deprecated, has no effect. Pacing is now owned by the - server-side long-poll of the Retrieve_Endpoint. """ logger.info('[upload_to_hub] submitting upload, waiting for completion...') self._submit( diff --git a/tests/infra/test_ray_registry.py b/tests/infra/test_ray_registry.py new file mode 100644 index 000000000..f3b4c6f2e --- /dev/null +++ b/tests/infra/test_ray_registry.py @@ -0,0 +1,62 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Registry isolation must not reserve a CPU for ordinary non-server use.""" +import subprocess +import sys +import textwrap + + +def test_registry_cpu_reservations_and_instance_isolation(): + # Own a subprocess and cluster so the CPU assertions do not count actors + # belonging to other tests or reset their process-global RayHelper cache. + code = textwrap.dedent(''' + import os + import time + import uuid + import ray + from twinkle.infra._ray.ray_helper import RayHelper + + namespace = 'registry-' + uuid.uuid4().hex[:8] + ray.init(address='local', num_cpus=2, namespace=namespace, + include_dashboard=False, log_to_driver=False, + _temp_dir='/tmp/twreg-' + uuid.uuid4().hex[:6]) + + def registry(name, managed): + RayHelper._registry = None + os.environ['TWINKLE_RAY_REGISTRY_NAME'] = name + if managed: + os.environ['TWINKLE_RAY_NAMESPACE'] = namespace + else: + os.environ.pop('TWINKLE_RAY_NAMESPACE', None) + RayHelper.init_registry() + actor = RayHelper._registry + ray.get(actor.add_config.remote('owner', name)) + return actor + + def assert_free_cpus(expected): + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + available = ray.available_resources().get('CPU', 0) + if abs(available - expected) < 1e-6: + return + time.sleep(.01) + raise AssertionError((available, expected)) + + try: + ordinary = registry('ordinary', False) + assert_free_cpus(2) + ray.kill(ordinary, no_restart=True) + left = registry('left', True) + right = registry('right', True) + assert_free_cpus(1.98) + assert ray.get(left.get_config.remote('owner')) == 'left' + assert ray.get(right.get_config.remote('owner')) == 'right' + ray.kill(left, no_restart=True) + assert_free_cpus(1.99) + assert ray.get(right.get_config.remote('owner')) == 'right' + ray.kill(right, no_restart=True) + assert_free_cpus(2) + finally: + ray.shutdown() + ''') + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/loss/test_dpo.py b/tests/loss/test_dpo.py index 89b9ddac4..62af61237 100644 --- a/tests/loss/test_dpo.py +++ b/tests/loss/test_dpo.py @@ -75,12 +75,12 @@ def test_dpo_reference_free(self): result = loss_fn(inputs, outputs) assert result['loss'].dim() == 0 - def test_dpo_no_ref_returns_zero(self): - """Without ref_logps and not reference_free, loss should be zero.""" + def test_dpo_no_ref_is_rejected_unless_reference_free(self): + """Missing references must fail before producing a misleading zero loss.""" loss_fn = DPOLoss(beta=0.1) inputs, outputs, _ = _make_preference_batch(with_ref=False) - result = loss_fn(inputs, outputs) - assert result['loss'].item() == pytest.approx(0.0, abs=1e-6) + with pytest.raises(ValueError, match='requires ref_logps'): + loss_fn(inputs, outputs) def test_dpo_with_label_smoothing(self): loss_fn = DPOLoss(beta=0.1, label_smoothing=0.1, loss_type='sigmoid') diff --git a/tests/loss/test_grpo_gkd.py b/tests/loss/test_grpo_gkd.py index 1f9daa707..9623df67e 100644 --- a/tests/loss/test_grpo_gkd.py +++ b/tests/loss/test_grpo_gkd.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Tests for GRPO family losses and GKD loss.""" +import numpy as np import pytest import torch import torch.nn.functional as F @@ -286,3 +287,73 @@ def test_gkd_no_ignored_tokens(self): result = loss_fn({'labels': labels}, {'logits': student_logits}, teacher_logits=teacher_logits) assert torch.isfinite(result['loss']) + + +class TestGRPOValidateInputs: + """Pre-forward validation must match the in-loss alignment rules.""" + + @staticmethod + def _batch(): + # Two trainable positions per row, at indices 2 and 3, so a length-3 + # row cannot be read as the full-sequence form (index 3 is missing). + labels = torch.tensor([[-100, -100, 0, 0], [-100, -100, 0, 0]]) + return [{'labels': labels}, {'labels': labels}] + + def test_accepts_response_only_scalar_and_full_length(self): + loss_fn = GRPOLoss() + inputs = self._batch() + loss_fn.validate_inputs(inputs, old_logps=[[0.1, 0.2], [0.3, 0.4]]) + loss_fn.validate_inputs(inputs, old_logps=[[0.1], [0.2]]) + loss_fn.validate_inputs(inputs, old_logps=torch.zeros(2, 4)) + + def test_rejects_ragged_length_before_forward(self): + loss_fn = GRPOLoss() + inputs = self._batch() + with pytest.raises(AssertionError, match='data/mask length mismatch at sample 1'): + loss_fn.validate_inputs(inputs, old_logps=[[0.1, 0.2], [0.3, 0.4, 0.5]]) + + def test_rejects_ragged_advantages(self): + loss_fn = GRPOLoss() + inputs = self._batch() + with pytest.raises(AssertionError, match='data/mask length mismatch'): + loss_fn.validate_inputs(inputs, advantages=[1.0, [1.0, 2.0, 3.0]]) + + @pytest.mark.parametrize('field', ['old_logps', 'ref_logps', 'advantages']) + @pytest.mark.parametrize('representation', ['list', 'tuple', 'tensor', 'numpy']) + @pytest.mark.parametrize('labels,n_sample,supported', [ + ([-100, -100, 0, 0, -100], 0, False), + ([-100, -100, 0, 0, -100], 1, True), + ([-100, -100, 0, 0, -100], 2, True), + ([-100, -100, 0, 0, -100], 3, False), + ([-100, -100, 0, 0, -100], 4, True), + ([-100, -100, 0, 0, -100], 5, True), + ([-100, -100, 0, 0, -100], 6, True), + ([-100, -100, -100], 0, True), + ([-100, -100, -100], 1, True), + ([-100, -100, -100], 2, True), + ]) + def test_preflight_and_loss_agree_on_alignment(self, labels, n_sample, supported, representation, field): + loss_fn = GRPOLoss(beta=0.1) + inputs = {'labels': torch.tensor(labels)} + outputs = {'logps': torch.zeros(1, len(labels), requires_grad=True)} + data = [[0.1] * n_sample] + if representation == 'tuple': + data = tuple(tuple(row) for row in data) + elif representation == 'tensor': + data = torch.tensor(data) + elif representation == 'numpy': + data = np.asarray(data, dtype=np.float32) + kwargs = {'advantages': [1.0], field: data} + + if not supported: + with pytest.raises(AssertionError, match='data/mask length mismatch at sample 0'): + loss_fn.validate_inputs([inputs], **kwargs) + with pytest.raises(AssertionError, match='data/mask length mismatch at sample 0'): + loss_fn(inputs, outputs, **kwargs) + return + + loss_fn.validate_inputs([inputs], **kwargs) + result = loss_fn(inputs, outputs, **kwargs) + assert torch.isfinite(result['loss']) + result['loss'].backward() + assert torch.isfinite(outputs['logps'].grad).all() diff --git a/tests/model/test_optimizer_group.py b/tests/model/test_optimizer_group.py index 45a1a0e0e..518f1aa38 100644 --- a/tests/model/test_optimizer_group.py +++ b/tests/model/test_optimizer_group.py @@ -7,7 +7,7 @@ and ``Loss`` objects through the real class on CPU (no model, no GPU, no distributed): * ``do_grad_sync`` gradient-accumulation gating and its state side effect, -* the ``__setattr__`` seam that wraps any assigned ``loss_instance`` in ``safe_loss``, +* loss assignment / selection and numerical transparency, * ``accumulate_metrics`` / ``calculate_metrics`` against a real ``LossMetric``. """ import pytest @@ -17,7 +17,6 @@ from twinkle.loss.base import Loss from twinkle.metric import LossMetric from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus -from twinkle.utils.nccl_safe import SafeLossWrapper # --------------------------------------------------------------------------- @@ -90,47 +89,41 @@ def test_none_argument_uses_state_without_mutation(self): # --------------------------------------------------------------------------- -# __setattr__ safe_loss wrapping seam +# Loss assignment / selection # --------------------------------------------------------------------------- -class TestSafeLossWrapping: +class TestLossSelection: - def test_assigning_loss_wraps_in_safe_loss_wrapper(self): + def test_assigned_loss_is_selected_without_changing_identity(self): group = BaseOptimizerGroup() raw = CrossEntropyLoss(reduction='sum') group.loss_instance = raw - assert isinstance(group.loss_instance, SafeLossWrapper) - assert isinstance(group.loss_instance, Loss) - assert group.loss_instance._loss_instance is raw + assert group.get_loss() is raw + assert isinstance(group.get_loss(), Loss) - def test_wrapping_is_idempotent(self): + def test_reassignment_is_idempotent(self): group = BaseOptimizerGroup() raw = CrossEntropyLoss(reduction='sum') group.loss_instance = raw - wrapped_once = group.loss_instance - # Re-assigning the already-wrapped instance must not nest a second wrapper. - group.loss_instance = wrapped_once - assert group.loss_instance is wrapped_once - assert group.loss_instance._loss_instance is raw + group.loss_instance = group.get_loss() + assert group.get_loss() is raw - def test_assigning_none_is_not_wrapped(self): + def test_assigning_none_clears_default_loss(self): group = BaseOptimizerGroup() group.loss_instance = CrossEntropyLoss(reduction='sum') group.loss_instance = None - assert group.loss_instance is None + assert group.get_loss() is None - def test_wrapped_loss_delegates_transparently(self, monkeypatch): - # Default fail-fast mode makes the wrapper a pass-through, so the seam must not - # change the numeric contract of the underlying loss. - monkeypatch.setenv('TWINKLE_FAIL_FAST', '1') + def test_selected_loss_keeps_numeric_contract_and_propagates_errors(self): group = BaseOptimizerGroup() - raw = CrossEntropyLoss(reduction='sum') - group.loss_instance = raw + group.loss_instance = CrossEntropyLoss(reduction='sum') inputs = {'labels': torch.tensor([[1, 2, -100]])} outputs = {'logps': torch.tensor([[-1.0, -2.0, -3.0]])} - wrapped_out = group.loss_instance(inputs, outputs) - raw_out = CrossEntropyLoss(reduction='sum')(inputs, outputs) - assert wrapped_out['loss'].item() == pytest.approx(raw_out['loss'].item()) + selected = group.get_loss()(inputs, outputs) + raw = CrossEntropyLoss(reduction='sum')(inputs, outputs) + assert selected['loss'].item() == pytest.approx(raw['loss'].item()) + with pytest.raises(KeyError): + group.get_loss()(inputs, {}) # --------------------------------------------------------------------------- diff --git a/tests/pluger/test_loader.py b/tests/pluger/test_loader.py index 76735a697..9d688c4fb 100644 --- a/tests/pluger/test_loader.py +++ b/tests/pluger/test_loader.py @@ -35,7 +35,7 @@ def test_load_plugin_invalid_id(self): def test_load_plugin_safe_mode(self): """Test loading plugin when trust_remote_code is False.""" - with patch('twinkle.utils.loader.MSHub.download_model', return_value='/tmp/fake'): + with patch('twinkle.hub.hub.MSHub.download_model', return_value='/tmp/fake'): with patch('twinkle.utils.loader.trust_remote_code', return_value=False): with pytest.raises(ValueError, match='Twinkle does not support plugin in safe mode'): Plugin.load_plugin('ms://test/plugin', BasePlugin) @@ -43,7 +43,7 @@ def test_load_plugin_safe_mode(self): def test_load_plugin_missing_init_file(self): """Test loading plugin when __init__.py is missing.""" with tempfile.TemporaryDirectory() as tmpdir: - with patch('twinkle.utils.loader.MSHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.MSHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with pytest.raises(AssertionError, match='does not exist'): Plugin.load_plugin('ms://test/plugin', BasePlugin) @@ -58,7 +58,7 @@ def test_load_plugin_no_subclass(self): other_class = type('OtherClass', (), {}) other_class.__module__ = str(init_file) - with patch('twinkle.utils.loader.MSHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.MSHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with patch('twinkle.utils.loader.importlib.import_module') as mock_import: mock_module = MagicMock() @@ -87,7 +87,7 @@ def test_load_plugin_ms_hub(self): test_plugin_class.__module__ = '__init__' mock_module.TestPlugin = test_plugin_class - with patch('twinkle.utils.loader.MSHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.MSHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with patch('twinkle.utils.loader.importlib.import_module', return_value=mock_module): plugin_cls = Plugin.load_plugin('ms://test/plugin', BasePlugin) @@ -111,7 +111,7 @@ def test_load_plugin_hf_hub(self): test_plugin_class.__module__ = '__init__' mock_module.TestPlugin = test_plugin_class - with patch('twinkle.utils.loader.HFHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.HFHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with patch('twinkle.utils.loader.importlib.import_module', return_value=mock_module): plugin_cls = Plugin.load_plugin('hf://test/plugin', BasePlugin) @@ -137,7 +137,7 @@ def test_load_plugin_sys_path_management(self): assert tmpdir not in sys.path - with patch('twinkle.utils.loader.MSHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.MSHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with patch('twinkle.utils.loader.importlib.import_module', return_value=mock_module): Plugin.load_plugin('ms://test/plugin', BasePlugin) @@ -186,7 +186,7 @@ def test_construct_class_with_plugin_id(self): test_plugin_class.__module__ = '__init__' mock_module.TestPlugin = test_plugin_class - with patch('twinkle.utils.loader.MSHub.download_model', return_value=tmpdir): + with patch('twinkle.hub.hub.MSHub.download_model', return_value=tmpdir): with patch('twinkle.utils.loader.trust_remote_code', return_value=True): with patch('twinkle.utils.loader.importlib.import_module', return_value=mock_module): result = construct_class('ms://test/plugin', BasePlugin, [], name='test') diff --git a/tests/processor/test_processor.py b/tests/processor/test_processor.py index 6840a9364..eca28b879 100644 --- a/tests/processor/test_processor.py +++ b/tests/processor/test_processor.py @@ -5,6 +5,7 @@ import torch import twinkle +from twinkle import Platform from twinkle.processor import InputProcessor twinkle.initialize(mode='local') @@ -139,6 +140,48 @@ def test_micro_batch_variable_length(self): assert b['input_ids'].shape[0] == 2 +class TestMegatronCausalMask: + """The dense 4D mask must not be built when the backend derives its own.""" + + def _collate(self, monkeypatch, device_prefix, attention_mask_type): + monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda: device_prefix)) + proc = InputProcessor(padding_free=False, framework='megatron') + proc._create_4d_attention_mask = lambda _: pytest.fail('dense 4D mask allocated') + return proc.collate_fn( + _make_text_batch(4, seq_len=8), + micro_batch_size=2, + variable_seq_lengths=False, + attention_mask_type=attention_mask_type, + ) + + def test_causal_npu_omits_dense_mask(self, monkeypatch): + outputs = self._collate(monkeypatch, 'npu', 'causal') + assert len(outputs) == 2 + assert all(output['attention_mask'] is None for output in outputs) + + def test_non_causal_still_builds_dense_mask(self, monkeypatch): + monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda: 'npu')) + proc = InputProcessor(padding_free=False, framework='megatron') + outputs = proc.collate_fn( + _make_text_batch(4, seq_len=8), + micro_batch_size=2, + variable_seq_lengths=False, + attention_mask_type=None, + ) + assert all(output['attention_mask'].dim() == 4 for output in outputs) + + def test_other_backends_still_build_dense_mask(self, monkeypatch): + monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda: 'cuda')) + proc = InputProcessor(padding_free=False, framework='megatron') + outputs = proc.collate_fn( + _make_text_batch(4, seq_len=8), + micro_batch_size=2, + variable_seq_lengths=False, + attention_mask_type='causal', + ) + assert all(output['attention_mask'].dim() == 4 for output in outputs) + + class TestMultimodalMode: """Multimodal: pixel_values, image_grid_thw.""" diff --git a/tests/sampler/test_vllm_stream_lifetime.py b/tests/sampler/test_vllm_stream_lifetime.py new file mode 100644 index 000000000..d7542bafd --- /dev/null +++ b/tests/sampler/test_vllm_stream_lifetime.py @@ -0,0 +1,57 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The vLLM stream iterator owns its engine task until actual completion.""" +import ast +import asyncio +import threading +from pathlib import Path +from types import SimpleNamespace + + +def test_stream_iterator_close_joins_the_real_engine_loop_task(): + # Execute the repository method without constructing a GPU-backed vLLM engine. + path = Path('src/twinkle/sampler/vllm_sampler/vllm_sampler.py') + tree = ast.parse(path.read_text()) + method = next(node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == '_iter_in_loop') + namespace = {'asyncio': asyncio} + exec(compile(ast.Module(body=[method], type_ignores=[]), str(path), 'exec'), namespace) + loop = asyncio.new_event_loop() + running, finalizing, release, closed = [threading.Event() for _ in range(4)] + + def run_loop(): + asyncio.set_event_loop(loop) + running.set() + loop.run_forever() + + thread = threading.Thread(target=run_loop) + thread.start() + running.wait(1) + + async def generation(): + try: + yield 'first' + finally: + finalizing.set() + await asyncio.to_thread(release.wait) + + iterator = namespace['_iter_in_loop'](SimpleNamespace(_async_loop=loop), generation()) + closer = None + try: + assert next(iterator) == 'first' + assert finalizing.wait(1) + + def close(): + iterator.close() + closed.set() + + closer = threading.Thread(target=close) + closer.start() + assert not closed.wait(0.05) + release.set() + assert closed.wait(2) + finally: + release.set() + if closer: + closer.join(2) + loop.call_soon_threadsafe(loop.stop) + thread.join(2) + loop.close() diff --git a/tests/server/utils/__init__.py b/tests/server/checkpoint/__init__.py similarity index 100% rename from tests/server/utils/__init__.py rename to tests/server/checkpoint/__init__.py diff --git a/tests/server/utils/test_checkpoint_manager.py b/tests/server/checkpoint/test_checkpoint_manager.py similarity index 94% rename from tests/server/utils/test_checkpoint_manager.py rename to tests/server/checkpoint/test_checkpoint_manager.py index 88a40e0d0..b342d2935 100644 --- a/tests/server/utils/test_checkpoint_manager.py +++ b/tests/server/checkpoint/test_checkpoint_manager.py @@ -414,9 +414,28 @@ def write(path, data): @pytest.mark.asyncio async def test_reserved_sampler_name_is_rejected_before_backend_or_file_io(manager): from unittest.mock import AsyncMock - from twinkle.server.task_queue.types import UserTaskError + from twinkle.server.exceptions import RequestRejectedError save = AsyncMock() - with pytest.raises(UserTaskError, match='reserved'): + with pytest.raises(RequestRejectedError, match='reserved'): await manager.save_sampler('run', 'latest', save) save.assert_not_awaited() assert not Path(manager.get_save_dir('run', is_sampler=True)).exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('manager', ['tinker', 'twinkle'], indirect=True) +async def test_delete_rejects_non_checkpoint_paths(manager): + """Only a concrete checkpoint may be deleted, never a kind root, alias, or the run root.""" + await manager.save_training('run', 'step', _writer('training')) + await manager.save_sampler('run', None, _writer('live')) + run_dir = manager.training_run_manager.get_model_dir('run') + + def snapshot(): + return sorted(path.relative_to(run_dir).as_posix() for path in run_dir.rglob('*')) + + before = snapshot() + for checkpoint_id in ('weights', 'sampler_weights', '.', 'sampler_weights/latest', 'weights/../weights/step'): + assert manager.delete('run', checkpoint_id) is False, checkpoint_id + assert snapshot() == before + assert manager.delete('run', 'weights/step') is True + assert not (run_dir / 'weights/step').exists() diff --git a/tests/server/cli/test_cli.py b/tests/server/cli/test_cli.py index a49ae372e..ef859abca 100644 --- a/tests/server/cli/test_cli.py +++ b/tests/server/cli/test_cli.py @@ -89,5 +89,5 @@ def test_env_var_overrides_when_flag_omitted(monkeypatch) -> None: def test_example_config_loads_via_server_config() -> None: cfg = ServerConfig.from_yaml(EXAMPLE) assert isinstance(cfg, ServerConfig) - assert any(a.import_path == 'model' for a in cfg.applications) + assert cfg.training_config.models[0].backend == 'mock' assert any(a.import_path == 'sampler' for a in cfg.applications) diff --git a/tests/server/config/server_config_4b_e2e.yaml b/tests/server/config/server_config_4b_e2e.yaml index 6fc4d658f..e4b06ea67 100644 --- a/tests/server/config/server_config_4b_e2e.yaml +++ b/tests/server/config/server_config_4b_e2e.yaml @@ -32,6 +32,8 @@ applications: args: backend: transformers model_id: "ms://Qwen/Qwen3.5-4B" + min_replicas: 1 + max_replicas: 1 max_length: 10240 max_loras: 10 nproc_per_node: 2 diff --git a/tests/server/config/server_config_4b_e2e_megatron.yaml b/tests/server/config/server_config_4b_e2e_megatron.yaml index ab71f6bb5..1c83ddac9 100644 --- a/tests/server/config/server_config_4b_e2e_megatron.yaml +++ b/tests/server/config/server_config_4b_e2e_megatron.yaml @@ -32,6 +32,8 @@ applications: args: backend: megatron model_id: "ms://Qwen/Qwen3.5-4B" + min_replicas: 1 + max_replicas: 1 max_length: 10240 max_loras: 10 nproc_per_node: 4 diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 927089ff2..cb858b625 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -33,38 +33,26 @@ }), ) -_MODEL_APP = st.fixed_dictionaries({ - 'name': - st.just('m'), - 'route_prefix': - st.just('/api/v1/m'), - 'import_path': - st.just('model'), - 'args': - st.fixed_dictionaries({ +_TRAINING_MODEL = st.fixed_dictionaries({ + 'name': st.just('m'), 'import_path': st.just('model'), 'route_prefix': st.just('/api/v1/model/m'), + 'args': st.fixed_dictionaries({ 'model_id': st.just('model-id'), - 'device_group': st.just({ - 'name': 'g', - 'ranks': 1, - 'device_type': 'CPU' - }), - 'device_mesh': st.just({ - 'device_type': 'CPU', - 'dp_size': 1 - }), 'backend': st.sampled_from(['mock', 'transformers', 'megatron']), + 'device_group': st.just({'ranks': 2, 'device_type': 'cuda'}), + 'device_mesh': st.just({'dp_size': 2}), + 'train_mode': st.sampled_from(['lora', 'full']), }), }) _VALID_CONFIG = st.fixed_dictionaries({ 'persistence': _PERSISTENCE_VARIANTS, - 'applications': st.lists(_MODEL_APP, min_size=0, max_size=3), + 'applications': st.lists(_TRAINING_MODEL, min_size=0, max_size=1), }) # ---------- valid → fully validated ---------------------------------------- # -@settings(max_examples=100) +@settings(max_examples=100, deadline=None) @given(payload=_VALID_CONFIG) def test_valid_payload_yields_full_instance(payload: dict) -> None: cfg = ServerConfig.model_validate(payload) @@ -118,7 +106,7 @@ def test_nested_field_constraint_violation_named(bad_max_input_tokens: int) -> N assert any('max_input_tokens' in err['loc'] for err in errors) -def test_torch_sampler_is_rejected_during_config_validation() -> None: +def test_unknown_sampler_is_rejected_during_config_validation() -> None: with pytest.raises(ValidationError): ApplicationSpec.model_validate({ 'name': 'sampler', @@ -127,21 +115,29 @@ def test_torch_sampler_is_rejected_during_config_validation() -> None: 'model_id': 'm', 'device_group': {}, 'device_mesh': {}, - 'sampler_type': 'torch', + 'sampler_type': 'unknown_sampler', }, }) def test_sampler_docstring_values_match_selector() -> None: doc = build_sampler_app.__doc__ or '' - line = next(line for line in doc.splitlines() if 'sampler_type:' in line) + line = doc[doc.index('sampler_type:'):].split('Validated up front')[0] assert set(re.findall(r'``(\w+)``', line)) == set(SAMPLER_SELECTOR.builders) +@pytest.mark.parametrize('sampler_type', ['torch', 'sglang_async']) +def test_dev_sampler_extensions_remain_available(sampler_type): + app = ApplicationSpec.model_validate({ + 'name': 'sampler', 'import_path': 'sampler', 'args': { + 'model_id': 'm', 'device_group': {}, 'device_mesh': {}, 'sampler_type': sampler_type}}) + assert app.args.sampler_type == sampler_type + + # ---------- round-trip fidelity ----------------------------------------- # -@settings(max_examples=100) +@settings(max_examples=100, deadline=None) @given(payload=_VALID_CONFIG) def test_round_trip_fidelity(payload: dict) -> None: cfg = ServerConfig.model_validate(payload) diff --git a/tests/server/utils/task_queue/test_config.py b/tests/server/config/test_task_queue_config.py similarity index 98% rename from tests/server/utils/task_queue/test_config.py rename to tests/server/config/test_task_queue_config.py index 60787f567..511d6bde5 100644 --- a/tests/server/utils/task_queue/test_config.py +++ b/tests/server/config/test_task_queue_config.py @@ -13,7 +13,7 @@ from hypothesis import strategies as st from pydantic import ValidationError -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig # ---------- defaults snapshot used by the default-value test -------------- # diff --git a/tests/server/config/test_training_config.py b/tests/server/config/test_training_config.py new file mode 100644 index 000000000..d288bcb4e --- /dev/null +++ b/tests/server/config/test_training_config.py @@ -0,0 +1,419 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import copy +import math + +import pytest +from pydantic import ValidationError + +from twinkle.server.config import ServerConfig +from twinkle.server.config.training import TrainingConfig +from twinkle.server.exceptions import RequestRejectedError + + +def model_config(mode='lora'): + devices = 2 if mode == 'lora' else 4 + return { + 'name': 'Qwen/small' if mode == 'lora' else 'Qwen/small-full', + 'model_id': 'ms://Qwen/small', 'backend': 'transformers', 'device_type': 'cuda', + 'mode': mode, 'config': {'num_devices': devices, 'device_mesh': {'fsdp_size': devices}}, + } + + +def model_application(): + return {'name': 'model-small', 'route_prefix': '/api/v1/model/Qwen/small', + 'import_path': 'model', 'args': { + 'model_id': 'ms://Qwen/small', 'backend': 'transformers', + 'device_group': {'ranks': 2, 'device_type': 'cuda'}, + 'device_mesh': {'fsdp_size': 2}, + }} + + +@pytest.mark.asyncio +async def test_custom_model_specs_survive_catalog_roundtrip_and_worker_construction(monkeypatch): + from types import SimpleNamespace + from unittest.mock import AsyncMock, Mock + from twinkle.server.config.training import TrainingModelConfig + from twinkle.server.model import app as model_app + from twinkle.server.model.pool.driver import RayInstanceDriver + from twinkle.server.state.training_records import InstanceRecord + + app = model_application() + # These modules deliberately do not exist on the gateway. Resolution belongs to workers. + app['args'].update(model_cls='worker_only.encoder:Model', model_loader='worker_only.loader:load', + gradient_checkpointing=False, mixed_precision='no') + catalog = ServerConfig(applications=[app]).training_config + spec = TrainingModelConfig.model_validate_json(catalog.models[0].model_dump_json()) + captured = Mock(return_value=object()) + monkeypatch.setattr(model_app, 'build_model_app', captured) + import importlib + monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', + lambda name: object()) + driver = RayInstanceDriver(catalog, None, 'custom-config-test', 'controller') + driver.runtime = SimpleNamespace( + pool=SimpleNamespace(name=lambda *args: 'owned-instance'), + deployment_options=lambda *args: {}, deploy=AsyncMock()) + await driver.deploy(InstanceRecord(instance_id='instance', controller_epoch='epoch', + base_model=spec.name, mode='lora', startup_deadline=0), spec) + assert captured.call_args.kwargs['model_cls'] == app['args']['model_cls'] + assert captured.call_args.kwargs['model_loader'] == app['args']['model_loader'] + assert captured.call_args.kwargs['gradient_checkpointing'] is False + assert captured.call_args.kwargs['mixed_precision'] == 'no' + driver.runtime.deploy.assert_awaited_once() + + +@pytest.mark.parametrize('mutation', ['full', 'hybrid', 'megatron']) +def test_custom_model_class_rejects_incompatible_training_modes(mutation): + app = model_application() + app['args']['model_cls'] = 'worker_only.encoder:Model' + if mutation == 'full': + app['args']['train_mode'] = 'full' + elif mutation == 'hybrid': + app['args']['hybrid'] = {'allocation_path': '/tmp/allocation.json'} + else: + app['args']['backend'] = 'megatron' + with pytest.raises(ValidationError, match='model_cls'): + ServerConfig(applications=[app]) + + +def test_st_example_is_a_cold_on_demand_custom_model(): + config = ServerConfig.from_yaml('cookbook/client/swift_model/sentence_transformer_config.yaml') + model = config.training_config.models[0] + assert model.config.model_cls.endswith(':MultiLoraSentenceTransformerModel') + assert model.config.gradient_checkpointing is False + assert model.config.mixed_precision == 'no' + assert (model.config.min_replicas, model.config.max_replicas, model.capacity) == (0, 1, 2) + + +def test_catalog_defaults_and_mode_selection(): + config = TrainingConfig(models=[model_config(), model_config('full')]) + catalog = config + shared = catalog.get('Qwen/small', 'lora') + exclusive = catalog.get('Qwen/small-full', 'full') + assert shared.capacity == 5 + assert exclusive.capacity == 1 + assert shared.config.min_replicas == exclusive.config.min_replicas == 0 + assert shared.config.idle_timeout_s == 1800 + assert config.allocation_timeout_s == 3 + with pytest.raises(RequestRejectedError, match='Unknown training model'): + catalog.get('not-in-catalog', 'lora') + + +def test_existing_yaml_uses_explicit_instance_limits_and_business_idle(): + from pathlib import Path + paths = sorted(Path('cookbook/client/server').glob('**/server_config*.yaml')) + assert paths + for path in paths: + config = ServerConfig.from_yaml(path) + applications = [app for app in config.applications if app.import_path == 'model'] + for app, model in zip(applications, config.training_config.models): + assert model.config.idle_timeout_s == app.args.idle_timeout_s + assert model.config.min_replicas == app.args.min_replicas + assert model.config.max_replicas == app.args.max_replicas + assert model.config.adapter_config == app.args.adapter_config + + +@pytest.mark.parametrize('field,value', [ + ('num_devices', 0), ('num_devices', True), ('num_devices', 3), + ('min_replicas', -1), ('idle_timeout_s', math.inf), ('idle_timeout_s', math.nan), + ('idle_timeout_s', 0), ('ranks', [0, 1]), ('max_loras', 0), +]) +def test_invalid_mode_cannot_reach_allocation(field, value): + model = model_config() + model['config'][field] = value + with pytest.raises(ValidationError): + TrainingConfig(models=[model]) + + +def test_ep_and_vpp_do_not_multiply_physical_device_count(): + model = model_config() + model['backend'] = 'megatron' + model['config'] = { + 'num_devices': 8, + 'device_mesh': {'tp_size': 2, 'dp_size': 2, 'pp_size': 2, 'ep_size': 2, 'vpp_size': 2}, + } + assert TrainingConfig(models=[model]).models[0].config.num_devices == 8 + + +@pytest.mark.parametrize('mutation', ['duplicate', 'missing_mode', 'full_lora_field', 'hybrid_megatron', 'mesh_type', 'min_max']) +def test_cross_field_constraints(mutation): + model = model_config() + models = [model] + if mutation == 'duplicate': + models.append(copy.deepcopy(model)) + elif mutation == 'missing_mode': + del model['mode'] + elif mutation == 'full_lora_field': + model = model_config('full') + model['config']['max_loras'] = 2 + models = [model] + elif mutation == 'hybrid_megatron': + model['backend'] = 'megatron' + model['config']['hybrid'] = {'allocation_path': '/tmp/allocation.json'} + elif mutation == 'mesh_type': + model['config']['device_mesh']['device_type'] = 'npu' + elif mutation == 'min_max': + model['config'].update(min_replicas=2, max_replicas=1) + with pytest.raises(ValidationError): + TrainingConfig(models=models) + + +def test_full_only_is_explicit_and_lora_is_unavailable(): + model = model_config('full') + catalog = TrainingConfig(models=[model]) + assert catalog.get('Qwen/small-full', 'full').capacity == 1 + with pytest.raises(RequestRejectedError, match='does not provide lora'): + catalog.get('Qwen/small-full', 'lora') + + +def test_training_namespace_conflicts_are_rejected(): + with pytest.raises(ValidationError, match='unique'): + ServerConfig(applications=[model_application(), model_application()]) + with pytest.raises(ValidationError, match='route conflicts'): + ServerConfig(applications=[model_application(), { + 'name': 'processor', 'import_path': 'processor', 'route_prefix': '/api/v1/model/Qwen/small', + }]) + with pytest.raises(ValidationError, match='models'): + ServerConfig(training={'models': [model_config()]}) + + +def test_existing_yaml_examples_keep_their_model_entries(monkeypatch): + from pathlib import Path + monkeypatch.setenv('TWINKLE_LOCAL_MODEL_PATH', '/tmp/model-cache') + paths = list(Path('cookbook/client/server').rglob('server_config*.yaml')) + paths += list(Path('cookbook/client/async_rl').glob('server_config.yaml')) + paths += list(Path('tests/server/config').glob('server_config*e2e*.yaml')) + paths += list(Path('tests/server/fixtures').glob('server_config_mock*.yaml')) + assert len(paths) == 9 + for path in paths: + config = ServerConfig.from_yaml(path) + models = [app for app in config.applications if app.import_path == 'model'] + assert len(models) == len(config.training_config.models) == 1 + model = config.training_config.models[0] + assert model.model_id == models[0].args.model_id + assert model.route_prefix == models[0].route_prefix + mode = model.config + original = models[0] + expected_min = original.args.min_replicas + assert mode.min_replicas == expected_min + assert ServerConfig.model_validate(config.to_yaml_dict()).training_config == config.training_config + + +def test_configured_lifecycle_overrides_and_full_instances(): + app = {'name': 'model', 'route_prefix': '/api/v1/model/small', 'import_path': 'model', 'args': { + 'model_id': 'weights', 'backend': 'mock', 'train_mode': 'full', + 'device_group': {'ranks': 2, 'device_type': 'CPU'}, 'device_mesh': {'dp_size': 2}, + 'min_replicas': 0, 'max_replicas': 3, 'idle_timeout_s': 17, + }, 'deployments': [{'num_replicas': 1}]} + config = ServerConfig(applications=[app]) + mode = config.training_config.models[0].config + assert (mode.num_devices, mode.min_replicas, mode.max_replicas, mode.idle_timeout_s) == (2, 0, 3, 17) + assert config.training_config.models[0].mode == 'full' + + +def test_separate_model_applications_select_lora_and_full(): + shared = model_application() + exclusive = copy.deepcopy(shared) + exclusive.update(name='model-small-full', route_prefix='/api/v1/model/Qwen/small-full') + exclusive['args'].update(train_mode='full', min_replicas=0, max_replicas=2, idle_timeout_s=60) + config = ServerConfig(applications=[shared, exclusive]) + catalog = config.training_config + assert catalog.get('Qwen/small', 'lora').capacity == 5 + full = catalog.get('Qwen/small-full', 'full') + assert full.capacity == 1 + assert (full.config.min_replicas, full.config.max_replicas, full.config.idle_timeout_s) == (0, 2, 60) + with pytest.raises(RequestRejectedError, match='does not provide full'): + catalog.get('Qwen/small', 'full') + assert ServerConfig.model_validate(config.to_yaml_dict()).training_config == config.training_config + + +def test_launcher_keeps_training_and_sampling_capabilities_separate(monkeypatch): + from ray import serve + from twinkle.server.launcher.server_launcher import ServerLauncher + config = ServerConfig(applications=[model_application(), { + 'name': 'gateway', 'import_path': 'server', 'route_prefix': '/api/v1', + }, { + 'name': 'sampler', 'import_path': 'sampler', 'route_prefix': '/api/v1/sampler/teacher', + 'args': {'model_id': 'teacher', 'device_group': {}, 'device_mesh': {}, 'sampler_type': 'mock'}, + }]) + captured = {} + def builder(**kwargs): + captured.update(kwargs) + return 'app' + launcher = ServerLauncher(config) + monkeypatch.setattr(launcher, '_resolve_builder', lambda path: builder) + monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) + launcher._deploy_application(config.applications[1]) + assert captured['training_models'] == ['Qwen/small'] + assert captured['supported_models'] == ['teacher'] + + +def configured_application(): + return {'name': 'model', 'route_prefix': '/api/v1/model/small', 'import_path': 'model', 'args': { + 'model_id': 'weights', 'backend': 'mock', 'train_mode': 'lora', 'max_length': 128, + 'device_group': {'ranks': 2, 'device_type': 'cpu'}, 'device_mesh': {'dp_size': 2}, + 'env_vars': {'USER_OPTION': 'args-wins'}, 'min_replicas': 1, 'max_replicas': 2, + }, 'deployments': [{ + 'name': 'ModelManagement', 'max_ongoing_requests': 17, + 'autoscaling_config': {'min_replicas': 1, 'max_replicas': 2, 'target_ongoing_requests': 16}, + 'ray_actor_options': {'num_cpus': 2, 'runtime_env': { + 'env_vars': {'USER_OPTION': 'deployment', 'OTHER_OPTION': 'retained'}, 'pip': ['example==1.0'] + }}, + }]} + + +@pytest.mark.parametrize('extra', [ + {'modes': None}, + {'modes': {'lora': {'num_devices': 2, 'device_mesh': {'dp_size': 2}}}}, + {'device_type': 'cpu'}, +]) +def test_model_yaml_rejects_alternative_resource_and_mode_syntax(extra): + app = configured_application() + app['args'].update(extra) + with pytest.raises(ValidationError, match='Extra inputs are not permitted'): + ServerConfig(applications=[app]) + + +@pytest.mark.parametrize('mode', ['lora', 'full']) +def test_adapter_lifecycle_yaml_reaches_controller_and_round_trips(mode): + app = model_application() + policy = {'adapter_timeout': 30, 'adapter_max_lifetime': 36000} + app['args'].update(train_mode=mode, adapter_config=policy, idle_timeout_s=600) + config = ServerConfig(applications=[app]) + resolved = config.training_config.models[0].config + assert resolved.adapter_config.model_dump() == policy + assert resolved.idle_timeout_s == 600 + restored = ServerConfig.model_validate(config.to_yaml_dict()) + assert restored.training_config == config.training_config + + +@pytest.mark.parametrize('policy', [ + {'adapter_timeout': 0}, {'adapter_timeout': math.inf}, {'adapter_timeout': math.nan}, + {'adapter_max_lifetime': -1}, {'adapter_max_lifetime': math.inf}, {'unknown': 30}, +]) +def test_invalid_adapter_lifecycle_is_rejected_before_allocation(policy): + app = model_application() + app['args']['adapter_config'] = policy + with pytest.raises(ValidationError): + ServerConfig(applications=[app]) + + +def test_configured_options_reach_cpu_router_without_changing_yaml(monkeypatch): + from types import SimpleNamespace + from twinkle.server.launcher.server_launcher import ServerLauncher + from twinkle.server.model.pool.actor import TrainingControllerActor + import ray + from ray import serve + config = ServerConfig(applications=[configured_application()]) + remote = SimpleNamespace(initialize=SimpleNamespace(remote=lambda: None)) + monkeypatch.setattr(TrainingControllerActor, 'options', + lambda **kwargs: SimpleNamespace(remote=lambda *args: remote)) + monkeypatch.setattr(ray, 'get', lambda *args, **kwargs: None) + captured = {} + + def router(*args): + captured.update(args[-1]) + return 'app' + + monkeypatch.setattr('twinkle.server.model.pool.router.build_training_router', router) + monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) + ServerLauncher(config)._deploy_training() + assert captured['max_ongoing_requests'] == 17 + assert captured['autoscaling_config']['target_ongoing_requests'] == 16 + actor = captured['ray_actor_options'] + assert actor['num_cpus'] == 2 and actor['runtime_env']['pip'] == ['example==1.0'] + assert actor['runtime_env']['env_vars']['USER_OPTION'] == 'args-wins' + assert actor['runtime_env']['env_vars']['OTHER_OPTION'] == 'retained' + assert config.to_yaml_dict()['applications'][0]['deployments'][0] == configured_application()['deployments'][0] + + +@pytest.mark.asyncio +async def test_configured_options_reach_runtime_with_one_controller_owned_replica(monkeypatch): + import importlib + from twinkle.server.state.training_records import InstanceRecord + from twinkle.server.model.pool.driver import RayInstanceDriver + from ray import serve + config = ServerConfig(applications=[configured_application()]) + model = config.training_config.models[0] + driver = RayInstanceDriver(config.training_config, None, 'test', 'controller') + instance = InstanceRecord(instance_id='instance', controller_epoch='epoch', base_model=model.name, + mode='lora', startup_deadline=0, devices=['0', '1']) + captured = {} + + def builder(**kwargs): + captured.update(kwargs) + return 'app' + + monkeypatch.setattr('twinkle.server.model.app.build_model_app', builder) + monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', lambda name: 'pg') + monkeypatch.setattr(serve, 'run', lambda *args, **kwargs: None) + await driver.deploy(instance, driver.config.get(model.name, 'lora')) + options = captured['deploy_options'] + assert options['num_replicas'] == 1 and 'autoscaling_config' not in options + assert options['max_ongoing_requests'] == 17 + assert options['ray_actor_options']['num_cpus'] == 2 + assert options['ray_actor_options']['runtime_env']['pip'] == ['example==1.0'] + assert captured['max_length'] == 128 + assert captured['nproc_per_node'] == 2 + assert captured['queue_config'] == model.config.queue_config + assert 'adapter_config' not in captured, 'Controller lifecycle policy must not reach model constructors' + assert not driver.runtime.deployments + + +def test_mutated_configuration_does_not_reuse_a_stale_catalog(): + config = ServerConfig(applications=[configured_application()]) + assert config.training_config.models[0].config.min_replicas == 1 + config.applications[0].args.min_replicas = 0 + assert config.training_config.models[0].config.min_replicas == 0 + + +def test_router_replica_limits_do_not_implicitly_set_training_instance_limits(): + app = configured_application() + app['args'].pop('min_replicas') + app['args'].pop('max_replicas') + model = ServerConfig(applications=[app]).training_config.models[0] + assert model.config.min_replicas == 0 and model.config.max_replicas is None + assert model.deploy_options['autoscaling_config']['min_replicas'] == 1 + + +def test_l20d_candidate_is_lora_only_and_bounds_total_devices(): + config = ServerConfig.from_yaml('cookbook/client/server/megatron/server_config.yaml') + model = config.training_config.models[0] + mode = model.config + assert model.mode == 'lora' + assert (mode.num_devices, mode.min_replicas, mode.max_replicas, mode.max_loras) == (2, 0, 2, 3) + assert mode.device_mesh.world_size() == 2 and mode.device_mesh.pp_size == 2 and mode.idle_timeout_s == 600 + sampler = next(app for app in config.applications if app.import_path == 'sampler') + assert sampler.args.device_group['ranks'] == 1 + assert sampler.args.device_mesh['dp_size'] == 1 and 'tp_size' not in sampler.args.device_mesh + assert 'tensor_parallel_size' not in sampler.args.engine_args + assert (sampler.args.min_replicas, sampler.args.max_replicas, sampler.args.target_ongoing_requests) == (2, 4, 16) + assert mode.num_devices * mode.max_replicas + sampler.args.device_group['ranks'] * sampler.args.max_replicas == 8 + assert mode.max_loras * mode.max_replicas == 6 + with pytest.raises(RequestRejectedError, match='does not provide full'): + config.training_config.get(model.name, 'full') + + +@pytest.mark.parametrize('mode', ['lora', 'full']) +@pytest.mark.parametrize('yaml_entrypoint', [False, True]) +def test_single_device_megatron_is_rejected_before_allocation(mode, yaml_entrypoint): + if yaml_entrypoint: + app = configured_application() + app['args'].update(backend='megatron', train_mode=mode, device_group={'ranks': 1, 'device_type': 'cuda'}, + device_mesh={'dp_size': 1}) + else: + model = model_config(mode) + model['backend'] = 'megatron' + model['config'] = {'num_devices': 1, 'device_mesh': {'dp_size': 1}} + with pytest.raises(ValidationError, match='at least 2 devices'): + if yaml_entrypoint: + ServerConfig(applications=[app]) + else: + TrainingConfig(models=[model]) + + +def test_transformers_still_accepts_single_device(): + app = model_application() + app['args'].update(device_group={'ranks': 1, 'device_type': 'cuda'}, device_mesh={'dp_size': 1}) + assert ServerConfig(applications=[app]).training_config.models[0].config.num_devices == 1 diff --git a/tests/server/contract/client_api_routes.json b/tests/server/contract/client_api_routes.json index ec8f176e7..5a376875b 100644 --- a/tests/server/contract/client_api_routes.json +++ b/tests/server/contract/client_api_routes.json @@ -45,6 +45,11 @@ "response": "GetServerCapabilitiesResponse", "statusCode": 200 }, + "GET /health": { + "body": [], + "response": null, + "statusCode": 200 + }, "GET /healthz": { "body": [], "response": "HealthResponse", @@ -55,6 +60,11 @@ "response": null, "statusCode": 200 }, + "GET /ping": { + "body": [], + "response": null, + "statusCode": 200 + }, "GET /training_runs": { "body": [], "response": "TrainingRunsResponse", @@ -127,6 +137,11 @@ "response": null, "statusCode": 200 }, + "POST /completions": { + "body": [], + "response": null, + "statusCode": 200 + }, "POST /create_model": { "body": [ "CreateModelRequest" @@ -148,6 +163,11 @@ "response": "CreateSessionResponse", "statusCode": 200 }, + "POST /embeddings": { + "body": [], + "response": null, + "statusCode": 200 + }, "POST /forward": { "body": [ "ForwardRequest" @@ -169,6 +189,11 @@ "response": "Any", "statusCode": 200 }, + "POST /infer": { + "body": [], + "response": null, + "statusCode": 200 + }, "POST /load_weights": { "body": [ "LoadWeightsRequest" @@ -183,6 +208,11 @@ "response": "Any", "statusCode": 200 }, + "POST /ping": { + "body": [], + "response": null, + "statusCode": 200 + }, "POST /retrieve_future": { "body": [ "FutureRetrieveRequest" @@ -589,6 +619,13 @@ "response": "SamplerCreateResponse", "statusCode": 200 }, + "POST /twinkle/encode": { + "body": [ + "EncodeRequest" + ], + "response": "TaskEnvelope", + "statusCode": 200 + }, "POST /twinkle/sample": { "body": [ "SampleRequest" diff --git a/tests/server/contract/test_client_api_contract.py b/tests/server/contract/test_client_api_contract.py index f543ee436..89d9fe00a 100644 --- a/tests/server/contract/test_client_api_contract.py +++ b/tests/server/contract/test_client_api_contract.py @@ -66,7 +66,7 @@ def test_new_client_types_importable(): for symbol in ('StrictRequest', 'ResponseModel', 'DataModel', 'backend_only'): assert hasattr(base, symbol) - for symbol in ('ErrorPayload', 'ErrorCategory', 'QueueStateLiteral'): + for symbol in ('ErrorPayload', 'ErrorCategory', 'QueueState'): assert hasattr(errors, symbol) diff --git a/tests/server/contract/test_error_wire.py b/tests/server/contract/test_error_wire.py index 41b394e42..d2ae84862 100644 --- a/tests/server/contract/test_error_wire.py +++ b/tests/server/contract/test_error_wire.py @@ -9,6 +9,9 @@ class _State: + async def assert_future_owner(self, request_id, token): + return None + async def get_future(self, request_id: str): return { 'status': 'failed', diff --git a/tests/server/contract/test_protocol_migration.py b/tests/server/contract/test_protocol_migration.py index 44bab4116..2d5e8a24b 100644 --- a/tests/server/contract/test_protocol_migration.py +++ b/tests/server/contract/test_protocol_migration.py @@ -10,7 +10,7 @@ CheckpointsListResponse ClientFeatures ClipGradAndStepRequest ClipGradAndStepResponse ClipGradNormRequest ClipGradNormResponse CreateModelRequest CreateRequest CreateResponse CreateSessionRequest CreateSessionResponse Cursor DataAppendRequest DataGetRequest DataModel DataPlaneForwardOnlyRequest DataPlaneForwardRequest DataPlaneSampleRequest -DataPutRequest DataRef DataReleaseRequest DataRowsResponse DeleteCheckpointResponse FieldRole ForwardBackwardResponse +DataPutRequest DataRef DataReleaseRequest EncodeRequest DataRowsResponse DeleteCheckpointResponse FieldRole ForwardBackwardResponse ForwardBackwardTaskRequest ForwardOnlyRequest ForwardRequest ForwardResponse GetServerCapabilitiesResponse GetTrainConfigsResponse HealthResponse LoadRequest LoadResponse LoraConfig LrStepRequest LrStepResponse ModelResult OkResponse ParsedCheckpointTwinklePath ProcessorCallRequest ProcessorCallResponse ProcessorCreateRequest @@ -27,9 +27,9 @@ """.split()) -def test_protocol_exports_match_pre_migration_snapshot() -> None: +def test_protocol_exports_match_committed_snapshot() -> None: import twinkle.protocol.types as types assert set(types.__all__) == _MIGRATED_NAMES - assert len(types.__all__) == 120 + assert len(types.__all__) == 121 assert 'ErrorResponse' not in types.__all__ diff --git a/tests/server/fixtures/server_config_mock.yaml b/tests/server/fixtures/server_config_mock.yaml index 67303ed5c..2a7711062 100644 --- a/tests/server/fixtures/server_config_mock.yaml +++ b/tests/server/fixtures/server_config_mock.yaml @@ -42,6 +42,8 @@ applications: args: backend: mock model_id: mock-model + min_replicas: 1 + max_replicas: 1 nproc_per_node: 1 device_group: name: model diff --git a/tests/server/fixtures/server_config_mock_redis.yaml b/tests/server/fixtures/server_config_mock_redis.yaml index 4c7297fe3..59c2ab436 100644 --- a/tests/server/fixtures/server_config_mock_redis.yaml +++ b/tests/server/fixtures/server_config_mock_redis.yaml @@ -41,6 +41,8 @@ applications: args: backend: mock model_id: mock-model + min_replicas: 1 + max_replicas: 1 nproc_per_node: 1 device_group: name: model diff --git a/tests/server/gateway/test_openai_handlers.py b/tests/server/gateway/test_openai_handlers.py index 678e83d09..4c43085f8 100644 --- a/tests/server/gateway/test_openai_handlers.py +++ b/tests/server/gateway/test_openai_handlers.py @@ -6,9 +6,11 @@ """ from __future__ import annotations +import asyncio import json +import httpx import pytest -from fastapi import FastAPI +from fastapi import FastAPI, Request from fastapi.testclient import TestClient from unittest.mock import AsyncMock, MagicMock @@ -38,6 +40,7 @@ def mock_gateway(): # Per-instance template cache; a fresh mock per test isolates it, so the # former module-global clear fixture is no longer needed. mock_self._template_initialized = set() + mock_self.owned_by = 'twinkle' app = FastAPI() _register_openai_routes(app, lambda: mock_self) @@ -50,29 +53,39 @@ def mock_gateway(): class TestChatCompletions: - def _make_sample_response(self, text='Hello!', stop_reason='stop'): + def _make_sample_response(self, text='Hello!', stop_reason='stop', status='completed'): """Build a mock Response matching what proxy_request returns. - The sampler replies with a completed TaskEnvelope: the payload rides under ``result``, - never at the top level. Mocking any other shape here would make the tests pass against a - contract the real server never sends. + ``/twinkle/sample`` always answers with a TaskEnvelope whose completed + result carries the samples; the raw samples dict is not the wire shape. """ from fastapi import Response - body = json.dumps({ - 'request_id': 'req-test', - 'status': 'completed', - 'result': { - 'samples': [{ - 'sequences': [{ - 'stop_reason': stop_reason, - 'tokens': list(range(len(text))), - 'decoded': text, - }] + result = { + 'samples': [{ + 'sequences': [{ + 'stop_reason': stop_reason, + 'tokens': list(range(len(text))), + 'decoded': text, }] - }, - }).encode() + }] + } + payload = {'request_id': 'req-mock', 'status': status} + if status == 'completed': + payload['result'] = result + body = json.dumps(payload).encode() return Response(content=body, status_code=200, media_type='application/json') + def _samples_result(self, text='Hello!'): + return { + 'samples': [{ + 'sequences': [{ + 'stop_reason': 'stop', + 'tokens': list(range(len(text))), + 'decoded': text, + }] + }] + } + def test_basic_chat_completion(self, mock_gateway): mock_self, app = mock_gateway mock_self.proxy.proxy_request.return_value = self._make_sample_response('Hi there!') @@ -220,6 +233,270 @@ def test_adapter_resolves_via_metadata(self, mock_gateway): call_kwargs = mock_self.proxy.proxy_request.call_args assert call_kwargs.kwargs.get('base_model') or call_kwargs.args[2] == 'Qwen/Qwen3.5-4B' + def test_non_stream_polls_queued_task_to_completion(self, mock_gateway): + """A queued sampler envelope is polled until completed, then translated.""" + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='queued') + mock_self.state.get_future = AsyncMock(return_value={ + 'status': 'completed', + 'result': self._samples_result('Polled!'), + }) + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 200 + assert resp.json()['choices'][0]['message']['content'] == 'Polled!' + mock_self.state.get_future.assert_awaited() + + def test_non_stream_failed_task_returns_error(self, mock_gateway): + """A failed sampler envelope becomes an OpenAI error, never an empty 200.""" + from fastapi import Response + mock_self, app = mock_gateway + envelope = { + 'request_id': 'req-failed', + 'status': 'failed', + 'result': None, + 'error': { + 'error': 'data/mask length mismatch at sample 0', + 'category': 'server', + 'error_code': 500, + 'request_id': 'req-failed', + }, + } + mock_self.proxy.proxy_request.return_value = Response( + content=json.dumps(envelope).encode(), status_code=200, media_type='application/json') + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 500 + assert 'mismatch' in resp.json()['error']['message'] + assert resp.json()['error']['type'] == 'server_error' + + def test_non_stream_wait_is_bounded(self, mock_gateway, monkeypatch): + """A task that never terminates returns 504 instead of polling forever.""" + monkeypatch.setenv('TWINKLE_OPENAI_SAMPLE_TIMEOUT', '0.01') + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='running') + mock_self.state.get_future = AsyncMock(return_value={'status': 'running'}) + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 504 + assert resp.json()['error']['type'] == 'timeout' + assert 'did not finish' in resp.json()['error']['message'] + + @pytest.mark.asyncio + @pytest.mark.parametrize('blocked_step', ['disconnect_check', 'initial_read', 'long_poll']) + async def test_non_stream_timeout_covers_each_wait_step(self, mock_gateway, monkeypatch, blocked_step): + from twinkle.server.gateway import openai_handlers + + monkeypatch.setenv('TWINKLE_OPENAI_SAMPLE_TIMEOUT', '0.01') + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='running') + mock_self.state.get_future = AsyncMock(return_value={'status': 'running'}) + started = asyncio.Event() + cancelled = asyncio.Event() + + async def blocked(*args, **kwargs): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + if blocked_step == 'disconnect_check': + monkeypatch.setattr(Request, 'is_disconnected', blocked) + elif blocked_step == 'initial_read': + mock_self.state.get_future.side_effect = blocked + else: + monkeypatch.setattr(openai_handlers, 'poll_future', blocked) + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client: + resp = await asyncio.wait_for( + client.post('/chat/completions', json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{'role': 'user', 'content': 'Hello'}], + }), timeout=1.0) + + assert started.is_set() + assert cancelled.is_set() + assert resp.status_code == 504 + assert resp.json()['error']['type'] == 'timeout' + assert 'did not finish' in resp.json()['error']['message'] + + @pytest.mark.asyncio + async def test_non_stream_polls_share_one_total_budget(self, mock_gateway, monkeypatch): + from twinkle.server.gateway import openai_handlers + from twinkle.server.gateway.use_cases import FuturePollResult + + monkeypatch.setenv('TWINKLE_OPENAI_SAMPLE_TIMEOUT', '0.05') + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='running') + mock_self.state.get_future = AsyncMock(return_value={'status': 'running'}) + calls = 0 + + async def slow_poll(*args, **kwargs): + nonlocal calls + calls += 1 + await asyncio.sleep(0.02) + status = 'completed' if calls == 10 else 'running' + return FuturePollResult(record={ + 'status': status, + 'result': self._samples_result('Late!') if status == 'completed' else None, + }, timed_out=status != 'completed') + + monkeypatch.setattr(openai_handlers, 'poll_future', slow_poll) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client: + resp = await asyncio.wait_for( + client.post('/chat/completions', json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{'role': 'user', 'content': 'Hello'}], + }), timeout=1.0) + + assert resp.status_code == 504 + assert resp.json()['error']['type'] == 'timeout' + assert calls < 10 + + @pytest.mark.asyncio + async def test_non_stream_disconnect_stops_before_state_reads(self, mock_gateway, monkeypatch): + monkeypatch.setattr(Request, 'is_disconnected', AsyncMock(return_value=True)) + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='running') + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client: + resp = await client.post('/chat/completions', json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{'role': 'user', 'content': 'Hello'}], + }) + + assert resp.status_code == 504 + assert 'disconnected' in resp.json()['error']['message'] + mock_self.state.get_future.assert_not_awaited() + + def test_non_stream_completed_task_needs_no_wait_budget(self, mock_gateway, monkeypatch): + monkeypatch.setenv('TWINKLE_OPENAI_SAMPLE_TIMEOUT', '0') + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response('Done!') + + with TestClient(app) as client: + resp = client.post('/chat/completions', json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{'role': 'user', 'content': 'Hello'}], + }) + + assert resp.status_code == 200 + assert resp.json()['choices'][0]['message']['content'] == 'Done!' + mock_self.state.get_future.assert_not_awaited() + + def test_non_stream_lost_record_is_a_server_error(self, mock_gateway): + """A vanished future record is the server's loss, not a client 404.""" + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.return_value = self._make_sample_response(status='queued') + mock_self.state.get_future = AsyncMock(return_value=None) + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 500 + assert 'lost' in resp.json()['error']['message'] + assert resp.json()['error']['type'] == 'server_error' + + def test_template_request_uses_only_declared_sampler_fields(self, mock_gateway): + """The sampler's strict schema rejects top-level model_id; do not send it.""" + from fastapi import Response + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.side_effect = [ + Response(content=b'{"status":"ok"}', status_code=200, media_type='application/json'), + self._make_sample_response('Hi'), + ] + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 200 + template_call = mock_self.proxy.proxy_request.call_args_list[0] + template_body = json.loads(template_call.kwargs['body_override']) + assert set(template_body) == {'template_cls', 'adapter_name'} + assert 'Qwen/Qwen3.5-4B' in mock_self._template_initialized + + def test_failed_template_call_is_not_cached(self, mock_gateway): + """A rejected set_template must not be marked initialized.""" + from fastapi import Response + mock_self, app = mock_gateway + mock_self.proxy.proxy_request.side_effect = [ + Response(content=b'{}', status_code=422, media_type='application/json'), + self._make_sample_response('Hi'), + ] + + client = TestClient(app) + resp = client.post( + '/chat/completions', + json={ + 'model': 'Qwen/Qwen3.5-4B', + 'messages': [{ + 'role': 'user', + 'content': 'Hello' + }], + }, + headers={'Authorization': 'Bearer key'}, + ) + + assert resp.status_code == 200 + assert 'Qwen/Qwen3.5-4B' not in mock_self._template_initialized + # ---------- /chat/completions streaming tests ------------------------------ # diff --git a/tests/server/utils/task_queue/__init__.py b/tests/server/integration/cs_tasks/__init__.py similarity index 100% rename from tests/server/utils/task_queue/__init__.py rename to tests/server/integration/cs_tasks/__init__.py diff --git a/tests/server/integration/cs_tasks/task_models.py b/tests/server/integration/cs_tasks/task_models.py new file mode 100644 index 000000000..4884fe6b3 --- /dev/null +++ b/tests/server/integration/cs_tasks/task_models.py @@ -0,0 +1,188 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Fixed bases and batches for the CS task acceptance (embedding, reranker, seq_cls). + +From the repository root: + python -m tests.server.integration.cs_tasks.task_models --tiny + python -m tests.server.integration.cs_tasks.task_models --model ms://Qwen/Qwen3.5-4B + python -m tests.server.integration.cs_tasks.task_models --model ms://Qwen/Qwen3.5-4B --mcore-native + +Checkpoints land in output/cs-models/{embedding,reranker,seq_cls}. Cold workers +load these complete bases, never a newly randomized classification head. +--mcore-native also writes output/cs-models-mcore/* with Qwen3.5's original +configuration and model.language_model.* keys for Megatron text-only training. +""" +import argparse +import json +from pathlib import Path +import shutil + +import torch +import transformers +from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer + +# Each group is (query, positive, negative). Labels mark group starts, not relevance. +EMBEDDING_TRIPLETS = [ + ('cat animal', 'cat pet', 'plane sky'), + ('dog animal', 'dog pet', 'stock market'), + ('apple fruit', 'apple sweet', 'car vehicle'), + ('boat vehicle', 'boat water', 'banana fruit'), +] +# Wire labels are integer 0/1; PointwiseRerankerLoss casts them to the score dtype. +RERANKER_PAIRS = [ + ('cat animal', 'cat pet', 1), ('cat animal', 'plane sky', 0), + ('dog animal', 'dog pet', 1), ('dog animal', 'stock market', 0), + ('apple fruit', 'apple sweet', 1), ('apple fruit', 'car vehicle', 0), + ('boat vehicle', 'boat water', 1), ('boat vehicle', 'banana fruit', 0), + ('cat animal', 'cat pet', 1), ('cat animal', 'plane sky', 0), + ('dog animal', 'dog pet', 1), ('dog animal', 'stock market', 0), +] +# Class IDs match the three-output head: animal=0, fruit=1, vehicle=2. +SEQ_CLS_SAMPLES = [ + ('cat animal', 0), ('dog pet', 0), ('apple fruit', 1), ('banana sweet', 1), + ('car vehicle', 2), ('boat water', 2), ('cat pet', 0), ('dog animal', 0), + ('orange fruit', 1), ('apple sweet', 1), ('train vehicle', 2), ('boat vehicle', 2), +] + + +def encoded_batch(task, tokenizer): + """Return twelve right-paddable features; Megatron DP=2 keeps whole embedding triplets per rank.""" + if task == 'embedding': + texts = [(text, int(index == 0)) for group in EMBEDDING_TRIPLETS for index, text in enumerate(group)] + elif task == 'reranker': + texts = [('Query: ' + query + '\nDocument: ' + document, label) for query, document, label in RERANKER_PAIRS] + else: + texts = SEQ_CLS_SAMPLES + rows = [] + for text, label in texts: + row = dict(tokenizer(text, truncation=True, max_length=128)) + row['labels'] = [label] + row['position_ids'] = list(range(len(row['input_ids']))) + rows.append(row) + return rows + + +def tiny_backbone(directory, tokenizer_source='ms://Qwen/Qwen3.5-4B', seed=42): + from transformers import Qwen2Config, Qwen2ForCausalLM + from twinkle.hub import HubOperation + + tokenizer_path = HubOperation.download_model(tokenizer_source) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=False) + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + config = Qwen2Config(vocab_size=len(tokenizer), hidden_size=64, intermediate_size=128, + num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, + max_position_embeddings=512, pad_token_id=tokenizer.pad_token_id, + bos_token_id=tokenizer.bos_token_id, eos_token_id=tokenizer.eos_token_id, + attention_dropout=0., tie_word_embeddings=False) + Qwen2ForCausalLM(config).save_pretrained(directory) + tokenizer.save_pretrained(directory) + return str(directory) + + +def _mcore_native_key(key): + if key.startswith('model.') and not key.startswith('model.language_model.'): + return 'model.language_model.' + key[len('model.'):] + return key + + +def _save_mcore_native(source_dir, native_dir, config, width): + from safetensors.torch import load_file, save_file + + source_dir, native_dir = Path(source_dir), Path(native_dir) + shutil.copytree(source_dir, native_dir, dirs_exist_ok=True) + for file in native_dir.glob('model*.safetensors'): + state = load_file(str(file)) + save_file({_mcore_native_key(key): value for key, value in state.items()}, str(file)) + index_path = native_dir / 'model.safetensors.index.json' + if index_path.is_file(): + index = json.loads(index_path.read_text()) + index['weight_map'] = {_mcore_native_key(key): value for key, value in index['weight_map'].items()} + index_path.write_text(json.dumps(index, indent=2)) + native_config = type(config).from_dict(config.to_dict()) + if width is not None: + native_config.architectures = ['Qwen3_5ForSequenceClassification'] + native_config.num_labels = width + native_config.problem_type = 'regression' if width == 1 else 'single_label_classification' + native_config.tie_word_embeddings = False + native_config.text_config.tie_word_embeddings = False + native_config.language_model_only = True + native_config.save_pretrained(native_dir) + + +def prepare_models(model_id, output_dir, *, seed=42, mcore_native=False): + from twinkle.hub import HubOperation + + directory = Path(output_dir) + source = HubOperation.download_model(str(model_id)) + config = AutoConfig.from_pretrained(source, trust_remote_code=False) + if config.model_type == 'qwen3_5': + from packaging.version import Version + if Version(transformers.__version__) < Version('5.8'): + raise RuntimeError('Qwen3.5 preparation requires transformers>=5.8; use --tiny or upgrade') + if mcore_native and config.model_type != 'qwen3_5': + raise ValueError('--mcore-native currently supports Qwen3.5 checkpoints') + tokenizer = AutoTokenizer.from_pretrained(source, trust_remote_code=False) + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + if tokenizer.pad_token_id is None: + raise ValueError('The tokenizer needs a pad or EOS token') + task_directories = [directory / task for task in ('embedding', 'reranker', 'seq_cls')] + if mcore_native: + task_directories += [Path(str(directory) + '-mcore')] + for path in task_directories: + if path.exists() and any(path.iterdir()): + raise ValueError(f'Output directory is not empty: {path}; choose a new directory or clear it first') + base = getattr(transformers, config.architectures[0]).from_pretrained(source, torch_dtype='auto') + backbone = base.model + text_config = getattr(config, 'text_config', config) + if hasattr(backbone, 'language_model'): + backbone = backbone.language_model + text_config.pad_token_id = tokenizer.pad_token_id + native_directory = Path(str(directory) + '-mcore') + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + causal = transformers.AutoModelForCausalLM.from_config(text_config) + causal.to(next(backbone.parameters()).dtype) + causal.model.load_state_dict(backbone.state_dict(), strict=True) + causal.lm_head.load_state_dict(base.lm_head.state_dict(), strict=True) + causal.save_pretrained(directory / 'embedding') + tokenizer.save_pretrained(directory / 'embedding') + if mcore_native: + _save_mcore_native(directory / 'embedding', native_directory / 'embedding', config, None) + del causal + for task, width in [('reranker', 1), ('seq_cls', 3)]: + head_config = type(text_config).from_dict(text_config.to_dict()) + head_config.num_labels = width + head_config.problem_type = 'regression' if width == 1 else 'single_label_classification' + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + model = AutoModelForSequenceClassification.from_config(head_config) + model.to(next(backbone.parameters()).dtype) + model.model.load_state_dict(backbone.state_dict(), strict=True) + model.save_pretrained(directory / task) + tokenizer.save_pretrained(directory / task) + if mcore_native: + _save_mcore_native(directory / task, native_directory / task, config, width) + del model + return {task: str((directory / task).resolve()) for task in ('embedding', 'reranker', 'seq_cls')} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--model', help='Source model; omit it to build a tiny model with the Qwen3.5 tokenizer') + parser.add_argument('--output-dir', default='output/cs-models') + parser.add_argument('--seed', type=int, default=42) + parser.add_argument('--tiny', action='store_true', help='Build a small Qwen2 backbone with the source tokenizer') + parser.add_argument('--mcore-native', action='store_true', + help='Also write Qwen3.5 Megatron-native text-only checkpoints') + args = parser.parse_args() + tokenizer_source = args.model or 'ms://Qwen/Qwen3.5-4B' + source = tiny_backbone(Path(args.output_dir) / 'source', tokenizer_source, args.seed) \ + if args.tiny or args.model is None else args.model + prepare_models(source, args.output_dir, seed=args.seed, mcore_native=args.mcore_native) + + +if __name__ == '__main__': + main() diff --git a/tests/server/integration/cs_tasks/task_server_megatron.yaml b/tests/server/integration/cs_tasks/task_server_megatron.yaml new file mode 100644 index 000000000..c695fd417 --- /dev/null +++ b/tests/server/integration/cs_tasks/task_server_megatron.yaml @@ -0,0 +1,64 @@ +# Prepare the persisted bases first. Each Megatron instance uses two DP devices. +# Forward tasks select losses/pooling; task_type selects the persisted model head. +http_options: + host: 127.0.0.1 + port: 8000 +telemetry: + enabled: false +persistence: + mode: memory +resources: + startup_timeout_s: 600 + shutdown_timeout_s: 120 +applications: +- name: server + import_path: server + route_prefix: /api/v1 + args: + supported_models: [cs-embedding, cs-reranker, cs-seq_cls] + deployments: + - name: GatewayServer + ray_actor_options: {num_cpus: 0.1} +- name: embedding + import_path: model + route_prefix: /api/v1/model/cs-embedding + args: &model_args + backend: megatron + model_id: output/cs-models/embedding + train_mode: lora + task_type: causal_lm + language_model_only: true + max_loras: 2 + max_r: 8 + max_length: 256 + min_replicas: 0 + max_replicas: 1 + idle_timeout_s: 60 + nproc_per_node: 2 + device_group: {name: model, ranks: 2, device_type: cuda} + device_mesh: {device_type: cuda, dp_size: 2} + target_modules: [linear_qkv] + queue_config: {execution_timeout: 600} + deployments: &deployments + - name: ModelManagement + ray_actor_options: {num_cpus: 0.1} +- name: reranker + import_path: model + route_prefix: /api/v1/model/cs-reranker + args: + <<: *model_args + model_id: output/cs-models/reranker + task_type: seq_cls + num_labels: 1 + target_modules: [linear_qkv, output_layer] + deployments: *deployments +- name: seq-cls + import_path: model + route_prefix: /api/v1/model/cs-seq_cls + args: + <<: *model_args + model_id: output/cs-models/seq_cls + task_type: seq_cls + num_labels: 3 + target_modules: [linear_qkv, output_layer] + deployments: *deployments diff --git a/tests/server/integration/cs_tasks/task_server_transformers.yaml b/tests/server/integration/cs_tasks/task_server_transformers.yaml new file mode 100644 index 000000000..8f16566b5 --- /dev/null +++ b/tests/server/integration/cs_tasks/task_server_transformers.yaml @@ -0,0 +1,65 @@ +# Prepare the persisted bases with task_models.py before starting this server from the repository root. +# Each public alias has one frozen base and two independent LoRA tenant slots. +http_options: + host: 127.0.0.1 + port: 8000 +telemetry: + enabled: false +persistence: + mode: memory +resources: + startup_timeout_s: 600 + shutdown_timeout_s: 120 +applications: +- name: server + import_path: server + route_prefix: /api/v1 + args: + supported_models: [cs-embedding, cs-reranker, cs-seq_cls] + deployments: + - name: GatewayServer + ray_actor_options: {num_cpus: 0.1} +- name: embedding + import_path: model + route_prefix: /api/v1/model/cs-embedding + args: &model_args + backend: transformers + model_id: output/cs-models/embedding + train_mode: lora + task_type: causal_lm + gradient_checkpointing: false + mixed_precision: 'no' + max_loras: 2 + max_r: 8 + max_length: 256 + min_replicas: 0 + max_replicas: 1 + idle_timeout_s: 60 + nproc_per_node: 1 + device_group: {name: model, ranks: 1, device_type: cuda} + device_mesh: {device_type: cuda, dp_size: 1} + target_modules: [q_proj, v_proj] + queue_config: {execution_timeout: 600} + deployments: &deployments + - name: ModelManagement + ray_actor_options: {num_cpus: 0.1} +- name: reranker + import_path: model + route_prefix: /api/v1/model/cs-reranker + args: + <<: *model_args + model_id: output/cs-models/reranker + task_type: seq_cls + num_labels: 1 + target_modules: [q_proj, v_proj, score] + deployments: *deployments +- name: seq-cls + import_path: model + route_prefix: /api/v1/model/cs-seq_cls + args: + <<: *model_args + model_id: output/cs-models/seq_cls + task_type: seq_cls + num_labels: 3 + target_modules: [q_proj, v_proj, score] + deployments: *deployments diff --git a/tests/server/integration/e2e_helpers.py b/tests/server/integration/e2e_helpers.py index 2eca95f0d..8ac1c6905 100644 --- a/tests/server/integration/e2e_helpers.py +++ b/tests/server/integration/e2e_helpers.py @@ -290,3 +290,34 @@ def assert_metrics_valid(metrics: Any, label: str) -> None: assert math.isfinite(val), ( f'[{label}] Metric {key}={val} is not finite!') log(f'[{label}] Metrics valid: {metrics}') + + +def sampled_datum(prompt_tokens, sample, advantage): + from tinker import types + + tokens = prompt_tokens + list(sample.tokens) + prompt_positions = len(prompt_tokens) - 1 + size = len(tokens) - 1 + + def tensor(values, dtype): + return types.TensorData(data=values, dtype=dtype, shape=[size]) + + return types.Datum( + model_input=types.ModelInput.from_ints(tokens[:-1]), + loss_fn_inputs={ + 'target_tokens': tensor(tokens[1:], 'int64'), + 'weights': tensor([0.0] * prompt_positions + [1.0] * len(sample.tokens), 'float32'), + 'logprobs': tensor([0.0] * prompt_positions + list(sample.logprobs), 'float32'), + 'advantages': tensor([0.0] * prompt_positions + [advantage] * len(sample.tokens), 'float32'), + }, + ) + + +def completion_logprobs(client, datum, timeout): + result = client.forward([datum], 'cross_entropy').result(timeout=timeout) + values = np.asarray(result.loss_fn_outputs[0]['logprobs'].tolist(), dtype=np.float64) + weights = np.asarray(datum.loss_fn_inputs['weights'].tolist()) + assert len(values) >= len(weights) + values = values[:len(weights)][weights != 0] + assert values.size and np.isfinite(values).all() + return values diff --git a/tests/server/integration/on_demand_training/__init__.py b/tests/server/integration/on_demand_training/__init__.py new file mode 100644 index 000000000..0bb7387ac --- /dev/null +++ b/tests/server/integration/on_demand_training/__init__.py @@ -0,0 +1 @@ +"""On-demand training integration probes.""" diff --git a/tests/server/integration/on_demand_training/acceptance_support.py b/tests/server/integration/on_demand_training/acceptance_support.py new file mode 100644 index 000000000..c14827c6c --- /dev/null +++ b/tests/server/integration/on_demand_training/acceptance_support.py @@ -0,0 +1,171 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real SDK operations and bounded resource assertions for acceptance probes.""" +import json +import time + +import requests + +from tests.server.integration.on_demand_training.on_demand_sdk import unload_tinker_model + + +def checked_logps(values): + """Reject empty/default metric values before numerical checkpoint comparisons.""" + import numpy as np + rows = [np.asarray(row, dtype=float).reshape(-1) for row in values] + values = np.concatenate(rows) if rows else np.array([], dtype=float) + assert values.size and np.isfinite(values).all(), 'Empty or non-finite checkpoint probe' + assert (values < 0).any(), 'Checkpoint probe contains no nontrivial token log probabilities' + return values + + +class TrainingHandle: + """A small fixed-token workload with optimizer checkpoint support.""" + + def __init__(self, base_url, model, protocol, label, mode='lora'): + from twinkle_client import init_tinker_client + from twinkle_client.http import ClientTransport + from twinkle_client.http.context import ClientContext + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + + self.base_url, self.model, self.protocol = base_url, model, protocol + self.label, self.mode = label, mode + self.token = 'acceptance-' + protocol + self.transport = None + self.closed = False + self.consumed_samples = 0 + if protocol == 'twinkle': + from peft import LoraConfig + response = requests.post(base_url + '/api/v1/twinkle/create_session', json={}, timeout=30) + response.raise_for_status() + self.transport = ClientTransport(ClientContext( + base_url=base_url, api_key=self.token, session_id=response.json()['session_id'])) + self.client = MultiLoraTransformersModel(model, transport=self.transport) + self.client.add_adapter_to_model( + label, LoraConfig(r=8, target_modules='all-linear') if mode == 'lora' else None, + gradient_accumulation_steps=2) + self.identity = self.client._training_ids[label] + self.configure() + else: + init_tinker_client() + from tinker import ServiceClient + self.service = ServiceClient(base_url=base_url, api_key=self.token) + self.client = (self.service.create_lora_training_client(base_model=model, rank=8) + if mode == 'lora' else self.service.create_full_training_client(base_model=model)) + self.identity = self.client.model_id + self.inputs = [dict(input_ids=[100, 200, 300, 400], labels=[200, 300, 400, 500], + position_ids=list(range(4)), attention_mask=[1] * 4) for _ in range(2)] + + def configure(self): + if self.protocol == 'twinkle': + # Qwen3.5 is multimodal: this hook connects the visual branch even + # for text-only batches, as in the established SFT helpers. + self.client.set_template('Qwen3_5Template', max_length=512) + self.client.set_loss('CrossEntropyLoss') + self.client.set_optimizer('Adam', lr=1e-4) + self.client.set_processor('InputProcessor', padding_side='right') + + def datums(self): + from twinkle.server.model.tinker_datum import input_feature_to_datum + return [input_feature_to_datum(feature) for feature in self.inputs] + + def step(self): + if self.protocol == 'twinkle': + for _ in range(3): + self.client.forward_backward(self.inputs, gradient_accumulation_steps=2) + self.client.clip_grad_and_step(gradient_accumulation_steps=2) + self.consumed_samples += len(self.inputs) + else: + from tinker import types + self.client.forward_backward(self.datums(), 'cross_entropy').result() + self.client.optim_step(types.AdamParams(learning_rate=1e-4)).result() + self.consumed_samples += len(self.inputs) + + def evaluate(self): + if self.protocol == 'twinkle': + result = self.client.forward_only(self.inputs).result + values = result['logps'] + else: + result = self.client.forward(self.datums(), 'cross_entropy').result() + values = [output['logprobs'].tolist() for output in result.loss_fn_outputs] + return checked_logps(values) + + def save(self, name): + if self.protocol == 'twinkle': + path = self.client.save( + name=name, save_optimizer=True, consumed_train_samples=self.consumed_samples).twinkle_path + else: + path = self.client.save_state(name).result().path + assert path, 'Checkpoint path is empty' + return path + + def resume(self, path): + if self.protocol == 'twinkle': + progress = self.client.resume_from_checkpoint(path) + self.consumed_samples = progress['consumed_train_samples'] + return progress + old = self.identity + self.client = self.service.create_training_client_from_state_with_optimizer(path=path) + self.identity = self.client.model_id + unload_tinker_model(self.service, old) + + def assert_stale(self): + try: + if self.protocol == 'twinkle': + self.client.get_train_configs().result + else: + self.client.get_info() + except Exception as exc: + assert any(term in str(exc).lower() for term in ('404', 'not found', 'handle', 'active', 'unavailable')), exc + else: + raise AssertionError('An old training handle remained usable') + + def close(self, stale=False): + if self.closed: + return + try: + if not stale: + if self.protocol == 'twinkle': + self.client.remove_adapter() + else: + unload_tinker_model(self.service, self.identity) + finally: + if self.transport is not None: + self.transport.close() + elif hasattr(self, 'service'): + self.service.holder.close() + self.closed = True + + +def wait_until(predicate, timeout, label): + deadline = time.monotonic() + timeout + while True: + result = predicate() + if result: + return result + if time.monotonic() >= deadline: + raise TimeoutError(label) + time.sleep(0.5) + + +def assert_released(controller, namespace, redis_url=None, prefix=None): + import ray + from ray import serve + from ray.util.placement_group import placement_group_table + diagnostic = ray.get(controller.diagnostics.remote(), timeout=30) + assert all(row['state'] == 'STOPPED' for row in diagnostic['instances']), diagnostic + # Training teardown must not require a healthy sampling pool to disappear. + owned_prefix = 'resources-' + namespace + '-controller_' + assert not [name for name in serve.status().applications if name.startswith(owned_prefix)] + assert not [row for row in placement_group_table().values() + if row.get('name', '').startswith(owned_prefix) and row.get('state') != 'REMOVED'] + if redis_url: + import redis + client = redis.Redis.from_url(redis_url, decode_responses=True, socket_timeout=5) + try: + assert not list(client.scan_iter(match=prefix + 'model::training_*')) + assert not list(client.scan_iter(match=prefix + 'seq_ledger::training_*')) + for key in client.scan_iter(match=prefix + 'token_count::*'): + assert int(json.loads(client.get(key)) or 0) == 0, 'Quota drift' + finally: + client.close() + return True diff --git a/tests/server/integration/on_demand_training/http_probe.py b/tests/server/integration/on_demand_training/http_probe.py new file mode 100644 index 000000000..ebb01d86c --- /dev/null +++ b/tests/server/integration/on_demand_training/http_probe.py @@ -0,0 +1,297 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""HTTP integration with both installed SDKs and a CPU mock model.""" +import json +import os +import ray +import requests +import socket +import tempfile +import time +import uuid +from pathlib import Path +from on_demand_sdk import unload_tinker_model +from ray import serve + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher +from twinkle_client import init_tinker_client, init_twinkle_client +from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + + +def main(): + started = time.monotonic() + + def phase(name, **details): + print(json.dumps({'phase': name, 'elapsed_s': round(time.monotonic() - started, 2), **details}), flush=True) + + namespace = 'http_' + uuid.uuid4().hex[:8] + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + os.environ['TWINKLE_STATE_RAY_NAMESPACE'] = namespace + os.environ['TWINKLE_DEFAULT_SAVE_DIR'] = tempfile.mkdtemp(prefix='twinkle-training-http-') + os.environ['TWINKLE_LONG_POLL_TIMEOUT'] = '2' + os.environ['TWINKLE_PERSISTENCE_MODE'] = 'memory' + os.environ['PYTHONPATH'] = os.pathsep.join(filter(None, ( + str(Path.cwd() / 'src'), os.environ.get('PYTHONPATH')))) + ray.init( + address='local', + namespace=namespace, + num_cpus=8, + _temp_dir='/tmp/twh-' + namespace[-8:], + include_dashboard=False, + log_to_driver=False) + config = ServerConfig.model_validate({ + 'ray_namespace': + namespace, + 'http_options': { + 'host': '127.0.0.1', + 'port': port + }, + 'resources': { + 'startup_timeout_s': 600, + 'shutdown_timeout_s': 180 + }, + 'applications': [{ + 'name': 'small', + 'import_path': 'model', + 'route_prefix': '/api/v1/model/small', + 'args': { + 'model_id': 'small', + 'backend': 'mock', + 'train_mode': 'lora', + 'device_group': {'ranks': 1, 'device_type': 'cpu'}, + 'device_mesh': {'dp_size': 1}, + 'max_loras': 2, + 'max_replicas': 1, + 'adapter_config': {'adapter_timeout': 30}, + }, + }, { + 'name': 'other', + 'import_path': 'model', + 'route_prefix': '/api/v1/model/other', + 'args': { + 'model_id': 'small', + 'backend': 'mock', + 'train_mode': 'full', + 'device_group': {'ranks': 1, 'device_type': 'cpu'}, + 'device_mesh': {'dp_size': 1}, + 'min_replicas': 0, + 'max_replicas': 1, + 'adapter_config': {'adapter_timeout': 30}, + }, + 'deployments': [{ + 'max_ongoing_requests': 8, + 'autoscaling_config': {'min_replicas': 1, 'max_replicas': 1, 'target_ongoing_requests': 4}, + 'ray_actor_options': {'num_cpus': 0.1, 'runtime_env': {'env_vars': {'LEGACY_CONFIG_PROBE': '1'}}}, + }], + }, { + 'name': 'routed', 'import_path': 'model', 'route_prefix': '/api/v1/model/routed', + 'args': { + 'model_id': 'small', 'backend': 'mock', 'train_mode': 'lora', + 'device_group': {'ranks': 1, 'device_type': 'cpu'}, 'device_mesh': {'dp_size': 1}, + 'max_loras': 1, 'max_replicas': 2, + 'adapter_config': {'adapter_timeout': 30}, + }, + }, { + 'name': 'gateway', + 'import_path': 'server', + 'route_prefix': '/api/v1', + 'deployments': [{ + 'ray_actor_options': { + 'num_cpus': 0.1 + } + }] + }], + }) + launcher = ServerLauncher(config) + launcher._start_serve() + base = f'http://127.0.0.1:{port}' + native = fresh_native = service = None + + def wait_epoch(controller, old_epoch): + deadline = time.monotonic() + 90 + last_error = None + while time.monotonic() < deadline: + try: + diagnostic = ray.get(controller.diagnostics.remote(), timeout=2) + if diagnostic['epoch'] != old_epoch: + return diagnostic + except ray.exceptions.RayError as exc: + last_error = exc + time.sleep(0.1) + raise AssertionError(f'Controller did not recover: {last_error}') + + try: + for spec in config.applications: + if spec.import_path != 'model': + launcher._deploy_application(spec) + launcher._deploy_training() + controller = launcher._training_controller + assert ray.get(controller.diagnostics.remote())['instances'] == [] + phase('server_ready') + native = init_twinkle_client(base_url=base, api_key='probe-owner') + transport = native.transport + first = MultiLoraTransformersModel('small', transport=transport) + second = MultiLoraTransformersModel('small', transport=transport) + from peft import LoraConfig + + def assert_separate_instances(identities): + diagnostic = ray.get(controller.diagnostics.remote()) + owners = {identity: item['instance_id'] for item in diagnostic['instances'] + for identity in item['training_ids'] if identity in identities} + assert set(owners) == set(identities) and len(set(owners.values())) == 2, owners + return owners + + routed = [MultiLoraTransformersModel('routed', transport=transport) for _ in range(2)] + for client in routed: + client.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) + native_owners = assert_separate_instances([client._training_ids['same'] for client in routed]) + for _ in range(3): + for client in reversed(routed): + assert json.loads(client.get_train_configs().result) == {} + client.zero_grad() + for client in routed: + client.remove_adapter() + phase('native_routing_verified') + first.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) + instance = next( + item for item in ray.get(controller.diagnostics.remote())['instances'] + if first._training_ids['same'] in item['training_ids']) + bypass = requests.post( + base + f"/api/v1/training-runtime/{instance['instance_id']}/twinkle/add_adapter_to_model", + json={ + 'adapter_name': 'bypass', + 'config': None + }, + headers={'Authorization': 'Bearer probe-owner'}, + timeout=10) + assert bypass.status_code == 400 and 'catalog endpoint' in bypass.text, bypass.text + second.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) + old_id = first._training_ids['same'] + assert old_id != second._training_ids['same'] + assert json.loads(first.get_train_configs().result) == {} + first.zero_grad() + first.remove_adapter() + first.add_adapter_to_model('same', LoraConfig(r=8, target_modules='all-linear')) + assert old_id != first._training_ids['same'] + first.remove_adapter() + second.remove_adapter() + native_lora = MultiLoraTransformersModel('small', transport=transport) + native_lora.add_adapter_to_model('native-lora', LoraConfig(r=8, target_modules='all-linear')) + native_lora.zero_grad() + native_lora.remove_adapter() + phase('native_lifecycle_verified') + init_tinker_client() + from tinker import ServiceClient + service = ServiceClient(base_url=base, api_key='probe-tinker') + routed_tinker = [service.create_lora_training_client(base_model='routed', rank=8) for _ in range(2)] + tinker_owners = assert_separate_instances([client.model_id for client in routed_tinker]) + for _ in range(3): + for client in reversed(routed_tinker): + assert client.get_info().model_id == client.model_id + for client in routed_tinker: + assert unload_tinker_model(service, client.model_id).model_id == client.model_id + phase('tinker_routing_verified') + shared1 = service.create_lora_training_client(base_model='small', rank=8) + shared2 = service.create_lora_training_client(base_model='small', rank=8) + assert shared1.model_id != shared2.model_id + assert shared1.get_info().is_lora + assert shared2.get_info().model_name == 'small' + from tinker import types + datum = types.Datum( + model_input=types.ModelInput.from_ints([1, 2, 3]), + loss_fn_inputs={ + 'target_tokens': types.TensorData(data=[2, 3, 4], dtype='int64', shape=[3]), + 'weights': types.TensorData(data=[1.0, 1.0, 1.0], dtype='float32', shape=[3]), + }) + assert shared1.forward_backward([datum], 'cross_entropy').result().metrics + shared1.optim_step(types.AdamParams(learning_rate=1e-4)).result() + exclusive = service.create_full_training_client(base_model='other') + assert not exclusive.get_info().is_lora + assert exclusive.forward_backward([datum], 'cross_entropy').result().metrics + exclusive.optim_step(types.AdamParams(learning_rate=1e-4)).result() + assert unload_tinker_model(service, exclusive.model_id).model_id == exclusive.model_id + assert shared1.get_info().is_lora and shared2.get_info().is_lora + try: + service.create_lora_training_client(base_model='small', rank=8) + raise AssertionError('A third LoRA user bypassed configured capacity') + except Exception as exc: + assert 'replica limit' in str(exc).lower(), exc + assert unload_tinker_model(service, shared2.model_id).model_id == shared2.model_id + replacement = service.create_lora_training_client(base_model='small', rank=8) + assert replacement.model_id != shared2.model_id + assert shared1.get_info().is_lora + old_epoch = ray.get(controller.diagnostics.remote())['epoch'] + ray.kill(controller, no_restart=False) + diagnostic = wait_epoch(controller, old_epoch) + assert not [item for item in diagnostic['instances'] if item['state'] != 'STOPPED'] + phase('controller_recovered') + stale = requests.post( + base + '/api/v1/get_info', + json={ + 'model_id': shared1.model_id, + 'type': 'get_info' + }, + headers={ + 'Authorization': 'Bearer probe-tinker', + 'x-request-id': uuid.uuid4().hex + }, + timeout=10) + assert stale.status_code == 404, stale.text + before = ray.available_resources().get('CPU', 0) + for index in range(3): + model = MultiLoraTransformersModel('small' if index % 2 == 0 else 'other', transport=transport) + model.add_adapter_to_model('cycle', LoraConfig(r=8, target_modules='all-linear') + if index % 2 == 0 else None) + model.zero_grad() + model.remove_adapter() + phase('cycles', completed=index + 1) + assert ray.available_resources().get('CPU', 0) >= before - 0.02 + final = ray.get(controller.diagnostics.remote()) + assert all(item['state'] == 'STOPPED' for item in final['instances']) + live = MultiLoraTransformersModel('other', transport=transport) + live.add_adapter_to_model('state-loss', None) + native.close() + ray.kill(ray.get_actor('twinkle_state_actor', namespace=namespace), no_restart=False) + after_loss = wait_epoch(controller, final['epoch']) + assert not [item for item in after_loss['instances'] if item['state'] != 'STOPPED'] + fresh_native = init_twinkle_client(base_url=base, api_key='probe-owner') + fresh_transport = fresh_native.transport + fresh = MultiLoraTransformersModel('other', transport=fresh_transport) + fresh.add_adapter_to_model('state-loss', None) + assert fresh._training_ids['state-loss'] != live._training_ids['state-loss'] + fresh.remove_adapter() + phase('state_loss_recovered') + evidence = { + 'passed': True, + 'native_recreation': True, + 'native_stable_routing': native_owners, + 'tinker_stable_routing': tinker_owners, + 'native_lora': True, + 'tinker_shared': True, + 'tinker_full': True, + 'restart': True, + 'state_loss': True, + 'runtime_lifecycle_rejected': True, + 'cycles': 3, + 'instances': len(final['instances']) + } + if result_path := os.environ.get('TWINKLE_TEST_HTTP_RESULT'): + Path(result_path).write_text(json.dumps(evidence)) + print(json.dumps(evidence), flush=True) + finally: + for client in (native, fresh_native): + if client is not None: + client.close() + if service is not None: + service.holder.close() + if getattr(launcher, '_training_controller', None) is not None: + ray.get(launcher._training_controller.shutdown.remote(), timeout=180) + ray.kill(launcher._training_controller, no_restart=True) + serve.shutdown() + ray.shutdown() + + +if __name__ == '__main__': + main() diff --git a/tests/server/integration/on_demand_training/managed_redis.py b/tests/server/integration/on_demand_training/managed_redis.py new file mode 100644 index 000000000..c832bbaee --- /dev/null +++ b/tests/server/integration/on_demand_training/managed_redis.py @@ -0,0 +1,65 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""An owned Redis subprocess for safe, durable restart fault injection.""" +import socket +import subprocess +import time +from pathlib import Path + + +class ManagedRedis: + """Never stop an externally supplied Redis server.""" + + def __init__(self, directory): + self.directory = Path(directory).resolve() + self.directory.mkdir(parents=True, exist_ok=True) + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + self.port = listener.getsockname()[1] + self.url = f'redis://127.0.0.1:{self.port}/0' + self.process = None + self.log = None + + def start(self): + import redis + assert self.process is None + self.log = (self.directory / 'redis.log').open('a') + self.process = subprocess.Popen([ + 'redis-server', '--bind', '127.0.0.1', '--port', str(self.port), + '--dir', str(self.directory), '--appendonly', 'yes', '--appendfsync', 'always', + '--save', '', '--daemonize', 'no'], stdout=self.log, stderr=subprocess.STDOUT) + client = redis.Redis.from_url(self.url, socket_timeout=1, socket_connect_timeout=1) + try: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + if self.process.poll() is not None: + raise RuntimeError('Owned test Redis exited during startup') + try: + if client.ping(): + return + except redis.exceptions.ConnectionError: + pass + time.sleep(0.1) + raise TimeoutError('Owned test Redis did not start') + finally: + client.close() + + def stop(self): + if self.process is not None: + self.process.terminate() + try: + self.process.wait(timeout=30) + except subprocess.TimeoutExpired: + self.process.kill() + self.process.wait(timeout=5) + raise + finally: + self.log.close() + self.process = None + self.log = None + + def fault(self, probe): + self.stop() + try: + probe() + finally: + self.start() diff --git a/tests/server/integration/on_demand_training/matrix_validation.py b/tests/server/integration/on_demand_training/matrix_validation.py new file mode 100644 index 000000000..c5f9513b5 --- /dev/null +++ b/tests/server/integration/on_demand_training/matrix_validation.py @@ -0,0 +1,513 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Run the full training matrix on an isolated on-demand server.""" +import argparse +import atexit +import json +import math +import os +import ray +import socket +import subprocess +import sys +import time +import uuid +import yaml +from pathlib import Path +from ray import serve +from xml.etree import ElementTree + +_ROOT = Path(__file__).resolve().parents[4] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher + + +def validate_controller_restart(base_url, controller, model_name): + """Train, restart, restore optimizer checkpoints and recover a failed worker.""" + import numpy as np + from ray.util import list_named_actors + from tests.server.integration.on_demand_training.acceptance_support import ( + TrainingHandle, assert_released, wait_until) + + namespace = ray.get_runtime_context().namespace + old = [TrainingHandle(base_url, model_name, protocol, 'restart-' + protocol) + for protocol in ('twinkle', 'tinker')] + snapshots = [] + try: + for handle in old: + # Two calls end the native GA=2 workload on an optimizer boundary. + handle.step() + handle.step() + path = handle.save('restart-state') + values, consumed = handle.evaluate(), handle.consumed_samples + handle.step() + handle.step() + continued = handle.evaluate() + assert np.max(np.abs(continued - values)) > 2e-3, 'The reference optimizer did not change logprobs' + snapshots.append((path, values, consumed, continued)) + old_epoch = ray.get(controller.diagnostics.remote())['epoch'] + ray.kill(controller, no_restart=False) + + def restarted(): + try: + diagnostic = ray.get(controller.diagnostics.remote(), timeout=10) + return diagnostic if diagnostic['epoch'] != old_epoch else None + except ray.exceptions.RayError: + return None + + diagnostic = wait_until(restarted, 600, 'Controller restart timed out') + assert_released(controller, namespace) + for handle in old: + handle.assert_stale() + restored = [] + try: + # Keep restored peers on one runtime: checkpoint/optimizer isolation + # is checked without reloading the base model for each protocol. + optimizer_errors = {} + for protocol, (path, values, consumed, continued) in zip(('twinkle', 'tinker'), snapshots): + fresh = TrainingHandle(base_url, model_name, protocol, 'restored-' + protocol) + restored.append(fresh) + progress = fresh.resume(path) + if protocol == 'twinkle': + assert progress['consumed_train_samples'] == consumed + np.testing.assert_allclose(fresh.evaluate(), values, rtol=0, atol=2e-3) + fresh.step() + fresh.step() + resumed = fresh.evaluate() + np.testing.assert_allclose(resumed, continued, rtol=0, atol=2e-3) + optimizer_errors[protocol] = { + 'max_error': float(np.max(np.abs(resumed - continued))), + 'checkpoint_logps': values.tolist(), 'continued_logps': continued.tolist(), + 'resumed_logps': resumed.tolist(), + } + active = next(row for row in ray.get(controller.diagnostics.remote())['instances'] + if restored[0].identity in row['training_ids']) + assert set(active['training_ids']) == {handle.identity for handle in restored} + owned = [row for row in list_named_actors(all_namespaces=True) + if row['namespace'] == namespace and active['instance_id'] in row['name'] + and row['name'].startswith('model-') and not row['name'].endswith('-registry')] + assert owned, 'No model worker available for fault injection' + ray.kill(ray.get_actor(owned[0]['name'], namespace=namespace), no_restart=True) + wait_until(lambda: all(row['state'] == 'STOPPED' + for row in ray.get(controller.diagnostics.remote())['instances']), + 600, 'Worker failure cleanup timed out') + for handle in restored: + handle.assert_stale() + assert_released(controller, namespace) + finally: + for handle in restored: + handle.close(stale=True) + fresh = TrainingHandle(base_url, model_name, 'twinkle', 'after-worker-fault') + try: + fresh.step() + finally: + fresh.close() + finally: + for handle in old: + handle.close(stale=True) + return {'epoch': diagnostic['epoch'], 'worker_fault_recovered': True, 'checkpoint_restored': True, + 'optimizer_continuation': optimizer_errors} + + +def validate_idempotent_create(base_url, controller, model_name): + """Concurrent HTTP retries share one receipt, handle and allocation.""" + from concurrent.futures import ThreadPoolExecutor + import requests + from peft import LoraConfig + from twinkle.protocol.serialize import serialize_object + from twinkle.protocol.types.model import TrainingHandleResponse + from twinkle_client._future import resolve_response + from twinkle_client.http import ClientTransport + from twinkle_client.http.context import ClientContext + from tests.server.integration.on_demand_training.acceptance_support import assert_released + + session_reply = requests.post(base_url + '/api/v1/twinkle/create_session', json={}, timeout=30) + session_reply.raise_for_status() + session = session_reply.json()['session_id'] + transport = ClientTransport(ClientContext(base_url=base_url, api_key='idempotency', session_id=session)) + headers = {'Authorization': 'Bearer idempotency', 'X-Twinkle-Session-Id': session} + payload = {'client_create_id': uuid.uuid4().hex, 'adapter_name': 'duplicate', + 'config': serialize_object(LoraConfig(r=8, target_modules='all-linear')), + 'gradient_accumulation_steps': 2} + endpoint = base_url + '/api/v1/model/' + model_name + '/twinkle/add_adapter_to_model' + + def submit(_): + response = requests.post(endpoint, json=payload, + headers={**headers, 'x-request-id': uuid.uuid4().hex}, timeout=60) + response.raise_for_status() + return response + + try: + with ThreadPoolExecutor(max_workers=8) as executor: + replies = list(executor.map(submit, range(8))) + assert len({reply.json()['request_id'] for reply in replies}) == 1 + handle = resolve_response(replies[0], TrainingHandleResponse, transport=transport) + live = [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] + assert len(live) == 1 and live[0]['training_ids'] == [handle.training_id] + reply = requests.post(base_url + '/api/v1/model/' + model_name + '/twinkle/remove_adapter', + json={'training_id': handle.training_id, 'adapter_name': 'duplicate'}, + headers={**headers, 'x-request-id': uuid.uuid4().hex}, timeout=60) + reply.raise_for_status() + resolve_response(reply, None, transport=transport) + assert_released(controller, ray.get_runtime_context().namespace) + finally: + transport.close() + return True + + +def validate_capacity(base_url, controller, model_name, full_model_name, capacity): + """Assert shared LoRA capacity and full-training exclusivity on real devices.""" + from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle, assert_released + handles = [] + try: + for index in range(capacity): + handles.append(TrainingHandle(base_url, model_name, ('twinkle', 'tinker')[index % 2], + 'capacity-' + str(index))) + live = [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] + assert len(live) == 1 and len(live[0]['training_ids']) == capacity + try: + overflow = TrainingHandle(base_url, model_name, 'tinker', 'overflow') + except Exception as exc: + assert 'replica limit' in str(exc).lower(), exc + else: + overflow.close() + raise AssertionError('LoRA capacity was exceeded') + handles[0].step() + finally: + for handle in reversed(handles): + handle.close() + assert_released(controller, ray.get_runtime_context().namespace) + exclusive = TrainingHandle(base_url, full_model_name, 'twinkle', 'exclusive', mode='full') + try: + exclusive.step() + live = [row for row in ray.get(controller.diagnostics.remote())['instances'] if row['state'] != 'STOPPED'] + assert len(live) == 1 and live[0]['mode'] == 'full' and len(live[0]['training_ids']) == 1 + try: + overflow = TrainingHandle(base_url, full_model_name, 'tinker', 'exclusive-overflow', mode='full') + except Exception as exc: + assert 'replica limit' in str(exc).lower(), exc + else: + overflow.close() + raise AssertionError('Full-training instance was shared') + finally: + exclusive.close() + assert_released(controller, ray.get_runtime_context().namespace) + return True + + +def validate_redis_restart(base_url, controller, model_name, managed): + """Only interrupt a Redis process started and owned by this harness.""" + import requests + from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle + handle = TrainingHandle(base_url, model_name, 'twinkle', 'redis-fault') + try: + handle.step() + before = ray.get(controller.diagnostics.remote()) + + def unavailable(): + response = requests.post( + base_url + '/api/v1/model/' + model_name + '/twinkle/get_train_configs', + json={'training_id': handle.identity, 'adapter_name': handle.label}, + headers={'Authorization': 'Bearer ' + handle.token, 'x-request-id': uuid.uuid4().hex, + 'X-Twinkle-Session-Id': handle.transport.context.session_id}, timeout=15) + assert response.status_code >= 500, response.text + + managed.fault(unavailable) + after = ray.get(controller.diagnostics.remote(), timeout=30) + assert after['epoch'] == before['epoch'] + assert {row['instance_id'] for row in before['instances']} == { + row['instance_id'] for row in after['instances']} + handle.step() + finally: + handle.close() + return True + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--backend', required=True, choices=['transformers', 'megatron']) + parser.add_argument('--model', required=True) + parser.add_argument('--training-devices', type=int, default=2) + parser.add_argument('--training-pp-size', type=int, default=1) + parser.add_argument('--run-dir', required=True) + parser.add_argument('--persistence', choices=['memory', 'redis'], default='memory') + parser.add_argument('--managed-redis', action='store_true', help='Start an owned Redis for restart injection') + parser.add_argument('--redis-url', default=os.environ.get('TWINKLE_TEST_REDIS_URL')) + parser.add_argument('--startup-timeout-s', type=float, default=1800) + parser.add_argument('--shutdown-timeout-s', type=float, default=180) + parser.add_argument('--recovery-rounds', type=int, default=1) + parser.add_argument('--skip-full-cycle', action='store_true') + parser.add_argument('--recovery-only', action='store_true') + parser.add_argument('--full-cycle-only', action='store_true') + parser.add_argument('--stress-iterations', type=int, default=0) + parser.add_argument('--soak-seconds', type=float, default=0) + args = parser.parse_args() + if args.persistence == 'redis' and not args.redis_url and not args.managed_redis: + parser.error('--redis-url or --managed-redis is required for Redis persistence') + if args.managed_redis and args.persistence != 'redis': + parser.error('--managed-redis requires --persistence redis') + if any(not math.isfinite(value) or value <= 0 + for value in (args.startup_timeout_s, args.shutdown_timeout_s)): + parser.error('startup and shutdown timeouts must be finite and positive') + if args.training_devices < 1: + parser.error('--training-devices must be positive') + if args.backend == 'megatron' and args.training_devices < 2: + parser.error('megatron training instances require at least 2 devices') + if args.training_pp_size < 1 or args.training_devices % args.training_pp_size: + parser.error('--training-pp-size must be positive and divide --training-devices') + if args.recovery_rounds < 0 or args.stress_iterations < 0 or args.soak_seconds < 0: + parser.error('validation budgets must be non-negative') + root = Path(__file__).resolve().parents[4] + run_dir = Path(args.run_dir) + run_dir.mkdir(parents=True, exist_ok=True) + (run_dir / 'checkpoints').mkdir(parents=True, exist_ok=True) + managed = None + if args.managed_redis: + from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + managed = ManagedRedis(run_dir / 'redis') + atexit.register(managed.stop) + managed.start() + args.redis_url = managed.url + namespace = 'matrix_' + uuid.uuid4().hex[:8] + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + source = root / 'tests/server/config' / ('server_config_4b_e2e.yaml' if args.backend == 'transformers' else + 'server_config_4b_e2e_megatron.yaml') + raw = yaml.safe_load(source.read_text()) + raw.setdefault('resources', {}).update({ + 'startup_timeout_s': args.startup_timeout_s, + 'shutdown_timeout_s': args.shutdown_timeout_s, + }) + redis_prefix = f'on-demand-test::{namespace}::' + raw['persistence'] = ({ + 'mode': 'redis', + 'redis_url': args.redis_url, + 'key_prefix': redis_prefix, + } if args.persistence == 'redis' else {'mode': 'memory'}) + raw['ray_namespace'] = namespace + raw['http_options'] = {'host': '127.0.0.1', 'port': port} + full_applications = [] + for app in raw['applications']: + params = app['args'] + if app['import_path'] == 'model': + common = { + 'device_group': {'ranks': args.training_devices, 'device_type': 'cuda'}, + 'device_mesh': {'device_type': 'cuda', 'dp_size': args.training_devices // args.training_pp_size, + 'pp_size': args.training_pp_size}, + 'min_replicas': 0, 'max_replicas': 1, 'idle_timeout_s': 3600, + 'max_length': params.get('max_length'), + 'queue_config': {**params.get('queue_config', {}), 'execution_timeout': 600}, + } + app['args'] = { + **common, 'model_id': args.model, 'backend': args.backend, 'train_mode': 'lora', + 'max_loras': 10, 'max_r': 32, + 'env_vars': {'TWINKLE_TRUST_REMOTE_CODE': '1'}, + } + app.pop('deployments', None) + full_applications.append({ + 'name': app['name'] + '-full', + 'route_prefix': app['route_prefix'] + '-full', + 'import_path': 'model', + 'args': {**common, 'model_id': args.model, 'backend': args.backend, 'train_mode': 'full', + 'env_vars': app['args']['env_vars']}, + }) + elif app['import_path'] == 'sampler': + params['model_id'] = args.model + params['engine_args']['attention_config'] = {'flash_attn_version': 3} + raw['applications'].extend(full_applications) + if args.recovery_only: + raw['applications'] = [app for app in raw['applications'] if app['import_path'] in ('server', 'model')] + config = ServerConfig.model_validate(raw) + evidence_config = config.to_yaml_dict() + if args.persistence == 'redis': + evidence_config['persistence']['redis_url'] = '' + (run_dir / 'server.yaml').write_text(yaml.safe_dump(evidence_config, sort_keys=False)) + os.environ.update({ + 'TWINKLE_STATE_RAY_NAMESPACE': namespace, + 'TWINKLE_PERSISTENCE_MODE': args.persistence, + 'TWINKLE_DEFAULT_SAVE_DIR': str(run_dir / 'checkpoints'), + 'TWINKLE_MODEL_ID_ALIASES': json.dumps({model.name: model.model_id + for model in config.training_config.models}), + 'TWINKLE_TRUST_REMOTE_CODE': '1', + 'TWINKLE_TEST_GPU_E2E': '1', + 'TWINKLE_TEST_BACKEND': args.backend, + 'TWINKLE_SERVER_URL': f'http://127.0.0.1:{port}', + 'TWINKLE_TEST_OPERATION_TIMEOUT': '600', + 'PYTHONPATH': os.pathsep.join(filter(None, (str(root / 'src'), str(root), os.environ.get('PYTHONPATH')))), + }) + if args.persistence == 'redis': + os.environ.update({ + 'TWINKLE_PERSISTENCE_REDIS_URL': args.redis_url, + 'TWINKLE_PERSISTENCE_KEY_PREFIX': redis_prefix, + }) + else: + os.environ.pop('TWINKLE_PERSISTENCE_REDIS_URL', None) + os.environ.pop('TWINKLE_PERSISTENCE_KEY_PREFIX', None) + evidence = { + 'backend': args.backend, + 'persistence': args.persistence, + 'namespace': namespace, + 'cases': [], + 'lazy_start': False, + } + started = time.monotonic() + ray.init( + address='local', + namespace=namespace, + num_cpus=24, + num_gpus=args.training_devices + (0 if args.recovery_only else 1), + _temp_dir='/tmp/twm-' + namespace[-8:], + include_dashboard=False) + evidence['ray_startup_seconds'] = time.monotonic() - started + launcher = ServerLauncher(config) + try: + started = time.monotonic() + launcher._start_serve() + for app in config.applications: + if app.import_path != 'sampler': + launcher._deploy_application(app) + launcher._deploy_training() + assert ray.get(launcher._training_controller.diagnostics.remote())['instances'] == [] + evidence['lazy_start'] = True + evidence['server_startup_seconds'] = time.monotonic() - started + print(json.dumps({'phase': 'training_server_ready', 'seconds': evidence['server_startup_seconds'], + 'sampler_running': False}), flush=True) + filenames = ['test_sft_e2e.py', 'test_dpo_e2e.py', 'test_grpo_e2e.py'] + if not args.skip_full_cycle: + filenames.append('test_full_cycle_e2e.py') + if args.recovery_only: + filenames = [] + elif args.full_cycle_only: + filenames = ['test_full_cycle_e2e.py'] + for filename in filenames: + collection = subprocess.run([ + sys.executable, '-m', 'pytest', '--noconftest', '--collect-only', '-q', + str(root / 'tests/server/integration' / filename), + ], capture_output=True, text=True, env=os.environ.copy(), timeout=180) + (run_dir / (filename + '.collect.log')).write_text(collection.stdout + collection.stderr) + expected = sum('::test_' in line for line in collection.stdout.splitlines()) + assert collection.returncode == 0 and expected > 0, collection.stdout + collection.stderr + if filename in ('test_grpo_e2e.py', 'test_full_cycle_e2e.py') and not launcher._sampler_controllers: + started = time.monotonic() + print(json.dumps({'phase': 'sampler_starting', 'needed_by': filename}), flush=True) + for app in config.applications: + if app.import_path == 'sampler': + launcher._deploy_application(app) + evidence['sampler_startup_seconds'] = time.monotonic() - started + evidence['sampler_started_for'] = filename + print(json.dumps({'phase': 'sampler_ready', 'seconds': evidence['sampler_startup_seconds'], + 'needed_by': filename}), flush=True) + start = time.monotonic() + print(json.dumps({'phase': 'case_start', 'file': filename, + 'sampler_running': bool(launcher._sampler_controllers)}), flush=True) + with (run_dir / (filename + '.log')).open('w') as output: + result = subprocess.run([ + sys.executable, + '-m', + 'pytest', + '--noconftest', # HTTP clients do not need a second local Ray fixture. + str(root / 'tests/server/integration' / filename), + '-vv', + '-s', + '--junitxml=' + str(run_dir / (filename + '.xml')), + ], + stdout=output, + stderr=subprocess.STDOUT, + env=os.environ.copy(), + timeout=7200) + junit = run_dir / (filename + '.xml') + cases = ElementTree.parse(junit).getroot().findall('.//testcase') if junit.exists() else [] + skipped = sum(case.find('skipped') is not None for case in cases) + evidence['cases'].append({ + 'file': filename, + 'returncode': result.returncode, + 'tests': len(cases), + 'skipped': skipped, + 'sampler_running': bool(launcher._sampler_controllers), + 'seconds': time.monotonic() - start + }) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + assert result.returncode == 0, f'{filename} failed; see {run_dir / (filename + ".log")}' + assert len(cases) == expected and skipped == 0 + assert not any(case.find('failure') is not None or case.find('error') is not None for case in cases) + base_url = f'http://127.0.0.1:{port}' + model_name = config.training_config.models[0].name + evidence['recovery'] = [] + for _ in range(args.recovery_rounds): + recovery = validate_controller_restart(base_url, launcher._training_controller, model_name) + for sampler in launcher._sampler_controllers: + assert ray.get(sampler.diagnostics.remote(), timeout=30)['ready'], 'Training restart disturbed sampler' + if launcher._sampler_controllers: + recovery['sampler_remained_ready'] = True + evidence['recovery'].append(recovery) + evidence['restart_epoch'] = recovery['epoch'] + evidence['worker_fault_recovered'] = recovery['worker_fault_recovered'] + if managed is not None: + evidence['redis_fault_rounds'] = evidence.get('redis_fault_rounds', 0) + int( + validate_redis_restart(base_url, launcher._training_controller, model_name, managed)) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + if args.recovery_rounds: + evidence['idempotency_checked'] = validate_idempotent_create( + base_url, launcher._training_controller, model_name) + evidence['capacity_checked'] = validate_capacity( + base_url, launcher._training_controller, model_name, + next(model.name for model in config.training_config.models if model.mode == 'full'), + config.training_config.models[0].config.max_loras) + if args.stress_iterations or args.soak_seconds: + from tests.server.integration.on_demand_training.acceptance_support import assert_released + controller = launcher._training_controller + assert_released(controller, namespace, + args.redis_url if args.persistence == 'redis' else None, redis_prefix) + from tests.server.integration.on_demand_training.stress_validation import run_stress + evidence['stress'] = [] + budgets = ([('short-' + str(index), args.stress_iterations, 0) for index in range(3)] + if args.stress_iterations else []) + if args.soak_seconds: + budgets.append(('soak', 0, args.soak_seconds)) + for name, iterations, duration in budgets: + evidence['stress'].append(run_stress( + f'http://127.0.0.1:{port}', config.training_config.models[0].name, + concurrency=2, iterations=iterations, duration=duration, + controller=launcher._training_controller, namespace=namespace, + redis_url=args.redis_url if args.persistence == 'redis' else None, prefix=redis_prefix, + progress_path=run_dir / (name + '.json'))) + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + finally: + cleanup_started = time.monotonic() + try: + controller = getattr(launcher, '_training_controller', None) + if controller is not None: + ray.get(controller.shutdown.remote(), timeout=args.shutdown_timeout_s) + diagnostic = ray.get(controller.diagnostics.remote(), timeout=30) + assert all(row['state'] == 'STOPPED' for row in diagnostic['instances']) + from tests.server.integration.on_demand_training.acceptance_support import assert_released + assert_released(controller, namespace, + args.redis_url if args.persistence == 'redis' else None, redis_prefix) + evidence['training_resources_released'] = True + launcher._shutdown() + finally: + serve.shutdown() + ray.shutdown() + if args.persistence == 'redis' and evidence.get('training_resources_released'): + import redis + client = redis.Redis.from_url(args.redis_url) + try: + keys = list(client.scan_iter(match=redis_prefix + '*', count=500)) + if keys: + client.delete(*keys) + finally: + client.close() + if managed is not None: + managed.stop() + evidence['cleanup_seconds'] = time.monotonic() - cleanup_started + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + evidence['passed'] = True + (run_dir / 'results.json').write_text(json.dumps(evidence, indent=2)) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/server/integration/on_demand_training/on_demand_sdk.py b/tests/server/integration/on_demand_training/on_demand_sdk.py new file mode 100644 index 000000000..d8e30f8aa --- /dev/null +++ b/tests/server/integration/on_demand_training/on_demand_sdk.py @@ -0,0 +1,18 @@ +"""SDK operations shared by the mock and real-device acceptance harnesses.""" +import time + + +def unload_tinker_model(service, model_id): + from tinker import types + from tinker.lib.api_future_impl import _APIFuture + from tinker.lib.client_connection_pool_type import ClientConnectionPoolType + + async def unload(): + started = time.time() + with service.holder.aclient(ClientConnectionPoolType.TRAIN) as client: + future = await client.models.unload(request=types.UnloadModelRequest(model_id=model_id)) + return await _APIFuture( + types.UnloadModelResponse, service.holder, future, request_start_time=started, + request_type='UnloadModel').result_async() + + return service.holder.run_coroutine_threadsafe(unload()).result(timeout=180) diff --git a/tests/server/integration/on_demand_training/stress_validation.py b/tests/server/integration/on_demand_training/stress_validation.py new file mode 100644 index 000000000..3dc2a3156 --- /dev/null +++ b/tests/server/integration/on_demand_training/stress_validation.py @@ -0,0 +1,102 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Concurrent and soak validation against an already running single-node server.""" +import argparse +import concurrent.futures +import json +import requests +import sys +import time +from pathlib import Path + +_ROOT = Path(__file__).resolve().parents[4] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle, assert_released + + +def run_stress(base_url, model, concurrency=2, iterations=100, duration=0, controller=None, + namespace=None, redis_url=None, prefix=None, progress_path=None, protocols=('twinkle', 'tinker')): + """Run actual training; duration mode is never truncated by an iteration limit.""" + started, completed = time.monotonic(), 0 + evidence = {'completed': 0, 'failures': [], 'passed': False, 'duration_target': duration} + + def cycle(protocol, index): + handle = TrainingHandle(base_url, model, protocol, 'stress-' + str(index)) + try: + handle.step() + if (index // len(protocols)) % 10 == 0: + path = handle.save('soak-' + str(index)) + handle.resume(path) + handle.step() + finally: + handle.close() + + # Keep a peer alive so cycles test slot reuse instead of spending hours + # repeatedly loading the same base model. Its final release must free the PG. + anchor = TrainingHandle(base_url, model, protocols[0], 'soak-anchor') + try: + anchor.step() + # Measure sustained traffic time, not model cold-start/JIT time. + started = time.monotonic() + with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor: + while (time.monotonic() - started < duration) if duration else (completed < iterations): + count = concurrency if duration else min(concurrency, iterations - completed) + jobs = [executor.submit(cycle, protocols[(completed + slot) % len(protocols)], completed + slot) + for slot in range(count)] + for job in jobs: + job.result() # Do not convert incomplete cycles into successful evidence. + completed += count + # Session heartbeats do not renew the business-idle lease. + # Keep the warm peer active during a multi-hour soak. + if completed % 10 == 0: + anchor.step() + if controller is not None: + import ray + live = [row for row in ray.get(controller.diagnostics.remote(), timeout=30)['instances'] + if row['state'] != 'STOPPED'] + assert len(live) == 1 and live[0]['state'] == 'READY' + assert live[0]['training_ids'] == [anchor.identity], live + assert len(live[0]['devices']) == 2 + evidence.update(completed=completed, seconds=time.monotonic() - started) + if progress_path: + Path(progress_path).write_text(json.dumps(evidence, indent=2)) + print(json.dumps({'phase': 'stress', **evidence}), flush=True) + finally: + anchor.close() + if controller is not None: + assert_released(controller, namespace, redis_url, prefix) + evidence['passed'] = True + if duration: + assert evidence['seconds'] >= duration + if progress_path: + Path(progress_path).write_text(json.dumps(evidence, indent=2)) + return evidence + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--backend', required=True, choices=['transformers', 'megatron']) + parser.add_argument('--client', choices=['twinkle', 'tinker', 'both'], default='both') + parser.add_argument('--server-url', required=True) + parser.add_argument('--model', required=True) + parser.add_argument('--concurrency', type=int, default=2) + parser.add_argument('--iterations', type=int, default=100) + parser.add_argument('--duration', type=float, default=0, help='Run for this many seconds instead of fixed iterations') + parser.add_argument('--results', required=True) + args = parser.parse_args() + if args.concurrency <= 0 or args.iterations <= 0 or args.duration < 0: + parser.error('concurrency and iterations must be positive; duration must be non-negative') + + protocols = ['twinkle', 'tinker'] if args.client == 'both' else [args.client] + evidence = run_stress(args.server_url, args.model, args.concurrency, args.iterations, args.duration, + progress_path=args.results, protocols=protocols) + response = requests.get(args.server_url + '/api/v1/model/' + args.model + '/healthz', timeout=30) + response.raise_for_status() + assert response.json()['healthy'] and not response.json()['instances'] + print(json.dumps(evidence), flush=True) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/server/integration/on_demand_training/test_resident_release_e2e.py b/tests/server/integration/on_demand_training/test_resident_release_e2e.py new file mode 100644 index 000000000..a7879000f --- /dev/null +++ b/tests/server/integration/on_demand_training/test_resident_release_e2e.py @@ -0,0 +1,190 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real resident reset ACK racing worker death; needs two dedicated accelerators.""" +import asyncio +import json +import os +import socket +import subprocess +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np +import pytest + +pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_RESIDENT_RELEASE') != '1', + reason='Requires two dedicated devices and a real cached model') + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + yield + + +@pytest.fixture(autouse=True) +def _reset_canonical_state_actor(): + yield + + +def release_gate(actor, action): + """Delay only the ACK after the real model reset; never substitute computation.""" + if action == 'install': + entered, resume = asyncio.Event(), asyncio.Event() + original = actor.driver.release_training + + async def delayed(instance, training): + await original(instance, training) + entered.set() + await resume.wait() + + actor.driver.release_training = delayed + actor._release_test_gate = (entered, resume, original) + elif action == 'entered': + return actor._release_test_gate[0].is_set() + elif action == 'resume': + actor._release_test_gate[1].set() + actor.driver.release_training = actor._release_test_gate[2] + + +def worker_origin(actor): + import twinkle + from twinkle.server.model.pool import controller + return {'python': sys.executable, 'twinkle': twinkle.__file__, 'controller': controller.__file__, + 'visible_devices': os.environ.get('CUDA_VISIBLE_DEVICES')} + + +def test_resident_release_ack_after_worker_failure(tmp_path): + import ray + from ray import serve + from ray.util import list_named_actors + from twinkle.server.config import ServerConfig + from twinkle.server.launcher.server_launcher import ServerLauncher + from tests.server.integration.on_demand_training.acceptance_support import TrainingHandle, wait_until, assert_released + from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + + root = Path(__file__).resolve().parents[4] + directory = Path(os.environ.get('TWINKLE_TEST_RUN_DIR', str(tmp_path))) + directory.mkdir(parents=True, exist_ok=True) + backend = os.environ['TWINKLE_TEST_BACKEND'] + model = os.environ['TWINKLE_TEST_MODEL'] + namespace = 'resident_' + uuid.uuid4().hex[:8] + persistence = 'redis' if backend == 'transformers' else 'memory' + managed = ManagedRedis(directory / 'redis') if persistence == 'redis' else None + if managed: + managed.start() + prefix = namespace + '::' + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + model_name = 'Qwen/Qwen3.5-4B' + raw = { + 'ray_namespace': namespace, 'http_options': {'host': '127.0.0.1', 'port': port}, + 'resources': {'startup_timeout_s': 1800, 'shutdown_timeout_s': 180, 'sweep_interval_s': 1}, + 'persistence': ({'mode': persistence, 'redis_url': managed.url, 'key_prefix': prefix} + if managed else {'mode': persistence}), + 'applications': [ + {'name': 'server', 'route_prefix': '/api/v1', 'import_path': 'server', + 'args': {'supported_models': [model_name]}}, + {'name': 'model', 'route_prefix': '/api/v1/model/' + model_name, 'import_path': 'model', + 'args': {'model_id': model, 'backend': backend, 'train_mode': 'full', + 'device_group': {'ranks': 2, 'device_type': 'cuda'}, + 'min_replicas': 1, 'max_replicas': 1, + 'device_mesh': {'dp_size': 2 if backend == 'transformers' else 1, + 'pp_size': 1 if backend == 'transformers' else 2}, + 'max_length': 512, 'idle_timeout_s': 3600, + 'queue_config': {'execution_timeout': 600}, + 'env_vars': {'TWINKLE_TRUST_REMOTE_CODE': '1'}}}, + ], + } + (directory / 'server.json').write_text(json.dumps(raw, indent=2)) + os.environ.update(TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE=persistence, + TWINKLE_DEFAULT_SAVE_DIR=str(directory / 'checkpoints'), TWINKLE_TRUST_REMOTE_CODE='1', + TWINKLE_MODEL_ID_ALIASES=json.dumps({model_name: model}), + PYTHONPATH=os.pathsep.join((str(root / 'src'), str(root)))) + if managed: + os.environ.update(TWINKLE_PERSISTENCE_REDIS_URL=managed.url, TWINKLE_PERSISTENCE_KEY_PREFIX=prefix) + else: + os.environ.pop('TWINKLE_PERSISTENCE_REDIS_URL', None) + os.environ.pop('TWINKLE_PERSISTENCE_KEY_PREFIX', None) + evidence = {'backend': backend, 'namespace': namespace, + 'sha': subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=root, text=True).strip()} + ray.init(address='local', namespace=namespace, num_cpus=24, num_gpus=2, + _temp_dir='/tmp/twr-' + namespace[-8:], include_dashboard=False) + launcher = ServerLauncher(ServerConfig.model_validate(raw)) + handles, controller = [], None + try: + launcher._start_serve() + for app in launcher.config.applications: + launcher._deploy_application(app) + launcher._deploy_training() + controller = launcher._training_controller + url = f'http://127.0.0.1:{port}' + old = TrainingHandle(url, model_name, 'twinkle', 'before-fault', mode='full') + handles.append(old) + baseline = old.evaluate() + old.step() + old.step() + trained = old.evaluate() + assert np.max(np.abs(trained - baseline)) > 2e-3, 'Training did not update real model logprobs' + diagnostic = ray.get(controller.diagnostics.remote()) + instance = next(row for row in diagnostic['instances'] if old.identity in row['training_ids']) + owned = [row for row in list_named_actors(all_namespaces=True) + if row['namespace'] == namespace and instance['instance_id'] in row['name'] + and row['name'].startswith('model-') and not row['name'].endswith('-registry')] + assert len(owned) == 2 + evidence['worker_origins'] = [ray.get(ray.get_actor(row['name'], namespace=namespace).__ray_call__.remote( + worker_origin)) for row in owned] + for origin in evidence['worker_origins']: + assert Path(origin['twinkle']).is_relative_to(root / 'src') + assert origin['python'] == sys.executable + ray.get(controller.__ray_call__.remote(release_gate, 'install')) + with ThreadPoolExecutor(max_workers=1) as executor: + releasing = executor.submit(old.close) + try: + wait_until(lambda: ray.get(controller.__ray_call__.remote(release_gate, 'entered')), + 600, 'Real reset did not finish') + ray.kill(ray.get_actor(owned[0]['name'], namespace=namespace), no_restart=True) + + def stopped(): + rows = ray.get(controller.diagnostics.remote())['instances'] + return next((row for row in rows if row['instance_id'] == instance['instance_id'] + and row['state'] == 'STOPPED'), None) + + wait_until(stopped, 600, 'Health monitor did not retire the failed resident') + finally: + ray.get(controller.__ray_call__.remote(release_gate, 'resume')) + releasing.result(timeout=600) + assert stopped(), 'A late reset ACK resurrected the dead resident' + fresh = TrainingHandle(url, model_name, 'tinker', 'after-fault', mode='full') + handles.append(fresh) + restored = fresh.evaluate() + np.testing.assert_allclose(restored, baseline, rtol=0, atol=2e-3) + fresh.step() + continued = fresh.evaluate() + assert np.max(np.abs(continued - restored)) > 2e-3 + fresh.close() + assert stopped() + rows = ray.get(controller.diagnostics.remote())['instances'] + assert sum(row['state'] == 'READY' for row in rows) == 1 + evidence.update(late_ack_fenced=True, recovered=True, baseline_logps=baseline.tolist(), + trained_logps=trained.tolist(), restored_logps=restored.tolist(), + continued_logps=continued.tolist(), instances=rows) + finally: + if controller is not None: + ray.get(controller.__ray_call__.remote( + lambda actor: release_gate(actor, 'resume') if hasattr(actor, '_release_test_gate') else None)) + for handle in handles: + handle.close(stale=True) + try: + if controller is not None: + ray.get(controller.shutdown.remote(), timeout=240) + assert_released(controller, namespace, managed.url if managed else None, prefix) + evidence['released'] = True + launcher._shutdown() + finally: + serve.shutdown() + ray.shutdown() + if managed: + managed.stop() + (directory / 'result.json').write_text(json.dumps(evidence, indent=2)) diff --git a/tests/server/integration/resources/__init__.py b/tests/server/integration/resources/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/integration/resources/allocation_probe.py b/tests/server/integration/resources/allocation_probe.py new file mode 100644 index 000000000..9d057d215 --- /dev/null +++ b/tests/server/integration/resources/allocation_probe.py @@ -0,0 +1,69 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Run in its own process: test production allocation against logical Ray GPUs.""" +import asyncio +import json +import os +import ray +import uuid + +from twinkle import DeviceGroup +from twinkle.infra._ray.resource_manager import ResourceManager +from twinkle.server.config.training import TrainingConfig +from twinkle.server.resources.allocation import CapacityUnavailableError, DevicePool + + +async def main(): + # These are synthetic Ray GPU reservations, not CUDA workloads. + os.environ.pop('CUDA_VISIBLE_DEVICES', None) + namespace = 'alloc_' + uuid.uuid4().hex + ray.init( + address='local', + namespace=namespace, + num_cpus=4, + num_gpus=8, + _temp_dir='/tmp/twa-' + namespace[-8:], + include_dashboard=False, + dashboard_port=0, + log_to_driver=False) + pool = DevicePool(namespace, allocation_timeout=3, shutdown_timeout=10) + allocations = [] + try: + for index, count in enumerate((2, 2, 4)): + config = TrainingConfig(models=[{'name': str(index), 'model_id': str(index), 'backend': 'mock', 'device_type': 'cuda', 'mode': 'full', 'config': {'num_devices': count, 'device_mesh': {'fsdp_size': count}}}]) + spec = config.get(str(index), 'full') + allocation = await pool.allocate(spec.config.num_devices, spec.device_type, 'epoch', str(index)) + allocations.append(allocation) + devices = [device for allocation in allocations for device in allocation.devices] + assert len(set(devices)) == 8, devices + assert len({allocation.node_id for allocation in allocations}) == 1 + try: + await pool.allocate(spec.config.num_devices, spec.device_type, 'epoch', 'overflow') + raise AssertionError('over-allocation was admitted') + except CapacityUnavailableError: + pass + from ray.util.placement_group import get_placement_group + group = DeviceGroup(name='owned', ranks=2, device_type='cuda') + owned_pg = get_placement_group(allocations[0].name) + view = ResourceManager.from_allocation(group, owned_pg, list(allocations[0].devices)) + assert view.get_config('owned').ranks == [0, 1] + assert [entry['gpu_rank'][0] for entry in view.get_group('owned')] == list(allocations[0].devices) + # A ResourceManager view cannot release a controller-owned allocation. + view.destroy_placement_group() + assert get_placement_group(allocations[0].name).id == owned_pg.id + freed = allocations[0] + await pool.remove(freed.name) + two = TrainingConfig(models=[{'name': 'replacement', 'model_id': 'replacement', 'backend': 'mock', 'device_type': 'cuda', 'mode': 'full', 'config': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}}}]).get('replacement', 'full') + replacement = await pool.allocate(two.config.num_devices, two.device_type, 'epoch', 'replacement') + allocations.append(replacement) + assert set(replacement.devices) == set(freed.devices) + print(json.dumps({'passed': True, 'devices': devices, 'reused': list(replacement.devices)})) + finally: + try: + for allocation in allocations: + await pool.remove(allocation.name) + finally: + ray.shutdown() + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/tests/server/integration/resources/fragmentation_probe.py b/tests/server/integration/resources/fragmentation_probe.py new file mode 100644 index 000000000..f684c51c3 --- /dev/null +++ b/tests/server/integration/resources/fragmentation_probe.py @@ -0,0 +1,73 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Exercise single-node admission with real Ray nodes and logical devices.""" +import asyncio +import json +import os +import ray +import uuid +from ray.cluster_utils import Cluster +from ray.util.placement_group import placement_group, placement_group_table, remove_placement_group + +from twinkle.server.config.training import TrainingConfig +from twinkle.server.resources.allocation import CapacityUnavailableError, DevicePool + + +async def main(): + # Synthetic reservations must not inherit a hidden CUDA-device mask. + os.environ.pop('CUDA_VISIBLE_DEVICES', None) + cluster = Cluster() + occupied = [] + allocated = None + pool = None + try: + cluster.add_node(num_cpus=4, num_gpus=4, resources={'left': 1}, include_dashboard=False, dashboard_port=0) + cluster.add_node(num_cpus=4, num_gpus=4, resources={'right': 1}) + ray.init(address=cluster.address, namespace='fragmentation_' + uuid.uuid4().hex, log_to_driver=False) + # Keep fragmentation rejection fast, but allow a cold Ray worker to + # import the allocation probe before confirming a successful placement. + pool = DevicePool(ray.get_runtime_context().namespace, allocation_timeout=1, shutdown_timeout=60) + spec = TrainingConfig(models=[{'name': 'two', 'model_id': 'two', 'backend': 'mock', 'device_type': 'cuda', 'mode': 'full', 'config': {'num_devices': 2, 'device_mesh': {'fsdp_size': 2}}}]).get('two', 'full') + for marker in ('left', 'right'): + pg = placement_group([{'GPU': 3, marker: 0.01}], strategy='STRICT_PACK') + occupied.append(pg) + await asyncio.wait_for(pg.ready(), 20) + assert ray.available_resources()['GPU'] == 2 + try: + await pool.allocate(spec.config.num_devices, spec.device_type, 'epoch', 'fragmented') + except CapacityUnavailableError: + pass + else: + raise AssertionError('Two nodes with one free GPU each admitted a two-GPU instance') + groups = placement_group_table().values() + assert not any(group['name'] == pool.name('epoch', 'fragmented') and group['state'] != 'REMOVED' + for group in groups) + remove_placement_group(occupied[0]) + for _ in range(200): + if ray.available_resources().get('GPU', 0) == 5: + break + await asyncio.sleep(0.05) + else: + raise AssertionError('Owned GPU reservation was not released') + allocated = await pool.allocate(spec.config.num_devices, spec.device_type, 'epoch', 'available') + assert len(set(allocated.devices)) == 2 + print( + json.dumps({ + 'passed': True, + 'nodes': len(ray.nodes()), + 'fragmented_rejected': True, + 'pending_removed': True, + 'devices': allocated.devices + })) + finally: + try: + if allocated is not None: + await pool.remove(allocated.name) + for pg in occupied: + remove_placement_group(pg) + finally: + ray.shutdown() + cluster.shutdown() + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/tests/server/integration/resources/test_allocation.py b/tests/server/integration/resources/test_allocation.py new file mode 100644 index 000000000..54d8f8b11 --- /dev/null +++ b/tests/server/integration/resources/test_allocation.py @@ -0,0 +1,31 @@ +import json +import subprocess +import sys +from pathlib import Path + + +def test_production_allocator_uses_disjoint_devices_and_reuses_released_group(): + result = subprocess.run([sys.executable, str(Path(__file__).with_name('allocation_probe.py'))], + capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr + evidence = json.loads(result.stdout.strip().splitlines()[-1]) + assert evidence['passed'] and len(set(evidence['devices'])) == 8 + assert len(evidence['reused']) == 2 + + +def test_fragmented_devices_cannot_be_combined_across_nodes(): + result = subprocess.run([sys.executable, str(Path(__file__).with_name('fragmentation_probe.py'))], + capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr + evidence = json.loads(result.stdout.strip().splitlines()[-1]) + assert evidence['passed'] and evidence['nodes'] == 2 + assert evidence['fragmented_rejected'] and evidence['pending_removed'] + assert len(evidence['devices']) == 2 + + +def test_sampler_uses_real_cpu_services_for_scaling_and_release(): + script = Path(__file__).resolve().parents[1] / 'sampler' / 'sampling_probe.py' + result = subprocess.run([sys.executable, str(script)], capture_output=True, text=True, timeout=1800) + assert result.returncode == 0, result.stdout + result.stderr + evidence = [json.loads(line) for line in result.stdout.splitlines() if line.startswith('{"sizes":')] + assert evidence == [{'sizes': [2, 4, 2], 'requests': 8, 'confirmed_release': True}] diff --git a/tests/server/integration/sampler/__init__.py b/tests/server/integration/sampler/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/integration/sampler/sampling_probe.py b/tests/server/integration/sampler/sampling_probe.py new file mode 100644 index 000000000..6e32c6ba4 --- /dev/null +++ b/tests/server/integration/sampler/sampling_probe.py @@ -0,0 +1,108 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Exercise the actual sampler driver and mock compute service on an owned CPU cluster.""" +import asyncio +import json +import os +import uuid + +import ray +from ray import serve +from ray.util import list_named_actors +from ray.util.placement_group import placement_group_table + +from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.sampler.pool.controller import SamplerController +from twinkle.server.sampler.pool.driver import SamplerDriver + + +async def run(namespace): + args = SamplerArgs( + model_id='mock', sampler_type='mock', + device_group={'name': 'sampler', 'ranks': 1, 'device_type': 'cpu'}, + device_mesh={'device_type': 'cpu', 'dp_size': 1}, + min_replicas=2, max_replicas=4, target_ongoing_requests=1, idle_timeout_s=1) + budgets = {'allocation_timeout_s': 10, 'startup_timeout_s': 600, + 'shutdown_timeout_s': 60, 'sweep_interval_s': 1} + driver = SamplerDriver(args, namespace, 'cpu-allocation', {}, budgets) + clock = [0] + pool = SamplerController(args, driver, budgets, clock=lambda: clock[0]) + # Hold replies after the actual compute route runs. The pool must retain + # their tickets until the caller can confirm backend completion. + invoke = driver.invoke + replies = 0 + dispatched, release = asyncio.Event(), asyncio.Event() + + async def delayed_reply(*args, **kwargs): + nonlocal replies + response = await invoke(*args, **kwargs) + replies += 1 + if replies == 8: + dispatched.set() + await release.wait() + return response + + driver.invoke = delayed_reply + calls = [] + try: + await pool.initialize() + sizes = [len(pool._ready_instances())] + payload = {'endpoint': 'twinkle/sample', 'method': 'POST', + 'headers': [(b'content-type', b'application/json'), + (b'authorization', b'Bearer cpu-probe'), + (b'x-request-id', b'cpu-probe')], 'query_string': b'', + 'body': json.dumps({'inputs': [{'input_ids': [1, 2]}], + 'sampling_params': {'max_tokens': 2}}).encode()} + calls = [asyncio.create_task(pool.request(payload)) for _ in range(8)] + await asyncio.wait_for(dispatched.wait(), 600) + await asyncio.gather(*tuple(pool.lifecycle.startups.values())) + sizes.append(len(pool._ready_instances())) + clock[0] = 10 + await pool.sweep() + assert pool._owned_count() == 4 + release.set() + responses = await asyncio.gather(*calls) + assert all(response['status'] == 200 for response in responses), responses + await pool.lifecycle.join() + from twinkle.server.state import get_server_state + state = get_server_state() + for response in responses: + record = await state.get_future(json.loads(response['body'])['request_id']) + assert record['status'] == 'completed', record + assert len(record['result']['samples'][0]['sequences'][0]['tokens']) == 2 + health = await asyncio.gather(*(driver.health(instance) for instance in pool._ready_instances())) + assert all(not result['unhealthy'] for result in health) + clock[0] = 12 + await pool.sweep() + await pool.lifecycle.join() + sizes.append(len(pool._ready_instances())) + assert sizes == [2, 4, 2], sizes + await pool.shutdown() + await driver.runtime.wait_for_cleanup() + assert not any(name.startswith(driver.prefix) for name in serve.status().applications) + assert not any(row.get('name', '').startswith(driver.prefix) and row['state'] != 'REMOVED' + for row in placement_group_table().values()) + assert not any(row['namespace'] == namespace and 'sampler_' in row['name'] + for row in list_named_actors(all_namespaces=True)) + return {'sizes': sizes, 'requests': len(responses), 'confirmed_release': True} + finally: + release.set() + await asyncio.gather(*calls, return_exceptions=True) + await pool.shutdown() + await driver.runtime.wait_for_cleanup() + + +def main(): + namespace = 'sampling_probe_' + uuid.uuid4().hex[:8] + os.environ.update(TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE='memory') + ray.init(address='local', namespace=namespace, num_cpus=24, num_gpus=0, + _temp_dir='/tmp/tws-' + namespace[-8:], include_dashboard=False) + try: + serve.start(http_options={'host': '127.0.0.1', 'port': 0}) + print(json.dumps(asyncio.run(run(namespace))), flush=True) + finally: + serve.shutdown() + ray.shutdown() + + +if __name__ == '__main__': + main() diff --git a/tests/server/integration/sampler/test_real_pool_e2e.py b/tests/server/integration/sampler/test_real_pool_e2e.py new file mode 100644 index 000000000..614796770 --- /dev/null +++ b/tests/server/integration/sampler/test_real_pool_e2e.py @@ -0,0 +1,385 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real vLLM pool acceptance. Run only on two or four dedicated devices with the PPU SDK loaded.""" +import asyncio +import json +import math +import os +import shutil +import subprocess +import time +import traceback +import uuid +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_REAL_POOL') != '1', + reason='Requires two or four dedicated devices and a real cached model') + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + yield + + +@pytest.fixture(autouse=True) +def _reset_canonical_state_actor(): + yield + + +def payload(endpoint, body): + return {'endpoint': endpoint, 'method': 'POST', 'body': json.dumps(body).encode(), + 'headers': [(b'content-type', b'application/json'), (b'x-request-id', uuid.uuid4().hex.encode()), + (b'authorization', b'Bearer pool-acceptance')], 'query_string': b''} + + +def checkpoint_alias(model, target): + target.mkdir() + for source in model.iterdir(): + if source.is_file(): + (target / source.name).symlink_to(source) + return str(target) + + +def make_lora(model, target): + import torch + from peft import LoraConfig + from safetensors import safe_open + from safetensors.torch import save_file + rank, weights = 8, {} + generator = torch.Generator().manual_seed(312) + for shard in model.glob('*.safetensors'): + with safe_open(str(shard), framework='pt', device='cpu') as source: + for name in source.keys(): + if name.endswith('.self_attn.q_proj.weight'): + output_size, input_size = source.get_slice(name).get_shape() + prefix = 'base_model.model.' + name[:-len('.weight')] + weights[prefix + '.lora_A.weight'] = torch.randn(rank, input_size, generator=generator) * .05 + weights[prefix + '.lora_B.weight'] = torch.randn(output_size, rank, generator=generator) * .05 + assert weights, 'No q_proj weights found; this is not a skipped acceptance case' + target.mkdir() + LoraConfig(r=rank, lora_alpha=rank, target_modules=['q_proj'], bias='none', task_type='CAUSAL_LM', + inference_mode=True, base_model_name_or_path=str(model)).save_pretrained(target) + save_file(weights, str(target / 'adapter_model.safetensors')) + return str(target) + + +async def exercise(model, directory, namespace, evidence, *, extended=True): + import numpy as np + import ray + from ray import serve + from ray.util import list_named_actors + from ray.util.placement_group import placement_group_table + from transformers import AutoTokenizer + from twinkle.server.config.application_spec import SamplerArgs + from twinkle.server.sampler.pool.controller import SamplerController + from twinkle.server.sampler.pool.driver import SamplerDriver + from twinkle.server.state import get_server_state + + devices_count = int(os.environ.get('TWINKLE_TEST_POOL_DEVICES', '4')) + assert devices_count in (2, 4) + minimum = devices_count // 2 + args = SamplerArgs(model_id=str(model), sampler_type='vllm', + device_group={'name': 'sampler', 'ranks': 1, 'device_type': 'cuda'}, + device_mesh={'device_type': 'cuda', 'dp_size': 1}, + engine_args={'max_model_len': 2048, 'gpu_memory_utilization': .4, 'enable_lora': True, + 'enforce_eager': True, 'logprobs_mode': 'processed_logprobs', + 'attention_config': {'flash_attn_version': 3}}, + min_replicas=minimum, max_replicas=devices_count, target_ongoing_requests=1, idle_timeout_s=2, + queue_config={'execution_timeout': 600, 'queue_timeout': 600, 'inline_fast_path_timeout': 0.01}) + budgets = {'allocation_timeout_s': 60, 'startup_timeout_s': 900, + 'shutdown_timeout_s': 180, 'sweep_interval_s': 1} + driver = SamplerDriver(args, namespace, 'real-pool', {}, budgets) + clock_offset = [0] + pool = SamplerController(args, driver, budgets, clock=lambda: time.monotonic() + clock_offset[0]) + pools = [pool] + state = get_server_state() + tokenizer = AutoTokenizer.from_pretrained(str(model), local_files_only=True) + ids = tokenizer.apply_chat_template([{'role': 'user', 'content': 'Count the integers from one to one hundred thousand.'}], + tokenize=True, add_generation_prompt=True, enable_thinking=False, + return_dict=False) + calls = [] + futures = [] + invoke = driver.invoke + + async def observed(instance, ticket, body, queue=None): + calls.append({'instance': instance.identity, 'endpoint': body['endpoint'], 'devices': instance.devices}) + return await invoke(instance, ticket, body, queue) + + driver.invoke = observed + + async def sample(target=pool, path=None, full=None, max_tokens=64): + body = {'inputs': [{'input_ids': ids}], + 'sampling_params': {'max_tokens': max_tokens, 'temperature': 0, 'logprobs': 1, 'ignore_eos': True}} + if path: + body['adapter_uri'] = path + reply = await target.request(payload('twinkle/sample', body), full_path=full) + assert reply['status'] == 200, reply + result = json.loads(reply['body']) + request_id = result['request_id'] + futures.append(request_id) + deadline = time.monotonic() + 600 + while result['status'] not in ('completed', 'failed', 'cancelled'): + assert time.monotonic() < deadline, result + await asyncio.sleep(.05) + result = await state.get_future(request_id) + assert result['status'] == 'completed', result + seq = result['result']['samples'][0]['sequences'][0] + assert seq['tokens'] and len(seq['tokens']) == len(seq['logprobs']) + assert all(math.isfinite(position[0][1]) for position in seq['logprobs']) + return seq + + async def settled(target=pool): + await asyncio.gather(*tuple(target.lifecycle.startups.values())) + await target.lifecycle.join() + + def save(): + snapshots = evidence.setdefault('device_snapshots', []) + name = f'devices-stage-{len(snapshots):02d}.txt' + (directory / name).write_text(subprocess.check_output(['ppu-smi'], text=True)) + snapshots.append(name) + (directory / 'results.json').write_text(json.dumps(evidence, indent=2)) + + try: + await pool.initialize() + evidence['sizes'] = [len(pool._ready_instances())] + source_root = str(Path(__file__).resolve().parents[4]) + + async def verify_worker_origins(): + workers = [row for row in list_named_actors(all_namespaces=True) + if row['namespace'] == namespace and row['name'].startswith('sampler-') + and not row['name'].endswith('-registry')] + assert len(workers) == len(pool._ready_instances()) + origins = {} + for row in workers: + worker = ray.get_actor(row['name'], namespace=namespace) + origin = await worker.__ray_call__.remote(lambda _: __import__('twinkle').__file__) + assert Path(origin).resolve().is_relative_to(source_root), origin + origins[row['name']] = origin + evidence['worker_import_origins'] = origins + + await verify_worker_origins() + save() + await asyncio.gather(*(sample(max_tokens=256) for _ in range(2 * devices_count))) + await settled() + assert len(pool._ready_instances()) == devices_count + devices = [device for instance in pool._ready_instances() for device in instance.devices] + assert len(devices) == len(set(devices)) == devices_count + calls.clear() + samples = await asyncio.gather(*(sample(max_tokens=128) for _ in range(2 * devices_count))) + await settled() + await verify_worker_origins() + evidence['sample_result'] = samples[0] + assert {row['instance'] for row in calls} == {i.identity for i in pool._ready_instances()} + for instance in pool._ready_instances(): + assert not (await driver.health(instance))['unhealthy'] + evidence.update(device_ids=devices, request_placements=list(calls)) + evidence['sizes'].append(devices_count) + save() + + # Observe real running compute before moving the idle clock past its deadline. + futures.clear() + active = [asyncio.create_task(sample(max_tokens=1024)) for _ in range(devices_count)] + try: + deadline = time.monotonic() + 30 + while True: + records = [await state.get_future(identity) for identity in futures] + if (len(records) == devices_count and all(row['status'] == 'running' for row in records) + and all(instance.requests for instance in pool._ready_instances())): + break + assert time.monotonic() < deadline and not any(task.done() for task in active), records + await asyncio.sleep(.01) + clock_offset[0] += 10 + await pool.sweep() + assert len(pool._ready_instances()) == devices_count + evidence['active_compute_sweep_safe'] = True + finally: + await asyncio.gather(*active) + await settled() + + # Pause the ACK after a real backend update, then sweep across the idle window. + entered, resume = asyncio.Event(), asyncio.Event() + original_wait = driver.wait_operation + + async def delayed_ack(instance, ticket): + await original_wait(instance, ticket) + entered.set() + await resume.wait() + + driver.wait_operation = delayed_ack + update = asyncio.create_task(pool.request(payload('twinkle/set_template', {'template_cls': 'Qwen3_5Template'}))) + try: + await asyncio.wait_for(entered.wait(), 120) + await asyncio.sleep(2.1) + await pool.sweep() + assert len(pool._ready_instances()) == devices_count + finally: + resume.set() + await update + driver.wait_operation = original_wait + evidence['control_sweep_safe'] = True + + if not extended: + await asyncio.sleep(2.1) + await pool.sweep() + await settled() + evidence['sizes'].append(len(pool._ready_instances())) + assert evidence['sizes'] == [minimum, devices_count, minimum] + evidence['scaling_complete'] = True + return + + old = checkpoint_alias(model, directory / 'checkpoint-A') + new = checkpoint_alias(model, directory / 'checkpoint-B') + await sample(full=old) + await settled() + original_load = driver.load_weights + entered.clear(); resume.clear() + + async def delayed_load(instance, path): + await original_load(instance, path) + entered.set() + await resume.wait() + + driver.load_weights = delayed_load + update = asyncio.create_task(sample(full=new)) + try: + await asyncio.wait_for(entered.wait(), 180) + await asyncio.sleep(2.1) + await pool.sweep() + assert len(pool._ready_instances()) == devices_count + finally: + resume.set() + await update + driver.load_weights = original_load + await settled() + shutil.rmtree(old) # Only the test-owned symlink directory; never the model cache. + evidence['full_update_sweep_safe'] = True + await asyncio.sleep(2.1) + await pool.sweep() + await settled() + evidence['sizes'].append(len(pool._ready_instances())) + assert evidence['sizes'] == [minimum, devices_count, minimum] + save() + await asyncio.gather(*(sample(max_tokens=128) for _ in range(2 * devices_count))) + await settled() + assert len(pool._ready_instances()) == devices_count + evidence['superseded_checkpoint_rebuild'] = True + + lora = make_lora(model, directory / 'lora') + before = await sample(path=lora) + await settled() + response = await pool.request(payload('twinkle/unload_adapter_paths', {'adapter_paths': [lora]})) + assert response['status'] == 200 + after = await sample(path=lora) + await settled() + assert before['tokens'] == after['tokens'] + np.testing.assert_allclose([x[0][1] for x in before['logprobs']], + [x[0][1] for x in after['logprobs']], rtol=0, atol=1e-5) + evidence['lora_results'] = {'before_eviction': before, 'after_reload': after} + evidence['lora_eviction_reload'] = True + save() + + # Lose a real control ACK; all four real participants must be stopped. + entered.clear() + async def missing_ack(instance, ticket): + await original_wait(instance, ticket) + entered.set() + await asyncio.Event().wait() + driver.wait_operation = missing_ack + args.queue_config.execution_timeout = 5 + try: + with pytest.raises(TimeoutError): + await pool.request(payload('twinkle/set_template', {'template_cls': 'Qwen3_5Template'})) + assert entered.is_set() + assert not pool._ready_instances() and pool._owned_count() == 0 + finally: + driver.wait_operation = original_wait + args.queue_config.execution_timeout = 600 + await pool.sweep() + await settled() + await sample(full=new) + await settled() + evidence['control_timeout_recovered'] = True + save() + + victim = pool._ready_instances()[0] + workers = [row for row in list_named_actors(all_namespaces=True) + if row['namespace'] == namespace and victim.identity in row['name'] + and row['name'].startswith('sampler-') and not row['name'].endswith('-registry')] + assert workers, 'No owned sampler worker for fault injection' + ray.kill(ray.get_actor(workers[0]['name'], namespace=namespace), no_restart=True) + deadline = time.monotonic() + 60 + while not (await driver.health(victim))['unhealthy']: + assert time.monotonic() < deadline + await asyncio.sleep(.2) + await pool._close(victim) + await pool.sweep() + await settled() + await sample(full=new) + await settled() + evidence['worker_fault_recovered'] = True + save() + + # The new controller's production recovery finds and retires the old pool's resources. + old_names = {i.application for i in pool._ready_instances()} + replacement_driver = SamplerDriver(args, namespace, 'real-pool', {}, budgets) + replacement = SamplerController(args, replacement_driver, budgets) + pools.append(replacement) + await replacement.initialize() + assert replacement.epoch != pool.epoch + assert not old_names.intersection(serve.status().applications) + await sample(target=replacement) + await settled(replacement) + evidence['controller_orphan_recovery'] = True + save() + except BaseException: + evidence['failure'] = traceback.format_exc() + save() + raise + finally: + for target in reversed(pools): + await target.shutdown() + await target.driver.runtime.wait_for_cleanup() + assert not any(name.startswith(driver.prefix) for name in serve.status().applications) + assert not any(row.get('name', '').startswith(driver.prefix) and row['state'] != 'REMOVED' + for row in placement_group_table().values()) + assert not any(row['namespace'] == namespace and 'sampler_' in row['name'] + for row in list_named_actors(all_namespaces=True)) + evidence['confirmed_release'] = True + if not extended and evidence.get('scaling_complete'): + evidence['passed'] = True + save() + evidence['passed'] = True + save() + + +def _run_real_pool(*, extended): + import ray + import twinkle + from ray import serve + model = Path(os.environ['TWINKLE_TEST_MODEL_PATH']) + directory = Path(os.environ['TWINKLE_TEST_ARTIFACT_DIR']) + directory.mkdir(parents=True, exist_ok=True) + namespace = 'pool_acceptance_' + uuid.uuid4().hex[:8] + os.environ.update(TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE='memory', + TWINKLE_GPU_PG_CPUS_PER_PROC='1', TWINKLE_TRUST_REMOTE_CODE='1', + TWINKLE_DEFAULT_SAVE_DIR=str(directory / 'checkpoints')) + evidence = {'namespace': namespace, 'model': str(model), 'import': twinkle.__file__, + 'sha': subprocess.check_output(['git', 'rev-parse', 'HEAD'], text=True).strip()} + ray.init(address='local', namespace=namespace, num_cpus=24, + num_gpus=int(os.environ.get('TWINKLE_TEST_POOL_DEVICES', '4')), + _temp_dir='/tmp/twp-' + namespace[-8:], include_dashboard=False, dashboard_port=0) + try: + serve.start(http_options={'host': '127.0.0.1', 'port': 0}) + asyncio.run(exercise(model, directory, namespace, evidence, extended=extended)) + finally: + serve.shutdown() + ray.shutdown() + + +def test_real_pool_scaling(): + _run_real_pool(extended=False) + + +def test_real_pool_lifecycle(): + _run_real_pool(extended=True) diff --git a/tests/server/integration/st_cs/__init__.py b/tests/server/integration/st_cs/__init__.py new file mode 100644 index 000000000..85b3e739d --- /dev/null +++ b/tests/server/integration/st_cs/__init__.py @@ -0,0 +1 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. diff --git a/tests/server/integration/st_cs/test_http_e2e.py b/tests/server/integration/st_cs/test_http_e2e.py new file mode 100644 index 000000000..aadd83b19 --- /dev/null +++ b/tests/server/integration/st_cs/test_http_e2e.py @@ -0,0 +1,30 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Opt-in ST HTTP acceptance on two reserved devices (CS=0, bare reference=1).""" +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest + +pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_ST_CS') != '1', + reason='Set TWINKLE_TEST_ST_CS=1 on two reserved devices') + + +def test_st_http_training_isolation_and_resume(tmp_path): + model = os.environ['TWINKLE_TEST_ST_MODEL'] + directory = Path(os.environ.get('TWINKLE_TEST_ARTIFACT_DIR', str(tmp_path))) + directory.mkdir(parents=True, exist_ok=True) + command = [sys.executable, '-u', str(Path(__file__).with_name('validation.py')), + '--model', model, '--run-dir', str(directory)] + root = Path(__file__).resolve().parents[4] + env = {**os.environ, 'PYTHONPATH': os.pathsep.join(filter(None, ( + str(root / 'src'), str(root), os.environ.get('PYTHONPATH'))))} + with (directory / 'server-and-client.log').open('w') as output: + result = subprocess.run(command, stdout=output, stderr=subprocess.STDOUT, env=env, timeout=1800) + assert result.returncode == 0, (directory / 'server-and-client.log').read_text()[-6000:] + evidence = json.loads((directory / 'results.json').read_text()) + assert evidence['passed'] and evidence['slot_reuse'] + assert evidence['bare_parameter_error'] < 2e-5 and evidence['embedding_delta'] > 1e-6 + assert evidence['cold_progress'] == evidence['progress'] diff --git a/tests/server/integration/st_cs/validation.py b/tests/server/integration/st_cs/validation.py new file mode 100644 index 000000000..5f116d59c --- /dev/null +++ b/tests/server/integration/st_cs/validation.py @@ -0,0 +1,209 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Owned ST HTTP acceptance and matching bare-library numerical reference.""" +import argparse +import copy +import json +import os +from pathlib import Path +import socket +import subprocess +import sys +import uuid + +import numpy as np +from peft import LoraConfig +from safetensors.torch import load_file +from transformers import AutoTokenizer + +from cookbook.client.swift_model.sentence_transformer_client import MODEL, TEXTS, configure, encoded_batch, step, vectors + + +def reference(args): + import torch + from sentence_transformers import SentenceTransformer + from swift.dev.model.multi_lora_sentence_transformer_model import MultiLoraSentenceTransformerModel + from twinkle import initialize + initialize(mode='local') + batch = encoded_batch(AutoTokenizer.from_pretrained(args.model)) + model = MultiLoraSentenceTransformerModel(model_id=args.model, mixed_precision='no', + gradient_checkpointing=False, max_loras=2, max_r=8, + target_modules=['query', 'value']) + for name, lr in [('a', 1e-3), ('b', 2e-3)]: + model.add_adapter_to_model(name, LoraConfig(r=4, target_modules=['query', 'value'])) + model.set_template('Template', max_length=128 if name == 'a' else 96, adapter_name=name) + model.set_processor('InputProcessor', padding_side='right', adapter_name=name) + model.set_loss('InfonceLoss', temperature=0.1, adapter_name=name) + model.set_optimizer('AdamW', lr=lr, adapter_name=name) + # Restore the initial checkpoint's RNG and optimizer as well as weights. + # This aligns dropout for a numerical update comparison across processes/devices. + model.resume_from_checkpoint(args.initial, adapter_name='a') + before = model.forward_only(inputs=copy.deepcopy(batch), adapter_name='a')['embeddings'].detach().cpu().numpy() + for _ in range(2): + model.forward_backward(inputs=copy.deepcopy(batch), adapter_name='a', task='embedding') + model.clip_grad_and_step(adapter_name='a') + after = model.forward_only(inputs=copy.deepcopy(batch), adapter_name='a')['embeddings'].detach().cpu().numpy() + saved = model.save('reference', output_dir=args.run_dir, adapter_name='a', save_optimizer=True, + consumed_train_samples=2 * len(batch)) + native = SentenceTransformer(args.model, device='cuda').eval() + with torch.inference_mode(): + raw = native.encode(TEXTS, convert_to_numpy=True, show_progress_bar=False) + Path(args.run_dir, 'reference.json').write_text(json.dumps({ + 'before': before.tolist(), 'after': after.tolist(), 'raw': raw.tolist(), 'checkpoint': saved})) + del native, model + + +def validate_http(base_url, controller, args): + from twinkle import init_twinkle_client + import ray + from tests.server.integration.on_demand_training.acceptance_support import wait_until + + directory = Path(args.run_dir) + batch = encoded_batch(AutoTokenizer.from_pretrained(args.model)) + clients, live = [], [] + + def create(token, label, lr=1e-3, max_length=128): + client = init_twinkle_client(base_url=base_url, api_key=token) + clients.append(client) + model = client.model(MODEL) + model.add_adapter_to_model(label, LoraConfig(r=4, target_modules=['query', 'value'])) + live.append(model) + configure(model, lr=lr, max_length=max_length) + return model + + def release(model): + model.remove_adapter() + live.remove(model) + + def save(model, name, samples): + return model.save(name, save_optimizer=True, consumed_train_samples=samples) + + def parameters(path): + return load_file(str(Path(path) / 'adapter_model.safetensors')) + + def error(a, b): + return float(np.max(np.abs(np.asarray(a) - np.asarray(b)))) + + try: + assert ray.get(controller.diagnostics.remote())['instances'] == [] + a = create('st-owner-a', 'tenant') + b = create('st-owner-b', 'tenant', lr=2e-3, max_length=96) + aid, bid = a._training_ids['tenant'], b._training_ids['tenant'] + diagnostic = ray.get(controller.diagnostics.remote()) + active = [row for row in diagnostic['instances'] if row['state'] == 'READY'] + assert len(active) == 1 and set(active[0]['training_ids']) == {aid, bid} + assert 'Learning rate: 0.001' in a.get_train_configs().result + assert 'Learning rate: 0.002' in b.get_train_configs().result + initial_a, initial_b = vectors(a, batch), vectors(b, batch) + initial = save(a, 'initial', 0) + for _ in range(2): + step(a, batch) + trained = vectors(a, batch) + assert error(trained, initial_a) > 1e-6 + np.testing.assert_allclose(vectors(b, batch), initial_b, rtol=0, atol=0) + snapshot = save(a, 'trained', 2 * len(batch)) + # Independent bare-library process uses the other, explicitly reserved PPU. + subprocess.run([sys.executable, str(Path(__file__).resolve()), '--reference', '--model', args.model, + '--run-dir', args.run_dir, '--initial', initial.checkpoint_dir], + env={**os.environ, 'CUDA_VISIBLE_DEVICES': args.reference_device}, check=True, timeout=600) + bare = json.loads((directory / 'reference.json').read_text()) + np.testing.assert_allclose(initial_a, bare['before'], rtol=0, atol=2e-5) + np.testing.assert_allclose(trained, bare['after'], rtol=0, atol=2e-5) + np.testing.assert_allclose(a.forward_only(inputs=batch, disable_lora=True).result['embeddings'], + bare['raw'], rtol=0, atol=2e-5) + left, right = parameters(snapshot.checkpoint_dir), parameters(bare['checkpoint']) + assert left.keys() == right.keys() + parameter_error = max(float((left[key] - right[key]).abs().max()) for key in left) + assert parameter_error < 2e-5, parameter_error + initial_weights = parameters(initial.checkpoint_dir) + parameter_delta = max(float((left[key] - initial_weights[key]).abs().max()) for key in left) + assert parameter_delta > 1e-6 + step(b, batch) + np.testing.assert_allclose(vectors(a, batch), trained, rtol=0, atol=0) + a.resume_from_checkpoint(snapshot.twinkle_path) + step(a, batch) + continued = vectors(a, batch) + release(a) + c = create('st-owner-a', 'reused') + fresh = save(c, 'fresh', 0) + fresh_progress = json.loads((Path(fresh.checkpoint_dir) / 'trainer_state.json').read_text()) + assert fresh_progress['cur_step'] == 0 + np.testing.assert_allclose(vectors(c, batch), initial_a, rtol=0, atol=2e-5) + progress = c.resume_from_checkpoint(snapshot.twinkle_path) + assert progress['consumed_train_samples'] == 2 * len(batch) and progress['cur_step'] == 2 + np.testing.assert_allclose(vectors(c, batch), trained, rtol=0, atol=2e-5) + step(c, batch) + np.testing.assert_allclose(vectors(c, batch), continued, rtol=0, atol=2e-5) + release(c) + release(b) + wait_until(lambda: all(row['state'] == 'STOPPED' for row in ray.get( + controller.diagnostics.remote())['instances']), 180, 'ST resources were not released') + d = create('st-owner-a', 'cold-restored') + cold_progress = d.resume_from_checkpoint(snapshot.twinkle_path) + assert cold_progress == progress + np.testing.assert_allclose(vectors(d, batch), trained, rtol=0, atol=2e-5) + step(d, batch) + release(d) + return {'passed': True, 'shared_instance': active[0]['instance_id'], 'initial_ids': [aid, bid], + 'embedding_delta': error(trained, initial_a), 'bare_embedding_error': error(trained, bare['after']), + 'bare_parameter_error': parameter_error, 'parameter_delta': parameter_delta, + 'progress': progress, 'cold_progress': cold_progress, 'slot_reuse': True} + finally: + for model in reversed(live): + try: + model.remove_adapter() + except Exception: + pass + for client in clients: + client.close() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--model', required=True, help='Original ST model cache directory') + parser.add_argument('--run-dir', required=True) + parser.add_argument('--reference', action='store_true') + parser.add_argument('--initial') + parser.add_argument('--reference-device', default='1') + args = parser.parse_args() + Path(args.run_dir).mkdir(parents=True, exist_ok=True) + if args.reference: + reference(args) + return + import ray + from ray import serve + from twinkle.server.config import ServerConfig + from twinkle.server.launcher.server_launcher import ServerLauncher + root = Path(__file__).resolve().parents[4] + config = ServerConfig.from_yaml(root / 'cookbook/client/swift_model/sentence_transformer_config.yaml') + with socket.socket() as sock: + sock.bind(('127.0.0.1', 0)) + config.http_options.port = sock.getsockname()[1] + namespace = 'st_cs_' + uuid.uuid4().hex[:8] + config.ray_namespace = namespace + model_app = next(app for app in config.applications if app.import_path == 'model') + model_app.args.model_id = args.model + model_app.args.idle_timeout_s = 900 + os.environ.update(TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE='memory', + TWINKLE_DEFAULT_SAVE_DIR=str(Path(args.run_dir).resolve() / 'checkpoints'), + TWINKLE_TRUST_REMOTE_CODE='1', + TWINKLE_MODEL_ID_ALIASES=json.dumps({MODEL: args.model})) + ray.init(address='local', namespace=namespace, num_cpus=16, num_gpus=1, + _temp_dir='/tmp/twst-' + namespace[-8:], include_dashboard=False) + launcher = ServerLauncher(config) + try: + launcher._start_serve() + for app in config.applications: + if app.import_path != 'model': + launcher._deploy_application(app) + launcher._deploy_training() + result = validate_http(f'http://127.0.0.1:{config.http_options.port}', launcher._training_controller, args) + result['validation_idle_timeout_s'] = model_app.args.idle_timeout_s + Path(args.run_dir, 'results.json').write_text(json.dumps(result, indent=2)) + finally: + launcher._shutdown() + serve.shutdown() + ray.shutdown() + + +if __name__ == '__main__': + main() diff --git a/tests/server/integration/task_http_validation.py b/tests/server/integration/task_http_validation.py new file mode 100644 index 000000000..7a051db22 --- /dev/null +++ b/tests/server/integration/task_http_validation.py @@ -0,0 +1,229 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Owned real-device CS acceptance: training progress, tenant isolation, checkpoint and release.""" +import argparse +import copy +import json +import os +from pathlib import Path +import socket +import uuid +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from peft import LoraConfig +from safetensors.torch import load_file +from transformers import AutoTokenizer + + +def validate(base_url, controller, args): + import ray + from twinkle import init_twinkle_client + from tests.server.integration.cs_tasks.task_models import encoded_batch + from tests.server.integration.on_demand_training.acceptance_support import wait_until + + base_path = str(Path(args.model_dir, args.task).resolve()) + batch = encoded_batch(args.task, AutoTokenizer.from_pretrained(base_path)) + targets = (['q_proj', 'v_proj'] if args.backend == 'transformers' else ['linear_qkv']) + if args.task != 'embedding': + targets += ['score' if args.backend == 'transformers' else 'output_layer'] + forward_kwargs = {'task': args.task, 'return_logits': True, 'compute_loss': False} + if args.backend == 'megatron': + forward_kwargs['micro_batch_size'] = 6 + + def stopped(): + diagnostics = ray.get(controller.diagnostics.remote(base_model='cs-' + args.task)) + return all(row['state'] == 'STOPPED' for row in diagnostics['instances']) + + clients, live = [], [] + + def create(token, label): + client = init_twinkle_client(base_url=base_url, api_key=token) + clients.append(client) + model = client.model('cs-' + args.task) + model.add_adapter_to_model(label, LoraConfig(r=4, lora_alpha=8, target_modules=targets), + gradient_accumulation_steps=2) + live.append(model) + model.set_template('Template', max_length=128) + model.set_processor('InputProcessor', padding_side='right') + if args.task == 'embedding': + model.set_loss('InfonceLoss', temperature=0.1, use_batch=True) + elif args.task == 'reranker': + model.set_loss('PointwiseRerankerLoss') + else: + model.set_loss('SeqClsLoss', problem_type='single_label_classification', num_labels=3) + model.set_optimizer('Adam', lr=1e-3) + return model + + def release(model): + model.remove_adapter() + live.remove(model) + + def read(model): + result = model.forward_only(inputs=copy.deepcopy(batch), **forward_kwargs).result + value = np.asarray(result['embeddings' if args.task == 'embedding' else 'logits'], dtype=float) + assert value.shape[0] == len(batch) and np.isfinite(value).all() + if args.task != 'embedding': + assert value.shape == (len(batch), 1 if args.task == 'reranker' else 3) + return value + + def step(model): + kwargs = {'task': args.task, 'gradient_accumulation_steps': 2} + if args.backend == 'megatron': + kwargs['micro_batch_size'] = 6 + result = model.forward_backward(inputs=copy.deepcopy(batch), **kwargs).result + assert np.isfinite(result['loss']) + # Call after every backward; the backend gates updates at its GA boundary. + model.clip_grad_and_step() + return float(result['loss']) + + def save(model, name, samples): + return model.save(name, save_optimizer=True, consumed_train_samples=samples) + + def weights(saved): + return load_file(str(Path(saved.checkpoint_dir, 'adapter_model.safetensors'))) + + try: + # Longer single-tenant run: the loss must clearly fall and the checkpoint must record progress. + trainer = create('task-owner-train', 'train') + before = read(trainer) + losses = [step(trainer) for _ in range(args.steps)] + assert np.max(np.abs(read(trainer) - before)) > 1e-7 + first_loss = float(np.mean(losses[:3])) + last_loss = float(np.mean(losses[-3:])) + loss_drop = (first_loss - last_loss) / first_loss + assert loss_drop >= args.min_loss_drop, (losses, loss_drop) + trained_progress = save(trainer, 'train-final', args.steps * len(batch)) + state = json.loads(Path(trained_progress.checkpoint_dir, 'trainer_state.json').read_text()) + assert state['cur_step'] == args.steps and state['consumed_train_samples'] == args.steps * len(batch) + release(trainer) + wait_until(stopped, 180, 'Training run did not release the instance') + if args.training_only: + return {'passed': True, 'task': args.task, 'backend': args.backend, 'slot_reuse': None, + 'losses': losses, 'loss_drop': loss_drop} + + a, b = create('task-owner-a', 'tenant'), create('task-owner-b', 'tenant') + diagnostics = ray.get(controller.diagnostics.remote(base_model='cs-' + args.task)) + ready = [row for row in diagnostics['instances'] if row['state'] == 'READY'] + assert len(ready) == 1 and len(ready[0]['training_ids']) == 2 + initial, peer = read(a), read(b) + initial_weights = weights(save(a, 'initial', 0)) + for _ in range(3): + step(a) + trained = read(a) + np.testing.assert_allclose(read(b), peer, rtol=0, atol=0) + snapshot = save(a, 'trained', 3 * len(batch)) + changed = weights(snapshot) + assert changed.keys() == initial_weights.keys() + updates = [key for key in changed if not np.array_equal(changed[key].float().numpy(), + initial_weights[key].float().numpy())] + assert updates, 'No adapter parameter changed' + if args.task != 'embedding': + assert any('score.' in key or 'output_layer.' in key for key in updates), 'Head LoRA not saved/updated' + step(a) + step(a) + continued = read(a) + release(a) + c = create('task-owner-a', 'reused') + fresh = save(c, 'fresh', 0) + fresh_state = json.loads(Path(fresh.checkpoint_dir, 'trainer_state.json').read_text()) + assert fresh_state['cur_step'] == 0 + np.testing.assert_allclose(read(c), initial, rtol=0, atol=2e-5) + progress = c.resume_from_checkpoint(snapshot.twinkle_path) + assert progress['cur_step'] == 3 and progress['consumed_train_samples'] == 3 * len(batch) + np.testing.assert_allclose(read(c), trained, rtol=0, atol=2e-5) + step(c) + step(c) + np.testing.assert_allclose(read(c), continued, rtol=0, atol=2e-5) + np.testing.assert_allclose(read(b), peer, rtol=0, atol=0) + release(c) + release(b) + wait_until(stopped, 180, 'Tenant instances did not stop') + cold = create('task-owner-a', 'cold') + assert cold.resume_from_checkpoint(snapshot.twinkle_path) == progress + np.testing.assert_allclose(read(cold), trained, rtol=0, atol=2e-5) + step(cold) + step(cold) + np.testing.assert_allclose(read(cold), continued, rtol=0, atol=2e-5) + release(cold) + wait_until(stopped, 180, 'Cold instance did not stop') + return {'passed': True, 'task': args.task, 'backend': args.backend, + 'slot_reuse': True, 'progress': progress, 'head_update_keys': updates, + 'losses': losses, 'loss_drop': loss_drop} + finally: + try: + for model in reversed(live): + model.remove_adapter() + finally: + for client in clients: + client.close() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--backend', choices=['transformers', 'megatron'], required=True) + parser.add_argument('--task', choices=['embedding', 'reranker', 'seq_cls'], nargs='+', required=True) + parser.add_argument('--model-dir', required=True, + help='Bases from cs_tasks/task_models.py; Megatron Qwen3.5 needs the -mcore output') + parser.add_argument('--run-dir', required=True) + parser.add_argument('--steps', type=int, default=31) + parser.add_argument('--reranker-steps', type=int, default=63) + parser.add_argument('--min-loss-drop', type=float, default=0.1) + parser.add_argument('--training-only', action='store_true') + args = parser.parse_args() + directory = Path(args.run_dir).resolve() + directory.mkdir(parents=True, exist_ok=True) + import ray + from ray import serve + from twinkle.server.config import ServerConfig + from twinkle.server.launcher.server_launcher import ServerLauncher + + config = ServerConfig.from_yaml(Path(__file__).with_name('cs_tasks') / f'task_server_{args.backend}.yaml') + model_apps = [] + aliases = {} + for task in args.task: + app = next(app for app in config.applications if app.name == ('seq-cls' if task == 'seq_cls' else task)) + app.args.model_id = str(Path(args.model_dir, task).resolve()) + model_apps.append(app) + aliases['cs-' + task] = app.args.model_id + config.applications = [app for app in config.applications if app.import_path == 'server'] + model_apps + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + config.http_options.port = listener.getsockname()[1] + namespace = 'cs_tasks_' + uuid.uuid4().hex[:8] + config.ray_namespace = namespace + os.environ.update(TWINKLE_STATE_RAY_NAMESPACE=namespace, TWINKLE_PERSISTENCE_MODE='memory', + TWINKLE_DEFAULT_SAVE_DIR=str(directory / 'checkpoints'), TWINKLE_TRUST_REMOTE_CODE='1', + TWINKLE_MODEL_ID_ALIASES=json.dumps(aliases)) + ray.init(address='local', namespace=namespace, num_cpus=8 * len(args.task), + num_gpus=(2 if args.backend == 'megatron' else 1) * len(args.task), + _temp_dir='/tmp/twtask-' + namespace[-8:], include_dashboard=False) + launcher = ServerLauncher(config) + task_args = [] + for task in args.task: + current = copy.copy(args) + current.task = task + current.steps = args.reranker_steps if task == 'reranker' else args.steps + current.run_dir = str(directory / (args.backend + '-' + task)) + Path(current.run_dir).mkdir(parents=True, exist_ok=True) + task_args.append(current) + try: + launcher.launch(blocking=False) + base_url = f'http://127.0.0.1:{config.http_options.port}' + with ThreadPoolExecutor(max_workers=len(task_args)) as executor: + results = list(executor.map(lambda current: validate(base_url, launcher._training_controller, current), + task_args)) + for current, result in zip(task_args, results): + (Path(current.run_dir) / 'results.json').write_text(json.dumps(result, indent=2)) + summary = {'backend': args.backend, 'steps': args.steps, 'min_loss_drop': args.min_loss_drop, + 'tasks': results} + (directory / 'summary.json').write_text(json.dumps(summary, indent=2)) + finally: + try: + launcher._shutdown() + finally: + serve.shutdown() + ray.shutdown() + + +if __name__ == '__main__': + main() diff --git a/tests/server/integration/test_actor_recovery.py b/tests/server/integration/test_actor_recovery.py index 16527bbae..dfc03a74b 100644 --- a/tests/server/integration/test_actor_recovery.py +++ b/tests/server/integration/test_actor_recovery.py @@ -40,7 +40,7 @@ async def call_backend(self, fn, /, *args, admit: bool = True, **kwargs): @pytest.mark.asyncio -async def test_timeout_probe_marks_unhealthy_then_recovers(): +async def test_timeout_probe_keeps_runtime_unhealthy_even_if_actor_recovers(): model = _MockModel() h = _HealthHarness(model) @@ -54,11 +54,12 @@ async def test_timeout_probe_marks_unhealthy_then_recovers(): await h._probe_after_timeout() assert h._model_unhealthy is True # /healthz would return 503 - # Actor recovers; one successful probe clears the bit (no restart needed). + # Actor liveness cannot prove that timed-out training stopped mutating state. + # The controller must replace this runtime before accepting more work. model.alive = True result = await h.check_model_health() - assert result['healthy'] is True - assert h._model_unhealthy is False + assert result['healthy'] is False + assert h._model_unhealthy is True @pytest.mark.asyncio @@ -76,11 +77,11 @@ async def test_health_route_returns_503_when_probe_fails(): @pytest.mark.asyncio -async def test_mark_unhealthy_is_cleared_by_successful_probe(): +async def test_mark_unhealthy_is_not_cleared_by_successful_probe(): h = _HealthHarness(_MockModel()) h.mark_unhealthy() assert h._model_unhealthy is True result = await h.check_model_health() - assert result['healthy'] is True - assert h._model_unhealthy is False + assert result['healthy'] is False + assert h._model_unhealthy is True diff --git a/tests/server/integration/test_blocking_boundary.py b/tests/server/integration/test_blocking_boundary.py index 7d35e229d..7fb895993 100644 --- a/tests/server/integration/test_blocking_boundary.py +++ b/tests/server/integration/test_blocking_boundary.py @@ -21,7 +21,7 @@ from twinkle.server.task_queue.backend_gate import BackendGate # noqa: E402 from twinkle.server.task_queue.mixin import TaskQueueMixin # noqa: E402 -from twinkle.server.task_queue.types import BackendBusyError # noqa: E402 +from twinkle.server.exceptions import BackendBusyError # noqa: E402 class _Harness(TaskQueueMixin): diff --git a/tests/server/integration/test_full_cycle_e2e.py b/tests/server/integration/test_full_cycle_e2e.py index 6fa31f006..ec7a4c52d 100644 --- a/tests/server/integration/test_full_cycle_e2e.py +++ b/tests/server/integration/test_full_cycle_e2e.py @@ -6,13 +6,12 @@ Backend selection via env var TWINKLE_TEST_BACKEND: - "transformers" (default): all 6 phases run strictly - - "megatron": Phase C/D skipped (known multi-LoRA strict-load bug), - Phase E/F best-effort (GPU OOM possible) + - "megatron": all 6 phases run strictly (requires a working optimizer-resume backend) Phase A — initial training (STEPS_PHASE_A steps) Phase B — keep training STEPS_PHASE_B more steps, save again Phase C — RELOAD VERIFY: load() ckpt_a, forward_only on fixed batch, - assert loss matches Phase A end + assert token log probabilities match Phase A end Phase D — RESUME VERIFY: resume_from_checkpoint(ckpt_b), train more Phase E + F — vLLM LoRA-effect greedy probe @@ -58,13 +57,13 @@ assert BACKEND in ('transformers', 'megatron'), f'Invalid TWINKLE_TEST_BACKEND={BACKEND!r}' BASE_MODEL = 'Qwen/Qwen3.5-4B' -BASE_URL = 'http://localhost:9000' +BASE_URL = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:9000') API_KEY = 'EMPTY_API_KEY' -SAVE_DIR = '/mnt/nas2/yunlin.myl/twinkle/output/twinkle_e2e_full_cycle' +SAVE_DIR = os.environ.get('TWINKLE_DEFAULT_SAVE_DIR', 'output/twinkle_e2e_full_cycle') STEPS_PHASE_A = 60 STEPS_PHASE_B = 4 STEPS_PHASE_D = 4 -RELOAD_LOSS_TOLERANCE = 0.05 # |reloaded - original| / original +RELOAD_LOGPROB_ATOL = 2e-3 # Same absolute tolerance as controller-restart acceptance. RESUME_LOSS_BAND = 3.0 # resumed step's loss must be within this factor of Phase-B's last SAMPLE_MAX_TOKENS = 32 @@ -125,21 +124,13 @@ def _train_n_steps(model, dataloader, n: int, *, label: str, start_step: int = 0 return losses -def _record_fixed_batch_loss(model, batch, *, label: str) -> float: - """Run forward_only on a fixed batch and report the loss for reload comparison. - - Uses ``forward_only`` + ``calculate_metric(is_training=False)`` to retrieve - loss. This is compatible with both Transformers and Megatron backends - (Megatron does not support standalone ``calculate_loss``). - """ - model.forward_only(inputs=batch) - metric = model.calculate_metric(is_training=False) - try: - val = float(metric.result.get('loss')) if hasattr(metric.result, 'get') else float(metric.result['loss']) - except (TypeError, KeyError): - val = float(metric.result) - logger.info(f'[{label}] fixed-batch loss = {val:.4f}') - return val +def _record_fixed_batch_logps(model, batch, *, label: str): + """Read actual forward outputs; forward_only need not populate loss metrics.""" + from tests.server.integration.on_demand_training.acceptance_support import checked_logps + values = checked_logps(model.forward_only(inputs=batch).result['logps']) + logger.info('[%s] fixed-batch logps: count=%d min=%.6f max=%.6f', + label, values.size, values.min(), values.max()) + return values # Probe prompts for vLLM LoRA-effect verification. Mixing different prompt @@ -217,7 +208,7 @@ def main() -> int: # Capture a fixed batch we'll re-use after reload to assert state survived. fixed_batch = next(iter(_build_dataset_loader())) - loss_a_fixed = _record_fixed_batch_loss(model_a, fixed_batch, label='A-fixed') + logps_a_fixed = _record_fixed_batch_logps(model_a, fixed_batch, label='A-fixed') ckpt_a_resp = model_a.save( name='e2e-phase-a', @@ -249,11 +240,11 @@ def main() -> int: logger.info('Phase C: reload-verify (load ckpt_a, fixed batch)') logger.info('=' * 60) model_a.load(ckpt_a) - loss_c_fixed = _record_fixed_batch_loss(model_a, fixed_batch, label='C-fixed') - delta = abs(loss_c_fixed - loss_a_fixed) / max(abs(loss_a_fixed), 1e-6) - logger.info(f'Phase C reload delta: |{loss_c_fixed:.4f} - {loss_a_fixed:.4f}| / {loss_a_fixed:.4f} = {delta:.4f}') - assert delta <= RELOAD_LOSS_TOLERANCE, ( - f'Phase C FAILED: reload delta {delta:.4f} > tolerance {RELOAD_LOSS_TOLERANCE}') + import numpy as np + logps_c_fixed = _record_fixed_batch_logps(model_a, fixed_batch, label='C-fixed') + np.testing.assert_allclose(logps_c_fixed, logps_a_fixed, rtol=0, atol=RELOAD_LOGPROB_ATOL) + delta = float(np.max(np.abs(logps_c_fixed - logps_a_fixed))) + logger.info('Phase C reload max logprob error: %.6f', delta) # --------------------------------------------------------------------- # Phase D — RESUME VERIFY @@ -307,8 +298,7 @@ def main() -> int: logger.info('=' * 60) logger.info(' Phase A losses (%d steps): first=%.3f last=%.3f', len(losses_a), losses_a[0][1], losses_a[-1][1]) logger.info(' Phase B losses (%d steps): %s', len(losses_b), [f'{l:.3f}' for _, l in losses_b]) - logger.info(' Phase C reload: |%.4f - %.4f| / %.4f = %.4f (tol %.2f)', loss_c_fixed, loss_a_fixed, - loss_a_fixed, delta, RELOAD_LOSS_TOLERANCE) + logger.info(' Phase C reload max logprob error: %.6f (atol %.4f)', delta, RELOAD_LOGPROB_ATOL) logger.info(' Phase D resume losses (%d steps): %s', len(losses_d), [f'{l:.3f}' for _, l in losses_d]) logger.info(' Phase F LoRA-effect probes (%d/%d diverged):', n_diverged, len(PROBE_PROMPTS)) for prompt, _, _, div in probe_results: diff --git a/tests/server/integration/test_full_param_e2e.py b/tests/server/integration/test_full_param_e2e.py index 857fbc711..32d605191 100644 --- a/tests/server/integration/test_full_param_e2e.py +++ b/tests/server/integration/test_full_param_e2e.py @@ -117,7 +117,7 @@ def main() -> int: raise except Exception as e: # noqa: BLE001 msg = str(e) - assert 'full-parameter' in msg or 'lora_config' in msg, f'Phase B FAIL: unexpected error: {msg[:500]}' + assert 'lora' in msg.lower(), f'Phase B FAIL: unexpected error: {msg[:500]}' logger.info('Phase B OK: rejected with mode-mismatch error') # ── Phase C: second full tenant must be rejected (exclusive) ── @@ -130,8 +130,8 @@ def main() -> int: raise except Exception as e: # noqa: BLE001 msg = str(e) - assert 'exclusive' in msg or 'already' in msg, f'Phase C FAIL: unexpected error: {msg[:500]}' - logger.info('Phase C OK: rejected with exclusive-busy error') + assert 'replica limit' in msg.lower() or 'allocation' in msg.lower(), f'Phase C FAIL: unexpected error: {msg[:500]}' + logger.info('Phase C OK: rejected at configured instance capacity') # ── Phase D: save full weights for sampler + sample ── logger.info('Phase D: save_weights_and_get_sampling_client + sample') diff --git a/tests/server/integration/test_mock_mode_startup.py b/tests/server/integration/test_mock_mode_startup.py index d5f706061..e377afff2 100644 --- a/tests/server/integration/test_mock_mode_startup.py +++ b/tests/server/integration/test_mock_mode_startup.py @@ -3,11 +3,11 @@ Launches the all-mock cookbook config inside the test process via Ray Serve, then asserts: -- Both Model and Sampler deployments report RUNNING within 30 seconds. +- Managed training and sampler routes become ready within the 120-second startup budget. - Repeated calls to the mock model and mock sampler over HTTP return byte-identical responses for identical input. -- The launch path imports cleanly even when ``transformers`` / ``vllm`` / - ``megatron`` would not be available — the mock branches don't pull them. +- The actual launcher normalizes the original applications YAML and deploys + CPU routes and managed mock instances without loading real model weights. This test is heavier than the property suite (boots a full Ray Serve cluster) and is gated behind ``TWINKLE_TEST_INTEGRATION=1`` so plain @@ -18,8 +18,6 @@ import httpx import os import pytest -import subprocess -import sys import time from tests.server.fixtures import MOCK_SERVER_CONFIG, MOCK_SERVER_CONFIG_REDIS @@ -37,23 +35,10 @@ SELECTED_CONFIG = ( MOCK_SERVER_CONFIG_REDIS if os.environ.get('TWINKLE_TEST_REDIS_PERSISTENCE', '0') == '1' else MOCK_SERVER_CONFIG) -READY_BUDGET_SECONDS = 30.0 +READY_BUDGET_SECONDS = 120.0 RAY_NODE_CPUS = 8 -def _run_ray_command(*args: str) -> None: - ray_bin = os.path.join(os.path.dirname(sys.executable), 'ray') - result = subprocess.run( - [ray_bin, *args], - check=False, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - if result.returncode != 0: - raise RuntimeError(f'ray {" ".join(args)} failed with code {result.returncode}:\n{result.stdout}') - - @pytest.fixture(scope='module') def ray_cluster(): """Start a fresh local Ray head node for the duration of the module. @@ -78,18 +63,8 @@ def ray_cluster(): if ray.is_initialized(): ray.shutdown() - _run_ray_command('stop', '--force') - _run_ray_command( - 'start', - '--head', - '--port=0', - f'--num-cpus={RAY_NODE_CPUS}', - '--num-gpus=0', - '--include-dashboard=false', - '--disable-usage-stats', - ) ray.init( - address='auto', + address='local', num_cpus=RAY_NODE_CPUS, num_gpus=0, include_dashboard=False, runtime_env={'env_vars': persistence_env} if persistence_env else None, ) yield @@ -101,10 +76,6 @@ def ray_cluster(): ray.shutdown() except Exception: pass - try: - _run_ray_command('stop', '--force') - except Exception: - pass def _wait_until_healthy(serve_module, timeout: float) -> dict: @@ -124,12 +95,10 @@ def _http(url: str, method: str = 'GET', json: dict | None = None) -> httpx.Resp return httpx.request(method, url, json=json, timeout=10.0) -def test_mock_mode_reaches_ready_under_30s_and_is_deterministic(ray_cluster) -> None: +def test_mock_mode_reaches_ready_and_is_deterministic(ray_cluster, monkeypatch) -> None: + import ray from ray import serve - - from twinkle.server.gateway import build_gateway_app - from twinkle.server.model import build_model_app - from twinkle.server.sampler import build_sampler_app + from twinkle.server.launcher.server_launcher import ServerLauncher cfg = ServerConfig.from_yaml(SELECTED_CONFIG) @@ -139,44 +108,14 @@ def test_mock_mode_reaches_ready_under_30s_and_is_deterministic(ray_cluster) -> serve.start(http_options={'host': host, 'port': port}) started = time.monotonic() - deploys: list[tuple[str, str]] = [] - builders = { - 'server': build_gateway_app, - 'model': build_model_app, - 'sampler': build_sampler_app, - } + namespace = ray.get_runtime_context().namespace + monkeypatch.setenv('TWINKLE_STATE_RAY_NAMESPACE', namespace) + cfg = cfg.model_copy(update={'ray_namespace': namespace, + 'http_options': cfg.http_options.model_copy(update={'host': host, 'port': port})}) + launcher = ServerLauncher(cfg) for app_spec in cfg.applications: - builder = builders[app_spec.import_path] - # Shallow-dump so nested typed models (notably ``queue_config``) stay as - # ``TaskQueueConfig`` instances instead of being serialized back to a - # dict — this mirrors what the production launcher does at - # ``launcher/server_launcher.py:161`` and is required since the - # builders + deployment ``__init__`` accept ``TaskQueueConfig`` directly - # (the ``from_dict`` revival path was removed). - args = {k: v for k, v in dict(app_spec.args).items() if v is not None} - if app_spec.import_path == 'server': - # Gateway's ServiceProxy reads http_options.port to build internal - # proxy targets; the fixture hard-codes 8000 but the test runs on - # a randomized port, so override before passing in. - http_opts = cfg.http_options.model_dump() - http_opts['host'] = host - http_opts['port'] = port - args.setdefault('http_options', http_opts) - # Strip ray_actor_options runtime_env to keep the test light, BUT keep - # ``num_cpus`` low so the three deployments leave room for the mock - # distributed runtime workers in the dedicated local Ray node. - deploy_options: dict = {'ray_actor_options': {'num_cpus': 0.1}} - for raw in app_spec.deployments: - if isinstance(raw, dict): - deploy_options = { - k: v - for k, v in raw.items() if k not in ('name', 'ray_actor_options', 'autoscaling_config') - } - deploy_options['ray_actor_options'] = {'num_cpus': 0.1} - break - bound = builder(deploy_options=deploy_options, **args) - serve.run(bound, name=app_spec.name, route_prefix=app_spec.route_prefix) - deploys.append((app_spec.name, app_spec.route_prefix)) + launcher._deploy_application(app_spec) + launcher._deploy_training() statuses = _wait_until_healthy(serve, READY_BUDGET_SECONDS) elapsed = time.monotonic() - started @@ -207,6 +146,16 @@ def test_mock_mode_reaches_ready_under_30s_and_is_deterministic(ray_cluster) -> # separately because tinker's polling can stall the test budget on flaky # cross-replica state. if os.environ.get('TWINKLE_TEST_TINKER', '0') == '1': + # The native patch dirties the resident instance; release replaces it. + # A create during that replacement is correctly rejected, not queued. + deadline = time.monotonic() + READY_BUDGET_SECONDS + while time.monotonic() < deadline: + health = _http(f'{base}/api/v1/model/mock-model/healthz').json() + if health.get('healthy'): + break + time.sleep(0.2) + else: + raise AssertionError('Resident replacement did not become ready') _exercise_tinker_client(base) @@ -257,20 +206,23 @@ def _exercise_twinkle_clients(base: str) -> None: save_resp = model.save(name='step-1') assert save_resp.twinkle_path and save_resp.twinkle_path.startswith('twinkle://') + # Mock save creates an empty directory; identify it as the LoRA checkpoint + # expected by the sampler's checkpoint classification introduced on main. + assert save_resp.checkpoint_dir + lora_cfg.save_pretrained(save_resp.checkpoint_dir) model.load(name=save_resp.twinkle_path) model.resume_from_checkpoint(name=save_resp.twinkle_path) model.apply_patch('NoopPatch') - model.upload_to_hub( - checkpoint_dir=save_resp.twinkle_path, - hub_model_id='mock/dummy', - hub_token='EMPTY_TOKEN', - async_upload=False, - poll_interval=0.5, - ) + # Hub upload is checked with a local mocked HubOperation; this smoke + # must not publish files to an external account. # --- sampler service --- - sampler = vLLMSampler(model_id='mock-model') + from twinkle_client.http import ClientTransport + from twinkle_client.http.context import capture_transport + # Cold PEFT imports on NAS can exceed the default 90-second HTTP timeout. + sampler_transport = ClientTransport(capture_transport().context, timeout=600) + sampler = vLLMSampler(model_id='mock-model', transport=sampler_transport) # Sampler /add_adapter_to_sampler reuses the model-side ``AddAdapterRequest`` # shape (``config: str``), so send JSON. Skip a real ``LoraConfig`` here # because its ``runtime_config`` member isn't JSON-serializable. @@ -287,8 +239,7 @@ def _exercise_twinkle_clients(base: str) -> None: ) assert samples and samples[0].sequences and len(samples[0].sequences[0].tokens) == 4 - # adapter_uri triggers reset_prefix_cache; reuse the training-weights - # path because mock ``save()`` materialized that directory on disk. + # Exercise adapter_uri resolution using the mock LoRA checkpoint directory. sampler.sample( inputs=[{ 'input_ids': [1, 2, 3] @@ -299,6 +250,9 @@ def _exercise_twinkle_clients(base: str) -> None: ) sampler.apply_patch('NoopPatch') + sampler_transport.close() + model.remove_adapter() + twinkle.close() def _exercise_tinker_client(base: str) -> None: @@ -339,6 +293,12 @@ def _exercise_tinker_client(base: str) -> None: sampler_ckpt = training.save_weights_for_sampler(name='step-1').result() # path mode: save_weights_for_sampler(name) must return path != None assert sampler_ckpt.path is not None + from pathlib import Path + from peft import LoraConfig + from twinkle.server.checkpoint import create_checkpoint_manager + checkpoint_manager = create_checkpoint_manager('EMPTY_TOKEN', client_type='tinker') + sampler_dir = Path(checkpoint_manager.resolve_checkpoint_dir(sampler_ckpt.path)).resolve() + LoraConfig(r=4, target_modules='all-linear').save_pretrained(sampler_dir) # Gateway's /asample resolves ``base_model`` from ``body.base_model`` or # ``sampling_session_id``; pass it explicitly because the SDK only sets # ``model_path`` and the gateway doesn't parse ``twinkle://`` URIs. @@ -352,6 +312,11 @@ def _exercise_tinker_client(base: str) -> None: # sampling_session_seq_id mode: save_weights_and_get_sampling_client() # must return path == None and sampling_session_id != None (asserted by SDK internally) sampling_client = training.save_weights_and_get_sampling_client() + # Mock saves contain no tensors. Mark the newly versioned sampler checkpoint + # as LoRA, just as the native smoke does, rather than pretending it is full weights. + for checkpoint in sampler_dir.parent.iterdir(): + if checkpoint.is_dir(): + LoraConfig(r=4, target_modules='all-linear').save_pretrained(checkpoint) sampling_client.sample( prompt=types.ModelInput.from_ints([1, 2, 3]), num_samples=1, @@ -362,3 +327,5 @@ def _exercise_tinker_client(base: str) -> None: # the wire-level handler is /tinker/save_weights either way. ckpt = training.save_state(name='step-2').result() training.load_state(ckpt.path).result() + from tests.server.integration.on_demand_training.on_demand_sdk import unload_tinker_model + unload_tinker_model(client, training.model_id) diff --git a/tests/server/integration/test_on_demand_gpu_redis_e2e.py b/tests/server/integration/test_on_demand_gpu_redis_e2e.py new file mode 100644 index 000000000..69fa201a6 --- /dev/null +++ b/tests/server/integration/test_on_demand_gpu_redis_e2e.py @@ -0,0 +1,53 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Opt-in single-node GPU matrix using the real Redis state backend.""" +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.skipif( + os.environ.get('TWINKLE_TEST_GPU_E2E') != '1' or not os.environ.get('TWINKLE_TEST_REDIS_URL'), + reason='Set TWINKLE_TEST_GPU_E2E=1 and TWINKLE_TEST_REDIS_URL for the GPU Redis matrix', +) + + +@pytest.mark.parametrize('backend', ['transformers', 'megatron']) +def test_on_demand_training_gpu_redis_matrix(tmp_path, backend): + model = os.environ.get('TWINKLE_TEST_MODEL') or os.environ.get('TWINKLE_DEVICE_MODEL') + if not model: + pytest.skip('Set TWINKLE_TEST_MODEL or TWINKLE_DEVICE_MODEL') + root = Path(__file__).resolve().parents[3] + run_dir = tmp_path / backend + script = Path(__file__).with_name('on_demand_training') / 'matrix_validation.py' + result = subprocess.run([ + sys.executable, + str(script), + '--backend', + backend, + '--model', + model, + '--run-dir', + str(run_dir), + '--persistence', + 'redis', + '--startup-timeout-s', + '1800', + '--shutdown-timeout-s', + '180', + ], + cwd=root, + capture_output=True, + text=True, + timeout=4 * 60 * 60) + assert result.returncode == 0, result.stdout + result.stderr + evidence = json.loads((run_dir / 'results.json').read_text()) + assert evidence['passed'] + assert evidence['persistence'] == 'redis' + assert evidence['lazy_start'] and evidence['training_resources_released'] + assert evidence['restart_epoch'] and evidence['worker_fault_recovered'] + assert all(row['checkpoint_restored'] for row in evidence['recovery']) + assert evidence['capacity_checked'] and evidence['idempotency_checked'] + assert all(case['returncode'] == 0 and not case['skipped'] for case in evidence['cases']) diff --git a/tests/server/integration/test_sampler_checkpoints_e2e.py b/tests/server/integration/test_sampler_checkpoints_e2e.py index e95858118..4affee9a3 100644 --- a/tests/server/integration/test_sampler_checkpoints_e2e.py +++ b/tests/server/integration/test_sampler_checkpoints_e2e.py @@ -8,7 +8,7 @@ import numpy as np import pytest -from tests.server.integration.test_tinker_rl_correctness_e2e import _completion_logprobs, _datum +from tests.server.integration.e2e_helpers import completion_logprobs, sampled_datum pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_GPU_E2E') != '1', reason='Requires a real trainer and vLLM sampler') @@ -58,15 +58,15 @@ def same(actual, expected): sampler_x = service.create_sampling_client(base_model=model_id, model_path=x) first = sample(sampler_x) same(sample(sampler_x), first) - datum = _datum(prompt, first, 1.0) - before = _completion_logprobs(client, datum, timeout) + datum = sampled_datum(prompt, first, 1.0) + before = completion_logprobs(client, datum, timeout) losses = [] for _ in range(3): result = client.forward_backward([datum], 'cross_entropy').result(timeout=timeout) losses.append(float(result.metrics['loss:mean'])) assert np.isfinite(losses[-1]) client.optim_step(types.AdamParams(learning_rate=1e-4)).result(timeout=timeout) - after = _completion_logprobs(client, datum, timeout) + after = completion_logprobs(client, datum, timeout) assert float(np.mean(after - before)) > 1e-4 y = client.save_weights_for_sampler(name='y').result(timeout=timeout).path diff --git a/tests/server/integration/test_task_http_e2e.py b/tests/server/integration/test_task_http_e2e.py new file mode 100644 index 000000000..0b557b892 --- /dev/null +++ b/tests/server/integration/test_task_http_e2e.py @@ -0,0 +1,38 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Opt-in task HTTP acceptance. Run with --noconftest on reserved devices. + +Set TWINKLE_TEST_TASK_CS=1 and TWINKLE_TEST_TASK_MODEL_DIR to the prepared bases. +TWINKLE_TEST_BACKEND selects transformers or megatron; run once for each backend. +""" +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest + +pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_TASK_CS') != '1', + reason='Set TWINKLE_TEST_TASK_CS=1 on reserved devices') + + +def test_task_cli_isolation_checkpoint_and_release(tmp_path): + import ray + assert not ray.is_initialized(), 'Use --noconftest; this test owns an isolated Ray cluster' + backend = os.environ.get('TWINKLE_TEST_BACKEND', 'transformers') + directory = Path(os.environ.get('TWINKLE_TEST_ARTIFACT_DIR', str(tmp_path))) + directory.mkdir(parents=True, exist_ok=True) + script = Path(__file__).with_name('task_http_validation.py') + root = Path(__file__).resolve().parents[3] + command = [sys.executable, '-u', str(script), '--backend', backend, + '--task', 'embedding', 'reranker', 'seq_cls', + '--model-dir', os.environ['TWINKLE_TEST_TASK_MODEL_DIR'], '--run-dir', str(directory), + '--steps', '31', '--min-loss-drop', '0.1'] + env = {**os.environ, 'PYTHONPATH': str(root / 'src') + os.pathsep + str(root)} + with (directory / 'server.log').open('w') as output: + completed = subprocess.run(command, cwd=root, env=env, stdout=output, stderr=subprocess.STDOUT, timeout=5400) + assert completed.returncode == 0, (directory / 'server.log').read_text()[-6000:] + for task in ['embedding', 'reranker', 'seq_cls']: + result = json.loads((directory / (backend + '-' + task) / 'results.json').read_text()) + assert result['passed'] and result['slot_reuse'] + assert result['loss_drop'] >= 0.1 diff --git a/tests/server/integration/test_tinker_rl_correctness_e2e.py b/tests/server/integration/test_tinker_rl_correctness_e2e.py index 925841dbe..5f3fc5a66 100644 --- a/tests/server/integration/test_tinker_rl_correctness_e2e.py +++ b/tests/server/integration/test_tinker_rl_correctness_e2e.py @@ -14,6 +14,8 @@ import numpy as np import pytest +from tests.server.integration.e2e_helpers import completion_logprobs, sampled_datum + pytestmark = pytest.mark.skipif( os.environ.get('TWINKLE_TEST_GPU_E2E', '0') != '1', reason='Requires a running training + vLLM server on real accelerators', @@ -83,50 +85,20 @@ def live_clients(): return service, model_id, prompt_tokens, samples[0], timeout -def _datum(prompt_tokens, sample, advantage): - from tinker import types - - tokens = prompt_tokens + list(sample.tokens) - prompt_positions = len(prompt_tokens) - 1 - size = len(tokens) - 1 - - def tensor(values, dtype): - return types.TensorData(data=values, dtype=dtype, shape=[size]) - - return types.Datum( - model_input=types.ModelInput.from_ints(tokens[:-1]), - loss_fn_inputs={ - 'target_tokens': tensor(tokens[1:], 'int64'), - 'weights': tensor([0.0] * prompt_positions + [1.0] * len(sample.tokens), 'float32'), - 'logprobs': tensor([0.0] * prompt_positions + list(sample.logprobs), 'float32'), - 'advantages': tensor([0.0] * prompt_positions + [advantage] * len(sample.tokens), 'float32'), - }, - ) - - -def _completion_logprobs(client, datum, timeout): - result = client.forward([datum], 'cross_entropy').result(timeout=timeout) - values = np.asarray(result.loss_fn_outputs[0]['logprobs'].tolist(), dtype=np.float64) - weights = np.asarray(datum.loss_fn_inputs['weights'].tolist()) - assert len(values) >= len(weights) - values = values[:len(weights)][weights != 0] - assert values.size and np.isfinite(values).all() - return values - def test_live_unsupported_losses_are_rejected_without_changing_policy(live_clients): from tinker import BadRequestError service, model_id, prompt, sample, timeout = live_clients client = service.create_lora_training_client(base_model=model_id, rank=16, seed=2026) - datum = _datum(prompt, sample, -1.0) - before = _completion_logprobs(client, datum, timeout) + datum = sampled_datum(prompt, sample, -1.0) + before = completion_logprobs(client, datum, timeout) for loss_fn in ('ppo', 'cispo', 'dro'): with pytest.raises(BadRequestError) as exc: client.forward_backward([datum], loss_fn).result(timeout=timeout) assert exc.value.status_code == 400 assert f"Unsupported Tinker loss_fn '{loss_fn}'" in str(exc.value) - after = _completion_logprobs(client, datum, timeout) + after = completion_logprobs(client, datum, timeout) np.testing.assert_allclose(after, before, rtol=0, atol=1e-5) _report(check='unsupported_loss', loss_fn=loss_fn, http_status=400, policy_unchanged=True) @@ -138,15 +110,15 @@ def test_live_importance_sampling_updates_policy_in_reward_direction(live_client service, model_id, prompt, sample, timeout = live_clients client = service.create_lora_training_client(base_model=model_id, rank=16, seed=2026) - datum = _datum(prompt, sample, advantage) - before = _completion_logprobs(client, datum, timeout) + datum = sampled_datum(prompt, sample, advantage) + before = completion_logprobs(client, datum, timeout) sampler_difference = float(np.mean(np.abs(before - np.asarray(sample.logprobs)))) assert sampler_difference < 0.05, f'Sampler/trainer logprob mismatch: {sampler_difference}' result = client.forward_backward([datum] * batch_size, 'importance_sampling').result(timeout=timeout) assert len(result.loss_fn_outputs) == batch_size assert np.isfinite(float(result.metrics['loss:mean'])) client.optim_step(types.AdamParams(learning_rate=1e-4)).result(timeout=timeout) - after = _completion_logprobs(client, datum, timeout) + after = completion_logprobs(client, datum, timeout) delta = float(np.mean(after) - np.mean(before)) _report( check='reward_direction', @@ -166,7 +138,7 @@ def test_live_dpo_still_requires_pairs(live_clients): service, model_id, prompt, sample, timeout = live_clients client = service.create_lora_training_client(base_model=model_id, rank=16, seed=2026) - datum = _datum(prompt, sample, 1.0) + datum = sampled_datum(prompt, sample, 1.0) reference = client.forward([datum], 'cross_entropy').result(timeout=timeout) datum.loss_fn_inputs['ref_logps'] = reference.loss_fn_outputs[0]['logprobs'] with pytest.raises(UnprocessableEntityError) as exc: diff --git a/tests/server/lifecycle/test_envelope.py b/tests/server/lifecycle/test_envelope.py index fdc588bd8..1dd368a94 100644 --- a/tests/server/lifecycle/test_envelope.py +++ b/tests/server/lifecycle/test_envelope.py @@ -12,12 +12,18 @@ def test_wire_maps_cover_exactly_the_canonical_reason_codes(): """The Twinkle and Tinker failure maps must both stay in lockstep with the canonical reason-code set, so a newly added domain reason cannot silently fall back to 500 on one protocol face.""" - from twinkle.server.gateway.tinker_handlers import _TINKER_FAILURE_WIRE - from twinkle.server.lifecycle.envelope import _FAILURE_WIRE - from twinkle.server.state.models import FAILURE_REASON_CODES - - assert set(_FAILURE_WIRE) == FAILURE_REASON_CODES - assert set(_TINKER_FAILURE_WIRE) == FAILURE_REASON_CODES + from twinkle.server.gateway.tinker_handlers import _tinker_error_from_failure + from twinkle.server.lifecycle.envelope import error_payload_from_failure + from twinkle.server.task_errors import FailureResponse + + for reason in FailureResponse.reasons(): + policy = FailureResponse.for_reason(reason) + failure = {'reason_code': reason, 'message': 'failed', 'attribution': policy.category.value} + native = error_payload_from_failure(failure, request_id='request') + assert native.error_code == policy.error_code + assert native.category is policy.category + assert _tinker_error_from_failure(failure, request_id='request') == native.model_dump( + mode='json', exclude_none=True) def test_completed_with_none_result_is_a_success_not_a_failure(): diff --git a/tests/server/lifecycle/test_preflight_rejection.py b/tests/server/lifecycle/test_preflight_rejection.py index 74639c137..eb997cd8a 100644 --- a/tests/server/lifecycle/test_preflight_rejection.py +++ b/tests/server/lifecycle/test_preflight_rejection.py @@ -14,7 +14,7 @@ from twinkle.server.deployment import twinkle_server_error_handler from twinkle.server.exceptions import (BatchSizeError, InputTokensExceededError, RateLimitExceededError, RequestRejectedError, TwinkleServerError) -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin diff --git a/tests/server/lifecycle/test_retrieve_endpoint.py b/tests/server/lifecycle/test_retrieve_endpoint.py index 84b8767e4..704ea2f71 100644 --- a/tests/server/lifecycle/test_retrieve_endpoint.py +++ b/tests/server/lifecycle/test_retrieve_endpoint.py @@ -19,6 +19,9 @@ class _State: """A fake ServerState whose get_future returns a fixed record (or None).""" + async def assert_future_owner(self, request_id, token): + return None + def __init__(self, record): self._record = record diff --git a/tests/server/lifecycle/test_run_submit_dedup.py b/tests/server/lifecycle/test_run_submit_dedup.py index f9e91d76c..b5f7fe710 100644 --- a/tests/server/lifecycle/test_run_submit_dedup.py +++ b/tests/server/lifecycle/test_run_submit_dedup.py @@ -20,25 +20,28 @@ class _FakeState: def __init__(self, record_after_claim): self._record_after_claim = record_after_claim self.claimed = {} + self.claim_ttls = [] self.released = [] - async def claim_seq(self, dedup_key, request_id, ttl): - # Unseen -> claim it and let the caller proceed (returns None). - self.claimed[dedup_key] = request_id - return None + async def claim_training_seq(self, training_id, stream, seq_id, request_id, retention): + key = (training_id, stream, seq_id) + prior = self.claimed.get(key) + self.claimed[key] = prior or request_id + self.claim_ttls.append(retention) + return prior async def get_future(self, request_id): # Simulates whether a PENDING/QUEUED record was written (task enqueued). return self._record_after_claim - async def release_seq(self, dedup_key): - self.released.append(dedup_key) + async def release_training_seq(self, training_id, stream, seq_id, request_id): + self.released.append((training_id, stream, seq_id)) class _FakeManagement: def __init__(self, record_after_claim): self.state = _FakeState(record_after_claim) - self.task_queue_config = SimpleNamespace(effective_execution_timeout=60.0) + self.sequence_retention_s = 3600 # A real deployment declares its backend; preflight reads it from here. self.backend = 'transformers' self.data_world_size = 1 @@ -61,7 +64,7 @@ def _body(adapter_name: str = 'ad', seq_id: int = 7) -> ForwardBackwardTaskReque ``run_submit`` now reads field roles off the body to build the backend kwargs and to run preflight, so a ``SimpleNamespace`` would exercise a shape production never sees. """ - return ForwardBackwardTaskRequest(inputs=[{'input_ids': [1, 2]}], adapter_name=adapter_name, seq_id=seq_id) + return ForwardBackwardTaskRequest(inputs=[{'input_ids': [1, 2]}], adapter_name=adapter_name, seq_id=seq_id, training_id='training-' + adapter_name) async def _call(self, body, adapter_name, token): # pragma: no cover - never invoked @@ -75,7 +78,7 @@ async def test_release_kept_when_task_already_enqueued(): with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) assert mgmt.state.released == [], 'claim wrongly released for an already-enqueued task' - assert 'seq::sess-1::sess-1-ad::7' in mgmt.state.claimed + assert ('training-ad', 'native::sess-1::forward_backward', 7) in mgmt.state.claimed @pytest.mark.asyncio @@ -84,7 +87,16 @@ async def test_release_when_never_enqueued(): mgmt = _FakeManagement(record_after_claim=None) with pytest.raises(RuntimeError): await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) - assert mgmt.state.released == ['seq::sess-1::sess-1-ad::7'], 'claim should be released when nothing enqueued' + assert mgmt.state.released == [('training-ad', 'native::sess-1::forward_backward', 7)], 'claim should be released when nothing enqueued' + + +@pytest.mark.asyncio +async def test_dedup_uses_training_sequence_retention(): + """The claim must outlive queue wait plus execution, not just execution.""" + mgmt = _FakeManagement(record_after_claim={'status': 'queued'}) + with pytest.raises(RuntimeError): + await run_submit(mgmt, _request(), _body(), task_type='forward_backward', backend_call=_call) + assert mgmt.state.claim_ttls == [mgmt.sequence_retention_s] @pytest.mark.asyncio @@ -107,5 +119,5 @@ async def test_dedup_key_is_scoped_per_adapter(): backend_call=_call) claimed = set(mgmt.state.claimed) - assert claimed == {'seq::sess-1::sess-1-lora-A::1', 'seq::sess-1::sess-1-lora-B::1'}, ( + assert claimed == {('training-lora-A', 'native::sess-1::forward_backward', 1), ('training-lora-B', 'native::sess-1::forward_backward', 1)}, ( f'adapters collided on one dedup key: {claimed}') diff --git a/tests/server/lifecycle/test_static_guards.py b/tests/server/lifecycle/test_static_guards.py index 730c4ef7f..3e635a294 100644 --- a/tests/server/lifecycle/test_static_guards.py +++ b/tests/server/lifecycle/test_static_guards.py @@ -5,7 +5,6 @@ - ``TaskEnvelope`` has exactly one construction site. - The client-side invariant: the client HTTP timeout is <= 120 and strictly greater than the server Long_Poll_Window. -- The task status set has two independent declarations that must not drift. These exist because the spec states most of its guarantees in prose. A prose claim that nothing checks decays into a false claim -- as happened with "a consistency test asserts @@ -71,30 +70,6 @@ def test_task_envelope_has_exactly_one_construction_site(): assert sites, 'expected to find the construction sites inside envelope.py' -def test_server_task_status_enum_matches_client_literal(): - """The two independent declarations of the task status set must not drift. - - ``twinkle.protocol.types.lifecycle.TaskStatus`` (a Literal on the wire model) and the - server's ``TaskStatus`` enum are declared separately. ``envelope_from_record`` copies - ``record['status']`` straight into ``TaskEnvelope.status``, so a value the server can - write but the Literal does not list would fail pydantic validation *while serialising - the response* -- i.e. a 500 from retrieve for a task that actually finished. - - The client module's comment claimed such a test existed; it did not. This is it. - """ - from typing import get_args - - from twinkle.server.task_queue.types import TaskStatus as ServerTaskStatus - from twinkle.protocol.types.lifecycle import TERMINAL_STATUSES - from twinkle.protocol.types.lifecycle import TaskStatus as WireTaskStatus - - server_values = {member.value for member in ServerTaskStatus} - wire_values = set(get_args(WireTaskStatus)) - assert server_values == wire_values, (f'task status sets drifted: server-only={server_values - wire_values}, ' - f'wire-only={wire_values - server_values}') - assert TERMINAL_STATUSES <= wire_values, 'TERMINAL_STATUSES must be a subset of the declared statuses' - - def test_client_future_layer_is_not_imported_by_the_server(): """``_future.py`` carries an underscore because the dependency runs one way only. diff --git a/tests/server/lifecycle/test_submit_peek_e2e.py b/tests/server/lifecycle/test_submit_peek_e2e.py index c87aba599..b3de7aa6c 100644 --- a/tests/server/lifecycle/test_submit_peek_e2e.py +++ b/tests/server/lifecycle/test_submit_peek_e2e.py @@ -13,7 +13,7 @@ ray = pytest.importorskip('ray') from twinkle.server.state import ServerState # noqa: E402 -from twinkle.server.task_queue.config import TaskQueueConfig # noqa: E402 +from twinkle.server.config.task_queue import TaskQueueConfig # noqa: E402 from twinkle.server.task_queue.mixin import TaskQueueMixin # noqa: E402 from twinkle_client import _future # noqa: E402 from twinkle_client.exceptions import TaskFailedError # noqa: E402 diff --git a/tests/server/lifecycle/test_timing_bounds.py b/tests/server/lifecycle/test_timing_bounds.py index b2f5f5464..7664b5bb4 100644 --- a/tests/server/lifecycle/test_timing_bounds.py +++ b/tests/server/lifecycle/test_timing_bounds.py @@ -21,7 +21,7 @@ from twinkle.server.lifecycle.poll_config import long_poll_window, retrieve_poll_interval # noqa: E402 from twinkle.server.state import ServerState # noqa: E402 -from twinkle.server.task_queue.config import TaskQueueConfig # noqa: E402 +from twinkle.server.config.task_queue import TaskQueueConfig # noqa: E402 from twinkle.server.task_queue.mixin import TaskQueueMixin # noqa: E402 _WINDOW = 0.05 diff --git a/tests/server/lifecycle/test_tinker_retrieve_regression.py b/tests/server/lifecycle/test_tinker_retrieve_regression.py index 4af4611f7..504ec294c 100644 --- a/tests/server/lifecycle/test_tinker_retrieve_regression.py +++ b/tests/server/lifecycle/test_tinker_retrieve_regression.py @@ -15,6 +15,9 @@ class _State: + async def assert_future_owner(self, request_id, token): + return None + def __init__(self, record): self._record = record diff --git a/tests/server/model/test_full_mode_exclusivity.py b/tests/server/model/test_full_mode_exclusivity.py new file mode 100644 index 000000000..2c6ad3635 --- /dev/null +++ b/tests/server/model/test_full_mode_exclusivity.py @@ -0,0 +1,104 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Full-mode exclusivity must be one atomic state transition, not check-then-register.""" +from __future__ import annotations + +import asyncio + +import pytest +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock +from starlette.requests import Request + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.helpers import request, setup +from twinkle.server.config.training import TrainingConfig +from twinkle.server.model.pool.router import TrainingRouter + + +async def _setup_full_instance(): + """Full-mode slots belong to the controller, not the retired model create route.""" + controller, state, runtime, _ = await setup(backend=FakeBackend()) + raw = controller.config.model_dump() + raw['models'][1]['config']['max_replicas'] = 1 + controller.config = TrainingConfig.model_validate(raw) + runtime.start_gate.clear() + return controller, state, runtime + + +@pytest.mark.asyncio +async def test_concurrent_full_mode_claims_have_exactly_one_winner(): + controller, state, runtime = await _setup_full_instance() + try: + refs = await asyncio.gather(*(controller.create(request(name, mode='full')) + for name in ('tenant-a', 'tenant-b'))) + assert len(controller.instances) == 1 + assert len(next(iter(controller.instances.values())).training_ids) == 1 + runtime.start_gate.set() + await controller.wait_for_jobs() + records = [await state.get_future(ref['request_id']) for ref in refs] + assert sum(record['status'] == 'completed' for record in records) == 1 + assert sum(record['status'] == 'failed' and record['failure']['reason_code'] == 'capacity_unavailable' + for record in records) == 1 + assert await state._backend.get('token_count::owner') == 1 + finally: + runtime.start_gate.set() + await controller.wait_for_jobs() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_same_tenant_may_reclaim_its_own_resource(): + controller, _, runtime = await _setup_full_instance() + try: + refs = await asyncio.gather(*(controller.create(request('same', mode='full')) for _ in range(2))) + assert refs[0] == refs[1] + assert len(controller.instances) == 1 + assert len(next(iter(controller.instances.values())).training_ids) == 1 + finally: + runtime.start_gate.set() + await controller.wait_for_jobs() + await controller.shutdown() + + +def _request(body): + async def receive(): + return {'type': 'http.request', 'body': json.dumps(body).encode(), 'more_body': False} + incoming = Request({'type': 'http', 'method': 'POST', 'path': '/twinkle/add_adapter_to_model', 'headers': []}, receive) + incoming.state.session_id = 'session' + incoming.state.token = 'owner' + incoming.state.request_id = 'req' + return incoming + + +@pytest.mark.asyncio +async def test_concurrent_handler_creates_have_exactly_one_winner(): + """Drive the current catalog handler into real controller reservations.""" + controller, state, runtime = await _setup_full_instance() + router = TrainingRouter.__new__(TrainingRouter) + router.state = state + router.config = controller.config + router.model = SimpleNamespace(name='small-full') + router.controller = SimpleNamespace(create=SimpleNamespace(remote=AsyncMock(side_effect=controller.create))) + + async def create_remote(payload): + from twinkle.server.state.training_records import CreateTraining + return await controller.create(CreateTraining.model_validate(payload)) + router.controller.create.remote.side_effect = create_remote + try: + refs = await asyncio.gather(*(router.native_request(_request({ + 'client_create_id': name, 'adapter_name': name, + }), 'add_adapter_to_model') for name in ('adapter-a', 'adapter-b'))) + assert len(controller.instances) == 1 + assert len(next(iter(controller.instances.values())).training_ids) == 1 + runtime.start_gate.set() + await controller.wait_for_jobs() + records = [await state.get_future(ref.request_id) for ref in refs] + assert sum(record['status'] == 'completed' for record in records) == 1 + assert sum(record['status'] == 'failed' and record['failure']['reason_code'] == 'capacity_unavailable' + for record in records) == 1 + assert await state._backend.get('token_count::owner') == 1 + finally: + runtime.start_gate.set() + await controller.wait_for_jobs() + await controller.shutdown() diff --git a/tests/server/model/test_preflight_loss_inputs.py b/tests/server/model/test_preflight_loss_inputs.py new file mode 100644 index 000000000..6c76c9354 --- /dev/null +++ b/tests/server/model/test_preflight_loss_inputs.py @@ -0,0 +1,103 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Both backend entry points share one full-batch preflight implementation.""" +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from twinkle.loss.grpo import GRPOLoss +from twinkle.server.model.backends import common +from twinkle.server.model.backends.common import TwinkleCompatModelBase +from twinkle.server.model.backends.megatron_model import (TwinkleCompatFullMegatronModel, TwinkleCompatMegatronModel) +from twinkle.server.model.backends.transformers_model import (TwinkleCompatFullTransformersModel, + TwinkleCompatTransformersModel) + + +class _FakeBackend(TwinkleCompatModelBase): + """Just the state preflight touches; no Ray actor, no device.""" + + def __init__(self, loss=GRPOLoss, with_loss=True): + self.optimizer_group = { + 'adapter': SimpleNamespace(loss_instance=loss(epsilon=0.2) if with_loss and loss else None), + } + self.set_loss_calls = [] + + def _get_default_group(self): + return 'adapter' + + def set_loss(self, loss_cls, **kwargs): + kwargs.pop('adapter_name', None) + self.set_loss_calls.append((loss_cls, kwargs)) + self.optimizer_group['adapter'].loss_instance = GRPOLoss(**{ + key: value + for key, value in kwargs.items() if key in ('epsilon', 'beta') + }) + + def get_template(self, adapter_name): + return None + + +_BAD_SAMPLE = { + 'labels': [-100, -100, 0, 0], +} +_BAD_OLD_LOGPS = [[0.1, 0.2, 0.3]] + + +@pytest.mark.parametrize('backend', [ + TwinkleCompatTransformersModel, + TwinkleCompatFullTransformersModel, + TwinkleCompatMegatronModel, + TwinkleCompatFullMegatronModel, +]) +def test_every_backend_inherits_the_preflight_entry_point(backend): + """A backend without it silently skips the check and can still deadlock.""" + assert hasattr(backend, 'preflight_loss_inputs') + + +def test_preflight_rejects_a_ragged_batch_on_the_twinkle_route(): + backend = _FakeBackend() + with pytest.raises(AssertionError, match='mismatch'): + TwinkleCompatModelBase.preflight_loss_inputs( + backend, + inputs=[_BAD_SAMPLE], + adapter_name='adapter', + old_logps=_BAD_OLD_LOGPS, + ) + assert backend.set_loss_calls == [] + + +def test_preflight_accepts_a_response_only_row(): + backend = _FakeBackend() + TwinkleCompatModelBase.preflight_loss_inputs( + backend, + inputs=[_BAD_SAMPLE], + adapter_name='adapter', + old_logps=[[0.1, 0.2]], + ) + + +def test_preflight_tinker_route_sets_up_the_loss_before_validating(monkeypatch): + """The tinker payload hides its RL features in the datums, so they must be extracted.""" + monkeypatch.setattr(common, 'datum_to_input_feature', lambda inputs, template: [_BAD_SAMPLE]) + monkeypatch.setattr(common, 'extract_rl_features_for_loss', lambda inputs: {'old_logps': _BAD_OLD_LOGPS, 'ref_outputs': {'logps': []}}) + backend = _FakeBackend(with_loss=False) + + with pytest.raises(AssertionError, match='mismatch'): + TwinkleCompatModelBase.preflight_loss_inputs( + backend, + inputs=[SimpleNamespace(loss_fn_inputs={})], + adapter_name='adapter', + loss_fn='importance_sampling', + epsilon=0.3, + ) + + assert backend.set_loss_calls == [('GRPOLoss', { + 'epsilon': 0.3, + 'beta': 0.0 + })] + + +def test_preflight_is_a_no_op_without_a_loss_instance(): + backend = _FakeBackend(with_loss=False) + TwinkleCompatModelBase.preflight_loss_inputs(backend, inputs=[_BAD_SAMPLE], adapter_name='adapter', old_logps=_BAD_OLD_LOGPS) diff --git a/tests/server/model/test_replica_lifecycle.py b/tests/server/model/test_replica_lifecycle.py index 84b961391..63b9ccaa8 100644 --- a/tests/server/model/test_replica_lifecycle.py +++ b/tests/server/model/test_replica_lifecycle.py @@ -34,6 +34,7 @@ def _make_lifecycle_manager(state: _CapacityState, replica_id: str, max_loras: i manager.state = state manager.replica_id = replica_id manager.max_loras = max_loras + manager.train_mode = 'lora' manager._replica_registered = False manager.data_plane = SimpleNamespace(close=AsyncMock()) manager.shutdown_task_queue = AsyncMock() @@ -54,24 +55,26 @@ async def test_replica_lifecycle_updates_shared_capacity() -> None: assert await state.get_capacity_info() == {'max_loras': 6, 'used_loras': 0, 'free_loras': 6} await second.shutdown() + # Only the owner retires capacity after confirming the GPU workers stopped. + assert await state.get_capacity_info() == {'max_loras': 6, 'used_loras': 0, 'free_loras': 6} + await state.unregister_replica(second.replica_id) assert await state.get_capacity_info() == {'max_loras': 3, 'used_loras': 0, 'free_loras': 3} @pytest.mark.asyncio async def test_async_constructor_registers_replica_before_ready() -> None: state = SimpleNamespace(register_replica=AsyncMock(), touch_replica_last_seen=AsyncMock()) - replica_context = SimpleNamespace(replica_id=SimpleNamespace(unique_id='replica-1')) manager = ModelManagement.__new__(ModelManagement) manager._task_queue_config = SimpleNamespace(effective_execution_timeout=1800.0) with patch('twinkle.server.model.app.DeviceGroup', return_value=SimpleNamespace(name='group')), \ patch('twinkle.server.model.app.init_twinkle_runtime', return_value=None), \ - patch('twinkle.server.model.app.serve.get_replica_context', return_value=replica_context), \ patch.object(MODEL_SELECTOR, 'construct', return_value=MagicMock()), \ patch('twinkle.server.model.app.get_server_state', return_value=state), \ patch('twinkle.server.data_plane.DataPlaneProxy', return_value=MagicMock()), \ patch.object(ModelManagement, '_init_task_queue'), \ - patch.object(ModelManagement, '_init_adapter_manager'): + patch.object(ModelManagement, 'initialize_training_runtime', new_callable=AsyncMock), \ + patch.object(ModelManagement, 'wait_for_training_workers', new_callable=AsyncMock): await ModelManagement.__init__( manager, model_id='model', @@ -80,6 +83,10 @@ async def test_async_constructor_registers_replica_before_ready() -> None: device_mesh={}, backend='mock', max_loras=3, + training_context={'instance_id': 'replica-1', 'epoch': 'epoch', 'namespace': 'test', + 'registry_name': 'test', 'controller_name': 'test'}, + placement_group=None, + allocation_devices=[], ) state.register_replica.assert_awaited_once_with('replica-1', 3) diff --git a/tests/server/model/test_task_head.py b/tests/server/model/test_task_head.py new file mode 100644 index 000000000..681d6e52e --- /dev/null +++ b/tests/server/model/test_task_head.py @@ -0,0 +1,101 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""CS model head configuration contracts.""" +from pathlib import Path +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError + +from tests.server.integration.cs_tasks.task_models import _mcore_native_key, encoded_batch +from twinkle.server.config import ServerConfig +from twinkle.server.config.training import TrainingArguments +from twinkle.server.model.app import _validate_model_task + +TASK_CONFIGS = Path(__file__).resolve().parents[1] / 'integration' / 'cs_tasks' + + +@pytest.mark.parametrize('options', [ + {'task_type': 'seq_cls'}, {'num_labels': 3}, {'task_type': 'causal_lm', 'num_labels': 1}, + {'task_type': 'seq_cls', 'num_labels': 0}, {'task_type': 'seq_cls', 'num_labels': True}, +]) +def test_task_head_arguments_fail_fast(options): + with pytest.raises(ValidationError): + TrainingArguments(**options) + + +@pytest.mark.parametrize('backend', ['transformers', 'megatron']) +def test_task_recipe_yaml_preserves_head_contract(backend): + config = ServerConfig.from_yaml(TASK_CONFIGS / f'task_server_{backend}.yaml') + models = {model.name: model for model in config.training_config.models} + assert set(models) == {'cs-embedding', 'cs-reranker', 'cs-seq_cls'} + assert models['cs-reranker'].config.num_labels == 1 + assert models['cs-seq_cls'].config.num_labels == 3 + assert models['cs-embedding'].config.num_labels is None + head = 'score' if backend == 'transformers' else 'output_layer' + assert head in models['cs-reranker'].config.target_modules + assert models['cs-reranker'].capacity == 2 + if backend == 'megatron': + app = next(app for app in config.applications if app.name == 'reranker') + assert app.args.language_model_only is True + + +def test_cookbook_seq_cls_yaml_declares_heads(): + root = Path(__file__).resolve().parents[3] + config = ServerConfig.from_yaml(root / 'cookbook/client/server/transformer/server_config_seq_cls.yaml') + models = {model.name: model.config for model in config.training_config.models} + assert models['Qwen3.5-4B-reranker'].num_labels == 1 + assert models['Qwen3.5-4B-seq-cls'].num_labels == 3 + assert all('score' in model.target_modules for model in models.values()) + + +def test_mcore_native_weights_keep_the_language_model_prefix(): + assert _mcore_native_key('model.layers.0.self_attn.q_proj.weight') == \ + 'model.language_model.layers.0.self_attn.q_proj.weight' + assert _mcore_native_key('score.weight') == 'score.weight' + assert _mcore_native_key('lm_head.weight') == 'lm_head.weight' + + +@pytest.mark.parametrize('task,width,architecture,actual,success', [ + ('seq_cls', 3, 'Qwen2ForSequenceClassification', 3, True), + ('seq_cls', 1, 'Qwen2ForSequenceClassification', 3, False), + ('seq_cls', 3, 'Qwen2ForCausalLM', 3, False), + ('causal_lm', None, 'Qwen2ForSequenceClassification', 3, False), + ('causal_lm', None, 'Qwen2ForCausalLM', None, True), +]) +def test_persisted_head_matches_deployment(tmp_path, monkeypatch, task, width, architecture, actual, success): + import torch + from safetensors.torch import save_file + from transformers import AutoConfig + save_file({'score.weight': torch.zeros(actual or 1, 8)}, str(tmp_path / 'model.safetensors')) + from twinkle.hub import HubOperation + monkeypatch.setattr(HubOperation, 'download_model', lambda path: path) + monkeypatch.setattr(AutoConfig, 'from_pretrained', lambda *a, **kw: SimpleNamespace( + architectures=[architecture], num_labels=actual)) + kwargs = {'model_id': str(tmp_path), 'task_type': task, 'num_labels': width} + if success: + _validate_model_task(kwargs) + else: + with pytest.raises(ValueError): + _validate_model_task(kwargs) + + +def test_default_model_path_does_not_load_head_metadata(monkeypatch): + from twinkle.hub import HubOperation + + def unexpected(*args, **kwargs): + raise AssertionError('Default deployments must keep their existing loading path') + monkeypatch.setattr(HubOperation, 'download_model', unexpected) + _validate_model_task({'model_id': 'unused'}) + + +@pytest.mark.parametrize('task', ['embedding', 'reranker', 'seq_cls']) +def test_task_batches_match_the_wire_contract(task): + from pydantic import TypeAdapter + from twinkle.protocol.types.data import WireInputBatch + + def tokenizer(text, **kwargs): + return {'input_ids': [4, 5, 6], 'attention_mask': [1, 1, 1]} + batch = encoded_batch(task, tokenizer) + parsed = TypeAdapter(WireInputBatch).validate_python(batch) + assert len(parsed) == 12 + assert all(type(label) is int for row in parsed for label in row.labels) diff --git a/tests/server/model/test_tinker_handlers.py b/tests/server/model/test_tinker_handlers.py index 6daa6d566..2df46bb7e 100644 --- a/tests/server/model/test_tinker_handlers.py +++ b/tests/server/model/test_tinker_handlers.py @@ -9,8 +9,9 @@ from twinkle.server.model.tinker_handlers import _register_model_tinker_routes from twinkle.server.deployment import twinkle_server_error_handler from twinkle.server.exceptions import BatchSizeError, TwinkleServerError -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin +from twinkle.server.model.runtime import TrainingExecutionState class _DummyManagement: @@ -190,6 +191,7 @@ class _SaveWeightsDummyManagement: is_full_mode = ModelManagement.is_full_mode def __init__(self): + self.training_state = TrainingExecutionState(grad_ready=True) self.model = MagicMock() self.state = MagicMock() self.state.get_model_metadata = AsyncMock(return_value={'base_model': 'test-model'}) @@ -198,8 +200,8 @@ def __init__(self): async def _on_request_start(self, request): return 'token1' - def get_adapter_name(self, adapter_name=None): - return adapter_name + def get_training_state(self, training_id): + return self.training_state resolve_model_adapter_name = ModelManagement.resolve_model_adapter_name @@ -269,7 +271,6 @@ async def test_load_weights_resolves_location_and_preserves_optimizer_flag(monke output_dir = manager.get_save_dir('model1') monkeypatch.setattr(handlers, 'create_checkpoint_manager', lambda *args, **kwargs: manager) management = _SaveWeightsDummyManagement() - management.set_resource_state = MagicMock() app = FastAPI() _register_model_tinker_routes(app, lambda: management) endpoint = next(r.endpoint for r in app.routes if getattr(r, 'path', '') == '/tinker/load_weights') @@ -278,7 +279,7 @@ async def test_load_weights_resolves_location_and_preserves_optimizer_flag(monke management.model.tinker_load.assert_called_once_with( checkpoint_name='step' if output_dir else path, output_dir=output_dir, load_optimizer=True, adapter_name='model1') - management.set_resource_state.assert_called_once_with('model1', 'grad_ready', False) + assert management.training_state.grad_ready is False assert response.path == path and response.type == 'load_weights' @@ -289,7 +290,6 @@ async def test_tinker_loss_metric_survives_sdk_reduction(backward): management = _SaveWeightsDummyManagement() management.data_world_size = 1 - management.set_resource_state = MagicMock() outputs = [{'logprobs': types.TensorData(data=[-0.25, -0.5], dtype='float32', shape=[2])}] management.model.tinker_forward_only.return_value = (outputs, 0.375) management.model.tinker_forward_backward.return_value = (outputs, 0.375) diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py index 9f2ee1e64..aa8f0b2c1 100644 --- a/tests/server/model/test_twinkle_async_inputs.py +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -101,7 +101,7 @@ async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() - kwarg_fields={ 'old_logps': 'sampled_logprobs', 'advantages': 'advantage', - }, + }, training_id='training-a', ) await route.endpoint(request, body, management) @@ -109,7 +109,7 @@ async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() - assert management.scheduled[-1]['batch_size'] == 8 assert management.scheduled[-1]['data_world_size'] == 2 inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] - assert adapter_name == 'session-adapter' + assert adapter_name == 'training-a' assert [row['input_ids'] for row in inputs] == [[index] for index in range(8)] assert forwarded_kwargs['old_logps'] == [[-0.1]] * 4 + [[-0.2]] * 4 assert forwarded_kwargs['advantages'] == [1.0] * 4 + [-1.0] * 4 @@ -139,13 +139,13 @@ async def test_forward_backward_binds_nested_dpo_ref_logps_without_coercion() -> body = types.DataPlaneForwardRequest( adapter_name='adapter', input_refs=[types.DataRef(ref_id='dpo', size=2, num_tokens=6)], - kwarg_fields={'ref_outputs.logps': 'ref_logps'}, + kwarg_fields={'ref_outputs.logps': 'ref_logps'}, training_id='training-a', ) await route.endpoint(request, body, management) inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] - assert adapter_name == 'session-adapter' + assert adapter_name == 'training-a' assert [row['input_ids'] for row in inputs] == [[1, 2, 3], [1, 4, 5]] assert forwarded_kwargs['ref_outputs']['logps'] == [ [-0.1, -0.2, -0.3], @@ -177,7 +177,7 @@ def save(self, name, output_dir, **kwargs): monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) runs = TwinkleTrainingRunManager('token') - runs._write_info('session-adapter', {'base_model': 'test-model', 'is_lora': True}) + runs._write_info('training-a', {'base_model': 'test-model', 'is_lora': True}) manager = TwinkleCheckpointManager('token', runs) monkeypatch.setattr(handlers, 'create_checkpoint_manager', lambda *args, **kwargs: manager) service = Management() @@ -189,14 +189,14 @@ def save(self, name, output_dir, **kwargs): request.state.session_id = 'session' body = types.SaveRequest(adapter_name='adapter', name='step/one', is_sampler=is_sampler, - save_optimizer=save_optimizer, consumed_train_samples=7) + save_optimizer=save_optimizer, consumed_train_samples=7, training_id='training-a') response = await endpoint(request, body, service) weights_type = 'sampler_weights' if is_sampler else 'weights' - directory = manager.get_ckpt_dir('session-adapter', f'{weights_type}/step_one') - assert response.result == {'twinkle_path': f'twinkle://session-adapter/{weights_type}/step_one', + directory = manager.get_ckpt_dir('training-a', f'{weights_type}/step_one') + assert response.result == {'twinkle_path': f'twinkle://training-a/{weights_type}/step_one', 'checkpoint_dir': str(directory)} assert (directory / 'weights').read_text() == '1' - assert service.model_calls == [('step_one', {'adapter_name': '' if full_mode else 'session-adapter', + assert service.model_calls == [('step_one', {'adapter_name': '' if full_mode else 'training-a', 'save_optimizer': save_optimizer, 'consumed_train_samples': 7})] @@ -211,9 +211,9 @@ def save(self, **kwargs): monkeypatch.setattr('twinkle.server.checkpoint.training_run_manager.TWINKLE_DEFAULT_SAVE_DIR', str(tmp_path)) runs = TwinkleTrainingRunManager('token') - runs._write_info('session-adapter', {'base_model': 'test-model'}) + runs._write_info('training-a', {'base_model': 'test-model'}) manager = TwinkleCheckpointManager('token', runs) - before = runs._read_info('session-adapter') + before = runs._read_info('training-a') monkeypatch.setattr(handlers, 'create_checkpoint_manager', lambda *args, **kwargs: manager) service = Management() app = FastAPI() @@ -221,12 +221,12 @@ def save(self, **kwargs): endpoint = next(r.endpoint for r in app.routes if getattr(r, 'path', '') == '/twinkle/save') request = Request({'type': 'http', 'headers': []}) request.state.session_id = 'session' - body = types.SaveRequest(adapter_name='adapter', name='failed') + body = types.SaveRequest(adapter_name='adapter', name='failed', training_id='training-a') with pytest.raises(RuntimeError, match='backend save failed'): await endpoint(request, body, service) - assert manager.list_checkpoints('session-adapter').checkpoints == [] - assert runs._read_info('session-adapter') == before + assert manager.list_checkpoints('training-a').checkpoints == [] + assert runs._read_info('training-a') == before @pytest.mark.asyncio @@ -253,22 +253,94 @@ async def test_checkpoint_consumers_keep_backend_path_and_options(monkeypatch, t request.state.session_id = 'session' if operation == 'load': - body = types.LoadRequest(adapter_name='adapter', name=path, load_optimizer=True, strict=False) + body = types.LoadRequest(adapter_name='adapter', name=path, load_optimizer=True, strict=False, training_id='training-a') elif operation == 'resume_from_checkpoint': - body = types.ResumeFromCheckpointRequest(adapter_name='adapter', name=path, resume_only_model=True) + body = types.ResumeFromCheckpointRequest(adapter_name='adapter', name=path, resume_only_model=True, training_id='training-a') service.model.resume_from_checkpoint.return_value = {'step': 7} else: body = types.UploadToHubRequest(checkpoint_dir=path, hub_model_id='org/output', hub_token='hub-token') + if not local: + # A hub model ID is accepted by load/resume; upload requires an owned local checkpoint. + with pytest.raises(ValueError, match='Path escapes the permitted directory'): + await endpoint(request, body, service) + service.model.upload_to_hub.assert_not_called() + return response = await endpoint(request, body, service) if operation == 'load': service.model.load.assert_called_once_with( name='step' if local else path, output_dir=manager.get_save_dir('run') if local else None, - load_optimizer=True, adapter_name='session-adapter', token='token', strict=False) + load_optimizer=True, adapter_name='training-a', token='token', strict=False) elif operation == 'resume_from_checkpoint': service.model.resume_from_checkpoint.assert_called_once_with( - directory, resume_only_model=True, adapter_name='session-adapter') + directory, resume_only_model=True, adapter_name='training-a') assert response.result == {'result': {'step': 7}} else: service.model.upload_to_hub.assert_called_once_with( checkpoint_dir=directory, hub_model_id='org/output', hub_token='hub-token', async_upload=False) + + +class _PreflightManagement(_SchedulingManagement): + + def __init__(self): + super().__init__() + self.preflight_calls = [] + + def preflight_loss_inputs(self, *, inputs, adapter_name, **kwargs): + self.preflight_calls.append((inputs, adapter_name, kwargs)) + self.model_calls.append(('preflight', adapter_name)) + + +@pytest.mark.asyncio +async def test_forward_backward_preflights_loss_inputs_before_backend() -> None: + """The full batch is validated before the backend call can slice it.""" + management = _PreflightManagement() + app = FastAPI() + _register_model_twinkle_routes(app, lambda: management) + route = next(route for route in app.routes if getattr(route, 'path', None) == '/twinkle/forward_backward') + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.ForwardBackwardTaskRequest( + inputs=[{ + 'input_ids': [index, index + 1], + 'labels': [-100, 5], + } for index in range(3)], + adapter_name='adapter', + loss_kwargs={'old_logps': [[-0.1], [-0.2], [-0.3]]}, training_id='training-a', + ) + + await route.endpoint(request, body, management) + + assert len(management.preflight_calls) == 1 + preflight_inputs, preflight_adapter, preflight_kwargs = management.preflight_calls[0] + # Slice-free: a per-rank slice would carry at most one of the three samples. + assert len(preflight_inputs) == 3 + assert preflight_adapter == 'training-a' + assert preflight_kwargs['old_logps'] == [[-0.1], [-0.2], [-0.3]] + assert management.model_calls[0][0] == 'preflight' + assert management.model_calls[-1][0] != 'preflight' + assert management.model_calls[-1][1] == 'training-a' + + +@pytest.mark.asyncio +async def test_forward_backward_skips_preflight_without_rl_inputs() -> None: + """Plain SFT batches carry no ragged targets, so the extra actor round trip is skipped.""" + management = _PreflightManagement() + app = FastAPI() + _register_model_twinkle_routes(app, lambda: management) + route = next(route for route in app.routes if getattr(route, 'path', None) == '/twinkle/forward_backward') + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.ForwardBackwardTaskRequest( + inputs=[{ + 'input_ids': [1, 2], + 'labels': [-100, 5], + }], + adapter_name='adapter', training_id='training-a', + ) + + await route.endpoint(request, body, management) + + assert management.preflight_calls == [] + assert len(management.model_calls) == 1 + assert management.model_calls[0][1] == 'training-a' diff --git a/tests/server/resources/__init__.py b/tests/server/resources/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/resources/test_allocation.py b/tests/server/resources/test_allocation.py new file mode 100644 index 000000000..9e89ef403 --- /dev/null +++ b/tests/server/resources/test_allocation.py @@ -0,0 +1,120 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Resource RPC deadlines must preserve ownership of late placement groups.""" +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +from twinkle.server.resources.allocation import DevicePool + + +def test_stale_resource_summary_does_not_reject_a_schedulable_group(monkeypatch): + import importlib + import ray + placement = importlib.import_module('ray.util.placement_group') + monkeypatch.setattr(ray, 'nodes', lambda: [{'Alive': True, 'Resources': {'CPU': 4, 'GPU': 2}}]) + monkeypatch.setattr(ray, 'available_resources', lambda: {'CPU': 4}) + created = [] + group = object() + + def create(bundles, **kwargs): + created.append((bundles, kwargs)) + return group + + monkeypatch.setattr(placement, 'placement_group', create) + assert DevicePool._create_group('GPU', 2, 'replacement') is group + assert created[0][0][0]['GPU'] == 2 + assert created[0][1] == {'strategy': 'STRICT_PACK', 'name': 'replacement'} + + +@pytest.mark.asyncio +@pytest.mark.parametrize('blocked', ['nodes', 'placement_group']) +async def test_cancelled_allocation_keeps_late_creation_owned(monkeypatch, blocked): + import importlib + import ray + placement = importlib.import_module('ray.util.placement_group') + entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + created, removed = [], [] + + async def ready(): + return None + + pg = SimpleNamespace(id=SimpleNamespace(hex=lambda: 'pg-id'), ready=ready) + + def call(name, result): + if name == blocked: + entered.set() + try: + if not resume.wait(2): + raise TimeoutError('blocked allocation RPC') + finally: + returned.set() + return result + + def create(*args, **kwargs): + result = call('placement_group', pg) + created.append(result) + return result + + def get(name): + if not created: + raise ValueError('not created yet') + return pg + + monkeypatch.setattr(ray, 'nodes', lambda: call('nodes', [{'Alive': True, 'Resources': {'CPU': 2}}])) + monkeypatch.setattr(placement, 'placement_group', create) + monkeypatch.setattr(placement, 'get_placement_group', get) + monkeypatch.setattr(placement, 'remove_placement_group', lambda group: removed.append(group)) + monkeypatch.setattr(placement, 'placement_group_table', lambda group: {'state': 'REMOVED'}) + pool = DevicePool('test', allocation_timeout=.02, shutdown_timeout=.02) + operation = asyncio.create_task(pool.allocate(1, 'cpu', 'epoch', 'instance')) + try: + assert await asyncio.to_thread(entered.wait, 3) + with pytest.raises(TimeoutError): + await operation + assert not returned.is_set(), 'Allocation RPC blocked its deadline' + with pytest.raises(TimeoutError): + await asyncio.wait_for(pool.remove(pool.name('epoch', 'instance')), .02) + assert not created, 'The blocked RPC should still be pending' + finally: + resume.set() + await asyncio.gather(operation, return_exceptions=True) + await asyncio.to_thread(returned.wait, 3) + await asyncio.wait_for(pool.remove(pool.name('epoch', 'instance')), 1) + assert created and removed and all(group is pg for group in removed) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('blocked', ['get_placement_group', 'remove_placement_group', 'placement_group_table']) +async def test_removal_rpc_cannot_block_quarantine_deadline(monkeypatch, blocked): + import importlib + placement = importlib.import_module('ray.util.placement_group') + entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + pg = object() + + def call(name, result): + if name == blocked: + entered.set() + try: + if not resume.wait(2): + raise TimeoutError('blocked placement group RPC') + finally: + returned.set() + return result + + monkeypatch.setattr(placement, 'get_placement_group', lambda name: call('get_placement_group', pg)) + monkeypatch.setattr(placement, 'remove_placement_group', lambda group: call('remove_placement_group', None)) + monkeypatch.setattr( + placement, 'placement_group_table', lambda group: call('placement_group_table', {'state': 'REMOVED'})) + pool = DevicePool('test', allocation_timeout=1, shutdown_timeout=.02) + operation = asyncio.create_task(pool.remove('owned')) + try: + assert await asyncio.to_thread(entered.wait, 3) + with pytest.raises(TimeoutError): + await operation + assert not returned.is_set(), 'Placement group RPC blocked the close deadline' + finally: + resume.set() + await asyncio.gather(operation, return_exceptions=True) + await asyncio.to_thread(returned.wait, 3) diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py new file mode 100644 index 000000000..fb5c01b5f --- /dev/null +++ b/tests/server/resources/test_instances.py @@ -0,0 +1,378 @@ +import asyncio +from types import SimpleNamespace +import logging +import pytest +from twinkle.server.resources.instances import InstanceLifecycle + + +@pytest.mark.asyncio +async def test_spawn_logs_exception_with_traceback(caplog, monkeypatch): + # Capture one path regardless of the parent logger's runtime configuration. + target = logging.getLogger('twinkle.server.resources.instances') + monkeypatch.setattr(target, 'handlers', [caplog.handler]) + monkeypatch.setattr(target, 'propagate', False) + lifecycle = InstanceLifecycle('Training') + + async def broken(): + raise RuntimeError('background sentinel') + + with caplog.at_level(logging.ERROR, logger=target.name): + lifecycle.spawn(broken()) + await lifecycle.join() + records = [record for record in caplog.records if record.getMessage() == 'Training background job failed'] + assert records and records[0].exc_info[1].args == ('background sentinel',) + assert records[0].exc_info[2] is not None + assert not lifecycle.jobs + + +@pytest.mark.asyncio +async def test_concurrent_close_shares_one_stop_and_retries_quarantine(caplog, monkeypatch): + import asyncio + target = logging.getLogger('twinkle.server.resources.instances') + monkeypatch.setattr(target, 'handlers', [caplog.handler]) + monkeypatch.setattr(target, 'propagate', False) + caplog.set_level(logging.WARNING, logger=target.name) + lifecycle = InstanceLifecycle('test') + entered, resume = asyncio.Event(), asyncio.Event() + calls, events = [], [] + + async def stop(): + calls.append('stop') + entered.set() + await resume.wait() + if len(calls) == 1: + raise RuntimeError('death unconfirmed') + + async def stopped(): + events.append('stopped') + + async def quarantined(error): + events.append(str(error)) + + first = asyncio.create_task(lifecycle.close('one', stop, 1, stopped, quarantined)) + await entered.wait() + second = asyncio.create_task(lifecycle.close('one', stop, 1, stopped, quarantined)) + await asyncio.sleep(0) + resume.set() + await asyncio.gather(first, second) + assert calls == ['stop'] and events == ['death unconfirmed'] + records = [record for record in caplog.records if record.getMessage() == 'test instance one shutdown is unconfirmed'] + assert len(records) == 1 and records[0].exc_info[1].args == ('death unconfirmed',) + assert records[0].exc_info[2] is not None + await lifecycle.close('one', stop, 1, stopped, quarantined) + assert calls == ['stop', 'stop'] and events[-1] == 'stopped' + await lifecycle.join() + assert not lifecycle.jobs and not lifecycle.closing + +@pytest.mark.asyncio +async def test_shared_health_rule_tolerates_transient_probe_timeout(): + from twinkle.server.resources.instances import InstanceLifecycle + clock = SimpleNamespace(now=0) + lifecycle = InstanceLifecycle('health', clock=lambda: clock.now) + + async def unreachable(): + raise TimeoutError('control-plane probe timed out') + + async def healthy(): + return {'unhealthy': False} + + async def failed(): + return {'unhealthy': True} + + assert await lifecycle.health('instance', unreachable) is None + clock.now = 29 + assert await lifecycle.health('instance', unreachable) is None + clock.now = 30 + assert await lifecycle.health('instance', unreachable) is False + assert await lifecycle.health('instance', healthy) is True + assert not lifecycle.unreachable_since + assert await lifecycle.health('instance', unreachable) is None + assert await lifecycle.health('instance', failed) is False + + +def _patch_actor_cleanup_rpc(monkeypatch, rows): + from unittest.mock import AsyncMock, Mock + import ray + from ray import serve + from ray._private import gcs_utils, worker + from ray.core.generated import gcs_service_pb2, gcs_service_pb2_grpc + from ray.core.generated.gcs_pb2 import ActorTableData + + events = [] + + async def lookup(request, **kwargs): + events.append(('lookup', request.name, request.ray_namespace)) + return gcs_service_pb2.GetNamedActorInfoReply(actor_table_data=ActorTableData( + actor_id=request.name.encode(), name=request.name, ray_namespace=request.ray_namespace)) + + async def kill(request, **kwargs): + events.append(('kill', request.actor_id, request.force_kill, request.no_restart)) + return gcs_service_pb2.KillActorViaGcsReply() + + async def state(request, **kwargs): + events.append(('state', request.actor_id)) + return gcs_service_pb2.GetActorInfoReply( + actor_table_data=ActorTableData(actor_id=request.actor_id, state=ActorTableData.DEAD)) + + rpc = SimpleNamespace(GetNamedActorInfo=AsyncMock(side_effect=lookup), + KillActorViaGcs=AsyncMock(side_effect=kill), GetActorInfo=AsyncMock(side_effect=state)) + channel = AsyncMock() + monkeypatch.setattr(gcs_utils, 'create_gcs_channel', lambda *args, **kwargs: channel) + monkeypatch.setattr(gcs_service_pb2_grpc, 'ActorInfoGcsServiceStub', lambda unused: rpc) + monkeypatch.setattr(worker.global_worker, 'gcs_client', SimpleNamespace(address='test-gcs'), raising=False) + monkeypatch.setattr(ray.util, 'list_named_actors', lambda **kwargs: rows) + monkeypatch.setattr(ray, 'get_actor', Mock(side_effect=AssertionError('Worker handles must not be deserialized'))) + monkeypatch.setattr(ray, 'kill', Mock(side_effect=AssertionError('Cleanup must use actor IDs'))) + monkeypatch.setattr(serve, 'delete', lambda name: events.append(('delete', name))) + rpc.events, rpc.channel = events, channel + return rpc + + +@pytest.mark.asyncio +async def test_shutdown_actor_lookup_cannot_block_the_close_deadline(monkeypatch): + from unittest.mock import AsyncMock + from twinkle.server.resources.runtime import RayRuntime + + rpc = _patch_actor_cleanup_rpc(monkeypatch, [{'namespace': 'test', 'name': 'owned-worker'}]) + entered = asyncio.Event() + + async def lookup(*args, **kwargs): + entered.set() + await asyncio.Event().wait() + + rpc.GetNamedActorInfo.side_effect = lookup + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + runtime.pool.remove = AsyncMock() + stopped, quarantined = AsyncMock(), AsyncMock() + lifecycle = InstanceLifecycle('test') + await lifecycle.close('owned', lambda: runtime.shutdown('owned', 'app'), .02, stopped, quarantined) + assert entered.is_set() + quarantined.assert_awaited_once() + stopped.assert_not_awaited() + runtime.pool.remove.assert_not_awaited() + rpc.KillActorViaGcs.assert_not_awaited() + rpc.channel.__aexit__.assert_awaited_once() + await lifecycle.join() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('kind', ['sampler', 'training']) +async def test_serve_handle_lookup_cannot_block_driver_deadlines(monkeypatch, kind): + import threading + from types import SimpleNamespace + from ray import serve + from twinkle.server.resources.runtime import RayRuntime + from twinkle.server.sampler.pool.driver import SamplerDriver + from twinkle.server.model.pool.driver import RayInstanceDriver + + entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + + def lookup(name): + entered.set() + try: + if not resume.wait(2): + raise TimeoutError('blocked Serve controller RPC') + raise RuntimeError('lookup released after cancellation') + finally: + returned.set() + + monkeypatch.setattr(serve, 'get_app_handle', lookup) + driver_type = SamplerDriver if kind == 'sampler' else RayInstanceDriver + driver = driver_type.__new__(driver_type) + driver.runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + instance = SimpleNamespace(application='app', application_name='app') + operation = asyncio.create_task(driver.health(instance)) + try: + assert await asyncio.to_thread(entered.wait, 3) + with pytest.raises(TimeoutError): + await asyncio.wait_for(operation, .02) + assert not returned.is_set(), 'Serve handle lookup blocked the controller event loop' + finally: + resume.set() + await asyncio.gather(operation, return_exceptions=True) + await asyncio.to_thread(returned.wait, 3) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('kind', ['applications', 'placement_groups']) +async def test_orphan_inventory_cannot_block_recovery_deadlines(monkeypatch, kind): + import threading + from types import SimpleNamespace + from ray import serve + from twinkle.server.resources.runtime import RayRuntime + + # Import the module rather than ray.util's same-named public function. + import importlib + placement_module = importlib.import_module('ray.util.placement_group') + entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + + def inventory(): + entered.set() + try: + if not resume.wait(2): + raise TimeoutError('blocked resource inventory RPC') + return SimpleNamespace(applications={}) if kind == 'applications' else {} + finally: + returned.set() + + monkeypatch.setattr( + serve, 'status', inventory if kind == 'applications' else lambda: SimpleNamespace(applications={})) + monkeypatch.setattr( + placement_module, 'placement_group_table', inventory if kind == 'placement_groups' else lambda: {}) + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + operation = asyncio.create_task(runtime.recover_orphans(lambda name: None, None)) + try: + assert await asyncio.to_thread(entered.wait, 3) + with pytest.raises(TimeoutError): + await asyncio.wait_for(operation, .02) + assert not returned.is_set(), 'Resource inventory blocked the controller event loop' + finally: + resume.set() + await asyncio.gather(operation, return_exceptions=True) + await asyncio.to_thread(returned.wait, 3) + + +@pytest.mark.asyncio +async def test_shutdown_uses_actor_ids_and_keeps_namespace_and_instance_scope(monkeypatch): + from unittest.mock import AsyncMock + from twinkle.server.resources.runtime import RayRuntime + + rows = [{'namespace': 'test', 'name': 'owned-worker'}, + {'namespace': 'test', 'name': 'owned-registry'}, + {'namespace': 'other', 'name': 'owned-worker'}, + {'namespace': 'test', 'name': 'unrelated-worker'}] + rpc = _patch_actor_cleanup_rpc(monkeypatch, rows) + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + runtime.pool.remove = AsyncMock(side_effect=lambda name: rpc.events.append(('remove', name))) + await runtime.shutdown('owned', 'app') + assert rpc.events == [ + ('lookup', 'owned-worker', 'test'), ('kill', b'owned-worker', True, True), + ('lookup', 'owned-registry', 'test'), ('kill', b'owned-registry', True, True), + ('delete', 'app'), ('state', b'owned-worker'), ('state', b'owned-registry'), ('remove', 'app')] + rpc.channel.__aexit__.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('failure', ['namespace', 'name', 'missing_id', 'lookup_status', 'kill_status', 'state_status']) +async def test_shutdown_requires_trusted_identity_and_successful_gcs_replies(monkeypatch, failure): + from unittest.mock import AsyncMock + from ray.core.generated.gcs_pb2 import ActorTableData + from ray.core.generated.gcs_service_pb2 import GetNamedActorInfoReply, KillActorViaGcsReply, GetActorInfoReply + from twinkle.server.exceptions import ResourceReleaseError + from twinkle.server.resources.runtime import RayRuntime + + rpc = _patch_actor_cleanup_rpc(monkeypatch, [{'namespace': 'test', 'name': 'owned-worker'}]) + actor = ActorTableData(actor_id=b'owned-worker', name='owned-worker', ray_namespace='test') + if failure == 'namespace': + actor.ray_namespace = 'other' + elif failure == 'name': + actor.name = 'unrelated-worker' + elif failure == 'missing_id': + actor.actor_id = b'' + lookup = GetNamedActorInfoReply(actor_table_data=actor) + if failure == 'lookup_status': + lookup.status.code = 1 + rpc.GetNamedActorInfo.side_effect = None + rpc.GetNamedActorInfo.return_value = lookup + if failure == 'kill_status': + reply = KillActorViaGcsReply() + reply.status.code = 1 + rpc.KillActorViaGcs.side_effect = None + rpc.KillActorViaGcs.return_value = reply + if failure == 'state_status': + reply = GetActorInfoReply(actor_table_data=ActorTableData(actor_id=b'owned-worker', state=ActorTableData.DEAD)) + reply.status.code = 1 + rpc.GetActorInfo.side_effect = None + rpc.GetActorInfo.return_value = reply + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + runtime.pool.remove = AsyncMock() + with pytest.raises(ResourceReleaseError): + await runtime.shutdown('owned', 'app') + runtime.pool.remove.assert_not_awaited() + if failure not in ('kill_status', 'state_status'): + rpc.KillActorViaGcs.assert_not_awaited() + rpc.channel.__aexit__.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('blocked_rpc', ['KillActorViaGcs', 'GetActorInfo']) +async def test_shutdown_remote_actor_rpc_remains_within_close_deadline(monkeypatch, blocked_rpc): + from unittest.mock import AsyncMock + from twinkle.server.resources.runtime import RayRuntime + + rpc = _patch_actor_cleanup_rpc(monkeypatch, [{'namespace': 'test', 'name': 'owned-worker'}]) + entered = asyncio.Event() + + async def blocked(*args, **kwargs): + entered.set() + await asyncio.Event().wait() + + getattr(rpc, blocked_rpc).side_effect = blocked + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + runtime.pool.remove = AsyncMock() + stopped, quarantined = AsyncMock(), AsyncMock() + lifecycle = InstanceLifecycle('test') + await lifecycle.close('owned', lambda: runtime.shutdown('owned', 'app'), .02, stopped, quarantined) + assert entered.is_set() + quarantined.assert_awaited_once() + stopped.assert_not_awaited() + runtime.pool.remove.assert_not_awaited() + await lifecycle.join() + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_dead_state_not_only_kill_ack(monkeypatch): + from unittest.mock import AsyncMock + from ray.core.generated.gcs_pb2 import ActorTableData + from ray.core.generated.gcs_service_pb2 import GetActorInfoReply + from twinkle.server.resources.runtime import RayRuntime + + rpc = _patch_actor_cleanup_rpc(monkeypatch, [{'namespace': 'test', 'name': 'owned-worker'}]) + rpc.GetActorInfo.side_effect = [ + GetActorInfoReply(actor_table_data=ActorTableData(actor_id=b'owned-worker', state=state)) + for state in (ActorTableData.ALIVE, ActorTableData.DEAD)] + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + runtime.pool.remove = AsyncMock() + await runtime.shutdown('owned', 'app') + assert rpc.GetActorInfo.await_count == 2 + runtime.pool.remove.assert_awaited_once_with('app') + + +@pytest.mark.asyncio +@pytest.mark.parametrize('outcome', [ + 'unconfirmed_worker', 'unavailable_worker', 'dead_worker', 'removed', 'removal_failed' +]) +async def test_worker_death_precedes_group_release_and_removal_timing(monkeypatch, outcome): + from unittest.mock import AsyncMock, Mock + from ray.core.generated.gcs_service_pb2 import GetActorInfoReply + from twinkle.server.exceptions import ResourceReleaseError + from twinkle.server.resources.runtime import RayRuntime + + rows = ([{'namespace': 'test', 'name': 'owned-worker'}] + if outcome in ('unconfirmed_worker', 'unavailable_worker', 'dead_worker') else []) + rpc = _patch_actor_cleanup_rpc(monkeypatch, rows) + if outcome == 'unavailable_worker': + rpc.GetActorInfo.side_effect = TimeoutError('Control-plane connection lost') + elif outcome == 'unconfirmed_worker': + rpc.GetActorInfo.side_effect = None + rpc.GetActorInfo.return_value = GetActorInfoReply() # No identity or death witness. + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + runtime.pool.remove = AsyncMock(side_effect=RuntimeError('PG removal failed') + if outcome == 'removal_failed' else None) + observe = Mock() + if outcome in ('unconfirmed_worker', 'unavailable_worker'): + with pytest.raises(ResourceReleaseError) as failure: + await runtime.shutdown('owned', 'app', observe_removal=observe) + assert failure.value.phase == 'worker_death' + runtime.pool.remove.assert_not_awaited() + observe.assert_not_called() + else: + if outcome == 'removal_failed': + with pytest.raises(RuntimeError, match='PG removal failed'): + await runtime.shutdown('owned', 'app', observe_removal=observe) + else: + await runtime.shutdown('owned', 'app', observe_removal=observe) + runtime.pool.remove.assert_awaited_once_with('app') + assert observe.call_count == 1 + duration, status = observe.call_args.args + assert duration >= 0 and status == ('error' if outcome == 'removal_failed' else 'success') diff --git a/tests/server/sampler/test_pool.py b/tests/server/sampler/test_pool.py new file mode 100644 index 000000000..12592bef4 --- /dev/null +++ b/tests/server/sampler/test_pool.py @@ -0,0 +1,515 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Pool decisions against a controllable driver, without Ray or GPU side effects.""" +import asyncio +import json +from types import SimpleNamespace + +import pytest + +from twinkle.server.config.application_spec import SamplerArgs +from twinkle.server.sampler.pool.controller import SamplerController + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + yield + + +class Driver: + + def __init__(self): + self.started = [] + self.stopped = [] + self.calls = [] + self.finish = asyncio.Event() + self.start_gate = asyncio.Event() + self.start_gate.set() + self.fail_start = False + self.fail_close = False + + async def start(self, instance): + self.started.append(instance.identity) + await self.start_gate.wait() + if self.fail_start: + raise RuntimeError('allocation failed') + instance.devices = [str(len(self.started))] + + async def invoke(self, instance, ticket, payload, queue=None): + self.calls.append((instance.identity, payload)) + return {'status': 200, 'headers': [], 'body': b'{"status":"pending"}'} + + async def wait_operation(self, instance, ticket): + await self.finish.wait() + + async def load_weights(self, instance, path): + self.calls.append((instance.identity, {'base_path': path})) + + async def shutdown(self, instance): + if self.fail_close: + raise RuntimeError('death unconfirmed') + self.stopped.append(instance.identity) + + async def recover_orphans(self): + pass + + +def setup(minimum=2, maximum=4, target=1): + args = SamplerArgs( + model_id='mock', + sampler_type='mock', + device_group={ + 'name': 'sampler', + 'ranks': 1, + 'device_type': 'cpu' + }, + device_mesh={'dp_size': 1}, + min_replicas=minimum, + max_replicas=maximum, + target_ongoing_requests=target, + idle_timeout_s=10) + driver = Driver() + clock = SimpleNamespace(now=0) + pool = SamplerController( + args, driver, { + 'startup_timeout_s': 2, + 'shutdown_timeout_s': 1, + 'sweep_interval_s': 1 + }, clock=lambda: clock.now) + return pool, driver, clock + + +def request(endpoint='twinkle/sample', body=None): + return { + 'endpoint': endpoint, + 'body': json.dumps(body or {}).encode(), + 'method': 'POST', + 'headers': [(b'x-request-id', b'test')], + 'query_string': b'' + } + + +async def settle_starts(pool): + await asyncio.gather(*tuple(pool.lifecycle.startups.values()), return_exceptions=True) + + +async def settle_jobs(pool): + await asyncio.gather(*tuple(pool.lifecycle.jobs), return_exceptions=True) + + +@pytest.mark.asyncio +async def test_shutdown_rejects_unconfirmed_release_and_allows_retry(): + pool, driver, _ = setup() + await pool.initialize() + driver.fail_close = True + try: + with pytest.raises(RuntimeError, match='shutdown.*unconfirmed'): + await pool.shutdown() + assert pool._owned_count() == 2 and not driver.stopped + assert all(instance.state == 'QUARANTINED' for instance in pool.instances.values()) + finally: + driver.fail_close = False + await pool.shutdown() + assert pool._owned_count() == 0 and len(driver.stopped) == 2 + + +@pytest.mark.asyncio +async def test_minimum_expand_and_reclaim_only_idle_extras(): + pool, driver, clock = setup() + await pool.initialize() + assert len(driver.started) == 2 + try: + for _ in range(8): + await pool.request(request()) + await settle_starts(pool) + assert len(driver.started) == 4 and pool._owned_count() == 4 + live = pool._ready_instances() + assert all(instance.requests for instance in live) + assert max(len(instance.requests) for instance in live) - min(len(instance.requests) for instance in live) <= 1 + clock.now = 30 + await pool.sweep() + assert not driver.stopped + driver.finish.set() + await settle_jobs(pool) + clock.now += 11 + await pool.sweep() + await settle_jobs(pool) + assert len(driver.stopped) == 2 and len(pool._ready_instances()) == 2 + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_waiter_cancellation_preserves_generation_and_defers_control(): + pool, driver, _ = setup() + await pool.initialize() + await pool.request(request()) + control = asyncio.create_task(pool.request(request('twinkle/set_template', {'template_cls': 'Template'}))) + try: + await asyncio.sleep(0) + control.cancel() + await asyncio.gather(control, return_exceptions=True) + assert pool._guard._readers == 1 + assert sum(len(instance.requests) for instance in pool.instances.values()) == 1 + assert len(driver.calls) == 1 + driver.finish.set() + await settle_jobs(pool) + assert any(payload['endpoint'] == 'twinkle/set_template' for _, payload in driver.calls) + assert pool._actions[0]['endpoint'] == 'twinkle/set_template' + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_new_instance_replays_controls_not_lora_weight_cache(): + pool, driver, _ = setup(minimum=1, target=1) + driver.finish.set() + await pool.initialize() + try: + await pool.request(request('twinkle/set_template', {'template_cls': 'Template'})) + driver.finish.clear() + await pool.request(request(body={'adapter_uri': 'checkpoint-A'})) + await pool.request(request(body={'adapter_uri': 'checkpoint-B'})) + driver.finish.set() + await settle_starts(pool) + await settle_jobs(pool) + assert len(driver.started) == 2 + new = driver.started[-1] + replay = [payload for identity, payload in driver.calls if identity == new] + assert replay and all(payload['endpoint'] == 'twinkle/set_template' for payload in replay) + assert all('adapter_uri' not in json.loads(payload['body']) for payload in pool._actions) + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_quarantine_counts_against_maximum_until_confirmed_death(): + pool, driver, _ = setup(minimum=1, maximum=1) + await pool.initialize() + instance = pool._ready_instances()[0] + driver.fail_close = True + await pool._close(instance) + try: + assert instance.state == 'QUARANTINED' + for _ in range(3): + await pool.sweep() + await settle_jobs(pool) + assert len(driver.started) == 1 and pool._owned_count() == 1 + driver.fail_close = False + await pool.sweep() + await settle_jobs(pool) + await pool.sweep() + await settle_starts(pool) + assert len(driver.started) == 2 and len(pool._ready_instances()) == 1 + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_startup_failure_rolls_back_loaded_minimum(): + pool, driver, _ = setup() + original = driver.start + + async def fail_second(instance): + if len(driver.started) == 1: + driver.fail_start = True + await original(instance) + + driver.start = fail_second + with pytest.raises(RuntimeError, match='allocation'): + await pool.initialize() + assert not pool._ready and pool._owned_count() == 0 + + +@pytest.mark.asyncio +async def test_lazy_first_requests_share_one_startup(): + pool, driver, _ = setup(minimum=0, maximum=1) + await pool.initialize() + driver.start_gate.clear() + calls = [asyncio.create_task(pool.request(request())) for _ in range(8)] + try: + for _ in range(10): + await asyncio.sleep(0) + assert len(driver.started) == 1 + driver.start_gate.set() + await asyncio.wait_for(asyncio.gather(*calls), 2) + assert len(driver.started) == 1 + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault', ['raise', 'http_error']) +async def test_partial_control_broadcast_fences_all_participants(fault): + pool, driver, _ = setup() + driver.finish.set() + await pool.initialize() + invoke = driver.invoke + changed = [] + + async def update(instance, ticket, payload, queue=None): + changed.append(instance.identity) + if len(changed) == 2: + if fault == 'raise': + raise ConnectionError('update ACK unknown') + return {'status': 500, 'headers': [], 'body': b'error'} + return await invoke(instance, ticket, payload, queue) + + driver.invoke = update + try: + if fault == 'raise': + with pytest.raises(ConnectionError): + await pool.request(request('twinkle/set_template')) + else: + reply = await pool.request(request('twinkle/set_template')) + assert reply['status'] == 500 + await settle_jobs(pool) + assert len(driver.stopped) == 2 + assert not pool._ready_instances() and not pool._actions + finally: + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_unknown_generation_exit_keeps_quarantined_capacity_reserved(): + pool, driver, clock = setup(minimum=1, maximum=1) + await pool.initialize() + driver.fail_close = True + + async def unknown(*args): + raise TimeoutError('generation completion unknown') + + driver.wait_operation = unknown + await pool.request(request()) + await settle_jobs(pool) + instance = next(iter(pool.instances.values())) + assert instance.state == 'QUARANTINED' and instance.requests + clock.now = 100 + await pool.sweep() + await settle_jobs(pool) + assert len(driver.started) == 1 and instance.requests + driver.fail_close = False + await pool.sweep() + await settle_jobs(pool) + assert instance.state == 'STOPPED' and not instance.requests + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_replaying_config_has_the_same_startup_deadline(): + pool, driver, _ = setup(minimum=0, maximum=1) + await pool.initialize() + pool._actions.append(request('twinkle/set_template')) + pool.budgets = pool.budgets.model_copy(update={'startup_timeout_s': 0.05}) + instance = pool._new_instance() + with pytest.raises(TimeoutError): + await pool.lifecycle.startups[instance.identity] + assert instance.state == 'STOPPED' and driver.stopped == [instance.identity] + await pool.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('kind', ['control', 'weights']) +async def test_mutation_reserves_every_replica_against_idle_sweep(kind): + pool, driver, clock = setup(minimum=1, maximum=2) + driver.finish.set() + await pool.initialize() + pool._new_instance() + await settle_starts(pool) + clock.now = 20 + entered, resume = asyncio.Event(), asyncio.Event() + invoke, load = driver.invoke, driver.load_weights + + async def update(instance, *args): + entered.set() + await resume.wait() + assert instance.state == 'READY' + return await (invoke(instance, *args) if kind == 'control' else load(instance, *args)) + + if kind == 'control': + driver.invoke = update + task = asyncio.create_task(pool.request(request('twinkle/set_template'))) + else: + driver.load_weights = update + task = asyncio.create_task(pool.request(request(), full_path='B')) + try: + await asyncio.wait_for(entered.wait(), 1) + await pool.sweep() + # Only wait for the close jobs, never for the deliberately blocked mutation. + await asyncio.gather(*tuple(pool.lifecycle.closing.values())) + assert len(pool._ready_instances()) == 2 + assert not driver.stopped + finally: + resume.set() + await asyncio.gather(task, return_exceptions=True) + await pool.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('stage', ['invoke', 'wait_operation', 'load_weights']) +async def test_control_deadline_fences_all_participants(stage): + pool, driver, _ = setup() + driver.finish.set() + pool.args.queue_config.execution_timeout = 0.03 + entered, resume = asyncio.Event(), asyncio.Event() + + async def blocked(*args): + entered.set() + await resume.wait() + + setattr(driver, stage, blocked) + task = asyncio.create_task(pool.request( + request() if stage == 'load_weights' else request('twinkle/set_template'), + full_path='B' if stage == 'load_weights' else None)) + try: + await asyncio.wait_for(entered.wait(), 1) + with pytest.raises(TimeoutError): + await asyncio.wait_for(asyncio.shield(task), 0.3) + assert task.done(), 'Configured execution deadline was not enforced' + await settle_jobs(pool) + assert not pool._ready_instances() + assert len(driver.stopped) == 2 + assert not pool._actions + finally: + resume.set() + await asyncio.gather(task, return_exceptions=True) + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_superseded_checkpoint_is_not_needed_to_start_new_replica(): + pool, driver, _ = setup(minimum=1) + driver.finish.set() + valid = {'A', 'B'} + original = driver.load_weights + + async def load(instance, path): + if path not in valid: + raise FileNotFoundError(path) + await original(instance, path) + + driver.load_weights = load + try: + for path in ('A', 'B'): + await pool.request(request(), full_path=path) + await settle_jobs(pool) + valid.remove('A') + instance = pool._new_instance() + await pool.lifecycle.startups[instance.identity] + assert instance.state == 'READY' + assert [payload['base_path'] for identity, payload in driver.calls + if identity == instance.identity and 'base_path' in payload] == ['B'] + finally: + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_queued_control_expires_without_mutating_or_fencing_generation(): + pool, driver, _ = setup(minimum=1, maximum=1) + pool.args.queue_config.queue_timeout = 0.02 + await pool.request(request()) + try: + with pytest.raises(TimeoutError): + await asyncio.wait_for(pool.request(request('twinkle/set_template')), 1) + assert len(driver.calls) == 1 and not driver.stopped + driver.finish.set() + await settle_jobs(pool) + await pool.request(request('twinkle/set_template')) + assert len(driver.calls) == 2 + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_cancelled_control_waiter_keeps_replica_occupied_until_completion(): + pool, driver, clock = setup(minimum=1, maximum=2) + await pool.initialize() + pool._new_instance() + await settle_starts(pool) + entered = asyncio.Event() + invoke = driver.invoke + + async def update(*args): + entered.set() + return await invoke(*args) + + driver.invoke = update + task = asyncio.create_task(pool.request(request('twinkle/set_template'))) + try: + await entered.wait() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + clock.now = 30 + await pool.sweep() + assert len(pool._ready_instances()) == 2 and not driver.stopped + driver.finish.set() + await settle_jobs(pool) + clock.now += 11 + await pool.sweep() + await settle_jobs(pool) + assert len(pool._ready_instances()) == 1 + finally: + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_control_timeout_retains_unknown_allocations_then_recovers(): + pool, driver, _ = setup(minimum=2, maximum=2) + pool.args.queue_config.execution_timeout = 0.02 + driver.fail_close = True + try: + with pytest.raises(TimeoutError): + await pool.request(request('twinkle/set_template')) + await settle_jobs(pool) + assert all(i.state == 'QUARANTINED' and i.requests for i in pool.instances.values()) + await pool.sweep() + await settle_jobs(pool) + assert len(driver.started) == 2 + driver.fail_close = False + driver.finish.set() + await pool.sweep() + await settle_jobs(pool) + await pool.sweep() + await settle_starts(pool) + await pool.request(request()) + assert len(driver.started) == 4 + assert len(pool._ready_instances()) == 2 + finally: + driver.fail_close = False + driver.finish.set() + await pool.shutdown() + + +@pytest.mark.asyncio +async def test_replay_compacts_setters_per_owner_without_crossing_patch(): + pool, driver, _ = setup(minimum=1) + driver.finish.set() + try: + for index in range(12): + await pool.request(request('twinkle/set_template', {'template_cls': f'T{index}'})) + for owner in ('one', 'two'): + for rank in (4, 8): + payload = request('twinkle/add_adapter_to_sampler', {'adapter_name': 'same', 'config': {'r': rank}}) + payload['headers'].append((b'twinkle-authorization', owner.encode())) + await pool.request(payload) + assert len(pool._actions) == 3 + await pool.request(request('twinkle/apply_patch', {'patch_cls': 'P'})) + await pool.request(request('twinkle/set_template', {'template_cls': 'AfterPatch'})) + await pool.request(request('twinkle/unload_adapter_paths', {'adapter_paths': ['unused']})) + instance = pool._new_instance() + await pool.lifecycle.startups[instance.identity] + replay = [p for identity, p in driver.calls if identity == instance.identity] + assert [p['endpoint'] for p in replay] == [ + 'twinkle/set_template', 'twinkle/add_adapter_to_sampler', 'twinkle/add_adapter_to_sampler', + 'twinkle/apply_patch', 'twinkle/set_template'] + assert json.loads(replay[0]['body'])['template_cls'] == 'T11' + assert all(json.loads(p['body'])['config']['r'] == 8 for p in replay[1:3]) + assert json.loads(replay[-1]['body'])['template_cls'] == 'AfterPatch' + finally: + await pool.shutdown() diff --git a/tests/server/sampler/test_request_runtime.py b/tests/server/sampler/test_request_runtime.py new file mode 100644 index 000000000..26b619a97 --- /dev/null +++ b/tests/server/sampler/test_request_runtime.py @@ -0,0 +1,271 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sampler runtime initialization, completion ownership and cancellation.""" +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +from tests.server.task_queue.helpers import _DummyQueue +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.sampler.concurrency import SamplerConcurrency +from twinkle.server.task_queue.types import ExecutionScope + + +@pytest.mark.asyncio +async def test_encode_retains_execution_ownership_and_blocks_weight_mutation(): + import json + from unittest.mock import AsyncMock + from twinkle.server.sampler.runtime import SamplerRuntime + + entered, release, mutated = asyncio.Event(), asyncio.Event(), asyncio.Event() + owner = _DummyQueue() + owner._task_queue_config = TaskQueueConfig(inline_fast_path_timeout=0.01) + owner.enable_compute_worker() + owner.sampling_guard = SamplerConcurrency() + owner._on_request_start = AsyncMock(return_value='token') + + async def encode(*args, **kwargs): + entered.set() + await release.wait() + return [SimpleNamespace(data=[0.25, 0.75])] + + async def call_backend(fn, *args, **kwargs): + return await fn(*args, **kwargs) + + owner.sampler = SimpleNamespace(encode=encode) + owner.call_backend = call_backend + runtime = SamplerRuntime(owner) + mutation = completion = None + try: + reply = await runtime.invoke('encode', { + 'endpoint': 'twinkle/encode', 'method': 'POST', + 'headers': [(b'content-type', b'application/json'), + (b'authorization', b'Bearer encode-test'), (b'x-request-id', b'encode-test')], + 'body': json.dumps({'inputs': [{'input_ids': [1, 2]}]}).encode()}) + assert reply['status'] == 200, reply['body'] + await asyncio.wait_for(entered.wait(), 1) + scope = runtime._operations['encode'].scope + assert scope.request_ids and scope.jobs + completion = asyncio.create_task(runtime.wait_operation('encode')) + + async def update(): + mutated.set() + + mutation = asyncio.create_task(owner.sampling_guard.run(update, exclusive=True)) + await asyncio.sleep(0.02) + assert not completion.done() and not mutated.is_set() + release.set() + await asyncio.wait_for(asyncio.gather(completion, mutation), 1) + result = await owner.state.get_future(next(iter(scope.request_ids))) + assert result['status'] == 'completed' + assert result['result'] == {'data': [[0.25, 0.75]]} + assert mutated.is_set() + finally: + release.set() + await asyncio.gather(*(task for task in (completion, mutation) if task), return_exceptions=True) + await owner._compute_worker.stop() + + +@pytest.mark.asyncio +async def test_allocated_sampler_initializes_with_typed_owner_context(monkeypatch): + from unittest.mock import Mock + from twinkle.infra._ray.ray_helper import RayHelper + from twinkle.infra._ray.resource_manager import ResourceManager + from twinkle.server.sampler import app as module + from twinkle.server.state.models import RuntimeContext + + context = RuntimeContext(namespace='owned', registry_name='registry', instance_id='sampler-owned') + allocation = object() + manager = object() + allocate = Mock(return_value=manager) + sampler = SimpleNamespace() + monkeypatch.setenv('TWINKLE_RAY_NAMESPACE', 'before') + monkeypatch.setenv('TWINKLE_RAY_REGISTRY_NAME', 'before') + monkeypatch.setattr(RayHelper, 'resource_manager', None) + monkeypatch.setattr(ResourceManager, 'from_allocation', allocate) + monkeypatch.setattr(module, 'init_twinkle_runtime', lambda **kwargs: None) + monkeypatch.setattr(module, '_construct_sampler_backend', lambda *args: sampler) + monkeypatch.setattr(module.SamplerManagement, '_init_task_queue', lambda *args, **kwargs: None) + monkeypatch.setattr(module, 'get_server_state', lambda: SimpleNamespace()) + monkeypatch.setattr(module.serve, 'get_replica_context', lambda: SimpleNamespace( + replica_id=SimpleNamespace(unique_id='serve-replica'))) + owner = module.SamplerManagement.__new__(module.SamplerManagement) + await owner.__init__( + model_id='mock', nproc_per_node=1, device_group={'name': 'sampler', 'ranks': 1, 'device_type': 'cpu'}, + device_mesh={'dp_size': 1}, sampler_type='mock', sampling_context=context, + placement_group=allocation, allocation_devices=['0']) + assert owner.replica_id == context.instance_id + assert RayHelper.resource_manager is manager + allocate.assert_called_once_with(owner.device_group, allocation, ['0']) + await owner.data_plane.close() + + +@pytest.mark.asyncio +async def test_worker_exports_actual_completion_handle_without_private_task_lookup(): + queue = _DummyQueue() + queue._task_queue_config = TaskQueueConfig(execution_timeout=0.05) + queue.enable_compute_worker() + guard = SamplerConcurrency() + entered, exit_gate = asyncio.Event(), asyncio.Event() + scope = ExecutionScope() + + async def backend(): + entered.set() + await exit_gate.wait() + return 'late result' + + ref = await queue.schedule_concurrent_task(lambda: guard.run(backend), execution_scope=scope) + try: + await asyncio.wait_for(entered.wait(), 1) + completion = asyncio.create_task(scope.wait()) + for _ in range(100): + record = await queue.state.get_future(ref['request_id']) + if record['status'] == 'failed': + break + await asyncio.sleep(0.01) + assert record['status'] == 'failed' + assert not completion.done() and guard._readers == 1 + assert scope.request_ids == {ref['request_id']} and len(scope.jobs) == 1 + exit_gate.set() + await asyncio.wait_for(completion, 1) + assert guard._readers == 0 + finally: + exit_gate.set() + await queue._compute_worker.stop() + + + +@pytest.mark.asyncio +async def test_runtime_does_not_certify_an_unconfirmed_cancel_ack(): + from twinkle.server.sampler.runtime import SamplerRuntime + from twinkle.server.exceptions import BackendBusyError + guard = SamplerConcurrency() + runtime = SamplerRuntime(SimpleNamespace(sampling_guard=guard)) + task = asyncio.create_task(asyncio.sleep(0)) + await task + from twinkle.server.sampler.runtime import SamplingOperation + runtime._operations['cancelled'] = SamplingOperation(ExecutionScope(), task) + guard.invalidate_weights(RuntimeError('cancel ACK preceded actual generation exit')) + with pytest.raises(BackendBusyError): + await runtime.wait_operation('cancelled') + assert 'cancelled' in runtime._operations + + +@pytest.mark.asyncio +async def test_cold_lora_config_does_not_block_sampler_health(monkeypatch): + import json + import sys + from twinkle.server.sampler.runtime import SamplerRuntime + entered, release = threading.Event(), threading.Event() + calls = [] + + class ColdConfig: + + def __init__(self, **kwargs): + entered.set() + assert release.wait(2) + + monkeypatch.setitem(sys.modules, 'peft', SimpleNamespace(LoraConfig=ColdConfig)) + + async def call_backend(fn, *args, **kwargs): + return fn(*args, **kwargs) + + def add_adapter(name, config): + calls.append((name, config)) + + owner = SimpleNamespace( + model_id='mock', + sampling_guard=SamplerConcurrency(), + sampler=SimpleNamespace(add_adapter_to_sampler=add_adapter), + call_backend=call_backend) + runtime = SamplerRuntime(owner) + payload = { + 'endpoint': 'twinkle/add_adapter_to_sampler', + 'method': 'POST', + 'headers': [(b'content-type', b'application/json'), (b'x-request-id', b'configuration')], + 'body': json.dumps({ + 'adapter_name': 'a', + 'config': { + 'r': 4 + } + }).encode() + } + invocation = asyncio.create_task(runtime.invoke('configuration', payload)) + try: + began = await asyncio.to_thread(entered.wait, 5) + assert began, invocation.result() if invocation.done() else 'Configuration route did not start' + assert (await asyncio.wait_for(runtime.runtime_identity(), 0.1))['unhealthy'] is False + assert not invocation.done() and not calls + release.set() + assert (await invocation)['status'] == 200 + await runtime.wait_operation('configuration') + assert len(calls) == 1 + finally: + release.set() + await asyncio.gather(invocation, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_managed_compute_http_cannot_bypass_pool_admission(monkeypatch): + import httpx + from fastapi import FastAPI + from twinkle.server.exceptions import RequestRejectedError + from twinkle.server.sampler import app as module + + def scaffold(title, register_routes, **kwargs): + app = FastAPI() + register_routes(app, lambda: SimpleNamespace()) + return app + + monkeypatch.setattr(module, 'build_deployment_app', scaffold) + monkeypatch.setattr(module, 'bind_deployment', lambda app, *args, **kwargs: app) + args = { + 'model_id': 'mock', + 'nproc_per_node': 1, + 'device_group': { + 'name': 'sampler', + 'ranks': 1, + 'device_type': 'cpu' + }, + 'device_mesh': { + 'dp_size': 1 + }, + 'deploy_options': {}, + 'sampler_type': 'mock' + } + managed = module.build_sampler_app(**args, sampling_context={'instance_id': 'owned'}) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=managed), base_url='http://test') as client: + with pytest.raises(RequestRejectedError, match='stable sampler endpoint'): + await client.post('/twinkle/create', json={}) + standalone = module.build_sampler_app(**args) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=standalone), base_url='http://test') as client: + assert (await client.post('/twinkle/create', json={})).status_code == 200 + + +@pytest.mark.asyncio +async def test_scope_waits_for_jobs_registered_by_an_earlier_job(): + scope = ExecutionScope() + child_started, child_release, parent_release = asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def child(): + child_started.set() + await child_release.wait() + + async def parent(): + await parent_release.wait() + scope.track(asyncio.create_task(child())) + + scope.track(asyncio.create_task(parent())) + waiter = asyncio.create_task(scope.wait()) + try: + await asyncio.sleep(0) + parent_release.set() + await child_started.wait() + with pytest.raises(TimeoutError): + await asyncio.wait_for(asyncio.shield(waiter), 0.05) + child_release.set() + await waiter + finally: + child_release.set() + await asyncio.gather(waiter, return_exceptions=True) diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py index 34a08d665..e9cba0ed1 100644 --- a/tests/server/sampler/test_sampling_concurrency.py +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -22,8 +22,8 @@ from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes from twinkle.server.task_queue.backend_gate import BackendGate -from twinkle.server.task_queue.config import TaskQueueConfig -from twinkle.server.task_queue.types import BackendBusyError +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.exceptions import BackendBusyError from twinkle.server.task_queue.worker import ComputeWorker @@ -41,6 +41,13 @@ def __init__(self): async def store_future_status(self, request_id, status, model_id=None, **kwargs): self.records[request_id] = {'status': status, **kwargs} + async def try_start_future(self, request_id): + record = self.records.get(request_id) + if record is None or record['status'] not in ('pending', 'queued'): + return False + record['status'] = 'running' + return True + async def get_future(self, request_id): return self.records.get(request_id) diff --git a/tests/server/sampler/test_stream_guarantees.py b/tests/server/sampler/test_stream_guarantees.py index 370dbe8c1..d2b72b159 100644 --- a/tests/server/sampler/test_stream_guarantees.py +++ b/tests/server/sampler/test_stream_guarantees.py @@ -10,7 +10,7 @@ from fastapi import FastAPI from twinkle.server.sampler.app import SamplerManagement -from twinkle.server.sampler.twinkle_handlers import _stream_queue +from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes, _stream_queue class _BlockingQueue: diff --git a/tests/server/sampler/test_streaming_bridge.py b/tests/server/sampler/test_streaming_bridge.py new file mode 100644 index 000000000..bba725b4f --- /dev/null +++ b/tests/server/sampler/test_streaming_bridge.py @@ -0,0 +1,45 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Consumer disconnection never shortens backend ownership.""" +from types import SimpleNamespace + +import pytest + +from twinkle.server.sampler.backends.streaming import STREAM_SENTINEL, stream_to_queue + + +@pytest.mark.parametrize('disconnect', [False, True]) +def test_stream_drains_and_closes_after_consumer_disconnect(disconnect): + consumed, closed, published = [], [], [] + + def generate(*args): + try: + for index in range(3): + consumed.append(index) + yield index, None + finally: + closed.append(True) + + def put(value): + if disconnect: + raise RuntimeError('queue actor is gone') + published.append(value) + + stream_to_queue(SimpleNamespace(sample_stream=generate), SimpleNamespace(put=put), []) + assert consumed == [0, 1, 2] and closed == [True] + assert published == ([] if disconnect else [(0, None), (1, None), (2, None), STREAM_SENTINEL]) + + +def test_backend_failure_is_reported_and_iterator_closed(): + closed, published = [], [] + + def generate(*args): + try: + yield 1, None + raise ValueError('generation failed') + finally: + closed.append(True) + + stream_to_queue(SimpleNamespace(sample_stream=generate), SimpleNamespace(put=published.append), []) + assert closed == [True] + assert isinstance(published[-2], ValueError) + assert published[-1] == STREAM_SENTINEL diff --git a/tests/server/sampler/test_twinkle_async_rows.py b/tests/server/sampler/test_twinkle_async_rows.py index 05ccdc3b5..b83b043be 100644 --- a/tests/server/sampler/test_twinkle_async_rows.py +++ b/tests/server/sampler/test_twinkle_async_rows.py @@ -77,7 +77,7 @@ def __init__(self): async def _on_request_start(self, _request): return 'token' - async def submit_concurrent_and_peek(self, coro_factory, *, model_id=None, task_type=None): + async def submit_concurrent_and_peek(self, coro_factory, *, model_id=None, task_type=None, execution_scope=None): self.scheduled.append({'model_id': model_id, 'task_type': task_type}) result = await coro_factory() from twinkle.protocol.types.lifecycle import TaskEnvelope diff --git a/tests/server/session_resource/test_contract.py b/tests/server/session_resource/test_contract.py index 477b509cf..843b7d611 100644 --- a/tests/server/session_resource/test_contract.py +++ b/tests/server/session_resource/test_contract.py @@ -7,7 +7,6 @@ import pytest from twinkle.server.exceptions import RequestRejectedError -from twinkle.server.session_resource.adapter import AdapterManagerMixin from twinkle.server.session_resource.base import SessionResourceMixin from twinkle.server.session_resource.processor import ProcessorManagerMixin @@ -39,16 +38,12 @@ class _MissingBaseHook(SessionResourceMixin): pass -class _MissingAdapterHook(AdapterManagerMixin): - pass - - class _MissingProcessorHook(ProcessorManagerMixin): pass def test_all_resource_expiry_hooks_are_abstract() -> None: - for cls in (_MissingBaseHook, _MissingAdapterHook, _MissingProcessorHook): + for cls in (_MissingBaseHook, _MissingProcessorHook): with pytest.raises(TypeError): cls() diff --git a/tests/server/start_e2e_server.py b/tests/server/start_e2e_server.py index 4af2e2b85..21995f969 100644 --- a/tests/server/start_e2e_server.py +++ b/tests/server/start_e2e_server.py @@ -38,10 +38,10 @@ MEGATRON_BASE = "tests/server/config/server_config_4b_e2e_megatron.yaml" _FULL = { "train_mode": "full", + "max_replicas": 1, # Full-parameter training holds the deployment exclusively; use long # timeouts so the tenant is not expired mid-test. - "adapter_config.adapter_timeout": 3600, - "adapter_config.adapter_max_lifetime": 36000, + "idle_timeout_s": 3600, } VARIANTS = { "lora": (TRANSFORMERS_BASE, {}), diff --git a/tests/server/state/fake_backend.py b/tests/server/state/fake_backend.py index eaac98833..10f5c9e56 100644 --- a/tests/server/state/fake_backend.py +++ b/tests/server/state/fake_backend.py @@ -67,13 +67,15 @@ async def update_atomic( key: str, transform: Callable[[Any | None], Any | None], ttl: int | None = None, + *, + return_changed: bool = False, ) -> Any | None: current = await self.get(key) updated = transform(current) if updated is None: - return current + return False if return_changed else current await self.set(key, updated, ttl) - return updated + return True if return_changed else updated async def flush_all(self) -> None: self._store.clear() diff --git a/tests/server/state/test_future_lifecycle.py b/tests/server/state/test_future_lifecycle.py index 9728780f3..e596391de 100644 --- a/tests/server/state/test_future_lifecycle.py +++ b/tests/server/state/test_future_lifecycle.py @@ -7,86 +7,20 @@ """ from __future__ import annotations +import asyncio import time -from collections.abc import Callable -from fnmatch import fnmatch -from typing import Any from unittest import mock import pytest -from twinkle.server.state.backend.base import StateBackend +from tests.server.state.fake_backend import FakeBackend from twinkle.server.state.future_manager import FutureManager from twinkle.server.state.models import FutureFailureRecord -class _FakeBackend(StateBackend): - """Dict-backed StateBackend. Single event loop, so ``update_atomic`` is atomic - simply by not awaiting between read and write -- the real backends' guarantee.""" - - def __init__(self) -> None: - self._store: dict[str, tuple[Any, float | None]] = {} - - def _is_expired(self, key: str) -> bool: - entry = self._store.get(key) - if entry is None: - return True - _, expire_at = entry - if expire_at is not None and time.time() >= expire_at: - del self._store[key] - return True - return False - - async def set(self, key: str, value: Any, ttl: int | None = None) -> None: - self._store[key] = (value, (time.time() + ttl) if ttl is not None else None) - - async def get(self, key: str) -> Any | None: - return None if self._is_expired(key) else self._store[key][0] - - async def mget(self, keys: list[str]) -> list[Any | None]: - return [await self.get(key) for key in keys] - - async def delete(self, key: str) -> None: - self._store.pop(key, None) - - async def exists(self, key: str) -> bool: - return not self._is_expired(key) - - async def keys(self, pattern: str) -> list[str]: - return [k for k in list(self._store) if not self._is_expired(k) and fnmatch(k, pattern)] - - async def count(self, pattern: str) -> int: - return len(await self.keys(pattern)) - - async def set_nx(self, key: str, value: Any, ttl: int | None = None) -> bool: - if not self._is_expired(key): - return False - await self.set(key, value, ttl) - return True - - async def update_atomic( - self, - key: str, - transform: Callable[[Any | None], Any | None], - ttl: int | None = None, - ) -> Any | None: - current = await self.get(key) - updated = transform(current) - if updated is None: - return current - await self.set(key, updated, ttl) - return updated - - async def close(self) -> None: - pass - - async def health_check(self) -> bool: - return True - - @pytest.fixture def manager(): - return FutureManager(_FakeBackend(), expiration_timeout=300.0) + return FutureManager(FakeBackend(), expiration_timeout=300.0) async def _store(manager, request_id, status, *, replica_id=None, absolute_deadline=None): @@ -203,7 +137,7 @@ async def test_stored_timestamps_align_with_wall_clock_regardless_of_host_tz(man @pytest.mark.asyncio async def test_claim_seq_dedups_then_release_readmits(): from twinkle.server.state.server_state import ServerState - state = ServerState(backend=_FakeBackend()) + state = ServerState(backend=FakeBackend()) # First claim of a (session, seq_id) is unseen -> None, caller proceeds to enqueue. assert await state.claim_seq('seq::s1::1', 'reqA', ttl=60) is None # A duplicate claim returns the original request_id -> caller returns its envelope. @@ -218,7 +152,7 @@ async def test_claim_seq_dedups_then_release_readmits(): @pytest.mark.asyncio async def test_cancel_drops_pending_but_never_running(): from twinkle.server.state.server_state import ServerState - state = ServerState(backend=_FakeBackend()) + state = ServerState(backend=FakeBackend()) # pending -> cancel drops it to a terminal domain failure. await state.store_future_status('rp', 'pending', 'm1') assert await state.cancel_future('rp') == {'cancelled': True, 'state': 'cancelled'} @@ -247,3 +181,68 @@ def test_cancelled_record_maps_to_error_envelope(): env = envelope_from_record('rc', rec) assert env.status == 'cancelled' assert env.error is not None and env.error.error_code == 499 + + +@pytest.mark.asyncio +async def test_late_queued_and_repeated_running_cannot_regress_or_renew(manager): + await _store(manager, 'race', 'pending') + assert await manager.try_start('race') + before = (await manager.get('race')).model_dump() + await _store(manager, 'race', 'queued') + await _store(manager, 'race', 'running') + assert (await manager.get('race')).model_dump() == before + assert await manager.cancel_if_pending('race') == 'running' + assert not await manager.try_start('race') + assert not await manager.try_start('missing') + + +@pytest.mark.asyncio +@pytest.mark.parametrize('cancel_first', [False, True]) +async def test_start_and_cancel_have_one_atomic_winner(manager, cancel_first): + await _store(manager, 'race', 'queued') + operations = [manager.try_start('race'), manager.cancel_if_pending('race')] + if cancel_first: + operations.reverse() + results = await asyncio.gather(*operations) + started = results[1] if cancel_first else results[0] + status = (await manager.get('race')).status + assert started == (status == 'running') + assert status in ('running', 'cancelled') + assert not await manager.try_start('race') + + +@pytest.mark.asyncio +async def test_concurrent_starters_only_one_obtains_execution(manager): + await _store(manager, 'race', 'pending') + assert sum(await asyncio.gather(*(manager.try_start('race') for _ in range(20)))) == 1 + + +@pytest.mark.asyncio +async def test_future_owner_outlives_admission_retention_until_result_cleanup(monkeypatch): + from twinkle.server.exceptions import ResourceNotFoundError + from twinkle.server.state import ServerState + state = ServerState(backend=FakeBackend(), expiration_timeout=10) + await state.store_future_status('owned', 'pending', 'training_test', owner_token='owner') + await state.store_future_status('owned', 'completed', 'training_test', result={'value': 42}) + now = time.time() + monkeypatch.setattr(time, 'time', lambda: now + 11) + assert await state.get_future('owned') is not None + await state.assert_future_owner('owned', 'owner') + with pytest.raises(ResourceNotFoundError): + await state.assert_future_owner('owned', 'other') + await state._future_mgr.cleanup_expired(now + 11) + assert await state.get_future('owned') is None + assert not await state._backend.keys('future*') + + +@pytest.mark.asyncio +async def test_future_owner_is_immutable_and_not_part_of_wire_envelope(): + from twinkle.server.state import ServerState + from twinkle.server.lifecycle.envelope import envelope_from_record + state = ServerState(backend=FakeBackend()) + await state.store_future_status('owned', 'pending', None, owner_token='secret-owner') + await state.store_future_status('owned', 'completed', None, result=42, owner_token='replacement') + record = await state.get_future('owned') + assert record['owner_token'] == 'secret-owner' + assert 'owner_token' not in envelope_from_record('owned', record).model_dump() + assert 'secret-owner' not in envelope_from_record('owned', record).model_dump_json() diff --git a/tests/server/state/test_leader_election.py b/tests/server/state/test_leader_election.py index 792989c30..216105e99 100644 --- a/tests/server/state/test_leader_election.py +++ b/tests/server/state/test_leader_election.py @@ -11,10 +11,14 @@ import asyncio import pytest +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from tests.server.state.fake_backend import FakeBackend from twinkle.server.state import ServerState from twinkle.server.state.backend.memory_backend import RayActorBackend -from twinkle.server.state.cleanup_coordinator import LEADER_KEY, LEASE_RENEW +from twinkle.server.state.cleanup_coordinator import LEADER_KEY, LEASE_RENEW, ResourceCleanupCoordinator from twinkle.server.telemetry import MetricsRegistry @@ -100,13 +104,7 @@ async def test_renew_keeps_leader() -> None: @pytest.mark.asyncio async def test_leader_recovers_after_renewal_failure() -> None: - """Regression: when a leader's renewal raises, it releases - the lease best-effort so the very next election tick re-acquires leadership - without waiting LEASE_TTL. - - On unfixed code ``_is_leader`` flips to False while the lease value lingers - in the backend, so ``set_nx`` keeps returning False for up to LEASE_TTL. - """ + """An unconfirmed renewal fences cleanup until the original lease expires.""" backend = RayActorBackend() state = ServerState(backend=backend, cleanup_interval=600.0) @@ -128,11 +126,13 @@ async def flaky_update_atomic(*args, **kwargs): from unittest import mock with mock.patch.object(backend, 'update_atomic', side_effect=flaky_update_atomic): await state._try_acquire_or_renew() - # Renewal failed → no longer leader, and the stale lease was released. assert state._is_leader is False - # The lease key was deleted best-effort, so the next tick re-acquires - # immediately (no LEASE_TTL wait, no lingering self-owned lease). + assert await backend.get(LEADER_KEY) == state._leader_id + await state._try_acquire_or_renew() + assert state._is_leader is False + await backend.set(LEADER_KEY, state._leader_id, ttl=1) + await asyncio.sleep(1.1) await state._try_acquire_or_renew() assert state._is_leader is True @@ -159,6 +159,75 @@ async def test_renewal_failure_does_not_steal_other_leader_lease() -> None: assert await backend.get(LEADER_KEY) == leader_value, 'follower stole the real leader lease' +def coordinator(backend=None): + models = SimpleNamespace(rebuild_indexes=AsyncMock()) + return ResourceCleanupCoordinator(backend or FakeBackend(), {'models': models}, + cleanup_interval=600, expiration_timeout=300, + sweep_processor_quotas=AsyncMock()) + + +@pytest.mark.asyncio +async def test_stale_leader_renewal_failure_preserves_successor(monkeypatch): + backend = FakeBackend() + cleanup = coordinator(backend) + await cleanup._try_acquire_or_renew() + assert cleanup.is_leader + + async def successor_then_fail(*args, **kwargs): + await backend.set(LEADER_KEY, 'successor', ttl=30) + raise ConnectionError('renewal acknowledgement lost') + + monkeypatch.setattr(backend, 'update_atomic', successor_then_fail) + await cleanup._try_acquire_or_renew() + assert not cleanup.is_leader + assert await backend.get(LEADER_KEY) == 'successor' + + +@pytest.mark.asyncio +async def test_cleanup_concurrent_start_rebuilds_once_and_stop_joins(): + cleanup = coordinator() + entered, proceed = asyncio.Event(), asyncio.Event() + + async def rebuild(): + entered.set() + await proceed.wait() + + cleanup._managers['models'].rebuild_indexes.side_effect = rebuild + first = asyncio.create_task(cleanup.start()) + await entered.wait() + second = asyncio.create_task(cleanup.start()) + proceed.set() + assert await asyncio.gather(first, second) == [True, False] + tasks = (cleanup._cleanup_task, cleanup._leader_task) + assert await cleanup.stop() + assert all(task.done() for task in tasks) + assert not await cleanup.stop() + cleanup._managers['models'].rebuild_indexes.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_failed_cleanup_start_can_retry(): + cleanup = coordinator() + cleanup._managers['models'].rebuild_indexes.side_effect = [RuntimeError('offline'), None] + with pytest.raises(RuntimeError, match='offline'): + await cleanup.start() + assert cleanup._cleanup_task is None and cleanup._leader_task is None + assert await cleanup.start() + await cleanup.stop() + + +@pytest.mark.asyncio +async def test_lazy_cleanup_retries_after_failed_start(): + from twinkle.server.deployment import LazyCleanupMixin + client = LazyCleanupMixin() + client.state = SimpleNamespace(start_cleanup_task=AsyncMock(side_effect=[RuntimeError('offline'), True])) + await client._ensure_state_cleanup_started() + assert not client._state_cleanup_started + await client._ensure_state_cleanup_started() + assert client._state_cleanup_started + assert client.state.start_cleanup_task.await_count == 2 + + # ============================================================ # Metrics ObservableGauge (merged from test_metrics_observable_gauge) # ============================================================ diff --git a/tests/server/state/test_managers.py b/tests/server/state/test_managers.py index 242dc3508..6c2dfd47b 100644 --- a/tests/server/state/test_managers.py +++ b/tests/server/state/test_managers.py @@ -534,3 +534,20 @@ async def get_all_then_touch(): session_alive = await state.get_session_last_heartbeat('s-race') is not None model_alive = await state.get_model_metadata('m-race') is not None assert session_alive == model_alive + + +@pytest.mark.parametrize('environment_limit', [None, '7']) +def test_state_initializes_without_or_with_a_partial_environment_policy(monkeypatch, environment_limit): + """A missing optional policy must resolve to defaults before field access.""" + import uuid + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.config.application_spec import SERVER_STATE_ENV_KEYS + from twinkle.server.state.server_state import get_server_state + + for key in SERVER_STATE_ENV_KEYS: + monkeypatch.delenv(key, raising=False) + if environment_limit is not None: + monkeypatch.setenv('TWINKLE_SERVER_STATE_PER_TOKEN_MODEL_LIMIT', environment_limit) + state = get_server_state(cache_key='policy-' + uuid.uuid4().hex, backend=FakeBackend()) + assert state.expiration_timeout == 86400 + assert state.cleanup_interval == 3600 diff --git a/tests/server/state/test_redis_read_retry.py b/tests/server/state/test_redis_read_retry.py new file mode 100644 index 000000000..c5cac36e3 --- /dev/null +++ b/tests/server/state/test_redis_read_retry.py @@ -0,0 +1,49 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Read retry is bounded, request-local, and never replays writes.""" +import asyncio +from unittest.mock import AsyncMock + +import pytest +from redis.exceptions import ConnectionError + +from twinkle.server.state.backend.redis_backend import RedisBackend + + +@pytest.mark.asyncio +async def test_stale_read_retries_without_blocking_another_reader(monkeypatch): + backend = RedisBackend('redis://127.0.0.1:1') + entered, resume = asyncio.Event(), asyncio.Event() + attempts = [] + + async def read(key): + attempts.append(key) + if key == 'stale' and attempts.count(key) == 1: + entered.set() + await resume.wait() + raise ConnectionError('stale connection') + return '"' + key + '"' + + backend._client.get = read + monkeypatch.setattr('twinkle.server.state.backend.redis_backend.asyncio.sleep', AsyncMock()) + try: + first = asyncio.create_task(backend.get('stale')) + await entered.wait() + assert await backend.get('healthy') == 'healthy' + resume.set() + assert await first == 'stale' + assert attempts == ['stale', 'healthy', 'stale'] + finally: + resume.set() + await backend.close() + + +@pytest.mark.asyncio +async def test_cancellation_is_not_retried(): + backend = RedisBackend('redis://127.0.0.1:1') + backend._client.get = AsyncMock(side_effect=asyncio.CancelledError()) + try: + with pytest.raises(asyncio.CancelledError): + await backend.get('key') + backend._client.get.assert_awaited_once() + finally: + await backend.close() diff --git a/tests/server/state/test_redis_recovery.py b/tests/server/state/test_redis_recovery.py new file mode 100644 index 000000000..bfde03295 --- /dev/null +++ b/tests/server/state/test_redis_recovery.py @@ -0,0 +1,114 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Read-only Redis reconnects are bounded; uncertain writes are never replayed.""" +import shutil +from unittest.mock import AsyncMock, Mock + +import pytest + +redis = pytest.importorskip('redis') +ConnectionError = redis.exceptions.ConnectionError + +from twinkle.server.state.backend.redis_backend import RedisBackend +from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + yield # These backend tests need no Ray cluster. + + +@pytest.mark.asyncio +async def test_stale_read_is_retried_but_write_is_not(): + backend = RedisBackend('redis://127.0.0.1:1/0') + try: + backend._client.get = AsyncMock(side_effect=[ConnectionError('stale socket'), '{"value": 1}']) + assert await backend.get('owned') == {'value': 1} + assert backend._client.get.await_count == 2 + backend._client.set = AsyncMock(side_effect=ConnectionError('lost acknowledgement')) + with pytest.raises(ConnectionError): + await backend.set('owned', {'value': 2}) + assert backend._client.set.await_count == 1 + finally: + await backend.close() + + +@pytest.mark.asyncio +async def test_persistent_outage_does_not_look_like_missing_state(): + backend = RedisBackend('redis://127.0.0.1:1/0') + try: + backend._client.get = AsyncMock(side_effect=ConnectionError('offline')) + with pytest.raises(ConnectionError): + await backend.get('controller-active') + assert backend._client.get.await_count == 3 + finally: + await backend.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif(not shutil.which('redis-server'), reason='redis-server executable required') +async def test_real_restart_recovers_old_pooled_readers(tmp_path): + managed = ManagedRedis(tmp_path) + managed.start() + first, stale, atomic = (RedisBackend(managed.url) for _ in range(3)) + try: + await first.set('owner', {'epoch': 'original'}) + assert await stale.get('owner') == {'epoch': 'original'} + assert await atomic.get('owner') == {'epoch': 'original'} + managed.stop() + with pytest.raises(ConnectionError): + await first.get('owner') + managed.start() + assert await stale.get('owner') == {'epoch': 'original'} + assert await stale.exists('owner') + assert await first.mget(['owner', 'missing']) == [{'epoch': 'original'}, None] + assert await stale.health_check() + assert await atomic.update_atomic('counter', lambda value: int(value or 0) + 1) == 1 + assert await atomic.get('counter') == 1 + finally: + await first.close() + await stale.close() + await atomic.close() + managed.stop() + + +class PipelineStub: + def __init__(self): + self.watch = AsyncMock() + self.get = AsyncMock(return_value='0') + self.multi = Mock() + self.set = Mock() + self.execute = AsyncMock(return_value=[]) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + +@pytest.mark.asyncio +async def test_transport_watch_error_after_exec_is_never_replayed(): + backend = RedisBackend('redis://127.0.0.1:1/0') + pipe = PipelineStub() + pipe.execute.side_effect = redis.exceptions.WatchError( + 'A ConnectionError occurred while watching one or more keys') + backend._client.pipeline = lambda: pipe + try: + with pytest.raises(ConnectionError, match='outcome could not be confirmed'): + await backend.update_atomic('counter', lambda value: value + 1) + assert pipe.execute.await_count == 1 + finally: + await backend.close() + + +@pytest.mark.asyncio +async def test_connection_loss_before_exec_is_safe_to_retry(): + backend = RedisBackend('redis://127.0.0.1:1/0') + pipe = PipelineStub() + pipe.watch.side_effect = [ConnectionError('stale'), None] + backend._client.pipeline = lambda: pipe + try: + assert await backend.update_atomic('counter', lambda value: value + 1) == 1 + assert pipe.watch.await_count == 2 and pipe.execute.await_count == 1 + finally: + await backend.close() diff --git a/tests/server/state/test_sequence.py b/tests/server/state/test_sequence.py new file mode 100644 index 000000000..053ef23ef --- /dev/null +++ b/tests/server/state/test_sequence.py @@ -0,0 +1,70 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import time +import uuid + +import pytest +import pytest_asyncio + +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.state import ServerState +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.state.backend.redis_backend import RedisBackend + + +@pytest_asyncio.fixture(params=['memory', 'redis']) +async def state(request): + if request.param == 'redis': + url = os.environ.get('TWINKLE_TEST_REDIS_URL') + if not url: + pytest.skip('TWINKLE_TEST_REDIS_URL is required for the Redis regression') + backend = RedisBackend(url, key_prefix='sequence-test-' + uuid.uuid4().hex + ':') + else: + backend = RayActorBackend() + yield ServerState(backend=backend) + if request.param == 'redis': + for key in await backend.keys('*'): + await backend.delete(key) + await backend.close() + + +@pytest.mark.asyncio +async def test_expired_retry_receipts_cannot_execute_old_sequences_again(state, monkeypatch): + backend = state._backend + now = time.time() + monkeypatch.setattr('twinkle.server.state.server_state.time.time', lambda: now) + assert await state.claim_training_seq('training', 'native::session::step', 1, 'first', 10) is None + assert await state.claim_training_seq('training', 'native::session::step', 1, 'retry', 10) == 'first' + now += 11 + assert await state.claim_training_seq('training', 'native::session::step', 2, 'second', 10) is None + with pytest.raises(RequestRejectedError, match='retry window expired'): + await state.claim_training_seq('training', 'native::session::step', 1, 'late', 10) + ledger = await backend.get('seq_ledger::training::native::session::step') + assert list(ledger['receipts']) == ['2'] + assert len(await backend.keys('seq_ledger::*')) == 1 + + +@pytest.mark.asyncio +async def test_out_of_order_sequences_and_preflight_release_remain_retryable(state): + assert await state.claim_training_seq('training', 'tinker::step', 2, 'second', 10) is None + assert await state.claim_training_seq('training', 'tinker::step', 1, 'first', 10) is None + await state.release_training_seq('training', 'tinker::step', 1, 'another-request') + assert await state.claim_training_seq('training', 'tinker::step', 1, 'retry', 10) == 'first' + await state.release_training_seq('training', 'tinker::step', 1, 'first') + assert await state.claim_training_seq('training', 'tinker::step', 1, 'retry', 10) is None + assert await state.claim_training_seq('other-training', 'tinker::step', 1, 'other', 10) is None + + +@pytest.mark.asyncio +async def test_retiring_a_handle_removes_native_and_tinker_stream_ledgers(state): + from tests.server.training.helpers import create, setup + controller, state, _, _ = await setup(backend=state._backend) + try: + identity = await create(controller, state, 'dedup-retirement') + for stream in ('native::session::step', 'tinker::optim_step'): + await state.claim_training_seq(identity, stream, 1, stream, 10) + assert len(await state._backend.keys(f'seq_ledger::{identity}::*')) == 2 + await controller.release(identity, 'owner', 'session') + assert await state._backend.keys(f'seq_ledger::{identity}::*') == [] + finally: + await controller.shutdown() diff --git a/tests/server/state/test_update_atomic.py b/tests/server/state/test_update_atomic.py index d03b3a79e..2bcc78f29 100644 --- a/tests/server/state/test_update_atomic.py +++ b/tests/server/state/test_update_atomic.py @@ -129,6 +129,43 @@ async def test_key_prefix_is_hidden_from_logical_keys(backend_pair) -> None: assert await first.keys('session::*') == ['session::one'] +@pytest.mark.asyncio +async def test_atomic_update_reports_this_call_not_the_existing_value(backend_pair): + first, second = backend_pair + assert await first.update_atomic('claim', functools.partial(_replace_with, value=True), + return_changed=True) is True + assert await second.update_atomic('claim', _no_op, return_changed=True) is False + assert await second.get('claim') is True + assert await second.update_atomic('missing', _no_op, return_changed=True) is False + + +@pytest.mark.asyncio +async def test_future_execution_is_claimed_once_across_backend_handles(backend_pair): + from twinkle.server.state.future_manager import FutureManager + first, second = [FutureManager(backend, 300) for backend in backend_pair] + await first.store_status('start', 'pending', 'm') + starts = await asyncio.gather(*(manager.try_start('start') for manager in [first, second] * 8)) + assert sum(starts) == 1 + await second.store_status('start', 'queued', 'm') + assert await first.cancel_if_pending('start') == 'running' + assert (await first.get('start')).status == 'running' + + +@pytest.mark.asyncio +@pytest.mark.parametrize('cancel_first', [False, True]) +async def test_future_start_cancel_race_across_backend_handles(backend_pair, cancel_first): + from twinkle.server.state.future_manager import FutureManager + first, second = [FutureManager(backend, 300) for backend in backend_pair] + await first.store_status('start', 'queued', 'm') + calls = [first.try_start('start'), second.cancel_if_pending('start')] + if cancel_first: + calls.reverse() + results = await asyncio.gather(*calls) + started = results[1] if cancel_first else results[0] + assert started == ((await first.get('start')).status == 'running') + assert not await second.try_start('start') + + # ---------- Memory backend ----------------------------------------------- # @@ -258,3 +295,19 @@ def force_watch_conflict(current: Any | None) -> int: 'error_code': 503, 'request_id': 'req-contention', } + + +@pytest.mark.asyncio +async def test_future_owner_is_shared_with_result_and_removed_with_it(backend_pair): + from twinkle.server.exceptions import ResourceNotFoundError + from twinkle.server.state import ServerState + first, second = (ServerState(backend=backend) for backend in backend_pair) + await first.store_future_status('owned', 'pending', None, owner_token='owner') + await second.store_future_status('owned', 'completed', None, result=42) + await first.assert_future_owner('owned', 'owner') + with pytest.raises(ResourceNotFoundError): + await second.assert_future_owner('owned', 'other') + assert (await second.get_future('owned'))['owner_token'] == 'owner' + await first._future_mgr.cleanup_expired(float('inf')) + assert await second.get_future('owned') is None + assert not await backend_pair[0].keys('future*') diff --git a/tests/server/static/backend_call_exemptions.py b/tests/server/static/backend_call_exemptions.py index 1d0307ffe..3b1ac1142 100644 --- a/tests/server/static/backend_call_exemptions.py +++ b/tests/server/static/backend_call_exemptions.py @@ -12,12 +12,15 @@ The allowed exemptions are: +- ``_submit_group`` submits a batch of dev async futures from within its caller's + single ``call_backend(_submit_group, ...)``; it is already inside that boundary; + - the ray ``Queue.get`` inside ``sample_stream``'s ``_stream_queue``: it bridges the sampler actor's process boundary and is bounded by a dedicated double-timeout, not by ``call_backend``; -- the ``.sample_stream_to_queue.remote(...)`` call inside ``sample_stream`` - itself: streaming generation must keep producing while the HTTP response streams, - so it cannot use ``call_backend`` as-is and carries its own double timeout. This is +- the ``.sample_stream_to_queue.remote(...)`` call inside ``_stream_generation``: + the HTTP queue has its own deadline, while the actor completion is retained by + ExecutionScope and the weight guard until real engine-loop generation exits. This is a ``remote_function`` bypass that the guard now *detects* (via backend-derived local tracking) and that is *explicitly* accepted here — replacing the previous "No remote_function call is exempt" claim, which was true only because the guard @@ -27,6 +30,7 @@ # (module_relpath under src/twinkle/server, innermost enclosing function name) BACKEND_CALL_EXEMPTIONS: frozenset[tuple[str, str]] = frozenset({ + ('sampler/twinkle_handlers.py', '_submit_group'), ('sampler/twinkle_handlers.py', '_stream_queue'), - ('sampler/twinkle_handlers.py', 'sample_stream'), + ('sampler/twinkle_handlers.py', '_stream_generation'), }) diff --git a/tests/server/static/test_adapter_name_mapping.py b/tests/server/static/test_adapter_name_mapping.py index 6856f099e..708bef399 100644 --- a/tests/server/static/test_adapter_name_mapping.py +++ b/tests/server/static/test_adapter_name_mapping.py @@ -2,7 +2,7 @@ """Static check: model backend calls map the tenant adapter name (F002 / P002). ``run_submit`` hands each ``backend_call`` a *tenant-scoped* adapter name -(``owner_id-``). Before that reaches a backend model method it must be +(``training_id``). Before that reaches a backend model method it must be translated by ``ModelManagement.resolve_model_adapter_name`` — which returns ``''`` in full-parameter mode (the empty-string default optimizer group) and the name unchanged in LoRA mode. The inline forward/backward/step endpoints did this; diff --git a/tests/server/task_queue/__init__.py b/tests/server/task_queue/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/task_queue/helpers.py b/tests/server/task_queue/helpers.py new file mode 100644 index 000000000..b3e9bb5de --- /dev/null +++ b/tests/server/task_queue/helpers.py @@ -0,0 +1,59 @@ +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.task_queue.mixin import TaskQueueMixin +from twinkle.server.task_queue.worker import ComputeWorker + + +class _DummyState: + + def __init__(self): + self.records = [] + self._latest = {} + + async def store_future_status(self, *args, **kwargs): + self.records.append((args, kwargs)) + request_id, status = args[0], args[1] + self._latest[request_id] = { + 'status': status, + 'result': kwargs.get('result'), + 'failure': (kwargs['failure'].model_dump() if kwargs.get('failure') is not None else None), + 'queue_state': kwargs.get('queue_state'), + 'queue_state_reason': kwargs.get('queue_state_reason'), + } + + async def try_start_future(self, request_id): + record = self._latest.get(request_id) + if record is None or record['status'] not in ('pending', 'queued'): + return False + await self.store_future_status(request_id, 'running', None, queue_state='active') + return self._latest[request_id]['status'] == 'running' + + async def get_future(self, request_id): + return self._latest.get(request_id) + + +class _AllowingRateLimiter: + + async def check_and_record(self, token, input_tokens): + return True, None + + +class _DummyQueue(TaskQueueMixin): + + def __init__(self): + self.state = _DummyState() + # A generous Inline_Fast_Path window keeps "task settles inside submit" + # deterministic for the trivial in-process coroutines used here. + self._task_queue_config = TaskQueueConfig(inline_fast_path_timeout=5.0) + self._rate_limiter = _AllowingRateLimiter() + self._task_metrics = None + self._deployment_name = 'test' + + def enable_compute_worker(self): + self._compute_worker = ComputeWorker( + state=self.state, + config=self._task_queue_config, + task_metrics=None, + deployment_name=self._deployment_name, + ) + self._event_loop = None + self._queue_tasks = set() diff --git a/tests/server/utils/test_rate_limiter.py b/tests/server/task_queue/test_rate_limiter.py similarity index 100% rename from tests/server/utils/test_rate_limiter.py rename to tests/server/task_queue/test_rate_limiter.py diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/task_queue/test_task_queue_mixin.py similarity index 69% rename from tests/server/utils/test_task_queue_mixin.py rename to tests/server/task_queue/test_task_queue_mixin.py index 033c372b4..e2f111e66 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/task_queue/test_task_queue_mixin.py @@ -1,59 +1,52 @@ +from tests.server.task_queue.helpers import _DummyState, _AllowingRateLimiter, _DummyQueue import asyncio import pytest import pytest_asyncio -from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.config.task_queue import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin -from twinkle.server.task_queue.types import UserTaskError -from twinkle.server.task_queue.worker import ComputeWorker +from twinkle.server.exceptions import RequestRejectedError +from twinkle.server.task_queue.worker import ComputeWorker, _exception_summary -class _DummyState: +def test_exception_summary_prefers_the_wrapped_cause(): + from ray.exceptions import RayTaskError - def __init__(self): - self.records = [] - self._latest = {} + err = RayTaskError( + 'TwinkleCompatTransformersModel.preflight_loss_inputs', + 'ray::TwinkleCompatTransformersModel.preflight_loss_inputs() (pid=1, ip=10.0.0.1)', + AssertionError('data/mask length mismatch at sample 0'), + ) - async def store_future_status(self, *args, **kwargs): - self.records.append((args, kwargs)) - request_id, status = args[0], args[1] - self._latest[request_id] = { - 'status': status, - 'result': kwargs.get('result'), - 'failure': (kwargs['failure'].model_dump() if kwargs.get('failure') is not None else None), - 'queue_state': kwargs.get('queue_state'), - 'queue_state_reason': kwargs.get('queue_state_reason'), - } + summary = _exception_summary(err) + assert summary.startswith('AssertionError:') + assert 'mismatch' in summary - async def get_future(self, request_id): - return self._latest.get(request_id) +def test_exception_summary_keeps_a_deliberate_wrapper_message(): + """``raise X from e`` is a deliberate message; only Ray's wrapper is unwrapped.""" -class _AllowingRateLimiter: + class _Rejected(Exception): + pass - async def check_and_record(self, token, input_tokens): - return True, None + try: + try: + raise ValueError('internal detail') + except ValueError as inner: + raise _Rejected('adapter_name is required') from inner + except _Rejected as outer: + summary = _exception_summary(outer) + assert summary == '_Rejected: adapter_name is required' -class _DummyQueue(TaskQueueMixin): - def __init__(self): - self.state = _DummyState() - # A generous Inline_Fast_Path window keeps "task settles inside submit" - # deterministic for the trivial in-process coroutines used here. - self._task_queue_config = TaskQueueConfig(inline_fast_path_timeout=5.0) - self._rate_limiter = _AllowingRateLimiter() - self._task_metrics = None - self._deployment_name = 'test' +def test_exception_summary_falls_back_to_the_ray_error_itself(): + from ray.exceptions import RayTaskError - def enable_compute_worker(self): - self._compute_worker = ComputeWorker( - state=self.state, - config=self._task_queue_config, - task_metrics=None, - deployment_name=self._deployment_name, - ) - self._event_loop = None + err = RayTaskError('Model.step', 'ray::Model.step() (pid=1)', None) + summary = _exception_summary(err) + assert summary.startswith('RayTaskError') + assert 'Model.step' in summary @pytest.mark.asyncio @@ -100,7 +93,7 @@ async def wait_for_with_completion_race(coro, timeout): @pytest.mark.asyncio -async def test_worker_can_restart_after_stopping_while_idle(): +async def test_idle_worker_stop_is_terminal_for_its_replica(): queue = _DummyQueue() queue.enable_compute_worker() worker = queue._compute_worker @@ -112,9 +105,11 @@ async def work(): await worker.ensure_started() await asyncio.sleep(0) # Let the idle worker wait for new work. await asyncio.wait_for(worker.stop(), timeout=1) - env = await queue.submit_and_peek(work, model_id='model1', token='token1') - assert env.status == 'completed' - assert env.result == {'ok': True} + from twinkle.server.exceptions import ResourceExhaustedError + with pytest.raises(ResourceExhaustedError): + await queue.submit_and_peek(work, model_id='model1', token='token1') + assert not queue.state.records + assert not worker._started finally: await worker.stop() @@ -181,6 +176,7 @@ async def test_preflight_accepts_batch_with_per_dp_multiple(): @pytest.mark.asyncio async def test_background_task_tracks_status(): queue = _DummyQueue() + queue.enable_compute_worker() async def work(): return {'ok': True} @@ -189,9 +185,13 @@ async def work(): work, model_id='model1', ) - await asyncio.sleep(0) + for _ in range(100): + if not queue._compute_worker._detached_tasks: + break + await asyncio.sleep(0.001) + await queue._compute_worker.stop() - assert [args[1] for args, _ in queue.state.records] == ['running', 'completed'] + assert [args[1] for args, _ in queue.state.records] == ['pending', 'running', 'completed'] assert queue.state.records[0][1]['absolute_deadline'] > 0 assert queue.state.records[-1][1]['result'] == {'ok': True} @@ -279,7 +279,7 @@ async def test_user_task_error_is_stored_as_user_failure(): queue.enable_compute_worker() async def work(): - raise UserTaskError('invalid request') + raise RequestRejectedError('invalid request') try: await queue.schedule_task(work, model_id='model1', token='token1') @@ -415,7 +415,7 @@ async def test_concurrent_task_keeps_queue_failure_mapping(make_concurrent_queue queue = make_concurrent_queue() async def work(): - raise UserTaskError('invalid request') + raise RequestRejectedError('invalid request') await queue.schedule_concurrent_task(work, token='token1') for _ in range(100): @@ -574,6 +574,108 @@ async def write(*args, **kwargs): assert [args[1] for args, _ in queue.state.records] == ['pending', 'failed'] +@pytest.mark.asyncio +async def test_background_capacity_and_shutdown_are_supervised(make_concurrent_queue): + from twinkle.server.exceptions import ResourceExhaustedError + queue = make_concurrent_queue(TaskQueueConfig(concurrent_task_limit=1)) + entered, exited = asyncio.Event(), asyncio.Event() + + async def work(): + entered.set() + try: + await asyncio.Event().wait() + finally: + exited.set() + + ref = await queue.schedule_background_task(work) + await asyncio.wait_for(entered.wait(), 1) + with pytest.raises(ResourceExhaustedError): + await queue.schedule_background_task(work) + assert len(queue.state._latest) == 1 + await queue._compute_worker.stop() + assert exited.is_set() + assert queue.state._latest[ref['request_id']]['status'] == 'failed' + assert queue._compute_worker._detached_reserved == 0 + + +@pytest.mark.asyncio +async def test_stopped_replica_rejects_serial_submission_without_a_future(): + from twinkle.server.exceptions import ResourceExhaustedError + queue = _DummyQueue() + queue.enable_compute_worker() + await queue._compute_worker.stop() + with pytest.raises(ResourceExhaustedError): + await queue.schedule_task(lambda: asyncio.sleep(0)) + assert not queue.state.records + assert not queue._compute_worker._started + + +@pytest.mark.asyncio +async def test_stop_during_pending_write_cannot_restart_the_serial_worker(): + from twinkle.server.exceptions import ResourceExhaustedError + queue = _DummyQueue() + queue.enable_compute_worker() + original = queue.state.store_future_status + + async def write_then_stop(*args, **kwargs): + await original(*args, **kwargs) + if args[1] == 'pending': + await queue._compute_worker.stop() + + queue.state.store_future_status = write_then_stop + with pytest.raises(ResourceExhaustedError): + await queue.schedule_task(lambda: asyncio.sleep(0), request_id='late-serial') + assert queue.state._latest['late-serial']['status'] == 'failed' + assert not queue._compute_worker._started + assert queue._compute_worker.total_queued() == 0 + + +@pytest.mark.asyncio +async def test_shutdown_drains_serial_tasks_that_never_started(): + queue = _DummyQueue() + queue.enable_compute_worker() + executed = [] + + async def work(): + executed.append(True) + + ref = await queue.schedule_task(work) + await queue._compute_worker.stop() + assert not executed + assert queue.state._latest[ref['request_id']]['status'] == 'failed' + assert queue._compute_worker.total_queued() == 0 + assert not queue._compute_worker.has_task(ref['request_id']) + + +@pytest.mark.asyncio +async def test_unknown_start_ack_never_executes_backend(make_concurrent_queue, monkeypatch): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.state.server_state import ServerState + queue = make_concurrent_queue() + backend = FakeBackend() + queue.state = queue._compute_worker._state = ServerState(backend=backend) + original = backend.update_atomic + executed = [] + + async def lost_ack(*args, **kwargs): + result = await original(*args, **kwargs) + if kwargs.get('return_changed'): + raise ConnectionError('start acknowledgement lost') + return result + + async def work(): + executed.append(True) + + monkeypatch.setattr(backend, 'update_atomic', lost_ack) + ref = await queue.schedule_concurrent_task(work) + for _ in range(100): + if not queue._compute_worker._detached_tasks: + break + await asyncio.sleep(0.001) + assert not executed + assert (await queue.state.get_future(ref['request_id']))['status'] == 'failed' + + @pytest.mark.asyncio async def test_cancel_winning_running_write_race_skips_execution(make_concurrent_queue): queue = make_concurrent_queue() @@ -603,3 +705,53 @@ async def work(): assert queue.state._latest[ref['request_id']]['status'] == 'cancelled' assert queue._compute_worker._detached_reserved == 0 assert not queue._compute_worker._detached_tasks + + +@pytest.mark.asyncio +@pytest.mark.parametrize('method', ['schedule_task', 'schedule_concurrent_task', 'schedule_background_task']) +async def test_admission_persists_future_and_owner_together(method): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.state import ServerState + queue = TaskQueueMixin() + queue.state = ServerState(backend=FakeBackend()) + queue.replica_id = 'ownership-test' + queue._init_task_queue() + entered, finish = asyncio.Event(), asyncio.Event() + + async def operation(): + entered.set() + await finish.wait() + return 42 + + try: + ref = await getattr(queue, method)(operation, token='owner') + record = await queue.state.get_future(ref['request_id']) + assert record['owner_token'] == 'owner' + assert not await queue.state._backend.keys('future_owner::*') + await asyncio.wait_for(entered.wait(), 1) + from twinkle.server.exceptions import ResourceNotFoundError + with pytest.raises(ResourceNotFoundError): + await queue.state.assert_future_owner(ref['request_id'], 'other') + finally: + finish.set() + await queue.shutdown_task_queue() + + +@pytest.mark.asyncio +async def test_backend_cancellation_is_an_explicit_server_failure(make_concurrent_queue): + from twinkle.server.lifecycle.envelope import envelope_from_record + from twinkle.server.task_queue.types import ExecutionScope + queue = make_concurrent_queue() + scope = ExecutionScope() + + async def work(): + raise asyncio.CancelledError() + + ref = await queue.schedule_concurrent_task(work, token='token1', execution_scope=scope) + await scope.wait() + record = await queue.state.get_future(ref['request_id']) + assert record['status'] == 'failed' + assert record['failure']['reason_code'] == 'backend_cancelled' + envelope = envelope_from_record(ref['request_id'], record) + assert envelope.error.error_code == 503 + assert envelope.error.category == 'server' diff --git a/tests/server/telemetry/test_training_metrics.py b/tests/server/telemetry/test_training_metrics.py new file mode 100644 index 000000000..7f4214593 --- /dev/null +++ b/tests/server/telemetry/test_training_metrics.py @@ -0,0 +1,259 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Managed training metrics stay low-cardinality and survive registry resets.""" +from collections import defaultdict +from types import SimpleNamespace + +import pytest + +from twinkle.server.telemetry import metrics +from twinkle.server.model.pool.controller import TrainingController +from twinkle.server.telemetry.metrics import TrainingInstanceMetrics + +_ALLOWED_ATTRIBUTES = {'backend', 'mode', 'state', 'status', 'reason', 'device_type', 'resident'} + + +@pytest.fixture(autouse=True) +def reset_registry(): + yield + metrics.MetricsRegistry.reset() + + +class FakeInstrument: + + def __init__(self, name): + self.name = name + self.calls = [] + + def add(self, value, attributes=None): + self.calls.append((value, attributes or {})) + + def record(self, value, attributes=None): + self.calls.append((value, attributes or {})) + + +class FakeMeter: + + def __init__(self): + self.instruments = {} + + def _create(self, name, **kwargs): + instrument = FakeInstrument(name) + self.instruments[name] = instrument + return instrument + + create_counter = _create + create_histogram = _create + create_up_down_counter = _create + + def create_observable_gauge(self, name, **kwargs): + return self._create(name, **kwargs) + + +def build_registry(monkeypatch): + meter = FakeMeter() + monkeypatch.setattr(metrics, 'get_meter', lambda name: meter) + metrics.MetricsRegistry.reset() + return metrics.MetricsRegistry.get(), meter + + +def test_training_and_redis_instruments_are_registered(monkeypatch): + registry, meter = build_registry(monkeypatch) + training = registry.training_metrics() + state = registry.state_backend_metrics() + + training.instances_active.set(2, tags={'backend': 'transformers', 'mode': 'lora', 'state': 'READY', + 'resident': 'false'}) + training.startup_seconds.observe(1.5, tags={'backend': 'transformers', 'mode': 'lora', 'status': 'success'}) + training.quarantined_total.inc(tags={'backend': 'megatron', 'mode': 'full', 'reason': 'timeout'}) + state.redis_health.set(1) + state.update_atomic_conflicts.inc() + state.update_atomic_retries.observe(2) + + expected = { + 'twinkle.training.instances.active', + 'twinkle.training.instance.startup.duration_seconds', + 'twinkle.training.instance.shutdown.duration_seconds', + 'twinkle.training.placement_group.allocation.duration_seconds', + 'twinkle.training.placement_group.removal.duration_seconds', + 'twinkle.training.quarantined.total', + 'twinkle.training.sweep.total', + 'twinkle.training.sweep.duration_seconds', + 'twinkle.training.devices.allocated', + 'twinkle.training.placement_groups.active', + 'twinkle.state.redis.health', + 'twinkle.state.redis.update_atomic.conflicts.total', + 'twinkle.state.redis.update_atomic.retries', + } + assert expected <= meter.instruments.keys() + for instrument in meter.instruments.values(): + for _, attributes in instrument.calls: + assert set(attributes) <= _ALLOWED_ATTRIBUTES + assert not {'token', 'session_id', 'training_id', 'instance_id', 'request_id'} & set(attributes) + + +def test_training_metric_adapters_rebind_after_registry_reset(monkeypatch): + first, _ = build_registry(monkeypatch) + first_metrics = first.training_metrics() + metrics.MetricsRegistry.reset() + second_metrics = metrics.MetricsRegistry.get().training_metrics() + assert first_metrics is not second_metrics + assert first_metrics.startup_seconds._instrument is not second_metrics.startup_seconds._instrument + + +@pytest.mark.parametrize('backend', ['transformers', 'megatron']) +def test_controller_metrics_clear_old_states_and_released_resources(monkeypatch, backend): + registry, meter = build_registry(monkeypatch) + spec = SimpleNamespace(backend=backend, device_type='cuda') + controller = TrainingController.__new__(TrainingController) + controller.config = SimpleNamespace(get=lambda base_model, mode: spec) + controller._metrics = registry.training_metrics() + controller._instance_metrics = TrainingInstanceMetrics(controller._metrics) + controller.instances = { + str(index): SimpleNamespace( + instance_id=f'instance_{index}', base_model='private-model', mode='lora', + state='READY', resident=False, devices=[index * 2, index * 2 + 1], + application_name=f'private-app-{index}', training_ids=[f'training_{index}']) + for index in range(2) + } + + def totals(name): + values = defaultdict(float) + for delta, attributes in meter.instruments[name].calls: + assert set(attributes) <= _ALLOWED_ATTRIBUTES + assert not {'instance_id', 'base_model', 'application_name', 'training_id'} & set(attributes) + values[tuple(sorted(attributes.items()))] += delta + return values + + instances_name = 'twinkle.training.instances.active' + devices_name = 'twinkle.training.devices.allocated' + groups_name = 'twinkle.training.placement_groups.active' + controller._report_instance_metrics() + first_calls = {name: len(meter.instruments[name].calls) + for name in (instances_name, devices_name, groups_name)} + controller._report_instance_metrics() + assert first_calls == {name: len(meter.instruments[name].calls) for name in first_calls} + assert sum(totals(instances_name).values()) == 2 + assert sum(totals(devices_name).values()) == 4 + assert sum(totals(groups_name).values()) == 2 + + controller.instances['0'].state = 'DRAINING' + controller._report_instance_metrics() + state_totals = {dict(key)['state']: value for key, value in totals(instances_name).items()} + assert state_totals == {'READY': 1, 'DRAINING': 1} + assert sum(totals(devices_name).values()) == 4 + assert sum(totals(groups_name).values()) == 2 + + for instance in controller.instances.values(): + instance.state = 'STOPPED' + controller._report_instance_metrics() + for name in (instances_name, devices_name, groups_name): + assert all(value == 0 for value in totals(name).values()) + + +@pytest.mark.asyncio +async def test_repeated_failed_cleanup_is_one_quarantine_transition(): + from unittest.mock import MagicMock + from tests.server.state.fake_backend import FakeBackend + from tests.server.training.helpers import create, setup + + controller, state, driver, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'metric-retry') + instruments = MagicMock() + controller._metrics = instruments + driver.fail_shutdown = True + try: + with pytest.raises(RuntimeError, match='quarantined'): + await controller.release(identity, 'owner', 'session') + for _ in range(2): + await controller.sweep() + await controller.wait_for_jobs() + assert instruments.quarantined_total.inc.call_count == 1 + assert instruments.shutdown_seconds.observe.call_count == 3 + finally: + driver.fail_shutdown = False + await controller.shutdown() + + +@pytest.mark.parametrize('phase', ['worker_death', 'pg_removal']) +def test_quarantine_classification_uses_typed_phase(phase): + from twinkle.server.exceptions import ResourceReleaseError + assert TrainingInstanceMetrics.quarantine_reason(ResourceReleaseError('unconfirmed', phase=phase)) == phase + assert TrainingInstanceMetrics.quarantine_reason(RuntimeError('user text: worker death')) == 'other' + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fail_start', [False, True]) +async def test_startup_and_shutdown_histograms_follow_real_attempts(fail_start): + from unittest.mock import MagicMock + from tests.server.state.fake_backend import FakeBackend + from tests.server.training.helpers import setup, request + + controller, state, driver, _ = await setup(backend=FakeBackend()) + controller._metrics = MagicMock() + driver.fail_start = fail_start + try: + ref = await controller.create(request('timed-start')) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == ('failed' if fail_start else 'completed') + call = controller._metrics.startup_seconds.observe.call_args + assert call.args[0] >= 0 + assert call.kwargs['tags'] == {'backend': 'mock', 'mode': 'lora', + 'status': 'error' if fail_start else 'success'} + if not fail_start: + await controller.release(future['result']['training_id'], 'owner', 'session') + call = controller._metrics.shutdown_seconds.observe.call_args + assert call.args[0] >= 0 and call.kwargs['tags']['status'] == 'success' + finally: + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('failed', [False, True]) +async def test_sweep_instruments_one_iteration_and_its_outcome(monkeypatch, failed): + import asyncio + from unittest.mock import AsyncMock, MagicMock + from twinkle.server.model.pool.actor import TrainingControllerActor + + async def stop_after_iteration(delay): + raise asyncio.CancelledError + + actor = SimpleNamespace(core=SimpleNamespace(sweep=AsyncMock( + side_effect=RuntimeError('state unavailable') if failed else None)), + config=SimpleNamespace(sweep_interval_s=5), _metrics=MagicMock()) + monkeypatch.setattr('twinkle.server.model.pool.actor.asyncio.sleep', stop_after_iteration) + with pytest.raises(asyncio.CancelledError): + await TrainingControllerActor.__ray_metadata__.modified_class._sweep(actor) + tags = {'status': 'error' if failed else 'success'} + actor._metrics.sweep_total.inc.assert_called_once_with(tags=tags) + call = actor._metrics.sweep_seconds.observe.call_args + assert call.args[0] >= 0 and call.kwargs['tags'] == tags + + +@pytest.mark.asyncio +@pytest.mark.parametrize('failed', [False, True]) +async def test_driver_records_allocation_attempts(failed): + from unittest.mock import AsyncMock, MagicMock + from twinkle.server.config.training import TrainingConfig + from twinkle.server.resources.allocation import DeviceAllocation + from twinkle.server.model.pool.driver import RayInstanceDriver + + config = TrainingConfig(models=[{'name': 'tiny', 'model_id': 'tiny', 'backend': 'mock', + 'device_type': 'cpu', 'mode': 'full', 'config': {'num_devices': 1, 'device_mesh': {'device_type': 'cpu'}}}]) + store = SimpleNamespace(models=SimpleNamespace(unregister_replica_capacity=AsyncMock())) + driver = RayInstanceDriver(config, store, 'test', 'controller') + driver._metrics = MagicMock() + allocation = DeviceAllocation('owned', 'pg', 'node', ('0',)) + driver.runtime.pool.allocate = AsyncMock(return_value=allocation, + side_effect=RuntimeError('allocation failed') if failed else None) + instance = SimpleNamespace(instance_id='instance', controller_epoch='epoch', application_name='owned', + base_model='tiny', mode='full') + if failed: + with pytest.raises(RuntimeError, match='allocation failed'): + await driver.allocate(instance, config.models[0]) + else: + assert (await driver.allocate(instance, config.models[0]))['devices'] == ['0'] + call = driver._metrics.allocation_seconds.observe.call_args + assert call.args[0] >= 0 + assert call.kwargs['tags'] == {'backend': 'mock', 'mode': 'full', 'device_type': 'cpu', + 'status': 'error' if failed else 'success'} diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py new file mode 100644 index 000000000..03d795da0 --- /dev/null +++ b/tests/server/test_dependencies.py @@ -0,0 +1,115 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Executable dependency architecture for every ``twinkle.server`` component.""" +import grimp +import subprocess +import sys + +import pytest + +_ROOT = 'twinkle.server' +_ROOT_COMPONENT = '__root__' + +# Dependencies must point from a later layer to an earlier layer. Components in +# the same layer are independent and therefore may not import one another. +_LAYERS = ( + frozenset({'exceptions'}), + frozenset({'config', 'middleware', 'model.tinker_datum', 'runtime', 'task_errors', 'utils', 'validation'}), + frozenset({'telemetry'}), + frozenset({'state'}), + frozenset({'resources'}), + frozenset({'deployment', 'lifecycle', 'session_resource'}), + frozenset({'task_queue'}), + frozenset({'checkpoint', 'data_plane'}), + frozenset({'gateway', 'model', 'sampler', 'processor'}), + frozenset({'model.pool', 'sampler.pool'}), + frozenset({'launcher'}), + frozenset({'cli'}), + frozenset({'__main__', _ROOT_COMPONENT}), +) + + +def _component(module: str) -> str: + if module == _ROOT: + return _ROOT_COMPONENT + relative = module.removeprefix(f'{_ROOT}.') + for component in ('model.pool', 'sampler.pool', 'model.tinker_datum'): + if relative == component or relative.startswith(component + '.'): + return component + return relative.split('.', 1)[0] + + +@pytest.mark.parametrize('statement', [ + 'import twinkle.server.exceptions', + 'from twinkle.server.config import ServerConfig', + 'from twinkle.server.state.models import FutureRecord', + 'from twinkle.server.state import SessionRecord', +]) +def test_lightweight_imports_do_not_load_runtime(statement): + code = statement + ''' +import sys +assert not any(name in sys.modules for name in ( + 'torch', 'ray', 'twinkle.server.launcher', 'twinkle.server.model.pool.actor', 'twinkle.hub.hub' +)), sorted(sys.modules) +''' + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=60) + assert result.returncode == 0, result.stderr + + +def test_hub_first_import_does_not_reenter_plugin_loader(): + code = ''' +from twinkle.hub.hub import HFHub +from twinkle.utils.loader import Plugin +assert HFHub.__module__ == 'twinkle.hub.hub' +assert Plugin.__module__ == 'twinkle.utils.loader' +''' + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=120) + assert result.returncode == 0, result.stderr + + +def test_public_launcher_is_the_actual_implementation(): + import twinkle.server + from twinkle.server.launcher.server_launcher import ServerLauncher, launch_server + assert twinkle.server.ServerLauncher is ServerLauncher + assert twinkle.server.launch_server is launch_server + assert twinkle.server.__all__ == ['ServerLauncher', 'launch_server'] + + +def test_server_components_form_declared_dag(): + """Reject cycles, upward imports, same-layer coupling, and unclassified components.""" + graph = grimp.build_graph('twinkle', include_external_packages=False) + server_modules = { + module for module in graph.modules if module == _ROOT or module.startswith(f'{_ROOT}.') + } + observed_components = {_component(module) for module in server_modules} + declared_components = set().union(*_LAYERS) + assert observed_components == declared_components, ( + f'Update the server DAG when adding or removing components: ' + f'undeclared={sorted(observed_components - declared_components)}, ' + f'missing={sorted(declared_components - observed_components)}') + + layer_by_component = { + component: layer_index for layer_index, layer in enumerate(_LAYERS) for component in layer + } + component_edges = { + (_component(importer), _component(imported)) + for importer in server_modules + for imported in graph.find_modules_directly_imported_by(importer) + if imported in server_modules and _component(importer) != _component(imported) + } + violations = sorted( + (importer, imported) for importer, imported in component_edges + if layer_by_component[importer] <= layer_by_component[imported]) + assert not violations, ( + 'Server dependencies must point strictly down the declared DAG; ' + f'invalid edges={violations}') + + +@pytest.mark.parametrize('component', ['model.pool', 'sampler.pool']) +def test_controller_import_does_not_load_ray_composition(component): + code = f'''import twinkle.server.{component}.controller +import sys +assert 'ray' not in sys.modules or not sys.modules['ray'].is_initialized() +assert not any('twinkle.server.{component}.' + name in sys.modules for name in ('actor', 'driver', 'app')) +''' + result = subprocess.run([sys.executable, '-c', code], capture_output=True, text=True, timeout=60) + assert result.returncode == 0, result.stderr diff --git a/tests/server/test_gateway_services.py b/tests/server/test_gateway_services.py index 5a5d032d8..e61dfd631 100644 --- a/tests/server/test_gateway_services.py +++ b/tests/server/test_gateway_services.py @@ -10,6 +10,9 @@ class _State: + async def assert_future_owner(self, request_id, token): + return None + def __init__(self, records): self.records = list(records) @@ -43,7 +46,7 @@ def test_gateway_services_do_not_import_protocol_models(): class _FakeGateway: """Minimal ``GatewayServer`` stand-in; ``poll_future`` is patched so state is unused.""" - state = None + state = _State([]) def _parity_client(monkeypatch, canonical_record): diff --git a/tests/server/gateway/test_proxy.py b/tests/server/test_service_proxy.py similarity index 61% rename from tests/server/gateway/test_proxy.py rename to tests/server/test_service_proxy.py index 9b03e4a98..dc9eb4359 100644 --- a/tests/server/gateway/test_proxy.py +++ b/tests/server/test_service_proxy.py @@ -1,19 +1,17 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Directed tests for ``gateway/proxy.py``. - -Covers the route-URL construction now sourced from ``gateway.routes``, the -``H_MULTIPLEX`` header compatibility in ``_prepare_headers``, and the 502 ``ErrorPayload`` -fallback when the upstream is unreachable. -""" +"""ServiceProxy URL construction, headers, bounded waits, and upstream error mapping.""" from __future__ import annotations +import asyncio import json +import httpx + import pytest from starlette.requests import Request from unittest.mock import AsyncMock -from twinkle.server.gateway.proxy import ServiceProxy +from twinkle.server.lifecycle.proxy import ServiceProxy from twinkle.protocol.headers import H_MULTIPLEX, H_MULTIPLEX_LEGACY, H_REQUEST_ID @@ -51,6 +49,37 @@ def test_prepare_headers_sets_multiplex_from_request_id(): assert 'host' not in {k.lower() for k in headers} +@pytest.mark.asyncio +async def test_proxy_waits_are_bounded_and_timeout_keeps_existing_wire(): + closed = asyncio.Event() + + async def stall(reader, writer): + try: + await reader.read(4096) + await reader.read() + finally: + writer.close() + await writer.wait_closed() + closed.set() + + server = await asyncio.start_server(stall, '127.0.0.1', 0) + port = server.sockets[0].getsockname()[1] + proxy = ServiceProxy(http_options={'host': '127.0.0.1', 'port': port}) + assert all(value is not None and value > 0 for value in proxy.client.timeout.as_dict().values()) + proxy.client.timeout = httpx.Timeout(0.05) + request = _make_request(headers=[(H_REQUEST_ID.encode(), b'req-timeout')]) + try: + reply = await asyncio.wait_for(proxy.proxy_request(request, 'tinker/forward', 'small', 'model', + body_override=b'{}'), 1) + assert reply.status_code == 502 + assert json.loads(reply.body)['request_id'] == 'req-timeout' + finally: + await proxy.close() + server.close() + await server.wait_closed() + await asyncio.wait_for(closed.wait(), 1) + + @pytest.mark.asyncio async def test_proxy_request_502_fallback_returns_error_payload(): proxy = ServiceProxy(http_options={'host': 'localhost', 'port': 8000}, route_prefix='/api/v1') diff --git a/tests/server/training/helpers.py b/tests/server/training/helpers.py new file mode 100644 index 000000000..ba4a5fb9d --- /dev/null +++ b/tests/server/training/helpers.py @@ -0,0 +1,174 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.config.training import TrainingConfig +from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.state import ServerState +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.state.training_records import CreateTraining +from twinkle.server.task_queue import TaskQueueMixin +from twinkle.server.model.pool.controller import TrainingController +from twinkle.server.model.runtime import TrainingRuntime + + +class Clock: + now = 0.0 + + def __call__(self): + return self.now + + +class Runtime: + """Allocates a finite pool and stores actual per-user state for assertions.""" + + def __init__(self, devices=16): + self.devices = devices + self.allocations = {} + self.users = {} + self.starts = 0 + self.stopped = [] + self.start_gate = asyncio.Event() + self.start_gate.set() + self.fail_start = False + self.fail_shutdown = False + + async def allocate(self, instance, spec): + self.starts += 1 + if self.fail_start: + raise RuntimeError('injected startup failure') + used = sum(self.allocations.values()) + if used + spec.config.num_devices > self.devices: + raise CapacityUnavailableError('No devices available') + self.allocations[instance.instance_id] = spec.config.num_devices + return {} + + async def deploy(self, instance, spec): + await self.start_gate.wait() + + async def create_training(self, instance, training, request): + assert instance.instance_id in self.allocations + assert training.training_id not in self.users + self.users[training.training_id] = { + 'instance': instance.instance_id, + 'label': training.adapter_label, + 'steps': 0 + } + + async def release_training(self, instance, training): + self.users.pop(training.training_id, None) + + async def shutdown(self, instance): + if self.fail_shutdown: + raise RuntimeError('cannot confirm process death') + self.allocations.pop(instance.instance_id, None) + self.users = {key: value for key, value in self.users.items() if value['instance'] != instance.instance_id} + self.stopped.append(instance.instance_id) + + async def verify_operation(self, instance, request_id): + return {'state': 'pending'} + + async def recover_orphans(self): + pass + + +async def setup(*, devices=16, min_replicas=0, max_replicas=None, quota=100, backend=None, adapter_config=None): + mode = { + 'num_devices': 2, + 'device_mesh': { + 'fsdp_size': 2 + }, + 'max_loras': 5, + 'idle_timeout_s': 10, + 'min_replicas': min_replicas, + 'max_replicas': max_replicas + } + if adapter_config is not None: + mode['adapter_config'] = adapter_config + config = TrainingConfig(models=[{ + 'name': 'small', + 'model_id': 'ms://small', + 'backend': 'mock', + 'device_type': 'cuda', + 'mode': 'lora', + 'config': mode, + }, { + 'name': 'small-full', + 'model_id': 'ms://small', + 'backend': 'mock', + 'device_type': 'cuda', + 'mode': 'full', + 'config': { + 'num_devices': 2, + 'device_mesh': { + 'fsdp_size': 2 + }, + 'idle_timeout_s': 10 + }, + }]) + state = ServerState(backend=backend or RayActorBackend(), per_token_model_limit=quota) + await state.create_session({'session_id': 'session'}) + runtime, clock = Runtime(devices), Clock() + controller = TrainingController(config, state.create_training_store('unit'), runtime, clock=clock) + await controller.initialize() + return controller, state, runtime, clock + + +def request(key, *, mode='lora', label='same', token='owner'): + return CreateTraining( + base_model='small-full' if mode == 'full' else 'small', + mode=mode, + token=token, + session_id='session', + create_key=str(key), + client_type='twinkle', + payload={}, + adapter_label=label) + + +async def create(controller, state, key, **kwargs): + ref = await controller.create(request(key, **kwargs)) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == 'completed', future + return future['result']['training_id'] + + +class RPC: + + def __init__(self, method): + self.method = method + + def remote(self, *args, **kwargs): + return asyncio.create_task(self.method(*args, **kwargs)) + + +class RuntimeQueue(TrainingRuntime, TaskQueueMixin): + + def __init__(self, controller, state, instance): + self.state, self.replica_id = state, instance.instance_id + from twinkle.server.state.training_records import TrainingRuntimeContext + self.training_context = TrainingRuntimeContext( + instance_id=instance.instance_id, + epoch=controller.epoch, + namespace='test', + registry_name='test', + controller_name='test') + self._model_unhealthy = False + self._training_controller = SimpleNamespace( + **{name: RPC(getattr(controller, name)) + for name in ('admit', 'finish_operation', 'fail_instance')}) + self._init_task_queue( + TaskQueueConfig(execution_timeout=0.1, inline_fast_path_timeout=0.01), enable_admission_gate=True) + + +async def settle(state, request_id): + for _ in range(300): + future = await state.get_future(request_id) + if future and future['status'] in ('completed', 'failed', 'cancelled'): + return future + await asyncio.sleep(0.01) + raise AssertionError('Future did not settle') diff --git a/tests/server/training/test_acceptance_harness.py b/tests/server/training/test_acceptance_harness.py new file mode 100644 index 000000000..58d46e7de --- /dev/null +++ b/tests/server/training/test_acceptance_harness.py @@ -0,0 +1,124 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Acceptance harness budgets and Redis ownership are independently testable.""" +import shutil + +import pytest + +from tests.server.integration.on_demand_training import stress_validation +from tests.server.integration.on_demand_training.managed_redis import ManagedRedis + + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + # These tests deliberately use no Ray runtime. + yield + + +class FakeHandle: + calls = [] + + def __init__(self, base, model, protocol, label): + self.calls.append(('create', protocol)) + + def step(self): + self.calls.append(('step', None)) + + def save(self, name): + return 'checkpoint' + + def resume(self, path): + assert path == 'checkpoint' + + def close(self): + self.calls.append(('close', None)) + + +def test_iteration_budget_does_not_round_up(monkeypatch): + FakeHandle.calls = [] + monkeypatch.setattr(stress_validation, 'TrainingHandle', FakeHandle) + evidence = stress_validation.run_stress('unused', 'unused', concurrency=2, iterations=3) + assert evidence['completed'] == 3 + assert len([call for call in FakeHandle.calls if call[0] == 'close']) == 4 # Includes the warm peer. + + +def test_duration_is_not_truncated_by_iterations(monkeypatch): + now = [0] + + def clock(): + now[0] += 1 + return now[0] + + monkeypatch.setattr(stress_validation.time, 'monotonic', clock) + monkeypatch.setattr(stress_validation, 'TrainingHandle', FakeHandle) + evidence = stress_validation.run_stress('unused', 'unused', concurrency=1, iterations=1, duration=10) + assert evidence['completed'] > 1 and evidence['seconds'] >= 10 + + +def test_failed_training_never_reports_success(monkeypatch): + + class Broken(FakeHandle): + + def step(self): + raise RuntimeError('injected failure') + + monkeypatch.setattr(stress_validation, 'TrainingHandle', Broken) + with pytest.raises(RuntimeError, match='injected failure'): + stress_validation.run_stress('unused', 'unused', iterations=1) + + +@pytest.mark.skipif(not shutil.which('redis-server'), reason='redis-server executable required') +def test_owned_redis_restart_preserves_state(tmp_path): + redis = pytest.importorskip('redis') + managed = ManagedRedis(tmp_path) + managed.start() + client = redis.Redis.from_url(managed.url, socket_timeout=1) + try: + client.set('checkpoint', 'durable') + + def unavailable(): + with pytest.raises(redis.exceptions.ConnectionError): + client.get('checkpoint') + + managed.fault(unavailable) + assert client.get('checkpoint') == b'durable' + finally: + client.close() + managed.stop() + + +def test_training_release_assertion_excludes_but_does_not_remove_sampling(monkeypatch): + import importlib + import ray + from ray import serve + from types import SimpleNamespace + from tests.server.integration.on_demand_training.acceptance_support import assert_released + namespace = 'owned' + sampling_name = 'resources-owned-sampling_pool_epoch-sampler_id' + applications = {sampling_name: object()} + groups = {'sampler': {'name': sampling_name, 'state': 'CREATED'}} + controller = SimpleNamespace(diagnostics=SimpleNamespace(remote=lambda: {'instances': [{'state': 'STOPPED'}]})) + monkeypatch.setattr(ray, 'get', lambda value, **kwargs: value) + monkeypatch.setattr(serve, 'status', lambda: SimpleNamespace(applications=applications)) + monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'placement_group_table', lambda: groups) + assert assert_released(controller, namespace) + assert sampling_name in applications and groups['sampler']['state'] == 'CREATED' + training_name = 'resources-owned-controller_epoch-instance_id' + applications[training_name] = object() + with pytest.raises(AssertionError): + assert_released(controller, namespace) + applications.pop(training_name) + groups['training'] = {'name': training_name, 'state': 'CREATED'} + with pytest.raises(AssertionError): + assert_released(controller, namespace) + + +@pytest.mark.parametrize('values', [[], [0.0, 0.0], [float('nan')], [float('-inf')]]) +def test_checkpoint_probe_rejects_vacuous_comparisons(values): + from tests.server.integration.on_demand_training.acceptance_support import checked_logps + with pytest.raises(AssertionError): + checked_logps(values) + + +def test_checkpoint_probe_accepts_variable_length_sequences(): + from tests.server.integration.on_demand_training.acceptance_support import checked_logps + assert checked_logps([[-0.25, -0.5], [-0.75]]).tolist() == [-0.25, -0.5, -0.75] diff --git a/tests/server/training/test_actor_initialization.py b/tests/server/training/test_actor_initialization.py new file mode 100644 index 000000000..dd9a3f530 --- /dev/null +++ b/tests/server/training/test_actor_initialization.py @@ -0,0 +1,129 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Initialization failures and backoff during automatic controller restarts.""" +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from twinkle.server.resources.instances import InstanceLifecycle +from twinkle.server.model.pool.actor import TrainingControllerActor + + +@pytest.mark.asyncio +async def test_failed_initialization_restarts_only_after_backoff(monkeypatch): + cls = TrainingControllerActor.__ray_metadata__.modified_class + gate, sleeping = asyncio.Event(), asyncio.Event() + delays = [] + handle = object() + kill = Mock() + get_actor = Mock(return_value=handle) + original_sleep = asyncio.sleep + + async def sleep(delay): + delays.append(delay) + sleeping.set() + await gate.wait() + + monkeypatch.setattr('twinkle.server.model.pool.actor.asyncio.sleep', sleep) + monkeypatch.setattr('twinkle.server.model.pool.actor.ray.get_actor', get_actor) + monkeypatch.setattr('twinkle.server.model.pool.actor.ray.kill', kill) + actor = SimpleNamespace(core=SimpleNamespace(initialize=AsyncMock(side_effect=RuntimeError('allocation failed'))), + config=SimpleNamespace(sweep_interval_s=7), namespace='initialization-test', _monitors=[]) + actor._restart_after_failure = lambda: cls._restart_after_failure(actor) + with pytest.raises(RuntimeError, match='allocation failed'): + await cls._initialize(actor) + await sleeping.wait() + assert delays == [7] and not actor._monitors + kill.assert_not_called() + gate.set() + await actor._restart_task + await original_sleep(0) + get_actor.assert_called_once_with('training-controller-initialization-test', namespace='initialization-test') + kill.assert_called_once_with(handle, no_restart=False) + + +@pytest.mark.asyncio +async def test_shutdown_cancels_pending_restart(): + cls = TrainingControllerActor.__ray_metadata__.modified_class + wait = asyncio.Event() + restart = asyncio.create_task(wait.wait()) + initialization = asyncio.get_running_loop().create_future() + initialization.set_result(None) + actor = SimpleNamespace(driver=SimpleNamespace(runtime=SimpleNamespace(wait_for_cleanup=AsyncMock())), + core=SimpleNamespace(shutdown=AsyncMock(), wait_for_jobs=AsyncMock()), + _monitors=[], _restart_task=restart, _initialization=initialization, _failures={}) + await cls.shutdown(actor) + assert restart.cancelled() + actor.core.shutdown.assert_awaited_once() + actor.driver.runtime.wait_for_cleanup.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_heartbeat_keeps_healthy_peer_alive_while_probe_and_close_block(): + cls = TrainingControllerActor.__ray_metadata__.modified_class + good, bad = [SimpleNamespace(instance_id=name, state='READY') for name in ('good', 'bad')] + probe_gate, close_gate = asyncio.Event(), asyncio.Event() + healthy_seen, close_started = asyncio.Event(), asyncio.Event() + + async def health(instance): + if instance is bad: + await probe_gate.wait() + return {'unhealthy': instance is bad} + + async def fail(*args): + close_started.set() + await close_gate.wait() + + async def touch(identity): + if identity == 'good': + healthy_seen.set() + + actor = SimpleNamespace( + store=SimpleNamespace(backend=SimpleNamespace(get=AsyncMock(return_value='active')), + key=lambda *args: 'active', models=SimpleNamespace(touch_replica_last_seen=touch)), + core=SimpleNamespace(epoch='controller', instances={'bad': bad, 'good': good}, fail_instance=fail, + lifecycle=InstanceLifecycle('Training')), + driver=SimpleNamespace(health=health), _unreachable_since={}, _failures={}) + actor._check_instance_health = lambda instance: cls._check_instance_health(actor, instance) + actor._fail_unhealthy = lambda instance, message: cls._fail_unhealthy(actor, instance, message) + monitor = asyncio.create_task(cls._heartbeat(actor)) + try: + await asyncio.wait_for(healthy_seen.wait(), 1) + assert not probe_gate.is_set() + probe_gate.set() + await asyncio.wait_for(close_started.wait(), 1) + # Closing a bad peer does not occupy the heartbeat monitor. + healthy_seen.clear() + await cls._check_instance_health(actor, good) + assert healthy_seen.is_set() + cls._fail_unhealthy(actor, bad, 'duplicate') + assert len(actor.core.lifecycle.named_jobs) == 1 + finally: + probe_gate.set() + close_gate.set() + monitor.cancel() + await asyncio.gather(monitor, *actor.core.lifecycle.named_jobs.values(), return_exceptions=True) + assert not actor.core.lifecycle.named_jobs + + +@pytest.mark.asyncio +async def test_shutdown_cancels_and_joins_initialization_before_closing_core(): + cls = TrainingControllerActor.__ray_metadata__.modified_class + started, exited = asyncio.Event(), asyncio.Event() + + async def initialize(): + started.set() + try: + await asyncio.Event().wait() + finally: + exited.set() + + initialization = asyncio.create_task(initialize()) + await started.wait() + actor = SimpleNamespace(driver=SimpleNamespace(runtime=SimpleNamespace(wait_for_cleanup=AsyncMock())), + core=SimpleNamespace(shutdown=AsyncMock(), wait_for_jobs=AsyncMock()), + _initialization=initialization, _monitors=[], _failures={}) + await cls.shutdown(actor) + assert initialization.cancelled() and exited.is_set() + actor.core.shutdown.assert_awaited_once() diff --git a/tests/server/training/test_controller.py b/tests/server/training/test_controller.py new file mode 100644 index 000000000..f4c3babad --- /dev/null +++ b/tests/server/training/test_controller.py @@ -0,0 +1,799 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Exercise real controller/storage decisions against a controllable runtime.""" +from __future__ import annotations +from tests.server.training.helpers import Clock, Runtime, setup, request, create + +import asyncio + +import pytest + +from twinkle.server.config.training import TrainingConfig +from twinkle.server.exceptions import RequestRejectedError, ResourceNotFoundError +from twinkle.server.state import ServerState +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.model.pool.controller import TrainingController +from twinkle.server.exceptions import TrainingExpiredError +from twinkle.server.resources.allocation import CapacityUnavailableError +from twinkle.server.state.training_records import CreateTraining + + +@pytest.mark.asyncio +async def test_initialization_timeout_only_retires_the_failing_tenant(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend()) + active = await create(controller, state, 'active') + runtime.users[active]['steps'] = 7 + instance = controller.instances[controller.trainings[active].instance_id] + async def timeout(instance, training, request): + raise TimeoutError('User initialization timed out') + runtime.create_training = timeout + ref = await controller.create(request('failing')) + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['failure']['reason_code'] == 'startup_timeout' + assert controller.trainings[active].state == 'ACTIVE' + assert runtime.users[active]['steps'] == 7 + assert instance.state == 'READY' and instance.training_ids == [active] + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_late_initializer_keeps_its_slot_reserved_until_serial_cleanup(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend(), max_replicas=1) + active = await create(controller, state, 'active') + raw = controller.config.model_dump() + raw['startup_timeout_s'] = 0.03 + raw['models'][0]['config']['max_loras'] = 2 + controller.config = TrainingConfig.model_validate(raw) + entered, unblock = asyncio.Event(), asyncio.Event() + original_create, original_release = runtime.create_training, runtime.release_training + initializer = None + async def initialize(instance, training, req): + nonlocal initializer + async def late(): + entered.set() + await unblock.wait() + await original_create(instance, training, req) + initializer = asyncio.create_task(late()) + await asyncio.shield(initializer) + async def release(instance, training): + await initializer + await original_release(instance, training) + runtime.create_training, runtime.release_training = initialize, release + ref = await controller.create(request('late')) + await entered.wait() + for _ in range(100): + future = await state.get_future(ref['request_id']) + if future['status'] == 'failed': + break + await asyncio.sleep(0.005) + assert future['failure']['reason_code'] == 'startup_timeout' + instance = controller.instances[controller.trainings[active].instance_id] + assert len(instance.training_ids) == 2 and controller.trainings[active].state == 'ACTIVE' + refused = await controller.create(request('no-free-slot')) + assert (await state.get_future(refused['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + unblock.set() + await controller.wait_for_jobs() + assert instance.training_ids == [active] and set(runtime.users) == {active} + runtime.create_training, runtime.release_training = original_create, original_release + replacement = await create(controller, state, 'replacement') + assert controller.trainings[replacement].instance_id == instance.instance_id + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_user_gets_full_initialization_window_after_base_model_startup(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + runtime.start_gate.clear() + ref = await controller.create(request('near-startup-deadline')) + instance = next(iter(controller.instances.values())) + clock.now = instance.startup_deadline - 1 + original = runtime.create_training + async def initialize(instance, training, req): + clock.now += 2 + await original(instance, training, req) + runtime.create_training = initialize + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'completed' + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_resident_rebuild_does_not_block_expiry_or_another_mode(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + expired = await create(controller, state, 'expired-full', mode='full') + raw = controller.config.model_dump() + raw['models'][0]['config']['min_replicas'] = 1 + controller.config = TrainingConfig.model_validate(raw) + runtime.start_gate.clear() + await asyncio.wait_for(controller.sweep(), 0.1) + for _ in range(20): + await asyncio.sleep(0) + if runtime.starts == 2: + break + clock.now = 11 + await asyncio.wait_for(controller.sweep(), 0.1) + for _ in range(20): + await asyncio.sleep(0) + if controller.trainings[expired].state == 'EXPIRED': + break + assert controller.trainings[expired].state == 'EXPIRED' + ref = await controller.create(request('unrelated-full', mode='full')) + assert (await state.get_future(ref['request_id']))['status'] == 'pending' + resident_ids = [item.instance_id for item in controller.instances.values() if item.resident] + assert len(resident_ids) == 1 + await controller.sweep() + assert [item.instance_id for item in controller.instances.values() if item.resident] == resident_ids + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'completed' + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_reconciliation_of_committed_completion_does_not_extend_idle_time(): + import time + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'active') + await controller.admit(identity, 'owner', 'session', 'lost-finish') + async def finished(instance, request_id): + return {'state': 'finished', 'executed': True, 'finished_at': time.time() - 31} + runtime.verify_operation = finished + clock.now = 31 + await controller.sweep() + await controller.wait_for_jobs() + assert not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].last_activity_at < 0.01 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_disabled_max_lifetime_and_live_heartbeat_do_not_override_business_idle(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup( + backend=FakeBackend(), adapter_config={'adapter_max_lifetime': None}) + identity = await create(controller, state, 'idle-only') + controller.trainings[identity].created_at = '2000-01-01T00:00:00+00:00' + clock.now = 9 + await state.touch_session('session') + await controller.sweep() + assert controller.trainings[identity].state == 'ACTIVE' + clock.now = 10 + await state.touch_session('session') + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[identity].state == 'EXPIRED' + assert not runtime.allocations + + +@pytest.mark.asyncio +@pytest.mark.parametrize('reason', ['heartbeat', 'max_lifetime']) +@pytest.mark.parametrize('busy', [False, True]) +async def test_adapter_lifecycle_expires_independently_and_waits_for_running_work(reason, busy): + import time + from datetime import datetime, timedelta, timezone + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup( + backend=FakeBackend(), adapter_config={'adapter_timeout': 30, 'adapter_max_lifetime': 60}) + identity = await create(controller, state, 'lifecycle') + if busy: + await controller.admit(identity, 'owner', 'session', 'running') + if reason == 'heartbeat': + session = await state.sessions.get('session') + session.last_heartbeat = time.time() - 31 + await state.sessions.add('session', session) + else: + controller.trainings[identity].created_at = (datetime.now(timezone.utc) - timedelta(seconds=61)).isoformat() + await controller.sweep() + await controller.wait_for_jobs() + if busy: + assert controller.trainings[identity].state == 'ACTIVE' + assert identity in runtime.users + if reason == 'heartbeat': + with pytest.raises(TrainingExpiredError): + await controller.admit(identity, 'owner', 'session', 'new-operation') + await controller.finish_operation(identity, 'running', controller.epoch, executed=True) + await controller.sweep() + await controller.wait_for_jobs() + assert controller.trainings[identity].state == 'EXPIRED' + assert identity not in runtime.users + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_max_lifetime_stops_new_work_while_admitted_work_finishes(): + from datetime import datetime, timedelta, timezone + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup( + backend=FakeBackend(), adapter_config={'adapter_max_lifetime': 60}) + identity = await create(controller, state, 'lifetime-admission') + await controller.admit(identity, 'owner', 'session', 'running') + training = controller.trainings[identity] + training.created_at = (datetime.now(timezone.utc) - timedelta(seconds=61)).isoformat() + with pytest.raises(TrainingExpiredError): + await controller.admit(identity, 'owner', 'session', 'new-operation') + assert set(training.pending_operation_ids) == {'running'} + assert identity in runtime.users + await controller.finish_operation(identity, 'running', controller.epoch, executed=True) + await controller.wait_for_jobs() + assert training.state == 'EXPIRED' + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_shutdown_rejects_unconfirmed_release_and_allows_retry(): + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup(backend=FakeBackend()) + await create(controller, state, 'shutdown-ack') + runtime.fail_shutdown = True + try: + with pytest.raises(RuntimeError, match='shutdown.*unconfirmed'): + await controller.shutdown() + assert runtime.allocations + assert all(instance.state == 'QUARANTINED' for instance in controller.instances.values()) + finally: + runtime.fail_shutdown = False + await controller.shutdown() + assert not runtime.allocations + assert all(instance.state == 'STOPPED' for instance in controller.instances.values()) + + +@pytest.mark.asyncio +async def test_shutdown_attempts_healthy_peer_after_an_instance_metadata_failure(): + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup(backend=FakeBackend()) + first = await create(controller, state, 'shutdown-first') + peer = await create(controller, state, 'shutdown-peer', mode='full') + first_instance = controller.instances[controller.trainings[first].instance_id] + peer_instance = controller.instances[controller.trainings[peer].instance_id] + save = controller.store.save_instance + + async def unavailable(record): + if record.instance_id == first_instance.instance_id and record.state == 'DRAINING': + raise RuntimeError('Instance metadata unavailable') + await save(record) + + controller.store.save_instance = unavailable + try: + with pytest.raises(RuntimeError, match='Instance metadata unavailable'): + await controller.shutdown() + assert peer_instance.state == 'STOPPED' + assert peer_instance.instance_id not in runtime.allocations + assert first_instance.instance_id in runtime.allocations + with pytest.raises(ResourceNotFoundError): + await controller.route(first, 'owner') + finally: + controller.store.save_instance = save + await controller.shutdown() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_lazy_queries_allocate_nothing(): + controller, state, runtime, _ = await setup() + assert runtime.starts == 0 + assert controller.config.get('small-full', 'full').capacity == 1 + with pytest.raises(ResourceNotFoundError): + await controller.route('unknown', 'owner') + assert not runtime.allocations + assert await state.get_capacity_info() == {'max_loras': 0, 'used_loras': 0, 'free_loras': 0} + + +@pytest.mark.asyncio +async def test_32_concurrent_creates_coalesce_even_before_allocation_ready(): + controller, state, runtime, _ = await setup() + runtime.start_gate.clear() + refs = await asyncio.gather(*(controller.create(request(index)) for index in range(32))) + assert len(controller.instances) == 7 + assert sorted(len(instance.training_ids) for instance in controller.instances.values()) == [2, 5, 5, 5, 5, 5, 5] + runtime.start_gate.set() + await controller.wait_for_jobs() + records = [await state.get_future(ref['request_id']) for ref in refs] + assert all(record['status'] == 'completed' for record in records) + assert len({record['result']['training_id'] for record in records}) == 32 + assert runtime.starts == 7 + assert len(runtime.users) == 32 + + +@pytest.mark.asyncio +async def test_full_and_lora_isolation_and_device_reuse(): + controller, state, runtime, _ = await setup(devices=6) + shared = await create(controller, state, 1) + full1 = await create(controller, state, 2, mode='full') + full2 = await create(controller, state, 3, mode='full') + identities = {runtime.users[key]['instance'] for key in (shared, full1, full2)} + assert len(identities) == 3 + old_full = runtime.users[full1]['instance'] + await controller.release(full1, 'owner') + replacement = await create(controller, state, 4, mode='full') + assert runtime.users[replacement]['instance'] != old_full + assert shared in runtime.users and full2 in runtime.users + assert len(runtime.allocations) == 3 + + +@pytest.mark.asyncio +async def test_capacity_failure_is_terminal_and_does_not_hold_quota(): + controller, state, runtime, _ = await setup(devices=2) + await create(controller, state, 1, mode='full') + ref = await controller.create(request(2, mode='full')) + await controller.wait_for_jobs() + record = await state.get_future(ref['request_id']) + assert record['status'] == 'failed' + assert record['failure']['reason_code'] == 'capacity_unavailable' + assert await state._backend.get('token_count::owner') == 1 + assert len(runtime.allocations) == 1 + assert await controller.create(request(2, mode='full')) == ref + + +@pytest.mark.asyncio +async def test_create_idempotency_and_changed_payload_conflict(): + controller, state, runtime, _ = await setup() + refs = await asyncio.gather(*(controller.create(request('one')) for _ in range(16))) + assert len({ref['request_id'] for ref in refs}) == 1 + await controller.wait_for_jobs() + assert len(runtime.users) == 1 + changed = request('one') + changed.payload = {'different': True} + with pytest.raises(RequestRejectedError, match='different parameters'): + await controller.create(changed) + + +@pytest.mark.asyncio +async def test_pending_creates_count_toward_quota(): + controller, state, runtime, _ = await setup(quota=1) + runtime.start_gate.clear() + first = await controller.create(request(1)) + second = await controller.create(request(2)) + failure = await state.get_future(second['request_id']) + assert failure['failure']['reason_code'] == 'resource_quota_exceeded' + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(first['request_id']))['status'] == 'completed' + assert len(runtime.users) == 1 + + +@pytest.mark.asyncio +async def test_heartbeat_does_not_prevent_idle_expiry_or_affect_busy_user(): + controller, state, runtime, clock = await setup() + idle = await create(controller, state, 1) + busy = await create(controller, state, 2) + await controller.admit(busy, 'owner', 'session', 'op') + clock.now = 11 + await state.touch_session('session') + await controller.sweep() + await controller.wait_for_jobs() + assert idle not in runtime.users and busy in runtime.users + with pytest.raises(TrainingExpiredError): + await controller.route(idle, 'owner') + await controller.finish_operation(busy, 'op', controller.epoch, executed=True) + assert controller.trainings[busy].last_activity_at == 11 + clock.now = 20 + await controller.sweep() + await controller.wait_for_jobs() + assert busy in runtime.users + clock.now = 22 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_release_drains_accepted_operation_before_cleanup(): + controller, state, runtime, _ = await setup() + identity = await create(controller, state, 1) + await controller.admit(identity, 'owner', 'session', 'op') + releasing = asyncio.create_task(controller.release(identity, 'owner')) + while controller.trainings[identity].state == 'ACTIVE': + await asyncio.sleep(0) + assert identity in runtime.users + with pytest.raises(ResourceNotFoundError): + await controller.admit(identity, 'owner', 'session', 'next') + await controller.finish_operation(identity, 'op', controller.epoch, executed=True) + await controller.finish_operation(identity, 'op', controller.epoch, executed=True) + await controller.wait_for_jobs() + await releasing + assert not runtime.users and not runtime.allocations + assert len(runtime.stopped) == 1 + + +@pytest.mark.asyncio +async def test_same_label_recreation_and_ownership_do_not_reattach_old_handle(): + controller, state, runtime, _ = await setup() + old = await create(controller, state, 1) + with pytest.raises(ResourceNotFoundError): + await controller.route(old, 'other') + with pytest.raises(ResourceNotFoundError): + await controller.route(old, 'owner', 'other-session') + await controller.release(old, 'owner') + new = await create(controller, state, 2) + assert new != old and runtime.users[new]['label'] == 'same' + with pytest.raises(ResourceNotFoundError): + await controller.route(old, 'owner') + + +@pytest.mark.asyncio +async def test_resident_survives_empty_but_extra_replica_does_not(): + controller, state, runtime, _ = await setup(min_replicas=1) + refs = [await controller.create(request(index)) for index in range(6)] + await controller.wait_for_jobs() + identities = [(await state.get_future(ref['request_id']))['result']['training_id'] for ref in refs] + for identity in identities: + await controller.release(identity, 'owner') + assert len(runtime.allocations) == 1 + resident = next(instance for instance in controller.instances.values() if instance.resident) + assert resident.state == 'READY' + + +@pytest.mark.asyncio +async def test_startup_failure_cleans_all_waiters_and_quota(): + controller, state, runtime, _ = await setup() + runtime.fail_start = True + refs = await asyncio.gather(*(controller.create(request(index)) for index in range(5))) + await controller.wait_for_jobs() + records = [await state.get_future(ref['request_id']) for ref in refs] + assert all(record['status'] == 'failed' for record in records) + assert (await state._backend.get('token_count::owner') or 0) == 0 + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_managed_records_are_not_deleted_by_generic_training_ttl(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + await state._model_mgr.cleanup_expired(10**12, ['session']) + assert await state.get_model_metadata(identity) is not None + assert identity in runtime.users + + +@pytest.mark.asyncio +async def test_restart_invalidates_handles_and_cleans_old_allocation(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + replacement = TrainingController(controller.config, state.create_training_store('unit'), runtime, clock=clock) + await replacement.initialize() + assert replacement.epoch != controller.epoch + assert not runtime.allocations and not runtime.users + assert await state.get_model_metadata(identity) is None + assert (await state._backend.get('token_count::owner') or 0) == 0 + with pytest.raises(ResourceNotFoundError): + await replacement.route(identity, 'owner') + + +@pytest.mark.asyncio +async def test_cancelled_create_never_initializes_user_and_keeps_receipt(): + controller, state, runtime, _ = await setup() + runtime.start_gate.clear() + ref = await controller.create(request('cancel')) + assert (await state.cancel_future(ref['request_id']))['cancelled'] + runtime.start_gate.set() + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'cancelled' + assert not runtime.users and not runtime.allocations + assert await controller.create(request('cancel')) == ref + + +@pytest.mark.asyncio +async def test_unconfirmed_worker_shutdown_holds_devices_and_quota_until_retry(): + from twinkle.server.state.models import FutureFailureRecord + controller, state, runtime, _ = await setup(devices=2) + identity = await create(controller, state, 1, mode='full') + instance = controller.instances[controller.trainings[identity].instance_id] + await controller.admit(identity, 'owner', 'session', 'running') + runtime.fail_shutdown = True + await controller.fail_instance(instance.instance_id, FutureFailureRecord( + reason_code='execution_timeout', message='RPC timed out', attribution='server')) + assert instance.state == 'QUARANTINED' + assert runtime.allocations and controller.trainings[identity].pending_operation_ids + assert await state.get_model_metadata(identity) is not None + failed = await controller.create(request(2, mode='full')) + await controller.wait_for_jobs() + assert (await state.get_future(failed['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + runtime.fail_shutdown = False + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + assert instance.state == 'STOPPED' + assert (await state._backend.get('token_count::owner') or 0) == 0 + + +@pytest.mark.asyncio +async def test_uncertain_enqueue_is_failed_before_ticket_is_released(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + await controller.admit(identity, 'owner', 'session', 'lost-submit') + await controller.store.future('lost-submit', 'pending', identity) + async def missing(instance, request_id): + return {'state': 'missing'} + runtime.verify_operation = missing + clock.now = 31 + await controller.sweep() + await controller.wait_for_jobs() + assert (await state.get_future('lost-submit'))['status'] == 'failed' + assert not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].last_activity_at == 0 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_restart_fails_pending_operations_and_fences_old_controller(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + await controller.admit(identity, 'owner', 'session', 'queued') + await controller.store.future('queued', 'queued', identity) + replacement = TrainingController(controller.config, state.create_training_store('unit'), runtime, clock=clock) + await replacement.initialize() + assert (await state.get_future('queued'))['failure']['reason_code'] == 'orphaned_replica' + with pytest.raises(ResourceNotFoundError, match='epoch'): + await controller.create(request(2)) + with pytest.raises(ResourceNotFoundError, match='epoch'): + await controller.route(identity, 'owner') + assert await replacement.create(request(1)) == {'request_id': controller.trainings[identity].create_request_id} + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_info_does_not_renew_and_expiry_wins_before_next_sweep(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + clock.now = 9 + await controller.admit(identity, 'owner', 'session', 'info', renew=False) + await controller.finish_operation(identity, 'info', controller.epoch, executed=True) + assert controller.trainings[identity].last_activity_at == 0 + clock.now = 10 + with pytest.raises(TrainingExpiredError): + await controller.admit(identity, 'owner', 'session', 'next') + await controller.wait_for_jobs() + assert not runtime.users + + +@pytest.mark.asyncio +async def test_full_resident_cannot_accept_user_while_resetting(): + controller, state, runtime, _ = await setup() + original = controller.config.model_dump() + original['models'][1]['config']['min_replicas'] = 1 + original['models'][1]['config']['max_replicas'] = 1 + controller.config = TrainingConfig.model_validate(original) + await controller.sweep() + await controller.wait_for_jobs() + identity = await create(controller, state, 1, mode='full') + entered, resume = asyncio.Event(), asyncio.Event() + previous = runtime.release_training + async def resetting(instance, training): + entered.set() + await resume.wait() + await previous(instance, training) + runtime.release_training = resetting + releasing = asyncio.create_task(controller.release(identity, 'owner')) + await entered.wait() + instance = controller.instances[controller.trainings[identity].instance_id] + assert instance.state == 'RESETTING' + ref = await controller.create(request(2, mode='full')) + assert (await state.get_future(ref['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + resume.set() + await releasing + assert instance.state == 'READY' + next_id = await create(controller, state, 3, mode='full') + assert controller.trainings[next_id].instance_id == instance.instance_id + + +@pytest.mark.asyncio +async def test_dirty_resident_is_rebuilt_after_last_user_leaves(): + controller, state, runtime, _ = await setup(min_replicas=1) + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + instance.dirty = True + await controller.release(identity, 'owner') + assert instance.state == 'STOPPED' + await controller.sweep() + await controller.wait_for_jobs() + ready = [item for item in controller.instances.values() if item.state == 'READY'] + assert len(ready) == 1 and ready[0].instance_id != instance.instance_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize('late_phase', ['runtime', 'user']) +async def test_late_startup_result_cannot_publish_success(late_phase): + controller, state, runtime, clock = await setup() + if late_phase == 'runtime': + original = runtime.deploy + async def late(instance, spec): + result = await original(instance, spec) + clock.now = instance.startup_deadline + 1 + return result + runtime.deploy = late + else: + original = runtime.create_training + async def late(instance, training, request): + await original(instance, training, request) + clock.now = training.creation_deadline + 1 + runtime.create_training = late + ref = await controller.create(request('late')) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == 'failed' and future['failure']['reason_code'] == 'startup_timeout' + assert not runtime.allocations and not runtime.users + assert all(instance.state == 'STOPPED' for instance in controller.instances.values()) + + +@pytest.mark.asyncio +async def test_failed_resident_startup_rolls_back_previously_loaded_residents(): + controller, state, runtime, clock = await setup(devices=2) + raw = controller.config.model_dump() + raw['models'][0]['config']['min_replicas'] = 2 + replacement = TrainingController(TrainingConfig.model_validate(raw), state.create_training_store('unit'), runtime, clock=clock) + with pytest.raises(CapacityUnavailableError): + await replacement.initialize() + assert not replacement._ready and not runtime.allocations + assert all(item.state == 'STOPPED' for item in replacement.instances.values()) + await replacement.wait_for_jobs() + + +@pytest.mark.asyncio +async def test_retired_diagnostics_are_bounded_without_discarding_live_training(): + controller, state, runtime, clock = await setup() + old = await create(controller, state, 'old') + await controller.release(old, 'owner', 'session') + retired_instance = controller.trainings[old].instance_id + live = await create(controller, state, 'live') + controller.trainings[old].terminal_at -= state.expiration_timeout + 1 + controller.instances[retired_instance].stopped_at -= state.expiration_timeout + 1 + await controller.sweep() + await controller.wait_for_jobs() + assert old not in controller.trainings and retired_instance not in controller.instances + assert retired_instance not in controller.lifecycle.startups + assert (await controller.route(live, 'owner', 'session'))['instance_id'] in runtime.allocations + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('startup_blocked', [False, True]) +async def test_creation_record_write_failure_closes_new_empty_instance(monkeypatch, startup_blocked): + controller, state, runtime, clock = await setup() + if startup_blocked: + runtime.start_gate.clear() + original = state._backend.set + failed = False + async def failing_set(key, value, *args, **kwargs): + nonlocal failed + if (not failed and key.startswith('model::') and isinstance(value, dict) + and value.get('instance_id') and value.get('state') == 'PENDING'): + failed = True + raise RuntimeError('injected training record write failure') + return await original(key, value, *args, **kwargs) + monkeypatch.setattr(state._backend, 'set', failing_set) + ref = await controller.create(request('write-failure')) + await controller.wait_for_jobs() + assert (await state.get_future(ref['request_id']))['status'] == 'failed' + assert not runtime.allocations + assert await state._backend.get('token_count::owner') == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault', ['cancel_after_read', 'unknown_start_ack']) +async def test_create_requires_its_own_atomic_start_ack(fault, monkeypatch): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend()) + initialized = [] + original_create = runtime.create_training + + async def initialize(*args): + initialized.append(True) + await original_create(*args) + + runtime.create_training = initialize + if fault == 'cancel_after_read': + save = controller.store.save_training + + async def cancel_during_save(record): + await save(record) + instance = controller.instances.get(record.instance_id) + if record.state == 'PENDING' and instance is not None and instance.state == 'READY': + assert (await state.cancel_future(record.create_request_id))['cancelled'] + + monkeypatch.setattr(controller.store, 'save_training', cancel_during_save) + else: + start = state.futures.try_start + + async def lose_ack(request_id): + assert await start(request_id) + raise ConnectionError('start acknowledgement lost') + + monkeypatch.setattr(state.futures, 'try_start', lose_ack) + try: + ref = await controller.create(request(fault)) + await controller.wait_for_jobs() + future = await state.get_future(ref['request_id']) + assert future['status'] == ('cancelled' if fault == 'cancel_after_read' else 'failed') + assert not initialized and not runtime.users and not runtime.allocations + assert (await state._backend.get('token_count::owner') or 0) == 0 + finally: + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_quarantined_resident_counts_towards_maximum_until_death_confirmed(): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.state.models import FutureFailureRecord + controller, _, runtime, _ = await setup(backend=FakeBackend(), min_replicas=1, max_replicas=1) + old = next(iter(controller.instances.values())) + runtime.fail_shutdown = True + try: + await controller.fail_instance(old.instance_id, FutureFailureRecord( + reason_code='orphaned_replica', attribution='server', message='worker lost')) + for _ in range(3): + await controller.sweep() + await controller.wait_for_jobs() + assert old.state == 'QUARANTINED' + assert runtime.starts == 1 and list(runtime.allocations) == [old.instance_id] + runtime.fail_shutdown = False + await controller.sweep() + await controller.wait_for_jobs() + await controller.sweep() + await controller.wait_for_jobs() + live = [instance for instance in controller.instances.values() if instance.state != 'STOPPED'] + assert old.state == 'STOPPED' and len(live) == 1 + assert live[0].resident and live[0].state == 'READY' + assert live[0].instance_id != old.instance_id + finally: + runtime.fail_shutdown = False + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_last_lora_release_does_not_report_unconfirmed_device_release(): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'last-user') + instance = controller.instances[controller.trainings[identity].instance_id] + runtime.fail_shutdown = True + try: + with pytest.raises(RuntimeError, match='quarantined'): + await controller.release(identity, 'owner', 'session') + assert instance.state == 'QUARANTINED' and runtime.allocations + runtime.fail_shutdown = False + await controller.sweep() + await controller.wait_for_jobs() + assert instance.state == 'STOPPED' and not runtime.allocations + finally: + runtime.fail_shutdown = False + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_two_device_lora_capacity_shares_two_instances_and_releases_them(): + from tests.server.state.fake_backend import FakeBackend + from twinkle.server.config import ServerConfig + controller, state, runtime, _ = await setup(backend=FakeBackend(), devices=4, quota=3) + controller.config = ServerConfig.from_yaml('cookbook/client/server/megatron/server_config.yaml').training_config + name = controller.config.models[0].name + try: + requests = [request(index, token=f'owner-{index // 3}').model_copy(update={'base_model': name}) + for index in range(6)] + refs = await asyncio.gather(*(controller.create(incoming) for incoming in requests)) + await controller.wait_for_jobs() + futures = [await state.get_future(ref['request_id']) for ref in refs] + assert all(future['status'] == 'completed' for future in futures) + assert len(controller.instances) == 2 + assert all(len(instance.training_ids) == 3 for instance in controller.instances.values()) + assert sum(runtime.allocations.values()) == 4 and set(runtime.allocations.values()) == {2} + excess = request('excess', token='new-owner').model_copy(update={'base_model': name}) + refused = await controller.create(excess) + assert (await state.get_future(refused['request_id']))['failure']['reason_code'] == 'capacity_unavailable' + for incoming, future in zip(requests, futures): + await controller.release(future['result']['training_id'], incoming.token, incoming.session_id) + assert not runtime.allocations and not runtime.users + assert all(instance.state == 'STOPPED' for instance in controller.instances.values()) + finally: + await controller.shutdown() diff --git a/tests/server/training/test_failure_recovery.py b/tests/server/training/test_failure_recovery.py new file mode 100644 index 000000000..115a3c8fc --- /dev/null +++ b/tests/server/training/test_failure_recovery.py @@ -0,0 +1,408 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Fault injection ensures persistence and notification failures do not leak resources.""" +import asyncio +import logging +from unittest.mock import AsyncMock + +import pytest + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.helpers import create, request, setup +from tests.server.training.helpers import RPC, RuntimeQueue, settle +from twinkle.protocol.types.errors import ErrorCategory +from twinkle.server.exceptions import ResourceExhaustedError, ResourceNotFoundError +from twinkle.server.state.models import FutureFailureRecord + + +async def wait_idle(queue): + worker = queue._compute_worker + for _ in range(100): + if not worker._request_ids and not worker.new_task_event.is_set(): + await asyncio.sleep(0) + return + await asyncio.sleep(0.001) + raise AssertionError('Worker did not enter idle wait') + + +@pytest.mark.asyncio +async def test_warm_idle_worker_wakes_after_enqueue_failure(monkeypatch): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'queue-failure') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + executed = [] + + async def operation(): + executed.append(True) + return {'ok': True} + + try: + warmup = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step') + await settle(state, warmup['request_id']) + await wait_idle(queue) + executed.clear() + previous = state.futures.store_status + + async def broken(request_id, status, *args, **kwargs): + if status == 'queued': + raise RuntimeError('injected queued write failure') + return await previous(request_id, status, *args, **kwargs) + + monkeypatch.setattr(state.futures, 'store_status', broken) + with pytest.raises(RuntimeError, match='queued write'): + await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step', + request_id='queued-failure') + ref = {'request_id': 'queued-failure'} + assert (await settle(state, ref['request_id']))['status'] == 'completed' + await wait_idle(queue) + assert executed == [True] + assert not controller.trainings[identity].pending_operation_ids + assert queue._compute_worker.total_queued() == 0 + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault,mode', [ + ('retire', 'lora'), ('save_draining', 'full'), ('save_stopped', 'full'), + ('close_retire', 'full'), ('post_retire_save', 'lora') +]) +async def test_cleanup_storage_failure_is_retryable(fault, mode, monkeypatch): + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, fault, mode=mode) + training = controller.trainings[identity] + instance = controller.instances[training.instance_id] + previous_retire = controller.store.retire + previous_save = controller.store.save_instance + failed = False + + async def retire(record): + nonlocal failed + if fault in ('retire', 'close_retire') and not failed: + failed = True + raise RuntimeError('injected retirement failure') + return await previous_retire(record) + + async def save(record): + nonlocal failed + matches = ((fault == 'save_draining' and record.state == 'DRAINING') + or (fault == 'save_stopped' and record.state == 'STOPPED') + or (fault == 'post_retire_save' and not record.training_ids)) + if matches and not failed: + failed = True + raise RuntimeError('injected instance write failure') + return await previous_save(record) + + monkeypatch.setattr(controller.store, 'retire', retire) + monkeypatch.setattr(controller.store, 'save_instance', save) + with pytest.raises(RuntimeError, match='injected'): + await controller.release(identity, 'owner', 'session') + assert failed + assert identity not in controller._cleaning + assert instance.instance_id not in controller.lifecycle.closing + assert training.state == 'CLOSING' + assert instance.state != 'STOPPED' + for _ in range(3): + await controller.sweep() + await controller.wait_for_jobs() + if training.state == 'RELEASED': + break + assert training.state == 'RELEASED' + assert not instance.training_ids + assert instance.state == 'STOPPED' + assert not runtime.allocations + assert await state.get_model_metadata(identity) is None + assert (await state._backend.get('token_count::owner') or 0) == 0 + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_failure_transition_write_is_retried_and_admission_fenced(monkeypatch): + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'failure-transition') + training = controller.trainings[identity] + instance = controller.instances[training.instance_id] + await controller.admit(identity, 'owner', 'session', 'pending-op') + await controller.store.future('pending-op', 'pending', identity) + original = controller.store.save_instance + failed = False + + async def save(record): + nonlocal failed + if record.state == 'DRAINING' and not failed: + failed = True + raise RuntimeError('injected draining write failure') + return await original(record) + + monkeypatch.setattr(controller.store, 'save_instance', save) + failure = FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message='Worker died') + with pytest.raises(RuntimeError, match='draining write'): + await controller.fail_instance(instance.instance_id, failure) + assert instance.state == 'READY' + with pytest.raises(ResourceNotFoundError): + await controller.route(identity, 'owner') + await controller.sweep() + await controller.wait_for_jobs() + assert instance.state == 'STOPPED' and training.state == 'FAILED' + assert not runtime.allocations and not training.pending_operation_ids + assert (await state.get_future('pending-op'))['failure']['reason_code'] == 'orphaned_replica' + assert not controller._failed_instances + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_failed_resident_is_excluded_when_rebuilding_minimum_replicas(monkeypatch): + controller, _, runtime, _ = await setup(backend=FakeBackend(), min_replicas=1) + failed_instance = next(iter(controller.instances.values())) + original = controller.store.save_instance + failed = False + + async def save(record): + nonlocal failed + if record.state == 'DRAINING' and not failed: + failed = True + raise RuntimeError('injected draining write failure') + return await original(record) + + monkeypatch.setattr(controller.store, 'save_instance', save) + failure = FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message='Worker died') + with pytest.raises(RuntimeError, match='draining write'): + await controller.fail_instance(failed_instance.instance_id, failure) + + assert failed_instance.state == 'READY' + assert failed_instance.instance_id in controller._failed_instances + await controller.sweep() + await controller.wait_for_jobs() + + ready_residents = [ + instance for instance in controller.instances.values() if instance.resident and instance.state == 'READY' + ] + assert failed_instance.state == 'STOPPED' + assert len(ready_residents) == 1 + assert ready_residents[0].instance_id != failed_instance.instance_id + assert set(runtime.allocations) == {ready_residents[0].instance_id} + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_an_in_flight_instance_close(): + controller, _, runtime, _ = await setup(backend=FakeBackend(), min_replicas=1) + instance = next(iter(controller.instances.values())) + original_shutdown = runtime.shutdown + entered, resume = asyncio.Event(), asyncio.Event() + + async def slow_shutdown(target): + entered.set() + await resume.wait() + await original_shutdown(target) + + runtime.shutdown = slow_shutdown + closing = asyncio.create_task(controller._close_instance(instance)) + await entered.wait() + shutdown = asyncio.create_task(controller.shutdown()) + for _ in range(100): + if instance.state == 'DRAINING' or shutdown.done(): + break + await asyncio.sleep(0) + + assert instance.state == 'DRAINING' + assert not shutdown.done() + resume.set() + await asyncio.gather(closing, shutdown) + assert instance.state == 'STOPPED' + assert instance.instance_id not in controller.lifecycle.closing + + +@pytest.mark.asyncio +async def test_cleanup_error_is_not_classified_as_startup_failure(): + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'release-failure') + instance = controller.instances[controller.trainings[identity].instance_id] + runtime.release_training = AsyncMock(side_effect=TimeoutError('release timed out')) + await controller.release(identity, 'owner') + assert instance.failure == 'Training cleanup failed: release timed out' + assert controller.trainings[identity].terminal_reason == 'internal_error' + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_unready_service_is_503_server(): + controller, _, _, _ = await setup(backend=FakeBackend()) + controller._ready = False + with pytest.raises(ResourceExhaustedError) as error: + await controller.create(request('unready')) + assert error.value.error_code == 503 and error.value.category is ErrorCategory.Server + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_sweep_heartbeat_io_does_not_hold_route_lock(monkeypatch): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'slow-heartbeat') + entered, resume = asyncio.Event(), asyncio.Event() + original = state.sessions.get_last_heartbeats + + async def slow(identities): + entered.set() + await resume.wait() + return await original(identities) + + monkeypatch.setattr(state.sessions, 'get_last_heartbeats', slow) + sweeping = asyncio.create_task(controller.sweep()) + try: + await entered.wait() + assert await asyncio.wait_for(controller.route(identity, 'owner'), 0.1) + finally: + resume.set() + await sweeping + await controller.wait_for_jobs() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_explicit_tinker_request_id_is_preserved(): + from twinkle.server.lifecycle.submit import submit_training_task + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'explicit-id') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + try: + ref = await submit_training_task(queue, AsyncMock(return_value='ok'), training_id=identity, seq_id=10, + stream='tinker::step', token='owner', + task_type='step', request_id='explicit-id') + assert ref['request_id'] == 'explicit-id' + assert (await settle(state, 'explicit-id'))['status'] == 'completed' + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_late_serve_deployment_before_confirming_release(monkeypatch): + import importlib + import threading + from types import SimpleNamespace + from twinkle.server.config.training import TrainingConfig + from twinkle.server.state.training_records import InstanceRecord + from twinkle.server.model.pool.driver import RayInstanceDriver + from ray import serve + config = TrainingConfig(models=[{'name': 'small', 'model_id': 'weights', 'backend': 'mock', 'device_type': 'cpu', 'mode': 'lora', 'config': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}]) + store = SimpleNamespace(models=SimpleNamespace(unregister_replica_capacity=AsyncMock())) + driver = RayInstanceDriver(config, store, 'test-late', 'controller') + instance = InstanceRecord(instance_id='instance', controller_epoch='epoch', base_model='small', + mode='lora', startup_deadline=0) + entered, release = threading.Event(), threading.Event() + remove = AsyncMock() + + def deploy(*args, **kwargs): + entered.set() + if not release.wait(10): + raise TimeoutError('test did not release Serve') + + monkeypatch.setattr('twinkle.server.model.app.build_model_app', lambda **kwargs: 'app') + monkeypatch.setattr(importlib.import_module('ray.util.placement_group'), 'get_placement_group', lambda name: 'pg') + monkeypatch.setattr('ray.util.list_named_actors', lambda **kwargs: []) + monkeypatch.setattr(serve, 'run', deploy) + monkeypatch.setattr(serve, 'delete', lambda name: None) + monkeypatch.setattr(driver.runtime.pool, 'remove', remove) + deploying = asyncio.create_task(driver.deploy(instance, driver.config.get('small', 'lora'))) + closing = None + try: + assert await asyncio.to_thread(entered.wait, 2) + deploying.cancel() + await asyncio.gather(deploying, return_exceptions=True) + closing = asyncio.create_task(driver.shutdown(instance)) + await asyncio.sleep(0.01) + assert not closing.done() + remove.assert_not_awaited() + store.models.unregister_replica_capacity.assert_not_awaited() + finally: + release.set() + await asyncio.gather(deploying, return_exceptions=True) + if closing is not None: + await closing + await driver.runtime.wait_for_cleanup() + assert not driver.runtime.deployments and not driver.runtime.lifecycle.jobs + store.models.unregister_replica_capacity.assert_awaited() + + +@pytest.mark.asyncio +async def test_capacity_unregister_does_not_retire_managed_model(): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'capacity-only') + instance = controller.instances[controller.trainings[identity].instance_id] + await state.register_replica(instance.instance_id, 5) + await state.models.unregister_replica_capacity(instance.instance_id) + assert await state.get_model_metadata(identity) is not None + assert await state._backend.get('token_count::owner') == 1 + await controller.release(identity, 'owner') + assert await state.get_model_metadata(identity) is None + assert await state._backend.get('token_count::owner') == 0 + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('mode', ['lora', 'full']) +@pytest.mark.parametrize('failure_stage', ['stopped', 'quarantined', 'persist_draining']) +async def test_late_release_ack_cannot_retire_a_failed_instance_user(mode, failure_stage, monkeypatch): + from twinkle.server.config.training import TrainingConfig + + controller, state, runtime, _ = await setup(backend=FakeBackend()) + config = controller.config.model_dump() + next(model['config'] for model in config['models'] if model['mode'] == mode).update(min_replicas=1, max_replicas=1) + controller.config = TrainingConfig.model_validate(config) + await controller.sweep() + await controller.wait_for_jobs() + identity = await create(controller, state, 'late-release', mode=mode) + training = controller.trainings[identity] + instance = controller.instances[training.instance_id] + entered, resume = asyncio.Event(), asyncio.Event() + release = runtime.release_training + + async def delayed_ack(target, user): + await release(target, user) + entered.set() + await resume.wait() + + runtime.release_training = delayed_ack + training.state, training.terminal_reason = 'CLOSING', 'released' + await controller.store.save_training(training) + cleanup = asyncio.create_task(controller._cleanup_user(training)) + save_instance = controller.store.save_instance + + async def unavailable(record): + if record.state == 'DRAINING': + raise ConnectionError('draining write unavailable') + await save_instance(record) + + try: + await asyncio.wait_for(entered.wait(), 1) + runtime.fail_shutdown = failure_stage == 'quarantined' + failure = FutureFailureRecord(reason_code='orphaned_replica', attribution='server', message='Worker died') + if failure_stage == 'persist_draining': + monkeypatch.setattr(controller.store, 'save_instance', unavailable) + with pytest.raises(ConnectionError): + await controller.fail_instance(instance.instance_id, failure) + else: + await controller.fail_instance(instance.instance_id, failure) + state_before_ack = instance.state + resume.set() + await asyncio.wait_for(cleanup, 1) + assert instance.state == state_before_ack + if failure_stage == 'stopped': + assert training.state == 'FAILED' + assert not runtime.allocations + else: + assert training.state == 'CLOSING' + assert identity in instance.training_ids + assert await state.get_model_metadata(identity) is not None + assert await state._backend.get('token_count::owner') == 1 + with pytest.raises(ResourceNotFoundError): + await controller.route(identity, 'owner') + finally: + resume.set() + await asyncio.gather(cleanup, return_exceptions=True) + runtime.fail_shutdown = False + monkeypatch.setattr(controller.store, 'save_instance', save_instance) + await controller.shutdown() diff --git a/tests/server/training/test_http_integration.py b/tests/server/training/test_http_integration.py new file mode 100644 index 000000000..164a1562d --- /dev/null +++ b/tests/server/training/test_http_integration.py @@ -0,0 +1,39 @@ +import json +import os +import subprocess +import sys +import uuid +from pathlib import Path + + +def test_installed_native_and_tinker_sdks_use_managed_http_routes(tmp_path): + # This probe owns a separate Ray cluster. Stop the session fixture's cluster + # first so two local GCS/raylet pairs do not evict each other on a slow host. + import ray + ray.shutdown() + try: + # Both SDKs, recovery and release/recreate cycles start fresh + # workers. Allow cold imports from a shared environment; each service + # operation still enforces its own startup/execution deadline. + evidence_path = tmp_path / 'http-result.json' + result = subprocess.run([ + sys.executable, + str(Path(__file__).resolve().parents[3] + / 'tests/server/integration/on_demand_training/http_probe.py') + ], env={**os.environ, 'TWINKLE_TEST_HTTP_RESULT': str(evidence_path)}, timeout=2 * 60 * 60) + finally: + ray.init( + ignore_reinit_error=True, + num_cpus=2, + log_to_driver=False, + namespace='twinkle_test', + _temp_dir='/tmp/twinkle-test-' + uuid.uuid4().hex[:8]) + assert result.returncode == 0, 'HTTP probe failed; see subprocess output above' + evidence = json.loads(evidence_path.read_text()) + assert evidence['passed'] and evidence['native_recreation'] and evidence['tinker_shared'] + assert len(set(evidence['native_stable_routing'].values())) == 2 + assert len(set(evidence['tinker_stable_routing'].values())) == 2 + assert evidence['native_lora'] + assert evidence['runtime_lifecycle_rejected'] + assert evidence['tinker_full'] + assert evidence['restart'] and evidence['state_loss'] and evidence['cycles'] == 3 diff --git a/tests/server/training/test_launcher.py b/tests/server/training/test_launcher.py new file mode 100644 index 000000000..bf8b73be5 --- /dev/null +++ b/tests/server/training/test_launcher.py @@ -0,0 +1,205 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import socket +import asyncio +import threading +import uuid +from unittest.mock import Mock + +import pytest +import ray + +from twinkle.server.config import ServerConfig +from twinkle.server.launcher.server_launcher import ServerLauncher +from twinkle.server.model.pool.actor import controller_name + + +def test_unresponsive_controller_does_not_block_other_cleanup(): + @ray.remote(num_cpus=0) + class Controller: + def __init__(self): + self.entered = asyncio.Event() + self.resume = asyncio.Event() + + async def shutdown(self): + self.entered.set() + await self.resume.wait() + + async def wait_until_entered(self): + await self.entered.wait() + + async def unblock(self): + self.resume.set() + + controller = Controller.remote() + launcher = ServerLauncher(ServerConfig(resources={'shutdown_timeout_s': 0.3})) + launcher._training_controller = controller + finished = threading.Event() + errors, cleaned = [], [] + + def stop(): + try: + launcher._run_cleanup([ + ('blocked controller', launcher._stop_training_controller), + ('other owner', lambda: cleaned.append(True))], raise_errors=True) + except Exception as error: + errors.append(error) + finally: + finished.set() + + worker = threading.Thread(target=stop) + # Ensure actor startup is outside the deadline under test. + ray.get(controller.__ray_ready__.remote(), timeout=10) + entered = controller.wait_until_entered.remote() + worker.start() + try: + ray.get(entered, timeout=10) + assert finished.wait(2), 'Controller shutdown exceeded its budget and blocked the next owner' + assert cleaned == [True] + assert len(errors) == 1 and isinstance(errors[0].__cause__, ray.exceptions.GetTimeoutError) + assert launcher._training_controller is controller + ray.get(controller.unblock.remote(), timeout=5) + launcher._stop_training_controller() + assert launcher._training_controller is None + finally: + if not finished.is_set(): + ray.get(controller.unblock.remote(), timeout=5) + worker.join(5) + ray.kill(controller, no_restart=True) + + +def test_failed_initialization_does_not_leave_an_actor_that_blocks_a_retry(): + namespace = 'launcher-retry-' + uuid.uuid4().hex + with socket.socket() as listener: + listener.bind(('127.0.0.1', 0)) + port = listener.getsockname()[1] + devices = int(ray.cluster_resources()['CPU']) + 1 + payload = { + 'ray_namespace': namespace, 'http_options': {'host': '127.0.0.1', 'port': port}, + 'applications': [{'name': 'retry', 'route_prefix': '/api/v1/model/retry', 'import_path': 'model', + 'args': {'model_id': 'mock', 'backend': 'mock', 'train_mode': 'full', + 'device_group': {'ranks': devices, 'device_type': 'cpu'}, + 'device_mesh': {'dp_size': devices}, 'min_replicas': 1}}], + } + launcher = ServerLauncher(ServerConfig.model_validate(payload)) + launcher._start_serve() + try: + with pytest.raises(ray.exceptions.RayTaskError, match='No single-node allocation'): + launcher._deploy_training() + with pytest.raises(ValueError): + ray.get_actor(controller_name(namespace), namespace=namespace) + payload['applications'][0]['args']['min_replicas'] = 0 + launcher.config = ServerConfig.model_validate(payload) + launcher._deploy_training() + diagnostic = ray.get(launcher._training_controller.diagnostics.remote()) + assert diagnostic['ready'] and diagnostic['instances'] == [] + finally: + launcher._shutdown() + + +def test_cleanup_attempts_all_owners_and_reports_the_original_failure(): + calls = [] + error = ConnectionError('shutdown ACK unknown') + + def failed(): + calls.append('failed') + raise error + + def healthy(): + calls.append('healthy') + + with pytest.raises(RuntimeError, match='cleanup is incomplete') as caught: + ServerLauncher._run_cleanup([('first', failed), ('second', healthy)], raise_errors=True) + assert calls == ['failed', 'healthy'] + assert caught.value.__cause__ is error + + +def test_startup_rollback_preserves_startup_error_and_cleans_all_owners(monkeypatch): + config = ServerConfig.model_validate({'applications': [ + {'name': 'gateway', 'import_path': 'server', 'route_prefix': '/api/v1'}]}) + launcher = ServerLauncher(config) + calls = [] + error = ValueError('invalid application') + monkeypatch.setattr('twinkle.server.launcher.server_launcher.apply_ray_serve_patches', lambda: None) + monkeypatch.setattr(launcher, '_configure_environment', lambda: None) + monkeypatch.setattr(launcher, '_init_ray', lambda: None) + monkeypatch.setattr(launcher, '_start_serve', lambda: None) + + def deploy(app): + raise error + + def failed_cleanup(): + calls.append('training') + raise ConnectionError('cleanup pending') + + monkeypatch.setattr(launcher, '_deploy_application', deploy) + monkeypatch.setattr(launcher, '_stop_training_controller', failed_cleanup) + monkeypatch.setattr('twinkle.server.telemetry.flush_telemetry_safely', lambda: calls.append('telemetry')) + with pytest.raises(ValueError) as caught: + launcher.launch() + assert caught.value is error + assert calls == ['training', 'telemetry'] + + +@pytest.mark.parametrize('blocking', [False, True]) +def test_launch_preserves_blocking_policy_and_deploys_on_demand_training(monkeypatch, blocking): + config = ServerConfig.model_validate({'applications': [ + {'name': 'gateway', 'import_path': 'server', 'route_prefix': '/api/v1'}, + {'name': 'train', 'import_path': 'model', 'route_prefix': '/api/v1/model/train', + 'args': {'model_id': 'mock', 'backend': 'mock', 'train_mode': 'full', + 'device_group': {'ranks': 1, 'device_type': 'cpu'}, + 'device_mesh': {'dp_size': 1}, 'min_replicas': 0}}, + ]}) + launcher = ServerLauncher(config) + calls = [] + monkeypatch.setattr('twinkle.server.launcher.server_launcher.apply_ray_serve_patches', lambda: None) + monkeypatch.setattr(launcher, '_configure_environment', lambda: calls.append('environment')) + monkeypatch.setattr(launcher, '_init_ray', lambda: calls.append('ray')) + monkeypatch.setattr(launcher, '_start_serve', lambda: calls.append('serve')) + monkeypatch.setattr(launcher, '_deploy_application', lambda app: calls.append(app.name)) + monkeypatch.setattr(launcher, '_deploy_training', lambda: calls.append('training')) + monkeypatch.setattr(launcher, '_wait_for_shutdown', lambda: calls.append('wait')) + monkeypatch.setattr(launcher, '_shutdown', lambda **kwargs: calls.append('shutdown')) + launcher.launch(blocking=blocking) + expected = ['environment', 'ray', 'serve', 'gateway', 'train', 'training'] + assert calls == expected + (['wait', 'shutdown'] if blocking else []) + + +def test_start_serve_reuses_existing_instance_without_restart(monkeypatch): + launcher = ServerLauncher(ServerConfig()) + client = Mock(return_value=object()) + start, shutdown = Mock(), Mock() + monkeypatch.setattr('ray.serve.context._get_global_client', client) + monkeypatch.setattr('ray.serve.start', start) + monkeypatch.setattr('ray.serve.shutdown', shutdown) + launcher._start_serve() + launcher._start_serve() + client.assert_called_once_with() + start.assert_not_called() + shutdown.assert_not_called() + assert launcher._serve_started + + +def test_start_serve_starts_missing_instance_with_configured_listener(monkeypatch): + launcher = ServerLauncher(ServerConfig(proxy_location='HeadOnly')) + client = Mock(side_effect=RuntimeError('Serve is not running')) + start, shutdown = Mock(), Mock() + monkeypatch.setattr('ray.serve.context._get_global_client', client) + monkeypatch.setattr('ray.serve.start', start) + monkeypatch.setattr('ray.serve.shutdown', shutdown) + launcher._start_serve() + start.assert_called_once_with(http_options=launcher.config.http_options.model_dump(), proxy_location='HeadOnly') + shutdown.assert_not_called() + assert launcher._serve_started + + +@pytest.mark.parametrize('blocking', [False, True]) +def test_launch_server_forwards_blocking_policy(monkeypatch, blocking): + from twinkle.server.launcher.server_launcher import launch_server + + config = ServerConfig() + launcher = Mock() + factory = Mock(return_value=launcher) + monkeypatch.setattr('twinkle.server.launcher.server_launcher.ServerLauncher', factory) + launch_server(config=config, blocking=blocking) + factory.assert_called_once_with(config=config, ray_namespace=config.ray_namespace) + launcher.launch.assert_called_once_with(blocking=blocking) diff --git a/tests/server/training/test_router.py b/tests/server/training/test_router.py new file mode 100644 index 000000000..52e4b1c7c --- /dev/null +++ b/tests/server/training/test_router.py @@ -0,0 +1,97 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +from unittest.mock import Mock +from types import SimpleNamespace + +import pytest +from starlette.requests import Request + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.helpers import settle +from twinkle.hub import HubOperation +from twinkle.server.state import ServerState +from twinkle.server.model.pool.router import TrainingRouter +from twinkle.server.config.training import TrainingConfig + + +def http_request(body): + async def receive(): + return {'type': 'http.request', 'body': json.dumps(body).encode(), 'more_body': False} + request = Request({'type': 'http', 'method': 'POST', 'path': '/upload_to_hub', 'headers': []}, receive) + request.state.token = 'owner' + request.state.session_id = '' + return request + + +@pytest.mark.asyncio +async def test_upload_after_recycle_is_cpu_only_and_preserves_future_ownership(monkeypatch, tmp_path): + router = TrainingRouter.__new__(TrainingRouter) + router.state = ServerState(backend=FakeBackend()) + router.store = router.state.create_training_store('upload-test') + router.replica_id = 'router' + await router.state.register_replica(router.replica_id, 0) + await router.state.touch_replica_last_seen(router.replica_id) + router._init_task_queue() + router.controller = Mock(side_effect=AssertionError('Upload must not create or route a training')) + upload = Mock(return_value=None) + monkeypatch.setattr(HubOperation, 'push_to_hub', upload) + try: + response = await router.native_request(http_request({ + 'checkpoint_dir': str(tmp_path), 'hub_model_id': 'owner/saved', 'hub_token': 'hub-token', + }), 'upload_to_hub') + future = await settle(router.state, response.request_id) + assert future['status'] == 'completed' + upload.assert_called_once_with(repo_id='owner/saved', folder_path=str(tmp_path), token='hub-token', private=True) + assert (await router.state.get_future(response.request_id))['owner_token'] == 'owner' + assert not await router.state._backend.keys('model::*') + finally: + await router.shutdown_task_queue() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('rank', [8, 0, 33, True]) +async def test_native_lora_admission_does_not_import_peft(monkeypatch, rank): + import builtins + from unittest.mock import AsyncMock + from twinkle.server.exceptions import RequestRejectedError + router = TrainingRouter.__new__(TrainingRouter) + router.model = SimpleNamespace(name='small') + router.state = ServerState(backend=FakeBackend()) + await router.state.create_session({'session_id': 'session'}) + router.config = TrainingConfig(models=[{'name': 'small', 'model_id': 'weights', 'backend': 'transformers', 'device_type': 'cuda', 'mode': 'lora', 'config': {'num_devices': 1, 'device_mesh': {'dp_size': 1}}}]) + async def create(payload): + await router.state.store_future_status('created', 'pending', None) + return {'request_id': 'created'} + router.controller = SimpleNamespace(create=SimpleNamespace(remote=AsyncMock(side_effect=create))) + original_import = builtins.__import__ + def cpu_import(name, *args, **kwargs): + if name == 'peft' or name.startswith('peft.'): + raise AssertionError('CPU admission must not import PEFT') + return original_import(name, *args, **kwargs) + monkeypatch.setattr(builtins, '__import__', cpu_import) + request = http_request({'client_create_id': 'create', 'adapter_name': 'same', + 'config': json.dumps({'_TWINKLE_TYPE_': 'LoraConfig', 'r': rank})}) + request.state.session_id = 'session' + if type(rank) is int and rank == 8: + result = await router.native_request(request, 'add_adapter_to_model') + assert result.request_id == 'created' + router.controller.create.remote.assert_awaited_once() + else: + with pytest.raises(RequestRejectedError, match='rank'): + await router.native_request(request, 'add_adapter_to_model') + router.controller.create.remote.assert_not_awaited() + + +def test_router_schemas_match_model_runtime_routes(): + from fastapi import FastAPI + from twinkle.server.model.tinker_handlers import _register_model_tinker_routes + from twinkle.server.model.twinkle_handlers import _register_model_twinkle_routes + from twinkle.server.model.pool.router import NATIVE_REQUESTS, TINKER_REQUESTS + + app = FastAPI() + _register_model_tinker_routes(app, lambda: None) + _register_model_twinkle_routes(app, lambda: None) + for protocol, schemas in [('twinkle', NATIVE_REQUESTS), ('tinker', TINKER_REQUESTS)]: + endpoints = {route.path.rsplit('/', 1)[-1] for route in app.routes + if route.path.startswith(f'/{protocol}/') and 'POST' in route.methods} + assert endpoints == set(schemas) diff --git a/tests/server/training/test_runtime_initialization.py b/tests/server/training/test_runtime_initialization.py new file mode 100644 index 000000000..b136d4fc8 --- /dev/null +++ b/tests/server/training/test_runtime_initialization.py @@ -0,0 +1,49 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The runtime event loop remains responsive during the first dynamic config import.""" +import asyncio +import threading +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tests.server.state.fake_backend import FakeBackend +from tests.server.training.helpers import create, setup +from twinkle.server.model.runtime import TrainingRuntime +from twinkle.server.state.training_records import CreateTraining + + +@pytest.mark.asyncio +async def test_configuration_deserialization_does_not_block_event_loop(monkeypatch): + controller, state, _, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'configuration-import', mode='full') + training = controller.trainings[identity] + entered, resume, finished = threading.Event(), threading.Event(), threading.Event() + + def deserialize(config): + entered.set() + resume.wait(3) + finished.set() + return None + + monkeypatch.setattr('twinkle.server.model.runtime.deserialize_object', deserialize) + manager = SimpleNamespace(save=Mock()) + monkeypatch.setattr('twinkle.server.model.runtime.create_training_run_manager', lambda *args, **kw: manager) + + async def serial(factory, task_type): + return await factory() + + runtime = SimpleNamespace(_run_lifecycle_operation=serial, resolve_model_adapter_name=lambda _: '', + _training_states={}, is_full_mode=True) + request = CreateTraining(base_model='small', mode='full', token='owner', session_id='session', create_key='key', + client_type='twinkle', payload={'adapter_name': 'same', 'config': None}) + initializing = asyncio.create_task(TrainingRuntime.initialize_training( + runtime, training.model_dump(), request.model_dump())) + try: + assert await asyncio.to_thread(entered.wait, 2) + assert not finished.is_set(), 'Synchronous import blocked the event loop' + finally: + resume.set() + await initializing + await controller.shutdown() + manager.save.assert_called_once() diff --git a/tests/server/training/test_runtime_queue.py b/tests/server/training/test_runtime_queue.py new file mode 100644 index 000000000..ac89f4b25 --- /dev/null +++ b/tests/server/training/test_runtime_queue.py @@ -0,0 +1,367 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Lifecycle protection through the actual compute worker and backend thread.""" +import asyncio +import threading +from types import SimpleNamespace +from tests.server.training.helpers import RPC, RuntimeQueue, settle + +import pytest + +from twinkle.server.config.task_queue import TaskQueueConfig +from twinkle.server.task_queue import TaskQueueMixin +from twinkle.server.model.runtime import TrainingRuntime +from twinkle.server.lifecycle.submit import submit_training_task +from tests.server.training.helpers import create, setup + + +@pytest.mark.asyncio +async def test_execution_timeout_keeps_live_thread_and_allocation_quarantined(): + controller, state, runtime, _ = await setup(devices=2) + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + entered, unblock = threading.Event(), threading.Event() + runtime.fail_shutdown = True + def blocking(): + entered.set() + unblock.wait(3) + async def operation(): + await queue.call_backend(blocking) + try: + ref = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='forward_backward') + future = await settle(state, ref['request_id']) + assert entered.is_set() and not unblock.is_set() + assert future['failure']['reason_code'] == 'execution_timeout' + for _ in range(100): + if instance.state == 'QUARANTINED':break + await asyncio.sleep(0.01) + assert instance.state == 'QUARANTINED' and runtime.allocations + assert controller.trainings[identity].pending_operation_ids + assert queue._model_unhealthy + finally: + unblock.set() + await queue.shutdown_task_queue() + runtime.fail_shutdown = False + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + + +@pytest.mark.asyncio +async def test_tinker_seq_retry_runs_once_and_different_operations_do_not_collide(): + controller, state, runtime, _ = await setup() + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + counter = [] + async def operation(): + counter.append('executed') + return {'ok': True} + try: + first, retry = await asyncio.gather(*( + submit_training_task(queue, operation, training_id=identity, seq_id=1, + stream='tinker::forward_backward', token='owner', task_type='forward_backward') + for _ in range(2))) + assert first == retry + assert (await settle(state, first['request_id']))['status'] == 'completed' + step = await submit_training_task(queue, operation, training_id=identity, seq_id=1, + stream='tinker::optim_step', token='owner', task_type='optim_step') + assert step['request_id'] != first['request_id'] + await settle(state, step['request_id']) + assert len(counter) == 2 + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_cancelled_queued_operation_releases_ticket_without_business_renewal(monkeypatch): + import time + # Freeze wall time as well as the injected business clock; RPC latency is not activity. + now = time.time() + monkeypatch.setattr('twinkle.server.task_queue.worker.time.time', lambda: now) + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 1) + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + queue._task_queue_config = TaskQueueConfig(execution_timeout=2) + queue._compute_worker._config = queue._task_queue_config + entered, unblock = asyncio.Event(), asyncio.Event() + executions = [] + async def busy(): + entered.set() + await unblock.wait() + async def cancelled(): + executions.append('must not run') + try: + first = await queue.schedule_task(busy, model_id=identity, token='owner', task_type='forward_backward') + await entered.wait() + second = await queue.schedule_task(cancelled, model_id=identity, token='owner', task_type='step') + assert (await state.cancel_future(second['request_id']))['cancelled'] + clock.now = 9 + unblock.set() + await settle(state, first['request_id']) + for _ in range(100): + if not controller.trainings[identity].pending_operation_ids:break + await asyncio.sleep(0.01) + assert not executions and not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].last_activity_at == 9 + assert (await state.get_future(second['request_id']))['status'] == 'cancelled' + finally: + unblock.set() + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_readiness_waits_for_real_ray_worker_constructor(): + import ray + import time + @ray.remote(num_cpus=0.1) + class Worker: + def __init__(self): + time.sleep(0.3) + self.loaded = True + def is_loaded(self): + return self.loaded + worker = Worker.remote() + runtime = TrainingRuntime() + runtime.model = SimpleNamespace(_actors=[worker]) + try: + await runtime.wait_for_training_workers() + assert await worker.is_loaded.remote() + assert ray._private.state.actors(worker._actor_id.hex())['State'] == 'ALIVE' + finally: + ray.kill(worker, no_restart=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('reason', ['heartbeat', 'max_lifetime']) +async def test_adapter_expiry_preserves_already_queued_work(reason): + import time + from datetime import datetime, timedelta, timezone + from tests.server.state.fake_backend import FakeBackend + + controller, state, runtime, _ = await setup( + backend=FakeBackend(), adapter_config={'adapter_timeout': 30, 'adapter_max_lifetime': 60}) + identity = await create(controller, state, 'queued-expiry') + training = controller.trainings[identity] + queue = RuntimeQueue(controller, state, controller.instances[training.instance_id]) + queue._task_queue_config = TaskQueueConfig(execution_timeout=5) + queue._compute_worker._config = queue._task_queue_config + entered, resume, finished = asyncio.Event(), asyncio.Event(), asyncio.Event() + executions = [] + + async def first(): + entered.set() + await resume.wait() + executions.append('first') + + async def second(): + executions.append('second') + + async def notify(*args, **kwargs): + await controller.finish_operation(*args, **kwargs) + if not training.pending_operation_ids: + finished.set() + return True + + queue._training_controller.finish_operation = RPC(notify) + try: + active = await queue.schedule_task(first, model_id=identity, token='owner', task_type='forward_backward') + await asyncio.wait_for(entered.wait(), 2) + queued = await queue.schedule_task(second, model_id=identity, token='owner', task_type='step') + if reason == 'heartbeat': + session = await state.sessions.get('session') + session.last_heartbeat = time.time() - 31 + await state.sessions.add('session', session) + else: + training.created_at = (datetime.now(timezone.utc) - timedelta(seconds=61)).isoformat() + await controller.sweep() + await controller.wait_for_jobs() + assert identity in runtime.users and len(training.pending_operation_ids) == 2 + resume.set() + assert (await settle(state, active['request_id']))['status'] == 'completed' + assert (await settle(state, queued['request_id']))['status'] == 'completed' + await asyncio.wait_for(finished.wait(), 2) + await controller.sweep() + await controller.wait_for_jobs() + assert executions == ['first', 'second'] + assert training.state == 'EXPIRED' and not runtime.allocations + finally: + resume.set() + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_lost_finish_notification_is_reconciled_without_stopping_compute(): + controller, state, runtime, clock = await setup() + identity = await create(controller, state, 'lost-notification') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + async def lost(*args, **kwargs): + raise RuntimeError('Completion RPC lost') + queue._training_controller.finish_operation = RPC(lost) + async def status(instance, request_id): + return await queue.verify_operation(request_id) + runtime.verify_operation = status + async def compute(): + return {'updated': True} + try: + ref = await queue.schedule_task(compute, model_id=identity, token='owner', task_type='step') + assert (await settle(state, ref['request_id']))['status'] == 'completed' + assert controller.trainings[identity].pending_operation_ids + clock.now = 6 + await controller.sweep() + await controller.wait_for_jobs() + assert not controller.trainings[identity].pending_operation_ids + assert controller.trainings[identity].state == 'ACTIVE' and not queue._model_unhealthy + clock.now += 11 + await controller.sweep() + await controller.wait_for_jobs() + assert not runtime.allocations + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fault,executed', [('queue_timeout', False), ('unknown_start_ack', False), + ('backend_error', True)]) +async def test_lost_failure_notification_uses_actual_execution_facts(fault, executed, monkeypatch): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + identity = await create(controller, state, fault) + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + + async def lost(*args, **kwargs): + raise ConnectionError('Completion RPC lost') + + queue._training_controller.finish_operation = RPC(lost) + runtime.verify_operation = lambda instance, request_id: queue.verify_operation(request_id) + ran = [] + + async def operation(): + ran.append(True) + raise ValueError('executed backend error') + + if fault == 'queue_timeout': + queue._compute_worker._config = TaskQueueConfig(queue_timeout=0) + elif fault == 'unknown_start_ack': + start = state.try_start_future + + async def lose_ack(request_id): + assert await start(request_id) + raise ConnectionError('start acknowledgement lost') + + monkeypatch.setattr(state, 'try_start_future', lose_ack) + try: + clock.now = 7 + ref = await queue.schedule_task(operation, model_id=identity, token='owner', task_type='step') + assert (await settle(state, ref['request_id']))['status'] == 'failed' + for _ in range(100): + if not queue._compute_worker.has_task(ref['request_id']): + break + await asyncio.sleep(0.01) + outcome = await queue.verify_operation(ref['request_id']) + assert outcome['state'] == 'finished' and outcome['executed'] is executed + assert bool(ran) is executed + assert controller.trainings[identity].pending_operation_ids + clock.now = 13 # Admission at 7 must age past the five-second reconciliation interval. + await controller.sweep() + await controller.wait_for_jobs() + training = controller.trainings[identity] + assert not training.pending_operation_ids + assert (training.last_activity_at > 6) is executed + if not executed: + clock.now = 14 + await controller.sweep() + await controller.wait_for_jobs() + assert training.state == 'EXPIRED' + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_execution_facts_expire_without_future_status_guessing(monkeypatch): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, _ = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'facts-expire') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + + async def lost(*args, **kwargs): + raise ConnectionError('Completion RPC lost') + + queue._training_controller.finish_operation = RPC(lost) + try: + ref = await queue.schedule_task(lambda: asyncio.sleep(0), model_id=identity, token='owner', task_type='step') + await settle(state, ref['request_id']) + worker = queue._compute_worker + for _ in range(100): + if not worker.has_task(ref['request_id']): + break + await asyncio.sleep(0.01) + outcome = worker.get_task_outcome(ref['request_id']) + assert outcome['executed'] + monkeypatch.setattr('twinkle.server.task_queue.worker.time.time', + lambda: outcome['finished_at'] + worker._config.absolute_future_ttl(1) + 1) + assert (await queue.verify_operation(ref['request_id']))['state'] == 'missing' + assert not worker._task_outcomes + finally: + await queue.shutdown_task_queue() + await controller.shutdown() + + +@pytest.mark.asyncio +async def test_completion_write_failure_retains_ticket_and_execution_fact(monkeypatch): + from tests.server.state.fake_backend import FakeBackend + controller, state, runtime, clock = await setup(backend=FakeBackend()) + identity = await create(controller, state, 'finish-write') + instance = controller.instances[controller.trainings[identity].instance_id] + queue = RuntimeQueue(controller, state, instance) + runtime.verify_operation = lambda instance, request_id: queue.verify_operation(request_id) + save = controller.store.save_training + failed = False + + async def fail_completion(record): + nonlocal failed + if record.state == 'ACTIVE' and not record.pending_operation_ids and not failed: + failed = True + raise ConnectionError('completion write failed') + await save(record) + + monkeypatch.setattr(controller.store, 'save_training', fail_completion) + try: + clock.now = 7 + ref = await queue.schedule_task(lambda: asyncio.sleep(0), model_id=identity, token='owner', task_type='step') + await settle(state, ref['request_id']) + for _ in range(100): + if not queue._compute_worker.has_task(ref['request_id']): + break + await asyncio.sleep(0.01) + for _ in range(100): + if failed: + break + await asyncio.sleep(0.01) + training = controller.trainings[identity] + assert failed and ref['request_id'] in training.pending_operation_ids + assert training.last_activity_at == 0 + assert queue._compute_worker.get_task_outcome(ref['request_id'])['executed'] + clock.now = 13 + await controller.sweep() + await controller.wait_for_jobs() + assert not training.pending_operation_ids and 6 < training.last_activity_at <= 13 + followup = await queue.schedule_task(lambda: asyncio.sleep(0), model_id=identity, token='owner', task_type='step') + await settle(state, followup['request_id']) + for _ in range(100): + if queue._compute_worker.get_task_outcome(followup['request_id']) is None: + break + await asyncio.sleep(0.01) + assert queue._compute_worker.get_task_outcome(followup['request_id']) is None + finally: + await queue.shutdown_task_queue() + await controller.shutdown() diff --git a/tests/server/training/test_storage_recovery.py b/tests/server/training/test_storage_recovery.py new file mode 100644 index 000000000..a0eefdb9f --- /dev/null +++ b/tests/server/training/test_storage_recovery.py @@ -0,0 +1,66 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The lifecycle uses the same recovery semantics on real memory and Redis backends.""" +import os +import uuid + +import pytest + +from twinkle.server.state.backend.memory_backend import RayActorBackend +from twinkle.server.state.backend.redis_backend import RedisBackend +from twinkle.server.model.pool.controller import TrainingController +from tests.server.training.helpers import create, request, setup + + +@pytest.mark.asyncio +@pytest.mark.parametrize('mode', ['memory', 'redis']) +async def test_pending_operation_recovery_preserves_creation_identity_and_quota(mode): + if mode == 'redis' and not os.environ.get('TWINKLE_TEST_REDIS_URL'): + pytest.skip('Set TWINKLE_TEST_REDIS_URL to a dedicated local Redis') + backend = (RayActorBackend() if mode == 'memory' else + RedisBackend(os.environ['TWINKLE_TEST_REDIS_URL'], key_prefix='training-test-' + uuid.uuid4().hex + ':')) + try: + controller, state, runtime, clock = await setup(backend=backend) + identity = await create(controller, state, 'retry') + original = await controller.create(request('retry')) + await controller.admit(identity, 'owner', 'session', 'queued') + await controller.store.future('queued', 'queued', identity) + replacement = TrainingController(controller.config, state.create_training_store('unit'), runtime, clock=clock) + await replacement.initialize() + assert (await state.get_future('queued'))['failure']['reason_code'] == 'orphaned_replica' + assert await replacement.create(request('retry')) == original + assert not runtime.allocations + assert await state.get_model_metadata(identity) is None + assert (await backend.get('token_count::owner') or 0) == 0 + new = await create(replacement, state, 'new') + assert new != identity + assert await backend.get('token_count::owner') == 1 + await replacement.shutdown() + finally: + if mode == 'redis': + for key in await backend.keys('*'): + await backend.delete(key) + await backend.close() + + +@pytest.mark.asyncio +async def test_managed_memory_actor_restart_reconnects_existing_handles(): + import asyncio + import ray + backend = RayActorBackend(namespace='memory-recovery-' + uuid.uuid4().hex) + try: + await backend.set('before', 'old') + ray.kill(backend._actor, no_restart=False) + for _ in range(100): + try: + if await backend.health_check() and await backend.get('before') is None: + break + except ray.exceptions.RayError: + pass + await asyncio.sleep(0.1) + else: + raise AssertionError('State actor did not restart') + assert await backend.get('before') is None + await backend.set('after', 'new') + assert await backend.get('after') == 'new' + finally: + ray.kill(backend._actor, no_restart=True) diff --git a/tests/server/validation/test_preflight.py b/tests/server/validation/test_preflight.py index 2912b1e8e..f6ffe539c 100644 --- a/tests/server/validation/test_preflight.py +++ b/tests/server/validation/test_preflight.py @@ -28,22 +28,23 @@ def __init__(self, backend: str = 'transformers'): self.futures: dict[str, dict] = {} self.backend_calls: list = [] self.claimed: list = [] + self.sequence_retention_s = 3600 self._task_queue_config = SimpleNamespace(effective_execution_timeout=60.0) self.state = SimpleNamespace( - claim_seq=self._claim_seq, + claim_training_seq=self._claim_seq, get_future=self._get_future, - release_seq=self._release_seq, + release_training_seq=self._release_seq, ) - async def _claim_seq(self, dedup_key, request_id, ttl): - self.claimed.append(dedup_key) + async def _claim_seq(self, training_id, stream, sequence, request_id, retention): + self.claimed.append((training_id, stream, sequence)) return None async def _get_future(self, request_id): return self.futures.get(request_id) - async def _release_seq(self, dedup_key): - self.claimed.remove(dedup_key) + async def _release_seq(self, training_id, stream, sequence, request_id): + self.claimed.remove((training_id, stream, sequence)) async def _on_request_start(self, request): return 'token' @@ -82,7 +83,7 @@ def test_megatron_rejects_the_split_gradient_endpoints_with_501(): for capability in (BackendCapability.Forward, BackendCapability.Backward, BackendCapability.CalculateLoss): with pytest.raises(EndpointUnavailableError) as raised: assert_request_supported( - _Deployment('megatron'), model_types.AdapterRequest(adapter_name='a'), capability=capability) + _Deployment('megatron'), model_types.AdapterRequest(adapter_name='a', training_id='training-a'), capability=capability) assert raised.value.error_code == 501 assert 'forward_backward' in str(raised.value), 'the alternative endpoints must be named' @@ -90,14 +91,14 @@ def test_megatron_rejects_the_split_gradient_endpoints_with_501(): def test_transformers_serves_the_split_gradient_endpoints(): assert_request_supported( _Deployment('transformers'), - model_types.AdapterRequest(adapter_name='a'), + model_types.AdapterRequest(adapter_name='a', training_id='training-a'), capability=BackendCapability.Forward) def test_mock_backend_serves_every_endpoint(): """The mock is a test double; restricting it would only break tests.""" assert_request_supported( - _Deployment('mock'), model_types.AdapterRequest(adapter_name='a'), capability=BackendCapability.Forward) + _Deployment('mock'), model_types.AdapterRequest(adapter_name='a', training_id='training-a'), capability=BackendCapability.Forward) @pytest.mark.asyncio @@ -107,7 +108,7 @@ async def test_capability_rejection_writes_no_future_and_calls_no_backend(): await run_submit( deployment, _request(), - model_types.AdapterRequest(adapter_name='a'), + model_types.AdapterRequest(adapter_name='a', training_id='training-a'), task_type='backward', backend_call=_call, capability=BackendCapability.Backward) @@ -123,7 +124,7 @@ async def test_capability_rejection_writes_no_future_and_calls_no_backend(): def test_a_megatron_only_parameter_is_rejected_on_transformers(): with pytest.raises(RequestRejectedError) as raised: - assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a', merge_lora=True)) + assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a', merge_lora=True, training_id='training-a')) assert raised.value.error_code == 422 assert 'merge_lora' in str(raised.value) assert 'megatron' in str(raised.value) @@ -132,7 +133,7 @@ def test_a_megatron_only_parameter_is_rejected_on_transformers(): def test_a_transformers_only_parameter_is_rejected_on_megatron(): with pytest.raises(RequestRejectedError): assert_request_supported( - _Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt', strict=True)) + _Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt', strict=True, training_id='training-a')) def test_an_unset_backend_only_parameter_is_not_rejected(): @@ -142,8 +143,8 @@ def test_an_unset_backend_only_parameter_is_not_rejected(): backend's own default, it would look sent on every request and the other half of the fleet would reject everything. """ - assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a')) - assert_request_supported(_Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt')) + assert_request_supported(_Deployment('transformers'), model_types.SaveRequest(adapter_name='a', training_id='training-a')) + assert_request_supported(_Deployment('megatron'), model_types.LoadRequest(adapter_name='a', name='ckpt', training_id='training-a')) def test_every_restricted_field_is_optional_with_a_none_default(): @@ -178,14 +179,14 @@ def test_a_real_parameter_read_from_kwargs_is_not_rejected(): assert_request_supported( _Deployment(), model_types.SetProcessorRequest( - processor_cls='InputProcessor', adapter_name='a', init_kwargs={'padding_side': 'right'})) + processor_cls='InputProcessor', adapter_name='a', init_kwargs={'padding_side': 'right'}, training_id='training-a')) def test_an_unrecognised_plugin_argument_is_forwarded(): """Plugins accept ``**kwargs``, so "unknown" cannot mean "wrong".""" assert_request_supported( _Deployment(), - model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'my_custom_knob': 1})) + model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'my_custom_knob': 1}, training_id='training-a')) def test_no_plugin_download_happens_during_validation(monkeypatch): @@ -195,7 +196,7 @@ def test_no_plugin_download_happens_during_validation(monkeypatch): lambda *a, **k: pytest.fail('validation must not download a plugin')) assert_request_supported( _Deployment(), - model_types.SetLossRequest(loss_cls='ms://someone/MyLoss', adapter_name='a', init_kwargs={'beta': 0.1})) + model_types.SetLossRequest(loss_cls='ms://someone/MyLoss', adapter_name='a', init_kwargs={'beta': 0.1}, training_id='training-a')) # --------------------------------------------------------------------------- # @@ -210,32 +211,32 @@ def test_control_fields_are_never_forwarded_to_the_backend(): into a backend signature. """ body = model_types.ForwardBackwardTaskRequest( - inputs=[{'input_ids': [1, 2]}], adapter_name='a', seq_id=3, task='embedding') + inputs=[{'input_ids': [1, 2]}], adapter_name='a', seq_id=3, task='embedding', training_id='training-a') assert backend_kwargs(body) == {'task': 'embedding'} def test_unset_backend_parameters_are_not_forwarded(): - body = model_types.ForwardRequest(inputs=[{'input_ids': [1]}], adapter_name='a') + body = model_types.ForwardRequest(inputs=[{'input_ids': [1]}], adapter_name='a', training_id='training-a') assert backend_kwargs(body) == {} def test_passthrough_contents_are_flattened(): body = model_types.ForwardBackwardTaskRequest( - inputs=[{'input_ids': [1]}], adapter_name='a', loss_kwargs={'advantages': [0.5]}) + inputs=[{'input_ids': [1]}], adapter_name='a', loss_kwargs={'advantages': [0.5]}, training_id='training-a') assert backend_kwargs(body) == {'advantages': [0.5]} def test_a_passthrough_key_shadowing_a_declared_parameter_is_an_error(): """Silently letting one win would make the effective value depend on merge order.""" body = model_types.ForwardRequest(inputs=[{'input_ids': [1]}], adapter_name='a', task='causal_lm', - loss_kwargs={'task': 'embedding'}) + loss_kwargs={'task': 'embedding'}, training_id='training-a') with pytest.raises(ValueError, match='collides'): backend_kwargs(body) def test_plugin_identifiers_are_not_forwarded_as_kwargs(): """The handler passes ``loss_cls`` positionally; forwarding it too would duplicate it.""" - body = model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}) + body = model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, training_id='training-a') assert backend_kwargs(body) == {'beta': 0.1} diff --git a/tests/server/validation/test_request_wire.py b/tests/server/validation/test_request_wire.py index 73759d25d..5e0ab0acf 100644 --- a/tests/server/validation/test_request_wire.py +++ b/tests/server/validation/test_request_wire.py @@ -55,7 +55,7 @@ def _post(client: TestClient, path: str, body: dict): def _valid_forward(**overrides) -> dict: - return {'inputs': [{'input_ids': [1, 2, 3]}], 'adapter_name': 'a', **overrides} + return {'training_id': 'training-a', 'inputs': [{'input_ids': [1, 2, 3]}], 'adapter_name': 'a', **overrides} def test_a_valid_body_reaches_the_handler(client): @@ -140,6 +140,7 @@ def test_a_dynamic_plugin_argument_is_accepted_inside_its_region(client): response = _post(client, '/set_loss', { 'loss_cls': 'DPOLoss', 'adapter_name': 'a', + 'training_id': 'training-a', 'init_kwargs': {'beta': 0.1, 'anything_at_all': [1, 2]}, }) assert response.status_code == 200 @@ -154,7 +155,7 @@ def test_a_non_json_value_fails_in_the_client_before_any_request(): """ from pydantic import ValidationError with pytest.raises(ValidationError) as raised: - model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': object()}) + model_types.SetLossRequest(loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': object()}, training_id='training-a') assert 'init_kwargs' in str(raised.value), 'the error must name the field that holds the bad value' @@ -246,3 +247,12 @@ def test_sampler_add_adapter_accepts_a_dict_config_where_model_rejects_it(): if __name__ == '__main__': raise SystemExit(pytest.main([__file__, '-v'])) + + +def test_missing_training_handle_is_rejected_before_handler(client): + body = _valid_forward() + body.pop('training_id') + response = _post(client, '/forward', body) + assert response.status_code == 422 + assert any(detail['field'] == 'training_id' for detail in response.json()['details']) + assert 'reached_handler' not in response.json() diff --git a/tests/transformers/test_dev_task_training.py b/tests/transformers/test_dev_task_training.py new file mode 100644 index 000000000..34f196522 --- /dev/null +++ b/tests/transformers/test_dev_task_training.py @@ -0,0 +1,136 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Exercise dev task entry points on real tiny decoder models without downloads.""" +import copy + +import pytest +import torch + + +@pytest.mark.parametrize('task', ['embedding', 'reranker', 'seq_cls']) +def test_dev_task_forward_backward_and_update(tmp_path, monkeypatch, task): + from tokenizers import Tokenizer + from tokenizers.models import WordPiece + from transformers import (PreTrainedTokenizerFast, Qwen2Config, Qwen2ForCausalLM, + Qwen2ForSequenceClassification) + from twinkle import Platform, initialize + from twinkle.model.transformers.transformers import TransformersModel + monkeypatch.setenv('ACCELERATE_USE_CPU', 'true') + monkeypatch.setattr(Platform, 'get_local_device', lambda *args, **kwargs: 'cpu') + initialize(mode='local') + config = Qwen2Config(vocab_size=32, hidden_size=16, intermediate_size=32, num_hidden_layers=1, + num_attention_heads=2, num_key_value_heads=2, pad_token_id=0, + num_labels=3 if task == 'seq_cls' else 1) + model_cls = Qwen2ForCausalLM if task == 'embedding' else Qwen2ForSequenceClassification + model_cls(config).save_pretrained(tmp_path) + vocab = {'[PAD]': 0, '[UNK]': 1, '[CLS]': 2, '[SEP]': 3, '[MASK]': 4} + vocab.update({f't{i}': i for i in range(5, 32)}) + PreTrainedTokenizerFast(tokenizer_object=Tokenizer(WordPiece(vocab, unk_token='[UNK]')), + pad_token='[PAD]', unk_token='[UNK]').save_pretrained(tmp_path) + model = TransformersModel(model_id=str(tmp_path), mixed_precision='no', gradient_checkpointing=False) + model.set_processor('InputProcessor', padding_side='right') + if task == 'embedding': + model.set_loss('InfonceLoss', temperature=0.1) + labels = [1, 0, 0, 1, 0, 0] + elif task == 'reranker': + model.set_loss('PointwiseRerankerLoss') + labels = [1., 0., 0., 1., 0., 0.] + else: + model.set_loss('SeqClsLoss', problem_type='single_label_classification', num_labels=3) + labels = [0, 1, 2, 0, 1, 2] + model.set_optimizer('AdamW', lr=1e-3) + inputs = [{'input_ids': [2, 5+i, 3], 'attention_mask': [1] * 3, 'labels': [label]} + for i, label in enumerate(labels)] + before = {name: p.detach().clone() for name, p in model.model.named_parameters()} + result = model.forward_backward(inputs=copy.deepcopy(inputs), task=task) + assert torch.isfinite(torch.tensor(result['loss'])) + model.clip_grad_and_step() + assert any(not torch.equal(before[name], p) for name, p in model.model.named_parameters()) + result = model.forward_only(inputs=copy.deepcopy(inputs), task=task, return_logits=True) + output = result['embeddings' if task == 'embedding' else 'logits'] + assert output.shape[0] == len(inputs) and torch.isfinite(output).all() + + +@pytest.mark.parametrize('task', ['embedding', 'reranker', 'seq_cls']) +def test_head_lora_tenants_and_cold_optimizer_resume(tmp_path, monkeypatch, task): + from peft import LoraConfig + from safetensors.torch import load_file + from transformers import AutoTokenizer + from tests.server.integration.cs_tasks.task_models import encoded_batch, prepare_models, tiny_backbone + from twinkle import Platform, initialize + from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel + + monkeypatch.setenv('ACCELERATE_USE_CPU', 'true') + monkeypatch.setattr(Platform, 'get_local_device', lambda *args, **kwargs: 'cpu') + initialize(mode='local') + paths = prepare_models(tiny_backbone(tmp_path / 'source'), tmp_path / 'bases') + targets = ['q_proj', 'v_proj'] + ([] if task == 'embedding' else ['score']) + batch = encoded_batch(task, AutoTokenizer.from_pretrained(paths[task])) + assert all(row['input_ids'] for row in batch), batch + + def build(): + return MultiLoraTransformersModel(model_id=paths[task], mixed_precision='no', + gradient_checkpointing=False, max_loras=2, max_r=8, + target_modules=targets) + + def add(model, name): + model.add_adapter_to_model(name, LoraConfig(r=4, lora_alpha=8, target_modules=targets), + gradient_accumulation_steps=2) + model.set_processor('InputProcessor', padding_side='right', adapter_name=name) + if task == 'embedding': + model.set_loss('InfonceLoss', temperature=0.1, use_batch=True, adapter_name=name) + elif task == 'reranker': + model.set_loss('PointwiseRerankerLoss', adapter_name=name) + else: + model.set_loss('SeqClsLoss', problem_type='single_label_classification', + num_labels=3, adapter_name=name) + model.set_optimizer('Adam', lr=1e-3, adapter_name=name) + + def read(model, name): + return model.forward_only(inputs=copy.deepcopy(batch), task=task, + return_logits=True, adapter_name=name)[ + 'embeddings' if task == 'embedding' else 'logits'].detach().clone() + + def step(model, name): + model.forward_backward(inputs=copy.deepcopy(batch), task=task, adapter_name=name) + model.clip_grad_and_step(adapter_name=name) + + model = build() + add(model, 'a') + add(model, 'b') + base = {k: p.detach().clone() for k, p in model.model.named_parameters() if 'lora_' not in k} + initial, peer = read(model, 'a'), read(model, 'b') + for _ in range(3): + step(model, 'a') + trained = read(model, 'a') + assert not torch.equal(initial, trained) + torch.testing.assert_close(read(model, 'b'), peer, rtol=0, atol=0) + for key, parameter in model.model.named_parameters(): + if key in base: + assert not parameter.requires_grad + torch.testing.assert_close(parameter, base[key], rtol=0, atol=0) + saved = model.save('trained', output_dir=str(tmp_path), adapter_name='a', + save_optimizer=True, consumed_train_samples=3 * len(batch)) + if task != 'embedding': + head = {k: v for k, v in load_file(saved + '/adapter_model.safetensors').items() if 'score.' in k} + assert head and any(torch.count_nonzero(v) for k, v in head.items() if 'lora_B' in k) + step(model, 'a') + step(model, 'a') + continued = read(model, 'a') + model.remove_adapter('a') + add(model, 'c') + assert not model.optimizer_group['c'].optimizer.state + assert model.optimizer_group['c'].cur_step == 0 + torch.testing.assert_close(read(model, 'c'), initial, rtol=0, atol=0) + model.resume_from_checkpoint(saved, adapter_name='c') + step(model, 'c') + step(model, 'c') + torch.testing.assert_close(read(model, 'c'), continued, rtol=0, atol=1e-6) + cold = build() + add(cold, 'a') + progress = cold.resume_from_checkpoint(saved, adapter_name='a') + assert progress['cur_step'] == 3 + step(cold, 'a') + step(cold, 'a') + torch.testing.assert_close(read(cold, 'a'), continued, rtol=0, atol=1e-6) + with pytest.raises(ValueError, match='modules_to_save'): + cold.add_adapter_to_model('bad', LoraConfig(target_modules=targets, modules_to_save=['score'])) diff --git a/tests/transformers/test_full_checkpoint_resume.py b/tests/transformers/test_full_checkpoint_resume.py index 63ae4b99e..d3153a8a2 100644 --- a/tests/transformers/test_full_checkpoint_resume.py +++ b/tests/transformers/test_full_checkpoint_resume.py @@ -30,7 +30,9 @@ def test_full_resume_restores_weights_and_optional_training_state(tmp_path, weig cur_step=0, gradient_accumulation_steps=1)} result = wrapper.resume_from_checkpoint(str(tmp_path), resume_only_model=weights_only) assert model(torch.ones(1, 1)).item() == 2 - assert result['cur_step'] == 3 + # dev's weights-only contract starts fresh progress as well as fresh optimizer state. + assert result['cur_step'] == (0 if weights_only else 3) + assert result['consumed_train_samples'] == (0 if weights_only else 12) if weights_only: assert not optimizer.state assert wrapper.optimizer_group[''].cur_step == 0 @@ -38,3 +40,76 @@ def test_full_resume_restores_weights_and_optional_training_state(tmp_path, weig torch.testing.assert_close(optimizer.state[model.weight]['exp_avg'], torch.full((1, 1), 0.5)) assert wrapper.optimizer_group[''].cur_step == 3 assert wrapper.optimizer_group[''].gradient_accumulation_steps == 2 + + +@pytest.mark.parametrize('offload_model', [False, True]) +@pytest.mark.parametrize('offload_optimizer', [False, True]) +def test_offload_handle_selects_model_and_optimizer_independently(offload_model, offload_optimizer): + from unittest.mock import Mock + from twinkle.model.offload import offload_record + + wrapper = object.__new__(TransformersModel) + torch.nn.Module.__init__(wrapper) + wrapper.model = torch.nn.Linear(2, 2) + optimizer = torch.optim.AdamW(wrapper.model.parameters()) + wrapper.optimizer_group = {'': SimpleNamespace(optimizer=optimizer)} + wrapper.strategy = SimpleNamespace(offload_to_cpu=Mock()) + wrapper.offload_to_cpu(offload_model=offload_model, offload_optimizer=offload_optimizer) + if not offload_model and not offload_optimizer: + wrapper.strategy.offload_to_cpu.assert_not_called() + assert offload_record(wrapper.strategy) == (False, False) + elif offload_model: + wrapper.strategy.offload_to_cpu.assert_called_once_with( + wrapper.model, optimizer if offload_optimizer else None) + else: + wrapper.strategy.offload_to_cpu.assert_called_once_with(wrapper.model, optimizer, offload_model=False) + + +@pytest.mark.parametrize('strategy_name', ['accelerate', 'native_fsdp']) +@pytest.mark.parametrize('offload_model', [False, True]) +@pytest.mark.parametrize('offload_optimizer', [False, True]) +def test_independent_offload_preserves_optimizer_continuation( + strategy_name, offload_model, offload_optimizer, monkeypatch): + import copy + from unittest.mock import Mock + import twinkle.model.offload as offload + from twinkle.model.transformers.strategy.accelerate import AccelerateStrategy + + torch.manual_seed(42) + model = torch.nn.Linear(2, 2) + optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) + inputs = torch.ones(2, 2) + model(inputs).sum().backward() + optimizer.step() + optimizer.zero_grad() + reference = copy.deepcopy(model) + reference_optimizer = torch.optim.AdamW(reference.parameters(), lr=0.01) + reference_optimizer.load_state_dict(copy.deepcopy(optimizer.state_dict())) + if strategy_name == 'accelerate': + strategy = object.__new__(AccelerateStrategy) + strategy.accelerator = SimpleNamespace(device=torch.device('cpu'), unwrap_model=lambda module, **kw: module) + else: + strategy = NativeFSDPStrategy(device_mesh=None, enable_ep=False) + model.to = Mock(wraps=model.to) + move_state = Mock(wraps=offload.move_optimizer_state) + monkeypatch.setattr(offload, 'move_optimizer_state', move_state) + monkeypatch.setattr(torch.cuda, 'is_available', lambda: False) + strategy.offload_to_cpu(model, optimizer if offload_optimizer else None, offload_model=offload_model) + assert offload.offload_record(strategy) == (offload_model, offload_optimizer) + assert model.to.call_count == int(offload_model) + assert move_state.call_count == int(offload_optimizer) + strategy.reload_to_gpu(model, optimizer) + assert offload.offload_record(strategy) == (False, False) + assert model.to.call_count == 2 * int(offload_model) + assert move_state.call_count == 2 * int(offload_optimizer) + strategy.reload_to_gpu(model, optimizer) + assert model.to.call_count == 2 * int(offload_model) + assert move_state.call_count == 2 * int(offload_optimizer) + for module, optim in [(model, optimizer), (reference, reference_optimizer)]: + module(inputs).sum().backward() + optim.step() + for name, parameter in model.named_parameters(): + torch.testing.assert_close(parameter, dict(reference.named_parameters())[name], rtol=0, atol=0) + for parameter, expected_parameter in zip(model.parameters(), reference.parameters()): + for key, value in optimizer.state[parameter].items(): + torch.testing.assert_close(value, reference_optimizer.state[expected_parameter][key], rtol=0, atol=0) diff --git a/tests/transformers/test_multi_lora_checkpoint_context.py b/tests/transformers/test_multi_lora_checkpoint_context.py index 1b89e3d51..39100bba6 100644 --- a/tests/transformers/test_multi_lora_checkpoint_context.py +++ b/tests/transformers/test_multi_lora_checkpoint_context.py @@ -21,6 +21,7 @@ def test_checkpoint_activates_target_tenant_before_optimizer_state_initializatio wrapper = object.__new__(MultiLoraTransformersModel) wrapper.__dict__['model'] = model wrapper.multi_adapter = manager + wrapper._model_wrapped = True wrapper.strategy = SimpleNamespace(unwrap_model=lambda value: value) with manager.adapter('first'): params = [param for name, param in model.named_parameters() if '.lora_0.' in name] @@ -46,3 +47,39 @@ def checkpoint_optimizer(*args, **kwargs): wrapper.load(str(tmp_path), adapter_name='first', load_optimizer=True) assert model.active_adapter == 'lora_0' assert set(model.peft_config) == {'lora_0', 'lora_1'} + + +def _make_two_tenant_wrapper(): + manager = MultiLora(max_loras=2, max_r=4) + config = LoraConfig(r=4, lora_alpha=8, target_modules=['0']) + model = manager.patch(torch.nn.Sequential(torch.nn.Linear(4, 4)), target_modules=['0'], lora_config=config) + manager.reset_adapter_status() + manager.acquire_lora('first', config) + manager.acquire_lora('second', config) + wrapper = object.__new__(MultiLoraTransformersModel) + wrapper.multi_adapter = manager + return wrapper, model + + +def _lora_params(model): + return {name: param for name, param in model.named_parameters() if 'lora_' in name} + + +@pytest.mark.parametrize('disable_lora', [False, True]) +def test_adapter_context_keeps_every_slot_trainable(disable_lora): + wrapper, model = _make_two_tenant_wrapper() + for tenant in ('second', 'first', 'second'): + with wrapper._adapter_context(tenant, disable_lora=disable_lora): + assert all(param.requires_grad for param in _lora_params(model).values()) + assert all(param.requires_grad for param in _lora_params(model).values()) + + +def test_adapter_context_only_active_slot_receives_gradient(): + wrapper, model = _make_two_tenant_wrapper() + for param in _lora_params(model).values(): + torch.nn.init.normal_(param) + with wrapper._adapter_context('second'): + model(torch.randn(2, 4)).sum().backward() + grads = {name: param.grad for name, param in _lora_params(model).items()} + assert all(grad is not None for name, grad in grads.items() if '.lora_1.' in name) + assert all(grad is None for name, grad in grads.items() if '.lora_0.' in name) diff --git a/tests/transformers/test_multi_lora_distributed_slots.py b/tests/transformers/test_multi_lora_distributed_slots.py new file mode 100644 index 000000000..1542585fd --- /dev/null +++ b/tests/transformers/test_multi_lora_distributed_slots.py @@ -0,0 +1,176 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real two-process CPU regression for resident Multi-LoRA slot switching.""" +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from peft import LoraConfig +from safetensors.torch import save_file + +from twinkle.model.multi_lora import MultiLora +from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel + + +@pytest.mark.parametrize('enabled', [False, True]) +def test_gradient_checkpointing_option_reaches_the_backbone(enabled, monkeypatch): + from transformers import AutoModel, BertConfig + from twinkle import DeviceMesh, Platform, initialize + monkeypatch.setenv('ACCELERATE_USE_CPU', 'true') + monkeypatch.setattr(Platform, 'get_local_device', lambda *args, **kwargs: 'cpu') + initialize(mode='local') + model = MultiLoraTransformersModel( + model_cls=AutoModel, config=BertConfig(vocab_size=16, hidden_size=8, num_hidden_layers=1, + num_attention_heads=2, intermediate_size=16), + device_mesh=DeviceMesh.from_sizes(device_type='cpu', dp_size=1), mixed_precision='no', + max_loras=1, max_r=2, target_modules=['query', 'value'], gradient_checkpointing=enabled) + assert model.model.base_model.model.is_gradient_checkpointing is enabled + + +def _wrapper(module, manager, wrap_model): + wrapper = object.__new__(MultiLoraTransformersModel) + torch.nn.Module.__init__(wrapper) + wrapper.model = module + wrapper.multi_adapter = manager + wrapper._model_wrapped = False + wrapper.optimizer_group = {} + wrapper._enable_expert_parallel = False + wrapper.sp_strategy = None + wrapper.strategy = SimpleNamespace( + capture_pre_ep_state_if_needed=lambda *args, **kwargs: None, + wrap_model=wrap_model, + unwrap_model=lambda model: model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model) + # This linear model does not use expert/sequence parallelism or GDN patches. + for name in ('_maybe_apply_expert_parallel', '_ensure_sp_strategy', '_ensure_gdn_padding_free_patch'): + object.__setattr__(wrapper, name, lambda: None) + return wrapper + + +def _distributed_slots(rank, rendezvous, mode): + dist.init_process_group('gloo', init_method='file://' + rendezvous, rank=rank, world_size=2, + timeout=timedelta(seconds=60)) + try: + torch.set_num_threads(1) + # DDP must synchronize initialization even when workers have different RNG states. + torch.manual_seed(42 + rank) + config = LoraConfig(r=2, lora_alpha=2, target_modules=['0']) + manager = MultiLora(max_loras=2, max_r=2) + module = manager.patch(torch.nn.Sequential(torch.nn.Linear(4, 4)), target_modules=['0'], lora_config=config) + manager.save_initial_weights() + manager.reset_adapter_status() + manager.acquire_lora('first', config) + wrapper = _wrapper(module, manager, lambda model: torch.nn.parallel.DistributedDataParallel( + model, find_unused_parameters=True)) + wrapper._lazy_wrap_model() + assert not any(param.requires_grad for name, param in module.named_parameters() if '.base_layer.' in name) + + def synchronized_weights(tenant): + state = manager.get_state_dict(tenant) + flat = torch.cat([value.detach().flatten() for value in state.values()]) + copies = [torch.empty_like(flat) for _ in range(2)] + dist.all_gather(copies, flat) + torch.testing.assert_close(copies[0], copies[1], rtol=0, atol=1e-7) + return {key: value.detach().clone() for key, value in state.items()} + + for tenant in (('first', 'second') if mode == 'switch' else ('first', 'reused')): + if tenant == 'second': + manager.acquire_lora(tenant, config) + elif tenant == 'reused': + manager.release_lora('first') + manager.acquire_lora(tenant, config) + initial = synchronized_weights(tenant) + with manager.adapter(tenant): + parameters = [p for p in module.parameters() if p.requires_grad] + optimizer = torch.optim.SGD(parameters, lr=.1) + for _ in range(3): + optimizer.zero_grad() + # Different local batches expose missing gradient synchronization. + wrapper.model(torch.ones(2, 4) * (rank + 1)).square().mean().backward() + optimizer.step() + trained = synchronized_weights(tenant) + assert any(not torch.equal(initial[name], value) for name, value in trained.items()) + with manager.adapter(tenant), torch.no_grad(): + fixed = wrapper.model(torch.ones(2, 4)).detach().clone() + # Rank 0 checkpoint must preserve predictions on every DP replica. + checkpoint = [trained if rank == 0 else None] + dist.broadcast_object_list(checkpoint, src=0) + manager.set_state_dict(tenant, checkpoint[0]) + torch.testing.assert_close(wrapper.model(torch.ones(2, 4)), fixed, rtol=0, atol=1e-7) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize('mode', ['switch', 'reuse']) +def test_all_slots_sync_across_dp_ranks_after_switch_and_reuse(tmp_path, mode): + mp.spawn(_distributed_slots, args=(str(tmp_path / 'rendezvous'), mode), nprocs=2, join=True) + + +@pytest.mark.parametrize('options', [{'find_unused_parameters': False}, {'static_graph': True}]) +def test_multi_lora_rejects_static_ddp_reducers(options, monkeypatch): + monkeypatch.setattr(MultiLoraTransformersModel, '_try_init_process_group', lambda _: None) + with pytest.raises(ValueError, match='Multi-LoRA DDP requires'): + MultiLoraTransformersModel(ddp_config=options) + + +@pytest.mark.parametrize('reshard_after_forward', [None, False, 2]) +def test_multi_lora_rejects_retained_or_partial_fsdp_views(reshard_after_forward, monkeypatch): + monkeypatch.setattr(MultiLoraTransformersModel, '_try_init_process_group', lambda _: None) + + def unexpected_strategy(*args): + raise AssertionError('Invalid FSDP config reached strategy initialization') + + monkeypatch.setattr(MultiLoraTransformersModel, '_decide_strategy', unexpected_strategy) + with pytest.raises(ValueError, match='Multi-LoRA FSDP requires reshard_after_forward=True'): + MultiLoraTransformersModel(fsdp_config={'reshard_after_forward': reshard_after_forward}) + + +def test_checkpoint_as_first_operation_preserves_fresh_slot_reset_basis(tmp_path): + config = LoraConfig(r=2, lora_alpha=2, target_modules=['0']) + manager = MultiLora(max_loras=1, max_r=2) + module = manager.patch(torch.nn.Sequential(torch.nn.Linear(4, 4)), target_modules=['0'], lora_config=config) + manager.save_initial_weights() + manager.acquire_lora('first', config) + initial = {name: value.detach().clone() for name, value in manager.get_state_dict('first').items()} + checkpoint = {name: torch.full_like(value, 2) for name, value in initial.items()} + save_file(checkpoint, str(tmp_path / 'adapter_model.safetensors')) + wrapper = _wrapper(module, manager, lambda model: model) + wrapper.optimizer_group = {'first': SimpleNamespace(optimizer=None)} + wrapper.load(str(tmp_path), adapter_name='first') + for name, value in manager.get_state_dict('first').items(): + torch.testing.assert_close(value, checkpoint[name]) + # Further operations must not replace the initial reset basis with trained weights. + wrapper._lazy_wrap_model() + manager.release_lora('first') + manager.acquire_lora('replacement', config) + for name, value in manager.get_state_dict('replacement').items(): + torch.testing.assert_close(value, initial[name]) + + +def test_slot_registration_never_unfreezes_base_parameters_with_lora_in_their_name(): + from collections import OrderedDict + config = LoraConfig(r=2, target_modules=['lora_projection']) + manager = MultiLora(max_loras=2, max_r=2) + module = manager.patch(torch.nn.Sequential(OrderedDict([ + ('lora_projection', torch.nn.Linear(4, 4)), + ])), target_modules=['lora_projection'], lora_config=config) + base = {name: parameter for name, parameter in module.named_parameters() if '.base_layer.' in name} + assert base + # Respect an explicitly frozen base even when its module name contains lora_. + for parameter in base.values(): + parameter.requires_grad_(False) + manager.reset_adapter_status() + mask = {name: parameter.requires_grad for name, parameter in module.named_parameters()} + with pytest.raises(RuntimeError, match='wrapper failed'): + with manager.trainable_slots(): + assert all(parameter.requires_grad for _, parameter in manager._named_slot_parameters(module)) + assert not any(parameter.requires_grad for parameter in base.values()) + raise RuntimeError('wrapper failed') + assert {name: parameter.requires_grad for name, parameter in module.named_parameters()} == mask + manager.acquire_lora('first', config) + wrapper = _wrapper(module, manager, lambda model: model) + for disable_lora in (False, True): + with wrapper._adapter_context('first', disable_lora=disable_lora): + assert not any(parameter.requires_grad for parameter in base.values()) + assert not any(parameter.requires_grad for parameter in base.values()) diff --git a/tests/transformers/test_optimizer_group_regression.py b/tests/transformers/test_optimizer_group_regression.py new file mode 100644 index 000000000..d03d16927 --- /dev/null +++ b/tests/transformers/test_optimizer_group_regression.py @@ -0,0 +1,77 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Optimizer groups must not depend on DDP wrapping or repeated GA kwargs.""" +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus +from twinkle.model.transformers.transformers import OptimizerGroup, TransformersModel + + +@pytest.mark.parametrize('group_cls', [BaseOptimizerGroup, OptimizerGroup]) +@pytest.mark.parametrize('is_training', [True, False]) +def test_forward_kwargs_do_not_duplicate_metric_control_parameters(group_cls, is_training): + metric = Mock(_twinkle_loss_key=None) + kwargs = {'gradient_accumulation_steps': 99, 'lr': 99, 'step': 99, 'grad_norm': 99, + 'loss_reduction': 'mean', 'custom_metric_arg': 7} + original_kwargs = dict(kwargs) + status = TrainStatus(inputs=[{}], outputs={'loss': 1.0}, metrics=[metric], forward_kwargs=kwargs, + loss_reduction='sum') + group = group_cls(cur_step=3, gradient_accumulation_steps=2, _last_grad_norm=1.25) + if is_training: + group.train_status = status + else: + group.eval_status = status + group.accumulate_metrics(is_training) + metric.accumulate.assert_called_once_with( + status.inputs, status.outputs, gradient_accumulation_steps=2, lr=[], step=2, grad_norm=1.25, + loss_reduction='sum', custom_metric_arg=7) + assert kwargs == original_kwargs + + +@pytest.mark.parametrize('group_cls', [BaseOptimizerGroup, OptimizerGroup]) +@pytest.mark.parametrize('loss_reduction', [None, 'sum']) +def test_shared_metric_arguments_preserve_keyed_loss_dispatch(group_cls, loss_reduction): + selected, unrelated, default = [Mock(_twinkle_loss_key=key) for key in ['selected', 'other', None]] + kwargs = {'_metric_loss_key': 'selected', 'step': 99, 'loss_reduction': 'none'} + status = TrainStatus(inputs=[{}], outputs={'loss': 1.0}, metrics=[selected, unrelated, default], + forward_kwargs=kwargs, loss_reduction=loss_reduction) + group = group_cls(cur_step=3) + group.set_loss(SimpleNamespace(reduction='mean'), loss_key='selected') + group.train_status = status + group.accumulate_metrics(True, include_default=False, metric_step=17) + selected.accumulate.assert_called_once_with( + status.inputs, status.outputs, step=17, loss_reduction=loss_reduction or 'mean', + lr=[], gradient_accumulation_steps=1, grad_norm=0.0) + unrelated.accumulate.assert_not_called() + default.accumulate.assert_not_called() + assert kwargs == {'_metric_loss_key': 'selected', 'step': 99, 'loss_reduction': 'none'} + + +def test_parameter_groups_and_optimizer_restore_are_stable_after_wrapping(): + model = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.LayerNorm(4)) + + class Wrapper(torch.nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + harness = SimpleNamespace( + model=model, + strategy=SimpleNamespace(unwrap_model=lambda wrapped: getattr(wrapped, 'module', wrapped)), + _get_trainable_parameters=lambda adapter: dict(model.named_parameters()), + ) + original_groups = TransformersModel._create_param_group(harness, '', weight_decay=0.1) + optimizer = torch.optim.Adam(original_groups, lr=1e-4) + model(torch.ones(2, 4)).sum().backward() + optimizer.step() + state = optimizer.state_dict() + harness.model = Wrapper(model) + late_groups = TransformersModel._create_param_group(harness, '', weight_decay=0.1) + assert [g['param_names'] for g in late_groups] == [g['param_names'] for g in original_groups] + assert original_groups[0]['param_names'] == ['0.weight'] + restored = torch.optim.Adam(late_groups, lr=1e-4) + restored.load_state_dict(state) + assert len(restored.state) == len(optimizer.state) diff --git a/tests/transformers/test_resident_training_reset.py b/tests/transformers/test_resident_training_reset.py new file mode 100644 index 000000000..885e4bd8d --- /dev/null +++ b/tests/transformers/test_resident_training_reset.py @@ -0,0 +1,233 @@ +import ast +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from twinkle.model.transformers.transformers import TransformersModel + + +def test_resident_reset_removes_previous_templates_forward_hook(): + class ScalingTemplate: + def __init__(self, scale): + self.scale = scale + + def pre_forward_hook(self, model, args, kwargs): + return (args[0] * self.scale,), kwargs + + def group(scale): + return SimpleNamespace(template=ScalingTemplate(scale), _handler=None) + + model = torch.nn.Linear(1, 1, bias=False) + model.weight.data.fill_(1) + wrapper = object.__new__(TransformersModel) + wrapper.__dict__['model'] = model + wrapper.model_id = 'local-reset-test' + wrapper._model_wrapped = True + wrapper.strategy = SimpleNamespace(unwrap_model=lambda value: value, + load_full_state_dict=lambda value, state: value.load_state_dict(state)) + wrapper.optimizer_group = {'': group(3)} + object.__setattr__(wrapper, '_construct_default_optimizer_group', lambda: group(1)) + wrapper.register_mm_forward_hook(wrapper.optimizer_group['']) + assert model(torch.ones(1, 1)).item() == 3 + initial = {'weight': torch.ones(1, 1)} + with patch('twinkle.model.transformers.transformers._read_hf_state_dict', return_value=initial): + for _ in range(2): + model.weight.grad = torch.ones_like(model.weight) + wrapper.reload_initial_weights(reset_training_state=True) + assert model(torch.ones(1, 1)).item() == 1 + assert model.weight.grad is None + assert len(model._forward_pre_hooks) == 1 + + +def test_blank_model_still_resets_training_state(): + wrapper = object.__new__(TransformersModel) + wrapper.model_id = '' + wrapper._model_wrapped = False + previous = SimpleNamespace(template=None, _handler=None, optimizer=object()) + fresh = SimpleNamespace(template=None, _handler=None, optimizer=None) + wrapper.optimizer_group = {'': previous} + object.__setattr__(wrapper, '_construct_default_optimizer_group', lambda: fresh) + wrapper.reload_initial_weights(reset_training_state=True) + assert wrapper.optimizer_group == {'': fresh} + assert wrapper.active_group == '' and fresh.adapter_name == '' + + +@pytest.mark.parametrize('reload_kwargs', [{}, {'reset_training_state': False}], ids=['default', 'explicit-false']) +def test_transformers_reload_preserves_training_progress_by_default(reload_kwargs): + model = torch.nn.Linear(1, 1, bias=False) + model.weight.data.fill_(3) + model.weight.grad = torch.ones_like(model.weight) + previous = SimpleNamespace(optimizer=object(), cur_step=7) + wrapper = object.__new__(TransformersModel) + wrapper.__dict__['model'] = model + wrapper.model_id = 'local-reset-test' + wrapper.optimizer_group = {'': previous} + wrapper.strategy = SimpleNamespace(load_full_state_dict=lambda value, state: value.load_state_dict(state)) + initial = {'weight': torch.ones_like(model.weight)} + with patch('twinkle.model.transformers.transformers._read_hf_state_dict', return_value=initial): + wrapper.reload_initial_weights(**reload_kwargs) + torch.testing.assert_close(model.weight, initial['weight']) + assert model.weight.grad is None + assert wrapper.optimizer_group[''] is previous and previous.cur_step == 7 + + +def test_transformers_reload_propagates_checkpoint_failure(): + wrapper = object.__new__(TransformersModel) + wrapper.model_id = 'unreadable-checkpoint' + with patch('twinkle.model.transformers.transformers._read_hf_state_dict', + side_effect=OSError('checkpoint unavailable')): + with pytest.raises(OSError, match='checkpoint unavailable'): + wrapper.reload_initial_weights() + + +@pytest.mark.parametrize('reload_kwargs', [{}, {'reset_training_state': False}, {'reset_training_state': True}], + ids=['default', 'explicit-false', 'reset']) +def test_megatron_reload_resets_optimizer_state_only_when_requested(reload_kwargs): + # Exercise the repository method without importing the optional Megatron stack. + path = Path('src/twinkle/model/megatron/megatron.py') + model_class = next(node for node in ast.parse(path.read_text()).body + if isinstance(node, ast.ClassDef) and node.name == 'MegatronModel') + method = next(node for node in model_class.body + if isinstance(node, ast.FunctionDef) and node.name == 'reload_initial_weights') + method.decorator_list = [] + namespace = {'_default_adapter_name': '', 'dist': SimpleNamespace(is_initialized=lambda: False)} + exec(compile(ast.Module(body=[method], type_ignores=[]), str(path), 'exec'), namespace) + model = torch.nn.Linear(1, 1, bias=False) + model.weight.grad = torch.ones_like(model.weight) + original_weights = {'weight': torch.ones_like(model.weight)} + + def load_weights(chunks, checkpoint, *, peft_format): + assert checkpoint == 'initial-checkpoint' and peft_format is False + for chunk in chunks: + chunk.load_state_dict(original_weights) + + previous = SimpleNamespace(optimizer=object(), cur_step=7) + fresh = SimpleNamespace(optimizer=None, cur_step=0) + wrapper = SimpleNamespace( + model=[model], _model_path='initial-checkpoint', + optimizer_group={'': previous}, + strategy=SimpleNamespace(unwrap_model=lambda value: value, + bridge=SimpleNamespace(load_weights=load_weights)), + _is_model_ddp_wrapped=lambda: False, + _construct_default_optimizer_group=lambda: fresh, + ) + namespace['reload_initial_weights'](wrapper, **reload_kwargs) + assert model.weight.grad is None + torch.testing.assert_close(model.weight, original_weights['weight']) + assert wrapper.optimizer_group[''] is (fresh if reload_kwargs.get('reset_training_state') else previous) + if reload_kwargs.get('reset_training_state'): + assert fresh.adapter_name == '' and fresh.optimizer is None and fresh.cur_step == 0 + + +def _megatron_dpo_loss_callback(*, forward_only, disable_lora, kwargs=None, ref_logps=None, return_logits=False): + from twinkle.loss import DPOLoss + + path = Path(__file__).resolve().parents[2] / 'src/twinkle/model/megatron/megatron.py' + model_class = next(node for node in ast.parse(path.read_text()).body + if isinstance(node, ast.ClassDef) and node.name == 'MegatronModel') + method = next(node for node in model_class.body + if isinstance(node, ast.FunctionDef) and node.name == 'forward_backward') + compute_loss = next(node for node in method.body if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == 'compute_loss' for target in node.targets)) + callback = next(node for node in method.body + if isinstance(node, ast.FunctionDef) and node.name == 'post_loss_function') + namespace = { + 'torch': torch, 'ModelOutput': dict, + 'Platform': SimpleNamespace(get_local_device=lambda: 'cpu'), + 'forward_only': forward_only, 'disable_lora': disable_lora, 'kwargs': dict(kwargs or {}), + '_mb_counter': [0], 'return_logits': return_logits, 'loss_instance': DPOLoss(), + 'loss_extra_kwargs_per_mb': [{} if ref_logps is None else {'ref_logps': ref_logps}], + 'self': SimpleNamespace(strategy=SimpleNamespace( + reduce_loss=lambda loss, count, logits, logps: [loss, count, {'logits': logits, 'logps': logps}])), + } + exec(compile(ast.Module(body=[compute_loss, callback], type_ignores=[]), str(path), 'exec'), namespace) + return namespace['post_loss_function'] + + +def test_megatron_adapter_disabled_reference_does_not_require_dpo_references(): + callback = _megatron_dpo_loss_callback(forward_only=True, disable_lora=True) + logps = torch.tensor([[-1., -2.], [-3., -4.]]) + result = callback(torch.zeros(2, 2, 4), {'labels': torch.ones(2, 2, dtype=torch.long)}, logps) + assert result[0].item() == 0 and result[1].item() == 0 + assert result[2]['logps'] is logps and result[2]['logits'] is None + + +@pytest.mark.parametrize('forward_only,disable_lora,kwargs', [ + (False, False, {}), (False, True, {}), (True, False, {}), (True, True, {'compute_loss': True}), +]) +def test_megatron_training_and_explicit_loss_probes_still_require_dpo_references( + forward_only, disable_lora, kwargs): + callback = _megatron_dpo_loss_callback(forward_only=forward_only, disable_lora=disable_lora, kwargs=kwargs) + with pytest.raises(ValueError, match='requires ref_logps'): + callback(torch.zeros(2, 2, 4), {'labels': torch.ones(2, 2, dtype=torch.long)}, + torch.tensor([[-1., -2.], [-3., -4.]])) + + +@pytest.mark.parametrize('forward_only', [False, True]) +def test_megatron_non_reference_dpo_still_computes_a_nonzero_loss(forward_only): + callback = _megatron_dpo_loss_callback(forward_only=forward_only, disable_lora=False, + ref_logps=torch.zeros(2, 2)) + result = callback(torch.zeros(2, 2, 4), {'labels': torch.ones(2, 2, dtype=torch.long)}, + torch.tensor([[-1., -2.], [-3., -4.]])) + assert result[0].item() > 0 and result[1].item() == 1 + + +def test_megatron_empty_probe_reports_pooled_scores_not_token_logits(): + callback = _megatron_dpo_loss_callback(forward_only=True, disable_lora=True, return_logits=True) + pooled = torch.randn(2, 3) + result = callback(torch.randn(2, 7, 3), {}, None, unpacked_logits=pooled) + torch.testing.assert_close(result[2]['logits'], pooled) + assert result[0].item() == 0 and result[1].item() == 0 + + +def test_megatron_empty_probe_reports_canonical_embeddings(): + callback = _megatron_dpo_loss_callback(forward_only=True, disable_lora=True) + embeddings = torch.randn(2, 8, requires_grad=True) + result = callback(embeddings, {}, None, embeddings=embeddings) + torch.testing.assert_close(result[2]['embeddings'], embeddings) + assert not result[2]['embeddings'].requires_grad + + +def test_megatron_empty_logits_are_omitted_before_dp_collection(): + path = Path(__file__).resolve().parents[2] / 'src/twinkle/model/megatron/megatron.py' + model_class = next(node for node in ast.parse(path.read_text()).body + if isinstance(node, ast.ClassDef) and node.name == 'MegatronModel') + method = next(node for node in model_class.body + if isinstance(node, ast.FunctionDef) and node.name == 'forward_backward') + normalizer = next(node for node in reversed(method.body) if isinstance(node, ast.If) + and any(isinstance(child, ast.Name) and child.id == 'return_logits' + for child in ast.walk(node.test))) + namespace = {'logits': [], 'return_logits': True, 'torch': torch} + exec(compile(ast.Module(body=[normalizer], type_ignores=[]), str(path), 'exec'), namespace) + assert namespace['logits'] is None + + logits = torch.randn(2, 3) + namespace['logits'] = logits + exec(compile(ast.Module(body=[normalizer], type_ignores=[]), str(path), 'exec'), namespace) + assert namespace['logits'] is logits + + +def test_megatron_seq_cls_lora_checkpoint_opens_bridge_head_load_gate(): + from twinkle.model.megatron.multi_lora_megatron import _seq_cls_head_lora_load_gate + + class Bridge: + _peft_format = True + config = SimpleNamespace(task_type='seq_cls') + hf_score_key = 'score.weight' + + def __init__(self): + self.loaded_keys = None + + def _convert_post_process(self, mg_model, hf_state_dict, hf_prefix, to_mcore): + self.loaded_keys = set(hf_state_dict) + + bridge = Bridge() + checkpoint = {'score.lora_A.weight': object(), 'score.lora_B.weight': object()} + with _seq_cls_head_lora_load_gate(bridge): + bridge._convert_post_process(None, checkpoint, '', True) + assert bridge.loaded_keys == {'score.weight', 'score.lora_A.weight', 'score.lora_B.weight'} + assert 'score.weight' not in checkpoint + assert bridge._convert_post_process.__self__ is bridge diff --git a/tests/transformers/test_sentence_transformer_tenants.py b/tests/transformers/test_sentence_transformer_tenants.py new file mode 100644 index 000000000..d13cf6e99 --- /dev/null +++ b/tests/transformers/test_sentence_transformer_tenants.py @@ -0,0 +1,108 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Real offline ST pipeline: tenant isolation, optimizer continuation and slot reuse.""" +import copy + +import pytest +import torch +from peft import LoraConfig, PeftModel + +st = pytest.importorskip('sentence_transformers') +pytest.importorskip('swift.dev.model.multi_lora_sentence_transformer_model') + + +def test_st_pipeline_tenants_checkpoint_and_reuse(tmp_path, monkeypatch): + from sentence_transformers.base.modules import Transformer, Dense, Normalize + from sentence_transformers.sentence_transformer.modules import Pooling + from tokenizers import Tokenizer + from tokenizers.models import WordPiece + from transformers import BertConfig, BertModel, PreTrainedTokenizerFast + from swift.dev.model.multi_lora_sentence_transformer_model import MultiLoraSentenceTransformerModel + from twinkle import Platform, initialize + + monkeypatch.setenv('ACCELERATE_USE_CPU', 'true') + monkeypatch.setattr(Platform, 'get_local_device', lambda *args, **kwargs: 'cpu') + initialize(mode='local') + torch.manual_seed(42) + backbone = tmp_path / 'bert' + BertModel(BertConfig(vocab_size=32, hidden_size=16, num_hidden_layers=1, num_attention_heads=2, + intermediate_size=32, hidden_dropout_prob=0, attention_probs_dropout_prob=0)).save_pretrained(backbone) + vocab = {'[PAD]': 0, '[UNK]': 1, '[CLS]': 2, '[SEP]': 3, '[MASK]': 4} + vocab.update({f't{i}': i for i in range(5, 32)}) + PreTrainedTokenizerFast(tokenizer_object=Tokenizer(WordPiece(vocab, unk_token='[UNK]')), + pad_token='[PAD]', unk_token='[UNK]').save_pretrained(backbone) + original = tmp_path / 'st' + st.SentenceTransformer(modules=[Transformer(str(backbone)), Pooling(16), + Dense(16, 8), Normalize()], device='cpu').save(str(original)) + module_config = (original / 'modules.json').read_bytes() + model = MultiLoraSentenceTransformerModel(model_id=str(original), mixed_precision='no', + gradient_checkpointing=False, max_loras=2, max_r=4, + target_modules=['query', 'value']) + batch = [{'input_ids': [2, 5 + i, 3], 'attention_mask': [1] * 3, 'labels': [int(i % 3 == 0)]} + for i in range(6)] + + def add(name, lr): + model.add_adapter_to_model(name, LoraConfig(r=4, target_modules=['query', 'value'])) + model.set_processor('InputProcessor', padding_side='right', adapter_name=name) + model.set_loss('InfonceLoss', temperature=0.1, adapter_name=name) + model.set_optimizer('AdamW', lr=lr, adapter_name=name) + + def state(name): + return {k: v.detach().cpu().clone() for k, v in model.get_state_dict(adapter_name=name).items()} + + def vectors(name, **kwargs): + return model.forward_only(inputs=copy.deepcopy(batch), adapter_name=name, **kwargs)['embeddings'].cpu() + + def step(name): + model.forward_backward(inputs=copy.deepcopy(batch), adapter_name=name, task='embedding') + model.clip_grad_and_step(adapter_name=name) + + add('a', 1e-3) + add('b', 2e-3) + group_a, group_b = model.optimizer_group['a'], model.optimizer_group['b'] + template_a, template_b = group_a.template, group_b.template + assert template_a is not template_b + model.set_template('Template', max_length=17, adapter_name='b') + assert group_a.template is template_a and group_b.template is not template_b + template_b = group_b.template + a0, b0, b_vectors = state('a'), state('b'), vectors('b') + reference = st.SentenceTransformer(str(original), device='cpu').eval() + features = {'input_ids': torch.tensor([row['input_ids'] for row in batch]), + 'attention_mask': torch.ones(6, 3, dtype=torch.long)} + with torch.no_grad(): + raw = reference(dict(features))['sentence_embedding'] + torch.testing.assert_close(vectors('a', disable_lora=True), raw, rtol=0, atol=1e-6) + step('a') + assert any(not torch.equal(a0[k], v) for k, v in state('a').items()) + for key, value in state('b').items(): + torch.testing.assert_close(value, b0[key], rtol=0, atol=0) + torch.testing.assert_close(vectors('b'), b_vectors, rtol=0, atol=0) + assert not group_b.optimizer.state and group_b.cur_step == 0 + assert group_a.template is template_a and group_b.template is template_b + saved_vectors = vectors('a') + checkpoint = model.save('a', output_dir=str(tmp_path), adapter_name='a', save_optimizer=True, + consumed_train_samples=6) + step('a') + continued = vectors('a') + model.remove_adapter('a') + add('c', 1e-3) + group_c = model.optimizer_group['c'] + assert not group_c.optimizer.state and group_c.cur_step == 0 + assert group_c.template is not template_a and group_c.template is not template_b + for key, value in state('c').items(): + torch.testing.assert_close(value, a0[key], rtol=0, atol=0) + progress = model.resume_from_checkpoint(checkpoint, adapter_name='c') + assert progress['consumed_train_samples'] == 6 and progress['cur_step'] == 1 + assert model.optimizer_group['c'].optimizer.state + torch.testing.assert_close(vectors('c'), saved_vectors, rtol=0, atol=1e-6) + step('c') + torch.testing.assert_close(vectors('c'), continued, rtol=0, atol=1e-6) + torch.testing.assert_close(vectors('b'), b_vectors, rtol=0, atol=0) + assert (original / 'modules.json').read_bytes() == module_config + # A native ST pipeline plus this tenant adapter preserves Dense/Pooling/Normalize. + reference[0].model = PeftModel.from_pretrained(reference[0].model, checkpoint) + reference.eval() + with torch.no_grad(): + reloaded = reference(dict(features))['sentence_embedding'] + torch.testing.assert_close(reloaded, saved_vectors, rtol=0, atol=1e-6) + model.remove_adapter('b') + model.remove_adapter('c') diff --git a/tests/transformers/test_spectral_hybrid_lora.py b/tests/transformers/test_spectral_hybrid_lora.py index 9e9714249..684c60708 100644 --- a/tests/transformers/test_spectral_hybrid_lora.py +++ b/tests/transformers/test_spectral_hybrid_lora.py @@ -759,6 +759,24 @@ def track_set_adapter(adapter_name, inference_mode=False): assert fft_wrapper.active_adapters == [] +def test_hybrid_restores_lora_gradients_if_fft_restore_fails(monkeypatch): + from twinkle.model.transformers.hybrid import SpectralHybridTransformersModel + + _, model, manager, spectral = _make_multi_tenant_hybrid() + service = object.__new__(SpectralHybridTransformersModel) + service.multi_adapter = manager + service.fft_slots = spectral + + def fail_restore(_adapter): + raise RuntimeError('FFT restore failed') + + monkeypatch.setattr(spectral, 'activate_fft_slot', fail_restore) + with pytest.raises(RuntimeError, match='FFT restore failed'): + with service._adapter_context('hybrid', disable_lora=True): + pass + assert all(parameter.requires_grad for name, parameter in model.named_parameters() if 'lora_' in name) + + def test_hybrid_context_switches_between_two_fft_tenants(): from twinkle.model.transformers.hybrid import SpectralHybridTransformersModel @@ -871,11 +889,15 @@ def test_disable_lora_restores_slots_after_operation_fails(): assert all(parameter.requires_grad for name, parameter in model.named_parameters() if 'lora_' in name) -def test_multi_tenant_hybrid_isolation_and_non_destructive_merge(): +@pytest.mark.parametrize('fft_disabled', [False, True]) +def test_multi_tenant_hybrid_isolation_and_non_destructive_merge(fft_disabled): base, model, manager, spectral = _make_multi_tenant_hybrid() inputs = torch.randn(2, 8) with manager.adapter('regular'): regular_before = model(inputs).detach().clone() + if fft_disabled: + # A normal LoRA selection can disable auxiliary FFT wrappers in PEFT. + spectral._get_fft_wrapper('layers.0.self_attn.q_proj').enable_adapters(False) hybrid = manager.find_lora_by_tenant('hybrid') modules = dict(model.named_modules()) @@ -1091,14 +1113,14 @@ def needs_wrapped_optimizer_state(): @staticmethod def save_optimizer_checkpoint(_model, optimizer, output_path, *, param_name_mapping=None): - from twinkle.model.transformers.strategy.optimizer_state import remap_optimizer_state_names + from twinkle.model.state_dict_utils import remap_optimizer_state_names state = optimizer.state_dict() remap_optimizer_state_names(state, param_name_mapping or {}) torch.save(state, output_path) @staticmethod def load_optimizer_checkpoint(_model, optimizer, input_path, *, param_name_mapping=None): - from twinkle.model.transformers.strategy.optimizer_state import remap_optimizer_state_names + from twinkle.model.state_dict_utils import remap_optimizer_state_names state = torch.load(input_path, map_location='cpu', weights_only=False) remap_optimizer_state_names(state, param_name_mapping or {}) optimizer.load_state_dict(state) @@ -1163,6 +1185,7 @@ def load_optimizer_checkpoint(_model, optimizer, input_path, *, param_name_mappi 'cur_step': 3, 'consumed_train_samples': 17, 'gradient_accumulation_steps': 2, + 'component_states': {}, } hybrid.tenant_config.lora_alpha += 1 @@ -1212,14 +1235,14 @@ def needs_wrapped_optimizer_state(): @staticmethod def save_optimizer_checkpoint(_model, optimizer, output_path, *, param_name_mapping=None): - from twinkle.model.transformers.strategy.optimizer_state import remap_optimizer_state_names + from twinkle.model.state_dict_utils import remap_optimizer_state_names state = optimizer.state_dict() remap_optimizer_state_names(state, param_name_mapping or {}) torch.save(state, output_path) @staticmethod def load_optimizer_checkpoint(_model, optimizer, input_path, *, param_name_mapping=None): - from twinkle.model.transformers.strategy.optimizer_state import remap_optimizer_state_names + from twinkle.model.state_dict_utils import remap_optimizer_state_names state = torch.load(input_path, map_location='cpu', weights_only=False) remap_optimizer_state_names(state, param_name_mapping or {}) optimizer.load_state_dict(state) @@ -1303,7 +1326,7 @@ def train_step(wrapper, inputs): uninterrupted_optimizer = uninterrupted.optimizer_group['hybrid'].optimizer.state_dict() resumed_optimizer = resumed.optimizer_group['hybrid'].optimizer.state_dict() # Normalized optimizer groups must match despite different runtime slots. - from twinkle.model.transformers.strategy.optimizer_state import remap_optimizer_state_names + from twinkle.model.state_dict_utils import remap_optimizer_state_names remap_optimizer_state_names( uninterrupted_optimizer, uninterrupted._optimizer_param_name_mapping( diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index 0a092627c..932f1446e 100644 --- a/tests/twinkle_client/test_async_components.py +++ b/tests/twinkle_client/test_async_components.py @@ -41,6 +41,10 @@ def _recorder(calls, result_factory): def post(url, headers=None, data=None, timeout=None, **kwargs): body = json.loads(data) if data else kwargs.get('json') or {} calls.append((url, body)) + if '/model/' in url and url.endswith('/create'): + return _Response({'training_handles': True}) + if url.endswith('/add_adapter_to_model'): + return _Response(_completed({'training_id': 'training-test', 'adapter_name': body['adapter_name']})) return _Response(result_factory(url)) return post @@ -65,7 +69,7 @@ def test_model_forward_backward_sends_multiple_data_refs(monkeypatch) -> None: }})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) refs = [ DataRef(ref_id='data-1', size=2, fields=['train_input']), DataRef(ref_id='data-2', size=2, fields=['train_input']), @@ -91,7 +95,7 @@ def test_model_inline_forward_methods_keep_the_original_endpoints(monkeypatch) - _patch_transport(monkeypatch, calls, lambda url: {} if url.endswith('/create') else _completed({'result': {}})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) inputs = [{'input_ids': [1, 2]}] model.forward(inputs, return_logits=True) model.forward_only(inputs, disable_lora=True) @@ -110,6 +114,35 @@ def test_model_inline_forward_methods_keep_the_original_endpoints(monkeypatch) - assert calls[-1][1]['micro_batch_size'] == 1 +def test_st_model_factory_keeps_owner_transport_and_lazy_creation(monkeypatch) -> None: + from twinkle_client.http import ClientContext, ClientTransport + from twinkle_client.http.context import clear_default_transport, set_default_transport + from twinkle_client.manager import TwinkleClient + + calls = [] + transport = ClientTransport(ClientContext(base_url='http://st-server', api_key='owner', session_id='st-session')) + other = ClientTransport(ClientContext(base_url='http://other-server', api_key='other')) + monkeypatch.setattr(transport._session, 'post', _recorder(calls, lambda url: _completed({'result': {}}))) + # A later default-client change must not redirect this client's model factory. + set_default_transport(other) + try: + with TwinkleClient(transport=transport) as client: + model = client.model('sentence-transformers/all-MiniLM-L6-v2') + assert model._transport is transport and calls == [] + model.add_adapter_to_model('tenant', {'r': 4, 'target_modules': ['query', 'value']}) + assert [url.rsplit('/', 1)[-1] for url, _ in calls] == ['create', 'add_adapter_to_model'] + assert all(url.startswith(transport.context.base_url + '/model/sentence-transformers/all-MiniLM-L6-v2/') + for url, _ in calls) + assert transport.context.api_key == 'owner' and transport.context.session_id == 'st-session' + assert model._training_ids == {'tenant': 'training-test'} + restored = client.model('sentence-transformers/all-MiniLM-L6-v2') + assert restored._transport is transport and restored._training_ids == {} + assert len(calls) == 2 + finally: + clear_default_transport(other) + other.close() + + def test_model_add_adapter_serializes_dict_lora_config(monkeypatch) -> None: """Dict input remains compatible with the strict serialized LoRA wire field.""" from twinkle_client.model import multi_lora_transformers as module @@ -140,7 +173,7 @@ def test_undeclared_forward_arguments_are_routed_to_loss_kwargs(monkeypatch) -> _patch_transport(monkeypatch, calls, lambda url: {} if url.endswith('/create') else _completed({'result': {}})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) model.forward_backward([{'input_ids': [1, 2]}], advantages=[0.5], old_logps=[[-1.0, -2.0]]) _, body = calls[-1] @@ -158,7 +191,7 @@ def test_model_data_plane_forward_uses_a_separate_api(monkeypatch) -> None: }})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) ref = DataRef(ref_id='data-1', size=2, fields=['train_input']) model.forward_from_data_plane(ref, input_field='train_input') @@ -180,7 +213,7 @@ def test_model_data_plane_forward_only_can_append_selected_outputs(monkeypatch) if url.endswith('/create') else _completed({'result': updated_ref.model_dump()})) model = module.MultiLoraTransformersModel('ms://base') - model.adapter_name = 'adapter' + model.add_adapter_to_model('adapter', None) result = model.forward_only_from_data_plane( ref, output_ref=ref, @@ -223,3 +256,169 @@ def test_sampler_async_data_plane_path_returns_reference_without_materializing(m assert url.endswith('/sampler/base/twinkle/sample_to_data_plane') assert body['num_samples'] == 4 assert body['group_ids'] == ['group-1'] + + +def test_old_remove_completion_does_not_clear_new_same_name_binding(): + import threading + from twinkle_client.http import ClientContext, ClientTransport + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + entered, resume = threading.Event(), threading.Event() + class Transport(ClientTransport): + def __init__(self): + super().__init__(ClientContext(base_url='http://server', api_key='owner')) + self.created = 0 + self.operations = [] + def post(self, url, json_data): + return _Response({'training_handles': True}) + def post_model(self, url, body): + self.operations.append((url.rsplit('/', 1)[-1], body.model_dump())) + if url.endswith('/add_adapter_to_model'): + self.created += 1 + return _Response(_completed({'training_id': f'training-{self.created}', 'adapter_name': 'same'})) + if url.endswith('/remove_adapter'): + if not entered.is_set(): + entered.set() + assert resume.wait(3) + return _Response(_completed(None)) + transport = Transport() + model = MultiLoraTransformersModel('small', transport=transport) + model.add_adapter_to_model('same', None) + errors = [] + def remove(): + try: + model.remove_adapter() + except Exception as exc: + errors.append(exc) + releasing = threading.Thread(target=remove) + releasing.start() + try: + assert entered.wait(3) + model.remove_adapter('same') + model.add_adapter_to_model('same', None) + finally: + resume.set() + releasing.join(3) + assert not releasing.is_alive() and not errors + model.zero_grad() + assert transport.operations[-1][1]['training_id'] == 'training-2' + removal = next(body for endpoint, body in transport.operations if endpoint == 'remove_adapter') + assert removal['training_id'] == 'training-1' + + +def test_duplicate_adapter_is_rejected_before_rpc_and_release_allows_recreate(monkeypatch): + import pytest + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed(None)) + first = MultiLoraTransformersModel('base') + first.add_adapter_to_model('same', None) + count = len(calls) + with pytest.raises(RuntimeError, match='already'): + first.add_adapter_to_model('same', None) + assert len(calls) == count + second = MultiLoraTransformersModel('base') + second.add_adapter_to_model('same', None) + first.remove_adapter('same') + first.add_adapter_to_model('same', None) + + +def test_concurrent_duplicate_and_failed_creation_leave_reusable_name(monkeypatch): + import threading + from concurrent.futures import ThreadPoolExecutor + import pytest + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed(None)) + model = MultiLoraTransformersModel('base') + original = model._submit + entered, resume = threading.Event(), threading.Event() + + def fail(*args, **kwargs): + entered.set() + assert resume.wait(3) + raise RuntimeError('create failed') + + monkeypatch.setattr(model, '_submit', fail) + with ThreadPoolExecutor(1) as executor: + future = executor.submit(model.add_adapter_to_model, 'same', None) + try: + assert entered.wait(3) + with pytest.raises(RuntimeError, match='already being created'): + model.add_adapter_to_model('same', None) + finally: + resume.set() + with pytest.raises(RuntimeError, match='create failed'): + future.result() + monkeypatch.setattr(model, '_submit', original) + model.add_adapter_to_model('same', None) + assert model._training_ids['same'] == 'training-test' + + +def test_unknown_training_rejected_before_handshake(monkeypatch): + import pytest + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed({})) + model = MultiLoraTransformersModel('base') + with pytest.raises(RuntimeError, match='add_adapter_to_model'): + model.step() + assert not calls + + +def test_expired_handle_can_be_removed_and_recreated_on_same_object(monkeypatch): + import pytest + from twinkle_client.exceptions import TwinkleHTTPError + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed(None)) + model = MultiLoraTransformersModel('base') + model.add_adapter_to_model('same', None) + submit = model._submit + + def expired(*args, **kwargs): + raise TwinkleHTTPError('Training handle expired', status_code=404, error_code=404) + + monkeypatch.setattr(model, '_submit', expired) + with pytest.raises(TwinkleHTTPError) as error: + model.remove_adapter('same') + assert error.value.status_code == 404 + monkeypatch.setattr(model, '_submit', submit) + model.add_adapter_to_model('same', None) + model.zero_grad() + assert calls[-1][1]['training_id'] == model._training_ids['same'] + + +def test_failed_release_keeps_handle_when_outcome_is_unknown(monkeypatch): + import pytest + from twinkle_client.exceptions import TwinkleHTTPError + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed(None)) + model = MultiLoraTransformersModel('base') + model.add_adapter_to_model('same', None) + identity = model._training_ids['same'] + for error in (TwinkleHTTPError('Unavailable', status_code=503), TimeoutError('ACK unknown')): + def failed(*args, **kwargs): + raise error + monkeypatch.setattr(model, '_submit', failed) + with pytest.raises(type(error)): + model.remove_adapter('same') + assert model._training_ids['same'] == identity + assert model.adapter_name == 'same' + + +def test_sequence_and_handle_selection_are_object_local(monkeypatch): + from twinkle_client.model.multi_lora_transformers import MultiLoraTransformersModel + calls = [] + _patch_transport(monkeypatch, calls, lambda url: _completed({})) + first, second = MultiLoraTransformersModel('base'), MultiLoraTransformersModel('base') + for model in (first, second): + model._created = True + model.adapter_name = 'default' + first._training_ids = {'default': 'first-default', 'other': 'first-other'} + second._training_ids = {'default': 'second-default'} + first.step() + first.step(adapter_name='other') + second.step() + assert [(body['adapter_name'], body['training_id'], body['seq_id']) for _, body in calls] == [ + ('default', 'first-default', 1), ('other', 'first-other', 2), ('default', 'second-default', 1)] diff --git a/tests/twinkle_client/test_auto_agent_tools.py b/tests/twinkle_client/test_auto_agent_tools.py index dcfc871e5..9d5ecc59d 100644 --- a/tests/twinkle_client/test_auto_agent_tools.py +++ b/tests/twinkle_client/test_auto_agent_tools.py @@ -81,6 +81,15 @@ def test_generated_server_config_matches_launcher_schema(tmp_path, monkeypatch, gateway = next(app for app in config.applications if app.import_path == 'server') assert gateway.deployments[0]['name'] == 'GatewayServer' assert gateway.args.supported_models == ['Qwen/student'] + [s['model_id'] for s in samplers] + application = next(app for app in config.applications if app.import_path == 'model') + assert application.name == 'models-student' + assert application.args.train_mode == 'lora' + assert application.args.adapter_config.adapter_timeout == 600 + model = config.training_config.models[0] + assert model.env_vars['TWINKLE_TRUST_REMOTE_CODE'] == '1' + assert model.mode == 'lora' + assert model.config.min_replicas == 0 and model.config.max_replicas == 1 + assert model.config.num_devices == 2 and model.config.idle_timeout_s == 600 if samplers: sampler = next(app for app in config.applications if app.import_path == 'sampler') assert sampler.args.nproc_per_node == 4 @@ -102,6 +111,14 @@ def test_removed_sampler_is_not_advertised_or_written(tmp_path, monkeypatch): assert not (tmp_path / '.cache' / 'twinkle' / 'server_config.yaml').exists() +def test_single_device_megatron_config_is_not_written(tmp_path, monkeypatch): + from pathlib import Path + monkeypatch.setattr(Path, 'home', lambda: tmp_path) + with pytest.raises(ValueError, match='at least 2 devices'): + ToolExecutor._generate_server_config('Qwen/student', train_gpus=1, backend='megatron') + assert not (tmp_path / '.cache' / 'twinkle' / 'server_config.yaml').exists() + + @pytest.mark.parametrize('backend', ['torch', 'missing']) def test_unsupported_training_backend_is_not_written(tmp_path, monkeypatch, backend): from pathlib import Path diff --git a/tests/twinkle_client/test_future_layer.py b/tests/twinkle_client/test_future_layer.py index e078b5388..4569dcbe8 100644 --- a/tests/twinkle_client/test_future_layer.py +++ b/tests/twinkle_client/test_future_layer.py @@ -62,6 +62,15 @@ def test_terminal_submit_failure_raises_taskfailed_with_payload(monkeypatch): assert not isinstance(exc.value, requests.HTTPError) +def test_terminal_failure_keeps_remote_traceback(monkeypatch): + """The client exception must carry the server diagnostic, not only the summary.""" + monkeypatch.setattr(_future, '_post_retrieve', lambda _r, _transport: pytest.fail('no retrieve')) + diagnostic = 'Traceback (most recent call last):\nAssertionError: data/mask length mismatch at sample 0' + with pytest.raises(TaskFailedError) as exc: + _future.resolve(_failed(traceback=diagnostic), model_cls=_Model) + assert exc.value.traceback == diagnostic + + def test_model_cls_none_returns_none_result(monkeypatch): """A method that returned None before still returns None (not swallowed).""" monkeypatch.setattr(_future, '_post_retrieve', lambda _r, _transport: pytest.fail('no retrieve')) diff --git a/tests/twinkle_client/test_request_builder.py b/tests/twinkle_client/test_request_builder.py index 0c7a15bce..f629f108c 100644 --- a/tests/twinkle_client/test_request_builder.py +++ b/tests/twinkle_client/test_request_builder.py @@ -32,7 +32,7 @@ def no_network(monkeypatch): def test_a_declared_name_goes_to_its_field(): - body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', task='embedding') + body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', task='embedding', training_id='training-a') assert body.task == 'embedding' assert body.loss_kwargs == {} @@ -46,7 +46,7 @@ def test_an_undeclared_name_goes_to_the_single_passthrough_region(): }], adapter_name='a', advantages=[0.5], - old_logps=[[-1.0]]) + old_logps=[[-1.0]], training_id='training-a') assert body.loss_kwargs == {'advantages': [0.5], 'old_logps': [[-1.0]]} @@ -71,7 +71,7 @@ class _TwoRegions(StrictRequest): def test_a_model_without_a_region_rejects_an_unknown_name(): with pytest.raises(TwinkleClientValidationError) as raised: - build_request(model_types.SaveRequest, adapter_name='a', name='ckpt', typo=1) + build_request(model_types.SaveRequest, adapter_name='a', name='ckpt', typo=1, training_id='training-a') assert 'typo' in str(raised.value) @@ -81,7 +81,7 @@ def test_an_explicit_region_and_a_routed_key_are_merged(): loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, - loss_type='sigmoid') + loss_type='sigmoid', training_id='training-a') assert body.init_kwargs == {'beta': 0.1, 'loss_type': 'sigmoid'} @@ -89,7 +89,7 @@ def test_a_key_passed_twice_is_an_error(): """Silently letting one win would make the effective value depend on merge order.""" with pytest.raises(TwinkleClientValidationError, match='both directly and inside'): build_request( - model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, beta=0.2) + model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', init_kwargs={'beta': 0.1}, beta=0.2, training_id='training-a') def test_an_omitted_optional_argument_is_not_routed_into_the_region(): @@ -98,7 +98,7 @@ def test_an_omitted_optional_argument_is_not_routed_into_the_region(): Routing an explicit ``None`` would hand the backend a null argument it never received before, changing behaviour for callers who simply did not pass anything. """ - body = build_request(model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', unused=None) + body = build_request(model_types.SetLossRequest, loss_cls='DPOLoss', adapter_name='a', unused=None, training_id='training-a') assert body.init_kwargs == {} @@ -109,28 +109,28 @@ def test_an_omitted_optional_argument_is_not_routed_into_the_region(): def test_a_wrongly_typed_field_fails_in_process(): with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature='hot') + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature='hot', training_id='training-a') def test_an_out_of_range_value_fails_in_process(): with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature=0) + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', temperature=0, training_id='training-a') def test_malformed_inputs_fail_in_process(): """The client shares the server's schema, so it catches this without asking.""" with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [True]}], adapter_name='a') + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [True]}], adapter_name='a', training_id='training-a') def test_a_missing_required_field_fails_in_process(): with pytest.raises(ValidationError): - build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}]) + build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], training_id='training-a') def test_forward_only_requires_an_adapter_context(): with pytest.raises(ValidationError): - build_request(model_types.ForwardOnlyRequest, inputs=[{'input_ids': [1]}]) + build_request(model_types.ForwardOnlyRequest, inputs=[{'input_ids': [1]}], training_id='training-a') # --------------------------------------------------------------------------- # @@ -141,9 +141,9 @@ def test_forward_only_requires_an_adapter_context(): def test_unset_optionals_stay_off_the_wire(): """Absent and "not requested" must look the same, or the server needs its own defaults.""" import json - body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a') + body = build_request(model_types.ForwardRequest, inputs=[{'input_ids': [1]}], adapter_name='a', training_id='training-a') payload = json.loads(request_json(body)) - assert payload == {'inputs': [{'input_ids': [1]}], 'adapter_name': 'a', 'loss_kwargs': {}} + assert payload == {'training_id': 'training-a', 'inputs': [{'input_ids': [1]}], 'adapter_name': 'a', 'loss_kwargs': {}} def test_a_lora_config_is_serialized_to_the_form_the_server_decodes(): diff --git a/tests/twinkle_client/test_types_contract.py b/tests/twinkle_client/test_types_contract.py index 7dbaa25d1..a2d35bbed 100644 --- a/tests/twinkle_client/test_types_contract.py +++ b/tests/twinkle_client/test_types_contract.py @@ -1,7 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Contract-base consistency and naming-disambiguation tests. -- ``QueueStateLiteral`` value set equals the server ``QueueState`` enum. - naming disambiguation guard. The two SDKs already share public names. The contract freezes that legacy set and @@ -12,12 +11,9 @@ import ast import pathlib -import typing import twinkle -from twinkle.server.task_queue.types import QueueState from twinkle.protocol.types import model as model_types -from twinkle.protocol.types.errors import QueueStateLiteral from twinkle.protocol.types.server import GetServerCapabilitiesResponse _TWINKLE_SRC = pathlib.Path(twinkle.__file__).resolve().parent @@ -43,11 +39,6 @@ }) -def test_queue_state_literal_matches_server_enum(): - literal_values = set(typing.get_args(QueueStateLiteral)) - enum_values = {state.value for state in QueueState} - assert literal_values == enum_values, (f'QueueStateLiteral {literal_values} != QueueState {enum_values}') - def test_old_capabilities_response_gets_conservative_defaults(): response = GetServerCapabilitiesResponse.model_validate({'supported_models': []}) diff --git a/tests/utils/test_parallel.py b/tests/utils/test_parallel.py new file mode 100644 index 000000000..6b9c0a0ee --- /dev/null +++ b/tests/utils/test_parallel.py @@ -0,0 +1,146 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Exercise the lock utility with real files, including independent processes.""" +import importlib.util +import os +from pathlib import Path +import subprocess +import sys +import textwrap + +import pytest + +MODULE_PATH = Path(__file__).resolve().parents[2] / 'src/twinkle/utils/parallel.py' + + +def _load_parallel(): + # This utility does not require the training dependencies imported by utils. + spec = importlib.util.spec_from_file_location('twinkle_parallel_test', MODULE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _run(cwd, script, env): + loader = f""" +import importlib.util +spec = importlib.util.spec_from_file_location('twinkle_parallel_test', {str(MODULE_PATH)!r}) +parallel = importlib.util.module_from_spec(spec) +spec.loader.exec_module(parallel) +""" + return subprocess.run( + [sys.executable, '-c', textwrap.dedent(loader) + textwrap.dedent(script)], + cwd=cwd, env=env, capture_output=True, text=True, timeout=20) + + +@pytest.mark.skipif(os.name != 'posix' or os.geteuid() == 0, reason='Requires enforced POSIX permissions') +@pytest.mark.parametrize('custom_home', [False, True]) +def test_locks_work_from_read_only_directory(tmp_path, custom_home): + cwd = tmp_path / 'read_only' + cwd.mkdir() + env = os.environ.copy() + env.pop('TWINKLE_HOME', None) + # A fresh interpreter honors TMPDIR when choosing the default location. + env['TMPDIR'] = str(tmp_path) + if custom_home: + env['TWINKLE_HOME'] = str(tmp_path / 'runtime') + env['TWINKLE_SESSION_ID'] = tmp_path.name + cwd.chmod(0o555) + try: + result = _run(cwd, """ +with parallel.processing_lock('dataset/name'): + pass +assert parallel.try_claim_once('read-only-notification') is True +assert parallel.try_claim_once('read-only-notification') is False +""", env) + assert result.returncode == 0, result.stderr + assert not (cwd / '.locks').exists() + finally: + cwd.chmod(0o755) + + +def test_import_does_not_create_lock_directory(tmp_path, monkeypatch): + runtime = tmp_path / 'runtime' + monkeypatch.setenv('TWINKLE_HOME', str(runtime)) + monkeypatch.chdir(tmp_path) + _load_parallel() + assert not runtime.exists() + assert not (tmp_path / '.locks').exists() + + +def test_claim_is_shared_across_working_directories(tmp_path): + env = os.environ.copy() + env['TWINKLE_HOME'] = str(tmp_path / 'runtime') + env['TWINKLE_SESSION_ID'] = 'first-session' + first = tmp_path / 'first' + second = tmp_path / 'second' + first.mkdir() + second.mkdir() + script = "print(parallel.try_claim_once('same-key', payload='payload', namespace='notify'))" + result = _run(first, script, env) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == 'True' + result = _run(second, script, env) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == 'False' + env['TWINKLE_SESSION_ID'] = 'next-session' + result = _run(second, script, env) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == 'True' + + +def test_processing_lock_waits_across_working_directories(tmp_path, monkeypatch): + monkeypatch.setenv('TWINKLE_HOME', str(tmp_path / 'runtime')) + first = tmp_path / 'first' + second = tmp_path / 'second' + first.mkdir() + second.mkdir() + monkeypatch.chdir(first) + parallel = _load_parallel() + child = None + with parallel.processing_lock('same-dataset'): + loader = f""" +import importlib.util +spec = importlib.util.spec_from_file_location('twinkle_parallel_test', {str(MODULE_PATH)!r}) +parallel = importlib.util.module_from_spec(spec) +spec.loader.exec_module(parallel) +print('ready', flush=True) +with parallel.processing_lock('same-dataset'): + print('entered', flush=True) +""" + child = subprocess.Popen( + [sys.executable, '-c', textwrap.dedent(loader)], cwd=second, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + try: + # Wait for module import before checking that the child blocks. + import selectors + with selectors.DefaultSelector() as selector: + selector.register(child.stdout, selectors.EVENT_READ) + assert selector.select(timeout=10), 'Child did not finish importing' + assert child.stdout.readline().strip() == 'ready' + with selectors.DefaultSelector() as selector: + selector.register(child.stdout, selectors.EVENT_READ) + assert not selector.select(timeout=0.25), 'Child bypassed the held lock' + except BaseException: + child.kill() + child.communicate() + raise + try: + stdout, stderr = child.communicate(timeout=10) + assert child.returncode == 0, stderr + assert stdout.strip() == 'entered' + finally: + if child.poll() is None: + child.kill() + child.communicate() + + +def test_processing_lock_releases_after_exception(tmp_path, monkeypatch): + monkeypatch.setenv('TWINKLE_HOME', str(tmp_path / 'runtime')) + monkeypatch.chdir(tmp_path) + parallel = _load_parallel() + with pytest.raises(ValueError): + with parallel.processing_lock('dataset/name'): + raise ValueError('processing failed') + lock = parallel.FileLock(os.path.join(parallel._LOCK_DIR, 'dataset_name.lock')) + assert parallel.acquire_lock(lock, blocking=False) + parallel.release_lock(lock) diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index 556b8dc83..86744f9b8 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -144,6 +144,106 @@ def test_bfloat16_fallback(self): # bfloat16 resolution rather than float32's. assert torch.allclose(result.float(), expected, atol=5e-2) + def test_bfloat16_does_not_quantize_large_negative_logps(self): + # A peaked row pushes the selected log-prob to ~-100.3, where the + # bfloat16 ULP is 0.5. Computing the softmax in float32 and only then + # selecting keeps the result at float32 resolution. + vocab_size = 4096 + logits = torch.zeros(2, vocab_size, dtype=torch.bfloat16) + logits[:, 0] = 100.0 + logits[:, 1] = -0.3 + index = torch.ones(2, dtype=torch.long) + + result = selective_log_softmax(logits, index) + expected = torch.gather( + logits.double().log_softmax(-1), -1, index.unsqueeze(-1)).squeeze(-1) + + # Computing wide only helps if the result stays wide: narrowing it back to + # bfloat16 would re-quantize to the 0.5 ULP this test exists to catch, and + # no other test in the suite asserts the output dtype. + assert result.dtype == torch.float32 + assert torch.allclose(result.double(), expected, atol=1e-3) + + @pytest.mark.parametrize('dtype', [torch.float32, torch.bfloat16]) + @pytest.mark.parametrize('return_entropy', [False, True]) + def test_backward_saves_no_extra_full_vocab_activation(self, dtype, return_entropy): + # Every tensor handed to save_for_backward stays resident from the end of + # forward until backward runs, so a per-row copy of the vocabulary is as + # expensive as materializing the whole log_softmax. Recomputing in + # backward must keep nothing wider than the labels. + rows, vocab_size = 8, 512 + logits = torch.randn(rows, vocab_size, dtype=dtype, requires_grad=True) + index = torch.randint(0, vocab_size, (rows, )) + logits_storage = logits.untyped_storage().data_ptr() + extra_storages = {} + + def pack(tensor): + storage = tensor.untyped_storage() + if storage.data_ptr() != logits_storage: + extra_storages[storage.data_ptr()] = storage.nbytes() + return tensor + + with torch.autograd.graph.saved_tensors_hooks(pack, lambda tensor: tensor): + result = selective_log_softmax(logits, index, return_entropy=return_entropy) + + outputs = result if return_entropy else (result, ) + torch.autograd.backward(outputs, [torch.ones_like(output) for output in outputs]) + + vocab_bytes = vocab_size * logits.element_size() + assert all(size < vocab_bytes for size in extra_storages.values()) + + def test_scales_with_chunks_not_with_rows(self): + # The per-row Python loop issued one save per row per intermediate, i.e. + # thousands of tiny kernel launches at production sequence lengths. + rows, vocab_size = 512, 512 + logits = torch.randn(rows, vocab_size, requires_grad=True) + index = torch.randint(0, vocab_size, (rows, )) + save_calls = [] + + def pack(tensor): + save_calls.append(tensor) + return tensor + + with torch.autograd.graph.saved_tensors_hooks(pack, lambda tensor: tensor): + result = selective_log_softmax(logits, index) + result.sum().backward() + + assert len(save_calls) <= 32 + + @pytest.mark.parametrize('return_entropy', [False, True]) + def test_gradient_matches_naive_implementation(self, return_entropy): + # The backward is hand-written -- recompute-in-backward for the logps, plus + # a separate term for the entropy -- so it needs a numerical reference. The + # two tests above both run backward but assert only on saved-storage size + # and save count, so they would still pass with a wrong gradient, and a + # wrong gradient is silent: training proceeds and just converges somewhere + # else. + torch.manual_seed(0) + rows, vocab_size = 6, 37 + index = torch.randint(0, vocab_size, (rows, )) + + # float64, so gradcheck's finite differences are not the limiting error. + logits = torch.randn(rows, vocab_size, dtype=torch.float64, requires_grad=True) + assert torch.autograd.gradcheck( + lambda t: selective_log_softmax(t, index, return_entropy=return_entropy), + (logits, ), eps=1e-6, atol=1e-5, rtol=1e-4) + + def naive(t): + log_probs = t.log_softmax(-1) + gathered = torch.gather(log_probs, -1, index.unsqueeze(-1)).squeeze(-1) + if not return_entropy: + return (gathered, ) + return gathered, -(log_probs.exp() * log_probs).sum(-1) + + x_new = torch.randn(rows, vocab_size, requires_grad=True) + x_ref = x_new.detach().clone().requires_grad_(True) + out_new = selective_log_softmax(x_new, index, return_entropy=return_entropy) + out_ref = naive(x_ref) + out_new = out_new if return_entropy else (out_new, ) + torch.autograd.backward(list(out_new), [torch.ones_like(o) for o in out_new]) + torch.autograd.backward(list(out_ref), [torch.ones_like(o) for o in out_ref]) + assert torch.allclose(x_new.grad, x_ref.grad, atol=1e-5) + class TestPadAndStackTensors: