diff --git a/angelslim/utils/vllm_calibration_dp.py b/angelslim/utils/vllm_calibration_dp.py index 0d81a7a5..d55c3b9c 100644 --- a/angelslim/utils/vllm_calibration_dp.py +++ b/angelslim/utils/vllm_calibration_dp.py @@ -55,8 +55,6 @@ def validate_vllm_calibration_dp_args(parser, args): parser.error("--tp-size must be >= 1") if args.dp_size < 1: parser.error("--dp-size must be >= 1") - if args.dp_timeout < 1: - parser.error("--dp-timeout must be >= 1") if args.dp_size > 1 and args.distributed_executor_backend != "ray": parser.error( @@ -193,19 +191,10 @@ def _run_ray_actor_calibration_dp(args, worker_fn: Callable[[Any], Any]) -> list # Wait for all placement groups to be ready ready_refs = [pg.ready() for pg in placement_groups] - ready, pending = ray.wait( + ray.wait( ready_refs, num_returns=len(ready_refs), - timeout=args.dp_timeout, ) - if pending: - raise TimeoutError( - "Timed out waiting for calibration DP placement groups: " - f"ready={len(ready)}, pending={len(pending)}, " - f"timeout={args.dp_timeout}s. " - "Check whether enough free GPUs are available for " - "dp_size * tp_size." - ) # Create actors for each DP rank ReplicaActor = ray.remote( @@ -234,18 +223,11 @@ def _run_ray_actor_calibration_dp(args, worker_fn: Callable[[Any], Any]) -> list actor.run.remote(args_dict, dp_rank) for dp_rank, actor in enumerate(actors) ] - # Wait for results with timeout - ready, pending = ray.wait( + # Wait for results + ray.wait( result_refs, num_returns=len(result_refs), - timeout=args.dp_timeout, ) - if pending: - raise TimeoutError( - "Timed out waiting for calibration DP actors: " - f"finished={len(ready)}, pending={len(pending)}, " - f"timeout={args.dp_timeout}s." - ) results = ray.get(result_refs) merged_activation_stats = _merge_dp_payloads(args.output_dir, results) @@ -271,17 +253,10 @@ def _run_ray_actor_calibration_dp(args, worker_fn: Callable[[Any], Any]) -> list actor.run_kv_search.remote(args_dict, dp_rank, activation_stats_ref) for dp_rank, actor in enumerate(actors) ] - ready, pending = ray.wait( + ray.wait( kv_refs, num_returns=len(kv_refs), - timeout=args.dp_timeout, ) - if pending: - raise TimeoutError( - "Timed out waiting for DP KV search actors: " - f"finished={len(ready)}, pending={len(pending)}, " - f"timeout={args.dp_timeout}s." - ) kv_results = ray.get(kv_refs) _merge_dp_kv_search_payloads( args.output_dir, diff --git a/tools/run_vllm_calibrate.py b/tools/run_vllm_calibrate.py index b0111c58..7bdc6c5b 100644 --- a/tools/run_vllm_calibrate.py +++ b/tools/run_vllm_calibrate.py @@ -245,12 +245,6 @@ def parse_args(): help="Calibration data parallel size (default: 1). When > 1, Ray Data " "creates one vLLM calibration actor per DP replica.", ) - parser.add_argument( - "--dp-timeout", - type=int, - default=7200, - help="Seconds to wait for Ray Data-managed DP workers before failing " "(default: 7200).", - ) # MTP (Multi-Token Prediction) configuration parser.add_argument(