From 15f8c4b1a9c1ac1e934e0c253dcd532369eacf00 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 27 Oct 2026 09:57:10 +0000 Subject: [PATCH 1/6] feat: integrate on-demand training and ST CS on dev-swift-v5 --- .dev_scripts/ci_container_test.sh | 2 +- cookbook/client/async_rl/server_config.yaml | 6 +- .../client/server/megatron/server_config.yaml | 57 +- .../server/megatron/server_config_4b.yaml | 4 +- .../server/transformer/server_config.yaml | 4 +- cookbook/client/swift_model/README-st.md | 68 + .../swift_model/reproduce_st_dependency.py | 24 + .../client/swift_model/requirements-st.txt | 4 + .../sentence_transformer_client.py | 87 ++ .../sentence_transformer_config.yaml | 146 +- cookbook/client/tinker/upload_to_hub.py | 2 - cookbook/client/twinkle/upload_to_hub.py | 2 - .../Usage Guide/Server and Client/Server.md | 61 +- .../\346\234\215\345\212\241\347\253\257.md" | 68 +- pyproject.toml | 7 +- src/twinkle/hub/__init__.py | 20 +- src/twinkle/infra/_ray/ray_helper.py | 251 +-- src/twinkle/infra/_ray/resource_manager.py | 192 ++- src/twinkle/loss/base.py | 15 +- src/twinkle/loss/grpo.py | 489 +++--- src/twinkle/model/megatron/megatron.py | 1104 +++++++------ src/twinkle/model/multi_lora.py | 346 +++-- src/twinkle/model/optimizer_group.py | 107 +- .../model/transformers/hybrid/fft_slots.py | 134 +- .../model/transformers/hybrid/model.py | 341 ++-- .../transformers/multi_lora_transformers.py | 276 ++-- .../model/transformers/transformers.py | 1369 +++++++++-------- src/twinkle/processor/base.py | 33 +- src/twinkle/protocol/types/base.py | 42 +- src/twinkle/protocol/types/errors.py | 29 +- src/twinkle/protocol/types/lifecycle.py | 27 +- src/twinkle/protocol/types/model.py | 108 +- .../sampler/vllm_sampler/vllm_sampler.py | 400 ++--- src/twinkle/server/__init__.py | 12 +- .../server/checkpoint/checkpoint_manager.py | 207 +-- src/twinkle/server/config/__init__.py | 36 +- src/twinkle/server/config/application_spec.py | 180 +-- src/twinkle/server/config/resources.py | 59 + src/twinkle/server/config/server_config.py | 51 +- .../config.py => config/task_queue.py} | 21 +- src/twinkle/server/config/training.py | 276 ++++ src/twinkle/server/deployment.py | 84 +- src/twinkle/server/exceptions.py | 37 +- src/twinkle/server/gateway/app.py | 60 +- src/twinkle/server/gateway/openai_handlers.py | 344 +++-- src/twinkle/server/gateway/proxy.py | 242 --- src/twinkle/server/gateway/routes.py | 26 - src/twinkle/server/gateway/tinker_handlers.py | 245 ++- .../server/gateway/twinkle_handlers.py | 156 +- src/twinkle/server/gateway/use_cases.py | 10 +- src/twinkle/server/launcher/__init__.py | 4 +- .../server/launcher/env_propagation.py | 52 - .../server/launcher/server_launcher.py | 454 +++--- src/twinkle/server/lifecycle/envelope.py | 49 +- src/twinkle/server/lifecycle/protocols.py | 26 +- src/twinkle/server/lifecycle/proxy.py | 137 ++ src/twinkle/server/lifecycle/submit.py | 97 +- src/twinkle/server/model/app.py | 245 ++- src/twinkle/server/model/backends/common.py | 145 +- .../server/model/backends/mock_model.py | 82 +- src/twinkle/server/model/pool/__init__.py | 2 + src/twinkle/server/model/pool/actor.py | 184 +++ src/twinkle/server/model/pool/controller.py | 875 +++++++++++ src/twinkle/server/model/pool/driver.py | 150 ++ src/twinkle/server/model/pool/router.py | 265 ++++ src/twinkle/server/model/routing.py | 75 - src/twinkle/server/model/runtime.py | 221 +++ src/twinkle/server/model/tinker_handlers.py | 457 +++--- src/twinkle/server/model/twinkle_handlers.py | 537 +++---- src/twinkle/server/resources/__init__.py | 2 + src/twinkle/server/resources/allocation.py | 144 ++ src/twinkle/server/resources/instances.py | 123 ++ src/twinkle/server/resources/runtime.py | 128 ++ src/twinkle/server/sampler/app.py | 187 ++- .../server/sampler/backends/__init__.py | 11 +- .../server/sampler/backends/streaming.py | 55 +- src/twinkle/server/sampler/concurrency.py | 24 +- src/twinkle/server/sampler/pool/__init__.py | 2 + src/twinkle/server/sampler/pool/actor.py | 53 + src/twinkle/server/sampler/pool/controller.py | 379 +++++ src/twinkle/server/sampler/pool/driver.py | 121 ++ src/twinkle/server/sampler/pool/router.py | 107 ++ src/twinkle/server/sampler/runtime.py | 159 ++ src/twinkle/server/sampler/tinker_handlers.py | 40 +- .../server/sampler/twinkle_handlers.py | 328 ++-- src/twinkle/server/sampler/weights.py | 23 +- .../server/session_resource/__init__.py | 3 +- .../server/session_resource/adapter.py | 86 -- .../server/session_resource/processor.py | 18 +- src/twinkle/server/state/__init__.py | 65 +- src/twinkle/server/state/backend/base.py | 7 +- src/twinkle/server/state/backend/factory.py | 16 +- .../server/state/backend/memory_backend.py | 40 +- .../server/state/backend/redis_backend.py | 76 +- .../server/state/cleanup_coordinator.py | 120 +- src/twinkle/server/state/future_manager.py | 260 ++-- src/twinkle/server/state/model_manager.py | 45 +- src/twinkle/server/state/models.py | 57 +- src/twinkle/server/state/server_state.py | 256 +-- src/twinkle/server/state/session_manager.py | 14 +- src/twinkle/server/state/training_records.py | 81 + src/twinkle/server/state/training_store.py | 171 ++ src/twinkle/server/task_errors.py | 54 +- src/twinkle/server/task_queue/__init__.py | 21 +- src/twinkle/server/task_queue/backend_gate.py | 20 +- src/twinkle/server/task_queue/mixin.py | 277 ++-- src/twinkle/server/task_queue/types.py | 69 +- src/twinkle/server/task_queue/worker.py | 349 +++-- src/twinkle/server/telemetry/metrics.py | 220 ++- src/twinkle/utils/parallel.py | 89 +- src/twinkle/utils/torch_utils.py | 132 +- src/twinkle_client/_future.py | 1 + src/twinkle_client/auto/agent/server_tools.py | 615 ++++---- src/twinkle_client/exceptions.py | 2 + .../model/multi_lora_transformers.py | 327 ++-- tests/infra/test_ray_registry.py | 62 + tests/loss/test_channel_loss.py | 1 - tests/loss/test_dpo.py | 8 +- tests/loss/test_grpo_gkd.py | 71 + tests/model/test_optimizer_group.py | 43 +- tests/pluger/test_loader.py | 14 +- tests/processor/test_processor.py | 43 + tests/sampler/test_vllm_stream_lifetime.py | 57 + .../server/{utils => checkpoint}/__init__.py | 0 .../test_checkpoint_manager.py | 23 +- tests/server/cli/test_cli.py | 2 +- tests/server/config/server_config_4b_e2e.yaml | 2 + .../config/server_config_4b_e2e_megatron.yaml | 2 + tests/server/config/test_server_config.py | 44 +- .../test_task_queue_config.py} | 2 +- tests/server/config/test_training_config.py | 419 +++++ tests/server/contract/client_api_routes.json | 37 + .../contract/test_client_api_contract.py | 2 +- tests/server/contract/test_error_wire.py | 3 + .../contract/test_protocol_migration.py | 6 +- tests/server/fixtures/server_config_mock.yaml | 2 + .../fixtures/server_config_mock_redis.yaml | 2 + tests/server/gateway/test_openai_handlers.py | 311 +++- tests/server/integration/e2e_helpers.py | 31 + .../on_demand_training/__init__.py | 1 + .../on_demand_training/acceptance_support.py | 171 ++ .../on_demand_training/http_probe.py | 297 ++++ .../on_demand_training/managed_redis.py | 65 + .../on_demand_training/matrix_validation.py | 513 ++++++ .../on_demand_training/on_demand_sdk.py | 18 + .../on_demand_training/stress_validation.py | 102 ++ .../test_resident_release_e2e.py | 190 +++ .../resources}/__init__.py | 0 .../integration/resources/allocation_probe.py | 69 + .../resources/fragmentation_probe.py | 73 + .../integration/resources/test_allocation.py | 31 + tests/server/integration/sampler/__init__.py | 0 .../integration/sampler/sampling_probe.py | 108 ++ .../integration/sampler/test_real_pool_e2e.py | 385 +++++ tests/server/integration/st_cs/__init__.py | 1 + .../server/integration/st_cs/test_http_e2e.py | 30 + tests/server/integration/st_cs/validation.py | 210 +++ .../server/integration/test_actor_recovery.py | 15 +- .../integration/test_blocking_boundary.py | 2 +- .../server/integration/test_full_cycle_e2e.py | 48 +- .../server/integration/test_full_param_e2e.py | 6 +- .../integration/test_mock_mode_startup.py | 139 +- .../test_on_demand_gpu_redis_e2e.py | 53 + .../test_sampler_checkpoints_e2e.py | 8 +- .../test_tinker_rl_correctness_e2e.py | 46 +- tests/server/lifecycle/test_envelope.py | 18 +- .../lifecycle/test_preflight_rejection.py | 2 +- .../lifecycle/test_retrieve_endpoint.py | 3 + .../server/lifecycle/test_run_submit_dedup.py | 34 +- tests/server/lifecycle/test_static_guards.py | 25 - .../server/lifecycle/test_submit_peek_e2e.py | 2 +- tests/server/lifecycle/test_timing_bounds.py | 2 +- .../test_tinker_retrieve_regression.py | 3 + .../model/test_full_mode_exclusivity.py | 104 ++ .../model/test_preflight_loss_inputs.py | 103 ++ tests/server/model/test_replica_lifecycle.py | 13 +- tests/server/model/test_tinker_handlers.py | 12 +- .../server/model/test_twinkle_async_inputs.py | 108 +- tests/server/resources/__init__.py | 0 tests/server/resources/test_allocation.py | 120 ++ tests/server/resources/test_instances.py | 254 +++ tests/server/sampler/test_pool.py | 515 +++++++ tests/server/sampler/test_request_runtime.py | 271 ++++ .../sampler/test_sampling_concurrency.py | 11 +- .../server/sampler/test_stream_guarantees.py | 2 +- tests/server/sampler/test_streaming_bridge.py | 45 + .../server/sampler/test_twinkle_async_rows.py | 2 +- .../server/session_resource/test_contract.py | 7 +- tests/server/start_e2e_server.py | 4 +- tests/server/state/fake_backend.py | 6 +- tests/server/state/test_future_lifecycle.py | 141 +- tests/server/state/test_leader_election.py | 91 +- tests/server/state/test_managers.py | 17 + tests/server/state/test_redis_read_retry.py | 49 + tests/server/state/test_redis_recovery.py | 114 ++ tests/server/state/test_sequence.py | 70 + tests/server/state/test_update_atomic.py | 53 + .../server/static/backend_call_exemptions.py | 12 +- .../static/test_adapter_name_mapping.py | 2 +- tests/server/task_queue/__init__.py | 0 tests/server/task_queue/helpers.py | 59 + .../test_rate_limiter.py | 0 .../test_task_queue_mixin.py | 246 ++- .../server/telemetry/test_training_metrics.py | 259 ++++ tests/server/test_dependencies.py | 115 ++ tests/server/test_gateway_services.py | 5 +- .../test_proxy.py => test_service_proxy.py} | 43 +- tests/server/training/helpers.py | 174 +++ .../training/test_acceptance_harness.py | 124 ++ .../training/test_actor_initialization.py | 129 ++ tests/server/training/test_controller.py | 799 ++++++++++ .../server/training/test_failure_recovery.py | 408 +++++ .../server/training/test_http_integration.py | 39 + tests/server/training/test_launcher.py | 205 +++ tests/server/training/test_router.py | 97 ++ .../training/test_runtime_initialization.py | 49 + tests/server/training/test_runtime_queue.py | 367 +++++ .../server/training/test_storage_recovery.py | 66 + tests/server/validation/test_preflight.py | 45 +- tests/server/validation/test_request_wire.py | 14 +- tests/transformers/test_dev_task_training.py | 50 + .../test_full_checkpoint_resume.py | 77 +- .../test_multi_lora_checkpoint_context.py | 37 + .../test_multi_lora_distributed_slots.py | 176 +++ .../test_optimizer_group_regression.py | 77 + .../test_resident_training_reset.py | 175 +++ .../test_sentence_transformer_tenants.py | 108 ++ .../transformers/test_spectral_hybrid_lora.py | 35 +- tests/twinkle_client/test_async_components.py | 180 ++- tests/twinkle_client/test_auto_agent_tools.py | 17 + tests/twinkle_client/test_future_layer.py | 9 + tests/twinkle_client/test_request_builder.py | 26 +- tests/twinkle_client/test_types_contract.py | 9 - tests/utils/test_parallel.py | 146 ++ tests/utils/test_utils.py | 100 ++ 235 files changed, 20930 insertions(+), 6823 deletions(-) create mode 100644 cookbook/client/swift_model/README-st.md create mode 100644 cookbook/client/swift_model/reproduce_st_dependency.py create mode 100644 cookbook/client/swift_model/requirements-st.txt create mode 100644 cookbook/client/swift_model/sentence_transformer_client.py create mode 100644 src/twinkle/server/config/resources.py rename src/twinkle/server/{task_queue/config.py => config/task_queue.py} (77%) create mode 100644 src/twinkle/server/config/training.py delete mode 100644 src/twinkle/server/gateway/proxy.py delete mode 100644 src/twinkle/server/gateway/routes.py delete mode 100644 src/twinkle/server/launcher/env_propagation.py create mode 100644 src/twinkle/server/lifecycle/proxy.py create mode 100644 src/twinkle/server/model/pool/__init__.py create mode 100644 src/twinkle/server/model/pool/actor.py create mode 100644 src/twinkle/server/model/pool/controller.py create mode 100644 src/twinkle/server/model/pool/driver.py create mode 100644 src/twinkle/server/model/pool/router.py delete mode 100644 src/twinkle/server/model/routing.py create mode 100644 src/twinkle/server/model/runtime.py create mode 100644 src/twinkle/server/resources/__init__.py create mode 100644 src/twinkle/server/resources/allocation.py create mode 100644 src/twinkle/server/resources/instances.py create mode 100644 src/twinkle/server/resources/runtime.py create mode 100644 src/twinkle/server/sampler/pool/__init__.py create mode 100644 src/twinkle/server/sampler/pool/actor.py create mode 100644 src/twinkle/server/sampler/pool/controller.py create mode 100644 src/twinkle/server/sampler/pool/driver.py create mode 100644 src/twinkle/server/sampler/pool/router.py create mode 100644 src/twinkle/server/sampler/runtime.py delete mode 100644 src/twinkle/server/session_resource/adapter.py create mode 100644 src/twinkle/server/state/training_records.py create mode 100644 src/twinkle/server/state/training_store.py create mode 100644 tests/infra/test_ray_registry.py create mode 100644 tests/sampler/test_vllm_stream_lifetime.py rename tests/server/{utils => checkpoint}/__init__.py (100%) rename tests/server/{utils => checkpoint}/test_checkpoint_manager.py (94%) rename tests/server/{utils/task_queue/test_config.py => config/test_task_queue_config.py} (98%) create mode 100644 tests/server/config/test_training_config.py create mode 100644 tests/server/integration/on_demand_training/__init__.py create mode 100644 tests/server/integration/on_demand_training/acceptance_support.py create mode 100644 tests/server/integration/on_demand_training/http_probe.py create mode 100644 tests/server/integration/on_demand_training/managed_redis.py create mode 100644 tests/server/integration/on_demand_training/matrix_validation.py create mode 100644 tests/server/integration/on_demand_training/on_demand_sdk.py create mode 100644 tests/server/integration/on_demand_training/stress_validation.py create mode 100644 tests/server/integration/on_demand_training/test_resident_release_e2e.py rename tests/server/{utils/task_queue => integration/resources}/__init__.py (100%) create mode 100644 tests/server/integration/resources/allocation_probe.py create mode 100644 tests/server/integration/resources/fragmentation_probe.py create mode 100644 tests/server/integration/resources/test_allocation.py create mode 100644 tests/server/integration/sampler/__init__.py create mode 100644 tests/server/integration/sampler/sampling_probe.py create mode 100644 tests/server/integration/sampler/test_real_pool_e2e.py create mode 100644 tests/server/integration/st_cs/__init__.py create mode 100644 tests/server/integration/st_cs/test_http_e2e.py create mode 100644 tests/server/integration/st_cs/validation.py create mode 100644 tests/server/integration/test_on_demand_gpu_redis_e2e.py create mode 100644 tests/server/model/test_full_mode_exclusivity.py create mode 100644 tests/server/model/test_preflight_loss_inputs.py create mode 100644 tests/server/resources/__init__.py create mode 100644 tests/server/resources/test_allocation.py create mode 100644 tests/server/resources/test_instances.py create mode 100644 tests/server/sampler/test_pool.py create mode 100644 tests/server/sampler/test_request_runtime.py create mode 100644 tests/server/sampler/test_streaming_bridge.py create mode 100644 tests/server/state/test_redis_read_retry.py create mode 100644 tests/server/state/test_redis_recovery.py create mode 100644 tests/server/state/test_sequence.py create mode 100644 tests/server/task_queue/__init__.py create mode 100644 tests/server/task_queue/helpers.py rename tests/server/{utils => task_queue}/test_rate_limiter.py (100%) rename tests/server/{utils => task_queue}/test_task_queue_mixin.py (69%) create mode 100644 tests/server/telemetry/test_training_metrics.py create mode 100644 tests/server/test_dependencies.py rename tests/server/{gateway/test_proxy.py => test_service_proxy.py} (61%) create mode 100644 tests/server/training/helpers.py create mode 100644 tests/server/training/test_acceptance_harness.py create mode 100644 tests/server/training/test_actor_initialization.py create mode 100644 tests/server/training/test_controller.py create mode 100644 tests/server/training/test_failure_recovery.py create mode 100644 tests/server/training/test_http_integration.py create mode 100644 tests/server/training/test_launcher.py create mode 100644 tests/server/training/test_router.py create mode 100644 tests/server/training/test_runtime_initialization.py create mode 100644 tests/server/training/test_runtime_queue.py create mode 100644 tests/server/training/test_storage_recovery.py create mode 100644 tests/transformers/test_dev_task_training.py create mode 100644 tests/transformers/test_multi_lora_distributed_slots.py create mode 100644 tests/transformers/test_optimizer_group_regression.py create mode 100644 tests/transformers/test_resident_training_reset.py create mode 100644 tests/transformers/test_sentence_transformer_tenants.py create mode 100644 tests/utils/test_parallel.py 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/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/swift_model/README-st.md b/cookbook/client/swift_model/README-st.md new file mode 100644 index 000000000..49cd1454e --- /dev/null +++ b/cookbook/client/swift_model/README-st.md @@ -0,0 +1,68 @@ +# ST CS 依赖与验证入口 + +本分支将 on-demand 生命周期迁移到 `dev-swift-v5`，并保留 `model_cls` / `model_loader` 从 YAML、训练目录、持久化配置到 worker 的链路。 +ST 使用 Swift 的 `MultiLoraSentenceTransformerModel`，依赖固定在 `requirements-st.txt` 的完整 SHA；不另建训练后端。 + +## 固定版本的梯度检查点限制 + +Swift `7e4203a961416952e1c8ea5b8fb1763139d56bfa` 的多租户 ST 构造函数调用： + +```python +self._enable_gradient_checkpointing(gradient_checkpointing_kwargs) +``` + +其父类 `SentenceTransformerModel._enable_gradient_checkpointing(self)` 不接受该参数，因此默认构造在创建租户前失败： + +```text +TypeError: SentenceTransformerModel._enable_gradient_checkpointing() takes 1 positional argument but 2 were given +``` + +裸库在迁移前 dev 基线和本分支均可复现。最小入口（无需启动 CS）： + +```bash +python cookbook/client/swift_model/reproduce_st_dependency.py \ + --model sentence-transformers/all-MiniLM-L6-v2 +# 离线时 --model 指向已有的 ST 模型目录。 +``` + +ST 首版在服务配置中显式使用 `gradient_checkpointing: false`；这是 Swift 构造函数已有的参数，不改 Swift，也不修改其方法。该选择会增加训练显存占用。将来启用梯度检查点前，需要先修复 Swift 封装并更新固定 SHA。 + +服务示例使用 FP32（`mixed_precision: "no"`）、一个 on-demand 实例及两个租户槽位，模型和采样接口不混用：这里通过训练 handle 的 `forward_only` 读取原生 ST 流水线向量。 + +## 运行示例 + +在已就绪的 Twinkle 设备环境中，用 `--no-deps` 安装 `requirements-st.txt` 的固定 Swift/ST 版本，缺失的 ST 常规依赖单独补齐，然后启动： + +```bash +python -m twinkle.server launch --config cookbook/client/swift_model/sentence_transformer_config.yaml +python cookbook/client/swift_model/sentence_transformer_client.py --base-url http://127.0.0.1:8000 +``` + +服务端离线时将 YAML 的 `model_id` 改为完整 ST 模型缓存目录，保留公开路由名；客户端以 `--tokenizer` 指向相同原始模型目录。恢复必须使用同一原始 ST 模型和模块配置，adapter checkpoint 仅保存租户权重、optimizer 与训练进度。 + +## 验收范围 + +先完成 on-demand 普通模型的设备验收，再运行 ST 链路：冷启动 → 创建两个 adapter → 设置 embedding processor、loss 和 optimizer → 训练与读取向量 → 保存 adapter → 释放并重建 handle → 恢复 optimizer 和进度 → 继续训练 → 释放并复用槽位。 +对照同一原始 ST 模型、模块配置和精度，比较池化、投影、归一化后的向量及参数更新；同时检查租户间 optimizer、进度、模板与权重不串用。 +首版限定 LoRA，不扩展全参、SP 或 padding-free。 + +## 验收状态（2026-10-27） + +合并 `fix/model-state-cleanup` 后，普通 Transformers 模型的真实 PPU CS 专项已重新验证双 SDK 训练、optimizer 保存恢复、Controller 重启、worker 故障恢复、创建去重、共享容量、Full 独占和资源释放；Twinkle、Tinker 恢复后继续训练的 logprob 最大误差均为 0。Redis 重启另有迁移阶段验收记录，本轮未据此宣称重新验收。 +合并后的 Transformers/Megatron × Twinkle/Tinker × SFT/DPO/GRPO 核心矩阵 12/12 通过，无跳过；SFT 双 SDK 保存与 optimizer 恢复通过，两后端均确认训练资源释放。Megatron 的 adapter-disabled 参考前向默认不计算已配置的 DPO loss；普通评估、训练和显式要求计算 loss 的前向语义保持不变。 +本地 ST 回归覆盖 Pooling、Dense、Normalize 流水线与租户 optimizer、进度、模板和槽位隔离。all-MiniLM-L6-v2 的 CPU 裸库训练、adapter 恢复和继续训练也已通过。 + +**真实 ST HTTP / PPU 验收已通过。** 覆盖双租户隔离、训练、保存、槽位复用、optimizer 与进度恢复、实例释放后的冷重建及继续训练。相同原始 MiniLM 模型、模块配置、FP32 精度及 checkpoint/RNG 起点下，CS 与裸库的向量和 adapter 参数最大误差均为 0，向量与参数均发生实际更新。 +设备验收入口单独将 `idle_timeout_s` 设置为 900 秒，以覆盖裸库对照进程的冷导入和设备初始化；服务示例的 60 秒闲置配置保持不变。session 心跳不会续期训练业务闲置期限，长时间离开 CS 做本地计算需要相应配置业务闲置预算。 + +采样池独立 `1→2→1` 设备验收仍待完成，不能用普通训练或 ST 的通过替代。扩展用例中的 Full checkpoint IPC 热更新超时继续单独记录，不纳入本次修复。 + +独立设备验收入口（需要两张已确认空闲的卡；CS 使用物理卡 0，裸库对照进程使用物理卡 1）： + +```bash +TWINKLE_TEST_ST_CS=1 TWINKLE_TEST_ST_MODEL=/absolute/ST/model/cache \ +CUDA_VISIBLE_DEVICES=0 python -m pytest --noconftest \ + tests/server/integration/st_cs/test_http_e2e.py -q +``` + +采样扩缩容入口使用 `TWINKLE_TEST_REAL_POOL=1`、`TWINKLE_TEST_POOL_DEVICES=2`、`TWINKLE_TEST_MODEL_PATH` 和 `TWINKLE_TEST_ARTIFACT_DIR`，选择 `tests/server/integration/sampler/test_real_pool_e2e.py::test_real_pool_scaling`。原有 `test_real_pool_lifecycle` 继续保留 Full 热更新、LoRA 驱逐及故障恢复等扩展断言。 diff --git a/cookbook/client/swift_model/reproduce_st_dependency.py b/cookbook/client/swift_model/reproduce_st_dependency.py new file mode 100644 index 000000000..4ce5df2d2 --- /dev/null +++ b/cookbook/client/swift_model/reproduce_st_dependency.py @@ -0,0 +1,24 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Reproduce the pinned Swift ST constructor failure without HTTP or Ray. + +Pass an existing SentenceTransformer directory to --model for an offline run. +The default downloads the same model used by the planned CS acceptance. +""" + +import argparse + +import twinkle + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", default="sentence-transformers/all-MiniLM-L6-v2") + args = parser.parse_args() + from swift.dev.model.multi_lora_sentence_transformer_model import MultiLoraSentenceTransformerModel + + twinkle.initialize(mode="local") + MultiLoraSentenceTransformerModel(model_id=args.model, mixed_precision="no", max_loras=2, max_r=4) + + +if __name__ == "__main__": + main() diff --git a/cookbook/client/swift_model/requirements-st.txt b/cookbook/client/swift_model/requirements-st.txt new file mode 100644 index 000000000..95497e022 --- /dev/null +++ b/cookbook/client/swift_model/requirements-st.txt @@ -0,0 +1,4 @@ +# Pinned dependency for the dev on-demand/ST CS investigation. +# The constructor blocker is documented in README-st.md. +ms-swift @ git+https://github.com/modelscope/ms-swift.git@7e4203a961416952e1c8ea5b8fb1763139d56bfa +sentence-transformers==6.2.0 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..5a55bd13b --- /dev/null +++ b/cookbook/client/swift_model/sentence_transformer_client.py @@ -0,0 +1,87 @@ +# 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 +from twinkle_client.model import MultiLoraTransformersModel + +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 = MultiLoraTransformersModel(args.model, transport=client.transport) + 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 = MultiLoraTransformersModel(args.model, transport=client.transport) + 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..2977173b0 100644 --- a/cookbook/client/swift_model/sentence_transformer_config.yaml +++ b/cookbook/client/swift_model/sentence_transformer_config.yaml @@ -1,98 +1,62 @@ -# 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. +# Install requirements-st.txt; see README-st.md for the pinned Swift constraint. +# 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/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/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/\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/pyproject.toml b/pyproject.toml index 5ec67ed3c..c8458b017 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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..426576b1c 100644 --- a/src/twinkle/hub/__init__.py +++ b/src/twinkle/hub/__init__.py @@ -1,3 +1,19 @@ # 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..4575a3014 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -1,13 +1,14 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import os -from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Type, TypeVar, Union +from collections.abc import Callable +from typing import Any, Literal, Optional, TypeVar, Union from twinkle import DeviceGroup, Platform, find_free_port, find_node_ip, requires from twinkle.utils.loader import run_worker_setup_hook from .. import _torch_setup from .resource_manager import ResourceManager -T = TypeVar('T') +T = TypeVar("T") class LazyCollect: @@ -34,6 +35,7 @@ def __init__(self, futures, method, collect_func, device_mesh, timeout=None): def _get_result(self): """Internal method to lazily collect and cache results""" import ray + if self._result is None: result = [] for future in self._futures: @@ -79,6 +81,7 @@ def ready(self) -> bool: exception is raised later, at the real :meth:`__call__` collect. """ import ray + refs = [future for future in self._futures if isinstance(future, ray.ObjectRef)] if not refs: return True @@ -112,37 +115,36 @@ def __contains__(self, key): class RayHelper: - - resource_manager: Optional[ResourceManager] = None + resource_manager: ResourceManager | None = None _registry = None - _remote_components: Dict[str, Any] = {} + _remote_components: dict[str, Any] = {} @staticmethod - def _get_ray_custom_resources(device_groups: List[DeviceGroup]) -> Dict[str, float]: + def _get_ray_custom_resources(device_groups: list[DeviceGroup]) -> dict[str, float]: """Return custom accelerator resources needed for local Ray startup. Ray discovers CUDA GPUs itself. Other accelerators, such as Ascend NPUs, need to be registered as Ray custom resources so that placement groups can request them. """ - device_types = {group.device_type.upper() for group in device_groups} - {'CPU'} + device_types = {group.device_type.upper() for group in device_groups} - {"CPU"} # ResourceManager supports one accelerator type per run. Only NPU # currently needs an explicit Ray custom-resource registration. - if device_types != {'NPU'}: + if device_types != {"NPU"}: return {} try: import torch - npu = getattr(torch, 'npu', None) + npu = getattr(torch, "npu", None) npu_count = npu.device_count() if npu is not None and npu.is_available() else 0 except (ImportError, AttributeError, RuntimeError): return {} - return {'NPU': float(npu_count)} if npu_count > 0 else {} + return {"NPU": float(npu_count)} if npu_count > 0 else {} @staticmethod def init_registry(): @@ -161,7 +163,7 @@ def __init__(self): def add_config(self, key: str, value: Any): self.config[key] = value - def add_or_get(self, key: str, value: Any) -> Tuple[bool, Any]: + def add_or_get(self, key: str, value: Any) -> tuple[bool, Any]: """Add or get config, because ray is single threaded.""" if key in self.config: return self.config[key] @@ -174,17 +176,25 @@ 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 - def initialize(nproc_per_node: int, ncpu_proc_per_node: int, device_groups: List[DeviceGroup]): + def initialize(nproc_per_node: int, ncpu_proc_per_node: int, device_groups: list[DeviceGroup]): """Initialize RayHelper. Args: @@ -195,13 +205,14 @@ def initialize(nproc_per_node: int, ncpu_proc_per_node: int, device_groups: List Returns: None """ - requires('ray') + requires("ray") import ray + RayHelper.device_groups = device_groups if not RayHelper.ray_inited(): # Connecting to an existing cluster (e.g. RAY_ADDRESS set for multi-node): custom # resources are declared at ``ray start`` time, and Ray rejects ``resources`` on connect. - connecting = bool(os.environ.get('RAY_ADDRESS')) + connecting = bool(os.environ.get("RAY_ADDRESS")) resources = None if connecting else RayHelper._get_ray_custom_resources(device_groups) ray.init(ignore_reinit_error=True, resources=resources or None) @@ -219,6 +230,7 @@ def teardown(): if RayHelper._registry is not None: import ray + try: ray.get(RayHelper._registry.clear.remote()) ray.kill(RayHelper._registry) @@ -240,11 +252,17 @@ def ray_inited(): def is_worker(): """Check if this process is the worker""" import ray + return RayHelper.ray_inited() and ray._private.worker.global_worker.mode == ray._private.worker.WORKER_MODE @staticmethod - def execute_all_sync(method_name: str, workers_and_args: List[Tuple[Any, List[Any], Dict[str, Any]]], *, - on_submit: Optional[Callable] = None, timeout=None): + def execute_all_sync( + method_name: str, + workers_and_args: list[tuple[Any, list[Any], dict[str, Any]]], + *, + on_submit: Callable | None = None, + timeout=None, + ): """Execute method and return results. ``timeout`` is passed to ``ray.get(list, timeout=)``, whose semantics are @@ -254,11 +272,16 @@ def execute_all_sync(method_name: str, workers_and_args: List[Tuple[Any, List[An are strictly tighter. """ import ray + return ray.get(RayHelper.execute_all_async(method_name, workers_and_args, on_submit=on_submit), timeout=timeout) @staticmethod - def execute_all_async(method_name: str, workers_and_args: List[Tuple[Any, List[Any], Dict[str, Any]]], *, - on_submit: Optional[Callable] = None): + def execute_all_async( + method_name: str, + workers_and_args: list[tuple[Any, list[Any], dict[str, Any]]], + *, + on_submit: Callable | None = None, + ): """Execute method and return futures.""" output = [] for worker_and_args in workers_and_args: @@ -274,16 +297,19 @@ def execute_all_async(method_name: str, workers_and_args: List[Tuple[Any, List[A @staticmethod def add_or_get_config(key: str, value: Any): import ray + return ray.get(RayHelper._registry.add_or_get.remote(key, value)) @staticmethod def add_config(key: str, value: Any): import ray + ray.get(RayHelper._registry.add_config.remote(key, value)) @staticmethod def get_config(key: str): import ray + return ray.get(RayHelper._registry.get_config.remote(key)) @staticmethod @@ -291,6 +317,7 @@ def _get_remote_component(component): """Avoid create remote component twice.""" if component not in RayHelper._remote_components: import ray + RayHelper._remote_components[component] = ray.remote(component) return RayHelper._remote_components[component] @@ -306,11 +333,7 @@ def get_node_address(): return ip, port @staticmethod - def do_get_and_collect_func(collect_func: Callable, - method: Union[str, Callable], - futures, - device_mesh, - timeout=None): + def do_get_and_collect_func(collect_func: Callable, method: str | Callable, futures, device_mesh, timeout=None): """Return the Future-protocol handle over one dispatched call's refs.""" return LazyCollect(futures, method, collect_func, device_mesh, timeout=timeout) @@ -319,14 +342,14 @@ def do_get_and_collect(args, kwargs): """Collect `LazyCollect` in each arg.""" new_args = [] for arg in args: - if isinstance(arg, Callable) and getattr(arg, '_is_lazy_collect', False): + if isinstance(arg, Callable) and getattr(arg, "_is_lazy_collect", False): arg = arg() new_args.append(arg) new_kwargs = {} for key in list(kwargs.keys()): value = kwargs[key] - if isinstance(value, Callable) and getattr(value, '_is_lazy_collect', False): + if isinstance(value, Callable) and getattr(value, "_is_lazy_collect", False): value = value() new_kwargs[key] = value return new_args, new_kwargs @@ -334,24 +357,26 @@ def do_get_and_collect(args, kwargs): @staticmethod def has_ref(args, kwargs) -> bool: for arg in args: - if isinstance(arg, Callable) and getattr(arg, '_is_lazy_collect', False): + if isinstance(arg, Callable) and getattr(arg, "_is_lazy_collect", False): return True for key in list(kwargs.keys()): value = kwargs[key] - if isinstance(value, Callable) and getattr(value, '_is_lazy_collect', False): + if isinstance(value, Callable) and getattr(value, "_is_lazy_collect", False): return True return False @staticmethod - def create_workers(worker_cls: Type[T], - group: str, - execute: Literal['all', 'peer', 'first'], - *args, - instance_id, - extra_torch_kwargs: Optional[Dict[str, Any]] = None, - max_concurrency: Optional[int] = None, - worker_slots: Optional[List[int]] = None, - **kwargs) -> List[T]: + def create_workers( + worker_cls: type[T], + group: str, + execute: Literal["all", "peer", "first"], + *args, + instance_id, + extra_torch_kwargs: dict[str, Any] | None = None, + max_concurrency: int | None = None, + worker_slots: list[int] | None = None, + **kwargs, + ) -> list[T]: # TODO when will remote create remote? # Should it peer create peer? or peer create all? # Whether the input data of each remote is independent, or they are a part of the whole device mesh? @@ -369,14 +394,16 @@ def create_workers(worker_cls: Type[T], assert len(placement_groups) == len(ranks) if worker_slots is not None: slots = tuple(worker_slots) - if any(type(slot) is not int or not 0 <= slot < len(ranks) for slot in slots) or len(set(slots)) != len(slots): - raise ValueError(f'worker_slots must be unique indices in [0, {len(ranks)}), got {slots!r}.') + if any(type(slot) is not int or not 0 <= slot < len(ranks) for slot in slots) or len(set(slots)) != len( + slots + ): + raise ValueError(f"worker_slots must be unique indices in [0, {len(ranks)}), got {slots!r}.") if not slots: return [] placement_groups = [placement_groups[slot] for slot in slots] ranks = [ranks[slot] for slot in slots] - key = f'{group}-{worker_cls.__class__.__name__}-{instance_id}' - if execute == 'peer': + key = f"{group}-{worker_cls.__class__.__name__}-{instance_id}" + if execute == "peer": # Create the peer worker # 0 1 2 3 # | | | | @@ -384,46 +411,42 @@ def create_workers(worker_cls: Type[T], _slice = Platform.get_peer_index(len(ranks)) placement_groups = placement_groups[_slice] ranks = ranks[_slice] - ip, port = RayHelper.get_master_id_port(placement_groups[0]['placement_group']) - ip = RayHelper.add_or_get_config(key + '-ip', ip) - port = RayHelper.add_or_get_config(key + '-port', port) - elif execute == 'first': + ip, port = RayHelper.get_master_id_port(placement_groups[0]["placement_group"]) + ip = RayHelper.add_or_get_config(key + "-ip", ip) + port = RayHelper.add_or_get_config(key + "-port", port) + elif execute == "first": placement_groups = placement_groups[:1] ranks = ranks[:1] - ip, port = RayHelper.get_master_id_port(placement_groups[0]['placement_group']) + ip, port = RayHelper.get_master_id_port(placement_groups[0]["placement_group"]) else: - ip, port = RayHelper.get_master_id_port(placement_groups[0]['placement_group']) + ip, port = RayHelper.get_master_id_port(placement_groups[0]["placement_group"]) - device_type_upper = (device_config.device_type or '').upper() - if device_type_upper != 'CPU': + device_type_upper = (device_config.device_type or "").upper() + if device_type_upper != "CPU": world_size = len(ranks) device_type = Platform.get_platform(device_type_upper).__name__ - for pg_idx, (deploy_pg, gpu) in enumerate(zip(placement_groups, ranks)): - deploy_pg: Dict + for pg_idx, (deploy_pg, _gpu) in enumerate(zip(placement_groups, ranks, strict=False)): + deploy_pg: dict cluster_name = group - worker_name = key + '-' + str(pg_idx) + worker_name = key + "-" + str(pg_idx) env_vars = os.environ.copy() - env_vars.update({ - 'WORLD_SIZE': - str(world_size), - 'RANK': - str(pg_idx), - 'LOCAL_RANK': - str(0), - 'CLUSTER_NAME': - cluster_name, - 'WORKER_NAME': - worker_name, - Platform.get_platform(device_type_upper).visible_device_env(): - ','.join([str(r) for r in deploy_pg['gpu_rank']]), - 'TWINKLE_MODE': - 'ray', - _torch_setup.EXTRA_TORCH_KWARGS_ENV: - _torch_setup.to_env_value(extra_torch_kwargs), - }) - - env_vars['MASTER_ADDR'] = ip - env_vars['MASTER_PORT'] = str(port) + env_vars.update( + { + "WORLD_SIZE": str(world_size), + "RANK": str(pg_idx), + "LOCAL_RANK": str(0), + "CLUSTER_NAME": cluster_name, + "WORKER_NAME": worker_name, + Platform.get_platform(device_type_upper).visible_device_env(): ",".join( + [str(r) for r in deploy_pg["gpu_rank"]] + ), + "TWINKLE_MODE": "ray", + _torch_setup.EXTRA_TORCH_KWARGS_ENV: _torch_setup.to_env_value(extra_torch_kwargs), + } + ) + + env_vars["MASTER_ADDR"] = ip + env_vars["MASTER_PORT"] = str(port) # Prevent Ray from overriding CUDA_VISIBLE_DEVICES set in runtime_env # This is critical for multi-GPU workers (gpus_per_worker > 1) @@ -432,21 +455,22 @@ def create_workers(worker_cls: Type[T], if max_concurrency is not None: # Read back in the worker to size the thread pool that runs # blocking continuous-work methods off the actor's event loop. - env_vars['TWINKLE_ACTOR_MAX_CONCURRENCY'] = str(max_concurrency) + env_vars["TWINKLE_ACTOR_MAX_CONCURRENCY"] = str(max_concurrency) # Ray serialises runtime_env to JSON; a callable setup hook is rejected by Ray 2.58 # ("Object of type function is not JSON serializable"). Pass the import path instead. runtime_env = RuntimeEnv( - env_vars=env_vars, - worker_process_setup_hook='twinkle.utils.loader.run_worker_setup_hook') + env_vars=env_vars, worker_process_setup_hook="twinkle.utils.loader.run_worker_setup_hook" + ) worker_options = { - 'scheduling_strategy': - PlacementGroupSchedulingStrategy(placement_group=deploy_pg['placement_group']), - 'name': worker_name, - 'namespace': 'default', - 'runtime_env': runtime_env, - 'num_cpus': 0.01, + "scheduling_strategy": PlacementGroupSchedulingStrategy( + placement_group=deploy_pg["placement_group"] + ), + "name": worker_name, + "namespace": os.environ.get("TWINKLE_RAY_NAMESPACE", "default"), + "runtime_env": runtime_env, + "num_cpus": 0.01, } # The placement group already reserves a GPU worker's full @@ -460,14 +484,16 @@ def create_workers(worker_cls: Type[T], # the worker in ``set_visible_accelerator_ids`` before its # constructor runs. The CPU claim and placement-group strategy # are sufficient to place the actor in the GPU-owning bundle. - if device_type != 'GPU': + if device_type != "GPU": # Use a custom resource key for non-GPU accelerators # (for example, NPU). - worker_options['resources'] = {device_type: 0.01} + worker_options["resources"] = {device_type: 0.01} if max_concurrency is not None: - worker_options['max_concurrency'] = max_concurrency + 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: @@ -480,44 +506,49 @@ def create_workers(worker_cls: Type[T], workers = [] # For CPU case, don't set visible device environment variables _visible_device_env = {} - for rank, (deploy_pg, index) in enumerate(zip(placement_groups, list(range(world_size)))): - deploy_pg: Dict + for rank, (deploy_pg, _index) in enumerate(zip(placement_groups, list(range(world_size)), strict=False)): + deploy_pg: dict cluster_name = group - worker_name = key + '-' + str(rank) + worker_name = key + "-" + str(rank) env_vars = os.environ.copy() - env_vars.update({ - 'WORLD_SIZE': str(world_size), - 'RANK': str(rank), - 'LOCAL_RANK': str(0), - 'CLUSTER_NAME': cluster_name, - 'WORKER_NAME': worker_name, - 'TWINKLE_MODE': 'ray', - _torch_setup.EXTRA_TORCH_KWARGS_ENV: _torch_setup.to_env_value(extra_torch_kwargs), - **_visible_device_env - }) + env_vars.update( + { + "WORLD_SIZE": str(world_size), + "RANK": str(rank), + "LOCAL_RANK": str(0), + "CLUSTER_NAME": cluster_name, + "WORKER_NAME": worker_name, + "TWINKLE_MODE": "ray", + _torch_setup.EXTRA_TORCH_KWARGS_ENV: _torch_setup.to_env_value(extra_torch_kwargs), + **_visible_device_env, + } + ) if max_concurrency is not None: # Read back in the worker to size the thread pool that runs # blocking continuous-work methods off the actor's event loop. - env_vars['TWINKLE_ACTOR_MAX_CONCURRENCY'] = str(max_concurrency) + env_vars["TWINKLE_ACTOR_MAX_CONCURRENCY"] = str(max_concurrency) # Ray serialises runtime_env to JSON; a callable setup hook is rejected by Ray 2.58 # ("Object of type function is not JSON serializable"). Pass the import path instead. runtime_env = RuntimeEnv( - env_vars=env_vars, - worker_process_setup_hook='twinkle.utils.loader.run_worker_setup_hook') + env_vars=env_vars, worker_process_setup_hook="twinkle.utils.loader.run_worker_setup_hook" + ) worker_options = { - 'scheduling_strategy': - PlacementGroupSchedulingStrategy(placement_group=deploy_pg['placement_group']), - 'name': worker_name, - 'namespace': 'default', - 'runtime_env': runtime_env, - 'num_cpus': 0.01, + "scheduling_strategy": PlacementGroupSchedulingStrategy( + placement_group=deploy_pg["placement_group"] + ), + "name": worker_name, + "namespace": os.environ.get("TWINKLE_RAY_NAMESPACE", "default"), + "runtime_env": runtime_env, + "num_cpus": 0.01, } if max_concurrency is not None: - worker_options['max_concurrency'] = max_concurrency + 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..e229c6916 100644 --- a/src/twinkle/infra/_ray/resource_manager.py +++ b/src/twinkle/infra/_ray/resource_manager.py @@ -1,7 +1,6 @@ # Copyright (c) Alibaba, Inc. and its affiliates. import math import os -from typing import Dict, List from twinkle import DeviceGroup, Platform from twinkle.utils import get_logger @@ -11,7 +10,7 @@ def _gpu_placement_group_cpus(node_cpu: int, nproc_per_node: int) -> int: """Return the logical CPUs reserved by one GPU placement group.""" - cpus_per_proc = int(os.environ.get('TWINKLE_GPU_PG_CPUS_PER_PROC', 4)) + cpus_per_proc = int(os.environ.get("TWINKLE_GPU_PG_CPUS_PER_PROC", 4)) cpus_per_proc = max(cpus_per_proc, 1) required_cpus = max(nproc_per_node * cpus_per_proc, 1) node_cap = max(node_cpu // 4, 1) @@ -19,8 +18,44 @@ 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. - def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[DeviceGroup]): + 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, # this creates a huge PG (e.g., 640 CPU node requests 160 CPU PG). @@ -28,41 +63,43 @@ def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[De # such large CPU PGs may stay pending forever, causing Serve replica's __init__ to hang. # - New strategy: request based on "actual CPU processes needed on node * CPUs per process", # with node_cpu//4 as the upper bound. - cpu_pg_cpus_per_proc = int(os.environ.get('TWINKLE_CPU_PG_CPUS_PER_PROC', 1)) + cpu_pg_cpus_per_proc = int(os.environ.get("TWINKLE_CPU_PG_CPUS_PER_PROC", 1)) cpu_pg_cpus_per_proc = max(cpu_pg_cpus_per_proc, 1) import ray from ray.util.placement_group import PlacementGroup + all_ranks = [] shared_groups = [] last_rank = -1 cpu_proc_count = 0 group_names = [group.name for group in groups] if len(group_names) != len(set(group_names)): - raise ValueError(f'DeviceGroup names must be unique, got {group_names}.') - device_types = {group.device_type.upper() for group in groups} - {'CPU'} + raise ValueError(f"DeviceGroup names must be unique, got {group_names}.") + device_types = {group.device_type.upper() for group in groups} - {"CPU"} assert len(device_types) <= 1 if not device_types: # Pure cpu task - device_type = 'CPU' + device_type = "CPU" else: device_type = next(iter(device_types)) device_type = Platform.get_platform(device_type).__name__ for group in groups: ranks = group.ranks - if (group.device_type or '').upper() == 'CPU': + if (group.device_type or "").upper() == "CPU": if group.shared: raise ValueError(f"CPU DeviceGroup '{group.name}' cannot be shared.") - assert isinstance(ranks, int), 'CPU group only supports integer ranks' + assert isinstance(ranks, int), "CPU group only supports integer ranks" cpu_proc_count += ranks continue if isinstance(ranks, int): if group.shared: raise ValueError( - f"Shared DeviceGroup '{group.name}' requires explicit physical ranks; got ranks={ranks}.") + f"Shared DeviceGroup '{group.name}' requires explicit physical ranks; got ranks={ranks}." + ) ranks = list(range(last_rank + 1, last_rank + 1 + ranks)) ranks = list(ranks) if not ranks: @@ -75,15 +112,16 @@ def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[De last_rank = max(last_rank, max(ranks)) if len(set(all_ranks)) != len(all_ranks): - raise ValueError('Non-shared DeviceGroup ranks must not overlap.') + raise ValueError("Non-shared DeviceGroup ranks must not overlap.") allocated_ranks = set(all_ranks) for group in shared_groups: missing = sorted(set(group.ranks) - allocated_ranks) if missing: raise ValueError( f"Shared DeviceGroup '{group.name}' references unallocated physical ranks {missing}; " - 'shared ranks must be a subset of non-shared DeviceGroup ranks.') - if device_type != 'CPU': + "shared ranks must be a subset of non-shared DeviceGroup ranks." + ) + if device_type != "CPU": # Calculate required nodes based on actual node indices spanned by all_ranks if all_ranks: node_indices = [rank // nproc_per_node for rank in all_ranks] @@ -99,28 +137,29 @@ def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[De self.nodes = [] self.cpu_nodes = [] for node in ray.nodes(): - if not node.get('Alive', True): + if not node.get("Alive", True): continue - resource = node['Resources'] - if device_type != 'CPU': + resource = node["Resources"] + if device_type != "CPU": node_device_num = int(resource.get(device_type, 0)) if node_device_num >= nproc_per_node: self.nodes.append(node) else: - if int(resource.get('CPU', 0)) // 4 >= ncpu_proc_per_node: + if int(resource.get("CPU", 0)) // 4 >= ncpu_proc_per_node: self.nodes.append(node) # Independent CPU node pool: any alive node with enough CPU - node_cpu = int(resource.get('CPU', 0)) + node_cpu = int(resource.get("CPU", 0)) if node_cpu // 4 >= ncpu_proc_per_node: self.cpu_nodes.append(node) - assert self.nnodes <= len( - self.nodes), f'Not enough resources, required nodes: {self.nnodes}, available: {len(self.nodes)}' + assert self.nnodes <= len(self.nodes), ( + f"Not enough resources, required nodes: {self.nnodes}, available: {len(self.nodes)}" + ) bundles = [] cpu_bundles = [] - if device_type != 'CPU': + if device_type != "CPU": for i in range(self.nnodes): # TODO not accurate, because placement_group cannot distribute to node same ordered with self.nodes node_idx = self.min_node_idx + i @@ -128,22 +167,25 @@ def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[De node = self.nodes[node_idx] except IndexError: node = self.nodes[0] - node_cpu = int(node['Resources']['CPU']) - bundles.append({ - device_type: nproc_per_node, - 'CPU': _gpu_placement_group_cpus(node_cpu, nproc_per_node), - }) + node_cpu = int(node["Resources"]["CPU"]) + bundles.append( + { + device_type: nproc_per_node, + "CPU": _gpu_placement_group_cpus(node_cpu, nproc_per_node), + } + ) # CPU placement groups: only create when there are actual CPU processes to allocate. if cpu_proc_count > 0: cpu_nnodes = math.ceil(cpu_proc_count / ncpu_proc_per_node) - assert cpu_nnodes <= len( - self.cpu_nodes), (f'Not enough nodes for CPU processes, required nodes: {cpu_nnodes}, ' - f'available CPU nodes: {len(self.cpu_nodes)}, ensure your ray cluster ' - f'has enough nodes with sufficient CPU resources.') + assert cpu_nnodes <= len(self.cpu_nodes), ( + f"Not enough nodes for CPU processes, required nodes: {cpu_nnodes}, " + f"available CPU nodes: {len(self.cpu_nodes)}, ensure your ray cluster " + f"has enough nodes with sufficient CPU resources." + ) for i in range(cpu_nnodes): node = self.cpu_nodes[i] - node_cpu = int(node['Resources']['CPU']) + node_cpu = int(node["Resources"]["CPU"]) # How many CPU processes will actually be placed on this node # (last node may have fewer than ncpu_proc_per_node) procs_on_node = min( @@ -154,7 +196,7 @@ def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[De # but don't request 160 CPU for just 1~2 processes. node_cap = max(node_cpu // 4, 1) need = max(procs_on_node * cpu_pg_cpus_per_proc, 1) - cpu_bundles.append({'CPU': min(node_cap, need)}) + cpu_bundles.append({"CPU": min(node_cap, need)}) self.cpu_node_map = {} for i in range(cpu_proc_count): @@ -172,9 +214,12 @@ def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[De self.node_ranks = [] if self.placement_groups: - self.node_ranks = ray.get([ - ray.remote(Platform.get_node_rank).options(placement_group=pg).remote() for pg in self.placement_groups - ]) + self.node_ranks = ray.get( + [ + ray.remote(Platform.get_node_rank).options(placement_group=pg).remote() + for pg in self.placement_groups + ] + ) if self.node_ranks.count(0) > 1: self.node_ranks = list(range(len(self.placement_groups))) @@ -187,48 +232,48 @@ def get_visible_devices(): if self.placement_groups: visible_device_futures = [] for pg in self.placement_groups: - probe_options = {'placement_group': pg} - if device_type == 'GPU': + probe_options = {"placement_group": pg} + if device_type == "GPU": # Ask Ray for the GPUs owned by this placement group. A # no-GPU probe with NOSET_* only sees the node-global # CUDA_VISIBLE_DEVICES list, so independently initialized # Model and Sampler groups both select its first entry. # The probe is short-lived and releases the bundle before # component workers are created. - probe_options['num_gpus'] = nproc_per_node + probe_options["num_gpus"] = nproc_per_node else: - probe_options['resources'] = {device_type: nproc_per_node} + probe_options["resources"] = {device_type: nproc_per_node} visible_device_futures.append(get_visible_devices.options(**probe_options).remote()) self.visible_devices = ray.get(visible_device_futures) visible_devices = [] for visible_device in self.visible_devices: if visible_device: - visible_device = [int(device) for device in visible_device.split(',')] + visible_device = [int(device) for device in visible_device.split(",")] else: visible_device = list(range(nproc_per_node)) visible_devices.append(visible_device) self.visible_devices = visible_devices - self.node2pg: Dict[int, PlacementGroup] = {} + self.node2pg: dict[int, PlacementGroup] = {} # Map actual node indices to placement groups # For GPU/NPU groups, node indices start from self.min_node_idx - if device_type != 'CPU': + if device_type != "CPU": for i, placement_group in enumerate(self.placement_groups): actual_node_idx = self.min_node_idx + i self.node2pg[actual_node_idx] = placement_group else: # For CPU-only or when using default node_ranks - for node_rank, placement_group in zip(self.node_ranks, self.placement_groups): + for node_rank, placement_group in zip(self.node_ranks, self.placement_groups, strict=False): self.node2pg[node_rank] = placement_group self.device_groups = {} ray_address = str(ray.get_runtime_context().gcs_address) global_cpu_proc_idx = 0 for group in groups: - if (group.device_type or '').upper() != 'CPU': + if (group.device_type or "").upper() != "CPU": ranks = group.ranks - gpus_per_worker = getattr(group, 'gpus_per_worker', 1) + gpus_per_worker = getattr(group, "gpus_per_worker", 1) local_device_groups = [] # Use original ranks for GPU mapping so each DeviceGroup maps to # the correct physical devices. E.g. ranks=[2,3] with @@ -237,13 +282,15 @@ def get_visible_devices(): if gpus_per_worker > 1: if len(normalized_ranks) % gpus_per_worker != 0: - raise ValueError(f"DeviceGroup '{group.name}': number of ranks ({len(normalized_ranks)}) " - f'must be divisible by gpus_per_worker ({gpus_per_worker})') + raise ValueError( + f"DeviceGroup '{group.name}': number of ranks ({len(normalized_ranks)}) " + f"must be divisible by gpus_per_worker ({gpus_per_worker})" + ) num_workers = len(normalized_ranks) // gpus_per_worker for worker_idx in range(num_workers): start_idx = worker_idx * gpus_per_worker - worker_ranks = normalized_ranks[start_idx:start_idx + gpus_per_worker] + worker_ranks = normalized_ranks[start_idx : start_idx + gpus_per_worker] # All GPUs for a worker should be on the same node gpu_ranks_local = [] @@ -255,21 +302,26 @@ def get_visible_devices(): gpu_ranks_local.append(gpu_ranks) if len(set(node_ranks)) > 1: - raise ValueError(f"DeviceGroup '{group.name}': GPUs {worker_ranks} span multiple nodes. " - f"Each worker's GPUs must be on the same node.") + raise ValueError( + f"DeviceGroup '{group.name}': GPUs {worker_ranks} span multiple nodes. " + f"Each worker's GPUs must be on the same node." + ) node_rank = node_ranks[0] local_device_groups.append( dict( gpu_rank=gpu_ranks_local, placement_group=self.node2pg[node_rank], - ray_address=ray_address)) + ray_address=ray_address, + ) + ) else: for alloc_rank in normalized_ranks: node_rank = alloc_rank // nproc_per_node gpu_rank = self.visible_devices[node_rank - self.min_node_idx][alloc_rank % nproc_per_node] local_device_groups.append( - dict(gpu_rank=[gpu_rank], placement_group=self.node2pg[node_rank], ray_address=ray_address)) + dict(gpu_rank=[gpu_rank], placement_group=self.node2pg[node_rank], ray_address=ray_address) + ) self.device_groups[group.name] = local_device_groups @@ -278,48 +330,54 @@ def get_visible_devices(): # Create virtual ranks for workers (not GPUs) group.ranks = list(range(len(local_device_groups))) else: - assert getattr(group, 'gpus_per_worker', 1) == 1 + assert getattr(group, "gpus_per_worker", 1) == 1 ranks = group.ranks local_device_groups = [] for _ in range(ranks): local_device_groups.append( dict( placement_group=self.cpu_placement_groups[self.cpu_node_map[global_cpu_proc_idx][0]], - ray_address=ray_address)) + ray_address=ray_address, + ) + ) global_cpu_proc_idx += 1 self.device_groups[group.name] = local_device_groups self.group_configs = groups logger.info(f"nodes: {[n['NodeID'][:8] for n in self.nodes]}") - logger.info(f'node_ranks: {self.node_ranks}') - logger.info(f'node2pg keys: {list(self.node2pg.keys())}') + logger.info(f"node_ranks: {self.node_ranks}") + logger.info(f"node2pg keys: {list(self.node2pg.keys())}") @staticmethod def noset_env(): return { - 'RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES': '1', - 'RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES': '1', - 'RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES': '1', - 'RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES': '1', - 'RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES': '1', - 'RAY_EXPERIMENTAL_NOSET_NEURON_RT_VISIBLE_CORES': '1', - 'RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS': '1', - 'RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR': '1', + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES": "1", + "RAY_EXPERIMENTAL_NOSET_NEURON_RT_VISIBLE_CORES": "1", + "RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS": "1", + "RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR": "1", } def get_config(self, group: str): for config in self.group_configs: if config.name == group: return config - assert False, f'No group {group} found in group list: {[group.name for group in self.group_configs]}' + raise AssertionError(f"No group {group} found in group list: {[group.name for group in self.group_configs]}") def get_group(self, group: str): - assert group in self.device_groups, (f'No group {group} found in group ' - f'list: {[group.name for group in self.group_configs]}') + assert group in self.device_groups, ( + f"No group {group} found in group list: {[group.name for group in self.group_configs]}" + ) 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) for pg in self.cpu_placement_groups: diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py index cf2c3b282..24b7f523c 100644 --- a/src/twinkle/loss/base.py +++ b/src/twinkle/loss/base.py @@ -3,20 +3,27 @@ class Loss: - require_logits = False require_entropy = False require_logps = True require_values = False - def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput: - ... + 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 raise NotImplementedError( - f'{self.__class__.__name__} does not support micro-batching, including dynamic batching') + f"{self.__class__.__name__} does not support micro-batching, including dynamic batching" + ) def batch_normalizer(self, inputs: list[InputFeature]): """Return this batch's additive reduction denominator, or ``None`` for mean-reduced losses.""" diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index d0459fa4c..8c596b98e 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -1,6 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +from typing import TYPE_CHECKING, Literal, Optional, Union + import numpy as np -from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union from twinkle.data_format import LossOutput from twinkle.loss.base import Loss @@ -28,7 +29,7 @@ class GRPOLoss(Loss): def __init__( self, epsilon: float = 0.2, - epsilon_high: Optional[float] = None, + epsilon_high: float | None = None, beta: float = 0.0, entropy_coef: float = 0.0, ignore_index: int = -100, @@ -40,20 +41,20 @@ def __init__( # legacy ``GRPOTrainer`` applied inline (and dev carried in ``_ConfiguredGRPOLoss`` / # ``_AdvancedGRPOLoss``): sequence-level IS, dual-clip, entropy filter/logging, truncated-IS # rollout correction, off-policy sequence masking, FIPO future-KL weighting and SDAR distillation. - importance_sampling_level: str = 'token', - delta: Optional[float] = None, + importance_sampling_level: str = "token", + delta: float | None = None, top_entropy_quantile: float = 1.0, log_entropy: bool = False, - rollout_importance_sampling_mode: Optional[str] = None, + rollout_importance_sampling_mode: str | None = None, rollout_importance_sampling_threshold: float = 2.0, log_rollout_offpolicy_metrics: bool = False, - off_policy_sequence_mask_delta: Optional[float] = None, + off_policy_sequence_mask_delta: float | None = None, overlong_filter: bool = False, - loss_type: str = 'grpo', + loss_type: str = "grpo", fipo_decay_rate: float = 32.0, - fipo_clip_range: Optional[float] = 0.2, + fipo_clip_range: float | None = 0.2, fipo_clip_high_only: bool = True, - fipo_safety_threshold: Optional[float] = 4.0, + fipo_safety_threshold: float | None = 4.0, sdar_loss_coef: float = 0.0, sdar_gate_beta: float = 5.0, **kwargs, @@ -64,9 +65,9 @@ def __init__( self.entropy_coef = entropy_coef self.enable_sampling_replay = enable_sampling_replay if enable_sampling_replay and beta != 0.0: - raise ValueError('sampling replay does not support a GRPO KL penalty (beta must be 0)') + raise ValueError("sampling replay does not support a GRPO KL penalty (beta must be 0)") if enable_sampling_replay and entropy_coef != 0.0: - raise ValueError('sampling replay does not support a GRPO entropy bonus') + raise ValueError("sampling replay does not support a GRPO entropy bonus") self.ignore_index = ignore_index self.importance_sampling_level = importance_sampling_level @@ -79,7 +80,7 @@ def __init__( self.off_policy_sequence_mask_delta = off_policy_sequence_mask_delta self.overlong_filter = overlong_filter self.loss_type = loss_type - self.fipo_gamma = 2.0**(-1.0 / fipo_decay_rate) + self.fipo_gamma = 2.0 ** (-1.0 / fipo_decay_rate) self.fipo_clip_range = fipo_clip_range self.fipo_clip_high_only = fipo_clip_high_only self.fipo_safety_threshold = fipo_safety_threshold @@ -93,9 +94,16 @@ def __init__( # so a plain GSPO run stays on the ±5-clamped base path exactly as before; only when one of the # controls below is active does the ±20-clamped configured pipeline take over (which then also # reads ``importance_sampling_level``). - self._advanced = (loss_type == 'fipo' or delta is not None or top_entropy_quantile < 1.0 or log_entropy - or overlong_filter or rollout_importance_sampling_mode is not None - or log_rollout_offpolicy_metrics or off_policy_sequence_mask_delta is not None) + self._advanced = ( + loss_type == "fipo" + or delta is not None + or top_entropy_quantile < 1.0 + or log_entropy + or overlong_filter + or rollout_importance_sampling_mode is not None + or log_rollout_offpolicy_metrics + or off_policy_sequence_mask_delta is not None + ) # ``_distill_wrap`` marks the SDAR auxiliary objective, which reads the rollout logits/logps. self._distill_wrap = sdar_loss_coef > 0.0 # The configured pipeline recomputes the surrogate and aggregation itself and never calls @@ -106,11 +114,12 @@ def __init__( # only ``_compute_per_token_loss`` (e.g. GPG) still compose, since the pipeline calls that hook. if self._advanced and type(self)._reduce_loss is not GRPOLoss._reduce_loss: raise ValueError( - f'{type(self).__name__} is a group-/sequence-level objective (it overrides `_reduce_loss`) ' - 'that the advanced GRPO ratio/entropy/rollout controls would bypass. Drop the conflicting ' - 'control (importance_sampling_level / delta / top_entropy_quantile / log_entropy / ' - 'overlong_filter / rollout_importance_sampling_mode / log_rollout_offpolicy_metrics / ' - 'off_policy_sequence_mask_delta / fipo) to use this loss_type.') + f"{type(self).__name__} is a group-/sequence-level objective (it overrides `_reduce_loss`) " + "that the advanced GRPO ratio/entropy/rollout controls would bypass. Drop the conflicting " + "control (importance_sampling_level / delta / top_entropy_quantile / log_entropy / " + "overlong_filter / rollout_importance_sampling_mode / log_rollout_offpolicy_metrics / " + "off_policy_sequence_mask_delta / fipo) to use this loss_type." + ) # Gate the expensive forward paths: entropy for the bonus OR the advanced entropy # logging/filtering; logits for the SDAR auxiliary term. self.require_entropy = entropy_coef > 0.0 or log_entropy or top_entropy_quantile < 1.0 @@ -121,10 +130,10 @@ def micro_batch_scale(self, inputs, indices): def _compute_log_importance_weights( self, - per_token_logps: 'torch.Tensor', - per_token_old_logps: 'torch.Tensor', - loss_mask: 'torch.Tensor', - ) -> 'torch.Tensor': + per_token_logps: "torch.Tensor", + per_token_old_logps: "torch.Tensor", + loss_mask: "torch.Tensor", + ) -> "torch.Tensor": """ Compute log importance sampling weights at ``self.importance_sampling_level``. @@ -144,25 +153,28 @@ def _compute_log_importance_weights( log_weights: [batch, seq_len] log importance weights """ import torch + log_ratio = per_token_logps - per_token_old_logps # Clamp for numerical stability (±5 bounds ratio to [exp(-5), exp(5)] ≈ [0.007, 148]) log_ratio = torch.clamp(log_ratio, min=-5.0, max=5.0) - if self.importance_sampling_level == 'token': + if self.importance_sampling_level == "token": return log_ratio sequence = ((log_ratio * loss_mask).sum(-1) / loss_mask.sum(-1).clamp(min=1.0)).unsqueeze(-1) - if self.importance_sampling_level == 'sequence': + if self.importance_sampling_level == "sequence": return sequence - if self.importance_sampling_level == 'sequence_token': + if self.importance_sampling_level == "sequence_token": return per_token_logps - per_token_logps.detach() + sequence.detach() - raise ValueError(f'Unknown importance_sampling_level={self.importance_sampling_level!r}. ' - "Expected 'token', 'sequence' or 'sequence_token'.") + raise ValueError( + f"Unknown importance_sampling_level={self.importance_sampling_level!r}. " + "Expected 'token', 'sequence' or 'sequence_token'." + ) def _compute_per_token_loss( self, - ratio: 'torch.Tensor', - advantages: 'torch.Tensor', - per_token_logps: 'torch.Tensor', - ) -> 'torch.Tensor': + ratio: "torch.Tensor", + advantages: "torch.Tensor", + per_token_logps: "torch.Tensor", + ) -> "torch.Tensor": """ Compute per-token loss with PPO clipping. @@ -177,6 +189,7 @@ def _compute_per_token_loss( per_token_loss: [batch, seq_len] loss for each token """ import torch + clipped_ratio = torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon_high) loss1 = ratio * advantages loss2 = clipped_ratio * advantages @@ -184,10 +197,10 @@ def _compute_per_token_loss( def _aggregate_loss( self, - per_token_loss: 'torch.Tensor', - loss_mask: 'torch.Tensor', + per_token_loss: "torch.Tensor", + loss_mask: "torch.Tensor", **kwargs, - ) -> 'torch.Tensor': + ) -> "torch.Tensor": """ Aggregate per-token loss to scalar. @@ -206,7 +219,7 @@ def _aggregate_loss( # Each sequence contributes equally regardless of length. return ((per_token_loss * loss_mask).sum(-1) / loss_mask.sum(-1).clamp(min=1.0)).mean() - def _loss_num_tokens(self, loss_mask: 'torch.Tensor'): + def _loss_num_tokens(self, loss_mask: "torch.Tensor"): """Token denominator reported in ``LossOutput.num_tokens``. 0 (default) => framework uses the PER-TOKEN-MEAN accumulation path, where each @@ -219,12 +232,12 @@ def _loss_num_tokens(self, loss_mask: 'torch.Tensor'): def _pad_and_align_to_batch( self, - data: 'Union[torch.Tensor, List, np.ndarray]', - mask: 'torch.Tensor', - device: 'torch.device', - dtype: 'torch.dtype', + data: "torch.Tensor | list | np.ndarray", + mask: "torch.Tensor", + device: "torch.device", + dtype: "torch.dtype", fill_value: float = 0.0, - ) -> 'torch.Tensor': + ) -> "torch.Tensor": """Align data to mask: scalars broadcast, sequences scatter. Two valid per-sample sequence forms are supported and disambiguated @@ -232,10 +245,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. @@ -255,8 +268,9 @@ def _pad_and_align_to_batch( data = data.unsqueeze(1) if data.shape[1] == 1: assert data.shape[0] == batch_size, ( - f'scalar broadcast expects data.shape[0]==batch_size, ' - f'got data.shape={tuple(data.shape)} mask.shape={(batch_size, seq_len)}') + f"scalar broadcast expects data.shape[0]==batch_size, " + f"got data.shape={tuple(data.shape)} mask.shape={(batch_size, seq_len)}" + ) fill = torch.full((batch_size, seq_len), fill_value, dtype=dtype, device=device) expanded = data.expand(batch_size, seq_len) return torch.where(mask, expanded, fill) @@ -266,7 +280,8 @@ def _pad_and_align_to_batch( if isinstance(data, (list, tuple)): if all(isinstance(x, (int, float)) for x in data): # Flat scalars return self._pad_and_align_to_batch( - torch.tensor(data, dtype=dtype, device=device), mask, device, dtype, fill_value) + torch.tensor(data, dtype=dtype, device=device), mask, device, dtype, fill_value + ) data = [torch.as_tensor(s, dtype=dtype, device=device) for s in data] # Scatter sequences @@ -276,26 +291,99 @@ 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 - def _resolve_loss_mask(self, inputs: Dict, labels: 'torch.Tensor') -> 'torch.Tensor': + @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, strict=False)): + 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. ``labels`` alone answers "should this token be scored", which is all SFT @@ -309,8 +397,9 @@ def _resolve_loss_mask(self, inputs: Dict, labels: 'torch.Tensor') -> 'torch.Ten ``labels != ignore_index`` and old trajectories train exactly as before. """ import torch + trainable = (labels != self.ignore_index).bool() - completion_mask = inputs.get('completion_mask') + completion_mask = inputs.get("completion_mask") if completion_mask is None: return trainable if not torch.is_tensor(completion_mask): @@ -319,23 +408,27 @@ def _resolve_loss_mask(self, inputs: Dict, labels: 'torch.Tensor') -> 'torch.Ten if completion_mask.dim() == 1: completion_mask = completion_mask.unsqueeze(0) if completion_mask.shape != trainable.shape: - raise ValueError(f'completion_mask shape {tuple(completion_mask.shape)} does not match labels shape ' - f'{tuple(trainable.shape)}. A misaligned mask would apply importance ratios to ' - 'the wrong tokens, so it is refused rather than broadcast.') + raise ValueError( + f"completion_mask shape {tuple(completion_mask.shape)} does not match labels shape " + f"{tuple(trainable.shape)}. A misaligned mask would apply importance ratios to " + "the wrong tokens, so it is refused rather than broadcast." + ) loss_mask = trainable & completion_mask.bool() if self.enable_sampling_replay and not bool((loss_mask == trainable).all()): - raise ValueError('sampling replay does not support turns generated outside the sampled policy: ' - 'they are trainable but have no sampling mask to replay against.') + raise ValueError( + "sampling replay does not support turns generated outside the sampled policy: " + "they are trainable but have no sampling mask to replay against." + ) return loss_mask def _base_grpo( self, - inputs: Dict, - outputs: Dict, + inputs: dict, + outputs: dict, *, - old_logps: Optional[Union['torch.Tensor', List[List[float]]]] = None, - ref_logps: Optional['torch.Tensor'] = None, - advantages: Optional[Union['torch.Tensor', List[float], np.ndarray]] = None, + old_logps: Union["torch.Tensor", list[list[float]]] | None = None, + ref_logps: Optional["torch.Tensor"] = None, + advantages: Union["torch.Tensor", list[float], np.ndarray] | None = None, **kwargs, ): """ @@ -357,25 +450,26 @@ def _base_grpo( **kwargs: Additional arguments """ import torch + if self.enable_sampling_replay: if old_logps is None: - raise ValueError('old_logps are required when sampling replay is enabled') - labels = inputs.get('labels') + raise ValueError("old_logps are required when sampling replay is enabled") + labels = inputs.get("labels") assert labels is not None, "inputs must contain 'labels'" if not torch.is_tensor(labels): labels = torch.as_tensor(labels) if labels.dim() == 1: labels = labels.unsqueeze(0) - logps = outputs.get('logps') + logps = outputs.get("logps") if self.enable_sampling_replay and logps is None: - raise RuntimeError('sampling replay logps must be computed by the model forward') + raise RuntimeError("sampling replay logps must be computed by the model forward") loss_mask = self._resolve_loss_mask(inputs, labels) if logps is None: - logits = outputs.get('logits') + logits = outputs.get("logits") if logits.shape[1] != labels.shape[1]: # some mllm return logits with image tokens, exclude here - logits = logits[:, -labels.shape[1]:] + logits = logits[:, -labels.shape[1] :] masked_labels = labels.clone() masked_labels[~loss_mask] = 0 logps = selective_log_softmax(logits, masked_labels) @@ -423,18 +517,19 @@ def _base_grpo( # ── Compute loss ──────────────────────────────────────────────── log_importance_weights = self._compute_log_importance_weights(logps, old_logps, loss_mask) - return self._reduce_loss(logps, old_logps, ref_logps, advantages, loss_mask, log_importance_weights, - outputs, **kwargs) + return self._reduce_loss( + logps, old_logps, ref_logps, advantages, loss_mask, log_importance_weights, outputs, **kwargs + ) def _reduce_loss( self, - logps: 'torch.Tensor', - old_logps: 'torch.Tensor', - ref_logps: Optional['torch.Tensor'], - advantages: 'torch.Tensor', - loss_mask: 'torch.Tensor', - log_importance_weights: 'torch.Tensor', - outputs: Dict, + logps: "torch.Tensor", + old_logps: "torch.Tensor", + ref_logps: Optional["torch.Tensor"], + advantages: "torch.Tensor", + loss_mask: "torch.Tensor", + log_importance_weights: "torch.Tensor", + outputs: dict, **kwargs, ) -> LossOutput: """Reduce the per-token importance weights to the scalar ``LossOutput``. @@ -445,18 +540,21 @@ def _reduce_loss( -- byte-for-byte identical to inlining it at the tail of ``__call__``. """ import torch + ratio = torch.exp(log_importance_weights) per_token_loss = self._compute_per_token_loss(ratio, advantages, logps) if self.beta > 0.0 and ref_logps is not None: - per_token_kl = (torch.exp(ref_logps - logps) - (ref_logps - logps) - 1) + per_token_kl = torch.exp(ref_logps - logps) - (ref_logps - logps) - 1 per_token_loss = per_token_loss + self.beta * per_token_kl if self.entropy_coef > 0.0: - entropies = outputs.get('entropies') - assert entropies is not None, ('entropy_coef > 0 requires outputs[\'entropies\'] — make sure the ' - "loss instance's require_entropy flag was set before the forward call.") + entropies = outputs.get("entropies") + assert entropies is not None, ( + "entropy_coef > 0 requires outputs['entropies'] — make sure the " + "loss instance's require_entropy flag was set before the forward call." + ) # entropies may come in fp32 from the kernel; cast to match logps dtype # so the final per_token_loss stays consistent (bf16 under amp). per_token_loss = per_token_loss - self.entropy_coef * entropies.to(per_token_loss.dtype) @@ -474,15 +572,15 @@ def _reduce_loss( def __call__( self, - inputs: Dict, - outputs: Dict, + inputs: dict, + outputs: dict, *, - old_logps: Optional[Union['torch.Tensor', List[List[float]]]] = None, - ref_logps: Optional['torch.Tensor'] = None, - advantages: Optional[Union['torch.Tensor', List[float], np.ndarray]] = None, - rollout_logps: Optional[Union['torch.Tensor', List[List[float]]]] = None, - truncated: Optional[List[bool]] = None, - teacher_logps: Optional[Union['torch.Tensor', List[List[float]]]] = None, + old_logps: Union["torch.Tensor", list[list[float]]] | None = None, + ref_logps: Optional["torch.Tensor"] = None, + advantages: Union["torch.Tensor", list[float], np.ndarray] | None = None, + rollout_logps: Union["torch.Tensor", list[list[float]]] | None = None, + truncated: list[bool] | None = None, + teacher_logps: Union["torch.Tensor", list[list[float]]] | None = None, **kwargs, ): """Compute the GRPO loss, applying any advanced controls this instance was built with. @@ -494,32 +592,41 @@ def __call__( is exactly :meth:`_base_grpo`. """ if not (self._advanced or self._distill_wrap): - return self._base_grpo(inputs, outputs, old_logps=old_logps, ref_logps=ref_logps, - advantages=advantages, **kwargs) + return self._base_grpo( + inputs, outputs, old_logps=old_logps, ref_logps=ref_logps, advantages=advantages, **kwargs + ) if self._advanced: result = self._configured_grpo( - inputs, outputs, old_logps=old_logps, ref_logps=ref_logps, advantages=advantages, - rollout_logps=rollout_logps, truncated=truncated, **kwargs) + inputs, + outputs, + old_logps=old_logps, + ref_logps=ref_logps, + advantages=advantages, + rollout_logps=rollout_logps, + truncated=truncated, + **kwargs, + ) else: result = self._base_grpo( - inputs, outputs, old_logps=old_logps, ref_logps=ref_logps, advantages=advantages, **kwargs) + inputs, outputs, old_logps=old_logps, ref_logps=ref_logps, advantages=advantages, **kwargs + ) - loss = result['loss'] - channel_loss = result.get('channel_loss') + loss = result["loss"] + channel_loss = result.get("channel_loss") if self._distill_wrap: loss = self._add_sdar(loss, inputs, outputs, teacher_logps) result = LossOutput(loss=loss, num_tokens=0) else: - result = LossOutput(loss=loss, num_tokens=result['num_tokens']) + result = LossOutput(loss=loss, num_tokens=result["num_tokens"]) if channel_loss is not None: - result['channel_loss'] = channel_loss + result["channel_loss"] = channel_loss return result def _configured_grpo( self, - inputs: Dict, - outputs: Dict, + inputs: dict, + outputs: dict, *, old_logps=None, ref_logps=None, @@ -535,34 +642,36 @@ def _configured_grpo( intersect with ``completion_mask`` the way the base :meth:`_resolve_loss_mask` does. """ import torch - labels = inputs['labels'] + + labels = inputs["labels"] if not torch.is_tensor(labels): labels = torch.as_tensor(labels) if labels.dim() == 1: labels = labels.unsqueeze(0) - logps = outputs.get('logps') + logps = outputs.get("logps") if logps is None: raise RuntimeError('The advanced GRPO controls require outputs["logps"].') alignment_mask = labels.ne(self.ignore_index) device, dtype = logps.device, logps.dtype - old = (logps.detach() if old_logps is None else self._pad_and_align_to_batch( - old_logps, alignment_mask, device, dtype)) + old = ( + logps.detach() + if old_logps is None + else self._pad_and_align_to_batch(old_logps, alignment_mask, device, dtype) + ) advantages = self._pad_and_align_to_batch(advantages, alignment_mask, device, dtype) mask = alignment_mask.clone() - channel_loss: Dict[str, 'torch.Tensor'] = {} + channel_loss: dict[str, torch.Tensor] = {} entropy_mask = None if self.require_entropy: - entropies = outputs.get('entropies') + entropies = outputs.get("entropies") if entropies is None: raise RuntimeError('Entropy logging/filtering requires outputs["entropies"].') if self.log_entropy: count = alignment_mask.sum().detach().float() - channel_loss['entropy'] = torch.stack( - ((entropies * alignment_mask).sum().detach().float(), count)) + channel_loss["entropy"] = torch.stack(((entropies * alignment_mask).sum().detach().float(), count)) if self.top_entropy_quantile < 1.0 and alignment_mask.any(): - threshold = torch.quantile( - entropies[alignment_mask].float(), 1.0 - self.top_entropy_quantile) + threshold = torch.quantile(entropies[alignment_mask].float(), 1.0 - self.top_entropy_quantile) entropy_mask = entropies.ge(threshold) if truncated is not None: @@ -571,7 +680,7 @@ def _configured_grpo( log_ratio, ratio = self._advanced_importance_weights(logps, old, mask) per_token_loss = self._dual_clip_per_token_loss(ratio, advantages, logps) - if self.loss_type == 'fipo': + if self.loss_type == "fipo": per_token_loss = per_token_loss * self._fipo_weights(log_ratio, ratio, advantages, mask) if entropy_mask is not None: per_token_loss = per_token_loss * entropy_mask @@ -588,12 +697,13 @@ def _configured_grpo( rollout_log_ratio = old - rollout if self.log_rollout_offpolicy_metrics: count = mask.sum().detach().float() - channel_loss['rollout_log_ratio'] = torch.stack( - ((rollout_log_ratio.abs() * mask).sum().detach().float(), count)) + channel_loss["rollout_log_ratio"] = torch.stack( + ((rollout_log_ratio.abs() * mask).sum().detach().float(), count) + ) if self.rollout_importance_sampling_mode is not None: per_token_loss = per_token_loss * self._rollout_is_weights(rollout_log_ratio, mask) elif self.rollout_importance_sampling_mode is not None: - raise ValueError('rollout_importance_sampling_mode requires rollout_logps from the sampler.') + raise ValueError("rollout_importance_sampling_mode requires rollout_logps from the sampler.") if self.off_policy_sequence_mask_delta is not None: old_policy = rollout if rollout is not None else old @@ -602,13 +712,13 @@ def _configured_grpo( keep = ~((sequence_delta > self.off_policy_sequence_mask_delta) & (sequence_advantage < 0)) mask = mask & keep.unsqueeze(-1) - if self.loss_type == 'fipo': + if self.loss_type == "fipo": loss = (per_token_loss * mask).sum() / mask.sum().clamp(min=1.0) else: loss = self._aggregate_loss(per_token_loss, mask, **kwargs) result = LossOutput(loss=loss, num_tokens=0) if channel_loss: - result['channel_loss'] = channel_loss + result["channel_loss"] = channel_loss return result def _advanced_importance_weights(self, logps, old_logps, mask): @@ -618,24 +728,26 @@ def _advanced_importance_weights(self, logps, old_logps, mask): consumes it directly. The ±20 clamp is the legacy stability bound for this pipeline. """ import torch + log_ratio = torch.clamp(logps - old_logps, min=-20.0, max=20.0) - if self.importance_sampling_level == 'token': + if self.importance_sampling_level == "token": log_weights = log_ratio else: sequence = ((log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).unsqueeze(-1) - if self.importance_sampling_level == 'sequence': + if self.importance_sampling_level == "sequence": log_weights = sequence - elif self.importance_sampling_level == 'sequence_token': + elif self.importance_sampling_level == "sequence_token": # GSPO-token: grad flows through the per-token logp, the sequence weight is a constant. log_weights = logps - logps.detach() + sequence.detach() else: - raise ValueError(f'Unknown importance_sampling_level={self.importance_sampling_level!r}.') + raise ValueError(f"Unknown importance_sampling_level={self.importance_sampling_level!r}.") return log_ratio, torch.exp(log_weights) def _dual_clip_per_token_loss(self, ratio, advantages, logps): """PPO clip with an optional dual-clip floor ``delta`` on the raw ratio (legacy ordering).""" import torch - if self.delta is None or self.loss_type not in {'grpo', 'dapo', 'fipo', 'bnpo', 'dr_grpo'}: + + if self.delta is None or self.loss_type not in {"grpo", "dapo", "fipo", "bnpo", "dr_grpo"}: return self._compute_per_token_loss(ratio, advantages, logps) clipped_ratio = torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon_high) dual_clipped_ratio = torch.clamp(ratio, max=self.delta) @@ -644,24 +756,25 @@ def _dual_clip_per_token_loss(self, ratio, advantages, logps): def _rollout_is_weights(self, log_ratio, mask): """Truncated importance-sampling weights correcting the policy-vs-sampler (rollout) gap.""" import torch + ratio = torch.exp(torch.clamp(log_ratio, min=-20.0, max=20.0)) mode = self.rollout_importance_sampling_mode threshold = self.rollout_importance_sampling_threshold - if mode == 'token_truncate': + if mode == "token_truncate": return torch.clamp(ratio, max=threshold) - if mode == 'token_mask': + if mode == "token_mask": return torch.where(ratio <= threshold, ratio, torch.zeros_like(ratio)) - sequence = torch.exp((torch.log(ratio.clamp(min=1e-10)) * mask).sum(-1) - / mask.sum(-1).clamp(min=1.0)) - if mode == 'sequence_truncate': + sequence = torch.exp((torch.log(ratio.clamp(min=1e-10)) * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)) + if mode == "sequence_truncate": return torch.clamp(sequence, max=threshold).unsqueeze(-1).expand_as(ratio) - if mode == 'sequence_mask': + if mode == "sequence_mask": return ratio * (sequence <= threshold).unsqueeze(-1) return ratio def _fipo_weights(self, log_ratio, ratio, advantages, mask): """FIPO's detached future-KL influence weight (block-matmul of the decayed future log-ratio).""" import torch + future_delta = log_ratio.masked_fill(~mask, 0.0) if self.delta is not None: # Dual-clip participation: high-ratio tokens do not contribute to the future KL. @@ -674,8 +787,8 @@ def _fipo_weights(self, log_ratio, ratio, advantages, mask): block_positions = torch.arange(start, end, device=log_ratio.device).unsqueeze(0) distance = block_positions - positions decay = torch.pow( - torch.as_tensor(self.fipo_gamma, dtype=log_ratio.dtype, device=log_ratio.device), - distance.clamp(min=0)) + torch.as_tensor(self.fipo_gamma, dtype=log_ratio.dtype, device=log_ratio.device), distance.clamp(min=0) + ) decay = decay * (distance >= 0).to(log_ratio.dtype) future_kl += torch.matmul(future_delta[:, start:end], decay.t()) weights = torch.exp(future_kl.masked_fill(~mask, 0.0)) @@ -692,10 +805,10 @@ def _add_sdar(self, loss, rl_inputs, rl_outputs, teacher_logps): if teacher_logps is None or self.sdar_loss_coef <= 0: return loss from twinkle.loss.sdar import compute_sdar_loss - student_logps = rl_outputs['logps'] - response_mask = rl_inputs['labels'].ne(self.ignore_index) - teacher = self._pad_and_align_to_batch( - teacher_logps, response_mask, student_logps.device, student_logps.dtype) + + student_logps = rl_outputs["logps"] + response_mask = rl_inputs["labels"].ne(self.ignore_index) + teacher = self._pad_and_align_to_batch(teacher_logps, response_mask, student_logps.device, student_logps.dtype) sdar_loss, _ = compute_sdar_loss(student_logps, teacher, response_mask, self.sdar_gate_beta) return loss + self.sdar_loss_coef * sdar_loss @@ -713,22 +826,23 @@ class PPOLoss(GRPOLoss): sequence and then averages the sequences. """ - _LOSS_AGG_MODES = {'token-mean', 'seq-mean-token-mean'} + _LOSS_AGG_MODES = {"token-mean", "seq-mean-token-mean"} - def __init__(self, loss_agg_mode: Literal['token-mean', 'seq-mean-token-mean'] = 'token-mean', **kwargs): + def __init__(self, loss_agg_mode: Literal["token-mean", "seq-mean-token-mean"] = "token-mean", **kwargs): super().__init__(**kwargs) if loss_agg_mode not in self._LOSS_AGG_MODES: - raise ValueError(f'Unsupported PPO loss_agg_mode: {loss_agg_mode}. ' - f'Expected one of {sorted(self._LOSS_AGG_MODES)}.') + raise ValueError( + f"Unsupported PPO loss_agg_mode: {loss_agg_mode}. Expected one of {sorted(self._LOSS_AGG_MODES)}." + ) self.loss_agg_mode = loss_agg_mode def _aggregate_loss( self, - per_token_loss: 'torch.Tensor', - loss_mask: 'torch.Tensor', + per_token_loss: "torch.Tensor", + loss_mask: "torch.Tensor", **kwargs, - ) -> 'torch.Tensor': - if self.loss_agg_mode == 'token-mean': + ) -> "torch.Tensor": + if self.loss_agg_mode == "token-mean": return (per_token_loss * loss_mask).sum() / loss_mask.sum().clamp(min=1.0) return super()._aggregate_loss(per_token_loss, loss_mask, **kwargs) @@ -742,7 +856,7 @@ class GSPOLoss(GRPOLoss): """ def __init__(self, *args, **kwargs): - kwargs.setdefault('importance_sampling_level', 'sequence') + kwargs.setdefault("importance_sampling_level", "sequence") super().__init__(*args, **kwargs) @@ -768,12 +882,13 @@ def __init__( def _compute_per_token_loss( self, - ratio: 'torch.Tensor', - advantages: 'torch.Tensor', - per_token_logps: 'torch.Tensor', - ) -> 'torch.Tensor': + ratio: "torch.Tensor", + advantages: "torch.Tensor", + per_token_logps: "torch.Tensor", + ) -> "torch.Tensor": """Soft-gated loss.""" import torch + gate_pos = torch.sigmoid(self.tau_pos * (ratio - 1)) * (4.0 / self.tau_pos) gate_neg = torch.sigmoid(self.tau_neg * (ratio - 1)) * (4.0 / self.tau_neg) is_positive = advantages > 0 @@ -791,8 +906,8 @@ class CISPOLoss(GRPOLoss): def micro_batch_scale(self, inputs, indices): token_counts = [] for model_input in inputs: - labels = model_input['labels'] - if hasattr(labels, 'ne'): + labels = model_input["labels"] + if hasattr(labels, "ne"): token_counts.append(int(labels.ne(self.ignore_index).sum().item())) else: token_counts.append(sum(int(token != self.ignore_index) for token in labels)) @@ -803,10 +918,10 @@ def micro_batch_scale(self, inputs, indices): def _compute_per_token_loss( self, - ratio: 'torch.Tensor', - advantages: 'torch.Tensor', - per_token_logps: 'torch.Tensor', - ) -> 'torch.Tensor': + ratio: "torch.Tensor", + advantages: "torch.Tensor", + per_token_logps: "torch.Tensor", + ) -> "torch.Tensor": """Clamped ratio * advantage * log_prob.""" import torch @@ -815,13 +930,13 @@ def _compute_per_token_loss( def _aggregate_loss( self, - per_token_loss: 'torch.Tensor', - loss_mask: 'torch.Tensor', + per_token_loss: "torch.Tensor", + loss_mask: "torch.Tensor", **kwargs, - ) -> 'torch.Tensor': + ) -> "torch.Tensor": """Sum over all tokens, divide by total token count.""" # Use provided num_items_in_batch if available, otherwise use mask sum - num_items = kwargs.get('num_items_in_batch', loss_mask.sum()) + num_items = kwargs.get("num_items_in_batch", loss_mask.sum()) return (per_token_loss * loss_mask).sum() / num_items @@ -857,15 +972,16 @@ class BNPOLoss(GRPOLoss): sequence-mean as the micro size approaches 1). """ - def __init__(self, *args, token_mean_scope: str = 'micro', **kwargs): + def __init__(self, *args, token_mean_scope: str = "micro", **kwargs): super().__init__(*args, **kwargs) - assert token_mean_scope in ('global', 'micro'), \ - f'token_mean_scope must be global|micro, got {token_mean_scope!r}' + assert token_mean_scope in ("global", "micro"), ( + f"token_mean_scope must be global|micro, got {token_mean_scope!r}" + ) self.token_mean_scope = token_mean_scope # 'global' 返回的是 token 和（梯度在下游按 num_tokens=Σmask 归一）。必须同步告诉展示层 # 这是 sum-reduction，否则 LossMetric（metric/loss.py）不会除以 num_tokens，会把每个 micro 的 # token 和当均值直接平均，展示出一个被 token 数放大的巨大 loss（梯度不受影响，纯展示失真）。 - self.reduction = 'sum' if token_mean_scope == 'global' else 'mean' + self.reduction = "sum" if token_mean_scope == "global" else "mean" def micro_batch_scale(self, inputs, indices): """The weight one micro-batch carries, which has to follow ``token_mean_scope``. @@ -879,26 +995,26 @@ def micro_batch_scale(self, inputs, indices): the global token-mean and erase the distinction this scope exists to make -- which is the +3.2e-4 vs +0.031 per-token pg_loss measured above. """ - if self.token_mean_scope == 'global': + if self.token_mean_scope == "global": return 1.0 return super().micro_batch_scale(inputs, indices) def _aggregate_loss( self, - per_token_loss: 'torch.Tensor', - loss_mask: 'torch.Tensor', + per_token_loss: "torch.Tensor", + loss_mask: "torch.Tensor", **kwargs, - ) -> 'torch.Tensor': + ) -> "torch.Tensor": """global: return the token SUM (the global division is done downstream via num_tokens=Σmask). micro (legacy): local token-mean, later equal-weighted across micro/dp groups.""" summed = (per_token_loss * loss_mask).sum() - if self.token_mean_scope == 'global': + if self.token_mean_scope == "global": return summed return summed / loss_mask.sum().clamp(min=1.0) - def _loss_num_tokens(self, loss_mask: 'torch.Tensor'): - if self.token_mean_scope == 'global': + def _loss_num_tokens(self, loss_mask: "torch.Tensor"): + if self.token_mean_scope == "global": return loss_mask.sum().clamp(min=1.0) return 0 @@ -923,10 +1039,10 @@ def __init__( def _aggregate_loss( self, - per_token_loss: 'torch.Tensor', - loss_mask: 'torch.Tensor', + per_token_loss: "torch.Tensor", + loss_mask: "torch.Tensor", **kwargs, - ) -> 'torch.Tensor': + ) -> "torch.Tensor": """Normalize by batch_size * max_completion_length.""" batch_size = loss_mask.shape[0] return (per_token_loss * loss_mask).sum() / (batch_size * self.max_completion_length) @@ -954,20 +1070,21 @@ def __init__(self, real_tau: float = 0.5, num_generations: int = 1, **kwargs): def _reduce_loss( self, - logps: 'torch.Tensor', - old_logps: 'torch.Tensor', - ref_logps: Optional['torch.Tensor'], - advantages: 'torch.Tensor', - loss_mask: 'torch.Tensor', - log_importance_weights: 'torch.Tensor', - outputs: Dict, + logps: "torch.Tensor", + old_logps: "torch.Tensor", + ref_logps: Optional["torch.Tensor"], + advantages: "torch.Tensor", + loss_mask: "torch.Tensor", + log_importance_weights: "torch.Tensor", + outputs: dict, **kwargs, ) -> LossOutput: import torch + batch_size = loss_mask.shape[0] assert self.num_generations > 0 and batch_size % self.num_generations == 0, ( - f'REAL needs num_generations ({self.num_generations}) to divide the batch ({batch_size}) ' - 'into whole groups.') + f"REAL needs num_generations ({self.num_generations}) to divide the batch ({batch_size}) into whole groups." + ) mask_sum = loss_mask.sum(-1).clamp(min=1.0) # Per-sequence score = mean token log importance ratio. Legacy REAL uses the UNCLAMPED token log @@ -996,10 +1113,10 @@ def _reduce_loss( zeros = torch.zeros(batch_scores.size(0), 1, device=batch_scores.device, dtype=batch_scores.dtype) # Negative loss: log(1 + sum(exp(S_neg))) - neg_input = scaled_scores.masked_fill(~batch_neg_mask, float('-inf')) + neg_input = scaled_scores.masked_fill(~batch_neg_mask, float("-inf")) neg_loss = torch.logsumexp(torch.cat([neg_input, zeros], dim=1), dim=1) # Positive loss: log(1 + sum(exp(-S_pos))) - pos_input = (-scaled_scores).masked_fill(~batch_pos_mask, float('-inf')) + pos_input = (-scaled_scores).masked_fill(~batch_pos_mask, float("-inf")) pos_loss = torch.logsumexp(torch.cat([pos_input, zeros], dim=1), dim=1) loss = (neg_loss + pos_loss).sum() / group_rewards.size(0) diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 915bf38c1..22898f9e6 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -1,22 +1,24 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import contextlib import inspect -import json -import numpy as np import os import random import re +from collections.abc import Callable, Generator +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Literal, Optional, Union + +import json +import numpy as np import torch import torch.distributed as dist import torch.nn as nn -from contextlib import contextmanager -from dataclasses import dataclass from peft import LoraConfig, PeftConfig, PeftModel, get_peft_model from peft.tuners.lora import Linear as LoraLinear from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler from transformers import PretrainedConfig -from typing import Any, Callable, Dict, Generator, List, Literal, Optional, Tuple, Type, Union import twinkle import twinkle.metric @@ -25,10 +27,15 @@ from twinkle.checkpoint_engine.delta.trainer.backends import MegatronDeltaBackend from twinkle.checkpoint_engine.delta.trainer.mixin import DeltaTrainerMixin from twinkle.checkpoint_engine.mixin import CheckpointEngineMixin -from twinkle.checkpoint_engine.weight_keys import add_base_layer_suffix, is_lora_key, normalize_param_name, print_weight_example +from twinkle.checkpoint_engine.weight_keys import ( + add_base_layer_suffix, + is_lora_key, + normalize_param_name, + print_weight_example, +) from twinkle.data_format import InputFeature, ModelOutput, Trajectory, is_encoded from twinkle.hub import HubOperation -from twinkle.infra import collect_tensor_dict, remote_function +from twinkle.infra import collect_tensor_dict from twinkle.infra._torch_setup import seed_from_env from twinkle.loss import CrossEntropyLoss, Loss from twinkle.metric import LossMetric, Metric, TrainMetric @@ -42,9 +49,9 @@ from twinkle.utils import construct_class, get_logger, selective_log_softmax from .moe.router_replay import RouterReplay from .mtp.mtp import MtpTrainingState +from .strategy import MegatronStrategy from .utils._mindspeed_runtime import ensure_mindspeed_adaptor_patched from .utils.utils import _find_embedding, _find_layers -from .strategy import MegatronStrategy logger = get_logger() @@ -53,13 +60,15 @@ def _resolve_task_context(model, task): """The per-forward Patch for ``task``: 'embedding' swaps the output_layer for identity (with TP/SP gather) and pools to ``[n_seqs, hidden]``; causal_lm/seq_cls/reranker/value need no patch (mcore-bridge already builds the right head; pooling happens in forward_step).""" - if task in (None, 'causal_lm', 'seq_cls', 'reranker', 'generative_reranker', 'value'): + if task in (None, "causal_lm", "seq_cls", "reranker", "generative_reranker", "value"): return contextlib.nullcontext() - if task == 'embedding': + if task == "embedding": from twinkle.patch.megatron_emb import MegatronEmbeddingPatch + return apply_context(model, MegatronEmbeddingPatch()) - raise ValueError(f'Unknown task={task!r}; expected one of: causal_lm, embedding, seq_cls, reranker, ' - 'generative_reranker, value.') + raise ValueError( + f"Unknown task={task!r}; expected one of: causal_lm, embedding, seq_cls, reranker, generative_reranker, value." + ) @dataclass @@ -68,12 +77,13 @@ class MegatronOptimizerGroup(BaseOptimizerGroup): Similar to OptimizerGroup but adapted for Megatron's distributed training. """ + # Megatron-specific fields _last_step_success: bool = True def __post_init__(self): if self._device_mesh.data_world_size > 1: - self._dp_group = self._device_mesh.create_process_group(['dp', 'fsdp']) + self._dp_group = self._device_mesh.create_process_group(["dp", "fsdp"]) train_metrics = [ LossMetric(self._device_mesh, self._dp_group), TrainMetric(self._device_mesh, self._dp_group), @@ -92,94 +102,100 @@ def _get_lr(self): return _lrs _default_lr = self.optimizer.chained_optimizers[0].config.lr for param_group in self.optimizer.param_groups: - _lrs.append(param_group.get('lr', _default_lr)) + _lrs.append(param_group.get("lr", _default_lr)) return _lrs -_default_adapter_name = '' +_default_adapter_name = "" -@remote_class(execute='all') -class MegatronModel(TrainableModel, nn.Module, CheckpointEngineMixin, DeltaTrainerMixin, - MegatronDeltaBackend): - +@remote_class(execute="all") +class MegatronModel(TrainableModel, nn.Module, CheckpointEngineMixin, DeltaTrainerMixin, MegatronDeltaBackend): _mcore_async_queue = None def __init__( self, model_id: str, - config: Optional[PretrainedConfig] = None, - ddp_config: Optional[Dict[str, Any]] = None, - device_mesh: Optional[DeviceMesh] = None, - mixed_precision: Literal['no', 'fp16', 'bf16'] = 'bf16', + config: PretrainedConfig | None = None, + ddp_config: dict[str, Any] | None = None, + device_mesh: DeviceMesh | None = None, + mixed_precision: Literal["no", "fp16", "bf16"] = "bf16", load_weights: bool = True, - recompute_granularity: Optional[str] = 'full', # Activation checkpointing - recompute_method: Optional[str] = 'uniform', - recompute_num_layers: Optional[int] = 1, - recompute_modules: Optional[list] = None, # Modules to recompute - vit_gradient_checkpointing: Optional[bool] = None, - vit_gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None, - new_special_tokens: Optional[List[str]] = None, - bridge_backend='mcore-bridge', + recompute_granularity: str | None = "full", # Activation checkpointing + recompute_method: str | None = "uniform", + recompute_num_layers: int | None = 1, + recompute_modules: list | None = None, # Modules to recompute + vit_gradient_checkpointing: bool | None = None, + vit_gradient_checkpointing_kwargs: dict[str, Any] | None = None, + new_special_tokens: list[str] | None = None, + bridge_backend="mcore-bridge", align_grad_reduce: bool = True, nccl_comm_warmup: bool = False, - attn_impl: Optional[str] = None, + attn_impl: str | None = None, **kwargs, ): - requires('megatron_core') + requires("megatron_core") nn.Module.__init__(self) from twinkle.patch.megatron_peft import MegatronPeft self.model_id = model_id self.device_mesh = device_mesh self.mixed_precision = mixed_precision - self._model_path = HubOperation.download_model(model_id, revision=kwargs.pop('revision', None)) - self.tokenizer_id = kwargs.pop('tokenizer_id', self.model_id) + self._model_path = HubOperation.download_model(model_id, revision=kwargs.pop("revision", None)) + self.tokenizer_id = kwargs.pop("tokenizer_id", self.model_id) self._default_tokenizer = None - self.use_distributed_optimizer = kwargs.pop('use_distributed_optimizer', True) - self.variable_seq_lengths = kwargs.get('variable_seq_lengths', True) - self._router_replay_enabled = bool(kwargs.pop('enable_router_replay', False)) + self.use_distributed_optimizer = kwargs.pop("use_distributed_optimizer", True) + self.variable_seq_lengths = kwargs.get("variable_seq_lengths", True) + self._router_replay_enabled = bool(kwargs.pop("enable_router_replay", False)) self._router_replay_applied = False if self._router_replay_enabled: - kwargs['moe_enable_routing_replay'] = True + kwargs["moe_enable_routing_replay"] = True if MegatronStrategy.uses_fsdp(ddp_config): - os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '32' + os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "32" torch_util.set_device() self._try_init_process_group() # MindSpeed must patch before mcore_bridge imports its patcher, otherwise # mcore_bridge pulls in megatron.core/TE too early on NPU. ensure_mindspeed_adaptor_patched() - requires('mcore_bridge') + requires("mcore_bridge") # Route batched pipeline P2P through the WORLD group to avoid subgroup-P2P PP hangs. from twinkle.patch.megatron_p2p import MegatronBatchedP2PGroupPatch + apply_patch(None, MegatronBatchedP2PGroupPatch()) # Parallelize torch DCP checkpoint reads (mcore optimizer/RNG resume) for faster loading. from twinkle.patch.torch_dcp_reader import TorchDCPParallelReaderPatch + apply_patch(None, TorchDCPParallelReaderPatch()) # Skip slow, redundant torch DCP shard validation during mcore checkpoint save/load. from twinkle.patch.torch_dcp_validation import MegatronDCPValidationPatch + apply_patch(None, MegatronDCPValidationPatch()) - kwargs.update({ - 'recompute_granularity': recompute_granularity, - 'recompute_modules': recompute_modules, - 'recompute_method': recompute_method, - 'recompute_num_layers': recompute_num_layers, - 'variable_seq_lengths': self.variable_seq_lengths, - }) - seed = kwargs.pop('seed', None) or seed_from_env(42) - model_loader: Optional[ModelLoaderProtocol] = kwargs.pop('model_loader', None) + kwargs.update( + { + "recompute_granularity": recompute_granularity, + "recompute_modules": recompute_modules, + "recompute_method": recompute_method, + "recompute_num_layers": recompute_num_layers, + "variable_seq_lengths": self.variable_seq_lengths, + } + ) + seed = kwargs.pop("seed", None) or seed_from_env(42) + model_loader: ModelLoaderProtocol | None = kwargs.pop("model_loader", None) if config is not None: self.hf_config = config elif model_loader is not None: self.hf_config = model_loader.process_config(model_loader.build_config(self._model_path)) else: from transformers import AutoConfig + # Fallback path (no ModelLoader): this config load may execute the checkpoint's own # Python. State it loudly -- the same advisory class as a hardcoded trust_remote_code. - logger.warning_once('loading %s with trust_remote_code=True: the checkpoint\'s own Python code ' - 'will execute inside this process. Only use checkpoints you trust.', - self._model_path) + logger.warning_once( + "loading %s with trust_remote_code=True: the checkpoint's own Python code " + "will execute inside this process. Only use checkpoints you trust.", + self._model_path, + ) self.hf_config = AutoConfig.from_pretrained(self._model_path, trust_remote_code=True) if new_special_tokens: self._extend_vocab_for_new_special_tokens(new_special_tokens) @@ -195,8 +211,9 @@ def __init__( align_grad_reduce=align_grad_reduce, nccl_comm_warmup=nccl_comm_warmup, attn_impl=attn_impl, - **kwargs) - self.model: List[nn.Module] = self.strategy.create_megatron_model(load_weights) + **kwargs, + ) + self.model: list[nn.Module] = self.strategy.create_megatron_model(load_weights) if vit_gradient_checkpointing is not None: self._apply_vit_gradient_checkpointing(vit_gradient_checkpointing, vit_gradient_checkpointing_kwargs) if new_special_tokens: @@ -210,14 +227,14 @@ def __init__( # lives in the MtpTrainingState instance; the model keeps thin same-named shells below. self._mtp = MtpTrainingState() # This correctly handles vocab sharding in Tensor Parallelism - self.optimizer_group: Dict[str, MegatronOptimizerGroup] = { + self.optimizer_group: dict[str, MegatronOptimizerGroup] = { _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 MegatronPeft().__call__() - def _apply_vit_gradient_checkpointing(self, enabled: bool, gc_kwargs: Optional[Dict[str, Any]]) -> None: + def _apply_vit_gradient_checkpointing(self, enabled: bool, gc_kwargs: dict[str, Any] | None) -> None: """Checkpoint each HF vision tower of a multimodal mcore model. mcore builds the tower under ``visual`` on the first PP stage only; tower names come from @@ -227,15 +244,16 @@ def _apply_vit_gradient_checkpointing(self, enabled: bool, gc_kwargs: Optional[D tower when the input embedding is frozen (LoRA / quantized). """ from twinkle.utils import deep_getattr + for chunk in self.model: core = chunk - while hasattr(core, 'module'): # unwrap DDP / Float16Module down to the mcore model + while hasattr(core, "module"): # unwrap DDP / Float16Module down to the mcore model core = core.module - visual = getattr(core, 'visual', None) + visual = getattr(core, "visual", None) if visual is None: # not the first pipeline stage, or a plain LLM continue - visual_cls = getattr(getattr(core, 'model_meta', None), 'visual_cls', None) - for tower_name in (getattr(visual_cls, '_vision_tower', None) or []): + visual_cls = getattr(getattr(core, "model_meta", None), "visual_cls", None) + for tower_name in getattr(visual_cls, "_vision_tower", None) or []: tower = deep_getattr(visual, tower_name) if tower is None: continue @@ -249,9 +267,9 @@ def _apply_vit_gradient_checkpointing(self, enabled: bool, gc_kwargs: Optional[D else: tower.gradient_checkpointing_disable() except (NotImplementedError, AttributeError, ValueError) as exc: - logger.warning(f'vision-tower gradient checkpointing skipped for {tower_name}: {exc}') + logger.warning(f"vision-tower gradient checkpointing skipped for {tower_name}: {exc}") - def _extend_vocab_for_new_special_tokens(self, new_special_tokens: List[str]) -> None: + def _extend_vocab_for_new_special_tokens(self, new_special_tokens: list[str]) -> None: """Grow the tokenizer vocab and raise ``hf_config.vocab_size`` to match. mcore builds VocabParallelEmbedding FROM ``hf_config.vocab_size``, so the vocab must be @@ -261,10 +279,11 @@ def _extend_vocab_for_new_special_tokens(self, new_special_tokens: List[str]) -> processor = self._default_tokenizer if processor is None: from transformers import AutoTokenizer + processor = AutoTokenizer.from_pretrained(self.tokenizer_id, trust_remote_code=True) self._default_tokenizer = processor - tokenizer = getattr(processor, 'tokenizer', processor) - added = tokenizer.add_special_tokens({'additional_special_tokens': list(new_special_tokens)}) + tokenizer = getattr(processor, "tokenizer", processor) + added = tokenizer.add_special_tokens({"additional_special_tokens": list(new_special_tokens)}) if added: text_config = self.hf_config.get_text_config() new_vocab_size = ((len(tokenizer) + 127) // 128) * 128 @@ -279,15 +298,16 @@ def _initialize_new_special_token_embeddings(self) -> None: group snapshots the parameters. """ from twinkle.utils import deep_getattr + for chunk in self.model: model = chunk - while hasattr(model, 'module'): # unwrap DDP / Float16Module down to the mcore model + while hasattr(model, "module"): # unwrap DDP / Float16Module down to the mcore model model = model.module init_method = model.config.init_method - if hasattr(model, 'language_model'): # multimodal: the decoder carries the embeddings + if hasattr(model, "language_model"): # multimodal: the decoder carries the embeddings model = model.language_model - for key in ('embedding.word_embeddings', 'output_layer'): - if key == 'output_layer' and model.share_embeddings_and_output_weights: + for key in ("embedding.word_embeddings", "output_layer"): + if key == "output_layer" and model.share_embeddings_and_output_weights: continue # tied to the input embedding, already initialized above module = deep_getattr(model, key) if module is None: @@ -297,23 +317,24 @@ def _initialize_new_special_token_embeddings(self) -> None: if num_to_initialize == 0: continue from megatron.core import mpu + if mpu.get_data_parallel_rank() == 0: - logger.info(f'Initializing {num_to_initialize} zero rows of {key!r} for new special tokens.') + logger.info(f"Initializing {num_to_initialize} zero rows of {key!r} for new special tokens.") tensor = module.weight.new_empty(num_to_initialize, module.weight.shape[1]) module.weight.data[initialize_mask] = init_method(tensor) - if getattr(module.weight, 'main_param', None) is not None: + if getattr(module.weight, "main_param", None) is not None: module.weight.main_param.copy_(module.weight.view(-1)) def _should_bind_device_id_for_process_group(self, backend: str) -> bool: # Keep NCCL's device binding behavior, but avoid binding HCCL's default # PG so Megatron's later Gloo DP groups stay decoupled on NPU. - return backend == 'nccl' + return backend == "nccl" def _construct_default_optimizer_group(self): return MegatronOptimizerGroup( - loss_instance=CrossEntropyLoss(reduction='sum'), + loss_instance=CrossEntropyLoss(reduction="sum"), template=Template(self.tokenizer_id), - processor=InputProcessor(self.device_mesh, framework='megatron'), + processor=InputProcessor(self.device_mesh, framework="megatron"), _device_mesh=self.device_mesh, ) @@ -367,30 +388,31 @@ def _slice_value_for_microbatch(value, mb_start: int, mb_end: int, micro_batch_s return value @remote_function() - def forward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs): - raise NotImplementedError('Megatron only supports `forward_backward` and `forward_only`') + def forward(self, *, inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], **kwargs): + raise NotImplementedError("Megatron only supports `forward_backward` and `forward_only`") # MTP training: thin shells over the MtpTrainingState instance; the same-named methods stay # so subclass overrides and duck-typed callers (``getattr(model, 'pop_mtp_metrics')``) keep resolving. def _mtp_training_enabled(self, *, forward_only: bool, task: str, disable_lora: bool) -> bool: return self._mtp.training_enabled(self, forward_only=forward_only, task=task, disable_lora=disable_lora) - @remote_function(collect='last_pp_first', lazy_collect=False) - def pop_mtp_metrics(self, loss_scale: Optional[float] = None) -> dict[str, float]: + @remote_function(collect="last_pp_first", lazy_collect=False) + def pop_mtp_metrics(self, loss_scale: float | None = None) -> dict[str, float]: """Drain megatron's MTP loss tracker; see :meth:`MtpTrainingState.pop`.""" return self._mtp.pop(loss_scale) - # sync=True orders the workers for PP's NCCL collectives; lazy_collect follows the global default # (True) so the return keeps the Future protocol. Note the sync collect itself still materializes # on the calling side -- the handle is a value wrapper there, resolved like any other. - @remote_function(dispatch='slice_dp', collect=collect_tensor_dict, sync=True) + @remote_function(dispatch="slice_dp", collect=collect_tensor_dict, sync=True) @timed - def forward_only(self, - *, - inputs: Union[InputFeature, List[InputFeature], List[Trajectory]], - micro_batch_size: Optional[int] = None, - **kwargs): + def forward_only( + self, + *, + inputs: InputFeature | list[InputFeature] | list[Trajectory], + micro_batch_size: int | None = None, + **kwargs, + ): """Forward pass without gradient computation. Args: @@ -411,9 +433,10 @@ def generate(self, *args, **kwargs): engine understands. Failing loudly beats silently generating from a partial, renamed view. """ raise NotImplementedError( - 'MegatronModel cannot generate in place: its weights are sharded across TP/PP ranks under ' - 'Megatron parameter names, which no inference engine reads. Use a dedicated sampler ' - '(vLLM/SGLang/Transformers) that loads its own copy of the weights.') + "MegatronModel cannot generate in place: its weights are sharded across TP/PP ranks under " + "Megatron parameter names, which no inference engine reads. Use a dedicated sampler " + "(vLLM/SGLang/Transformers) that loads its own copy of the weights." + ) def generate_stream(self, *args, **kwargs): """Streaming counterpart of :meth:`generate`, refused for the same reason. @@ -422,29 +445,32 @@ def generate_stream(self, *args, **kwargs): is called, not only when a returned iterator is consumed, so a facade surfaces it at the call site. """ raise NotImplementedError( - 'MegatronModel cannot generate in place: its weights are sharded across TP/PP ranks under ' - 'Megatron parameter names, which no inference engine reads. Use a dedicated sampler ' - '(vLLM/SGLang/Transformers) that loads its own copy of the weights.') + "MegatronModel cannot generate in place: its weights are sharded across TP/PP ranks under " + "Megatron parameter names, which no inference engine reads. Use a dedicated sampler " + "(vLLM/SGLang/Transformers) that loads its own copy of the weights." + ) - @remote_function(collect='mean') + @remote_function(collect="mean") @timed def calculate_loss(self, **kwargs): - raise NotImplementedError('Megatron only supports `forward_backward` and `forward_only`') + raise NotImplementedError("Megatron only supports `forward_backward` and `forward_only`") @remote_function() def backward(self, **kwargs): - raise NotImplementedError('Megatron only supports `forward_backward` and `forward_only`') + raise NotImplementedError("Megatron only supports `forward_backward` and `forward_only`") - @remote_function(dispatch='slice_dp', collect=collect_tensor_dict, sync=True) + @remote_function(dispatch="slice_dp", collect=collect_tensor_dict, sync=True) @timed - def forward_backward(self, - *, - inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], - micro_batch_size: Optional[int] = None, - loss_key: Optional[str] = None, - loss_weight: float = 1.0, - finalize_gradients: bool = True, - **kwargs): + def forward_backward( + self, + *, + inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], + micro_batch_size: int | None = None, + loss_key: str | None = None, + loss_weight: float = 1.0, + finalize_gradients: bool = True, + **kwargs, + ): """Combined forward and backward pass using Megatron's scheduler. sync=True is required in Ray mode: pipeline parallel's NCCL P2P needs all ranks to enter @@ -462,50 +488,54 @@ def forward_backward(self, self._lazy_wrap_model() self._lazy_finish_param_config() from functools import partial + from megatron.core import parallel_state as mpu from megatron.core.pipeline_parallel import get_forward_backward_func - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) - disable_lora = kwargs.pop('disable_lora', False) - temperature = float(kwargs.pop('temperature', 1.0)) - forward_only = kwargs.pop('forward_only', False) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) + disable_lora = kwargs.pop("disable_lora", False) + temperature = float(kwargs.pop("temperature", 1.0)) + forward_only = kwargs.pop("forward_only", False) # Per-call opt-in to skip the configured loss and return logps only (an empty-loss probe: # 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) - return_logits = kwargs.pop('return_logits', False) - task = kwargs.pop('task', 'causal_lm') + 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 # backend's action. Popped before loss_extra_kwargs_per_mb so it never reaches the loss. - router_replay_mode = kwargs.pop('router_replay_mode', None) - router_replay_manual_cleanup = kwargs.pop('router_replay_manual_cleanup', False) + router_replay_mode = kwargs.pop("router_replay_mode", None) + router_replay_manual_cleanup = kwargs.pop("router_replay_manual_cleanup", False) optimizer_config = self.optimizer_group[adapter_name] loss_instance = optimizer_config.get_loss(loss_key) if not inputs: - raise ValueError('inputs empty, check your DataLoader outputs') + raise ValueError("inputs empty, check your DataLoader outputs") if forward_only and not finalize_gradients: - raise ValueError('forward_only does not support finalize_gradients=False') + raise ValueError("forward_only does not support finalize_gradients=False") if not forward_only and not finalize_gradients and self.strategy.uses_fsdp(self.strategy.ddp_config): - ddp_config = getattr(self.model[0], 'ddp_config', None) + ddp_config = getattr(self.model[0], "ddp_config", None) sharding_strategy = getattr( ddp_config, - 'data_parallel_sharding_strategy', - self.strategy.ddp_config.get('data_parallel_sharding_strategy'), + "data_parallel_sharding_strategy", + self.strategy.ddp_config.get("data_parallel_sharding_strategy"), ) - if sharding_strategy in {'optim_grads', 'optim_grads_params'}: + if sharding_strategy in {"optim_grads", "optim_grads_params"}: raise ValueError( - 'Deferred backward is not supported by Megatron-FSDP with gradient sharding ' - f'({sharding_strategy}); use data_parallel_sharding_strategy="no_shard" or "optim".') + "Deferred backward is not supported by Megatron-FSDP with gradient sharding " + f'({sharding_strategy}); use data_parallel_sharding_strategy="no_shard" or "optim".' + ) completing_pending = False if not forward_only: completing_pending = optimizer_config.prepare_backward_transaction(finalize_gradients, loss_weight) transactional = not forward_only and (not finalize_gradients or completing_pending) - if (isinstance(inputs, dict) and self._not_encoded(inputs)) or (isinstance(inputs, list) - and self._not_encoded(inputs[0])): + if (isinstance(inputs, dict) and self._not_encoded(inputs)) or ( + isinstance(inputs, list) and self._not_encoded(inputs[0]) + ): # Trajectory or List[Trajectory] - assert optimizer_config.template is not None, \ - 'Use set_template to add a template when trying to input `List[Trajectory]`' + assert optimizer_config.template is not None, ( + "Use set_template to add a template when trying to input `List[Trajectory]`" + ) if isinstance(inputs, dict): inputs = [inputs] inputs = optimizer_config.template.batch_encode(inputs) # noqa @@ -513,32 +543,39 @@ def forward_backward(self, inputs = [inputs] local_batch_normalizer = loss_instance.batch_normalizer(inputs) if transactional else None processor: InputProcessor = optimizer_config.processor - assert isinstance(processor, InputProcessor), 'Set InputProcessor correctly before forwarding' + assert isinstance(processor, InputProcessor), "Set InputProcessor correctly before forwarding" if micro_batch_size is None: # Compatible with DPO micro_batch_size = min(2, len(inputs)) unwrapped_model = self.strategy.unwrap_model(self.model)[0] + 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 cp_size = self.device_mesh.cp_world_size # MoE routing replay orchestration of this call (a null object when disabled or mode-less); # the per-microbatch dance lives in forward_step_func (pre_mb / post_mb). - rr = RouterReplay(enabled=self._router_replay_enabled, mode=router_replay_mode, - forward_only=forward_only, cp_size=cp_size, - apply_patch=self._maybe_apply_router_replay) + rr = RouterReplay( + enabled=self._router_replay_enabled, + mode=router_replay_mode, + forward_only=forward_only, + cp_size=cp_size, + apply_patch=self._maybe_apply_router_replay, + ) # Check actual sequence_parallel setting from model config # Bridge may auto-enable sequence_parallel for MoE models if self.variable_seq_lengths: - seq_length = max(inp['input_ids'].shape[-1] for inp in inputs) + seq_length = max(inp["input_ids"].shape[-1] for inp in inputs) else: - original_seq_length = inputs[0]['input_ids'].shape[1] * (cp_size or 1) + original_seq_length = inputs[0]["input_ids"].shape[1] * (cp_size or 1) if cp_size > 1: divisor = 2 * cp_size elif self.strategy.sequence_parallel and self.device_mesh.tp_world_size > 1: @@ -592,18 +629,19 @@ def post_loss_function(output_tensor, inputs, logps, unpacked_logits=None, entro logits = unpacked_logits if unpacked_logits is not None else output_tensor outputs = ModelOutput(logits=logits, logps=logps) if entropies is not None: - outputs['entropies'] = entropies + 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). zero = torch.zeros((), device=Platform.get_local_device(), dtype=torch.float32) - return self.strategy.reduce_loss(zero, zero.clone(), output_tensor, logps) + reported_logits = output_tensor if return_logits else None + return self.strategy.reduce_loss(zero, zero.clone(), reported_logits, logps) result = loss_instance(inputs, outputs, **current_kwargs) if unpacked_logits is not None: - outputs.pop('logits', None) + outputs.pop("logits", None) del unpacked_logits - losses = result['loss'] - counts = result['num_tokens'] + losses = result["loss"] + counts = result["num_tokens"] if not counts: # Later will gather this value, so it becomes: # 1. SUM loss: gather_sum(local_num_tokens) = global_num_tokens @@ -614,31 +652,32 @@ 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) - if result.get('channel_loss') is not None: - reduced[2]['channel_loss'] = result['channel_loss'] + reported_logits = output_tensor if return_logits else None + reduced = self.strategy.reduce_loss(losses, counts, reported_logits, logps) + if result.get("channel_loss") is not None: + reduced[2]["channel_loss"] = result["channel_loss"] return reduced # Define forward step function for Megatron # forward_step_func(data_iterator, model) -> (output_tensor, partial(loss_func)) def forward_step_func(data_iterator, model): batch = next(data_iterator) - labels = batch.pop('labels', None) - loss_scale = batch.pop('loss_scale', None) - channel = batch.pop('channel', None) + labels = batch.pop("labels", None) + loss_scale = batch.pop("loss_scale", None) + channel = batch.pop("channel", None) # Not a model argument; restored below so the loss can read it. - completion_mask = batch.pop('completion_mask', None) + completion_mask = batch.pop("completion_mask", None) # The whole-model routing to replay (R3, delivered by the sampler): popped so it never # reaches model(**model_kwargs); the replay dance below slices it to the local shard. - routed_experts_mb = batch.pop('routed_experts', None) + routed_experts_mb = batch.pop("routed_experts", None) # Per-token mask selecting which tokens replay vs recompute natively; not a model argument. - replay_mask_mb = batch.pop('replay_mask', None) + replay_mask_mb = batch.pop("replay_mask", None) # MTP joint training: labels stay out of the model call so the main loss stays external; # the MTP heads get next-token targets on a separate keyword, never left in ``batch``. model_kwargs = batch if feed_mtp and labels is not None: model_kwargs = dict(batch) - model_kwargs['mtp_labels'] = labels + model_kwargs["mtp_labels"] = labels unwrapped_model = self.strategy.unwrap_model([model])[0] rr.pre_mb(unwrapped_model, batch, routed_experts_mb, replay_mask_mb) if disable_lora and isinstance(unwrapped_model, PeftModel): @@ -648,13 +687,13 @@ def forward_step_func(data_iterator, model): output_tensor = model(**model_kwargs) rr.post_mb(unwrapped_model, batch) - batch['labels'] = labels + batch["labels"] = labels if loss_scale is not None: - batch['loss_scale'] = loss_scale + batch["loss_scale"] = loss_scale if channel is not None: - batch['channel'] = channel + batch["channel"] = channel if completion_mask is not None: - batch['completion_mask'] = completion_mask + batch["completion_mask"] = completion_mask logps = None unpacked_logits = None entropies = None @@ -662,58 +701,61 @@ def forward_step_func(data_iterator, model): _loss_instance = loss_instance is_last_pp = mpu.is_pipeline_last_stage(False, unwrapped_model.vp_stage) try: - if task == 'embedding': + if task == "embedding": # MegatronEmbeddingPatch already pooled output to [n_seqs, hidden] on last PP stage. if is_last_pp: embeddings = output_tensor - elif task in ('seq_cls', 'reranker', 'generative_reranker'): + elif task in ("seq_cls", "reranker", "generative_reranker"): # mcore-bridge emits per-token head output; pool the last valid token per sequence # here (transformers does this inside *ForSequenceClassification). These losses read # logits (require_logps=False); the pool runs regardless of labels because a # preference row carries no next-token labels at all. if is_last_pp: - _packed = batch.get('packed_seq_params') + _packed = batch.get("packed_seq_params") # CP reconstruct first (like the causal_lm branch), then take the last token. - cu_seqlens_q = getattr(_packed, 'cu_seqlens_q', None) if _packed is not None else None + cu_seqlens_q = getattr(_packed, "cu_seqlens_q", None) if _packed is not None else None output_tensor = processor.postprocess_tensor_cp(output_tensor, cu_seqlens=cu_seqlens_q) if _packed is None: - _am = batch.get('attention_mask') + _am = batch.get("attention_mask") if _am is None: - _am = batch.get('attention_mask_2d') + _am = batch.get("attention_mask_2d") if _am is None: # No mask in the batch (the collator right-pads without emitting one): pool # the final position, exactly what transformers does when handed no mask. - _last_idx = torch.full((output_tensor.shape[0], ), - output_tensor.shape[1] - 1, - dtype=torch.long, - device=output_tensor.device) + _last_idx = torch.full( + (output_tensor.shape[0],), + output_tensor.shape[1] - 1, + dtype=torch.long, + device=output_tensor.device, + ) else: # attention_mask_2d compatibility: a 4D mask collapses to per-row validity. if _am.dim() > 2: _am = (~_am).sum(dim=(1, 2)) > 0 _last_idx = processor._last_valid_indices(_am.long()) unpacked_logits = output_tensor[ - torch.arange(output_tensor.shape[0], device=output_tensor.device), _last_idx] + torch.arange(output_tensor.shape[0], device=output_tensor.device), _last_idx + ] else: # Packed: one row [1, total, C]; each segment ends at cu_seqlens_q[i]+seq_lens[i]-1. _n = _packed.seq_lens.shape[0] _last_idx = _packed.cu_seqlens_q[:_n] + _packed.seq_lens - 1 unpacked_logits = output_tensor[0, _last_idx] - elif task == 'value': + elif task == "value": # PPO critic: the same num_labels=1 seq_cls head, but keep its per-token V(s_t) # riding the logits channel ([b, s]); PPOValueLoss masks to response tokens. # Non-packed / non-SP only, mirroring the transformers TransformersValuePatch scope. if is_last_pp and labels is not None: - _packed = batch.get('packed_seq_params') - cu_seqlens_q = getattr(_packed, 'cu_seqlens_q', None) if _packed is not None else None + _packed = batch.get("packed_seq_params") + cu_seqlens_q = getattr(_packed, "cu_seqlens_q", None) if _packed is not None else None output_tensor = processor.postprocess_tensor_cp(output_tensor, cu_seqlens=cu_seqlens_q) # [b, s, 1] -> [b, s] (num_labels=1); already per-token, no pooling. unpacked_logits = output_tensor.squeeze(-1) if output_tensor.dim() == 3 else output_tensor elif labels is not None and is_last_pp: - _loss_require_logps = getattr(_loss_instance, 'require_logps', True) - _loss_require_entropy = getattr(_loss_instance, 'require_entropy', False) - _packed = batch.get('packed_seq_params') - cu_seqlens_q = getattr(_packed, 'cu_seqlens_q', None) if _packed is not None else None + _loss_require_logps = getattr(_loss_instance, "require_logps", True) + _loss_require_entropy = getattr(_loss_instance, "require_entropy", False) + _packed = batch.get("packed_seq_params") + cu_seqlens_q = getattr(_packed, "cu_seqlens_q", None) if _packed is not None else None if _loss_require_logps: loss_mask = (labels != -100).bool() masked_labels = labels.clone() @@ -727,36 +769,36 @@ def forward_step_func(data_iterator, model): logps = processor.postprocess_tensor_cp(logps, cu_seqlens=cu_seqlens_q) if entropies is not None: entropies = processor.postprocess_tensor_cp(entropies, cu_seqlens=cu_seqlens_q) - batch['labels'] = processor.postprocess_tensor_cp(labels, cu_seqlens=cu_seqlens_q) - if batch.get('loss_scale') is not None: - batch['loss_scale'] = processor.postprocess_tensor_cp( - batch['loss_scale'], cu_seqlens=cu_seqlens_q) + batch["labels"] = processor.postprocess_tensor_cp(labels, cu_seqlens=cu_seqlens_q) + if batch.get("loss_scale") is not None: + batch["loss_scale"] = processor.postprocess_tensor_cp( + batch["loss_scale"], cu_seqlens=cu_seqlens_q + ) if completion_mask is not None: # Same index space as labels, so it needs the same CP reassembly. - batch['completion_mask'] = processor.postprocess_tensor_cp( - completion_mask, cu_seqlens=cu_seqlens_q) - if 'position_ids' in batch: - pos = batch['position_ids'] + batch["completion_mask"] = processor.postprocess_tensor_cp( + completion_mask, cu_seqlens=cu_seqlens_q + ) + if "position_ids" in batch: + pos = batch["position_ids"] if pos.dim() == 3: pos = pos[0] # [2/3, 1, seq] → [1, seq] - batch['position_ids'] = processor.postprocess_tensor_cp(pos, cu_seqlens=cu_seqlens_q) + batch["position_ids"] = processor.postprocess_tensor_cp(pos, cu_seqlens=cu_seqlens_q) # Unpack packed sequences into per-sequence batch format - _outputs = {'logps': logps} + _outputs = {"logps": logps} if entropies is not None: - _outputs['entropies'] = entropies - if hasattr(_loss_instance, 'require_logits') and _loss_instance.require_logits: - _outputs['logits'] = output_tensor + _outputs["entropies"] = entropies + if hasattr(_loss_instance, "require_logits") and _loss_instance.require_logits: + _outputs["logits"] = output_tensor batch, _outputs = processor.unpack_packed_sequences(batch, _outputs) - logps = _outputs['logps'] - entropies = _outputs.get('entropies', None) - unpacked_logits = _outputs.get('logits', None) + logps = _outputs["logps"] + entropies = _outputs.get("entropies", None) + unpacked_logits = _outputs.get("logits", None) except Exception as e: # Log the rank-attributed failure, then re-raise: swallowing it would let this DP # rank skip the PP collectives while the others wait, and a degraded value would # mask the failure. - logger.warning('[nccl_safe] forward_step_func data processing error on rank: ' - '%s: %s', - type(e).__name__, e) + logger.warning("[forward_step] data processing error on rank: %s: %s", type(e).__name__, e) raise return output_tensor, partial( post_loss_function, @@ -776,7 +818,7 @@ def forward_step_func(data_iterator, model): if vpp_size in (0, 1): vpp_size = None - micro_batch_size = inputs[0]['input_ids'].shape[0] + micro_batch_size = inputs[0]["input_ids"].shape[0] if vpp_size is None: data_iter = iter(inputs) else: @@ -788,14 +830,15 @@ def forward_step_func(data_iterator, model): # transactional / deferred-finalize pass is owned by the strategy. rr.arm() with self.strategy.transactional_grad_schedule( - self.model, - forward_only=forward_only, - transactional=transactional, - loss_weight=loss_weight, - finalize_gradients=finalize_gradients, - completing_pending=completing_pending, - global_normalizer=global_normalizer, - **kwargs): + self.model, + forward_only=forward_only, + transactional=transactional, + loss_weight=loss_weight, + finalize_gradients=finalize_gradients, + completing_pending=completing_pending, + global_normalizer=global_normalizer, + **kwargs, + ): with _resolve_task_context(self.model, task): losses = forward_backward_func( forward_step_func=forward_step_func, @@ -817,22 +860,22 @@ def forward_step_func(data_iterator, model): if losses: for loss_dict in losses: if isinstance(loss_dict, dict): - if 'loss' in loss_dict: - loss += loss_dict['loss'] + if "loss" in loss_dict: + loss += loss_dict["loss"] n_loss_entries += 1 - if 'logits' in loss_dict: - logits.append(loss_dict['logits']) - if 'logps' in loss_dict: - logps.append(loss_dict['logps']) - if 'num_tokens' in loss_dict: - count += loss_dict['num_tokens'] - for channel, stats in loss_dict.get('channel_loss', {}).items(): + if loss_dict.get("logits") is not None: + logits.append(loss_dict["logits"]) + if loss_dict.get("logps") is not None: + logps.append(loss_dict["logps"]) + if "num_tokens" in loss_dict: + count += loss_dict["num_tokens"] + for channel, stats in loss_dict.get("channel_loss", {}).items(): if channel in channel_loss: channel_loss[channel] = channel_loss[channel] + stats else: channel_loss[channel] = stats elif isinstance(loss_dict, torch.Tensor): - raise ValueError('Expected loss dict, got tensor') + raise ValueError("Expected loss dict, got tensor") if n_loss_entries > 1: # Backend equivalence: the transformers backend reports ONE mean loss for the whole batch, so # a run split into several microbatches must average the per-microbatch losses, not SUM them @@ -852,8 +895,11 @@ def forward_step_func(data_iterator, model): if mpu.get_pipeline_model_parallel_world_size() > 1: src_rank = mpu.get_pipeline_model_parallel_last_rank() pp_group = mpu.get_pipeline_model_parallel_group() - loss_f = loss.detach().float() if isinstance(loss, torch.Tensor) else torch.tensor( - float(loss), device=Platform.get_local_device()) + loss_f = ( + loss.detach().float() + if isinstance(loss, torch.Tensor) + else torch.tensor(float(loss), device=Platform.get_local_device()) + ) count_f = torch.as_tensor(float(num_tokens), device=loss_f.device, dtype=torch.float32) # One broadcast carries both scalars so the collective count stays identical on all stages. packed = torch.stack([loss_f, count_f]) @@ -885,11 +931,10 @@ def forward_step_func(data_iterator, model): inputs = processor.unpack_inputs(inputs, task=task) model_output = ModelOutput(logits=logits, loss=loss, logps=logps, num_tokens=num_tokens) if channel_loss: - model_output['channel_loss'] = channel_loss + model_output["channel_loss"] = channel_loss # R2 RECORD: hand the captured whole-model routing back to the driver (symmetric with # TransformersModel.forward's return_outputs['routed_experts']). - rr.finish(model_output, unwrapped_model.config, num_microbatches, vpp_size, - router_replay_manual_cleanup) + rr.finish(model_output, unwrapped_model.config, num_microbatches, vpp_size, router_replay_manual_cleanup) if forward_only: optimizer_config.eval_status.inputs = inputs optimizer_config.eval_status.outputs = model_output @@ -898,61 +943,60 @@ def forward_step_func(data_iterator, model): status = optimizer_config.train_status status.inputs = inputs status.outputs = model_output - status.forward_kwargs = {**kwargs, '_metric_loss_key': loss_key} - status.loss_reduction = 'mean' if transactional else getattr(loss_instance, 'reduction', 'mean') + status.forward_kwargs = {**kwargs, "_metric_loss_key": loss_key} + status.loss_reduction = "mean" if transactional else getattr(loss_instance, "reduction", "mean") if transactional: metric_step = optimizer_config.cur_step - 1 if finalize_gradients else optimizer_config.cur_step optimizer_config.accumulate_metrics(True, include_default=False, metric_step=metric_step) status.inputs = None status.outputs = None - total_loss = optimizer_config.record_transaction_loss(model_output.get('loss', 0.0), - finalize_gradients) + total_loss = optimizer_config.record_transaction_loss(model_output.get("loss", 0.0), finalize_gradients) if finalize_gradients: status.inputs = {} status.outputs = ModelOutput(loss=total_loss) status.forward_kwargs = {} - status.loss_reduction = 'mean' + status.loss_reduction = "mean" return ModelOutput(loss=total_loss) return model_output - @remote_function(dispatch='all') + @remote_function(dispatch="all") def clip_grad_norm(self, max_grad_norm: float = 1.0, norm_type: int = 2, **kwargs): - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) - self.optimizer_group[adapter_name].require_finalized('clip gradients') + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) + self.optimizer_group[adapter_name].require_finalized("clip gradients") # Megatron optimizer will cover this function. return 0 - @remote_function(dispatch='all') + @remote_function(dispatch="all") def step(self, **kwargs): """Optimizer step (dispatch='all' so every worker executes it). For non-DDP models (e.g. PEFT/LoRA) gradients are NOT synchronized across DP ranks; each replica trains independently on its own data.""" - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.require_finalized('step the optimizer') + optimizer_config.require_finalized("step the optimizer") - if not optimizer_config.do_grad_sync(kwargs.pop('gradient_accumulation_steps', None)): + if not optimizer_config.do_grad_sync(kwargs.pop("gradient_accumulation_steps", None)): return optimizer = optimizer_config.optimizer - assert optimizer is not None, 'Set optimizer correctly before stepping' + assert optimizer is not None, "Set optimizer correctly before stepping" # Megatron optimizer step() returns (success, grad_norm, num_zeros) - optim_params = kwargs.pop('optim_params', {}) + optim_params = kwargs.pop("optim_params", {}) if optim_params: for group in optimizer.param_groups: - group['lr'] = optim_params['lr'] - if group['weight_decay'] > 0.0 and optim_params.get('weight_decay', None) is not None: - group['weight_decay'] = optim_params['weight_decay'] - if optim_params.get('eps') is not None: - group['eps'] = optim_params['eps'] - if optim_params.get('betas') is not None: - group['betas'] = optim_params['betas'] + group["lr"] = optim_params["lr"] + if group["weight_decay"] > 0.0 and optim_params.get("weight_decay", None) is not None: + group["weight_decay"] = optim_params["weight_decay"] + if optim_params.get("eps") is not None: + group["eps"] = optim_params["eps"] + if optim_params.get("betas") is not None: + group["betas"] = optim_params["betas"] success, grad_norm, num_zeros = optimizer.step() # Store grad_norm for later retrieval if grad_norm is not None: - grad_norm = grad_norm.item() if hasattr(grad_norm, 'item') else grad_norm + grad_norm = grad_norm.item() if hasattr(grad_norm, "item") else grad_norm optimizer_config._last_grad_norm = float(grad_norm) if grad_norm is not None else 0.0 optimizer_config._last_step_success = success @@ -964,15 +1008,16 @@ def _is_model_ddp_wrapped(self) -> bool: """ from megatron.core.distributed import DistributedDataParallel as MegatronDDP from torch.nn.parallel import DistributedDataParallel as TorchDDP + return isinstance(self.model[0], (MegatronDDP, TorchDDP)) - @remote_function(dispatch='all') + @remote_function(dispatch="all") def zero_grad(self, **kwargs): """Zero gradients. For DDP-wrapped models the gradient buffer is ALWAYS zeroed (the schedule tracks accumulation through it); do_grad_sync only gates optimizer.zero_grad().""" - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.require_finalized('zero gradients') + optimizer_config.require_finalized("zero gradients") # For DDP-wrapped models, ALWAYS zero the gradient buffer # This is essential because Megatron's forward_backward_func uses @@ -980,10 +1025,10 @@ def zero_grad(self, **kwargs): # (self.model is a list of chunks; zero_grad_buffer lives on each chunk) if self._is_model_ddp_wrapped(): for model_chunk in self.model: - if hasattr(model_chunk, 'zero_grad_buffer'): + if hasattr(model_chunk, "zero_grad_buffer"): model_chunk.zero_grad_buffer() - if not optimizer_config.do_grad_sync(kwargs.pop('gradient_accumulation_steps', None)): + if not optimizer_config.do_grad_sync(kwargs.pop("gradient_accumulation_steps", None)): return optimizer = optimizer_config.optimizer @@ -997,46 +1042,47 @@ def lr_step(self, **kwargs): Args: **kwargs: Additional arguments. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.require_finalized('step the learning-rate scheduler') + optimizer_config.require_finalized("step the learning-rate scheduler") - if not optimizer_config.do_grad_sync(kwargs.pop('gradient_accumulation_steps', None)): + if not optimizer_config.do_grad_sync(kwargs.pop("gradient_accumulation_steps", None)): return lr_scheduler = optimizer_config.lr_scheduler if lr_scheduler is not None: # Megatron's OptimizerParamScheduler.step() requires increment argument - increment = kwargs.pop('increment', 1) + increment = kwargs.pop("increment", 1) lr_scheduler.step(increment=increment) - @remote_function(dispatch='all') - def set_loss(self, loss_cls: Union[Loss, Type[Loss], str, Callable[[InputFeature, ModelOutput, ...], torch.Tensor]], - **kwargs): + @remote_function(dispatch="all") + def set_loss( + self, loss_cls: Loss | type[Loss] | str | Callable[[InputFeature, ModelOutput, ...], torch.Tensor], **kwargs + ): """Register a loss used by Megatron's scheduler-side loss callback. Args: loss_cls: Loss class, instance, or registered name. **kwargs: Constructor arguments plus optional ``loss_key`` and ``adapter_name``. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) - loss_key = kwargs.pop('loss_key', None) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) + loss_key = kwargs.pop("loss_key", None) optimizer_config = self.optimizer_group[adapter_name] loss = construct_class(loss_cls, Loss, twinkle.loss, **kwargs) # Gather over the DP group, not WORLD: under PP the loss runs only on the last # pipeline stage, so a WORLD-group all-gather would deadlock on earlier-stage ranks. - if getattr(loss, 'process_group', None) is None and getattr(optimizer_config, '_dp_group', None) is not None: + if getattr(loss, "process_group", None) is None and getattr(optimizer_config, "_dp_group", None) is not None: loss.process_group = optimizer_config._dp_group optimizer_config.set_loss(loss, loss_key=loss_key) - @remote_function(dispatch='all') + @remote_function(dispatch="all") def empty_cache(self, **kwargs): """Release the caching allocator's unused blocks on every worker; backend-agnostic through torch_util, mirrors TransformersModel.empty_cache.""" torch_util.empty_cache() @remote_function() - def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] = None, **kwargs): + def add_metric(self, metric_cls: Metric | str, is_training: bool | None = None, **kwargs): """Add an eval metric Args: @@ -1046,11 +1092,11 @@ def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] adapter_name: Lora adapter name. Any parameters needed to construct the metric_cls instance. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) - loss_key = kwargs.pop('loss_key', None) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) + loss_key = kwargs.pop("loss_key", None) optimizer_config = self.optimizer_group[adapter_name] - kwargs['device_mesh'] = self.device_mesh - kwargs['process_group'] = optimizer_config._dp_group + kwargs["device_mesh"] = self.device_mesh + kwargs["process_group"] = optimizer_config._dp_group if is_training is None or is_training is True: metric = construct_class(metric_cls, Metric, twinkle.metric, **kwargs) metric._twinkle_loss_key = loss_key @@ -1060,8 +1106,8 @@ def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] metric._twinkle_loss_key = loss_key optimizer_config.eval_status.metrics.append(metric) - @remote_function(dispatch='all') - def set_optimizer(self, optimizer_cls: Union[Optimizer, Type[Optimizer], str], **kwargs): + @remote_function(dispatch="all") + def set_optimizer(self, optimizer_cls: Optimizer | type[Optimizer] | str, **kwargs): """Set optimizer. Args: @@ -1072,25 +1118,26 @@ def set_optimizer(self, optimizer_cls: Union[Optimizer, Type[Optimizer], str], * - For standard optimizers: lr, weight_decay, etc. - For MegatronDistributed: use_distributed_optimizer, clip_grad, etc. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] if not self._model_wrapped: self.model = self.strategy.wrap_model(self.model) self._model_wrapped = True # Check if requesting Megatron distributed optimizer - if not optimizer_cls or optimizer_cls in ('MegatronOptimizer', 'default', 'Adam'): + if not optimizer_cls or optimizer_cls in ("MegatronOptimizer", "default", "Adam"): optimizer_config.optimizer = self._create_megatron_optimizer(**kwargs) # noqa else: raise NotImplementedError( - f'Unsupported optimizer: {optimizer_cls}, only support MegatronOptimizer currently.') + f"Unsupported optimizer: {optimizer_cls}, only support MegatronOptimizer currently." + ) - @remote_function(collect='last_pp_first', lazy_collect=False) + @remote_function(collect="last_pp_first", lazy_collect=False) @timed def calculate_metric(self, is_training, **kwargs): - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.require_finalized('calculate metrics') + optimizer_config.require_finalized("calculate metrics") return optimizer_config.calculate_metrics(is_training) def _create_megatron_optimizer(self, **kwargs): @@ -1112,26 +1159,26 @@ def _create_megatron_optimizer(self, **kwargs): from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer # Build optimizer config - lr = kwargs.pop('lr', 1e-4) - self.use_distributed_optimizer: bool = kwargs.pop('use_distributed_optimizer', self.use_distributed_optimizer) + lr = kwargs.pop("lr", 1e-4) + self.use_distributed_optimizer: bool = kwargs.pop("use_distributed_optimizer", self.use_distributed_optimizer) # Pop before building OptimizerConfig: these are get_megatron_optimizer args, not config fields. - config_overrides = kwargs.pop('config_overrides', None) - param_groups_spec = kwargs.pop('param_groups_spec', None) + config_overrides = kwargs.pop("config_overrides", None) + param_groups_spec = kwargs.pop("param_groups_spec", None) opt_config = OptimizerConfig( - optimizer=kwargs.pop('optimizer', 'adam'), + optimizer=kwargs.pop("optimizer", "adam"), lr=lr, - min_lr=kwargs.pop('min_lr', 0.0), - weight_decay=kwargs.pop('weight_decay', 0.01), - adam_beta1=kwargs.pop('adam_beta1', 0.9), - adam_beta2=kwargs.pop('adam_beta2', 0.999), - adam_eps=kwargs.pop('adam_eps', 1e-8), - clip_grad=kwargs.pop('clip_grad', 1.0), - bf16=kwargs.pop('bf16', self.mixed_precision == 'bf16'), - fp16=kwargs.pop('fp16', self.mixed_precision == 'fp16'), + min_lr=kwargs.pop("min_lr", 0.0), + weight_decay=kwargs.pop("weight_decay", 0.01), + adam_beta1=kwargs.pop("adam_beta1", 0.9), + adam_beta2=kwargs.pop("adam_beta2", 0.999), + adam_eps=kwargs.pop("adam_eps", 1e-8), + clip_grad=kwargs.pop("clip_grad", 1.0), + bf16=kwargs.pop("bf16", self.mixed_precision == "bf16"), + fp16=kwargs.pop("fp16", self.mixed_precision == "fp16"), use_distributed_optimizer=self.use_distributed_optimizer, - overlap_param_gather=kwargs.pop('overlap_param_gather', False), - log_num_zeros_in_grad=kwargs.pop('log_num_zeros_in_grad', False), + overlap_param_gather=kwargs.pop("overlap_param_gather", False), + log_num_zeros_in_grad=kwargs.pop("log_num_zeros_in_grad", False), **kwargs, ) @@ -1139,21 +1186,25 @@ def _create_megatron_optimizer(self, **kwargs): # Layer the spec's rules on top of the standard defaults (bias / length-1 wd skip) -- # or a caller-supplied config_overrides -- so per-tower / LoRA+ LRs do not drop the skip. from megatron.core.optimizer import get_standard_config_overrides + from twinkle.module.optimizer import build_param_group_overrides, resolve_megatron_tower_prefixes - base_overrides = config_overrides if config_overrides is not None else get_standard_config_overrides( - opt_config) + + base_overrides = ( + config_overrides if config_overrides is not None else get_standard_config_overrides(opt_config) + ) spec_overrides = build_param_group_overrides( param_groups_spec, model_chunks=self.model, tower_prefixes=resolve_megatron_tower_prefixes(self.model), base_lr=opt_config.lr, - base_min_lr=opt_config.min_lr) + base_min_lr=opt_config.min_lr, + ) config_overrides = {**base_overrides, **spec_overrides} # Ensure each model chunk has ddp_config attached (required by Megatron optimizer) model_chunks = self.model for model_chunk in model_chunks: - assert hasattr(model_chunk, 'ddp_config') + assert hasattr(model_chunk, "ddp_config") optimizer = get_megatron_optimizer( config=opt_config, model_chunks=model_chunks, @@ -1163,22 +1214,23 @@ def _create_megatron_optimizer(self, **kwargs): def _create_megatron_scheduler(self, optimizer, lr_decay_steps, max_lr=1e-4, **kwargs): from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler + return OptimizerParamScheduler( optimizer, - init_lr=kwargs.pop('init_lr', 0.0), + init_lr=kwargs.pop("init_lr", 0.0), max_lr=max_lr, - min_lr=kwargs.pop('min_lr', 0.0), - lr_warmup_steps=kwargs.pop('lr_warmup_steps', 0), + min_lr=kwargs.pop("min_lr", 0.0), + lr_warmup_steps=kwargs.pop("lr_warmup_steps", 0), lr_decay_steps=lr_decay_steps, - lr_decay_style=kwargs.pop('lr_decay_style', 'cosine'), - start_wd=kwargs.pop('start_wd', 0.01), - end_wd=kwargs.pop('end_wd', 0.01), + lr_decay_style=kwargs.pop("lr_decay_style", "cosine"), + start_wd=kwargs.pop("start_wd", 0.01), + end_wd=kwargs.pop("end_wd", 0.01), wd_incr_steps=lr_decay_steps, - wd_incr_style=kwargs.pop('wd_incr_style', 'constant'), + wd_incr_style=kwargs.pop("wd_incr_style", "constant"), **kwargs, ) - def _get_trainable_parameters(self, adapter_name: str = _default_adapter_name) -> Dict[str, nn.Parameter]: + def _get_trainable_parameters(self, adapter_name: str = _default_adapter_name) -> dict[str, nn.Parameter]: """Get trainable parameters. Args: @@ -1190,7 +1242,7 @@ def _get_trainable_parameters(self, adapter_name: str = _default_adapter_name) - is_default = adapter_name == _default_adapter_name # Same pattern family as the transformers backend: LoRA deltas plus the fully trained # modules_to_save copies -- missing it here would freeze those copies out of the optimizer. - pattern = re.compile(rf'\.(?:lora_\w+|modules_to_save)\.{re.escape(adapter_name)}') + pattern = re.compile(rf"\.(?:lora_\w+|modules_to_save)\.{re.escape(adapter_name)}") params = {} model = self.strategy.unwrap_model(self.model) @@ -1200,39 +1252,42 @@ def _get_trainable_parameters(self, adapter_name: str = _default_adapter_name) - params[name] = param return params - @remote_function(dispatch='all') - def set_lr_scheduler(self, scheduler_cls: Union[LRScheduler, Type[LRScheduler], str], **kwargs): + @remote_function(dispatch="all") + def set_lr_scheduler(self, scheduler_cls: LRScheduler | type[LRScheduler] | str, **kwargs): """Set learning rate scheduler. Args: scheduler_cls: Scheduler class or string name. **kwargs: Additional arguments. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] optimizer = optimizer_config.optimizer - if not scheduler_cls or scheduler_cls in ('OptimizerParamScheduler', 'default'): + if not scheduler_cls or scheduler_cls in ("OptimizerParamScheduler", "default"): optimizer_config.lr_scheduler = self._create_megatron_scheduler(optimizer, **kwargs) # noqa else: raise NotImplementedError( - f'Unsupported scheduler: {scheduler_cls}, only support OptimizerParamScheduler currently.') + f"Unsupported scheduler: {scheduler_cls}, only support OptimizerParamScheduler currently." + ) - @remote_function(dispatch='all') + @remote_function(dispatch="all") @timed def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type=2, **kwargs): self.step(**kwargs) self.zero_grad(**kwargs) self.lr_step(**kwargs) - @remote_function(dispatch='all', collect='first', sync=True, timeout=3600) + @remote_function(dispatch="all", collect="first", sync=True, timeout=3600) @timed - def save(self, - name: Optional[str] = None, - output_dir: Optional[str] = None, - interval: int = 1, - save_optimizer: bool = False, - merge_lora: bool = False, - **kwargs): + def save( + self, + name: str | None = None, + output_dir: str | None = None, + interval: int = 1, + save_optimizer: bool = False, + merge_lora: bool = False, + **kwargs, + ): """Save model checkpoint. Always saves HF-format model weights. When ``save_optimizer`` is True, @@ -1251,29 +1306,31 @@ def save(self, **kwargs: Additional arguments forwarded to the underlying save methods (e.g. ``adapter_name``, ``safe_serialization``, and ``max_shard_size``). """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) - safe_serialization = kwargs.pop('safe_serialization', True) - max_shard_size = kwargs.pop('max_shard_size', '5GB') - save_total_limit = kwargs.pop('save_total_limit', None) - async_save = kwargs.pop('async_save', False) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) + safe_serialization = kwargs.pop("safe_serialization", True) + max_shard_size = kwargs.pop("max_shard_size", "5GB") + save_total_limit = kwargs.pop("save_total_limit", None) + async_save = kwargs.pop("async_save", False) if async_save and save_total_limit is not None: # The rolling delete cannot tell whether an in-flight async write has finished, so # the two would race on the same directory. - raise ValueError('async_save is incompatible with save_total_limit: the rotation cannot tell ' - 'whether a background write has completed. Disable one of the two.') + raise ValueError( + "async_save is incompatible with save_total_limit: the rotation cannot tell " + "whether a background write has completed. Disable one of the two." + ) if not safe_serialization: - raise NotImplementedError('Megatron HF-format checkpoints require safe_serialization=True.') + raise NotImplementedError("Megatron HF-format checkpoints require safe_serialization=True.") optimizer_config = self.optimizer_group[adapter_name] if optimizer_config.cur_step % interval != 0: return if name is None: - name = f'checkpoint-step-{optimizer_config.cur_step}' + name = f"checkpoint-step-{optimizer_config.cur_step}" if output_dir is None: - output_dir = 'output' + output_dir = "output" checkpoint_dir = os.path.join(output_dir, name) - is_lora = (optimizer_config.adapter_name != _default_adapter_name) + is_lora = optimizer_config.adapter_name != _default_adapter_name if merge_lora and is_lora: self._merge_lora_adapters(optimizer_config.adapter_name) @@ -1294,23 +1351,23 @@ def save(self, **kwargs, ) trainer_state = { - 'checkpoint_version': 1, - 'cur_step': optimizer_config.cur_step, - 'consumed_train_samples': kwargs.get('consumed_train_samples', 0), - 'gradient_accumulation_steps': optimizer_config.gradient_accumulation_steps, - 'component_states': kwargs.get('component_states', {}), + "checkpoint_version": 1, + "cur_step": optimizer_config.cur_step, + "consumed_train_samples": kwargs.get("consumed_train_samples", 0), + "gradient_accumulation_steps": optimizer_config.gradient_accumulation_steps, + "component_states": kwargs.get("component_states", {}), } - state_path = os.path.join(checkpoint_dir, 'trainer_state.json') + state_path = os.path.join(checkpoint_dir, "trainer_state.json") if async_request is not None: # trainer_state.json + the iteration tracker are what a resume reads, so they ride the # async finalize and are written only after every rank's shard write landed. def _finalize_async_checkpoint() -> None: rank = dist.get_rank() if rank == 0: - with open(state_path, 'w') as f: + with open(state_path, "w") as f: json.dump(trainer_state, f, indent=2) - tracker_path = os.path.join(checkpoint_dir, 'latest_checkpointed_iteration.txt') - with open(tracker_path, 'w') as f: + tracker_path = os.path.join(checkpoint_dir, "latest_checkpointed_iteration.txt") + with open(tracker_path, "w") as f: f.write(str(optimizer_config.cur_step)) copy_checkpoint_args(output_dir, checkpoint_dir) @@ -1321,7 +1378,7 @@ def _finalize_async_checkpoint() -> None: return checkpoint_dir rank = dist.get_rank() if dist.is_initialized() else 0 if rank == 0 and save_optimizer: - with open(state_path, 'w') as f: + with open(state_path, "w") as f: json.dump(trainer_state, f, indent=2) # Copy the run metadata and rotate only after every checkpoint component has completed. Rank 0 @@ -1333,9 +1390,9 @@ def _finalize_async_checkpoint() -> None: return checkpoint_dir - @remote_function(dispatch='all') + @remote_function(dispatch="all") @timed - def load(self, name: str, output_dir: Optional[str] = None, **kwargs): + def load(self, name: str, output_dir: str | None = None, **kwargs): """Load model weights, and optionally optimizer / scheduler / RNG state. Args: @@ -1347,19 +1404,19 @@ def load(self, name: str, output_dir: Optional[str] = None, **kwargs): **kwargs: Additional arguments (``adapter_name``, ``no_load_optim``, ``no_load_rng``, etc.). """ - resume = kwargs.pop('load_optimizer', False) + resume = kwargs.pop("load_optimizer", False) if output_dir is not None: checkpoint_dir = os.path.join(output_dir, name) elif os.path.exists(name): checkpoint_dir = name elif not resume: # load from hub - token = kwargs.pop('token', None) + token = kwargs.pop("token", None) checkpoint_dir = HubOperation.download_model(name, token=token) else: - checkpoint_dir = os.path.join('output', name) + checkpoint_dir = os.path.join("output", name) - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) if resume: self._load_mcore_optimizer( @@ -1381,7 +1438,7 @@ def load(self, name: str, output_dir: Optional[str] = None, **kwargs): if dist.is_initialized(): dist.barrier() - @remote_function(dispatch='all') + @remote_function(dispatch="all") def reload_initial_weights(self, **kwargs): """Reload the base model weights from the original model path, used by full-parameter server deployments after a tenant releases the (exclusive) model.""" @@ -1395,57 +1452,60 @@ def reload_initial_weights(self, **kwargs): # leak into the next tenant's first optimizer step. if self._is_model_ddp_wrapped(): for model_chunk in self.model: - if hasattr(model_chunk, 'zero_grad_buffer'): + if hasattr(model_chunk, "zero_grad_buffer"): model_chunk.zero_grad_buffer() 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() # lazy_collect=False: the driver consumes the returned training-state dict immediately # (loop.resume subscripts state['cur_step']); a lazy callable default would break it. - @remote_function(dispatch='all', collect='first', sync=True, lazy_collect=False, timeout=3600) + @remote_function(dispatch="all", collect="first", sync=True, lazy_collect=False, timeout=3600) @timed def resume_from_checkpoint(self, checkpoint_dir, *, resume_only_model=False, **kwargs): - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] self.load(checkpoint_dir, load_optimizer=not resume_only_model, adapter_name=adapter_name, **kwargs) if resume_only_model: return { - 'cur_step': 0, - 'consumed_train_samples': 0, - 'gradient_accumulation_steps': optimizer_config.gradient_accumulation_steps, - 'component_states': {}, + "cur_step": 0, + "consumed_train_samples": 0, + "gradient_accumulation_steps": optimizer_config.gradient_accumulation_steps, + "component_states": {}, } - trainer_state_path = os.path.join(checkpoint_dir, 'trainer_state.json') + trainer_state_path = os.path.join(checkpoint_dir, "trainer_state.json") with open(trainer_state_path) as f: trainer_state = json.load(f) result = { - 'cur_step': trainer_state['cur_step'], - 'consumed_train_samples': trainer_state['consumed_train_samples'], - 'gradient_accumulation_steps': trainer_state['gradient_accumulation_steps'], + "cur_step": trainer_state["cur_step"], + "consumed_train_samples": trainer_state["consumed_train_samples"], + "gradient_accumulation_steps": trainer_state["gradient_accumulation_steps"], } - if 'component_states' in trainer_state: - result['component_states'] = trainer_state['component_states'] + if "component_states" in trainer_state: + result["component_states"] = trainer_state["component_states"] return result @staticmethod - def _get_rng_state() -> 'ShardedObject': + def _get_rng_state() -> "ShardedObject": from megatron.core import parallel_state as mpu from megatron.core import tensor_parallel from megatron.core.dist_checkpointing.mapping import ShardedObject rng_state = { - 'random_rng_state': random.getstate(), - 'np_rng_state': np.random.get_state(), - 'torch_rng_state': torch.get_rng_state(), + "random_rng_state": random.getstate(), + "np_rng_state": np.random.get_state(), + "torch_rng_state": torch.get_rng_state(), # Backend-agnostic device RNG (CUDA / NPU / MPS); key kept as # 'cuda_rng_state' for backward compatibility with existing checkpoints. - 'cuda_rng_state': Platform.get_device_rng_state(), - 'rng_tracker_states': tensor_parallel.get_cuda_rng_tracker().get_states(), + "cuda_rng_state": Platform.get_device_rng_state(), + "rng_tracker_states": tensor_parallel.get_cuda_rng_tracker().get_states(), } rng_state_list = [rng_state] @@ -1455,7 +1515,7 @@ def _get_rng_state() -> 'ShardedObject': tp_size = mpu.get_tensor_model_parallel_world_size() return ShardedObject( - 'rng_state', + "rng_state", rng_state_list, (pp_size, tp_size), (pp_rank, tp_rank), @@ -1468,9 +1528,9 @@ def _generate_state_dict( optimizer=None, opt_param_scheduler=None, rng_state=None, - iteration: Optional[int] = None, - model_sd_kwargs: Optional[dict] = None, - optim_sd_kwargs: Optional[dict] = None, + iteration: int | None = None, + model_sd_kwargs: dict | None = None, + optim_sd_kwargs: dict | None = None, save_optim: bool = True, save_rng: bool = True, ) -> dict: @@ -1478,28 +1538,28 @@ def _generate_state_dict( optim_sd_kwargs = optim_sd_kwargs or {} state_dict: dict = { - 'checkpoint_version': 3.0, + "checkpoint_version": 3.0, } if iteration is not None: - state_dict['iteration'] = iteration + state_dict["iteration"] = iteration # Model sharded state dict for i, m in enumerate(model): - key = 'model' if len(model) == 1 else f'model{i}' + key = "model" if len(model) == 1 else f"model{i}" state_dict[key] = m.sharded_state_dict(**model_sd_kwargs) # Optimizer + scheduler if save_optim and optimizer is not None: - state_dict['optimizer'] = optimizer.sharded_state_dict( + state_dict["optimizer"] = optimizer.sharded_state_dict( state_dict, **optim_sd_kwargs, ) if opt_param_scheduler is not None: - state_dict['opt_param_scheduler'] = opt_param_scheduler.state_dict() + state_dict["opt_param_scheduler"] = opt_param_scheduler.state_dict() # RNG if save_rng and rng_state is not None: - state_dict['rng_state'] = rng_state + state_dict["rng_state"] = rng_state return state_dict @@ -1507,16 +1567,16 @@ def _mcore_async_calls_queue(self): """The mcore-native async write queue, owned per model instance (no megatron global args).""" if self._mcore_async_queue is None: from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue + self._mcore_async_queue = AsyncCallsQueue(persistent=False) return self._mcore_async_queue def _maybe_finalize_mcore_async_saves(self, blocking: bool = False) -> None: """Join completed async shard writes; runs on all ranks (the finalize holds a barrier).""" if self._mcore_async_queue is not None: - self._mcore_async_queue.maybe_finalize_async_calls( - blocking, no_dist=not dist.is_initialized()) + self._mcore_async_queue.maybe_finalize_async_calls(blocking, no_dist=not dist.is_initialized()) - @remote_function(dispatch='all', collect='none', lazy_collect=False) + @remote_function(dispatch="all", collect="none", lazy_collect=False) def wait_async_saves(self) -> None: """Join every in-flight save on both async channels (thread-pool tail + mcore queue).""" super().wait_async_saves() @@ -1525,7 +1585,7 @@ def wait_async_saves(self) -> None: def _save_mcore_optimizer( self, checkpoint_dir: str, - optimizer_config: 'MegatronOptimizerGroup', + optimizer_config: "MegatronOptimizerGroup", **kwargs, ): """Write the optimizer/scheduler/RNG state in mcore dist-ckpt format. @@ -1539,14 +1599,14 @@ def _save_mcore_optimizer( from megatron.core.dist_checkpointing.strategies.fully_parallel import FullyParallelSaveStrategyWrapper iteration = optimizer_config.cur_step - iter_dir = os.path.join(checkpoint_dir, f'iter_{iteration:07d}') + iter_dir = os.path.join(checkpoint_dir, f"iter_{iteration:07d}") os.makedirs(iter_dir, exist_ok=True) # The sharding type must match the data-parallel wrapper the strategy chose at wrap time # (Megatron-FSDP stores optimizer state as DTensors, plain DDP does not). sharded_sd_metadata = self.strategy.get_sharded_sd_metadata() - no_save_rng = kwargs.pop('no_save_rng', False) + no_save_rng = kwargs.pop("no_save_rng", False) rng_state = None if no_save_rng else self._get_rng_state() model = self.model @@ -1556,18 +1616,18 @@ def _save_mcore_optimizer( opt_param_scheduler=optimizer_config.lr_scheduler, rng_state=rng_state, iteration=iteration, - model_sd_kwargs={'metadata': sharded_sd_metadata}, - optim_sd_kwargs={'metadata': sharded_sd_metadata}, + model_sd_kwargs={"metadata": sharded_sd_metadata}, + optim_sd_kwargs={"metadata": sharded_sd_metadata}, ) # thread_count is read at save time, so setting it post-construction is safe; # non-torch_dist backends do not expose it and are skipped. - thread_count = kwargs.pop('thread_count', 2) + thread_count = kwargs.pop("thread_count", 2) # megatron-core < 0.19 exposed get_default_save_sharded_strategy(); 0.19 made # save(sharded_strategy=None) default to TorchDistSaveShardedStrategy itself -- build that # explicitly, the thread_count tweak and the wrapper below need a concrete instance. save_strategy = TorchDistSaveShardedStrategy() - if hasattr(save_strategy, 'thread_count'): + if hasattr(save_strategy, "thread_count"): save_strategy.thread_count = thread_count if mpu.get_data_parallel_world_size(with_context_parallel=True) > 1: save_strategy = FullyParallelSaveStrategyWrapper( @@ -1577,7 +1637,7 @@ def _save_mcore_optimizer( # Async staging is only sound under an initialized process group (the metadata finalize # holds a barrier); a single-process run has no skew to hide anyway. - async_save = kwargs.pop('async_save', False) and dist.is_initialized() + async_save = kwargs.pop("async_save", False) and dist.is_initialized() async_request = dist_checkpointing.save( state_dict, iter_dir, @@ -1599,51 +1659,53 @@ def _save_mcore_optimizer( if rank == 0: tracker_path = os.path.join( checkpoint_dir, - 'latest_checkpointed_iteration.txt', + "latest_checkpointed_iteration.txt", ) - with open(tracker_path, 'w') as f: + with open(tracker_path, "w") as f: f.write(str(iteration)) - logger.info(f'Saved mcore optimizer state at iteration {iteration} ' - f'to {checkpoint_dir}') + logger.info(f"Saved mcore optimizer state at iteration {iteration} to {checkpoint_dir}") def _load_mcore_optimizer( self, checkpoint_dir: str, - adapter_name: str = '', + adapter_name: str = "", **kwargs, ): - from megatron.core import dist_checkpointing + from megatron.core import dist_checkpointing, tensor_parallel from megatron.core import parallel_state as mpu - from megatron.core import tensor_parallel from megatron.core.dist_checkpointing.serialization import get_default_load_sharded_strategy from megatron.core.dist_checkpointing.strategies.fully_parallel import FullyParallelLoadStrategyWrapper - no_load_optim = kwargs.pop('no_load_optim', False) - no_load_rng = kwargs.pop('no_load_rng', False) - data_parallel_random_init = kwargs.pop('data_parallel_random_init', False) + no_load_optim = kwargs.pop("no_load_optim", False) + no_load_rng = kwargs.pop("no_load_rng", False) + data_parallel_random_init = kwargs.pop("data_parallel_random_init", False) - optimizer_config = self.optimizer_group.get(adapter_name or self._get_default_group(), ) + optimizer_config = self.optimizer_group.get( + adapter_name or self._get_default_group(), + ) # Read iteration from tracker file. tracker_path = os.path.join( checkpoint_dir, - 'latest_checkpointed_iteration.txt', + "latest_checkpointed_iteration.txt", ) iteration = self._read_iteration(tracker_path) if iteration == 0: - logger.warning(f'No checkpoint found in {checkpoint_dir}') + logger.warning(f"No checkpoint found in {checkpoint_dir}") return - iter_dir = os.path.join(checkpoint_dir, f'iter_{iteration:07d}') + iter_dir = os.path.join(checkpoint_dir, f"iter_{iteration:07d}") # Load common (non-sharded) state to inspect content metadata. common_state = dist_checkpointing.load_common_state_dict(iter_dir) - sharded_sd_metadata = dist_checkpointing.load_content_metadata(preloaded_state_dict=common_state, ) + sharded_sd_metadata = dist_checkpointing.load_content_metadata( + preloaded_state_dict=common_state, + ) # Build optimizer / scheduler references for the sharded state dict. optimizer = optimizer_config.optimizer if not no_load_optim else None - opt_param_scheduler = (optimizer_config.lr_scheduler if not no_load_optim else None) + opt_param_scheduler = optimizer_config.lr_scheduler if not no_load_optim else None rng_state = self._get_rng_state() if not no_load_rng else None optim_sd_kwargs = dict(metadata=sharded_sd_metadata, is_loading=True) @@ -1674,45 +1736,49 @@ def _load_mcore_optimizer( # Restore model weights. if len(self.model) == 1: - self.model[0].load_state_dict(state_dict['model'], strict=False) + self.model[0].load_state_dict(state_dict["model"], strict=False) else: for i, m in enumerate(self.model): - key = f'model{i}' + key = f"model{i}" if key in state_dict: m.load_state_dict(state_dict[key], strict=False) # Restore optimizer + LR scheduler. - if not no_load_optim and optimizer is not None and 'optimizer' in state_dict: + if not no_load_optim and optimizer is not None and "optimizer" in state_dict: with torch.no_grad(): - optimizer.load_state_dict(state_dict['optimizer']) - if (opt_param_scheduler is not None and 'opt_param_scheduler' in state_dict): - opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler'], ) + optimizer.load_state_dict(state_dict["optimizer"]) + if opt_param_scheduler is not None and "opt_param_scheduler" in state_dict: + opt_param_scheduler.load_state_dict( + state_dict["opt_param_scheduler"], + ) - if not no_load_rng and 'rng_state' in state_dict: - rng = state_dict['rng_state'] + if not no_load_rng and "rng_state" in state_dict: + rng = state_dict["rng_state"] if isinstance(rng, (list, tuple)): rng_index = mpu.get_data_parallel_rank() if data_parallel_random_init else 0 if rng_index >= len(rng): raise ValueError( - f'Checkpoint contains {len(rng)} RNG states, but data-parallel rank {rng_index} was requested.') + f"Checkpoint contains {len(rng)} RNG states, but data-parallel rank {rng_index} was requested." + ) rng = rng[rng_index] - random.setstate(rng['random_rng_state']) - np.random.set_state(rng['np_rng_state']) - torch.set_rng_state(rng['torch_rng_state']) + random.setstate(rng["random_rng_state"]) + np.random.set_state(rng["np_rng_state"]) + torch.set_rng_state(rng["torch_rng_state"]) # Backend-agnostic restore: tolerates ckpt produced on different backend # (returns None) and avoids hard-coded torch.cuda which crashes on NPU. - Platform.set_device_rng_state(rng.get('cuda_rng_state')) - tensor_parallel.get_cuda_rng_tracker().set_states(rng['rng_tracker_states'], ) + Platform.set_device_rng_state(rng.get("cuda_rng_state")) + tensor_parallel.get_cuda_rng_tracker().set_states( + rng["rng_tracker_states"], + ) # Restore iteration counter. - if optimizer_config is not None and 'iteration' in state_dict: - optimizer_config.cur_step = state_dict['iteration'] + if optimizer_config is not None and "iteration" in state_dict: + optimizer_config.cur_step = state_dict["iteration"] if dist.is_initialized(): dist.barrier() - logger.info(f'Resumed from mcore checkpoint at iteration {iteration} ' - f'from {checkpoint_dir}') + logger.info(f"Resumed from mcore checkpoint at iteration {iteration} from {checkpoint_dir}") @staticmethod def _read_iteration(tracker_path: str) -> int: @@ -1731,9 +1797,10 @@ def _read_iteration(tracker_path: str) -> int: iteration = int(iters_dev[0].item()) return iteration - def _merge_lora_adapters(self, adapter_name: str = 'default'): + def _merge_lora_adapters(self, adapter_name: str = "default"): """Merge LoRA adapters into base model weights.""" from mcore_bridge import LoraParallelLinear + with torch.no_grad(): for model in self.strategy.unwrap_model(self.model): for module in model.modules(): @@ -1743,6 +1810,7 @@ def _merge_lora_adapters(self, adapter_name: str = 'default'): def _unmerge_lora_adapters(self): """Unmerge LoRA adapters to restore training state.""" from mcore_bridge import LoraParallelLinear + with torch.no_grad(): for model in self.strategy.unwrap_model(self.model): for module in model.modules(): @@ -1754,28 +1822,25 @@ def _export_hf_config(self): mcore model config, and drop a stale ``quantization_config`` (exported safetensors are plain bf16/fp16 weights). The live ``self.hf_config`` stays untouched for training.""" from copy import deepcopy + hf_config = deepcopy(self.hf_config) - mtp_num_layers = getattr(self.strategy.config, 'mtp_num_layers', None) + mtp_num_layers = getattr(self.strategy.config, "mtp_num_layers", None) if mtp_num_layers: text_config = hf_config.get_text_config() - for key in ('num_nextn_predict_layers', 'mtp_num_hidden_layers'): + for key in ("num_nextn_predict_layers", "mtp_num_hidden_layers"): if hasattr(text_config, key): setattr(text_config, key, mtp_num_layers) break else: text_config.num_nextn_predict_layers = mtp_num_layers - if getattr(hf_config, 'quantization_config', None) is not None: + if getattr(hf_config, "quantization_config", None) is not None: try: - delattr(hf_config, 'quantization_config') + delattr(hf_config, "quantization_config") except AttributeError: hf_config.quantization_config = None return hf_config - def _save_hf_format(self, - output_dir: str, - adapter_name: str, - lora_converter=None, - max_shard_size: str = '5GB'): + def _save_hf_format(self, output_dir: str, adapter_name: str, lora_converter=None, max_shard_size: str = "5GB"): """Save in HuggingFace format using bridge adapter. For distributed training: @@ -1787,7 +1852,7 @@ def _save_hf_format(self, - Saves in PEFT format (adapter_model.safetensors + adapter_config.json) """ # Check if this is LoRA training - is_peft_format = (adapter_name != _default_adapter_name) + is_peft_format = adapter_name != _default_adapter_name is_global_zero = (not dist.is_initialized()) or dist.get_rank() == 0 if is_global_zero: @@ -1801,19 +1866,19 @@ def _save_hf_format(self, model = self.strategy.unwrap_model(self.model) save_weights = self.strategy.bridge.save_weights save_kwargs = { - 'peft_format': is_peft_format, - 'adapter_name': adapter_name, - 'converter': lora_converter, - 'max_shard_size': max_shard_size, + "peft_format": is_peft_format, + "adapter_name": adapter_name, + "converter": lora_converter, + "max_shard_size": max_shard_size, } # Newer mcore-bridge releases can preserve source-checkpoint tensors for modules Megatron # does not materialize; adapter checkpoints contain deltas only and ignore the flag. try: - supports_missing_weights = 'save_missing_weights' in inspect.signature(save_weights).parameters + supports_missing_weights = "save_missing_weights" in inspect.signature(save_weights).parameters except (TypeError, ValueError): supports_missing_weights = False if supports_missing_weights: - save_kwargs['save_missing_weights'] = True + save_kwargs["save_missing_weights"] = True save_weights(model, output_dir, **save_kwargs) # Save config on global rank 0 only (avoid concurrent writers). @@ -1825,7 +1890,7 @@ def _save_hf_format(self, if is_peft_format and isinstance(model[0], PeftModel): config = model[0].peft_config[adapter_name] target_modules = config.target_modules - config.target_modules = 'all-linear' + config.target_modules = "all-linear" model[0].peft_config[adapter_name].save_pretrained(output_dir) config.target_modules = target_modules @@ -1849,7 +1914,7 @@ def _save_megatron_format(self, output_dir: str, adapter_name: str, lora_convert # Save with rank info for distributed checkpointing rank = dist.get_rank() if dist.is_initialized() else 0 - checkpoint_path = os.path.join(output_dir, f'model_rank{rank}.pt') + checkpoint_path = os.path.join(output_dir, f"model_rank{rank}.pt") torch.save(cpu_state_dict, checkpoint_path) # Save shared config on global rank 0 only (avoid concurrent writers). model = self.strategy.unwrap_model(self.model) @@ -1865,10 +1930,11 @@ def _save_megatron_format(self, output_dir: str, adapter_name: str, lora_convert def _save_tokenizer(self, output_dir: str, **kwargs): from twinkle.utils import is_last_rank + if not is_last_rank(): return - adapter_name = kwargs.pop('adapter_name', _default_adapter_name) + adapter_name = kwargs.pop("adapter_name", _default_adapter_name) optimizer_config = self.optimizer_group[adapter_name] template_ins = optimizer_config.template if template_ins is not None: @@ -1876,7 +1942,7 @@ def _save_tokenizer(self, output_dir: str, **kwargs): else: self._default_tokenizer.save_pretrained(output_dir) - @remote_function(execute='first') + @remote_function(execute="first") def get_state_dict(self, **kwargs): """Get trainable state dict. @@ -1886,10 +1952,10 @@ def get_state_dict(self, **kwargs): Returns: State dict of trainable parameters. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) return self._get_trainable_parameters(adapter_name) - def get_hf_state_dict(self, adapter_name: str = '') -> Generator[Tuple[str, torch.Tensor], None, None]: + def get_hf_state_dict(self, adapter_name: str = "") -> Generator[tuple[str, torch.Tensor], None, None]: """Yield model weights in HuggingFace format, one by one (generator to avoid OOM). Preferred for weight synchronization to vLLM: converts Megatron->HF on the fly, avoids @@ -1902,19 +1968,19 @@ def get_hf_state_dict(self, adapter_name: str = '') -> Generator[Tuple[str, torc only_master_rank=False, # All ranks participate in weight sync peft_format=bool(adapter_name), adapter_name=adapter_name if adapter_name else None, - tqdm_desc='Weight sync: ', + tqdm_desc="Weight sync: ", ) - def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str, Dict[str, Any]], **kwargs): - assert adapter_name, 'Use a non-empty adapter_name' + def _patch_adapter(self, adapter_name: str, config_or_dir: PeftConfig | str | dict[str, Any], **kwargs): + assert adapter_name, "Use a non-empty adapter_name" # freeze_* drive the multimodal 'all-linear' target expansion in get_target_modules; popped so they # do not leak into get_peft_model. Defaults match dev's TunerConfig defaults. - freeze_llm = kwargs.pop('freeze_llm', False) - freeze_vit = kwargs.pop('freeze_vit', True) - freeze_aligner = kwargs.pop('freeze_aligner', True) + freeze_llm = kwargs.pop("freeze_llm", False) + freeze_vit = kwargs.pop("freeze_vit", True) + freeze_aligner = kwargs.pop("freeze_aligner", True) model = self.strategy.unwrap_model(self.model) - original_target_modules: List[str] = [] - original_modules_to_save: List[str] = [] + original_target_modules: list[str] = [] + original_modules_to_save: list[str] = [] if isinstance(config_or_dir, str): config_or_dir = HubOperation.download_model(config_or_dir) else: @@ -1934,7 +2000,8 @@ def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str for _model in model: if isinstance(config_or_dir, str): _model = PeftModel.from_pretrained( - _model, config_or_dir, adapter_name=adapter_name, is_trainable=kwargs.get('is_trainable', True)) + _model, config_or_dir, adapter_name=adapter_name, is_trainable=kwargs.get("is_trainable", True) + ) config = _model.peft_config else: config = config_or_dir @@ -1946,7 +2013,8 @@ def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str original_target_modules, freeze_llm=freeze_llm, freeze_vit=freeze_vit, - freeze_aligner=freeze_aligner) + freeze_aligner=freeze_aligner, + ) # Expand the modules_to_save shorthands per chunk too ('all-embedding'; 'all-norm' is # rejected inside get_modules_to_save -- mcore fuses norm into its linear layers). if original_modules_to_save: @@ -1960,16 +2028,16 @@ def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str self.optimizer_group[adapter_name] = self._construct_default_optimizer_group() self.optimizer_group[adapter_name].adapter_name = adapter_name self.optimizer_group[adapter_name].adapter_config = config # noqa - self.optimizer_group[adapter_name].gradient_accumulation_steps = kwargs.get('gradient_accumulation_steps', 1) + self.optimizer_group[adapter_name].gradient_accumulation_steps = kwargs.get("gradient_accumulation_steps", 1) # Fix: use .processor instead of .tokenizer - Template class uses self.processor self._default_tokenizer = self.optimizer_group[adapter_name].template.processor self.active_group = adapter_name - @remote_function(dispatch='all', sync=True, timeout=3600) + @remote_function(dispatch="all", sync=True, timeout=3600) def add_adapter_to_model( self, adapter_name: str, - config_or_dir: Union[Dict[str, Any], LoraConfig, str], + config_or_dir: dict[str, Any] | LoraConfig | str, **kwargs, ): """Add LoRA adapter to model. @@ -1982,44 +2050,46 @@ def add_adapter_to_model( self._patch_adapter(adapter_name, config_or_dir, **kwargs) @remote_function() - def apply_patch(self, patch_cls: Union[Patch, Type[Patch], str], **kwargs): + def apply_patch(self, patch_cls: Patch | type[Patch] | str, **kwargs): apply_patch(self.model, patch_cls, **kwargs) - @remote_function(dispatch='all') - def set_template(self, template_cls: Union[Template, Type[Template], str], **kwargs): + @remote_function(dispatch="all") + def set_template(self, template_cls: Template | type[Template] | str, **kwargs): """Set template for input encoding. Args: template_cls: Template class or string name. **kwargs: Additional arguments. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - if 'template_spec' not in kwargs: - kwargs['model_id'] = self.tokenizer_id + if "template_spec" not in kwargs: + kwargs["model_id"] = self.tokenizer_id optimizer_config.template = construct_template(template_cls, **kwargs) - @remote_function(dispatch='all') - def set_processor(self, processor_cls: Union[InputProcessor, Type[InputProcessor], str, Callable], **kwargs): + @remote_function(dispatch="all") + def set_processor(self, processor_cls: InputProcessor | type[InputProcessor] | str | Callable, **kwargs): """Set input processor. Args: processor_cls: Processor class or string name. **kwargs: Additional arguments. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - kwargs['framework'] = 'megatron' + kwargs["framework"] = "megatron" # processor/base.py: self.device_mesh.cp_world_size - kwargs['device_mesh'] = kwargs.get('device_mesh', self.device_mesh) + kwargs["device_mesh"] = kwargs.get("device_mesh", self.device_mesh) processor = construct_class(processor_cls, InputProcessor, twinkle.processor, **kwargs) if processor.padding_free and not self.variable_seq_lengths: - raise ValueError('padding_free=True requires variable_seq_lengths=True in MegatronModel. ' - 'Padding-free packing merges sequences into batch=1, making fixed-length ' - 'microbatch slicing impossible.') + raise ValueError( + "padding_free=True requires variable_seq_lengths=True in MegatronModel. " + "Padding-free packing merges sequences into batch=1, making fixed-length " + "microbatch slicing impossible." + ) optimizer_config.processor = processor - @remote_function(execute='first', lazy_collect=False) + @remote_function(execute="first", lazy_collect=False) def get_train_configs(self, **kwargs): """Get training configuration summary. @@ -2029,29 +2099,29 @@ def get_train_configs(self, **kwargs): Returns: Configuration summary string. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - expr = 'Backend: Megatron-Core\n' - expr += f'DP size: {self.device_mesh.dp_world_size}\n' - expr += f'TP size: {self.device_mesh.tp_world_size}\n' - expr += f' - VPP size: {self.device_mesh.vpp_size}\n' - expr += f'PP size: {self.device_mesh.pp_world_size}\n' - expr += f'CP size: {self.device_mesh.cp_world_size}\n' - expr += f'EP size: {self.device_mesh.ep_size}\n' - expr += f'Sequence Parallel: {self.strategy.sequence_parallel}\n' + expr = "Backend: Megatron-Core\n" + expr += f"DP size: {self.device_mesh.dp_world_size}\n" + expr += f"TP size: {self.device_mesh.tp_world_size}\n" + expr += f" - VPP size: {self.device_mesh.vpp_size}\n" + expr += f"PP size: {self.device_mesh.pp_world_size}\n" + expr += f"CP size: {self.device_mesh.cp_world_size}\n" + expr += f"EP size: {self.device_mesh.ep_size}\n" + expr += f"Sequence Parallel: {self.strategy.sequence_parallel}\n" if optimizer_config.adapter_config is not None: config = optimizer_config.adapter_config.__dict__ config = {key: str(value) for key, value in config.items() if value is not None} - expr += f'Adapter config:\n{json.dumps(config, indent=2, ensure_ascii=False)}\n' + expr += f"Adapter config:\n{json.dumps(config, indent=2, ensure_ascii=False)}\n" if optimizer_config.optimizer: - expr += f'Optimizer: {optimizer_config.optimizer.__class__.__name__}\n' - expr += f'Learning rate: {optimizer_config.optimizer.chained_optimizers[0].config.lr}\n' + expr += f"Optimizer: {optimizer_config.optimizer.__class__.__name__}\n" + expr += f"Learning rate: {optimizer_config.optimizer.chained_optimizers[0].config.lr}\n" if optimizer_config.lr_scheduler: - expr += f'LR scheduler: {optimizer_config.lr_scheduler.__class__.__name__}\n' - expr += f'Gradient accumulation steps: {optimizer_config.gradient_accumulation_steps}\n' + expr += f"LR scheduler: {optimizer_config.lr_scheduler.__class__.__name__}\n" + expr += f"Gradient accumulation steps: {optimizer_config.gradient_accumulation_steps}\n" return expr @@ -2065,7 +2135,7 @@ def _get_weight_generator( adapter_name: str = None, base_sync_done: bool = False, merge_and_sync: bool = False, - model_keys: List[str] = None, + model_keys: list[str] = None, **kwargs, ): if adapter_name is None: @@ -2088,7 +2158,7 @@ def merge_lora(): def weight_generator(): with merge_lora(): names = [] - for name, tensor in self.get_hf_state_dict(adapter_name=''): + for name, tensor in self.get_hf_state_dict(adapter_name=""): if name is None or tensor is None: continue # Skip LoRA-specific weights for base model sync @@ -2106,7 +2176,7 @@ def weight_generator(): for name, tensor in self.get_hf_state_dict(adapter_name=adapter_name): if name is None or tensor is None: continue - if 'lora' not in name: + if "lora" not in name: continue name = normalize_param_name(name, keep_base_layer=True) names.append(name) @@ -2117,7 +2187,7 @@ def weight_generator(): # First full base-model sync. def _raw_weights(with_base_layer=False): names = [] - for name, tensor in self.get_hf_state_dict(adapter_name=''): + for name, tensor in self.get_hf_state_dict(adapter_name=""): if name is None or tensor is None: continue # Skip LoRA-specific weights for base model sync @@ -2156,7 +2226,7 @@ def _get_local_weight_generator(self): models = self.strategy.unwrap_model(self.model) for chunk_idx, chunk in enumerate(models): for name, param in chunk.named_parameters(): - yield f'chunk{chunk_idx}.{name}', param.data + yield f"chunk{chunk_idx}.{name}", param.data def _get_local_state_dict(self): """Return {name: param_data} for this rank's local Megatron shard.""" @@ -2164,11 +2234,11 @@ def _get_local_state_dict(self): state = {} for chunk_idx, chunk in enumerate(models): for name, param in chunk.named_parameters(): - state[f'chunk{chunk_idx}.{name}'] = param.data + state[f"chunk{chunk_idx}.{name}"] = param.data return state - @remote_function(collect='first', lazy_collect=False) - def get_peft_config_dict(self, adapter_name: str = None) -> Optional[Dict[str, Any]]: + @remote_function(collect="first", lazy_collect=False) + def get_peft_config_dict(self, adapter_name: str = None) -> dict[str, Any] | None: """Return the PEFT config as a dict for vLLM's PEFTHelper. Used by CheckpointEngineManager for LoRA-only weight sync. @@ -2185,18 +2255,20 @@ def get_peft_config_dict(self, adapter_name: str = None) -> Optional[Dict[str, A if isinstance(config, dict): config = config.get(adapter_name, next(iter(config.values()))) target_modules = config.target_modules - config.target_modules = 'all-linear' - _peft_config = config.to_dict() if hasattr(config, 'to_dict') else dict(config) - _peft_config['target_modules'] = target_modules + config.target_modules = "all-linear" + _peft_config = config.to_dict() if hasattr(config, "to_dict") else dict(config) + _peft_config["target_modules"] = target_modules return _peft_config @staticmethod - def get_target_modules(model: 'torch.nn.Module', - target_modules: List[str], - *, - freeze_llm: bool = False, - freeze_vit: bool = True, - freeze_aligner: bool = True) -> List[str]: + def get_target_modules( + model: "torch.nn.Module", + target_modules: list[str], + *, + freeze_llm: bool = False, + freeze_vit: bool = True, + freeze_aligner: bool = True, + ) -> list[str]: """Expand the ``all-*`` keywords into concrete module names, honouring the freeze_* partition. On a multimodal model an ``all-*`` keyword must respect freeze_llm / freeze_vit / @@ -2205,12 +2277,15 @@ def get_target_modules(model: 'torch.nn.Module', """ import torch - def find_all_linears(model: torch.nn.Module) -> List[str]: - from megatron.core.extensions.transformer_engine import (TEGroupedLinear, TELayerNormColumnParallelLinear, - TELinear) + def find_all_linears(model: torch.nn.Module) -> list[str]: + from megatron.core.extensions.transformer_engine import ( + TEGroupedLinear, + TELayerNormColumnParallelLinear, + TELinear, + ) def _cond(name: str, module: torch.nn.Module) -> bool: - if name == 'output_layer' or 'lora' in name: + if name == "output_layer" or "lora" in name: return False if isinstance(module, (TELinear, TELayerNormColumnParallelLinear, TEGroupedLinear, torch.nn.Linear)): return True @@ -2218,34 +2293,40 @@ def _cond(name: str, module: torch.nn.Module) -> bool: return _find_layers(model, _cond) - def find_router(model: torch.nn.Module) -> List[str]: + def find_router(model: torch.nn.Module) -> list[str]: from megatron.core.transformer.moe.router import TopKRouter - return _find_layers(model, lambda name, module: isinstance(module, TopKRouter) and 'lora' not in name) - keywords = ('all-linear', 'all-embedding', 'all-router') + return _find_layers(model, lambda name, module: isinstance(module, TopKRouter) and "lora" not in name) + + keywords = ("all-linear", "all-embedding", "all-router") # Explicit names are kept verbatim; only the all-* keywords are expanded (and, on a # multimodal model, filtered by tower). explicit = [t for t in target_modules if t not in keywords] - expanded: List[str] = [] - if 'all-linear' in target_modules: + expanded: list[str] = [] + if "all-linear" in target_modules: expanded += find_all_linears(model) - if 'all-embedding' in target_modules: + if "all-embedding" in target_modules: expanded += _find_embedding(model) - if 'all-router' in target_modules: + if "all-router" in target_modules: expanded += find_router(model) from twinkle.module.optimizer import resolve_megatron_tower_prefixes + tower_prefixes = resolve_megatron_tower_prefixes([model]) - is_multimodal = bool(tower_prefixes['vision_tower'] or tower_prefixes['aligner']) + is_multimodal = bool(tower_prefixes["vision_tower"] or tower_prefixes["aligner"]) if is_multimodal and any(keyword in target_modules for keyword in keywords): if freeze_llm and freeze_vit and freeze_aligner: - raise ValueError('freeze_llm / freeze_vit / freeze_aligner are all True, so a multimodal ' - 'LoRA has no tower left to adapt. Unfreeze at least one of them.') - frozen = {'language_model': freeze_llm, 'vision_tower': freeze_vit, 'aligner': freeze_aligner} + raise ValueError( + "freeze_llm / freeze_vit / freeze_aligner are all True, so a multimodal " + "LoRA has no tower left to adapt. Unfreeze at least one of them." + ) + frozen = {"language_model": freeze_llm, "vision_tower": freeze_vit, "aligner": freeze_aligner} # Longest prefix first, so the aligner (nested under the vision tower) is classified # as the aligner rather than the tower; a module in no declared tower is dropped. - ranked = sorted(((prefix, role) for role, prefixes in tower_prefixes.items() for prefix in prefixes), - key=lambda item: -len(item[0])) + ranked = sorted( + ((prefix, role) for role, prefixes in tower_prefixes.items() for prefix in prefixes), + key=lambda item: -len(item[0]), + ) kept = [] for name in expanded: role = next((r for prefix, r in ranked if name.startswith(prefix)), None) @@ -2255,20 +2336,21 @@ def find_router(model: torch.nn.Module) -> List[str]: return list(set(explicit + expanded)) @staticmethod - def get_modules_to_save(model: 'torch.nn.Module', modules_to_save: List[str]) -> List[str]: + def get_modules_to_save(model: "torch.nn.Module", modules_to_save: list[str]) -> list[str]: """Expand the ``modules_to_save`` shorthands into concrete mcore module names. Only ``'all-embedding'`` is expanded; ``'all-norm'`` is rejected (mcore fuses the normalization into its linear layers, so there is no standalone norm module). """ - if 'all-norm' in modules_to_save: + if "all-norm" in modules_to_save: raise ValueError( "modules_to_save 'all-norm' is not supported on the megatron backend: mcore fuses the " - 'normalization into its linear layers, so there is no standalone norm module to save. ' - 'List the concrete module names to save instead.') - if 'all-embedding' not in modules_to_save: + "normalization into its linear layers, so there is no standalone norm module to save. " + "List the concrete module names to save instead." + ) + if "all-embedding" not in modules_to_save: return list(modules_to_save) - explicit = [m for m in modules_to_save if m != 'all-embedding'] + explicit = [m for m in modules_to_save if m != "all-embedding"] found = _find_embedding(model) if not found: logger.warning("modules_to_save 'all-embedding' matched no mcore embedding module; ignored.") diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 0aa8e62fe..88f829a1c 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -1,46 +1,46 @@ import re -import torch +from collections.abc import Callable from contextlib import contextmanager, nullcontext from copy import deepcopy from dataclasses import dataclass, field +from types import MethodType +from typing import Any, Optional, Union + +import torch from peft import LoraConfig, PeftModel, get_peft_model from peft.tuners.lora import Embedding, Linear, LoraLayer from torch.distributed.tensor import distribute_tensor -from types import MethodType -from typing import Any, Callable, Dict, List, Optional, Union 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() @dataclass class LoraTenant: - index: int adapter_name: str config: LoraConfig - tenant_adapter_name: Optional[str] = None - tenant_config: Optional[LoraConfig] = None - lora_A_weights: Dict[str, torch.Tensor] = field(default_factory=lambda: {}) + tenant_adapter_name: str | None = None + tenant_config: LoraConfig | None = None + lora_A_weights: dict[str, torch.Tensor] = field(default_factory=lambda: {}) class MultiLora: - def __init__(self, max_loras=5, max_r=32, max_length: int = 8192): self.max_loras = max_loras self.max_r = max_r - self.loras: List[LoraTenant] = [] + self.loras: list[LoraTenant] = [] self.module: PeftModel self._active_adapters = [] self.max_length = max_length self.target_parameter_manager = TargetParameterLoraManager(max_loras=max_loras, max_r=max_r) - self.lora_layer_names: List[str] = [] + self.lora_layer_names: list[str] = [] - def _get_available_lora(self) -> Optional[LoraTenant]: + def _get_available_lora(self) -> LoraTenant | None: for _lora in self.loras: if _lora.tenant_adapter_name is None: return _lora @@ -51,11 +51,11 @@ def _read_param_tensor(self, parameter): @staticmethod def _is_distributed_param(parameter): - return hasattr(parameter, 'device_mesh') and hasattr(parameter, 'placements') + return hasattr(parameter, "device_mesh") and hasattr(parameter, "placements") @staticmethod def _is_target_parameter_lora_name(name: str) -> bool: - return '._twinkle_lora_' in name + return "._twinkle_lora_" in name def _write_param_tensor(self, parameter, value): if value is None: @@ -66,7 +66,7 @@ def _write_param_tensor(self, parameter, value): parameter.data.copy_(value.to(parameter.device)) return - local_param = parameter.to_local() if hasattr(parameter, 'to_local') else None + local_param = parameter.to_local() if hasattr(parameter, "to_local") else None parameter_shape = tuple(parameter.shape) value_shape = tuple(value.shape) if local_param is not None: @@ -76,8 +76,9 @@ def _write_param_tensor(self, parameter, value): return if value_shape != parameter_shape: raise ValueError( - f'Cannot write tensor with shape {value_shape} to distributed parameter with global shape ' - f'{parameter_shape} and local shape {local_shape}') + f"Cannot write tensor with shape {value_shape} to distributed parameter with global shape " + f"{parameter_shape} and local shape {local_shape}" + ) value = distribute_tensor(value.to(parameter.device), parameter.device_mesh, parameter.placements) else: value = value.to(parameter.device) @@ -87,13 +88,13 @@ def _write_param_tensor(self, parameter, value): def _slice_rank_tensor(name: str, tensor, rank: int): if tensor is None: return None - if 'embedding_A' in name: + if "embedding_A" in name: return tensor[:, :rank] - if 'embedding_B' in name: + if "embedding_B" in name: return tensor[:rank, :] - if '_A' in name: + if "_A" in name: return tensor[:rank, :] - if '_B' in name: + if "_B" in name: return tensor[:, :rank] return tensor @@ -101,14 +102,14 @@ def _slice_rank_tensor(name: str, tensor, rank: int): def _copy_rank_tensor(name: str, target, value): if target is None or value is None: return None - if 'embedding_A' in name: - target[:, :value.shape[1]].copy_(value) - elif 'embedding_B' in name: - target[:value.shape[0], :].copy_(value) - elif '_A' in name: - target[:value.shape[0], :].copy_(value) - elif '_B' in name: - target[:, :value.shape[1]].copy_(value) + if "embedding_A" in name: + target[:, : value.shape[1]].copy_(value) + elif "embedding_B" in name: + target[: value.shape[0], :].copy_(value) + elif "_A" in name: + target[: value.shape[0], :].copy_(value) + elif "_B" in name: + target[:, : value.shape[1]].copy_(value) else: target.copy_(value) return target @@ -120,13 +121,61 @@ def reset_adapter_status(self): """Force lora_0 require_grad, disable others""" if isinstance(self.module, list): for _module in self.module: - _module.set_adapter('lora_0') + _module.set_adapter("lora_0") else: - self.module.set_adapter('lora_0') + 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, strict=False): + 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') + raise ValueError(f"Adapter {tenant_adapter_name} does not exist") adapter_name = self.find_lora_by_tenant(tenant_adapter_name).adapter_name if isinstance(self.module, list): for _module in self.module: @@ -161,8 +210,8 @@ def _autocast_adapter_dtype(peft_model, adapter_name: str) -> None: it explicitly keeps every preallocated slot on the same dtype PEFT would produce, instead of forcing fp32 ourselves. It is a no-op for an fp32 base. """ - base = getattr(peft_model, 'base_model', None) - if base is not None and hasattr(base, '_cast_adapter_dtype'): + base = getattr(peft_model, "base_model", None) + if base is not None and hasattr(base, "_cast_adapter_dtype"): base._cast_adapter_dtype(adapter_name=adapter_name, autocast_adapter_dtype=True) @contextmanager @@ -221,44 +270,44 @@ def _after(_module): def check_length( self, - inputs: Union[InputFeature, List[InputFeature]], + inputs: InputFeature | list[InputFeature], ): if isinstance(inputs, dict): inputs = [inputs] for index, item in enumerate(inputs): - if 'input_ids' not in item: + if "input_ids" not in item: continue - length = len(item['input_ids']) + length = len(item["input_ids"]) if length > self.max_length: - raise ValueError(f'Input length {length} exceeds max_length {self.max_length} at sample {index}') + raise ValueError(f"Input length {length} exceeds max_length {self.max_length} at sample {index}") def acquire_lora(self, tenant_adapter_name: str, config: LoraConfig) -> str: if self.has_lora(tenant_adapter_name): - raise ValueError(f'Lora {tenant_adapter_name} already exists') + raise ValueError(f"Lora {tenant_adapter_name} already exists") _available_lora = self._get_available_lora() if _available_lora is None: - raise RuntimeError(f'No lora available for tenant {tenant_adapter_name}. Max loras: {self.max_loras}') + raise RuntimeError(f"No lora available for tenant {tenant_adapter_name}. Max loras: {self.max_loras}") if config.r > self.max_r: - raise RuntimeError(f'Too big rank for lora: {config.r}') + raise RuntimeError(f"Too big rank for lora: {config.r}") _available_lora.tenant_config = config _available_lora.tenant_adapter_name = tenant_adapter_name - if getattr(config, 'target_parameters', None): + if getattr(config, "target_parameters", None): self.target_parameter_manager.acquire( tenant_adapter_name=tenant_adapter_name, slot_name=_available_lora.adapter_name, config=config, ) - logger.info(f'Lora count: {len(self.loras)}, available lora: {self._count_available_loras()}') + logger.info(f"Lora count: {len(self.loras)}, available lora: {self._count_available_loras()}") return _available_lora.adapter_name - def release_lora(self, tenant_adapter_name: str) -> Optional[str]: + def release_lora(self, tenant_adapter_name: str) -> str | None: try: _lora = self.find_lora_by_tenant(tenant_adapter_name) _lora.tenant_config = None _lora.tenant_adapter_name = None self._load_initial_weights(_lora.adapter_name) self.target_parameter_manager.release(tenant_adapter_name) - logger.info(f'Lora count: {len(self.loras)}, available lora: {self._count_available_loras()}') + logger.info(f"Lora count: {len(self.loras)}, available lora: {self._count_available_loras()}") except ValueError: return @@ -270,36 +319,37 @@ def find_lora_by_tenant(self, tenant_adapter_name): if len(_loras) > 0: return _loras[0] else: - raise ValueError(f'No lora found for tenant {tenant_adapter_name}') + raise ValueError(f"No lora found for tenant {tenant_adapter_name}") def find_lora(self, adapter_name): _loras = [_lora for _lora in self.loras if _lora.adapter_name == adapter_name] if len(_loras) > 0: return _loras[0] else: - raise ValueError(f'No lora found for real adapter_name {adapter_name}') + raise ValueError(f"No lora found for real adapter_name {adapter_name}") def validate_tenant_target_modules(self, target_modules, target_parameters=None) -> 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)] - if target_modules == 'all-linear' or (isinstance(target_modules, - (list, set)) and 'all-linear' in target_modules): + if target_modules == "all-linear" or ( + isinstance(target_modules, (list, set)) and "all-linear" in target_modules + ): return if isinstance(target_modules, (list, set)): missing = [target for target in target_modules if not any(name.endswith(target) for name in layers)] if missing: - raise ValueError(f'LoRA target_modules are outside the preallocated range: {sorted(missing)}') + raise ValueError(f"LoRA target_modules are outside the preallocated range: {sorted(missing)}") return if isinstance(target_modules, str) and any(self.match_target_modules(name, target_modules) for name in layers): return - raise ValueError(f'LoRA target_modules do not resolve inside the preallocated range: {target_modules!r}') + raise ValueError(f"LoRA target_modules do not resolve inside the preallocated range: {target_modules!r}") @staticmethod def match_target_modules( module_name: str, - target_modules: Optional[Union[List[str], str]], + target_modules: list[str] | str | None, ) -> bool: if target_modules is None: return False @@ -307,18 +357,18 @@ def match_target_modules( if isinstance(target_modules, (list, set)) and len(target_modules) == 0: return False - if isinstance(target_modules, (list, set)) and 'all-linear' in target_modules: + if isinstance(target_modules, (list, set)) and "all-linear" in target_modules: return True - if target_modules == 'all-linear': + if target_modules == "all-linear": return True # Strip LoRA-specific suffixes (e.g. ".lora_A.default.weight") so that # a full parameter name like "model.layers.0.attn.proj.lora_A.default.weight" # can be matched against target_modules like ["attn.proj"]. cleaned = module_name - if '.lora_' in module_name: - cleaned = re.sub(r'\.lora_\w+(\.[\w-]+)*$', '', module_name) + if ".lora_" in module_name: + cleaned = re.sub(r"\.lora_\w+(\.[\w-]+)*$", "", module_name) if isinstance(target_modules, str): return re.fullmatch(target_modules, cleaned) is not None @@ -358,13 +408,16 @@ def _linear_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: lora_B = self.lora_B[active_adapter] dropout = self.lora_dropout[_lora.adapter_name] scaling = _lora.tenant_config.lora_alpha / ( - _lora.tenant_config.r**0.5 if _lora.tenant_config.use_rslora else _lora.tenant_config.r) + _lora.tenant_config.r**0.5 if _lora.tenant_config.use_rslora else _lora.tenant_config.r + ) x = self._cast_input_dtype(x, lora_A.weight.dtype) dropout_x = dropout(x) lora_A_out = torch.nn.functional.linear( - dropout_x, lora_A.weight[:_lora.tenant_config.r, :], bias=None) + dropout_x, lora_A.weight[: _lora.tenant_config.r, :], bias=None + ) lora_B_out = torch.nn.functional.linear( - lora_A_out, lora_B.weight[:, :_lora.tenant_config.r], bias=None) + lora_A_out, lora_B.weight[:, : _lora.tenant_config.r], bias=None + ) result = result + lora_B_out * scaling result = result.to(torch_result_dtype) return result @@ -392,10 +445,11 @@ def _embedding_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: embedding_B = self.lora_embedding_B[active_adapter] scaling = _lora.tenant_config.lora_alpha / ( - _lora.tenant_config.r**0.5 if _lora.tenant_config.use_rslora else _lora.tenant_config.r) + _lora.tenant_config.r**0.5 if _lora.tenant_config.use_rslora else _lora.tenant_config.r + ) - embedding_A_T = embedding_A.T[:, :_lora.tenant_config.r] - embedding_B_T = embedding_B.T[:_lora.tenant_config.r, :] + embedding_A_T = embedding_A.T[:, : _lora.tenant_config.r] + embedding_B_T = embedding_B.T[: _lora.tenant_config.r, :] after_A = self._embed(x, embedding_A_T.T) lora_out = after_A @ embedding_B_T.T @@ -411,8 +465,11 @@ def _embedding_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: elif isinstance(base_layer, _LoraParallelLinear): def _megatron_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any): - from megatron.core.extensions.transformer_engine import (TEGroupedLinear, - TELayerNormColumnParallelLinear, TELinear) + from megatron.core.extensions.transformer_engine import ( + TEGroupedLinear, + TELayerNormColumnParallelLinear, + TELinear, + ) from megatron.core.transformer.moe.router import TopKRouter previous_dtype = x.dtype @@ -435,7 +492,7 @@ def _megatron_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any): with self._patch_router_gating(): result, bias = self.base_layer(x, *args, **kwargs) else: - raise ValueError(f'Unsupported base layer type: {type(self.base_layer)}') + raise ValueError(f"Unsupported base layer type: {type(self.base_layer)}") if not isinstance(self.base_layer, TopKRouter) and not self.disable_adapters and not self.merged: for active_adapter in self.active_adapters: @@ -451,29 +508,36 @@ def _megatron_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any): lora_B = self.lora_B[active_adapter] dropout = self.lora_dropout[_lora.adapter_name] scaling = _lora.tenant_config.lora_alpha / ( - _lora.tenant_config.r**0.5 if _lora.tenant_config.use_rslora else _lora.tenant_config.r) + _lora.tenant_config.r**0.5 if _lora.tenant_config.use_rslora else _lora.tenant_config.r + ) dtype = lora_A.weight0.dtype if isinstance(lora_A, TEGroupedLinear) else lora_A.weight.dtype x = x.to(dtype) fp8_context = nullcontext() - if getattr(self.config, 'fp8_param', False): + if getattr(self.config, "fp8_param", False): import transformer_engine.pytorch as te + fp8_context = te.fp8_autocast(enabled=False) with fp8_context: # Keep the backend layer's TP/SP communication and parameter # identities. Mask inactive ranks between A and B instead of # replacing parallel layers with local linear operations. lora_result = ( - lora_A(dropout(x), *args, **kwargs) if isinstance(lora_A, TEGroupedLinear) else lora_A( - dropout(x))) + lora_A(dropout(x), *args, **kwargs) + if isinstance(lora_A, TEGroupedLinear) + else lora_A(dropout(x)) + ) if isinstance(lora_result, tuple): lora_result = lora_result[0] - lora_result = torch.nn.functional.pad(lora_result[..., :_lora.tenant_config.r], - (0, _self.max_r - _lora.tenant_config.r)) + lora_result = torch.nn.functional.pad( + lora_result[..., : _lora.tenant_config.r], (0, _self.max_r - _lora.tenant_config.r) + ) lora_result = ( lora_B(lora_result, *args, **kwargs) - if isinstance(lora_B, TEGroupedLinear) else lora_B(lora_result)) + if isinstance(lora_B, TEGroupedLinear) + else lora_B(lora_result) + ) if isinstance(lora_result, tuple): lora_result = lora_result[0] lora_result = lora_result * scaling @@ -486,20 +550,16 @@ def _megatron_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any): base_layer.forward = MethodType(_megatron_forward, base_layer) base_layer.layer_name = name - def patch(self, - module: Union[torch.nn.Module, List[torch.nn.Module]], - target_modules='all-linear', - *args, - **kwargs): + def patch(self, module: torch.nn.Module | list[torch.nn.Module], target_modules="all-linear", *args, **kwargs): # ``module`` may be a list of model chunks (Megatron); probe the first # chunk for the device in that case. first_module = module[0] if isinstance(module, (list, tuple)) else module - module_device = getattr(first_module, 'device', None) + module_device = getattr(first_module, "device", None) if module_device is None: module_device = next(first_module.parameters()).device - low_cpu_mem_usage = module_device.type == 'meta' + low_cpu_mem_usage = module_device.type == "meta" - base_config = kwargs.get('lora_config', None) + base_config = kwargs.get("lora_config", None) for i in range(self.max_loras): config = deepcopy(base_config) if base_config is not None else None if config is None: @@ -507,18 +567,19 @@ def patch(self, r=self.max_r, target_modules=target_modules, lora_alpha=32, - exclude_modules=['o_a_proj'], + exclude_modules=["o_a_proj"], ) - lora_tenant = LoraTenant(index=i, adapter_name=f'lora_{i}', config=config) + lora_tenant = LoraTenant(index=i, adapter_name=f"lora_{i}", config=config) self.loras.append(lora_tenant) - def _patch_peft(_module): + def _patch_peft(_module, lora_tenant=lora_tenant, config=config): if isinstance(_module, PeftModel): _module.add_adapter(lora_tenant.adapter_name, config, low_cpu_mem_usage=low_cpu_mem_usage) self._autocast_adapter_dtype(_module, lora_tenant.adapter_name) else: _peft_model: PeftModel = get_peft_model( - _module, config, lora_tenant.adapter_name, low_cpu_mem_usage=low_cpu_mem_usage) + _module, config, lora_tenant.adapter_name, low_cpu_mem_usage=low_cpu_mem_usage + ) _module.active_adapters = _peft_model.active_adapters _module = _peft_model @@ -527,7 +588,7 @@ def _patch_peft(_module): self._patch_lora_forward(name, submodule) return _module - def _patch_megatron(_module): + def _patch_megatron(_module, lora_tenant=lora_tenant, config=config): # Expand target_modules (e.g., 'all-linear' -> actual module names) _config = deepcopy(config) if isinstance(_module, PeftModel): @@ -542,9 +603,11 @@ def _patch_megatron(_module): target_modules = list(_config.target_modules) from .megatron import MegatronModel + _config.target_modules = MegatronModel.get_target_modules(_module, target_modules) _module = get_peft_model( - _module, _config, lora_tenant.adapter_name, low_cpu_mem_usage=low_cpu_mem_usage) + _module, _config, lora_tenant.adapter_name, low_cpu_mem_usage=low_cpu_mem_usage + ) for name, submodule in _module.named_modules(): if isinstance(submodule, LoraLayer): @@ -559,18 +622,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 @@ -578,14 +635,14 @@ def _enable_all_lora_grad(_module): def save_initial_weights(self): for i in range(self.max_loras): lora_tenant = self.loras[i] - pattern = re.compile(rf'\.lora_(?:A|embedding_A)\.{re.escape(lora_tenant.adapter_name)}\.') + pattern = re.compile(rf"\.lora_(?:A|embedding_A)\.{re.escape(lora_tenant.adapter_name)}\.") - def _store_weights(_module): + def _store_weights(_module, pattern=pattern, lora_tenant=lora_tenant): for name, parameter in _module.named_parameters(): if self._is_target_parameter_lora_name(name): continue if pattern.search(name): - lora_tenant.lora_A_weights[name] = self._read_param_tensor(parameter).clone().to('cpu') + lora_tenant.lora_A_weights[name] = self._read_param_tensor(parameter).clone().to("cpu") if isinstance(self.module, list): for _module in self.module: @@ -594,29 +651,32 @@ def _store_weights(_module): _store_weights(self.module) def load_lora_converter(self, name, parameter, **kwargs): - def convert_param(name, parameter): - if 'embedding_A' in name: + if "embedding_A" in name: r_saved = parameter.shape[1] parameter = torch.cat( - (parameter, torch.zeros(parameter.shape[0], self.max_r - r_saved).to(parameter.dtype)), dim=1) - elif 'embedding_B' in name: + (parameter, torch.zeros(parameter.shape[0], self.max_r - r_saved).to(parameter.dtype)), dim=1 + ) + elif "embedding_B" in name: r_saved = parameter.shape[0] parameter = torch.cat( - (parameter, torch.zeros(self.max_r - r_saved, parameter.shape[1]).to(parameter.dtype)), dim=0) - elif '_A' in name: + (parameter, torch.zeros(self.max_r - r_saved, parameter.shape[1]).to(parameter.dtype)), dim=0 + ) + elif "_A" in name: r_saved = parameter.shape[0] parameter = torch.cat( - (parameter, torch.zeros(self.max_r - r_saved, parameter.shape[1]).to(parameter.dtype)), dim=0) - elif '_B' in name: + (parameter, torch.zeros(self.max_r - r_saved, parameter.shape[1]).to(parameter.dtype)), dim=0 + ) + elif "_B" in name: r_saved = parameter.shape[1] parameter = torch.cat( - (parameter, torch.zeros(parameter.shape[0], self.max_r - r_saved).to(parameter.dtype)), dim=1) + (parameter, torch.zeros(parameter.shape[0], self.max_r - r_saved).to(parameter.dtype)), dim=1 + ) return name, parameter if isinstance(parameter, torch.Tensor): return convert_param(name, parameter) - elif 'lazytensor' in parameter.__class__.__name__.lower(): + elif "lazytensor" in parameter.__class__.__name__.lower(): def _loader(self): tensor = self.loader_origin() @@ -651,14 +711,15 @@ def save_hf_key_context(self, adapter_name): target_modules = _lora.tenant_config.target_modules # 'all-linear' matches everything — no patching needed - if target_modules == 'all-linear' or (isinstance(target_modules, - (list, set)) and 'all-linear' in target_modules): + if target_modules == "all-linear" or ( + isinstance(target_modules, (list, set)) and "all-linear" in target_modules + ): yield return # Create a sibling class with compatible memory layout that is NOT # recognised as LoraParallelLinear by isinstance. - _ExcludedLinear = type('_ExcludedLinear', _LoraParallelLinear.__bases__, {}) + _ExcludedLinear = type("_ExcludedLinear", _LoraParallelLinear.__bases__, {}) patched = [] # (module, original_class) modules_list = self.module if isinstance(self.module, list) else [self.module] @@ -679,30 +740,30 @@ def save_hf_key_context(self, adapter_name): def save_lora_converter(self, name, parameter, adapter_name): _lora = self.find_lora(adapter_name) # Skip weights belonging to OTHER adapters - if re.search(r'\.lora_\w+\.\w+\.', name) and not re.search(rf'\.lora_\w+\.{adapter_name}\.', name): + if re.search(r"\.lora_\w+\.\w+\.", name) and not re.search(rf"\.lora_\w+\.{adapter_name}\.", name): return None # target_modules filtering is handled by save_hf_key_context (class # patching makes the bridge skip non-target modules entirely), so we # only check the adapter-name / weight pattern here. - if re.search(rf'\.lora_\w+\.({adapter_name}|weight)', name): + if re.search(rf"\.lora_\w+\.({adapter_name}|weight)", name): _param = self._slice_rank_tensor(name, self._read_param_tensor(parameter), _lora.tenant_config.r) if _param is not None: _param = _param.clone() - name = name.replace(f'.{_lora.adapter_name}.', '.') + name = name.replace(f".{_lora.adapter_name}.", ".") return name, _param else: return None def set_state_dict(self, tenant_adapter_name, state_dict): _lora = self.find_lora_by_tenant(tenant_adapter_name) - pattern = re.compile(rf'\.lora_\w+\.{re.escape(_lora.adapter_name)}\.') + pattern = re.compile(rf"\.lora_\w+\.{re.escape(_lora.adapter_name)}\.") def _load_weights(_module): for name, parameter in _module.named_parameters(): if self._is_target_parameter_lora_name(name): continue if pattern.search(name) and self.match_target_modules(name, _lora.tenant_config.target_modules): - state_key = name.replace(f'.{_lora.adapter_name}.', '.') + state_key = name.replace(f".{_lora.adapter_name}.", ".") target_tensor = self._read_param_tensor(parameter) if target_tensor is None: continue @@ -716,13 +777,13 @@ def _load_weights(_module): _load_weights(_module) else: _load_weights(self.module) - if getattr(_lora.tenant_config, 'target_parameters', None): + if getattr(_lora.tenant_config, "target_parameters", None): self.target_parameter_manager.set_state_dict(tenant_adapter_name, state_dict) def get_state_dict(self, tenant_adapter_name): state_dict = {} _lora = self.find_lora_by_tenant(tenant_adapter_name) - pattern = re.compile(rf'\.lora_\w+\.{re.escape(_lora.adapter_name)}\.') + pattern = re.compile(rf"\.lora_\w+\.{re.escape(_lora.adapter_name)}\.") def _get_weights(_module): state_dict = {} @@ -733,7 +794,7 @@ def _get_weights(_module): _param = self._slice_rank_tensor(name, self._read_param_tensor(parameter), _lora.tenant_config.r) if _param is None: continue - name = name.replace(f'.{_lora.adapter_name}.', '.') + name = name.replace(f".{_lora.adapter_name}.", ".") state_dict[name] = _param return state_dict @@ -742,18 +803,18 @@ def _get_weights(_module): state_dict.update(_get_weights(_module)) else: state_dict = _get_weights(self.module) - if getattr(_lora.tenant_config, 'target_parameters', None): + if getattr(_lora.tenant_config, "target_parameters", None): target_state_dict = self.target_parameter_manager.get_state_dict(tenant_adapter_name) overlap = state_dict.keys() & target_state_dict.keys() if overlap: - raise ValueError(f'Duplicate LoRA state keys: {sorted(overlap)[:5]}') + raise ValueError(f"Duplicate LoRA state keys: {sorted(overlap)[:5]}") state_dict.update(target_state_dict) return state_dict def _load_initial_weights(self, origin_adapter_name): _lora = self.find_lora(origin_adapter_name) - pattern_A = re.compile(rf'\.lora_(?:A|embedding_A)\.{origin_adapter_name}\.') - pattern_B = re.compile(rf'\.lora_(?:B|embedding_B)\.{origin_adapter_name}\.') + pattern_A = re.compile(rf"\.lora_(?:A|embedding_A)\.{origin_adapter_name}\.") + pattern_B = re.compile(rf"\.lora_(?:B|embedding_B)\.{origin_adapter_name}\.") def _load_initial_weights(_module): for name, parameter in _module.named_parameters(): @@ -781,11 +842,11 @@ def _active_lora_numel(name, parameter, r: int) -> int: shape = parameter.shape if len(shape) != 2: return parameter.numel() - if 'embedding_A' in name: + if "embedding_A" in name: rank_dim = 1 - elif 'embedding_B' in name or '_A' in name: + elif "embedding_B" in name or "_A" in name: rank_dim = 0 - elif '_B' in name: + elif "_B" in name: rank_dim = 1 else: return parameter.numel() @@ -800,13 +861,13 @@ def get_nb_trainable_parameters(self, tenant_adapter_name) -> tuple[int, int]: """ _lora = self.find_lora_by_tenant(tenant_adapter_name) adapter_name = _lora.adapter_name - pattern = re.compile(rf'\.lora_\w+\.{re.escape(adapter_name)}\.') + pattern = re.compile(rf"\.lora_\w+\.{re.escape(adapter_name)}\.") def _count_trainable_parameters(_module): trainable_params = 0 all_param = 0 for name, param in _module.named_parameters(): - if not pattern.search(name) and 'lora_' in name: + if not pattern.search(name) and "lora_" in name: # Other lora continue if pattern.search(name) and not self.match_target_modules(name, _lora.tenant_config.target_modules): @@ -817,13 +878,13 @@ def _count_trainable_parameters(_module): num_params = self._active_lora_numel(name, param, _lora.tenant_config.r) else: num_params = param.numel() - if num_params == 0 and hasattr(param, 'ds_numel'): + if num_params == 0 and hasattr(param, "ds_numel"): num_params = param.ds_numel - if param.__class__.__name__ == 'Params4bit': - if hasattr(param, 'element_size'): + if param.__class__.__name__ == "Params4bit": + if hasattr(param, "element_size"): num_bytes = param.element_size() - elif not hasattr(param, 'quant_storage'): + elif not hasattr(param, "quant_storage"): num_bytes = 1 else: num_bytes = param.quant_storage.itemsize @@ -850,14 +911,17 @@ def get_trainable_parameters_example(self, tenant_adapter_name): trainable_param_names = [] _lora = self.find_lora_by_tenant(tenant_adapter_name) adapter_name = _lora.adapter_name - pattern = re.compile(rf'\.lora_\w+\.{re.escape(adapter_name)}\.') + pattern = re.compile(rf"\.lora_\w+\.{re.escape(adapter_name)}\.") def _get_parameters(_module): for name, parameter in _module.named_parameters(): - if parameter.requires_grad and pattern.search(name) and self.match_target_modules( - name, _lora.tenant_config.target_modules): - name = name.replace(f'A.{adapter_name}', f'A.{tenant_adapter_name}') - name = name.replace(f'B.{adapter_name}', f'B.{tenant_adapter_name}') + if ( + parameter.requires_grad + and pattern.search(name) + and self.match_target_modules(name, _lora.tenant_config.target_modules) + ): + name = name.replace(f"A.{adapter_name}", f"A.{tenant_adapter_name}") + name = name.replace(f"B.{adapter_name}", f"B.{tenant_adapter_name}") trainable_param_names.append(name) if isinstance(self.module, list): @@ -866,8 +930,8 @@ def _get_parameters(_module): else: _get_parameters(self.module) - trainable_param_names = trainable_param_names[:5] + ['...'] + trainable_param_names[-5:] - trainable_param_names = '\n'.join(trainable_param_names) + trainable_param_names = trainable_param_names[:5] + ["..."] + trainable_param_names[-5:] + trainable_param_names = "\n".join(trainable_param_names) return trainable_param_names def get_target_parameter_trainable_parameters(self, tenant_adapter_name): diff --git a/src/twinkle/model/optimizer_group.py b/src/twinkle/model/optimizer_group.py index cced80a67..e087f27fb 100644 --- a/src/twinkle/model/optimizer_group.py +++ b/src/twinkle/model/optimizer_group.py @@ -1,10 +1,11 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import math -import torch from dataclasses import dataclass, field +from typing import Any, Optional + +import torch from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler -from typing import Any, Dict, List, Optional from twinkle import DeviceMesh from twinkle.data_format import InputFeature, ModelOutput @@ -20,13 +21,14 @@ class TrainStatus: Encapsulates inputs, outputs, loss, tokens count and metrics for a training/eval step. """ - inputs: List[InputFeature] = None + + inputs: list[InputFeature] = None outputs: ModelOutput = None loss_value: Any = None num_tokens: int = 0 - metrics: List[Metric] = field(default_factory=list) - forward_kwargs: Dict[str, Any] = field(default_factory=dict) - loss_reduction: Optional[str] = None + metrics: list[Metric] = field(default_factory=list) + forward_kwargs: dict[str, Any] = field(default_factory=dict) + loss_reduction: str | None = None @dataclass @@ -35,12 +37,13 @@ class BaseOptimizerGroup: Subclasses: OptimizerGroup (Transformers), MegatronOptimizerGroup (Megatron) """ + adapter_name: str = None adapter_config: Any = None optimizer: Optimizer = None lr_scheduler: LRScheduler = None loss_instance: Loss = None - loss_registry: Dict[str, Loss] = field(default_factory=dict) + loss_registry: dict[str, Loss] = field(default_factory=dict) train_status: TrainStatus = None eval_status: TrainStatus = None template: Template = None @@ -54,22 +57,23 @@ class BaseOptimizerGroup: pending_failed: bool = False pending_loss: float = 0.0 - def set_loss(self, loss: Loss, loss_key: Optional[str] = None) -> None: + def set_loss(self, loss: Loss, loss_key: str | None = None) -> None: if loss_key is None: self.loss_registry.clear() self.loss_instance = loss return if not isinstance(loss_key, str) or not loss_key: - raise ValueError('loss_key must be a non-empty string') + raise ValueError("loss_key must be a non-empty string") self.loss_registry[loss_key] = loss - def get_loss(self, loss_key: Optional[str] = None) -> Loss: + def get_loss(self, loss_key: str | None = None) -> Loss: if loss_key is not None: try: return self.loss_registry[loss_key] except KeyError as exc: raise ValueError( - f'Unknown loss_key={loss_key!r}; registered keys are {list(self.loss_registry)}') from exc + f"Unknown loss_key={loss_key!r}; registered keys are {list(self.loss_registry)}" + ) from exc if len(self.loss_registry) == 1: return next(iter(self.loss_registry.values())) if len(self.loss_registry) > 1: @@ -78,7 +82,7 @@ def get_loss(self, loss_key: Optional[str] = None) -> Loss: # carry a key must keep working when CHORD registers its second, keyed loss. if self.loss_instance is not None: return self.loss_instance - raise ValueError(f'loss_key is required when multiple losses are registered: {list(self.loss_registry)}') + raise ValueError(f"loss_key is required when multiple losses are registered: {list(self.loss_registry)}") return self.loss_instance def prepare_backward_transaction(self, finalize_gradients: bool, loss_weight: float) -> bool: @@ -87,15 +91,16 @@ def prepare_backward_transaction(self, finalize_gradients: bool, loss_weight: fl Returns whether this call completes an already pending backward. """ if not math.isfinite(loss_weight) or loss_weight < 0: - raise ValueError(f'loss_weight must be finite and non-negative, got {loss_weight}') + raise ValueError(f"loss_weight must be finite and non-negative, got {loss_weight}") if self.pending_failed: raise RuntimeError( - 'The deferred backward transaction was interrupted; its gradients cannot be reused safely. ' - 'Reinitialize the model/optimizer state before continuing.') + "The deferred backward transaction was interrupted; its gradients cannot be reused safely. " + "Reinitialize the model/optimizer state before continuing." + ) completing_pending = self.pending_backward if not finalize_gradients: if completing_pending: - raise RuntimeError('A deferred backward is already pending; finalize it before starting another one.') + raise RuntimeError("A deferred backward is already pending; finalize it before starting another one.") self.pending_backward = True self.pending_failed = True elif completing_pending: @@ -116,55 +121,53 @@ def record_transaction_loss(self, local_loss: float, finalize_gradients: bool) - def require_finalized(self, operation: str) -> None: if self.pending_failed: raise RuntimeError( - f'Cannot {operation}: the deferred backward transaction was interrupted and its gradients ' - 'cannot be reused safely. Reinitialize the model/optimizer state before continuing.') + f"Cannot {operation}: the deferred backward transaction was interrupted and its gradients " + "cannot be reused safely. Reinitialize the model/optimizer state before continuing." + ) if self.pending_backward: raise RuntimeError( - f'Cannot {operation} while a deferred backward is pending; call forward_backward with ' - 'finalize_gradients=True first.') + f"Cannot {operation} while a deferred backward is pending; call forward_backward with " + "finalize_gradients=True first." + ) - def do_grad_sync(self, gradient_accumulation_steps: Optional[int] = None) -> bool: + def do_grad_sync(self, gradient_accumulation_steps: int | None = None) -> bool: if gradient_accumulation_steps is None: gradient_accumulation_steps = self.gradient_accumulation_steps else: self.gradient_accumulation_steps = gradient_accumulation_steps - return gradient_accumulation_steps == 1 or ((self.cur_step - 1) % gradient_accumulation_steps == 0 - and self.cur_step > 1) + return gradient_accumulation_steps == 1 or ( + (self.cur_step - 1) % gradient_accumulation_steps == 0 and self.cur_step > 1 + ) def _get_lr(self): """Get learning rates from optimizer. Override in subclass.""" return [] - 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.""" + def _metric_kwargs(self, status: TrainStatus, metric_step: int | None = 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: int | None = None): + """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..722f38c2f 100644 --- a/src/twinkle/model/transformers/hybrid/fft_slots.py +++ b/src/twinkle/model/transformers/hybrid/fft_slots.py @@ -1,9 +1,9 @@ # Copyright (c) ModelScope Contributors. All rights reserved. + import torch from peft.tuners.lora import LoraLayer from peft.utils import ModulesToSaveWrapper from torch import nn -from typing import Dict, List from twinkle.model.multi_lora import MultiLora @@ -16,20 +16,20 @@ class HybridFftSlots: from a tenant's existing LoRA slot index. """ - def __init__(self, multi_lora: MultiLora, s_fft: List[str]) -> None: + def __init__(self, multi_lora: MultiLora, s_fft: list[str]) -> None: """Bind the shared LoRA manager to its server-owned FFT allocation.""" if not s_fft: - raise ValueError('Hybrid requires at least one S_FFT module.') + raise ValueError("Hybrid requires at least one S_FFT module.") if len(set(s_fft)) != len(s_fft): - raise ValueError('Hybrid S_FFT contains the same module more than once.') + raise ValueError("Hybrid S_FFT contains the same module more than once.") normalized_s_fft = [self._canonical_module_name(name) for name in s_fft] if len(set(normalized_s_fft)) != len(normalized_s_fft): - raise ValueError('Hybrid S_FFT aliases resolve to the same layer.') + raise ValueError("Hybrid S_FFT aliases resolve to the same layer.") self.multi_lora = multi_lora self.s_fft = normalized_s_fft self.hybrid_adapters: set[str] = set() - self.allocated_to_layer_name: Dict[str, str] = {} - self.allocated_to_wrapper_name: Dict[str, str] = {} + self.allocated_to_layer_name: dict[str, str] = {} + self.allocated_to_wrapper_name: dict[str, str] = {} @property def module(self): @@ -39,38 +39,39 @@ def module(self): @staticmethod def _canonical_module_name(name: str) -> str: """Remove the module name prefix introduced by PEFT.""" - prefix = 'base_model.model.' - return name[len(prefix):] if name.startswith(prefix) else name + prefix = "base_model.model." + return name[len(prefix) :] if name.startswith(prefix) else name def install_fft_slots(self) -> None: """Install PEFT ``ModulesToSaveWrapper`` slots before DDP/FSDP wrapping.""" if isinstance(self.module, list): - raise NotImplementedError('Hybrid FFT slots currently require the Transformers backend.') + raise NotImplementedError("Hybrid FFT slots currently require the Transformers backend.") named_modules = dict(self.module.named_modules()) resolved_layer_names = set() for allocated_name in self.s_fft: matches = [ - (name, layer) for name, layer in named_modules.items() + (name, layer) + for name, layer in named_modules.items() if self._canonical_module_name(name) == allocated_name and isinstance(layer, (LoraLayer, nn.Linear)) ] if len(matches) != 1: - raise ValueError(f'Hybrid S_FFT module {allocated_name!r} resolved to {len(matches)} layers.') + raise ValueError(f"Hybrid S_FFT module {allocated_name!r} resolved to {len(matches)} layers.") layer_name, layer = matches[0] if layer_name in resolved_layer_names: - raise ValueError(f'Hybrid S_FFT aliases resolve to the same layer {layer_name!r}.') + raise ValueError(f"Hybrid S_FFT aliases resolve to the same layer {layer_name!r}.") resolved_layer_names.add(layer_name) if isinstance(layer, LoraLayer): original_module = layer.base_layer - wrapper_name = f'{layer_name}.base_layer' + wrapper_name = f"{layer_name}.base_layer" else: original_module = layer wrapper_name = layer_name if any(parameter.is_meta for parameter in original_module.parameters()): - raise ValueError('Hybrid FFT slots require materialized base weights.') - wrapper = ModulesToSaveWrapper(original_module, 'fft_0') + raise ValueError("Hybrid FFT slots require materialized base weights.") + wrapper = ModulesToSaveWrapper(original_module, "fft_0") for slot in range(1, self.multi_lora.max_loras): - wrapper.update(f'fft_{slot}') + wrapper.update(f"fft_{slot}") # Keep FFT copies in the base dtype so FSDP2 sees the same original # dtype when another tenant activates a different slot. base_dtype = next((parameter.dtype for parameter in original_module.parameters()), torch.float32) @@ -81,7 +82,7 @@ def install_fft_slots(self) -> None: if isinstance(layer, LoraLayer): layer.base_layer = wrapper else: - parent_name, _, child_name = layer_name.rpartition('.') + parent_name, _, child_name = layer_name.rpartition(".") parent = self.module.get_submodule(parent_name) if parent_name else self.module setattr(parent, child_name, wrapper) self.allocated_to_layer_name[allocated_name] = layer_name @@ -105,7 +106,7 @@ def _tenant(self, adapter_name: str): def _fft_adapter_name(self, adapter_name: str) -> str: """Map a tenant's LoRA slot index to its PEFT FFT adapter name.""" - return f'fft_{self._tenant(adapter_name).index}' + return f"fft_{self._tenant(adapter_name).index}" def _get_fft_wrapper(self, allocated_name: str) -> ModulesToSaveWrapper: """Resolve one allocated module to its installed PEFT wrapper.""" @@ -120,6 +121,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) @@ -137,7 +142,7 @@ def _enable_preallocated_fft_grad(wrapper: ModulesToSaveWrapper) -> None: if not parameter.requires_grad: parameter.requires_grad_(True) - def resolve_lora_targets(self, target_modules) -> List[str]: + def resolve_lora_targets(self, target_modules) -> list[str]: """Resolve tenant LoRA targets while reserving S_FFT for full tuning.""" lora_layers = self.multi_lora.lora_layer_names fft_layers = set(self.allocated_to_layer_name.values()) @@ -147,12 +152,14 @@ def resolve_lora_targets(self, target_modules) -> List[str]: selected = [name for name in lora_layers if self.multi_lora.match_target_modules(name, target_modules)] return sorted(name for name in selected if name not in fft_layers) - def _tenant_lora_layer_names(self, adapter_name: str) -> List[str]: + def _tenant_lora_layer_names(self, adapter_name: str) -> list[str]: """Return this Hybrid tenant's LoRA layers, excluding its FFT layers.""" tenant = self._tenant(adapter_name) fft_layers = set(self.allocated_to_layer_name.values()) return [ - name for name in self.multi_lora.lora_layer_names if name not in fft_layers + name + for name in self.multi_lora.lora_layer_names + if name not in fft_layers and self.multi_lora.match_target_modules(name, tenant.tenant_config.target_modules) ] @@ -170,87 +177,93 @@ def reset_adapter_slot(self, adapter_name: str) -> None: for allocated_name, _, _, parameter_name, target in self._iter_fft_slot_parameters(adapter_name): wrapper = self._get_fft_wrapper(allocated_name) original_parameters = dict(wrapper.original_module.named_parameters()) - self.multi_lora._write_param_tensor(target, - self.multi_lora._read_param_tensor(original_parameters[parameter_name])) + self.multi_lora._write_param_tensor( + target, self.multi_lora._read_param_tensor(original_parameters[parameter_name]) + ) @staticmethod def _checkpoint_key(allocated_name: str, parameter_name: str) -> str: """Build the plain Transformers checkpoint key for an FFT tensor.""" - return f'base_model.model.{allocated_name}.{parameter_name}' + return f"base_model.model.{allocated_name}.{parameter_name}" - def get_fft_state_dict(self, adapter_name: str) -> Dict[str, torch.Tensor]: + def get_fft_state_dict(self, adapter_name: str) -> dict[str, torch.Tensor]: """Return an independent snapshot of one tenant's FFT state.""" if not self.is_hybrid(adapter_name): return {} state = {} for allocated_name, _, _, parameter_name, parameter in self._iter_fft_slot_parameters(adapter_name): state[self._checkpoint_key(allocated_name, parameter_name)] = ( - self.multi_lora._read_param_tensor(parameter).detach().clone()) + self.multi_lora._read_param_tensor(parameter).detach().clone() + ) return state - def set_fft_state_dict(self, adapter_name: str, state_dict: Dict[str, torch.Tensor]) -> None: + def set_fft_state_dict(self, adapter_name: str, state_dict: dict[str, torch.Tensor]) -> None: """Restore one tenant's FFT copies from lossless training state.""" if not self.is_hybrid(adapter_name): return for allocated_name, _, _, parameter_name, parameter in self._iter_fft_slot_parameters(adapter_name): key = self._checkpoint_key(allocated_name, parameter_name) if key not in state_dict: - raise ValueError(f'Hybrid training state is missing {key!r}.') + raise ValueError(f"Hybrid training state is missing {key!r}.") self.multi_lora._write_param_tensor(parameter, state_dict[key]) def named_fft_parameters(self, adapter_name: str): """Return trainable FFT parameters with their PEFT state-dict names.""" if not self.is_hybrid(adapter_name): return [] - return [(f'{wrapper_name}.modules_to_save.{fft_adapter_name}.{parameter_name}', parameter) for _, wrapper_name, - fft_adapter_name, parameter_name, parameter in self._iter_fft_slot_parameters(adapter_name)] + return [ + (f"{wrapper_name}.modules_to_save.{fft_adapter_name}.{parameter_name}", parameter) + for _, wrapper_name, fft_adapter_name, parameter_name, parameter in self._iter_fft_slot_parameters( + adapter_name + ) + ] @staticmethod def _normalize_base_state_key(name: str) -> str: """Convert a PEFT base-module key to its plain Transformers key.""" - prefix = 'base_model.model.' + prefix = "base_model.model." if name.startswith(prefix): - name = name[len(prefix):] - name = name.replace('.base_layer.original_module.', '.') - name = name.replace('.original_module.', '.') - return name.replace('.base_layer.', '.') + name = name[len(prefix) :] + name = name.replace(".base_layer.original_module.", ".") + name = name.replace(".original_module.", ".") + return name.replace(".base_layer.", ".") - def iter_merged_state_dict(self, adapter_name: str, full_state_dict: Dict[str, torch.Tensor]): + def iter_merged_state_dict(self, adapter_name: str, full_state_dict: dict[str, torch.Tensor]): """Yield a non-destructively merged plain Transformers state dict.""" if not self.is_hybrid(adapter_name): - raise ValueError(f'Adapter {adapter_name!r} is not a Hybrid adapter.') + raise ValueError(f"Adapter {adapter_name!r} is not a Hybrid adapter.") tenant = self._tenant(adapter_name) slot = tenant.adapter_name - replacements: Dict[str, torch.Tensor] = {} + replacements: dict[str, torch.Tensor] = {} for layer_name in self._tenant_lora_layer_names(adapter_name): - base_key = f'{layer_name}.base_layer.weight' - a_key = f'{layer_name}.lora_A.{slot}.weight' - b_key = f'{layer_name}.lora_B.{slot}.weight' + base_key = f"{layer_name}.base_layer.weight" + a_key = f"{layer_name}.lora_A.{slot}.weight" + b_key = f"{layer_name}.lora_B.{slot}.weight" missing = [key for key in (base_key, a_key, b_key) if key not in full_state_dict] if missing: - raise ValueError(f'Cannot export Hybrid module {layer_name!r}; missing {missing}.') + raise ValueError(f"Cannot export Hybrid module {layer_name!r}; missing {missing}.") base = full_state_dict[base_key] rank = tenant.tenant_config.r a = full_state_dict[a_key][:rank, :] b = full_state_dict[b_key][:, :rank] scaling = tenant.tenant_config.lora_alpha / (rank**0.5 if tenant.tenant_config.use_rslora else rank) delta = b.to(torch.float32) @ a.to(torch.float32) - if getattr(tenant.tenant_config, 'fan_in_fan_out', False): + if getattr(tenant.tenant_config, "fan_in_fan_out", False): delta = delta.transpose(0, 1) replacements[base_key] = base + delta.to(dtype=base.dtype) * scaling for allocated_name, wrapper_name, fft_adapter_name, parameter_name, _ in self._iter_fft_slot_parameters( - adapter_name): - base_key = f'{wrapper_name}.original_module.{parameter_name}' - fft_key = f'{wrapper_name}.modules_to_save.{fft_adapter_name}.{parameter_name}' + adapter_name + ): + base_key = f"{wrapper_name}.original_module.{parameter_name}" + fft_key = f"{wrapper_name}.modules_to_save.{fft_adapter_name}.{parameter_name}" if base_key not in full_state_dict or fft_key not in full_state_dict: - raise ValueError(f'Cannot export Hybrid FFT layer {allocated_name!r}.') + raise ValueError(f"Cannot export Hybrid FFT layer {allocated_name!r}.") replacements[base_key] = full_state_dict[fft_key] emitted = set() for name, value in full_state_dict.items(): - if ('.lora_' in name or '.modules_to_save.' in name - or self.multi_lora._is_target_parameter_lora_name(name)): + if ".lora_" in name or ".modules_to_save." in name or self.multi_lora._is_target_parameter_lora_name(name): continue output_name = self._normalize_base_state_key(name) if output_name in emitted: @@ -258,27 +271,28 @@ def iter_merged_state_dict(self, adapter_name: str, full_state_dict: Dict[str, t emitted.add(output_name) yield output_name, replacements.get(name, value).detach().cpu() - def build_merged_state_dict(self, adapter_name: str, full_state_dict: Dict[str, torch.Tensor]): + def build_merged_state_dict(self, adapter_name: str, full_state_dict: dict[str, torch.Tensor]): """Materialize the merged state dict for deployment-oriented saving.""" return dict(self.iter_merged_state_dict(adapter_name, full_state_dict)) - def build_training_state_dict(self, adapter_name: str, full_state_dict: Dict[str, torch.Tensor]): + def build_training_state_dict(self, adapter_name: str, full_state_dict: dict[str, torch.Tensor]): """Extract lossless LoRA and FFT state for a Hybrid tenant.""" if not self.is_hybrid(adapter_name): - raise ValueError(f'Adapter {adapter_name!r} is not a Hybrid adapter.') + raise ValueError(f"Adapter {adapter_name!r} is not a Hybrid adapter.") tenant = self._tenant(adapter_name) state = {} for layer_name in self._tenant_lora_layer_names(adapter_name): - for kind in ('A', 'B'): - source = f'{layer_name}.lora_{kind}.{tenant.adapter_name}.weight' + for kind in ("A", "B"): + source = f"{layer_name}.lora_{kind}.{tenant.adapter_name}.weight" if source not in full_state_dict: - raise ValueError(f'Hybrid training state is missing {source!r}.') + raise ValueError(f"Hybrid training state is missing {source!r}.") value = self.multi_lora._slice_rank_tensor(source, full_state_dict[source], tenant.tenant_config.r) - state[source.replace(f'.{tenant.adapter_name}.', '.')] = value.detach().cpu() + state[source.replace(f".{tenant.adapter_name}.", ".")] = value.detach().cpu() for allocated_name, wrapper_name, fft_adapter_name, parameter_name, _ in self._iter_fft_slot_parameters( - adapter_name): - source = f'{wrapper_name}.modules_to_save.{fft_adapter_name}.{parameter_name}' + adapter_name + ): + source = f"{wrapper_name}.modules_to_save.{fft_adapter_name}.{parameter_name}" if source not in full_state_dict: - raise ValueError(f'Hybrid training state is missing {source!r}.') + raise ValueError(f"Hybrid training state is missing {source!r}.") state[self._checkpoint_key(allocated_name, parameter_name)] = full_state_dict[source].detach().cpu() return state diff --git a/src/twinkle/model/transformers/hybrid/model.py b/src/twinkle/model/transformers/hybrid/model.py index 26d25f674..fa2ee6d07 100644 --- a/src/twinkle/model/transformers/hybrid/model.py +++ b/src/twinkle/model/transformers/hybrid/model.py @@ -1,14 +1,15 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -import json import os import shutil +from contextlib import contextmanager +from typing import Any, Optional, Union + +import json import torch import torch.distributed as dist -from contextlib import contextmanager from peft import PeftConfig from safetensors.torch import load_file, save_file from torch.optim import Optimizer -from typing import Any, Dict, Optional, Type, Union from twinkle import DeviceMesh, Platform, remote_class, remote_function from twinkle.utils.safetensors import StreamingSafetensorSaver @@ -17,41 +18,44 @@ from .spectral_allocation import load_spectral_allocation _HYBRID_CONFIG_FIELDS = ( - 'r', - 'lora_alpha', - 'lora_dropout', - 'use_rslora', - 'fan_in_fan_out', - 'use_dora', - 'bias', - 'rank_pattern', - 'alpha_pattern', - 'target_modules', - 'modules_to_save', + "r", + "lora_alpha", + "lora_dropout", + "use_rslora", + "fan_in_fan_out", + "use_dora", + "bias", + "rank_pattern", + "alpha_pattern", + "target_modules", + "modules_to_save", ) -HYBRID_ADAPTER_MODE = 'hybrid' +HYBRID_ADAPTER_MODE = "hybrid" @remote_class() class SpectralHybridTransformersModel(MultiLoraTransformersModel): """Transformers MultiLoRA service extended with server-owned FFT slots.""" - def __init__(self, - hybrid: Dict[str, Any], - memory_efficient_init: bool = False, - device_mesh: Optional[DeviceMesh] = None, - **kwargs): + def __init__( + self, + hybrid: dict[str, Any], + memory_efficient_init: bool = False, + device_mesh: DeviceMesh | None = None, + **kwargs, + ): if memory_efficient_init: raise ValueError( - 'Spectral Hybrid does not support memory_efficient_init because FFT slots require materialized ' - 'base weights before FSDP wrapping.') + "Spectral Hybrid does not support memory_efficient_init because FFT slots require materialized " + "base weights before FSDP wrapping." + ) config = dict(hybrid or {}) - allocation_path = config.get('allocation_path') + allocation_path = config.get("allocation_path") if not allocation_path: - raise ValueError('Spectral Hybrid requires allocation_path.') + raise ValueError("Spectral Hybrid requires allocation_path.") s_fft = load_spectral_allocation(allocation_path) - self.default_lr_lora = float(config.get('default_lr_lora', 2.5e-5)) - self.default_lr_fft = float(config.get('default_lr_fft', 1e-6)) + self.default_lr_lora = float(config.get("default_lr_lora", 2.5e-5)) + self.default_lr_fft = float(config.get("default_lr_fft", 1e-6)) super().__init__(memory_efficient_init=False, device_mesh=device_mesh, **kwargs) self.fft_slots = HybridFftSlots(self.multi_adapter, s_fft) self.fft_slots.install_fft_slots() @@ -64,40 +68,30 @@ 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): - adapter_mode = kwargs.pop('adapter_mode', 'lora') - if adapter_mode == 'lora': + def add_adapter_to_model(self, adapter_name: str, config_or_dir: PeftConfig | str, **kwargs): + adapter_mode = kwargs.pop("adapter_mode", "lora") + if adapter_mode == "lora": return super().add_adapter_to_model(adapter_name, config_or_dir, **kwargs) if adapter_mode != HYBRID_ADAPTER_MODE: raise ValueError(f'Unsupported adapter_mode {adapter_mode!r}; expected "lora" or {HYBRID_ADAPTER_MODE!r}.') config = self._copy_lora_config(config_or_dir) if config.modules_to_save: - raise ValueError('Hybrid modules_to_save is controlled by the server allocation.') - if getattr(config, 'target_parameters', None): - raise ValueError('Hybrid target_parameters is not supported.') + raise ValueError("Hybrid modules_to_save is controlled by the server allocation.") + if getattr(config, "target_parameters", None): + raise ValueError("Hybrid target_parameters is not supported.") config.target_modules = set(self.fft_slots.resolve_lora_targets(config.target_modules)) config.modules_to_save = list(self.fft_slots.s_fft) self._register_adapter(adapter_name, config, **kwargs) @@ -107,39 +101,43 @@ def _create_param_group(self, adapter_name: str, lr: float = 1e-5, weight_decay: if not self.fft_slots.is_hybrid(adapter_name): return super()._create_param_group(adapter_name=adapter_name, lr=lr, weight_decay=weight_decay, **kwargs) params = self._get_trainable_parameters(adapter_name) - fft_token = '.modules_to_save.fft_' + fft_token = ".modules_to_save.fft_" lora_names = [name for name in params if fft_token not in name] fft_names = [name for name in params if fft_token in name] groups = [] if lora_names: - groups.append({ - 'params': [params[name] for name in lora_names], - 'param_names': lora_names, - 'lr': kwargs.get('lr_lora', self.default_lr_lora), - 'weight_decay': weight_decay, - }) + groups.append( + { + "params": [params[name] for name in lora_names], + "param_names": lora_names, + "lr": kwargs.get("lr_lora", self.default_lr_lora), + "weight_decay": weight_decay, + } + ) if fft_names: - groups.append({ - 'params': [params[name] for name in fft_names], - 'param_names': fft_names, - 'lr': kwargs.get('lr_fft', self.default_lr_fft), - 'weight_decay': weight_decay, - }) + groups.append( + { + "params": [params[name] for name in fft_names], + "param_names": fft_names, + "lr": kwargs.get("lr_fft", self.default_lr_fft), + "weight_decay": weight_decay, + } + ) if not groups: - raise ValueError(f'Spectral Hybrid adapter {adapter_name!r} has no trainable parameters.') + raise ValueError(f"Spectral Hybrid adapter {adapter_name!r} has no trainable parameters.") return groups @remote_function() - def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str], **kwargs): - adapter_name = kwargs.get('adapter_name') + def set_optimizer(self, optimizer_cls: type[Optimizer] | str, **kwargs): + adapter_name = kwargs.get("adapter_name") if not self.fft_slots.is_hybrid(adapter_name): return super().set_optimizer(optimizer_cls, **kwargs) - if 'params' not in kwargs: - lr_lora = kwargs.pop('lr_lora', kwargs.get('lr', self.default_lr_lora)) - lr_fft = kwargs.pop('lr_fft', self.default_lr_fft) - kwargs['params'] = self._create_param_group( + if "params" not in kwargs: + lr_lora = kwargs.pop("lr_lora", kwargs.get("lr", self.default_lr_lora)) + lr_fft = kwargs.pop("lr_fft", self.default_lr_fft) + kwargs["params"] = self._create_param_group( adapter_name, - weight_decay=kwargs.get('weight_decay', 0.01), + weight_decay=kwargs.get("weight_decay", 0.01), lr_lora=lr_lora, lr_fft=lr_fft, ) @@ -153,7 +151,8 @@ def _get_trainable_parameters(self, adapter_name): # MultiLoRA preallocates every slot; only collect S_LORA targets enabled for this tenant. params = { name: parameter - for name, parameter in params.items() if self.multi_adapter.match_target_modules(name, target_modules) + for name, parameter in params.items() + if self.multi_adapter.match_target_modules(name, target_modules) } known_parameter_ids = {id(parameter) for parameter in params.values()} for name, parameter in self.fft_slots.named_fft_parameters(adapter_name): @@ -162,22 +161,22 @@ def _get_trainable_parameters(self, adapter_name): known_parameter_ids.add(id(parameter)) return params - def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) -> Dict[str, str]: + def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) -> dict[str, str]: """Extend the LoRA mapping with the Hybrid FFT ``default`` slot.""" mapping = super()._optimizer_param_name_mapping(adapter_name, optimizer) tenant = self.multi_adapter.find_lora_by_tenant(adapter_name) - physical_token = f'.modules_to_save.fft_{tenant.index}.' - checkpoint_token = '.modules_to_save.default.' + physical_token = f".modules_to_save.fft_{tenant.index}." + checkpoint_token = ".modules_to_save.default." for group in optimizer.param_groups: - for name in group.get('param_names', []): + for name in group.get("param_names", []): checkpoint_name = name.replace(physical_token, checkpoint_token) if checkpoint_name != name: mapping[name] = checkpoint_name return mapping - @remote_function(collect='first') + @remote_function(collect="first") def get_state_dict(self, **kwargs): - adapter_name = kwargs.get('adapter_name') + adapter_name = kwargs.get("adapter_name") self._check_adapter_valid(adapter_name) state = self.multi_adapter.get_state_dict(adapter_name) if self.fft_slots.is_hybrid(adapter_name): @@ -188,37 +187,42 @@ def _validate_hybrid_training_checkpoint_boundary(self, adapter_name: str) -> No """Require a quiescent optimizer-step boundary for a lossless checkpoint.""" optimizer_group = self.optimizer_group[adapter_name] if optimizer_group.optimizer is None: - raise ValueError('Spectral Hybrid optimizer must be configured before save_optimizer=True.') + raise ValueError("Spectral Hybrid optimizer must be configured before save_optimizer=True.") for group in optimizer_group.optimizer.param_groups: - if len(group.get('param_names', [])) != len(group['params']): + if len(group.get("param_names", [])) != len(group["params"]): raise ValueError( - 'Spectral Hybrid lossless checkpoints require optimizer param_names for every parameter.') + "Spectral Hybrid lossless checkpoints require optimizer param_names for every parameter." + ) train_status = optimizer_group.train_status if train_status.loss_value is not None or train_status.num_tokens != 0: - raise ValueError('Spectral Hybrid training state can only be saved after the optimizer step and zero_grad.') + raise ValueError("Spectral Hybrid training state can only be saved after the optimizer step and zero_grad.") if any(parameter.grad is not None for parameter in self._get_trainable_parameters(adapter_name).values()): raise ValueError( - 'Spectral Hybrid training state can only be saved after zero_grad cleared accumulated gradients.') - if (optimizer_group.cur_step > 0 and optimizer_group.gradient_accumulation_steps > 1 - and not optimizer_group.do_grad_sync()): - raise ValueError('Spectral Hybrid training state cannot be saved in the middle of gradient accumulation.') + "Spectral Hybrid training state can only be saved after zero_grad cleared accumulated gradients." + ) + if ( + optimizer_group.cur_step > 0 + and optimizer_group.gradient_accumulation_steps > 1 + and not optimizer_group.do_grad_sync() + ): + raise ValueError("Spectral Hybrid training state cannot be saved in the middle of gradient accumulation.") @staticmethod - def _class_identity(instance) -> Optional[str]: + def _class_identity(instance) -> str | None: if instance is None: return None cls = instance.__class__ - return f'{cls.__module__}.{cls.__qualname__}' + return f"{cls.__module__}.{cls.__qualname__}" @staticmethod def _normalize_hybrid_config(config) -> dict: - raw = config.to_dict() if hasattr(config, 'to_dict') else dict(config) + raw = config.to_dict() if hasattr(config, "to_dict") else dict(config) normalized = {} for field in _HYBRID_CONFIG_FIELDS: value = raw.get(field) - if field in ('target_modules', 'modules_to_save'): + if field in ("target_modules", "modules_to_save"): value = sorted(value or []) - elif field in ('rank_pattern', 'alpha_pattern'): + elif field in ("rank_pattern", "alpha_pattern"): value = dict(sorted((value or {}).items())) normalized[field] = value return normalized @@ -227,66 +231,70 @@ def _write_hybrid_training_manifest(self, training_dir: str, adapter_name: str) if not Platform.is_master(): return optimizer_group = self.optimizer_group[adapter_name] - trainer_state_path = os.path.join(training_dir, 'trainer_state.json') - with open(trainer_state_path, encoding='utf-8') as handle: + trainer_state_path = os.path.join(training_dir, "trainer_state.json") + with open(trainer_state_path, encoding="utf-8") as handle: trainer_state = json.load(handle) - trainer_state.update({ - 'checkpoint_boundary': 'optimizer_step', - 'scheduler_class': self._class_identity(optimizer_group.lr_scheduler), - 'has_scaler': optimizer_group.scaler is not None, - }) - with open(trainer_state_path, 'w', encoding='utf-8') as handle: + trainer_state.update( + { + "checkpoint_boundary": "optimizer_step", + "scheduler_class": self._class_identity(optimizer_group.lr_scheduler), + "has_scaler": optimizer_group.scaler is not None, + } + ) + with open(trainer_state_path, "w", encoding="utf-8") as handle: json.dump(trainer_state, handle, ensure_ascii=False, indent=2, sort_keys=True) - handle.write('\n') + handle.write("\n") - def _validate_hybrid_resume_state(self, training_dir: str, adapter_name: str, saved_config: dict, - trainer_state: dict) -> None: + def _validate_hybrid_resume_state( + self, training_dir: str, adapter_name: str, saved_config: dict, trainer_state: dict + ) -> None: optimizer_group = self.optimizer_group[adapter_name] current_config = self._normalize_hybrid_config(optimizer_group.adapter_config) checkpoint_config = self._normalize_hybrid_config(saved_config) differences = { field: (checkpoint_config[field], current_config[field]) - for field in _HYBRID_CONFIG_FIELDS if checkpoint_config[field] != current_config[field] + for field in _HYBRID_CONFIG_FIELDS + if checkpoint_config[field] != current_config[field] } if differences: - raise ValueError(f'Spectral Hybrid adapter config does not match checkpoint: {differences}') - if trainer_state.get('checkpoint_boundary') != 'optimizer_step': - raise ValueError('Spectral Hybrid checkpoint was not saved at a supported optimizer-step boundary.') + raise ValueError(f"Spectral Hybrid adapter config does not match checkpoint: {differences}") + if trainer_state.get("checkpoint_boundary") != "optimizer_step": + raise ValueError("Spectral Hybrid checkpoint was not saved at a supported optimizer-step boundary.") if optimizer_group.optimizer is None: - raise ValueError('Spectral Hybrid optimizer must be configured before resuming training.') - optimizer_path = os.path.join(training_dir, 'optimizer.pt') + raise ValueError("Spectral Hybrid optimizer must be configured before resuming training.") + optimizer_path = os.path.join(training_dir, "optimizer.pt") if not os.path.isfile(optimizer_path): - raise ValueError('Spectral Hybrid training state is missing optimizer.pt.') - saved_scheduler = trainer_state.get('scheduler_class') + raise ValueError("Spectral Hybrid training state is missing optimizer.pt.") + saved_scheduler = trainer_state.get("scheduler_class") if saved_scheduler != self._class_identity(optimizer_group.lr_scheduler): - raise ValueError('Spectral Hybrid scheduler configuration does not match the checkpoint.') - if saved_scheduler and not os.path.isfile(os.path.join(training_dir, 'scheduler.pt')): - raise ValueError('Spectral Hybrid training state is missing scheduler.pt.') - saved_scaler = bool(trainer_state.get('has_scaler')) + raise ValueError("Spectral Hybrid scheduler configuration does not match the checkpoint.") + if saved_scheduler and not os.path.isfile(os.path.join(training_dir, "scheduler.pt")): + raise ValueError("Spectral Hybrid training state is missing scheduler.pt.") + saved_scaler = bool(trainer_state.get("has_scaler")) if saved_scaler != (optimizer_group.scaler is not None): - raise ValueError('Spectral Hybrid grad scaler configuration does not match the checkpoint.') - if saved_scaler and not os.path.isfile(os.path.join(training_dir, 'scaler.pt')): - raise ValueError('Spectral Hybrid training state is missing scaler.pt.') + raise ValueError("Spectral Hybrid grad scaler configuration does not match the checkpoint.") + if saved_scaler and not os.path.isfile(os.path.join(training_dir, "scaler.pt")): + raise ValueError("Spectral Hybrid training state is missing scaler.pt.") rank = dist.get_rank() if dist.is_initialized() else 0 - rank_rng_path = os.path.join(training_dir, f'rng_state_rank_{rank}.pt') + rank_rng_path = os.path.join(training_dir, f"rng_state_rank_{rank}.pt") if not os.path.isfile(rank_rng_path): - raise ValueError(f'Spectral Hybrid training state is missing rank RNG state: {rank_rng_path}') + raise ValueError(f"Spectral Hybrid training state is missing rank RNG state: {rank_rng_path}") - def _save_spectral_hybrid(self, name, output_dir: Optional[str], interval: int, adapter_name: str, **kwargs): + def _save_spectral_hybrid(self, name, output_dir: str | None, interval: int, adapter_name: str, **kwargs): optimizer_group = self.optimizer_group[adapter_name] - save_optimizer = kwargs.get('save_optimizer', False) - save_only_training_state = kwargs.get('save_only_training_state', False) + save_optimizer = kwargs.get("save_optimizer", False) + save_only_training_state = kwargs.get("save_only_training_state", False) if save_only_training_state and not save_optimizer: - raise ValueError('save_only_training_state=True requires save_optimizer=True.') + raise ValueError("save_only_training_state=True requires save_optimizer=True.") if name is None: - name = f'checkpoint-step-{optimizer_group.cur_step}' - output_dir = output_dir or 'output' + name = f"checkpoint-step-{optimizer_group.cur_step}" + output_dir = output_dir or "output" checkpoint_dir = os.path.join(output_dir, name) if optimizer_group.cur_step % interval != 0: return None if save_optimizer: self._validate_hybrid_training_checkpoint_boundary(adapter_name) - training_dir = os.path.join(checkpoint_dir, 'twinkle_training_state') + training_dir = os.path.join(checkpoint_dir, "twinkle_training_state") if Platform.is_master() and os.path.isdir(training_dir): shutil.rmtree(training_dir) @@ -295,8 +303,8 @@ def _save_spectral_hybrid(self, name, output_dir: Optional[str], interval: int, full_state = self.strategy.get_full_state_dict(self.model) saver = StreamingSafetensorSaver( checkpoint_dir, - max_shard_size=kwargs.get('max_shard_size', '5GB'), - save_rank='master', + max_shard_size=kwargs.get("max_shard_size", "5GB"), + save_rank="master", ) if Platform.is_master(): for key, value in self.fft_slots.iter_merged_state_dict(adapter_name, full_state): @@ -306,11 +314,11 @@ def _save_spectral_hybrid(self, name, output_dir: Optional[str], interval: int, model = self.strategy.unwrap_model(self.model) if Platform.is_master(): self.hf_config.save_pretrained(checkpoint_dir) - generation_config = getattr(model, 'generation_config', None) + generation_config = getattr(model, "generation_config", None) if generation_config is not None: generation_config.save_pretrained(checkpoint_dir) else: - generation_config_path = os.path.join(checkpoint_dir, 'generation_config.json') + generation_config_path = os.path.join(checkpoint_dir, "generation_config.json") if os.path.exists(generation_config_path): os.unlink(generation_config_path) self._save_tokenizer(checkpoint_dir, adapter_name=adapter_name) @@ -319,45 +327,42 @@ def _save_spectral_hybrid(self, name, output_dir: Optional[str], interval: int, if save_only_training_state: tenant = self.multi_adapter.find_lora_by_tenant(adapter_name) full_state = self.strategy.get_adapter_state_dict(self.model, tenant.adapter_name) - full_state.update(self.strategy.get_adapter_state_dict(self.model, f'fft_{tenant.index}')) + full_state.update(self.strategy.get_adapter_state_dict(self.model, f"fft_{tenant.index}")) if Platform.is_master(): adapter_state = self.fft_slots.build_training_state_dict(adapter_name, full_state) os.makedirs(training_dir, exist_ok=True) optimizer_group.adapter_config.save_pretrained(training_dir) - config_path = os.path.join(training_dir, 'adapter_config.json') - with open(config_path, encoding='utf-8') as handle: + config_path = os.path.join(training_dir, "adapter_config.json") + with open(config_path, encoding="utf-8") as handle: adapter_config = json.load(handle) - adapter_config['twinkle_adapter_mode'] = HYBRID_ADAPTER_MODE - with open(config_path, 'w', encoding='utf-8') as handle: + adapter_config["twinkle_adapter_mode"] = HYBRID_ADAPTER_MODE + with open(config_path, "w", encoding="utf-8") as handle: json.dump(adapter_config, handle, ensure_ascii=False, indent=2, sort_keys=True) - handle.write('\n') + handle.write("\n") save_file( - { - key: value.contiguous() - for key, value in adapter_state.items() - }, - os.path.join(training_dir, 'adapter_model.safetensors'), + {key: value.contiguous() for key, value in adapter_state.items()}, + os.path.join(training_dir, "adapter_model.safetensors"), ) if dist.is_initialized(): dist.barrier() self._save_training_state( training_dir, adapter_name=adapter_name, - consumed_train_samples=kwargs.get('consumed_train_samples', 0), - component_states=kwargs.get('component_states', {}), + consumed_train_samples=kwargs.get("consumed_train_samples", 0), + component_states=kwargs.get("component_states", {}), ) self._write_hybrid_training_manifest(training_dir, adapter_name) rank = dist.get_rank() if dist.is_initialized() else 0 - torch.save(self._get_training_rng_state(), os.path.join(training_dir, f'rng_state_rank_{rank}.pt')) + torch.save(self._get_training_rng_state(), os.path.join(training_dir, f"rng_state_rank_{rank}.pt")) if dist.is_initialized(): dist.barrier() return checkpoint_dir # lazy_collect=False: see TransformersModel.save -- a checkpoint write the caller acts on immediately # must block until it lands, or the run's final save is torn down with the actor before writing. - @remote_function(collect='first', lazy_collect=False) - def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): - adapter_name = kwargs.pop('adapter_name', None) + @remote_function(collect="first", lazy_collect=False) + def save(self, name, output_dir: str | None = None, interval=1, **kwargs): + adapter_name = kwargs.pop("adapter_name", None) self._check_adapter_valid(adapter_name) if not self.fft_slots.is_hybrid(adapter_name): return super().save(name, output_dir, interval, adapter_name=adapter_name, **kwargs) @@ -367,20 +372,22 @@ def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): return checkpoint_dir def _resume_spectral_hybrid(self, checkpoint_dir: str, adapter_name: str, resume_only_model: bool): - training_dir = os.path.join(checkpoint_dir, 'twinkle_training_state') - adapter_path = os.path.join(training_dir, 'adapter_model.safetensors') - trainer_state_path = os.path.join(training_dir, 'trainer_state.json') + training_dir = os.path.join(checkpoint_dir, "twinkle_training_state") + adapter_path = os.path.join(training_dir, "adapter_model.safetensors") + trainer_state_path = os.path.join(training_dir, "trainer_state.json") if not os.path.isfile(adapter_path) or not os.path.isfile(trainer_state_path): - raise ValueError('Cannot resume Spectral Hybrid training from a merged-only checkpoint. ' - 'Save the checkpoint with save_optimizer=True to create twinkle_training_state.') - adapter_config_path = os.path.join(training_dir, 'adapter_config.json') + raise ValueError( + "Cannot resume Spectral Hybrid training from a merged-only checkpoint. " + "Save the checkpoint with save_optimizer=True to create twinkle_training_state." + ) + adapter_config_path = os.path.join(training_dir, "adapter_config.json") if not os.path.isfile(adapter_config_path): - raise ValueError('Spectral Hybrid training state is missing adapter_config.json.') - with open(adapter_config_path, encoding='utf-8') as handle: + raise ValueError("Spectral Hybrid training state is missing adapter_config.json.") + with open(adapter_config_path, encoding="utf-8") as handle: saved_config = json.load(handle) - if saved_config.get('twinkle_adapter_mode') != HYBRID_ADAPTER_MODE: - raise ValueError('Checkpoint training state is not a Spectral Hybrid adapter.') - with open(trainer_state_path, encoding='utf-8') as handle: + if saved_config.get("twinkle_adapter_mode") != HYBRID_ADAPTER_MODE: + raise ValueError("Checkpoint training state is not a Spectral Hybrid adapter.") + with open(trainer_state_path, encoding="utf-8") as handle: trainer_state = json.load(handle) if not resume_only_model: self._validate_hybrid_resume_state(training_dir, adapter_name, saved_config, trainer_state) @@ -388,32 +395,32 @@ def _resume_spectral_hybrid(self, checkpoint_dir: str, adapter_name: str, resume current_config = self._normalize_hybrid_config(self.optimizer_group[adapter_name].adapter_config) checkpoint_config = self._normalize_hybrid_config(saved_config) if current_config != checkpoint_config: - raise ValueError('Spectral Hybrid adapter config does not match checkpoint.') + raise ValueError("Spectral Hybrid adapter config does not match checkpoint.") rank = dist.get_rank() if dist.is_initialized() else 0 - adapter_state = load_file(adapter_path, device='cpu') + adapter_state = load_file(adapter_path, device="cpu") self.multi_adapter.set_state_dict(adapter_name, adapter_state) self.fft_slots.set_fft_state_dict(adapter_name, adapter_state) if not resume_only_model: trainer_state = self._restore_training_state(training_dir, adapter_name=adapter_name) - self._load_rng_state(os.path.join(training_dir, f'rng_state_rank_{rank}.pt')) + self._load_rng_state(os.path.join(training_dir, f"rng_state_rank_{rank}.pt")) result = { - 'cur_step': trainer_state['cur_step'], - 'consumed_train_samples': trainer_state['consumed_train_samples'], - 'gradient_accumulation_steps': trainer_state['gradient_accumulation_steps'], + "cur_step": trainer_state["cur_step"], + "consumed_train_samples": trainer_state["consumed_train_samples"], + "gradient_accumulation_steps": trainer_state["gradient_accumulation_steps"], } if resume_only_model: - result['component_states'] = {} - elif 'component_states' in trainer_state: - result['component_states'] = trainer_state['component_states'] + result["component_states"] = {} + elif "component_states" in trainer_state: + result["component_states"] = trainer_state["component_states"] return result # lazy_collect=False: see TransformersModel.resume_from_checkpoint -- the driver consumes the returned # training-state dict immediately, so this override must stay non-lazy too (it either delegates to # super() or returns _resume_spectral_hybrid's dict). Preserves backend equivalence. - @remote_function(dispatch='all', collect='first', sync=True, lazy_collect=False) + @remote_function(dispatch="all", collect="first", sync=True, lazy_collect=False) def resume_from_checkpoint(self, checkpoint_dir, *, resume_only_model=False, **kwargs): - adapter_name = kwargs.get('adapter_name', '') + adapter_name = kwargs.get("adapter_name", "") self._check_adapter_valid(adapter_name) if not self.fft_slots.is_hybrid(adapter_name): return super().resume_from_checkpoint(checkpoint_dir, resume_only_model=resume_only_model, **kwargs) diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index 3c8e57dea..3461b6e2f 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -1,14 +1,16 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import os -import torch.distributed as dist -import transformers +from collections.abc import Callable from contextlib import contextmanager from copy import deepcopy +from typing import Any, Literal, Optional, Union + +import torch.distributed as dist +import transformers from peft import LoraConfig, PeftConfig, PeftModel, load_peft_weights from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel -from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union from twinkle import DeviceMesh, remote_class, remote_function, template, torch_util from twinkle.data_format import InputFeature, Trajectory @@ -24,42 +26,49 @@ @remote_class() class MultiLoraTransformersModel(TransformersModel, PreTrainedModel): - def __init__( - self, # noqa - model_cls=None, - model_id: Optional[str] = None, - config: Optional[PretrainedConfig] = None, - device_mesh: Optional[DeviceMesh] = None, - mixed_precision: Literal['no', 'fp8', 'fp16', 'bf16'] = 'bf16', - strategy: Literal['accelerate', 'native_fsdp'] = 'accelerate', - ddp_config: Dict[str, Any] = None, - fsdp_config: Dict[str, Any] = None, - lora_config: Dict[str, Any] = None, - grad_scaler_config: Dict[str, Any] = None, - memory_efficient_init: bool = False, - max_loras: int = 5, - max_r: int = 32, - max_length: int = 8192, - target_modules: Union[List[str], str] = 'all-linear', - model_loader: Optional[ModelLoaderProtocol] = None, - acc_strategy: str = 'token', - model_arch: Optional[Dict[str, Any]] = None, - **kwargs): - new_special_tokens = kwargs.pop('new_special_tokens', None) + self, # noqa + model_cls=None, + model_id: str | None = None, + config: PretrainedConfig | None = None, + device_mesh: DeviceMesh | None = None, + mixed_precision: Literal["no", "fp8", "fp16", "bf16"] = "bf16", + strategy: Literal["accelerate", "native_fsdp"] = "accelerate", + ddp_config: dict[str, Any] = None, + fsdp_config: dict[str, Any] = None, + lora_config: dict[str, Any] = None, + grad_scaler_config: dict[str, Any] = None, + memory_efficient_init: bool = False, + max_loras: int = 5, + max_r: int = 32, + max_length: int = 8192, + target_modules: list[str] | str = "all-linear", + model_loader: ModelLoaderProtocol | None = None, + gradient_checkpointing: bool = True, + gradient_checkpointing_kwargs: dict[str, Any] | None = None, + acc_strategy: str = "token", + model_arch: dict[str, Any] | None = None, + **kwargs, + ): + new_special_tokens = kwargs.pop("new_special_tokens", None) if new_special_tokens: # Tenants share one frozen base tokenizer/embedding, so the vocab cannot be grown here; # pre-grow the tokenizer or use the single-tenant TransformersModel. - raise ValueError('new_special_tokens is not supported on the multi-LoRA model.') - os.environ['TOKENIZERS_PARALLELISM'] = 'true' + raise ValueError("new_special_tokens is not supported on the multi-LoRA model.") + os.environ["TOKENIZERS_PARALLELISM"] = "true" self._try_init_process_group() super(PreTrainedModel, self).__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_enabled = bool(kwargs.pop("enable_router_replay", False)) self._router_replay_applied = False # ``_acc_strategy`` and ``model_arch`` are read by inherited ``TransformersModel`` code -- the # former by ``_construct_default_optimizer_group`` (reached from ``_register_adapter``, i.e. every @@ -69,7 +78,7 @@ def __init__( # same role normalization the base applies. self._acc_strategy = acc_strategy self.model_arch = dict(model_arch or {}) - for _role in ('language_model', 'vision_tower', 'aligner'): + for _role in ("language_model", "vision_tower", "aligner"): _value = self.model_arch.get(_role) self.model_arch[_role] = [_value] if isinstance(_value, str) else list(_value or []) self._decide_strategy(strategy) @@ -88,8 +97,9 @@ def __init__( self.hf_config = model_loader.process_config(model_loader.build_config(model_id)) else: from transformers import AutoConfig + self.hf_config = AutoConfig.from_pretrained(model_id, trust_remote_code=True) - if model_cls is None and hasattr(self.hf_config, 'architectures'): + if model_cls is None and hasattr(self.hf_config, "architectures"): model_cls = self.hf_config.architectures[0] if model_cls is None: model_cls = AutoModelForCausalLM @@ -112,13 +122,16 @@ def __init__( self.model = model_cls.from_pretrained(model_id, config=self.hf_config, **kwargs) # pop: tokenizer_id is not a from_pretrained argument, so it must not leak into the model # construction kwargs above. - self.tokenizer_id = kwargs.pop('tokenizer_id', self.model_id) + self.tokenizer_id = kwargs.pop("tokenizer_id", self.model_id) self._model_wrapped = False self.sp_strategy = None # Initialize expert parallel attributes (required by set_optimizer in TransformersModel) - self.optimizer_group: Dict[str, OptimizerGroup] = {} + 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 @@ -131,11 +144,11 @@ def __init__( self.multi_adapter.reset_adapter_status() def _check_adapter_valid(self, adapter_name: str): - assert adapter_name and adapter_name in self.optimizer_group, (f'Use a valid adapter_name first, ' - f'current is: {adapter_name}') + assert adapter_name and adapter_name in self.optimizer_group, ( + f"Use a valid adapter_name first, current is: {adapter_name}" + ) def register_global_mm_forward_hook(self): - def forward_hook(model, args, kwargs): active_adapter = model.active_adapters[0] active_adapter = self.multi_adapter.find_lora(active_adapter).tenant_adapter_name @@ -153,18 +166,29 @@ 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') + raise NotImplementedError("Expert parallel is not supported with memory_efficient_init") return super()._maybe_apply_expert_parallel() def _ensure_target_parameter_lora_installed(self, config: LoraConfig) -> None: - target_parameters = getattr(config, 'target_parameters', None) + target_parameters = getattr(config, "target_parameters", None) if not target_parameters: return if self._model_wrapped: - raise RuntimeError('target_parameters LoRA must be installed before FSDP/DDP wrapping') - if getattr(self, '_enable_expert_parallel', False): + raise RuntimeError("target_parameters LoRA must be installed before FSDP/DDP wrapping") + if getattr(self, "_enable_expert_parallel", False): self.strategy.capture_pre_ep_state_if_needed(self.model, enable_ep=True) # Do not shard the MoE layers here: sharding must wait until the rank0 weight broadcast # during wrap_model -- ranks without materialized weights cannot be sharded yet. @@ -172,37 +196,46 @@ 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 - - @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) - def forward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) - optimizer_config = self.optimizer_group[kwargs.get('adapter_name')] - if (isinstance(inputs, dict) and self._not_encoded(inputs)) or (isinstance(inputs, list) - and self._not_encoded(inputs[0])): + """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: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], **kwargs): + self._check_adapter_valid(kwargs.get("adapter_name")) + optimizer_config = self.optimizer_group[kwargs.get("adapter_name")] + if (isinstance(inputs, dict) and self._not_encoded(inputs)) or ( + isinstance(inputs, list) and self._not_encoded(inputs[0]) + ): # Trajectory or List[Trajectory] - assert optimizer_config.template is not None, \ - 'Use set_template to add a template when trying to input `List[Trajectory]`' + assert optimizer_config.template is not None, ( + "Use set_template to add a template when trying to input `List[Trajectory]`" + ) if isinstance(inputs, dict): inputs = [inputs] inputs = optimizer_config.template.batch_encode(inputs) # noqa self.multi_adapter.check_length(inputs) - with self._adapter_context(kwargs.get('adapter_name')): + with self._adapter_context(kwargs.get("adapter_name")): return super().forward(inputs=inputs, **kwargs) - @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) - def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajectory]], **kwargs): - adapter_name = kwargs.get('adapter_name') - disable_lora = kwargs.get('disable_lora', False) + @remote_function(dispatch="slice_dp", collect=collect_tensor_dict) + def forward_only(self, *, inputs: InputFeature | list[InputFeature] | list[Trajectory], **kwargs): + adapter_name = kwargs.get("adapter_name") + disable_lora = kwargs.get("disable_lora", False) self._check_adapter_valid(adapter_name) optimizer_config = self.optimizer_group[adapter_name] - if (isinstance(inputs, dict) and self._not_encoded(inputs)) or (isinstance(inputs, list) - and self._not_encoded(inputs[0])): + if (isinstance(inputs, dict) and self._not_encoded(inputs)) or ( + isinstance(inputs, list) and self._not_encoded(inputs[0]) + ): # Trajectory or List[Trajectory] - assert optimizer_config.template is not None, \ - 'Use set_template to add a template when trying to input `List[Trajectory]`' + assert optimizer_config.template is not None, ( + "Use set_template to add a template when trying to input `List[Trajectory]`" + ) if isinstance(inputs, dict): inputs = [inputs] inputs = optimizer_config.template.batch_encode(inputs) # noqa @@ -210,88 +243,89 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T with self._adapter_context(adapter_name, disable_lora=disable_lora): return super().forward_only(inputs=inputs, **kwargs) - @remote_function(collect='mean') + @remote_function(collect="mean") def calculate_loss(self, **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) - with self._adapter_context(kwargs.get('adapter_name')): + self._check_adapter_valid(kwargs.get("adapter_name")) + with self._adapter_context(kwargs.get("adapter_name")): return super().calculate_loss(**kwargs) @remote_function() def backward(self, **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) - with self._adapter_context(kwargs.get('adapter_name')): + self._check_adapter_valid(kwargs.get("adapter_name")) + with self._adapter_context(kwargs.get("adapter_name")): super().backward(**kwargs) @remote_function() def clip_grad_norm(self, max_grad_norm: float = 1.0, norm_type=2, **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) - with self._adapter_context(kwargs.get('adapter_name')): + self._check_adapter_valid(kwargs.get("adapter_name")) + with self._adapter_context(kwargs.get("adapter_name")): return super().clip_grad_norm(max_grad_norm, norm_type=norm_type, **kwargs) @remote_function() def step(self, **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) - with self._adapter_context(kwargs.get('adapter_name')): + self._check_adapter_valid(kwargs.get("adapter_name")) + with self._adapter_context(kwargs.get("adapter_name")): super().step(**kwargs) @remote_function() def zero_grad(self, **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) - with self._adapter_context(kwargs.get('adapter_name')): + self._check_adapter_valid(kwargs.get("adapter_name")) + with self._adapter_context(kwargs.get("adapter_name")): super().zero_grad(**kwargs) @remote_function() def lr_step(self, **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) - with self._adapter_context(kwargs.get('adapter_name')): + self._check_adapter_valid(kwargs.get("adapter_name")) + with self._adapter_context(kwargs.get("adapter_name")): super().lr_step(**kwargs) @remote_function() - def set_loss(self, loss_cls: Union[Type[Loss], str], **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) + def set_loss(self, loss_cls: type[Loss] | str, **kwargs): + self._check_adapter_valid(kwargs.get("adapter_name")) super().set_loss(loss_cls, **kwargs) @remote_function() - def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str], **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) - with self._adapter_context(kwargs.get('adapter_name')): + def set_optimizer(self, optimizer_cls: type[Optimizer] | str, **kwargs): + self._check_adapter_valid(kwargs.get("adapter_name")) + with self._adapter_context(kwargs.get("adapter_name")): super().set_optimizer(optimizer_cls, **kwargs) @remote_function() - def add_adapter_to_model(self, adapter_name: str, config_or_dir: Union[PeftConfig, str], **kwargs): - adapter_mode = kwargs.pop('adapter_mode', 'lora') - if adapter_mode != 'lora': - raise ValueError(f'MultiLoraTransformersModel only supports LoRA adapters, got {adapter_mode!r}.') + def add_adapter_to_model(self, adapter_name: str, config_or_dir: PeftConfig | str, **kwargs): + adapter_mode = kwargs.pop("adapter_mode", "lora") + if adapter_mode != "lora": + raise ValueError(f"MultiLoraTransformersModel only supports LoRA adapters, got {adapter_mode!r}.") config = self._copy_lora_config(config_or_dir) if config.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.target_modules, - getattr(config, 'target_parameters', None)) + raise ValueError("modules_to_save is not supported for a multi-tenant LoRA adapter.") + self.multi_adapter.validate_tenant_target_modules( + config.target_modules, getattr(config, "target_parameters", None) + ) self._register_adapter(adapter_name, config, **kwargs) @staticmethod - def _copy_lora_config(config_or_dir: Union[PeftConfig, str]) -> LoraConfig: + def _copy_lora_config(config_or_dir: PeftConfig | str) -> LoraConfig: """Validate and copy a client-provided LoRA config before normalizing it.""" if isinstance(config_or_dir, str): - raise ValueError('Loading an adapter config from a model path or hub is not supported.') + raise ValueError("Loading an adapter config from a model path or hub is not supported.") if not isinstance(config_or_dir, LoraConfig): - raise TypeError('config_or_dir must be a LoraConfig instance.') + raise TypeError("config_or_dir must be a LoraConfig instance.") return deepcopy(config_or_dir) def _register_adapter(self, adapter_name: str, config: LoraConfig, **kwargs) -> None: """Register an already validated LoRA configuration in a free slot.""" config = self.strategy.prepare_adapter_config( config, - enable_ep=getattr(self, '_enable_expert_parallel', False), + enable_ep=getattr(self, "_enable_expert_parallel", False), ) # Limit the max peft version in pyproject.toml, in case any newer version opens some untested module grad. - config.bias = 'none' + config.bias = "none" config.init_lora_weights = False config.trainable_token_indices = None self.optimizer_group[adapter_name] = self._construct_default_optimizer_group() self.optimizer_group[adapter_name].adapter_name = adapter_name self.optimizer_group[adapter_name].adapter_config = config - _gas_default = kwargs.get('gradient_accumulation_steps', 1) + _gas_default = kwargs.get("gradient_accumulation_steps", 1) self.optimizer_group[adapter_name].gradient_accumulation_steps = _gas_default self._default_tokenizer = self.optimizer_group[adapter_name].template.processor self._ensure_target_parameter_lora_installed(config) @@ -302,37 +336,37 @@ def _register_adapter(self, adapter_name: str, config: LoraConfig, **kwargs) -> raise @remote_function() - def set_lr_scheduler(self, scheduler_cls: Union[Type[LRScheduler], str], **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) + def set_lr_scheduler(self, scheduler_cls: type[LRScheduler] | str, **kwargs): + self._check_adapter_valid(kwargs.get("adapter_name")) super().set_lr_scheduler(scheduler_cls, **kwargs) @remote_function() - def set_template(self, template_cls: Union[Type[template.Template], str], **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) + def set_template(self, template_cls: type[template.Template] | str, **kwargs): + self._check_adapter_valid(kwargs.get("adapter_name")) super().set_template(template_cls, **kwargs) @remote_function() - def set_processor(self, processor_cls: Union[Type[InputProcessor], str, Callable], **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) + def set_processor(self, processor_cls: type[InputProcessor] | str | Callable, **kwargs): + self._check_adapter_valid(kwargs.get("adapter_name")) super().set_processor(processor_cls, **kwargs) - @remote_function(collect='first') + @remote_function(collect="first") def get_state_dict(self, **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) - return self.multi_adapter.get_state_dict(kwargs.get('adapter_name')) + self._check_adapter_valid(kwargs.get("adapter_name")) + return self.multi_adapter.get_state_dict(kwargs.get("adapter_name")) def _get_adapter_state_dict_for_save(self, adapter_name: str) -> dict: adapter_state = self.multi_adapter.get_state_dict(adapter_name) return {key: torch_util.to_local_tensor(value).cpu() for key, value in adapter_state.items()} - def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) -> Dict[str, str]: + def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) -> dict[str, str]: """Map the tenant's physical LoRA slot to the PEFT ``default`` name.""" tenant = self.multi_adapter.find_lora_by_tenant(adapter_name) - physical_token = f'.{tenant.adapter_name}.' - checkpoint_token = f'.{CHECKPOINT_ADAPTER_NAME}.' + physical_token = f".{tenant.adapter_name}." + checkpoint_token = f".{CHECKPOINT_ADAPTER_NAME}." mapping = {} for group in optimizer.param_groups: - for name in group.get('param_names', []): + for name in group.get("param_names", []): checkpoint_name = name.replace(physical_token, checkpoint_token) if checkpoint_name != name: mapping[name] = checkpoint_name @@ -340,10 +374,11 @@ def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) # lazy_collect=False: see TransformersModel.save -- a checkpoint write the caller acts on immediately # must block until it lands, or the run's final save is torn down with the actor before writing. - @remote_function(collect='first', lazy_collect=False) - def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): - adapter_name = kwargs.get('adapter_name') + @remote_function(collect="first", lazy_collect=False) + def save(self, name, output_dir: str | None = 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(): @@ -351,23 +386,26 @@ def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): return checkpoint_dir @remote_function() - def load(self, name: str, output_dir: Optional[str] = None, **kwargs): - adapter_name = kwargs.get('adapter_name') + def load(self, name: str, output_dir: str | None = 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) + load_optimizer = kwargs.get("load_optimizer", False) if output_dir is None: if os.path.exists(name): checkpoint_dir = name else: - token = kwargs.pop('token', None) + token = kwargs.pop("token", None) checkpoint_dir = HubOperation.download_model(name, token=token) else: checkpoint_dir = os.path.join(output_dir, name) model = self.strategy.unwrap_model(self.model) if isinstance(model, PeftModel): # Load to CPU to avoid safetensors device issues in Ray environment - adapter_weights = load_peft_weights(checkpoint_dir, device='cpu') + adapter_weights = load_peft_weights(checkpoint_dir, device="cpu") self.multi_adapter.set_state_dict(adapter_name, adapter_weights) if load_optimizer: @@ -377,17 +415,17 @@ def load(self, name: str, output_dir: Optional[str] = None, **kwargs): @remote_function() def set_grad_scaler(self, **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) + self._check_adapter_valid(kwargs.get("adapter_name")) super().set_grad_scaler(**kwargs) @remote_function() - def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] = None, **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) + def add_metric(self, metric_cls: Metric | str, is_training: bool | None = None, **kwargs): + self._check_adapter_valid(kwargs.get("adapter_name")) super().add_metric(metric_cls, is_training, **kwargs) - @remote_function(collect='first', lazy_collect=False) + @remote_function(collect="first", lazy_collect=False) def calculate_metric(self, is_training, **kwargs): - self._check_adapter_valid(kwargs.get('adapter_name')) + self._check_adapter_valid(kwargs.get("adapter_name")) return super().calculate_metric(is_training, **kwargs) @remote_function() diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 85fa19b37..c412f5dda 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -1,16 +1,19 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import contextlib -import json -import numpy as np import os import random import re import time +from collections.abc import Callable +from copy import copy, deepcopy +from dataclasses import dataclass, field +from typing import Any, Literal, Optional, Union, overload + +import json +import numpy as np import torch import torch.distributed as dist import transformers -from copy import copy, deepcopy -from dataclasses import dataclass, field from peft import PeftConfig, PeftModel, get_peft_model from peft.utils import load_peft_weights from safetensors.torch import save_file @@ -19,55 +22,70 @@ from torch.optim.lr_scheduler import LRScheduler from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel from transformers.models.auto.auto_factory import _BaseAutoModelClass -from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union, overload import twinkle import twinkle.module.optimizer import twinkle.module.scheduler from twinkle import DeviceMesh, Platform, remote_class, remote_function from twinkle.checkpoint_engine import CheckpointEngine -from twinkle.checkpoint_engine.mixin import CheckpointEngineMixin from twinkle.checkpoint_engine.delta.trainer.backends import FSDP2DeltaBackend from twinkle.checkpoint_engine.delta.trainer.mixin import DeltaTrainerMixin -from twinkle.checkpoint_engine.weight_keys import (add_base_layer_suffix, is_lora_key, normalize_param_name, - print_weight_example) +from twinkle.checkpoint_engine.mixin import CheckpointEngineMixin +from twinkle.checkpoint_engine.weight_keys import ( + add_base_layer_suffix, + is_lora_key, + normalize_param_name, + print_weight_example, +) from twinkle.data_format import InputFeature, ModelOutput, SampleResponse, SamplingParams, Trajectory from twinkle.hub import HubOperation from twinkle.infra import collect_dp_and_flatten, collect_tensor_dict from twinkle.loss import CrossEntropyLoss, Loss from twinkle.metric import Accuracy, LossMetric, Metric, TrainMetric from twinkle.model.base import ModelLoaderProtocol, TrainableModel, copy_checkpoint_args, rotate_checkpoints -from twinkle.model.state_dict_utils import normalize_adapter_keys from twinkle.model.micro_batch import MicroBatchConfig, merge_micro_batch_outputs, plan_micro_batches, select_batch from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus +from twinkle.model.state_dict_utils import normalize_adapter_keys from twinkle.model.transformers.moe import apply_expert_parallel from twinkle.model.transformers.strategy import AccelerateStrategy, DeepSpeedStrategy, NativeFSDPStrategy -from twinkle.module.optimizer import (GaLoreConfig, MuonConfig, apply_param_groups_spec, create_galore_param_groups, - create_muon_param_groups, resolve_tower_prefixes) +from twinkle.module.optimizer import ( + GaLoreConfig, + MuonConfig, + apply_param_groups_spec, + create_galore_param_groups, + create_muon_param_groups, + resolve_tower_prefixes, +) from twinkle.patch import Patch, apply_context, apply_patch from twinkle.perf import timed from twinkle.processor import InputProcessor from twinkle.template import Template from twinkle.template.utils import construct_template -from twinkle.utils import (construct_class, get_logger, prepare_replayed_selective_log_softmax, - replayed_selective_log_softmax, selective_log_softmax, torch_util) +from twinkle.utils import ( + construct_class, + get_logger, + prepare_replayed_selective_log_softmax, + replayed_selective_log_softmax, + selective_log_softmax, + torch_util, +) from twinkle.utils.framework import Torch from twinkle.utils.grad_clip import normalize_and_clip_grad_norm from twinkle.utils.transformers_utils import filter_from_config_kwargs logger = get_logger() -CHECKPOINT_ADAPTER_NAME = 'default' +CHECKPOINT_ADAPTER_NAME = "default" def _get_vocab_size(config: PretrainedConfig) -> int: """Resolve the LM vocabulary size without running a sharded model forward.""" - vocab_size = getattr(config, 'vocab_size', None) + vocab_size = getattr(config, "vocab_size", None) if vocab_size is None: - text_config = getattr(config, 'text_config', None) - vocab_size = getattr(text_config, 'vocab_size', None) + text_config = getattr(config, "text_config", None) + vocab_size = getattr(text_config, "vocab_size", None) if not isinstance(vocab_size, int) or vocab_size <= 0: - raise ValueError('sampling replay requires a positive vocab_size in the model config') + raise ValueError("sampling replay requires a positive vocab_size in the model config") return vocab_size @@ -80,7 +98,7 @@ def _prepare_sampling_replay( ): """Build the labels and metadata shared by training and evaluation replay.""" if labels is None: - raise ValueError('labels are required when sampling replay is enabled') + raise ValueError("labels are required when sampling replay is enabled") loss_mask = (labels != -100).bool() masked_labels = labels.masked_fill(~loss_mask, 0) metadata = prepare_replayed_selective_log_softmax( @@ -102,7 +120,7 @@ def _prepare_sampling_replay( #: so excluding them here keeps the two backends equivalent. This matters for a FROZEN seq_cls reward/value #: model: it carries no loss, ``require_logps`` defaults True, so without this guard its ``forward_only`` #: would try to gather vocab logps out of a ``[B, 1]`` score and raise a shape error. -_NON_LOGPS_TASKS = frozenset({'embedding', 'seq_cls', 'reranker', 'generative_reranker', 'value'}) +_NON_LOGPS_TASKS = frozenset({"embedding", "seq_cls", "reranker", "generative_reranker", "value"}) def _resolve_task_context(model, task, template=None): @@ -117,33 +135,40 @@ def _resolve_task_context(model, task, template=None): 'seq_cls' / 'reranker' capture per-token scores BEFORE HF's internal pooling so ``postprocess_tensor_sp`` can do packed-aware last-token reduction (same as embedding). """ - if task in (None, 'causal_lm'): + if task in (None, "causal_lm"): return contextlib.nullcontext() - if task in ('seq_cls', 'reranker'): + if task in ("seq_cls", "reranker"): # Skip HF's internal pooling; surface per-token [B, T, num_labels] so the processor's # _pool_last_valid_token handles packed/non-packed uniformly. from twinkle.patch.transformers_value import TransformersValuePatch + return apply_context(model, TransformersValuePatch(squeeze_single=False)) - if task == 'embedding': + if task == "embedding": from twinkle.patch.transformers_emb import TransformersEmbeddingPatch + return apply_context(model, TransformersEmbeddingPatch()) - if task == 'generative_reranker': + if task == "generative_reranker": from twinkle.patch.transformers_reranker import TransformersGenerativeRerankerPatch + pos_id, neg_id = _resolve_generative_reranker_token_ids(template) return apply_context(model, TransformersGenerativeRerankerPatch(pos_id, neg_id)) - if task == 'fused_lm_ce': + if task == "fused_lm_ce": # Skip the lm_head GEMM so a fused-linear-CE loss can take over (Liger). # Device-agnostic: Liger self-dispatches the fused kernel across CUDA/NPU. from twinkle.patch.transformers_fused_ce import TransformersFusedCEPatch + return apply_context(model, TransformersFusedCEPatch()) - if task == 'value': + if task == "value": # PPO critic: reuse the num_labels=1 seq_cls ``score`` head but expose its PER-TOKEN output # (skip the last-token pooling ``*ForSequenceClassification`` does), so forward returns a value # V(s_t) at every token in ``logits`` [B, T]. Symmetric with Megatron's task='value'. from twinkle.patch.transformers_value import TransformersValuePatch + return apply_context(model, TransformersValuePatch()) - raise ValueError(f'Unknown task={task!r}; expected one of: causal_lm, embedding, reranker, seq_cls, ' - 'generative_reranker, fused_lm_ce, value.') + raise ValueError( + f"Unknown task={task!r}; expected one of: causal_lm, embedding, reranker, seq_cls, " + "generative_reranker, fused_lm_ce, value." + ) def _resolve_mm_graph_context(model, template): @@ -153,8 +178,9 @@ def _resolve_mm_graph_context(model, template): otherwise returns a no-op context. The patch ensures vision-tower params receive gradients even on ranks that get pure-text batches, preventing DDP/FSDP hangs. """ - if template is not None and getattr(template, 'is_mm', False) and hasattr(model, 'get_image_features'): + if template is not None and getattr(template, "is_mm", False) and hasattr(model, "get_image_features"): from twinkle.patch.transformers_uneven_mm_graph import MultimodalUnevenGraphPatch + return apply_context(model, MultimodalUnevenGraphPatch(template)) return contextlib.nullcontext() @@ -167,24 +193,27 @@ def _resolve_generative_reranker_token_ids(template): scores identically under legacy inference. """ import os - assert template is not None and getattr(template, 'processor', None) is not None, ( - 'generative_reranker needs a template with a tokenizer; call set_template before forwarding.') + + assert template is not None and getattr(template, "processor", None) is not None, ( + "generative_reranker needs a template with a tokenizer; call set_template before forwarding." + ) tokenizer = template.processor - pos_token = os.environ.get('GENERATIVE_RERANKER_POSITIVE_TOKEN', 'yes') - neg_token = os.environ.get('GENERATIVE_RERANKER_NEGATIVE_TOKEN', 'no') + pos_token = os.environ.get("GENERATIVE_RERANKER_POSITIVE_TOKEN", "yes") + neg_token = os.environ.get("GENERATIVE_RERANKER_NEGATIVE_TOKEN", "no") return tokenizer.convert_tokens_to_ids(pos_token), tokenizer.convert_tokens_to_ids(neg_token) @dataclass class OptimizerGroup(BaseOptimizerGroup): """Optimizer group for Transformers training.""" + adapter_config: PeftConfig = None loss_instance: Loss = CrossEntropyLoss scaler: GradScaler = None scaler_has_nan: bool = False checkpoint_engine: CheckpointEngine = None _handler: Any = None - acc_strategy: str = 'token' + acc_strategy: str = "token" def __post_init__(self): self._ensure_dp_group() @@ -216,7 +245,7 @@ def _ensure_dp_group(self): if dist.get_world_size() < raw_world_size: # World size is smaller than the requested dp/fsdp group; skip to avoid crash. return - dims = [dim for dim in ('dp', 'fsdp') if self._device_mesh.has_dim(dim)] + dims = [dim for dim in ("dp", "fsdp") if self._device_mesh.has_dim(dim)] if not dims: return self._dp_group = self._device_mesh.create_process_group(dims) @@ -224,38 +253,19 @@ def _ensure_dp_group(self): def _get_lr(self): if self.optimizer is not None: _lrs = [] - _default_lr = self.optimizer.defaults.get('lr') + _default_lr = self.optimizer.defaults.get("lr") for param_group in self.optimizer.param_groups: - _lrs.append(param_group.get('lr', _default_lr)) + _lrs.append(param_group.get("lr", _default_lr)) return _lrs else: return [] - def accumulate_metrics(self, is_training, include_default: bool = True, metric_step: Optional[int] = None): + def accumulate_metrics(self, is_training, include_default: bool = True, metric_step: int | None = 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) - - -_default_adapter_name = '' + super().accumulate_metrics(is_training, include_default=include_default, metric_step=metric_step) + + +_default_adapter_name = "" DEFAULT_LEARNING_RATE = 1e-5 DEFAULT_WEIGHT_DECAY = 0.01 # How many prompts go into one padded generate() during an in-model rollout. Matches the transformers @@ -263,7 +273,7 @@ def accumulate_metrics(self, is_training, include_default: bool = True, metric_s _GENERATION_MAX_BATCH_SIZE = 8 -def _read_hf_state_dict(checkpoint_dir: str) -> Dict[str, torch.Tensor]: +def _read_hf_state_dict(checkpoint_dir: str) -> dict[str, torch.Tensor]: """Read a full HuggingFace checkpoint directory into a CPU state dict. Supports both single-file and sharded ``safetensors`` layouts, falling back @@ -271,40 +281,41 @@ def _read_hf_state_dict(checkpoint_dir: str) -> Dict[str, torch.Tensor]: """ import json - state_dict: Dict[str, torch.Tensor] = {} - st_index = os.path.join(checkpoint_dir, 'model.safetensors.index.json') - st_single = os.path.join(checkpoint_dir, 'model.safetensors') - bin_index = os.path.join(checkpoint_dir, 'pytorch_model.bin.index.json') - bin_single = os.path.join(checkpoint_dir, 'pytorch_model.bin') + state_dict: dict[str, torch.Tensor] = {} + st_index = os.path.join(checkpoint_dir, "model.safetensors.index.json") + st_single = os.path.join(checkpoint_dir, "model.safetensors") + bin_index = os.path.join(checkpoint_dir, "pytorch_model.bin.index.json") + bin_single = os.path.join(checkpoint_dir, "pytorch_model.bin") if os.path.exists(st_index) or os.path.exists(st_single): from safetensors.torch import load_file + if os.path.exists(st_index): with open(st_index) as f: - weight_map = json.load(f)['weight_map'] + weight_map = json.load(f)["weight_map"] shards = sorted(set(weight_map.values())) for shard in shards: - state_dict.update(load_file(os.path.join(checkpoint_dir, shard), device='cpu')) + state_dict.update(load_file(os.path.join(checkpoint_dir, shard), device="cpu")) else: - state_dict.update(load_file(st_single, device='cpu')) + state_dict.update(load_file(st_single, device="cpu")) elif os.path.exists(bin_index) or os.path.exists(bin_single): if os.path.exists(bin_index): with open(bin_index) as f: - weight_map = json.load(f)['weight_map'] + weight_map = json.load(f)["weight_map"] shards = sorted(set(weight_map.values())) for shard in shards: state_dict.update( - torch.load(os.path.join(checkpoint_dir, shard), map_location='cpu', weights_only=True)) + torch.load(os.path.join(checkpoint_dir, shard), map_location="cpu", weights_only=True) + ) else: - state_dict.update(torch.load(bin_single, map_location='cpu', weights_only=True)) + state_dict.update(torch.load(bin_single, map_location="cpu", weights_only=True)) else: - raise FileNotFoundError(f'No safetensors/bin weights found in {checkpoint_dir}') + raise FileNotFoundError(f"No safetensors/bin weights found in {checkpoint_dir}") return state_dict @remote_class() -class TransformersModel(TrainableModel, PreTrainedModel, CheckpointEngineMixin, DeltaTrainerMixin, - FSDP2DeltaBackend): +class TransformersModel(TrainableModel, PreTrainedModel, CheckpointEngineMixin, DeltaTrainerMixin, FSDP2DeltaBackend): """The transformers model wrapper. Args: @@ -323,44 +334,45 @@ class TransformersModel(TrainableModel, PreTrainedModel, CheckpointEngineMixin, """ @overload - def __init__(self, *, model_cls: Type[PreTrainedModel], config: PretrainedConfig, remote_group, **kwargs) -> None: - ... + def __init__( + self, *, model_cls: type[PreTrainedModel], config: PretrainedConfig, remote_group, **kwargs + ) -> None: ... @overload - def __init__(self, *, model_id: str, config: Optional[PretrainedConfig] = None, **kwargs) -> None: - ... + def __init__(self, *, model_id: str, config: PretrainedConfig | None = None, **kwargs) -> None: ... def __init__( - self, # noqa - model_cls: Optional[Union[Type[PreTrainedModel], str, Type[_BaseAutoModelClass]]] = None, - model_id: Optional[str] = None, - config: Optional[PretrainedConfig] = None, - device_mesh: Optional[DeviceMesh] = None, - mixed_precision: Literal['no', 'fp8', 'fp16', 'bf16'] = 'bf16', - strategy: Literal['accelerate', 'native_fsdp', 'deepspeed'] = 'accelerate', - ddp_config: Dict[str, Any] = None, - fsdp_config: Dict[str, Any] = None, - grad_scaler_config: Dict[str, Any] = None, - memory_efficient_init: bool = False, - deepspeed_config: Dict[str, Any] = None, - model_loader: Optional[ModelLoaderProtocol] = None, - gradient_checkpointing: bool = True, - gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None, - vit_gradient_checkpointing: Optional[bool] = None, - vit_gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None, - model_arch: Optional[Dict[str, Any]] = None, - neftune_noise_alpha: Optional[float] = None, - acc_strategy: str = 'token', - new_special_tokens: Optional[List[str]] = None, - init_strategy: Optional[str] = None, - use_liger_kernel: bool = False, - liger_kernel_config: Optional[Dict[str, Any]] = None, - **kwargs): - os.environ['TOKENIZERS_PARALLELISM'] = 'true' + self, # noqa + model_cls: type[PreTrainedModel] | str | type[_BaseAutoModelClass] | None = None, + model_id: str | None = None, + config: PretrainedConfig | None = None, + device_mesh: DeviceMesh | None = None, + mixed_precision: Literal["no", "fp8", "fp16", "bf16"] = "bf16", + strategy: Literal["accelerate", "native_fsdp", "deepspeed"] = "accelerate", + ddp_config: dict[str, Any] = None, + fsdp_config: dict[str, Any] = None, + grad_scaler_config: dict[str, Any] = None, + memory_efficient_init: bool = False, + deepspeed_config: dict[str, Any] = None, + model_loader: ModelLoaderProtocol | None = None, + gradient_checkpointing: bool = True, + gradient_checkpointing_kwargs: dict[str, Any] | None = None, + vit_gradient_checkpointing: bool | None = None, + vit_gradient_checkpointing_kwargs: dict[str, Any] | None = None, + model_arch: dict[str, Any] | None = None, + neftune_noise_alpha: float | None = None, + acc_strategy: str = "token", + new_special_tokens: list[str] | None = None, + init_strategy: str | None = None, + use_liger_kernel: bool = False, + liger_kernel_config: dict[str, Any] | None = None, + **kwargs, + ): + os.environ["TOKENIZERS_PARALLELISM"] = "true" # Opt-out of the cuDNN SDPA backend (falls back to flash/mem-efficient, numerically # equivalent): sporadic `mha_graph.execute` RuntimeError on Blackwell + CUDA 13 # (ablate12 E7 crashed at update 21 mid-forward). Env-gated to keep default behavior. - if os.environ.get('TWINKLE_DISABLE_CUDNN_SDP', '0') == '1': + if os.environ.get("TWINKLE_DISABLE_CUDNN_SDP", "0") == "1": torch.backends.cuda.enable_cudnn_sdp(False) self._try_init_process_group() super(PreTrainedModel, self).__init__() @@ -372,20 +384,20 @@ def __init__( self._ddp_config = ddp_config or {} self._acc_strategy = acc_strategy self.model_arch = dict(model_arch or {}) - for _role in ('language_model', 'vision_tower', 'aligner'): + for _role in ("language_model", "vision_tower", "aligner"): _value = self.model_arch.get(_role) self.model_arch[_role] = [_value] if isinstance(_value, str) else list(_value or []) self._memory_efficient_init = memory_efficient_init - self._router_replay_enabled = bool(kwargs.pop('enable_router_replay', False)) + self._router_replay_enabled = bool(kwargs.pop("enable_router_replay", False)) self._router_replay_applied = False self._decide_strategy(strategy, deepspeed_config=deepspeed_config) self.grad_scaler_config = grad_scaler_config if model_id is not None: - model_id = HubOperation.download_model(model_id, revision=kwargs.get('revision')) + model_id = HubOperation.download_model(model_id, revision=kwargs.get("revision")) self.model_id = model_id # pop: tokenizer_id is not a from_pretrained argument, so it must not leak into the model # construction kwargs below. - self.tokenizer_id = kwargs.pop('tokenizer_id', self.model_id) + self.tokenizer_id = kwargs.pop("tokenizer_id", self.model_id) if config is not None: self.hf_config = config elif model_loader is not None and model_id is not None: @@ -395,12 +407,16 @@ def __init__( self.hf_config = model_loader.process_config(model_loader.build_config(model_id)) else: from transformers import AutoConfig + # Fallback path (no ModelLoader): this config load may execute the checkpoint's own # Python. State it loudly -- the same advisory class as a hardcoded trust_remote_code. - logger.warning_once('loading %s with trust_remote_code=True: the checkpoint\'s own Python code ' - 'will execute inside this process. Only use checkpoints you trust.', model_id) + logger.warning_once( + "loading %s with trust_remote_code=True: the checkpoint's own Python code " + "will execute inside this process. Only use checkpoints you trust.", + model_id, + ) self.hf_config = AutoConfig.from_pretrained(model_id, trust_remote_code=True) - if model_cls is None and hasattr(self.hf_config, 'architectures'): + if model_cls is None and hasattr(self.hf_config, "architectures"): model_cls = self.hf_config.architectures[0] if model_cls is None: model_cls = AutoModelForCausalLM @@ -431,15 +447,17 @@ def __init__( # config reference aligned with the config that the model actually uses. self.hf_config = self.model.config logger.info_once( - f'[TransformersModel] attn_implementation: requested={kwargs.get("attn_implementation")!r}, ' - f'resolved={getattr(self.model.config, "_attn_implementation", None)!r}') + f"[TransformersModel] attn_implementation: requested={kwargs.get('attn_implementation')!r}, " + f"resolved={getattr(self.model.config, '_attn_implementation', None)!r}" + ) # Gated worker-side so the flag is honored under Ray too, where a driver-side post-construction # reach into ``model.model`` only touches the proxy stub and silently no-ops. if gradient_checkpointing: self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs) - vision_tower_names = self.model_arch.get('vision_tower') + vision_tower_names = self.model_arch.get("vision_tower") if vit_gradient_checkpointing is not None and vision_tower_names: from twinkle.utils import deep_getattr + for tower_name in vision_tower_names: tower = deep_getattr(self.model, tower_name) if tower is None: @@ -448,13 +466,14 @@ def __init__( if vit_gradient_checkpointing: if vit_gradient_checkpointing_kwargs: tower.gradient_checkpointing_enable( - gradient_checkpointing_kwargs=vit_gradient_checkpointing_kwargs) + gradient_checkpointing_kwargs=vit_gradient_checkpointing_kwargs + ) else: tower.gradient_checkpointing_enable() else: tower.gradient_checkpointing_disable() except (NotImplementedError, AttributeError, ValueError) as exc: - logger.warning(f'vision-tower gradient checkpointing skipped for {tower_name}: {exc}') + logger.warning(f"vision-tower gradient checkpointing skipped for {tower_name}: {exc}") if neftune_noise_alpha: self._activate_neftune(neftune_noise_alpha) @@ -471,7 +490,7 @@ def __init__( self._apply_init_strategy(init_strategy) self.sp_strategy = None self._model_wrapped = False - self.optimizer_group: Dict[str, OptimizerGroup] = { + self.optimizer_group: dict[str, OptimizerGroup] = { _default_adapter_name: self._construct_default_optimizer_group() } self.optimizer_group[_default_adapter_name].adapter_name = _default_adapter_name @@ -490,8 +509,7 @@ def _activate_neftune(self, neftune_noise_alpha: float) -> None: """ embeddings = self.model.get_input_embeddings() if embeddings is None: - logger.warning('neftune_noise_alpha is set but the model exposes no input embeddings; ' - 'skipping NEFTune.') + logger.warning("neftune_noise_alpha is set but the model exposes no input embeddings; skipping NEFTune.") return def neftune_post_forward_hook(module, input, output): @@ -503,7 +521,7 @@ def neftune_post_forward_hook(module, input, output): embeddings.register_forward_hook(neftune_post_forward_hook) - def _apply_liger_kernel(self, liger_kernel_config: Optional[Dict[str, Any]] = None) -> None: + def _apply_liger_kernel(self, liger_kernel_config: dict[str, Any] | None = None) -> None: """Patch the HF module's per-layer ops with Liger kernels via the mapping-driven ``kernelize``. ``liger_kernel_config`` (an ``{op: bool}`` dict) switches individual ops off on top of @@ -512,10 +530,11 @@ def _apply_liger_kernel(self, liger_kernel_config: Optional[Dict[str, Any]] = No before the optimizer is built. """ from twinkle.kernel import kernelize + kernelize(self.model, self._liger_kernel_mapping(liger_kernel_config)) @staticmethod - def _liger_kernel_mapping(liger_kernel_config: Optional[Dict[str, Any]]) -> Optional[dict]: + def _liger_kernel_mapping(liger_kernel_config: dict[str, Any] | None) -> dict | None: """DEFAULT_KERNEL_CONFIG minus any op the config switches off; None keeps the built-in default. Returning None rather than a copy of the default matters: ``kernelize`` logs @@ -526,14 +545,15 @@ def _liger_kernel_mapping(liger_kernel_config: Optional[Dict[str, Any]]) -> Opti if not liger_kernel_config: return None from twinkle.kernel.config import DEFAULT_KERNEL_CONFIG + disabled = {op for op, enabled in liger_kernel_config.items() if not enabled} return { target: choice for target, choice in DEFAULT_KERNEL_CONFIG.items() - if getattr(choice, 'op', None) not in disabled + if getattr(choice, "op", None) not in disabled } - def _apply_new_special_tokens(self, new_special_tokens: List[str]) -> None: + def _apply_new_special_tokens(self, new_special_tokens: list[str]) -> None: """Grow the tokenizer's vocabulary with ``new_special_tokens`` and resize the embeddings to match. ``new_special_tokens`` is an already-resolved token list -- any ``.txt`` file expansion happens on @@ -545,10 +565,11 @@ def _apply_new_special_tokens(self, new_special_tokens: List[str]) -> None: processor = self._default_tokenizer if processor is None: from transformers import AutoTokenizer + processor = AutoTokenizer.from_pretrained(self.tokenizer_id, trust_remote_code=True) self._default_tokenizer = processor - tokenizer = getattr(processor, 'tokenizer', processor) - added = tokenizer.add_special_tokens({'additional_special_tokens': list(new_special_tokens)}) + tokenizer = getattr(processor, "tokenizer", processor) + added = tokenizer.add_special_tokens({"additional_special_tokens": list(new_special_tokens)}) if added: vocab_size = ((len(tokenizer) + 127) // 128) * 128 self.model.resize_token_embeddings(vocab_size) @@ -569,20 +590,21 @@ def high_dim(param, init): nn.init.zeros_(param) initializers = { - 'zero': nn.init.zeros_, - 'uniform': lambda param: nn.init.uniform_(param, -0.1, 0.1), - 'normal': lambda param: nn.init.normal_(param, 0.0, 0.01), - 'xavier_uniform': lambda param: high_dim(param, nn.init.xavier_uniform_), - 'xavier_normal': lambda param: high_dim(param, nn.init.xavier_normal_), - 'kaiming_uniform': lambda param: high_dim( - param, lambda value: nn.init.kaiming_uniform_(value, mode='fan_out', nonlinearity='leaky_relu', - a=0.1)), - 'kaiming_normal': lambda param: high_dim( - param, lambda value: nn.init.kaiming_normal_(value, mode='fan_in', nonlinearity='relu')), - 'orthogonal': lambda param: high_dim(param, nn.init.orthogonal_), + "zero": nn.init.zeros_, + "uniform": lambda param: nn.init.uniform_(param, -0.1, 0.1), + "normal": lambda param: nn.init.normal_(param, 0.0, 0.01), + "xavier_uniform": lambda param: high_dim(param, nn.init.xavier_uniform_), + "xavier_normal": lambda param: high_dim(param, nn.init.xavier_normal_), + "kaiming_uniform": lambda param: high_dim( + param, lambda value: nn.init.kaiming_uniform_(value, mode="fan_out", nonlinearity="leaky_relu", a=0.1) + ), + "kaiming_normal": lambda param: high_dim( + param, lambda value: nn.init.kaiming_normal_(value, mode="fan_in", nonlinearity="relu") + ), + "orthogonal": lambda param: high_dim(param, nn.init.orthogonal_), } if init_strategy not in initializers: - raise ValueError(f'Unknown init_strategy={init_strategy!r}; expected one of {sorted(initializers)}.') + raise ValueError(f"Unknown init_strategy={init_strategy!r}; expected one of {sorted(initializers)}.") initialize = initializers[init_strategy] with torch.no_grad(): for param in self.model.parameters(): @@ -594,12 +616,12 @@ def high_dim(param, init): initialize(param) def _should_init_empty_pretrained_model_on_this_rank(self) -> bool: - use_rank0_broadcast = getattr(self.strategy, 'use_rank0_pretrained_broadcast', lambda: False) + use_rank0_broadcast = getattr(self.strategy, "use_rank0_pretrained_broadcast", lambda: False) if not (use_rank0_broadcast() and dist.is_available() and dist.is_initialized()): return False local_rank = Platform.get_local_rank() if local_rank < 0: - raise RuntimeError('Native FSDP memory_efficient_init requires LOCAL_RANK.') + raise RuntimeError("Native FSDP memory_efficient_init requires LOCAL_RANK.") return local_rank != 0 def _init_empty_model_from_config(self, model_cls, **kwargs): @@ -607,40 +629,42 @@ def _init_empty_model_from_config(self, model_cls, **kwargs): config_kwargs = filter_from_config_kwargs(kwargs) with init_empty_weights(include_buffers=False): - if hasattr(model_cls, 'from_config'): + if hasattr(model_cls, "from_config"): model = model_cls.from_config(self.hf_config, **config_kwargs) else: model = model_cls._from_config(self.hf_config, **config_kwargs) - if hasattr(model, 'tie_weights'): + if hasattr(model, "tie_weights"): model.tie_weights() return model - def _decide_strategy(self, - strategy: Literal['accelerate', 'native_fsdp', 'deepspeed'], - deepspeed_config: Dict[str, Any] = None): + def _decide_strategy( + self, strategy: Literal["accelerate", "native_fsdp", "deepspeed"], deepspeed_config: dict[str, Any] = None + ): if self.device_mesh is not None and self.device_mesh.pp_world_size > 1: raise ValueError( - f'The transformers backend has no pipeline-parallel runtime (pp_world_size=' - f'{self.device_mesh.pp_world_size}): its strategies wire dp/fsdp/tp/cp/sp only, and transformers ' - '5.x ships pp_plan as a declaration layer with no training scheduler. Run pp>1 with ' - '--backend megatron.') - self._expert_parallel_config = self._fsdp_config.pop('expert_parallel', None) - self._enable_expert_parallel = self._should_enable_expert_parallel(self._expert_parallel_config, - self.device_mesh) + f"The transformers backend has no pipeline-parallel runtime (pp_world_size=" + f"{self.device_mesh.pp_world_size}): its strategies wire dp/fsdp/tp/cp/sp only, and transformers " + "5.x ships pp_plan as a declaration layer with no training scheduler. Run pp>1 with " + "--backend megatron." + ) + self._expert_parallel_config = self._fsdp_config.pop("expert_parallel", None) + self._enable_expert_parallel = self._should_enable_expert_parallel( + self._expert_parallel_config, self.device_mesh + ) self._expert_parallel_applied = False - if strategy == 'deepspeed': + if strategy == "deepspeed": self.strategy = DeepSpeedStrategy( mixed_precision=self.mixed_precision, deepspeed_config=deepspeed_config, device_mesh=self.device_mesh, ) else: - use_native_fsdp = self._enable_expert_parallel or strategy == 'native_fsdp' + use_native_fsdp = self._enable_expert_parallel or strategy == "native_fsdp" if use_native_fsdp: - ep_size = (self._expert_parallel_config.get('ep_size') if self._expert_parallel_config else None) + ep_size = self._expert_parallel_config.get("ep_size") if self._expert_parallel_config else None if ep_size is None and self.device_mesh is not None: - ep_size = getattr(self.device_mesh, 'ep_size', None) + ep_size = getattr(self.device_mesh, "ep_size", None) self.strategy = NativeFSDPStrategy( mixed_precision=self.mixed_precision, fsdp_config=self._fsdp_config, @@ -655,17 +679,18 @@ def _decide_strategy(self, ddp_config=self._ddp_config, fsdp_config=self._fsdp_config, device_mesh=self.device_mesh, - memory_efficient_init=self._memory_efficient_init) + memory_efficient_init=self._memory_efficient_init, + ) # Sequence parallel ("ulysses") is derived from dp/fsdp ranks; it does not change world size. # We construct `sp_strategy` after the underlying HF model is initialized (see __init__). self._enable_sp = False if self.device_mesh is not None: - sp_size = getattr(self.device_mesh, 'ulysses_size', None) + sp_size = getattr(self.device_mesh, "ulysses_size", None) self._enable_sp = bool(sp_size and sp_size > 1) def _ensure_sp_strategy(self) -> None: - if not getattr(self, '_enable_sp', False): + if not getattr(self, "_enable_sp", False): return if self.sp_strategy is not None: return @@ -681,13 +706,13 @@ def _ensure_sp_strategy(self) -> None: def _ensure_gdn_padding_free_patch(self) -> None: """Install the non-SP GatedDeltaNet patch before wrapping the model.""" self._gdn_padding_free_enabled = False - if getattr(self, '_enable_sp', False): + if getattr(self, "_enable_sp", False): return - if not getattr(self.model, '_twinkle_gdn_padding_free_patched', False): + if not getattr(self.model, "_twinkle_gdn_padding_free_patched", False): from twinkle.patch.gdn_padding_free import GatedDeltaNetPaddingFreePatch apply_patch(self.model, GatedDeltaNetPaddingFreePatch, hf_config=self.hf_config, enable_sp=False) - self._gdn_padding_free_enabled = bool(getattr(self.model, '_twinkle_gdn_padding_free_patched', False)) + self._gdn_padding_free_enabled = bool(getattr(self.model, "_twinkle_gdn_padding_free_patched", False)) def _get_default_group(self): """Get the only group, else return the default one""" @@ -698,7 +723,7 @@ def _get_default_group(self): def __repr__(self) -> str: # nn.Module.__repr__ dumps the full tree; RayTaskError interpolates actor repr # into the exception string and that dump swallows the real traceback. - return f'{type(self).__name__}(model_id={getattr(self, "model_id", None)!r})' + return f"{type(self).__name__}(model_id={getattr(self, 'model_id', None)!r})" def _lazy_wrap_model(self): if not self._model_wrapped: @@ -731,7 +756,7 @@ def register_mm_forward_hook(self, optimizer_group: OptimizerGroup): assert template is not None def forward_hook(model, args, kwargs): - selected_group = getattr(self, '_forward_optimizer_group', None) + selected_group = getattr(self, "_forward_optimizer_group", None) if selected_group is None: selected_group = self.optimizer_group.get(self._get_default_group()) if selected_group is optimizer_group: @@ -741,7 +766,7 @@ def forward_hook(model, args, kwargs): @contextlib.contextmanager def _mm_forward_context(self, optimizer_group: OptimizerGroup): - previous = getattr(self, '_forward_optimizer_group', None) + previous = getattr(self, "_forward_optimizer_group", None) self._forward_optimizer_group = optimizer_group try: yield @@ -754,22 +779,23 @@ def unregister_mm_forward_hook(self, optimizer_group: OptimizerGroup): optimizer_group._handler = None @staticmethod - def _should_enable_expert_parallel(expert_parallel_config: Optional[Dict[str, Any]], - device_mesh: Optional[DeviceMesh]) -> bool: + def _should_enable_expert_parallel( + expert_parallel_config: dict[str, Any] | None, device_mesh: DeviceMesh | None + ) -> bool: if expert_parallel_config is None or device_mesh is None: return False # Check ep_size from config first, then from device_mesh.ep_size attribute - ep_size = expert_parallel_config.get('ep_size') or getattr(device_mesh, 'ep_size', None) or 1 + ep_size = expert_parallel_config.get("ep_size") or getattr(device_mesh, "ep_size", None) or 1 if ep_size <= 1: return False - return expert_parallel_config.get('enabled', True) + return expert_parallel_config.get("enabled", True) def _maybe_apply_expert_parallel(self): if not self._enable_expert_parallel or self._expert_parallel_applied: return self._ensure_optimizer_dp_groups() model = self.strategy.unwrap_model(self.model) - ep_fsdp_mesh = getattr(self.strategy, 'ep_fsdp_device_mesh', None) + ep_fsdp_mesh = getattr(self.strategy, "ep_fsdp_device_mesh", None) apply_expert_parallel( model, self.device_mesh, @@ -784,18 +810,19 @@ def _maybe_apply_router_replay(self): return model = self.strategy.unwrap_model(self.model) from .moe.router_replay import RouterReplay + RouterReplay.apply_router_replay_patch(model) self._router_replay_applied = True - def _router_replay(self, action, *, routed_experts=None, batch_size=1, manual_cleanup=False, - replay_mask=None): + def _router_replay(self, action, *, routed_experts=None, batch_size=1, manual_cleanup=False, replay_mask=None): """Set up routing replay before a model forward; returns the per-call :class:`RouterReplay`. A disabled/absent action yields the null object (all methods no-op); call ``finish()`` after the forward to gather recorded indices (RECORD) and clear global state. """ from .moe.router_replay import RouterReplay - if not getattr(self, '_router_replay_enabled', False) or action is None: + + if not getattr(self, "_router_replay_enabled", False) or action is None: return RouterReplay(action=None) self._maybe_apply_router_replay() return RouterReplay( @@ -819,7 +846,7 @@ def _ensure_optimizer_dp_groups(self): def _construct_default_optimizer_group(self): return OptimizerGroup( - loss_instance=CrossEntropyLoss(reduction='sum'), + loss_instance=CrossEntropyLoss(reduction="sum"), template=Template(self.tokenizer_id), processor=InputProcessor(self.device_mesh), _device_mesh=self.device_mesh, @@ -828,29 +855,30 @@ def _construct_default_optimizer_group(self): def _restore_loss_inputs(self, inputs, labels, loss_scale=None, channel=None, completion_mask=None): """Put the loss-related tensors popped before the model call back onto ``inputs``.""" - inputs['labels'] = labels + inputs["labels"] = labels if loss_scale is not None: - inputs['loss_scale'] = loss_scale + inputs["loss_scale"] = loss_scale if channel is not None: - inputs['channel'] = channel + inputs["channel"] = channel if completion_mask is not None: - inputs['completion_mask'] = completion_mask + inputs["completion_mask"] = completion_mask - def _postprocess_forward_outputs(self, outputs, *, labels, task, temperature, sampling_masks, return_logits, - loss_instance, replay=None): + def _postprocess_forward_outputs( + self, outputs, *, labels, task, temperature, sampling_masks, return_logits, loss_instance, replay=None + ): """Shared forward/forward_only tail: derive logps/entropy/values, then reclaim logits memory. ``replay`` is the ``(loss_mask, masked_labels, metadata)`` tuple from ``_prepare_sampling_replay`` when sampling replay is enabled. """ - loss_require_logits = getattr(loss_instance, 'require_logits', False) - loss_require_entropy = getattr(loss_instance, 'require_entropy', False) - loss_require_logps = getattr(loss_instance, 'require_logps', True) + loss_require_logits = getattr(loss_instance, "require_logits", False) + loss_require_entropy = getattr(loss_instance, "require_entropy", False) + loss_require_logps = getattr(loss_instance, "require_logps", True) if task not in _NON_LOGPS_TASKS and labels is not None and loss_require_logps: if replay is not None: loss_mask, masked_labels, replay_metadata = replay - outputs['logps'] = replayed_selective_log_softmax( - logits=outputs['logits'], + outputs["logps"] = replayed_selective_log_softmax( + logits=outputs["logits"], labels=masked_labels, loss_mask=loss_mask, sampling_masks=sampling_masks, @@ -860,23 +888,36 @@ def _postprocess_forward_outputs(self, outputs, *, labels, task, temperature, sa else: loss_mask = (labels != -100).bool() masked_labels = labels.masked_fill(~loss_mask, 0) - logits = outputs['logits'] + logits = outputs["logits"] logits.div_(temperature) if loss_require_entropy: - outputs['logps'], outputs['entropies'] = selective_log_softmax( - logits, masked_labels, return_entropy=True) + outputs["logps"], outputs["entropies"] = selective_log_softmax( + logits, masked_labels, return_entropy=True + ) else: - outputs['logps'] = selective_log_softmax(logits, masked_labels) + outputs["logps"] = selective_log_softmax(logits, masked_labels) del logits - if getattr(loss_instance, 'require_values', False): - values = outputs['logits'] - outputs['values'] = values.squeeze(-1) if values.shape[-1] == 1 else values - outputs['past_key_values'] = None + if getattr(loss_instance, "require_values", False): + values = outputs["logits"] + outputs["values"] = values.squeeze(-1) if values.shape[-1] == 1 else values + outputs["past_key_values"] = None if not (return_logits or loss_require_logits): - outputs['logits'] = None + outputs["logits"] = None - def _finalize_forward_outputs(self, inputs, outputs, *, processor, task, optimizer_config, is_training, - forward_kwargs, loss_instance, return_logits, recorded_routing): + def _finalize_forward_outputs( + self, + inputs, + outputs, + *, + processor, + task, + optimizer_config, + is_training, + forward_kwargs, + loss_instance, + return_logits, + recorded_routing, + ): """Shared forward/forward_only tail: SP postprocess, unpack, status recording, and the return copy.""" inputs, outputs = processor.postprocess_tensor_sp(inputs, outputs, sp_strategy=self.sp_strategy, task=task) inputs, outputs = processor.unpack_packed_sequences(inputs, outputs, task=task) @@ -885,17 +926,17 @@ def _finalize_forward_outputs(self, inputs, outputs, *, processor, task, optimiz status.outputs = outputs status.forward_kwargs = forward_kwargs if is_training: - status.loss_reduction = getattr(loss_instance, 'reduction', 'mean') - status.loss_value = outputs.get('aux_loss', 0) + status.loss_reduction = getattr(loss_instance, "reduction", "mean") + status.loss_value = outputs.get("aux_loss", 0) return_outputs = copy(outputs) if not return_logits: - return_outputs['logits'] = None + return_outputs["logits"] = None if recorded_routing is not None: - return_outputs['routed_experts'] = recorded_routing + return_outputs["routed_experts"] = recorded_routing return return_outputs - @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) - def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajectory]], **kwargs): + @remote_function(dispatch="slice_dp", collect=collect_tensor_dict) + def forward(self, *, inputs: InputFeature | list[InputFeature] | list[Trajectory], **kwargs): """Call forward function and record the inputs and outputs. Args: @@ -905,47 +946,49 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec Returns: The output of the model forward. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) - temperature = float(kwargs.pop('temperature', 1.0)) - sampling_masks = kwargs.pop('sampling_masks', None) - return_logits = kwargs.pop('return_logits', False) - task = kwargs.pop('task', 'causal_lm') - loss_key = kwargs.pop('loss_key', None) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) + temperature = float(kwargs.pop("temperature", 1.0)) + sampling_masks = kwargs.pop("sampling_masks", None) + return_logits = kwargs.pop("return_logits", False) + task = kwargs.pop("task", "causal_lm") + loss_key = kwargs.pop("loss_key", None) optimizer_config = self.optimizer_group[adapter_name] self._lazy_wrap_model() if not inputs: - raise ValueError('inputs empty, check your DataLoader outputs') + raise ValueError("inputs empty, check your DataLoader outputs") self.model.train() - if (isinstance(inputs, dict) and self._not_encoded(inputs)) or (isinstance(inputs, list) - and self._not_encoded(inputs[0])): + if (isinstance(inputs, dict) and self._not_encoded(inputs)) or ( + isinstance(inputs, list) and self._not_encoded(inputs[0]) + ): # Trajectory or List[Trajectory] - assert optimizer_config.template is not None, \ - 'Use set_template to add a template when trying to input `List[Trajectory]`' + assert optimizer_config.template is not None, ( + "Use set_template to add a template when trying to input `List[Trajectory]`" + ) if isinstance(inputs, dict): inputs = [inputs] inputs = optimizer_config.template.batch_encode(inputs) # noqa processor: InputProcessor = optimizer_config.processor loss_instance = optimizer_config.get_loss(loss_key) - enable_sampling_replay = getattr(loss_instance, 'enable_sampling_replay', False) + enable_sampling_replay = getattr(loss_instance, "enable_sampling_replay", False) if enable_sampling_replay: if sampling_masks is None: - raise ValueError('sampling_masks are required when sampling replay is enabled') + raise ValueError("sampling_masks are required when sampling replay is enabled") cp_world_size = self.device_mesh.cp_world_size if self.device_mesh is not None else 1 - if getattr(self, '_enable_sp', False) or cp_world_size > 1: - raise ValueError('sampling replay does not support sequence or context parallelism') - assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' - inputs: Dict[str, Any] = processor( + if getattr(self, "_enable_sp", False) or cp_world_size > 1: + raise ValueError("sampling replay does not support sequence or context parallelism") + assert isinstance(processor, InputProcessor), "Set a correct `InputProcessor` before forwarding" + inputs: dict[str, Any] = processor( inputs, sp_strategy=self.sp_strategy, hf_config=self.hf_config, - enable_sp=getattr(self, '_enable_sp', False), + enable_sp=getattr(self, "_enable_sp", False), gdn_padding_free_enabled=self._gdn_padding_free_enabled, ) - labels: torch.Tensor = inputs.pop('labels', None) - loss_scale = inputs.pop('loss_scale', None) - channel = inputs.pop('channel', None) + labels: torch.Tensor = inputs.pop("labels", None) + loss_scale = inputs.pop("loss_scale", None) + channel = inputs.pop("channel", None) # Not a model argument; the loss reads it back off `inputs` further down. - completion_mask = inputs.pop('completion_mask', None) + completion_mask = inputs.pop("completion_mask", None) replay_metadata = replay_loss_mask = replay_masked_labels = None if enable_sampling_replay: replay_loss_mask, replay_masked_labels, replay_metadata = _prepare_sampling_replay( @@ -953,32 +996,41 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec sampling_masks=sampling_masks, temperature=temperature, vocab_size=_get_vocab_size(self.hf_config), - allow_packed_masks=(processor.padding_free - or processor._is_packed_position_ids(inputs.get('position_ids'))), + allow_packed_masks=( + processor.padding_free or processor._is_packed_position_ids(inputs.get("position_ids")) + ), ) optimizer_config.accumulate_metrics(True) # Routing replay: the driver passes the MODE scalar; the model layer owns the mode->action mapping # (resolve_router_replay_mode -- the training forward always REPLAYS, so a missing routed_experts # still fails loudly in RouterReplay.__init__). - router_replay_mode = kwargs.pop('router_replay_mode', None) - router_replay_manual_cleanup = kwargs.pop('router_replay_manual_cleanup', False) - routed_experts = inputs.pop('routed_experts', None) - replay_mask = inputs.pop('replay_mask', None) - batch_size = labels.shape[0] if labels is not None else ( - inputs['input_ids'].shape[0] if 'input_ids' in inputs else 1) + router_replay_mode = kwargs.pop("router_replay_mode", None) + router_replay_manual_cleanup = kwargs.pop("router_replay_manual_cleanup", False) + routed_experts = inputs.pop("routed_experts", None) + replay_mask = inputs.pop("replay_mask", None) + batch_size = ( + labels.shape[0] if labels is not None else (inputs["input_ids"].shape[0] if "input_ids" in inputs else 1) + ) from .moe.router_replay import RouterReplay + rr = self._router_replay( - RouterReplay.resolve_router_replay_mode(router_replay_mode), routed_experts=routed_experts, - batch_size=batch_size, manual_cleanup=router_replay_manual_cleanup, replay_mask=replay_mask) + RouterReplay.resolve_router_replay_mode(router_replay_mode), + routed_experts=routed_experts, + batch_size=batch_size, + manual_cleanup=router_replay_manual_cleanup, + replay_mask=replay_mask, + ) # Per-forward task patches probe module attributes (score/classifier heads, # get_image_features) that DDP-wrapped models do not forward -- resolve them on the # unwrapped module. Hooks registered there still fire: the wrapper's forward calls # the inner module's __call__. unwrapped_model = self.strategy.unwrap_model(self.model) - with _resolve_task_context(unwrapped_model, task, template=optimizer_config.template), \ - _resolve_mm_graph_context(unwrapped_model, optimizer_config.template): + with ( + _resolve_task_context(unwrapped_model, task, template=optimizer_config.template), + _resolve_mm_graph_context(unwrapped_model, optimizer_config.template), + ): outputs = self.model(**inputs) recorded_routing = rr.finish() @@ -1001,7 +1053,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec task=task, optimizer_config=optimizer_config, is_training=True, - forward_kwargs={**kwargs, '_metric_loss_key': loss_key}, + forward_kwargs={**kwargs, "_metric_loss_key": loss_key}, loss_instance=loss_instance, return_logits=return_logits, recorded_routing=recorded_routing, @@ -1010,9 +1062,9 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec # lazy_collect follows the global default (True): a caller that reads the output itself calls # ``.result()`` first, while a ref-only pipeline hands the handle straight to a downstream remote # call, which resolves and collects it where the data is consumed (see ``LazyCollect``). - @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) + @remote_function(dispatch="slice_dp", collect=collect_tensor_dict) @timed - def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajectory]], **kwargs): + def forward_only(self, *, inputs: InputFeature | list[InputFeature] | list[Trajectory], **kwargs): """Call forward function without grad and record the inputs and outputs. Args: @@ -1028,26 +1080,28 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T The output of the model forward. """ config = MicroBatchConfig.from_kwargs(kwargs) - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] inputs = self._prepare_forward_only_inputs(inputs, optimizer_config) if config is None: return self._forward_only_single(inputs=inputs, optimizer_config=optimizer_config, **kwargs) return self._forward_only_micro_batches( - inputs=inputs, config=config, optimizer_config=optimizer_config, **kwargs) + inputs=inputs, config=config, optimizer_config=optimizer_config, **kwargs + ) def _prepare_forward_only_inputs(self, inputs, optimizer_config): """Normalize and encode the local batch once, before optional micro-batching.""" if isinstance(inputs, dict): inputs = [inputs] if not inputs: - raise ValueError('inputs empty, check your DataLoader outputs') + raise ValueError("inputs empty, check your DataLoader outputs") if self._not_encoded(inputs[0]): - assert optimizer_config.template is not None, \ - 'Use set_template to add a template when trying to input `List[Trajectory]`' + assert optimizer_config.template is not None, ( + "Use set_template to add a template when trying to input `List[Trajectory]`" + ) inputs = optimizer_config.template.batch_encode(inputs) # noqa if not inputs: - raise ValueError('encoded inputs empty, check your template outputs') + raise ValueError("encoded inputs empty, check your template outputs") return inputs def _forward_only_micro_batches(self, *, inputs, config, optimizer_config, **kwargs): @@ -1070,52 +1124,53 @@ def _forward_only_micro_batches(self, *, inputs, config, optimizer_config, **kwa ) batches.append(dict(outputs)) micro_batch_scales = None - if any(output.get('aux_loss') is not None for output in batches): - loss_instance = optimizer_config.get_loss(kwargs.get('loss_key')) + if any(output.get("aux_loss") is not None for output in batches): + loss_instance = optimizer_config.get_loss(kwargs.get("loss_key")) micro_batch_scales = [loss_instance.micro_batch_scale(inputs, indices) for indices in plan] return merge_micro_batch_outputs( - batches, plan, task=kwargs.get('task', 'causal_lm'), micro_batch_scales=micro_batch_scales) + batches, plan, task=kwargs.get("task", "causal_lm"), micro_batch_scales=micro_batch_scales + ) def _forward_only_single(self, *, inputs, optimizer_config, **kwargs): """Run one evaluation micro-batch (or, without micro-batching, the whole batch).""" - disable_lora = kwargs.pop('disable_lora', False) - temperature = float(kwargs.pop('temperature', 1.0)) - sampling_masks = kwargs.pop('sampling_masks', None) - return_logits = kwargs.pop('return_logits', False) - task = kwargs.pop('task', 'causal_lm') - loss_key = kwargs.pop('loss_key', None) + disable_lora = kwargs.pop("disable_lora", False) + temperature = float(kwargs.pop("temperature", 1.0)) + sampling_masks = kwargs.pop("sampling_masks", None) + return_logits = kwargs.pop("return_logits", False) + task = kwargs.pop("task", "causal_lm") + loss_key = kwargs.pop("loss_key", None) # Accepted for backend parity with Megatron's empty-loss probe entry (``compute_loss=False``). This # forward_only derives logps via selective_log_softmax and never calls the loss, so there is nothing # to skip here; popped so it does not linger in eval_status.forward_kwargs. - kwargs.pop('compute_loss', None) + kwargs.pop("compute_loss", None) self._lazy_wrap_model() self.model.eval() with torch.no_grad(): processor: InputProcessor = optimizer_config.processor - assert isinstance(processor, InputProcessor), 'Set InputProcessor correctly before forwarding' + assert isinstance(processor, InputProcessor), "Set InputProcessor correctly before forwarding" loss_instance = optimizer_config.get_loss(loss_key) - enable_sampling_replay = getattr(loss_instance, 'enable_sampling_replay', False) + enable_sampling_replay = getattr(loss_instance, "enable_sampling_replay", False) if enable_sampling_replay: if sampling_masks is None: - raise ValueError('sampling_masks are required when sampling replay is enabled') + raise ValueError("sampling_masks are required when sampling replay is enabled") cp_world_size = self.device_mesh.cp_world_size if self.device_mesh is not None else 1 - if getattr(self, '_enable_sp', False) or cp_world_size > 1: - raise ValueError('sampling replay does not support sequence or context parallelism') - inputs: Dict[str, Any] = processor( + if getattr(self, "_enable_sp", False) or cp_world_size > 1: + raise ValueError("sampling replay does not support sequence or context parallelism") + inputs: dict[str, Any] = processor( inputs, sp_strategy=self.sp_strategy, hf_config=self.hf_config, - enable_sp=getattr(self, '_enable_sp', False), + enable_sp=getattr(self, "_enable_sp", False), gdn_padding_free_enabled=self._gdn_padding_free_enabled, ) - labels = inputs.pop('labels', None) - loss_scale = inputs.pop('loss_scale', None) - channel = inputs.pop('channel', None) + labels = inputs.pop("labels", None) + loss_scale = inputs.pop("loss_scale", None) + channel = inputs.pop("channel", None) # Not a model argument; the loss reads it back off `inputs` further down. - completion_mask = inputs.pop('completion_mask', None) + completion_mask = inputs.pop("completion_mask", None) replay_metadata = replay_loss_mask = replay_masked_labels = None if enable_sampling_replay: - packed_position_ids = processor._is_packed_position_ids(inputs.get('position_ids')) + packed_position_ids = processor._is_packed_position_ids(inputs.get("position_ids")) replay_loss_mask, replay_masked_labels, replay_metadata = _prepare_sampling_replay( labels=labels, sampling_masks=sampling_masks, @@ -1128,21 +1183,30 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs): # Routing replay: the driver passes the MODE scalar; allow_record makes an R2 capture pass # (a forward_only whose inputs carry no routing yet) RECORD here. - router_replay_mode = kwargs.pop('router_replay_mode', None) - router_replay_manual_cleanup = kwargs.pop('router_replay_manual_cleanup', False) - routed_experts = inputs.pop('routed_experts', None) - replay_mask = inputs.pop('replay_mask', None) - batch_size = labels.shape[0] if labels is not None else ( - inputs['input_ids'].shape[0] if 'input_ids' in inputs else 1) + router_replay_mode = kwargs.pop("router_replay_mode", None) + router_replay_manual_cleanup = kwargs.pop("router_replay_manual_cleanup", False) + routed_experts = inputs.pop("routed_experts", None) + replay_mask = inputs.pop("replay_mask", None) + batch_size = ( + labels.shape[0] + if labels is not None + else (inputs["input_ids"].shape[0] if "input_ids" in inputs else 1) + ) from .moe.router_replay import RouterReplay + rr = self._router_replay( RouterReplay.resolve_router_replay_mode(router_replay_mode, allow_record=True), - routed_experts=routed_experts, batch_size=batch_size, - manual_cleanup=router_replay_manual_cleanup, replay_mask=replay_mask) + routed_experts=routed_experts, + batch_size=batch_size, + manual_cleanup=router_replay_manual_cleanup, + replay_mask=replay_mask, + ) lora_ctx = ( unwrapped_model.disable_adapter() - if disable_lora and isinstance(unwrapped_model, PeftModel) else contextlib.nullcontext()) + if disable_lora and isinstance(unwrapped_model, PeftModel) + else contextlib.nullcontext() + ) # unwrapped on purpose: task patches probe head/module attributes DDP does not forward. with _resolve_task_context(unwrapped_model, task, template=optimizer_config.template), lora_ctx: outputs = self.model(**inputs) @@ -1167,25 +1231,25 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs): task=task, optimizer_config=optimizer_config, is_training=False, - forward_kwargs={**kwargs, '_metric_loss_key': loss_key}, + forward_kwargs={**kwargs, "_metric_loss_key": loss_key}, loss_instance=loss_instance, return_logits=return_logits, recorded_routing=recorded_routing, ) - @remote_function(dispatch='slice_dp', collect=collect_dp_and_flatten, lazy_collect=False) + @remote_function(dispatch="slice_dp", collect=collect_dp_and_flatten, lazy_collect=False) def generate( self, - inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], - sampling_params: Optional[Union[SamplingParams, Dict[str, Any]]] = None, - adapter_name: str = '', - adapter_path: Optional[str] = None, + inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], + sampling_params: SamplingParams | dict[str, Any] | None = None, + adapter_name: str = "", + adapter_path: str | None = None, *, return_encoded: bool = False, use_base_model: bool = False, - adapter_paths: Optional[List[Optional[str]]] = None, + adapter_paths: list[str | None] | None = None, strict: bool = True, - ) -> List[SampleResponse]: + ) -> list[SampleResponse]: """Generate completions on the weights this model already holds, inside its own workers. This is the target a :class:`~twinkle.sampler.TransformersSampler` facade forwards to when it is @@ -1238,16 +1302,19 @@ def generate( # they are (or their base, via use_base_model); use a TransformersSampler(model_id=...) with # adapter_path to generate from a separate on-disk adapter. raise NotImplementedError( - 'TransformersModel.generate samples the weights this model already holds; it cannot load ' - 'an external adapter_path on top of a live training model. Drop adapter_path/adapter_paths ' - 'to sample the current (or, with use_base_model=True, the base) weights, or use a ' - 'TransformersSampler(model_id=...) for a separate adapter.') + "TransformersModel.generate samples the weights this model already holds; it cannot load " + "an external adapter_path on top of a live training model. Drop adapter_path/adapter_paths " + "to sample the current (or, with use_base_model=True, the base) weights, or use a " + "TransformersSampler(model_id=...) for a separate adapter." + ) adapter_name = adapter_name or self._get_default_group() template = self.optimizer_group[adapter_name].template if template is None: - raise ValueError('Set a template with set_template() before generate(): a Trajectory input ' - 'needs it to encode, and the engine needs its tokenizer to pad and decode.') + raise ValueError( + "Set a template with set_template() before generate(): a Trajectory input " + "needs it to encode, and the engine needs its tokenizer to pad and decode." + ) inputs_list = [inputs] if isinstance(inputs, dict) else list(inputs) if not inputs_list: return [] @@ -1268,12 +1335,13 @@ def generate( # FSDP2 fully_shard is in place, so self.model still exposes .generate and its all-gather hooks # must run -- generate on the wrapped model. A DDP/DeepSpeed wrapper has no .generate, so fall # back to the unwrapped module (the wrapper only exists for gradient sync, irrelevant here). - generation_target = self.model if hasattr(self.model, 'generate') else unwrapped_model + generation_target = self.model if hasattr(self.model, "generate") else unwrapped_model engine = TransformersEngine(self.tokenizer_id, model=generation_target, tokenizer=template.tokenizer) encoded, states, failures = generation.encode_all( - template, inputs_list, logprobs_only=logprobs_only, strict=strict) - results: List[Optional[SampleResponse]] = [None] * len(inputs_list) + template, inputs_list, logprobs_only=logprobs_only, strict=strict + ) + results: list[SampleResponse | None] = [None] * len(inputs_list) for index, response in failures.items(): results[index] = response pending = [index for index in range(len(inputs_list)) if results[index] is None] @@ -1283,22 +1351,28 @@ def call_batch(prompts, extra): # responsive and Ray forbids asyncio.run, so there is nothing to gain from batch_sample's # loop-and-thread hop (and its lock, which would be bound to the wrong loop). return engine.generate_batch( - prompts, params, num_samples=params.num_samples, logprobs=params.logprobs is not None, - extra_model_inputs=extra or None) + prompts, + params, + num_samples=params.num_samples, + logprobs=params.logprobs is not None, + extra_model_inputs=extra or None, + ) - lora_ctx = (unwrapped_model.disable_adapter() - if use_base_model and isinstance(unwrapped_model, PeftModel) else contextlib.nullcontext()) + lora_ctx = ( + unwrapped_model.disable_adapter() + if use_base_model and isinstance(unwrapped_model, PeftModel) + else contextlib.nullcontext() + ) with lora_ctx: for chunk in generation.chunks(pending, _GENERATION_MAX_BATCH_SIZE): for index, response in zip( - chunk, generation.run_chunk(chunk, encoded, states, call_batch, strict=strict)): + chunk, generation.run_chunk(chunk, encoded, states, call_batch, strict=strict), strict=False + ): results[index] = response if not logprobs_only: generation.attach_features(template, results, encoded, states) - return [ - r if r is not None else generation.error_response(encoded.get(i, {})) for i, r in enumerate(results) - ] + return [r if r is not None else generation.error_response(encoded.get(i, {})) for i, r in enumerate(results)] def generate_stream(self, *args, **kwargs): """Refuse streaming in-place generation rather than half-build it. @@ -1311,29 +1385,32 @@ def generate_stream(self, *args, **kwargs): consumed -- the same choice :meth:`MegatronModel.generate_stream` makes. """ raise NotImplementedError( - 'TransformersModel.generate does not stream: generation is sliced across data-parallel ' - 'workers and collected as a batch. Call generate() for the whole completion, or use a ' - 'TransformersSampler(model_id=...) whose sample_stream runs on a single engine.') + "TransformersModel.generate does not stream: generation is sliced across data-parallel " + "workers and collected as a batch. Call generate() for the whole completion, or use a " + "TransformersSampler(model_id=...) whose sample_stream runs on a single engine." + ) - @remote_function(collect='mean') + @remote_function(collect="mean") @timed - def calculate_loss(self, - loss_key: Optional[str] = None, - loss_weight: float = 1.0, - global_normalizer: Optional[torch.Tensor] = None, - **kwargs): + def calculate_loss( + self, + loss_key: str | None = None, + loss_weight: float = 1.0, + global_normalizer: torch.Tensor | None = None, + **kwargs, + ): """Calculate one selected loss, optionally pre-normalized for a multi-RPC transaction.""" - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] loss_instance: Loss = optimizer_config.get_loss(loss_key) - assert isinstance(loss_instance, Loss), 'Set a loss_instance before calculating loss' + assert isinstance(loss_instance, Loss), "Set a loss_instance before calculating loss" status = optimizer_config.train_status if self.model.training else optimizer_config.eval_status inputs = status.inputs outputs = status.outputs - assert inputs is not None and outputs is not None, 'Cannot calculate loss of empty inputs and outputs' + assert inputs is not None and outputs is not None, "Cannot calculate loss of empty inputs and outputs" result = loss_instance(inputs, outputs, **kwargs) - main_loss = result['loss'] - raw_counts = result['num_tokens'] + main_loss = result["loss"] + raw_counts = result["num_tokens"] counts = raw_counts if global_normalizer is not None: dp_world_size = self.device_mesh._get_dp_fsdp_world_size() if self.device_mesh is not None else 1 @@ -1350,40 +1427,41 @@ def calculate_loss(self, counts = counts / raw_dp_fsdp_world_size status.num_tokens += counts.item() status.loss_value = (status.loss_value + main_loss) * objective_scale - outputs['loss'] = status.loss_value - outputs['num_tokens'] = raw_counts.detach() if hasattr(raw_counts, 'detach') else raw_counts - status.loss_reduction = 'mean' if global_normalizer is not None else getattr(loss_instance, 'reduction', 'mean') - if result.get('channel_loss') is not None: - outputs['channel_loss'] = result['channel_loss'] + outputs["loss"] = status.loss_value + outputs["num_tokens"] = raw_counts.detach() if hasattr(raw_counts, "detach") else raw_counts + status.loss_reduction = "mean" if global_normalizer is not None else getattr(loss_instance, "reduction", "mean") + if result.get("channel_loss") is not None: + outputs["channel_loss"] = result["channel_loss"] return status.loss_value.item() @remote_function() - def backward(self, - increment_step: bool = True, - sync_gradients: Optional[bool] = None, - **kwargs): + def backward(self, increment_step: bool = True, sync_gradients: bool | None = None, **kwargs): """Backward propagation with explicit step and synchronization controls.""" - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] loss_value = optimizer_config.train_status.loss_value - assert loss_value is not None, 'Do forwarding and calculating loss before backward' + assert loss_value is not None, "Do forwarding and calculating loss before backward" scaler = optimizer_config.scaler - if scaler is None and self.mixed_precision == 'fp16': + if scaler is None and self.mixed_precision == "fp16": self.set_grad_scaler(adapter_name=adapter_name) scaler = optimizer_config.scaler if increment_step: optimizer_config.cur_step += 1 - should_sync = (optimizer_config.do_grad_sync(kwargs.get('gradient_accumulation_steps')) - if sync_gradients is None else bool(sync_gradients)) + should_sync = ( + optimizer_config.do_grad_sync(kwargs.get("gradient_accumulation_steps")) + if sync_gradients is None + else bool(sync_gradients) + ) no_sync_ctx = contextlib.nullcontext() - if not should_sync and hasattr(self.model, 'no_sync'): + if not should_sync and hasattr(self.model, "no_sync"): no_sync_ctx = self.model.no_sync() - router_replay_mode = kwargs.pop('router_replay_mode', None) + router_replay_mode = kwargs.pop("router_replay_mode", None) rr = None if router_replay_mode is not None: from .moe.router_replay import RouterReplay, RouterReplayAction + rr = self._router_replay(RouterReplayAction.REPLAY_BACKWARD) with no_sync_ctx: @@ -1404,7 +1482,7 @@ def _build_micro_batch_plan(self, inputs, config, optimizer_config): same number, and fails loudly when no common count exists. """ processor = optimizer_config.processor - assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' + assert isinstance(processor, InputProcessor), "Set a correct `InputProcessor` before forwarding" optimizer_config._ensure_dp_group() dp_group = optimizer_config._dp_group min_micro_batches = 1 @@ -1421,31 +1499,33 @@ def _build_micro_batch_plan(self, inputs, config, optimizer_config): if dp_group is None: raise plan = None - planning_error = f'{type(exc).__name__}: {exc}' + planning_error = f"{type(exc).__name__}: {exc}" if dp_group is None: return plan local_state = { - 'micro_batch_count': len(plan) if plan is not None else None, - 'input_count': len(inputs), - 'error': planning_error, + "micro_batch_count": len(plan) if plan is not None else None, + "input_count": len(inputs), + "error": planning_error, } states = [None] * dist.get_world_size(dp_group) dist.all_gather_object(states, local_state, group=dp_group) errors = [ - f'rank {rank}: {state["error"]}' for rank, state in enumerate(states) if state['error'] is not None + f"rank {rank}: {state['error']}" for rank, state in enumerate(states) if state["error"] is not None ] if errors: - raise RuntimeError('micro-batch planning failed on one or more model DP ranks: ' + '; '.join(errors)) + raise RuntimeError("micro-batch planning failed on one or more model DP ranks: " + "; ".join(errors)) - counts = [state['micro_batch_count'] for state in states] + counts = [state["micro_batch_count"] for state in states] if all(count == len(plan) for count in counts): return plan min_micro_batches = max(counts) - if any(min_micro_batches > state['input_count'] for state in states): - raise ValueError('model DP ranks cannot execute the same number of non-empty micro-batches; ' - 'make the input batch divisible by the model data-parallel size') + if any(min_micro_batches > state["input_count"] for state in states): + raise ValueError( + "model DP ranks cannot execute the same number of non-empty micro-batches; " + "make the input batch divisible by the model data-parallel size" + ) def _global_loss_normalizer(self, inputs, loss_instance, optimizer_config): local_normalizer = loss_instance.batch_normalizer(inputs) @@ -1490,7 +1570,7 @@ def _forward_backward_micro_batch( normalizer_delta = optimizer_config.train_status.num_tokens - previous_normalizer optimizer_config.train_status.loss_value = optimizer_config.train_status.loss_value * loss_scale optimizer_config.train_status.num_tokens = previous_normalizer + normalizer_delta * loss_scale - outputs['loss'] = loss * loss_scale + outputs["loss"] = loss * loss_scale self.backward( increment_step=increment_step, sync_gradients=sync_gradients, @@ -1531,8 +1611,8 @@ def _forward_backward_micro_batches( increment_step=increment_step and is_last_micro_batch, **micro_kwargs, ) - total_loss += float(outputs.get('loss', 0.0)) - outputs['loss'] = total_loss + total_loss += float(outputs.get("loss", 0.0)) + outputs["loss"] = total_loss return outputs def _forward_backward_single( @@ -1546,20 +1626,22 @@ def _forward_backward_single( increment_step, **kwargs, ): - adapter_name = kwargs.get('adapter_name', self._get_default_group()) + adapter_name = kwargs.get("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] if isinstance(inputs, dict): inputs = [inputs] if not inputs: - raise ValueError('inputs empty, check your DataLoader outputs') + raise ValueError("inputs empty, check your DataLoader outputs") if self._not_encoded(inputs[0]): - assert optimizer_config.template is not None, \ - 'Use set_template to add a template when trying to input `List[Trajectory]`' + assert optimizer_config.template is not None, ( + "Use set_template to add a template when trying to input `List[Trajectory]`" + ) inputs = optimizer_config.template.batch_encode(inputs) loss_instance = optimizer_config.get_loss(loss_key) global_normalizer = ( - self._global_loss_normalizer(inputs, loss_instance, optimizer_config) if normalize_globally else None) + self._global_loss_normalizer(inputs, loss_instance, optimizer_config) if normalize_globally else None + ) micro_batch_config = MicroBatchConfig.from_kwargs(kwargs) if micro_batch_config is not None: return self._forward_backward_micro_batches( @@ -1588,23 +1670,23 @@ def _forward_backward_single( **kwargs, ) - @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) + @remote_function(dispatch="slice_dp", collect=collect_tensor_dict) @timed def forward_backward( self, *, - inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], - loss_key: Optional[str] = None, + inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], + loss_key: str | None = None, loss_weight: float = 1.0, finalize_gradients: bool = True, **kwargs, ): """Run one homogeneous loss RPC and optionally finalize a two-RPC backward transaction.""" - adapter_name = kwargs.get('adapter_name', self._get_default_group()) + adapter_name = kwargs.get("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] completing_pending = optimizer_config.prepare_backward_transaction(finalize_gradients, loss_weight) transactional = not finalize_gradients or completing_pending - requested_sync = kwargs.pop('sync_gradients', None) + requested_sync = kwargs.pop("sync_gradients", None) outputs = self._forward_backward_single( inputs=inputs, loss_key=loss_key, @@ -1622,7 +1704,7 @@ def forward_backward( status = optimizer_config.train_status status.inputs = None status.outputs = None - total_loss = optimizer_config.record_transaction_loss(outputs.get('loss', 0.0), finalize_gradients) + total_loss = optimizer_config.record_transaction_loss(outputs.get("loss", 0.0), finalize_gradients) if not finalize_gradients: return outputs @@ -1631,12 +1713,12 @@ def forward_backward( status.inputs = {} status.outputs = ModelOutput(loss=total_loss) status.forward_kwargs = {} - status.loss_reduction = 'mean' + status.loss_reduction = "mean" return ModelOutput(loss=total_loss) @remote_function() def clip_grad_norm(self, max_grad_norm: float = 1.0, norm_type=2, **kwargs): - """ Clip the gradient norm + """Clip the gradient norm Args: max_grad_norm: The maximum grad norm, default `1.0`. @@ -1646,10 +1728,10 @@ def clip_grad_norm(self, max_grad_norm: float = 1.0, norm_type=2, **kwargs): Returns: Total norm of the parameter gradients (viewed as a single vector). """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.require_finalized('clip gradients') - if not optimizer_config.do_grad_sync(kwargs.get('gradient_accumulation_steps')): + optimizer_config.require_finalized("clip gradients") + if not optimizer_config.do_grad_sync(kwargs.get("gradient_accumulation_steps")): return optimizer = optimizer_config.optimizer @@ -1657,6 +1739,7 @@ def clip_grad_norm(self, max_grad_norm: float = 1.0, norm_type=2, **kwargs): context = contextlib.nullcontext if self.device_mesh is not None and self.device_mesh.tp_world_size > 1: from torch.distributed.tensor.experimental import implicit_replication + context = implicit_replication with context(): @@ -1669,8 +1752,9 @@ def clip_grad_norm(self, max_grad_norm: float = 1.0, norm_type=2, **kwargs): num_tokens = sum(num_tokens) parameters = list(self._get_trainable_parameters(adapter_name).values()) - ep_clip_kwargs = self.strategy.get_ep_clip_kwargs(self.model) if hasattr( - self.strategy, 'get_ep_clip_kwargs') else {} + ep_clip_kwargs = ( + self.strategy.get_ep_clip_kwargs(self.model) if hasattr(self.strategy, "get_ep_clip_kwargs") else {} + ) grad_norm = normalize_and_clip_grad_norm( parameters, @@ -1684,7 +1768,7 @@ def clip_grad_norm(self, max_grad_norm: float = 1.0, norm_type=2, **kwargs): optimizer_config.train_status.num_tokens = 0 return grad_norm - @remote_function(dispatch='all') + @remote_function(dispatch="all") @timed def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type=2, **kwargs): self.clip_grad_norm(max_grad_norm, norm_type, **kwargs) @@ -1692,49 +1776,54 @@ def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type=2, **kwargs): self.zero_grad(**kwargs) self.lr_step(**kwargs) - def _create_param_group(self, - adapter_name: str, - lr: float = DEFAULT_LEARNING_RATE, - weight_decay: float = DEFAULT_WEIGHT_DECAY, - **kwargs): + def _create_param_group( + self, adapter_name: str, lr: float = DEFAULT_LEARNING_RATE, weight_decay: float = DEFAULT_WEIGHT_DECAY, **kwargs + ): # Some code borrowed from transformers def get_parameter_names(model, forbidden_layer_types, forbidden_layer_names=None): - forbidden_layer_patterns = ([re.compile(pattern) for pattern in forbidden_layer_names] - if forbidden_layer_names is not None else []) + forbidden_layer_patterns = ( + [re.compile(pattern) for pattern in forbidden_layer_names] if forbidden_layer_names is not None else [] + ) result = [] for name, child in model.named_children(): child_params = get_parameter_names(child, forbidden_layer_types, forbidden_layer_names) result += [ - f'{name}.{n}' for n in child_params - if not isinstance(child, tuple(forbidden_layer_types)) and not any( - pattern.search(f'{name}.{n}'.lower()) for pattern in forbidden_layer_patterns) + f"{name}.{n}" + for n in child_params + if not isinstance(child, tuple(forbidden_layer_types)) + and not any(pattern.search(f"{name}.{n}".lower()) for pattern in forbidden_layer_patterns) ] # Add model specific parameters that are not in any child result += [ - k for k in model._parameters + k + for k in model._parameters if not any(pattern.search(k.lower()) for pattern in forbidden_layer_patterns) ] 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) + forbidden_name_patterns = [r"bias", r"layernorm", r"rmsnorm", r"(?:^|\.)norm(?:$|\.)", r"_norm(?:$|\.)"] + # 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)] optimizer_grouped_parameters = [ { - 'params': [params[n] for n in decay_param_names], - 'param_names': decay_param_names, - 'weight_decay': weight_decay, - 'lr': lr + "params": [params[n] for n in decay_param_names], + "param_names": decay_param_names, + "weight_decay": weight_decay, + "lr": lr, }, { - 'params': [params[n] for n in no_decay_param_names], - 'param_names': no_decay_param_names, - 'weight_decay': 0.0, - 'lr': lr + "params": [params[n] for n in no_decay_param_names], + "param_names": no_decay_param_names, + "weight_decay": 0.0, + "lr": lr, }, ] return optimizer_grouped_parameters @@ -1748,43 +1837,46 @@ def step(self, **kwargs): adapter_name: Lora adapter name. Any parameters needed for `optimizer.step`. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.require_finalized('step the optimizer') - grad_accum_steps = kwargs.pop('gradient_accumulation_steps', None) + optimizer_config.require_finalized("step the optimizer") + grad_accum_steps = kwargs.pop("gradient_accumulation_steps", None) if not optimizer_config.do_grad_sync(grad_accum_steps): return optimizer = optimizer_config.optimizer scaler = optimizer_config.scaler - assert isinstance(optimizer, Optimizer), 'Set optimizer correctly before forwarding' + assert isinstance(optimizer, Optimizer), "Set optimizer correctly before forwarding" context = contextlib.nullcontext if self.device_mesh is not None and self.device_mesh.tp_world_size > 1: from torch.distributed.tensor.experimental import implicit_replication + context = implicit_replication - optim_params = kwargs.pop('optim_params', {}) + optim_params = kwargs.pop("optim_params", {}) if optim_params: # After _lazy_wrap_model the optimizer may be wrapped (e.g. # accelerate's AcceleratedOptimizer); check the inner instance. - inner_optimizer = getattr(optimizer, 'optimizer', optimizer) - assert isinstance(inner_optimizer, (AdamW, Adam)), \ - f'optim_params is only supported for Adam/AdamW, got {type(inner_optimizer).__name__}' + inner_optimizer = getattr(optimizer, "optimizer", optimizer) + assert isinstance(inner_optimizer, (AdamW, Adam)), ( + f"optim_params is only supported for Adam/AdamW, got {type(inner_optimizer).__name__}" + ) for group in optimizer.param_groups: - group['lr'] = optim_params['lr'] - if group['weight_decay'] > 0.0 and optim_params.get('weight_decay', None) is not None: - group['weight_decay'] = optim_params['weight_decay'] - if optim_params.get('eps') is not None: - group['eps'] = optim_params['eps'] - if optim_params.get('betas') is not None: - group['betas'] = optim_params['betas'] + group["lr"] = optim_params["lr"] + if group["weight_decay"] > 0.0 and optim_params.get("weight_decay", None) is not None: + group["weight_decay"] = optim_params["weight_decay"] + if optim_params.get("eps") is not None: + group["eps"] = optim_params["eps"] + if optim_params.get("betas") is not None: + group["betas"] = optim_params["betas"] with context(): if scaler is not None: scaler.step(optimizer, **kwargs) scaler.update() - optimizer_config.scaler_has_nan = sum(v.item() - for v in scaler._found_inf_per_device(optimizer).values()) > 0 + optimizer_config.scaler_has_nan = ( + sum(v.item() for v in scaler._found_inf_per_device(optimizer).values()) > 0 + ) else: optimizer.step(**kwargs) @@ -1797,13 +1889,13 @@ def zero_grad(self, **kwargs): adapter_name: Lora adapter name. Any parameters needed for `optimizer.zero_grad`. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.require_finalized('zero gradients') - if not optimizer_config.do_grad_sync(kwargs.pop('gradient_accumulation_steps', None)): + optimizer_config.require_finalized("zero gradients") + if not optimizer_config.do_grad_sync(kwargs.pop("gradient_accumulation_steps", None)): return optimizer = optimizer_config.optimizer - assert isinstance(optimizer, Optimizer), 'Set optimizer correctly before forwarding' + assert isinstance(optimizer, Optimizer), "Set optimizer correctly before forwarding" optimizer.zero_grad(set_to_none=True) @remote_function() @@ -1815,10 +1907,10 @@ def lr_step(self, **kwargs): adapter_name: Lora adapter name. Any parameters needed for `lr_scheduler.step`. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.require_finalized('step the learning-rate scheduler') - if not optimizer_config.do_grad_sync(kwargs.pop('gradient_accumulation_steps', None)): + optimizer_config.require_finalized("step the learning-rate scheduler") + if not optimizer_config.do_grad_sync(kwargs.pop("gradient_accumulation_steps", None)): return if optimizer_config.scaler_has_nan: return @@ -1827,8 +1919,9 @@ def lr_step(self, **kwargs): lr_scheduler.step(**kwargs) @remote_function() - def set_loss(self, loss_cls: Union[Loss, Type[Loss], str, Callable[[InputFeature, ModelOutput, ...], torch.Tensor]], - **kwargs): + def set_loss( + self, loss_cls: Loss | type[Loss] | str | Callable[[InputFeature, ModelOutput, ...], torch.Tensor], **kwargs + ): """Set the loss instance. Args: @@ -1837,13 +1930,13 @@ def set_loss(self, loss_cls: Union[Loss, Type[Loss], str, Callable[[InputFeature adapter_name: Lora adapter name. Any parameters needed to construct the loss instance. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) - loss_key = kwargs.pop('loss_key', None) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) + loss_key = kwargs.pop("loss_key", None) optimizer_config = self.optimizer_group[adapter_name] optimizer_config.set_loss(construct_class(loss_cls, Loss, twinkle.loss, **kwargs), loss_key=loss_key) @remote_function() - def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str, Optimizer], **kwargs): + def set_optimizer(self, optimizer_cls: type[Optimizer] | str | Optimizer, **kwargs): """Set the optimizer. Args: @@ -1866,23 +1959,23 @@ def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str, Optimizer], * passing it is how `MuonClip` gets more than plain momentum SGD. Any parameters needed to construct the optimizer instance. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) # Metric copies the dp group at construction, and OptimizerGroup builds metrics before # dist init -- so rebuild here (first path that runs post-init on every backend) to get # dp-wide token-weighted logging. Logging only; gradients are unaffected. self._ensure_optimizer_dp_groups() optimizer_config = self.optimizer_group[adapter_name] - galore_config = kwargs.pop('galore_config', None) - muon_config = kwargs.pop('muon_config', None) - param_groups_spec = kwargs.pop('param_groups_spec', None) + galore_config = kwargs.pop("galore_config", None) + muon_config = kwargs.pop("muon_config", None) + param_groups_spec = kwargs.pop("param_groups_spec", None) if isinstance(optimizer_cls, Optimizer): optimizer_config.optimizer = optimizer_cls return - params = kwargs.pop('params', None) + params = kwargs.pop("params", None) if params is None: - lr = kwargs.get('lr', DEFAULT_LEARNING_RATE) - weight_decay = kwargs.get('weight_decay', DEFAULT_WEIGHT_DECAY) + lr = kwargs.get("lr", DEFAULT_LEARNING_RATE) + weight_decay = kwargs.get("weight_decay", DEFAULT_WEIGHT_DECAY) params = self._create_param_group(adapter_name, lr=lr, weight_decay=weight_decay) if param_groups_spec: # Split the decay / no-decay groups by the spec before GaLore / Muon reshape them, so a @@ -1892,10 +1985,13 @@ def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str, Optimizer], * params, param_groups_spec, model=self.strategy.unwrap_model(self.model), - tower_prefixes=resolve_tower_prefixes(self.model_arch)) + tower_prefixes=resolve_tower_prefixes(self.model_arch), + ) if galore_config is not None and muon_config is not None: - raise ValueError('galore_config and muon_config are mutually exclusive: the Muon grouping ' - 'would stamp GaLore keys (rank, proj_type, ...) onto every sub-group.') + raise ValueError( + "galore_config and muon_config are mutually exclusive: the Muon grouping " + "would stamp GaLore keys (rank, proj_type, ...) onto every sub-group." + ) if galore_config is not None: if isinstance(galore_config, dict): galore_config = GaLoreConfig(**galore_config) @@ -1906,12 +2002,11 @@ def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str, Optimizer], * # Families whose output head is not called 'lm_head' (internlm2 'output', chatglm # 'transformer.output_layer', ...) would otherwise take the Muon update; the embedding # side needs no key because create_muon_param_groups excludes it by type. - lm_head = self.model_arch.get('lm_head') + lm_head = self.model_arch.get("lm_head") if lm_head and lm_head not in muon_config.exclude_keys: muon_config.exclude_keys = list(muon_config.exclude_keys) + [lm_head] - params = create_muon_param_groups( - params, muon_config, model=self.strategy.unwrap_model(self.model)) - if hasattr(self.strategy, 'adjust_optimizer_kwargs'): + params = create_muon_param_groups(params, muon_config, model=self.strategy.unwrap_model(self.model)) + if hasattr(self.strategy, "adjust_optimizer_kwargs"): kwargs = self.strategy.adjust_optimizer_kwargs(optimizer_cls, kwargs) optimizer_config.optimizer = construct_class( optimizer_cls, @@ -1925,7 +2020,7 @@ def _get_trainable_parameters(self, adapter_name=_default_adapter_name): is_default = adapter_name == _default_adapter_name # Previously the pattern did not match nn.Parameter() weights, causing EP LoRA # parameters to be missed by the optimizer. - pattern = re.compile(rf'\.(?:lora_\w+|modules_to_save)\.{re.escape(adapter_name)}') + pattern = re.compile(rf"\.(?:lora_\w+|modules_to_save)\.{re.escape(adapter_name)}") params = {} model = self.strategy.unwrap_model(self.model) for name, param in model.named_parameters(): @@ -1934,7 +2029,7 @@ def _get_trainable_parameters(self, adapter_name=_default_adapter_name): return params @remote_function() - def set_lr_scheduler(self, scheduler_cls: Union[Type[LRScheduler], str, LRScheduler], **kwargs): + def set_lr_scheduler(self, scheduler_cls: type[LRScheduler] | str | LRScheduler, **kwargs): """Set the lr_scheduler. Args: @@ -1943,20 +2038,21 @@ def set_lr_scheduler(self, scheduler_cls: Union[Type[LRScheduler], str, LRSchedu adapter_name: Lora adapter name. Any parameters needed to construct the lr_scheduler instance. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] optimizer = optimizer_config.optimizer - assert isinstance(optimizer, Optimizer), 'Set optimizer correctly before setting lr_scheduler' - kwargs['optimizer'] = optimizer - scheduler = construct_class(scheduler_cls, LRScheduler, [torch.optim.lr_scheduler, twinkle.module.scheduler], - **kwargs) + assert isinstance(optimizer, Optimizer), "Set optimizer correctly before setting lr_scheduler" + kwargs["optimizer"] = optimizer + scheduler = construct_class( + scheduler_cls, LRScheduler, [torch.optim.lr_scheduler, twinkle.module.scheduler], **kwargs + ) optimizer_config.lr_scheduler = scheduler @remote_function() - def apply_patch(self, patch_cls: Union[Patch, Type[Patch], str], **kwargs): + def apply_patch(self, patch_cls: Patch | type[Patch] | str, **kwargs): apply_patch(self.model, patch_cls, **kwargs) - @remote_function(dispatch='all') + @remote_function(dispatch="all") def empty_cache(self, **kwargs): """Release the caching allocator's unused blocks on every worker. @@ -1979,9 +2075,9 @@ def __del__(self): # would be torn down with the actor before writing, leaving an empty checkpoint-final. Megatron's save # is sync=True for the same reason; keeping the transformers backend non-lazy preserves that backend # equivalence. - @remote_function(collect='first', lazy_collect=False) + @remote_function(collect="first", lazy_collect=False) @timed - def save_initial_adapter(self, initial_dir: str, adapter_name: Optional[str] = None) -> bool: + def save_initial_adapter(self, initial_dir: str, adapter_name: str | None = None) -> bool: """Save the adapter's *initial* (pre-training) weights as a plain-LoRA checkpoint. An init strategy that rewrites the base weight at injection (peft: pissa / olora / corda ...) @@ -2005,10 +2101,11 @@ def save_initial_adapter(self, initial_dir: str, adapter_name: Optional[str] = N config.init_lora_weights = init_lora_weights return True - @remote_function(collect='first', lazy_collect=False) + @remote_function(collect="first", lazy_collect=False) @timed - def save_converted_adapter(self, converted_dir: str, initial_model_dir: str, - adapter_name: Optional[str] = None) -> None: + def save_converted_adapter( + self, converted_dir: str, initial_model_dir: str, adapter_name: str | None = None + ) -> None: """Save a *converted* copy of the adapter: its delta re-expressed against the ORIGINAL base. peft's ``path_initial_model_for_weight_conversion`` folds the init factor back out using the @@ -2017,7 +2114,7 @@ def save_converted_adapter(self, converted_dir: str, initial_model_dir: str, """ model = self.strategy.unwrap_model(self.model) if not isinstance(model, PeftModel): - raise TypeError('save_converted_adapter requires a peft-wrapped model.') + raise TypeError("save_converted_adapter requires a peft-wrapped model.") adapter_name = adapter_name or self._get_default_group() saved_config = copy(model.peft_config[adapter_name]) # The conversion loads the initial snapshot through peft's inject_adapter, which freezes the @@ -2025,16 +2122,17 @@ def save_converted_adapter(self, converted_dir: str, initial_model_dir: str, requires_grad_state = {name: p.requires_grad for name, p in model.named_parameters()} os.makedirs(os.path.dirname(converted_dir), exist_ok=True) model.save_pretrained( - converted_dir, path_initial_model_for_weight_conversion=os.path.abspath(initial_model_dir)) + converted_dir, path_initial_model_for_weight_conversion=os.path.abspath(initial_model_dir) + ) # peft replaces the adapter's config entry during conversion; restore the live one. model.peft_config[adapter_name] = saved_config for name, p in model.named_parameters(): if name in requires_grad_state: p.requires_grad = requires_grad_state[name] - @remote_function(collect='first', lazy_collect=False) + @remote_function(collect="first", lazy_collect=False) @timed - def save(self, name: Optional[str] = None, output_dir: Optional[str] = None, interval: int = 1, **kwargs): + def save(self, name: str | None = None, output_dir: str | None = None, interval: int = 1, **kwargs): """Save model. Args: @@ -2050,23 +2148,23 @@ def save(self, name: Optional[str] = None, output_dir: Optional[str] = None, int async_save: Offload the staged weight/tokenizer/trainer-state file I/O to a background writer (the state itself is always gathered synchronously). """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) - safe_serialization = kwargs.pop('safe_serialization', True) - max_shard_size = kwargs.pop('max_shard_size', '5GB') - save_total_limit = kwargs.pop('save_total_limit', None) - async_save = kwargs.pop('async_save', False) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) + safe_serialization = kwargs.pop("safe_serialization", True) + max_shard_size = kwargs.pop("max_shard_size", "5GB") + save_total_limit = kwargs.pop("save_total_limit", None) + async_save = kwargs.pop("async_save", False) optimizer_config = self.optimizer_group[adapter_name] if name is None: - name = f'checkpoint-step-{optimizer_config.cur_step}' + name = f"checkpoint-step-{optimizer_config.cur_step}" if output_dir is None: - output_dir = 'output' + output_dir = "output" checkpoint_dir = os.path.join(output_dir, name) if optimizer_config.cur_step % interval != 0: return model = self.strategy.unwrap_model(self.model) save_kwargs = { - 'safe_serialization': safe_serialization, - 'max_shard_size': max_shard_size, + "safe_serialization": safe_serialization, + "max_shard_size": max_shard_size, } if adapter_name == _default_adapter_name: # Full model save @@ -2078,33 +2176,34 @@ def save(self, name: Optional[str] = None, output_dir: Optional[str] = None, int # engine.save_checkpoint / FSDP2's get_optimizer_state_dict perform collectives), so it stays on # the synchronous path; async only offloads the staged weight/tokenizer/trainer-state file I/O. self._deferred_trainer_state_writes = [] - if kwargs.get('save_optimizer', False): + if kwargs.get("save_optimizer", False): self._save_training_state( checkpoint_dir, adapter_name=adapter_name, - consumed_train_samples=kwargs.get('consumed_train_samples', 0), - component_states=kwargs.get('component_states', {}), + consumed_train_samples=kwargs.get("consumed_train_samples", 0), + component_states=kwargs.get("component_states", {}), defer_trainer_state=async_save, ) def _write_weights() -> None: """Pure file I/O over the staged CPU payloads; no collectives, no device tensors.""" if isinstance(model, PeftModel): - if max_shard_size != '5GB': + if max_shard_size != "5GB": raise NotImplementedError( - 'max_shard_size is not supported for PEFT adapter checkpoints; remove the option or save a ' - 'full-model checkpoint.') + "max_shard_size is not supported for PEFT adapter checkpoints; remove the option or save a " + "full-model checkpoint." + ) if Platform.is_master(): model.peft_config[adapter_name].save_pretrained(checkpoint_dir) contiguous_dict = {k: v.contiguous() for k, v in processed_state_dict.items()} if safe_serialization: - save_file(contiguous_dict, os.path.join(checkpoint_dir, 'adapter_model.safetensors')) + save_file(contiguous_dict, os.path.join(checkpoint_dir, "adapter_model.safetensors")) else: - torch.save(contiguous_dict, os.path.join(checkpoint_dir, 'adapter_model.bin')) + torch.save(contiguous_dict, os.path.join(checkpoint_dir, "adapter_model.bin")) else: model.save_pretrained( - checkpoint_dir, state_dict=processed_state_dict, is_main_process=Platform.is_master(), - **save_kwargs) + checkpoint_dir, state_dict=processed_state_dict, is_main_process=Platform.is_master(), **save_kwargs + ) if async_save: # trainer_state.json doubles as the completion marker a resume reads; the deferred writes run @@ -2135,18 +2234,18 @@ def _get_adapter_state_dict_for_save(self, adapter_name: str) -> dict: adapter_state = self.strategy.get_adapter_state_dict(self.model, adapter_name) return normalize_adapter_keys(adapter_state, adapter_name) - def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) -> Dict[str, str]: + def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer) -> dict[str, str]: """Map runtime optimizer parameter names to checkpoint names.""" return {} def _save_optimizer(self, output_dir, **kwargs): - adapter_name = kwargs.pop('adapter_name', _default_adapter_name) + adapter_name = kwargs.pop("adapter_name", _default_adapter_name) optimizer_config = self.optimizer_group[adapter_name] optimizer = optimizer_config.optimizer lr_scheduler = optimizer_config.lr_scheduler if optimizer is not None: - optimizer_path = os.path.join(output_dir, 'optimizer.pt') + optimizer_path = os.path.join(output_dir, "optimizer.pt") name_mapping = self._optimizer_param_name_mapping(adapter_name, optimizer) self.strategy.save_optimizer_checkpoint( self.model, @@ -2156,11 +2255,11 @@ def _save_optimizer(self, output_dir, **kwargs): ) if Platform.is_master(): if lr_scheduler is not None: - torch.save(lr_scheduler.state_dict(), os.path.join(output_dir, 'scheduler.pt')) + torch.save(lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt")) def _save_training_state(self, output_dir, **kwargs): - adapter_name = kwargs.pop('adapter_name', _default_adapter_name) - defer_trainer_state = kwargs.pop('defer_trainer_state', False) + adapter_name = kwargs.pop("adapter_name", _default_adapter_name) + defer_trainer_state = kwargs.pop("defer_trainer_state", False) # The optimizer/scheduler files are written before _write_weights() gets a chance to # create the checkpoint dir via save_pretrained; torch.save then raises # "Parent directory ... does not exist". This lives HERE (not at the top of save()) so it @@ -2174,15 +2273,15 @@ def _save_training_state(self, output_dir, **kwargs): return trainer_state = { - 'checkpoint_version': 1, - 'cur_step': optimizer_config.cur_step, - 'gradient_accumulation_steps': optimizer_config.gradient_accumulation_steps, - 'consumed_train_samples': kwargs.get('consumed_train_samples', 0), - 'component_states': kwargs.get('component_states', {}), + "checkpoint_version": 1, + "cur_step": optimizer_config.cur_step, + "gradient_accumulation_steps": optimizer_config.gradient_accumulation_steps, + "consumed_train_samples": kwargs.get("consumed_train_samples", 0), + "component_states": kwargs.get("component_states", {}), } def _write_trainer_state() -> None: - with open(os.path.join(output_dir, 'trainer_state.json'), 'w', encoding='utf-8') as f: + with open(os.path.join(output_dir, "trainer_state.json"), "w", encoding="utf-8") as f: json.dump(trainer_state, f) # Under async_save the trainer_state.json is the resume completion marker: defer it to the @@ -2195,16 +2294,16 @@ def _write_trainer_state() -> None: if optimizer_config.scaler is not None: torch.save( { - 'scaler_state_dict': optimizer_config.scaler.state_dict(), - 'scaler_has_nan': optimizer_config.scaler_has_nan, + "scaler_state_dict": optimizer_config.scaler.state_dict(), + "scaler_has_nan": optimizer_config.scaler_has_nan, }, - os.path.join(output_dir, 'scaler.pt'), + os.path.join(output_dir, "scaler.pt"), ) - torch.save(self._get_training_rng_state(), os.path.join(output_dir, 'rng_state.pt')) + torch.save(self._get_training_rng_state(), os.path.join(output_dir, "rng_state.pt")) def _save_tokenizer(self, output_dir, **kwargs): - adapter_name = kwargs.pop('adapter_name', _default_adapter_name) + adapter_name = kwargs.pop("adapter_name", _default_adapter_name) optimizer_config = self.optimizer_group[adapter_name] template_ins = optimizer_config.template if Platform.is_master(): @@ -2215,7 +2314,7 @@ def _save_tokenizer(self, output_dir, **kwargs): @remote_function() @timed - def load(self, name: str, output_dir: Optional[str] = None, **kwargs): + def load(self, name: str, output_dir: str | None = None, **kwargs): """Load model state and optionally optimizer state from a checkpoint. Args: @@ -2225,21 +2324,21 @@ def load(self, name: str, output_dir: Optional[str] = None, **kwargs): adapter_name: Adapter to load. load_optimizer: Whether to load optimizer and scheduler states. """ - load_optimizer = kwargs.get('load_optimizer', False) - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + load_optimizer = kwargs.get("load_optimizer", False) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) if output_dir is None: if os.path.exists(name): checkpoint_dir = name else: # load from hub - token = kwargs.pop('token', None) + token = kwargs.pop("token", None) checkpoint_dir = HubOperation.download_model(name, token=token) else: checkpoint_dir = os.path.join(output_dir, name) model = self.strategy.unwrap_model(self.model) if isinstance(model, PeftModel): - adapter_weights = load_peft_weights(checkpoint_dir, device='cpu') + adapter_weights = load_peft_weights(checkpoint_dir, device="cpu") self.strategy.load_peft_weights(model, adapter_weights, adapter_name) else: # Full-parameter model: load a plain HF checkpoint in-place. @@ -2257,8 +2356,17 @@ 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).') + logger.warning("reload_initial_weights skipped: model_id is not set (blank model).") return state_dict = _read_hf_state_dict(self.model_id) self.strategy.load_full_state_dict(self.model, state_dict) @@ -2268,19 +2376,20 @@ def reload_initial_weights(self, **kwargs): param.grad = None def _load_optimizer(self, checkpoint_dir, **kwargs): - adapter_name = kwargs.pop('adapter_name', _default_adapter_name) - strict = kwargs.pop('strict', False) + adapter_name = kwargs.pop("adapter_name", _default_adapter_name) + strict = kwargs.pop("strict", False) # assume optimizer and lr_scheduler are created optimizer_config = self.optimizer_group[adapter_name] - optimizer_path = os.path.join(checkpoint_dir, 'optimizer.pt') - scheduler_path = os.path.join(checkpoint_dir, 'scheduler.pt') + optimizer_path = os.path.join(checkpoint_dir, "optimizer.pt") + scheduler_path = os.path.join(checkpoint_dir, "scheduler.pt") if strict and not os.path.exists(optimizer_path): raise FileNotFoundError(optimizer_path) if strict and optimizer_config.lr_scheduler is not None and not os.path.exists(scheduler_path): logger.warning( - f'Missing scheduler checkpoint {scheduler_path}; resuming without restoring lr scheduler state.', ) + f"Missing scheduler checkpoint {scheduler_path}; resuming without restoring lr scheduler state.", + ) if os.path.exists(optimizer_path) and optimizer_config.optimizer is not None: if self.strategy.needs_wrapped_optimizer_state() and not self._model_wrapped: @@ -2296,14 +2405,14 @@ def _load_optimizer(self, checkpoint_dir, **kwargs): ) if os.path.exists(scheduler_path) and optimizer_config.lr_scheduler is not None: - state_dict = torch.load(scheduler_path, map_location='cpu', weights_only=True) + state_dict = torch.load(scheduler_path, map_location="cpu", weights_only=True) optimizer_config.lr_scheduler.load_state_dict(state_dict) @staticmethod def _ensure_lora_dtype(model): """Force LoRA parameters to use the same dtype as base model for FSDP2 compatibility.""" base_dtype = None - is_npu_device = Platform.device_prefix() == 'npu' + is_npu_device = Platform.device_prefix() == "npu" for param in model.parameters(): if base_dtype is None and param.dtype in (torch.float16, torch.bfloat16, torch.float32): base_dtype = param.dtype @@ -2315,18 +2424,18 @@ def _ensure_lora_dtype(model): # Temporary workaround: NPU requires all parameters to align with the base dtype. with torch.no_grad(): for name, param in model.named_parameters(): - if (is_npu_device or 'lora_' in name.lower()) and param.dtype != base_dtype: + if (is_npu_device or "lora_" in name.lower()) and param.dtype != base_dtype: param.data = param.data.to(base_dtype) def _load_scaler_state(self, scaler_path, **kwargs): - adapter_name = kwargs.pop('adapter_name', _default_adapter_name) + adapter_name = kwargs.pop("adapter_name", _default_adapter_name) optimizer_config = self.optimizer_group[adapter_name] if optimizer_config.scaler is None: - raise ValueError(f'Grad scaler is not configured for adapter {adapter_name!r}') + raise ValueError(f"Grad scaler is not configured for adapter {adapter_name!r}") - scaler_state = torch.load(scaler_path, map_location='cpu', weights_only=True) - optimizer_config.scaler.load_state_dict(scaler_state['scaler_state_dict']) - optimizer_config.scaler_has_nan = scaler_state.get('scaler_has_nan', False) + scaler_state = torch.load(scaler_path, map_location="cpu", weights_only=True) + optimizer_config.scaler.load_state_dict(scaler_state["scaler_state_dict"]) + optimizer_config.scaler_has_nan = scaler_state.get("scaler_has_nan", False) def _get_training_rng_state(self): # numpy's state tuple carries int64 scalars and an ndarray, which torch's @@ -2334,60 +2443,66 @@ def _get_training_rng_state(self): # whole file loadable with weights_only=True (np.random.set_state accepts it). numpy_name, numpy_key, numpy_pos, numpy_has_gauss, numpy_cached = np.random.get_state() state = { - 'python_rng_state': random.getstate(), - 'numpy_rng_state': (numpy_name, numpy_key.tolist(), int(numpy_pos), int(numpy_has_gauss), - float(numpy_cached)), - 'torch_rng_state': torch.get_rng_state(), + "python_rng_state": random.getstate(), + "numpy_rng_state": ( + numpy_name, + numpy_key.tolist(), + int(numpy_pos), + int(numpy_has_gauss), + float(numpy_cached), + ), + "torch_rng_state": torch.get_rng_state(), } device_prefix = Platform.device_prefix() device_module = getattr(torch, device_prefix, None) - if device_module and hasattr(device_module, 'is_available') and device_module.is_available(): - state['device_type'] = device_prefix - state['device_rng_state'] = device_module.get_rng_state() + if device_module and hasattr(device_module, "is_available") and device_module.is_available(): + state["device_type"] = device_prefix + state["device_rng_state"] = device_module.get_rng_state() else: - state['device_type'] = 'cpu' - state['device_rng_state'] = None + state["device_type"] = "cpu" + state["device_rng_state"] = None return state def _load_rng_state(self, rng_path): try: - rng_state = torch.load(rng_path, map_location='cpu', weights_only=True) + rng_state = torch.load(rng_path, map_location="cpu", weights_only=True) except Exception: # Pre-hardening checkpoints stored numpy scalars/arrays, which the safe # unpickler refuses. RNG state only affects run-to-run reproducibility, so # the recovery is to drop it rather than fall back to unrestricted pickle. raise ValueError( - f'{rng_path} is not loadable in weights-only mode (likely a pre-hardening checkpoint ' - 'carrying numpy rng objects). Delete the file and resume: only exact RNG replay is ' - 'lost, training correctness is unaffected.') from None - random.setstate(rng_state['python_rng_state']) - np.random.set_state(rng_state['numpy_rng_state']) - torch.set_rng_state(rng_state['torch_rng_state']) - - device_type = rng_state.get('device_type') - device_rng_state = rng_state.get('device_rng_state') - if device_type != 'cpu' and device_rng_state is not None: + f"{rng_path} is not loadable in weights-only mode (likely a pre-hardening checkpoint " + "carrying numpy rng objects). Delete the file and resume: only exact RNG replay is " + "lost, training correctness is unaffected." + ) from None + random.setstate(rng_state["python_rng_state"]) + np.random.set_state(rng_state["numpy_rng_state"]) + torch.set_rng_state(rng_state["torch_rng_state"]) + + device_type = rng_state.get("device_type") + device_rng_state = rng_state.get("device_rng_state") + if device_type != "cpu" and device_rng_state is not None: device_module = getattr(torch, device_type, None) - if device_module and hasattr(device_module, 'is_available') and device_module.is_available(): + if device_module and hasattr(device_module, "is_available") and device_module.is_available(): device_module.set_rng_state(device_rng_state) - def _restore_training_state(self, checkpoint_dir, *, adapter_name=''): - trainer_state_path = os.path.join(checkpoint_dir, 'trainer_state.json') + def _restore_training_state(self, checkpoint_dir, *, adapter_name=""): + trainer_state_path = os.path.join(checkpoint_dir, "trainer_state.json") with open(trainer_state_path) as f: trainer_state = json.load(f) adapter_name = adapter_name or self._get_default_group() optimizer_config = self.optimizer_group[adapter_name] self._load_optimizer(checkpoint_dir, adapter_name=adapter_name) - scaler_path = os.path.join(checkpoint_dir, 'scaler.pt') + scaler_path = os.path.join(checkpoint_dir, "scaler.pt") if os.path.exists(scaler_path) and optimizer_config.scaler is not None: self._load_scaler_state(scaler_path, adapter_name=adapter_name) - rng_path = os.path.join(checkpoint_dir, 'rng_state.pt') + rng_path = os.path.join(checkpoint_dir, "rng_state.pt") if os.path.exists(rng_path): self._load_rng_state(rng_path) - optimizer_config.cur_step = trainer_state['cur_step'] - optimizer_config.gradient_accumulation_steps = trainer_state['gradient_accumulation_steps'] + optimizer_config.cur_step = trainer_state["cur_step"] + optimizer_config.gradient_accumulation_steps = trainer_state["gradient_accumulation_steps"] return trainer_state @@ -2397,37 +2512,37 @@ def _restore_training_state(self, checkpoint_dir, *, adapter_name=''): # raises "'function' object is not subscriptable". Mirrors forward_only/calculate_metric, the other # methods whose return value the driver reads. sync=True orders the workers but does not materialize # the result; the megatron/hybrid/multi_lora variants are non-lazy for the same backend equivalence. - @remote_function(dispatch='all', collect='first', sync=True, lazy_collect=False, timeout=3600) + @remote_function(dispatch="all", collect="first", sync=True, lazy_collect=False, timeout=3600) @timed def resume_from_checkpoint(self, checkpoint_dir, *, resume_only_model=False, **kwargs): - adapter_name = kwargs.get('adapter_name', '') + adapter_name = kwargs.get("adapter_name", "") self.load(checkpoint_dir, adapter_name=adapter_name) if resume_only_model: optimizer_config = self.optimizer_group[adapter_name] return { - 'cur_step': 0, - 'consumed_train_samples': 0, - 'gradient_accumulation_steps': optimizer_config.gradient_accumulation_steps, - 'component_states': {}, + "cur_step": 0, + "consumed_train_samples": 0, + "gradient_accumulation_steps": optimizer_config.gradient_accumulation_steps, + "component_states": {}, } trainer_state = self._restore_training_state(checkpoint_dir, adapter_name=adapter_name) result = { - 'cur_step': trainer_state['cur_step'], - 'consumed_train_samples': trainer_state['consumed_train_samples'], - 'gradient_accumulation_steps': trainer_state['gradient_accumulation_steps'], + "cur_step": trainer_state["cur_step"], + "consumed_train_samples": trainer_state["consumed_train_samples"], + "gradient_accumulation_steps": trainer_state["gradient_accumulation_steps"], } - if 'component_states' in trainer_state: - result['component_states'] = trainer_state['component_states'] + if "component_states" in trainer_state: + result["component_states"] = trainer_state["component_states"] return result - @remote_function(collect='first') + @remote_function(collect="first") def get_state_dict(self, **kwargs): - return self._get_trainable_parameters(kwargs.pop('adapter_name', self._get_default_group())) + return self._get_trainable_parameters(kwargs.pop("adapter_name", self._get_default_group())) - @remote_function(collect='first', lazy_collect=False) + @remote_function(collect="first", lazy_collect=False) def get_peft_config_dict(self, adapter_name: str = None) -> dict: """Return the PEFT config as a dict for vLLM's PEFTHelper. @@ -2447,27 +2562,28 @@ def get_peft_config_dict(self, adapter_name: str = None) -> dict: # or a single LoraConfig. Normalize to a single config. if isinstance(config, dict): config = config.get(adapter_name, next(iter(config.values()))) - return config.to_dict() if hasattr(config, 'to_dict') else dict(config) + return config.to_dict() if hasattr(config, "to_dict") else dict(config) - @remote_function(collect='first', lazy_collect=False) + @remote_function(collect="first", lazy_collect=False) @timed def calculate_metric(self, is_training, **kwargs): - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.require_finalized('calculate metrics') + optimizer_config.require_finalized("calculate metrics") return optimizer_config.calculate_metrics(is_training) - def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str], **kwargs): - assert adapter_name, 'Use a different adapter_name, current is empty.' - freeze_llm = kwargs.pop('freeze_llm', False) - freeze_vit = kwargs.pop('freeze_vit', True) - freeze_aligner = kwargs.pop('freeze_aligner', True) + def _patch_adapter(self, adapter_name: str, config_or_dir: PeftConfig | str, **kwargs): + assert adapter_name, "Use a different adapter_name, current is empty." + freeze_llm = kwargs.pop("freeze_llm", False) + freeze_vit = kwargs.pop("freeze_vit", True) + freeze_aligner = kwargs.pop("freeze_aligner", True) unwrapped_model = self.strategy.unwrap_model(self.model) if not isinstance(config_or_dir, str): config_or_dir = self.strategy.prepare_adapter_config( - config_or_dir, enable_ep=getattr(self, '_enable_expert_parallel', False)) + config_or_dir, enable_ep=getattr(self, "_enable_expert_parallel", False) + ) - if getattr(self, '_enable_expert_parallel', False): + if getattr(self, "_enable_expert_parallel", False): self.strategy.capture_pre_ep_state_if_needed(self.model, enable_ep=self._enable_expert_parallel) self._maybe_apply_expert_parallel() unwrapped_model = self.strategy.unwrap_model(self.model) @@ -2478,7 +2594,8 @@ def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str unwrapped_model, model_id=config_or_dir, adapter_name=adapter_name, - is_trainable=kwargs.get('is_trainable', True)) + is_trainable=kwargs.get("is_trainable", True), + ) if unwrapped_model is self.model: self.model = _adapted_model else: @@ -2491,14 +2608,11 @@ def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str else: config = config_or_dir self._expand_multimodal_lora_target( - config, - unwrapped_model, - freeze_llm=freeze_llm, - freeze_vit=freeze_vit, - freeze_aligner=freeze_aligner) + config, unwrapped_model, freeze_llm=freeze_llm, freeze_vit=freeze_vit, freeze_aligner=freeze_aligner + ) self._expand_modules_to_save(config, unwrapped_model) if not isinstance(unwrapped_model, PeftModel): - assert unwrapped_model is self.model, 'Cannot wrap model with peft after DDP/FSDP!' + assert unwrapped_model is self.model, "Cannot wrap model with peft after DDP/FSDP!" self.model = get_peft_model(unwrapped_model, config, adapter_name=adapter_name) else: unwrapped_model.add_adapter(adapter_name, config) @@ -2507,13 +2621,14 @@ def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str self.optimizer_group[adapter_name] = self._construct_default_optimizer_group() self.optimizer_group[adapter_name].adapter_name = adapter_name self.optimizer_group[adapter_name].adapter_config = config - _gas_default = kwargs.get('gradient_accumulation_steps', 1) + _gas_default = kwargs.get("gradient_accumulation_steps", 1) self.optimizer_group[adapter_name].gradient_accumulation_steps = _gas_default self._default_tokenizer = self.optimizer_group[adapter_name].template.processor self.active_group = adapter_name - def _expand_multimodal_lora_target(self, config, model, *, freeze_llm: bool, freeze_vit: bool, - freeze_aligner: bool) -> None: + def _expand_multimodal_lora_target( + self, config, model, *, freeze_llm: bool, freeze_vit: bool, freeze_aligner: bool + ) -> None: """Rewrite a multimodal ``'all-linear'`` LoRA target into the freeze_*-aware regex, in place. PEFT expands ``'all-linear'`` to every linear in the model, which on a multimodal model would @@ -2525,26 +2640,28 @@ def _expand_multimodal_lora_target(self, config, model, *, freeze_llm: bool, fre explicit user regex or name list is respected verbatim -- or when ``model_arch`` declares no vision tower / aligner, i.e. a plain LLM where ``'all-linear'`` already means the right thing. """ - target = getattr(config, 'target_modules', None) - is_all_linear = target == 'all-linear' or (isinstance(target, (list, tuple, set)) and 'all-linear' in target) + target = getattr(config, "target_modules", None) + is_all_linear = target == "all-linear" or (isinstance(target, (list, tuple, set)) and "all-linear" in target) if not is_all_linear: return - if not (self.model_arch.get('vision_tower') or self.model_arch.get('aligner')): + if not (self.model_arch.get("vision_tower") or self.model_arch.get("aligner")): return from twinkle.model.transformers.target_modules import get_multimodal_target_regex + # Walk the base model: target prefixes in model_arch are base-relative, and when an adapter is # added to an already-peft model the towers live under the wrapped base, not the PeftModel. base_model = model while isinstance(base_model, PeftModel): base_model = base_model.get_base_model() - include_embedding = isinstance(target, (list, tuple, set)) and 'all-embedding' in target + include_embedding = isinstance(target, (list, tuple, set)) and "all-embedding" in target config.target_modules = get_multimodal_target_regex( base_model, self.model_arch, freeze_llm=freeze_llm, freeze_vit=freeze_vit, freeze_aligner=freeze_aligner, - include_embedding=include_embedding) + include_embedding=include_embedding, + ) def _expand_modules_to_save(self, config, model) -> None: """Expand the ``'all-norm'`` / ``'all-embedding'`` shorthands in ``modules_to_save``, in place. @@ -2557,31 +2674,32 @@ def _expand_modules_to_save(self, config, model) -> None: target expansion, ``modules_to_save`` is NOT tower / freeze filtered: every norm / embedding the shorthand names is saved, regardless of which towers are LoRA-adapted (legacy does the same). """ - modules_to_save = getattr(config, 'modules_to_save', None) + modules_to_save = getattr(config, "modules_to_save", None) if not modules_to_save: return - shorthands = ('all-norm', 'all-embedding') + shorthands = ("all-norm", "all-embedding") if not any(name in shorthands for name in modules_to_save): return from twinkle.model.transformers.target_modules import find_embedding, find_norm + # Walk to the base model: when an adapter is added to an already-peft model the real modules # live under the wrapped base, not the PeftModel. base_model = model while isinstance(base_model, PeftModel): base_model = base_model.get_base_model() expanded = [name for name in modules_to_save if name not in shorthands] - for shorthand, finder in (('all-embedding', find_embedding), ('all-norm', find_norm)): + for shorthand, finder in (("all-embedding", find_embedding), ("all-norm", find_norm)): if shorthand not in modules_to_save: continue found = finder(base_model) if not found: - logger.warning(f'modules_to_save {shorthand!r} matched no module on this model; ignored.') + logger.warning(f"modules_to_save {shorthand!r} matched no module on this model; ignored.") expanded += found # de-dup while preserving order (a shorthand may overlap an explicitly-named module). config.modules_to_save = list(dict.fromkeys(expanded)) @remote_function() - def add_adapter_to_model(self, adapter_name: str, config_or_dir: Union[PeftConfig, str], **kwargs): + def add_adapter_to_model(self, adapter_name: str, config_or_dir: PeftConfig | str, **kwargs): """Add adapter to model. Args: @@ -2594,7 +2712,7 @@ def add_adapter_to_model(self, adapter_name: str, config_or_dir: Union[PeftConfi self._patch_adapter(adapter_name, config_or_dir, **kwargs) @remote_function() - def set_template(self, template_cls: Union[Type[Template], str, Template], **kwargs): + def set_template(self, template_cls: type[Template] | str | Template, **kwargs): """Set template. This is optional, if you need to input `Trajectory`, you need to set the template to encode them. @@ -2604,14 +2722,14 @@ def set_template(self, template_cls: Union[Type[Template], str, Template], **kwa adapter_name: Lora adapter name. Any parameters needed to construct the template_cls instance. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - if 'template_spec' not in kwargs: - kwargs['model_id'] = self.tokenizer_id + if "template_spec" not in kwargs: + kwargs["model_id"] = self.tokenizer_id optimizer_config.template = construct_template(template_cls, **kwargs) @remote_function() - def set_processor(self, processor_cls: Union[Type[InputProcessor], str, InputProcessor, Callable], **kwargs): + def set_processor(self, processor_cls: type[InputProcessor] | str | InputProcessor | Callable, **kwargs): """Set task processor to prepare the task inputs. Args: processor_cls: A processor_cls class name, a processor_cls plugin id, @@ -2620,9 +2738,9 @@ def set_processor(self, processor_cls: Union[Type[InputProcessor], str, InputPro adapter_name: Lora adapter name. Any parameters needed to construct the processor_cls instance. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - kwargs['device_mesh'] = self.device_mesh + kwargs["device_mesh"] = self.device_mesh processor = construct_class(processor_cls, InputProcessor, twinkle.processor, **kwargs) optimizer_config.processor = processor @@ -2634,15 +2752,16 @@ def set_grad_scaler(self, **kwargs): adapter_name: Lora adapter name. Any parameters needed to construct the GradScaler instance. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] from torch.amp.grad_scaler import GradScaler + grad_scaler_config = self.grad_scaler_config.copy() grad_scaler_config.update(kwargs) optimizer_config.scaler = GradScaler(**grad_scaler_config) @remote_function() - def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] = None, **kwargs): + def add_metric(self, metric_cls: Metric | str, is_training: bool | None = None, **kwargs): """Add an eval metric Args: @@ -2652,11 +2771,11 @@ def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] adapter_name: Lora adapter name. Any parameters needed to construct the metric_cls instance. """ - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) - loss_key = kwargs.pop('loss_key', None) + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) + loss_key = kwargs.pop("loss_key", None) optimizer_config = self.optimizer_group[adapter_name] - kwargs['device_mesh'] = self.device_mesh - kwargs['process_group'] = optimizer_config._dp_group + kwargs["device_mesh"] = self.device_mesh + kwargs["process_group"] = optimizer_config._dp_group if is_training is None or is_training is True: metric = construct_class(metric_cls, Metric, twinkle.metric, **kwargs) metric._twinkle_loss_key = loss_key @@ -2674,16 +2793,16 @@ def _get_nb_trainable_parameters(self, adapter_name, model): # everything under an `_ep_patched` module's `experts` subtree (raw # expert tensors + any PEFT LoRA stacked on them) is undercounted by # ep_size. Compensate here. - strategy = getattr(self, 'strategy', None) - ep_mesh = getattr(strategy, 'ep_fsdp_device_mesh', None) if strategy is not None else None - ep_size = ep_mesh['ep'].size() if ep_mesh is not None else 1 + strategy = getattr(self, "strategy", None) + ep_mesh = getattr(strategy, "ep_fsdp_device_mesh", None) if strategy is not None else None + ep_size = ep_mesh["ep"].size() if ep_mesh is not None else 1 ep_param_ids: set = set() if ep_size > 1: for module in model.modules(): - if not getattr(module, '_ep_patched', False): + if not getattr(module, "_ep_patched", False): continue - experts = getattr(module, 'experts', None) + experts = getattr(module, "experts", None) if experts is None: continue for p in experts.parameters(): @@ -2693,7 +2812,7 @@ def _get_nb_trainable_parameters(self, adapter_name, model): total = 0 for _, p in model.named_parameters(): n = p.numel() - if n == 0 and hasattr(p, 'ds_numel'): + if n == 0 and hasattr(p, "ds_numel"): n = p.ds_numel if id(p) in ep_param_ids: n *= ep_size @@ -2707,43 +2826,47 @@ def _get_trainable_parameters_example(self, adapter_name, model): for name, parameter in self.model.named_parameters(): if parameter.requires_grad: trainable_param_names.append(name) - trainable_param_names = trainable_param_names[:5] + ['...'] + trainable_param_names[-5:] - trainable_param_names = '\n'.join(trainable_param_names) + trainable_param_names = trainable_param_names[:5] + ["..."] + trainable_param_names[-5:] + trainable_param_names = "\n".join(trainable_param_names) return trainable_param_names - @remote_function(dispatch='all', collect='first', lazy_collect=False) + @remote_function(dispatch="all", collect="first", lazy_collect=False) def get_train_configs(self, **kwargs) -> str: - expr = '' - adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + expr = "" + adapter_name = kwargs.pop("adapter_name", self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] # adapter_config is always a single PeftConfig (see _patch_adapter); to_dict covers both # the from-scratch and the from-checkpoint load paths. if optimizer_config.adapter_config is not None: config = optimizer_config.adapter_config - config = config.to_dict() if hasattr(config, 'to_dict') else dict(config) + config = config.to_dict() if hasattr(config, "to_dict") else dict(config) else: config = {} config = {key: str(value) for key, value in config.items() if value is not None} trainable_params, all_param = self._get_nb_trainable_parameters(adapter_name, self.model) trainable_param_names = self._get_trainable_parameters_example(adapter_name, self.model) if optimizer_config.optimizer is not None: - expr += (f'Adapter config:\n' - f'{json.dumps(config, indent=2, ensure_ascii=False)}\n' - f'Trainable parameters examples:\n' - f'{trainable_param_names}\n' - f'Trainable params: {trainable_params:,d} || all params: {all_param:,d} || ' - f'trainable%: {100 * trainable_params / all_param:.4f}\n' - f'Optimizer: {optimizer_config.optimizer.__class__.__name__}\n' - f'Learning rate: {optimizer_config.optimizer.defaults.get("lr", "No default lr")}\n' - f'Lr scheduler: {optimizer_config.lr_scheduler.__class__.__name__}\n' - f'Gradient accumulation steps: {optimizer_config.gradient_accumulation_steps}\n') + expr += ( + f"Adapter config:\n" + f"{json.dumps(config, indent=2, ensure_ascii=False)}\n" + f"Trainable parameters examples:\n" + f"{trainable_param_names}\n" + f"Trainable params: {trainable_params:,d} || all params: {all_param:,d} || " + f"trainable%: {100 * trainable_params / all_param:.4f}\n" + f"Optimizer: {optimizer_config.optimizer.__class__.__name__}\n" + f"Learning rate: {optimizer_config.optimizer.defaults.get('lr', 'No default lr')}\n" + f"Lr scheduler: {optimizer_config.lr_scheduler.__class__.__name__}\n" + f"Gradient accumulation steps: {optimizer_config.gradient_accumulation_steps}\n" + ) else: - expr += (f'Adapter config:\n' - f'{json.dumps(config, indent=2, ensure_ascii=False)}\n' - f'Trainable parameters examples:\n' - f'{trainable_param_names}\n' - f'Trainable params: {trainable_params:,d} || all params: {all_param:,d} || ' - f'trainable%: {100 * trainable_params / all_param:.4f}%\n') + expr += ( + f"Adapter config:\n" + f"{json.dumps(config, indent=2, ensure_ascii=False)}\n" + f"Trainable parameters examples:\n" + f"{trainable_param_names}\n" + f"Trainable params: {trainable_params:,d} || all params: {all_param:,d} || " + f"trainable%: {100 * trainable_params / all_param:.4f}%\n" + ) return expr # ========================================================================= @@ -2758,7 +2881,7 @@ def _get_weight_generator( adapter_name: str = None, base_sync_done: bool = False, merge_and_sync: bool = False, - model_keys: List[str] = None, + model_keys: list[str] = None, **kwargs, ): if adapter_name is None: @@ -2770,6 +2893,7 @@ def _get_weight_generator( def _checkpoint_tensor(tensor): # Stage a (possibly CPU-offloaded) DTensor onto the device before the full_tensor all-gather. from torch.distributed.tensor import DTensor + if isinstance(tensor, DTensor): tensor = tensor.to(Torch.get_device(None), non_blocking=True) return Torch.to_local_tensor(tensor) @@ -2789,17 +2913,17 @@ def _tied_duplicate_names() -> set: own tensor), which would make a storage probe stop matching after the first sync. The flag can sit on a nested ``text_config`` for multimodal shells, so check both. """ - config = getattr(self, 'hf_config', None) + config = getattr(self, "hf_config", None) if config is None: return set() - tied = getattr(config, 'tie_word_embeddings', None) + tied = getattr(config, "tie_word_embeddings", None) if tied is None: - tied = getattr(getattr(config, 'text_config', None), 'tie_word_embeddings', None) + tied = getattr(getattr(config, "text_config", None), "tie_word_embeddings", None) if not tied: return set() # Drop the head, keep the embedding: vLLM's traversal reaches the nested embedding first, # making it the canonical name that ``_check_skipped_aliases`` wants to see loaded. - return {name for name in model.state_dict() if name.split('.')[-2:-1] == ['lm_head']} + return {name for name in model.state_dict() if name.split(".")[-2:-1] == ["lm_head"]} _tied_duplicates = _tied_duplicate_names() @@ -2826,6 +2950,7 @@ def weight_generator(): # LoRA-only mode: send only adapter weights. # Use PEFT's get_peft_model_state_dict for clean LoRA extraction from peft.utils import get_peft_model_state_dict + lora_state_dict = get_peft_model_state_dict(model, adapter_name=adapter_name) def weight_generator(): @@ -2873,6 +2998,7 @@ def _get_local_weight_generator(self): parallel_spec), so names and shapes match one-to-one without HF conversion. """ from twinkle.utils.framework import Torch + model = self.strategy.unwrap_model(self.model) for name, tensor in model.state_dict().items(): yield name, Torch.to_local_tensor(tensor) @@ -2880,5 +3006,6 @@ def _get_local_weight_generator(self): def _get_local_state_dict(self): """Return {name: param} for this rank's local shard (used by EMA receiver and naive sync).""" from twinkle.utils.framework import Torch + model = self.strategy.unwrap_model(self.model) return {name: Torch.to_local_tensor(tensor) for name, tensor in model.state_dict().items()} 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..36e650f6a 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 ----------- @@ -45,24 +45,25 @@ The role lives in the field's ``json_schema_extra`` so that one declaration site carries it -- there is deliberately no second per-endpoint parameter table. """ -from __future__ import annotations +from __future__ import annotations from enum import StrEnum +from typing import Any, Optional + from pydantic import BaseModel, ConfigDict, Field from pydantic.fields import FieldInfo -from typing import Any, Optional class StrictRequest(BaseModel): """Request bodies. Typos fail loudly.""" - model_config = ConfigDict(frozen=True, extra='forbid') + model_config = ConfigDict(frozen=True, extra="forbid") class ResponseModel(BaseModel): """Response bodies. An old client tolerates new server fields.""" - model_config = ConfigDict(frozen=True, extra='ignore') + model_config = ConfigDict(frozen=True, extra="ignore") class DataModel(BaseModel): @@ -81,24 +82,24 @@ class DataModel(BaseModel): which carry strict types; it does not belong to the field set. """ - model_config = ConfigDict(frozen=True, extra='allow') + model_config = ConfigDict(frozen=True, extra="allow") class FieldRole(StrEnum): """How a declared request field relates to the backend call.""" - Control = 'control' - BackendKwarg = 'backend_kwarg' - Passthrough = 'passthrough' + Control = "control" + BackendKwarg = "backend_kwarg" + Passthrough = "passthrough" # Keys under which field metadata is stored in a field's ``json_schema_extra``. -FIELD_ROLE_KEY = 'twinkle_field_role' -BACKEND_ONLY_KEY = 'twinkle_backend_only' +FIELD_ROLE_KEY = "twinkle_field_role" +BACKEND_ONLY_KEY = "twinkle_backend_only" def _with_extra(field_kwargs: dict[str, Any], **extra: Any) -> FieldInfo: - merged = dict(field_kwargs.pop('json_schema_extra', None) or {}) + merged = dict(field_kwargs.pop("json_schema_extra", None) or {}) merged.update(extra) return Field(json_schema_extra=merged, **field_kwargs) @@ -114,10 +115,13 @@ def backend_kwarg(*backends: str, **field_kwargs: Any) -> FieldInfo: constant default would make it carry a non-``None`` value on every deployment, so the "non-``None`` on the wrong backend" check would reject every request. """ - return _with_extra(field_kwargs, **{ - FIELD_ROLE_KEY: FieldRole.BackendKwarg.value, - BACKEND_ONLY_KEY: tuple(backends) or None, - }) + return _with_extra( + field_kwargs, + **{ + FIELD_ROLE_KEY: FieldRole.BackendKwarg.value, + BACKEND_ONLY_KEY: tuple(backends) or None, + }, + ) def backend_only(*backends: str, **field_kwargs: Any) -> FieldInfo: @@ -128,7 +132,7 @@ def backend_only(*backends: str, **field_kwargs: Any) -> FieldInfo: :func:`backend_kwarg` with a non-empty backend tuple, not a second mechanism. """ if not backends: - raise ValueError('backend_only() requires at least one backend; use backend_kwarg() for an unrestricted field') + raise ValueError("backend_only() requires at least one backend; use backend_kwarg() for an unrestricted field") return backend_kwarg(*backends, **field_kwargs) @@ -141,12 +145,12 @@ def passthrough(**field_kwargs: Any) -> FieldInfo: and ``inspect.signature`` cannot see such a read -- so any such check rejects valid requests. A misspelled plugin argument therefore still surfaces from the plugin. """ - field_kwargs.setdefault('default_factory', dict) + field_kwargs.setdefault("default_factory", dict) return _with_extra(field_kwargs, **{FIELD_ROLE_KEY: FieldRole.Passthrough.value}) def _read_extra(field_info: FieldInfo, key: str) -> Any: - extra = getattr(field_info, 'json_schema_extra', None) + extra = getattr(field_info, "json_schema_extra", None) if isinstance(extra, dict): return extra.get(key) return None diff --git a/src/twinkle/protocol/types/errors.py b/src/twinkle/protocol/types/errors.py index b896e4eed..cee969762 100644 --- a/src/twinkle/protocol/types/errors.py +++ b/src/twinkle/protocol/types/errors.py @@ -12,26 +12,31 @@ class can be lined up against its tinker counterpart): ``twinkle_client/utils/patch_tinker.py`` shows two SDKs can coexist in one process, so a semantically-equal but differently-named exception must be lookup-able. """ -from __future__ import annotations +from __future__ import annotations from enum import StrEnum -from pydantic import Field, field_validator, model_validator from typing import Any, Literal, Optional +from pydantic import Field, field_validator, model_validator + 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): """Error attribution. Matches tinker's ``RequestErrorCategory``.""" - Unknown = 'unknown' - Server = 'server' - User = 'user' + Unknown = "unknown" + Server = "server" + User = "user" class ErrorPayload(ResponseModel): @@ -54,15 +59,15 @@ class ErrorPayload(ResponseModel): traceback: str | None = Field(default=None, max_length=65536) details: list[dict[str, Any]] | None = None - @field_validator('category', mode='before') + @field_validator("category", mode="before") @classmethod def normalize_legacy_category(cls, value: Any) -> Any: if isinstance(value, str): return value.lower() return value - @model_validator(mode='after') + @model_validator(mode="after") def traceback_is_server_only(self) -> ErrorPayload: if self.traceback is not None and self.category is not ErrorCategory.Server: - raise ValueError('traceback is only valid for server errors') + raise ValueError("traceback is only valid for server errors") return self diff --git a/src/twinkle/protocol/types/lifecycle.py b/src/twinkle/protocol/types/lifecycle.py index 32555e24b..446f3618b 100644 --- a/src/twinkle/protocol/types/lifecycle.py +++ b/src/twinkle/protocol/types/lifecycle.py @@ -5,20 +5,27 @@ reverse-imports ``twinkle.protocol.types``); per the naming rulings in ``base.py`` it therefore intentionally carries **no** underscore prefix. """ -from __future__ import annotations -from typing import Any, Literal, Optional +from __future__ import annotations +from enum import StrEnum +from typing import Any, Optional from .base import ResponseModel, StrictRequest -from .errors import ErrorPayload, QueueStateLiteral +from .errors import ErrorPayload, QueueState + + +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" -# 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'}) +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..12d2675c1 100644 --- a/src/twinkle/protocol/types/model.py +++ b/src/twinkle/protocol/types/model.py @@ -17,10 +17,11 @@ home -- the passthrough regions -- so strictness never blocks a legitimate user-supplied argument. """ + from __future__ import annotations +from typing import Any, Optional, Union from pydantic import Field, JsonValue, field_validator, model_validator -from typing import Any, Dict, List, Optional, Union from .base import ResponseModel, StrictRequest, backend_kwarg, backend_only, passthrough from .component import DataRef @@ -28,7 +29,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 +39,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 +56,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) @@ -62,7 +71,7 @@ class LrStepRequest(AdapterRequest): # ``OptimizerParamScheduler.step(increment=...)``; the transformers scheduler has # no equivalent knob. - increment: int | None = backend_only('megatron', default=None, ge=0) + increment: int | None = backend_only("megatron", default=None, ge=0) class ClipGradNormRequest(AdapterRequest): @@ -83,10 +92,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 +113,7 @@ class CalculateMetricRequest(StrictRequest): # --------------------------------------------------------------------------- # -class _InlineForwardBase(StrictRequest): +class _InlineForwardBase(TrainingRequest): """Fields common to the three inline forward endpoints.""" inputs: WireInputBatch @@ -115,8 +123,8 @@ class _InlineForwardBase(StrictRequest): micro_batch_size: int | None = backend_kwarg(default=None, ge=1) gradient_accumulation_steps: int | None = backend_kwarg(default=None, ge=1) # Read only by the transformers backend. - sampling_masks: JsonValue | None = backend_only('transformers', default=None) - router_replay_action: str | None = backend_only('transformers', default=None) + sampling_masks: JsonValue | None = backend_only("transformers", default=None) + router_replay_action: str | None = backend_only("transformers", default=None) # Loss inputs (``advantages`` / ``old_logps`` / ``ref_outputs`` / ...). Their key # set is decided by the configured Loss, so they get a declared dict rather than # top-level fields; the flattening in ``backend_kwargs`` keeps the backend call @@ -127,7 +135,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 +146,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 +166,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 +174,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 +184,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 +195,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) @@ -202,10 +205,10 @@ class DataPlaneForwardOnlyRequest(StrictRequest): output_ref: DataRef | None = None output_fields: dict[str, str] = Field(default_factory=dict) - @model_validator(mode='after') + @model_validator(mode="after") def validate_output(self) -> DataPlaneForwardOnlyRequest: if (self.output_ref is None) != (len(self.output_fields) == 0): - raise ValueError('output_ref and output_fields must be configured together') + raise ValueError("output_ref and output_fields must be configured together") return self @@ -222,25 +225,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 +250,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,33 +274,32 @@ 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 consumed_train_samples: int | None = backend_kwarg(default=None, ge=0) - merge_lora: bool | None = backend_only('megatron', default=None) + 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) - no_load_rng: bool | None = backend_only('megatron', default=None) - strict: bool | None = backend_only('transformers', default=None) + no_load_optim: bool | None = backend_only("megatron", default=None) + no_load_rng: bool | None = backend_only("megatron", default=None) + strict: bool | None = backend_only("transformers", default=None) -class ResumeFromCheckpointRequest(StrictRequest): +class ResumeFromCheckpointRequest(TrainingRequest): """Body of ``POST /twinkle/resume_from_checkpoint``.""" name: str - adapter_name: str = '' + adapter_name: str = "" resume_only_model: bool = False 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. @@ -325,7 +320,7 @@ class UploadToHubRequest(StrictRequest): hub_model_id: str hub_token: str | None = None - @field_validator('checkpoint_dir', mode='before') + @field_validator("checkpoint_dir", mode="before") @classmethod def extract_checkpoint_dir(cls, v): """Accept a ``save`` response dict and take its twinkle path. @@ -334,9 +329,9 @@ def extract_checkpoint_dir(cls, v): wrong-shaped dict is a 422 naming the missing key instead of a 500. """ if isinstance(v, dict): - if 'twinkle_path' not in v: + if "twinkle_path" not in v: raise ValueError("checkpoint_dir dict must contain 'twinkle_path'") - return v['twinkle_path'] + return v["twinkle_path"] return v @@ -347,11 +342,13 @@ def extract_checkpoint_dir(cls, v): class OkResponse(ResponseModel): """Response for endpoints whose underlying method returns None.""" - status: str = 'ok' + + status: str = "ok" class ModelResult(ResponseModel): """Generic result wrapper; ``ModelResult`` is the retained historical public name.""" + result: Any @@ -360,42 +357,50 @@ class ModelResult(ResponseModel): class ForwardResponse(ResponseModel): """Response for /forward and /forward_only endpoints (returns ModelOutput).""" + result: Any class ForwardBackwardResponse(ResponseModel): """Response for /forward_backward endpoint (returns ModelOutput).""" + result: Any class CalculateLossResponse(ResponseModel): """Response for /calculate_loss endpoint (returns float).""" + result: float class ClipGradNormResponse(ResponseModel): """Response for /clip_grad_norm endpoint (returns float as str).""" + result: str class GetTrainConfigsResponse(ResponseModel): """Response for /get_train_configs endpoint (returns str).""" + result: str class CalculateMetricResponse(ResponseModel): """Response for /calculate_metric endpoint (returns Dict).""" + result: dict[str, Any] class SaveResponse(ResponseModel): """Response for /save endpoint (returns twinkle path + checkpoint dir).""" + twinkle_path: str checkpoint_dir: str | None = None class TrainingProgressResponse(ResponseModel): """Response for /resume_from_checkpoint endpoint.""" + result: dict[str, Any] @@ -418,6 +423,13 @@ class TrainingProgressResponse(ResponseModel): # --- Other responses --- +class TrainingHandleResponse(ResponseModel): + training_id: str + adapter_name: str + + class CreateResponse(ResponseModel): """Response for /create endpoint.""" - status: str = 'ok' + + 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 b7c6a2288..5888dd84d 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -1,18 +1,26 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import asyncio import atexit -import numpy as np import os import threading import uuid from copy import copy -from typing import Any, Dict, List, Optional, Type, Union +from typing import Any, Optional, Union + +import numpy as np from twinkle import DeviceMesh, get_logger, remote_class, remote_function, requires from twinkle.checkpoint_engine import CheckpointEngineMixin from twinkle.checkpoint_engine.delta.rollout.mixin import DeltaSamplerMixin -from twinkle.data_format import (InputFeature, PoolingParams, PoolingResponse, SampledSequence, SampleResponse, - SamplingParams, Trajectory) +from twinkle.data_format import ( + InputFeature, + PoolingParams, + PoolingResponse, + SampledSequence, + SampleResponse, + SamplingParams, + Trajectory, +) from twinkle.exceptions import FullWeightLoadError from twinkle.hub import HubOperation from twinkle.patch import Patch, apply_patch @@ -53,8 +61,8 @@ def _unwrap_mm_ref(entry: Any) -> Any: template's ``preprocess_image`` does -- yields a form the engine can load; any other entry (already a PIL image / path / bytes) passes through untouched, so text-only and pre-loaded media are unaffected. """ - if isinstance(entry, dict) and ('bytes' in entry or 'path' in entry): - return entry.get('bytes') or entry.get('path') + if isinstance(entry, dict) and ("bytes" in entry or "path" in entry): + return entry.get("bytes") or entry.get("path") return entry @@ -69,6 +77,7 @@ def _load_vllm_image(entry: Any) -> Any: image) passes through. Video is not a single image, so it only needs :func:`_unwrap_mm_ref`. """ from twinkle.utils.vision_tools import load_image + unwrapped = _unwrap_mm_ref(entry) if isinstance(unwrapped, (str, bytes, bytearray)): return load_image(unwrapped) @@ -78,7 +87,7 @@ def _load_vllm_image(entry: Any) -> Any: #: The actor's method-concurrency ceiling (Ray runs a decorated actor's methods one at a time #: unless this is raised). Shared with the server-side async backend, which decorates its own #: actor with the same ceiling, hence the public name. -SAMPLER_MAX_CONCURRENCY = max(1, int(os.environ.get('TWINKLE_SAMPLER_MAX_CONCURRENCY') or 24)) +SAMPLER_MAX_CONCURRENCY = max(1, int(os.environ.get("TWINKLE_SAMPLER_MAX_CONCURRENCY") or 24)) @remote_class(max_concurrency=SAMPLER_MAX_CONCURRENCY) @@ -94,7 +103,7 @@ class vLLMSampler(Sampler, CheckpointEngineMixin, PartialRolloutMixin, DeltaSamp # feature without one), and _generate_inputs already reads state.feature / engine_input_ids / media. accepts_rollout_state = True - def __init__(self, model_id: str, engine_args: Dict[str, Any] = None, device_mesh: DeviceMesh = None, **kwargs): + def __init__(self, model_id: str, engine_args: dict[str, Any] = None, device_mesh: DeviceMesh = None, **kwargs): """Initialize vLLMSampler. Args: @@ -105,10 +114,10 @@ def __init__(self, model_id: str, engine_args: Dict[str, Any] = None, device_mes device_mesh: Parallel configuration for data parallelism. **kwargs: Additional arguments. """ - os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn' - os.environ['VLLM_ENGINE_ITERATION_TIMEOUT_S'] = '86400' + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + os.environ["VLLM_ENGINE_ITERATION_TIMEOUT_S"] = "86400" super().__init__() - requires('vllm') + requires("vllm") self.model_id = model_id self.device_mesh = device_mesh @@ -122,7 +131,7 @@ def __init__(self, model_id: str, engine_args: Dict[str, Any] = None, device_mes # 3. By creating engine in the background thread's event loop, all async # operations stay in the same loop context self._async_loop = asyncio.new_event_loop() - self._async_thread = threading.Thread(target=self._run_event_loop, daemon=True, name='vLLMSampler-EventLoop') + self._async_thread = threading.Thread(target=self._run_event_loop, daemon=True, name="vLLMSampler-EventLoop") self._async_thread.start() # Engine request ids currently mid-generation (PartialRolloutMixin), so an in-place weight republish @@ -130,34 +139,35 @@ def __init__(self, model_id: str, engine_args: Dict[str, Any] = None, device_mes self._inflight_request_ids: set = set() from .vllm_engine import VLLMEngine + engine_kwargs = engine_args.copy() if engine_args else {} # Auto-detect tensor_parallel_size from CUDA_VISIBLE_DEVICES - if 'tensor_parallel_size' not in engine_kwargs: + if "tensor_parallel_size" not in engine_kwargs: tp_size = 1 - visible_devices = os.environ.get(Platform.visible_device_env(), '') + visible_devices = os.environ.get(Platform.visible_device_env(), "") if visible_devices: - num_gpus = len([d for d in visible_devices.split(',') if d.strip()]) + num_gpus = len([d for d in visible_devices.split(",") if d.strip()]) if num_gpus > 0: tp_size = num_gpus - logger.info(f'vLLM TP size: {tp_size}') - engine_kwargs['tensor_parallel_size'] = tp_size + logger.info(f"vLLM TP size: {tp_size}") + engine_kwargs["tensor_parallel_size"] = tp_size # Set unique seed per engine based on rank for diverse sampling across DP workers # User can override by passing 'seed' in engine_args - engine_seed = engine_kwargs.get('seed', None) + engine_seed = engine_kwargs.get("seed", None) if engine_seed is None: rank = Platform.get_rank() engine_seed = 42 + rank # set different seed to get different results - engine_kwargs['seed'] = engine_seed + engine_kwargs["seed"] = engine_seed # Create engine in the background event loop so all async operations # (including vLLM's internal background tasks) run in the same loop self.engine: VLLMEngine = self._run_in_loop(self._create_engine_async(VLLMEngine, model_id, engine_kwargs)) # fix: On NPU, monkey_patch_model can trigger Triton compatibility errors and abort sampler init. # The patch is still applied unconditionally; a real NPU skip needs a platform branch here. - self._run_in_loop(self.engine.engine.collective_rpc('monkey_patch_model')) + self._run_in_loop(self.engine.engine.collective_rpc("monkey_patch_model")) VLLMLoraWeights()(self) @@ -182,6 +192,7 @@ def _iter_in_loop(self, async_gen_coro): via a thread-safe queue. """ import queue as stdlib_queue + q: stdlib_queue.Queue = stdlib_queue.Queue() _SENTINEL = object() @@ -192,26 +203,32 @@ async def _drain(): except Exception as exc: q.put(exc) finally: - q.put(_SENTINEL) - - asyncio.run_coroutine_threadsafe(_drain(), self._async_loop) + try: + await async_gen_coro.aclose() + finally: + q.put(_SENTINEL) - 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.""" return engine_cls(model_id=model_id, **engine_kwargs) - def encode_trajectory_for_vllm(self, - trajectory: Trajectory, - adapter_name: str = '', - add_generation_prompt=True) -> Dict[str, Any]: + def encode_trajectory_for_vllm( + self, trajectory: Trajectory, adapter_name: str = "", add_generation_prompt=True + ) -> dict[str, Any]: """Encode trajectory for vLLM. Messages should already use transformers standard format (content is List[Dict]). @@ -229,47 +246,47 @@ def encode_trajectory_for_vllm(self, encoded[key] = encoded[key].tolist() return encoded - @remote_function(dispatch='all', collect='first', lazy_collect=False, exclusive=True, replay=True) - def apply_patch(self, patch_cls: Union[Patch, Type[Patch], str], **kwargs) -> None: + @remote_function(dispatch="all", collect="first", lazy_collect=False, exclusive=True, replay=True) + def apply_patch(self, patch_cls: Patch | type[Patch] | str, **kwargs) -> None: apply_patch(self, patch_cls, **kwargs) @staticmethod - def _extract_multi_modal_data(feat: Any) -> Optional[Dict[str, Any]]: + def _extract_multi_modal_data(feat: Any) -> dict[str, Any] | None: """Convert a canonical media bundle (or a legacy feature) to vLLM's request shape.""" if not isinstance(feat, dict): - images = getattr(feat, 'images', ()) - videos = getattr(feat, 'videos', ()) - audios = getattr(feat, 'audios', ()) + images = getattr(feat, "images", ()) + videos = getattr(feat, "videos", ()) + audios = getattr(feat, "audios", ()) mm_data = {} if images: - mm_data['image'] = [_load_vllm_image(image) for image in images] + mm_data["image"] = [_load_vllm_image(image) for image in images] if videos: - mm_data['video'] = [_unwrap_mm_ref(video) for video in videos] + mm_data["video"] = [_unwrap_mm_ref(video) for video in videos] if audios: - mm_data['audio'] = [_unwrap_mm_ref(audio) for audio in audios] + mm_data["audio"] = [_unwrap_mm_ref(audio) for audio in audios] return mm_data or None - images = feat.get('images') - videos = feat.get('videos') + images = feat.get("images") + videos = feat.get("videos") if not images and not videos: - for msg in feat.get('messages', []): - content = msg.get('content') + for msg in feat.get("messages", []): + content = msg.get("content") if not isinstance(content, list): continue for block in content: if not isinstance(block, dict): continue - btype = block.get('type') - if btype == 'image': - for key in ('image', 'url', 'path'): + btype = block.get("type") + if btype == "image": + for key in ("image", "url", "path"): if key in block and block[key] is not None: if images is None: images = [] images.append(block[key]) break - elif btype == 'video': - for key in ('video', 'url', 'path'): + elif btype == "video": + for key in ("video", "url", "path"): if key in block and block[key] is not None: if videos is None: videos = [] @@ -279,24 +296,24 @@ def _extract_multi_modal_data(feat: Any) -> Optional[Dict[str, Any]]: mm_data = {} if images: image_list = images if isinstance(images, (list, tuple)) else [images] - mm_data['image'] = [_load_vllm_image(im) for im in image_list] + mm_data["image"] = [_load_vllm_image(im) for im in image_list] if videos: video_list = videos if isinstance(videos, (list, tuple)) else [videos] - mm_data['video'] = [_unwrap_mm_ref(video) for video in video_list] + mm_data["video"] = [_unwrap_mm_ref(video) for video in video_list] return mm_data or None async def _sample_single( self, - feat: Dict[str, Any], + feat: dict[str, Any], sampling_params: SamplingParams, - lora_request: Optional[Any] = None, + lora_request: Any | None = None, *, - multi_modal_data: Optional[Dict[str, Any]] = None, + multi_modal_data: dict[str, Any] | None = None, logprobs_only: bool = False, disable_lora: bool = False, token_only: bool = False, - request_id: Optional[str] = None, - state: Optional[RolloutState] = None, + request_id: str | None = None, + state: RolloutState | None = None, ) -> SampleResponse: """Sample a single input asynchronously. @@ -307,8 +324,8 @@ async def _sample_single( if request_id is None: request_id = uuid.uuid4().hex if state is None and feature_has_media(feat): - raise NotImplementedError('multimodal sampling requires a Trajectory or RolloutState.') - prompt_ids = list(state.engine_input_ids) if state is not None else feat['input_ids'] + raise NotImplementedError("multimodal sampling requires a Trajectory or RolloutState.") + prompt_ids = list(state.engine_input_ids) if state is not None else feat["input_ids"] response = await self._sample_engine_request( request_id, self.engine.sample( @@ -316,32 +333,34 @@ async def _sample_single( sampling_params=sampling_params, lora_request=lora_request, multi_modal_data=multi_modal_data, - mm_processor_kwargs=feat.get('mm_processor_kwargs'), + mm_processor_kwargs=feat.get("mm_processor_kwargs"), disable_lora=disable_lora, request_id=request_id, ), ) - if 'input_ids' not in feat or multi_modal_data: - if 'input_ids' in feat: - if len(feat['input_ids']) != len(response.prompt_token_ids): - raise RuntimeError(f'Input ids length {len(feat["input_ids"])} does not ' - f'match prompt_token_ids length {len(response.prompt_token_ids)}') + if "input_ids" not in feat or multi_modal_data: + if "input_ids" in feat: + if len(feat["input_ids"]) != len(response.prompt_token_ids): + raise RuntimeError( + f"Input ids length {len(feat['input_ids'])} does not " + f"match prompt_token_ids length {len(response.prompt_token_ids)}" + ) else: - feat['input_ids'] = response.prompt_token_ids - feat['labels'] = [-100] * len(response.prompt_token_ids) + feat["input_ids"] = response.prompt_token_ids + feat["labels"] = [-100] * len(response.prompt_token_ids) # A sampling prompt (e.g. a tinker ModelInput) carries input_ids but no labels; # concat_input_feature would then derive a zero-length prefix completion_mask # and raise. The prompt is pure context, so materialise aligned all-context # labels when they are missing or length-mismatched. Present, aligned labels are # left untouched (preserving provenance), and the logprobs-only path -- which # never concatenates -- is not touched. - if not token_only and not logprobs_only and 'input_ids' in feat: + if not token_only and not logprobs_only and "input_ids" in feat: # ``labels`` may be an ndarray (np.roll output), whose truth value is ambiguous; # compare lengths directly instead of ``feat.get('labels') or []``. - labels = feat.get('labels') - if labels is None or len(labels) != len(feat['input_ids']): - feat['labels'] = [-100] * len(feat['input_ids']) + labels = feat.get("labels") + if labels is None or len(labels) != len(feat["input_ids"]): + feat["labels"] = [-100] * len(feat["input_ids"]) sequences = [] for seq in response.sequences: if token_only: @@ -354,7 +373,7 @@ async def _sample_single( elif logprobs_only: new_input_feature = _convert_ndarray_to_list(feat) if seq.routed_experts is not None: - new_input_feature['routed_experts'] = _convert_ndarray_to_list(seq.routed_experts) + new_input_feature["routed_experts"] = _convert_ndarray_to_list(seq.routed_experts) sampled_seq = SampledSequence( tokens=[], stop_reason=seq.stop_reason, @@ -364,7 +383,7 @@ async def _sample_single( completion = self.template.append_completion(state or feat, seq.tokens) new_input_feature = _convert_ndarray_to_list(completion.state.feature) if seq.routed_experts is not None: - new_input_feature['routed_experts'] = _convert_ndarray_to_list(seq.routed_experts) + new_input_feature["routed_experts"] = _convert_ndarray_to_list(seq.routed_experts) sampled_seq = SampledSequence( stop_reason=seq.stop_reason, tokens=seq.tokens, @@ -379,10 +398,11 @@ async def _sample_single( prompt_token_ids=response.prompt_token_ids, sequences=sequences, prompt_logprobs=response.prompt_logprobs, - topk_prompt_logprobs=response.topk_prompt_logprobs) + topk_prompt_logprobs=response.topk_prompt_logprobs, + ) @staticmethod - async def _gather_samples(samples) -> List[SampleResponse]: + async def _gather_samples(samples) -> list[SampleResponse]: """Finish every prompt before exposing the batch's result or failure.""" tasks = [asyncio.create_task(sample) for sample in samples] batch = asyncio.gather(*tasks) @@ -397,22 +417,27 @@ async def _gather_samples(samples) -> List[SampleResponse]: batch.exception() raise - @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False, enable_continous_work=True, - worker_selector='_sampling_worker_ids') + @remote_function( + dispatch="slice_dp", + collect="flatten", + lazy_collect=False, + enable_continous_work=True, + worker_selector="_sampling_worker_ids", + ) @timed def sample( self, - inputs: Union[InputFeature, RolloutState, Trajectory, List[Union[InputFeature, RolloutState, Trajectory]]], - sampling_params: Optional[Union[SamplingParams, Dict[str, Any]]] = None, - adapter_name: str = '', - adapter_path: Optional[str] = None, + inputs: InputFeature | RolloutState | Trajectory | list[InputFeature | RolloutState | Trajectory], + sampling_params: SamplingParams | dict[str, Any] | None = None, + adapter_name: str = "", + adapter_path: str | None = None, *, return_encoded: bool = False, use_base_model: bool = False, - adapter_paths: Optional[List[Optional[str]]] = None, + adapter_paths: list[str | None] | None = None, allow_partial_rollout: bool = False, token_only: bool = False, - ) -> List[SampleResponse]: + ) -> list[SampleResponse]: """Sample responses for given inputs. Args: @@ -470,21 +495,23 @@ def sample( use_base_model=use_base_model, adapter_paths=adapter_paths, allow_partial_rollout=allow_partial_rollout, - token_only=token_only)) + token_only=token_only, + ) + ) async def _generate_inputs( self, - inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], - sampling_params: Optional[Union[SamplingParams, Dict[str, Any]]] = None, - adapter_name: str = '', - adapter_path: Optional[str] = None, + inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], + sampling_params: SamplingParams | dict[str, Any] | None = None, + adapter_name: str = "", + adapter_path: str | None = None, *, return_encoded: bool = False, use_base_model: bool = False, - adapter_paths: Optional[List[Optional[str]]] = None, + adapter_paths: list[str | None] | None = None, allow_partial_rollout: bool = False, token_only: bool = False, - ) -> List[SampleResponse]: + ) -> list[SampleResponse]: """The async generation core shared by the blocking :meth:`sample` and the non-blocking ``submit``. ``sample`` drives this coroutine to completion on the background loop; ``submit`` schedules the @@ -507,21 +534,27 @@ async def _generate_inputs( if not raw_inputs: return [] input_states = [item if isinstance(item, RolloutState) else None for item in raw_inputs] - inputs_list = [state.feature if state is not None else item for item, state in zip(raw_inputs, input_states)] + inputs_list = [ + state.feature if state is not None else item for item, state in zip(raw_inputs, input_states, strict=False) + ] if adapter_paths is not None: if adapter_path is not None: - raise ValueError('Pass either adapter_path (one adapter for the whole call) or adapter_paths ' - '(one per input), not both.') + raise ValueError( + "Pass either adapter_path (one adapter for the whole call) or adapter_paths " + "(one per input), not both." + ) if len(adapter_paths) != len(inputs_list): - raise ValueError(f'adapter_paths has {len(adapter_paths)} entries but there are ' - f'{len(inputs_list)} inputs; they must correspond one-to-one so that DP ' - 'slicing keeps them aligned.') + raise ValueError( + f"adapter_paths has {len(adapter_paths)} entries but there are " + f"{len(inputs_list)} inputs; they must correspond one-to-one so that DP " + "slicing keeps them aligned." + ) # Check if inputs are Trajectory (not encoded) - aligned with Model.forward logic - is_trajectory = 'input_ids' not in inputs_list[0] + is_trajectory = "input_ids" not in inputs_list[0] if token_only and is_trajectory: - raise ValueError('Token-only sampling requires encoded input_ids.') + raise ValueError("Token-only sampling requires encoded input_ids.") logprobs_only = False if sampling_params.max_tokens == 0: sampling_params = copy(sampling_params) @@ -529,10 +562,9 @@ async def _generate_inputs( logprobs_only = True if is_trajectory: - assert self.template is not None, \ - 'Use set_template to add a template when trying to input Trajectory' + assert self.template is not None, "Use set_template to add a template when trying to input Trajectory" prompts = [ - self.template.encode_prompt(traj, backend='vllm', add_generation_prompt=not logprobs_only) + self.template.encode_prompt(traj, backend="vllm", add_generation_prompt=not logprobs_only) for traj in inputs_list ] states = [prompt.state for prompt in prompts] @@ -543,7 +575,7 @@ async def _generate_inputs( states = input_states multi_modal_data_list = [ self._extract_multi_modal_data(state.media if state is not None else feat) - for feat, state in zip(inputs_list, states) + for feat, state in zip(inputs_list, states, strict=False) ] if adapter_paths is not None: @@ -554,7 +586,7 @@ async def _generate_inputs( # ``request_id`` is the caller-side handle for async control (submit/pause/stop): taken from # the raw inputs before encoding (the encoder does not carry it) and handed to the engine as # the request id, so an abort can address this exact generation. - request_ids = [input.get('request_id') for input in inputs_list] + request_ids = [input.get("request_id") for input in inputs_list] async def _generate_one(feat, state, multi_modal_data, lora_request, request_id): if not allow_partial_rollout: @@ -599,42 +631,42 @@ async def run_attempt(current_input, current_state, attempt_params, attempt): return await self._gather_generation_tasks( _generate_one(feat, state, multi_modal_data, lora_request, request_id) - for feat, state, multi_modal_data, lora_request, request_id in zip(encoded_inputs, states, - multi_modal_data_list, lora_requests, - request_ids)) + for feat, state, multi_modal_data, lora_request, request_id in zip( + encoded_inputs, states, multi_modal_data_list, lora_requests, request_ids, strict=False + ) + ) async def _encode_single( self, - feat: Dict[str, Any], + feat: dict[str, Any], pooling_params: PoolingParams, - lora_request: Optional[Any] = None, + lora_request: Any | None = None, *, - multi_modal_data: Optional[Dict[str, Any]] = None, - state: Optional[RolloutState] = None, + multi_modal_data: dict[str, Any] | None = None, + state: RolloutState | None = None, ) -> PoolingResponse: """Run a single pooling forward asynchronously, the counterpart of ``_sample_single``.""" if state is None and feature_has_media(feat): - raise NotImplementedError('multimodal pooling requires a Trajectory or RolloutState.') + raise NotImplementedError("multimodal pooling requires a Trajectory or RolloutState.") return await self.engine.encode( - prompt=list(state.engine_input_ids) if state is not None else feat['input_ids'], + prompt=list(state.engine_input_ids) if state is not None else feat["input_ids"], pooling_params=pooling_params, lora_request=lora_request, multi_modal_data=multi_modal_data, - mm_processor_kwargs=feat.get('mm_processor_kwargs'), + mm_processor_kwargs=feat.get("mm_processor_kwargs"), ) - @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False, - worker_selector='_sampling_worker_ids') + @remote_function(dispatch="slice_dp", collect="flatten", lazy_collect=False, worker_selector="_sampling_worker_ids") @timed def encode( self, - inputs: Union[InputFeature, RolloutState, Trajectory, List[Union[InputFeature, RolloutState, Trajectory]]], - pooling_params: Optional[Union[PoolingParams, Dict[str, Any]]] = None, - adapter_name: str = '', - adapter_path: Optional[str] = None, + inputs: InputFeature | RolloutState | Trajectory | list[InputFeature | RolloutState | Trajectory], + pooling_params: PoolingParams | dict[str, Any] | None = None, + adapter_name: str = "", + adapter_path: str | None = None, *, - adapter_paths: Optional[List[Optional[str]]] = None, - ) -> List[PoolingResponse]: + adapter_paths: list[str | None] | None = None, + ) -> list[PoolingResponse]: """Run a pooling forward (embedding / classify) for each input. The pooling counterpart of :meth:`sample`, sharing its dispatch, trajectory encoding, LoRA @@ -661,25 +693,30 @@ def encode( raw_inputs = [inputs] if isinstance(inputs, RolloutState) else self._normalize_inputs(inputs) input_states = [item if isinstance(item, RolloutState) else None for item in raw_inputs] - inputs_list = [state.feature if state is not None else item for item, state in zip(raw_inputs, input_states)] + inputs_list = [ + state.feature if state is not None else item for item, state in zip(raw_inputs, input_states, strict=False) + ] if adapter_paths is not None: if adapter_path is not None: - raise ValueError('Pass either adapter_path (one adapter for the whole call) or adapter_paths ' - '(one per input), not both.') + raise ValueError( + "Pass either adapter_path (one adapter for the whole call) or adapter_paths " + "(one per input), not both." + ) if len(adapter_paths) != len(inputs_list): - raise ValueError(f'adapter_paths has {len(adapter_paths)} entries but there are ' - f'{len(inputs_list)} inputs; they must correspond one-to-one so that DP ' - 'slicing keeps them aligned.') + raise ValueError( + f"adapter_paths has {len(adapter_paths)} entries but there are " + f"{len(inputs_list)} inputs; they must correspond one-to-one so that DP " + "slicing keeps them aligned." + ) if not inputs_list: return [] - is_trajectory = 'input_ids' not in inputs_list[0] + is_trajectory = "input_ids" not in inputs_list[0] if is_trajectory: - assert self.template is not None, \ - 'Use set_template to add a template when trying to input Trajectory' + assert self.template is not None, "Use set_template to add a template when trying to input Trajectory" prompts = [ - self.template.encode_prompt(traj, backend='vllm', add_generation_prompt=False) for traj in inputs_list + self.template.encode_prompt(traj, backend="vllm", add_generation_prompt=False) for traj in inputs_list ] states = [prompt.state for prompt in prompts] encoded_inputs = [state.feature for state in states] @@ -688,7 +725,7 @@ def encode( encoded_inputs = inputs_list multi_modal_data_list = [ self._extract_multi_modal_data(state.media if state is not None else feat) - for feat, state in zip(encoded_inputs, states) + for feat, state in zip(encoded_inputs, states, strict=False) ] if adapter_paths is not None: @@ -704,14 +741,16 @@ async def _encode_all(): lora_request=lora_request, multi_modal_data=multi_modal_data, state=state, - ) for feat, state, multi_modal_data, lora_request in zip(encoded_inputs, states, multi_modal_data_list, - lora_requests) + ) + for feat, state, multi_modal_data, lora_request in zip( + encoded_inputs, states, multi_modal_data_list, lora_requests, strict=False + ) ] return await self._gather_samples(tasks) return self._run_in_loop(_encode_all()) - async def _aload_lora(self, adapter_path: Optional[str]): + async def _aload_lora(self, adapter_path: str | None): """Async :meth:`_load_lora`: resolve an adapter path to a vLLM ``LoRARequest`` on the event loop. Split out so the async generation core (``_generate_inputs``) can ``await`` it directly; calling @@ -719,11 +758,11 @@ async def _aload_lora(self, adapter_path: Optional[str]): """ if adapter_path is None: return None - logger.info(f'Loading LoRA from {adapter_path}') + logger.info(f"Loading LoRA from {adapter_path}") local_path = HubOperation.download_model(model_id_or_path=adapter_path) return await self.engine._get_or_load_lora(local_path) - def _load_lora(self, adapter_path: Optional[str]): + def _load_lora(self, adapter_path: str | None): """Resolve an adapter path to a vLLM ``LoRARequest``, or None for the base model. ``_get_or_load_lora`` caches by path, so repeating the same path across a batch of per-input @@ -733,10 +772,10 @@ def _load_lora(self, adapter_path: Optional[str]): def sample_stream( self, - inputs: Union[InputFeature, Trajectory, Dict[str, Any]], - sampling_params: Optional[Union[SamplingParams, Dict[str, Any]]] = None, - adapter_name: str = '', - adapter_path: Optional[str] = None, + inputs: InputFeature | Trajectory | dict[str, Any], + sampling_params: SamplingParams | dict[str, Any] | None = None, + adapter_name: str = "", + adapter_path: str | None = None, ): """Stream token deltas as they are generated by the vLLM engine. @@ -751,33 +790,35 @@ def sample_stream( state = inputs if isinstance(inputs, RolloutState) else None if state is not None: feat = state.feature - elif 'input_ids' not in inputs: + elif "input_ids" not in inputs: if self.template is None: raise ValueError(f"Template not set for adapter '{adapter_name}'. Use set_template() first.") - state = self.template.encode_prompt(inputs, backend='vllm').state + state = self.template.encode_prompt(inputs, backend="vllm").state feat = state.feature else: feat = inputs if state is None and feature_has_media(feat): - raise NotImplementedError('multimodal streaming requires a Trajectory or RolloutState.') + raise NotImplementedError("multimodal streaming requires a Trajectory or RolloutState.") lora_request = self._load_lora(adapter_path) - prompt = list(state.engine_input_ids) if state is not None else feat['input_ids'] + prompt = list(state.engine_input_ids) if state is not None else feat["input_ids"] yield from self._iter_in_loop( self.engine.generate_stream( prompt=prompt, sampling_params=sampling_params, lora_request=lora_request, - )) + ) + ) - def sample_stream_to_queue(self, queue, inputs, sampling_params=None, adapter_name='', adapter_path=None): + 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') + @remote_function(dispatch="all", collect="first") @timed def sleep(self, level: int = 1) -> None: """ @@ -785,28 +826,28 @@ def sleep(self, level: int = 1) -> None: """ self._run_in_loop(self.engine.sleep(level)) - @remote_function(dispatch='all', collect='first') + @remote_function(dispatch="all", collect="first") @timed - def wake_up(self, tags: List[str] = None) -> None: + def wake_up(self, tags: list[str] = None) -> None: self._run_in_loop(self.engine.wake_up(tags=tags)) - @remote_function(dispatch='all', collect='first') + @remote_function(dispatch="all", collect="first") def reset_prefix_cache(self): self._run_in_loop(self.engine.reset_prefix_cache()) - @remote_function(dispatch='all', collect='first') + @remote_function(dispatch="all", collect="first") def reset_mm_cache(self): self._run_in_loop(self.engine.reset_mm_cache()) - @remote_function(dispatch='all', collect='first') + @remote_function(dispatch="all", collect="first") def reset_encoder_cache(self): self._run_in_loop(self.engine.reset_encoder_cache()) - @remote_function(dispatch='all', collect='first', lazy_collect=False) + @remote_function(dispatch="all", collect="first", lazy_collect=False) def get_state_keys(self): return self._run_in_loop(self.engine.get_state_keys()) - @remote_function(dispatch='all', lazy_collect=True, exclusive=True) + @remote_function(dispatch="all", lazy_collect=True, exclusive=True) @timed def receive_weights( self, @@ -863,14 +904,14 @@ async def _receive_and_load(): self._run_in_loop(_receive_and_load()) - @remote_function(dispatch='all', collect='first', lazy_collect=False) + @remote_function(dispatch="all", collect="first", lazy_collect=False) def unload_adapter_paths(self, adapter_paths: list[str]) -> None: """Unload policy snapshots from vLLM and clear cached requests.""" self._run_in_loop(self.engine.unload_lora_paths(adapter_paths)) - @remote_function(dispatch='all', collect='first', lazy_collect=False, timeout=3600) + @remote_function(dispatch="all", collect="first", lazy_collect=False, timeout=3600) @timed - def load_full_weights_from_path(self, path: Optional[str] = None) -> int: + def load_full_weights_from_path(self, path: str | None = None) -> int: """Load a full (non-LoRA) HF checkpoint into the engine's base model. Unlike :meth:`receive_weights`, this does **not** involve the training model: @@ -900,36 +941,37 @@ def load_full_weights_from_path(self, path: Optional[str] = None) -> int: 1 if weights were (re)loaded, 0 if that path was already loaded. """ import glob + import json from safetensors import safe_open path = path or self.model_id resolved = path if os.path.exists(path) else HubOperation.download_model(path) - if getattr(self, '_loaded_full_weights_path', None) == resolved: + if getattr(self, "_loaded_full_weights_path", None) == resolved: return 0 # Resolve the shard list eagerly so a bad path fails here rather than # part-way through streaming tensors into a live engine. - index_path = os.path.join(resolved, 'model.safetensors.index.json') + index_path = os.path.join(resolved, "model.safetensors.index.json") if os.path.exists(index_path): - with open(index_path, encoding='utf-8') as f: - weight_map = json.load(f)['weight_map'] + with open(index_path, encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] shards = [os.path.join(resolved, s) for s in sorted(set(weight_map.values()))] else: - shards = sorted(glob.glob(os.path.join(resolved, '*.safetensors'))) + shards = sorted(glob.glob(os.path.join(resolved, "*.safetensors"))) if not shards: - raise FileNotFoundError(f'No .safetensors weights found under {resolved}') + raise FileNotFoundError(f"No .safetensors weights found under {resolved}") # Check every shard header before changing any live weights. for shard in shards: - with safe_open(shard, framework='pt', device='cpu') as f: + with safe_open(shard, framework="pt", device="cpu") as f: if not list(f.keys()): - raise ValueError(f'Empty checkpoint shard: {shard}') + raise ValueError(f"Empty checkpoint shard: {shard}") def _iter_weights(): # safe_open + get_tensor reads one tensor at a time (mmap-backed), so peak # host memory is a single tensor rather than the whole shard. for shard in shards: - with safe_open(shard, framework='pt', device='cpu') as f: + with safe_open(shard, framework="pt", device="cpu") as f: for name in f.keys(): yield name, f.get_tensor(name) @@ -939,17 +981,17 @@ async def _load(): # mirroring the `not base_sync_done` branch of receive_weights(). self.engine.invalidate_synced_lora() - logger.info(f'Loading full-parameter weights into sampler base model from {resolved}') + logger.info(f"Loading full-parameter weights into sampler base model from {resolved}") try: self._run_in_loop(_load()) self.reset_prefix_cache() except Exception as error: - raise FullWeightLoadError(f'Failed to install checkpoint {resolved}') from error + raise FullWeightLoadError(f"Failed to install checkpoint {resolved}") from error self._loaded_full_weights_path = resolved - logger.info(f'Reloaded base weights from {resolved} ({len(shards)} shard(s))') + logger.info(f"Reloaded base weights from {resolved} ({len(shards)} shard(s))") return 1 - @remote_function(dispatch='all', collect='first', lazy_collect=False) + @remote_function(dispatch="all", collect="first", lazy_collect=False) def shutdown(self): """Gracefully shutdown the vLLM engine and background event loop. @@ -963,16 +1005,16 @@ def shutdown(self): # 1. Shutdown vLLM engine (stops EngineCore process and output_handler) try: - if hasattr(self, 'engine') and self.engine is not None: + if hasattr(self, "engine") and self.engine is not None: self._run_in_loop(self.engine.shutdown()) except Exception as e: - logger.warning(f'vLLMSampler engine shutdown error: {e}') + logger.warning(f"vLLMSampler engine shutdown error: {e}") # 2. Stop the background event loop and join thread try: - if hasattr(self, '_async_loop') and self._async_loop.is_running(): + if hasattr(self, "_async_loop") and self._async_loop.is_running(): self._async_loop.call_soon_threadsafe(self._async_loop.stop) - if hasattr(self, '_async_thread') and self._async_thread.is_alive(): + if hasattr(self, "_async_thread") and self._async_thread.is_alive(): self._async_thread.join(timeout=5) except Exception as e: - logger.warning(f'vLLMSampler event loop shutdown error: {e}') + logger.warning(f"vLLMSampler event loop shutdown error: {e}") diff --git a/src/twinkle/server/__init__.py b/src/twinkle/server/__init__.py index 5bdb1f97c..d84e77426 100644 --- a/src/twinkle/server/__init__.py +++ b/src/twinkle/server/__init__.py @@ -1,7 +1,9 @@ # 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..34cdb46f7 100644 --- a/src/twinkle/server/checkpoint/checkpoint_manager.py +++ b/src/twinkle/server/checkpoint/checkpoint_manager.py @@ -9,10 +9,9 @@ weights, publish a version link, commit run metadata, then remove obsolete weights; failures before the commit restore the previous public path. """ -from __future__ import annotations +from __future__ import annotations import asyncio -import json import os import re import shutil @@ -22,12 +21,14 @@ from collections.abc import Awaitable from datetime import datetime from pathlib import Path -from typing import Any, Dict, List, Optional, Protocol +from typing import Any, Optional, Protocol + +import json 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 @@ -38,8 +39,7 @@ class _CheckpointWeightSaver(Protocol): """Backend callback writes weights into ``output_dir / name``.""" - def __call__(self, *, name: str, output_dir: str) -> Awaitable[str | None]: - ... + def __call__(self, *, name: str, output_dir: str) -> Awaitable[str | None]: ... def _remove_checkpoint_path(path: Path) -> None: @@ -63,11 +63,11 @@ def __init__(self, save_dir: Path, staging_dir: Path, name: str, version_id: str self.staged_checkpoint = staging_dir / name self.checkpoint_path = save_dir / name self._save_dir = save_dir - self._version_dir = save_dir / '.versions' / version_id + self._version_dir = save_dir / ".versions" / version_id self._version_dir.parent.mkdir(exist_ok=True) self._previous_version = self.checkpoint_path.resolve() if self.checkpoint_path.exists() else None - self._backup_path = staging_dir / '.previous' - self._new_link = staging_dir / '.new' + self._backup_path = staging_dir / ".previous" + self._new_link = staging_dir / ".new" self._live_alias = save_dir / live_alias_name if live_alias_name else None self._link_published = False @@ -86,7 +86,7 @@ def publish(self, metadata: dict[str, Any]) -> None: self._new_link.replace(self.checkpoint_path) self._link_published = True if self._live_alias is not None: - self._live_alias.symlink_to('latest', target_is_directory=True) + self._live_alias.symlink_to("latest", target_is_directory=True) def rollback(self) -> None: """Restore the previous public path and discard the uncommitted version.""" @@ -103,8 +103,9 @@ def cleanup_obsolete(self) -> None: previous = self._previous_version versions_dir = self._version_dir.parent.resolve() if previous is not None and previous.parent == versions_dir: - still_referenced = any(path.is_symlink() and path.resolve() == previous - for path in self._save_dir.iterdir()) + still_referenced = any( + path.is_symlink() and path.resolve() == previous for path in self._save_dir.iterdir() + ) if not still_referenced: shutil.rmtree(previous, ignore_errors=True) @@ -143,19 +144,21 @@ def path_prefix(self) -> str: pass @abstractmethod - def _create_checkpoint(self, - checkpoint_id: str, - checkpoint_type: str, - path: str, - size_bytes: int, - public: bool, - base_model: str | None = None, - is_lora: bool = False, - lora_rank: int | None = None, - train_unembed: bool | None = None, - train_mlp: bool | None = None, - train_attn: bool | None = None, - user_metadata: dict[str, Any] | None = None) -> dict[str, Any]: + def _create_checkpoint( + self, + checkpoint_id: str, + checkpoint_type: str, + path: str, + size_bytes: int, + public: bool, + base_model: str | None = None, + is_lora: bool = False, + lora_rank: int | None = None, + train_unembed: bool | None = None, + train_mlp: bool | None = None, + train_attn: bool | None = None, + user_metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: """ Create checkpoint data. @@ -252,7 +255,7 @@ def get_save_dir(self, model_id: str, is_sampler: bool = False) -> str: Returns: String path to save directory """ - weights_type = 'sampler_weights' if is_sampler else 'weights' + weights_type = "sampler_weights" if is_sampler else "weights" save_path = self.training_run_manager.get_model_dir(model_id) / weights_type return save_path.as_posix() @@ -261,9 +264,9 @@ def get_ckpt_name(name: str | None) -> str: """Generate or normalize checkpoint name.""" if name: # Normalize name to avoid issues with filesystem - name = re.sub(r'[^\w\-]', '_', name) + name = re.sub(r"[^\w\-]", "_", name) return name - return datetime.now().strftime('%Y%m%d_%H%M%S') + return datetime.now().strftime("%Y%m%d_%H%M%S") def _read_ckpt_info(self, model_id: str, checkpoint_id: str) -> dict[str, Any] | None: """ @@ -297,21 +300,18 @@ def _write_ckpt_info(self, model_id: str, checkpoint_id: str, data: dict[str, An ckpt_dir = self.get_ckpt_dir(model_id, checkpoint_id) ckpt_dir.mkdir(parents=True, exist_ok=True) meta_path = ckpt_dir / CHECKPOINT_INFO_FILENAME - with open(meta_path, 'w') as f: + with open(meta_path, "w") as f: json.dump(data, f, indent=2) - def _build_checkpoint_metadata(self, - model_id: str, - name: str, - checkpoint_path: Path, - is_sampler: bool, - public: bool = False) -> dict[str, Any]: + def _build_checkpoint_metadata( + self, model_id: str, name: str, checkpoint_path: Path, is_sampler: bool, public: bool = False + ) -> dict[str, Any]: if not validate_user_path(self.token, name): - raise ValueError(f'Invalid checkpoint name: {name}') - weights_type = 'sampler_weights' if is_sampler else 'weights' - checkpoint_type = 'sampler' if is_sampler else 'training' - checkpoint_id = f'{weights_type}/{name}' - path = f'{self.path_prefix}{model_id}/{checkpoint_id}' + raise ValueError(f"Invalid checkpoint name: {name}") + weights_type = "sampler_weights" if is_sampler else "weights" + checkpoint_type = "sampler" if is_sampler else "training" + checkpoint_id = f"{weights_type}/{name}" + path = f"{self.path_prefix}{model_id}/{checkpoint_id}" run_info = self.training_run_manager._read_info(model_id) return self._create_checkpoint( checkpoint_id=checkpoint_id, @@ -319,13 +319,14 @@ def _build_checkpoint_metadata(self, path=path, size_bytes=self.get_dir_size(checkpoint_path), public=public, - base_model=run_info.get('base_model'), - is_lora=run_info.get('is_lora', False), - lora_rank=run_info.get('lora_rank'), - train_unembed=run_info.get('train_unembed'), - train_mlp=run_info.get('train_mlp'), - train_attn=run_info.get('train_attn'), - user_metadata=run_info.get('user_metadata')) + base_model=run_info.get("base_model"), + is_lora=run_info.get("is_lora", False), + lora_rank=run_info.get("lora_rank"), + train_unembed=run_info.get("train_unembed"), + train_mlp=run_info.get("train_mlp"), + train_attn=run_info.get("train_attn"), + user_metadata=run_info.get("user_metadata"), + ) def save(self, model_id: str, name: str, is_sampler: bool = False, public: bool = False) -> str: """Record metadata and the corresponding last-checkpoint pointer. @@ -333,21 +334,23 @@ def save(self, model_id: str, name: str, is_sampler: bool = False, public: bool This compatibility API does not write weights. Server handlers use ``save_training`` or ``save_sampler`` to coordinate backend writes. """ - weights_type = 'sampler_weights' if is_sampler else 'weights' - checkpoint_id = f'{weights_type}/{name}' + weights_type = "sampler_weights" if is_sampler else "weights" + checkpoint_id = f"{weights_type}/{name}" ckpt_data = self._build_checkpoint_metadata( model_id=model_id, name=name, checkpoint_path=self.get_ckpt_dir(model_id, checkpoint_id), is_sampler=is_sampler, - public=public) + public=public, + ) self._write_ckpt_info(model_id, checkpoint_id, ckpt_data) - field = 'last_sampler_checkpoint' if is_sampler else 'last_checkpoint' + field = "last_sampler_checkpoint" if is_sampler else "last_checkpoint" self.training_run_manager.update(model_id, {field: ckpt_data}) - return f'{self.path_prefix}{model_id}/{checkpoint_id}' + return f"{self.path_prefix}{model_id}/{checkpoint_id}" - async def save_training(self, model_id: str, name: str | None, - save_weights: _CheckpointWeightSaver) -> tuple[str, str | None]: + async def save_training( + self, model_id: str, name: str | None, save_weights: _CheckpointWeightSaver + ) -> tuple[str, str | None]: """Write training weights before recording a successful checkpoint. The backend retains ownership of training files and optimizer state. @@ -360,8 +363,9 @@ async def save_training(self, model_id: str, name: str | None, uri = self.save(model_id, checkpoint_name) return uri, checkpoint_dir - async def save_sampler(self, model_id: str, name: str | None, - save_weights: _CheckpointWeightSaver) -> tuple[str, str]: + async def save_sampler( + self, model_id: str, name: str | None, save_weights: _CheckpointWeightSaver + ) -> tuple[str, str]: """Save weights, publish their link, then commit the run metadata. Named saves persist; unnamed saves replace ``latest`` and return a unique @@ -370,35 +374,37 @@ async def save_sampler(self, model_id: str, name: str | None, Both client dialects share the lock and the failure rollback below. """ 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.') + if explicit_name == "latest": + 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}' + checkpoint_name = explicit_name or "latest" + uri_name = explicit_name or f"live_{version_id}" live_alias_name = None if explicit_name else uri_name save_dir = Path(self.get_save_dir(model_id, is_sampler=True)) save_dir.mkdir(parents=True, exist_ok=True) - lock = FileLock(str(save_dir / '.save.lock')) + lock = FileLock(str(save_dir / ".save.lock")) # Use the shared helpers; blocking acquisition would stall async saves. while not acquire_lock(lock, blocking=False): await asyncio.sleep(0.05) try: - with tempfile.TemporaryDirectory(prefix='.pending-', dir=save_dir, ignore_cleanup_errors=True) as staging: + with tempfile.TemporaryDirectory(prefix=".pending-", dir=save_dir, ignore_cleanup_errors=True) as staging: files = _SamplerCheckpointFiles( save_dir=save_dir, staging_dir=Path(staging), name=checkpoint_name, version_id=version_id, - live_alias_name=live_alias_name) + live_alias_name=live_alias_name, + ) try: await save_weights(name=checkpoint_name, output_dir=staging) if not files.staged_checkpoint.is_dir(): - raise RuntimeError('Backend did not materialize the sampler checkpoint.') + raise RuntimeError("Backend did not materialize the sampler checkpoint.") metadata = self._build_checkpoint_metadata( - model_id=model_id, name=uri_name, checkpoint_path=files.staged_checkpoint, is_sampler=True) + model_id=model_id, name=uri_name, checkpoint_path=files.staged_checkpoint, is_sampler=True + ) files.publish(metadata) # Commit point: backups must survive until this atomic write succeeds. - self.training_run_manager.update(model_id, {'last_sampler_checkpoint': metadata}) + self.training_run_manager.update(model_id, {"last_sampler_checkpoint": metadata}) except BaseException: # Cancellation must restore the previous version too. files.rollback() @@ -406,7 +412,7 @@ async def save_sampler(self, model_id: str, name: str | None, files.cleanup_obsolete() finally: release_lock(lock) - uri = f'{self.path_prefix}{model_id}/sampler_weights/{uri_name}' + uri = f"{self.path_prefix}{model_id}/sampler_weights/{uri_name}" return uri, files.checkpoint_path.as_posix() def get(self, model_id: str, checkpoint_id: str) -> Any | None: @@ -437,13 +443,13 @@ def list_checkpoints(self, model_id: str) -> Any | None: checkpoints = [] # Iterate over weights and sampler_weights directories - for weights_type in ['weights', 'sampler_weights']: + for weights_type in ["weights", "sampler_weights"]: type_dir = run_dir / weights_type if not type_dir.exists() or not type_dir.is_dir(): continue for d in type_dir.iterdir(): - if not d.name.startswith('.') and d.is_dir() and (d / CHECKPOINT_INFO_FILENAME).exists(): - checkpoint_id = f'{weights_type}/{d.name}' + if not d.name.startswith(".") and d.is_dir() and (d / CHECKPOINT_INFO_FILENAME).exists(): + checkpoint_id = f"{weights_type}/{d.name}" ckpt = self.get(model_id, checkpoint_id) if ckpt and ckpt.checkpoint_id == checkpoint_id: checkpoints.append(ckpt) @@ -457,10 +463,22 @@ 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: + 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) @@ -471,10 +489,10 @@ def delete(self, model_id: str, checkpoint_id: str) -> bool: if ckpt_dir.is_symlink(): target = ckpt_dir.resolve() ckpt_dir.unlink() - latest = self.get_ckpt_dir(model_id, 'sampler_weights/latest') + latest = self.get_ckpt_dir(model_id, "sampler_weights/latest") if target == latest.resolve() and latest.exists(): _remove_checkpoint_path(latest) - if target.parent == (ckpt_dir.parent / '.versions').resolve(): + if target.parent == (ckpt_dir.parent / ".versions").resolve(): shutil.rmtree(target) elif ckpt_dir.is_dir(): shutil.rmtree(ckpt_dir) @@ -483,11 +501,11 @@ def delete(self, model_id: str, checkpoint_id: str) -> bool: # Refresh only the pointer for the deleted checkpoint's kind. all_ckpts = self.list_checkpoints(model_id) - kind = 'sampler' if checkpoint_id.startswith('sampler_weights/') else 'training' + kind = "sampler" if checkpoint_id.startswith("sampler_weights/") else "training" remaining = [ckpt for ckpt in all_ckpts.checkpoints if ckpt.checkpoint_type == kind] if all_ckpts else [] last_ckpt = remaining[-1] if remaining else None - field = 'last_sampler_checkpoint' if kind == 'sampler' else 'last_checkpoint' - self.training_run_manager.update(model_id, {field: last_ckpt.model_dump(mode='json') if last_ckpt else None}) + field = "last_sampler_checkpoint" if kind == "sampler" else "last_checkpoint" + self.training_run_manager.update(model_id, {field: last_ckpt.model_dump(mode="json") if last_ckpt else None}) return True def parse_path(self, path: str) -> Any | None: @@ -502,17 +520,17 @@ def parse_path(self, path: str) -> Any | None: """ if not path.startswith(self.path_prefix): return None - parts = path[len(self.path_prefix):].split('/') + parts = path[len(self.path_prefix) :].split("/") if len(parts) != 3: return None - if parts[1] not in ['weights', 'sampler_weights']: + if parts[1] not in ["weights", "sampler_weights"]: return None - checkpoint_type = 'training' if parts[1] == 'weights' else 'sampler' + checkpoint_type = "training" if parts[1] == "weights" else "sampler" return self._create_parsed_path( path=path, training_run_id=parts[0], checkpoint_type=checkpoint_type, - checkpoint_id='/'.join(parts[1:]), + checkpoint_id="/".join(parts[1:]), ) def get_weights_info(self, checkpoint_path: str) -> Any | None: @@ -526,7 +544,7 @@ def get_weights_info(self, checkpoint_path: str) -> Any | None: return self._get_weights_info_from_hub(checkpoint_path) ckpt_data = self._read_ckpt_info(resolved.training_run_id, resolved.checkpoint_id) - if not ckpt_data or not ckpt_data.get('base_model'): + if not ckpt_data or not ckpt_data.get("base_model"): return None return self._create_weights_info(ckpt_data) @@ -543,7 +561,8 @@ def _get_weights_info_from_hub(self, hub_model_id: str) -> Any | None: try: # Download only the checkpoint_metadata.json file from hub local_dir = HubOperation.download_file( - repo_id=hub_model_id, allow_patterns=[CHECKPOINT_INFO_FILENAME], token=self.token) + repo_id=hub_model_id, allow_patterns=[CHECKPOINT_INFO_FILENAME], token=self.token + ) # Read and parse the metadata metadata_path = os.path.join(local_dir, CHECKPOINT_INFO_FILENAME) @@ -553,7 +572,7 @@ def _get_weights_info_from_hub(self, hub_model_id: str) -> Any | None: with open(metadata_path) as f: ckpt_data = json.load(f) - if not ckpt_data.get('base_model'): + if not ckpt_data.get("base_model"): return None return self._create_weights_info(ckpt_data) @@ -569,14 +588,14 @@ def parse_adapter_uri(self, adapter_uri: str) -> tuple[str, str]: for nonstandard service URIs are preserved. """ if not adapter_uri.startswith(self.path_prefix): - return 'default', adapter_uri + return "default", adapter_uri parsed = self.parse_path(adapter_uri) if not parsed: - return 'default', adapter_uri[len(self.path_prefix):] + return "default", adapter_uri[len(self.path_prefix) :] checkpoint_dir = self.get_ckpt_dir(parsed.training_run_id, parsed.checkpoint_id) - if parsed.checkpoint_type == 'sampler': + if parsed.checkpoint_type == "sampler": checkpoint_dir = checkpoint_dir.resolve() return parsed.training_run_id, str(checkpoint_dir) @@ -594,30 +613,32 @@ def resolve_load_path(self, path: str, validate_exists: bool = True) -> Resolved checkpoint_dir=None, is_twinkle_path=False, training_run_id=None, - checkpoint_id=None) + checkpoint_id=None, + ) parsed = self.parse_path(path) if not parsed: - raise ValueError(f'Invalid {self.path_prefix} path format: {path}') + raise ValueError(f"Invalid {self.path_prefix} path format: {path}") training_run_id = parsed.training_run_id checkpoint_id = parsed.checkpoint_id - checkpoint_name = checkpoint_id.split('/')[-1] + checkpoint_name = checkpoint_id.split("/")[-1] if validate_exists: checkpoint = self.get(training_run_id, checkpoint_id) if not checkpoint: - raise ValueError(f'Checkpoint not found or access denied: {path}') + raise ValueError(f"Checkpoint not found or access denied: {path}") checkpoint_dir = self.get_ckpt_dir(training_run_id, checkpoint_id).parent if validate_exists and not checkpoint_dir.exists(): - raise ValueError(f'Checkpoint directory not found: {checkpoint_dir}') + raise ValueError(f"Checkpoint directory not found: {checkpoint_dir}") return ResolvedLoadPath( checkpoint_name=checkpoint_name, checkpoint_dir=checkpoint_dir.as_posix(), is_twinkle_path=True, training_run_id=training_run_id, - checkpoint_id=checkpoint_id) + checkpoint_id=checkpoint_id, + ) def resolve_checkpoint_dir(self, path: str) -> str: """Resolve a full directory for resume/upload; preserve external paths. diff --git a/src/twinkle/server/config/__init__.py b/src/twinkle/server/config/__init__.py index 3e3fe2f71..88ddc00b0 100644 --- a/src/twinkle/server/config/__init__.py +++ b/src/twinkle/server/config/__init__.py @@ -1,22 +1,32 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Server configuration package — aggregate root and per-deployment specs.""" -from .application_spec import (ApplicationSpec, DataPlaneArgs, HttpOptions, ModelArgs, ProcessorArgs, SamplerArgs, - ServerArgs, SpectralHybridArgs) +from .application_spec import ( + ApplicationSpec, + DataPlaneArgs, + HttpOptions, + ModelArgs, + ProcessorArgs, + SamplerArgs, + 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', - 'DataPlaneArgs', - 'HttpOptions', - 'ModelArgs', - 'PersistenceConfig', - 'ProcessorArgs', - 'SamplerArgs', - 'ServerArgs', - 'SpectralHybridArgs', - 'ServerConfig', - 'TelemetryConfig', + "ApplicationSpec", + "DataPlaneArgs", + "HttpOptions", + "ModelArgs", + "PersistenceConfig", + "ProcessorArgs", + "SamplerArgs", + "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..993b8a378 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -9,14 +9,15 @@ The schemas use ``extra='forbid'`` so unknown args (typos, copy-paste from other deployments) fail at load time instead of being silently dropped. """ -from __future__ import annotations -import math +from __future__ import annotations import os -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing import Any, Literal -from twinkle.server.task_queue.config import TaskQueueConfig +from pydantic import BaseModel, ConfigDict, Field, model_validator + +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 @@ -24,10 +25,10 @@ # defaults. Only the four *policy* fields are propagated; ``actor_name`` is a # per-process cache key, not a cross-worker policy, so it is excluded. SERVER_STATE_ENV_KEYS: tuple[str, ...] = ( - 'TWINKLE_SERVER_STATE_EXPIRATION_TIMEOUT', - 'TWINKLE_SERVER_STATE_CLEANUP_INTERVAL', - 'TWINKLE_SERVER_STATE_PER_TOKEN_MODEL_LIMIT', - 'TWINKLE_SERVER_STATE_METRICS_UPDATE_INTERVAL', + "TWINKLE_SERVER_STATE_EXPIRATION_TIMEOUT", + "TWINKLE_SERVER_STATE_CLEANUP_INTERVAL", + "TWINKLE_SERVER_STATE_PER_TOKEN_MODEL_LIMIT", + "TWINKLE_SERVER_STATE_METRICS_UPDATE_INTERVAL", ) # ---------- shared helpers ------------------------------------------------- # @@ -36,7 +37,7 @@ class _ArgsBase(BaseModel): """Base class for every per-deployment args schema (extra='forbid').""" - model_config = ConfigDict(extra='forbid') + model_config = ConfigDict(extra="forbid") class HttpOptions(BaseModel): @@ -49,86 +50,23 @@ class HttpOptions(BaseModel): listener on plain HTTP. """ - model_config = ConfigDict(extra='forbid') + model_config = ConfigDict(extra="forbid") - host: str = 'localhost' + host: str = "localhost" port: int = 8000 ssl_keyfile: str | None = None ssl_certfile: str | None = None - @model_validator(mode='after') - def _ssl_pair_set_together(self) -> 'HttpOptions': + @model_validator(mode="after") + def _ssl_pair_set_together(self) -> HttpOptions: if (self.ssl_keyfile is None) != (self.ssl_certfile is None): - raise ValueError('ssl_keyfile and ssl_certfile must be set together to serve HTTPS') + raise ValueError("ssl_keyfile and ssl_certfile must be set together to serve HTTPS") return self # ---------- 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. @@ -145,11 +83,33 @@ class SamplerArgs(_ArgsBase): nproc_per_node: int = 1 device_group: dict[str, Any] device_mesh: dict[str, Any] - sampler_type: Literal['mock', 'vllm', 'vllm_async', 'sglang', 'sglang_async', 'torch'] + sampler_type: Literal["mock", "vllm", "vllm_async", "sglang", "sglang_async", "torch"] engine_args: dict[str, Any] | None = None 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): @@ -179,13 +139,13 @@ def to_env_vars(self) -> dict[str, str]: """ env: dict[str, str] = {} if self.expiration_timeout is not None: - env['TWINKLE_SERVER_STATE_EXPIRATION_TIMEOUT'] = str(self.expiration_timeout) + env["TWINKLE_SERVER_STATE_EXPIRATION_TIMEOUT"] = str(self.expiration_timeout) if self.cleanup_interval is not None: - env['TWINKLE_SERVER_STATE_CLEANUP_INTERVAL'] = str(self.cleanup_interval) + env["TWINKLE_SERVER_STATE_CLEANUP_INTERVAL"] = str(self.cleanup_interval) if self.per_token_model_limit is not None: - env['TWINKLE_SERVER_STATE_PER_TOKEN_MODEL_LIMIT'] = str(self.per_token_model_limit) + env["TWINKLE_SERVER_STATE_PER_TOKEN_MODEL_LIMIT"] = str(self.per_token_model_limit) if self.metrics_update_interval is not None: - env['TWINKLE_SERVER_STATE_METRICS_UPDATE_INTERVAL'] = str(self.metrics_update_interval) + env["TWINKLE_SERVER_STATE_METRICS_UPDATE_INTERVAL"] = str(self.metrics_update_interval) return env @classmethod @@ -197,10 +157,10 @@ def from_env(cls) -> ServerStateArgs | None: configured to a default value". Only policy fields are populated; ``actor_name`` stays ``None``. """ - exp = os.environ.get('TWINKLE_SERVER_STATE_EXPIRATION_TIMEOUT') - clean = os.environ.get('TWINKLE_SERVER_STATE_CLEANUP_INTERVAL') - limit = os.environ.get('TWINKLE_SERVER_STATE_PER_TOKEN_MODEL_LIMIT') - interval = os.environ.get('TWINKLE_SERVER_STATE_METRICS_UPDATE_INTERVAL') + exp = os.environ.get("TWINKLE_SERVER_STATE_EXPIRATION_TIMEOUT") + clean = os.environ.get("TWINKLE_SERVER_STATE_CLEANUP_INTERVAL") + limit = os.environ.get("TWINKLE_SERVER_STATE_PER_TOKEN_MODEL_LIMIT") + interval = os.environ.get("TWINKLE_SERVER_STATE_METRICS_UPDATE_INTERVAL") if exp is None and clean is None and limit is None and interval is None: return None return cls( @@ -227,7 +187,7 @@ class ServerArgs(_ArgsBase): http_options: HttpOptions | None = None route_prefix: str | None = None api_key: str | None = None - owned_by: str = 'twinkle' + owned_by: str = "twinkle" class ProcessorArgs(_ArgsBase): @@ -251,11 +211,11 @@ class DataPlaneArgs(_ArgsBase): _ARGS_SCHEMA: dict[str, type[_ArgsBase]] = { - 'server': ServerArgs, - 'model': ModelArgs, - 'sampler': SamplerArgs, - 'processor': ProcessorArgs, - 'data_plane': DataPlaneArgs, + "server": ServerArgs, + "model": ModelArgs, + "sampler": SamplerArgs, + "processor": ProcessorArgs, + "data_plane": DataPlaneArgs, } # ---------- ApplicationSpec ------------------------------------------------ # @@ -273,19 +233,19 @@ class ApplicationSpec(BaseModel): falling back to a different schema's default. """ - model_config = ConfigDict(extra='forbid') + model_config = ConfigDict(extra="forbid") name: str - route_prefix: str = '/' - import_path: Literal['server', 'model', 'sampler', 'processor', 'data_plane'] + route_prefix: str = "/" + import_path: Literal["server", "model", "sampler", "processor", "data_plane"] # ``args`` is always populated by the ``mode='before'`` validator below # (which validates the raw block against the schema selected by # ``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') + @model_validator(mode="before") @classmethod def _coerce_args_to_schema(cls, data: Any) -> Any: """Validate the raw ``args`` block against the schema for ``import_path``. @@ -296,36 +256,16 @@ def _coerce_args_to_schema(cls, data: Any) -> Any: """ if not isinstance(data, dict): return data - import_path = data.get('import_path') + import_path = data.get("import_path") if import_path not in _ARGS_SCHEMA: # Let Pydantic's Literal validator handle bad import_path values. return data schema = _ARGS_SCHEMA[import_path] - raw_args = data.get('args') + raw_args = data.get("args") if isinstance(raw_args, schema): return data if raw_args is None: raw_args = {} # ``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 + return {**data, "args": schema.model_validate(raw_args)} 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..88dbba214 100644 --- a/src/twinkle/server/config/server_config.py +++ b/src/twinkle/server/config/server_config.py @@ -12,23 +12,26 @@ ``extra='forbid'`` so a YAML that uses a legacy field is rejected with the offending name pointed at. """ -from __future__ import annotations +from __future__ import annotations from pathlib import Path -from pydantic import BaseModel, ConfigDict, Field, model_validator from typing import Any +from pydantic import BaseModel, ConfigDict, Field, model_validator + 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): """Top-level server configuration aggregate root.""" - model_config = ConfigDict(extra='forbid') + model_config = ConfigDict(extra="forbid") ray_namespace: str | None = None proxy_location: str | None = None @@ -36,6 +39,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 ---------------------------------------------------------- # @@ -54,27 +58,52 @@ def from_yaml(cls, path: str | Path) -> ServerConfig: p = Path(path) if not p.exists(): - raise FileNotFoundError(f'Config file not found: {p}') + raise FileNotFoundError(f"Config file not found: {p}") try: raw = OmegaConf.to_container(OmegaConf.load(p), resolve=True) except Exception as e: # malformed YAML / OmegaConf parse failure - raise ConfigParseError(f'Malformed YAML in {p}: {e}') from e + raise ConfigParseError(f"Malformed YAML in {p}: {e}") from e if raw is None: raw = {} if not isinstance(raw, dict): - raise ConfigParseError(f'Top-level YAML in {p} must be a mapping, got {type(raw).__name__}', ) + raise ConfigParseError( + f"Top-level YAML in {p} must be a mapping, got {type(raw).__name__}", + ) return cls.model_validate(raw) # ---- cross-field validation ------------------------------------------ # - @model_validator(mode='after') + @model_validator(mode="after") def _validate_cross_field(self) -> ServerConfig: - if self.persistence.mode == 'redis' and not self.persistence.redis_url: - raise ValueError("persistence.redis_url is required when persistence.mode == 'redis'", ) + 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]: """Return a JSON-mode dict suitable for ``yaml.safe_dump`` / round-trip.""" - return self.model_dump(mode='json') + return self.model_dump(mode="json") diff --git a/src/twinkle/server/task_queue/config.py b/src/twinkle/server/config/task_queue.py similarity index 77% rename from src/twinkle/server/task_queue/config.py rename to src/twinkle/server/config/task_queue.py index 76f1813e8..2a0e96bd4 100644 --- a/src/twinkle/server/task_queue/config.py +++ b/src/twinkle/server/config/task_queue.py @@ -1,12 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -""" -Task queue configuration. +"""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. -""" from __future__ import annotations from pydantic import BaseModel, ConfigDict, Field @@ -39,7 +33,7 @@ class TaskQueueConfig(BaseModel): Must remain < Long_Poll_Window. """ - model_config = ConfigDict(extra='forbid') + model_config = ConfigDict(extra="forbid") rps_limit: float = Field(default=100.0, ge=0) tps_limit: float = Field(default=16000.0, ge=0) @@ -55,20 +49,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..d6ce7b505 --- /dev/null +++ b/src/twinkle/server/config/training.py @@ -0,0 +1,276 @@ +# 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 + + @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") + 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..6b327518c 100644 --- a/src/twinkle/server/deployment.py +++ b/src/twinkle/server/deployment.py @@ -25,16 +25,17 @@ [cleanup] → handler``), with metrics wrapping tracing and the replica-id header wrapping the full response path when enabled. """ -from __future__ import annotations +from __future__ import annotations import traceback from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager +from typing import Any + from fastapi import FastAPI, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse from ray import serve -from typing import Any from twinkle.protocol.types.errors import ErrorCategory from twinkle.server.exceptions import TwinkleServerError @@ -59,14 +60,14 @@ async def twinkle_server_error_handler(request: Request, exc: TwinkleServerError level rather than nested under ``detail``. A Decision_Boundary-left rejection (``category=user``) carries no traceback. """ - request_id = getattr(request.state, 'request_id', None) or '' + request_id = getattr(request.state, "request_id", None) or "" payload = build_error_payload( str(exc) or exc.__class__.__name__, category=exc.category, error_code=exc.error_code, request_id=request_id, ) - return JSONResponse(status_code=exc.error_code, content=payload.model_dump(mode='json', exclude_none=True)) + return JSONResponse(status_code=exc.error_code, content=payload.model_dump(mode="json", exclude_none=True)) # A body can produce hundreds of errors (one per element of a mis-typed tensor), and a @@ -76,28 +77,28 @@ async def twinkle_server_error_handler(request: Request, exc: TwinkleServerError def _validation_detail(error: dict[str, Any]) -> dict[str, Any]: """One pydantic error as a JSON-safe detail entry.""" - location = [str(part) for part in error.get('loc', ())] + location = [str(part) for part in error.get("loc", ())] return { - 'field': location[-1] if location else '', - 'path': '.'.join(location), - 'type': error.get('type', ''), - 'message': error.get('msg', ''), + "field": location[-1] if location else "", + "path": ".".join(location), + "type": error.get("type", ""), + "message": error.get("msg", ""), } def _validation_summary(errors: list[dict[str, Any]]) -> str: fields = [] for error in errors: - path = '.'.join(str(part) for part in error.get('loc', ())) + path = ".".join(str(part) for part in error.get("loc", ())) if path and path not in fields: fields.append(path) - shown = ', '.join(fields[:_MAX_VALIDATION_DETAILS]) or 'request body' - suffix = '' if len(fields) <= _MAX_VALIDATION_DETAILS else f' (+{len(fields) - _MAX_VALIDATION_DETAILS} more)' - return f'Request body validation failed for: {shown}{suffix}' + shown = ", ".join(fields[:_MAX_VALIDATION_DETAILS]) or "request body" + suffix = "" if len(fields) <= _MAX_VALIDATION_DETAILS else f" (+{len(fields) - _MAX_VALIDATION_DETAILS} more)" + return f"Request body validation failed for: {shown}{suffix}" def _validation_mentions_unknown_field(errors: list[dict[str, Any]]) -> bool: - return any(error.get('type') == 'extra_forbidden' for error in errors) + return any(error.get("type") == "extra_forbidden" for error in errors) async def validation_error_handler(request: Request, exc: RequestValidationError) -> JSONResponse: @@ -119,16 +120,18 @@ async def validation_error_handler(request: Request, exc: RequestValidationError if _validation_mentions_unknown_field(errors): # An unknown top-level field is what an older client looks like against a newer # server, so say so instead of leaving the caller to infer it from a field list. - message += ('. Unknown fields are rejected; if this worked before, upgrade ' - 'twinkle-kit on the client to match the server version.') + message += ( + ". Unknown fields are rejected; if this worked before, upgrade " + "twinkle-kit on the client to match the server version." + ) payload = build_error_payload( message, category=ErrorCategory.User, error_code=422, - request_id=getattr(request.state, 'request_id', None) or '', + request_id=getattr(request.state, "request_id", None) or "", details=[_validation_detail(error) for error in errors[:_MAX_VALIDATION_DETAILS]], ) - return JSONResponse(status_code=422, content=payload.model_dump(mode='json', exclude_none=True)) + return JSONResponse(status_code=422, content=payload.model_dump(mode="json", exclude_none=True)) def get_servable() -> Any: @@ -192,6 +195,7 @@ def build_deployment_app( async def lifespan(app: FastAPI): # Initialize telemetry in the worker process (after deserialization). from twinkle.server.telemetry.worker_init import ensure_telemetry_initialized + ensure_telemetry_initialized() # NOTE: ``state.start_cleanup_task()`` cannot run here — Ray Serve binds # ``servable_object`` AFTER lifespan startup. It is lazy-started from the @@ -206,6 +210,7 @@ async def lifespan(app: FastAPI): import asyncio from twinkle.server.telemetry import flush_telemetry_safely + await asyncio.to_thread(flush_telemetry_safely) app = FastAPI(lifespan=lifespan, **(fastapi_kwargs or {})) @@ -222,7 +227,7 @@ async def lifespan(app: FastAPI): # and finally the optional replica-id header as the outermost layer. if attach_cleanup_middleware: - @app.middleware('http') + @app.middleware("http") async def ensure_state_cleanup_started(request: Request, call_next): # ``LazyCleanupMixin._ensure_state_cleanup_started`` already # warn-and-swallows backend errors itself, so no outer try is @@ -231,7 +236,7 @@ async def ensure_state_cleanup_started(request: Request, call_next): await get_servable()._ensure_state_cleanup_started() return await call_next(request) - @app.middleware('http') + @app.middleware("http") async def catch_unhandled_exceptions(request: Request, call_next): try: return await call_next(request) @@ -241,7 +246,7 @@ async def catch_unhandled_exceptions(request: Request, call_next): # Unify the last-resort 500 with the rest of the wire: an # ``ErrorPayload`` body (Server category keeps the traceback) instead # of the legacy ``{'detail': }`` shape. - request_id = getattr(request.state, 'request_id', None) or '' + request_id = getattr(request.state, "request_id", None) or "" payload = build_error_payload( str(exc) or exc.__class__.__name__, category=ErrorCategory.Server, @@ -249,23 +254,23 @@ async def catch_unhandled_exceptions(request: Request, call_next): request_id=request_id, traceback_text=tb, ) - return JSONResponse(status_code=500, content=payload.model_dump(mode='json', exclude_none=True)) + return JSONResponse(status_code=500, content=payload.model_dump(mode="json", exclude_none=True)) - @app.middleware('http') + @app.middleware("http") async def verify_token(request: Request, call_next): return await verify_request_token(request=request, call_next=call_next) - app.middleware('http')(create_tracing_middleware(component)) - app.middleware('http')(create_metrics_middleware(component)) + app.middleware("http")(create_tracing_middleware(component)) + app.middleware("http")(create_metrics_middleware(component)) if attach_replica_id_header: - @app.middleware('http') + @app.middleware("http") async def inject_replica_id(request: Request, call_next): response = await call_next(request) try: ctx = serve.get_replica_context() - response.headers['X-Twinkle-Replica-Id'] = ctx.replica_id.unique_id + response.headers["X-Twinkle-Replica-Id"] = ctx.replica_id.unique_id except Exception: pass return response @@ -275,14 +280,14 @@ async def inject_replica_id(request: Request, call_next): def bind_deployment( - app: FastAPI, - servable_cls: type, - deploy_options: dict[str, Any], - *, - deployment_name: str, - request_router_config: Any | None = None, - bind_args: tuple = (), - bind_kwargs: dict[str, Any] | None = None, + app: FastAPI, + servable_cls: type, + deploy_options: dict[str, Any], + *, + deployment_name: str, + request_router_config: Any | None = None, + bind_args: tuple = (), + bind_kwargs: dict[str, Any] | None = None, ) -> Any: """Run ``serve.ingress(app)(cls)`` → ``serve.deployment(...)`` → ``.options().bind()``. @@ -296,9 +301,9 @@ def bind_deployment( bind_kwargs: Keyword args forwarded to ``.bind(...)``. """ ingress_cls = serve.ingress(app)(servable_cls) - deployment_kwargs: dict[str, Any] = {'name': deployment_name} + deployment_kwargs: dict[str, Any] = {"name": deployment_name} if request_router_config is not None: - deployment_kwargs['request_router_config'] = request_router_config + deployment_kwargs["request_router_config"] = request_router_config deployment_cls = serve.deployment(**deployment_kwargs)(ingress_cls) return deployment_cls.options(**deploy_options).bind(*bind_args, **(bind_kwargs or {})) @@ -328,5 +333,6 @@ async def _ensure_state_cleanup_started(self) -> None: # Idempotent via the ResourceCleanupCoordinator's own start guard. 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 + logger.warning(f"Failed to start ServerState cleanup task: {e}") + else: + self._state_cleanup_started = True diff --git a/src/twinkle/server/exceptions.py b/src/twinkle/server/exceptions.py index b8985879d..33252bdec 100644 --- a/src/twinkle/server/exceptions.py +++ b/src/twinkle/server/exceptions.py @@ -8,6 +8,7 @@ """ from __future__ import annotations +from typing import Literal from twinkle.protocol.types.errors import ErrorCategory @@ -24,7 +25,7 @@ class TwinkleServerError(Exception): def __init__( self, - message: str = '', + message: str = "", *, error_code: int | None = None, category: ErrorCategory | None = None, @@ -66,8 +67,8 @@ def __init__( self.value = value self.allowed = list(allowed) if allowed is not None else None if message is None: - allowed_part = f', allowed: {self.allowed}' if self.allowed is not None else '' - message = f'Invalid value for {field}: {value!r}{allowed_part}' + allowed_part = f", allowed: {self.allowed}" if self.allowed is not None else "" + message = f"Invalid value for {field}: {value!r}{allowed_part}" super().__init__(message) @@ -173,6 +174,30 @@ class FullModeBusyError(RequestRejectedError): def __init__(self, current_holder: str) -> None: self.current_holder = current_holder - 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.') + 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..082a411a6 100644 --- a/src/twinkle/server/gateway/app.py +++ b/src/twinkle/server/gateway/app.py @@ -5,18 +5,19 @@ A single Ray Serve deployment that serves both Tinker (/tinker/*) and Twinkle (/twinkle/*) management and proxy endpoints. """ -from __future__ import annotations +from __future__ import annotations import asyncio -from fastapi import FastAPI, HTTPException from typing import Any +from fastapi import FastAPI, HTTPException + 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 @@ -26,26 +27,30 @@ class GatewayServer(LazyCleanupMixin): """Unified gateway server handling both Tinker and Twinkle API clients.""" - def __init__(self, - supported_models: list | None = None, - server_config: Any = None, - http_options: dict[str, Any] | None = None, - **kwargs) -> None: + def __init__( + self, + supported_models: list | None = None, + server_config: Any = None, + http_options: dict[str, Any] | None = None, + **kwargs, + ) -> None: # ``server_config`` may be a typed ServerStateArgs (from the launcher), a # raw dict (direct callers), or None. Normalize to the keyword args # ``get_server_state`` expects, dropping unset (None) values so they fall # back to ServerState's own defaults. state_kwargs = self._normalize_server_state_args(server_config) self.state = get_server_state(**state_kwargs) - self.route_prefix = kwargs.get('route_prefix', '/api/v1') + self.route_prefix = kwargs.get("route_prefix", "/api/v1") 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. - self.api_key = kwargs.get('api_key') - self.owned_by = kwargs.get('owned_by') or 'twinkle' + self.api_key = kwargs.get("api_key") + self.owned_by = kwargs.get("owned_by") or "twinkle" self._modelscope_config_lock = asyncio.Lock() self._state_cleanup_started = False # Per-instance, not module-level: a process-global set never got cleared @@ -71,7 +76,7 @@ def _normalize_server_state_args(server_config: Any) -> dict[str, Any]: """ if server_config is None: return {} - if hasattr(server_config, 'model_dump'): + if hasattr(server_config, "model_dump"): data = server_config.model_dump() elif isinstance(server_config, dict): data = dict(server_config) @@ -93,24 +98,27 @@ 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) - if metadata and metadata.get('base_model'): - return metadata['base_model'] - raise HTTPException(status_code=404, detail=f'Model {model_id} not found') - - -def build_gateway_app(deploy_options: dict[str, Any], - supported_models: list | None = None, - server_config: Any = None, - http_options: dict[str, Any] | None = None, - **kwargs): + if metadata and metadata.get("base_model"): + return metadata["base_model"] + raise HTTPException(status_code=404, detail=f"Model {model_id} not found") + + +def build_gateway_app( + deploy_options: dict[str, Any], + supported_models: list | None = None, + server_config: Any = None, + http_options: dict[str, Any] | None = None, + **kwargs, +): """Build and configure the unified gateway server application. Serves Tinker endpoints at /* and Twinkle endpoints at /twinkle/*. @@ -142,7 +150,7 @@ async def _on_shutdown(servable: Any) -> None: await servable.proxy.close() app = build_deployment_app( - 'Gateway', + "Gateway", register_routes, on_shutdown=_on_shutdown, attach_cleanup_middleware=True, @@ -152,7 +160,7 @@ async def _on_shutdown(servable: Any) -> None: app, GatewayServer, deploy_options, - deployment_name='GatewayServer', + deployment_name="GatewayServer", bind_kwargs=dict( supported_models=supported_models, server_config=server_config, diff --git a/src/twinkle/server/gateway/openai_handlers.py b/src/twinkle/server/gateway/openai_handlers.py index 5495cb0c7..7e9d06e48 100644 --- a/src/twinkle/server/gateway/openai_handlers.py +++ b/src/twinkle/server/gateway/openai_handlers.py @@ -6,46 +6,111 @@ Translates OpenAI request/response shapes and proxies to the existing sampler /twinkle/sample (non-streaming) or /twinkle/sample_stream (streaming) routes. """ -from __future__ import annotations -import json +from __future__ import annotations +import asyncio +import os import secrets import time import uuid from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import json from fastapi import Depends, FastAPI, Request from fastapi.responses import JSONResponse, StreamingResponse -from typing import TYPE_CHECKING, Any 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, - translate_infer_request, translate_infer_response, translate_response, - translate_stream_chunk) +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, + translate_infer_request, + translate_infer_response, + translate_response, + translate_stream_chunk, +) from .use_cases import poll_future 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 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.""" - @app.post('/chat/completions') + @app.post("/chat/completions") async def chat_completions( - request: Request, - self: GatewayServer = Depends(self_fn), + request: Request, + self: GatewayServer = Depends(self_fn), ): """OpenAI-compatible chat completions endpoint. @@ -60,14 +125,14 @@ async def chat_completions( except ValueError as e: return _bad_request(e) - if not body.get('stream', False): - request_id = f'chatcmpl-{uuid.uuid4().hex[:24]}' + if not body.get("stream", False): + request_id = f"chatcmpl-{uuid.uuid4().hex[:24]}" return await _proxy_sample( self, request, sample_request, model, - endpoint='twinkle/sample', + endpoint="twinkle/sample", translate=lambda data, m: translate_response(data, m, request_id), ) @@ -77,12 +142,12 @@ async def chat_completions( return _model_not_found(model) sticky_headers = _build_sticky_headers(_sticky_key(request, model), request) await _ensure_template(self, base_model, sticky_headers, request) - request_id = f'chatcmpl-{uuid.uuid4().hex[:24]}' - first_chunk = {'is_first': True} + request_id = f"chatcmpl-{uuid.uuid4().hex[:24]}" + first_chunk = {"is_first": True} def _chat_chunk(delta_text: str, finish_reason: str | None) -> dict[str, Any]: - is_first = first_chunk['is_first'] - first_chunk['is_first'] = False + is_first = first_chunk["is_first"] + first_chunk["is_first"] = False return translate_stream_chunk( delta_text=delta_text, model=model, @@ -93,10 +158,10 @@ def _chat_chunk(delta_text: str, finish_reason: str | None) -> dict[str, Any]: return _streaming_response(self, request, sample_request, base_model, sticky_headers, _chat_chunk) - @app.post('/completions') + @app.post("/completions") async def completions( - request: Request, - self: GatewayServer = Depends(self_fn), + request: Request, + self: GatewayServer = Depends(self_fn), ): """OpenAI-compatible text completion endpoint. @@ -110,33 +175,38 @@ async def completions( except ValueError as e: return _bad_request(e) - if body.get('stream', False): + if body.get("stream", False): base_model = await _resolve_base_model(self, model) if base_model is None: return _model_not_found(model) sticky_headers = _build_sticky_headers(_sticky_key(request, model), request) await _ensure_template(self, base_model, sticky_headers, request) - request_id = f'cmpl-{uuid.uuid4().hex[:24]}' + request_id = f"cmpl-{uuid.uuid4().hex[:24]}" return _streaming_response( - self, request, sample_request, base_model, sticky_headers, + self, + request, + sample_request, + base_model, + sticky_headers, lambda delta, finish: translate_completion_stream_chunk( - delta_text=delta, model=model, finish_reason=finish, request_id=request_id), + delta_text=delta, model=model, finish_reason=finish, request_id=request_id + ), ) - request_id = f'cmpl-{uuid.uuid4().hex[:24]}' + request_id = f"cmpl-{uuid.uuid4().hex[:24]}" return await _proxy_sample( self, request, sample_request, model, - endpoint='twinkle/sample', + endpoint="twinkle/sample", translate=lambda data, m: translate_completion_response(data, m, request_id), ) - @app.post('/embeddings') + @app.post("/embeddings") async def embeddings( - request: Request, - self: GatewayServer = Depends(self_fn), + request: Request, + self: GatewayServer = Depends(self_fn), ): """OpenAI-compatible embeddings endpoint. @@ -153,14 +223,14 @@ async def embeddings( request, encode_request, model, - endpoint='twinkle/encode', + endpoint="twinkle/encode", translate=translate_embedding_response, ) - @app.post('/infer') + @app.post("/infer") async def infer( - request: Request, - self: GatewayServer = Depends(self_fn), + request: Request, + self: GatewayServer = Depends(self_fn), ): """Rollout inference: batched token ids + logprobs for RL training. @@ -179,35 +249,39 @@ async def infer( request, sample_request, model, - endpoint='twinkle/sample', + endpoint="twinkle/sample", translate=lambda data, m: translate_infer_response(data), ) - @app.get('/health') - @app.get('/ping') - @app.post('/ping') + @app.get("/health") + @app.get("/ping") + @app.post("/ping") async def health(): """Liveness probe. The gateway answering at all means the deployment is up.""" return JSONResponse(status_code=200, content={}) - @app.get('/models') + @app.get("/models") async def list_models( - request: Request, - self: GatewayServer = Depends(self_fn), + request: Request, + self: GatewayServer = Depends(self_fn), ): """OpenAI-compatible model listing endpoint.""" models = [] for m in self.supported_models: - models.append({ - 'id': m.model_name, - 'object': 'model', - 'created': int(time.time()), - 'owned_by': self.owned_by, - }) - return JSONResponse(content={ - 'object': 'list', - 'data': models, - }) + models.append( + { + "id": m.model_name, + "object": "model", + "created": int(time.time()), + "owned_by": self.owned_by, + } + ) + return JSONResponse( + content={ + "object": "list", + "data": models, + } + ) async def _resolve_base_model(gateway: GatewayServer, model: str) -> str | None: @@ -220,8 +294,8 @@ async def _resolve_base_model(gateway: GatewayServer, model: str) -> str | None: # Check if it's a registered adapter with metadata try: metadata = await gateway.state.get_model_metadata(model) - if metadata and metadata.get('base_model'): - return metadata['base_model'] + if metadata and metadata.get("base_model"): + return metadata["base_model"] except Exception: pass @@ -238,7 +312,7 @@ async def _resolve_base_model(gateway: GatewayServer, model: str) -> str | None: def _build_sticky_headers(sticky_key: str, request: Request) -> dict[str, str]: """Build the headers needed for sticky session routing.""" - auth = (request.headers.get(H_AUTH_TWINKLE) or request.headers.get(H_AUTH) or '') + auth = request.headers.get(H_AUTH_TWINKLE) or request.headers.get(H_AUTH) or "" return build_routing_headers(sticky_key, auth) @@ -247,14 +321,14 @@ def _sticky_key(request: Request, model: str) -> str: Using the model name alone would share slots across every caller of the same model. """ - return getattr(request.state, 'session_id', '') or getattr(request.state, 'token', '') or model + return getattr(request.state, "session_id", "") or getattr(request.state, "token", "") or model def _bad_request(exc: ValueError) -> JSONResponse: """A 400 for a translator's ``ValueError``; the exception text names the offending field.""" return JSONResponse( status_code=400, - content=make_error(message=f'Missing or invalid field: {exc}', param=str(exc) or None), + content=make_error(message=f"Missing or invalid field: {exc}", param=str(exc) or None), ) @@ -263,8 +337,8 @@ def _model_not_found(model: str) -> JSONResponse: status_code=404, content=make_error( message=f"Model '{model}' not found. Register it as an adapter or use a supported base model.", - error_type='model_not_found', - param='model', + error_type="model_not_found", + param="model", ), ) @@ -279,44 +353,20 @@ def _deny_without_key(request: Request, api_key: str | None) -> JSONResponse | N """ if api_key is None: return None - header = request.headers.get('authorization') or '' - supplied = header[len('Bearer '):] if header.startswith('Bearer ') else '' + header = request.headers.get("authorization") or "" + supplied = header[len("Bearer ") :] if header.startswith("Bearer ") else "" if not supplied or not secrets.compare_digest(supplied.encode(), api_key.encode()): return JSONResponse( status_code=401, content=make_error( message='Missing or invalid API key; send "Authorization: Bearer ".', - error_type='authentication_error', - code='invalid_api_key', + error_type="authentication_error", + code="invalid_api_key", ), ) 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, @@ -352,7 +402,7 @@ async def _proxy_sample( request, endpoint=endpoint, base_model=base_model, - service_type='sampler', + service_type="sampler", body_override=body_bytes, extra_headers=sticky_headers, ) @@ -360,33 +410,49 @@ async def _proxy_sample( return JSONResponse( status_code=response.status_code, content=make_error( - message=f'Sampler error: {response.body.decode()[:500]}', - error_type='server_error', + message=f"Sampler error: {response.body.decode()[:500]}", + error_type="server_error", ), ) - 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') + replica_id = response.headers.get("x-twinkle-replica-id") if replica_id: - resp_headers['X-Twinkle-Replica-Id'] = replica_id + resp_headers["X-Twinkle-Replica-Id"] = replica_id return JSONResponse(content=result, headers=resp_headers) @@ -409,34 +475,31 @@ def _streaming_response( async def _sse_generator(): try: async for line in gateway.proxy.proxy_request_stream( - request, - endpoint='twinkle/sample_stream', - base_model=base_model, - service_type='sampler', - body_override=json.dumps(sample_request).encode(), - extra_headers=sticky_headers, + request, + endpoint="twinkle/sample_stream", + base_model=base_model, + service_type="sampler", + body_override=json.dumps(sample_request).encode(), + extra_headers=sticky_headers, ): chunk_data = json.loads(line) - openai_chunk = chunk_translator(chunk_data.get('delta', ''), chunk_data.get('finish_reason')) - yield f'data: {json.dumps(openai_chunk)}\n\n' + openai_chunk = chunk_translator(chunk_data.get("delta", ""), chunk_data.get("finish_reason")) + yield f"data: {json.dumps(openai_chunk)}\n\n" - yield 'data: [DONE]\n\n' + yield "data: [DONE]\n\n" except httpx.HTTPStatusError as e: error_body = e.response.content.decode()[:500] error_chunk = make_error( - message=f'Streaming error: {error_body}', - error_type='server_error', + message=f"Streaming error: {error_body}", + error_type="server_error", ) - yield f'data: {json.dumps(error_chunk)}\n\n' - yield 'data: [DONE]\n\n' + yield f"data: {json.dumps(error_chunk)}\n\n" + yield "data: [DONE]\n\n" return StreamingResponse( _sse_generator(), - media_type='text/event-stream', - headers={ - 'Cache-Control': 'no-cache', - 'X-Accel-Buffering': 'no' - }, + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, ) @@ -450,9 +513,9 @@ def _register_api_key_gate(app: FastAPI, self_fn: Callable[[], GatewayServer]) - (and therefore wrapping the cleanup middleware ``build_deployment_app`` installs) so a rejected request short-circuits before any handler runs. """ - open_paths = frozenset({'/health', '/ping', '/models', '/twinkle/healthz', '/twinkle/healthz/deep'}) + open_paths = frozenset({"/health", "/ping", "/models", "/twinkle/healthz", "/twinkle/healthz/deep"}) - @app.middleware('http') + @app.middleware("http") async def require_api_key(request: Request, call_next): gateway = self_fn() if gateway.api_key is not None and request.url.path not in open_paths: @@ -477,24 +540,25 @@ async def _ensure_template( return template_cls = get_template_for_model(base_model) - set_template_body = json.dumps({ - 'template_cls': template_cls, - 'model_id': base_model, - 'adapter_name': '', - }).encode() + set_template_body = json.dumps( + { + "template_cls": template_cls, + "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', + service_type="sampler", body_override=set_template_body, extra_headers=sticky_headers, ) if resp.status_code == 200: gateway._template_initialized.add(base_model) else: - logger.warning('set_template failed: %s', resp.body.decode()[:200]) + logger.warning("set_template failed: %s", resp.body.decode()[:200]) except Exception as e: - logger.warning('set_template call failed: %s', e) + logger.warning("set_template call failed: %s", e) 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..3caf18100 100644 --- a/src/twinkle/server/gateway/tinker_handlers.py +++ b/src/twinkle/server/gateway/tinker_handlers.py @@ -5,64 +5,31 @@ All endpoints are prefixed /* and registered via _register_gateway_tinker_routes(app, self_fn). self_fn is injected via FastAPI Depends to obtain the GatewayServer instance at request time. """ -from __future__ import annotations +from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING, Any + from fastapi import Depends, FastAPI, HTTPException, Request, Response from tinker import types -from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from .app import GatewayServer 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 delete_checkpoint, get_weights_info, list_checkpoints, list_training_runs, poll_future from .use_cases import get_training_run as get_training_run_use_case -from .use_cases import get_weights_info, list_checkpoints, list_training_runs, poll_future 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: @@ -73,213 +40,219 @@ def _register_gateway_tinker_routes(app: FastAPI, self_fn: Callable[[], GatewayS It is wired in via ``Depends`` so it is resolved lazily at request time. """ - @app.get('/healthz') + @app.get("/healthz") async def healthz(request: Request) -> types.HealthResponse: - return types.HealthResponse(status='ok') + return types.HealthResponse(status="ok") - @app.get('/get_server_capabilities') + @app.get("/get_server_capabilities") async def get_server_capabilities( - request: Request, - self: GatewayServer = Depends(self_fn), + request: Request, + 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') + @app.post("/telemetry") async def telemetry(request: Request, body: types.TelemetrySendRequest) -> types.TelemetryResponse: - return types.TelemetryResponse(status='accepted') + return types.TelemetryResponse(status="accepted") - @app.post('/create_session') + @app.post("/create_session") async def create_session( - request: Request, - body: types.CreateSessionRequest, - self: GatewayServer = Depends(self_fn), + request: Request, + body: types.CreateSessionRequest, + self: GatewayServer = Depends(self_fn), ) -> types.CreateSessionResponse: session_id = await self.state.create_session(body.model_dump()) return types.CreateSessionResponse(session_id=session_id) - @app.post('/session_heartbeat') + @app.post("/session_heartbeat") async def session_heartbeat( request: Request, body: types.SessionHeartbeatRequest, self: GatewayServer = Depends(self_fn) ) -> types.SessionHeartbeatResponse: alive = await self.state.touch_session(body.session_id) if not alive: - raise HTTPException(status_code=404, detail='Unknown session') + raise HTTPException(status_code=404, detail="Unknown session") return types.SessionHeartbeatResponse() - @app.post('/create_sampling_session') + @app.post("/create_sampling_session") async def create_sampling_session( request: Request, body: types.CreateSamplingSessionRequest, self: GatewayServer = Depends(self_fn) ) -> types.CreateSamplingSessionResponse: sampling_session_id = await self.state.create_sampling_session(body.model_dump()) return types.CreateSamplingSessionResponse(sampling_session_id=sampling_session_id) - @app.post('/retrieve_future') - async def retrieve_future(request: Request, - body: types.FutureRetrieveRequest, - self: GatewayServer = Depends(self_fn)) -> Any: + @app.post("/retrieve_future") + async def retrieve_future( + request: Request, body: types.FutureRetrieveRequest, 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: - response_data: dict[str, Any] = {'type': 'try_again'} + response_data: dict[str, Any] = {"type": "try_again"} if record is not None: - if queue_state := record.get('queue_state'): - response_data['queue_state'] = queue_state - if queue_state_reason := record.get('queue_state_reason'): - response_data['queue_state_reason'] = queue_state_reason + if queue_state := record.get("queue_state"): + response_data["queue_state"] = queue_state + if queue_state_reason := record.get("queue_state_reason"): + response_data["queue_state_reason"] = queue_state_reason return response_data - status = record.get('status') - if status == 'failed': - return _tinker_error_from_failure(record.get('failure'), request_id=request_id) + status = record.get("status") + if status == "failed": + return _tinker_error_from_failure(record.get("failure"), request_id=request_id) - result = record.get('result') + result = record.get("result") if result is None: - raise HTTPException(status_code=500, detail='Task completed but no result found') - if hasattr(result, 'model_dump'): + raise HTTPException(status_code=500, detail="Task completed but no result found") + if hasattr(result, "model_dump"): return result.model_dump() return result # --- Training Runs Endpoints --- - @app.get('/training_runs') + @app.get("/training_runs") async def get_training_runs(request: Request, limit: int = 20, offset: int = 0) -> types.TrainingRunsResponse: token = get_token_from_request(request) - return list_training_runs(token, 'tinker', limit=limit, offset=offset) + return list_training_runs(token, "tinker", limit=limit, offset=offset) - @app.get('/training_runs/{run_id}') + @app.get("/training_runs/{run_id}") async def get_training_run(request: Request, run_id: str) -> types.TrainingRun: token = get_token_from_request(request) - run = get_training_run_use_case(token, 'tinker', run_id) + run = get_training_run_use_case(token, "tinker", run_id) if not run: - raise HTTPException(status_code=404, detail=f'Training run {run_id} not found') + raise HTTPException(status_code=404, detail=f"Training run {run_id} not found") return run - @app.get('/training_runs/{run_id}/checkpoints') + @app.get("/training_runs/{run_id}/checkpoints") async def get_run_checkpoints(request: Request, run_id: str) -> types.CheckpointsListResponse: token = get_token_from_request(request) - response = list_checkpoints(token, 'tinker', run_id) + response = list_checkpoints(token, "tinker", run_id) if not response: - raise HTTPException(status_code=404, detail=f'Training run {run_id} not found') + raise HTTPException(status_code=404, detail=f"Training run {run_id} not found") return response - @app.delete('/training_runs/{run_id}/checkpoints/{checkpoint_id:path}') + @app.delete("/training_runs/{run_id}/checkpoints/{checkpoint_id:path}") async def delete_run_checkpoint(request: Request, run_id: str, checkpoint_id: str) -> Any: token = get_token_from_request(request) - success = delete_checkpoint(token, 'tinker', run_id, checkpoint_id) + success = delete_checkpoint(token, "tinker", run_id, checkpoint_id) if not success: - raise HTTPException(status_code=404, detail=f'Checkpoint {checkpoint_id} not found for run {run_id}') + raise HTTPException(status_code=404, detail=f"Checkpoint {checkpoint_id} not found for run {run_id}") return None - @app.post('/weights_info') + @app.post("/weights_info") async def weights_info(request: Request, body: dict[str, Any]) -> types.WeightsInfoResponse: token = get_token_from_request(request) - tinker_path = body.get('tinker_path') - response = get_weights_info(token, 'tinker', tinker_path) + tinker_path = body.get("tinker_path") + response = get_weights_info(token, "tinker", tinker_path) if not response: - raise HTTPException(status_code=404, detail=f'Weights at {tinker_path} not found') + raise HTTPException(status_code=404, detail=f"Weights at {tinker_path} not found") return response - @app.post('/training_runs/{run_id}/checkpoints/{checkpoint_id:path}/publish') - async def publish_checkpoint(request: Request, - run_id: str, - checkpoint_id: str, - self: GatewayServer = Depends(self_fn)) -> Response: + @app.post("/training_runs/{run_id}/checkpoints/{checkpoint_id:path}/publish") + async def publish_checkpoint( + request: Request, run_id: str, checkpoint_id: str, self: GatewayServer = Depends(self_fn) + ) -> Response: token = get_token_from_request(request) - training_run_manager = create_training_run_manager(token, client_type='tinker') - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + training_run_manager = create_training_run_manager(token, client_type="tinker") + checkpoint_manager = create_checkpoint_manager(token, client_type="tinker") run = training_run_manager.get(run_id) if not run: - raise HTTPException(status_code=404, detail=f'Training run {run_id} not found or access denied') + raise HTTPException(status_code=404, detail=f"Training run {run_id} not found or access denied") checkpoint = checkpoint_manager.get(run_id, checkpoint_id) if not checkpoint: - raise HTTPException(status_code=404, detail=f'Checkpoint {checkpoint_id} not found') + raise HTTPException(status_code=404, detail=f"Checkpoint {checkpoint_id} not found") checkpoint_dir = str(checkpoint_manager.get_ckpt_dir(run_id, checkpoint_id)) async with self._modelscope_config_lock: try: from modelscope.hub.api import HubApi, ModelScopeConfig + hub_api = HubApi(token=token) hub_api.login() username = ModelScopeConfig.get_user_info()[0] except Exception as e: - logger.error(f'Failed to get username from ModelScope: {e}') + logger.error(f"Failed to get username from ModelScope: {e}") raise HTTPException( - status_code=401, - detail='Failed to get username from ModelScope. Please ensure your token is valid.') + status_code=401, detail="Failed to get username from ModelScope. Please ensure your token is valid." + ) from e - checkpoint_name = checkpoint_id.split('/')[-1] - hub_model_id = f'{username}/{run_id}_{checkpoint_name}' + checkpoint_name = checkpoint_id.split("/")[-1] + hub_model_id = f"{username}/{run_id}_{checkpoint_name}" HubOperation.push_to_hub(repo_id=hub_model_id, folder_path=checkpoint_dir, token=token, private=True) return Response(status_code=204) # --- Model Proxy Endpoints --- - @app.post('/create_model') - async def create_model(request: Request, body: types.CreateModelRequest, - self: GatewayServer = Depends(self_fn)) -> Any: + @app.post("/create_model") + async def create_model( + request: Request, body: types.CreateModelRequest, self: GatewayServer = Depends(self_fn) + ) -> Any: self._validate_base_model(body.base_model) - return await self.proxy.proxy_to_model(request, 'create_model', body.base_model) + return await self.proxy.proxy_to_model(request, "create_model", body.base_model) - @app.post('/get_info') + @app.post("/get_info") async def get_info(request: Request, body: types.GetInfoRequest, self: GatewayServer = Depends(self_fn)) -> Any: - return await self.proxy.proxy_to_model(request, 'get_info', await self._get_base_model(body.model_id)) + return await self.proxy.proxy_to_model(request, "get_info", await self._get_base_model(body.model_id)) - @app.post('/unload_model') - async def unload_model(request: Request, body: types.UnloadModelRequest, - self: GatewayServer = Depends(self_fn)) -> Any: - return await self.proxy.proxy_to_model(request, 'unload_model', await self._get_base_model(body.model_id)) + @app.post("/unload_model") + async def unload_model( + request: Request, body: types.UnloadModelRequest, self: GatewayServer = Depends(self_fn) + ) -> Any: + return await self.proxy.proxy_to_model(request, "unload_model", await self._get_base_model(body.model_id)) - @app.post('/forward') + @app.post("/forward") async def forward(request: Request, body: types.ForwardRequest, self: GatewayServer = Depends(self_fn)) -> Any: - return await self.proxy.proxy_to_model(request, 'forward', await self._get_base_model(body.model_id)) + return await self.proxy.proxy_to_model(request, "forward", await self._get_base_model(body.model_id)) - @app.post('/forward_backward') - async def forward_backward(request: Request, - body: types.ForwardBackwardRequest, - self: GatewayServer = Depends(self_fn)) -> Any: - return await self.proxy.proxy_to_model(request, 'forward_backward', await self._get_base_model(body.model_id)) + @app.post("/forward_backward") + async def forward_backward( + request: Request, body: types.ForwardBackwardRequest, self: GatewayServer = Depends(self_fn) + ) -> Any: + return await self.proxy.proxy_to_model(request, "forward_backward", await self._get_base_model(body.model_id)) - @app.post('/optim_step') + @app.post("/optim_step") async def optim_step(request: Request, body: types.OptimStepRequest, self: GatewayServer = Depends(self_fn)) -> Any: - return await self.proxy.proxy_to_model(request, 'optim_step', await self._get_base_model(body.model_id)) + return await self.proxy.proxy_to_model(request, "optim_step", await self._get_base_model(body.model_id)) - @app.post('/save_weights') - async def save_weights(request: Request, body: types.SaveWeightsRequest, - self: GatewayServer = Depends(self_fn)) -> Any: - return await self.proxy.proxy_to_model(request, 'save_weights', await self._get_base_model(body.model_id)) + @app.post("/save_weights") + async def save_weights( + request: Request, body: types.SaveWeightsRequest, self: GatewayServer = Depends(self_fn) + ) -> Any: + return await self.proxy.proxy_to_model(request, "save_weights", await self._get_base_model(body.model_id)) - @app.post('/load_weights') - async def load_weights(request: Request, body: types.LoadWeightsRequest, - self: GatewayServer = Depends(self_fn)) -> Any: - return await self.proxy.proxy_to_model(request, 'load_weights', await self._get_base_model(body.model_id)) + @app.post("/load_weights") + async def load_weights( + request: Request, body: types.LoadWeightsRequest, self: GatewayServer = Depends(self_fn) + ) -> Any: + return await self.proxy.proxy_to_model(request, "load_weights", await self._get_base_model(body.model_id)) # --- Sampler Proxy Endpoints --- - @app.post('/asample') + @app.post("/asample") async def asample(request: Request, body: types.SampleRequest, self: GatewayServer = Depends(self_fn)) -> Any: base_model = body.base_model if not base_model and body.sampling_session_id: session = await self.state.get_sampling_session(body.sampling_session_id) if session: - base_model = session.get('base_model') + base_model = session.get("base_model") if not base_model: - raise HTTPException(status_code=400, detail='base_model is required but could not be resolved') - return await self.proxy.proxy_to_sampler(request, 'asample', base_model) + raise HTTPException(status_code=400, detail="base_model is required but could not be resolved") + return await self.proxy.proxy_to_sampler(request, "asample", base_model) - @app.post('/save_weights_for_sampler') + @app.post("/save_weights_for_sampler") async def save_weights_for_sampler( - request: Request, - body: types.SaveWeightsForSamplerRequest, - self: GatewayServer = Depends(self_fn), + request: Request, + body: types.SaveWeightsForSamplerRequest, + self: GatewayServer = Depends(self_fn), ) -> Any: - return await self.proxy.proxy_to_model(request, 'save_weights_for_sampler', await - self._get_base_model(body.model_id)) + return await self.proxy.proxy_to_model( + request, "save_weights_for_sampler", await self._get_base_model(body.model_id) + ) diff --git a/src/twinkle/server/gateway/twinkle_handlers.py b/src/twinkle/server/gateway/twinkle_handlers.py index 3df44ffc8..b7eb557ba 100644 --- a/src/twinkle/server/gateway/twinkle_handlers.py +++ b/src/twinkle/server/gateway/twinkle_handlers.py @@ -4,12 +4,13 @@ All endpoints are prefixed /twinkle/* and registered via _register_gateway_twinkle_routes(app, self_fn). """ -from __future__ import annotations +from __future__ import annotations from collections.abc import Callable -from fastapi import Depends, FastAPI, Request from typing import TYPE_CHECKING +from fastapi import Depends, FastAPI, Request + if TYPE_CHECKING: from .app import GatewayServer @@ -20,9 +21,8 @@ from twinkle.server.lifecycle.poll_config import long_poll_window from twinkle.server.middleware.auth import get_token_from_request from twinkle.utils.logger import get_logger -from .use_cases import delete_checkpoint +from .use_cases import delete_checkpoint, get_weights_info, list_checkpoints, list_training_runs, poll_future from .use_cases import get_training_run as get_training_run_use_case -from .use_cases import get_weights_info, list_checkpoints, list_training_runs, poll_future logger = get_logger() @@ -30,22 +30,22 @@ def _register_gateway_twinkle_routes(app: FastAPI, self_fn: Callable[[], GatewayServer]) -> None: """Register all /twinkle/* routes on the given FastAPI app.""" - @app.get('/twinkle/capacity_info', response_model=types.CapacityInfoResponse) + @app.get("/twinkle/capacity_info", response_model=types.CapacityInfoResponse) async def get_capacity_info( - request: Request, - self: GatewayServer = Depends(self_fn), + request: Request, + self: GatewayServer = Depends(self_fn), ) -> types.CapacityInfoResponse: info = await self.state.get_capacity_info() return types.CapacityInfoResponse(**info) - @app.get('/twinkle/healthz', response_model=types.HealthResponse) + @app.get("/twinkle/healthz", response_model=types.HealthResponse) async def healthz(request: Request) -> types.HealthResponse: - return types.HealthResponse(status='ok') + return types.HealthResponse(status="ok") - @app.get('/twinkle/healthz/deep') + @app.get("/twinkle/healthz/deep") async def healthz_deep( - request: Request, - self: GatewayServer = Depends(self_fn), + request: Request, + self: GatewayServer = Depends(self_fn), ) -> dict: """Deep health check: verifies model actors are alive, not just the gateway. @@ -57,36 +57,40 @@ 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) + 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] = { - 'healthy': healthy, - 'status_code': resp.status_code, + "healthy": healthy, + "status_code": resp.status_code, } except Exception as e: all_healthy = False results[model_name] = { - 'healthy': False, - 'detail': str(e), + "healthy": False, + "detail": str(e), } - body = {'healthy': all_healthy, 'models': results} + body = {"healthy": all_healthy, "models": results} if not all_healthy: return JSONResponse(status_code=503, content=body) return body - @app.get('/twinkle/get_server_capabilities', response_model=types.GetServerCapabilitiesResponse) + @app.get("/twinkle/get_server_capabilities", response_model=types.GetServerCapabilitiesResponse) async def get_server_capabilities( - request: Request, - self: GatewayServer = Depends(self_fn), + request: Request, + 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, @@ -98,31 +102,31 @@ async def get_server_capabilities( limits=types.ProtocolLimits(long_poll_timeout_seconds=long_poll_window()), ) - @app.post('/twinkle/create_session', response_model=types.CreateSessionResponse) + @app.post("/twinkle/create_session", response_model=types.CreateSessionResponse) async def create_session( - request: Request, - body: types.CreateSessionRequest, - self: GatewayServer = Depends(self_fn), + request: Request, + body: types.CreateSessionRequest, + self: GatewayServer = Depends(self_fn), ) -> types.CreateSessionResponse: session_id = await self.state.create_session(body.model_dump()) return types.CreateSessionResponse(session_id=session_id) - @app.post('/twinkle/session_heartbeat', response_model=types.SessionHeartbeatResponse) + @app.post("/twinkle/session_heartbeat", response_model=types.SessionHeartbeatResponse) async def session_heartbeat( - request: Request, - body: types.SessionHeartbeatRequest, - self: GatewayServer = Depends(self_fn), + request: Request, + body: types.SessionHeartbeatRequest, + self: GatewayServer = Depends(self_fn), ) -> types.SessionHeartbeatResponse: alive = await self.state.touch_session(body.session_id) if not alive: - raise ResourceNotFoundError('Unknown session') + raise ResourceNotFoundError("Unknown session") return types.SessionHeartbeatResponse() - @app.post('/twinkle/retrieve_future', response_model=types.TaskEnvelope) + @app.post("/twinkle/retrieve_future", response_model=types.TaskEnvelope) async def retrieve_future( - request: Request, - body: types.RetrieveFutureRequest, - self: GatewayServer = Depends(self_fn), + request: Request, + body: types.RetrieveFutureRequest, + self: GatewayServer = Depends(self_fn), ) -> types.TaskEnvelope: """Long-poll a twinkle-native task to a terminal state. @@ -139,16 +143,17 @@ 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') + raise ResourceNotFoundError(f"request_id {request_id} not found or expired") return envelope_from_record(request_id, outcome.record) - @app.post('/twinkle/cancel', response_model=types.CancelResponse) + @app.post("/twinkle/cancel", response_model=types.CancelResponse) async def cancel_future( - request: Request, - body: types.CancelRequest, - self: GatewayServer = Depends(self_fn), + request: Request, + body: types.CancelRequest, + self: GatewayServer = Depends(self_fn), ) -> types.CancelResponse: """Best-effort cancel of a not-yet-started task. @@ -156,85 +161,88 @@ 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) - @app.get('/twinkle/training_runs', response_model=types.TrainingRunsResponse) + @app.get("/twinkle/training_runs", response_model=types.TrainingRunsResponse) async def get_training_runs(request: Request, limit: int = 20, offset: int = 0) -> types.TrainingRunsResponse: token = get_token_from_request(request) - return list_training_runs(token, 'twinkle', limit=limit, offset=offset) + return list_training_runs(token, "twinkle", limit=limit, offset=offset) - @app.get('/twinkle/training_runs/{run_id}', response_model=types.TrainingRun) + @app.get("/twinkle/training_runs/{run_id}", response_model=types.TrainingRun) async def get_training_run(request: Request, run_id: str) -> types.TrainingRun: token = get_token_from_request(request) - run = get_training_run_use_case(token, 'twinkle', run_id, check_permission=True) + run = get_training_run_use_case(token, "twinkle", run_id, check_permission=True) if not run: - raise ResourceNotFoundError(f'Training run {run_id} not found or access denied') + raise ResourceNotFoundError(f"Training run {run_id} not found or access denied") return run - @app.get('/twinkle/training_runs/{run_id}/checkpoints', response_model=types.CheckpointsListResponse) + @app.get("/twinkle/training_runs/{run_id}/checkpoints", response_model=types.CheckpointsListResponse) async def get_run_checkpoints(request: Request, run_id: str) -> types.CheckpointsListResponse: token = get_token_from_request(request) - response = list_checkpoints(token, 'twinkle', run_id) + response = list_checkpoints(token, "twinkle", run_id) if response is None: - raise ResourceNotFoundError(f'Training run {run_id} not found or access denied') + raise ResourceNotFoundError(f"Training run {run_id} not found or access denied") return response @app.delete( - '/twinkle/training_runs/{run_id}/checkpoints/{checkpoint_id:path}', - response_model=types.DeleteCheckpointResponse) - async def delete_run_checkpoint(request: Request, run_id: str, - checkpoint_id: str) -> types.DeleteCheckpointResponse: + "/twinkle/training_runs/{run_id}/checkpoints/{checkpoint_id:path}", + response_model=types.DeleteCheckpointResponse, + ) + async def delete_run_checkpoint( + request: Request, run_id: str, checkpoint_id: str + ) -> types.DeleteCheckpointResponse: token = get_token_from_request(request) if not validate_user_path(token, checkpoint_id): - raise RequestRejectedError('Invalid checkpoint path: path traversal not allowed') + raise RequestRejectedError("Invalid checkpoint path: path traversal not allowed") - success = delete_checkpoint(token, 'twinkle', run_id, checkpoint_id) + success = delete_checkpoint(token, "twinkle", run_id, checkpoint_id) if not success: - raise ResourceNotFoundError(f'Checkpoint {checkpoint_id} not found or access denied') + raise ResourceNotFoundError(f"Checkpoint {checkpoint_id} not found or access denied") - return types.DeleteCheckpointResponse(success=True, message=f'Checkpoint {checkpoint_id} deleted successfully') + return types.DeleteCheckpointResponse(success=True, message=f"Checkpoint {checkpoint_id} deleted successfully") - @app.post('/twinkle/weights_info', response_model=types.WeightsInfoResponse) + @app.post("/twinkle/weights_info", response_model=types.WeightsInfoResponse) async def weights_info(request: Request, body: types.WeightsInfoRequest) -> types.WeightsInfoResponse: token = get_token_from_request(request) - response = get_weights_info(token, 'twinkle', body.twinkle_path) + response = get_weights_info(token, "twinkle", body.twinkle_path) if response is None: - raise ResourceNotFoundError(f'Weights at {body.twinkle_path} not found or access denied') + raise ResourceNotFoundError(f"Weights at {body.twinkle_path} not found or access denied") return response - @app.get('/twinkle/checkpoint_path/{run_id}/{checkpoint_id:path}', response_model=types.CheckpointPathResponse) + @app.get("/twinkle/checkpoint_path/{run_id}/{checkpoint_id:path}", response_model=types.CheckpointPathResponse) async def get_checkpoint_path(request: Request, run_id: str, checkpoint_id: str) -> types.CheckpointPathResponse: token = get_token_from_request(request) if not validate_user_path(token, checkpoint_id): - raise RequestRejectedError('Invalid checkpoint path: path traversal not allowed') + raise RequestRejectedError("Invalid checkpoint path: path traversal not allowed") - training_run_manager = create_training_run_manager(token, client_type='twinkle') - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + training_run_manager = create_training_run_manager(token, client_type="twinkle") + checkpoint_manager = create_checkpoint_manager(token, client_type="twinkle") run = training_run_manager.get(run_id) if not run: - raise ResourceNotFoundError(f'Training run {run_id} not found or access denied') + raise ResourceNotFoundError(f"Training run {run_id} not found or access denied") checkpoint = checkpoint_manager.get(run_id, checkpoint_id) if not checkpoint: - raise ResourceNotFoundError(f'Checkpoint {checkpoint_id} not found') + raise ResourceNotFoundError(f"Checkpoint {checkpoint_id} not found") ckpt_dir = checkpoint_manager.get_ckpt_dir(run_id, checkpoint_id) return types.CheckpointPathResponse(path=str(ckpt_dir), twinkle_path=checkpoint.twinkle_path) - @app.get('/twinkle/status') + @app.get("/twinkle/status") async def status( - request: Request, - self: GatewayServer = Depends(self_fn), + request: Request, + self: GatewayServer = Depends(self_fn), ) -> dict: cleanup_stats = await self.state.get_cleanup_stats() return { - 'resources': cleanup_stats['resource_counts'], - 'cleanup': { - 'running': cleanup_stats['cleanup_running'], - 'expiration_timeout': cleanup_stats['expiration_timeout'], + "resources": cleanup_stats["resource_counts"], + "cleanup": { + "running": cleanup_stats["cleanup_running"], + "expiration_timeout": cleanup_stats["expiration_timeout"], }, } diff --git a/src/twinkle/server/gateway/use_cases.py b/src/twinkle/server/gateway/use_cases.py index 26691014b..5207806a0 100644 --- a/src/twinkle/server/gateway/use_cases.py +++ b/src/twinkle/server/gateway/use_cases.py @@ -9,8 +9,8 @@ ``cancel_future``, ``get_cleanup_stats``, ``get_model_metadata``) deliberately stay direct: wrapping them for symmetry would add forwarding, not structure. """ -from __future__ import annotations +from __future__ import annotations import asyncio from dataclasses import dataclass from typing import Any @@ -18,7 +18,9 @@ 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 -_TERMINAL_STATUSES = frozenset({'completed', 'failed', 'cancelled'}) +# 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"}) @dataclass(frozen=True, slots=True) @@ -34,7 +36,7 @@ async def poll_future(state: Any, request_id: str) -> FuturePollResult: record = None while True: record = await state.get_future(request_id) - if record is not None and record.get('status') in _TERMINAL_STATUSES: + if record is not None and record.get("status") in _TERMINAL_STATUSES: return FuturePollResult(record=record, timed_out=False) if asyncio.get_running_loop().time() >= deadline: return FuturePollResult(record=record, timed_out=True) @@ -47,7 +49,7 @@ def list_training_runs(token: str, client_type: str, *, limit: int, offset: int) def get_training_run(token: str, client_type: str, run_id: str, *, check_permission: bool = False) -> Any | None: manager = create_training_run_manager(token, client_type=client_type) - if check_permission and hasattr(manager, 'get_with_permission'): + if check_permission and hasattr(manager, "get_with_permission"): return manager.get_with_permission(run_id) return manager.get(run_id) diff --git a/src/twinkle/server/launcher/__init__.py b/src/twinkle/server/launcher/__init__.py index 4a2514ceb..1fd81aa12 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`` @@ -25,8 +24,9 @@ # CLI python -m twinkle.server launch --config server_config.yaml """ + from __future__ import annotations from .server_launcher import ServerLauncher, build_ray_runtime_env, ensure_ray_cluster, launch_server -__all__ = ['ServerLauncher', 'launch_server', 'build_ray_runtime_env', 'ensure_ray_cluster'] +__all__ = ["ServerLauncher", "launch_server", "build_ray_runtime_env", "ensure_ray_cluster"] 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 500feffd8..a719db0f6 100644 --- a/src/twinkle/server/launcher/server_launcher.py +++ b/src/twinkle/server/launcher/server_launcher.py @@ -1,28 +1,24 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""``ServerLauncher`` and the ``launch_server`` entry point. +"""Initialize server configuration, deploy applications, and own their shutdown.""" -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. -""" from __future__ import annotations - -import json import os import signal import threading from collections.abc import Callable +from copy import deepcopy from pathlib import Path from typing import Any +import json + 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, get_runtime_env_for_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() @@ -39,9 +35,9 @@ def build_ray_runtime_env() -> dict[str, Any]: runtime_env = get_runtime_env_for_patches() propagated_env_vars = build_propagated_env_vars() if propagated_env_vars: - merged_env_vars = dict(runtime_env.get('env_vars') or {}) + merged_env_vars = dict(runtime_env.get("env_vars") or {}) merged_env_vars.update(propagated_env_vars) - runtime_env['env_vars'] = merged_env_vars + runtime_env["env_vars"] = merged_env_vars return runtime_env @@ -60,72 +56,79 @@ def ensure_ray_cluster(*, namespace: str, address: str | None = None) -> None: if ray.is_initialized(): return init_kwargs: dict[str, Any] = { - 'namespace': namespace, - 'runtime_env': build_ray_runtime_env(), - 'ignore_reinit_error': True, + "namespace": namespace, + "runtime_env": build_ray_runtime_env(), + "ignore_reinit_error": True, } if address: - init_kwargs['address'] = address + init_kwargs["address"] = address ray.init(**init_kwargs) - logger.info(f'ray runtime ready (namespace={namespace}, address={address or "local head"})') + logger.info(f"ray runtime ready (namespace={namespace}, address={address or 'local head'})") 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 + raise TypeError( + "ServerLauncher requires a typed ServerConfig instance; " + f"got {type(config).__name__}. Build one with " + "ServerConfig.from_yaml(path) or ServerConfig(...)." + ) + 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: return - namespace = self.ray_namespace or self.config.ray_namespace or 'twinkle_cluster' - ensure_ray_cluster(namespace=namespace, address='auto') + namespace = self.ray_namespace or self.config.ray_namespace or "twinkle_cluster" + ensure_ray_cluster(namespace=namespace, address="auto") self._ray_initialized = True def _start_serve(self) -> None: @@ -139,7 +142,6 @@ def _start_serve(self) -> None: """ if self._serve_started: return - from ray import serve try: @@ -155,187 +157,203 @@ def _start_serve(self) -> None: self._serve_started = True return - http_options = self.config.http_options.model_dump() - serve_kwargs: dict[str, Any] = {'http_options': http_options} + options = {"http_options": self.config.http_options.model_dump()} # ``proxy_location`` controls where the Ray Serve HTTP proxy runs # (``EveryNode`` / ``HeadOnly`` / ``Disabled``). The example configs # set this field, so honour it here instead of silently ignoring. 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}') - + options["proxy_location"] = self.config.proxy_location + serve.start(**options) 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) + from twinkle.server.sampler.pool.actor import create_sampler_controller + from twinkle.server.sampler.pool.router import build_sampler_router - self._init_ray() - self._start_serve() + 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 + if app_spec.import_path == "sampler": + self._deploy_sampler(app_spec) + return + from ray import serve - for app_spec in applications: - self._deploy_application(app_spec) + 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 - host = self.config.http_options.host - port = self.config.http_options.port + from twinkle.server.model.pool.actor import TrainingControllerActor, controller_name + from twinkle.server.model.pool.router import build_training_router - logger.info('All applications started!') - for app_spec in applications: - logger.info(' Endpoint: http://%s:%s%s', host, port, app_spec.route_prefix) + 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 - if not blocking: - return + from ray import serve - # Graceful shutdown via signal handling - shutdown_event = threading.Event() + from twinkle.server.telemetry import flush_telemetry_safely - def _handle_signal(signum, frame): - sig_name = signal.Signals(signum).name - logger.info(f'Received {sig_name}, shutting down gracefully...') - shutdown_event.set() + 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 - signal.signal(signal.SIGINT, _handle_signal) - signal.signal(signal.SIGTERM, _handle_signal) + @staticmethod + def _wait_for_shutdown() -> None: + stopped = threading.Event() - # Block until a termination signal is received - shutdown_event.wait() + def signal_shutdown(signum, frame): + logger.info("Received %s", signal.Signals(signum).name) + stopped.set() - self._shutdown() + signal.signal(signal.SIGINT, signal_shutdown) + signal.signal(signal.SIGTERM, signal_shutdown) + stopped.wait() - def _shutdown(self) -> None: - """Tear down Serve and flush telemetry on graceful termination. + def launch(self, blocking: bool = True) -> None: + """Launch the server with all configured applications. - Telemetry flush runs AFTER ``serve.shutdown()`` and is wrapped so a - telemetry-shutdown failure cannot mask the user-facing shutdown path. + 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. """ - from ray import serve + 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: @@ -346,14 +364,16 @@ def _warn_public_unauthenticated(config: ServerConfig) -> None: free compute plus an internal-network pivot. Keeping the check here (the single launch choke point) covers YAML and programmatic configs alike. """ - if config.http_options.host != '0.0.0.0': + if config.http_options.host != "0.0.0.0": return - if any(app.import_path == 'server' and getattr(app.args, 'api_key', None) - for app in config.applications): + if any(app.import_path == "server" and getattr(app.args, "api_key", None) for app in config.applications): return - logger.warning('server binds 0.0.0.0:%s with no api_key configured on any gateway application: ' - 'every route is reachable unauthenticated by any network peer. Set ServerArgs.api_key ' - '(--api_key) or bind a non-public host.', config.http_options.port) + logger.warning( + "server binds 0.0.0.0:%s with no api_key configured on any gateway application: " + "every route is reachable unauthenticated by any network peer. Set ServerArgs.api_key " + "(--api_key) or bind a non-public host.", + config.http_options.port, + ) def launch_server( diff --git a/src/twinkle/server/lifecycle/envelope.py b/src/twinkle/server/lifecycle/envelope.py index 25343c9b0..33f84cb52 100644 --- a/src/twinkle/server/lifecycle/envelope.py +++ b/src/twinkle/server/lifecycle/envelope.py @@ -6,48 +6,25 @@ this mapping per endpoint is how a task that failed inside the Inline_Fast_Path window loses its payload. """ -from __future__ import annotations +from __future__ import annotations from typing import Any 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, @@ -58,7 +35,7 @@ def envelope_from_record( request_id: str, record: dict[str, Any] | None, *, - fallback_status: str = 'pending', + fallback_status: str = "pending", ) -> TaskEnvelope: """Map a stored ``FutureRecord`` dict to the wire ``TaskEnvelope``. @@ -77,21 +54,21 @@ def envelope_from_record( distinguish them by ``status`` while reading one field. """ record = record or {} - status = record.get('status', fallback_status) + status = record.get("status", fallback_status) common = dict( - queue_state=record.get('queue_state'), - queue_state_reason=record.get('queue_state_reason'), + queue_state=record.get("queue_state"), + queue_state_reason=record.get("queue_state_reason"), ) - if status in ('failed', 'cancelled'): + if status in ("failed", "cancelled"): return TaskEnvelope( request_id=request_id, status=status, - error=error_payload_from_failure(record.get('failure'), request_id=request_id), + error=error_payload_from_failure(record.get("failure"), request_id=request_id), **common, ) return TaskEnvelope( request_id=request_id, status=status, - result=record.get('result') if status == 'completed' else None, + result=record.get("result") if status == "completed" else None, **common, ) diff --git a/src/twinkle/server/lifecycle/protocols.py b/src/twinkle/server/lifecycle/protocols.py index c05ea12bb..8c762a6aa 100644 --- a/src/twinkle/server/lifecycle/protocols.py +++ b/src/twinkle/server/lifecycle/protocols.py @@ -14,14 +14,15 @@ cleanup middleware via ``attach_cleanup_middleware=False``; the gate below is what makes that fact visible statically instead of only via that boolean. """ + from __future__ import annotations +from typing import Any, Protocol, runtime_checkable 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 @@ -32,23 +33,17 @@ class QueuedDeployment(Protocol): replica_id: str @property - def task_queue_config(self) -> TaskQueueConfig: - ... + def task_queue_config(self) -> TaskQueueConfig: ... - async def _on_request_start(self, request: Request) -> str: - ... + async def _on_request_start(self, request: Request) -> str: ... - def assert_resource_exists(self, resource_id: str | None) -> None: - ... + def assert_resource_exists(self, resource_id: str | None) -> None: ... - async def _peek_terminal(self, request_id: str, *, fallback_status: str) -> TaskEnvelope: - ... + async def _peek_terminal(self, request_id: str, *, fallback_status: str) -> TaskEnvelope: ... - async def submit_and_peek(self, *args: Any, **kwargs: Any) -> TaskEnvelope: - ... + async def submit_and_peek(self, *args: Any, **kwargs: Any) -> TaskEnvelope: ... - async def call_backend(self, fn: Any, /, *args: Any, admit: bool = True, **kwargs: Any) -> Any: - ... + async def call_backend(self, fn: Any, /, *args: Any, admit: bool = True, **kwargs: Any) -> Any: ... @runtime_checkable @@ -56,5 +51,4 @@ class DataParallelDeployment(QueuedDeployment, Protocol): """A queued deployment that also shards a batch across data-parallel ranks.""" @property - def data_world_size(self) -> int: - ... + def data_world_size(self) -> int: ... 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..a1cd3a441 100644 --- a/src/twinkle/server/lifecycle/submit.py +++ b/src/twinkle/server/lifecycle/submit.py @@ -5,17 +5,24 @@ :class:`~twinkle.server.task_queue.mixin.TaskQueueMixin`, since it operates on queue state; this module owns the request-shaped pieces around it. """ -from __future__ import annotations +from __future__ import annotations import uuid from collections.abc import Callable, Coroutine -from fastapi import Request from typing import TYPE_CHECKING, Any +from fastapi import Request + from twinkle.data_format import InputFeature, Trajectory, is_encoded 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 @@ -50,7 +57,7 @@ def to_backend_inputs(inputs: Any, *, single: bool = False) -> Any: if single: if isinstance(entries, list): if len(entries) != 1: - raise ValueError('Streaming only supports a single input') + raise ValueError("Streaming only supports a single input") entries = entries[0] if isinstance(entries, dict): return _as_backend_entry(entries) @@ -85,14 +92,14 @@ 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))}') + raise ValueError(f"{name} collides with declared backend parameters: {', '.join(sorted(overlap))}") kwargs.update(region) return kwargs @@ -104,26 +111,26 @@ def input_metrics(self: DataParallelDeployment, body: Any, *, data_parallel: boo ``input_ids`` is either absent or a list of ints. """ inputs = body.inputs - input_tokens = sum(len(getattr(entry, 'input_ids', None) or ()) for entry in inputs) - metrics: dict[str, Any] = {'input_tokens': input_tokens} + input_tokens = sum(len(getattr(entry, "input_ids", None) or ()) for entry in inputs) + metrics: dict[str, Any] = {"input_tokens": input_tokens} if data_parallel: - metrics['batch_size'] = len(inputs) - metrics['data_world_size'] = self.data_world_size + metrics["batch_size"] = len(inputs) + metrics["data_world_size"] = self.data_world_size return metrics def resolve_twinkle_adapter_name(request: Request, adapter_name: str | None) -> str | None: """Build a stable per-session adapter name, falling back to request_id for older clients.""" - if adapter_name is None or adapter_name == '': + if adapter_name is None or adapter_name == "": return None owner_id = get_session_id_from_request(request) or request.state.request_id - return owner_id + '-' + adapter_name + return owner_id + "-" + adapter_name async def run_submit( self: QueuedDeployment, request: Request, - body: Any, + body: TrainingRequest, *, task_type: str, backend_call: Callable[..., Coroutine], @@ -152,13 +159,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 +172,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 +196,36 @@ 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..6ea1cd20a 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -5,47 +5,44 @@ Builds a single Ray Serve deployment (ModelManagement) that simultaneously handles both Tinker (/tinker/*) and Twinkle (/twinkle/*) model endpoints. """ -from __future__ import annotations +from __future__ import annotations import asyncio import importlib - -from fastapi import FastAPI, HTTPException, Request -from ray import serve -from ray.serve.config import RequestRouterConfig from typing import Any +from fastapi import FastAPI, Request + 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 logger = get_logger() -__all__ = ['ModelManagement', 'build_model_app'] +__all__ = ["ModelManagement", "build_model_app"] # Ctor kwargs consumed by the MultiLora wrappers' signatures but unknown to the # plain (full-parameter) model classes, where **kwargs flows into HF # ``from_pretrained``. Dropped when constructing a full-mode backend. -_LORA_ONLY_CTOR_KWARGS = ('max_loras', 'max_r', 'max_length', 'target_modules', 'lora_config') +_LORA_ONLY_CTOR_KWARGS = ("max_loras", "max_r", "max_length", "target_modules", "lora_config") def _make_mock_model(kw: dict[str, Any]) -> Any: from .backends.mock_model import TwinkleCompatMockModel - kw.pop('train_mode', None) # mock backend behaves the same for lora/full + kw.pop("train_mode", None) # mock backend behaves the same for lora/full return TwinkleCompatMockModel(**kw) @@ -63,25 +60,28 @@ def _make_model_class(base: type) -> type: """ from .backends.transformers_model import _TransformersTinkerCompatMixin - return remote_class()(type(f'TwinkleCompat{base.__name__}', (_TransformersTinkerCompatMixin, base), {})) + return remote_class()(type(f"TwinkleCompat{base.__name__}", (_TransformersTinkerCompatMixin, base), {})) def _make_transformers_model(kw: dict[str, Any]) -> Any: - train_mode = kw.pop('train_mode', 'lora') + train_mode = kw.pop("train_mode", "lora") # ``model_cls`` arrives as plain data and is resolved here, so the class object never reaches the # constructor -- it would either be rejected by the remote-class layer in safe mode # (``TWINKLE_TRUST_REMOTE_CODE=0``) or leak into every worker's constructor kwargs. The schema only # allows it for a LoRA-mode transformers deployment. - model_cls_spec = kw.pop('model_cls', None) - if train_mode == 'full': + model_cls_spec = kw.pop("model_cls", None) + if train_mode == "full": from .backends.transformers_model import TwinkleCompatFullTransformersModel + for key in _LORA_ONLY_CTOR_KWARGS: kw.pop(key, None) return TwinkleCompatFullTransformersModel(**kw) - from .backends.transformers_model import (TwinkleCompatSpectralHybridTransformersModel, - TwinkleCompatTransformersModel) + from .backends.transformers_model import ( + TwinkleCompatSpectralHybridTransformersModel, + TwinkleCompatTransformersModel, + ) - if train_mode == 'hybrid': + if train_mode == "hybrid": return TwinkleCompatSpectralHybridTransformersModel(**kw) if model_cls_spec is not None: return _make_model_class(_resolve_model_cls(model_cls_spec))(**kw) @@ -89,13 +89,15 @@ def _make_transformers_model(kw: dict[str, Any]) -> Any: def _make_megatron_model(kw: dict[str, Any]) -> Any: - train_mode = kw.pop('train_mode', 'lora') - if train_mode == 'full': + train_mode = kw.pop("train_mode", "lora") + if train_mode == "full": from .backends.megatron_model import TwinkleCompatFullMegatronModel + for key in _LORA_ONLY_CTOR_KWARGS: kw.pop(key, None) return TwinkleCompatFullMegatronModel(**kw) from .backends.megatron_model import TwinkleCompatMegatronModel + return TwinkleCompatMegatronModel(**kw) @@ -109,7 +111,7 @@ def _require_trusted_imports(spec: str) -> None: (``TWINKLE_TRUST_REMOTE_CODE=0``) would promise more than it delivers. """ if not trust_remote_code(): - raise ValueError(f'import spec {spec!r} requires TWINKLE_TRUST_REMOTE_CODE != 0') + raise ValueError(f"import spec {spec!r} requires TWINKLE_TRUST_REMOTE_CODE != 0") def _resolve_model_loader(spec: str, model_id: str) -> Any: @@ -121,17 +123,18 @@ def _resolve_model_loader(spec: str, model_id: str) -> Any: ``ModelLoaderProtocol`` object, or ``None`` to keep the generic AutoConfig/AutoModel path. """ _require_trusted_imports(spec) - module_name, _, attr = spec.partition(':') + module_name, _, attr = spec.partition(":") if not attr: - raise ValueError( - f'model_loader must be a "module:callable" import spec naming a loader factory, got {spec!r}') + raise ValueError(f'model_loader must be a "module:callable" import spec naming a loader factory, got {spec!r}') factory = getattr(importlib.import_module(module_name), attr) loader = factory(model_id) # Fail here, not at first use inside the model constructor: a misspelled factory that returns # something else would otherwise surface as an opaque AttributeError deep in the backend. - if loader is not None and not callable(getattr(loader, 'build_config', None)): - raise TypeError(f'model_loader factory {spec!r} returned {loader!r}, which is neither None nor a ' - 'ModelLoaderProtocol (missing callable build_config)') + if loader is not None and not callable(getattr(loader, "build_config", None)): + raise TypeError( + f"model_loader factory {spec!r} returned {loader!r}, which is neither None nor a " + "ModelLoaderProtocol (missing callable build_config)" + ) return loader @@ -144,100 +147,106 @@ def _resolve_model_cls(spec: str) -> type: ``swift.dev.model.multi_lora_sentence_transformer_model:MultiLoraSentenceTransformerModel``. """ _require_trusted_imports(spec) - module_name, _, attr = spec.partition(':') + module_name, _, attr = spec.partition(":") if not attr: raise ValueError(f'model_cls must be a "module:attr" import spec naming a model class, got {spec!r}') obj = getattr(importlib.import_module(module_name), attr) if not (isinstance(obj, type) and issubclass(obj, TrainableModel)): - raise TypeError(f'model_cls must name a TrainableModel subclass, got {spec!r} -> {obj!r}') + raise TypeError(f"model_cls must name a TrainableModel subclass, got {spec!r} -> {obj!r}") return obj # Single validate-then-dispatch selector for the model backend. MODEL_SELECTOR = BackendSelector( - 'backend', + "backend", { - 'mock': _make_mock_model, - 'transformers': _make_transformers_model, - 'megatron': _make_megatron_model, + "mock": _make_mock_model, + "transformers": _make_transformers_model, + "megatron": _make_megatron_model, }, ) -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 """ - async def __init__(self, - model_id: str, - nproc_per_node: int, - 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): + async def __init__( + self, + model_id: str, + nproc_per_node: int, + device_group: dict[str, Any], + device_mesh: dict[str, Any], + backend: str, + 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'), + 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.max_loras = kwargs.get('max_loras', 5) + 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 # full-parameter training. Popped here so the underlying model # constructor never receives it; re-injected into ctor_kwargs so the # backend builder can pick the full vs LoRA wrapper class. - self.train_mode = kwargs.pop('train_mode', 'lora') + self.train_mode = kwargs.pop("train_mode", "lora") # ``model_loader`` arrives as an import spec (plain data across the Ray boundary); resolve it to # the loader object here, in the replica, right before construction. ``None`` (unset, or no # family matched) leaves the generic AutoConfig/AutoModel path untouched. - model_loader_spec = kwargs.pop('model_loader', None) + model_loader_spec = kwargs.pop("model_loader", None) ctor_kwargs: dict[str, Any] = { - 'model_id': model_id, - 'remote_group': self.device_group.name, - 'instance_id': self.replica_id, - 'train_mode': self.train_mode, + "model_id": model_id, + "remote_group": self.device_group.name, + "instance_id": self.replica_id, + "train_mode": self.train_mode, **kwargs, } if model_loader_spec is not None: - ctor_kwargs['model_loader'] = _resolve_model_loader(model_loader_spec, model_id) + 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 + ctor_kwargs["device_mesh"] = self.device_mesh 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 + self.data_plane = DataPlaneProxy(data_plane_url) self._replica_registered = False self._model_unhealthy = False self._health_probe_task: asyncio.Task | None = None - actors = getattr(self.model, '_actors', None) + actors = getattr(self.model, "_actors", None) self._init_task_queue( queue_config, - deployment_name='Model', + deployment_name="Model", enable_admission_gate=True, on_backend_timeout=self._probe_after_timeout, 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 +261,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,49 +277,31 @@ 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}'} + return {"healthy": False, "detail": f"unexpected ping result: {result}"} except Exception as e: self._model_unhealthy = True - return {'healthy': False, 'detail': f'model actor unreachable: {e}'} + return {"healthy": False, "detail": f"model actor unreachable: {e}"} async def check_model_health(self) -> dict: """Run one coalesced actor probe outside the event loop.""" - current = getattr(self, '_health_probe_task', None) + current = getattr(self, "_health_probe_task", None) if current is None or current.done(): self._health_probe_task = asyncio.create_task(self._run_model_health_probe()) 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: """True when this deployment runs exclusive full-parameter training.""" - return self.train_mode == 'full' + return self.train_mode == "full" def resolve_model_adapter_name(self, adapter_name: str | None) -> str | None: """Map a per-tenant adapter name to the name the model expects. @@ -338,40 +310,21 @@ def resolve_model_adapter_name(self, adapter_name: str | None) -> str | None: optimizer group; in LoRA mode the per-tenant adapter name is used as-is. """ if self.is_full_mode: - return '' + 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, - device_group: dict[str, Any], - 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): +def build_model_app( + model_id: str, + nproc_per_node: int, + device_group: dict[str, Any], + device_mesh: dict[str, Any], + deploy_options: dict[str, Any], + backend: str, + queue_config: TaskQueueConfig | None = None, + data_plane_url: str | None = None, + **kwargs, +): """Build a unified model management application for distributed training. Supports both Tinker (polling-style) and Twinkle (synchronous) clients. @@ -385,7 +338,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 +350,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) @@ -407,16 +358,14 @@ def register_routes(app: FastAPI, get_self: Any) -> None: async def _on_shutdown(servable: Any) -> None: await servable.shutdown() - app = build_deployment_app('Model', register_routes, on_shutdown=_on_shutdown, attach_replica_id_header=True) + app = build_deployment_app("Model", register_routes, on_shutdown=_on_shutdown, attach_replica_id_header=True) return bind_deployment( app, 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), + deployment_name="ModelManagement", + 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..900f2c860 100644 --- a/src/twinkle/server/model/backends/common.py +++ b/src/twinkle/server/model/backends/common.py @@ -2,15 +2,18 @@ """ Shared helpers and base classes for backend model implementations. """ -import numpy as np + import re -import torch from collections.abc import Mapping from numbers import Number +from typing import Any + +import numpy as np +import torch 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 @@ -88,7 +91,6 @@ def to_cpu_safe_output(obj: Any) -> Any: def clean_metrics(metrics: dict) -> dict: - def _to_float(v): if isinstance(v, (float, int, Number, np.generic, str)): try: @@ -114,9 +116,9 @@ def _to_float(v): if s: try: head, unit = s.split(maxsplit=1) - cleaned[f'{key}/{unit}'] = float(head) + cleaned[f"{key}/{unit}"] = float(head) except Exception: - m = re.match(r'^([+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)', s) + m = re.match(r"^([+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)", s) if m: cleaned[key] = float(m.group(1)) @@ -129,11 +131,11 @@ class TwinkleCompatModelBase: @staticmethod def _normalize_ref_outputs(kwargs: dict) -> None: """Convert HTTP/Ray-serialized reference logps to one padded tensor.""" - ref_outputs = kwargs.get('ref_outputs') + ref_outputs = kwargs.get("ref_outputs") if not isinstance(ref_outputs, dict): return - logps = ref_outputs.get('logps') + logps = ref_outputs.get("logps") if not isinstance(logps, (list, tuple)) or not logps or isinstance(logps[0], torch.Tensor): return @@ -145,7 +147,8 @@ def _normalize_ref_outputs(kwargs: dict) -> None: rows.append(item) from twinkle.utils import pad_and_stack_tensors - ref_outputs['logps'] = pad_and_stack_tensors( + + ref_outputs["logps"] = pad_and_stack_tensors( [torch.as_tensor(row, dtype=torch.float32) for row in rows], pad_value=0.0, concat=False, @@ -157,21 +160,51 @@ def get_template(self, adapter_name: str) -> Template: def _tinker_setup_loss(self, loss_fn: str, inputs, adapter_name: str, kwargs: dict): """Set up loss function based on loss_fn; pops DPO/GRPO-specific params from kwargs in-place.""" is_dpo = validate_tinker_loss(loss_fn, inputs) - if loss_fn == 'cross_entropy': - self.set_loss('CrossEntropyLoss', adapter_name=adapter_name) - elif loss_fn == 'importance_sampling': + if loss_fn == "cross_entropy": + self.set_loss("CrossEntropyLoss", adapter_name=adapter_name) + elif loss_fn == "importance_sampling": if is_dpo: - beta = kwargs.pop('dpo_beta', 0.1) - loss_type = kwargs.pop('dpo_loss_type', 'sigmoid') - sft_weight = kwargs.pop('dpo_sft_weight', 0.0) + beta = kwargs.pop("dpo_beta", 0.1) + loss_type = kwargs.pop("dpo_loss_type", "sigmoid") + sft_weight = kwargs.pop("dpo_sft_weight", 0.0) self.set_loss( - 'DPOLoss', adapter_name=adapter_name, beta=beta, loss_type=loss_type, sft_weight=sft_weight) + "DPOLoss", adapter_name=adapter_name, beta=beta, loss_type=loss_type, sft_weight=sft_weight + ) # Only add DPOMetric if not already present for this adapter self._ensure_dpo_metric(adapter_name, beta) else: - epsilon = kwargs.pop('epsilon', 0.2) - grpo_beta = kwargs.pop('beta', 0.0) - self.set_loss('GRPOLoss', adapter_name=adapter_name, epsilon=epsilon, beta=grpo_beta) + epsilon = kwargs.pop("epsilon", 0.2) + 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. @@ -179,12 +212,13 @@ def _ensure_dpo_metric(self, adapter_name: str, beta: float): This prevents duplicate metric accumulation across training steps. """ from twinkle.metric.dpo import DPOMetric + optimizer_config = self.optimizer_group[adapter_name] # Check if DPOMetric already exists in training metrics for metric in optimizer_config.train_status.metrics: if isinstance(metric, DPOMetric): return - self.add_metric('DPOMetric', adapter_name=adapter_name, beta=beta) + self.add_metric("DPOMetric", adapter_name=adapter_name, beta=beta) def _apply_ref_outputs(self, loss_values: dict, loss_kwargs: dict, adapter_name: str) -> None: """Pop ref_outputs from loss_values into loss_kwargs and propagate to train_status. @@ -192,17 +226,17 @@ def _apply_ref_outputs(self, loss_values: dict, loss_kwargs: dict, adapter_name: DPOMetric reads ref_outputs from train_status.forward_kwargs during accumulate_metrics, so it must be set here before the subsequent loss calculation. """ - if 'ref_outputs' not in loss_values: + if "ref_outputs" not in loss_values: return - ref_outputs_dict = loss_values.pop('ref_outputs') - loss_kwargs['ref_outputs'] = ref_outputs_dict - self.optimizer_group[adapter_name].train_status.forward_kwargs['ref_outputs'] = ref_outputs_dict + ref_outputs_dict = loss_values.pop("ref_outputs") + loss_kwargs["ref_outputs"] = ref_outputs_dict + self.optimizer_group[adapter_name].train_status.forward_kwargs["ref_outputs"] = ref_outputs_dict def _tinker_build_output(self, inputs, outputs, return_full_logprobs: bool = False): """Extract logits/logps from model outputs and build per-datum output list.""" - seq_lens = [feature.loss_fn_inputs['target_tokens'].to_torch().long().view(-1).numel() for feature in inputs] - logits = self._tensor_output_to_rows(outputs.get('logits'), seq_lens, kind='logits') - logps = self._tensor_output_to_rows(outputs.get('logps'), seq_lens, kind='logps') + seq_lens = [feature.loss_fn_inputs["target_tokens"].to_torch().long().view(-1).numel() for feature in inputs] + logits = self._tensor_output_to_rows(outputs.get("logits"), seq_lens, kind="logits") + logps = self._tensor_output_to_rows(outputs.get("logps"), seq_lens, kind="logps") if logits is None and logps is None: # non-last PP stage: no outputs produced, collector will discard this return [] @@ -234,10 +268,11 @@ def _tensor_output_to_rows(value, seq_lens: list[int], *, kind: str) -> list[tor if isinstance(value, list) and value and isinstance(value[0], (list, tuple)): flat = [s for item in value for s in (item if isinstance(item[0], (list, tuple)) else [item])] from twinkle.utils import pad_and_stack_tensors + tensors = [ - pad_and_stack_tensors([torch.tensor(s, dtype=torch.float32) for s in flat], - pad_value=0.0, - concat=False) + pad_and_stack_tensors( + [torch.tensor(s, dtype=torch.float32) for s in flat], pad_value=0.0, concat=False + ) ] else: tensors = [torch.as_tensor(value, dtype=torch.float32)] @@ -248,13 +283,13 @@ def _tensor_output_to_rows(value, seq_lens: list[int], *, kind: str) -> list[tor rows = [] for tensor in tensors: - chunk_rows = TwinkleCompatModelBase._tensor_chunk_to_rows(tensor, seq_lens[len(rows):], kind=kind) + chunk_rows = TwinkleCompatModelBase._tensor_chunk_to_rows(tensor, seq_lens[len(rows) :], kind=kind) rows.extend(chunk_rows) return rows @staticmethod def _normalize_single_row_tensor(tensor: torch.Tensor, *, kind: str) -> torch.Tensor: - if kind == 'logits': + if kind == "logits": if tensor.dim() >= 3 and tensor.shape[0] == 1: return tensor.squeeze(0) return tensor @@ -272,7 +307,7 @@ def _split_packed_tensor(tensor: torch.Tensor, seq_lens: list[int]) -> list[torc for seq_len in seq_lens: if offset + seq_len > tensor.shape[0]: break - split_rows.append(tensor[offset:offset + seq_len]) + split_rows.append(tensor[offset : offset + seq_len]) offset += seq_len if offset == tensor.shape[0]: break @@ -283,7 +318,7 @@ def _tensor_chunk_to_rows(tensor: torch.Tensor, seq_lens: list[int], *, kind: st if tensor.dim() == 0 or len(seq_lens) <= 1: return [TwinkleCompatModelBase._normalize_single_row_tensor(tensor, kind=kind)] - if kind == 'logps': + if kind == "logps": while tensor.dim() > 1 and tensor.shape[0] == 1: tensor = tensor.squeeze(0) if tensor.dim() >= 2: @@ -297,7 +332,7 @@ def _tensor_chunk_to_rows(tensor: torch.Tensor, seq_lens: list[int], *, kind: st return packed_rows return [tensor] - if kind == 'logits': + if kind == "logits": if tensor.dim() >= 3 and tensor.shape[0] == 1: packed_rows = TwinkleCompatModelBase._split_packed_tensor(tensor.squeeze(0), seq_lens) if packed_rows is not None: @@ -314,13 +349,15 @@ def _tensor_chunk_to_rows(tensor: torch.Tensor, seq_lens: list[int], *, kind: st return packed_rows return [tensor] - raise ValueError(f'Unsupported tensor output kind: {kind}') + raise ValueError(f"Unsupported tensor output kind: {kind}") @staticmethod - def _get_forward_output(inputs: list[types.Datum], - logits: list[torch.Tensor] | None, - logps: list[torch.Tensor] | None, - return_full_logprobs: bool = False) -> list[dict]: + def _get_forward_output( + inputs: list[types.Datum], + logits: list[torch.Tensor] | None, + logps: list[torch.Tensor] | None, + return_full_logprobs: bool = False, + ) -> list[dict]: """Convert raw logits to the expected output format with logprobs and elementwise_loss. When return_full_logprobs is True (forward_only / reference pass), logprobs is returned @@ -333,28 +370,29 @@ def _get_forward_output(inputs: list[types.Datum], """ if logps is not None: if len(logps) != len(inputs): - raise ValueError(f'Expected {len(inputs)} logps rows, got {len(logps)}') + raise ValueError(f"Expected {len(inputs)} logps rows, got {len(logps)}") device = logps[0].device elif logits is not None: if len(logits) != len(inputs): - raise ValueError(f'Expected {len(inputs)} logits rows, got {len(logits)}') + raise ValueError(f"Expected {len(inputs)} logits rows, got {len(logits)}") device = logits[0].device else: - raise ValueError('At least one of logits or logps must be provided.') + raise ValueError("At least one of logits or logps must be provided.") results = [] if logits is None: logits = [None] * len(inputs) if logps is None: logps = [None] * len(inputs) - for feature, logit, feature_logps in zip(inputs, logits, logps): - labels = feature.loss_fn_inputs['target_tokens'].to_torch().long().view(-1).to(device) - weights = feature.loss_fn_inputs['weights'].to_torch().view(-1).to(device) + for feature, logit, feature_logps in zip(inputs, logits, logps, strict=False): + labels = feature.loss_fn_inputs["target_tokens"].to_torch().long().view(-1).to(device) + weights = feature.loss_fn_inputs["weights"].to_torch().view(-1).to(device) seq_len = labels.numel() # original unpadded length if feature_logps is None: - assert logit is not None, 'logit must not be None when logps is None' + assert logit is not None, "logit must not be None when logps is None" from twinkle.utils.torch_utils import selective_log_softmax + feature_logits = logit[:seq_len, :] token_log_probs_orig = selective_log_softmax(feature_logits, labels) if return_full_logprobs: @@ -363,6 +401,7 @@ def _get_forward_output(inputs: list[types.Datum], padded_len = logit.shape[0] if padded_len > seq_len: import torch.nn.functional as F + token_log_probs_full = F.pad(token_log_probs_orig, (0, padded_len - seq_len), value=0.0) else: token_log_probs_full = token_log_probs_orig @@ -379,8 +418,10 @@ def _get_forward_output(inputs: list[types.Datum], token_log_probs_orig = token_log_probs_orig.to(weights.device) elementwise_loss = -token_log_probs_orig * weights - results.append({ - 'logprobs': types.TensorData.from_torch(token_log_probs_full.cpu()), - 'elementwise_loss': types.TensorData.from_torch(elementwise_loss.cpu()) - }) + results.append( + { + "logprobs": types.TensorData.from_torch(token_log_probs_full.cpu()), + "elementwise_loss": types.TensorData.from_torch(elementwise_loss.cpu()), + } + ) return results diff --git a/src/twinkle/server/model/backends/mock_model.py b/src/twinkle/server/model/backends/mock_model.py index 22abfb64b..cbd31886e 100644 --- a/src/twinkle/server/model/backends/mock_model.py +++ b/src/twinkle/server/model/backends/mock_model.py @@ -10,12 +10,13 @@ The class is duck-typed against ``TwinkleCompatModelBase`` rather than subclassing it — the base class lives in a torch-importing module. """ -from __future__ import annotations -import numpy as np +from __future__ import annotations import os from typing import Any +import numpy as np + from twinkle import remote_class, remote_function from twinkle.utils.logger import get_logger from twinkle.utils.seed import stable_seed @@ -43,7 +44,7 @@ def __init__( # adapter_name -> arbitrary config payload self._adapters: dict[str, dict[str, Any]] = {} if kwargs: - logger.debug('MockModel ignoring unknown ctor kwargs: %s', sorted(kwargs)) + logger.debug("MockModel ignoring unknown ctor kwargs: %s", sorted(kwargs)) # ----- Forward family ------------------------------------------------ # @@ -65,11 +66,13 @@ def _build_forward_result( rng = np.random.default_rng(stable_seed(self.model_id, adapter_name, self._rng_seed, idx, seq_len)) logprobs = rng.uniform(-2.0, 0.0, size=seq_len).astype(np.float32) elementwise_loss = rng.uniform(0.0, 1.0, size=seq_len).astype(np.float32) - out.append({ - 'logprobs': logprobs.tolist(), - 'elementwise_loss': elementwise_loss.tolist(), - 'loss': float(loss_value), - }) + out.append( + { + "logprobs": logprobs.tolist(), + "elementwise_loss": elementwise_loss.tolist(), + "loss": float(loss_value), + } + ) return out @remote_function() @@ -78,22 +81,22 @@ def tinker_forward_only(self, *, inputs: Any, adapter_name: str | None = None, * @remote_function() def tinker_forward_backward(self, *, inputs: Any, adapter_name: str, loss_fn: str, **kwargs: Any) -> list[Any]: - loss_seed = stable_seed(self.model_id, adapter_name, self._rng_seed, 'loss', loss_fn) + loss_seed = stable_seed(self.model_id, adapter_name, self._rng_seed, "loss", loss_fn) loss = float(np.random.default_rng(loss_seed).uniform(0.0, 1.0)) return [_to_tinker_loss_outputs(self._build_forward_result(inputs, adapter_name, loss_value=loss)), loss] @remote_function() def forward(self, *, inputs: Any, **kwargs: Any) -> list[dict[str, Any]]: - return self._build_forward_result(inputs, kwargs.get('adapter_name')) + return self._build_forward_result(inputs, kwargs.get("adapter_name")) @remote_function() def forward_only(self, *, inputs: Any, **kwargs: Any) -> list[dict[str, Any]]: - return self._build_forward_result(inputs, kwargs.get('adapter_name')) + return self._build_forward_result(inputs, kwargs.get("adapter_name")) @remote_function() def forward_backward(self, *, inputs: Any, **kwargs: Any) -> list[Any]: loss = float(np.random.default_rng(self._rng_seed).uniform(0.0, 1.0)) - return [self._build_forward_result(inputs, kwargs.get('adapter_name'), loss_value=loss), loss] + return [self._build_forward_result(inputs, kwargs.get("adapter_name"), loss_value=loss), loss] @remote_function() def calculate_loss(self, *, inputs: Any = None, **kwargs: Any) -> float: @@ -133,11 +136,11 @@ def tinker_step(self, *, adam_params: Any = None, **kwargs: Any) -> None: @remote_function() def tinker_calculate_metric(self, is_training: bool, **kwargs: Any) -> dict[str, float]: - return {'loss': 0.5, 'grad_norm': 0.1} + return {"loss": 0.5, "grad_norm": 0.1} @remote_function() def calculate_metric(self, *args: Any, **kwargs: Any) -> dict[str, float]: - return {'loss': 0.5, 'grad_norm': 0.1} + return {"loss": 0.5, "grad_norm": 0.1} @remote_function() def tinker_load(self, *, checkpoint_name: str, output_dir: str | None = None, **kwargs: Any) -> None: @@ -176,15 +179,15 @@ def apply_patch(self, *args: Any, **kwargs: Any) -> None: # ----- Persistence stubs --------------------------------------------- # @remote_function() - def save(self, name: str = 'latest', output_dir: str | None = None, **kwargs: Any) -> str: + def save(self, name: str = "latest", output_dir: str | None = None, **kwargs: Any) -> str: """Materialize an empty checkpoint dir so downstream existence checks pass. 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) + target = os.path.join(output_dir or "/tmp/twinkle_mock_ckpt", name) os.makedirs(target, exist_ok=True) return target @@ -199,7 +202,7 @@ def reload_initial_weights(self, *args: Any, **kwargs: Any) -> None: @remote_function() def resume_from_checkpoint(self, *args: Any, **kwargs: Any) -> dict[str, Any]: - return {'status': 'ok', 'progress': {}} + return {"status": "ok", "progress": {}} @remote_function() def get_state_dict(self, *args: Any, **kwargs: Any) -> dict[str, Any]: @@ -209,7 +212,7 @@ def get_state_dict(self, *args: Any, **kwargs: Any) -> dict[str, Any]: def get_train_configs(self, *args: Any, **kwargs: Any) -> str: # Real backend returns a JSON-serialized config string; handler wraps it # into ``GetTrainConfigsResponse.result: str`` (types/model.py:213). - return '{}' + return "{}" @remote_function() def upload_to_hub(self, *args: Any, **kwargs: Any) -> None: @@ -226,21 +229,21 @@ def add_adapter(self, adapter_name: str, **cfg: Any) -> None: def add_adapter_to_model(self, adapter_name: str, config: Any = None, **cfg: Any) -> None: merged: dict[str, Any] = dict(cfg) if config is not None: - merged.setdefault('config', config) + merged.setdefault("config", config) self._adapters[adapter_name] = merged @remote_function() def remove_adapter(self, adapter_name: str) -> None: """Remove ``adapter_name``; raise ``KeyError`` if it was never added.""" if adapter_name not in self._adapters: - raise KeyError(f'adapter not present: {adapter_name}') + raise KeyError(f"adapter not present: {adapter_name}") del self._adapters[adapter_name] @remote_function() def has_adapter(self, adapter_name: str) -> bool: return adapter_name in self._adapters - @remote_function(collect='first', lazy_collect=False, sync=True, timeout=4) + @remote_function(collect="first", lazy_collect=False, sync=True, timeout=4) def ping(self) -> bool: """Lightweight liveness probe for watchdog health checks.""" return True @@ -260,17 +263,18 @@ def _to_tinker_loss_outputs(records: list[dict[str, Any]]) -> list[dict[str, Any out: list[dict[str, Any]] = [] for rec in records: - logprobs = rec['logprobs'] - elementwise_loss = rec['elementwise_loss'] - loss = float(rec['loss']) - out.append({ - 'logprobs': - TensorData(data=list(logprobs), dtype='float32', shape=[len(logprobs)]), - 'elementwise_loss': - TensorData(data=list(elementwise_loss), dtype='float32', shape=[len(elementwise_loss)]), - 'loss': - TensorData(data=[loss], dtype='float32', shape=[]), - }) + logprobs = rec["logprobs"] + elementwise_loss = rec["elementwise_loss"] + loss = float(rec["loss"]) + out.append( + { + "logprobs": TensorData(data=list(logprobs), dtype="float32", shape=[len(logprobs)]), + "elementwise_loss": TensorData( + data=list(elementwise_loss), dtype="float32", shape=[len(elementwise_loss)] + ), + "loss": TensorData(data=[loss], dtype="float32", shape=[]), + } + ) return out @@ -294,17 +298,17 @@ def _input_seq_lengths(inputs: Any) -> list[int]: def _seq_length_of(item: Any) -> int: # Datum-like: model_input.tokens or loss_fn_inputs['target_tokens'] - for attr in ('model_input', 'inputs', 'tokens'): + for attr in ("model_input", "inputs", "tokens"): v = getattr(item, attr, None) if v is None: continue - tokens = getattr(v, 'tokens', v) - if hasattr(tokens, '__len__'): + tokens = getattr(v, "tokens", v) + if hasattr(tokens, "__len__"): return max(1, len(tokens)) if isinstance(item, dict): - for k in ('input_ids', 'tokens', 'target_tokens'): - if k in item and hasattr(item[k], '__len__'): + for k in ("input_ids", "tokens", "target_tokens"): + if k in item and hasattr(item[k], "__len__"): return max(1, len(item[k])) - if hasattr(item, '__len__'): + if hasattr(item, "__len__"): return max(1, len(item)) return 1 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..ba4511612 100644 --- a/src/twinkle/server/model/tinker_handlers.py +++ b/src/twinkle/server/model/tinker_handlers.py @@ -7,22 +7,21 @@ /tinker/... and use schedule_task() returning UntypedAPIFuture. ``self_fn`` is injected via FastAPI Depends to obtain the ModelManagement instance at request time. """ -from __future__ import annotations -import traceback +from __future__ import annotations from collections.abc import Callable -from fastapi import Depends, FastAPI, Request from functools import partial -from tinker import types from typing import TYPE_CHECKING +from fastapi import Depends, FastAPI, Request +from tinker import types + if TYPE_CHECKING: 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 @@ -36,77 +35,25 @@ def _register_model_tinker_routes(app: FastAPI, self_fn: Callable[[], ModelManag replica instance. It is wired in via Depends so it is resolved lazily at request time. """ - @app.post('/tinker/create_model') + @app.post("/tinker/create_model") async def create_model( - request: Request, - body: types.CreateModelRequest, - self: ModelManagement = Depends(self_fn), + request: Request, + body: types.CreateModelRequest, + self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: - token = await self._on_request_start(request) + raise RequestRejectedError("Training lifecycle operations must use the catalog endpoint") - 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') - - @app.post('/tinker/get_info') - async def get_info(request: Request, body: types.GetInfoRequest, - self: ModelManagement = Depends(self_fn)) -> types.GetInfoResponse: + @app.post("/tinker/get_info") + async def get_info( + request: Request, body: types.GetInfoRequest, self: ModelManagement = Depends(self_fn) + ) -> types.GetInfoResponse: token = await self._on_request_start(request) - training_run_manager = create_training_run_manager(token, client_type='tinker') + training_run_manager = create_training_run_manager(token, client_type="tinker") metadata = training_run_manager.get(str(body.model_id)) model_name = metadata.base_model if metadata else self.base_model lora_rank = None is_lora = False - if metadata and hasattr(metadata, 'lora_rank') and metadata.lora_rank: + if metadata and hasattr(metadata, "lora_rank") and metadata.lora_rank: lora_rank = metadata.lora_rank is_lora = metadata.is_lora return types.GetInfoResponse( @@ -117,232 +64,258 @@ async def get_info(request: Request, body: types.GetInfoRequest, model_name=model_name, ) - @app.post('/tinker/unload_model') + @app.post("/tinker/unload_model") async def unload_model( - request: Request, - body: types.UnloadModelRequest, - self: ModelManagement = Depends(self_fn), + request: Request, + 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) + raise RequestRejectedError("Training lifecycle operations must use the catalog endpoint") - return await self.schedule_task(_do_unload, model_id=body.model_id, token=token, task_type='unload_model') - - @app.post('/tinker/forward') - async def forward(request: Request, body: types.ForwardRequest, - self: ModelManagement = Depends(self_fn)) -> types.UntypedAPIFuture: + @app.post("/tinker/forward") + async def forward( + request: Request, body: types.ForwardRequest, self: ModelManagement = Depends(self_fn) + ) -> types.UntypedAPIFuture: 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, data_world_size=self.data_world_size, - task_type='forward', + task_type="forward", ) - @app.post('/tinker/forward_backward') + @app.post("/tinker/forward_backward") async def forward_backward( - request: Request, - body: types.ForwardBackwardRequest, - self: ModelManagement = Depends(self_fn), + request: Request, + body: types.ForwardBackwardRequest, + self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: 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}, + 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 ) - except Exception: - logger.error(traceback.format_exc()) - raise + 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, data_world_size=self.data_world_size, batch_size_multiple=2 if is_dpo else None, - task_type='forward_backward', + task_type="forward_backward", ) - @app.post('/tinker/optim_step') + @app.post("/tinker/optim_step") async def optim_step( - request: Request, - body: types.OptimStepRequest, - self: ModelManagement = Depends(self_fn), + request: Request, + body: types.OptimStepRequest, + self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: 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') - - @app.post('/tinker/save_weights') + 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( - request: Request, - body: types.SaveWeightsRequest, - self: ModelManagement = Depends(self_fn), + request: Request, + body: types.SaveWeightsRequest, + self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: 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') - - @app.post('/tinker/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=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( - request: Request, - body: types.SaveWeightsForSamplerRequest, - self: ModelManagement = Depends(self_fn), + request: Request, + body: types.SaveWeightsForSamplerRequest, + self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: 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') - - @app.post('/tinker/load_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=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( - request: Request, - body: types.LoadWeightsRequest, - self: ModelManagement = Depends(self_fn), + request: Request, + body: types.LoadWeightsRequest, + self: ModelManagement = Depends(self_fn), ) -> types.UntypedAPIFuture: 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..c88cab994 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -9,27 +9,29 @@ resolves the envelope to a terminal state. ``self_fn`` is injected via FastAPI Depends to obtain the ModelManagement instance at request time. """ -from __future__ import annotations -import torch +from __future__ import annotations from collections.abc import Callable -from fastapi import Depends, FastAPI, HTTPException, Request from functools import partial from typing import TYPE_CHECKING +import torch +from fastapi import Depends, FastAPI, HTTPException, Request + if TYPE_CHECKING: from .app import ModelManagement 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.model.data_plane_inputs import (data_plane_request_shape, merge_forward_kwargs, - resolve_data_plane_model_inputs, select_output_rows) +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 from twinkle.utils.logger import get_logger @@ -49,7 +51,7 @@ def _tokens_only_metrics(self, body): def _data_plane_metrics(self, body): """Scheduling metrics derived from DataRef shape for *_from_data_plane endpoints.""" input_tokens, batch_size = data_plane_request_shape(body) - return {'input_tokens': input_tokens, 'batch_size': batch_size, 'data_world_size': self.data_world_size} + return {"input_tokens": input_tokens, "batch_size": batch_size, "data_world_size": self.data_world_size} def _register_model_twinkle_routes(app: FastAPI, self_fn: Callable[[], ModelManagement]) -> None: @@ -59,72 +61,74 @@ def _register_model_twinkle_routes(app: FastAPI, self_fn: Callable[[], ModelMana replica instance. It is wired in via Depends so it is resolved lazily at request time. """ - @app.get('/healthz') + @app.get("/healthz") async def model_healthz( - request: Request, - self: ModelManagement = Depends(self_fn), + request: Request, + self: ModelManagement = Depends(self_fn), ) -> dict: """Deep health probe: pings underlying model actors to verify liveness.""" result = await self.check_model_health() - if self._model_unhealthy or not result['healthy']: + if self._model_unhealthy or not result["healthy"]: from fastapi.responses import JSONResponse + return JSONResponse(status_code=503, content=result) return result - @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() + @app.post("/twinkle/create", response_model=types.CreateResponse) + async def create( + request: Request, body: types.CreateRequest, self: ModelManagement = Depends(self_fn) + ) -> types.CreateResponse: + raise RequestRejectedError("Training lifecycle operations must use the catalog endpoint") # ------------------------------------------------------------------ # # Inline data / forward family # ------------------------------------------------------------------ # - @app.post('/twinkle/forward', response_model=types.TaskEnvelope) - async def forward(request: Request, body: types.ForwardRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + @app.post("/twinkle/forward", response_model=types.TaskEnvelope) + async def forward( + request: Request, body: types.ForwardRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): ret = await self.call_backend( self.model.forward, inputs=to_backend_inputs(body.inputs), adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) - return {'result': ret} + **backend_kwargs(body), + ) + return {"result": ret} return await run_submit( self, request, body, - task_type='forward', + task_type="forward", backend_call=_call, metrics=_dp_metrics, - capability=BackendCapability.Forward) + capability=BackendCapability.Forward, + ) - @app.post('/twinkle/forward_only', response_model=types.TaskEnvelope) + @app.post("/twinkle/forward_only", response_model=types.TaskEnvelope) async def forward_only( - request: Request, body: types.ForwardOnlyRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.ForwardOnlyRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): ret = await self.call_backend( self.model.forward_only, inputs=to_backend_inputs(body.inputs), adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) - return {'result': ret} + **backend_kwargs(body), + ) + return {"result": ret} return await run_submit( - self, request, body, task_type='forward_only', backend_call=_call, metrics=_tokens_only_metrics) + self, request, body, task_type="forward_only", backend_call=_call, metrics=_tokens_only_metrics + ) - @app.post('/twinkle/forward_backward', response_model=types.TaskEnvelope) + @app.post("/twinkle/forward_backward", response_model=types.TaskEnvelope) async def forward_backward( - request: Request, body: types.ForwardBackwardTaskRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.ForwardBackwardTaskRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): - def first_element(data): while isinstance(data, list): if len(data) == 0: @@ -137,55 +141,69 @@ 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)) - return {'result': ret} + self.model.forward_backward, inputs=all_inputs, adapter_name=resolved_adapter_name, **kwargs + ) + return {"result": ret} return await run_submit( - self, request, body, task_type='forward_backward', backend_call=_call, metrics=_dp_metrics) + self, request, body, task_type="forward_backward", backend_call=_call, metrics=_dp_metrics + ) - @app.post('/twinkle/calculate_loss', response_model=types.TaskEnvelope) + @app.post("/twinkle/calculate_loss", response_model=types.TaskEnvelope) async def calculate_loss( - request: Request, body: types.AdapterRequest, self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.AdapterRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): ret = await self.call_backend( self.model.calculate_loss, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) - return {'result': ret} + **backend_kwargs(body), + ) + return {"result": ret} return await run_submit( self, request, body, - task_type='calculate_loss', + task_type="calculate_loss", backend_call=_call, - capability=BackendCapability.CalculateLoss) - - @app.post('/twinkle/backward', response_model=types.TaskEnvelope) - async def backward(request: Request, body: types.AdapterRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: + capability=BackendCapability.CalculateLoss, + ) + @app.post("/twinkle/backward", response_model=types.TaskEnvelope) + async def backward( + request: Request, body: types.AdapterRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): await self.call_backend( - self.model.backward, adapter_name=self.resolve_model_adapter_name(adapter_name), **backend_kwargs(body)) + self.model.backward, adapter_name=self.resolve_model_adapter_name(adapter_name), **backend_kwargs(body) + ) return await run_submit( - self, request, body, task_type='backward', backend_call=_call, capability=BackendCapability.Backward) + self, request, body, task_type="backward", backend_call=_call, capability=BackendCapability.Backward + ) # ------------------------------------------------------------------ # # Data-plane forward family (DataRef inputs; response only enters the contract) # ------------------------------------------------------------------ # - @app.post('/twinkle/forward_from_data_plane', response_model=types.TaskEnvelope) + @app.post("/twinkle/forward_from_data_plane", response_model=types.TaskEnvelope) async def forward_from_data_plane( - request: Request, body: types.DataPlaneForwardRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.DataPlaneForwardRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: 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) @@ -193,23 +211,24 @@ async def _call(self, body, adapter_name, token): self.model.forward, inputs=to_backend_inputs(raw_inputs), adapter_name=self.resolve_model_adapter_name(adapter_name), - **kwargs) - return {'result': ret} + **kwargs, + ) + return {"result": ret} return await run_submit( self, request, body, - task_type='forward_from_data_plane', + task_type="forward_from_data_plane", backend_call=_call, metrics=_data_plane_metrics, - capability=BackendCapability.Forward) + capability=BackendCapability.Forward, + ) - @app.post('/twinkle/forward_only_from_data_plane', response_model=types.TaskEnvelope) + @app.post("/twinkle/forward_only_from_data_plane", response_model=types.TaskEnvelope) async def forward_only_from_data_plane( - request: Request, body: types.DataPlaneForwardOnlyRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.DataPlaneForwardOnlyRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): raw_inputs, field_kwargs = await resolve_data_plane_model_inputs(body, self.data_plane) inputs = to_backend_inputs(raw_inputs) @@ -218,197 +237,210 @@ async def _call(self, body, adapter_name, token): self.model.forward_only, inputs=inputs, adapter_name=self.resolve_model_adapter_name(adapter_name), - **kwargs) + **kwargs, + ) if body.output_ref is not None: rows = select_output_rows(ret, batch_size=len(inputs), output_fields=body.output_fields) output_ref = await self.data_plane.append(body.output_ref, rows) - return {'result': output_ref.model_dump()} - return {'result': ret} + return {"result": output_ref.model_dump()} + return {"result": ret} return await run_submit( self, request, body, - task_type='forward_only_from_data_plane', + task_type="forward_only_from_data_plane", backend_call=_call, - metrics=_data_plane_metrics) + metrics=_data_plane_metrics, + ) - @app.post('/twinkle/forward_backward_from_data_plane', response_model=types.TaskEnvelope) + @app.post("/twinkle/forward_backward_from_data_plane", response_model=types.TaskEnvelope) async def forward_backward_from_data_plane( - request: Request, body: types.DataPlaneForwardRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.DataPlaneForwardRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: 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) - return {'result': ret} + self.model.forward_backward, inputs=inputs, adapter_name=resolved_adapter_name, **kwargs + ) + return {"result": ret} return await run_submit( self, request, body, - task_type='forward_backward_from_data_plane', + task_type="forward_backward_from_data_plane", backend_call=_call, - metrics=_data_plane_metrics) + metrics=_data_plane_metrics, + ) # ------------------------------------------------------------------ # # Optimizer / control plane # ------------------------------------------------------------------ # - @app.post('/twinkle/clip_grad_norm', response_model=types.TaskEnvelope) + @app.post("/twinkle/clip_grad_norm", response_model=types.TaskEnvelope) async def clip_grad_norm( - request: Request, body: types.ClipGradNormRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.ClipGradNormRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): ret = await self.call_backend( self.model.clip_grad_norm, max_grad_norm=body.max_grad_norm, norm_type=body.norm_type, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) - return {'result': str(ret)} - - return await run_submit(self, request, body, task_type='clip_grad_norm', backend_call=_call) + **backend_kwargs(body), + ) + return {"result": str(ret)} - @app.post('/twinkle/step', response_model=types.TaskEnvelope) - async def step(request: Request, body: types.StepRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: + return await run_submit(self, request, body, task_type="clip_grad_norm", backend_call=_call) + @app.post("/twinkle/step", response_model=types.TaskEnvelope) + async def step( + request: Request, body: types.StepRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): await self.call_backend( - self.model.step, adapter_name=self.resolve_model_adapter_name(adapter_name), **backend_kwargs(body)) - - return await run_submit(self, request, body, task_type='step', backend_call=_call) + self.model.step, adapter_name=self.resolve_model_adapter_name(adapter_name), **backend_kwargs(body) + ) - @app.post('/twinkle/zero_grad', response_model=types.TaskEnvelope) - async def zero_grad(request: Request, body: types.AdapterRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: + return await run_submit(self, request, body, task_type="step", backend_call=_call) + @app.post("/twinkle/zero_grad", response_model=types.TaskEnvelope) + async def zero_grad( + request: Request, body: types.AdapterRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): await self.call_backend( - self.model.zero_grad, - adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) - - return await run_submit(self, request, body, task_type='zero_grad', backend_call=_call) + self.model.zero_grad, adapter_name=self.resolve_model_adapter_name(adapter_name), **backend_kwargs(body) + ) - @app.post('/twinkle/lr_step', response_model=types.TaskEnvelope) - async def lr_step(request: Request, body: types.LrStepRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: + return await run_submit(self, request, body, task_type="zero_grad", backend_call=_call) + @app.post("/twinkle/lr_step", response_model=types.TaskEnvelope) + async def lr_step( + request: Request, body: types.LrStepRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): await self.call_backend( - self.model.lr_step, adapter_name=self.resolve_model_adapter_name(adapter_name), **backend_kwargs(body)) + self.model.lr_step, adapter_name=self.resolve_model_adapter_name(adapter_name), **backend_kwargs(body) + ) - return await run_submit(self, request, body, task_type='lr_step', backend_call=_call) + return await run_submit(self, request, body, task_type="lr_step", backend_call=_call) - @app.post('/twinkle/clip_grad_and_step', response_model=types.TaskEnvelope) + @app.post("/twinkle/clip_grad_and_step", response_model=types.TaskEnvelope) async def clip_grad_and_step( - request: Request, body: types.ClipGradAndStepRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.ClipGradAndStepRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): await self.call_backend( self.model.clip_grad_and_step, max_grad_norm=body.max_grad_norm, norm_type=body.norm_type, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) + **backend_kwargs(body), + ) - return await run_submit(self, request, body, task_type='clip_grad_and_step', backend_call=_call) + return await run_submit(self, request, body, task_type="clip_grad_and_step", backend_call=_call) - @app.post('/twinkle/get_train_configs', response_model=types.TaskEnvelope) + @app.post("/twinkle/get_train_configs", response_model=types.TaskEnvelope) async def get_train_configs( - request: Request, body: types.AdapterRequest, self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.AdapterRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): ret = await self.call_backend( self.model.get_train_configs, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) - return {'result': ret} - - return await run_submit(self, request, body, task_type='get_train_configs', backend_call=_call) + **backend_kwargs(body), + ) + return {"result": ret} - @app.post('/twinkle/set_loss', response_model=types.TaskEnvelope) - async def set_loss(request: Request, body: types.SetLossRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: + return await run_submit(self, request, body, task_type="get_train_configs", backend_call=_call) + @app.post("/twinkle/set_loss", response_model=types.TaskEnvelope) + async def set_loss( + request: Request, body: types.SetLossRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): await self.call_backend( self.model.set_loss, body.loss_cls, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) + **backend_kwargs(body), + ) - return await run_submit(self, request, body, task_type='set_loss', backend_call=_call) + return await run_submit(self, request, body, task_type="set_loss", backend_call=_call) - @app.post('/twinkle/set_optimizer', response_model=types.TaskEnvelope) + @app.post("/twinkle/set_optimizer", response_model=types.TaskEnvelope) async def set_optimizer( - request: Request, body: types.SetOptimizerRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.SetOptimizerRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): await self.call_backend( self.model.set_optimizer, body.optimizer_cls, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) + **backend_kwargs(body), + ) - return await run_submit(self, request, body, task_type='set_optimizer', backend_call=_call) + return await run_submit(self, request, body, task_type="set_optimizer", backend_call=_call) - @app.post('/twinkle/set_lr_scheduler', response_model=types.TaskEnvelope) + @app.post("/twinkle/set_lr_scheduler", response_model=types.TaskEnvelope) async def set_lr_scheduler( - request: Request, body: types.SetLrSchedulerRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.SetLrSchedulerRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): await self.call_backend( self.model.set_lr_scheduler, body.scheduler_cls, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) + **backend_kwargs(body), + ) - return await run_submit(self, request, body, task_type='set_lr_scheduler', backend_call=_call) + return await run_submit(self, request, body, task_type="set_lr_scheduler", backend_call=_call) - @app.post('/twinkle/set_template', response_model=types.TaskEnvelope) + @app.post("/twinkle/set_template", response_model=types.TaskEnvelope) async def set_template( - request: Request, body: types.SetTemplateRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.SetTemplateRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): await self.call_backend( self.model.set_template, body.template_cls, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) + **backend_kwargs(body), + ) - return await run_submit(self, request, body, task_type='set_template', backend_call=_call) + return await run_submit(self, request, body, task_type="set_template", backend_call=_call) - @app.post('/twinkle/set_processor', response_model=types.TaskEnvelope) + @app.post("/twinkle/set_processor", response_model=types.TaskEnvelope) async def set_processor( - request: Request, body: types.SetProcessorRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.SetProcessorRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): await self.call_backend( self.model.set_processor, body.processor_cls, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) - - return await run_submit(self, request, body, task_type='set_processor', backend_call=_call) + **backend_kwargs(body), + ) - @app.post('/twinkle/add_metric', response_model=types.TaskEnvelope) - async def add_metric(request: Request, body: types.AddMetricRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: + return await run_submit(self, request, body, task_type="set_processor", backend_call=_call) + @app.post("/twinkle/add_metric", response_model=types.TaskEnvelope) + async def add_metric( + request: Request, body: types.AddMetricRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): metric_cls = deserialize_object(body.metric_cls) await self.call_backend( @@ -416,68 +448,70 @@ async def _call(self, body, adapter_name, token): metric_cls, is_training=body.is_training, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) + **backend_kwargs(body), + ) - return await run_submit(self, request, body, task_type='add_metric', backend_call=_call) + return await run_submit(self, request, body, task_type="add_metric", backend_call=_call) - @app.post('/twinkle/apply_patch', response_model=types.TaskEnvelope) + @app.post("/twinkle/apply_patch", response_model=types.TaskEnvelope) async def apply_patch( - request: Request, body: types.ApplyPatchRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.ApplyPatchRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): patch_cls = deserialize_object(body.patch_cls) await self.call_backend( self.model.apply_patch, patch_cls, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) + **backend_kwargs(body), + ) - return await run_submit(self, request, body, task_type='apply_patch', backend_call=_call) + return await run_submit(self, request, body, task_type="apply_patch", backend_call=_call) - @app.post('/twinkle/calculate_metric', response_model=types.TaskEnvelope) + @app.post("/twinkle/calculate_metric", response_model=types.TaskEnvelope) async def calculate_metric( - request: Request, body: types.CalculateMetricRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.CalculateMetricRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): ret = await self.call_backend( self.model.calculate_metric, is_training=body.is_training, adapter_name=self.resolve_model_adapter_name(adapter_name), - **backend_kwargs(body)) - return {'result': ret} + **backend_kwargs(body), + ) + return {"result": ret} - return await run_submit(self, request, body, task_type='calculate_metric', backend_call=_call) + return await run_submit(self, request, body, task_type="calculate_metric", backend_call=_call) # ------------------------------------------------------------------ # # Checkpoint I/O (need the caller token) # ------------------------------------------------------------------ # - @app.post('/twinkle/save', response_model=types.TaskEnvelope) - async def save(request: Request, body: types.SaveRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + @app.post("/twinkle/save", response_model=types.TaskEnvelope) + async def save( + request: Request, body: types.SaveRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + checkpoint_manager = create_checkpoint_manager(token, client_type="twinkle") save_weights = partial( self.call_backend, self.model.save, adapter_name=self.resolve_model_adapter_name(adapter_name), save_optimizer=body.save_optimizer, - **backend_kwargs(body)) + **backend_kwargs(body), + ) save_checkpoint = checkpoint_manager.save_sampler if body.is_sampler else checkpoint_manager.save_training twinkle_path, checkpoint_dir = await save_checkpoint(adapter_name, body.name, save_weights) - return {'twinkle_path': twinkle_path, 'checkpoint_dir': checkpoint_dir} - - return await run_submit(self, request, body, task_type='save', backend_call=_call) + return {"twinkle_path": twinkle_path, "checkpoint_dir": checkpoint_dir} - @app.post('/twinkle/load', response_model=types.TaskEnvelope) - async def load(request: Request, body: types.LoadRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: + return await run_submit(self, request, body, task_type="save", backend_call=_call) + @app.post("/twinkle/load", response_model=types.TaskEnvelope) + async def load( + request: Request, body: types.LoadRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + checkpoint_manager = create_checkpoint_manager(token, client_type="twinkle") resolved = checkpoint_manager.resolve_load_path(body.name) await self.call_backend( self.model.load, @@ -486,48 +520,43 @@ async def _call(self, body, adapter_name, token): adapter_name=self.resolve_model_adapter_name(adapter_name), load_optimizer=body.load_optimizer, token=token, - **backend_kwargs(body)) + **backend_kwargs(body), + ) - return await run_submit(self, request, body, task_type='load', backend_call=_call) + return await run_submit(self, request, body, task_type="load", backend_call=_call) - @app.post('/twinkle/resume_from_checkpoint', response_model=types.TaskEnvelope) + @app.post("/twinkle/resume_from_checkpoint", response_model=types.TaskEnvelope) async def resume_from_checkpoint( - request: Request, body: types.ResumeFromCheckpointRequest, - self: ModelManagement = Depends(self_fn)) -> types.TaskEnvelope: - + request: Request, body: types.ResumeFromCheckpointRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: async def _call(self, body, adapter_name, token): - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + checkpoint_manager = create_checkpoint_manager(token, client_type="twinkle") checkpoint_dir = checkpoint_manager.resolve_checkpoint_dir(body.name) ret = await self.call_backend( self.model.resume_from_checkpoint, checkpoint_dir, resume_only_model=body.resume_only_model, - adapter_name=self.resolve_model_adapter_name(adapter_name)) - return {'result': ret} + adapter_name=self.resolve_model_adapter_name(adapter_name), + ) + return {"result": ret} - return await run_submit(self, request, body, task_type='resume', backend_call=_call) + return await run_submit(self, request, body, task_type="resume", backend_call=_call) # ------------------------------------------------------------------ # # Adapter lifecycle (create / drop the adapter itself: no resource assert) # ------------------------------------------------------------------ # - @app.post('/twinkle/remove_adapter', response_model=types.TaskEnvelope) + @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) + request: Request, body: types.AdapterRequest, self: ModelManagement = Depends(self_fn) + ) -> types.TaskEnvelope: + raise RequestRejectedError("Training lifecycle operations must use the catalog endpoint") - @app.post('/twinkle/add_adapter_to_model', response_model=types.TaskEnvelope) + @app.post("/twinkle/add_adapter_to_model", response_model=types.TaskEnvelope) async def add_adapter_to_model( - request: Request, - body: types.AddAdapterRequest, - self: ModelManagement = Depends(self_fn), + request: Request, + 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 @@ -538,85 +567,26 @@ async def add_adapter_to_model( # 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) # ------------------------------------------------------------------ # - @app.post('/twinkle/upload_to_hub', response_model=types.TaskEnvelope) + @app.post("/twinkle/upload_to_hub", response_model=types.TaskEnvelope) async def upload_to_hub( - request: Request, - body: types.UploadToHubRequest, - self: ModelManagement = Depends(self_fn), + request: Request, + body: types.UploadToHubRequest, + self: ModelManagement = Depends(self_fn), ) -> types.TaskEnvelope: token = await self._on_request_start(request) async def _task(): - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') - if body.checkpoint_dir.startswith('twinkle://'): + checkpoint_manager = create_checkpoint_manager(token, client_type="twinkle") + if body.checkpoint_dir.startswith("twinkle://"): parsed = checkpoint_manager.parse_twinkle_path(body.checkpoint_dir) if not parsed: - raise ValueError(f'Invalid twinkle path format: {body.checkpoint_dir}') + raise ValueError(f"Invalid twinkle path format: {body.checkpoint_dir}") checkpoint_dir = checkpoint_manager.resolve_checkpoint_dir(body.checkpoint_dir) else: # A raw local path is still confined to this token's directory: without the check an @@ -625,8 +595,9 @@ async def _task(): from twinkle.server.checkpoint.paths import ensure_inside - checkpoint_dir = str(ensure_inside( - checkpoint_manager.training_run_manager.get_base_dir(), Path(body.checkpoint_dir))) + checkpoint_dir = str( + ensure_inside(checkpoint_manager.training_run_manager.get_base_dir(), Path(body.checkpoint_dir)) + ) await self.call_backend( self.model.upload_to_hub, checkpoint_dir=checkpoint_dir, @@ -635,4 +606,4 @@ async def _task(): async_upload=False, ) - return await self.submit_background_and_peek(_task, task_type='upload_to_hub') + return await self.submit_background_and_peek(_task, task_type="upload_to_hub") 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..4134d3c5f --- /dev/null +++ b/src/twinkle/server/resources/runtime.py @@ -0,0 +1,128 @@ +# 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): + import ray + 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) + # Both features give workers the controller identity, not a Serve replica ID. + actors = [] + 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 + try: + # Ray may import the worker class when resolving its handle. + # Keep imports and GCS calls outside the controller event loop + # so its shutdown deadline and other instances remain live. + handle = await asyncio.to_thread(ray.get_actor, row["name"], namespace=self.namespace) + except ValueError: + continue + actors.append(handle) + await asyncio.to_thread(ray.kill, handle, no_restart=True) + await asyncio.to_thread(serve.delete, name) + for handle in actors: + try: + await asyncio.wait_for(handle.__ray_ready__.remote(), 5) + except ray.exceptions.ActorDiedError: + continue + except ray.exceptions.RayActorError as exc: + raise ResourceReleaseError("Owned worker death could not be confirmed", phase="worker_death") from exc + raise ResourceReleaseError("Owned worker death could not be confirmed", phase="worker_death") + + 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..89275e8da 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -5,23 +5,29 @@ Builds a single Ray Serve deployment (SamplerManagement) that simultaneously handles both Tinker (/tinker/asample) and Twinkle (/twinkle/*) sampler endpoints. """ + from __future__ import annotations +import asyncio +import os +from typing import Any 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 @@ -69,14 +75,14 @@ def _make_torch_sampler(kw: dict[str, Any]) -> Any: # Single validate-then-dispatch selector for the sampler backend. SAMPLER_SELECTOR = BackendSelector( - 'sampler_type', + "sampler_type", { - 'mock': _make_mock_sampler, - 'vllm': _make_vllm_sampler, - 'vllm_async': _make_vllm_async_sampler, - 'sglang': _make_sglang_sampler, - 'sglang_async': _make_sglang_async_sampler, - 'torch': _make_torch_sampler, + "mock": _make_mock_sampler, + "vllm": _make_vllm_sampler, + "vllm_async": _make_vllm_async_sampler, + "sglang": _make_sglang_sampler, + "sglang_async": _make_sglang_async_sampler, + "torch": _make_torch_sampler, }, ) @@ -101,55 +107,71 @@ class SamplerManagement(LazyCleanupMixin, TaskQueueMixin): - Template configuration for trajectory encoding """ - async def __init__(self, - model_id: str, - nproc_per_node: int, - device_group: dict[str, Any], - device_mesh: dict[str, Any], - sampler_type: str, - engine_args: dict[str, Any] | None = None, - queue_config: TaskQueueConfig | None = None, - data_plane_url: str | None = None, - template: dict[str, Any] | None = None, - **kwargs): + async def __init__( + self, + model_id: str, + nproc_per_node: int, + device_group: dict[str, Any], + device_mesh: dict[str, Any], + sampler_type: str, + engine_args: dict[str, Any] | None = None, + queue_config: TaskQueueConfig | None = None, + 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'), + is_mock=(sampler_type == "mock"), nproc_per_node=nproc_per_node, device_group=self.device_group, device_mesh_dict=device_mesh, ) 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, - 'remote_group': self.device_group.name, - 'instance_id': self.replica_id, + "model_id": model_id, + "remote_group": self.device_group.name, + "instance_id": self.replica_id, } - if sampler_type != 'mock': + if sampler_type != "mock": sampler_kwargs.update( engine_args=engine_args or {}, device_mesh=self.device_mesh, - **{ - k: v - for k, v in kwargs.items() if k not in ('engine_args', ) - }, + **{k: v for k, v in kwargs.items() if k not in ("engine_args",)}, ) 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 + self.data_plane = DataPlaneProxy(data_plane_url) - actors = getattr(self.sampler, '_actors', None) + actors = getattr(self.sampler, "_actors", None) self._init_task_queue( queue_config, - deployment_name='Sampler', + deployment_name="Sampler", collect_width=len(actors) if actors else 1, ) # ComputeWorker bounds the client future. Keep collecting the actual @@ -163,22 +185,23 @@ async def __init__(self, # (see the handler), which keeps that crude fallback from clobbering an explicitly chosen one. # Runs after _init_task_queue: call_backend needs the BackendGate it installs. self._template_pinned = False - if template and hasattr(self.sampler, 'set_template'): + if template and hasattr(self.sampler, "set_template"): template_kwargs = dict(template) - template_cls = template_kwargs.pop('template_cls', None) + template_cls = template_kwargs.pop("template_cls", None) if template_cls is None: raise ValueError( "template must carry a 'template_cls' key naming the template class; got keys: " - f'{sorted(template_kwargs) or "none"}') + f"{sorted(template_kwargs) or 'none'}" + ) await self.call_backend(self.sampler.set_template, template_cls, **template_kwargs) self._template_pinned = True - if sampler_type != 'mock' and not self._template_pinned: + if sampler_type != "mock" and not self._template_pinned: await self.call_backend(self.sampler.set_template, get_template_for_model(model_id), model_id=model_id) async def shutdown(self) -> None: self.sampling_guard.close() try: - cancel_all = getattr(self.sampler, 'abort_all_inflight', None) + cancel_all = getattr(self.sampler, "abort_all_inflight", None) if callable(cancel_all): await self.call_backend(cancel_all, admit=False) finally: @@ -187,37 +210,46 @@ 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) - @serve.multiplexed(max_num_models_per_replica=5) - async def _sticky_entry(self, sticky_key: str): - return sticky_key + async def pending_futures(self): + return await self.request_runtime.pending_futures() - async def _ensure_sticky(self): - sticky_key = serve.get_multiplexed_model_id() - await self._sticky_entry(sticky_key) + async def runtime_identity(self): + return await self.request_runtime.runtime_identity() + + 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) -def build_sampler_app(model_id: str, - nproc_per_node: int, - device_group: dict[str, Any], - device_mesh: dict[str, Any], - deploy_options: dict[str, Any], - sampler_type: str, - engine_args: dict[str, Any] | None = None, - queue_config: TaskQueueConfig | None = None, - data_plane_url: str | None = None, - template: dict[str, Any] | None = None, - **kwargs): +def build_sampler_app( + model_id: str, + nproc_per_node: int, + device_group: dict[str, Any], + device_mesh: dict[str, Any], + deploy_options: dict[str, Any], + sampler_type: str, + engine_args: dict[str, Any] | None = None, + queue_config: TaskQueueConfig | None = None, + data_plane_url: str | None = None, + template: dict[str, Any] | None = None, + **kwargs, +): """Build a unified sampler application for text generation inference. Supports both Tinker (polling-style /tinker/asample) and @@ -250,17 +282,28 @@ 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', + "Sampler", register_routes, fastapi_kwargs={ - 'title': 'Unified Sampler', - 'description': 'REST API for distributed text generation inference (Tinker + Twinkle)', - 'version': '1.0.0', + "title": "Unified Sampler", + "description": "REST API for distributed text generation inference (Tinker + Twinkle)", + "version": "1.0.0", }, attach_replica_id_header=True, on_shutdown=lambda servable: servable.shutdown(), @@ -270,12 +313,20 @@ def register_routes(app: FastAPI, get_self: Any) -> None: app, SamplerManagement, deploy_options, - deployment_name='SamplerManagement', - bind_args=(model_id, nproc_per_node, device_group, device_mesh, sampler_type, engine_args, queue_config, - data_plane_url), + deployment_name="SamplerManagement", + bind_args=( + model_id, + nproc_per_node, + device_group, + device_mesh, + sampler_type, + engine_args, + queue_config, + data_plane_url, + ), # ``template`` is an explicit named parameter, so it is not in ``kwargs``; forward it explicitly # or the pinned-template branch in ``SamplerManagement.__init__`` never runs. - bind_kwargs={**kwargs, 'template': template}, + bind_kwargs={**kwargs, "template": template}, ) 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..03eaea353 100644 --- a/src/twinkle/server/sampler/backends/streaming.py +++ b/src/twinkle/server/sampler/backends/streaming.py @@ -1,26 +1,45 @@ # 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__' +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. +def stream_to_queue(sampler, queue, inputs, sampling_params=None, adapter_name="", adapter_path=None): + """Drain and close the generator even when its HTTP consumer disconnects.""" + from contextlib import ExitStack - 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) + 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..6f3d1f9a9 100644 --- a/src/twinkle/server/sampler/concurrency.py +++ b/src/twinkle/server/sampler/concurrency.py @@ -1,17 +1,17 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Coordinate sampling with changes to shared sampler state.""" -from __future__ import annotations +from __future__ import annotations 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): @@ -25,9 +25,9 @@ def __init__(self): def check_available(self): if self._closed: - raise BackendBusyError('The sampler is shutting down.') + raise BackendBusyError("The sampler is shutting down.") if self._weights_error is not None: - raise BackendBusyError('A checkpoint write failed; the sampler is restarting.') + raise BackendBusyError("A checkpoint write failed; the sampler is restarting.") def invalidate_weights(self, error: Exception): self._weights_error = error @@ -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..37605a3b1 100644 --- a/src/twinkle/server/sampler/tinker_handlers.py +++ b/src/twinkle/server/sampler/tinker_handlers.py @@ -6,23 +6,25 @@ inheritance relationship with the deployment class. Provides POST /tinker/asample using schedule_concurrent_task() returning UntypedAPIFuture. """ -from __future__ import annotations +from __future__ import annotations import os import traceback from collections.abc import Callable +from typing import TYPE_CHECKING + from fastapi import Depends, FastAPI, Request from tinker import types -from typing import TYPE_CHECKING if TYPE_CHECKING: from .app import SamplerManagement 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() @@ -50,9 +52,10 @@ def _register_tinker_sampler_routes(app: FastAPI, self_fn: Callable[[], SamplerM It is wired in via Depends so it is resolved lazily at request time. """ - @app.post('/tinker/asample') - async def asample(request: Request, body: types.SampleRequest, - self: SamplerManagement = Depends(self_fn)) -> types.UntypedAPIFuture: + @app.post("/tinker/asample") + async def asample( + request: Request, body: types.SampleRequest, self: SamplerManagement = Depends(self_fn) + ) -> types.UntypedAPIFuture: """Execute text generation (inference) for Tinker clients. Args: @@ -67,33 +70,35 @@ async def asample(request: Request, body: types.SampleRequest, async def _do_sample(): try: # Extract prompt token IDs from ModelInput - prompt_inputs = {'input_ids': body.prompt.to_ints()} + prompt_inputs = {"input_ids": body.prompt.to_ints()} # Get model_path from body or sampling session model_path = body.model_path if not model_path and body.sampling_session_id: session = await self.state.get_sampling_session(body.sampling_session_id) if session: - model_path = session.get('model_path') + model_path = session.get("model_path") # Parse and resolve adapter URI from model_path adapter_uri = None if model_path: - checkpoint_manager = create_checkpoint_manager(token, client_type='tinker') + checkpoint_manager = create_checkpoint_manager(token, client_type="tinker") adapter_name, adapter_uri = checkpoint_manager.parse_adapter_uri(model_path) # 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 if body.sampling_params: sampling_params = SamplingParams( max_tokens=body.sampling_params.max_tokens - if body.sampling_params.max_tokens is not None else 256, + if body.sampling_params.max_tokens is not None + else 256, temperature=body.sampling_params.temperature - if body.sampling_params.temperature is not None else 1.0, + if body.sampling_params.temperature is not None + else 1.0, top_p=body.sampling_params.top_p, top_k=body.sampling_params.top_k, stop=body.sampling_params.stop, @@ -129,14 +134,16 @@ async def generate(lora_path): logprobs = flattened else: raise RuntimeError( - f'Sampler returned {len(flattened)} logprobs for {len(seq.tokens)} generated ' - 'tokens; refusing to return a misaligned Tinker SampledSequence.') + f"Sampler returned {len(flattened)} logprobs for {len(seq.tokens)} generated " + "tokens; refusing to return a misaligned Tinker SampledSequence." + ) tinker_sequences.append( _sampled_sequence( stop_reason=seq.stop_reason, tokens=list(seq.tokens), logprobs=logprobs, - )) + ) + ) return _sample_response( sequences=tinker_sequences, prompt_logprobs=responses[0].prompt_logprobs, @@ -152,5 +159,6 @@ async def generate(lora_path): _do_sample, token=token, input_tokens=input_tokens, - task_type='sample', + 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..7aee52fed 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -6,15 +6,16 @@ inheritance relationship with the deployment class. Provides /twinkle/* sampler endpoints. """ -from __future__ import annotations +from __future__ import annotations import asyncio -import json import uuid from collections.abc import Callable, Mapping +from typing import TYPE_CHECKING + +import json from fastapi import Depends, FastAPI, Request from fastapi.responses import StreamingResponse -from typing import TYPE_CHECKING if TYPE_CHECKING: from .app import SamplerManagement @@ -33,6 +34,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() @@ -65,7 +67,7 @@ def _json_safe_value(value): return [safe for safe in (_json_safe_value(item) for item in value) if safe is not _DROP] if isinstance(value, _JSON_PRIMITIVES): return value - tolist = getattr(value, 'tolist', None) + tolist = getattr(value, "tolist", None) if callable(tolist): try: return value.tolist() @@ -90,33 +92,36 @@ def _build_rollout_rows_and_tags( """Flatten sampler output to one TQ row per generated sequence.""" resolved_group_ids = group_ids or [uuid.uuid4().hex for _ in sample_models] if len(resolved_group_ids) != len(sample_models): - raise ValueError(f'group_ids contains {len(resolved_group_ids)} values for ' - f'{len(sample_models)} sampler inputs') + raise ValueError(f"group_ids contains {len(resolved_group_ids)} values for {len(sample_models)} sampler inputs") rows = [] tags = [] - for prompt_index, (response, group_id) in enumerate(zip(sample_models, resolved_group_ids)): + for prompt_index, (response, group_id) in enumerate(zip(sample_models, resolved_group_ids, strict=False)): for generation_idx, sequence in enumerate(response.sequences): sampled_logprobs = [ 0.0 if not position else float(position[0][1]) for position in (sequence.logprobs or []) ] - rows.append({ - 'train_input': sequence.new_input_feature, - 'sampled_logprobs': sampled_logprobs, - 'tokens': sequence.tokens, - 'decoded': sequence.decoded, - 'stop_reason': sequence.stop_reason, - 'prompt_logprobs': response.prompt_logprobs, - 'topk_prompt_logprobs': response.topk_prompt_logprobs, - }) - tags.append({ - 'record_type': 'sample', - 'group_id': group_id, - 'prompt_index': prompt_index, - 'generation_idx': generation_idx, - 'rollout_status': 'ROLLOUT_DONE', - 'rollout_policy_version': policy_version, - 'rollout_adapter_uri': adapter_uri, - }) + rows.append( + { + "train_input": sequence.new_input_feature, + "sampled_logprobs": sampled_logprobs, + "tokens": sequence.tokens, + "decoded": sequence.decoded, + "stop_reason": sequence.stop_reason, + "prompt_logprobs": response.prompt_logprobs, + "topk_prompt_logprobs": response.topk_prompt_logprobs, + } + ) + tags.append( + { + "record_type": "sample", + "group_id": group_id, + "prompt_index": prompt_index, + "generation_idx": generation_idx, + "rollout_status": "ROLLOUT_DONE", + "rollout_policy_version": policy_version, + "rollout_adapter_uri": adapter_uri, + } + ) return rows, tags @@ -130,9 +135,13 @@ def _to_sample_response_models(responses) -> list[types.SampleResponseModel]: tokens=list(sequence.tokens), logprobs=list(sequence.logprobs) if sequence.logprobs is not None else None, decoded=sequence.decoded, - new_input_feature=(_serialize_input_feature(sequence.new_input_feature) - if sequence.new_input_feature is not None else None), - ) for sequence in response.sequences + new_input_feature=( + _serialize_input_feature(sequence.new_input_feature) + if sequence.new_input_feature is not None + else None + ), + ) + for sequence in response.sequences ] sample_models.append( types.SampleResponseModel( @@ -144,7 +153,8 @@ def _to_sample_response_models(responses) -> list[types.SampleResponseModel]: prompt_token_ids=(list(response.prompt_token_ids) if response.prompt_token_ids is not None else None), prompt_logprobs=response.prompt_logprobs, topk_prompt_logprobs=response.topk_prompt_logprobs, - )) + ) + ) return sample_models @@ -156,25 +166,28 @@ async def _stream_queue(q, sentinel, request_id: str, total_timeout: float, sing remaining = total_timeout - (loop.time() - start) if remaining <= 0: payload = task_error_payload( - 'sample_stream exceeded the execution time bound', request_id=request_id, error_code=504) - yield json.dumps(payload) + '\n' + "sample_stream exceeded the execution time bound", request_id=request_id, error_code=504 + ) + yield json.dumps(payload) + "\n" break try: item = await asyncio.wait_for( - loop.run_in_executor(None, q.get), timeout=min(single_get_timeout, remaining)) - except asyncio.TimeoutError: + loop.run_in_executor(None, q.get), timeout=min(single_get_timeout, remaining) + ) + except TimeoutError: payload = task_error_payload( - 'sample_stream timed out waiting for the next token', request_id=request_id, error_code=504) - yield json.dumps(payload) + '\n' + "sample_stream timed out waiting for the next token", request_id=request_id, error_code=504 + ) + yield json.dumps(payload) + "\n" break if item == sentinel: break if isinstance(item, Exception): - payload = task_error_payload(f'{type(item).__name__}: {item}', request_id=request_id, error_code=500) - yield json.dumps(payload) + '\n' + payload = task_error_payload(f"{type(item).__name__}: {item}", request_id=request_id, error_code=500) + yield json.dumps(payload) + "\n" break delta, reason = item - yield json.dumps({'delta': delta, 'finish_reason': reason}) + '\n' + yield json.dumps({"delta": delta, "finish_reason": reason}) + "\n" finally: try: q.shutdown(force=True) @@ -192,20 +205,21 @@ async def _await_futures(service: SamplerManagement, futures: List[Any], request """ loop = asyncio.get_running_loop() try: - results = await asyncio.gather( - *[loop.run_in_executor(None, future.result) for future in futures]) + results = await asyncio.gather(*[loop.run_in_executor(None, future.result) for future in futures]) return [response for result in results for response in result] except BaseException: + service.sampling_guard.invalidate_weights(RuntimeError("Generation completion was not confirmed")) for request_id in request_ids: try: - await service.call_backend(service.sampler.stop, {'request_id': request_id}) + await service.call_backend(service.sampler.stop, {"request_id": request_id}) except Exception: - logger.warning('Failed to stop generation %s', request_id, exc_info=True) + logger.warning("Failed to stop generation %s", request_id, exc_info=True) raise -def _submit_group(service: SamplerManagement, inputs: List[Any], params: Any, *, - adapter_name: str, adapter_path: Optional[str]) -> Tuple[List[Any], List[str], List[str]]: +def _submit_group( + service: SamplerManagement, inputs: List[Any], params: Any, *, adapter_name: str, adapter_path: Optional[str] +) -> Tuple[List[Any], List[str], List[str]]: """Submit one generation per input; returns (futures, request_ids, trajectories). Synchronous on purpose: the caller wraps it in ``call_backend`` (the backend gate's executor), @@ -218,10 +232,9 @@ def _submit_group(service: SamplerManagement, inputs: List[Any], params: Any, *, trajectories: List[Dict[str, Any]] = [] for entry in inputs: trajectory = dict(entry) - future = service.sampler.submit(trajectory, params, adapter_name=adapter_name, - adapter_path=adapter_path) + future = service.sampler.submit(trajectory, params, adapter_name=adapter_name, adapter_path=adapter_path) futures.append(future) - request_ids.append(trajectory['request_id']) + request_ids.append(trajectory["request_id"]) trajectories.append(trajectory) return futures, request_ids, trajectories @@ -233,15 +246,17 @@ def _register_twinkle_sampler_routes(app: FastAPI, self_fn: Callable[[], Sampler It is wired in via Depends so it is resolved lazily at request time. """ - @app.post('/twinkle/create', response_model=sampler_types.SamplerCreateResponse) + @app.post("/twinkle/create", response_model=sampler_types.SamplerCreateResponse) async def create( - request: Request, self: SamplerManagement = Depends(self_fn)) -> sampler_types.SamplerCreateResponse: + request: Request, self: SamplerManagement = Depends(self_fn) + ) -> sampler_types.SamplerCreateResponse: """Health check / session creation endpoint.""" return sampler_types.SamplerCreateResponse() - @app.post('/twinkle/sample', response_model=types.TaskEnvelope) - async def sample(request: Request, body: types.SampleRequest, - self: SamplerManagement = Depends(self_fn)) -> types.TaskEnvelope: + @app.post("/twinkle/sample", response_model=types.TaskEnvelope) + async def sample( + request: Request, body: types.SampleRequest, self: SamplerManagement = Depends(self_fn) + ) -> types.TaskEnvelope: """Sample completions from the model. Supports Trajectory or InputFeature inputs, with optional LoRA adapter. @@ -251,12 +266,13 @@ async def sample(request: Request, body: types.SampleRequest, async def _task(): # Resolve adapter resolved_uri = None - adapter_name = body.adapter_name or '' - full_adapter_name = resolve_twinkle_adapter_name(request, adapter_name) or '' + adapter_name = body.adapter_name or "" + full_adapter_name = resolve_twinkle_adapter_name(request, adapter_name) or "" if body.adapter_uri: from twinkle.server.checkpoint import create_checkpoint_manager - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + + checkpoint_manager = create_checkpoint_manager(token, client_type="twinkle") _, resolved_uri = checkpoint_manager.parse_adapter_uri(body.adapter_uri) # Parse inputs (shared seam; batch form) @@ -282,35 +298,42 @@ 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') + 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", + execution_scope=request_execution_scope(request), + ) - @app.post('/twinkle/sample_to_data_plane', response_model=types.TaskEnvelope) + @app.post("/twinkle/sample_to_data_plane", response_model=types.TaskEnvelope) async def sample_to_data_plane( - request: Request, - body: types.DataPlaneSampleRequest, - self: SamplerManagement = Depends(self_fn), + request: Request, + body: types.DataPlaneSampleRequest, + self: SamplerManagement = Depends(self_fn), ) -> types.TaskEnvelope: """Generate a complete group, store it server-side, and return a Task_Envelope whose result is the stored group's DataRef.""" token = await self._on_request_start(request) if not self.data_plane.enabled: - raise EndpointUnavailableError('sample_to_data_plane requires data_plane_url') - if not callable(getattr(self.sampler, 'submit', None)): - raise EndpointUnavailableError('sampler_type must be vllm_async or sglang_async') + raise EndpointUnavailableError("sample_to_data_plane requires data_plane_url") + if not callable(getattr(self.sampler, "submit", None)): + raise EndpointUnavailableError("sampler_type must be vllm_async or sglang_async") adapter_path = None - full_adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) or '' + full_adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) or "" if body.adapter_uri: from twinkle.server.checkpoint import create_checkpoint_manager - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + + checkpoint_manager = create_checkpoint_manager(token, client_type="twinkle") _, adapter_path = checkpoint_manager.parse_adapter_uri(body.adapter_uri) - inputs = (await self.data_plane.get(body.input_ref) if body.input_ref is not None else body.inputs) + inputs = await self.data_plane.get(body.input_ref) if body.input_ref is not None else body.inputs inputs = to_backend_inputs(inputs) params_dict = dict(body.sampling_params or {}) - params_dict['num_samples'] = body.num_samples + params_dict["num_samples"] = body.num_samples params = SamplingParams.from_dict(params_dict) async def _generate_and_store(): @@ -318,30 +341,34 @@ async def _generate_and_store(): # sequence runs as one background future (outside the serial compute queue) and its # result is the DataRef. One submit per input (num_samples rides SamplingParams), # each handing back a LazyCollect the handler awaits locally -- no per-status RPC. - async def generate(): + async def generate(resolved_path): futures, request_ids, _ = await self.call_backend( - _submit_group, self, inputs, params, - adapter_name=full_adapter_name, adapter_path=adapter_path) + _submit_group, self, inputs, params, adapter_name=full_adapter_name, adapter_path=resolved_path + ) return await _await_futures(self, futures, request_ids) - responses = await self.sampling_guard.run(generate) + responses = await sample_with_weights(self, adapter_path, generate) rows, tags = _build_rollout_rows_and_tags( _to_sample_response_models(responses), group_ids=body.group_ids, policy_version=body.policy_version, adapter_uri=body.adapter_uri, ) - ref = await self.data_plane.put([json_safe(item) for item in rows], kind='rollout', tags=tags) + ref = await self.data_plane.put([json_safe(item) for item in rows], kind="rollout", tags=tags) 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) + @app.post("/twinkle/encode", response_model=types.TaskEnvelope) async def encode( - request: Request, - body: types.EncodeRequest, - self: SamplerManagement = Depends(self_fn), + request: Request, + body: types.EncodeRequest, + self: SamplerManagement = Depends(self_fn), ) -> types.TaskEnvelope: """Run a pooling forward (embedding / classify) and return the pooled output per input. @@ -356,38 +383,42 @@ async def encode( 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) - return {'data': [list(response.data) for response in responses]} + adapter_name = resolve_twinkle_adapter_name(request, body.adapter_name) or "" + 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') + @app.post("/twinkle/unload_adapter_paths") async def unload_adapter_paths( - request: Request, - body: types.UnloadAdapterPathsRequest, - self: SamplerManagement = Depends(self_fn), + request: Request, + body: types.UnloadAdapterPathsRequest, + self: SamplerManagement = Depends(self_fn), ) -> dict[str, str]: """Best-effort eviction of published LoRA snapshots from sampler caches.""" token = await self._on_request_start(request) resolved_paths = [] for adapter_path in body.adapter_paths: - if adapter_path.startswith('twinkle://'): + if adapter_path.startswith("twinkle://"): from twinkle.server.checkpoint import create_checkpoint_manager - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + + checkpoint_manager = create_checkpoint_manager(token, client_type="twinkle") _, adapter_path = checkpoint_manager.parse_adapter_uri(adapter_path) resolved_paths.append(adapter_path) - unload = getattr(self.sampler, 'unload_adapter_paths', None) + unload = getattr(self.sampler, "unload_adapter_paths", None) if unload is not None: await self.sampling_guard.run(lambda: self.call_backend(unload, resolved_paths), exclusive=True) - return {'status': 'ok'} + return {"status": "ok"} - @app.post('/twinkle/set_template', response_model=sampler_types.SamplerSetTemplateResponse) + @app.post("/twinkle/set_template", response_model=sampler_types.SamplerSetTemplateResponse) async def set_template( - request: Request, - body: sampler_types.SamplerSetTemplateRequest, - self: SamplerManagement = Depends(self_fn), + request: Request, + body: sampler_types.SamplerSetTemplateRequest, + self: SamplerManagement = Depends(self_fn), ) -> sampler_types.SamplerSetTemplateResponse: """Set the chat template for encoding Trajectory inputs. @@ -397,7 +428,7 @@ async def set_template( unaffected either way. """ if self._template_pinned: - logger.debug('set_template ignored: this sampler was configured with a pinned template.') + logger.debug("set_template ignored: this sampler was configured with a pinned template.") return sampler_types.SamplerSetTemplateResponse() # Multi-tenant safety: this sampler replica is shared, and set_template mutates a single # self.template (it is not keyed per adapter), so one client's call would otherwise change the @@ -406,52 +437,59 @@ async def set_template( # only reachable for a deployment that deliberately did NOT pin one. The supported model is one # fixed template per sampler, set once at construction; a per-adapter/per-LoRA template would need # set_template to key by adapter, which is future work. - with traced_operation('sampler.set_template'): + with traced_operation("sampler.set_template"): await self.sampling_guard.run( lambda: self.call_backend(self.sampler.set_template, body.template_cls, **backend_kwargs(body)), - exclusive=True) + exclusive=True, + ) return sampler_types.SamplerSetTemplateResponse() - @app.post('/twinkle/add_adapter_to_sampler', response_model=sampler_types.SamplerAddAdapterResponse) + @app.post("/twinkle/add_adapter_to_sampler", response_model=sampler_types.SamplerAddAdapterResponse) async def add_adapter_to_sampler( - request: Request, - body: sampler_types.SamplerAddAdapterRequest, - self: SamplerManagement = Depends(self_fn), + request: Request, + body: sampler_types.SamplerAddAdapterRequest, + self: SamplerManagement = Depends(self_fn), ) -> sampler_types.SamplerAddAdapterResponse: """Add a LoRA adapter to the sampler.""" # Raised, not asserted: decidable from the request body alone, so it owes the caller # a real 400 rather than an AssertionError surfacing as a 500 -- and a bare assert # would vanish under `python -O`, letting an empty adapter_name reach the backend. if not body.adapter_name: - raise RequestRejectedError('`adapter_name` is required and must be non-empty.') + 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 - with traced_operation('sampler.add_adapter_to_sampler', attrs={MODEL_ID: self.model_id}): + 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( lambda: self.call_backend(self.sampler.add_adapter_to_sampler, full_adapter_name, config), - exclusive=True) + exclusive=True, + ) return sampler_types.SamplerAddAdapterResponse(adapter_name=full_adapter_name) - @app.post('/twinkle/apply_patch') + @app.post("/twinkle/apply_patch") async def apply_patch( - request: Request, - body: types.ApplyPatchRequest, - self: SamplerManagement = Depends(self_fn), + request: Request, + body: types.ApplyPatchRequest, + self: SamplerManagement = Depends(self_fn), ) -> None: patch_cls = deserialize_object(body.patch_cls) - with traced_operation('sampler.apply_patch'): + with traced_operation("sampler.apply_patch"): await self.sampling_guard.run( - lambda: self.call_backend(self.sampler.apply_patch, patch_cls, **backend_kwargs(body)), exclusive=True) + lambda: self.call_backend(self.sampler.apply_patch, patch_cls, **backend_kwargs(body)), exclusive=True + ) - @app.post('/twinkle/sample_stream') + @app.post("/twinkle/sample_stream") async def sample_stream( - request: Request, - body: types.SampleRequest, - self: SamplerManagement = Depends(self_fn), + request: Request, + body: types.SampleRequest, + self: SamplerManagement = Depends(self_fn), ): """Stream token deltas as newline-delimited JSON. @@ -463,16 +501,15 @@ async def sample_stream( """ token = await self._on_request_start(request) - adapter_path = None - adapter_name = body.adapter_name or '' - full_adapter_name = resolve_twinkle_adapter_name(request, adapter_name) or '' + resolved_uri = None + adapter_name = body.adapter_name or "" + full_adapter_name = resolve_twinkle_adapter_name(request, adapter_name) or "" if body.adapter_uri: from twinkle.server.checkpoint import create_checkpoint_manager - checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + + 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 @@ -480,7 +517,7 @@ async def sample_stream( try: inputs_parsed = to_backend_inputs(body.inputs, single=True) except ValueError as e: - raise RequestRejectedError(str(e)) + raise RequestRejectedError(str(e)) from e params = None if body.sampling_params: @@ -488,27 +525,44 @@ 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}' + request_id = f"req_{uuid.uuid4().hex}" actors = self.sampler._actors if not actors: async def _no_actor_generator(): - payload = task_error_payload('No available sampler actor', request_id=request_id, error_code=503) - yield json.dumps(payload) + '\n' + payload = task_error_payload("No available sampler actor", request_id=request_id, error_code=503) + yield json.dumps(payload) + "\n" - return StreamingResponse(_no_actor_generator(), media_type='application/x-ndjson') + 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, @@ -516,5 +570,5 @@ async def _no_actor_generator(): request_id, self.task_queue_config.effective_execution_timeout, ), - media_type='application/x-ndjson', + media_type="application/x-ndjson", ) diff --git a/src/twinkle/server/sampler/weights.py b/src/twinkle/server/sampler/weights.py index f2d1077f7..9e3fa4930 100644 --- a/src/twinkle/server/sampler/weights.py +++ b/src/twinkle/server/sampler/weights.py @@ -1,25 +1,29 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Run checkpoint loading and generation as one owned sampling job.""" -from __future__ import annotations +from __future__ import annotations import glob import os + 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: if not resolved_uri: return False - if os.path.isfile(os.path.join(resolved_uri, 'adapter_config.json')): + 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}') - 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"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 RequestRejectedError(f"No .safetensors weights found under {resolved_uri}") return True @@ -29,7 +33,8 @@ async def resolve_sampler_weights(service, resolved_uri: str | None, *, full_che full_checkpoint = is_full_checkpoint(resolved_uri) if full_checkpoint: return await service.sampling_guard.run( - lambda: resolve_sampler_weights(service, resolved_uri, full_checkpoint=True), exclusive=True) + lambda: resolve_sampler_weights(service, resolved_uri, full_checkpoint=True), exclusive=True + ) if not full_checkpoint: return resolved_uri or None try: @@ -38,7 +43,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..6395492b5 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..9bb035a28 100644 --- a/src/twinkle/server/session_resource/processor.py +++ b/src/twinkle/server/session_resource/processor.py @@ -2,11 +2,11 @@ """ 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 +from __future__ import annotations from abc import abstractmethod from twinkle.utils.logger import get_logger @@ -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__ @@ -30,7 +30,7 @@ class ProcessorManagerMixin(SessionResourceMixin): """ # Set resource type for logging - _resource_type = 'Processor' + _resource_type = "Processor" def _init_processor_manager( self, @@ -61,7 +61,7 @@ async def _on_resource_liveness_confirmed(self, resource_id: str) -> bool: return False try: renewed = await self.state.renew_processor_quota( - info['token'], + info["token"], resource_id, lease_seconds=self._processor_quota_lease_seconds, ) @@ -69,10 +69,10 @@ async def _on_resource_liveness_confirmed(self, resource_id: str) -> bool: # Keep the local processor during a transient backend outage. Once the # backend recovers, a lost/expired reservation returns False and the # countdown loop removes the unaccounted local resource. - logger.warning('[ProcessorManager] Failed to renew quota lease for %s: %r', resource_id, exc) + logger.warning("[ProcessorManager] Failed to renew quota lease for %s: %r", resource_id, exc) return True if not renewed: - logger.warning('[ProcessorManager] Quota lease for %s was lost; expiring local processor', resource_id) + logger.warning("[ProcessorManager] Quota lease for %s was lost; expiring local processor", resource_id) return renewed async def _on_resource_expired(self, resource_id: str) -> None: @@ -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..70b441b73 100644 --- a/src/twinkle/server/state/__init__.py +++ b/src/twinkle/server/state/__init__.py @@ -1,35 +1,46 @@ # 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', - 'SamplingSessionRecord', - 'FutureRecord', + "SessionRecord", + "ModelRecord", + "SamplingSessionRecord", + "FutureRecord", # Base - 'BaseManager', + "BaseManager", # Resource managers - 'SessionManager', - 'ModelManager', - 'SamplingSessionManager', - 'FutureManager', + "SessionManager", + "ModelManager", + "SamplingSessionManager", + "FutureManager", # Server state - 'ServerState', - 'ReplicaRegistry', - 'ResourceCleanupCoordinator', - 'ResourceCountPublisher', - 'get_server_state', + "ServerState", + "ReplicaRegistry", + "ResourceCleanupCoordinator", + "ResourceCountPublisher", + "get_server_state", # Persistence backend factory - 'create_backend', + "create_backend", ] diff --git a/src/twinkle/server/state/backend/base.py b/src/twinkle/server/state/backend/base.py index b65f7d2b2..ecaa32616 100644 --- a/src/twinkle/server/state/backend/base.py +++ b/src/twinkle/server/state/backend/base.py @@ -1,5 +1,4 @@ from __future__ import annotations - from abc import ABC, abstractmethod from collections.abc import Callable from typing import Any @@ -74,6 +73,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 +88,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..e1e5451e3 100644 --- a/src/twinkle/server/state/backend/factory.py +++ b/src/twinkle/server/state/backend/factory.py @@ -1,5 +1,7 @@ """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 @@ -25,15 +27,19 @@ def create_backend(config: PersistenceConfig | None = None) -> StateBackend: config = PersistenceConfig() match config.mode: - case 'memory': + case "memory": # 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) - case 'redis': + + 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') + raise ValueError("redis_url is required for redis persistence mode") from .redis_backend import RedisBackend + return RedisBackend(config.redis_url, key_prefix=config.key_prefix) case _: - raise ValueError(f'Unknown persistence mode: {config.mode}') + raise ValueError(f"Unknown persistence mode: {config.mode}") diff --git a/src/twinkle/server/state/backend/memory_backend.py b/src/twinkle/server/state/backend/memory_backend.py index 433f8baa4..2d96b2d51 100644 --- a/src/twinkle/server/state/backend/memory_backend.py +++ b/src/twinkle/server/state/backend/memory_backend.py @@ -1,21 +1,21 @@ from __future__ import annotations - -import ray import time from collections.abc import Callable from fnmatch import fnmatch from typing import Any +import ray + from .base import StateBackend -_ACTOR_NAME_PREFIX = 'twinkle_state_actor' +_ACTOR_NAME_PREFIX = "twinkle_state_actor" def _actor_name(key_prefix: str) -> str: """Namespace the detached actor by ``key_prefix`` so multiple Twinkle deployments sharing one Ray cluster do not collide. """ - return _ACTOR_NAME_PREFIX if not key_prefix else f'{_ACTOR_NAME_PREFIX}::{key_prefix}' + return _ACTOR_NAME_PREFIX if not key_prefix else f"{_ACTOR_NAME_PREFIX}::{key_prefix}" @ray.remote @@ -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,29 @@ 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.') + 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 +173,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..0afe614c4 100644 --- a/src/twinkle/server/state/backend/redis_backend.py +++ b/src/twinkle/server/state/backend/redis_backend.py @@ -1,15 +1,18 @@ from __future__ import annotations - import asyncio -import json import random from collections.abc import Callable from typing import Any +import json + +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 @@ -33,19 +36,20 @@ class RedisBackend(StateBackend): TTL is managed by Redis native EXPIRE mechanism. """ - def __init__(self, redis_url: str, key_prefix: str = '') -> None: + def __init__(self, redis_url: str, key_prefix: str = "") -> None: if not _REDIS_AVAILABLE: - raise ImportError('redis package required. Install with: pip install redis') + 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.""" - return f'{self._prefix}{key}' if self._prefix else key + return f"{self._prefix}{key}" if self._prefix else key def _strip_prefix(self, key: str) -> str: """Remove namespace prefix from full key.""" - return key[len(self._prefix):] if self._prefix else key + return key[len(self._prefix) :] if self._prefix else key async def set(self, key: str, value: Any, ttl: int | None = None) -> None: """Store key-value pair with optional TTL in seconds.""" @@ -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,29 +150,48 @@ 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 - raise ConcurrencyError(f'update_atomic exhausted {_UPDATE_ATOMIC_MAX_RETRIES} retries on key {key!r}') + 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]: """Batch-read via native Redis MGET — single round-trip.""" 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..1a4be1e8f 100644 --- a/src/twinkle/server/state/cleanup_coordinator.py +++ b/src/twinkle/server/state/cleanup_coordinator.py @@ -13,8 +13,8 @@ is driven through the ``on_become_leader`` / ``on_lose_leader`` callbacks so the telemetry dependency stays isolated to ``count_publisher``. """ -from __future__ import annotations +from __future__ import annotations import asyncio import functools import time @@ -37,7 +37,7 @@ # inside the shared StateBackend: the winner runs cleanup + publishes metrics, # the others stay quiet. -LEADER_KEY = 'cleanup_leader' # actual backend key: 'cleanup_leader' +LEADER_KEY = "cleanup_leader" # actual backend key: 'cleanup_leader' LEASE_TTL = 30 # seconds — leader loses the lease after this without a renew LEASE_RENEW = 10 # seconds — must be < LEASE_TTL/2 so two missed renews still beat the TTL @@ -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 @@ -94,10 +95,10 @@ async def cleanup_expired_resources(self) -> dict[str, int]: current_time = time.time() cutoff_time = current_time - self._expiration_timeout - session_mgr = self._managers['sessions'] - model_mgr = self._managers['models'] - sampling_mgr = self._managers['sampling_sessions'] - future_mgr = self._managers['futures'] + session_mgr = self._managers["sessions"] + model_mgr = self._managers["models"] + sampling_mgr = self._managers["sampling_sessions"] + future_mgr = self._managers["futures"] # Determine expired sessions and remove them in a SINGLE pass, then cascade # the SAME set to dependent resources. Using one authoritative set closes the @@ -113,10 +114,10 @@ async def cleanup_expired_resources(self) -> dict[str, int]: await self._sweep_processor_quotas() return { - 'sessions': sessions_removed, - 'models': models_removed, - 'sampling_sessions': samplings_removed, - 'futures': futures_removed, + "sessions": sessions_removed, + "models": models_removed, + "sampling_sessions": samplings_removed, + "futures": futures_removed, } async def _cleanup_loop(self) -> None: @@ -132,27 +133,25 @@ async def _cleanup_loop(self) -> None: continue stats = await self.cleanup_expired_resources() if any(stats.values()): - logger.debug(f'[ServerState Cleanup] Removed expired resources: {stats}') + logger.debug(f"[ServerState Cleanup] Removed expired resources: {stats}") except asyncio.CancelledError: break except Exception as e: - logger.warning(f'[ServerState Cleanup] Error during cleanup: {e}') + logger.warning(f"[ServerState Cleanup] Error during cleanup: {e}") continue # ----- Leader election ----- 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 + logger.warning(f"[ServerState Leader] renew error: {e}") + await asyncio.sleep(LEASE_RENEW) async def _try_acquire_or_renew(self) -> None: was_leader = self._is_leader @@ -163,30 +162,21 @@ async def _try_acquire_or_renew(self) -> None: functools.partial(_renew_if_owner, owner=self._leader_id), ttl=LEASE_TTL, ) - self._is_leader = (val == self._leader_id) + self._is_leader = val == self._leader_id else: self._is_leader = await self._backend.set_nx(LEADER_KEY, self._leader_id, ttl=LEASE_TTL) except Exception as e: - logger.warning(f'[ServerState Leader] backend error during election: {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]})') + logger.info(f"[ServerState] became cleanup leader (id={self._leader_id[:8]})") if self._on_become_leader is not None: await self._on_become_leader() elif not self._is_leader and was_leader: - logger.warning(f'[ServerState] lost cleanup leadership (id={self._leader_id[:8]})') + logger.warning(f"[ServerState] lost cleanup leadership (id={self._leader_id[:8]})") if self._on_lose_leader is not None: await self._on_lose_leader() @@ -199,49 +189,45 @@ 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.""" return { - 'expiration_timeout': self._expiration_timeout, - 'cleanup_interval': self._cleanup_interval, - 'cleanup_running': self._cleanup_running, - 'is_leader': self._is_leader, - 'leader_id': self._leader_id, - 'resource_counts': { - 'sessions': await self._managers['sessions'].count(), - 'models': await self._managers['models'].count(), - 'sampling_sessions': await self._managers['sampling_sessions'].count(), - 'futures': await self._managers['futures'].count(), + "expiration_timeout": self._expiration_timeout, + "cleanup_interval": self._cleanup_interval, + "cleanup_running": self._cleanup_running, + "is_leader": self._is_leader, + "leader_id": self._leader_id, + "resource_counts": { + "sessions": await self._managers["sessions"].count(), + "models": await self._managers["models"].count(), + "sampling_sessions": await self._managers["sampling_sessions"].count(), + "futures": await self._managers["futures"].count(), }, } diff --git a/src/twinkle/server/state/future_manager.py b/src/twinkle/server/state/future_manager.py index 54bca4c70..61894b8f4 100644 --- a/src/twinkle/server/state/future_manager.py +++ b/src/twinkle/server/state/future_manager.py @@ -1,10 +1,10 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from __future__ import annotations - import functools 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,107 +12,121 @@ 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'}) - - -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 + +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} + + @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]): """Manage future state, terminal retention, and immutable task deadlines.""" def __init__(self, backend: StateBackend, expiration_timeout: float) -> None: - super().__init__(backend, 'future::', FutureRecord, expiration_timeout) + super().__init__(backend, "future::", FutureRecord, expiration_timeout) # ----- Future-specific operations ----- @@ -128,6 +142,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. @@ -138,14 +153,14 @@ async def store_status( If the result object has a ``model_dump`` method (i.e. it is a Pydantic model) it is serialized to a plain dict before storage. """ - is_failure = status in ('failed', 'cancelled') + is_failure = status in ("failed", "cancelled") if is_failure and failure is None: - raise ValueError(f'{status} future requires a FutureFailureRecord') + raise ValueError(f"{status} future requires a FutureFailureRecord") if is_failure and result is not None: - raise ValueError(f'{status} future cannot carry result') + raise ValueError(f"{status} future cannot carry result") if not is_failure and failure is not None: - raise ValueError(f'{status} future cannot carry failure') - if result is not None and hasattr(result, 'model_dump'): + raise ValueError(f"{status} future cannot carry failure") + if result is not None and hasattr(result, "model_dump"): result = result.model_dump() failure_data = failure.model_dump() if failure is not None else None @@ -153,7 +168,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 +179,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. @@ -177,15 +201,15 @@ async def cancel_if_pending(self, request_id: str) -> str | None: record for ``request_id``. """ failure = FutureFailureRecord( - reason_code='cancelled', - message='Task cancelled by client', - attribution='user', + reason_code="cancelled", + message="Task cancelled by client", + attribution="user", ) 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 + return result.get("status") if result else None # ----- Cleanup ----- @@ -219,25 +243,25 @@ 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) + replica_alive = replica_id is None or alive_replica_ids is None or replica_id in alive_replica_ids if not replica_alive: await self.store_status( request_id, - 'failed', + "failed", record.model_id, failure=FutureFailureRecord( - reason_code='orphaned_replica', - message='The replica that owned this task is no longer available.', - attribution='server', + reason_code="orphaned_replica", + message="The replica that owned this task is no longer available.", + attribution="server", ), replica_id=replica_id, ) @@ -248,12 +272,12 @@ async def cleanup_expired( if now > deadline: await self.store_status( request_id, - 'failed', + "failed", record.model_id, failure=FutureFailureRecord( - reason_code='deadline_exceeded', - message='Task exceeded the absolute survival bound without reaching a terminal state.', - attribution='server', + reason_code="deadline_exceeded", + message="Task exceeded the absolute survival bound without reaching a terminal state.", + attribution="server", ), replica_id=replica_id, ) diff --git a/src/twinkle/server/state/model_manager.py b/src/twinkle/server/state/model_manager.py index aaf520fca..d8b03a5ee 100644 --- a/src/twinkle/server/state/model_manager.py +++ b/src/twinkle/server/state/model_manager.py @@ -8,8 +8,8 @@ consistent view of the cluster's model registry without an extra coordination actor in front. """ -from __future__ import annotations +from __future__ import annotations import functools import time @@ -23,7 +23,7 @@ # picked up by ``model::*`` scans. The count is maintained atomically via # ``StateBackend.update_atomic`` so two concurrent adds with the same token # cannot both pass the limit check (the prior count-then-add race). -_TOKEN_COUNT_PREFIX = 'token_count::' +_TOKEN_COUNT_PREFIX = "token_count::" def _counter_delta_transform(existing: object, *, delta: int) -> int: @@ -60,14 +60,14 @@ class ModelManager(BaseManager[ModelRecord]): """ def __init__(self, backend: StateBackend, expiration_timeout: float, per_token_model_limit: int = 30) -> None: - super().__init__(backend, 'model::', ModelRecord, expiration_timeout) + super().__init__(backend, "model::", ModelRecord, expiration_timeout) self._per_token_model_limit = per_token_model_limit self._replicas = ReplicaRegistry(backend) # ----- Index Rebuild -------------------------------------------------- # def _token_count_key(self, token: str) -> str: - return f'{_TOKEN_COUNT_PREFIX}{token}' + return f"{_TOKEN_COUNT_PREFIX}{token}" async def rebuild_indexes(self) -> None: """Rebuild the per-token model counters from the persisted ``model::*`` records. @@ -83,7 +83,7 @@ async def rebuild_indexes(self) -> None: counts[record.token] = counts.get(record.token, 0) + 1 # Reset any stale counters, then write the recomputed ones. - stale_keys = await self._backend.keys(f'{_TOKEN_COUNT_PREFIX}*') + stale_keys = await self._backend.keys(f"{_TOKEN_COUNT_PREFIX}*") for key in stale_keys: await self._backend.delete(key) for token, count in counts.items(): @@ -93,14 +93,15 @@ 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) return { - 'max_loras': total_max, - 'used_loras': total_used, - 'free_loras': max(0, total_max - total_used), + "max_loras": total_max, + "used_loras": total_used, + "free_loras": max(0, total_max - total_used), } # ----- Replica Registration ------------------------------------------ # @@ -116,6 +117,26 @@ 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) @@ -196,7 +217,8 @@ async def add(self, model_id: str, record: ModelRecord) -> None: # is the count that was already present before this add. await self._backend.update_atomic(key, functools.partial(_counter_delta_transform, delta=-1)) raise ResourceQuotaExceededError( - f'Model quota exceeded for this token: {new_count - 1}/{self._per_token_model_limit} models') + f"Model quota exceeded for this token: {new_count - 1}/{self._per_token_model_limit} models" + ) try: await super().add(model_id, record) @@ -220,6 +242,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..89eb92e5c 100644 --- a/src/twinkle/server/state/models.py +++ b/src/twinkle/server/state/models.py @@ -1,17 +1,27 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from __future__ import annotations - import time -from datetime import datetime, timezone -from pydantic import BaseModel, Field, model_validator +from datetime import UTC, datetime, timezone from typing import Any, Literal +from pydantic import BaseModel, ConfigDict, Field, model_validator + def _now_iso() -> str: # UTC-aware so _parse_timestamp (which reads timestamps back as UTC) agrees with # it and with time.time(); a naive local string would be misread as UTC and skew # every expiry comparison by the host's UTC offset. - return datetime.now(timezone.utc).isoformat() + return datetime.now(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): @@ -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): @@ -52,37 +65,15 @@ class FutureFailureRecord(BaseModel): reason_code: str message: str - attribution: Literal['user', 'server'] + attribution: Literal["user", "server"] details: list[dict[str, Any]] | None = None 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 @@ -96,13 +87,13 @@ class FutureRecord(BaseModel): created_at: str = Field(default_factory=_now_iso) updated_at: str = Field(default_factory=_now_iso) - @model_validator(mode='after') + @model_validator(mode="after") def result_and_failure_match_status(self) -> FutureRecord: - is_failure = self.status in ('failed', 'cancelled') + is_failure = self.status in ("failed", "cancelled") if is_failure and self.failure is None: - raise ValueError(f'{self.status} future requires failure') + raise ValueError(f"{self.status} future requires failure") if is_failure and self.result is not None: - raise ValueError(f'{self.status} future cannot carry result') + raise ValueError(f"{self.status} future cannot carry result") if not is_failure and self.failure is not None: - raise ValueError(f'{self.status} future cannot carry failure') + raise ValueError(f"{self.status} future cannot carry failure") return self diff --git a/src/twinkle/server/state/server_state.py b/src/twinkle/server/state/server_state.py index dd926aa57..cd07be2c7 100644 --- a/src/twinkle/server/state/server_state.py +++ b/src/twinkle/server/state/server_state.py @@ -1,6 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from __future__ import annotations - import functools import math import re @@ -10,9 +9,15 @@ from typing import Any from twinkle.server.config.persistence import PersistenceConfig -from twinkle.server.exceptions import ResourceQuotaExceededError -from twinkle.server.telemetry.correlation import (BASE_MODEL, MODEL_ID, REPLICA_ID, SAMPLING_SESSION_ID, SESSION_ID, - TOKEN_ID) +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 from twinkle.utils.logger import get_logger from .backend import StateBackend @@ -24,10 +29,11 @@ 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() -_PROCESSOR_QUOTA_PREFIX = 'processor_quota::' +_PROCESSOR_QUOTA_PREFIX = "processor_quota::" def _clean_processor_reservations(existing: Any, *, now: float) -> dict[str, dict[str, Any]]: @@ -38,7 +44,7 @@ def _clean_processor_reservations(existing: Any, *, now: float) -> dict[str, dic for processor_id, reservation in existing.items(): if not isinstance(processor_id, str) or not isinstance(reservation, dict): continue - expires_at = reservation.get('lease_expires_at') + expires_at = reservation.get("lease_expires_at") if isinstance(expires_at, (int, float)) and float(expires_at) > now: active[processor_id] = dict(reservation) return active @@ -56,8 +62,8 @@ def _reserve_processor_transform( reservations = _clean_processor_reservations(existing, now=now) if processor_id in reservations or len(reservations) < limit: reservations[processor_id] = { - 'session_id': session_id, - 'lease_expires_at': now + lease_seconds, + "session_id": session_id, + "lease_expires_at": now + lease_seconds, } return reservations @@ -72,7 +78,7 @@ def _renew_processor_transform( reservations = _clean_processor_reservations(existing, now=now) reservation = reservations.get(processor_id) if reservation is not None: - reservation['lease_expires_at'] = now + lease_seconds + reservation["lease_expires_at"] = now + lease_seconds return reservations @@ -116,13 +122,14 @@ class ServerState: """ def __init__( - self, - backend: StateBackend | None = None, - persistence_config: PersistenceConfig | None = None, - expiration_timeout: float = 86400.0, # 24 hours in seconds - cleanup_interval: float = 3600.0, # 1 hour in seconds - per_token_model_limit: int = 30, - metrics_update_interval: float = 15.0) -> None: + self, + backend: StateBackend | None = None, + persistence_config: PersistenceConfig | None = None, + expiration_timeout: float = 86400.0, # 24 hours in seconds + cleanup_interval: float = 3600.0, # 1 hour in seconds + per_token_model_limit: int = 30, + metrics_update_interval: float = 15.0, + ) -> None: if backend is not None: self._backend: StateBackend = backend else: @@ -131,6 +138,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 @@ -142,17 +150,17 @@ def __init__( # the leader publishes, while keeping the coordinator itself free # of any telemetry dependency. _managers = { - 'sessions': self._session_mgr, - 'models': self._model_mgr, - 'sampling_sessions': self._sampling_mgr, - 'futures': self._future_mgr, + "sessions": self._session_mgr, + "models": self._model_mgr, + "sampling_sessions": self._sampling_mgr, + "futures": self._future_mgr, } self._count_publisher = ResourceCountPublisher( [ - ('active_sessions', self._session_mgr), - ('active_models', self._model_mgr), - ('active_sampling_sessions', self._sampling_mgr), - ('active_futures', self._future_mgr), + ("active_sessions", self._session_mgr), + ("active_models", self._model_mgr), + ("active_sampling_sessions", self._sampling_mgr), + ("active_futures", self._future_mgr), ], interval=float(metrics_update_interval), ) @@ -173,6 +181,34 @@ 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: @@ -184,15 +220,15 @@ async def create_session(self, payload: dict[str, Any]) -> str: Returns: The session_id for the created session. """ - session_id = payload.get('session_id') or f'session_{uuid.uuid4().hex}' + session_id = payload.get("session_id") or f"session_{uuid.uuid4().hex}" with traced_operation( - 'server_state.create_session', - attrs={SESSION_ID: session_id}, + "server_state.create_session", + attrs={SESSION_ID: session_id}, ): record = SessionRecord( - tags=list(payload.get('tags') or []), - user_metadata=payload.get('user_metadata') or {}, - sdk_version=payload.get('sdk_version'), + tags=list(payload.get("tags") or []), + user_metadata=payload.get("user_metadata") or {}, + sdk_version=payload.get("sdk_version"), ) await self._session_mgr.add(session_id, record) return session_id @@ -213,14 +249,19 @@ 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, - payload: dict[str, Any], - token: str, - model_id: str | None = None, - replica_id: str | None = None, - session_id: str | None = None) -> str: + async def register_model( + self, + payload: dict[str, Any], + token: str, + model_id: str | None = None, + replica_id: str | None = None, + session_id: str | None = None, + ) -> str: """Register a new model with the server state. Args: @@ -234,27 +275,30 @@ async def register_model(self, Returns: The model_id for the registered model. """ - _time = datetime.now().strftime('%Y%m%d_%H%M%S') - _model_id: str = model_id or payload.get( - 'model_id') or f"{_time}-{payload.get('base_model', 'model')}-{uuid.uuid4().hex[:8]}" - _model_id = re.sub(r'[^\w\-]', '_', _model_id) + _time = datetime.now().strftime("%Y%m%d_%H%M%S") + _model_id: str = ( + model_id + or payload.get("model_id") + or f"{_time}-{payload.get('base_model', 'model')}-{uuid.uuid4().hex[:8]}" + ) + _model_id = re.sub(r"[^\w\-]", "_", _model_id) with traced_operation( - 'server_state.register_model', - attrs={ - MODEL_ID: _model_id, - BASE_MODEL: payload.get('base_model'), - REPLICA_ID: replica_id, - TOKEN_ID: token, - SESSION_ID: session_id or payload.get('session_id'), - }, + "server_state.register_model", + attrs={ + MODEL_ID: _model_id, + BASE_MODEL: payload.get("base_model"), + REPLICA_ID: replica_id, + TOKEN_ID: token, + SESSION_ID: session_id or payload.get("session_id"), + }, ): record = ModelRecord( - session_id=session_id or payload.get('session_id'), - model_seq_id=payload.get('model_seq_id'), - base_model=payload.get('base_model'), - user_metadata=payload.get('user_metadata') or {}, - lora_config=payload.get('lora_config'), + session_id=session_id or payload.get("session_id"), + model_seq_id=payload.get("model_seq_id"), + base_model=payload.get("base_model"), + user_metadata=payload.get("user_metadata") or {}, + lora_config=payload.get("lora_config"), token=token, replica_id=replica_id, ) @@ -284,8 +328,8 @@ async def register_replica(self, replica_id: str, max_loras: int) -> None: max_loras: Maximum number of LoRA adapters the replica can hold. """ with traced_operation( - 'server_state.register_replica', - attrs={REPLICA_ID: replica_id}, + "server_state.register_replica", + attrs={REPLICA_ID: replica_id}, ): await self._model_mgr.register_replica(replica_id, max_loras) @@ -312,7 +356,7 @@ async def get_available_replica_ids(self, candidate_ids: list[str]) -> list[str] @staticmethod def _processor_quota_key(token: str) -> str: - return f'{_PROCESSOR_QUOTA_PREFIX}{token}' + return f"{_PROCESSOR_QUOTA_PREFIX}{token}" @staticmethod def _processor_quota_ttl(lease_seconds: float) -> int: @@ -344,7 +388,7 @@ async def reserve_processor_quota( ttl=self._processor_quota_ttl(lease_seconds), ) if not isinstance(reservations, dict) or processor_id not in reservations: - raise ResourceQuotaExceededError(f'Per-user processor quota ({limit}) reached for token {token[:8]}...') + raise ResourceQuotaExceededError(f"Per-user processor quota ({limit}) reached for token {token[:8]}...") async def renew_processor_quota( self, @@ -378,7 +422,7 @@ async def release_processor_quota(self, token: str, processor_id: str, *, lease_ async def sweep_processor_quotas(self, *, lease_seconds: float = 30.0) -> None: """Remove expired leases from every persisted processor quota map.""" now = time.time() - keys = await self._backend.keys(f'{_PROCESSOR_QUOTA_PREFIX}*') + keys = await self._backend.keys(f"{_PROCESSOR_QUOTA_PREFIX}*") for key in keys: await self._backend.update_atomic( key, @@ -398,21 +442,22 @@ async def create_sampling_session(self, payload: dict[str, Any], sampling_sessio Returns: The sampling_session_id. """ - _sampling_session_id: str = sampling_session_id or payload.get( - 'sampling_session_id') or f'sampling_{uuid.uuid4().hex}' + _sampling_session_id: str = ( + sampling_session_id or payload.get("sampling_session_id") or f"sampling_{uuid.uuid4().hex}" + ) with traced_operation( - 'server_state.create_sampling_session', - attrs={ - SAMPLING_SESSION_ID: _sampling_session_id, - SESSION_ID: payload.get('session_id'), - BASE_MODEL: payload.get('base_model'), - }, + "server_state.create_sampling_session", + attrs={ + SAMPLING_SESSION_ID: _sampling_session_id, + SESSION_ID: payload.get("session_id"), + BASE_MODEL: payload.get("base_model"), + }, ): record = SamplingSessionRecord( - session_id=payload.get('session_id'), - seq_id=payload.get('sampling_session_seq_id'), - base_model=payload.get('base_model'), - model_path=payload.get('model_path'), + session_id=payload.get("session_id"), + seq_id=payload.get("sampling_session_seq_id"), + base_model=payload.get("base_model"), + model_path=payload.get("model_path"), ) await self._sampling_mgr.add(_sampling_session_id, record) return _sampling_session_id @@ -429,7 +474,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 +497,18 @@ 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. @@ -455,8 +518,8 @@ async def cancel_future(self, request_id: str) -> dict[str, Any]: """ status = await self._future_mgr.cancel_if_pending(request_id) if status is None: - return {'cancelled': False, 'state': 'not_found'} - return {'cancelled': status == 'cancelled', 'state': status} + return {"cancelled": False, "state": "not_found"} + return {"cancelled": status == "cancelled", "state": status} async def store_future_status( self, @@ -470,6 +533,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 +565,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 ----- @@ -571,13 +636,15 @@ async def get_cleanup_stats(self) -> dict[str, Any]: _DEFAULT_METRICS_UPDATE_INTERVAL = 15.0 -def get_server_state(cache_key: str = 'twinkle_server_state', - backend: StateBackend | None = None, - persistence_config: PersistenceConfig | None = None, - expiration_timeout: float | None = None, - cleanup_interval: float | None = None, - per_token_model_limit: int | None = None, - metrics_update_interval: float | None = None) -> ServerState: +def get_server_state( + cache_key: str = "twinkle_server_state", + backend: StateBackend | None = None, + persistence_config: PersistenceConfig | None = None, + expiration_timeout: float | None = None, + cleanup_interval: float | None = None, + per_token_model_limit: int | None = None, + metrics_update_interval: float | None = None, +) -> ServerState: """Return a process-local :class:`ServerState` bound directly to the backend. Within one process the same ``cache_key`` returns the same cached instance @@ -613,7 +680,8 @@ 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,14 +690,14 @@ 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), - _DEFAULT_PER_TOKEN_MODEL_LIMIT) - metrics_update_interval = _resolve(metrics_update_interval, getattr(env_policy, 'metrics_update_interval', None), - _DEFAULT_METRICS_UPDATE_INTERVAL) + 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, env_policy.metrics_update_interval, _DEFAULT_METRICS_UPDATE_INTERVAL + ) state = ServerState( backend=backend, @@ -641,9 +709,13 @@ def _resolve(explicit, env_value, default): ) _PROCESS_STATE_CACHE[cache_key] = state logger.info( - 'ServerState policy in effect: per_token_model_limit=%s expiration_timeout=%s ' - 'cleanup_interval=%s metrics_update_interval=%s (resolution: explicit>env>default)', per_token_model_limit, - expiration_timeout, cleanup_interval, metrics_update_interval) + "ServerState policy in effect: per_token_model_limit=%s expiration_timeout=%s " + "cleanup_interval=%s metrics_update_interval=%s (resolution: explicit>env>default)", + per_token_model_limit, + expiration_timeout, + cleanup_interval, + metrics_update_interval, + ) # Cleanup task is started by the deployment's FastAPI ``lifespan`` hook # via ``await state.start_cleanup_task()`` — that's the single async # entry point each worker has, so we don't need any sync-context diff --git a/src/twinkle/server/state/session_manager.py b/src/twinkle/server/state/session_manager.py index 0e4beeb37..03d8be1b1 100644 --- a/src/twinkle/server/state/session_manager.py +++ b/src/twinkle/server/state/session_manager.py @@ -1,6 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from __future__ import annotations - import functools import time @@ -22,7 +21,7 @@ def _session_touch_transform(existing: dict | None, *, now: float) -> dict | Non if existing is None: return None updated = dict(existing) - updated['last_heartbeat'] = now + updated["last_heartbeat"] = now return updated @@ -34,7 +33,7 @@ class SessionManager(BaseManager[SessionRecord]): """ def __init__(self, backend: StateBackend, expiration_timeout: float) -> None: - super().__init__(backend, 'session::', SessionRecord, expiration_timeout) + super().__init__(backend, "session::", SessionRecord, expiration_timeout) # ----- Session-specific operations ----- @@ -56,7 +55,7 @@ async def touch(self, session_id: str) -> bool: except ConcurrencyError as e: # touch() is per-heartbeat best-effort — losing one round under # contention is recoverable as soon as the next heartbeat arrives. - logger.warning('SessionManager.touch dropped due to contention: %s', e) + logger.warning("SessionManager.touch dropped due to contention: %s", e) return False return result is not None @@ -67,6 +66,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..6194cffcf 100644 --- a/src/twinkle/server/task_errors.py +++ b/src/twinkle/server/task_errors.py @@ -1,14 +1,60 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Helpers for constructing protocol-layer ``ErrorPayload`` values.""" -from __future__ import annotations +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' +_TRUNCATION_MARKER = "...[traceback truncated, tail kept]...\n" def trim_traceback(text: str) -> str: @@ -33,7 +79,7 @@ def build_error_payload( category = ErrorCategory(category.lower()) tb = trim_traceback(traceback_text) if category is ErrorCategory.Server and traceback_text else None lines = str(error).splitlines() - summary = (lines[0] if lines else '')[:_ERROR_MAX] + summary = (lines[0] if lines else "")[:_ERROR_MAX] return ErrorPayload( error=summary, category=category, @@ -62,4 +108,4 @@ def task_error_payload( traceback_text=traceback_text, details=details, ) - return payload.model_dump(mode='json', exclude_none=True) + return payload.model_dump(mode="json", exclude_none=True) diff --git a/src/twinkle/server/task_queue/__init__.py b/src/twinkle/server/task_queue/__init__.py index bf5391611..2c0a33cae 100644 --- a/src/twinkle/server/task_queue/__init__.py +++ b/src/twinkle/server/task_queue/__init__.py @@ -5,23 +5,22 @@ 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', + "TaskStatus", + "QueueState", + "QueuedTask", + "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..54dc1ff4f 100644 --- a/src/twinkle/server/task_queue/backend_gate.py +++ b/src/twinkle/server/task_queue/backend_gate.py @@ -13,8 +13,8 @@ invisible to the AST scan in ``test_no_direct_backend_call.py`` and needs no exemption entry. """ -from __future__ import annotations +from __future__ import annotations import asyncio import contextlib import functools @@ -22,7 +22,12 @@ from concurrent.futures import ThreadPoolExecutor from typing import Any -from .types import BackendBusyError +from ray.exceptions import GetTimeoutError + +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: @@ -34,14 +39,14 @@ class BackendGate: """ def __init__(self, *, enable_admission_gate: bool = False) -> None: - self._executor = ThreadPoolExecutor(thread_name_prefix='twinkle-backend') - self._probe_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='twinkle-backend-probe') + self._executor = ThreadPoolExecutor(thread_name_prefix="twinkle-backend") + self._probe_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="twinkle-backend-probe") self._admission: asyncio.Lock | None = asyncio.Lock() if enable_admission_gate else None self._poisoned = asyncio.Event() async def _acquire(self, gate: asyncio.Lock) -> None: if self._poisoned.is_set(): - raise BackendBusyError('This replica is waiting for a timed-out backend call to exit.') + raise BackendBusyError("This replica is waiting for a timed-out backend call to exit.") if not gate.locked(): await gate.acquire() else: @@ -63,14 +68,14 @@ async def _acquire(self, gate: asyncio.Lock) -> None: acquire_task.cancel() with contextlib.suppress(asyncio.CancelledError): await acquire_task - raise BackendBusyError('This replica is waiting for a timed-out backend call to exit.') + raise BackendBusyError("This replica is waiting for a timed-out backend call to exit.") poison_task.cancel() with contextlib.suppress(asyncio.CancelledError): await poison_task await acquire_task if self._poisoned.is_set(): gate.release() - raise BackendBusyError('This replica is waiting for a timed-out backend call to exit.') + raise BackendBusyError("This replica is waiting for a timed-out backend call to exit.") async def call(self, fn: Callable[..., Any], /, *args: Any, admit: bool = True, **kwargs: Any) -> Any: """Run one backend call outside the event loop. @@ -97,7 +102,6 @@ async def call(self, fn: Callable[..., Any], /, *args: Any, admit: bool = True, if gate is not None: def release_gate(_future) -> None: - def release() -> None: self._poisoned.clear() if gate.locked(): diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index b06ef8525..a2b7c27b5 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -5,16 +5,17 @@ used for I/O and for engines such as vLLM that own their compute concurrency and continuous batching internally. """ -from __future__ import annotations +from __future__ import annotations 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 ------------ @@ -64,7 +63,7 @@ class TaskQueueMixin: def _init_task_queue( self, config: TaskQueueConfig | None = None, - deployment_name: str = '', + deployment_name: str = "", *, enable_admission_gate: bool = False, on_backend_timeout: Callable[[], Coroutine[Any, Any, None]] | None = None, @@ -88,8 +87,10 @@ def _init_task_queue( self._task_queue_config = config if config is not None else TaskQueueConfig() if self._task_queue_config.execution_timeout == 0: logger.warning( - '[TaskQueue] execution_timeout=0: a finite %.0fs bound has replaced unbounded waiting ' - '(deployment=%s).', self._task_queue_config.effective_execution_timeout, deployment_name or 'unknown') + "[TaskQueue] execution_timeout=0: a finite %.0fs bound has replaced unbounded waiting (deployment=%s).", + self._task_queue_config.effective_execution_timeout, + deployment_name or "unknown", + ) # The Inline_Fast_Path window must stay strictly under Long_Poll_Window: a # submit that peeks longer than a retrieve would wait makes no sense (D4). # Raised, not asserted -- `python -O` strips asserts and would drop this @@ -98,10 +99,13 @@ def _init_task_queue( _window = long_poll_window() if _inline >= _window: raise ConfigError( - 'inline_fast_path_timeout', + "inline_fast_path_timeout", _inline, - message=(f'inline_fast_path_timeout ({_inline}s) must be < Long_Poll_Window ' - f'({_window}s); lower it or raise TWINKLE_LONG_POLL_TIMEOUT.')) + message=( + f"inline_fast_path_timeout ({_inline}s) must be < Long_Poll_Window " + f"({_window}s); lower it or raise TWINKLE_LONG_POLL_TIMEOUT." + ), + ) self._deployment_name = deployment_name self._task_metrics = get_task_metrics(deployment_name) if deployment_name else None self._future_absolute_ttl = self._task_queue_config.absolute_future_ttl(collect_width) @@ -130,6 +134,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 +146,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. @@ -152,16 +161,16 @@ async def call_backend(self, fn: Callable[..., Any], /, *args: Any, admit: bool return await self._backend_gate.call(fn, *args, admit=admit, **kwargs) def _future_deadline(self) -> float: - ttl = getattr(self, '_future_absolute_ttl', self._task_queue_config.absolute_future_ttl(1)) + ttl = getattr(self, "_future_absolute_ttl", self._task_queue_config.absolute_future_ttl(1)) return time.time() + ttl @staticmethod def _queue_key(model_id: str | None, token: str | None) -> str: if model_id: - return f'model:{model_id}' + return f"model:{model_id}" if token: - return f'token:{token}' - return 'default' + return f"token:{token}" + return "default" async def _perform_preflight_checks( self, @@ -184,24 +193,27 @@ async def _perform_preflight_checks( return if input_tokens > self._task_queue_config.max_input_tokens: - raise InputTokensExceededError(f'Input tokens ({input_tokens}) exceed maximum allowed ' - f'({self._task_queue_config.max_input_tokens})') + raise InputTokensExceededError( + f"Input tokens ({input_tokens}) exceed maximum allowed ({self._task_queue_config.max_input_tokens})" + ) if batch_size is not None and data_world_size is not None: if batch_size < data_world_size: - raise BatchSizeError(f'Batch size {batch_size} must be >= data world size {data_world_size}') + raise BatchSizeError(f"Batch size {batch_size} must be >= data world size {data_world_size}") if batch_size_multiple is not None: required_multiple = data_world_size * batch_size_multiple if batch_size % required_multiple != 0: - raise BatchSizeError(f'Batch size {batch_size} must be divisible by {required_multiple} ' - f'so each data-parallel shard gets a multiple of ' - f'{batch_size_multiple} examples') + raise BatchSizeError( + f"Batch size {batch_size} must be divisible by {required_multiple} " + f"so each data-parallel shard gets a multiple of " + f"{batch_size_multiple} examples" + ) allowed, reason = await self._rate_limiter.check_and_record(token, input_tokens) if not allowed: if self._task_metrics: - self._task_metrics.rate_limit_rejections.inc(tags={'deployment': self._deployment_name}) - raise RateLimitExceededError(f'Rate limit exceeded: {reason}') + self._task_metrics.rate_limit_rejections.inc(tags={"deployment": self._deployment_name}) + raise RateLimitExceededError(f"Rate limit exceeded: {reason}") async def _schedule_task( self, @@ -217,7 +229,7 @@ async def _schedule_task( ) -> dict[str, Any]: """Common enqueue path. Always persists status: the future record is the single delivery channel for both result and failure.""" - request_id = request_id or f'req_{uuid.uuid4().hex}' + request_id = request_id or f"req_{uuid.uuid4().hex}" # Decision_Boundary: raises RequestRejectedError before any state write. await self._perform_preflight_checks( @@ -229,46 +241,82 @@ 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, model_id, queue_state=QueueState.ACTIVE.value, - replica_id=getattr(self, 'replica_id', None), + 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 +368,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. @@ -331,7 +380,7 @@ async def schedule_concurrent_task( Returns: {'request_id': str, 'model_id': str | None} """ - request_id = request_id or f'req_{uuid.uuid4().hex}' + request_id = request_id or f"req_{uuid.uuid4().hex}" # Decision_Boundary: raises RequestRejectedError before any state write. await self._perform_preflight_checks(model_id=model_id, token=token, input_tokens=input_tokens) @@ -344,24 +393,30 @@ 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()) - logger.info(f'[TaskQueue] Task {request_id} started concurrently, type={task_type or "unknown"}, ' - f'model_id={model_id}, input_tokens={input_tokens}') - return {'request_id': request_id, 'model_id': model_id} + replica_id=getattr(self, "replica_id", None), + absolute_deadline=self._future_deadline(), + ) + logger.info( + f"[TaskQueue] Task {request_id} started concurrently, type={task_type or 'unknown'}, " + f"model_id={model_id}, input_tokens={input_tokens}" + ) + return {"request_id": request_id, "model_id": model_id} # Poll cadence *inside* the Inline_Fast_Path window. Much smaller than the window # itself so a task that finishes early is noticed promptly; the window bound # (config.inline_fast_path_timeout) is what actually caps submit latency. _INLINE_FAST_PATH_POLL = 0.005 - async def submit_concurrent_and_peek(self, coro_factory: Callable[[], Coroutine], - **schedule_kwargs) -> TaskEnvelope: + async def submit_concurrent_and_peek( + self, coro_factory: Callable[[], Coroutine], **schedule_kwargs + ) -> TaskEnvelope: """Concurrent compute submission with the existing Twinkle future envelope.""" ref = await self.schedule_concurrent_task(coro_factory, **schedule_kwargs) - return await self._peek_terminal(ref['request_id'], fallback_status='pending') + return await self._peek_terminal(ref["request_id"], fallback_status="pending") async def _peek_terminal(self, request_id: str, *, fallback_status: str) -> TaskEnvelope: """Poll a record up to the Inline_Fast_Path window; return a terminal envelope @@ -370,7 +425,7 @@ async def _peek_terminal(self, request_id: str, *, fallback_status: str) -> Task record = None while time.monotonic() < deadline: record = await self.state.get_future(request_id) - if record is not None and record.get('status') in TERMINAL_STATUSES: + if record is not None and record.get("status") in TERMINAL_STATUSES: return envelope_from_record(request_id, record) await asyncio.sleep(self._INLINE_FAST_PATH_POLL) return envelope_from_record(request_id, record, fallback_status=fallback_status) @@ -396,8 +451,9 @@ async def submit_and_peek( still non-terminal returns a non-terminal envelope carrying ``queue_state``. """ ref = await self.schedule_task( - coro_factory, model_id=model_id, token=token, task_type=task_type, request_id=request_id, **schedule_kwargs) - return await self._peek_terminal(ref['request_id'], fallback_status='pending') + coro_factory, model_id=model_id, token=token, task_type=task_type, request_id=request_id, **schedule_kwargs + ) + return await self._peek_terminal(ref["request_id"], fallback_status="pending") async def submit_background_and_peek( self, @@ -414,80 +470,21 @@ async def submit_background_and_peek( non-terminal fallback is ``running``. """ ref = await self.schedule_background_task(coro_factory, model_id=model_id, task_type=task_type) - return await self._peek_terminal(ref['request_id'], fallback_status='running') + return await self._peek_terminal(ref["request_id"], fallback_status="running") async def schedule_background_task( self, 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. - - 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(). + """Run pure I/O outside serial compute, with the same supervised lifecycle. - 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) @@ -496,17 +493,27 @@ def fail_pending_tasks_for_model(self, model_id: str, reason: str) -> None: """Fail and drop all queued tasks for a model. Thread-safe.""" queue_key = self._queue_key(model_id=model_id, token=None) if self._event_loop is None: - logger.warning(f'[TaskQueue] fail_pending_tasks_for_model called without event loop: {queue_key}') + logger.warning(f"[TaskQueue] fail_pending_tasks_for_model called without event loop: {queue_key}") 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') + 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..b9a1f57ac 100644 --- a/src/twinkle/server/task_queue/types.py +++ b/src/twinkle/server/task_queue/types.py @@ -1,66 +1,39 @@ # 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 +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 class QueuedTask: """Dataclass representing a task waiting in the compute queue.""" + request_id: str coro_factory: Callable[[], Coroutine] model_id: str | None @@ -68,3 +41,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..6fdeaca54 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -6,27 +6,43 @@ GPU compute tasks serially across all per-adapter/per-token queues using a round-robin policy. """ -from __future__ import annotations +from __future__ import annotations import asyncio import time import traceback from collections import deque +from collections.abc import Callable from dataclasses import replace -from typing import TYPE_CHECKING, Any, Callable, Deque - -from twinkle.protocol.types.errors import ErrorCategory -from twinkle.server.exceptions import (BatchSizeError, EndpointUnavailableError, FullModeBusyError, - InputTokensExceededError, RateLimitExceededError, RequestRejectedError, - ResourceExhaustedError, ResourceNotFoundError, ResourceQuotaExceededError, - StateBackendError, TrainModeMismatchError, TwinkleServerError) +from typing import TYPE_CHECKING, Any + +from ray.exceptions import RayTaskError + +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,25 +51,46 @@ 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], ...] = ( - (TrainModeMismatchError, 'invalid_request'), - (ResourceNotFoundError, 'resource_not_found'), - (FullModeBusyError, 'full_mode_busy'), - (InputTokensExceededError, 'input_tokens_exceeded'), - (BatchSizeError, 'batch_size_invalid'), - (RateLimitExceededError, 'rate_limit_exceeded'), - (ResourceQuotaExceededError, 'resource_quota_exceeded'), - (EndpointUnavailableError, 'endpoint_unavailable'), - (ConcurrencyError, 'state_contention'), - (StateBackendError, 'backend_gate_unavailable'), - (RequestRejectedError, 'request_rejected'), + (TrainModeMismatchError, "invalid_request"), + (ResourceNotFoundError, "resource_not_found"), + (FullModeBusyError, "full_mode_busy"), + (InputTokensExceededError, "input_tokens_exceeded"), + (BatchSizeError, "batch_size_invalid"), + (RateLimitExceededError, "rate_limit_exceeded"), + (ResourceQuotaExceededError, "resource_quota_exceeded"), + (EndpointUnavailableError, "endpoint_unavailable"), + (ConcurrencyError, "state_contention"), + (StateBackendError, "backend_gate_unavailable"), + (RequestRejectedError, "request_rejected"), ) for error_type, reason_code in mappings: if isinstance(exc, error_type): return reason_code - return 'internal_error' + return "internal_error" class ComputeWorker: @@ -85,7 +122,9 @@ def __init__( self._on_backend_timeout = on_backend_timeout self.task_queues: dict[str, asyncio.Queue] = {} - self.queue_order: Deque[str] = deque() + 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() self._worker_task: asyncio.Task | None = None @@ -97,14 +136,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 @@ -125,13 +169,15 @@ async def stop(self) -> None: if running.cancelled() and not await self._is_record_terminal(task.request_id): await self._store_task_failed( task, - 'Sampler task stopped during replica shutdown.', + "Sampler task stopped during replica shutdown.", QueueState.ACTIVE.value, - reason_code='orphaned_replica') + 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 +195,38 @@ 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()) @@ -170,35 +248,41 @@ def notify_new_task(self) -> None: def _record_tasks_total(self, task_type: str, status: str) -> None: """Increment the tasks_total counter if metrics are enabled.""" if self._task_metrics: - self._task_metrics.tasks_total.inc(tags={ - 'deployment': self._deployment_name, - 'task_type': task_type, - 'status': status, - }) + self._task_metrics.tasks_total.inc( + tags={ + "deployment": self._deployment_name, + "task_type": task_type, + "status": status, + } + ) def _record_execution_time(self, task_type: str, exec_time: float) -> None: """Observe execution duration in the execution_seconds histogram if metrics are enabled.""" if self._task_metrics: self._task_metrics.execution_seconds.observe( - exec_time, tags={ - 'deployment': self._deployment_name, - 'task_type': task_type, - }) + exec_time, + tags={ + "deployment": self._deployment_name, + "task_type": task_type, + }, + ) def _record_queue_depth(self) -> None: """Single writer of the queue-depth gauge.""" if self._task_metrics: total_depth = sum(qq.qsize() for qq in self.task_queues.values()) - self._task_metrics.queue_depth.set(total_depth, tags={'deployment': self._deployment_name}) + self._task_metrics.queue_depth.set(total_depth, tags={"deployment": self._deployment_name}) def _record_queue_metrics(self, task_type: str, queue_wait: float) -> None: """Observe queue wait time and update current queue depth if metrics are enabled.""" if self._task_metrics: self._task_metrics.queue_wait_seconds.observe( - queue_wait, tags={ - 'deployment': self._deployment_name, - 'task_type': task_type, - }) + queue_wait, + tags={ + "deployment": self._deployment_name, + "task_type": task_type, + }, + ) self._record_queue_depth() # ------------------------------------------------------------------ @@ -210,15 +294,15 @@ async def _store_task_failed( queue_state: str, queue_state_reason: str | None = None, *, - reason_code: str = 'internal_error', - attribution: str = 'server', + reason_code: str = "internal_error", + attribution: str = "server", traceback_text: str | None = None, ) -> None: """Store FAILED status with protocol-independent failure details.""" failure = FutureFailureRecord( reason_code=reason_code, - message=(error.splitlines() or [''])[0][:1024], - attribution='user' if attribution == 'user' else 'server', + message=(error.splitlines() or [""])[0][:1024], + attribution="user" if attribution == "user" else "server", diagnostic=traceback_text, ) await self._state.store_future_status( @@ -246,6 +330,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: @@ -268,10 +353,11 @@ async def _wait_for_work(self) -> None: async def _fail_timed_out_task(self, task: QueuedTask, waited: float, q: asyncio.Queue) -> None: """Mark a queue-timed-out task as FAILED and release it from the queue.""" - error = f'Queue timeout exceeded: waited {waited:.2f}s' + error = f"Queue timeout exceeded: waited {waited:.2f}s" 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') + 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.""" @@ -280,8 +366,9 @@ async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue finally: q.task_done() - async def submit_detached(self, task: QueuedTask, queue_key: str, *, replica_id: str | None, - absolute_deadline: float) -> None: + async def submit_detached( + self, task: QueuedTask, queue_key: str, *, replica_id: str | None, absolute_deadline: float + ) -> None: """Own admission, future persistence and startup as one lifecycle.""" self.reserve_detached() wrote_status = False @@ -292,7 +379,9 @@ 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: @@ -300,15 +389,17 @@ async def submit_detached(self, task: QueuedTask, queue_key: str, *, replica_id: if wrote_status: await self._store_task_failed( task, - 'Replica stopped before the concurrent task could start.', + "Replica stopped before the concurrent task could start.", QueueState.ACTIVE.value, - reason_code='orphaned_replica') + reason_code="orphaned_replica", + ) raise 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: - raise ResourceExhaustedError('Concurrent compute capacity is unavailable on this replica.') + 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 def release_detached(self) -> None: @@ -322,7 +413,7 @@ def run_detached(self, task: QueuedTask, queue_key: str) -> None: queued one. """ if self._stopping: - raise ResourceExhaustedError('The replica is shutting down.') + raise ResourceExhaustedError("The replica is shutting down.") execution = None def start(): @@ -353,10 +444,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) @@ -365,46 +463,41 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: Handles execution timeout and general exceptions. """ - task_type = task.task_type or 'unknown' + task_type = task.task_type or "unknown" exec_start = time.monotonic() - task_status = 'completed' + 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 = { TOKEN_ID: task.token, MODEL_ID: task.model_id, - 'twinkle.task_type': task_type, - 'twinkle.deployment': self._deployment_name or 'unknown', - 'twinkle.queue_key': queue_key, + "twinkle.task_type": task_type, + "twinkle.deployment": self._deployment_name or "unknown", + "twinkle.queue_key": queue_key, } handler_attrs = {TOKEN_ID: task.token, MODEL_ID: task.model_id} - handler_span_name = f'{self._deployment_name or "deployment"}.{task_type}' + 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): - task_status = 'cancelled' - logger.info(f'[ComputeWorker] Task {task.request_id} became terminal before start; skipping.') + if not await self._state.try_start_future(task.request_id): + task_status = "cancelled" + 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("task_queue.execute", attrs=queue_attrs): + logger.debug( + f"[ComputeWorker] Task {task.request_id} executing, 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}') + logger.info(f"[ComputeWorker] Task {task.request_id} completed in {exec_time:.2f}s, type={task_type}") await self._state.store_future_status( task.request_id, TaskStatus.COMPLETED.value, @@ -413,85 +506,87 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: queue_state=QueueState.ACTIVE.value, ) except asyncio.CancelledError: - task_status = 'cancelled' + task_status = "cancelled" exec_time = time.monotonic() - exec_start if not await self._is_record_terminal(task.request_id): await self._store_task_failed( task, - 'Task stopped during replica shutdown.' if self._stopping else 'Backend generation was cancelled.', + "Task stopped during replica shutdown." if self._stopping else "Backend generation was cancelled.", QueueState.ACTIVE.value, - reason_code='orphaned_replica' if self._stopping else 'backend_cancelled') + reason_code="orphaned_replica" if self._stopping else "backend_cancelled", + ) raise except BACKEND_TIMEOUT_EXCEPTIONS: - task_status = 'timeout' + task_status = "timeout" exec_time = time.monotonic() - exec_start - error = (f'Backend call timed out (bound {self._config.effective_execution_timeout}s), ' - f'actual execution time: {exec_time:.2f}s') - logger.error(f'[ComputeWorker] Task {task.request_id} TIMEOUT after {exec_time:.2f}s, ' - f'type={task_type}, queue_key={queue_key}') + error = ( + f"Backend call timed out (bound {self._config.effective_execution_timeout}s), " + f"actual execution time: {exec_time:.2f}s" + ) + logger.error( + f"[ComputeWorker] Task {task.request_id} TIMEOUT after {exec_time:.2f}s, " + f"type={task_type}, queue_key={queue_key}" + ) # asyncio.TimeoutError and Ray_Get_Timeout are 504/Server. - await self._store_task_failed(task, error, QueueState.ACTIVE.value, reason_code='execution_timeout') + await self._store_task_failed(task, error, QueueState.ACTIVE.value, reason_code="execution_timeout") # Probe actor liveness after a timeout so an operator learns the replica's # state without waiting for a second request to also time out. if self._on_backend_timeout is not None: try: 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', - ) + logger.error(f"[ComputeWorker] backend-timeout probe failed:\n{traceback.format_exc(limit=3)}") except BackendBusyError as exc: - task_status = 'failed' + task_status = "failed" exec_time = time.monotonic() - exec_start error = str(exc) - logger.error(f'[ComputeWorker] Task {task.request_id} REFUSED (admission gate held) after ' - f'{exec_time:.2f}s, type={task_type}, queue_key={queue_key}') + logger.error( + f"[ComputeWorker] Task {task.request_id} REFUSED (admission gate held) after " + f"{exec_time:.2f}s, type={task_type}, queue_key={queue_key}" + ) # Gate held by a leaked timed-out call -> 503/Server. - await self._store_task_failed(task, error, QueueState.ACTIVE.value, reason_code='backend_gate_unavailable') + await self._store_task_failed(task, error, QueueState.ACTIVE.value, reason_code="backend_gate_unavailable") except TwinkleServerError as exc: # A typed server error carries its own status + category (e.g. # ResourceNotFoundError = 404/User from a deferred # assert_resource_exists). Honour them instead of collapsing every # such failure to 500/Server. Only Server-category errors keep a # traceback; a User rejection does not. - task_status = 'failed' + task_status = "failed" exec_time = time.monotonic() - exec_start is_server = exc.category is ErrorCategory.Server if is_server: - logger.error(f'[ComputeWorker] Task {task.request_id} FAILED after {exec_time:.2f}s, ' - f'type={task_type}:\n{traceback.format_exc(limit=3)}') + 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( task, - f'{type(exc).__name__}: {exc}', + f"{type(exc).__name__}: {exc}", QueueState.ACTIVE.value, reason_code=_reason_code_for_server_error(exc), attribution=exc.category.value, traceback_text=traceback.format_exc() if is_server else None, ) except Exception as exc: - task_status = 'failed' + task_status = "failed" 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}' - logger.error(f'[ComputeWorker] Task {task.request_id} FAILED after {exec_time:.2f}s, ' - f'type={task_type}:\n{traceback.format_exc(limit=3)}') + 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( task, error, QueueState.ACTIVE.value, - reason_code='internal_error', + reason_code="internal_error", 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) @@ -518,19 +613,12 @@ async def _try_run_one(self) -> bool: # Record queue-wait metrics and check queue-level timeout queue_wait = time.monotonic() - task.created_at - self._record_queue_metrics(task.task_type or 'unknown', queue_wait) + self._record_queue_metrics(task.task_type or "unknown", queue_wait) if queue_wait > self._config.queue_timeout: 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,13 +627,10 @@ 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: + record = await self._state.get_future(request_id) + if record is None: return False - if not record: - return False - return record.get('status') in (TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, TaskStatus.CANCELLED.value) + return record.get("status") in (TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, TaskStatus.CANCELLED.value) # ------------------------------------------------------------------ # Main worker loop @@ -553,8 +638,10 @@ async def _is_record_terminal(self, request_id: str) -> bool: async def _worker_loop(self) -> None: """Main worker loop: wait for work, run one task, repeat.""" - logger.info(f'[ComputeWorker] Started, queue_timeout={self._config.queue_timeout}, ' - f'execution_timeout={self._config.execution_timeout}') + logger.info( + f"[ComputeWorker] Started, queue_timeout={self._config.queue_timeout}, " + f"execution_timeout={self._config.execution_timeout}" + ) # Python 3.10/3.11 wait_for can swallow cancellation when its child # completes simultaneously. Do not wait for more work after stop(). while not self._stopping: @@ -567,6 +654,6 @@ async def _worker_loop(self) -> None: except asyncio.CancelledError: break except Exception: - logger.error(f'[ComputeWorker] Unexpected error:\n{traceback.format_exc(limit=3)}') + logger.error(f"[ComputeWorker] Unexpected error:\n{traceback.format_exc(limit=3)}") continue - logger.info('[ComputeWorker] Worker stopped') + logger.info("[ComputeWorker] Worker stopped") diff --git a/src/twinkle/server/telemetry/metrics.py b/src/twinkle/server/telemetry/metrics.py index a182d57b4..0a5233de8 100644 --- a/src/twinkle/server/telemetry/metrics.py +++ b/src/twinkle/server/telemetry/metrics.py @@ -16,10 +16,12 @@ """ from __future__ import annotations +from dataclasses import dataclass +from typing import Any from pydantic import BaseModel, ConfigDict -from typing import Any +from twinkle.server.exceptions import ResourceReleaseError from .provider import get_meter try: @@ -28,10 +30,10 @@ Observation = None # type: ignore[assignment, misc] _RESOURCE_GAUGES: tuple[tuple[str, str, str], ...] = ( - ('active_sessions', 'twinkle.sessions.active', 'Number of active client sessions'), - ('active_models', 'twinkle.models.active', 'Number of registered models'), - ('active_sampling_sessions', 'twinkle.sampling_sessions.active', 'Number of active sampling sessions'), - ('active_futures', 'twinkle.futures.active', 'Number of pending futures/tasks'), + ("active_sessions", "twinkle.sessions.active", "Number of active client sessions"), + ("active_models", "twinkle.models.active", "Number of registered models"), + ("active_sampling_sessions", "twinkle.sampling_sessions.active", "Number of active sampling sessions"), + ("active_futures", "twinkle.futures.active", "Number of pending futures/tasks"), ) # --------------------------------------------------------------------------- @@ -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).""" @@ -136,45 +165,98 @@ class MetricsRegistry: _instance: MetricsRegistry | None = None def __init__(self) -> None: - meter = get_meter('twinkle-server') + meter = get_meter("twinkle-server") # === HTTP Requests === self.requests_total = meter.create_counter( - 'twinkle.http.requests.total', - description='Total HTTP requests received', + "twinkle.http.requests.total", + description="Total HTTP requests received", ) self.request_duration_seconds = meter.create_histogram( - 'twinkle.http.request.duration_seconds', - description='HTTP request duration in seconds', - unit='s', + "twinkle.http.request.duration_seconds", + description="HTTP request duration in seconds", + unit="s", ) # === Task Queue === self.queue_depth = meter.create_up_down_counter( - 'twinkle.queue.depth', - description='Current task queue depth', + "twinkle.queue.depth", + description="Current task queue depth", ) self.task_execution_seconds = meter.create_histogram( - 'twinkle.task.execution_seconds', - description='Task execution duration in seconds', - unit='s', + "twinkle.task.execution_seconds", + description="Task execution duration in seconds", + unit="s", ) self.task_wait_seconds = meter.create_histogram( - 'twinkle.task.wait_seconds', - description='Task wait time in queue before execution', - unit='s', + "twinkle.task.wait_seconds", + description="Task wait time in queue before execution", + unit="s", ) self.rate_limit_rejections = meter.create_counter( - 'twinkle.rate_limit.rejections.total', - description='Total requests rejected by rate limiter', + "twinkle.rate_limit.rejections.total", + description="Total requests rejected by rate limiter", ) self.tasks_total = meter.create_counter( - 'twinkle.tasks.total', - description='Total task completions, partitioned by status', + "twinkle.tasks.total", + description="Total task completions, partitioned by status", ) self.rate_limiter_active_tokens = meter.create_up_down_counter( - 'twinkle.rate_limiter.active_tokens', - description='Number of tokens currently tracked by the rate limiter', + "twinkle.rate_limiter.active_tokens", + 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) === @@ -192,6 +274,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 +311,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 +379,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, strict=False)) + 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, strict=False)) + 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..be42bb2e7 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,14 +27,24 @@ 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 # error at store creation, not as silent misbehaviour. -_COORD_PORT_OFFSET = int(os.environ.get('TWINKLE_COORD_PORT_OFFSET', 10091)) -_COORD_TIMEOUT = timedelta(seconds=int(os.environ.get('TWINKLE_COORD_TIMEOUT', 21600))) +_COORD_PORT_OFFSET = int(os.environ.get("TWINKLE_COORD_PORT_OFFSET", 10091)) +_COORD_TIMEOUT = timedelta(seconds=int(os.environ.get("TWINKLE_COORD_TIMEOUT", 21600))) _store = None _store_failed = False # One sequence per lock name: the same key is locked repeatedly over a run (encode, then check, @@ -51,12 +63,12 @@ def _sanitize_lock_name(name: str) -> str: Windows does not allow : / \ * ? " < > | in file names. """ # Replace problematic characters with underscores - return re.sub(r'[:/\\*?"<>|]', '_', name) + return re.sub(r'[:/\\*?"<>|]', "_", name) def acquire_lock(lock: FileLock, blocking: bool): try: - if 'blocking' in inspect.signature(lock.acquire).parameters: + if "blocking" in inspect.signature(lock.acquire).parameters: lock.acquire(blocking=blocking) else: lock.acquire(timeout=(0 if not blocking else None)) @@ -71,10 +83,10 @@ def release_lock(lock: FileLock): def _get_session_token() -> str: """Return a stable token shared by all ranks in the same training run.""" - return os.environ.get('TWINKLE_SESSION_ID') or str(os.getppid()) + return os.environ.get("TWINKLE_SESSION_ID") or str(os.getppid()) -def try_claim_once(key: str, *, payload: str = '', namespace: str = 'claim') -> bool: +def try_claim_once(key: str, *, payload: str = "", namespace: str = "claim") -> bool: """Atomically claim a one-shot slot identified by ``key`` (single-winner). Stale claims left by a prior session (identified by a session token stored @@ -90,9 +102,9 @@ def try_claim_once(key: str, *, payload: str = '', namespace: str = 'claim') -> """ try: session = _get_session_token() - digest = hashlib.md5(_sanitize_lock_name(key).encode('utf-8')).hexdigest()[:16] + digest = hashlib.md5(_sanitize_lock_name(key).encode("utf-8")).hexdigest()[:16] os.makedirs(_LOCK_DIR, exist_ok=True) - path = os.path.join(_LOCK_DIR, f'{namespace}_{digest}.once') + path = os.path.join(_LOCK_DIR, f"{namespace}_{digest}.once") return _try_create_claim(path, session, payload) except Exception: return True @@ -104,13 +116,13 @@ def _try_create_claim(path: str, session: str, payload: str) -> bool: try: fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644) try: - os.write(fd, f'{session}\n{payload}'.encode()) + os.write(fd, f"{session}\n{payload}".encode()) finally: os.close(fd) return True except FileExistsError: try: - with open(path, encoding='utf-8') as f: + with open(path, encoding="utf-8") as f: stored = f.readline().strip() if stored == session: return False # same session, genuine loser @@ -131,16 +143,17 @@ class PosixFileLock: def __init__(self, path: str): import fcntl + self._path = path self._fcntl = fcntl - self._fd = open(path, 'w') + self._fd = open(path, "w") self._pid = os.getpid() def _ensure_fd(self): # After fork, child must reopen so it doesn't share parent's fd state. pid = os.getpid() if pid != self._pid: - self._fd = open(self._path, 'w') + self._fd = open(self._path, "w") self._pid = pid def acquire(self): @@ -161,19 +174,20 @@ def __exit__(self, *exc): self.release() def __getstate__(self): - return {'_path': self._path} + return {"_path": self._path} def __setstate__(self, state): import fcntl - self._path = state['_path'] + + self._path = state["_path"] self._fcntl = fcntl - self._fd = open(self._path, 'w') + self._fd = open(self._path, "w") self._pid = os.getpid() def _coord_ranks(): """Rank/world-size of the group sharing one coordination store.""" - return int(os.environ.get('RANK', 0)), int(os.environ.get('WORLD_SIZE', 1)) + return int(os.environ.get("RANK", 0)), int(os.environ.get("WORLD_SIZE", 1)) def _is_local_master() -> bool: @@ -183,17 +197,17 @@ def _is_local_master() -> bool: rank reports True and the node tier collapses into 'global master first, then everyone' -- which is the right degradation when the node layout is not knowable from the env. """ - return int(os.environ.get('LOCAL_RANK', 0)) == 0 + return int(os.environ.get("LOCAL_RANK", 0)) == 0 def _node_index() -> int: """Index of this rank's node. torchrun exports GROUP_RANK; derive it otherwise.""" ## GROUP_RANK和LOCAL_WORLD_SIZE谁设置的？会不会有无法使用的可能性，或者多机条件下 - group_rank = os.environ.get('GROUP_RANK') + group_rank = os.environ.get("GROUP_RANK") if group_rank is not None: return int(group_rank) rank, _ = _coord_ranks() - return rank // (int(os.environ.get('LOCAL_WORLD_SIZE', 1)) or 1) + return rank // (int(os.environ.get("LOCAL_WORLD_SIZE", 1)) or 1) def _get_store(): @@ -211,16 +225,18 @@ def _get_store(): rank, _ = _coord_ranks() try: from torch.distributed import TCPStore + # world_size=None keeps this off the startup path: nobody blocks waiting for peers, and a # client simply retries until rank 0 has its server up. _store = TCPStore( - os.environ.get('MASTER_ADDR', '127.0.0.1'), + os.environ.get("MASTER_ADDR", "127.0.0.1"), ## _COORD_PORT_OFFSET 干啥用的 - int(os.environ.get('MASTER_PORT', 29500)) + _COORD_PORT_OFFSET, + int(os.environ.get("MASTER_PORT", 29500)) + _COORD_PORT_OFFSET, world_size=None, is_master=(rank == 0), timeout=_COORD_TIMEOUT, - wait_for_workers=False) + wait_for_workers=False, + ) except Exception: # Port taken, torch too old, no network -- degrade to the file lock rather than fail a run # over a coordination detail. @@ -231,13 +247,13 @@ def _get_store(): def _use_store() -> bool: """Whether ranks can be ordered through the store instead of the file lock.""" _, world_size = _coord_ranks() - return world_size > 1 and 'MASTER_ADDR' in os.environ and _get_store() is not None + return world_size > 1 and "MASTER_ADDR" in os.environ and _get_store() is not None def _wait_for(store, flag): store.wait([flag], _COORD_TIMEOUT) - if store.get(flag) != b'1': - raise LockPeerError(f'The peer responsible for {flag} failed; refusing to read after it') + if store.get(flag) != b"1": + raise LockPeerError(f"The peer responsible for {flag} failed; refusing to read after it") @contextmanager @@ -249,11 +265,11 @@ def _ordered_by_store(key: str, sticky: bool): if sticky: # The key already identifies the result (a repo id), and the work is idempotent, so one # flag serves the whole run: a rank arriving after the flag is set just proceeds. - prefix = f'twinkle/lock/{name}' + prefix = f"twinkle/lock/{name}" else: - prefix = f'twinkle/lock/{name}/{_seq[name]}' + prefix = f"twinkle/lock/{name}/{_seq[name]}" _seq[name] += 1 - global_flag, node_flag = f'{prefix}/global', f'{prefix}/node{_node_index()}' + global_flag, node_flag = f"{prefix}/global", f"{prefix}/node{_node_index()}" rank, _ = _coord_ranks() is_global_master, is_local_master = rank == 0, _is_local_master() @@ -265,7 +281,7 @@ def _ordered_by_store(key: str, sticky: bool): if is_local_master: # Pass the failure down, instead of leaving our own node waiting on a flag that # will never be set. - store.set(node_flag, b'0') + store.set(node_flag, b"0") raise ok = False @@ -275,8 +291,8 @@ def _ordered_by_store(key: str, sticky: bool): finally: if is_local_master: # Publish even on failure: a crashed writer must not hang the readers. - for flag in ([global_flag, node_flag] if is_global_master else [node_flag]): - store.set(flag, b'1' if ok else b'0') + for flag in [global_flag, node_flag] if is_global_master else [node_flag]: + store.set(flag, b"1" if ok else b"0") @contextmanager @@ -310,11 +326,10 @@ def processing_lock(lock_file: str, sticky: bool = False): return lock_name = _sanitize_lock_name(lock_file) - # Imported here, not at module scope: ``datasets`` is only needed when this fallback path - # 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 - lock: FileLock = FileLock(os.path.join(_LOCK_DIR, f'{lock_name}.lock')) # noqa + + 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): try: 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..8cfa5e8a7 100644 --- a/src/twinkle_client/auto/agent/server_tools.py +++ b/src/twinkle_client/auto/agent/server_tools.py @@ -1,11 +1,12 @@ # Copyright (c) Twinkle Contributors. All rights reserved. """Private server lifecycle and cluster tools used by ``ToolExecutor``.""" -from __future__ import annotations +from __future__ import annotations import asyncio -import json import os +import json + class _ServerTools: """Server-side operations mixed into ``ToolExecutor``. @@ -23,7 +24,7 @@ async def _check_server_health(self, url: str) -> bool: def _probe(): try: - req = urllib.request.Request(f'{url}/api/v1/healthz', method='GET') + req = urllib.request.Request(f"{url}/api/v1/healthz", method="GET") urllib.request.urlopen(req, timeout=3) return True except (urllib.error.URLError, OSError): @@ -41,16 +42,16 @@ async def _tool_start_server( model_id: str, train_gpus: int | None = None, port: int = 8000, - backend: str = 'transformers', + backend: str = "transformers", samplers: list[dict] | None = None, ) -> dict: """Start Ray cluster + Twinkle Server. Idempotent. Supports multi-model.""" - server_url = self._server_url or os.environ.get('TWINKLE_SERVER_URL') or f'http://localhost:{port}' + server_url = self._server_url or os.environ.get("TWINKLE_SERVER_URL") or f"http://localhost:{port}" # Idempotent: skip if already running if await self._check_server_health(server_url): self._server_url = server_url - return {'status': 'already_running', 'server_url': server_url} + return {"status": "already_running", "server_url": server_url} def _start(): sampler_list = samplers or [] @@ -58,12 +59,12 @@ def _start(): # Step 1: Detect hardware & compute GPU partition total_hw_gpus = self._detect_gpu_count() if total_hw_gpus == 0: - return {'status': 'error', 'error': 'No GPUs detected. Cannot start training server.'} + return {"status": "error", "error": "No GPUs detected. Cannot start training server."} alloc = self._compute_gpu_allocation(sampler_list, train_gpus, total_hw_gpus) - if 'error' in alloc: - return {'status': 'error', 'error': alloc['error']} - t_gpus, sampler_gpu_total = alloc['train_gpus'], alloc['sampler_gpus'] + if "error" in alloc: + return {"status": "error", "error": alloc["error"]} + t_gpus, sampler_gpu_total = alloc["train_gpus"], alloc["sampler_gpus"] # Step 2: Generate server_config.yaml config_path = self._generate_server_config( @@ -77,12 +78,12 @@ def _start(): # Step 3: Start Ray cluster (multi-node GPU partitioning) ray_err = self._start_ray_cluster(t_gpus, sampler_gpu_total) if ray_err: - return {'status': 'error', 'error': ray_err} + return {"status": "error", "error": ray_err} # Step 4: Launch Twinkle Server process proc, log_path, err = self._launch_server_process(config_path) if err: - return {'status': 'error', 'error': err} + return {"status": "error", "error": err} # Step 5: Wait for readiness (healthz + sampler engine) return self._wait_server_ready( @@ -97,7 +98,7 @@ def _start(): ) result = await asyncio.get_event_loop().run_in_executor(None, _start) - if result.get('status') in ('started', 'already_running'): + if result.get("status") in ("started", "already_running"): self._server_url = server_url return result @@ -105,15 +106,16 @@ def _start(): def _detect_gpu_count() -> int: """Detect total hardware GPU count via nvidia-smi.""" import subprocess as _sp + try: r = _sp.run( - ['nvidia-smi', '--query-gpu=index', '--format=csv,noheader'], + ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], capture_output=True, text=True, timeout=10, ) if r.returncode == 0: - return len([ln for ln in r.stdout.strip().split('\n') if ln.strip()]) + return len([ln for ln in r.stdout.strip().split("\n") if ln.strip()]) except (FileNotFoundError, OSError): pass return 0 @@ -127,8 +129,8 @@ def _compute_gpu_allocation( """Compute GPU partition: {train_gpus, sampler_gpus} or {error}.""" sampler_gpu_total = 0 for s in sampler_list: - s_tp = s.get('tp', 1) - s_dp, s_gpus = s.get('dp'), s.get('gpus') + s_tp = s.get("tp", 1) + s_dp, s_gpus = s.get("dp"), s.get("gpus") if s_gpus is not None: sampler_gpu_total += s_gpus elif s_dp is not None: @@ -140,10 +142,12 @@ def _compute_gpu_allocation( needed = t_gpus + sampler_gpu_total if needed > total_hw_gpus: return { - 'error': (f'Requested {needed} GPUs (train={t_gpus}, samplers={sampler_gpu_total}) ' - f'but only {total_hw_gpus} available.'), + "error": ( + f"Requested {needed} GPUs (train={t_gpus}, samplers={sampler_gpu_total}) " + f"but only {total_hw_gpus} available." + ), } - return {'train_gpus': t_gpus, 'sampler_gpus': sampler_gpu_total} + return {"train_gpus": t_gpus, "sampler_gpus": sampler_gpu_total} @staticmethod def _start_ray_cluster(train_gpus: int, sampler_gpus: int) -> str | None: @@ -162,10 +166,10 @@ def _start_ray_cluster(train_gpus: int, sampler_gpus: int) -> str | None: import tempfile from pathlib import Path - _sp.run(['ray', 'stop', '--force'], capture_output=True, timeout=15) + _sp.run(["ray", "stop", "--force"], capture_output=True, timeout=15) # Create unique temp dirs so each `ray start` spawns a separate raylet - ray_base = Path(tempfile.gettempdir()) / 'twinkle_ray' + ray_base = Path(tempfile.gettempdir()) / "twinkle_ray" ray_base.mkdir(parents=True, exist_ok=True) def _ray_node( @@ -173,37 +177,37 @@ def _ray_node( num_gpus: int, *, head: bool = False, - node_name: str = 'worker', + node_name: str = "worker", ) -> str | None: env = os.environ.copy() - env['CUDA_VISIBLE_DEVICES'] = devices + env["CUDA_VISIBLE_DEVICES"] = devices temp_dir = str(ray_base / node_name) - cmd = ['ray', 'start', f'--temp-dir={temp_dir}'] + cmd = ["ray", "start", f"--temp-dir={temp_dir}"] if head: - cmd += ['--head', '--port=6379', '--disable-usage-stats', '--include-dashboard=false'] + cmd += ["--head", "--port=6379", "--disable-usage-stats", "--include-dashboard=false"] else: - cmd += ['--address=127.0.0.1:6379'] - cmd.append(f'--num-gpus={num_gpus}') + cmd += ["--address=127.0.0.1:6379"] + cmd.append(f"--num-gpus={num_gpus}") r = _sp.run(cmd, capture_output=True, text=True, timeout=30, env=env) - if r.returncode != 0 and 'already' not in r.stderr.lower(): + if r.returncode != 0 and "already" not in r.stderr.lower(): return r.stderr.strip() return None # Head node — training model GPUs - model_devices = ','.join(str(i) for i in range(train_gpus)) - err = _ray_node(model_devices, train_gpus, head=True, node_name='head') + model_devices = ",".join(str(i) for i in range(train_gpus)) + err = _ray_node(model_devices, train_gpus, head=True, node_name="head") if err: - return f'Ray head start failed: {err}' + return f"Ray head start failed: {err}" # GPU Worker node — sampler GPUs if sampler_gpus > 0: - sampler_devices = ','.join(str(i) for i in range(train_gpus, train_gpus + sampler_gpus)) - err = _ray_node(sampler_devices, sampler_gpus, node_name='gpu_worker') + sampler_devices = ",".join(str(i) for i in range(train_gpus, train_gpus + sampler_gpus)) + err = _ray_node(sampler_devices, sampler_gpus, node_name="gpu_worker") if err: - return f'Ray GPU worker start failed: {err}' + return f"Ray GPU worker start failed: {err}" # CPU Worker node — processor (no GPU) - _ray_node('', 0, node_name='cpu_worker') + _ray_node("", 0, node_name="cpu_worker") return None @staticmethod @@ -215,12 +219,12 @@ def _launch_server_process(config_path: str) -> tuple: import subprocess as _sp from pathlib import Path - log_dir = Path.home() / '.cache' / 'twinkle' + log_dir = Path.home() / ".cache" / "twinkle" log_dir.mkdir(parents=True, exist_ok=True) - log_path = str(log_dir / 'server.log') - log_file = open(log_path, 'w') + log_path = str(log_dir / "server.log") + log_file = open(log_path, "w") - cmd = ['python', '-m', 'twinkle.server', 'launch', '--config', config_path] + cmd = ["python", "-m", "twinkle.server", "launch", "--config", config_path] try: proc = _sp.Popen( cmd, @@ -230,7 +234,7 @@ def _launch_server_process(config_path: str) -> tuple: ) except OSError as e: log_file.close() - return None, log_path, f'Failed to start Twinkle server: {e}' + return None, log_path, f"Failed to start Twinkle server: {e}" return proc, log_path, None @staticmethod @@ -250,74 +254,72 @@ def _wait_server_ready( import urllib.request timeout_s = 120 if sampler_list else 60 - needed = t_gpus + sum(s.get('gpus') or (s.get('tp', 1) * s.get('dp', 1)) for s in sampler_list) + needed = t_gpus + sum(s.get("gpus") or (s.get("tp", 1) * s.get("dp", 1)) for s in sampler_list) for _ in range(timeout_s): time.sleep(1) if proc.poll() is not None: # Server died — read log tail to diagnose log_tail = _ServerTools._read_log_tail(log_path, max_chars=2000) - error_msg = (f'Server exited immediately (code={proc.returncode}). ' - f'Model: {model_id}, GPUs: {t_gpus}, Samplers: {len(sampler_list)}.\n' - f'--- server.log tail ---\n{log_tail}') + error_msg = ( + f"Server exited immediately (code={proc.returncode}). " + f"Model: {model_id}, GPUs: {t_gpus}, Samplers: {len(sampler_list)}.\n" + f"--- server.log tail ---\n{log_tail}" + ) return { - 'status': - 'error', - 'error': - error_msg, - 'log_path': - log_path, - 'hint': - 'Check required packages (pip install -e ".[client,server,rl]"; ' - 'add the megatron extra for the Megatron backend).', + "status": "error", + "error": error_msg, + "log_path": log_path, + "hint": 'Check required packages (pip install -e ".[client,server,rl]"; ' + "add the megatron extra for the Megatron backend).", } try: - urllib.request.urlopen(f'{server_url}/api/v1/healthz', timeout=2) + urllib.request.urlopen(f"{server_url}/api/v1/healthz", timeout=2) except (OSError, Exception): continue # healthz OK — additionally wait for sampler vLLM engines if sampler_list and not _ServerTools._probe_sampler_ready(server_url, sampler_list, model_id): return { - 'status': 'started', - 'warning': 'Server is up but sampler may still be loading.', - 'server_url': server_url, - 'server_pid': proc.pid, - 'model_id': model_id, - 'log_path': log_path, + "status": "started", + "warning": "Server is up but sampler may still be loading.", + "server_url": server_url, + "server_pid": proc.pid, + "model_id": model_id, + "log_path": log_path, } return { - 'status': 'started', - 'server_url': server_url, - 'server_pid': proc.pid, - 'model_id': model_id, - 'train_gpus': t_gpus, - 'backend': backend, - 'samplers': [s.get('model_id') for s in sampler_list], - 'total_gpus_used': needed, - 'config_path': config_path, - 'log_path': log_path, + "status": "started", + "server_url": server_url, + "server_pid": proc.pid, + "model_id": model_id, + "train_gpus": t_gpus, + "backend": backend, + "samplers": [s.get("model_id") for s in sampler_list], + "total_gpus_used": needed, + "config_path": config_path, + "log_path": log_path, } return { - 'status': 'timeout', - 'error': 'Health check did not pass within timeout. Models may still be loading.', - 'server_pid': proc.pid, - 'log_path': log_path, + "status": "timeout", + "error": "Health check did not pass within timeout. Models may still be loading.", + "server_pid": proc.pid, + "log_path": log_path, } @staticmethod def _read_log_tail(log_path: str, max_chars: int = 2000) -> str: """Read the tail of a log file for error diagnosis.""" try: - with open(log_path, errors='replace') as f: + with open(log_path, errors="replace") as f: content = f.read() if len(content) <= max_chars: return content.strip() return content[-max_chars:].strip() except OSError: - return '(could not read log file)' + return "(could not read log file)" @staticmethod def _probe_sampler_ready(server_url: str, sampler_list: list[dict], fallback_model_id: str) -> bool: @@ -326,16 +328,16 @@ def _probe_sampler_ready(server_url: str, sampler_list: list[dict], fallback_mod import urllib.error import urllib.request - s_mid = sampler_list[0].get('model_id', fallback_model_id) - probe_url = f'{server_url}/api/v1/sampler/{s_mid}/twinkle/create' + s_mid = sampler_list[0].get("model_id", fallback_model_id) + probe_url = f"{server_url}/api/v1/sampler/{s_mid}/twinkle/create" for _ in range(90): try: req = urllib.request.Request( probe_url, - method='POST', - data=b'{}', - headers={'Content-Type': 'application/json'}, + method="POST", + data=b"{}", + headers={"Content-Type": "application/json"}, ) urllib.request.urlopen(req, timeout=5) return True # non-error response = ready @@ -352,7 +354,7 @@ def _generate_server_config( model_id: str, train_gpus: int, port: int = 8000, - backend: str = 'transformers', + backend: str = "transformers", samplers: list[dict] | None = None, ) -> str: """Generate a server_config.yaml from template and return its path. @@ -362,52 +364,52 @@ def _generate_server_config( - N sampler/teacher models (for RL/OPD) - 1 processor service """ - import yaml from pathlib import Path + import yaml + sampler_list = samplers or [] - if backend not in ('transformers', 'megatron'): - raise ValueError(f'Unsupported training backend: {backend!r}') + 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: - return mid.split('/')[-1] if '/' in mid else mid + return mid.split("/")[-1] if "/" in mid else mid model_short = _short(model_id) # Collect all model IDs for supported_models - all_model_ids = [model_id] + [s['model_id'] for s in sampler_list] + all_model_ids = [model_id] + [s["model_id"] for s in sampler_list] # === Build applications list === applications = [] # 1. API Gateway - applications.append({ - 'name': - 'server', - 'route_prefix': - '/api/v1', - 'import_path': - 'server', - 'args': { - 'server_config': { - 'per_token_model_limit': 3 - }, - 'supported_models': all_model_ids, - }, - 'deployments': [{ - 'name': 'GatewayServer', - 'max_ongoing_requests': 50, - 'autoscaling_config': { - 'min_replicas': 1, - 'max_replicas': 1, - 'target_ongoing_requests': 128, - }, - 'ray_actor_options': { - 'num_cpus': 0.1 + applications.append( + { + "name": "server", + "route_prefix": "/api/v1", + "import_path": "server", + "args": { + "server_config": {"per_token_model_limit": 3}, + "supported_models": all_model_ids, }, - }], - }) + "deployments": [ + { + "name": "GatewayServer", + "max_ongoing_requests": 50, + "autoscaling_config": { + "min_replicas": 1, + "max_replicas": 1, + "target_ongoing_requests": 128, + }, + "ray_actor_options": {"num_cpus": 0.1}, + } + ], + } + ) # 2. Build GPU-requiring applications (model + samplers), # then sort by GPU count DESCENDING before appending. @@ -416,81 +418,84 @@ def _short(mid: str) -> str: gpu_apps: list[tuple[int, dict]] = [] # (gpu_count, app_config) # 2a. Training model worker (student) - gpu_apps.append(( - train_gpus, - { - 'name': - f'models-{model_short}', - 'route_prefix': - f'/api/v1/model/{model_id}', - 'import_path': - 'model', - 'args': { - 'backend': backend, - 'model_id': f'ms://{model_id}', - 'max_length': 500000, # total tokens per forward pass (must match max_input_tokens) - 'nproc_per_node': train_gpus, - 'device_group': { - 'name': 'model', - 'ranks': train_gpus, - 'device_type': 'cuda', - }, - 'device_mesh': { - 'device_type': 'cuda', - 'dp_size': train_gpus, - }, - 'queue_config': { - 'rps_limit': 100, - 'tps_limit': 100000, - 'max_input_tokens': 500000, - }, - 'adapter_config': { - 'adapter_timeout': 600, - }, - }, - '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' - }, + gpu_apps.append( + ( + train_gpus, + { + "name": f"models-{model_short}", + "route_prefix": f"/api/v1/model/{model_id}", + "import_path": "model", + "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": { + "name": "model", + "ranks": train_gpus, + "device_type": "cuda", + }, + "device_mesh": { + "device_type": "cuda", + "dp_size": train_gpus, + }, + "queue_config": { + "rps_limit": 100, + "tps_limit": 100000, + "max_input_tokens": 500000, + }, + "adapter_config": { + "adapter_timeout": 600, }, }, - }], - })) + "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"}, + }, + }, + } + ], + }, + ) + ) # 2b. Sampler/teacher models sampler_name_count: dict[str, int] = {} for sampler_cfg in sampler_list: - s_model_id = sampler_cfg['model_id'] + s_model_id = sampler_cfg["model_id"] s_short = _short(s_model_id) # Deduplicate names when multiple samplers share the same short name sampler_name_count[s_short] = sampler_name_count.get(s_short, 0) + 1 if sampler_name_count[s_short] > 1: - s_name = f'sampler-{s_short}-{sampler_name_count[s_short]}' + s_name = f"sampler-{s_short}-{sampler_name_count[s_short]}" else: - s_name = f'sampler-{s_short}' + s_name = f"sampler-{s_short}" - s_engine = sampler_cfg.get('engine', 'vllm') - if s_engine != 'vllm': - raise ValueError(f'Unsupported sampler engine: {s_engine!r}; supported engine: vllm') - s_max_len = sampler_cfg.get('max_model_len', 16000) + s_engine = sampler_cfg.get("engine", "vllm") + if s_engine != "vllm": + raise ValueError(f"Unsupported sampler engine: {s_engine!r}; supported engine: vllm") + s_max_len = sampler_cfg.get("max_model_len", 16000) # Compute tp / dp / total GPUs: # tp = tensor parallelism (GPUs per vLLM process, for large models) # dp = data parallelism (number of independent inference replicas) # total GPUs = tp * dp - s_tp = sampler_cfg.get('tp', 1) - s_dp = sampler_cfg.get('dp', None) - s_gpus = sampler_cfg.get('gpus', None) + s_tp = sampler_cfg.get("tp", 1) + s_dp = sampler_cfg.get("dp", None) + s_gpus = sampler_cfg.get("gpus", None) if s_dp is not None and s_gpus is not None: # Both specified: validate consistency @@ -511,63 +516,60 @@ def _short(mid: str) -> str: # Build device_mesh: include tp_size when tp>1 so that # world_size = tp*dp and slice_dp dispatch computes correct # rank_stride for DP data sharding. - mesh_config: dict = {'device_type': 'cuda', 'dp_size': s_dp} + mesh_config: dict = {"device_type": "cuda", "dp_size": s_dp} if s_tp > 1: - mesh_config['tp_size'] = s_tp + mesh_config["tp_size"] = s_tp sampler_app: dict = { - 'name': - s_name, - 'route_prefix': - f'/api/v1/sampler/{s_model_id}', - 'import_path': - 'sampler', - 'args': { - 'model_id': f'ms://{s_model_id}', - 'nproc_per_node': s_total_gpus, - 'sampler_type': s_engine, - 'device_group': { - 'name': s_name, - 'ranks': s_total_gpus, - 'device_type': 'cuda', - 'gpus_per_worker': s_tp, + "name": s_name, + "route_prefix": f"/api/v1/sampler/{s_model_id}", + "import_path": "sampler", + "args": { + "model_id": f"ms://{s_model_id}", + "nproc_per_node": s_total_gpus, + "sampler_type": s_engine, + "device_group": { + "name": s_name, + "ranks": s_total_gpus, + "device_type": "cuda", + "gpus_per_worker": s_tp, }, - 'device_mesh': mesh_config, - 'queue_config': { - 'rps_limit': 100, - 'tps_limit': 100000, + "device_mesh": mesh_config, + "queue_config": { + "rps_limit": 100, + "tps_limit": 100000, }, }, - 'deployments': [{ - 'name': 'SamplerManagement', - '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' + "deployments": [ + { + "name": "SamplerManagement", + "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"}, }, }, - }, - }], + } + ], } # Add engine-specific args - if s_engine == 'vllm': + if s_engine == "vllm": engine_args = { - 'max_model_len': s_max_len, - 'gpu_memory_utilization': 0.85, - 'enable_lora': True, - 'logprobs_mode': 'processed_logprobs', + "max_model_len": s_max_len, + "gpu_memory_utilization": 0.85, + "enable_lora": True, + "logprobs_mode": "processed_logprobs", } # Set tensor_parallel_size when tp > 1 if s_tp > 1: - engine_args['tensor_parallel_size'] = s_tp - sampler_app['args']['engine_args'] = engine_args + engine_args["tensor_parallel_size"] = s_tp + sampler_app["args"]["engine_args"] = engine_args gpu_apps.append((s_total_gpus, sampler_app)) @@ -578,53 +580,52 @@ def _short(mid: str) -> str: applications.append(app_cfg) # 4. Processor service - applications.append({ - 'name': - 'processor', - 'route_prefix': - '/api/v1/processor', - 'import_path': - 'processor', - 'args': { - 'ncpu_proc_per_node': 2, - 'device_group': { - 'name': 'processor', - 'ranks': 2, - 'device_type': 'CPU', - }, - 'device_mesh': { - 'device_type': 'CPU', - 'dp_size': 2, - }, - }, - 'deployments': [{ - 'name': 'ProcessorManagement', - 'autoscaling_config': { - 'min_replicas': 1, - 'max_replicas': 1, - 'target_ongoing_requests': 128, - }, - 'ray_actor_options': { - 'num_cpus': 0.1 + applications.append( + { + "name": "processor", + "route_prefix": "/api/v1/processor", + "import_path": "processor", + "args": { + "ncpu_proc_per_node": 2, + "device_group": { + "name": "processor", + "ranks": 2, + "device_type": "CPU", + }, + "device_mesh": { + "device_type": "CPU", + "dp_size": 2, + }, }, - }], - }) + "deployments": [ + { + "name": "ProcessorManagement", + "autoscaling_config": { + "min_replicas": 1, + "max_replicas": 1, + "target_ongoing_requests": 128, + }, + "ray_actor_options": {"num_cpus": 0.1}, + } + ], + } + ) # === Assemble final config === config = { - 'proxy_location': 'EveryNode', - 'http_options': { - 'host': '0.0.0.0', - 'port': port, + "proxy_location": "EveryNode", + "http_options": { + "host": "0.0.0.0", + "port": port, }, - 'applications': applications, + "applications": applications, } # Write to ~/.cache/twinkle/server_config.yaml - config_dir = Path.home() / '.cache' / 'twinkle' + config_dir = Path.home() / ".cache" / "twinkle" config_dir.mkdir(parents=True, exist_ok=True) - config_path = config_dir / 'server_config.yaml' - with open(config_path, 'w') as f: + config_path = config_dir / "server_config.yaml" + with open(config_path, "w") as f: yaml.dump(config, f, default_flow_style=False, allow_unicode=True) return str(config_path) @@ -638,33 +639,33 @@ def _shutdown(): # 1. Try `serve shutdown` to cleanly stop Ray Serve deployments try: - r = _sp.run(['serve', 'shutdown', '-y'], capture_output=True, text=True, timeout=30) - results['serve_shutdown'] = 'ok' if r.returncode == 0 else r.stderr.strip() + r = _sp.run(["serve", "shutdown", "-y"], capture_output=True, text=True, timeout=30) + results["serve_shutdown"] = "ok" if r.returncode == 0 else r.stderr.strip() except (FileNotFoundError, OSError) as e: - results['serve_shutdown'] = f'skipped: {e}' + results["serve_shutdown"] = f"skipped: {e}" # 2. Kill any remaining twinkle.server processes try: - _sp.run(['pkill', '-f', 'twinkle.server'], capture_output=True, timeout=5) + _sp.run(["pkill", "-f", "twinkle.server"], capture_output=True, timeout=5) except (FileNotFoundError, OSError): pass # 3. Stop Ray cluster try: - r = _sp.run(['ray', 'stop', '--force'], capture_output=True, text=True, timeout=15) - results['ray_stop'] = 'ok' if r.returncode == 0 else r.stderr.strip() + r = _sp.run(["ray", "stop", "--force"], capture_output=True, text=True, timeout=15) + results["ray_stop"] = "ok" if r.returncode == 0 else r.stderr.strip() except (FileNotFoundError, OSError) as e: - results['ray_stop'] = f'failed: {e}' + results["ray_stop"] = f"failed: {e}" - results['status'] = 'shutdown_complete' - results['warning'] = 'All GPU model state has been released.' + results["status"] = "shutdown_complete" + results["warning"] = "All GPU model state has been released." return results return await asyncio.get_event_loop().run_in_executor(None, _shutdown) async def _tool_list_supported_models(self, base_url: str | None = None) -> dict: """Query the Twinkle server for supported models.""" - url = base_url or self._server_url or os.environ.get('TWINKLE_SERVER_URL') or 'http://localhost:8000' + url = base_url or self._server_url or os.environ.get("TWINKLE_SERVER_URL") or "http://localhost:8000" def _query(): # Use a lightweight HTTP GET instead of init_twinkle_client() which @@ -673,27 +674,27 @@ def _query(): import urllib.error import urllib.request - endpoint = f'{url}/api/v1/twinkle/get_server_capabilities' - req = urllib.request.Request(endpoint, method='GET') + endpoint = f"{url}/api/v1/twinkle/get_server_capabilities" + req = urllib.request.Request(endpoint, method="GET") resp = urllib.request.urlopen(req, timeout=10) data = json.loads(resp.read().decode()) - models = data.get('supported_models', []) + models = data.get("supported_models", []) # Each model entry may be a dict with 'model_name' or a plain string model_names = [] for m in models: if isinstance(m, dict): - model_names.append(m.get('model_name', '')) + model_names.append(m.get("model_name", "")) else: model_names.append(str(m)) return { - 'base_url': url, - 'supported_models': model_names, + "base_url": url, + "supported_models": model_names, } try: return await asyncio.get_event_loop().run_in_executor(None, _query) except Exception as e: - return {'error': f'Failed to query {url}: {e}'} + return {"error": f"Failed to query {url}: {e}"} async def _tool_get_cluster_info(self) -> dict: """Query cluster resources: try Ray first, fall back to nvidia-smi.""" @@ -702,15 +703,17 @@ def _query(): # 1. Try connecting to an existing Ray cluster ray_info = self._try_ray_cluster() if ray_info is not None: - ray_info['ray_active'] = True + ray_info["ray_active"] = True return ray_info # 2. Ray not available — fall back to nvidia-smi nvidia_info = self._try_nvidia_smi() - nvidia_info['ray_active'] = False - nvidia_info['hint'] = ('Ray cluster is not running. To use distributed training, ' - 'start Ray first: `ray start --head --num-gpus=N` or use ' - 'the server mode run.sh script.') + nvidia_info["ray_active"] = False + nvidia_info["hint"] = ( + "Ray cluster is not running. To use distributed training, " + "start Ray first: `ray start --head --num-gpus=N` or use " + "the server mode run.sh script." + ) return nvidia_info return await asyncio.get_event_loop().run_in_executor(None, _query) @@ -728,7 +731,7 @@ def _try_ray_cluster() -> dict | None: try: if not ray.is_initialized(): ray.init( - address='auto', + address="auto", ignore_reinit_error=True, _timeout_s=5, logging_level=_logging.ERROR, @@ -738,24 +741,25 @@ def _try_ray_cluster() -> dict | None: resources = ray.cluster_resources() available = ray.available_resources() nodes = ray.nodes() - gpu_total = resources.get('GPU', 0) - gpu_available = available.get('GPU', 0) + gpu_total = resources.get("GPU", 0) + gpu_available = available.get("GPU", 0) gpu_types = set() for node in nodes: - for key in node.get('Resources', {}): - if key.startswith('accelerator_type:'): - gpu_types.add(key.split(':', 1)[1]) + for key in node.get("Resources", {}): + if key.startswith("accelerator_type:"): + gpu_types.add(key.split(":", 1)[1]) return { - 'num_nodes': len([n for n in nodes if n.get('Alive')]), - 'gpu_total': int(gpu_total), - 'gpu_available': int(gpu_available), - 'gpu_types': sorted(gpu_types) if gpu_types else ['unknown'], - 'cpu_total': resources.get('CPU', 0), - 'memory_bytes': resources.get('memory', 0), + "num_nodes": len([n for n in nodes if n.get("Alive")]), + "gpu_total": int(gpu_total), + "gpu_available": int(gpu_available), + "gpu_types": sorted(gpu_types) if gpu_types else ["unknown"], + "cpu_total": resources.get("CPU", 0), + "memory_bytes": resources.get("memory", 0), } except Exception: try: import ray as _ray + if _ray.is_initialized(): _ray.shutdown() except Exception: @@ -770,43 +774,46 @@ def _try_nvidia_smi() -> dict: try: result = _sp.run( [ - 'nvidia-smi', '--query-gpu=index,name,memory.total,memory.free,utilization.gpu', - '--format=csv,noheader,nounits' + "nvidia-smi", + "--query-gpu=index,name,memory.total,memory.free,utilization.gpu", + "--format=csv,noheader,nounits", ], capture_output=True, text=True, timeout=10, ) if result.returncode != 0: - return {'error': f'nvidia-smi failed: {result.stderr.strip()}', 'gpu_total': 0} + return {"error": f"nvidia-smi failed: {result.stderr.strip()}", "gpu_total": 0} gpus = [] - for line in result.stdout.strip().split('\n'): + for line in result.stdout.strip().split("\n"): if not line.strip(): continue - parts = [p.strip() for p in line.split(',')] + parts = [p.strip() for p in line.split(",")] if len(parts) >= 5: try: - gpus.append({ - 'index': int(parts[0]), - 'name': parts[1], - 'memory_total_mb': int(parts[2]), - 'memory_free_mb': int(parts[3]), - 'utilization_pct': int(parts[4]) if parts[4].isdigit() else 0, - }) + gpus.append( + { + "index": int(parts[0]), + "name": parts[1], + "memory_total_mb": int(parts[2]), + "memory_free_mb": int(parts[3]), + "utilization_pct": int(parts[4]) if parts[4].isdigit() else 0, + } + ) except (ValueError, IndexError): # Skip lines with unparseable values (e.g. [N/A]) continue - gpu_types = sorted({g['name'] for g in gpus}) + gpu_types = sorted({g["name"] for g in gpus}) return { - 'gpu_total': len(gpus), - 'gpu_available': len([g for g in gpus if g['utilization_pct'] < 10]), - 'gpu_types': gpu_types if gpu_types else ['none'], - 'gpus': gpus, - 'source': 'nvidia-smi', + "gpu_total": len(gpus), + "gpu_available": len([g for g in gpus if g["utilization_pct"] < 10]), + "gpu_types": gpu_types if gpu_types else ["none"], + "gpus": gpus, + "source": "nvidia-smi", } except FileNotFoundError: - return {'error': 'nvidia-smi not found (no NVIDIA GPU or driver not installed)', 'gpu_total': 0} + return {"error": "nvidia-smi not found (no NVIDIA GPU or driver not installed)", "gpu_total": 0} except Exception as e: - return {'error': f'nvidia-smi query failed: {e}', 'gpu_total': 0} + return {"error": f"nvidia-smi query failed: {e}", "gpu_total": 0} 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..f428fe1a7 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -1,11 +1,11 @@ from __future__ import annotations - 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,19 +13,20 @@ 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 -logger = logging.getLogger('twinkle_client') +logger = logging.getLogger("twinkle_client") def _data_refs(inputs: DataRef | list[DataRef]) -> list[dict[str, Any]]: """Encode one or more opaque references for a DataPlane model endpoint.""" refs = [inputs] if isinstance(inputs, DataRef) else list(inputs) if not refs: - raise ValueError('at least one DataRef is required') + raise ValueError("at least one DataRef is required") if not all(isinstance(item, DataRef) for item in refs): - raise TypeError('data-plane model inputs must contain only DataRef values') + raise TypeError("data-plane model inputs must contain only DataRef values") return [item.model_dump() for item in refs] @@ -33,7 +34,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 +48,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] + 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.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,14 +86,42 @@ 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) - response = self._transport.post_model(f'{self.server_url}/{endpoint}', body) + self._ensure_created() + response = self._transport.post_model(f"{self.server_url}/{endpoint}", body) return self._await_task(response, response_cls) def _await_task(self, response, model_cls): @@ -102,6 +132,7 @@ def _await_task(self, response, model_cls): Keeps every public method's synchronous signature unchanged. """ from twinkle_client._future import resolve_response + return resolve_response(response, model_cls, transport=self._transport) def _next_seq_id(self) -> int: @@ -129,25 +160,54 @@ def add_adapter_to_model( from peft import LoraConfig config = LoraConfig(**config) - save_dir = kwargs.pop('save_dir', None) + 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,62 +215,37 @@ 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.""" return self._submit( - 'forward_backward', + "forward_backward", model_types.ForwardBackwardTaskRequest, model_types.ForwardBackwardResponse, inputs=inputs, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - seq_id=self._next_seq_id(), - **kwargs) + **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) + "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 @@ -226,14 +261,14 @@ def forward_from_data_plane( ) -> model_types.ForwardResponse: """Execute forward using rows referenced from the server DataPlane.""" return self._submit( - 'forward_from_data_plane', + "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) + **kwargs, + ) def forward_only_from_data_plane( self, @@ -247,16 +282,16 @@ def forward_only_from_data_plane( ) -> model_types.ForwardResponse | DataRef: """Execute forward-only using DataPlane rows and optionally append outputs.""" result = self._submit( - 'forward_only_from_data_plane', + "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, output_fields=output_fields or {}, - **kwargs) + **kwargs, + ) if output_ref is not None: return DataRef(**result.result) return result @@ -271,15 +306,14 @@ def forward_backward_from_data_plane( ) -> model_types.ForwardBackwardResponse: """Execute forward/backward using rows referenced from the server DataPlane.""" return self._submit( - 'forward_backward_from_data_plane', + "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) + **kwargs, + ) # ------------------------------------------------------------------ # # Optimizer / scheduler steps @@ -287,58 +321,39 @@ 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) - - def clip_grad_norm(self, - max_grad_norm: float = 1.0, - norm_type: int = 2, - **kwargs) -> model_types.ClipGradNormResponse: + self._submit("lr_step", model_types.LrStepRequest, None, **kwargs) + + def clip_grad_norm( + self, max_grad_norm: float = 1.0, norm_type: int = 2, **kwargs + ) -> model_types.ClipGradNormResponse: """Clip gradient norm.""" return self._submit( - 'clip_grad_norm', + "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) + **kwargs, + ) def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type: int = 2, **kwargs) -> None: """Clip gradient norm and execute optimizer step in one call.""" self._submit( - 'clip_grad_and_step', + "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) + **kwargs, + ) # ------------------------------------------------------------------ # # Plugin setters @@ -346,33 +361,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,54 +379,31 @@ 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.""" return self._submit( - 'calculate_metric', + "calculate_metric", model_types.CalculateMetricRequest, model_types.CalculateMetricResponse, is_training=is_training, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), - **kwargs) + **kwargs, + ) # ------------------------------------------------------------------ # # Checkpoint I/O @@ -437,34 +411,22 @@ 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.""" progress = self._submit( - 'resume_from_checkpoint', + "resume_from_checkpoint", model_types.ResumeFromCheckpointRequest, model_types.TrainingProgressResponse, name=name, - adapter_name=kwargs.pop('adapter_name', self.adapter_name), resume_only_model=resume_only_model, - **kwargs) + **kwargs, + ) return progress.result def upload_to_hub( @@ -472,8 +434,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,17 +444,14 @@ 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...') + logger.info("[upload_to_hub] submitting upload, waiting for completion...") self._submit( - 'upload_to_hub', + "upload_to_hub", model_types.UploadToHubRequest, None, checkpoint_dir=checkpoint_dir, hub_model_id=hub_model_id, - hub_token=hub_token) - logger.info('[upload_to_hub] upload completed successfully.') + hub_token=hub_token, + ) + logger.info("[upload_to_hub] upload completed successfully.") 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_channel_loss.py b/tests/loss/test_channel_loss.py index 2d6c8f008..e4bb78cfe 100644 --- a/tests/loss/test_channel_loss.py +++ b/tests/loss/test_channel_loss.py @@ -160,7 +160,6 @@ def forward(self, **kwargs): train_status=status, accumulate_metrics=lambda _: None, ) - def get_loss(loss_key=None): return group.loss_instance 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/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/utils/task_queue/__init__.py b/tests/server/integration/resources/__init__.py similarity index 100% rename from tests/server/utils/task_queue/__init__.py rename to tests/server/integration/resources/__init__.py 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..181ef30bc --- /dev/null +++ b/tests/server/integration/st_cs/validation.py @@ -0,0 +1,210 @@ +# 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 + from twinkle_client.model import MultiLoraTransformersModel + 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 = MultiLoraTransformersModel(MODEL, transport=client.transport) + 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/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_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_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..556bf2879 --- /dev/null +++ b/tests/server/resources/test_instances.py @@ -0,0 +1,254 @@ +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 + + +@pytest.mark.asyncio +async def test_shutdown_actor_lookup_cannot_block_the_close_deadline(monkeypatch): + import threading + from unittest.mock import AsyncMock + import ray + from ray import serve + from twinkle.server.resources.runtime import RayRuntime + + runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) + entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + + def lookup(*args, **kwargs): + entered.set() + try: + # A bounded fallback makes the broken implementation fail rather than hang pytest. + if not resume.wait(2): + raise TimeoutError('blocked actor class import') + raise ValueError('actor already gone') + finally: + returned.set() + + monkeypatch.setattr(ray.util, 'list_named_actors', lambda **kwargs: [{'namespace': 'test', 'name': 'owned-worker'}]) + monkeypatch.setattr(ray, 'get_actor', lookup) + monkeypatch.setattr(serve, 'delete', lambda name: None) + runtime.pool.remove = AsyncMock() + stopped, quarantined = AsyncMock(), AsyncMock() + lifecycle = InstanceLifecycle('test') + try: + await lifecycle.close('owned', lambda: runtime.shutdown('owned', 'app'), .02, stopped, quarantined) + assert entered.is_set() + quarantined.assert_awaited_once() + stopped.assert_not_awaited() + assert not returned.is_set(), 'Synchronous actor lookup blocked the event loop past the close deadline' + finally: + resume.set() + await asyncio.to_thread(returned.wait, 3) + 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 +@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 types import SimpleNamespace + from unittest.mock import AsyncMock, Mock + import ray + from ray import serve + from twinkle.server.exceptions import ResourceReleaseError + from twinkle.server.resources.runtime import RayRuntime + + async def ready(): + if outcome == 'unavailable_worker': + raise ray.exceptions.ActorUnavailableError('Control-plane connection lost', None) + if outcome == 'dead_worker': + raise ray.exceptions.ActorDiedError() + return None # Worker can still answer: death has not been confirmed. + + handle = SimpleNamespace(__ray_ready__=SimpleNamespace(remote=ready)) + rows = ([{'namespace': 'test', 'name': 'owned-worker'}] + if outcome in ('unconfirmed_worker', 'unavailable_worker', 'dead_worker') else []) + monkeypatch.setattr(ray.util, 'list_named_actors', lambda **kwargs: rows) + monkeypatch.setattr(ray, 'get_actor', lambda *args, **kwargs: handle) + monkeypatch.setattr(ray, 'kill', lambda *args, **kwargs: None) + monkeypatch.setattr(serve, 'delete', lambda name: None) + 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..02301c0e0 --- /dev/null +++ b/tests/transformers/test_dev_task_training.py @@ -0,0 +1,50 @@ +# 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() 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..d13d93680 --- /dev/null +++ b/tests/transformers/test_resident_training_reset.py @@ -0,0 +1,175 @@ +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): + 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': False, '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 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..431ed66d7 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) @@ -140,7 +144,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 +162,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 +184,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 +227,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: From dc84b2c53cbd439ed96122ffc480e8046cc55337 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 27 Oct 2026 10:28:40 +0000 Subject: [PATCH 2/6] chore: keep ST dependency reproducers and acceptance notes local --- .gitignore | 5 ++ cookbook/client/swift_model/README-st.md | 68 ------------------- .../swift_model/reproduce_st_dependency.py | 24 ------- .../client/swift_model/requirements-st.txt | 4 -- 4 files changed, 5 insertions(+), 96 deletions(-) delete mode 100644 cookbook/client/swift_model/README-st.md delete mode 100644 cookbook/client/swift_model/reproduce_st_dependency.py delete mode 100644 cookbook/client/swift_model/requirements-st.txt 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/swift_model/README-st.md b/cookbook/client/swift_model/README-st.md deleted file mode 100644 index 49cd1454e..000000000 --- a/cookbook/client/swift_model/README-st.md +++ /dev/null @@ -1,68 +0,0 @@ -# ST CS 依赖与验证入口 - -本分支将 on-demand 生命周期迁移到 `dev-swift-v5`，并保留 `model_cls` / `model_loader` 从 YAML、训练目录、持久化配置到 worker 的链路。 -ST 使用 Swift 的 `MultiLoraSentenceTransformerModel`，依赖固定在 `requirements-st.txt` 的完整 SHA；不另建训练后端。 - -## 固定版本的梯度检查点限制 - -Swift `7e4203a961416952e1c8ea5b8fb1763139d56bfa` 的多租户 ST 构造函数调用： - -```python -self._enable_gradient_checkpointing(gradient_checkpointing_kwargs) -``` - -其父类 `SentenceTransformerModel._enable_gradient_checkpointing(self)` 不接受该参数，因此默认构造在创建租户前失败： - -```text -TypeError: SentenceTransformerModel._enable_gradient_checkpointing() takes 1 positional argument but 2 were given -``` - -裸库在迁移前 dev 基线和本分支均可复现。最小入口（无需启动 CS）： - -```bash -python cookbook/client/swift_model/reproduce_st_dependency.py \ - --model sentence-transformers/all-MiniLM-L6-v2 -# 离线时 --model 指向已有的 ST 模型目录。 -``` - -ST 首版在服务配置中显式使用 `gradient_checkpointing: false`；这是 Swift 构造函数已有的参数，不改 Swift，也不修改其方法。该选择会增加训练显存占用。将来启用梯度检查点前，需要先修复 Swift 封装并更新固定 SHA。 - -服务示例使用 FP32（`mixed_precision: "no"`）、一个 on-demand 实例及两个租户槽位，模型和采样接口不混用：这里通过训练 handle 的 `forward_only` 读取原生 ST 流水线向量。 - -## 运行示例 - -在已就绪的 Twinkle 设备环境中，用 `--no-deps` 安装 `requirements-st.txt` 的固定 Swift/ST 版本，缺失的 ST 常规依赖单独补齐，然后启动： - -```bash -python -m twinkle.server launch --config cookbook/client/swift_model/sentence_transformer_config.yaml -python cookbook/client/swift_model/sentence_transformer_client.py --base-url http://127.0.0.1:8000 -``` - -服务端离线时将 YAML 的 `model_id` 改为完整 ST 模型缓存目录，保留公开路由名；客户端以 `--tokenizer` 指向相同原始模型目录。恢复必须使用同一原始 ST 模型和模块配置，adapter checkpoint 仅保存租户权重、optimizer 与训练进度。 - -## 验收范围 - -先完成 on-demand 普通模型的设备验收，再运行 ST 链路：冷启动 → 创建两个 adapter → 设置 embedding processor、loss 和 optimizer → 训练与读取向量 → 保存 adapter → 释放并重建 handle → 恢复 optimizer 和进度 → 继续训练 → 释放并复用槽位。 -对照同一原始 ST 模型、模块配置和精度，比较池化、投影、归一化后的向量及参数更新；同时检查租户间 optimizer、进度、模板与权重不串用。 -首版限定 LoRA，不扩展全参、SP 或 padding-free。 - -## 验收状态（2026-10-27） - -合并 `fix/model-state-cleanup` 后，普通 Transformers 模型的真实 PPU CS 专项已重新验证双 SDK 训练、optimizer 保存恢复、Controller 重启、worker 故障恢复、创建去重、共享容量、Full 独占和资源释放；Twinkle、Tinker 恢复后继续训练的 logprob 最大误差均为 0。Redis 重启另有迁移阶段验收记录，本轮未据此宣称重新验收。 -合并后的 Transformers/Megatron × Twinkle/Tinker × SFT/DPO/GRPO 核心矩阵 12/12 通过，无跳过；SFT 双 SDK 保存与 optimizer 恢复通过，两后端均确认训练资源释放。Megatron 的 adapter-disabled 参考前向默认不计算已配置的 DPO loss；普通评估、训练和显式要求计算 loss 的前向语义保持不变。 -本地 ST 回归覆盖 Pooling、Dense、Normalize 流水线与租户 optimizer、进度、模板和槽位隔离。all-MiniLM-L6-v2 的 CPU 裸库训练、adapter 恢复和继续训练也已通过。 - -**真实 ST HTTP / PPU 验收已通过。** 覆盖双租户隔离、训练、保存、槽位复用、optimizer 与进度恢复、实例释放后的冷重建及继续训练。相同原始 MiniLM 模型、模块配置、FP32 精度及 checkpoint/RNG 起点下，CS 与裸库的向量和 adapter 参数最大误差均为 0，向量与参数均发生实际更新。 -设备验收入口单独将 `idle_timeout_s` 设置为 900 秒，以覆盖裸库对照进程的冷导入和设备初始化；服务示例的 60 秒闲置配置保持不变。session 心跳不会续期训练业务闲置期限，长时间离开 CS 做本地计算需要相应配置业务闲置预算。 - -采样池独立 `1→2→1` 设备验收仍待完成，不能用普通训练或 ST 的通过替代。扩展用例中的 Full checkpoint IPC 热更新超时继续单独记录，不纳入本次修复。 - -独立设备验收入口（需要两张已确认空闲的卡；CS 使用物理卡 0，裸库对照进程使用物理卡 1）： - -```bash -TWINKLE_TEST_ST_CS=1 TWINKLE_TEST_ST_MODEL=/absolute/ST/model/cache \ -CUDA_VISIBLE_DEVICES=0 python -m pytest --noconftest \ - tests/server/integration/st_cs/test_http_e2e.py -q -``` - -采样扩缩容入口使用 `TWINKLE_TEST_REAL_POOL=1`、`TWINKLE_TEST_POOL_DEVICES=2`、`TWINKLE_TEST_MODEL_PATH` 和 `TWINKLE_TEST_ARTIFACT_DIR`，选择 `tests/server/integration/sampler/test_real_pool_e2e.py::test_real_pool_scaling`。原有 `test_real_pool_lifecycle` 继续保留 Full 热更新、LoRA 驱逐及故障恢复等扩展断言。 diff --git a/cookbook/client/swift_model/reproduce_st_dependency.py b/cookbook/client/swift_model/reproduce_st_dependency.py deleted file mode 100644 index 4ce5df2d2..000000000 --- a/cookbook/client/swift_model/reproduce_st_dependency.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Reproduce the pinned Swift ST constructor failure without HTTP or Ray. - -Pass an existing SentenceTransformer directory to --model for an offline run. -The default downloads the same model used by the planned CS acceptance. -""" - -import argparse - -import twinkle - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--model", default="sentence-transformers/all-MiniLM-L6-v2") - args = parser.parse_args() - from swift.dev.model.multi_lora_sentence_transformer_model import MultiLoraSentenceTransformerModel - - twinkle.initialize(mode="local") - MultiLoraSentenceTransformerModel(model_id=args.model, mixed_precision="no", max_loras=2, max_r=4) - - -if __name__ == "__main__": - main() diff --git a/cookbook/client/swift_model/requirements-st.txt b/cookbook/client/swift_model/requirements-st.txt deleted file mode 100644 index 95497e022..000000000 --- a/cookbook/client/swift_model/requirements-st.txt +++ /dev/null @@ -1,4 +0,0 @@ -# Pinned dependency for the dev on-demand/ST CS investigation. -# The constructor blocker is documented in README-st.md. -ms-swift @ git+https://github.com/modelscope/ms-swift.git@7e4203a961416952e1c8ea5b8fb1763139d56bfa -sentence-transformers==6.2.0 From 383eea1eb155553dc7da26057b469d35843533e9 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 28 Oct 2026 02:37:11 +0000 Subject: [PATCH 3/6] fix(examples): use client.model for ST training and acceptance --- .../sentence_transformer_client.py | 5 ++-- .../sentence_transformer_config.yaml | 2 +- tests/server/integration/st_cs/validation.py | 3 +- tests/twinkle_client/test_async_components.py | 29 +++++++++++++++++++ 4 files changed, 33 insertions(+), 6 deletions(-) diff --git a/cookbook/client/swift_model/sentence_transformer_client.py b/cookbook/client/swift_model/sentence_transformer_client.py index 5a55bd13b..4ece80392 100644 --- a/cookbook/client/swift_model/sentence_transformer_client.py +++ b/cookbook/client/swift_model/sentence_transformer_client.py @@ -9,7 +9,6 @@ from transformers import AutoTokenizer from twinkle import init_twinkle_client -from twinkle_client.model import MultiLoraTransformersModel MODEL = "sentence-transformers/all-MiniLM-L6-v2" TEXTS = [ @@ -56,7 +55,7 @@ def main(): client = init_twinkle_client(base_url=args.base_url, api_key=args.api_key) active = None try: - active = MultiLoraTransformersModel(args.model, transport=client.transport) + active = client.model(args.model) active.add_adapter_to_model("first", LoraConfig(r=4, target_modules=["query", "value"])) configure(active) before = vectors(active, batch) @@ -67,7 +66,7 @@ def main(): 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 = MultiLoraTransformersModel(args.model, transport=client.transport) + 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) diff --git a/cookbook/client/swift_model/sentence_transformer_config.yaml b/cookbook/client/swift_model/sentence_transformer_config.yaml index 2977173b0..e5ee771af 100644 --- a/cookbook/client/swift_model/sentence_transformer_config.yaml +++ b/cookbook/client/swift_model/sentence_transformer_config.yaml @@ -1,5 +1,5 @@ # ST multi-tenant LoRA over the on-demand training pool. -# Install requirements-st.txt; see README-st.md for the pinned Swift constraint. +# Validated with Swift 7e4203a961416952e1c8ea5b8fb1763139d56bfa and sentence-transformers 6.2.0. # The original ST pipeline remains the base model; checkpoints contain tenant adapters. proxy_location: EveryNode http_options: diff --git a/tests/server/integration/st_cs/validation.py b/tests/server/integration/st_cs/validation.py index 181ef30bc..5f116d59c 100644 --- a/tests/server/integration/st_cs/validation.py +++ b/tests/server/integration/st_cs/validation.py @@ -54,7 +54,6 @@ def reference(args): def validate_http(base_url, controller, args): from twinkle import init_twinkle_client - from twinkle_client.model import MultiLoraTransformersModel import ray from tests.server.integration.on_demand_training.acceptance_support import wait_until @@ -65,7 +64,7 @@ def validate_http(base_url, controller, args): 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 = MultiLoraTransformersModel(MODEL, transport=client.transport) + 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) diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py index 431ed66d7..932f1446e 100644 --- a/tests/twinkle_client/test_async_components.py +++ b/tests/twinkle_client/test_async_components.py @@ -114,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 From b123d2f45a1d8eccb32ca900995287c20cb7353f Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 28 Oct 2026 05:44:17 +0000 Subject: [PATCH 4/6] fix(resources): terminate owned actors without loading model classes --- src/twinkle/server/resources/runtime.py | 78 ++++++--- tests/server/resources/test_instances.py | 206 ++++++++++++++++++----- 2 files changed, 219 insertions(+), 65 deletions(-) diff --git a/src/twinkle/server/resources/runtime.py b/src/twinkle/server/resources/runtime.py index 4134d3c5f..3a1c3cd20 100644 --- a/src/twinkle/server/resources/runtime.py +++ b/src/twinkle/server/resources/runtime.py @@ -72,7 +72,6 @@ def lookup(): return await asyncio.to_thread(lookup) async def shutdown(self, identity, name, *, observe_removal=None): - import ray from ray import serve from ray.util import list_named_actors @@ -80,29 +79,60 @@ async def shutdown(self, identity, name, *, observe_removal=None): if deployment is not None: await asyncio.gather(asyncio.shield(deployment), return_exceptions=True) self.deployments.pop(identity, None) - # Both features give workers the controller identity, not a Serve replica ID. - actors = [] - 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 - try: - # Ray may import the worker class when resolving its handle. - # Keep imports and GCS calls outside the controller event loop - # so its shutdown deadline and other instances remain live. - handle = await asyncio.to_thread(ray.get_actor, row["name"], namespace=self.namespace) - except ValueError: - continue - actors.append(handle) - await asyncio.to_thread(ray.kill, handle, no_restart=True) - await asyncio.to_thread(serve.delete, name) - for handle in actors: - try: - await asyncio.wait_for(handle.__ray_ready__.remote(), 5) - except ray.exceptions.ActorDiedError: - continue - except ray.exceptions.RayActorError as exc: - raise ResourceReleaseError("Owned worker death could not be confirmed", phase="worker_death") from exc - raise ResourceReleaseError("Owned worker death could not be confirmed", phase="worker_death") + 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: diff --git a/tests/server/resources/test_instances.py b/tests/server/resources/test_instances.py index 556bf2879..fb5c01b5f 100644 --- a/tests/server/resources/test_instances.py +++ b/tests/server/resources/test_instances.py @@ -90,43 +90,69 @@ async def failed(): 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): - import threading from unittest.mock import AsyncMock - import ray - from ray import serve from twinkle.server.resources.runtime import RayRuntime - runtime = RayRuntime('test', allocation_timeout=1, shutdown_timeout=1) - entered, resume, returned = threading.Event(), threading.Event(), threading.Event() + rpc = _patch_actor_cleanup_rpc(monkeypatch, [{'namespace': 'test', 'name': 'owned-worker'}]) + entered = asyncio.Event() - def lookup(*args, **kwargs): + async def lookup(*args, **kwargs): entered.set() - try: - # A bounded fallback makes the broken implementation fail rather than hang pytest. - if not resume.wait(2): - raise TimeoutError('blocked actor class import') - raise ValueError('actor already gone') - finally: - returned.set() + await asyncio.Event().wait() - monkeypatch.setattr(ray.util, 'list_named_actors', lambda **kwargs: [{'namespace': 'test', 'name': 'owned-worker'}]) - monkeypatch.setattr(ray, 'get_actor', lookup) - monkeypatch.setattr(serve, 'delete', lambda name: None) + 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') - try: - await lifecycle.close('owned', lambda: runtime.shutdown('owned', 'app'), .02, stopped, quarantined) - assert entered.is_set() - quarantined.assert_awaited_once() - stopped.assert_not_awaited() - assert not returned.is_set(), 'Synchronous actor lookup blocked the event loop past the close deadline' - finally: - resume.set() - await asyncio.to_thread(returned.wait, 3) - await lifecycle.join() + 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 @@ -206,32 +232,130 @@ def inventory(): 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 types import SimpleNamespace from unittest.mock import AsyncMock, Mock - import ray - from ray import serve + from ray.core.generated.gcs_service_pb2 import GetActorInfoReply from twinkle.server.exceptions import ResourceReleaseError from twinkle.server.resources.runtime import RayRuntime - async def ready(): - if outcome == 'unavailable_worker': - raise ray.exceptions.ActorUnavailableError('Control-plane connection lost', None) - if outcome == 'dead_worker': - raise ray.exceptions.ActorDiedError() - return None # Worker can still answer: death has not been confirmed. - - handle = SimpleNamespace(__ray_ready__=SimpleNamespace(remote=ready)) rows = ([{'namespace': 'test', 'name': 'owned-worker'}] if outcome in ('unconfirmed_worker', 'unavailable_worker', 'dead_worker') else []) - monkeypatch.setattr(ray.util, 'list_named_actors', lambda **kwargs: rows) - monkeypatch.setattr(ray, 'get_actor', lambda *args, **kwargs: handle) - monkeypatch.setattr(ray, 'kill', lambda *args, **kwargs: None) - monkeypatch.setattr(serve, 'delete', lambda name: None) + 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) From af435b70d610f31680dcc351cbafce3e3c6c45bc Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 30 Oct 2026 01:46:02 +0000 Subject: [PATCH 5/6] feat: support CS embedding and classification tasks --- .../sentence_transformer_config.yaml | 1 - cookbook/client/twinkle/embedding.py | 228 ++++++++--------- .../client/twinkle/prepare_sequence_model.py | 134 ++++++++++ cookbook/client/twinkle/reranker.py | 131 ++++++++++ cookbook/client/twinkle/seq_cls.py | 127 ++++++++++ .../client/twinkle/task_server_megatron.yaml | 64 +++++ .../twinkle/task_server_megatron_qwen35.yaml | 64 +++++ .../twinkle/task_server_transformers.yaml | 65 +++++ .../Usage Guide/Embedding-Training.md | 9 +- .../Embedding\350\256\255\347\273\203.md" | 9 +- src/twinkle/model/megatron/megatron.py | 24 +- .../model/megatron/multi_lora_megatron.py | 35 ++- src/twinkle/model/multi_lora.py | 3 +- src/twinkle/server/config/training.py | 13 + src/twinkle/server/model/app.py | 38 +++ .../integration/task_http_validation.py | 235 ++++++++++++++++++ .../server/integration/test_task_http_e2e.py | 38 +++ tests/server/model/test_task_head.py | 103 ++++++++ tests/transformers/test_dev_task_training.py | 88 +++++++ .../test_resident_training_reset.py | 62 ++++- 20 files changed, 1337 insertions(+), 134 deletions(-) create mode 100644 cookbook/client/twinkle/prepare_sequence_model.py create mode 100644 cookbook/client/twinkle/reranker.py create mode 100644 cookbook/client/twinkle/seq_cls.py create mode 100644 cookbook/client/twinkle/task_server_megatron.yaml create mode 100644 cookbook/client/twinkle/task_server_megatron_qwen35.yaml create mode 100644 cookbook/client/twinkle/task_server_transformers.yaml create mode 100644 tests/server/integration/task_http_validation.py create mode 100644 tests/server/integration/test_task_http_e2e.py create mode 100644 tests/server/model/test_task_head.py diff --git a/cookbook/client/swift_model/sentence_transformer_config.yaml b/cookbook/client/swift_model/sentence_transformer_config.yaml index e5ee771af..efc28a2e2 100644 --- a/cookbook/client/swift_model/sentence_transformer_config.yaml +++ b/cookbook/client/swift_model/sentence_transformer_config.yaml @@ -1,5 +1,4 @@ # ST multi-tenant LoRA over the on-demand training pool. -# Validated with Swift 7e4203a961416952e1c8ea5b8fb1763139d56bfa and sentence-transformers 6.2.0. # The original ST pipeline remains the base model; checkpoints contain tenant adapters. proxy_location: EveryNode http_options: diff --git a/cookbook/client/twinkle/embedding.py b/cookbook/client/twinkle/embedding.py index 1e8eb36dc..52f94b26e 100644 --- a/cookbook/client/twinkle/embedding.py +++ b/cookbook/client/twinkle/embedding.py @@ -1,122 +1,128 @@ -# Twinkle Client - Embedding (InfoNCE contrastive) Training Example -# -# HTTP client example described in the client section of -# ``docs/source_zh/使用指引/Embedding训练.md``. Instead of building a local -# ``TransformersModel``/``MegatronModel`` + Ray, it drives the SAME training flow -# over HTTP through the Twinkle client: -# -# 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) -# -# 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). - -import dotenv +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Train an embedding adapter over HTTP, save it and continue after cold recreation. + +From the repository root: + python cookbook/client/twinkle/prepare_sequence_model.py --tiny + twinkle-server launch --config cookbook/client/twinkle/task_server_transformers.yaml + python cookbook/client/twinkle/embedding.py + +For Megatron with the tiny model, use task_server_megatron.yaml and run this example +with --backend megatron. For Qwen3.5, prepare with --model ms://Qwen/Qwen3.5-4B +--mcore-native, launch with task_server_megatron_qwen35.yaml, and keep the tokenizer +in output/cs-models/embedding. No sampler service is required. +""" +import argparse +import json import os +from pathlib import Path + +import numpy as np from peft import LoraConfig -from typing import Any, Dict, List +from transformers import AutoTokenizer -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 ========== -BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') -MODEL_ID = f'ms://{BASE_MODEL}' -ADAPTER_NAME = 'emb_adapter' -TEMPERATURE = 0.07 -LEARNING_RATE = 1e-4 -MAX_STEPS = 10 -MAX_LENGTH = 512 - -# Small self-contained anchor/positive corpus. Each anchor's positive shares its -# meaning, so InfoNCE has a clear contrastive signal to learn from. -PAIRS: List[tuple] = [ - ('What is the capital of France?', 'Paris is the capital city of France.'), - ('How do plants make their food?', 'Plants use photosynthesis to produce food from sunlight.'), - ('What is the boiling point of water?', 'Water boils at 100 degrees Celsius at sea level.'), - ('Who wrote the play Hamlet?', 'William Shakespeare wrote the tragedy Hamlet.'), - ('What is the speed of light?', 'Light travels at about 300,000 kilometers per second.'), - ('What language is spoken in Brazil?', 'The main language spoken in Brazil is Portuguese.'), +# Each group is (query, positive, negative). Labels mark group starts, not relevance. +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'), ] -def build_minibatch(tokenizer) -> List[Dict[str, Any]]: - """Build one interleaved [anchor, positive, ...] embedding minibatch. - - Each sample is a plain-dict feature with: - * ``input_ids`` -- tokenized text (plain Python list, JSON-serializable) - * ``attention_mask`` -- all ones - * ``labels`` -- a single scalar: 1 for an anchor, 0 for its positive - * ``position_ids`` -- ``range(len)`` so the megatron mrope model gets valid - positions (transformers derives them internally) - """ - minibatch: List[Dict[str, Any]] = [] - for anchor_text, positive_text in PAIRS: - for text, is_anchor in ((anchor_text, True), (positive_text, False)): - input_ids = tokenizer.encode(text, add_special_tokens=True) - minibatch.append({ - 'input_ids': list(input_ids), - 'attention_mask': [1] * len(input_ids), - 'labels': [1 if is_anchor else 0], - 'position_ids': list(range(len(input_ids))), - }) - return minibatch - - -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: 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) - - # Step 3: configure the embedding-training pipeline (ORDER MATTERS). - # Pass class names as strings; the server resolves them to core-lib classes. - model.set_processor('InputProcessor') - model.set_loss('InfonceLoss', temperature=TEMPERATURE, use_batch=True, hard_negatives=None) - # 'Adam' works on both transformers and megatron backends ('AdamW' is - # transformers-only; megatron expects 'Adam'/'default'/'MegatronOptimizer'). - model.set_optimizer('Adam', lr=LEARNING_RATE) - model.add_metric('EmbeddingMetric', is_training=True) - - # Step 4: build a small contrastive minibatch with a local tokenizer. - template = Qwen3_5Template(model_id=MODEL_ID, max_length=MAX_LENGTH, enable_thinking=False) - minibatch = build_minibatch(template.tokenizer) - logger.info(f'Built minibatch: {len(minibatch)} samples ({len(PAIRS)} anchor/positive pairs)') - - # 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). - 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.') +def encoded_batch(tokenizer): + rows = [] + for query, positive, negative in TRIPLETS: + for text, boundary in [(query, 1), (positive, 0), (negative, 0)]: + row = dict(tokenizer(text, truncation=True, max_length=128)) + row['labels'] = [boundary] + row['position_ids'] = list(range(len(row['input_ids']))) + rows.append(row) + return rows + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--backend', choices=['transformers', 'megatron'], default='transformers') + parser.add_argument('--base-url', default=os.environ.get('TWINKLE_SERVER_URL', 'http://127.0.0.1:8000')) + parser.add_argument('--api-key', default=os.environ.get('TWINKLE_SERVER_TOKEN', 'embedding-example')) + parser.add_argument('--model', default='cs-embedding', help='Public model alias in the server YAML') + parser.add_argument('--tokenizer', default='output/cs-models/embedding') + parser.add_argument('--steps', type=int, default=7, help='Backward batches; save on a GA=2 update boundary') + parser.add_argument('--lr', type=float, default=1e-3) + parser.add_argument('--atol', type=float, default=2e-5) + parser.add_argument('--result-json', help='Optional local file for model outputs and training progress') + args = parser.parse_args() + if args.steps < 3 or args.steps % 2 == 0: + parser.error('--steps must be odd and at least 3 so GA=2 has no unsaved pending gradients') + if not np.isfinite(args.lr) or args.lr <= 0 or not np.isfinite(args.atol) or args.atol < 0: + parser.error('lr must be positive; atol must be finite and nonnegative') + batch = encoded_batch(AutoTokenizer.from_pretrained(args.tokenizer)) + client = init_twinkle_client(base_url=args.base_url, api_key=args.api_key) + model = None + targets = ['q_proj', 'v_proj'] if args.backend == 'transformers' else ['linear_qkv'] + forward_kwargs = {'task': 'embedding', 'return_logits': True} + if args.backend == 'megatron': + forward_kwargs.update(micro_batch_size=6, compute_loss=False) + try: + # Allocate a tenant handle on the server's frozen causal text backbone. + model = client.model(args.model) + model.add_adapter_to_model('embedding', LoraConfig(r=4, lora_alpha=8, target_modules=targets), + gradient_accumulation_steps=2) + model.set_template('Template', max_length=128) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('InfonceLoss', temperature=0.1, use_batch=True) + model.set_optimizer('Adam', lr=args.lr) + before = model.forward_only(inputs=batch, **forward_kwargs).result['embeddings'] + losses = [] + for _ in range(args.steps): + kwargs = {'task': 'embedding', 'gradient_accumulation_steps': 2} + if args.backend == 'megatron': + # DP=2 gives six rows per rank, preserving two complete triplets. + kwargs['micro_batch_size'] = 6 + result = model.forward_backward(inputs=batch, **kwargs).result + # Call after every backward; the backend gates updates at its GA boundary. + model.clip_grad_and_step(max_grad_norm=1.0) + losses.append(float(result['loss'])) + trained = model.forward_only(inputs=batch, **forward_kwargs).result['embeddings'] + saved = model.save('embedding-trained', save_optimizer=True, + consumed_train_samples=args.steps * len(batch)) + model.remove_adapter() + model = None + + # Reload the same persisted base, then restore adapter and optimizer together. + model = client.model(args.model) + model.add_adapter_to_model('restored', LoraConfig(r=4, lora_alpha=8, target_modules=targets), + gradient_accumulation_steps=2) + model.set_template('Template', max_length=128) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('InfonceLoss', temperature=0.1, use_batch=True) + model.set_optimizer('Adam', lr=args.lr) + progress = model.resume_from_checkpoint(saved.twinkle_path) + restored = model.forward_only(inputs=batch, **forward_kwargs).result['embeddings'] + np.testing.assert_allclose(restored, trained, rtol=0, atol=args.atol) + for _ in range(2): + kwargs = {'task': 'embedding', 'gradient_accumulation_steps': 2} + if args.backend == 'megatron': + kwargs['micro_batch_size'] = 6 + model.forward_backward(inputs=batch, **kwargs) + model.clip_grad_and_step() + continued = model.forward_only(inputs=batch, **forward_kwargs).result['embeddings'] + result = {'task': 'embedding', 'backend': args.backend, 'losses': losses, + 'checkpoint': saved.twinkle_path, 'checkpoint_dir': saved.checkpoint_dir, + 'progress': progress, 'before': before, 'trained': trained, + 'restored': restored, 'continued': continued} + if args.result_json: + Path(args.result_json).write_text(json.dumps(result, indent=2)) + print(json.dumps({'losses': losses, 'checkpoint': saved.twinkle_path, 'progress': progress})) # noqa: T201 + finally: + try: + if model is not None: + model.remove_adapter() + finally: + client.close() if __name__ == '__main__': - train() + main() diff --git a/cookbook/client/twinkle/prepare_sequence_model.py b/cookbook/client/twinkle/prepare_sequence_model.py new file mode 100644 index 000000000..b4561264a --- /dev/null +++ b/cookbook/client/twinkle/prepare_sequence_model.py @@ -0,0 +1,134 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Prepare fixed text backbones and classification heads for CS training. + + python cookbook/client/twinkle/prepare_sequence_model.py --model ms://Qwen/Qwen3.5-4B + python cookbook/client/twinkle/prepare_sequence_model.py --model ms://Qwen/Qwen3.5-4B --mcore-native + python cookbook/client/twinkle/prepare_sequence_model.py --tiny + +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 + + +def tiny_backbone(directory, seed=42): + from tokenizers.pre_tokenizers import ByteLevel + from transformers import Qwen2Config, Qwen2ForCausalLM, Qwen2Tokenizer + + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + config = Qwen2Config(vocab_size=512, 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=0, bos_token_id=2, eos_token_id=3, + attention_dropout=0., tie_word_embeddings=False) + Qwen2ForCausalLM(config).save_pretrained(directory) + vocab = {'[PAD]': 0, '[UNK]': 1, '[BOS]': 2, '[EOS]': 3} + # Use the model family's tokenizer so AutoTokenizer preserves its byte vocabulary. + vocab.update({token: index + 4 for index, token in enumerate(sorted(ByteLevel.alphabet()))}) + Qwen2Tokenizer(vocab=vocab, merges=[], pad_token='[PAD]', unk_token='[UNK]', + bos_token='[BOS]', eos_token='[EOS]').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 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') + 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', default='ms://Qwen/Qwen3.5-4B') + 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 offline Qwen2 backbone') + parser.add_argument('--mcore-native', action='store_true', + help='Also write Qwen3.5 Megatron-native text-only checkpoints') + args = parser.parse_args() + source = tiny_backbone(Path(args.output_dir) / 'source', args.seed) if args.tiny else args.model + prepare_models(source, args.output_dir, seed=args.seed, mcore_native=args.mcore_native) + + +if __name__ == '__main__': + main() diff --git a/cookbook/client/twinkle/reranker.py b/cookbook/client/twinkle/reranker.py new file mode 100644 index 000000000..6a4a8a03a --- /dev/null +++ b/cookbook/client/twinkle/reranker.py @@ -0,0 +1,131 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Train a pointwise reranker over HTTP with LoRA on the backbone and score head. + +From the repository root: + python cookbook/client/twinkle/prepare_sequence_model.py --tiny + twinkle-server launch --config cookbook/client/twinkle/task_server_transformers.yaml + python cookbook/client/twinkle/reranker.py + +For Megatron with the tiny model, use task_server_megatron.yaml and pass --backend megatron. +For Qwen3.5, prepare with --model ms://Qwen/Qwen3.5-4B --mcore-native, launch +with task_server_megatron_qwen35.yaml, and keep the tokenizer in output/cs-models/reranker. +The base checkpoint must have a fixed one-output score head. +Each input is one query/document sequence; labels are binary relevance targets. +""" +import argparse +import json +import os +from pathlib import Path + +import numpy as np +from peft import LoraConfig +from transformers import AutoTokenizer + +from twinkle import init_twinkle_client + +# Wire labels are integer 0/1; PointwiseRerankerLoss casts them to the score dtype. +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), +] + + +def encoded_batch(tokenizer): + rows = [] + for query, document, relevance in PAIRS: + row = dict(tokenizer('Query: ' + query + '\nDocument: ' + document, + truncation=True, max_length=128)) + row['labels'] = [relevance] + row['position_ids'] = list(range(len(row['input_ids']))) + rows.append(row) + return rows + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--backend', choices=['transformers', 'megatron'], default='transformers') + parser.add_argument('--base-url', default=os.environ.get('TWINKLE_SERVER_URL', 'http://127.0.0.1:8000')) + parser.add_argument('--api-key', default=os.environ.get('TWINKLE_SERVER_TOKEN', 'reranker-example')) + parser.add_argument('--model', default='cs-reranker', help='Public model alias in the server YAML') + parser.add_argument('--tokenizer', default='output/cs-models/reranker') + parser.add_argument('--steps', type=int, default=7, help='Backward batches; save on a GA=2 update boundary') + parser.add_argument('--lr', type=float, default=1e-3) + parser.add_argument('--atol', type=float, default=2e-5) + parser.add_argument('--result-json', help='Optional local file for model outputs and training progress') + args = parser.parse_args() + if args.steps < 3 or args.steps % 2 == 0: + parser.error('--steps must be odd and at least 3 so GA=2 has no unsaved pending gradients') + if not np.isfinite(args.lr) or args.lr <= 0 or not np.isfinite(args.atol) or args.atol < 0: + parser.error('lr must be positive; atol must be finite and nonnegative') + batch = encoded_batch(AutoTokenizer.from_pretrained(args.tokenizer)) + client = init_twinkle_client(base_url=args.base_url, api_key=args.api_key) + model = None + targets = (['q_proj', 'v_proj', 'score'] if args.backend == 'transformers' else + ['linear_qkv', 'output_layer']) + forward_kwargs = {'task': 'reranker', 'return_logits': True} + if args.backend == 'megatron': + forward_kwargs.update(micro_batch_size=6, compute_loss=False) + try: + # score/output_layer LoRA is explicitly selected, not a full shared trainable head. + model = client.model(args.model) + model.add_adapter_to_model('reranker', LoraConfig(r=4, lora_alpha=8, target_modules=targets), + gradient_accumulation_steps=2) + model.set_template('Template', max_length=128) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('PointwiseRerankerLoss') + model.set_optimizer('Adam', lr=args.lr) + before = model.forward_only(inputs=batch, **forward_kwargs).result['logits'] + losses = [] + for _ in range(args.steps): + kwargs = {'task': 'reranker', 'gradient_accumulation_steps': 2} + if args.backend == 'megatron': + kwargs['micro_batch_size'] = 6 + result = model.forward_backward(inputs=batch, **kwargs).result + # Call after every backward; the backend gates updates at its GA boundary. + model.clip_grad_and_step(max_grad_norm=1.0) + losses.append(float(result['loss'])) + trained = model.forward_only(inputs=batch, **forward_kwargs).result['logits'] + saved = model.save('reranker-trained', save_optimizer=True, + consumed_train_samples=args.steps * len(batch)) + model.remove_adapter() + model = None + + # A cold worker loads the same persisted base head before restoring its LoRA. + model = client.model(args.model) + model.add_adapter_to_model('restored', LoraConfig(r=4, lora_alpha=8, target_modules=targets), + gradient_accumulation_steps=2) + model.set_template('Template', max_length=128) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('PointwiseRerankerLoss') + model.set_optimizer('Adam', lr=args.lr) + progress = model.resume_from_checkpoint(saved.twinkle_path) + restored = model.forward_only(inputs=batch, **forward_kwargs).result['logits'] + np.testing.assert_allclose(restored, trained, rtol=0, atol=args.atol) + for _ in range(2): + kwargs = {'task': 'reranker', 'gradient_accumulation_steps': 2} + if args.backend == 'megatron': + kwargs['micro_batch_size'] = 6 + model.forward_backward(inputs=batch, **kwargs) + model.clip_grad_and_step() + continued = model.forward_only(inputs=batch, **forward_kwargs).result['logits'] + result = {'task': 'reranker', 'backend': args.backend, 'losses': losses, + 'checkpoint': saved.twinkle_path, 'checkpoint_dir': saved.checkpoint_dir, + 'progress': progress, 'before': before, 'trained': trained, + 'restored': restored, 'continued': continued} + if args.result_json: + Path(args.result_json).write_text(json.dumps(result, indent=2)) + print(json.dumps({'losses': losses, 'checkpoint': saved.twinkle_path, 'progress': progress})) # noqa: T201 + finally: + try: + if model is not None: + model.remove_adapter() + finally: + client.close() + + +if __name__ == '__main__': + main() diff --git a/cookbook/client/twinkle/seq_cls.py b/cookbook/client/twinkle/seq_cls.py new file mode 100644 index 000000000..ca34fa99d --- /dev/null +++ b/cookbook/client/twinkle/seq_cls.py @@ -0,0 +1,127 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Train a three-class CS model with backbone and classification-head LoRA. + +From the repository root: + python cookbook/client/twinkle/prepare_sequence_model.py --tiny + twinkle-server launch --config cookbook/client/twinkle/task_server_transformers.yaml + python cookbook/client/twinkle/seq_cls.py + +For Megatron with the tiny model, use task_server_megatron.yaml and pass --backend megatron. +For Qwen3.5, prepare with --model ms://Qwen/Qwen3.5-4B --mcore-native, launch +with task_server_megatron_qwen35.yaml, and keep the tokenizer in output/cs-models/seq_cls. +The base checkpoint must contain a fixed three-output head; +this recipe does not train a full head through modules_to_save. +""" +import argparse +import json +import os +from pathlib import Path + +import numpy as np +from peft import LoraConfig +from transformers import AutoTokenizer + +from twinkle import init_twinkle_client + +# Class IDs match the head width in the service YAML: animal=0, fruit=1, vehicle=2. +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(tokenizer): + rows = [] + for text, class_id in SAMPLES: + row = dict(tokenizer(text, truncation=True, max_length=128)) + row['labels'] = [class_id] + row['position_ids'] = list(range(len(row['input_ids']))) + rows.append(row) + return rows + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--backend', choices=['transformers', 'megatron'], default='transformers') + parser.add_argument('--base-url', default=os.environ.get('TWINKLE_SERVER_URL', 'http://127.0.0.1:8000')) + parser.add_argument('--api-key', default=os.environ.get('TWINKLE_SERVER_TOKEN', 'classification-example')) + parser.add_argument('--model', default='cs-seq_cls', help='Public model alias in the server YAML') + parser.add_argument('--tokenizer', default='output/cs-models/seq_cls') + parser.add_argument('--steps', type=int, default=7, help='Backward batches; save on a GA=2 update boundary') + parser.add_argument('--lr', type=float, default=1e-3) + parser.add_argument('--atol', type=float, default=2e-5) + parser.add_argument('--result-json', help='Optional local file for model outputs and training progress') + args = parser.parse_args() + if args.steps < 3 or args.steps % 2 == 0: + parser.error('--steps must be odd and at least 3 so GA=2 has no unsaved pending gradients') + if not np.isfinite(args.lr) or args.lr <= 0 or not np.isfinite(args.atol) or args.atol < 0: + parser.error('lr must be positive; atol must be finite and nonnegative') + batch = encoded_batch(AutoTokenizer.from_pretrained(args.tokenizer)) + client = init_twinkle_client(base_url=args.base_url, api_key=args.api_key) + model = None + targets = (['q_proj', 'v_proj', 'score'] if args.backend == 'transformers' else + ['linear_qkv', 'output_layer']) + forward_kwargs = {'task': 'seq_cls', 'return_logits': True} + if args.backend == 'megatron': + forward_kwargs.update(micro_batch_size=6, compute_loss=False) + try: + model = client.model(args.model) + model.add_adapter_to_model('classification', LoraConfig(r=4, lora_alpha=8, target_modules=targets), + gradient_accumulation_steps=2) + model.set_template('Template', max_length=128) + model.set_processor('InputProcessor', padding_side='right') + # Cross entropy consumes one integer class ID per input sequence. + model.set_loss('SeqClsLoss', problem_type='single_label_classification', num_labels=3) + model.set_optimizer('Adam', lr=args.lr) + before = model.forward_only(inputs=batch, **forward_kwargs).result['logits'] + losses = [] + for _ in range(args.steps): + kwargs = {'task': 'seq_cls', 'gradient_accumulation_steps': 2} + if args.backend == 'megatron': + kwargs['micro_batch_size'] = 6 + result = model.forward_backward(inputs=batch, **kwargs).result + # Call after every backward; the backend gates updates at its GA boundary. + model.clip_grad_and_step(max_grad_norm=1.0) + losses.append(float(result['loss'])) + trained = model.forward_only(inputs=batch, **forward_kwargs).result['logits'] + saved = model.save('classification-trained', save_optimizer=True, + consumed_train_samples=args.steps * len(batch)) + model.remove_adapter() + model = None + + # Restore both LoRA (including score-head LoRA) and optimizer, not a new random head. + model = client.model(args.model) + model.add_adapter_to_model('restored', LoraConfig(r=4, lora_alpha=8, target_modules=targets), + gradient_accumulation_steps=2) + model.set_template('Template', max_length=128) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('SeqClsLoss', problem_type='single_label_classification', num_labels=3) + model.set_optimizer('Adam', lr=args.lr) + progress = model.resume_from_checkpoint(saved.twinkle_path) + restored = model.forward_only(inputs=batch, **forward_kwargs).result['logits'] + np.testing.assert_allclose(restored, trained, rtol=0, atol=args.atol) + for _ in range(2): + kwargs = {'task': 'seq_cls', 'gradient_accumulation_steps': 2} + if args.backend == 'megatron': + kwargs['micro_batch_size'] = 6 + model.forward_backward(inputs=batch, **kwargs) + model.clip_grad_and_step() + continued = model.forward_only(inputs=batch, **forward_kwargs).result['logits'] + result = {'task': 'seq_cls', 'backend': args.backend, 'losses': losses, + 'checkpoint': saved.twinkle_path, 'checkpoint_dir': saved.checkpoint_dir, + 'progress': progress, 'before': before, 'trained': trained, + 'restored': restored, 'continued': continued} + if args.result_json: + Path(args.result_json).write_text(json.dumps(result, indent=2)) + print(json.dumps({'losses': losses, 'checkpoint': saved.twinkle_path, 'progress': progress})) # noqa: T201 + finally: + try: + if model is not None: + model.remove_adapter() + finally: + client.close() + + +if __name__ == '__main__': + main() diff --git a/cookbook/client/twinkle/task_server_megatron.yaml b/cookbook/client/twinkle/task_server_megatron.yaml new file mode 100644 index 000000000..c695fd417 --- /dev/null +++ b/cookbook/client/twinkle/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/cookbook/client/twinkle/task_server_megatron_qwen35.yaml b/cookbook/client/twinkle/task_server_megatron_qwen35.yaml new file mode 100644 index 000000000..0f3520225 --- /dev/null +++ b/cookbook/client/twinkle/task_server_megatron_qwen35.yaml @@ -0,0 +1,64 @@ +# Qwen3.5 text-only Megatron tasks. Prepare with --mcore-native first. +# Tokenizers can stay in output/cs-models; these bases use model.language_model.* keys. +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-mcore/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-mcore/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-mcore/seq_cls + task_type: seq_cls + num_labels: 3 + target_modules: [linear_qkv, output_layer] + deployments: *deployments diff --git a/cookbook/client/twinkle/task_server_transformers.yaml b/cookbook/client/twinkle/task_server_transformers.yaml new file mode 100644 index 000000000..39975be1d --- /dev/null +++ b/cookbook/client/twinkle/task_server_transformers.yaml @@ -0,0 +1,65 @@ +# Run prepare_sequence_model.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/docs/source_en/Usage Guide/Embedding-Training.md b/docs/source_en/Usage Guide/Embedding-Training.md index 4198f6991..fbd712b47 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 save, cold recreation, optimizer resume and final release; parameter meanings and recommended values are in "Key Parameters" above. --- 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..346e98876 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`，包含保存、冷重建、optimizer 恢复和最终释放。各参数含义与取值建议见上文「关键参数」。 --- diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 9efb68ef3..398597e63 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -600,18 +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) - reported_logits = output_tensor if return_logits else None - return self.strategy.reduce_loss(zero, zero.clone(), reported_logits, 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 ) @@ -621,8 +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) - reported_logits = output_tensor if return_logits else None 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 @@ -819,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 @@ -832,6 +834,8 @@ def forward_step_func(data_iterator, model): logits.append(loss_dict['logits']) 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(): @@ -876,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() @@ -888,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 diff --git a/src/twinkle/model/megatron/multi_lora_megatron.py b/src/twinkle/model/megatron/multi_lora_megatron.py index 90e29eb26..4e58ba4fa 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,33 @@ 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. + """ + 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 +405,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 +482,11 @@ 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) # 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 a692b74d3..4b583a8e5 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -331,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 diff --git a/src/twinkle/server/config/training.py b/src/twinkle/server/config/training.py index d6ce7b505..78723fe0b 100644 --- a/src/twinkle/server/config/training.py +++ b/src/twinkle/server/config/training.py @@ -58,6 +58,17 @@ class TrainingArguments(BaseModel): 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): @@ -110,6 +121,8 @@ def _validate_spectral_backend(self): 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 diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index 3868e0a00..908dddbd8 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -39,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 @@ -216,6 +252,8 @@ 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() diff --git a/tests/server/integration/task_http_validation.py b/tests/server/integration/task_http_validation.py new file mode 100644 index 000000000..fee3a8b94 --- /dev/null +++ b/tests/server/integration/task_http_validation.py @@ -0,0 +1,235 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Owned real-device CS acceptance, including directly executed cookbook CLIs.""" +import argparse +import copy +import importlib +import json +import os +from pathlib import Path +import socket +import subprocess +import sys +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.on_demand_training.acceptance_support import wait_until + + root = Path(__file__).resolve().parents[3] + directory = Path(args.run_dir) + base_path = str(Path(args.model_dir, args.task).resolve()) + recipe = importlib.import_module('cookbook.client.twinkle.' + args.task) + batch = recipe.encoded_batch(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']) + + # Execute the user-facing main rather than inferring coverage from imported helpers. + command = [sys.executable, '-u', str(root / f'cookbook/client/twinkle/{args.task}.py'), + '--backend', args.backend, '--base-url', base_url, '--tokenizer', base_path, + '--steps', str(args.steps), '--result-json', str(directory / 'cli-result.json')] + with (directory / 'cli.log').open('w') as output: + completed = subprocess.run(command, cwd=root, stdout=output, stderr=subprocess.STDOUT, timeout=1800) + assert completed.returncode == 0, (directory / 'cli.log').read_text()[-6000:] + cli = json.loads((directory / 'cli-result.json').read_text()) + assert np.isfinite(cli['losses']).all() + first_loss = float(np.mean(cli['losses'][:3])) + last_loss = float(np.mean(cli['losses'][-3:])) + loss_drop = (first_loss - last_loss) / first_loss + assert loss_drop >= args.min_loss_drop, (cli['losses'], loss_drop) + assert cli['progress']['cur_step'] == args.steps + assert cli['progress']['consumed_train_samples'] == args.steps * len(batch) + assert np.max(np.abs(np.asarray(cli['trained']) - np.asarray(cli['before']))) > 1e-7 + wait_until(stopped, 180, 'CLI did not release the instance') + if args.cli_only: + return {'passed': True, 'task': args.task, 'backend': args.backend, 'cli': True, + 'slot_reuse': None, 'progress': cli['progress'], 'losses': cli['losses'], + 'loss_drop': loss_drop} + 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']) + model.clip_grad_and_step() + + 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: + 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, + 'cli': True, 'slot_reuse': True, 'progress': progress, 'head_update_keys': updates, + 'losses': cli['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) + 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('--cli-only', action='store_true') + args = parser.parse_args() + root = Path(__file__).resolve().parents[3] + 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(root / f'cookbook/client/twinkle/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_task_http_e2e.py b/tests/server/integration/test_task_http_e2e.py new file mode 100644 index 000000000..3de01f143 --- /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['cli'] and result['slot_reuse'] + assert result['loss_drop'] >= 0.1 diff --git a/tests/server/model/test_task_head.py b/tests/server/model/test_task_head.py new file mode 100644 index 000000000..d9c9ecfeb --- /dev/null +++ b/tests/server/model/test_task_head.py @@ -0,0 +1,103 @@ +# 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 cookbook.client.twinkle.prepare_sequence_model import _mcore_native_key +from twinkle.server.config import ServerConfig +from twinkle.server.config.training import TrainingArguments +from twinkle.server.model.app import _validate_model_task + + +@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): + root = Path(__file__).resolve().parents[3] + config = ServerConfig.from_yaml(root / f'cookbook/client/twinkle/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_qwen35_megatron_yaml_uses_native_bases(): + root = Path(__file__).resolve().parents[3] + config = ServerConfig.from_yaml(root / 'cookbook/client/twinkle/task_server_megatron_qwen35.yaml') + apps = [app for app in config.applications if app.import_path == 'model'] + assert len(apps) == 3 + for app in apps: + assert app.args.model_id.startswith('output/cs-models-mcore/') + assert app.args.language_model_only is True + + +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_cookbook_batches_match_the_wire_contract(task): + import importlib + from pydantic import TypeAdapter + from twinkle.protocol.types.data import WireInputBatch + + recipe = importlib.import_module('cookbook.client.twinkle.' + task) + def tokenizer(text, **kwargs): + return {'input_ids': [4, 5, 6], 'attention_mask': [1, 1, 1]} + batch = recipe.encoded_batch(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/transformers/test_dev_task_training.py b/tests/transformers/test_dev_task_training.py index 02301c0e0..0ba708e1a 100644 --- a/tests/transformers/test_dev_task_training.py +++ b/tests/transformers/test_dev_task_training.py @@ -48,3 +48,91 @@ def test_dev_task_forward_backward_and_update(tmp_path, monkeypatch, task): 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 cookbook.client.twinkle.prepare_sequence_model import prepare_models, tiny_backbone + from cookbook.client.twinkle import embedding, reranker, seq_cls + 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']) + recipe = {'embedding': embedding, 'reranker': reranker, 'seq_cls': seq_cls}[task] + batch = recipe.encoded_batch(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_resident_training_reset.py b/tests/transformers/test_resident_training_reset.py index d13d93680..885e4bd8d 100644 --- a/tests/transformers/test_resident_training_reset.py +++ b/tests/transformers/test_resident_training_reset.py @@ -122,7 +122,7 @@ def load_weights(chunks, checkpoint, *, peft_format): 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): +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' @@ -138,7 +138,7 @@ def _megatron_dpo_loss_callback(*, forward_only, disable_lora, kwargs=None, ref_ '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': False, 'loss_instance': DPOLoss(), + '_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}])), @@ -173,3 +173,61 @@ def test_megatron_non_reference_dpo_still_computes_a_nonzero_loss(forward_only): 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 From 67f5a746940464703b78e83436c3b1a39d3a7472 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 30 Oct 2026 03:08:59 +0000 Subject: [PATCH 6/6] fix: make CS model preparation reproducible --- .../client/twinkle/prepare_sequence_model.py | 39 ++++++++++++------- .../model/megatron/multi_lora_megatron.py | 6 ++- 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/cookbook/client/twinkle/prepare_sequence_model.py b/cookbook/client/twinkle/prepare_sequence_model.py index b4561264a..1d1f5d4e1 100644 --- a/cookbook/client/twinkle/prepare_sequence_model.py +++ b/cookbook/client/twinkle/prepare_sequence_model.py @@ -20,22 +20,23 @@ from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer -def tiny_backbone(directory, seed=42): - from tokenizers.pre_tokenizers import ByteLevel - from transformers import Qwen2Config, Qwen2ForCausalLM, Qwen2Tokenizer +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=512, hidden_size=64, intermediate_size=128, + 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=0, bos_token_id=2, eos_token_id=3, + 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) - vocab = {'[PAD]': 0, '[UNK]': 1, '[BOS]': 2, '[EOS]': 3} - # Use the model family's tokenizer so AutoTokenizer preserves its byte vocabulary. - vocab.update({token: index + 4 for index, token in enumerate(sorted(ByteLevel.alphabet()))}) - Qwen2Tokenizer(vocab=vocab, merges=[], pad_token='[PAD]', unk_token='[UNK]', - bos_token='[BOS]', eos_token='[EOS]').save_pretrained(directory) + tokenizer.save_pretrained(directory) return str(directory) @@ -75,6 +76,10 @@ def prepare_models(model_id, output_dir, *, seed=42, mcore_native=False): 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) @@ -82,6 +87,12 @@ def prepare_models(model_id, output_dir, *, seed=42, mcore_native=False): 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) @@ -119,14 +130,16 @@ def prepare_models(model_id, output_dir, *, seed=42, mcore_native=False): def main(): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument('--model', default='ms://Qwen/Qwen3.5-4B') + 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 offline Qwen2 backbone') + 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() - source = tiny_backbone(Path(args.output_dir) / 'source', args.seed) if args.tiny else args.model + 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) diff --git a/src/twinkle/model/megatron/multi_lora_megatron.py b/src/twinkle/model/megatron/multi_lora_megatron.py index 4e58ba4fa..19c3db52c 100644 --- a/src/twinkle/model/megatron/multi_lora_megatron.py +++ b/src/twinkle/model/megatron/multi_lora_megatron.py @@ -44,6 +44,9 @@ def _seq_cls_head_lora_load_gate(bridge): 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): @@ -485,7 +488,8 @@ def add_adapter_to_model( 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) + 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.bias = 'none' config_or_dir.init_lora_weights = False