From 1c9290848e221112bf22143a325de96640a19f92 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Wed, 11 Nov 2026 01:19:17 +0800 Subject: [PATCH 01/44] Kimi K3.1 RL runtime: Miles adapters, engine samplers, NC pool, cost guard, run up, sandbox hosts and dashboard Adds the Kimi K3.1 async GRPO path to yeto: the Miles adapters (ports island, sampler island, state plugin), the engine sampler with SGLang and vLLM adapters (R3 routed experts, top-K log-probs, fp32 lm-head), the National Compute sampler pool with elastic join/leave, the cost guard and resource ledger, `yeto run up` with budget and silence backstops, sandbox hosts for agentic tasks, the task dispenser and breaker, and the dashboard views for async GRPO runs. Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 1 + sitecustomize.py | 7 + yeto/cli.py | 257 +- yeto/cloud/modal_ckpt_store.py | 5 +- yeto/cloud/modal_engine_sampler.py | 338 ++ yeto/cloud/oss_store.py | 36 +- yeto/cloud/s3_store.py | 790 ++++- yeto/codepack.py | 25 +- yeto/cost/__init__.py | 10 + yeto/cost/billing.py | 311 ++ yeto/cost/cleanup.py | 239 ++ yeto/cost/guard.py | 306 ++ yeto/cost/idle.py | 154 + yeto/cost/ledger.py | 214 ++ yeto/cost/reaper.py | 74 + yeto/cost/report.py | 181 ++ yeto/dashboard/alerts.py | 48 +- yeto/dashboard/async_rl.py | 454 +++ yeto/dashboard/cli.py | 21 +- yeto/dashboard/cost.py | 13 +- yeto/dashboard/prices.example.json | 6 +- yeto/dashboard/reducer.py | 393 ++- yeto/dashboard/serve.py | 9 +- yeto/dashboard/sources.py | 29 +- yeto/dashboard/static/app.js | 114 +- yeto/dashboard/static/index.html | 63 +- yeto/dashboard/static/sft.js | 155 +- yeto/dashboard/store_tapes.py | 83 + yeto/data.py | 25 +- yeto/gpu_spec.py | 19 +- yeto/hw/card_telemetry.py | 243 +- yeto/hw/card_tools.py | 85 +- yeto/hw/catalog.py | 8 +- yeto/island_credential_guard.py | 13 +- yeto/island_overrides.py | 2 +- yeto/launch/agentic_node_agent.py | 95 + yeto/launch/agentic_setup.py | 984 ++++++ yeto/launch/args.py | 6 +- yeto/launch/code_pack.py | 6 +- yeto/launch/convertible_pool.py | 30 +- yeto/launch/fleet.py | 36 + yeto/launch/images.py | 14 +- yeto/launch/k31_checks.py | 210 ++ yeto/launch/manifest.py | 18 +- yeto/launch/modal_islands.py | 63 +- yeto/launch/nc/__init__.py | 7 + yeto/launch/nc/__main__.py | 3 + yeto/launch/nc/cli.py | 976 ++++++ yeto/launch/nc/client.py | 346 ++ yeto/launch/nc/elastic.py | 483 +++ yeto/launch/nc/lease_pool.py | 954 ++++++ yeto/launch/nc/node_pool.py | 263 ++ yeto/launch/nc/pool_run.py | 719 +++++ yeto/launch/nc/port_open.mjs | 32 + yeto/launch/nc/provider.py | 100 + yeto/launch/nc/relay.mjs | 72 + yeto/launch/nc/sky_provider.py | 175 ++ yeto/launch/nc/ssh_sampler.py | 369 +++ yeto/launch/nc/tasks.py | 265 ++ yeto/launch/nc/tunnel.py | 287 ++ yeto/launch/node_inventory.py | 8 +- yeto/launch/peer_recovery.py | 13 +- yeto/launch/rl_args.py | 105 +- yeto/launch/rl_backend.py | 10 +- yeto/launch/rl_island_task.py | 113 +- yeto/launch/rl_ports.py | 14 +- yeto/launch/storage.py | 109 +- yeto/launch/trainer_island.py | 18 +- yeto/launch/trainer_loop.py | 4 +- yeto/launch_preflight.py | 65 +- yeto/launcher.py | 91 +- yeto/megatron/miles_value_island.py | 17 +- yeto/modal_runner.py | 92 +- yeto/models.py | 1 + yeto/provenance.py | 78 +- yeto/rl/__init__.py | 63 + yeto/rl/adapters/miles/algo_flag_rows.py | 36 + yeto/rl/adapters/miles/algorithm_flags.py | 8 + yeto/rl/adapters/miles/buffer_rollout.py | 453 ++- .../rl/adapters/miles/buffer_rollout_miles.py | 82 +- yeto/rl/adapters/miles/carry_over.py | 27 +- yeto/rl/adapters/miles/config.py | 436 ++- yeto/rl/adapters/miles/convertible.py | 6 +- yeto/rl/adapters/miles/cut_injection.py | 39 +- yeto/rl/adapters/miles/cut_plugin.py | 126 +- yeto/rl/adapters/miles/dp_invariant_state.py | 575 ++++ yeto/rl/adapters/miles/entry.py | 229 +- yeto/rl/adapters/miles/eval_wiring.py | 2 +- yeto/rl/adapters/miles/ext_args.py | 37 + .../adapters/miles/full_parameter/evidence.py | 23 +- .../miles/full_parameter/inbound_sink.py | 20 + .../miles/full_parameter/master_check.py | 187 ++ .../miles/full_parameter/rank_state.py | 168 + .../miles/full_parameter/rank_stats.py | 12 +- .../miles/full_parameter/role_stream.py | 4 + .../miles/full_parameter/sao_binding.py | 6 +- .../adapters/miles/full_parameter/session.py | 3 +- yeto/rl/adapters/miles/harbor_agent.py | 7 +- .../adapters/miles/harness_glue/tool_wait.py | 3 +- yeto/rl/adapters/miles/island_args.py | 341 ++ yeto/rl/adapters/miles/island_entry.py | 232 +- .../miles/models/chunked_full_parameter.py | 18 +- .../miles/models/deepseek_v4/bridge.py | 6 +- .../miles/models/deepseek_v4/expert_clone.py | 24 +- .../miles/models/deepseek_v4/expert_full.py | 89 +- .../models/deepseek_v4/expert_full_runtime.py | 235 +- .../miles/models/deepseek_v4/ports_runtime.py | 7 +- .../adapters/miles/models/full_parameter.py | 337 -- .../full_parameter_continuation_probe.py | 389 --- .../models/full_parameter_manifest_probe.py | 80 +- .../miles/models/full_parameter_probe.py | 266 -- .../miles/models/kimi_k31/__init__.py | 17 + .../miles/models/kimi_k31/first_stage_grad.py | 68 + .../adapters/miles/models/kimi_k31/model.py | 156 + .../miles/models/kimi_k31/nan_inject.py | 59 + .../miles/models/kimi_k31/ports_recipe.py | 248 ++ .../miles/models/kimi_k31/ports_runtime.py | 72 + .../miles/models/kimi_k31/provider.py | 130 + .../adapters/miles/models/kimi_k31/release.py | 89 + yeto/rl/adapters/miles/models/nemotron_h.py | 8 +- .../miles/models/partial_experts/__init__.py | 23 + .../miles/models/partial_experts/attest.py | 59 + .../miles/models/partial_experts/config.py | 111 + .../miles/models/partial_experts/export.py | 50 + .../miles/models/partial_experts/freeze.py | 94 + .../miles/models/partial_experts/hooks.py | 86 + .../miles/models/partial_experts/optimizer.py | 133 + .../miles/models/partial_experts/selection.py | 648 ++++ yeto/rl/adapters/miles/native_slots_state.py | 196 ++ .../rl/adapters/miles/offpolicy_correction.py | 26 +- yeto/rl/adapters/miles/overlay.py | 171 - .../miles/overlays/miles-critic-c357.patch | 2786 ----------------- yeto/rl/adapters/miles/pins.py | 423 ++- yeto/rl/adapters/miles/policy_age.py | 8 +- yeto/rl/adapters/miles/proc_state.py | 50 + yeto/rl/adapters/miles/publish.py | 188 +- yeto/rl/adapters/miles/records_generate.py | 10 +- yeto/rl/adapters/miles/release_plugin.py | 46 +- yeto/rl/adapters/miles/remote_code_prewarm.py | 4 +- yeto/rl/adapters/miles/reshard.py | 2 +- yeto/rl/adapters/miles/revision.py | 29 +- yeto/rl/adapters/miles/role_switch.py | 93 + yeto/rl/adapters/miles/rollout.py | 51 +- yeto/rl/adapters/miles/rollout_meta_hook.py | 81 +- yeto/rl/adapters/miles/run_config_rules.py | 36 +- yeto/rl/adapters/miles/run_flags.py | 92 +- yeto/rl/adapters/miles/runtime_manifest.py | 149 +- yeto/rl/adapters/miles/sampler_island.py | 184 +- yeto/rl/adapters/miles/state_plugin.py | 297 +- yeto/rl/adapters/miles/synth_rollout.py | 150 + yeto/rl/adapters/miles/train_gap_probe.py | 12 +- yeto/rl/adapters/miles/trainer.py | 8 +- yeto/rl/adapters/miles/trainer_rebuild.py | 44 +- yeto/rl/adapters/sampler_engines/__init__.py | 2 +- yeto/rl/adapters/sampler_engines/fake.py | 85 +- yeto/rl/adapters/sampler_engines/from_spec.py | 113 +- .../adapters/sampler_engines/k31_release.py | 293 ++ .../rl/adapters/sampler_engines/quant_tier.py | 16 +- .../sampler_engines/sglang_mxfp4_reload.py | 334 ++ .../adapters/sampler_engines/sglang_server.py | 865 ++++- .../sampler_engines/vllm_expert_reload.py | 338 ++ .../adapters/sampler_engines/vllm_server.py | 620 +++- yeto/rl/adapters/verl/buffer_feed.py | 28 +- yeto/rl/adapters/verl/image.py | 2 +- yeto/rl/adapters/verl/rollout_engine.py | 4 +- yeto/rl/agentic/finish.py | 51 +- yeto/rl/agentic/tokens.py | 59 +- yeto/rl/agentic/types.py | 46 +- yeto/rl/algos/__init__.py | 1 + yeto/rl/algos/critic_fork.py | 34 +- yeto/rl/algos/loss_variants.py | 16 +- yeto/rl/algos/mismatch_correction.py | 4 +- yeto/rl/algos/reward_pipeline.py | 19 +- yeto/rl/algos/sample_filters.py | 9 +- yeto/rl/algos/sao.py | 19 +- yeto/rl/algos/score_centering.py | 308 ++ yeto/rl/algos/seq_adv.py | 3 +- yeto/rl/algos/vendor/miles_mis.py | 2 +- yeto/rl/breaker.py | 286 ++ yeto/rl/core.py | 5 +- yeto/rl/critic_warmup.py | 2 +- yeto/rl/engine/agentic_round.py | 73 +- yeto/rl/engine/algorithm.py | 12 +- yeto/rl/engine/backends.py | 4 +- yeto/rl/engine/bridges.py | 28 +- yeto/rl/engine/buffer_payload.py | 234 +- yeto/rl/engine/cut.py | 122 +- yeto/rl/engine/driver.py | 183 +- yeto/rl/engine/engine_sampler.py | 159 +- yeto/rl/engine/engine_sampler_entry.py | 259 +- yeto/rl/engine/events.py | 6 +- yeto/rl/engine/ext_attrs.py | 150 + yeto/rl/engine/fake.py | 7 + yeto/rl/engine/incarnation_lease.py | 23 + yeto/rl/engine/init_from_release.py | 182 ++ yeto/rl/engine/kimi_k31_lora.py | 204 ++ yeto/rl/engine/loop_guard.py | 6 +- yeto/rl/engine/lora_release_merge.py | 168 + yeto/rl/engine/merged_release.py | 84 +- yeto/rl/engine/nan_skip.py | 74 + yeto/rl/engine/object_store.py | 219 +- yeto/rl/engine/policy_age.py | 70 +- yeto/rl/engine/policy_release.py | 242 +- yeto/rl/engine/ports.py | 13 + yeto/rl/engine/proc_state.py | 149 + yeto/rl/engine/release_overlay.py | 64 + yeto/rl/engine/resume.py | 85 +- yeto/rl/engine/role_switch.py | 572 ++++ yeto/rl/engine/round_cut.py | 32 +- yeto/rl/engine/run_config.py | 57 +- yeto/rl/engine/run_flags.py | 26 + yeto/rl/engine/runtime_manifest.py | 4 + yeto/rl/engine/sample_buffer.py | 83 +- yeto/rl/engine/sample_roles.py | 177 +- yeto/rl/engine/sampler_checks.py | 112 +- yeto/rl/engine/sampler_engine.py | 79 +- yeto/rl/engine/sampler_loop.py | 38 +- yeto/rl/engine/sampler_resume.py | 74 + yeto/rl/engine/selection.py | 17 +- yeto/rl/engine/store_dispense.py | 1093 +++++++ yeto/rl/engine/tape_upload.py | 172 + yeto/rl/engine/task_dispense.py | 40 +- yeto/rl/engine/xhw_guard.py | 357 +++ yeto/rl/filters.py | 24 +- yeto/rl/grad_audit.py | 4 +- yeto/rl/harness/chat_templates.py | 143 +- yeto/rl/harness/gateway/chains.py | 40 +- yeto/rl/harness/gateway/proxy.py | 143 +- yeto/rl/harness/gateway/recorder.py | 20 +- yeto/rl/harness/gateway/records.py | 10 + yeto/rl/harness/gateway/records_env.py | 1 + yeto/rl/harness/harbor/README.md | 23 + yeto/rl/harness/harbor/capacity.py | 148 + yeto/rl/harness/harbor/pi_agent.py | 59 + yeto/rl/harness/harbor/pi_config.py | 109 + yeto/rl/harness/harbor/pins.py | 8 + yeto/rl/harness/harbor/remote_docker_env.py | 86 + yeto/rl/harness/harbor/task_lists.py | 109 +- yeto/rl/harness/harbor/trial.py | 92 +- yeto/rl/harness/placeholder_logprob_mask.py | 6 +- yeto/rl/harness/reasoning_loss_mask.py | 11 +- yeto/rl/image_registry.py | 31 + yeto/rl/quant/release.py | 2 +- yeto/rl/sampler_pins.py | 158 + yeto/rl/spot/__init__.py | 3 +- yeto/rl/spot/peer_fanout.py | 251 ++ yeto/rl/spot/s3_delivery.py | 66 +- yeto/rl/ssh_harness.py | 131 +- yeto/rl/up.py | 585 ++++ yeto/rl/wandb_rl.py | 5 +- yeto/runs.py | 49 +- yeto/sandbox/__init__.py | 1 + yeto/sandbox/cli.py | 111 + yeto/sandbox/docker.py | 189 ++ yeto/sandbox/envs.py | 299 ++ yeto/sandbox/gate.py | 180 ++ yeto/sandbox/ops.py | 474 +++ yeto/sandbox/ready.py | 114 + yeto/sandbox/trial_labels.py | 98 + yeto/shape/catalog.py | 1 + 260 files changed, 31318 insertions(+), 5528 deletions(-) create mode 100644 yeto/cloud/modal_engine_sampler.py create mode 100644 yeto/cost/__init__.py create mode 100644 yeto/cost/billing.py create mode 100644 yeto/cost/cleanup.py create mode 100644 yeto/cost/guard.py create mode 100644 yeto/cost/idle.py create mode 100644 yeto/cost/ledger.py create mode 100644 yeto/cost/reaper.py create mode 100644 yeto/cost/report.py create mode 100644 yeto/dashboard/async_rl.py create mode 100644 yeto/dashboard/store_tapes.py create mode 100644 yeto/launch/agentic_node_agent.py create mode 100644 yeto/launch/agentic_setup.py create mode 100644 yeto/launch/k31_checks.py create mode 100644 yeto/launch/nc/__init__.py create mode 100644 yeto/launch/nc/__main__.py create mode 100644 yeto/launch/nc/cli.py create mode 100644 yeto/launch/nc/client.py create mode 100644 yeto/launch/nc/elastic.py create mode 100644 yeto/launch/nc/lease_pool.py create mode 100644 yeto/launch/nc/node_pool.py create mode 100644 yeto/launch/nc/pool_run.py create mode 100644 yeto/launch/nc/port_open.mjs create mode 100644 yeto/launch/nc/provider.py create mode 100644 yeto/launch/nc/relay.mjs create mode 100644 yeto/launch/nc/sky_provider.py create mode 100644 yeto/launch/nc/ssh_sampler.py create mode 100644 yeto/launch/nc/tasks.py create mode 100644 yeto/launch/nc/tunnel.py create mode 100644 yeto/rl/adapters/miles/dp_invariant_state.py create mode 100644 yeto/rl/adapters/miles/ext_args.py create mode 100644 yeto/rl/adapters/miles/full_parameter/master_check.py create mode 100644 yeto/rl/adapters/miles/full_parameter/rank_state.py create mode 100644 yeto/rl/adapters/miles/island_args.py delete mode 100644 yeto/rl/adapters/miles/models/full_parameter.py delete mode 100644 yeto/rl/adapters/miles/models/full_parameter_continuation_probe.py delete mode 100644 yeto/rl/adapters/miles/models/full_parameter_probe.py create mode 100644 yeto/rl/adapters/miles/models/kimi_k31/__init__.py create mode 100644 yeto/rl/adapters/miles/models/kimi_k31/first_stage_grad.py create mode 100644 yeto/rl/adapters/miles/models/kimi_k31/model.py create mode 100644 yeto/rl/adapters/miles/models/kimi_k31/nan_inject.py create mode 100644 yeto/rl/adapters/miles/models/kimi_k31/ports_recipe.py create mode 100644 yeto/rl/adapters/miles/models/kimi_k31/ports_runtime.py create mode 100644 yeto/rl/adapters/miles/models/kimi_k31/provider.py create mode 100644 yeto/rl/adapters/miles/models/kimi_k31/release.py create mode 100644 yeto/rl/adapters/miles/models/partial_experts/__init__.py create mode 100644 yeto/rl/adapters/miles/models/partial_experts/attest.py create mode 100644 yeto/rl/adapters/miles/models/partial_experts/config.py create mode 100644 yeto/rl/adapters/miles/models/partial_experts/export.py create mode 100644 yeto/rl/adapters/miles/models/partial_experts/freeze.py create mode 100644 yeto/rl/adapters/miles/models/partial_experts/hooks.py create mode 100644 yeto/rl/adapters/miles/models/partial_experts/optimizer.py create mode 100644 yeto/rl/adapters/miles/models/partial_experts/selection.py create mode 100644 yeto/rl/adapters/miles/native_slots_state.py delete mode 100644 yeto/rl/adapters/miles/overlay.py delete mode 100644 yeto/rl/adapters/miles/overlays/miles-critic-c357.patch create mode 100644 yeto/rl/adapters/miles/proc_state.py create mode 100644 yeto/rl/adapters/miles/role_switch.py create mode 100644 yeto/rl/adapters/miles/synth_rollout.py create mode 100644 yeto/rl/adapters/sampler_engines/k31_release.py create mode 100644 yeto/rl/adapters/sampler_engines/sglang_mxfp4_reload.py create mode 100644 yeto/rl/adapters/sampler_engines/vllm_expert_reload.py create mode 100644 yeto/rl/algos/score_centering.py create mode 100644 yeto/rl/breaker.py create mode 100644 yeto/rl/engine/ext_attrs.py create mode 100644 yeto/rl/engine/init_from_release.py create mode 100644 yeto/rl/engine/kimi_k31_lora.py create mode 100644 yeto/rl/engine/lora_release_merge.py create mode 100644 yeto/rl/engine/nan_skip.py create mode 100644 yeto/rl/engine/proc_state.py create mode 100644 yeto/rl/engine/release_overlay.py create mode 100644 yeto/rl/engine/role_switch.py create mode 100644 yeto/rl/engine/store_dispense.py create mode 100644 yeto/rl/engine/tape_upload.py create mode 100644 yeto/rl/engine/xhw_guard.py create mode 100644 yeto/rl/harness/harbor/capacity.py create mode 100644 yeto/rl/harness/harbor/pi_agent.py create mode 100644 yeto/rl/harness/harbor/pi_config.py create mode 100644 yeto/rl/harness/harbor/remote_docker_env.py create mode 100644 yeto/rl/image_registry.py create mode 100644 yeto/rl/sampler_pins.py create mode 100644 yeto/rl/spot/peer_fanout.py create mode 100644 yeto/rl/up.py create mode 100644 yeto/sandbox/__init__.py create mode 100644 yeto/sandbox/cli.py create mode 100644 yeto/sandbox/docker.py create mode 100644 yeto/sandbox/envs.py create mode 100644 yeto/sandbox/gate.py create mode 100644 yeto/sandbox/ops.py create mode 100644 yeto/sandbox/ready.py create mode 100644 yeto/sandbox/trial_labels.py diff --git a/pyproject.toml b/pyproject.toml index 40654bf7c..82195ac60 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,4 +74,5 @@ include = ["yeto*"] [tool.setuptools.package-data] "yeto.dashboard" = ["*.json", "static/*"] "yeto.cloud" = ["*.json"] +"yeto.launch.nc" = ["*.mjs"] "yeto.rl.adapters.miles" = ["*.json", "overlays/*.patch"] diff --git a/sitecustomize.py b/sitecustomize.py index 5f45e2cd1..51a539ce3 100644 --- a/sitecustomize.py +++ b/sitecustomize.py @@ -24,6 +24,13 @@ install() +if os.environ.get("YETO_SGLANG_MXFP4_INPLACE") == "1": + # In-place reload of MXFP4 routed experts on SGLang's SM100 trtllm-gen path + # (every SGLang process, including the spawned schedulers). + from yeto.rl.adapters.sampler_engines.sglang_mxfp4_reload import install as _install_mxfp4_inplace + + _install_mxfp4_inplace() + def _install_tms_post_pause_hold() -> None: """Hold a Miles trainer actor after TMS pause for crash attribution. diff --git a/yeto/cli.py b/yeto/cli.py index 8abe64209..4dd0452b6 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -47,9 +47,22 @@ "stop-run", "_worker", "_head", + "cost", + "cleanup", + "breaker", + "_reaper", ) + + +class _DefaultLR(float): + """The --inner-lr default (an SFT learning rate): its type tells an unset flag + from an explicit 3e-4 (yeto.launch.rl_args refuses it for full-parameter RL).""" + + +DEFAULT_INNER_LR = _DefaultLR(3e-4) + def _add_launch_args(p: argparse.ArgumentParser) -> None: p.add_argument( "--gpu", @@ -237,7 +250,8 @@ def loss_spec(value: str) -> str: default="auto", help="adapter placement: attention-only, expanded-V4 " "attention+routed-expert clones, every linear, or auto " - "(attention for ordinary MoE β€” router and routed experts stay frozen)", + "(attention for ordinary MoE and Kimi-K3.1 β€” router and routed experts " + "stay frozen; all-linear must be named explicitly)", ) parent = tune.add_mutually_exclusive_group() parent.add_argument( @@ -309,7 +323,7 @@ def int_or_auto(value: str): help="accumulation steps with an explicit micro batch; with 'auto', " "the requested per-rank effective sequence batch", ) - tune.add_argument("--inner-lr", type=float, default=3e-4) + tune.add_argument("--inner-lr", type=float, default=DEFAULT_INNER_LR) tune.add_argument("--max-rows", type=int, default=None, help="cap dataset rows per learner") tune.add_argument( "--tokenize", @@ -633,6 +647,31 @@ def int_or_auto(value: str): "of shipping the sandbox token from the launching environment; the island setup " "fails closed when either variable is missing. Refused for non-Modal islands", ) + infra.add_argument( + "--modal-store-secret", + default=None, + metavar="NAME", + help="Modal islands with an s3:// --rl-checkpoint-store: attach the named Modal Secret " + "(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, and AWS_ENDPOINT_URL / AWS_REGION for an " + "S3-compatible store such as R2) to the island function instead of shipping the store " + "credentials from the launching environment. Refused for non-Modal islands", + ) + infra.add_argument( + "--rl-object-store-endpoint", + default=None, + metavar="URL", + help="S3-compatible endpoint of an s3:// --rl-checkpoint-store (R2, MinIO), for the head " + "and every island; default: AWS_ENDPOINT_URL from the environment, the endpoint_url of " + "--rl-object-store-profile, else AWS S3", + ) + infra.add_argument( + "--rl-object-store-profile", + default=None, + metavar="NAME", + help="AWS profile (~/.aws) of an s3:// --rl-checkpoint-store: the head reads with it, and its " + "key, endpoint_url and region go to every island (unless --modal-store-secret); " + "default: YETO_RL_STORE_S3_PROFILE", + ) infra.add_argument( "--modal-timeout-s", type=int, @@ -746,6 +785,9 @@ def int_or_auto(value: str): from .launch import trainer_loop as _trainer_loop _trainer_loop.add_arguments(infra) + from .cost.report import add_launch_arguments as _add_cost_launch_arguments + + _add_cost_launch_arguments(p) # W&B telemetry. One group per fleet (this run's name), one run per # island plus one for the syncer's event tape; see docs/WANDB.md. @@ -862,7 +904,8 @@ def build_parser() -> argparse.ArgumentParser: ) sub = p.add_subparsers( dest="command", - metavar="{launch,shape,merge,rl,sample-diffusion,status,logs,down}", + metavar="{launch,shape,merge,rl,sample-diffusion,engine-sampler-modal,status," + "dashboard,logs,down,stop-run}", ) launch = sub.add_parser( @@ -1015,6 +1058,22 @@ def build_parser() -> argparse.ArgumentParser: ) _add_diffusion_sample_args(sample) + engine_sampler = sub.add_parser( + "engine-sampler-modal", + help="start one engine sampler island (SGLang) on Modal that writes its own buffer shard", + ) + _add_engine_sampler_modal_args(engine_sampler) + + from .cost.report import add_parsers as _add_cost_parsers + + _add_cost_parsers(sub) + from .rl.breaker import add_parser as _add_breaker_parser + + _add_breaker_parser(sub) + reaper = sub.add_parser("_reaper") + reaper.add_argument("run") + reaper.add_argument("pid", type=int) + status = sub.add_parser("status", help="table of known runs") status.add_argument( "--tape", @@ -1026,6 +1085,14 @@ def build_parser() -> argparse.ArgumentParser: _add_dashboard_parser(sub) + from .sandbox.cli import add_parser as _add_sandbox_parser + + _add_sandbox_parser(sub) + + from .rl.up import add_parser as _add_rl_up_parser + + _add_rl_up_parser(sub) + logs = sub.add_parser("logs", help="stream a run's launcher log (Ctrl-C detaches)") logs.add_argument("run", help="run name (its --cluster-prefix)") logs.add_argument( @@ -1064,10 +1131,14 @@ def _spawn_worker(name: str) -> subprocess.Popen: survives this CLI exiting and never touches the controlling TTY; `-u` keeps its output unbuffered so log tailing is live. """ + from .cost.ledger import LEDGER_ENV, ResourceLedger + env = dict(os.environ) env["YETO_RUNS_DIR"] = str(runs.RUNS_DIR) + # processes of this run (and the reaper) record / read the same resource ledger + env[LEDGER_ENV] = str(ResourceLedger.for_run(name).path) with open(runs.log_path(name), "ab") as log_f: - return subprocess.Popen( + proc = subprocess.Popen( [sys.executable, "-u", "-m", "yeto.cli", "_worker", name], stdout=log_f, stderr=subprocess.STDOUT, @@ -1075,6 +1146,21 @@ def _spawn_worker(name: str) -> subprocess.Popen: start_new_session=True, env=env, ) + _spawn_reaper(name, proc.pid, env, log_f) + return proc + + +def _spawn_reaper(name: str, worker_pid: int, env: dict, log_f) -> None: + """A detached watcher in its own session: when the worker dies (also by SIGKILL) + it releases whatever the run's resource ledger still lists (yeto.cost.reaper).""" + try: + subprocess.Popen( + [sys.executable, "-u", "-m", "yeto.cli", "_reaper", name, str(worker_pid)], + stdout=log_f, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, + start_new_session=True, env=env) + except OSError as e: + print(f"[yeto] WARNING: no resource reaper for {name} ({e}); after a crash run " + f"`yeto cleanup --run {name}`", file=sys.stderr) def _stream_log(name: str, follow: bool = True, alive=None) -> None: @@ -1222,14 +1308,22 @@ def cmd_launch(args) -> int: code = _resolve_auto_fleet(args) if code: return 0 if code == 2 else code # user said no: clean exit + if "modal:" in str(args.gpu or "") and not getattr(args, "dry_run", False): + from .launch_preflight import modal_sdk_problem + + problem = modal_sdk_problem() # the worker drives Modal through this interpreter + if problem is not None: + print(f"[yeto] --gpu {args.gpu}: {problem}", file=sys.stderr) + return 1 name = args.cluster_prefix existing = runs.load_run(name) - if existing is not None and runs.is_alive(existing.get("pid")): + live = sorted(set(runs.worker_pids(name)) | ( + {int(existing["pid"])} if existing is not None and runs.is_alive(existing.get("pid")) else set())) + if live: print( - f"[yeto] run '{name}' already has a live worker " - f"(pid {existing['pid']}). Use `yeto logs {name}` to attach, " - f"`yeto down {name}` to stop it, or pick a different " - f"--cluster-prefix.", + f"[yeto] run '{name}' already has a live worker (pid {', '.join(map(str, live))}). " + f"Run `yeto down {name}` to stop it and its clusters, `yeto logs {name}` to attach, " + f"or pick a different --cluster-prefix.", file=sys.stderr, ) return 1 @@ -1291,6 +1385,47 @@ def cmd_launch(args) -> int: return code +def _add_engine_sampler_modal_args(p: argparse.ArgumentParser) -> None: + p.description = ("Configuration from the environment (RUN, SHARD, SELECTION, PROMPTS, RELOAD, DROP_ZV, " + "N, MAX_TOKENS, GROUPS, BACKLOG, THINKING_EFFORT, DATA_SLOT, EP, ...; " + "yeto.cloud.modal_engine_sampler); the flags below override it.") + p.add_argument("--run", default=None, help="run name: the store is s3://BUCKET/runs/RUN") + p.add_argument("--shard", default=None, help="this sampler's buffer shard (the trainer lists it " + "in --rl-buffer-shards)") + p.add_argument("--data-slot", default=None, metavar="I/N", + help="read prompt rows i with i %% N == I (one slot per sampler)") + p.add_argument("--bucket", default=None) + p.add_argument("--selection", default=None, help="expert selection JSON (its sha256 gates releases)") + p.add_argument("--prompts", default=None, help="local prompts JSONL") + p.add_argument("--reload", choices=["hot", "restart"], default=None) + p.add_argument("--secret", default=None, help="Modal Secret with the store credentials and endpoint") + p.add_argument("--app-name", default=None) + p.add_argument("--image", default=None, help="sampler image /@sha256: " + "(required: no default; Kimi-K3.1 samplers run on A-card vLLM nodes, not here)") + p.add_argument("--dry-run", action="store_true", help="print the spec and checks; no Modal call") + p.add_argument("--no-wait", action="store_true", help="return after the spawn (print the call id)") + + +def cmd_engine_sampler_modal(args) -> int: + from .cloud import modal_engine_sampler as mes + + try: + config = mes.SamplerConfig.from_env( + run=args.run, shard=args.shard, data_slot=args.data_slot, bucket=args.bucket, + selection=args.selection, prompts=args.prompts, reload=args.reload, secret=args.secret, + app_name=args.app_name, image=getattr(args, "image", None)) + report = mes.dry_run(config) + except (ValueError, KeyError, OSError) as e: + print(f"[sampler] {e}", file=sys.stderr) + return 2 + if args.dry_run or report["problems"]: + print(json.dumps(report, indent=2)) + return 1 if report["problems"] else 0 + result = mes.launch(config, wait=not args.no_wait) + print(json.dumps(result) if isinstance(result, dict) else result) + return 0 if not isinstance(result, dict) or result.get("rc") == 0 else 1 + + def cmd_sample_diffusion(args) -> int: if bool(args.prompt) == bool(args.data): print("[yeto] pass exactly one of --prompt or --data", file=sys.stderr) @@ -1741,6 +1876,15 @@ def record_instance_ids(cluster, ids) -> None: from .launcher import run as launcher_run + def on_term(signum, _frame): + # `yeto down` / a supervisor: unwind through the launcher's teardown (finally) + raise SystemExit(128 + signum) + + for sig in (signal.SIGTERM, signal.SIGHUP): + try: + signal.signal(sig, on_term) + except (ValueError, OSError): + pass try: import inspect @@ -1755,8 +1899,10 @@ def record_instance_ids(cluster, ids) -> None: traceback.print_exc() runs.update_run(name, state=runs.FAILED, exit_code=1, finished_at=time.time()) + _backstop_cleanup(name, args) return 1 code = int(code or 0) + _backstop_cleanup(name, args) runs.update_run( name, state=runs.SUCCEEDED if code == 0 else runs.FAILED, @@ -1766,6 +1912,25 @@ def record_instance_ids(cluster, ids) -> None: return code +def _backstop_cleanup(name: str, args) -> None: + """After the launcher (its own teardown ran): release whatever the run's resource + ledger still lists (e.g. a Modal app the teardown could not stop, a sampler app a + script of this run started, an NC lease). Skipped with --keep.""" + if getattr(args, "keep", False): + return + try: + from .cost.cleanup import cleanup_run + from .cost.ledger import ResourceLedger + + if ResourceLedger.for_run(name).open_resources(): + print(f"[yeto] resource ledger of {name} still lists open resources; releasing them", + flush=True) + cleanup_run(name) + except Exception as e: # noqa: BLE001 + print(f"[yeto] WARNING: backstop cleanup failed ({e}); run `yeto cleanup --run {name}`", + file=sys.stderr, flush=True) + + # --------------------------------------------------------------------------- # status @@ -2050,7 +2215,8 @@ def _modal_stop_app(run_name: str) -> None: try: ModalOps(modal_app_name(run_name)).stop_app() - print(f"[yeto] Modal app {modal_app_name(run_name)}: stopped") + print(f"[yeto] Modal app {modal_app_name(run_name)}: stop requested (modal app stop --yes " + "exited 0); checking its state") except Exception as e: # the status check below decides; this is advisory if "already stopped" in str(e): print(f"[yeto] Modal app {modal_app_name(run_name)}: already stopped") @@ -2071,6 +2237,24 @@ def _signal_worker(pid: int, sig: int) -> None: def cmd_down(args) -> int: + code = _cmd_down_clusters(args) + meta = runs.load_run(args.run) + if meta is None: + return code + # then every resource of the run's ledger (by exact id), e.g. sampler apps / NC leases + # that are no cluster of the run record + from .cost.cleanup import cleanup_run + from .cost.ledger import ResourceLedger + + if ResourceLedger.for_run(args.run).open_resources(): + if not cleanup_run(args.run): + print(f"[yeto] resource ledger of {args.run}: not every resource confirmed released; " + f"rerun `yeto cleanup --run {args.run}`", file=sys.stderr) + return code or 1 + return code + + +def _cmd_down_clusters(args) -> int: name = args.run meta = runs.load_run(name) if meta is None: @@ -2079,15 +2263,28 @@ def cmd_down(args) -> int: return 1 pid = meta.get("pid") - if runs.is_alive(pid): - print(f"[yeto] stopping worker pid {pid} (SIGTERM)") - _signal_worker(int(pid), signal.SIGTERM) + # the recorded worker plus any other `_worker ` process (a launcher killed + # before it recorded the pid, or a second worker) -- each by PID + pids = sorted(set(runs.worker_pids(name)) | ({int(pid)} if runs.is_alive(pid) else set())) + if pids: + print(f"[yeto] stopping worker pid {', '.join(map(str, pids))} (SIGTERM)") + for p in pids: + _signal_worker(p, signal.SIGTERM) deadline = time.time() + 10 - while runs.is_alive(pid) and time.time() < deadline: + while any(runs.is_alive(p) for p in pids) and time.time() < deadline: time.sleep(0.2) - if runs.is_alive(pid): - print(f"[yeto] worker pid {pid} did not exit; sending SIGKILL") - _signal_worker(int(pid), signal.SIGKILL) + stuck = [p for p in pids if runs.is_alive(p)] + if stuck: + print(f"[yeto] worker pid {', '.join(map(str, stuck))} did not exit; sending SIGKILL") + for p in stuck: + _signal_worker(p, signal.SIGKILL) + deadline = time.time() + 5 + while any(runs.is_alive(p) for p in stuck) and time.time() < deadline: + time.sleep(0.2) + left = [p for p in pids if runs.is_alive(p)] + if left: + print(f"[yeto] worker pid {', '.join(map(str, left))} is still alive after SIGKILL", + file=sys.stderr) else: print("[yeto] worker is not running") @@ -2319,12 +2516,22 @@ def main(argv=None) -> int: return 0 if args.command == "sample-diffusion": return cmd_sample_diffusion(args) + if args.command == "engine-sampler-modal": + return cmd_engine_sampler_modal(args) if args.command == "status": return cmd_status(args) if args.command == "dashboard": from .dashboard.cli import main as dashboard_main return dashboard_main(args) + if args.command == "run": + from .rl.up import main as rl_up_main + + return rl_up_main(args) + if args.command == "sandbox": + from .sandbox.cli import main as sandbox_main + + return sandbox_main(args) if args.command == "logs": return cmd_logs(args) if args.command == "down": @@ -2334,6 +2541,22 @@ def main(argv=None) -> int: return 0 if args.command == "_worker": return cmd_worker(args.run) + if args.command == "breaker": + from .rl.breaker import cmd_breaker + + return cmd_breaker(args) + if args.command == "cost": + from .cost.report import cmd_cost + + return cmd_cost(args) + if args.command == "cleanup": + from .cost.report import cmd_cleanup + + return cmd_cleanup(args) + if args.command == "_reaper": + from .cost.reaper import reap + + return reap(args.run, args.pid) if args.command == "_head": return cmd_head(args.args_json) parser.print_help() diff --git a/yeto/cloud/modal_ckpt_store.py b/yeto/cloud/modal_ckpt_store.py index 9100d5cfb..d41e5b179 100644 --- a/yeto/cloud/modal_ckpt_store.py +++ b/yeto/cloud/modal_ckpt_store.py @@ -203,8 +203,9 @@ def _volume(self) -> Any: return modal.Volume.from_name(self.volume_name) - def commit(self) -> float: - """Explicit commit; a failing commit raises (the sync then reports store_synced + def commit(self, paths: Any = None) -> float: + """Explicit commit (``paths`` is accepted and ignored: a Volume commit makes + every write of the container durable); a failing commit raises (the sync then reports store_synced False). In the island container the learner's Python cannot import the Modal client, so the commit runs in the Modal runner's interpreter (``YETO_MODAL_PYTHON`` / ``YETO_MODAL_SYSPATH``, exported by modal_runner).""" diff --git a/yeto/cloud/modal_engine_sampler.py b/yeto/cloud/modal_engine_sampler.py new file mode 100644 index 000000000..46491062d --- /dev/null +++ b/yeto/cloud/modal_engine_sampler.py @@ -0,0 +1,338 @@ +"""Cloud layer: one engine sampler island on Modal (``yeto engine-sampler-modal``). + +One Modal GPU function runs ``python -m yeto.rl.engine.engine_sampler --spec`` +(the sglang-server adapter) against the run's object store: it writes its own +buffer shard (``SamplerConfig.shard``) and follows the trainer's policy releases. +The trainer reads it with ``--rl-buffer-shards`` naming this shard. Several +samplers split the prompts with ``data_slot`` ``"I/N"`` (``JsonlPrompts``). + +The configuration is read on the launching machine (:meth:`SamplerConfig.from_env` +takes the same variable names the K3.1 ops launcher used: RUN, SHARD, RELOAD, +DROP_ZV, N, MAX_TOKENS, THINKING_EFFORT, ...; flags override) and travels to the +container as an argument: Modal re-imports this module in the container without +the launcher's environment. + +Store credentials come from a Modal Secret (``secret``): either yeto's own names +(``YETO_RL_STORE_S3_ACCESS_KEY_ID`` / ``..._SECRET_ACCESS_KEY`` / ``..._ENDPOINT_URL``) +or the AWS names (``AWS_ACCESS_KEY_ID`` / ``AWS_SECRET_ACCESS_KEY`` / ``AWS_ENDPOINT_URL``), +read only because the island env says ``YETO_RL_STORE_S3_CREDENTIALS=aws-env``. + +Stop a sampler: ``modal.FunctionCall.from_id(id).cancel(terminate_containers=True)`` +(the id is printed at spawn), then ``modal app stop --yes ``. + +Not the Kimi-K3.1 production path: K3.1 samplers run on A-card nodes with the +vLLM ROCm image (``yeto.rl.sampler_pins.VLLM_SAMPLER_IMAGE_ROCM``). This module has +no default image (its former default was a stale trainer image) and refuses to +start until ``SAMPLER_IMAGE`` / ``--image`` names one; the bucket +(``STORE_BUCKET``) and the Modal Secret (``SAMPLER_SECRET``) are required too. + +Imports ``modal`` lazily; :func:`build_spec` and :func:`store_env` are CPU-testable. +""" + +from __future__ import annotations + +import dataclasses +import json +import os +import time +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from yeto.rl.engine.data_split import DataSlot + +REPO_ROOT = Path(__file__).resolve().parents[2] +PROMPTS_MOUNT = "/root/run/prompts.jsonl" +CODE_ROOT = "/root" # add_local_python_source("yeto") puts the package at /root/yeto +LOG_MOUNT = "/vol" +STAGE = "/tmp/stage" +DEFAULT_IMAGE = None # no default: see the module doc (not the K3.1 production path) +NO_IMAGE = ("the Modal engine sampler has no default image: Kimi-K3.1 samplers run on A-card " + "nodes with the vLLM ROCm image (yeto.rl.sampler_pins.VLLM_SAMPLER_IMAGE_ROCM), not " + "on Modal. To run this path anyway, name an SGLang sampler image with " + "SAMPLER_IMAGE=/@sha256: (or --image)") +CODE_IGNORE = ["**/__pycache__", "**/*.pyc"] +RELOADS = ("hot", "restart") + + +@dataclass(frozen=True) +class SamplerConfig: + run: str + shard: str + bucket: str + selection_sha256: str + prompts: str + data_slot: str | None = None + image: str | None = DEFAULT_IMAGE + engine_version: str = "sglang-miles-kimi-k31@35708f76" + model_volume: str = "k31-weights" + model_mount: str = "/mnt/yeto-models" + base_dir: str = "/mnt/yeto-models/hf/Kimi-K3.1-Pro/d9e99b07" + log_volume: str = "yeto-k31-sampler" + secret: str = "" + app_name: str = "yeto-k31-sampler-b300" + python: str = "/opt/sglang/bin/python3" + gpu: str = "B300" + tp: int = 8 + ep: int = 8 + mem_fraction: str = "0.85" + max_model_len: int = 4096 + max_tokens: int = 1024 + n: int = 4 + groups: int = 8 + backlog: int | None = None # None: 4 x groups + rounds: int = 1_000_000 + reload: str = "hot" + drop_zero_variance: bool = False + thinking_effort: str = "low" + reward: str = "yeto.rl.gsm8k_reward:score" + timeout_h: float = 6.0 + memory_gib: int = 2048 + cpu: int = 64 + prewarm: bool = True + extra_shards: tuple[str, ...] = field(default_factory=tuple) # informational: the trainer's list + + def __post_init__(self) -> None: + if self.reload not in RELOADS: + raise ValueError(f"reload {self.reload!r} is not one of {RELOADS}") + if self.data_slot: + DataSlot.parse(self.data_slot) # "I/N", refused early + if self.max_tokens + 1024 > self.max_model_len: + raise ValueError(f"max_tokens {self.max_tokens} + 1024 prompt tokens exceed max_model_len " + f"{self.max_model_len}") + if not self.shard or not self.run or not self.bucket: + raise ValueError("run, shard and bucket (STORE_BUCKET / --bucket) are required") + if not self.image: + raise ValueError(NO_IMAGE) + if not self.secret: + raise ValueError("the Modal Secret with the store credentials is required " + "(SAMPLER_SECRET / --secret)") + + @property + def max_backlog_groups(self) -> int: + return 4 * self.groups if self.backlog is None else int(self.backlog) + + @property + def trainer_buffer_shards(self) -> str: + """The ``--rl-buffer-shards`` value the trainer needs (with ``extra_shards``).""" + return ",".join([*self.extra_shards, self.shard]) + + def record(self) -> dict[str, Any]: + out = dataclasses.asdict(self) + out["extra_shards"] = list(self.extra_shards) + return out + + @classmethod + def from_record(cls, record: Mapping[str, Any]) -> "SamplerConfig": + data = dict(record) + data["extra_shards"] = tuple(data.get("extra_shards") or ()) + return cls(**data) + + @classmethod + def from_env(cls, environ: Mapping[str, str] | None = None, **overrides: Any) -> "SamplerConfig": + """The K3.1 ops launcher's variables (RUN, SHARD, SELECTION, PROMPTS, RELOAD, ...); + ``overrides`` (not None) win.""" + env = os.environ if environ is None else environ + selection = overrides.pop("selection", None) or env.get("SELECTION") + sha = overrides.pop("selection_sha256", None) or env.get("SELECTION_SHA256") + if not sha: + if not selection: + raise ValueError("SELECTION (expert selection file) or SELECTION_SHA256 is required") + sha = json.loads(Path(selection).read_text(encoding="utf-8"))["sha256"] + got = { + "run": env.get("RUN"), "shard": env.get("SHARD"), "bucket": env.get("STORE_BUCKET") or env.get("R2_BUCKET"), + "selection_sha256": sha, + "prompts": env.get("PROMPTS") or str(REPO_ROOT / "data" / "gsm8k" / "gsm8k-tasks512.jsonl"), + "data_slot": env.get("DATA_SLOT") or None, + "image": env.get("SAMPLER_IMAGE"), "engine_version": env.get("ENGINE_VERSION"), + "ep": _int(env.get("EP")), "mem_fraction": env.get("MEM_FRACTION"), + "max_model_len": _int(env.get("MAX_MODEL_LEN")), "max_tokens": _int(env.get("MAX_TOKENS")), + "n": _int(env.get("N")), "groups": _int(env.get("GROUPS")), "backlog": _int(env.get("BACKLOG")), + "rounds": _int(env.get("ROUNDS")), "reload": env.get("RELOAD"), + "drop_zero_variance": None if env.get("DROP_ZV") is None else env.get("DROP_ZV") == "1", + "thinking_effort": env.get("THINKING_EFFORT"), "reward": env.get("REWARD"), + "timeout_h": None if env.get("FN_TIMEOUT_H") is None else float(env["FN_TIMEOUT_H"]), + "memory_gib": _int(env.get("MEMORY_GIB")), + "prewarm": None if env.get("PREWARM") is None else env.get("PREWARM") == "1", + "secret": env.get("SAMPLER_SECRET"), "app_name": env.get("SAMPLER_APP"), + "extra_shards": tuple(s for s in (env.get("EXTRA_SHARDS") or "").split(",") if s) or None, + } + got.update({k: v for k, v in overrides.items() if v is not None}) + got.setdefault("bucket", None) + got["bucket"] = got["bucket"] or "" # refused in __post_init__ with the variable name + return cls(**{k: v for k, v in got.items() if v is not None}) + + +def _int(value: str | None) -> int | None: + return None if value in (None, "") else int(value) + + +def build_spec(c: SamplerConfig, *, out: str, stage: str = STAGE) -> dict[str, Any]: + """The engine sampler spec (``yeto.rl.engine.engine_sampler_entry``) of one island.""" + extra = ["--mem-fraction-static", c.mem_fraction, "--disable-flashinfer-autotune"] + if c.ep > 1: + extra += ["--ep-size", str(c.ep)] + hot = c.reload == "hot" + prompts: dict[str, Any] = { + "path": PROMPTS_MOUNT, "input_key": "messages", "label_key": "label", + "groups_per_round": c.groups, "format": "chat", "tokenizer": c.base_dir, "trust_remote_code": True, + "chat_template_kwargs": {"thinking_effort": c.thinking_effort}} + if c.data_slot: + prompts["data_slot"] = c.data_slot + return { + "sampler_id": c.shard, + "platform": "cuda", + "rollout_fields": {"routed_experts": True, "top_logprobs_num": 128, "fp32_lm_head": True}, + "expert_selection_sha256": c.selection_sha256, + "engine": { + "factory": "yeto.rl.adapters.sampler_engines.from_spec:build", + "adapter": "sglang-server", + "server": {"base_dir": c.base_dir, "port": 30000, "python": c.python, "card_model": c.gpu, + "device_family": "nvidia", "engine_version": c.engine_version, "precision": "mxfp4", + "quant_scheme": "MXFP4_EXPERTS_RTN", "topology": f"tp{c.tp}", + "max_model_len": c.max_model_len, "in_place_reload": hot, "start_timeout_s": 5400, + "extra_args": extra}, + "options": {"tp": c.tp, "expert_file_update": "auto" if hot else "off", "expert_update_probe": hot, + "live_update": False, "dp_attention_flag": "attn-dp-size", + "load_report_dir": f"{out}/load-reports"}, + }, + "sampling": {"max_tokens": c.max_tokens, "temperature": 1.0, "top_p": 1.0, "top_k": -1, "n": c.n}, + "num_rollouts": c.rounds, "max_backlog_groups": c.max_backlog_groups, "max_prompt_tokens": 1024, + "drop_zero_variance_groups": bool(c.drop_zero_variance), + "base_version": 0, "poll_s": 5.0, + "store": f"{stage}/{c.bucket}/runs/{c.run}", + "release_dir": "policy-releases", + "buffer_shards": [c.shard], + "prompts": prompts, + "reward": c.reward, + "reward_kind": "miles", + "checks": {"expected_expert_matrices": None, "allow_unverified_count": True, "require_fixed_answers": False}, + "incarnation_lease": False, + "max_consecutive_load_failures": 3, + "event_tape": f"{out}/events.jsonl", + } + + +def store_env(c: SamplerConfig, environ: Mapping[str, str], *, stage: str = STAGE) -> dict[str, str]: + """The store variables of the sampler process, from the container environment + (the Secret): yeto-named store credentials as they are, else the AWS names.""" + from yeto.cloud.s3_store import (CREDENTIALS_AWS_ENV, CREDENTIALS_ENV, ENDPOINT_ENV, KEY_ID_ENV, + REGION_ENV, STAGE_ENV, URI_ENV) + + env = {"YETO_RL_STORE_IMPL": "yeto.cloud.s3_store:store_from_env", + URI_ENV: f"s3://{c.bucket}/runs/{c.run}", STAGE_ENV: stage, "PYTHONPATH": CODE_ROOT} + endpoint = environ.get(ENDPOINT_ENV) or environ.get("AWS_ENDPOINT_URL") + if not endpoint: + raise ValueError(f"the store Secret has no endpoint ({ENDPOINT_ENV} or AWS_ENDPOINT_URL)") + env[ENDPOINT_ENV] = endpoint + env[REGION_ENV] = environ.get(REGION_ENV) or "auto" + if not environ.get(KEY_ID_ENV): + env[CREDENTIALS_ENV] = CREDENTIALS_AWS_ENV + return env + + +def dry_run(c: SamplerConfig) -> dict[str, Any]: + """What a launch would run (CPU only, no Modal call): the spec and the trainer flag.""" + rows = [line for line in Path(c.prompts).read_text(encoding="utf-8").splitlines() if line.strip()] + slot = DataSlot.parse(c.data_slot) + owned = sum(1 for i in range(len(rows)) if slot.owns(i)) + problems = [] + if c.groups > owned: + problems.append(f"groups_per_round {c.groups} exceeds the {owned} prompt rows of slot {slot.env()}") + return {"app": c.app_name, "store": f"s3://{c.bucket}/runs/{c.run}", "shard": c.shard, + "data_slot": slot.env(), "prompt_rows": owned, + "trainer_buffer_shards": c.trainer_buffer_shards, "image": c.image, "secret": c.secret, + "spec": build_spec(c, out=f"{LOG_MOUNT}/sampler/{c.run}/{c.shard}"), "problems": problems} + + +# ---------------------------------------------------------------------------- container +def sampler_body(record: Mapping[str, Any]) -> dict[str, Any]: + """Runs INSIDE the Modal container: write the spec, run the engine sampler.""" + import subprocess + import threading + + import modal + + c = SamplerConfig.from_record(record) + t0 = time.time() + logs = modal.Volume.from_name(c.log_volume) + out = f"{LOG_MOUNT}/sampler/{c.run}/{c.shard}-{time.strftime('%Y%m%d-%H%M%S')}" + Path(out).mkdir(parents=True, exist_ok=True) + spec = build_spec(c, out=out) + Path(out, "spec.json").write_text(json.dumps(spec, indent=2)) + env = dict(os.environ, **store_env(c, os.environ)) + pre = subprocess.run(f"nvidia-smi --query-gpu=name,compute_cap,memory.total --format=csv,noheader; " + f"{c.python} -c 'import sglang;print(sglang.__version__)'; " + f"head -c 300 {c.base_dir}/MANIFEST.json", shell=True, capture_output=True, text=True) + Path(out, "preflight.log").write_text(pre.stdout + pre.stderr) + print(pre.stdout, pre.stderr, flush=True) + stop = threading.Event() + + def smi() -> None: + with open(f"{out}/nvidia-smi.csv", "a", buffering=1) as fh: + while not stop.is_set(): + fh.write(subprocess.run(["nvidia-smi", "--query-gpu=timestamp,index,memory.used,utilization.gpu", + "--format=csv,noheader"], capture_output=True, text=True).stdout) + stop.wait(15) + + def commit_loop() -> None: + while not stop.wait(300): + logs.commit() + + def prewarm() -> None: + # read every weight file once, 32 at a time, so SGLang's own load hits page cache + t = time.time() + r = subprocess.run(f"ls {c.base_dir}/*.safetensors | xargs -P 32 -I{{}} cat {{}} > /dev/null", + shell=True, capture_output=True, text=True) + Path(out, "prewarm.log").write_text(f"rc={r.returncode} secs={time.time() - t:.1f}\n{r.stderr[-2000:]}") + + if c.prewarm: + threading.Thread(target=prewarm, daemon=True).start() + threading.Thread(target=smi, daemon=True).start() + threading.Thread(target=commit_loop, daemon=True).start() + with open(f"{out}/engine_sampler.log", "w") as log: + rc = subprocess.run([c.python, "-m", "yeto.rl.engine.engine_sampler", "--spec", f"{out}/spec.json"], + env=env, stdout=log, stderr=subprocess.STDOUT).returncode + stop.set() + logs.commit() + return {"rc": rc, "out": out, "secs": round(time.time() - t0, 1)} + + +def build_app(c: SamplerConfig) -> tuple[Any, Any]: + """(app, sampler GPU function running :func:`sampler_body`).""" + import modal + + image = (modal.Image.from_registry(c.image).entrypoint([]) # the image's version-assert entrypoint off + .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"}) # no PYTHONPATH: it shadows Modal's runtime + .add_local_file(c.prompts, PROMPTS_MOUNT, copy=False) + .add_local_python_source("yeto", copy=False, ignore=CODE_IGNORE)) + app = modal.App(c.app_name) + models = modal.Volume.from_name(c.model_volume) + logs = modal.Volume.from_name(c.log_volume, create_if_missing=True) + fn = app.function(image=image, gpu=f"{c.gpu}:{c.tp}", cpu=c.cpu, memory=c.memory_gib * 1024, + timeout=int(c.timeout_h * 3600), + volumes={c.model_mount: models.read_only(), LOG_MOUNT: logs}, + secrets=[modal.Secret.from_name(c.secret)], name="sampler")(sampler_body) + return app, fn + + +def launch(c: SamplerConfig, *, wait: bool = True, log=print) -> dict[str, Any] | str: + """Spawn the sampler; returns its result (``wait``) or the FunctionCall id.""" + from ..cost.ledger import record_safely + + app, fn = build_app(c) + record_safely("intent", "modal_app", c.app_name, gpu=c.gpu, gpus=c.tp) + with app.run(detach=True): + # a detached ephemeral app per launch: its own app id ($YETO_RESOURCE_LEDGER records it, + # so `yeto cleanup --run` stops exactly this app; the sampler outlives this process) + if getattr(app, "app_id", None): + record_safely("created", "modal_app", app.app_id, c.app_name, gpu=c.gpu, gpus=c.tp, + role="sampler", shard=c.shard) + call = fn.spawn(c.record()) + record_safely("created", "modal_call", str(call.object_id), "sampler", + app_id=getattr(app, "app_id", None), app_name=c.app_name) + log(f"[sampler] app {c.app_name} FunctionCall {call.object_id}; trainer --rl-buffer-shards " + f"{c.trainer_buffer_shards}; stop: modal.FunctionCall.from_id(id).cancel(terminate_containers=True) " + f"then modal app stop --yes {c.app_name}") + return call.get() if wait else call.object_id diff --git a/yeto/cloud/oss_store.py b/yeto/cloud/oss_store.py index b32f9fb3e..69f1e7c9c 100644 --- a/yeto/cloud/oss_store.py +++ b/yeto/cloud/oss_store.py @@ -44,8 +44,9 @@ from pathlib import Path from typing import Any -from yeto.cloud.s3_store import (PART_CONCURRENCY_ENV, PART_MB_ENV, POINTER_FILES, RECORD_FILES, - STAGED_SUBTREES, S3ObjectStore, TransferSettings, parse_s3_uri) +from yeto.cloud.s3_store import (FILE_CONCURRENCY_ENV, INDEPENDENT_DIRS, PART_CONCURRENCY_ENV, PART_MB_ENV, + POINTER_FILES, RECORD_FILES, STAGED_SUBTREES, S3ObjectStore, + TransferSettings, parse_s3_uri) from yeto.rl.engine.object_store import ObjectStoreError, TreeMirror from yeto.rl.engine.resume import STORE_IMPL_ENV, CheckpointStore @@ -62,7 +63,8 @@ DEFAULT_STAGE = "~/.cache/yeto/oss-store-stage" CREDENTIAL_ENVS = (KEY_ID_ENV, SECRET_ENV, TOKEN_ENV) SECRET_ENVS = frozenset({SECRET_ENV, TOKEN_ENV}) # shipped as task secrets, not plain envs -SETTING_ENVS = (ENDPOINT_ENV, REGION_ENV, CREDENTIALS_FILE_ENV, PART_MB_ENV, PART_CONCURRENCY_ENV) +SETTING_ENVS = (ENDPOINT_ENV, REGION_ENV, CREDENTIALS_FILE_ENV, PART_MB_ENV, PART_CONCURRENCY_ENV, + FILE_CONCURRENCY_ENV) FORBID_OVERWRITE_HEADER = "x-oss-forbid-overwrite" @@ -83,6 +85,7 @@ class OssObjectStore(S3ObjectStore): scheme = SCHEME buffer_shards = True + part_checksums = False # OSS's S3 API takes no x-amz-checksum-sha256 (its CRC64 is not composite) def create(self, key: str, data: bytes) -> bool: if not self.conditional_writes: @@ -196,11 +199,32 @@ def objects(self) -> OssObjectStore: # type: ignore[override] def _tree(self) -> TreeMirror: if self._mirror is None: self._mirror = TreeMirror(self.root, self.objects, subtrees=STAGED_SUBTREES, - pointers=POINTER_FILES, records=RECORD_FILES) + pointers=POINTER_FILES, records=RECORD_FILES, + independent=INDEPENDENT_DIRS, + file_concurrency=self.objects.transfer.file_concurrency) return self._mirror - def commit(self) -> float: - return self._tree().push() + def commit(self, paths: Any = None) -> float: + return self._tree().push(paths) + + @property + def cut_commits_beside_buffer_reads(self) -> bool: # type: ignore[override] + return True # always staged: the buffer goes by the object API + + def push_local_file(self, relative: str) -> int | None: + """Upload one file of the stage now (a trainer rank's cut shard, on whichever + node wrote it: a multi-node island has one stage per node); returns the size + the object store reports, None when this store has no stage.""" + return self._tree().push_file(relative) + + def local_file_checksum(self, relative: str) -> str | None: + """The checksum the object store keeps for a file :meth:`push_local_file` + uploaded (``S3ObjectStore.checksum``), None without one.""" + return self._tree().object_checksum(relative) + + def delete_remote_tree(self, relative: str) -> list[str]: + """Delete every object under ``relative``, including what other nodes uploaded.""" + return self._tree().delete_tree(relative) def reload(self) -> None: self._tree().pull() diff --git a/yeto/cloud/s3_store.py b/yeto/cloud/s3_store.py index 93b49f16c..5b9c5acef 100644 --- a/yeto/cloud/s3_store.py +++ b/yeto/cloud/s3_store.py @@ -32,16 +32,38 @@ from __future__ import annotations +import functools import os -from collections.abc import Mapping +import pickle +import subprocess +import sys +import time +from collections.abc import Callable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass +from dataclasses import asdict, dataclass from pathlib import Path from typing import Any from yeto.rl.engine.object_store import ObjectStoreError, TreeMirror, check_key from yeto.rl.engine.resume import STORE_IMPL_ENV, BucketMountStore +# The retryable failures of one streamed range read: the whole class of transfer +# errors, not one exception at a time (an R2 body cut short has surfaced as botocore +# ReadTimeoutError, ResponseStreamingError and urllib3 IncompleteRead so far). +import http.client as _http_client + +RANGE_READ_RETRYABLE: tuple[type[BaseException], ...] = ( + ConnectionError, TimeoutError, _http_client.HTTPException) # incl. ConnectionResetError, IncompleteRead +try: # botocore ships with boto3 + from botocore.exceptions import (ConnectionError as _BotoConnectionError, HTTPClientError, + IncompleteReadError, ResponseStreamingError) + from urllib3.exceptions import HTTPError as _Urllib3Error # ProtocolError, IncompleteRead, timeouts + RANGE_READ_RETRYABLE += (_BotoConnectionError, HTTPClientError, IncompleteReadError, + ResponseStreamingError, _Urllib3Error) +except ImportError: # pragma: no cover - a store without boto3 never reads + pass +RANGE_READ_ATTEMPTS = 5 + STORE_IMPL = "yeto.cloud.s3_store:store_from_env" # value of resume.STORE_IMPL_ENV URI_ENV = "YETO_RL_STORE_S3_URI" ENDPOINT_ENV = "YETO_RL_STORE_S3_ENDPOINT_URL" # S3-compatible endpoint (MinIO); unset: AWS @@ -49,12 +71,21 @@ KEY_ID_ENV = "YETO_RL_STORE_S3_ACCESS_KEY_ID" SECRET_ENV = "YETO_RL_STORE_S3_SECRET_ACCESS_KEY" TOKEN_ENV = "YETO_RL_STORE_S3_SESSION_TOKEN" +# a named AWS profile (~/.aws/config, ~/.aws/credentials) of a launcher-side client: +# its keys, endpoint_url and region; AWS_PROFILE is read too +PROFILE_ENV = "YETO_RL_STORE_S3_PROFILE" STAGE_ENV = "YETO_RL_STORE_S3_STAGE" # a local directory: cuts and releases by object API PART_MB_ENV = "YETO_RL_STORE_PART_MB" # multipart part size (MiB); shared by the S3 and OSS stores -PART_CONCURRENCY_ENV = "YETO_RL_STORE_PART_CONCURRENCY" # parts moved in parallel +PART_CONCURRENCY_ENV = "YETO_RL_STORE_PART_CONCURRENCY" # parts of one file moved in parallel +FILE_CONCURRENCY_ENV = "YETO_RL_STORE_FILE_CONCURRENCY" # files of one push tier moved in parallel +MAX_CONNECTIONS_ENV = "YETO_RL_STORE_MAX_CONNECTIONS" # HTTP connection pool of the client (default: the threads) +DOWNLOAD_PROCESSES_ENV = "YETO_RL_STORE_DOWNLOAD_PROCESSES" # processes of a many-file download (read_files) +PART_CHECKSUM_ENV = "YETO_RL_STORE_PART_CHECKSUM" # SHA-256 per uploaded part: 1 / 0 / auto (default) +UPLOAD_ATTEMPTS_ENV = "YETO_RL_STORE_UPLOAD_ATTEMPTS" # tries of one multipart upload, resumed between them CREDENTIAL_ENVS = (KEY_ID_ENV, SECRET_ENV, TOKEN_ENV) SECRET_ENVS = frozenset({SECRET_ENV, TOKEN_ENV}) # shipped as task secrets, not plain envs -SETTING_ENVS = (ENDPOINT_ENV, REGION_ENV, PART_MB_ENV, PART_CONCURRENCY_ENV) +SETTING_ENVS = (ENDPOINT_ENV, REGION_ENV, PART_MB_ENV, PART_CONCURRENCY_ENV, FILE_CONCURRENCY_ENV, + MAX_CONNECTIONS_ENV, DOWNLOAD_PROCESSES_ENV, PART_CHECKSUM_ENV, UPLOAD_ATTEMPTS_ENV) DEFAULT_STAGE = "~/.cache/yeto/store-stage" # what a staged store pulls at start (the newest committed cut: LATEST, its snapshot, the # cuts) and the pointer files it uploads last (resume.LATEST, policy_release.LATEST) @@ -63,6 +94,9 @@ # record files: uploaded after the data files next to them, before the pointers # (policy_release.RECORD lists the release's weight files) RECORD_FILES = frozenset({"RELEASE.json"}) +# top-level directories their writer commits by path (policy releases of every scheme, +# policy_release.PolicyReleaseStore): a cut's commit neither waits for nor uploads them +INDEPENDENT_DIRS = ("policy-releases",) # error codes of a conditional write that lost to an existing object (S3; OSS # answers FileAlreadyExists to x-oss-forbid-overwrite) @@ -71,23 +105,156 @@ _MIB = 1 << 20 MIN_PART_BYTES = 5 * _MIB # S3 and OSS: every part but the last is at least 5 MiB MAX_PARTS = 10_000 # S3 and OSS: parts of one multipart upload +# the resume journal of a multipart upload, next to its source file (a stage push +# skips it: object_store.TMP_SUFFIX); removed when the upload completes +UPLOAD_JOURNAL_SUFFIX = ".upload.tmp" +_FALSE = frozenset({"0", "false", "no", "off"}) +_TRUE = frozenset({"1", "true", "yes", "on"}) +# endpoints known to refuse x-amz-checksum-sha256 on UploadPart (NotImplemented): +# PART_CHECKSUM_ENV=auto sends no part checksums to them +NO_PART_CHECKSUM_HOSTS = (".r2.cloudflarestorage.com",) +_PART_CHECKSUM_WARNED = False + + +def part_checksum_default(environ: Mapping[str, str]) -> bool: + """:data:`PART_CHECKSUM_ENV`: ``1`` sends a SHA-256 with every part, ``0`` none; + unset or ``auto`` sends them except to an endpoint known not to take them + (:data:`NO_PART_CHECKSUM_HOSTS`, Cloudflare R2). An endpoint that still answers + ``NotImplemented`` turns them off for that store (:meth:`S3ObjectStore.write_file`).""" + raw = str(environ.get(PART_CHECKSUM_ENV) or "auto").strip().lower() + if raw in _FALSE: + return False + if raw in _TRUE: + return True + if raw != "auto": + raise ValueError(f"{PART_CHECKSUM_ENV}={raw!r}: expected 1, 0 or auto") + return not _is_r2(store_endpoint(environ)) + + +def _warn_part_checksums_off(uri: str, error: Exception) -> None: + global _PART_CHECKSUM_WARNED + if _PART_CHECKSUM_WARNED: + return + _PART_CHECKSUM_WARNED = True + print(f"yeto: warning: {uri} does not take SHA-256 part checksums ({error.__class__.__name__}: " + f"NotImplemented); uploading without them (set {PART_CHECKSUM_ENV}=0 to skip this probe)", + file=sys.stderr, flush=True) + + +def _b64(digest: bytes) -> str: + import base64 + + return base64.b64encode(digest).decode("ascii") + + +def composite_sha256(part_digests: Sequence[str]) -> str: + """The checksum S3 reports for a multipart object uploaded with SHA-256 part + checksums (a COMPOSITE checksum): base64 SHA-256 of the concatenated raw part + digests, then ``-``. ``part_digests``: base64, in part order.""" + import base64 + import hashlib + + raw = b"".join(base64.b64decode(d) for d in part_digests) + return f"{_b64(hashlib.sha256(raw).digest())}-{len(part_digests)}" + + +class _UploadJournal: + """What a multipart upload of one file has sent: its UploadId and, per part + number, the ETag and SHA-256 (base64) the store acknowledged. Saved atomically + after each part, next to the source (:data:`UPLOAD_JOURNAL_SUFFIX`), so an upload + cut short (a failed part, a killed process) continues where it stopped. It + belongs to one version of the source: a different key, size, mtime, part size or + checksum mode makes it stale.""" + + def __init__(self, source: str | os.PathLike[str], identity: dict[str, Any]) -> None: + import threading + + self.path = Path(str(source) + UPLOAD_JOURNAL_SUFFIX) + self.identity = identity + self.upload_id: str | None = None + self.parts: dict[int, dict[str, Any]] = {} + self._lock = threading.Lock() + self.writable = True + + def load(self) -> dict[str, Any] | None: + """The journal on disk: its contents when it belongs to this upload, a dict + with only ``stale_upload_id`` when it belongs to another, None when absent.""" + import json + + try: + data = json.loads(self.path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + if not isinstance(data, dict): + return None + if data.get("identity") != self.identity: + return {"stale_upload_id": data.get("upload_id")} + return data + + def save(self) -> None: + import json + + if not self.writable: + return + data = {"identity": self.identity, "upload_id": self.upload_id, + "parts": {str(n): v for n, v in sorted(self.parts.items())}} + tmp = self.path.with_name(self.path.name.removesuffix(".tmp") + ".w.tmp") # skipped by a push too + try: + tmp.write_text(json.dumps(data, sort_keys=True), encoding="utf-8") + os.replace(tmp, self.path) + except OSError: # a read-only source directory: the upload goes on without resume + self.writable = False + + def record(self, number: int, part: dict[str, Any]) -> None: + with self._lock: + self.parts[int(number)] = part + self.save() + + def remove(self) -> None: + self.path.unlink(missing_ok=True) @dataclass(frozen=True) class TransferSettings: """How large files move: a file larger than one part is uploaded in parts of ``part_bytes`` (raised so a file needs at most :data:`MAX_PARTS`) and - downloaded by ranges of the same size; ``concurrency`` parts move at once - (memory: up to ``concurrency`` parts).""" + downloaded by ranges of the same size; ``concurrency`` parts of a file move at + once, and a staged store's push moves ``file_concurrency`` files at once + (memory: up to ``file_concurrency * concurrency`` parts). + + Defaults from an R2 upload measurement (GPU-A, 2 GiB files, 64 MiB parts, one + 16-core container): one file tops out near 100 MB/s whatever its part + concurrency (76 / 109 / 102 MB/s at 8 / 32 / 64 parts); 4 and 8 files of 16 + parts each reach 255 and 273 MB/s. + + ``download_processes``: a many-file download (:meth:`S3ObjectStore.read_files`, + a policy release) splits its files over that many processes. One Python process + tops out near 700 MB/s of R2 range reads (MI355X node, 64 threads, nothing + written); 4 / 8 processes of 16 threads read 1.7 / 3.0 GB/s. 4 stays well under + the ~7 GB/s that dropped an MI355X node off the network (GPU-B section 4.1).""" part_bytes: int = 64 * _MIB - concurrency: int = 8 + concurrency: int = 16 + file_concurrency: int = 8 + download_processes: int = 4 + # SHA-256 per uploaded part: the store refuses a part whose body does not match + # it, and the object's composite checksum is checked after the upload + part_checksum: bool = True + # tries of one multipart upload in one call; each continues the previous one + # (list_parts), and after the last the upload is kept for the next call + upload_attempts: int = 3 def __post_init__(self) -> None: if self.part_bytes < MIN_PART_BYTES: raise ValueError(f"part size {self.part_bytes} B is below the {MIN_PART_BYTES} B minimum") if self.concurrency < 1: raise ValueError("part concurrency must be >= 1") + if self.file_concurrency < 1: + raise ValueError("file concurrency must be >= 1") + if self.download_processes < 1: + raise ValueError("download processes must be >= 1") + if self.upload_attempts < 1: + raise ValueError("upload attempts must be >= 1") def part_for(self, size: int) -> int: return max(self.part_bytes, -(-int(size) // MAX_PARTS)) @@ -96,7 +263,11 @@ def part_for(self, size: int) -> int: def from_env(cls, environ: Mapping[str, str]) -> "TransferSettings": default = cls() return cls(part_bytes=int(float(environ.get(PART_MB_ENV) or default.part_bytes / _MIB) * _MIB), - concurrency=int(environ.get(PART_CONCURRENCY_ENV) or default.concurrency)) + concurrency=int(environ.get(PART_CONCURRENCY_ENV) or default.concurrency), + file_concurrency=int(environ.get(FILE_CONCURRENCY_ENV) or default.file_concurrency), + download_processes=int(environ.get(DOWNLOAD_PROCESSES_ENV) or default.download_processes), + part_checksum=part_checksum_default(environ), + upload_attempts=int(environ.get(UPLOAD_ATTEMPTS_ENV) or default.upload_attempts)) def parse_s3_uri(uri: str, scheme: str = "s3") -> tuple[str, str]: @@ -108,24 +279,170 @@ def parse_s3_uri(uri: str, scheme: str = "s3") -> tuple[str, str]: return bucket, prefix.strip("/") +# The store credential as a Modal Secret carries it (the AWS names; +# yeto.launch.storage.modal_store_secret_name). An island reads that pair only when +# the launch says so (:data:`CREDENTIALS_ENV` = :data:`CREDENTIALS_AWS_ENV`): the +# launcher's own AWS_* variables are never a store credential (secret-handling +# hardening), and the default credential chain stays unused either way. +CREDENTIALS_ENV = "YETO_RL_STORE_S3_CREDENTIALS" +CREDENTIALS_AWS_ENV = "aws-env" # read AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY (a Modal Secret) +AWS_KEY_ID_ENV = "AWS_ACCESS_KEY_ID" +AWS_SECRET_ENV = "AWS_SECRET_ACCESS_KEY" +AWS_TOKEN_ENV = "AWS_SESSION_TOKEN" +AWS_ENDPOINT_ENV = "AWS_ENDPOINT_URL" +AWS_REGION_ENVS = ("AWS_REGION", "AWS_DEFAULT_REGION") + + +def store_credentials(environ: Mapping[str, str]) -> tuple[str, str, str | None] | None: + """``(key id, secret, session token or None)``: the store credential + (:data:`KEY_ID_ENV` / :data:`SECRET_ENV`); with :data:`CREDENTIALS_ENV` set to + :data:`CREDENTIALS_AWS_ENV` the AWS-named pair a Secret put in the environment. + None when no complete pair applies.""" + if environ.get(KEY_ID_ENV) and environ.get(SECRET_ENV): + return environ[KEY_ID_ENV], environ[SECRET_ENV], environ.get(TOKEN_ENV) or None + if (environ.get(CREDENTIALS_ENV) == CREDENTIALS_AWS_ENV and environ.get(AWS_KEY_ID_ENV) + and environ.get(AWS_SECRET_ENV)): + return environ[AWS_KEY_ID_ENV], environ[AWS_SECRET_ENV], environ.get(AWS_TOKEN_ENV) or None + return None + + +def missing_credentials_message() -> str: + return (f"s3 store: no store credentials. Set {KEY_ID_ENV} and {SECRET_ENV} (a key limited to the " + f"store's bucket), or {PROFILE_ENV} (or --profile) to an AWS profile that holds them; on " + f"an island with a Modal Secret, {CREDENTIALS_ENV}={CREDENTIALS_AWS_ENV} reads " + f"{AWS_KEY_ID_ENV} and {AWS_SECRET_ENV}; an island uses the store credentials only, never the default AWS chain") + + def has_store_credentials(environ: Mapping[str, str]) -> bool: - return bool(environ.get(KEY_ID_ENV) and environ.get(SECRET_ENV)) + return store_credentials(environ) is not None + + +AWS_PROFILE_ENV = "AWS_PROFILE" +AWS_ENDPOINT_S3_ENV = "AWS_ENDPOINT_URL_S3" + + +def store_profile(environ: Mapping[str, str]) -> str | None: + """The AWS profile a store client reads: :data:`PROFILE_ENV`, else ``AWS_PROFILE``.""" + return environ.get(PROFILE_ENV) or environ.get(AWS_PROFILE_ENV) or None + + +@functools.lru_cache(maxsize=16) +def _profile_config_cached(profile: str, config_file: str | None, credentials_file: str | None) -> dict[str, Any]: + import botocore.session + + try: + config = dict(botocore.session.Session(profile=profile).get_scoped_config()) + except Exception as error: # noqa: BLE001 - botocore ProfileNotFound + if type(error).__name__ == "ProfileNotFound": + return {"_missing": True} + raise + s3 = config.get("s3") if isinstance(config.get("s3"), dict) else {} + return {"endpoint_url": s3.get("endpoint_url") or config.get("endpoint_url"), "region": config.get("region")} + + +def profile_config(environ: Mapping[str, str]) -> dict[str, Any]: + """``endpoint_url`` and ``region`` of the store's AWS profile (:func:`store_profile`; + the ``s3`` section's endpoint first), ``{}`` without a profile, ``{"_missing": True}`` + when the profile is not configured on this host (an island).""" + profile = store_profile(environ) + if not profile: + return {} + try: + return _profile_config_cached(profile, environ.get("AWS_CONFIG_FILE"), + environ.get("AWS_SHARED_CREDENTIALS_FILE")) + except ImportError: # no botocore: nothing to read + return {} + + +def store_endpoint(environ: Mapping[str, str]) -> str | None: + """The S3-compatible endpoint: :data:`ENDPOINT_ENV`, else ``AWS_ENDPOINT_URL_S3`` / + ``AWS_ENDPOINT_URL``, else the ``endpoint_url`` of the store's AWS profile; None + for AWS S3.""" + return (environ.get(ENDPOINT_ENV) or environ.get(AWS_ENDPOINT_S3_ENV) or environ.get(AWS_ENDPOINT_ENV) + or profile_config(environ).get("endpoint_url") or None) + + +def _is_r2(endpoint: str | None) -> bool: + from urllib.parse import urlparse + + host = (urlparse(endpoint).hostname or "") if endpoint else "" + return any(host.endswith(suffix) for suffix in NO_PART_CHECKSUM_HOSTS) + + +def store_region(environ: Mapping[str, str]) -> str | None: + """:data:`REGION_ENV`, else the AWS region variables, else the store profile's + region, else ``auto`` for a Cloudflare R2 endpoint; None otherwise (boto3's + default). ``yeto.launch.storage.store_client_environ`` sets ``auto`` for an + endpoint given by ``--rl-object-store-endpoint`` (what R2 expects).""" + for name in (REGION_ENV, *AWS_REGION_ENVS): + if environ.get(name): + return environ[name] + region = profile_config(environ).get("region") + if region: + return str(region) + return "auto" if _is_r2(store_endpoint(environ)) else None + + +def _read_files_here(objects: "S3ObjectStore", items: Sequence[tuple[str, str]], threads: int) -> list[str]: + def one(item: tuple[str, str]) -> str | None: + return None if objects.read_file(*item) else item[0] + + with ThreadPoolExecutor(max(1, min(int(threads), len(items) or 1))) as pool: + return [key for key in pool.map(one, items) if key is not None] + + +def _read_files_child() -> None: + """One child process of :meth:`S3ObjectStore.read_files`: its own client; the + arguments come pickled on stdin (from the parent only), the missing keys go back + pickled on stdout.""" + uri, client_factory, transfer, threads, items = pickle.loads(sys.stdin.buffer.read()) + objects = S3ObjectStore(uri, client=client_factory(), transfer=TransferSettings(**transfer)) + missing = _read_files_here(objects, items, threads) + sys.stdout.buffer.write(pickle.dumps(missing)) + sys.stdout.buffer.flush() + + +READ_FILES_CHILD = "from yeto.cloud.s3_store import _read_files_child; _read_files_child()" + + +BOTOCORE_DEFAULT_POOL = 10 # botocore's max_pool_connections + + +def client_pool_size(environ: Mapping[str, str]) -> int: + """HTTP connections of the store's boto3 client: :data:`MAX_CONNECTIONS_ENV`, else + one per transfer thread (``file_concurrency * concurrency``). botocore's default + pool of 10 serialized the 8 x 16 range reads of a release download (MI355X: a + 207 GB Kimi-K3.1 release far below the node's R2 bandwidth).""" + raw = environ.get(MAX_CONNECTIONS_ENV) + if raw: + return max(1, int(raw)) + t = TransferSettings.from_env(environ) + return max(BOTOCORE_DEFAULT_POOL, t.file_concurrency * t.concurrency) def _client(environ: Mapping[str, str], *, require_credentials: bool) -> Any: import boto3 + from botocore.config import Config - if has_store_credentials(environ): - session = boto3.session.Session(aws_access_key_id=environ[KEY_ID_ENV], - aws_secret_access_key=environ[SECRET_ENV], - aws_session_token=environ.get(TOKEN_ENV) or None) + creds = store_credentials(environ) + profile = store_profile(environ) + if profile and profile_config(environ).get("_missing"): + if environ.get(PROFILE_ENV) or (creds is None and require_credentials): + raise ObjectStoreError(f"s3 store: AWS profile {profile!r} ({PROFILE_ENV if environ.get(PROFILE_ENV) else AWS_PROFILE_ENV}) " + "is not configured on this host (~/.aws/config, ~/.aws/credentials)") + profile = None # AWS_PROFILE of another host: not this store's + if creds is not None: + key, secret, token = creds + session = boto3.session.Session(aws_access_key_id=key, aws_secret_access_key=secret, + aws_session_token=token) + elif profile: + session = boto3.session.Session(profile_name=profile) elif require_credentials: - raise ObjectStoreError(f"s3 store: {KEY_ID_ENV} and {SECRET_ENV} are not set (an island " - "uses the store credentials only, never the default AWS chain)") + raise ObjectStoreError(missing_credentials_message()) else: session = boto3.session.Session() - return session.client("s3", endpoint_url=environ.get(ENDPOINT_ENV) or None, - region_name=environ.get(REGION_ENV) or None) + return session.client("s3", endpoint_url=store_endpoint(environ), region_name=store_region(environ), + config=Config(max_pool_connections=client_pool_size(environ))) class S3ObjectStore: @@ -143,13 +460,22 @@ class S3ObjectStore: scheme = "s3" buffer_shards = False + # the endpoint takes x-amz-checksum-sha256 on parts and reports a composite + # checksum (HEAD, ChecksumMode=ENABLED); OSS does not + part_checksums = True - def __init__(self, uri: str, *, client: Any, transfer: TransferSettings | None = None) -> None: + def __init__(self, uri: str, *, client: Any, transfer: TransferSettings | None = None, + client_factory: Callable[[], Any] | None = None) -> None: + """``client_factory``: a picklable callable that builds a client like ``client`` + in another process (:meth:`read_files`); None keeps every read in this process.""" self.bucket, self.prefix = parse_s3_uri(uri, self.scheme) self.uri = f"{self.scheme}://{self.bucket}" + (f"/{self.prefix}" if self.prefix else "") self.client = client self.transfer = transfer or TransferSettings() + self.client_factory = client_factory self.conditional_writes = True + # set when the endpoint answered a part checksum with NotImplemented + self.part_checksums_refused = False def _full(self, name: str) -> str: return f"{self.prefix}/{name}" if self.prefix else name @@ -207,34 +533,190 @@ def size(self, key: str) -> int | None: head = self._head(key) return None if head is None else int(head["ContentLength"]) - def write_file(self, key: str, source: str | os.PathLike[str]) -> None: + def _checksums(self) -> bool: + return bool(self.part_checksums and self.transfer.part_checksum and not self.part_checksums_refused) + + def checksum(self, key: str) -> str | None: + """The SHA-256 checksum the store keeps for object ``key`` (base64; a + multipart object's is composite, ``-``), None when the object is + missing or has none (uploaded without checksums, an endpoint without them).""" + if not self._checksums(): + return None + try: + head = self.client.head_object(Bucket=self.bucket, Key=self._key(key), ChecksumMode="ENABLED") + except Exception as error: # noqa: BLE001 - botocore ClientError, by code + if self._code(error) in _MISSING: + return None + raise + value = head.get("ChecksumSHA256") + return str(value) if value else None + + def write_file(self, key: str, source: str | os.PathLike[str]) -> str | None: """Upload a file; one larger than a part goes as a multipart upload whose parts are read from the file as they are sent (never the whole file in - memory). A failed upload is aborted, so no part is left billed.""" + memory). Returns the object's SHA-256 checksum (:meth:`checksum`), None + without checksums. + + Integrity: each part carries its SHA-256 (the store refuses a part whose body + does not match), and the composite checksum the store reports afterwards must + equal the one of the parts sent (:func:`composite_sha256`). + + Resume: the UploadId and every acknowledged part are journaled next to the + source (:class:`_UploadJournal`). A failed try continues from the parts + ``list_parts`` and the journal agree on (same ETag, same checksum), up to + ``upload_attempts`` tries; after the last the upload and its journal are kept, + so a later call (the next push, a restarted process) continues it. A journal + of another version of the file aborts its upload and starts again. + + An endpoint that answers a checksum with ``NotImplemented`` (Cloudflare R2 + on UploadPart) turns checksums off for this store, with one warning per + process, and the file is uploaded again without them.""" + try: + return self._write_file(key, source) + except Exception as error: # noqa: BLE001 - botocore ClientError, by code + if not self._checksums() or self._code(error) != "NotImplemented": + raise + self.part_checksums_refused = True + _warn_part_checksums_off(self.uri, error) + return self._write_file(key, source) + + def _write_file(self, key: str, source: str | os.PathLike[str]) -> str | None: + import hashlib + + checks = self._checksums() size = os.path.getsize(source) part = self.transfer.part_for(size) full = self._key(key) if size <= part: with open(source, "rb") as fh: - self.client.put_object(Bucket=self.bucket, Key=full, Body=fh.read()) + body = fh.read() + extra = {"ChecksumSHA256": _b64(hashlib.sha256(body).digest())} if checks else {} + self.client.put_object(Bucket=self.bucket, Key=full, Body=body, **extra) + return extra.get("ChecksumSHA256") + count = -(-size // part) + st = os.stat(source) + journal = _UploadJournal(source, {"bucket": self.bucket, "key": full, "size": int(size), + "mtime_ns": int(st.st_mtime_ns), "part_bytes": int(part), + "checksum": checks}) + self._open_upload(journal, full, part, size, checks) + for attempt in range(1, self.transfer.upload_attempts + 1): + try: + self._send_parts(journal, full, source, part, size, count, checks) + parts = [{"PartNumber": n, "ETag": journal.parts[n]["ETag"], + **({"ChecksumSHA256": journal.parts[n]["sha256"]} if checks else {})} + for n in range(1, count + 1)] + self.client.complete_multipart_upload(Bucket=self.bucket, Key=full, UploadId=journal.upload_id, + MultipartUpload={"Parts": parts}) + break + except Exception as error: # noqa: BLE001 - botocore ClientError, by code + if checks and self._code(error) == "NotImplemented": + try: # the endpoint takes no part checksums: write_file starts again without + self.client.abort_multipart_upload(Bucket=self.bucket, Key=full, + UploadId=journal.upload_id) + except Exception: # noqa: BLE001 - already gone + pass + journal.remove() + raise + if attempt == self.transfer.upload_attempts: + raise # the upload and its journal stay: the next call continues them + listed = self._listed_parts(full, journal.upload_id, part, size) + if listed is None: # the upload is gone (completed elsewhere, expired, aborted) + journal.upload_id, journal.parts = None, {} + self._open_upload(journal, full, part, size, checks, fresh=True) + else: + journal.parts = self._agreed(journal.parts, listed, checks) + journal.save() + journal.remove() + if not checks: + return None + expected = composite_sha256([journal.parts[n]["sha256"] for n in range(1, count + 1)]) + reported = self.checksum(key) + if reported is not None and reported != expected: + raise ObjectStoreError(f"{key}: the store reports checksum {reported}, the parts sent make {expected}") + return expected + + def _open_upload(self, journal: _UploadJournal, full: str, part: int, size: int, checks: bool, *, + fresh: bool = False) -> None: + """Continue the journaled upload of this file version, or create one.""" + state = None if fresh else journal.load() + if state is not None and state.get("upload_id") and "stale_upload_id" not in state: + listed = self._listed_parts(full, str(state["upload_id"]), part, size) + if listed is not None: + journal.upload_id = str(state["upload_id"]) + saved = {int(n): v for n, v in dict(state.get("parts") or {}).items()} + journal.parts = self._agreed(saved, listed, checks) + journal.save() + return + stale = (state or {}).get("stale_upload_id") + if stale: + try: # another version of the file: its parts are never completed + self.client.abort_multipart_upload(Bucket=self.bucket, Key=full, UploadId=str(stale)) + except Exception: # noqa: BLE001 - already gone + pass + extra = {"ChecksumAlgorithm": "SHA256"} if checks else {} + journal.upload_id = self.client.create_multipart_upload(Bucket=self.bucket, Key=full, **extra)["UploadId"] + journal.parts = {} + journal.save() + + def _send_parts(self, journal: _UploadJournal, full: str, source: str | os.PathLike[str], part: int, + size: int, count: int, checks: bool) -> None: + import hashlib + + missing = [n for n in range(1, count + 1) if n not in journal.parts] + if not missing: return - upload_id = self.client.create_multipart_upload(Bucket=self.bucket, Key=full)["UploadId"] - try: - with open(source, "rb") as fh: - def send(number: int) -> dict[str, Any]: - offset = (number - 1) * part - body = os.pread(fh.fileno(), min(part, size - offset), offset) - done = self.client.upload_part(Bucket=self.bucket, Key=full, UploadId=upload_id, - PartNumber=number, Body=body) - return {"PartNumber": number, "ETag": done["ETag"]} - - with ThreadPoolExecutor(self.transfer.concurrency) as pool: - parts = list(pool.map(send, range(1, -(-size // part) + 1))) - self.client.complete_multipart_upload(Bucket=self.bucket, Key=full, UploadId=upload_id, - MultipartUpload={"Parts": parts}) - except BaseException: - self.client.abort_multipart_upload(Bucket=self.bucket, Key=full, UploadId=upload_id) - raise + with open(source, "rb") as fh: + def send(number: int) -> None: + offset = (number - 1) * part + body = os.pread(fh.fileno(), min(part, size - offset), offset) + digest = _b64(hashlib.sha256(body).digest()) if checks else None + extra = {"ChecksumSHA256": digest} if checks else {} + done = self.client.upload_part(Bucket=self.bucket, Key=full, UploadId=journal.upload_id, + PartNumber=number, Body=body, **extra) + journal.record(number, {"ETag": done["ETag"], "sha256": digest, "size": len(body)}) + + with ThreadPoolExecutor(min(self.transfer.concurrency, len(missing))) as pool: + for future in [pool.submit(send, n) for n in missing]: + future.result() + + def _listed_parts(self, full: str, upload_id: str, part: int, size: int) -> dict[int, dict[str, Any]] | None: + """The parts the store holds for ``upload_id`` (number -> ETag, size, + checksum), None when the upload does not exist any more.""" + out: dict[int, dict[str, Any]] = {} + marker = 0 + while True: + try: + page = self.client.list_parts(Bucket=self.bucket, Key=full, UploadId=upload_id, + PartNumberMarker=marker) + except Exception as error: # noqa: BLE001 - botocore ClientError, by code + if self._code(error) in {"NoSuchUpload", *_MISSING}: + return None + raise + for p in page.get("Parts") or (): + out[int(p["PartNumber"])] = {"ETag": p.get("ETag"), "size": p.get("Size"), + "sha256": p.get("ChecksumSHA256")} + if not page.get("IsTruncated"): + return out + marker = int(page["NextPartNumberMarker"]) + + @staticmethod + def _agreed(journaled: Mapping[int, Mapping[str, Any]], listed: Mapping[int, Mapping[str, Any]], + checks: bool) -> dict[int, dict[str, Any]]: + """The journaled parts the store holds as journaled: same ETag, same size, + and (with checksums) the same SHA-256 when the store lists one. A part the + store holds but the journal never recorded is sent again.""" + out = {} + for n, mine in journaled.items(): + theirs = listed.get(int(n)) + if theirs is None or theirs.get("ETag") != mine.get("ETag"): + continue + if theirs.get("size") is not None and mine.get("size") is not None \ + and int(theirs["size"]) != int(mine["size"]): + continue + if checks and (not mine.get("sha256") or (theirs.get("sha256") and theirs["sha256"] != mine["sha256"])): + continue + out[int(n)] = dict(mine) + return out def read_file(self, key: str, target: str | os.PathLike[str]) -> bool: """Download object ``key`` to ``target``; a large one by parallel ranges, @@ -260,8 +742,18 @@ def read_file(self, key: str, target: str | os.PathLike[str]) -> bool: def fetch(offset: int) -> None: end = min(offset + part, size) - 1 - body = self.client.get_object(Bucket=self.bucket, Key=full, Range=f"bytes={offset}-{end}", - **pin)["Body"].read() + for attempt in range(RANGE_READ_ATTEMPTS): + try: + body = self.client.get_object(Bucket=self.bucket, Key=full, Range=f"bytes={offset}-{end}", + **pin)["Body"].read() + except RANGE_READ_RETRYABLE: + # botocore retries the request, not the streamed body: a stalled + # body read of one part would otherwise fail the whole download + if attempt + 1 == RANGE_READ_ATTEMPTS: + raise + time.sleep(min(30.0, 2.0 ** attempt)) + continue + break if len(body) != end - offset + 1: raise ObjectStoreError(f"{self.uri}/{key}: range {offset}-{end} returned {len(body)} B") os.pwrite(fh.fileno(), body, offset) @@ -270,6 +762,42 @@ def fetch(offset: int) -> None: list(pool.map(fetch, range(0, size, part))) return True + def read_files(self, items: Sequence[tuple[str, str | os.PathLike[str]]]) -> list[str]: + """Download each ``(key, target)`` like :meth:`read_file`; returns the keys + that are missing (nothing written for them). With a ``client_factory`` and + ``transfer.download_processes`` > 1 the files are split over that many + child interpreters, each moving ``ceil(file_concurrency / processes)`` files at + once; otherwise ``file_concurrency`` threads of this process move them. The + children are plain ``python -c`` subprocesses fed by stdin, not + ``multiprocessing``: spawn re-imports the caller's ``__main__`` (an unguarded + script runs again) and fork copies a process that may hold threads or a GPU + runtime.""" + items = [(key, str(target)) for key, target in items] + procs = min(int(self.transfer.download_processes), len(items)) + if self.client_factory is None or procs <= 1: + return _read_files_here(self, items, self.transfer.file_concurrency) + threads = max(1, -(-int(self.transfer.file_concurrency) // procs)) + groups = [items[i::procs] for i in range(procs)] + env = dict(os.environ, PYTHONPATH=os.pathsep.join(p for p in sys.path if p)) + children = [] + for group in groups: + child = subprocess.Popen([sys.executable, "-c", READ_FILES_CHILD], stdin=subprocess.PIPE, + stdout=subprocess.PIPE, env=env) + children.append(child) + child.stdin.write(pickle.dumps((self.uri, self.client_factory, asdict(self.transfer), threads, group))) + child.stdin.close() + missing, failed = [], [] + for child in children: + out = child.stdout.read() + if child.wait() != 0: + failed.append(child.returncode) + continue + missing += pickle.loads(out) + if failed: + raise ObjectStoreError(f"{self.uri}: {len(failed)} of {len(children)} download processes " + f"failed (exit {failed}); their stderr is above") + return missing + def list(self, prefix: str, *, start_after: str = "") -> list[str]: strip = len(self.prefix) + 1 if self.prefix else 0 kwargs = {"Bucket": self.bucket, "Prefix": self._full(prefix)} @@ -281,6 +809,27 @@ def list(self, prefix: str, *, start_after: str = "") -> list[str]: keys += [item["Key"][strip:] for item in page.get("Contents", ()) if not item["Key"].endswith("/")] return keys + def list_meta(self, prefix: str) -> tuple[list[tuple[str, float]], float]: + """``([(key, LastModified epoch s)], server time)`` under ``prefix``: the + object times and the clock they compare to both come from the server + (the ``Date`` header of the listing), never from this host's clock.""" + from email.utils import parsedate_to_datetime + + strip = len(self.prefix) + 1 if self.prefix else 0 + out: list[tuple[str, float]] = [] + now: float | None = None + for page in self.client.get_paginator("list_objects_v2").paginate(Bucket=self.bucket, + Prefix=self._full(prefix)): + date = (page.get("ResponseMetadata") or {}).get("HTTPHeaders", {}).get("date") + if date: + stamp = parsedate_to_datetime(date).timestamp() + now = stamp if now is None else max(now, stamp) + out += [(item["Key"][strip:], item["LastModified"].timestamp()) + for item in page.get("Contents", ()) if not item["Key"].endswith("/")] + if now is None: # no Date header: the newest object time is a lower bound of the server clock + now = max((t for _, t in out), default=0.0) + return out, now + def delete(self, key: str) -> None: self.client.delete_object(Bucket=self.bucket, Key=self._key(key)) @@ -308,22 +857,74 @@ def __init__(self, root: str | os.PathLike[str], uri: str, *, environ: Mapping[s def objects(self) -> S3ObjectStore: # type: ignore[override] if self._objects is None: client = self._client or _client(self._environ, require_credentials=True) + # another process builds its own client from the same settings (read_files); + # an injected client (tests) has no such factory + factory = None if self._client is not None else functools.partial( + _client, dict(self._environ), require_credentials=True) self._objects = S3ObjectStore(self.uri, client=client, - transfer=TransferSettings.from_env(self._environ)) + transfer=TransferSettings.from_env(self._environ), + client_factory=factory) return self._objects def _tree(self) -> TreeMirror: if self._mirror is None: self._mirror = TreeMirror(self.root, self.objects, subtrees=STAGED_SUBTREES, - pointers=POINTER_FILES, records=RECORD_FILES) + pointers=POINTER_FILES, records=RECORD_FILES, + independent=INDEPENDENT_DIRS, + file_concurrency=self.objects.transfer.file_concurrency) return self._mirror @property def put_renames(self) -> bool: # type: ignore[override] return self.staged # a stage directory is local disk; the bucket mount is not - def commit(self) -> float: - return self._tree().push() if self.staged else super().commit() + @property + def put_links(self) -> bool: # type: ignore[override] + return self.staged # a 207 GB release export is linked into the stage, not copied + + def commit(self, paths: Sequence[str] | None = None) -> float: + return self._tree().push(paths) if self.staged else super().commit(paths) + + @property + def cut_commits_beside_buffer_reads(self) -> bool: # type: ignore[override] + return self.staged # the buffer goes by the object API; a push takes the cut's paths only + + def push_local_file(self, relative: str) -> int | None: + """Upload one file of the stage now (a trainer rank's cut shard, on whichever + node wrote it: a multi-node island has one stage per node); returns the size + the object store reports, None when this store has no stage.""" + if not self.staged: + return None + return self._tree().push_file(relative) + + def local_file_checksum(self, relative: str) -> str | None: + """The checksum the object store keeps for a file :meth:`push_local_file` + uploaded (``S3ObjectStore.checksum``), None without one.""" + if not self.staged: + return None + return self._tree().object_checksum(relative) + + def evict_local_file(self, relative: str) -> int: + """Drop the stage copy of a file :meth:`push_local_file` uploaded (a cut shard: + a K3.1 cut is ~1.35 TB per node); the bucket keeps it. Returns the bytes freed.""" + return self._tree().evict_file(relative) if self.staged else 0 + + def delete_remote_tree(self, relative: str) -> list[str]: + """Delete every object under ``relative``, including what other nodes uploaded.""" + if not self.staged: + return [] + return self._tree().delete_tree(relative) + + def evict_local(self, relative: str) -> None: + """Drop the stage copy of a pushed directory (a 207 GB Kimi-K3.1 release + kept per release filled the trainer's disk); the bucket keeps it.""" + if self.staged: + self._tree().evict_tree(relative) + + def remove_tree(self, relative: str) -> None: + if self.staged: # an evicted directory has no local files left to prune + self.delete_remote_tree(relative) + super().remove_tree(relative) def reload(self) -> None: if self.staged: @@ -348,17 +949,104 @@ def head_objects(uri: str, *, environ: Mapping[str, str] | None = None, client: transfer=TransferSettings.from_env(env)) -def island_envs(uri: str, environ: Mapping[str, str] | None = None, *, stage: bool = False) -> dict[str, str]: +def island_envs(uri: str, environ: Mapping[str, str] | None = None, *, stage: bool = False, + credentials: bool = True) -> dict[str, str]: """The environment of an island that uses the store's object API: the store hook, the URI, the local stage directory when ``stage`` (one-node islands), and the store settings and credentials of the launcher's environment. Raises - ValueError without store credentials.""" - env = os.environ if environ is None else environ + ValueError without store credentials. ``credentials=False``: the island gets + its credential from a Modal Secret with the AWS names (:data:`CREDENTIALS_ENV` + tells it so); none is required here or copied.""" + env = dict(os.environ if environ is None else environ) parse_s3_uri(uri) - if not has_store_credentials(env): + if credentials and not (env.get(KEY_ID_ENV) and env.get(SECRET_ENV)) and env.get(PROFILE_ENV): + env.update(profile_credentials(env)) # the profile the launch names: its key goes to the island + if credentials and not (env.get(KEY_ID_ENV) and env.get(SECRET_ENV)): raise ValueError(f"an island that uses the s3 store {uri} needs store credentials: set " - f"{KEY_ID_ENV} and {SECRET_ENV} (a key limited to the store's bucket) on the " - "launcher; the island never uses the launcher's cloud credentials") + f"{KEY_ID_ENV} and {SECRET_ENV} (a key limited to the store's bucket), or " + f"{PROFILE_ENV} to an AWS profile holding one, on the launcher; the island " + "never uses the launcher's cloud credentials") out = {STORE_IMPL_ENV: STORE_IMPL, URI_ENV: uri, **({STAGE_ENV: DEFAULT_STAGE} if stage else {})} - out.update({name: env[name] for name in (*SETTING_ENVS, *CREDENTIAL_ENVS) if env.get(name)}) + out.update({name: env[name] for name in SETTING_ENVS if env.get(name)}) + # an endpoint / region found in the profile: the island has no profile to read them from + endpoint, region = store_endpoint(env), store_region(env) + if endpoint and ENDPOINT_ENV not in out: + out[ENDPOINT_ENV] = endpoint + if region and REGION_ENV not in out: + out[REGION_ENV] = region + if credentials: + out.update({name: env[name] for name in CREDENTIAL_ENVS if env.get(name)}) + else: + out[CREDENTIALS_ENV] = CREDENTIALS_AWS_ENV return out + + +def profile_credentials(environ: Mapping[str, str]) -> dict[str, str]: + """:data:`KEY_ID_ENV` / :data:`SECRET_ENV` (/ :data:`TOKEN_ENV`) from the store's + AWS profile (:data:`PROFILE_ENV`); ``{}`` when it has none.""" + profile = environ.get(PROFILE_ENV) + if not profile or profile_config(environ).get("_missing"): + return {} + import boto3 + + creds = boto3.session.Session(profile_name=profile).get_credentials() + if creds is None: + return {} + frozen = creds.get_frozen_credentials() + out = {KEY_ID_ENV: frozen.access_key, SECRET_ENV: frozen.secret_key} + if frozen.token: + out[TOKEN_ENV] = frozen.token + return out + + +# ------------------------------------------------------------------ CLI helpers +def add_store_client_arguments(parser: Any) -> None: + """``--endpoint-url`` and ``--profile`` of a CLI that opens an ``s3://`` store.""" + parser.add_argument("--endpoint-url", default=None, metavar="URL", + help=f"S3-compatible endpoint (R2, MinIO); default: {ENDPOINT_ENV}, AWS_ENDPOINT_URL, " + "the profile's endpoint_url, else AWS S3") + parser.add_argument("--profile", default=None, metavar="NAME", + help=f"AWS profile with the store's key, endpoint_url and region " + f"(default: {PROFILE_ENV}, AWS_PROFILE)") + + +def apply_store_client_arguments(args: Any, environ: Any = None) -> None: + """Put ``--endpoint-url`` / ``--profile`` into ``environ`` (default: the process + environment), where every store client of this process reads them.""" + env = os.environ if environ is None else environ + if getattr(args, "endpoint_url", None): + env[ENDPOINT_ENV] = str(args.endpoint_url) + if getattr(args, "profile", None): + env[PROFILE_ENV] = str(args.profile) + + +_AUTH_CODES = frozenset({"InvalidAccessKeyId", "SignatureDoesNotMatch", "AccessDenied", "Unauthorized", + "InvalidToken", "ExpiredToken", "403", "401"}) + + +def describe_store_error(error: BaseException, uri: str = "", environ: Mapping[str, str] | None = None) -> str | None: + """A message naming what to set for a store error of credentials or endpoint; + None for any other error.""" + env = os.environ if environ is None else environ + name = type(error).__name__ + code = str(getattr(error, "response", {}).get("Error", {}).get("Code", "")) + endpoint = store_endpoint(env) + where = f"endpoint {endpoint}" if endpoint else "AWS S3 (no endpoint set)" + hint_endpoint = (f"set {ENDPOINT_ENV} or AWS_ENDPOINT_URL (or --endpoint-url, or endpoint_url in the " + f"AWS profile) to the store's endpoint, e.g. https://.r2.cloudflarestorage.com") + target = f"s3 store {uri}" if uri else "s3 store" + if name in ("NoCredentialsError", "PartialCredentialsError"): + return f"{target}: {missing_credentials_message()}" + if name == "ProfileNotFound": + return f"{target}: AWS profile {store_profile(env)!r} not found; set {PROFILE_ENV} (or --profile) to a configured profile" + if name in ("EndpointConnectionError", "ConnectTimeoutError", "InvalidEndpointConfigurationError") \ + or isinstance(error, ValueError) and "endpoint" in str(error).lower(): + return f"{target}: cannot reach {where} ({error}); {hint_endpoint}" + if code in _AUTH_CODES: + extra = "" if endpoint else f"; a key of an S3-compatible store (R2) is refused by AWS: {hint_endpoint}" + return (f"{target}: {where} refused the credentials ({code}). Check {KEY_ID_ENV} / {SECRET_ENV} " + f"or the profile in {PROFILE_ENV} / AWS_PROFILE{extra}") + if code == "AuthorizationHeaderMalformed" or (code in ("InvalidRegionName", "InvalidRequest") + and endpoint and _is_r2(endpoint)): + return f"{target}: {where} refused the region ({code}); set {REGION_ENV}=auto for Cloudflare R2" + return None diff --git a/yeto/codepack.py b/yeto/codepack.py index bf0f1db89..7b74f9124 100644 --- a/yeto/codepack.py +++ b/yeto/codepack.py @@ -550,7 +550,7 @@ class RunPack: def resolve_run_pack(*, previous: Mapping[str, Any] | None, repo: str | os.PathLike[str], store: PackStore, cache_root: str | os.PathLike[str], semantic_sha256: str | None, upgrade_reason: str | None = None, - release_image: str | None = None, + release_image: str | None = None, keep_pinned: bool = False, log: Callable[[str], None] = print) -> RunPack: """The pack this start of the run uses. @@ -562,13 +562,25 @@ def resolve_run_pack(*, previous: Mapping[str, Any] | None, repo: str | os.PathL With ``release_image`` the candidate pack is that image (:func:`release_pack`) instead of the local tree, and nothing is stored. A run started on one pack - moves to another (another image, or between modes) only by an explicit upgrade.""" + moves to another (another image, or between modes) only by an explicit upgrade. + + Drift guard: when the local tree (commit + dirty patch, i.e. the pack digest) + differs from the pinned pack, a plain relaunch is refused instead of silently + running the old code. ``keep_pinned`` (a deliberate resume on the run's pack) + reuses it anyway; ``upgrade_reason`` moves to the local tree.""" old_record = (previous or {}).get("code_pack") if release_image is not None: return _resolve_release_pack(old_record, previous, release_image, semantic_sha256=semantic_sha256, upgrade_reason=upgrade_reason, log=log) + base = getattr(store, "base", None) or getattr(store, "uri_base", None) + old_uri = (old_record or {}).get("pack_uri") + if (old_uri and base and (old_record or {}).get("mode") != CODE_PACK_MODE_RELEASE and not str(old_uri).startswith(base.rstrip("/") + "/")): + # The manifest is keyed by cluster prefix: a previous run of the same prefix + # on another checkpoint store is a different run, not this run's resume. + log(f"[codepack] run manifest pack {old_uri} is not in store {base}: new run") + old_record = None if not old_record: pack, data = make_and_store(repo, store) log(f"[codepack] new run pack {pack.digest[:16]} (commit {pack.commit[:12]}" @@ -588,6 +600,15 @@ def resolve_run_pack(*, previous: Mapping[str, Any] | None, repo: str | os.PathL new_semantic=semantic_sha256, reason=upgrade_reason) log(f"[codepack] upgrade {old.digest[:16]} -> {new.digest[:16]}: {upgrade_reason}") return RunPack(new, unpack(data, new.digest, cache_root), reused=False) + if not keep_pinned: + local, _ = build_pack(repo) + if local.digest != old.digest: + raise CodePackError( + f"the run manifest pins code pack {old.digest[:16]} (commit {(old.commit or '?')[:12]}" + f"{', dirty' if old.dirty else ''}), but the local tree is {local.digest[:16]} " + f"(commit {local.commit[:12]}{', dirty' if local.dirty else ''}). Resume on the " + "pinned code with --rl-code-pack-keep, move the run to the local tree with " + "--rl-code-pack-upgrade REASON, or start a new run (another --cluster-prefix)") log(f"[codepack] reusing run pack {old.digest[:16]} from the run manifest") return RunPack(old, fetch_tree(old, store, cache_root), reused=True) diff --git a/yeto/cost/__init__.py b/yeto/cost/__init__.py new file mode 100644 index 000000000..44a239fa4 --- /dev/null +++ b/yeto/cost/__init__.py @@ -0,0 +1,10 @@ +"""Cost protection for unattended runs: resource ledger, budget breaker, +idle shutdown, exact cleanup and the cost report. + +:mod:`.ledger` crash-safe per-run list of every billed resource this run created +:mod:`.billing` what those resources cost (Modal billing API + live estimate, NC lease x rate) +:mod:`.idle` idle detection with a loading-phase exemption +:mod:`.guard` the breaker thread the launcher runs (warn, graceful stop, forced cleanup) +:mod:`.cleanup` ``yeto cleanup --run``: release exactly what the ledger lists +:mod:`.report` ``yeto cost --run``: per resource / per hour cost and GPU effective utilization +""" diff --git a/yeto/cost/billing.py b/yeto/cost/billing.py new file mode 100644 index 000000000..e129d6c8e --- /dev/null +++ b/yeto/cost/billing.py @@ -0,0 +1,311 @@ +"""What the resources of one run cost so far. + +* Modal app: billed cost from ``modal.billing.workspace_billing_report`` (hourly + rows, matched by exact app id from the ledger, only hours that overlap the + span the app belonged to this run), plus an estimate for the time after the + newest billed hour (billing lags): duration x rate, where the rate is the + ledger's ``usd_per_hour``, else the price table, else the newest complete + billed hour of that app. +* NC lease: (released or now) - granted_at, times the lease rate. ``lease_until`` + does NOT bound it: a lease past its end that was never deleted keeps billing. +* sky cluster: price table x wall clock (if the ledger carries gpu / gpus). + +A resource with no rate is "unpriced" (``usd=None``) - never $0 - and the +summary says so; the guard warns about unpriced resources. +""" + +from __future__ import annotations + +import math +import os +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import Any + +from .ledger import Resource + +HOUR = 3600.0 +NC_USD_PER_UCR_ENV = "YETO_NC_USD_PER_UCR" +NC_USD_PER_GPU_HOUR_ENV = "YETO_NC_USD_PER_GPU_HOUR" +NC_UCR_PER_GPU_HOUR_ENV = "YETO_NC_UCR_PER_GPU_HOUR" +# National Compute defaults: 1 credit (1e6 ucr) ~ 1 USD, and the 8-GPU MI355X node at +# 3e6 ucr per GPU-hour (24e6 ucr/h, ~24 USD/h); the env above overrides either +DEFAULT_NC_USD_PER_UCR = 1e-6 +DEFAULT_NC_UCR_PER_GPU_HOUR = 3_000_000 +# Modal list price per GPU-hour, used when neither the ledger nor a price table +# gives a rate (2x8 B300 ~ 136 USD/h). Keys are lower-case Modal gpu names. +DEFAULT_MODAL_USD_PER_GPU_HOUR = {"b300": 8.5, "b200": 6.25, "h200": 4.54, "h100": 3.95, + "a100-80gb": 2.5, "a100": 2.1, "l40s": 1.95} + + +def default_modal_rate(gpu: Any, gpus: Any) -> float | None: + unit = DEFAULT_MODAL_USD_PER_GPU_HOUR.get(str(gpu or "").lower().split(":")[0].rstrip("!")) + return None if unit is None or not gpus else unit * float(gpus) + + +def _ts(value: Any) -> float | None: + if value is None or value == "": + return None + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, datetime): + return (value if value.tzinfo else value.replace(tzinfo=UTC)).timestamp() + try: + return float(value) + except (TypeError, ValueError): + pass + dt = datetime.fromisoformat(str(value).replace("Z", "+00:00")) + return (dt if dt.tzinfo else dt.replace(tzinfo=UTC)).timestamp() + + +def floor_hour(t: float) -> float: + return math.floor(t / HOUR) * HOUR + + +def hour_label(t: float) -> str: + return datetime.fromtimestamp(floor_hour(t), UTC).strftime("%Y-%m-%dT%H:00Z") + + +def split_by_hour(start: float, end: float, usd_per_hour: float) -> dict[str, float]: + out: dict[str, float] = {} + t = start + while t < end: + nxt = min(end, floor_hour(t) + HOUR) + out[hour_label(t)] = out.get(hour_label(t), 0.0) + (nxt - t) / HOUR * usd_per_hour + t = nxt + return out + + +# -- Modal billing source ---------------------------------------------------------- + +@dataclass +class BillingRow: + object_id: str + description: str + interval_start: float + cost: float + + +def modal_billing_fetch(start: float, end: float) -> list[BillingRow]: + """Production fetch: ``modal.billing.workspace_billing_report`` at hourly resolution.""" + import modal + + rows = modal.billing.workspace_billing_report( + start=datetime.fromtimestamp(floor_hour(start), UTC), + end=datetime.fromtimestamp(end, UTC), resolution="h") + return [BillingRow(str(r["object_id"]), str(r.get("description") or ""), + _ts(r["interval_start"]), float(r["cost"])) for r in rows] + + +class ModalBilling: + """Caches the billing report; refreshed at most every ``refresh_s``. + A failing fetch keeps the previous rows (the estimate covers the gap) and + records the error.""" + + def __init__(self, fetch: Callable[[float, float], list[BillingRow]] = modal_billing_fetch, *, + refresh_s: float = 300.0) -> None: + self.fetch = fetch + self.refresh_s = refresh_s + self._rows: list[BillingRow] = [] + self._at: float | None = None + self._start: float | None = None + self.error: str | None = None + + def rows(self, start: float, now: float) -> list[BillingRow]: + if self._at is None or now - self._at >= self.refresh_s or ( + self._start is not None and start < self._start): + try: + self._rows = list(self.fetch(start, now)) + self._start = start + self.error = None + except Exception as exc: # noqa: BLE001 + self.error = f"{type(exc).__name__}: {exc}" + self._at = now + return self._rows + + +# -- per-resource cost --------------------------------------------------------------- + +@dataclass +class ItemCost: + key: str + kind: str + name: str | None + start: float + end: float + open: bool + billed_usd: float = 0.0 + estimated_usd: float = 0.0 + priced: bool = True + rate_usd_per_hour: float | None = None + hours: dict[str, float] = field(default_factory=dict) + note: str = "" + + @property + def usd(self) -> float | None: + if not self.priced and self.billed_usd == 0.0: + return None + return self.billed_usd + self.estimated_usd + + @property + def billable_s(self) -> float: + return max(0.0, self.end - self.start) + + +@dataclass +class CostSummary: + now: float + items: list[ItemCost] + billing_error: str | None = None + + @property + def total_usd(self) -> float: + return sum(i.usd or 0.0 for i in self.items) + + @property + def unpriced(self) -> list[ItemCost]: + return [i for i in self.items if i.usd is None and i.kind != "modal_call"] + + def to_json(self) -> dict[str, Any]: + return {"now": self.now, "total_usd": round(self.total_usd, 4), + "unpriced": [i.key for i in self.unpriced], "billing_error": self.billing_error, + "items": [{**{k: v for k, v in i.__dict__.items()}, "usd": i.usd, + "billable_s": i.billable_s} for i in self.items]} + + +@dataclass +class Rates: + """Where rates come from when the ledger entry carries none.""" + prices: Mapping[str, Any] | None = None # yeto.dashboard.cost price table + nc_usd_per_ucr: float | None = DEFAULT_NC_USD_PER_UCR + nc_usd_per_gpu_hour: float | None = None + nc_ucr_per_gpu_hour: float | None = DEFAULT_NC_UCR_PER_GPU_HOUR # when the ledger carries no rate + + @classmethod + def from_env(cls, prices_path: str | None = None, environ=None) -> "Rates": + env = os.environ if environ is None else environ + prices = None + path = prices_path or env.get("YETO_DASHBOARD_PRICES") + if path: + from ..dashboard.cost import load_prices + + prices = load_prices(path) + + def f(name): + v = env.get(name) + return float(v) if v else None + + usd_per_ucr = f(NC_USD_PER_UCR_ENV) + ucr = f(NC_UCR_PER_GPU_HOUR_ENV) + return cls(prices, DEFAULT_NC_USD_PER_UCR if usd_per_ucr is None else usd_per_ucr, + f(NC_USD_PER_GPU_HOUR_ENV), DEFAULT_NC_UCR_PER_GPU_HOUR if ucr is None else ucr) + + def table_rate(self, cloud: str, attrs: Mapping[str, Any]) -> float | None: + if self.prices is None or not attrs.get("gpu") or not attrs.get("gpus"): + return None + from ..dashboard.cost import price_per_gpu_h + + unit = price_per_gpu_h(dict(self.prices), cloud, attrs.get("gpu"), attrs.get("price_key")) + return None if unit is None else unit * float(attrs["gpus"]) + + +def _modal_app_cost(r: Resource, rows: Iterable[BillingRow], now: float, rates: Rates) -> ItemCost: + end = r.released_at if r.released_at is not None else now + item = ItemCost(r.key, r.kind, r.name, r.created_at, end, r.open) + lo, hi = floor_hour(r.created_at), end + mine = sorted((b for b in rows if b.object_id == r.id and lo <= b.interval_start < hi), + key=lambda b: b.interval_start) + for b in mine: + item.billed_usd += b.cost + lbl = hour_label(b.interval_start) + item.hours[lbl] = item.hours.get(lbl, 0.0) + b.cost + billed_until = (mine[-1].interval_start + HOUR) if mine else r.created_at + rate = r.attrs.get("usd_per_hour") + rate = float(rate) if rate is not None else rates.table_rate("modal", r.attrs) + if rate is None: + rate = default_modal_rate(r.attrs.get("gpu"), r.attrs.get("gpus")) + if rate is not None: + item.note = "default Modal list price" + if rate is None: + full = [b for b in mine if b.interval_start >= r.created_at and b.interval_start + HOUR <= end] + if full: + rate = full[-1].cost + item.note = "rate from the newest complete billed hour" + item.rate_usd_per_hour = rate + gap_start = max(r.created_at, billed_until) + if gap_start < end: + if rate is None: + item.priced = False + item.note = "no rate for the unbilled tail (set usd_per_hour or a price table)" + else: + for lbl, usd in split_by_hour(gap_start, end, rate).items(): + item.estimated_usd += usd + item.hours[lbl] = item.hours.get(lbl, 0.0) + usd + if r.attrs.get("shared"): + item.note = (item.note + "; " if item.note else "") + \ + "app shared with other runs: billed hours overlapping this run are counted in full" + return item + + +def _nc_rate(r: Resource, rates: Rates) -> float | None: + gpus = float(r.attrs.get("gpus") or 8) + if r.attrs.get("usd_per_gpu_hour") is not None: + return gpus * float(r.attrs["usd_per_gpu_hour"]) + if r.attrs.get("usd_per_hour") is not None: + return float(r.attrs["usd_per_hour"]) + if rates.nc_usd_per_gpu_hour is not None: + return gpus * rates.nc_usd_per_gpu_hour + ucr = r.attrs.get("rate_ucr_per_gpu_hour") + if ucr is None: + ucr = rates.nc_ucr_per_gpu_hour + if ucr is not None and rates.nc_usd_per_ucr is not None: + return gpus * float(ucr) * rates.nc_usd_per_ucr + return None + + +def _span_cost(r: Resource, start: float, now: float, rate: float | None, note: str = "") -> ItemCost: + end = r.released_at if r.released_at is not None else now + item = ItemCost(r.key, r.kind, r.name, start, max(start, end), r.open, note=note, + rate_usd_per_hour=rate) + if rate is None: + item.priced = False + item.note = note or "no rate (ledger attrs / price table / env)" + else: + item.hours = split_by_hour(start, max(start, end), rate) + item.estimated_usd = sum(item.hours.values()) + return item + + +def resource_costs(resources: Iterable[Resource], now: float, *, rates: Rates | None = None, + billing: ModalBilling | None = None) -> CostSummary: + rates = rates or Rates() + resources = [r for r in resources if not r.intent_only or r.open] + apps = [r for r in resources if r.kind == "modal_app" and r.id] + rows: list[BillingRow] = [] + if apps and billing is not None: + rows = billing.rows(min(r.created_at for r in apps), now) + items: list[ItemCost] = [] + for r in resources: + if r.kind == "modal_app": + if r.id is None: + items.append(_span_cost(r, r.created_at, now, None, "app id unknown (intent only)")) + else: + items.append(_modal_app_cost(r, rows, now, rates)) + elif r.kind == "nc_lease": + start = _ts(r.attrs.get("granted_at")) or r.created_at + note = "" + until = _ts(r.attrs.get("lease_until")) + end = r.released_at if r.released_at is not None else now + if until is not None and end > until: + note = f"{(end - until) / 60:.0f} min past lease_until" + items.append(_span_cost(r, start, now, _nc_rate(r, rates), note)) + elif r.kind == "sky_cluster": + rate = r.attrs.get("usd_per_hour") + rate = float(rate) if rate is not None else rates.table_rate(r.attrs.get("cloud") or "", r.attrs) + items.append(_span_cost(r, r.created_at, now, rate)) + else: # modal_call: its cost is in its app; listed for liveness / utilization + end = r.released_at if r.released_at is not None else now + items.append(ItemCost(r.key, r.kind, r.name, r.created_at, end, r.open, billed_usd=0.0, + note="cost counted in its Modal app")) + return CostSummary(now, items, billing_error=billing.error if billing is not None else None) diff --git a/yeto/cost/cleanup.py b/yeto/cost/cleanup.py new file mode 100644 index 000000000..e8476a662 --- /dev/null +++ b/yeto/cost/cleanup.py @@ -0,0 +1,239 @@ +"""``yeto cleanup --run ``: release exactly the resources the run's ledger lists. + +Rules: + +* Only ledger entries are touched, each by its exact provider id (Modal + FunctionCall id, Modal app id, NC ticket + account, sky cluster name). No + name pattern, no prefix match, no ``--all``. +* An ``intent`` entry (crash between "about to create" and "created") is + resolved only when the provider lists exactly one app whose name EQUALS the + recorded name and that was created after the intent; otherwise it is + reported unresolved and left alone. +* A Modal app marked ``shared`` (reused by other runs) is not stopped; only this + run's FunctionCalls in it are cancelled. +* NC leases go through the NC lease manager's own guarded release + (``yeto.launch.nc.cli.down``: refuses tickets it did not create). +* Order: FunctionCalls, then apps, then NC leases, then sky clusters. Each + success is recorded as ``release`` in the ledger (a rerun skips it); each + failure as ``release_failed`` (a rerun retries it). +""" + +from __future__ import annotations + +import json +import subprocess +import sys +from dataclasses import dataclass +from typing import Any, Protocol + +from .ledger import Resource, ResourceLedger + +ORDER = {"modal_call": 0, "modal_app": 1, "nc_lease": 2, "sky_cluster": 3} + + +class CleanupOps(Protocol): + def modal_cancel_call(self, call_id: str) -> None: ... + def modal_stop_app(self, app_id: str) -> None: ... + def modal_app_ids_named(self, name: str, created_after: float) -> list[str] | None: ... + def nc_release(self, account: int, ticket: str) -> dict[str, Any]: ... + def sky_down(self, cluster: str) -> bool: ... + + +class ProductionCleanupOps: + """Modal CLI/SDK by exact id, the NC lease manager's guarded down, sky down.""" + + def modal_cancel_call(self, call_id: str) -> None: + import modal + + modal.FunctionCall.from_id(call_id).cancel(terminate_containers=True) + + def modal_stop_app(self, app_id: str) -> None: + if not str(app_id).startswith("ap-"): + raise ValueError(f"refusing to stop {app_id!r}: not a Modal app id") + proc = subprocess.run([sys.executable, "-m", "modal", "app", "stop", "--yes", app_id], + capture_output=True, text=True, timeout=120) + out = (proc.stdout + proc.stderr).strip() + if proc.returncode != 0 and "already" not in out.lower() and "stopped" not in out.lower(): + raise RuntimeError(f"modal app stop {app_id}: {out or proc.returncode}") + + def modal_app_ids_named(self, name: str, created_after: float) -> list[str] | None: + proc = subprocess.run([sys.executable, "-m", "modal", "app", "list", "--json"], + capture_output=True, text=True, timeout=120) + if proc.returncode != 0: + return None + from .billing import _ts + + out = [] + for app in json.loads(proc.stdout or "[]"): + if app.get("description") != name: # exact equality, never a pattern + continue + if str(app.get("state", "")).lower() == "stopped": + continue + try: + created = _ts(app.get("created_at")) + except ValueError: + created = None + if created is not None and created >= created_after - 5: + out.append(str(app.get("app_id"))) + return out + + def nc_release(self, account: int, ticket: str) -> dict[str, Any]: + from ..launch.nc import cli as nc_cli + + return nc_cli.down(int(account), str(ticket)) + + def sky_down(self, cluster: str) -> bool: + import sky + + sky.stream_and_get(sky.down(cluster)) + return True + + +@dataclass +class Action: + key: str + action: str + ok: bool + detail: str = "" + + +def _release_one(r: Resource, ops: CleanupOps, ledger: ResourceLedger, dry_run: bool, + resolved_apps: dict[str, str]) -> Action: + if r.kind == "modal_call": + if dry_run: + return Action(r.key, f"would cancel FunctionCall {r.id}", True) + ops.modal_cancel_call(r.id) + ledger.released("modal_call", r.id, r.name, by="cleanup") + return Action(r.key, f"cancelled FunctionCall {r.id}", True) + if r.kind == "modal_app": + if r.attrs.get("shared"): + return Action(r.key, "shared app left running (its calls of this run were cancelled)", True) + app_id = r.id + if app_id is None: + ids = ops.modal_app_ids_named(r.name, r.created_at) + if ids is None: + return Action(r.key, "intent unresolved: app list unavailable", False) + if not ids: + if not dry_run: + ledger.released("modal_app", None, r.name, by="cleanup", never_created=True) + return Action(r.key, "intent never became an app: nothing to stop", True) + if len(ids) > 1: + return Action(r.key, f"intent unresolved: {len(ids)} live apps named exactly " + f"{r.name!r} created after the intent; left alone", False) + app_id = ids[0] + resolved_apps[r.key] = app_id + if dry_run: + return Action(r.key, f"would stop app {app_id}", True) + ops.modal_stop_app(app_id) + if r.id is None: + ledger.created("modal_app", app_id, r.name, resolved_by="cleanup") + ledger.released("modal_app", app_id, r.name, by="cleanup") + return Action(r.key, f"stopped app {app_id}", True) + if r.kind == "nc_lease": + account = r.attrs.get("account") + if account is None: + return Action(r.key, "NC lease without account in the ledger: not released", False) + if dry_run: + return Action(r.key, f"would release NC ticket {r.id} (account {account})", True) + res = ops.nc_release(int(account), r.id) + ledger.released("nc_lease", r.id, r.name, by="cleanup", result=res) + return Action(r.key, f"released NC ticket {r.id}: {res}", True) + if r.kind == "sky_cluster": + if dry_run: + return Action(r.key, f"would sky down {r.id}", True) + ok = ops.sky_down(r.id) + if ok: + ledger.released("sky_cluster", r.id, r.name, by="cleanup") + return Action(r.key, f"sky down {r.id}", bool(ok), "" if ok else "not confirmed") + return Action(r.key, "unknown kind", False) + + +def cleanup(ledger: ResourceLedger, ops: CleanupOps, *, dry_run: bool = False, + kinds: tuple[str, ...] | None = None, + keys: tuple[str, ...] | None = None) -> list[Action]: + actions = [] + resolved: dict[str, str] = {} + pending = sorted(ledger.open_resources(), key=lambda r: (ORDER.get(r.kind, 9), r.created_at)) + for r in pending: + if (kinds is not None and r.kind not in kinds) or (keys is not None and r.key not in keys): + continue + try: + actions.append(_release_one(r, ops, ledger, dry_run, resolved)) + except Exception as exc: # noqa: BLE001 - every other resource still gets its release + if not dry_run: + ledger.release_failed(r.kind, r.id, r.name, f"{type(exc).__name__}: {exc}") + actions.append(Action(r.key, "release failed", False, f"{type(exc).__name__}: {exc}")) + return actions + + +# NC leases belong to the NC sampler pool (python -m yeto.launch.nc pool / pool-stop): a trainer +# launch that fails or exits must never release them (2026-11-10: a failed trainer launch's +# backstop cleanup released a healthy, set-up sampler node). Only an explicit +# `yeto cleanup --run R --include-nc` (or pool-stop) releases them. +POOL_OWNED_KINDS = frozenset({"nc_lease"}) + + +def cleanup_run(run: str, ops: CleanupOps | None = None, *, dry_run: bool = False, + log=print, include_nc: bool = False) -> bool: + """Release the run's open resources except the NC pool's leases (``include_nc``); + True when nothing it may release is left open.""" + ledger = ResourceLedger.for_run(run) + if not ledger.path.exists(): + log(f"[cleanup] run {run}: no resource ledger at {ledger.path}; nothing to release") + return True + locks: list = [] + refused = False + if include_nc and any(r.kind in POOL_OWNED_KINDS for r in ledger.open_resources()): + held = _take_pool_locks(run, dry_run=dry_run, log=log) + if held is None: # a pool loop of this run runs: it owns its nodes (pool-stop asks it) + include_nc, refused = False, True + else: + locks = held + try: + return _cleanup_run(run, ledger, ops, dry_run=dry_run, log=log, include_nc=include_nc) and not refused + finally: + for fh in locks: + fh.close() + + +def _take_pool_locks(run: str, *, dry_run: bool, log) -> list | None: + """The pool locks (``pool_run.lock_path``) recording ``run``, all taken; None when a + live pool loop (or a pool-stop) holds one: NC leases are then not released here.""" + try: + from yeto.launch.nc.pool_run import acquire_lock, run_pool_locks # noqa: PLC0415 + except Exception as error: # noqa: BLE001 - a tree without the NC pool: nothing to share + log(f"[cleanup] NC pool lock unavailable ({type(error).__name__}); releasing without it") + return [] + held = [] + for path in run_pool_locks(run): + fh = acquire_lock(path, run=run, role="cleanup") + if fh is None: + for other in held: + other.close() + log(f"[cleanup] a pool loop of run {run} holds {path}: its NC leases are left to it " + f"(stop it with python -m yeto.launch.nc pool-stop --config )") + return None + held.append(fh) + return held + + +def _cleanup_run(run: str, ledger: ResourceLedger, ops, *, dry_run: bool, log, include_nc: bool) -> bool: + skipped = [] if include_nc else [r for r in ledger.open_resources() if r.kind in POOL_OWNED_KINDS] + for r in skipped: + log(f"[cleanup] keep {r.key}: an NC pool lease (release it with pool-stop, or " + f"yeto cleanup --run {run} --include-nc)") + actions = cleanup(ledger, ops or ProductionCleanupOps(), dry_run=dry_run, + kinds=None if include_nc else _kinds_except(ledger, POOL_OWNED_KINDS)) + if not actions: + log(f"[cleanup] run {run}: every ledger resource is already released") + for a in actions: + log(f"[cleanup] {'ok ' if a.ok else 'FAIL'} {a.key}: {a.action}" + + (f" ({a.detail})" if a.detail else "")) + remaining = [r for r in ledger.open_resources() + if not (r.kind == "modal_app" and r.attrs.get("shared")) + and (include_nc or r.kind not in POOL_OWNED_KINDS)] + return all(a.ok for a in actions) and (dry_run or not remaining) + + +def _kinds_except(ledger: ResourceLedger, skip: frozenset[str]) -> tuple[str, ...]: + return tuple(sorted({r.kind for r in ledger.open_resources()} - set(skip))) or ("",) diff --git a/yeto/cost/guard.py b/yeto/cost/guard.py new file mode 100644 index 000000000..aa4fcad2e --- /dev/null +++ b/yeto/cost/guard.py @@ -0,0 +1,306 @@ +"""The run's cost breaker: budget, idle shutdown, NC lease expiry. + +The launcher runs one :class:`CostGuard` thread per run. Every ``poll_s`` it + +1. reads the new lines of the run's event tapes (idle detection input); +2. totals the cost of every resource in the run's ledger (:mod:`.billing`); +3. at ``budget_warn_usd`` emits ``cost_budget_warn`` (once); +4. at ``budget_usd`` emits ``cost_budget_exceeded`` and calls ``graceful_stop`` + (default: the run's STOP flag - the trainer loop stops the trainer with a cut, + the fleet ends and the launcher tears down every resource); if the run has not + ended ``stop_grace_s`` later it calls ``force_stop`` (default: release every + ledger resource by exact id, :mod:`.cleanup`); +5. idle (see :mod:`.idle`): ``idle_warn`` after N minutes, ``idle_shutdown`` and + the same graceful -> forced stop after N + M minutes; +6. an NC lease still open ``nc_overdue_grace_s`` after its ``lease_until`` is + reported (``nc_lease_overdue``, once): it keeps billing past its end, but the + NC sampler pool owns its leases (its own lease-end drain, ``pool-stop``, + ``yeto cleanup --include-nc``); a trainer-side guard never releases one + unless ``release_overdue_nc`` is set. + +Every event goes to ``/cost-events.jsonl`` and to stdout +(``[cost-guard] ...``), so the launcher log and the dashboard see it. +""" + +from __future__ import annotations + +import glob +import json +import os +import threading +import time +from collections.abc import Callable, Iterable +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from .billing import CostSummary, ModalBilling, Rates, _ts, resource_costs +from .idle import IdleConfig, IdleDetector +from .ledger import Resource, ResourceLedger + +EVENTS_FILE = "cost-events.jsonl" +PREFIX = "YETO_RL_EVENT " + + +@dataclass +class GuardConfig: + budget_usd: float | None = None + budget_warn_usd: float | None = None # default: 80 % of budget_usd + poll_s: float = 60.0 + stop_grace_s: float = 20 * 60 + nc_overdue_grace_s: float = 5 * 60 + # the NC pool owns its leases: a trainer's guard only alerts on an overdue one + release_overdue_nc: bool = False + idle: IdleConfig = field(default_factory=IdleConfig) + + @property + def warn_usd(self) -> float | None: + if self.budget_warn_usd is not None: + return self.budget_warn_usd + return None if self.budget_usd is None else 0.8 * self.budget_usd + + +class TapeTail: + """New records of every file matching the globs, from where the last read stopped.""" + + def __init__(self, globs: Iterable[str]) -> None: + self.globs = list(globs) + self.offsets: dict[str, int] = {} + + def read(self) -> list[tuple[str, dict[str, Any]]]: + out = [] + for pattern in self.globs: + for path in sorted(set(glob.glob(pattern, recursive=True))): + if path.endswith(".gz"): + continue + try: + size = os.path.getsize(path) + off = self.offsets.get(path, 0) + if size < off: # rewritten file: start over + off = 0 + if size == off: + continue + with open(path, "rb") as fh: + fh.seek(off) + data = fh.read(size - off) + except OSError: + continue + cut = data.rfind(b"\n") + if cut < 0: + continue # no complete line yet + self.offsets[path] = off + cut + 1 + for raw in data[:cut].decode("utf-8", "replace").splitlines(): + i = raw.find(PREFIX) + text = raw[i + len(PREFIX):] if i >= 0 else raw + try: + rec = json.loads(text) + except ValueError: + continue + if isinstance(rec, dict) and "event" in rec: + out.append((path, rec)) + return out + + +def record_time(rec: dict[str, Any], default: float) -> float: + for k in ("time_unix", "ts", "time", "t", "wall_time"): + v = rec.get(k) + if v is None: + continue + try: + t = _ts(v) + except ValueError: + continue + if t is not None and t > 1e9: + return t + return default + + +class CostGuard: + def __init__(self, run: str, config: GuardConfig, *, ledger: ResourceLedger, + tapes: TapeTail | None = None, rates: Rates | None = None, + billing: ModalBilling | None = None, + graceful_stop: Callable[[str], None] | None = None, + force_stop: Callable[[str], None] | None = None, + release_resource: Callable[[Resource], None] | None = None, + run_ended: Callable[[], bool] = lambda: False, + events_path: str | os.PathLike | None = None, + clock: Callable[[], float] = time.time, log: Callable[[str], None] = print, + start: float | None = None) -> None: + self.run = run + self.config = config + self.ledger = ledger + self.tapes = tapes or TapeTail([]) + self.rates = rates or Rates() + self.billing = billing + self.graceful_stop = graceful_stop or self._default_graceful_stop + self.force_stop = force_stop or self._default_force_stop + self.release_resource = release_resource or self._default_release + self.run_ended = run_ended + self.clock = clock + self.log = log + self.start_ts = clock() if start is None else start + self.events_path = Path(events_path) if events_path else ledger.path.parent / EVENTS_FILE + self.idle = IdleDetector(config.idle, self.start_ts) + self.warned_budget = False + self.stop_reason: str | None = None + self.stop_at: float | None = None + self.forced = False + self.idle_warned = False + self.unpriced_warned: set[str] = set() + self.overdue_released: set[str] = set() + self.last_summary: CostSummary | None = None + self._stop = threading.Event() + self._thread: threading.Thread | None = None + + # -- defaults ------------------------------------------------------------------- + + def _default_graceful_stop(self, reason: str) -> None: + from .. import runs + + runs.request_stop(self.run) + + def _default_force_stop(self, reason: str) -> None: + from .cleanup import cleanup_run + + cleanup_run(self.run, log=self.log) + + def _default_release(self, r: Resource) -> None: + from .cleanup import ProductionCleanupOps, cleanup + + for a in cleanup(self.ledger, ProductionCleanupOps(), keys=(r.key,)): + self.log(f"[cost-guard] {a.key}: {a.action} {a.detail}") + + # -- events -------------------------------------------------------------------- + + def emit(self, event: str, **fields: Any) -> dict[str, Any]: + rec = {"event": event, "run": self.run, "time_unix": self.clock(), **fields} + try: + self.events_path.parent.mkdir(parents=True, exist_ok=True) + with open(self.events_path, "a", encoding="utf-8") as fh: + fh.write(json.dumps(rec, default=str) + "\n") + except OSError as exc: + self.log(f"[cost-guard] cannot write {self.events_path}: {exc}") + if event != "cost_tick": + self.log(f"[cost-guard] {event}: " + ", ".join(f"{k}={v}" for k, v in fields.items())) + return rec + + # -- one tick ------------------------------------------------------------------ + + def _request_stop(self, reason: str, now: float) -> None: + if self.stop_at is not None: + return + self.stop_reason, self.stop_at = reason, now + self.emit("cost_guard_stop", reason=reason, mode="graceful", + force_after_s=self.config.stop_grace_s) + try: + self.graceful_stop(reason) + except Exception as exc: # noqa: BLE001 - the forced stop still follows + self.log(f"[cost-guard] graceful stop failed ({exc}); forcing now") + self.stop_at = now - self.config.stop_grace_s + + def tick(self) -> CostSummary: + now = self.clock() + for path, rec in self.tapes.read(): + self.idle.observe(rec, record_time(rec, now), source=path) + resources = self.ledger.resources() + # A reused run name keeps its ledger: what an earlier launch of that name + # created and released before this guard started is not this launch's spend + # (``yeto cost --run`` still reports the whole ledger). A resource still open + # bills now and counts, whoever created it. + counted = [r for r in resources if r.released_at is None or r.released_at >= self.start_ts] + summary = resource_costs(counted, now, rates=self.rates, billing=self.billing) + self.last_summary = summary + total = summary.total_usd + self.emit("cost_tick", total_usd=round(total, 4), budget_usd=self.config.budget_usd, + open=[i.key for i in summary.items if i.open], + unpriced=[i.key for i in summary.unpriced], billing_error=summary.billing_error) + for i in summary.unpriced: + if i.open and i.key not in self.unpriced_warned: + self.unpriced_warned.add(i.key) + self.emit("cost_unpriced", resource=i.key, note=i.note, + detail="this resource is NOT counted against the budget") + cfg = self.config + if cfg.warn_usd is not None and total >= cfg.warn_usd and not self.warned_budget: + self.warned_budget = True + self.emit("cost_budget_warn", total_usd=round(total, 2), warn_usd=cfg.warn_usd, + budget_usd=cfg.budget_usd) + if cfg.budget_usd is not None and total >= cfg.budget_usd and self.stop_at is None: + self.emit("cost_budget_exceeded", total_usd=round(total, 2), budget_usd=cfg.budget_usd) + self._request_stop(f"budget ${cfg.budget_usd:.2f} reached (${total:.2f})", now) + verdict = self.idle.verdict(now) + if verdict.state in ("warn", "stop") and not self.idle_warned: + self.idle_warned = True + self.emit("idle_warn", reason=verdict.reason, idle_min=round(verdict.idle_s / 60, 1), + shutdown_in_min=round(max(0.0, cfg.idle.idle_after_s + cfg.idle.stop_after_s + - verdict.idle_s) / 60, 1)) + elif verdict.state in ("ok", "loading") and self.idle_warned and self.stop_at is None: + self.idle_warned = False + self.emit("idle_cleared", state=verdict.state) + if verdict.state == "stop" and self.stop_at is None: + self.emit("idle_shutdown", reason=verdict.reason, idle_min=round(verdict.idle_s / 60, 1)) + self._request_stop(f"idle: {verdict.reason}", now) + for r in resources: + if r.kind != "nc_lease" or not r.open or r.key in self.overdue_released: + continue + until = _ts(r.attrs.get("lease_until")) + if until is not None and now >= until + cfg.nc_overdue_grace_s: + self.overdue_released.add(r.key) + release = cfg.release_overdue_nc + self.emit("nc_lease_overdue", resource=r.key, lease_until=r.attrs.get("lease_until"), + overdue_min=round((now - until) / 60, 1), action="release" if release else "alert", + detail=None if release else "kept: the NC pool owns it (pool-stop / " + f"yeto cleanup --run {self.run} --include-nc)") + if not release: + continue + try: + self.release_resource(r) + except Exception as exc: # noqa: BLE001 + self.emit("nc_lease_release_failed", resource=r.key, error=str(exc)) + if self.stop_at is not None and not self.forced and not self.run_ended() \ + and now - self.stop_at >= cfg.stop_grace_s: + self.forced = True + self.emit("cost_guard_stop", reason=self.stop_reason, mode="forced", + detail=f"run still up {cfg.stop_grace_s / 60:.0f} min after the graceful stop") + try: + self.force_stop(self.stop_reason or "forced") + except Exception as exc: # noqa: BLE001 + self.emit("cost_guard_force_failed", error=str(exc)) + return summary + + # -- thread -------------------------------------------------------------------- + + def start(self) -> "CostGuard": + def loop() -> None: + while not self._stop.is_set(): + try: + self.tick() + except Exception as exc: # noqa: BLE001 - the guard never takes the run down + self.log(f"[cost-guard] tick failed: {type(exc).__name__}: {exc}") + self._stop.wait(self.config.poll_s) + + self._thread = threading.Thread(target=loop, name=f"yeto-cost-guard-{self.run}", daemon=True) + self._thread.start() + return self + + def stop(self) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=10) + + +def guard_config_from_args(args) -> GuardConfig: + g = lambda name, default=None: getattr(args, name, default) # noqa: E731 + m = lambda name, default: float(g(name, default) if g(name) is not None else default) * 60 # noqa: E731 + return GuardConfig( + budget_usd=g("budget_usd"), budget_warn_usd=g("budget_warn_usd"), + poll_s=float(g("cost_poll_seconds") or 60), stop_grace_s=m("budget_stop_grace_minutes", 20), + idle=IdleConfig(idle_after_s=m("idle_warn_minutes", 15), stop_after_s=m("idle_stop_minutes", 15), + load_grace_s=m("load_grace_minutes", 45), max_load_s=m("max_load_minutes", 120))) + + +def guard_enabled(args) -> bool: + """On for every CLI launch (the parser sets the idle flags) unless --no-cost-guard; + a programmatic launch without the cost flags runs it only when it sets a budget.""" + if getattr(args, "no_cost_guard", False): + return False + return hasattr(args, "idle_warn_minutes") or getattr(args, "budget_usd", None) is not None diff --git a/yeto/cost/idle.py b/yeto/cost/idle.py new file mode 100644 index 000000000..8843e0eda --- /dev/null +++ b/yeto/cost/idle.py @@ -0,0 +1,154 @@ +"""Idle detection for a run, with a loading-phase exemption. + +A run is idle when, outside a loading phase, either + +* no progress event (a trained round, a sampled round, a cut, a release ...) + arrived for ``idle_after_s``, or +* every GPU source that reports utilization has read 0 % for ``idle_after_s``. + +Heartbeats are not progress: a hung process still heartbeats. + +Loading phase: the run start, and every loading event (driver start-up steps, +weight loading, sampler start, trainer incarnation start ...), opens a window +of ``load_grace_s`` in which neither rule fires - loading K3.1 weights takes +~20 min with 0 % GPU and no progress. A progress event closes the window early. +The window is bounded: loading events alone cannot keep a run alive past +``max_load_s`` since the last progress (a loader stuck in a restart loop). + +Once idle for ``idle_after_s`` the detector says ``warn``; after a further +``stop_after_s`` it says ``stop``. Any progress resets it. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Mapping + +PROGRESS_EVENTS = frozenset({ + "rl_round_trained", "rl_local_round", "rl_sampler_round", "rl_sampler_admitted", + "rl_cut_saved", "rl_round_cut", "rl_publication", "rl_member_publication", + "rl_policy_release_delta", "rl_policy_release_role", "rl_sampler_policy_loaded", + "rl_fragment_push", "rl_policy_apply", "rl_dispense_commit", "rl_learner_finalized", + "rl_trajectory_reward", "sft_step", "step", +}) +LOADING_EVENTS = frozenset({ + "rl_driver_start", "rl_startup_step", "rl_driver_phase", "rl_sampler_start", + "rl_trainer_lifecycle", "rl_trainer_init_from_release", "rl_resume", "rl_data_cursor_restored", + "rl_sampler_release_not_loaded", "island_ready", "rl_verl_island_start", "loading", +}) +GPU_SAMPLE_EVENTS = frozenset({"modal_host_sample", "rl_resource_sample", "gpu_sample"}) +# rl_driver_phase / rl_trainer_lifecycle count as loading only in these phases. +LOADING_PHASES = frozenset({"startup", "loading", "load", "init", "starting", "start", "weights", + "load_weights", "submitted", "provisioning", "ready_wait"}) + + +def _gpu_util(rec: Mapping[str, Any]) -> float | None: + vals = rec.get("gpu_util_pct") + if vals is None: + vals = rec.get("util_pct") + if isinstance(vals, (int, float)): + return float(vals) + if isinstance(vals, list): + nums = [float(v) for v in vals if isinstance(v, (int, float))] + return max(nums) if nums else None + return None + + +def classify(rec: Mapping[str, Any]) -> str | None: + ev = rec.get("event") + if ev in GPU_SAMPLE_EVENTS: + return "gpu" + if ev in PROGRESS_EVENTS: + return "progress" + if ev in ("rl_driver_phase", "rl_trainer_lifecycle"): + phase = str(rec.get("phase") or rec.get("state") or "").lower() + return "loading" if phase in LOADING_PHASES or not phase else None + if ev in LOADING_EVENTS: + return "loading" + return None + + +@dataclass +class IdleConfig: + idle_after_s: float = 15 * 60 + stop_after_s: float = 15 * 60 + load_grace_s: float = 45 * 60 + max_load_s: float = 120 * 60 + + @property + def enabled(self) -> bool: + return self.idle_after_s > 0 + + +@dataclass +class IdleVerdict: + state: str # "ok" | "loading" | "warn" | "stop" + reason: str | None = None + idle_s: float = 0.0 + + +@dataclass +class IdleDetector: + config: IdleConfig + start: float + last_progress: float | None = None + loading_until: float | None = None + gpu: dict[str, tuple[float, float | None]] = field(default_factory=dict) # source -> (util, zero_since) + + def __post_init__(self) -> None: + self.loading_until = self.start + self.config.load_grace_s + + def observe(self, rec: Mapping[str, Any], ts: float, source: str = "") -> str | None: + kind = classify(rec) + if kind == "progress": + self.last_progress = max(self.last_progress or ts, ts) + self.loading_until = None + elif kind == "loading": + self.loading_until = max(self.loading_until or 0.0, ts + self.config.load_grace_s) + elif kind == "gpu": + util = _gpu_util(rec) + if util is not None: + key = str(rec.get("island") or rec.get("learner_id") or source) + ":" + \ + str(rec.get("node_rank") or rec.get("node") or 0) + _prev, zero_since = self.gpu.get(key, (None, None)) + if util > 0: + zero_since = None + elif zero_since is None: + zero_since = ts + self.gpu[key] = (util, zero_since) + return kind + + def _loading_end(self) -> float | None: + """End of the exempt loading window (bounded by ``max_load_s`` after the last progress).""" + if self.loading_until is None: + return None + anchor = self.last_progress if self.last_progress is not None else self.start + return min(self.loading_until, anchor + self.config.max_load_s) + + def _loading(self, now: float) -> bool: + end = self._loading_end() + return end is not None and now < end + + def verdict(self, now: float) -> IdleVerdict: + if not self.config.enabled: + return IdleVerdict("ok") + if self._loading(now): + return IdleVerdict("loading") + # the idle clocks start no earlier than the end of the last (exempt) loading window + end = self._loading_end() + quiet_from = self.start if end is None else min(end, now) + no_progress = now - max(self.last_progress or self.start, quiet_from) + reasons = [] + idle = 0.0 + if no_progress >= self.config.idle_after_s: + reasons.append(f"no progress event for {no_progress / 60:.0f} min") + idle = no_progress + if self.gpu and all(z is not None for _u, z in self.gpu.values()): + zero = now - max(max(z for _u, z in self.gpu.values()), quiet_from) + if zero >= self.config.idle_after_s: + reasons.append(f"GPU utilization 0 % on every reporting GPU for {zero / 60:.0f} min") + idle = max(idle, zero) + if not reasons: + return IdleVerdict("ok") + state = "stop" if idle >= self.config.idle_after_s + self.config.stop_after_s else "warn" + return IdleVerdict(state, "; ".join(reasons), idle) diff --git a/yeto/cost/ledger.py b/yeto/cost/ledger.py new file mode 100644 index 000000000..eaeb470ca --- /dev/null +++ b/yeto/cost/ledger.py @@ -0,0 +1,214 @@ +"""Crash-safe resource ledger of one run: ``/resources.jsonl``. + +Every billed resource a run creates is appended here the moment it exists +(Modal app, Modal FunctionCall, NC lease ticket, sky cluster), and a second +record when it is released. Lines are appended with ``O_APPEND`` and fsynced, +so a launcher that is SIGKILLed leaves a ledger that lists everything it had +created up to its last completed call. :mod:`.cleanup` releases exactly the +entries listed here - it never looks resources up by name pattern. + +Record shape (one JSON object per line):: + + {"op": "intent"|"create"|"release"|"release_failed", "kind": , + "id": , "name": , "ts": , "attrs": {...}} + +``intent`` is written right before a create whose id is only known +afterwards (``modal app deploy``); a crash between the two leaves an intent +that cleanup resolves by exact name AND creation time after the intent, or +reports as unresolved. + +Kinds and the attrs each writer fills in: + +* ``modal_app`` id = app id (``ap-...``), name = app name; attrs: gpu, gpus, + usd_per_hour (optional estimate rate), shared (bool: app is reused by other runs). +* ``modal_call`` id = FunctionCall id (``fc-...``), name = island; attrs: app_id, app_name. +* ``nc_lease`` id = ticket; attrs: account, granted_at, lease_until (ISO or unix), + gpus, rate_ucr_per_gpu_hour, usd_per_gpu_hour (optional). +* ``sky_cluster`` id = cluster name (sky names are exact ids). + +Other processes of the same run (sampler launch scripts, the NC pool) find the +ledger through ``$YETO_RESOURCE_LEDGER`` (see :func:`current`). +""" + +from __future__ import annotations + +import json +import os +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +LEDGER_ENV = "YETO_RESOURCE_LEDGER" +LEDGER_FILE = "resources.jsonl" + +KINDS = ("modal_app", "modal_call", "nc_lease", "sky_cluster") + + +@dataclass +class Resource: + kind: str + id: str | None + name: str | None + created_at: float + attrs: dict[str, Any] = field(default_factory=dict) + released_at: float | None = None + release_error: str | None = None + intent_only: bool = False + + @property + def key(self) -> str: + return _key(self.kind, self.id, self.name) + + @property + def open(self) -> bool: + return self.released_at is None + + def to_json(self) -> dict[str, Any]: + return dict(self.__dict__) + + +def _key(kind: str, rid: str | None, name: str | None) -> str: + return f"{kind}:{rid if rid is not None else 'intent:' + str(name)}" + + +class ResourceLedger: + def __init__(self, path: str | os.PathLike, *, clock=time.time) -> None: + self.path = Path(path) + self.clock = clock + self._lock = threading.Lock() + + @classmethod + def for_run(cls, run: str, **kw) -> "ResourceLedger": + from .. import runs + + return cls(runs.run_dir(run) / LEDGER_FILE, **kw) + + # -- writing --------------------------------------------------------------- + + def _append(self, record: dict[str, Any]) -> None: + line = json.dumps(record, default=str, sort_keys=True) + "\n" + with self._lock: + self.path.parent.mkdir(parents=True, exist_ok=True) + fd = os.open(self.path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600) + try: + os.write(fd, line.encode("utf-8")) + os.fsync(fd) + finally: + os.close(fd) + + def _rec(self, op: str, kind: str, rid: str | None, name: str | None, attrs: dict) -> None: + if kind not in KINDS: + raise ValueError(f"unknown resource kind {kind!r}; one of {KINDS}") + self._append({"op": op, "kind": kind, "id": rid, "name": name, "ts": self.clock(), + "attrs": attrs}) + + def intent(self, kind: str, name: str, **attrs: Any) -> None: + self._rec("intent", kind, None, name, attrs) + + def created(self, kind: str, rid: str, name: str | None = None, **attrs: Any) -> None: + if not rid: + raise ValueError("a created resource needs its provider id") + self._rec("create", kind, str(rid), name, attrs) + + def released(self, kind: str, rid: str | None, name: str | None = None, **attrs: Any) -> None: + self._rec("release", kind, None if rid is None else str(rid), name, attrs) + + def release_failed(self, kind: str, rid: str | None, name: str | None, error: str) -> None: + self._rec("release_failed", kind, rid, name, {"error": str(error)[:500]}) + + # -- reading --------------------------------------------------------------- + + def records(self) -> list[dict[str, Any]]: + try: + text = self.path.read_text(encoding="utf-8") + except FileNotFoundError: + return [] + out = [] + for line in text.splitlines(): + try: + rec = json.loads(line) + except ValueError: + continue # a torn last line after a crash + if isinstance(rec, dict) and rec.get("kind") in KINDS: + out.append(rec) + return out + + def resources(self) -> list[Resource]: + """Folded state, in creation order. A create resolves the newest open + intent of the same kind and name.""" + res: dict[str, Resource] = {} + for r in self.records(): + kind, rid, name, ts = r["kind"], r.get("id"), r.get("name"), float(r.get("ts") or 0) + attrs = dict(r.get("attrs") or {}) + op = r.get("op") + if op == "intent": + k = _key(kind, None, name) + res[k] = Resource(kind, None, name, ts, attrs, intent_only=True) + elif op == "create": + ik = _key(kind, None, name) + prev = res.pop(ik, None) if name is not None else None + k = _key(kind, rid, name) + if k in res: # a re-recorded create: keep the first creation time + res[k].attrs.update(attrs) + res[k].released_at = None + continue + created = prev.created_at if prev is not None else ts + res[k] = Resource(kind, rid, name, created, {**(prev.attrs if prev else {}), **attrs}) + elif op in ("release", "release_failed"): + k = _key(kind, rid, name) + if k not in res: + continue + if op == "release": + res[k].released_at = ts + res[k].release_error = None + else: + res[k].release_error = attrs.get("error") + # a stopped app ends every FunctionCall in it + stopped = [a for a in res.values() if a.kind == "modal_app" and a.released_at is not None] + for c in res.values(): + if c.kind != "modal_call" or c.released_at is not None: + continue + for a in stopped: + same = (c.attrs.get("app_id") and c.attrs.get("app_id") == a.id) or \ + (not c.attrs.get("app_id") and c.attrs.get("app_name") == a.name) + if same and a.released_at >= c.created_at: + c.released_at = a.released_at + break + return list(res.values()) + + def open_resources(self) -> list[Resource]: + return [r for r in self.resources() if r.open] + + +class NullLedger: + """Records nothing (a programmatic launch that opted out of the ledger).""" + + path = None + + def __getattr__(self, name): + if name in ("intent", "created", "released", "release_failed"): + return lambda *a, **k: None + if name in ("records", "resources", "open_resources"): + return lambda: [] + raise AttributeError(name) + + +def current(environ=None) -> ResourceLedger | None: + """The ledger of the run this process belongs to (``$YETO_RESOURCE_LEDGER``), or None.""" + env = os.environ if environ is None else environ + path = env.get(LEDGER_ENV) + return ResourceLedger(path) if path else None + + +def record_safely(fn_name: str, *args: Any, ledger: ResourceLedger | None = None, **kw: Any) -> None: + """Call ``ledger.(...)`` on the given or current ledger; a ledger + write error is printed loudly but never breaks the create it records.""" + led = ledger if ledger is not None else current() + if led is None: + return + try: + getattr(led, fn_name)(*args, **kw) + except Exception as exc: # noqa: BLE001 + print(f"[cost] WARNING: resource ledger {led.path} write failed: {exc}", flush=True) diff --git a/yeto/cost/reaper.py b/yeto/cost/reaper.py new file mode 100644 index 000000000..91a3fb4e2 --- /dev/null +++ b/yeto/cost/reaper.py @@ -0,0 +1,74 @@ +"""The resource reaper: a detached process per run that outlives the launcher. + +``yeto launch`` starts it in its own session next to the worker. It polls the +worker pid; once the worker is gone - a normal exit, an exception, SIGTERM, or +SIGKILL that skipped every ``finally`` - it releases whatever the run's +resource ledger still lists as open (:func:`.cleanup.cleanup_run`, exact ids +only). Runs launched with ``--keep`` and head-controller runs (the controller +lives on the head VM) are left alone. One reaper per run: a second exits. +""" + +from __future__ import annotations + +import fcntl +import time +from collections.abc import Callable + +POLL_S = 15.0 + + +def reap(run: str, worker_pid: int, *, alive: Callable[[int], bool] | None = None, + cleanup: Callable[[str], bool] | None = None, sleep: Callable[[float], None] = time.sleep, + poll_s: float = POLL_S, log=print) -> int: + from .. import runs + from .cleanup import cleanup_run + from .ledger import ResourceLedger + + alive = alive or runs.is_alive + cleanup = cleanup or cleanup_run + d = runs.run_dir(run) + d.mkdir(parents=True, exist_ok=True) + lock = open(d / ".reaper.lock", "w") + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError: + return 0 # another reaper watches this run + try: + while alive(worker_pid): + sleep(poll_s) + meta = runs.load_run(run) or {} + args = meta.get("args") or {} + if args.get("keep"): + log(f"[reaper] {run}: worker {worker_pid} gone; --keep: resources left up") + return 0 + if meta.get("controller") == "head": + return 0 + if meta.get("pid") not in (None, worker_pid) and alive(int(meta["pid"])): + return 0 # the run was relaunched under the same name by a newer worker + led = ResourceLedger.for_run(run) + open_res = [r for r in led.open_resources() + if not (r.kind == "modal_app" and r.attrs.get("shared"))] + if not open_res: + return 0 + log(f"[reaper] {run}: worker {worker_pid} gone with {len(open_res)} open resource(s) in " + f"{led.path}; releasing them") + ok = cleanup(run) + if not ok: + log(f"[reaper] {run}: not everything confirmed released; run `yeto cleanup --run {run}`") + try: + runs.update_run(run, state=runs.TEARDOWN_INCOMPLETE) + except OSError: + pass + return 0 if ok else 1 + finally: + try: + fcntl.flock(lock, fcntl.LOCK_UN) + lock.close() + except OSError: + pass + + +if __name__ == "__main__": # pragma: no cover + import sys + + sys.exit(reap(sys.argv[1], int(sys.argv[2]))) diff --git a/yeto/cost/report.py b/yeto/cost/report.py new file mode 100644 index 000000000..e22ef3d25 --- /dev/null +++ b/yeto/cost/report.py @@ -0,0 +1,181 @@ +"""``yeto cost --run ``: cost per resource and per hour, and GPU effective utilization. + +Effective utilization = productive time / billed time, where + +* billed time is the sum over the run's billed compute resources (Modal apps, + NC leases, sky clusters) of their billable span; and +* productive time is, per event source (one tape = one island / sampler), + the union of the gaps between consecutive progress events (trained round, + sampled round, cut, release ...) that are at most ``gap_s`` apart; a longer + gap is waiting / loading / idling, not work. Summed over sources. + +It is an estimate from event tapes, not a hardware counter; the hardware +utilization samples (``gpu_util_pct``), when present, are reported next to it. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +from .billing import CostSummary, ModalBilling, Rates, resource_costs +from .guard import TapeTail, record_time +from .idle import _gpu_util, classify +from .ledger import ResourceLedger + +DEFAULT_GAP_S = 10 * 60 +COMPUTE_KINDS = ("modal_app", "nc_lease", "sky_cluster") + + +def productive_seconds(records: Iterable[tuple[str, dict[str, Any]]], gap_s: float = DEFAULT_GAP_S + ) -> tuple[dict[str, float], list[float]]: + by_src: dict[str, list[float]] = {} + utils: list[float] = [] + for src, rec in records: + kind = classify(rec) + if kind == "progress": + t = record_time(rec, 0.0) + if t > 0: + by_src.setdefault(src, []).append(t) + elif kind == "gpu": + u = _gpu_util(rec) + if u is not None: + utils.append(u) + out = {} + for src, ts in by_src.items(): + ts.sort() + out[src] = sum(b - a for a, b in zip(ts, ts[1:]) if b - a <= gap_s) + return out, utils + + +def build_report(run: str, *, summary: CostSummary, tape_globs: Iterable[str], + gap_s: float = DEFAULT_GAP_S) -> dict[str, Any]: + recs = TapeTail(tape_globs).read() + prod, utils = productive_seconds(recs, gap_s) + billed = sum(i.billable_s for i in summary.items if i.kind in COMPUTE_KINDS) + productive = sum(prod.values()) + hours: dict[str, float] = {} + for i in summary.items: + for h, usd in i.hours.items(): + hours[h] = hours.get(h, 0.0) + usd + return { + "run": run, "total_usd": round(summary.total_usd, 2), + "unpriced": [i.key for i in summary.unpriced], "billing_error": summary.billing_error, + "resources": [{"resource": i.key, "name": i.name, "open": i.open, + "billable_h": round(i.billable_s / 3600, 3), + "billed_usd": round(i.billed_usd, 2), "estimated_usd": round(i.estimated_usd, 2), + "usd": None if i.usd is None else round(i.usd, 2), + "rate_usd_per_hour": i.rate_usd_per_hour, "note": i.note} + for i in summary.items], + "by_hour": {h: round(v, 2) for h, v in sorted(hours.items())}, + "billed_compute_h": round(billed / 3600, 3), + "productive_h": round(productive / 3600, 3), + "productive_by_source_h": {Path(k).name: round(v / 3600, 3) for k, v in sorted(prod.items())}, + "effective_utilization": None if billed <= 0 else round(productive / billed, 4), + "gpu_util_pct_mean": None if not utils else round(sum(utils) / len(utils), 1), + "gpu_util_samples": len(utils), + } + + +def default_tape_globs(run: str) -> list[str]: + from .. import runs + + d = runs.run_dir(run) + return [str(d / "events" / "*.jsonl"), str(d / "events" / "**" / "*.jsonl"), + str(d / "tapes" / "*.jsonl")] + + +def format_report(rep: dict[str, Any]) -> str: + lines = [f"run {rep['run']}: total ${rep['total_usd']:.2f}" + + (f" (UNPRICED, not counted: {', '.join(rep['unpriced'])})" if rep["unpriced"] else "")] + if rep.get("billing_error"): + lines.append(f" Modal billing API error: {rep['billing_error']} (estimates only)") + lines.append(f" {'resource':<48} {'open':<5} {'hours':>8} {'billed $':>10} {'est. $':>9} {'total $':>9}") + for r in rep["resources"]: + usd = "unpriced" if r["usd"] is None else f"{r['usd']:.2f}" + lines.append(f" {r['resource'][:48]:<48} {'yes' if r['open'] else 'no':<5} {r['billable_h']:>8.2f} " + f"{r['billed_usd']:>10.2f} {r['estimated_usd']:>9.2f} {usd:>9}" + + (f" [{r['note']}]" if r["note"] else "")) + lines.append(" by hour (UTC): " + ", ".join(f"{h} ${v:.2f}" for h, v in rep["by_hour"].items())) + eff = rep["effective_utilization"] + lines.append(f" GPU effective utilization: " + f"{'n/a' if eff is None else f'{100 * eff:.1f} %'} " + f"(productive {rep['productive_h']:.2f} h / billed {rep['billed_compute_h']:.2f} h)" + + ("" if rep["gpu_util_pct_mean"] is None else + f"; hardware util mean {rep['gpu_util_pct_mean']} % over {rep['gpu_util_samples']} samples")) + return "\n".join(lines) + + +def cmd_cost(args) -> int: + import time + + led = ResourceLedger.for_run(args.run) + if not led.path.exists(): + print(f"[cost] run {args.run}: no resource ledger at {led.path}") + return 1 + billing = None if args.no_billing else ModalBilling() + summary = resource_costs(led.resources(), time.time(), rates=Rates.from_env(args.prices), + billing=billing) + globs = default_tape_globs(args.run) + list(args.tapes or []) + rep = build_report(args.run, summary=summary, tape_globs=globs, gap_s=args.gap_minutes * 60) + print(json.dumps(rep, indent=2, default=str) if args.json else format_report(rep)) + return 0 + + +def cmd_cleanup(args) -> int: + from .cleanup import cleanup_run + + return 0 if cleanup_run(args.run, dry_run=args.dry_run, include_nc=getattr(args, "include_nc", False)) else 1 + + +def add_parsers(sub) -> None: + cost = sub.add_parser("cost", help="cost of a run per resource and per hour, and GPU effective " + "utilization (Modal billing API + estimates)") + cost.add_argument("--run", required=True, help="run name (its --cluster-prefix)") + cost.add_argument("--json", action="store_true") + cost.add_argument("--no-billing", action="store_true", + help="do not call the Modal billing API (estimates only)") + cost.add_argument("--prices", default=None, help="price table JSON (yeto.dashboard.cost format)") + cost.add_argument("--tapes", nargs="*", help="extra event tape globs (e.g. pulled sampler tapes)") + cost.add_argument("--gap-minutes", type=float, default=DEFAULT_GAP_S / 60, + help="progress events further apart than this are not productive time") + clean = sub.add_parser("cleanup", help="release exactly the resources in a run's resource ledger " + "(by id; never by name pattern, never --all)") + clean.add_argument("--run", required=True, help="run name (its --cluster-prefix)") + clean.add_argument("--dry-run", action="store_true", help="list what would be released") + clean.add_argument("--include-nc", action="store_true", + help="also release the NC pool's leases (normally pool-stop's job)") + + +def add_launch_arguments(p) -> None: + g = p.add_argument_group("cost guard (budget breaker, idle shutdown)") + g.add_argument("--budget-usd", type=float, default=None, + help="stop the run gracefully (cut, then teardown) when its cost reaches this " + "(default: no budget stop)") + g.add_argument("--budget-warn-usd", type=float, default=None, + help="emit cost_budget_warn at this cost (default: 80%% of --budget-usd)") + g.add_argument("--budget-stop-grace-minutes", type=float, default=20.0, + help="after a graceful stop, release every ledger resource by force when the " + "run is still up this much later (default 20)") + g.add_argument("--idle-warn-minutes", type=float, default=15.0, + help="idle_warn after this long without progress / with 0%% GPU (0 disables idle " + "shutdown; default 15)") + g.add_argument("--idle-stop-minutes", type=float, default=15.0, + help="stop the run this long after idle_warn if still idle (default 15)") + g.add_argument("--load-grace-minutes", type=float, default=45.0, + help="no idle verdict for this long after the start / a loading event (weight " + "loading is silent; default 45)") + g.add_argument("--max-load-minutes", type=float, default=120.0, + help="loading events alone cannot exempt a run longer than this since its last " + "progress (default 120)") + g.add_argument("--cost-poll-seconds", type=float, default=60.0, + help="cost guard tick (default 60)") + g.add_argument("--cost-prices", default=None, + help="price table JSON for estimates of not-yet-billed time") + g.add_argument("--cost-tapes", nargs="*", default=None, + help="extra event tape globs the idle detector reads (e.g. pulled NC sampler tapes)") + g.add_argument("--no-cost-guard", action="store_true", + help="run without the cost guard (the resource ledger is still written)") + diff --git a/yeto/dashboard/alerts.py b/yeto/dashboard/alerts.py index 3eb983ee8..ad2da77bd 100644 --- a/yeto/dashboard/alerts.py +++ b/yeto/dashboard/alerts.py @@ -19,6 +19,11 @@ # FN 2x8 measured 10-15 min). Tasks 8.1. "startup_warn_s": 1200.0, "startup_severe_s": 1800.0, + # the actor init stage of a trainer start (model build + weight load: trainer_start_phase + # actor_import_done -> actor_init_done, or rl_startup_step before weights_loaded) writes + # no event for its whole length; K3.1 mbridge load measured ~22 min + "startup_init_warn_s": 2700.0, + "startup_init_severe_s": 3600.0, # 8.4: one startup sub-step (Ray up / engines ready / weights loaded) taking # longer than this while heartbeats still arrive (process alive, step stuck). "startup_step_warn_s": 900.0, @@ -72,20 +77,26 @@ def _next_startup_step(done: str | None) -> str: return STARTUP_STEPS[i + 1] if i + 1 < len(STARTUP_STEPS) else "rl_driver_start" +ENDED = ("done", "stopped", "left", "lost", "unknown") +QUIET_OK = ("done", "stopped", "left", "unknown") # no heartbeat alert (a lost island keeps its alert) + + def heartbeat(cards: list[dict], th: dict) -> list[dict]: out = [] for c in cards: age = c["last_event_age_s"] - if c["status"] in ("done", "stopped") or c.get("finalized") or age is None: + if c["status"] in QUIET_OK or c.get("finalized") or age is None: continue if c.get("starting"): + stage = c.get("startup_stage") or {} + warn, severe = stage.get("warn_s", th["startup_warn_s"]), stage.get("severe_s", th["startup_severe_s"]) + step_warn = max(th["startup_step_warn_s"], warn) if stage.get("init") else th["startup_step_warn_s"] step_age = c.get("startup_step_age_s") - if step_age is not None and step_age > th["startup_step_warn_s"] and age <= th["startup_warn_s"]: + if step_age is not None and step_age > step_warn and age <= warn: pending = _next_startup_step(c.get("startup_step")) out.append(_alert(1, "startup_step", f"island {c['id']} startup step slow", f"waiting for {pending} for {step_age:.0f}s (previous step {c.get('startup_step') or 'start'}; " - f"threshold {th['startup_step_warn_s']:.0f}s)", island=c["id"])) - warn, severe = th["startup_warn_s"], th["startup_severe_s"] + f"threshold {step_warn:.0f}s)", island=c["id"])) if age > severe: sev = 0 elif age > warn: @@ -93,19 +104,27 @@ def heartbeat(cards: list[dict], th: dict) -> list[dict]: else: continue out.append(_alert(sev, "startup", f"island {c['id']} no events during startup", - f"still starting (no rl_driver_start yet), last event {age:.0f}s ago" - f" (startup threshold {warn:.0f}s/{severe:.0f}s)", island=c["id"], metric="reward")) + f"still starting (no rl_driver_start yet; stage: {stage.get('label') or 'n/a'}" + + (f" for {stage['age_s']:.0f}s" if stage.get("age_s") is not None else "") + + f"), last event {age:.0f}s ago (startup threshold {warn:.0f}s/{severe:.0f}s)", + island=c["id"], metric="reward")) continue - if age > th["heartbeat_severe_s"]: + if c.get("waiting_backlog"): + continue # a sampler idle at its backlog cap, by design (its trainers are behind) + # without heartbeats an island is quiet for as long as its own cadence says + # (rounds / generations / policy loads; the card's stale_after_s) + warn = max(th["heartbeat_warn_s"], c.get("stale_after_s") or 0.0) + severe = th["heartbeat_severe_s"] if warn <= th["heartbeat_warn_s"] else max(th["heartbeat_severe_s"], 2 * warn) + if age > severe: sev = 0 - elif age > th["heartbeat_warn_s"]: + elif age > warn: sev = 1 else: continue what = "heartbeat" if c["heartbeat_seen"] else "event (no heartbeat events, using the last event)" out.append(_alert(sev, "heartbeat", f"island {c['id']} heartbeat timeout", - f"last {what} {age:.0f}s ago (threshold {th['heartbeat_warn_s']:.0f}s/" - f"{th['heartbeat_severe_s']:.0f}s)", island=c["id"], round=c["round"], + f"last {what} {age:.0f}s ago (threshold {warn:.0f}s/" + f"{severe:.0f}s)", island=c["id"], round=c["round"], metric="reward")) return out @@ -224,17 +243,20 @@ def recovery(reducer) -> list[dict]: return out -def source_lost(cards: list[dict]) -> list[dict]: +def source_lost(cards: list[dict], torn_down: bool = False) -> list[dict]: return [_alert(1, "source_lost", f"island {c['id']} tape source lost", f"ssh follow disconnected: {(c['source_lost'].get('error') or '')[:120]}", island=c["id"]) - for c in cards if c["source_lost"] and not c["source_lost"].get("recovered")] + for c in cards if c["source_lost"] and not c["source_lost"].get("recovered") + and not torn_down and c["status"] not in ("done", "stopped", "left")] + # an island that ended (or a run torn down): its tape follow ending is expected def evaluate(reducer, *, now: float, rounds: list[dict], cost: dict, cards: list[dict]) -> list[dict]: th = reducer.thresholds alerts = (recovery(reducer) + heartbeat(cards, th) + consecutive_missed(rounds, th) + quorum(rounds, th) + grad_norm(reducer, th) + clip_kl(reducer, th) - + budget(cost, th) + source_lost(cards)) + + budget(cost, th) + + source_lost(cards, torn_down=(reducer.operator_stop() or {}).get("cause") == "teardown")) alerts.sort(key=lambda a: a["sev"]) for i, a in enumerate(alerts): a["id"] = i diff --git a/yeto/dashboard/async_rl.py b/yeto/dashboard/async_rl.py new file mode 100644 index 000000000..4f1a2d874 --- /dev/null +++ b/yeto/dashboard/async_rl.py @@ -0,0 +1,454 @@ +"""Async GRPO (sample buffer) view: samplers, buffer backlog, releases, NC pool. + +The buffer trainers and the samplers of an async run write their own event +types (``rl_sampler_*``, ``rl_group_policies``, ``rl_buffer_stale_skipped``, +``rl_policy_released`` ...); this module folds them into the tables and +series of the page's "Async GRPO" section. Pure observation of the tapes: +every number is read from a record, missing stays None. + +Island roles are inferred from the events an island writes: ``sampler`` (any +``rl_sampler_*`` event), ``trainer`` (``rl_round_trained`` / ``rl_driver_start``). +Health of each role is judged by its own cadence (see ``expected_gap_s``). +""" + +from __future__ import annotations + +import math +from typing import Any + +POOL_PREFIX = "nc_elastic_" +POOL_MAX = 2000 +ROUNDS_MAX = 5000 +CADENCE_WINDOW = 5 # recent rounds whose duration sets the expected gap between events +CADENCE_FACTOR = 2.0 +CADENCE_SLACK_S = 60.0 +# no event for LOST_FACTOR x the expected gap: the island is counted lost (not active) +LOST_FACTOR = 3.0 +SAMPLER_FLOOR_S = 300.0 # the shortest quiet time a sampler is allowed (3 x p50 round below it) +# events after which a sampler writes nothing by design until something else happens +SAMPLER_WAIT_EVENTS = {"rl_sampler_backlog_full": "backlog full (trainers behind)", + "rl_sampler_waiting_policy": "waiting for a policy release"} +SAMPLER_END_EVENTS = frozenset({"rl_sampler_stopped", "rl_sampler_data_end"}) +TRAINER_EVENTS = frozenset({"rl_round_trained", "rl_driver_start", "rl_group_policies", + "rl_buffer_stale_skipped", "rl_policy_released"}) +# sampler-reported reward/length summary keys (rl_sampler_engine_round) +# the sample-buffer reader's short-group handling per round (rl_buffer_stale_skipped) +FILL_KEYS = ("incomplete_groups", "dropped_small_groups", "carried_over_groups", "padded_samples") +SAMPLER_STAT_KEYS = ("reward_mean", "reward_std", "reward_p10", "reward_p50", "reward_p90", + "resp_len_mean", "resp_len_p95", "truncated", "zero_variance_seen") +# pool events that end a shard's sampler (left: drained / released / fenced; lost: the node went away) +POOL_LEFT = {"nc_elastic_released": "left", "nc_elastic_released_externally": "left", + "nc_elastic_shard_fenced": "left", "nc_elastic_sampler_finished": "done", + "nc_elastic_sampler_crash_loop": "left", "nc_elastic_drain_failed": "left", + "nc_elastic_evicted": "lost", "nc_elastic_expired": "lost", "nc_elastic_node_lost": "lost"} +POOL_UP = {"nc_elastic_member_ready", "nc_elastic_node_reachable_again", "nc_elastic_control_plane_recovered"} +POOL_STOPPED = "nc_elastic_pool_stopped" +POOL_KEYS = ("ticket", "account", "shard", "why", "reason", "detail", "error", "state", "after_s", + "waited_s", "restarts", "attempt", "data_slot") + + +def _f(v: Any) -> float | None: + if isinstance(v, bool) or not isinstance(v, (int, float)): + return None + v = float(v) + return v if math.isfinite(v) else None + + +def _int(v: Any) -> int | None: + return v if isinstance(v, int) and not isinstance(v, bool) else None + + +def quantile(vals: list[float], q: float) -> float | None: + xs = sorted(v for v in vals if v is not None) + if not xs: + return None + k = (len(xs) - 1) * q + lo, hi = math.floor(k), math.ceil(k) + return xs[lo] + (xs[hi] - xs[lo]) * (k - lo) + + +def pool_ts(r: dict) -> float | None: + """``nc_elastic_*`` records carry the wall clock in ``t`` (EventLog).""" + return _f(r.get("t")) + + +class AsyncState: + def __init__(self) -> None: + self.roles: dict[str, str] = {} + self.sampler: dict[str, dict] = {} + self.trainer: dict[str, dict] = {} + self.releases: list[dict] = [] + self.loads: list[dict] = [] + self.pool: list[dict] = [] + + # -- folding ---------------------------------------------------------------- + def _s(self, iid: str) -> dict: + s = self.sampler.get(iid) + if s is None: + s = self.sampler[iid] = {"rounds": {}, "start": None, "last_wait": None, "ended": False, + "backlog": [], "max_backlog_groups": None, "resource": None, + "first_round_ts": None, "last_type": None, "last_ts": None, + # one entry per sampler incarnation (an in-place restart is a new one) + "incarnations": []} + return s + + def _t(self, iid: str) -> dict: + t = self.trainer.get(iid) + if t is None: + t = self.trainer[iid] = {"rounds": {}, "left": None, "walls": []} + return t + + def fold_pool(self, r: dict) -> None: + if len(self.pool) < POOL_MAX: + self.pool.append({"ts": pool_ts(r), "event": r.get("event"), + **{k: r.get(k) for k in POOL_KEYS if r.get(k) is not None}}) + + def fold(self, r: dict, iid: str, ts: float | None) -> None: + ev = r.get("event") + if not isinstance(ev, str): + return + if ev.startswith("rl_sampler_"): + self.roles[iid] = "sampler" + self._fold_sampler(r, iid, ts, ev) + elif ev in TRAINER_EVENTS or ev in ("rl_rollout_carry_over", "offpolicy_gate", "rl_round_phases", + "rl_round_wall", "island_leave"): + if ev in TRAINER_EVENTS and self.roles.get(iid) != "sampler": + self.roles[iid] = "trainer" + self._fold_trainer(r, iid, ts, ev) + elif ev == "rl_resource_sample" and iid in self.sampler: + self.sampler[iid]["resource"] = {"available": r.get("available"), "reason": r.get("reason"), + "source": r.get("source")} + + def _fold_sampler(self, r: dict, iid: str, ts: float | None, ev: str) -> None: + s = self._s(iid) + s["last_type"], s["last_ts"] = ev, ts + rid = _int(r.get("rollout_id")) + inc = r.get("incarnation") + if inc is not None and ev in ("rl_sampler_join", "rl_sampler_start"): + inc = str(inc) + row = next((x for x in s["incarnations"] if x["incarnation"] == inc), None) + if row is None: + s["incarnations"].append({"incarnation": inc, "first_ts": ts, "last_ts": ts}) + elif ts is not None: + row["first_ts"] = ts if row["first_ts"] is None else min(row["first_ts"], ts) + s["incarnation"] = inc + if s["incarnations"] and ts is not None: + cur = next((x for x in s["incarnations"] if x["incarnation"] == s.get("incarnation")), None) + if cur is not None: + cur["last_ts"] = ts if cur["last_ts"] is None else max(cur["last_ts"], ts) + if ev == "rl_sampler_join": + s["shard"] = r.get("sampler") + if ev == "rl_sampler_start": + s["start"] = ts + s["max_backlog_groups"] = r.get("max_backlog_groups") + s["sampling"] = {k: (r.get("sampling") or {}).get(k) for k in ("n", "max_tokens", "temperature")} + s["provenance"] = {k: (r.get("provenance") or {}).get(k) + for k in ("card_model", "engine", "precision", "topology")} + if ev == "rl_sampler_round" and s["first_round_ts"] is None: + s["first_round_ts"] = ts if ts is not None else 0.0 + key = (s.get("incarnation"), rid) # a restarted sampler may count its rounds from 0 again + if ev == "rl_sampler_round" and rid is not None: + row = s["rounds"].setdefault(key, {"rollout_id": rid, "incarnation": s.get("incarnation")}) + row.update({"ts": ts, "policy_version": r.get("policy_version"), "version_lag": r.get("version_lag"), + "groups": r.get("groups"), "buffer_end": r.get("buffer_end"), + "generate_s": _f(r.get("generate_s"))}) + if s["first_round_ts"] is None: + s["first_round_ts"] = ts + elif ev == "rl_sampler_engine_round" and rid is not None: + row = s["rounds"].setdefault(key, {"rollout_id": rid, "incarnation": s.get("incarnation")}) + finish = r.get("finish") if isinstance(r.get("finish"), dict) else {} + if row.get("ts") is None: + row["ts"] = ts # rl_sampler_round (when it follows) sets the round's own time + samples = _f(r.get("samples")) + trunc = _f(finish.get("length")) + row.update({"samples": samples, "response_tokens": _f(r.get("response_tokens")), + "dropped_groups": r.get("dropped_groups"), "dropped_reason": r.get("dropped_reason"), + "zero_variance_dropped": r.get("zero_variance_groups"), + "truncated": trunc if trunc is not None else (0.0 if finish else None), + "truncated_frac": (trunc or 0.0) / samples if samples and finish else None}) + for k in SAMPLER_STAT_KEYS: + if r.get(k) is not None: + row[k] = r.get(k) + if samples and _f(r.get("response_tokens")) is not None and row.get("resp_len_mean") is None: + row["resp_len_mean"] = _f(r.get("response_tokens")) / samples + agentic = r.get("agentic") if isinstance(r.get("agentic"), dict) else None + for k in ("retried_trials", "lost_trials"): + if agentic and _f(agentic.get(k)) is not None: + row[k] = _f(agentic.get(k)) + if agentic and row.get("reward_mean") is None and _f(agentic.get("reward_mean")) is not None: + row["reward_mean"] = _f(agentic.get("reward_mean")) + elif ev == "rl_sampler_policy_loaded": + ft = r.get("follower_timings") if isinstance(r.get("follower_timings"), dict) else {} + seen = _f(r.get("seen_at")) + self.loads.append({"sampler": iid, "policy_version": r.get("policy_version"), "seq": r.get("seq"), + "released_at": _f(r.get("released_at")), "seen_at": seen, "loaded_at": ts, + "resolve_s": _f(ft.get("resolve_s")), "load_s": _f(ft.get("load_s")), + "pause_s": _f(r.get("pause_s")), "buffer_selected": r.get("buffer_selected"), + "load_wall_s": round(ts - seen, 3) if ts is not None and seen is not None else None}) + elif ev in ("rl_sampler_backlog_full", "rl_sampler_admitted"): + b = _f(r.get("backlog_groups")) + if ts is not None and b is not None: + s["backlog"].append([ts, b, ev == "rl_sampler_backlog_full"]) + if r.get("max_backlog_groups") is not None: + s["max_backlog_groups"] = r.get("max_backlog_groups") + if ev in SAMPLER_END_EVENTS: + s["ended"] = True + s["last_wait"] = SAMPLER_WAIT_EVENTS.get(ev) + + def _fold_trainer(self, r: dict, iid: str, ts: float | None, ev: str) -> None: + t = self._t(iid) + rid = _int(r.get("rollout_id")) + if ev == "island_leave": + t["left"] = {"ts": ts, "reason": r.get("reason_name") or r.get("reason"), + "error_type": r.get("error_type")} + return + if ev == "rl_policy_released": + origin = r.get("origin") if isinstance(r.get("origin"), dict) else {} + timings = r.get("release_timings") if isinstance(r.get("release_timings"), dict) else {} + self.releases.append({"trainer": iid, "policy_version": r.get("policy_version"), + "local_version": r.get("local_version"), "seq": r.get("seq"), + "released_at": _f(r.get("released_at")) or ts, "ts": ts, + "bytes": r.get("release_bytes"), "upload_s": _f(timings.get("upload_s")), + "origin": origin.get("kind"), "buffer_selected": r.get("buffer_selected")}) + return + if ev in ("rl_round_phases", "rl_round_wall"): + w = _f(r.get("wall_seconds")) if ev == "rl_round_phases" else _f(r.get("round_seconds")) + if w is not None: + t["walls"] = (t["walls"] + [w])[-CADENCE_WINDOW:] + return + if rid is None or len(t["rounds"]) >= ROUNDS_MAX and rid not in t["rounds"]: + return + row = t["rounds"].setdefault(rid, {"rollout_id": rid}) + if ev == "rl_round_trained": + row.update({"ts": ts, "trained_groups": r.get("trained_groups"), + "trained_samples": r.get("trained_samples"), + "policy_age_mean": r.get("policy_age_mean"), "policy_age_max": r.get("policy_age_max")}) + elif ev == "rl_buffer_stale_skipped": + row["stale_groups"] = r.get("stale_groups") + row["max_policy_age"] = r.get("max_policy_age") + for k in FILL_KEYS: + if r.get(k) is not None: + row[k] = r.get(k) + elif ev == "rl_group_policies": + ages = r.get("ages") if isinstance(r.get("ages"), dict) else {} + row["groups_read"] = len(r["groups"]) if isinstance(r.get("groups"), list) else None + row["ages"] = {str(k): v for k, v in sorted(ages.items(), key=lambda kv: _age_key(kv[0]))} + row["merged_version"] = r.get("merged_version") + elif ev == "offpolicy_gate": + c = r.get("counts") if isinstance(r.get("counts"), dict) else {} + row["gate_dropped"] = c.get("group_dropped") + row["unknown_version_dropped"] = c.get("unknown_version_dropped") + elif ev == "rl_rollout_carry_over": + row["over_age_discarded_groups"] = r.get("over_age_discarded_groups") + row["trained_response_tokens"] = r.get("trained_response_tokens") + if ev == "rl_buffer_stale_skipped" and row.get("ts") is None: + row["read_ts"] = ts # the read happens before the step; rl_round_trained's ts wins + + # -- derived per-island facts used by the reducer's island cards ---------- + def role(self, iid: str) -> str | None: + return self.roles.get(iid) + + def expected_gap_s(self, iid: str) -> float | None: + """Longest quiet time this island's own cadence explains: a sampler writes + one round event per generation and a load event per release (the load can + take minutes); a trainer one round per step. ``CADENCE_FACTOR`` times the + longest recent duration plus slack; None before the first round.""" + if self.roles.get(iid) == "sampler": + s = self.sampler.get(iid) or {} + rows = sorted((v for v in s.get("rounds", {}).values() if v.get("generate_s") is not None), + key=_row_order)[-CADENCE_WINDOW:] + if not rows: + return None + loads = [x["load_wall_s"] for x in self.loads if x["sampler"] == iid and x["load_wall_s"] is not None] + gens = sorted(v["generate_s"] for v in rows) + return max(SAMPLER_FLOOR_S, 3.0 * gens[len(gens) // 2], CADENCE_FACTOR * gens[-1]) \ + + (max(loads) if loads else 0.0) + CADENCE_SLACK_S + t = self.trainer.get(iid) or {} + if t.get("walls"): + return CADENCE_FACTOR * max(t["walls"]) + CADENCE_SLACK_S + return None + + def sampler_state(self, iid: str) -> dict: + s = self.sampler.get(iid) or {} + return {"started_rounds": s.get("first_round_ts") is not None, + "waiting": s.get("last_wait"), "ended": bool(s.get("ended"))} + + def pool_exit(self, iid: str, after: float | None = None) -> dict | None: + """The NC pool's word on sampler ``iid`` (its shard): the last exit event of + the shard, or of a ticket that served it, with no member_ready after it; a + pool_stopped ends every shard. None when the pool has not ended it (or the + exit is older than ``after``, the sampler's last event: a restart since).""" + tickets = {p.get("ticket") for p in self.pool if p.get("shard") is not None + and str(p.get("shard")) == str(iid) and p.get("ticket")} + state = None + for p in sorted(self.pool, key=lambda x: x.get("ts") or 0): + ev = p.get("event") + mine = (p.get("shard") is not None and str(p.get("shard")) == str(iid)) or ( + p.get("shard") is None and p.get("ticket") in tickets) + if ev == POOL_STOPPED: + state = {"status": "left", "event": ev, "ts": p.get("ts"), "why": p.get("why")} + elif mine and ev in POOL_LEFT: + state = {"status": POOL_LEFT[ev], "event": ev, "ts": p.get("ts"), + "why": p.get("why") or p.get("detail"), "ticket": p.get("ticket")} + elif mine and ev in POOL_UP: + state = None + if state and after is not None and state["ts"] is not None and state["ts"] < after - 5: + return None + return state + + def trainer_left(self, iid: str) -> dict | None: + return (self.trainer.get(iid) or {}).get("left") + + # -- view --------------------------------------------------------------------- + def active(self) -> bool: + return bool(self.sampler or self.releases or self.pool or any( + t["rounds"] and any("ages" in r or "stale_groups" in r for r in t["rounds"].values()) + for t in self.trainer.values())) + + def view(self) -> dict | None: + if not self.active(): + return None + samplers = {} + for iid, s in self.sampler.items(): + rows = sorted(s["rounds"].values(), key=_row_order) + # empty rounds (no group written: backlog / policy waits) would push every + # reward-bearing round out of the 500-row window of the charts + full = [r for r in rows if not _empty_round(r)] + lags = [_f(r.get("version_lag")) for r in rows if _f(r.get("version_lag")) is not None] + hist: dict[str, int] = {} + for lag in lags: + hist[str(int(lag))] = hist.get(str(int(lag)), 0) + 1 + gens = [r["generate_s"] for r in rows if r.get("generate_s") is not None] + + def tot(key: str, _rows: list = rows) -> float | None: + vals = [_f(r.get(key)) for r in _rows if _f(r.get(key)) is not None] + return sum(vals) if vals else None + + last = rows[-1] if rows else {} + rew = next((r for r in reversed(rows) if r.get("reward_mean") is not None), None) + incs = sorted(s["incarnations"], key=lambda x: (x["first_ts"] is None, x["first_ts"] or 0)) + samplers[iid] = { + "rounds": full[-500:], "n_rounds": len(rows), "empty_rounds": len(rows) - len(full), + "last_round": last or None, + "policy_version": last.get("policy_version"), + "lag_hist": dict(sorted(hist.items(), key=lambda kv: int(kv[0]))), + "lag_p50": quantile(lags, 0.5), "lag_max": max(lags) if lags else None, + "generate_p50_s": quantile(gens, 0.5), "generate_max_s": max(gens) if gens else None, + "groups": tot("groups"), "samples": tot("samples"), "truncated": tot("truncated"), + "zero_variance_dropped": tot("zero_variance_dropped"), "zero_variance_seen": tot("zero_variance_seen"), + "dropped_groups": tot("dropped_groups"), + "retried_trials": tot("retried_trials"), "lost_trials": tot("lost_trials"), + # the last round that carried a reward (empty rounds carry none) + "reward_mean_last": rew.get("reward_mean") if rew else None, + "reward_last_ts": rew.get("ts") if rew else None, + "reward_last_rollout_id": rew.get("rollout_id") if rew else None, + # incarnations in start order; every one after the first is a restart + "incarnations": incs, + "restarts": [{"ts": x["first_ts"], "incarnation": x["incarnation"]} for x in incs[1:]], + "backlog": s["backlog"][-500:], "max_backlog_groups": s["max_backlog_groups"], + "waiting": s["last_wait"], "ended": s["ended"], "resource": s["resource"], + "sampling": s.get("sampling"), "provenance": s.get("provenance"), + "incarnation": incs[-1]["incarnation"] if incs else s.get("incarnation"), "shard": s.get("shard"), + "first_ts": s.get("start"), "last_ts": s.get("last_ts"), + } + trainers = {} + for iid, t in self.trainer.items(): + rows = [t["rounds"][k] for k in sorted(t["rounds"])] + if not rows: + continue + ages: dict[str, float] = {} + for r in rows: + for k, v in (r.get("ages") or {}).items(): + if _f(v) is not None: + ages[k] = ages.get(k, 0) + v + + def tot(key: str, _rows: list = rows) -> float | None: + vals = [_f(r.get(key)) for r in _rows if _f(r.get(key)) is not None] + return sum(vals) if vals else None + + trainers[iid] = {"rounds": rows[-500:], "age_hist": dict(sorted(ages.items(), key=lambda kv: _age_key(kv[0]))), + "stale_groups": tot("stale_groups"), "gate_dropped": tot("gate_dropped"), + "over_age_discarded_groups": tot("over_age_discarded_groups"), + "trained_groups": tot("trained_groups"), "left": t["left"], + **{k: tot(k) for k in FILL_KEYS}} + return {"samplers": samplers, "trainers": trainers, "backlog": self.backlog_series(), + "releases": self.release_rows(), "pool": self.pool[-500:], + "roles": dict(self.roles)} + + def backlog_series(self) -> dict | None: + """Groups in the buffer not yet read by the trainers: groups written by all + samplers (``rl_sampler_round.groups``) minus the trainers' read position + (``buffer_selected`` of ``rl_policy_released`` / ``rl_sampler_policy_loaded``, + the position the samplers' own backlog rule starts from). Points at every + write and every read-position update; None without sampler rounds.""" + flow = [(r["ts"], 0, _f(r.get("groups"))) + for s in self.sampler.values() for r in s["rounds"].values() + if r.get("ts") is not None and _f(r.get("groups"))] + if not flow: + return None + reads = [(x["ts"], 1, _f(x["buffer_selected"])) for x in self.releases + if x["ts"] is not None and _f(x.get("buffer_selected")) is not None] + reads += [(x["loaded_at"], 1, _f(x["buffer_selected"])) for x in self.loads + if x["loaded_at"] is not None and _f(x.get("buffer_selected")) is not None] + source = "buffer_selected" if reads else None + if not reads: + # no buffer_selected on the tapes (e.g. no release yet): the trainers' read + # position is the groups they consumed (rl_group_policies + stale skipped) + done = 0.0 + for ts, n in sorted((r.get("ts") or r.get("read_ts"), (r.get("groups_read") or 0) + (_f(r.get("stale_groups")) or 0)) + for t in self.trainer.values() for r in t["rounds"].values() + if (r.get("ts") or r.get("read_ts")) is not None and r.get("groups_read") is not None): + done += n + reads.append((ts, 1, done)) + source = "trainer groups consumed" if reads else None + flow.sort(key=lambda x: x[0]) + written, read, pts = 0.0, None, [] + for ts, kind, v in sorted(flow + reads, key=lambda x: (x[0], x[1])): + if kind == 0: + written += v + else: + read = v if read is None else max(read, v) + pts.append([ts, written, read, None if read is None else max(0.0, written - read)]) + step = max(1, len(pts) // 400) + out = pts[::step] + if out[-1] is not pts[-1]: + out.append(pts[-1]) + reported = {iid: s["backlog"] for iid, s in self.sampler.items() if s["backlog"]} + return {"points": out, "written": written, "read": read, "has_reads": read is not None, "read_source": source, + "reported": reported} + + def release_rows(self) -> list[dict]: + """One row per released policy version: when the trainer released it and + when each sampler finished loading it (loaded_at - released_at = latency).""" + by_v: dict[Any, dict] = {} + for rel in self.releases: + v = rel["policy_version"] + row = by_v.setdefault(v, {"policy_version": v, "loads": {}}) + row.update({k: rel[k] for k in ("trainer", "released_at", "upload_s", "bytes", "origin", "local_version")}) + for ld in self.loads: + v = ld["policy_version"] + row = by_v.setdefault(v, {"policy_version": v, "loads": {}, "released_at": ld["released_at"]}) + if row.get("released_at") is None: + row["released_at"] = ld["released_at"] + rel = row.get("released_at") + row["loads"][ld["sampler"]] = { + "loaded_at": ld["loaded_at"], "seen_at": ld["seen_at"], "load_wall_s": ld["load_wall_s"], + "resolve_s": ld["resolve_s"], "load_s": ld["load_s"], + "latency_s": round(ld["loaded_at"] - rel, 3) if ld["loaded_at"] is not None and rel is not None else None} + return [by_v[k] for k in sorted(by_v, key=lambda v: (v is None, v if isinstance(v, int) else 0))] + + +def _row_order(r: dict) -> tuple: + ts = r.get("ts", r.get("read_ts")) + return (ts is None, ts if ts is not None else 0.0, r.get("rollout_id") or 0) + + +def _empty_round(r: dict) -> bool: + return r.get("reward_mean") is None and not _f(r.get("groups")) and not _f(r.get("samples")) + + +def _age_key(k: Any) -> int: + try: + return int(k) + except (TypeError, ValueError): + return 1 << 30 diff --git a/yeto/dashboard/cli.py b/yeto/dashboard/cli.py index 99510784c..c3b84788b 100644 --- a/yeto/dashboard/cli.py +++ b/yeto/dashboard/cli.py @@ -19,6 +19,9 @@ def common(p: argparse.ArgumentParser) -> None: p.add_argument("--tapes", nargs="*", default=[], metavar="PATH", help="extra tape files or directories (*.jsonl up to 2 levels deep): learner tapes, " "syncer tape, controller journal.jsonl, fleet.jsonl") + p.add_argument("--store", default=None, metavar="URI", + help="run store (s3://bucket/runs/ or a directory): read every island's tape " + "uploaded under tapes/ (samplers on leased nodes, replacements included)") p.add_argument("--prices", default=None, help="price table JSON (default: built-in example, not a bill)") p.add_argument("--budget", type=float, default=None, help="budget cap in $ for the cost bar/alerts") p.add_argument("--thresholds", default=None, help="alert threshold overrides JSON") @@ -61,11 +64,19 @@ def _paths(args) -> list[str]: paths = list(args.tapes) if args.run: paths = run_sources(args.run) + paths - if not paths: - raise SystemExit("[dashboard] no tapes: pass --run and/or --tapes PATH...") + if not paths and not getattr(args, "store", None): + raise SystemExit("[dashboard] no tapes: pass --run , --tapes PATH... and/or --store URI") return paths +def _store(args): + if not getattr(args, "store", None): + return None + from .store_tapes import StoreTapes, open_objects + + return StoreTapes(open_objects(args.store), label="store") + + def _attach_compare(reducer, args) -> None: """load the --compare run into its own reducer.""" if not getattr(args, "compare", None): @@ -99,7 +110,7 @@ def main(args) -> int: try: return serve(paths, reducer=reducer, host=args.host, port=args.port, - interval=args.poll, head=args.head) + interval=args.poll, head=args.head, store=_store(args)) except ValueError as exc: print(f"[dashboard] {exc}", file=sys.stderr) return 2 @@ -108,6 +119,10 @@ def main(args) -> int: from .sources import load_all sources = load_all(reducer, paths) + st = _store(args) + if st is not None: + st.pump(reducer) + sources = list(sources) + sorted(st.sources) out = export_html(reducer, args.out) print(f"[dashboard] {len(sources)} tape(s), {reducer.event_seq} record(s) -> {out}") return 0 diff --git a/yeto/dashboard/cost.py b/yeto/dashboard/cost.py index c42f19194..d01c51152 100644 --- a/yeto/dashboard/cost.py +++ b/yeto/dashboard/cost.py @@ -15,6 +15,7 @@ from pathlib import Path EXAMPLE_PRICES = Path(__file__).with_name("prices.example.json") +INFERRED_RUNNING_S = 900.0 # an inferred island with an event this recent is still billed LOCAL_CLOUDS = ("local", "\u672c\u5730") # "local" in Chinese, as older tapes wrote it @@ -74,13 +75,21 @@ def island_rate(prices: dict, isl: dict) -> dict: def island_cost(isl: dict, prices: dict, now: float) -> dict: rt = island_rate(prices, isl) unit = rt["unit_usd_gpu_h"] - if isl.get("ready_ts") is None: + running = isl.get("open_ts") is not None + if isl.get("ready_ts") is None and isl.get("hw_inferred"): + # no fleet.jsonl (store-only run): billed from each launch's first to last event + spans = [(v["first_ts"], v["last_ts"]) for v in (isl.get("launches") or {}).values() + if v.get("first_ts") is not None and v.get("last_ts") is not None] + if not spans and isl.get("first_ts") is not None and isl.get("last_event_ts") is not None: + spans = [(isl["first_ts"], isl["last_event_ts"])] + hours = sum(b - a for a, b in spans) / 3600.0 if spans else None + running = isl.get("last_event_ts") is not None and now - isl["last_event_ts"] < INFERRED_RUNNING_S + elif isl.get("ready_ts") is None: hours = None else: open_s = max(0.0, now - isl["open_ts"]) if isl.get("open_ts") is not None else 0.0 hours = (isl.get("closed_s", 0.0) + open_s) / 3600.0 rate = rt["rate_usd_h"] - running = isl.get("open_ts") is not None return {"id": isl["id"], "unit_usd_gpu_h": unit, "rate_usd_h": rate, "hours": hours, "gpu_rate_usd_h": rt["gpu_rate_usd_h"], "host_rate_usd_h": rt["host_rate_usd_h"], "gpu_only": rt["gpu_only"], "cpus": isl.get("cpus"), "memory_gib": isl.get("memory_gib"), diff --git a/yeto/dashboard/prices.example.json b/yeto/dashboard/prices.example.json index 76732e13c..48de520a0 100644 --- a/yeto/dashboard/prices.example.json +++ b/yeto/dashboard/prices.example.json @@ -1,5 +1,5 @@ { - "_note": "Example price table: illustrative values, not billing data; point --prices at your own table. host_prices are per core-hour and per GiB-hour rates (Modal converted from its public rates of $0.0000131/core-second and $0.00000222/GiB-second)", + "_note": "Example price table: illustrative values, not billing data; point --prices at your own table. host_prices are per core-hour and per GiB-hour rates (Modal converted from its public rates of $0.0000131/core-second and $0.00000222/GiB-second); modal:B300 = $136/h per 2x8 B300 island / 16, nc:MI355X = $24/h per 8-card NC machine / 8 (K3.1 run quotes)", "prices": { "modal:H100": 3.95, "modal:H200": 4.54, @@ -8,7 +8,9 @@ "aws:H100": 6.88, "verda:H100": 2.29, "verda:H200": 2.99, - "*:A100": 1.79 + "*:A100": 1.79, + "modal:B300": 8.5, + "nc:MI355X": 3.0 }, "host_prices": { "modal": { diff --git a/yeto/dashboard/reducer.py b/yeto/dashboard/reducer.py index 8552c6c49..9dcba48ec 100644 --- a/yeto/dashboard/reducer.py +++ b/yeto/dashboard/reducer.py @@ -22,15 +22,18 @@ from __future__ import annotations import math +import re import time from collections import deque from typing import Any, Iterable from . import alerts as alerts_mod +from . import async_rl from . import cost as cost_mod from . import sft as sft_mod -FLEET_EVENTS = frozenset({"island_launch", "island_ready", "island_lost", "island_stop", "cost_tick"}) +FLEET_EVENTS = frozenset({"island_launch", "island_ready", "island_lost", "island_stop", "island_done", + "cost_tick"}) POLICY_SWEEP_LEDGER_EVENT = "policy_sweep_ledger" SYNCER_ISLAND_EVENTS = frozenset({"rl_fragment_push", "rl_policy_apply", "rl_pull_resend", "rl_member_publication"}) @@ -50,6 +53,10 @@ ("reward_p90", "reward p90", ("reward_p90",)), ("pg_loss", "pg_loss", ("pg_loss", "train_metrics.pg_loss", "train/pg_loss")), ("kl", "KL", ("mean_kl", "train_metrics.ppo_kl", "train_metrics.kl_loss", "train/train_rollout_kl")), + # async GRPO: the train-vs-rollout KL (a buffer step's PPO KL is 0 by construction) + ("kl_tr", "train-rollout KL", ("mean_kl", "train_metrics.train_rollout_kl", "mismatch.train_rollout_kl", + "train/train_rollout_kl")), + ("tis_clip", "TIS clipfrac", ("mismatch.tis_clipfrac", "train_metrics.tis_clipfrac", "train/tis_clipfrac")), ("entropy", "entropy", ("entropy", "train_metrics.entropy", "train_metrics.entropy_loss")), ("grad_norm", "grad_norm", ("grad_norm", "train_metrics.grad_norm", "train/grad_norm")), ("clip", "clip total", ("clip_fraction", "train_metrics.pg_clipfrac", "train/pg_clipfrac")), @@ -57,6 +64,14 @@ ("clip_hi", "clip high side", ("train_metrics.pg_clipfrac_upper", "train_metrics.pg_clipfrac_high")), ("resp_len", "response len", ("resp_len_mean",)), ("tok_s", "tok/s", ("tok_per_s",)), + ("policy_age", "policy age mean", ("policy_age_mean",)), + ("policy_age_max", "policy age max", ("policy_age_max",)), + # score centering: the off-policy guard is the IS weight, not a PPO ratio clip + ("is_clip", "IS clipped tokens", ("train_metrics.sc_is_clipped", "train/sc_is_clipped")), + ("is_weight", "IS weight mean", ("train_metrics.sc_importance_weight", "train/sc_importance_weight")), + ("tail_ratio", "tail ratio", ("train_metrics.sc_tail_ratio", "train/sc_tail_ratio")), + # YETO_RL_NAN_SKIP: cumulative rounds skipped for a non-finite grad_norm (rl_round_trained) + ("nan_skipped", "NaN-skipped rounds", ("nan_skipped",)), ) EXTRA_SERIES: tuple[tuple[str, tuple[str, ...]], ...] = ( ("loss", ("loss", "train_metrics.loss", "train/loss")), @@ -68,6 +83,7 @@ ("resp_p95", ("resp_len_p95",)), ("logprob_diff", ("train_metrics.train_rollout_logprob_abs_diff", "train/train_rollout_logprob_abs_diff")), ("train_step", ("train_step",)), + ("step_s", ("step_seconds",)), ) @@ -117,6 +133,8 @@ def classify(record: dict) -> str: event = record.get("event") if isinstance(event, str) and event.startswith("sft_"): return "sft" + if isinstance(event, str) and event.startswith(async_rl.POOL_PREFIX): + return "pool" if event in FLEET_EVENTS: return "fleet" if "kind" in record and "seq" in record and event is None: @@ -136,6 +154,8 @@ def classify(record: dict) -> str: NODE_SERIES_POINTS = 300 # rl_timeline_span task -> page phase: R rollout / T training / S post-train sync / P publish SPAN_PHASES = {"generate": "R", "train": "T", "outer_sync": "S", "publish": "P"} +# rl_driver_phase name -> rl_timeline_span task (derived spans) +DRIVER_PHASE_TASK = {"sync": "outer_sync"} # rl_rollout_cutoff fields shown per round CUTOFF_KEYS = ("submitted_groups", "target_groups", "filtered_groups", "discarded_groups", "discarded_trajectories", "discarded_tokens", "discarded_unknown_groups", "mechanism") @@ -147,7 +167,9 @@ def classify(record: dict) -> str: UTIL_TRAJ_MAX = 400 # per-trajectory rows sent to the page (most recent) # Events only written once driver.run() is going (older tapes may lack rl_driver_start). DRIVER_RUNNING_EVENTS = frozenset({"rl_driver_phase", "rl_heartbeat", "rl_local_round", "rl_round_trained", - "rl_resource_sample", "rl_publication"}) + "rl_resource_sample", "rl_publication", + # a sampler island runs no driver: its loop's events + "rl_sampler_start", "rl_sampler_round"}) def _fmt_override_value(value: Any) -> str: @@ -204,8 +226,13 @@ def _new_island(iid: str) -> dict: "startup_steps": {}, "startup_step": None, "startup_step_ts": None, # cutoff per round, per-trajectory timing, load samples "cutoffs": {}, "trajs": [], "load": deque(maxlen=NODE_SERIES_MAX), + # rl_driver_phase marks (wall clock): phase spans when the tape has no rl_timeline_span + "phase_marks": [], # launch-preflight-guards 3.5: rl_island_override (negative-test island) "negative_test": None, + # store tape launches (tapes///: one per trainer restart) and + # trainer_start_phase marks of the current launch + "launches": {}, "launch": None, "launch_ts": None, "start_phase": None, "start_phase_ts": None, } @@ -240,10 +267,13 @@ def __init__(self, *, run: str | None = None, prices: dict | None = None, self.label: str | None = None # display label when compared (7.6) self.compare: list["Reducer"] = [] # self.sft = sft_mod.SftState() # ms-swift SFT islands + syncer side streams (sft.py) + self._source_island: dict[str, str] = {} # tape path -> the island its records name + self.async_ = async_rl.AsyncState() # samplers / buffer / releases / NC pool (async_rl.py) + self._superseded = False # the record being fed is from a superseded trainer launch # -- feeding ---------------------------------------------------------------- def feed(self, record: Any, *, source: str = "-", offset: int | None = None, - island: str | None = None, node: str | None = None) -> bool: + island: str | None = None, node: str | None = None, island_fixed: bool = False) -> bool: """Fold one record. Returns False when skipped as already consumed.""" if source not in self.sources_seen and source != "-": self.sources_seen.append(source) @@ -260,11 +290,24 @@ def feed(self, record: Any, *, source: str = "-", offset: int | None = None, self.spot.append(_spot_row(record)) kind = classify(record) self.counts[kind] = self.counts.get(kind, 0) + 1 - ts = record_ts(record) + ts = record_ts(record) if kind != "pool" else async_rl.pool_ts(record) if ts is not None: self.max_ts = ts if self.max_ts is None else max(self.max_ts, ts) self.min_ts = ts if self.min_ts is None else min(self.min_ts, ts) - iid = _iid(record.get("island_id", record.get("learner_id"))) or island + iid = (_iid(island) if island_fixed and island is not None and kind == "learner" else None) \ + or _iid(record.get("island_id", record.get("learner_id"))) or island \ + or (_iid(record.get("sampler_id")) if kind == "learner" else None) + if record.get("event") in ("dashboard_source_lost", "dashboard_source_restored") and source.endswith( + ".source.jsonl"): + # a mirror's status file names the island by the mirror's --island label; the + # records of the mirrored tape (the same path without .source.jsonl) name it + iid = self._source_island.get(source[: -len(".source.jsonl")], iid) + elif kind == "learner" and iid is not None and source != "-": + self._source_island[source] = iid + self._superseded = False + launch = _store_launch(source) + if kind == "learner" and iid is not None and launch is not None: + self._superseded = self._launch(self.island(iid), launch, ts) if kind != "sft" and sft_mod.is_syncer_tape(record): # The syncer writes the same tape for RL and SFT runs: the SFT page # reads it too, but it alone does not make a run an SFT run. @@ -348,6 +391,8 @@ def stopped_by_us(self, isl: dict) -> bool: stop = self.operator_stop() if stop is None or stop.get("time_unix") is None: return False + if stop.get("cause") == "teardown" and (isl["last_event_ts"] or 0) > stop["time_unix"] + 60: + return False # an older launch's teardown.json: this island ran after it fails = [x["ts"] for x in isl["recovery_required"] if x.get("ts") is not None] if isl["lost_ts"] is not None: fails.append(isl["lost_ts"]) @@ -357,6 +402,31 @@ def feed_many(self, records: Iterable[dict], **kw: Any) -> None: for r in records: self.feed(r, **kw) + def _launch(self, isl: dict, launch: str, ts: float | None) -> bool: + """Note a record of store tape ``launch`` (``tapes///``). A launch + that starts after the current one is a restart: the island's startup state starts + over. Returns True for a record of a superseded (older) launch, which must not + move the startup state of the current one.""" + row = isl["launches"].get(launch) + if row is None: + row = isl["launches"][launch] = {"launch": launch, "first_ts": ts, "last_ts": ts} + elif ts is not None: + row["first_ts"] = ts if row["first_ts"] is None else min(row["first_ts"], ts) + row["last_ts"] = ts if row["last_ts"] is None else max(row["last_ts"], ts) + cur = isl["launches"].get(isl["launch"]) if isl["launch"] is not None else None + if cur is None: + isl["launch"], isl["launch_ts"] = launch, row["first_ts"] + return False + if launch == isl["launch"]: + return False + if row["first_ts"] is not None and (cur["first_ts"] is None or row["first_ts"] > cur["first_ts"]): + isl["launch"] = launch # a restart: a new launch of the same island + isl.update({"driver_started": False, "startup_steps": {}, "startup_step": None, + "startup_step_ts": None, "start_phase": None, "start_phase_ts": None, + "launch_ts": row["first_ts"]}) + return False + return True + def island(self, iid: str) -> dict: isl = self.islands.get(iid) if isl is None: @@ -380,15 +450,29 @@ def _feed_learner(self, r: dict, iid: str | None, ts: float | None) -> None: if iid is None: return # no island identity: never fabricate a "?" island isl = self.island(iid) - self._touch(isl, ts) + if not event.startswith("dashboard_"): # the mirror's own status is not island activity + self._touch(isl, ts) isl["events_by_type"][event] = isl["events_by_type"].get(event, 0) + 1 isl["recent"].append({"ts": ts, "type": event, "summary": _summary(r)}) + self.async_.fold(r, iid, ts) if r.get("policy_version") is not None and event in ( "rl_policy_apply", "rl_publication", "rl_driver_phase", "rl_heartbeat", - "rl_round_cut", "rl_member_publication", "rl_resume"): + "rl_round_cut", "rl_member_publication", "rl_resume", + # a sampler's policy is the version it serves + "rl_sampler_round", "rl_sampler_policy_loaded"): isl["policy_version"] = r.get("policy_version") if event == "rl_island_hardware": # decoupling 7.7: recorded, never a refusal reason isl["hardware"] = {k: r.get(k) for k in ("compat_group", "driver_version", "cuda_version")} + if getattr(self, "_superseded", False) and ( + event in ("rl_driver_start", "rl_startup_step", "trainer_start_phase", "rl_heartbeat") + or event in DRIVER_RUNNING_EVENTS): + return # a stopped launch's late records: not this island's current state + if event == "remote_code_prewarm" and isinstance(r.get("nodes"), list): + isl["nodes_seen"] = max(len(r["nodes"]), isl.get("nodes_seen") or 0) + if event == "trainer_start_phase" and r.get("phase") is not None: + isl["start_phase"] = str(r.get("phase")) + if ts is not None: + isl["start_phase_ts"] = ts if event == "rl_driver_start" and isl.get("driver_start_ts") is None: isl["driver_start_ts"] = ts if event == "rl_driver_start": @@ -397,6 +481,8 @@ def _feed_learner(self, r: dict, iid: str | None, ts: float | None) -> None: isl["driver_started"] = True if event == "rl_driver_phase": isl["phase"] = r.get("phase") + if ts is not None and len(isl["phase_marks"]) < NODE_SERIES_MAX: + isl["phase_marks"].append((ts, r.get("phase"), r.get("rollout_id"))) elif event == "rl_startup_step": step = r.get("step") if step is not None: @@ -421,6 +507,26 @@ def _feed_learner(self, r: dict, iid: str | None, ts: float | None) -> None: isl["pubs"][pv] = ts elif event == "rl_resource_sample": isl["resource"] = _resource(r, ts) + gl = [g for g in (r.get("gpus") or []) if isinstance(g, dict)] + if gl: + isl["gpus_per_node"] = max(len(gl), isl.get("gpus_per_node") or 0) + isl["gpu_mem_total_mb"] = finite(gl[0].get("mem_total_mb")) or isl.get("gpu_mem_total_mb") + gm = [finite(g.get("mem_used_mb")) for g in isl["resource"]["gpus"]] + gm = [g for g in gm if g is not None] + if ts is not None and gm: # NVML samples of the driver's node (no host probe needed) + ser = isl["node_series"].setdefault("nvml", deque(maxlen=NODE_SERIES_MAX)) + ser.append([ts, round(sum(gm) / len(gm) / 1024, 2), isl["resource"]["util_pct"]]) + elif event == "rl_rollout_carry_over": + self._derived_resp_len(isl, r) + elif event == "rl_round_phases" and isinstance(r.get("rollout_id"), int): + self._point(isl, r["rollout_id"] + 1, {"round_s": finite(r.get("wall_seconds")), + "tokens": finite(r.get("trained_tokens"))}) + elif event == "rl_eval": + v = next((finite(r[k]) for k in sorted(r) if k.startswith("eval/") and finite(r.get(k)) is not None + and any(w in k for w in ("reward", "acc", "pass", "score"))), None) + if v is not None and ts is not None: + isl.setdefault("evals", []).append([ts, v, next(k for k in sorted(r) if finite(r.get(k)) == v + and k.startswith("eval/"))]) elif event == "rl_rollout_cutoff": rid = r.get("rollout_id") if isinstance(rid, int): @@ -476,6 +582,35 @@ def _feed_learner(self, r: dict, iid: str | None, ts: float | None) -> None: elif event == "rl_pull_resend": self.resends[step] = self.resends.get(step, 0) + 1 + def _derived_resp_len(self, isl: dict, r: dict) -> None: + """Mean response length of a round whose rl_round_trained has no batch + summary: the carry-over record's trained response tokens over the round's + trained samples (both counted on the trained batch).""" + rid = r.get("rollout_id") + if not isinstance(rid, int): + return + row = (self.async_.trainer.get(isl["id"]) or {}).get("rounds", {}).get(rid) or {} + toks, n = finite(r.get("trained_response_tokens")), finite(row.get("trained_samples")) + pt = isl["points"].get(rid + 1, {}) + if toks is not None and n and pt.get("resp_len") is None: + self._point(isl, rid + 1, {"resp_len": toks / n, "resp_len_derived": True}) + + def phase_spans(self, isl: dict) -> list[dict]: + """rl_timeline_span spans; without them (observation off) the wall-clock + spans between consecutive rl_driver_phase marks: a phase lasts until the next + mark. A publish mark has no rollout_id (round_records attaches it by time).""" + if isl["spans"]: + return isl["spans"] + out = [] + marks = isl["phase_marks"] + for (t, phase, rid), nxt in zip(marks, marks[1:]): + task = SPAN_PHASES.get(DRIVER_PHASE_TASK.get(phase, phase)) + if task is None or nxt[0] < t: + continue + out.append({"phase": task, "rollout_id": rid if isinstance(rid, int) else None, + "start": t, "end": nxt[0], "derived": True}) + return out + def _span(self, isl: dict, r: dict) -> None: """Keep phase spans on the wall clock. ``start``/``end`` are the driver's monotonic seconds; the record's ``time_unix`` is written at ``end``.""" @@ -491,8 +626,9 @@ def round_records(self, isl: dict) -> list[dict]: A publish span has no rollout_id: it belongs to the round whose sync it follows (it publishes policy rollout_id+1). Missing values stay None.""" rounds: dict[int, dict] = {} - pubs = sorted((s for s in isl["spans"] if s["phase"] == "P"), key=lambda s: s["start"]) - for sp in isl["spans"]: + spans = self.phase_spans(isl) + pubs = sorted((s for s in spans if s["phase"] == "P"), key=lambda s: s["start"]) + for sp in spans: rid = sp["rollout_id"] if sp["phase"] != "P" and isinstance(rid, int): rounds.setdefault(rid, {})[sp["phase"]] = [sp["start"], sp["end"]] @@ -513,7 +649,10 @@ def round_records(self, isl: dict) -> list[dict]: "p10": pt.get("reward_p10"), "p50": pt.get("reward_p50"), "p90": pt.get("reward_p90"), "trunc": pt.get("trunc"), "logprob_diff": pt.get("logprob_diff"), "resp_mean": pt.get("resp_len"), "resp_p95": pt.get("resp_p95"), "tok_s": pt.get("tok_s"), - "grad_norm": pt.get("grad_norm"), "kl": pt.get("kl"), + "grad_norm": pt.get("grad_norm"), "kl": pt.get("kl"), "kl_tr": pt.get("kl_tr"), + "tis_clip": pt.get("tis_clip"), "resp_derived": bool(pt.get("resp_len_derived")), + "ts": pt.get("t"), "lr": pt.get("lr"), "step_s": pt.get("step_s"), "round_s": pt.get("round_s"), + "tokens": pt.get("tokens"), "phases": ph, "dur": {k: round(v[1] - v[0], 1) for k, v in ph.items()}, "discarded_trainings": self.discarded_trainings(isl, rid), }) @@ -621,12 +760,18 @@ def page_view(self, *, live: bool = False, now: float | None = None) -> dict: "driver_start_ts": isl.get("driver_start_ts") or (starts[0] if starts else None), "cells": len(isl["cells"] or []), "hardware": isl.get("hardware"), "transactions": len(isl["tx_order"]), "ray_embed": self.ray_embed(isl, self.island_ids().index(iid)), + "role": self.async_.role(iid), "gpu_mem_note": gpu_mem_note(isl), + "evals": list(isl.get("evals") or [])[-500:], + "hbm_total_gib": _hbm_total_gib(isl), "last_ts": isl["last_event_ts"], } usage = {"syncer": self.counts["syncer"] > 0, "journal": self.counts["journal"] > 0, "cells": any(v["cells"] for v in isl_extra.values()), "transactions": any(v["transactions"] for v in isl_extra.values())} view = {"overview": ov, "islands": isl_extra, "rounds_syncer": self.rounds(), "usage": usage, "label": self.label, "spot": list(self.spot)} + asy = self.async_.view() + if asy is not None: + view["async"] = asy sft = self.sft_view(now) if sft is not None: view["sft"] = sft @@ -661,6 +806,9 @@ def _round_metrics(self, isl: dict, r: dict, event: str) -> None: values["tok_s_derived"] = True elif values.get("tok_s") is not None: values["tok_s_derived"] = False # measured tok_per_s wins over the derived estimate + if values.get("lr") is None and isinstance(r.get("applied_lrs"), list) and r["applied_lrs"]: + values["lr"] = finite(r["applied_lrs"][0]) + values["t"] = record_ts(r) # wall clock of the round (the metrics page plots over time) if nonfinite_seen(r, ("grad_norm", "train_metrics.grad_norm", "loss", "pg_loss")): isl["nonfinite"].append(x) self._point(isl, x, values) @@ -807,7 +955,7 @@ def _feed_fleet(self, r: dict, iid: str | None, ts: float | None) -> None: if isl["open_ts"] is None: isl["open_ts"] = ts isl["stop_ts"] = None - elif event in ("island_lost", "island_stop") and ts is not None: + elif event in ("island_lost", "island_stop", "island_done") and ts is not None: if event == "island_lost": isl["lost_ts"] = ts else: @@ -816,6 +964,9 @@ def _feed_fleet(self, r: dict, iid: str | None, ts: float | None) -> None: isl["closed_s"] += max(0.0, ts - isl["open_ts"]) isl["open_ts"] = None + def _feed_pool(self, r: dict, iid: str | None, ts: float | None) -> None: + self.async_.fold_pool(r) + def _feed_sft(self, r: dict, iid: str | None, ts: float | None) -> None: self.sft.fold(r, ts) @@ -884,7 +1035,14 @@ def _island_card(self, isl: dict, now: float) -> dict: age = (now - last_any) if last_any is not None else None hb_age = (now - isl["last_heartbeat_ts"]) if isl["last_heartbeat_ts"] is not None else None res = isl["resource"] or {} - starting = not isl["driver_started"] and not isl["finalized"] and isl["fleet_state"] not in ("lost", "stop") + role = self.async_.role(isl["id"]) + ended = isl["fleet_state"] in ("lost", "stop", "done") + starting = not isl["driver_started"] and not isl["finalized"] and not ended + samp = self.async_.sampler_state(isl["id"]) if role == "sampler" else None + if samp is not None: + # a sampler is starting until its first round (engine load + checks + the + # first generation write nothing for minutes) + starting = not samp["started_rounds"] and not samp["ended"] and not ended if starting and isl["ready_ts"] is not None: # A container that is still loading weights may have written nothing yet: # count the startup age from island_ready as well. @@ -892,7 +1050,34 @@ def _island_card(self, isl: dict, now: float) -> dict: age = now - ref startup_s = (now - min(t for t in (isl["ready_ts"], isl["first_ts"]) if t is not None) if starting and (isl["ready_ts"] is not None or isl["first_ts"] is not None) else None) + if starting and len(isl["launches"]) > 1 and isl["launch_ts"] is not None: + startup_s = now - isl["launch_ts"] # a restarted trainer: since its current launch + stage = self._startup_stage(isl, now) if starting else None + # the current step began at the last startup step / start phase mark (of this launch) + step_ref = max([t for t in (isl["startup_step_ts"], isl["start_phase_ts"]) if t is not None], default=None) \ + or (isl["launch_ts"] if len(isl["launches"]) > 1 else None) or isl["first_ts"] stopped = self.stopped_by_us(isl) + waiting_why = samp["waiting"] if samp else None + waiting = bool(waiting_why) or (bool(isl["recent"]) and isl["recent"][-1]["type"] == "rl_sampler_backlog_full") + left = self._trainer_left(isl) + pool = self.async_.pool_exit(isl["id"], after=isl["last_event_ts"]) if role == "sampler" else None + if left is None and pool is not None and pool["status"] != "done": + left = {"ts": pool["ts"], "reason": pool["event"].replace("nc_elastic_", "") + ( + f": {pool['why']}" if pool.get("why") else ""), "source": "nc pool", "lost": pool["status"] == "lost"} + sl = isl["source_lost"] + if left is None and sl and not sl.get("recovered") and sl.get("ts") is not None \ + and (isl["last_event_ts"] or 0) <= sl["ts"] + STOP_SLACK_S and not isl["finalized"]: + # the follow of its tape broke and nothing arrived since: the node is gone for this page + left = {"ts": sl["ts"], "reason": "tape source lost: " + str(sl.get("error") or "")[:120], + "source": "mirror", "lost": True} + role = role or self._guess_role(isl) + # no heartbeats: the island's own cadence (round / generation / load times) sets + # how long a quiet island is still healthy + gap = self.async_.expected_gap_s(isl["id"]) + warn_s = self.thresholds["heartbeat_warn_s"] + if not isl["heartbeat_seen"] and gap is not None: + warn_s = max(warn_s, gap) + trainers_done = role == "sampler" and self._trainers_done() if stopped and not isl["finalized"]: status = "stopped" starting = False @@ -900,11 +1085,32 @@ def _island_card(self, isl: dict, now: float) -> dict: status = "recovery" elif isl["fleet_state"] == "lost": status = "lost" - elif isl["finalized"] or isl["fleet_state"] == "stop": + elif left is not None and not isl["finalized"]: + # left the run before finishing (island_leave, or the NC pool released / lost its node) + status = "lost" if left.get("lost") else "left" + starting = False + elif pool is not None and pool["status"] == "done": + status = "done" + starting = False + elif role == "unrecognized": + status = "unknown" # a tape whose records name no trainer / sampler event + starting = False + elif isl["finalized"] or isl["fleet_state"] in ("stop", "done") or (samp and samp["ended"]): status = "done" + elif trainers_done: + status = "done" # a sampler that runs until its trainers end + starting = False + elif waiting: + status = "waiting" + starting = False elif starting and age is not None: - status = "stale" if age > self.thresholds["startup_warn_s"] else "starting" - elif age is not None and age > self.thresholds["heartbeat_warn_s"]: + # loading (weights, Ray, engine, first generation) is not stale: the + # startup alerts report a slow start instead + status = "starting" + elif age is not None and not isl["heartbeat_seen"] and gap is not None \ + and age > async_rl.LOST_FACTOR * warn_s: + status = "lost" # silent for 3x its own cadence: counted as gone + elif age is not None and age > warn_s: status = "stale" elif age is None: status = "unknown" @@ -914,14 +1120,23 @@ def _island_card(self, isl: dict, now: float) -> dict: "id": isl["id"], "name": isl["name"], "cloud": isl["cloud"], "region": isl["region"], "gpu": isl["gpu"], "gpus": isl["gpus"], "status": status, "finalized": bool(isl["finalized"]), "stopped_by_us": stopped, "starting": starting, "startup_s": _r(startup_s), "driver_started": isl["driver_started"], + # a sampler waiting at its backlog cap writes nothing until the trainers read + "waiting_backlog": waiting, "waiting_reason": waiting_why or ("backlog full" if waiting else None), + "role": role, "left": left, "stale_after_s": _r(warn_s), + "startup_slow": bool(starting and age is not None and age > (stage or {}).get( + "warn_s", self.thresholds["startup_warn_s"])), + # the current startup stage (trainer_start_phase / rl_startup_step) and its thresholds + "startup_stage": stage, + "launch": isl["launch"], + "launches": [{**v, "state": "current" if k == isl["launch"] else "superseded"} + for k, v in sorted(isl["launches"].items(), key=lambda kv: kv[1]["first_ts"] or 0)], "nodes": [isl["nodes"][k] for k in sorted(isl["nodes"], key=lambda x: int(x) if x.isdigit() else 0)], "last_event_age_s": _r(age), "heartbeat_age_s": _r(hb_age), "heartbeat_seen": isl["heartbeat_seen"], "round": isl["round"], "startup_steps": {k: {"seconds": v["seconds"], "step_s": v["step_s"]} for k, v in isl["startup_steps"].items()}, "startup_step": isl["startup_step"], - "startup_step_age_s": _r(now - (isl["startup_step_ts"] or isl["first_ts"])) - if starting and (isl["startup_step_ts"] or isl["first_ts"]) is not None else None, + "startup_step_age_s": _r(now - step_ref) if starting and step_ref is not None else None, "rollout_id": isl["rollout_id"], "policy_version": isl["policy_version"], "phase": isl["phase"], "staleness": isl["staleness"], "contribution": isl["contribution"], "gpu_util_pct": res.get("util_pct"), "mem_pct": res.get("mem_pct"), @@ -935,6 +1150,82 @@ def _island_card(self, isl: dict, now: float) -> dict: "negative_test": isl["negative_test"], } + def infer_hardware(self) -> None: + """Cloud / accelerator / count of an island that no fleet.jsonl record names (a + store-only run): a sampler from its rl_sampler_start provenance (card model, + tpN) and the NC pool naming its shard; a trainer from its tape launch dir + (``l0-modal-...``), its NVML samples (per-GPU memory -> model, GPUs per node) + and the node count of remote_code_prewarm. Marked ``hw_inferred``.""" + pool_shards = {str(p.get("shard")) for p in self.async_.pool if p.get("shard") is not None} + for iid, isl in self.islands.items(): + if isl["cloud"] is not None and isl["gpu"] is not None and not isl.get("hw_inferred"): + continue + if self.async_.role(iid) == "sampler": + prov = (self.async_.sampler.get(iid) or {}).get("provenance") or {} + m = re.search(r"tp(\d+)", str(prov.get("topology") or "")) + hw = {"cloud": "nc" if iid in pool_shards else None, "gpu": prov.get("card_model"), + "gpus": int(m.group(1)) if m else None} + elif isl["launch"] is not None or isl.get("gpu_mem_total_mb"): + per_node, mem = isl.get("gpus_per_node"), isl.get("gpu_mem_total_mb") + hw = {"cloud": "modal" if "-modal-" in f"-{isl['launch'] or ''}" else None, + "gpu": _gpu_from_mem(mem), + "gpus": per_node * max(1, isl.get("nodes_seen") or 1) if per_node else None} + else: + continue + for k, v in hw.items(): + if v is not None and (isl[k] is None or isl.get("hw_inferred")): + isl[k] = v + isl["hw_inferred"] = True + + def _startup_stage(self, isl: dict, now: float) -> dict: + """What a starting island is doing now and how long it may stay quiet: the actor + init stage (model build + weight load, K3.1 mbridge ~22 min with no log line) + gets ``startup_init_*`` thresholds instead of ``startup_*``.""" + th = self.thresholds + phase, step = isl["start_phase"], isl["startup_step"] + ts = max([t for t in (isl["start_phase_ts"], isl["startup_step_ts"]) if t is not None], default=None) + init = phase in INIT_STAGE_PHASES or step in INIT_STAGE_STEPS + if phase is not None and (isl["start_phase_ts"] or 0) >= (isl["startup_step_ts"] or 0): + label = START_STAGE.get(phase, f"after {phase}") + elif step is not None: + from .alerts import _next_startup_step # noqa: PLC0415 + label = f"waiting for {_next_startup_step(step)}" + else: + label = "starting" + return {"phase": phase, "step": step, "label": label, "init": init, + "since_ts": ts, "age_s": _r(now - ts) if ts is not None else None, + "warn_s": th["startup_init_warn_s"] if init else th["startup_warn_s"], + "severe_s": th["startup_init_severe_s"] if init else th["startup_severe_s"]} + + TRAINER_HINTS = ("rl_driver_", "trainer_start_", "rl_round_", "rl_local_round", "rl_publication", + "rl_heartbeat", "rl_policy_release", "rl_engine_selected", "rl_learner_") + + def _guess_role(self, isl: dict) -> str | None: + """Role from the tape's records, never from a file name: trainer events make a + trainer; a fleet.jsonl island with no tape yet stays None (a launcher island); + anything else (only dashboard / unknown events) is ``unrecognized``.""" + types = isl["events_by_type"] + if any(t.startswith(self.TRAINER_HINTS) for t in types): + return "trainer" + if isl["fleet_state"] is not None or not types: + return None + return "unrecognized" + + def _trainers_done(self) -> bool: + ids = [i for i, role in self.async_.roles.items() if role == "trainer" and i in self.islands] + return bool(ids) and all( + self.islands[i]["finalized"] or self.islands[i]["fleet_state"] in ("stop", "done", "lost") + or _clean_leave(self._trainer_left(self.islands[i])) for i in ids) + + def _trainer_left(self, isl: dict) -> dict | None: + """The island_leave of the island's current launch: a leave older than the + current launch belongs to a launch that a restart replaced.""" + left = self.async_.trainer_left(isl["id"]) + if left and left.get("ts") is not None and isl.get("launch_ts") is not None \ + and len(isl["launches"]) > 1 and left["ts"] < isl["launch_ts"]: + return None + return left + def series(self, isl: dict) -> dict: out: dict[str, list] = {} xs = sorted(isl["points"]) @@ -959,6 +1250,7 @@ def ray_embed(self, isl: dict, index: int) -> dict: def cost(self, now: float) -> dict: self.apply_run_host_shape() + self.infer_hardware() return cost_mod.cost_view(self, now) def overview(self, *, live: bool = False, now: float | None = None) -> dict: @@ -967,7 +1259,9 @@ def overview(self, *, live: bool = False, now: float | None = None) -> dict: cost = self.cost(now) cards = [self._island_card(self.islands[k], now) for k in self.island_ids()] alerts = alerts_mod.evaluate(self, now=now, rounds=rounds, cost=cost, cards=cards) - n_bad = sum(1 for c in cards if c["status"] in ("stale", "lost", "recovery")) + torn_down = (self.operator_stop() or {}).get("cause") == "teardown" + n_bad = sum(1 for c in cards if c["status"] in ("stale", "recovery") + or (c["status"] == "lost" and not torn_down)) # after teardown a lost node is expected sev0 = sum(1 for a in alerts if a["sev"] == 0) stop = self.operator_stop() if not cards: @@ -1103,18 +1397,79 @@ def modal_host_shape(args: dict) -> dict | None: return {"cpus": cpu * nodes, "memory_gib": mem * nodes} +GPU_MEM_HOW = ("launch with --rl-resource-sample-interval (trainer NVML, rl_resource_sample) " + "or set YETO_MODAL_HOSTMEM_SAMPLE_S= (Modal per-node probe, modal_host_sample)") + + +def _hbm_total_gib(isl: dict) -> float | None: + gpus = (isl["resource"] or {}).get("gpus") or [] + tot = [finite(g.get("mem_total_mb")) for g in gpus] + tot = [t for t in tot if t] + return round(max(tot) / 1024, 1) if tot else None + + +def gpu_mem_note(isl: dict) -> str | None: + """Why the GPU memory chart of an island is empty (None when it has data).""" + if isl["node_series"]: + return None + res = isl["resource"] + if res and res.get("available") is False: + return ("rl_resource_sample reported unavailable" + (f" ({res['reason']})" if res.get("reason") else "") + + ": the island's image lacks the card telemetry library (NVML: pynvml; ROCm: amdsmi)") + return "not sampled: " + GPU_MEM_HOW + + def _recovered_after(isl: dict) -> bool: last = isl["recovery_required"][-1]["ts"] or 0 return any(r.get("result") == "RECOVERED" and (r.get("ts") or 0) > last for r in isl["reconfig"]) +# per-GPU memory (NVML mem_total_mb) -> accelerator model, largest first +_GPU_BY_MEM = ((250_000, "B300"), (170_000, "B200"), (130_000, "H200"), (75_000, "H100")) + + +def _gpu_from_mem(mb: float | None) -> str | None: + if not mb: + return None + return next((name for floor, name in _GPU_BY_MEM if mb >= floor), None) + + +# island_leave reasons that end a run on purpose; a failed trainer (INFRA_FAILED ...) +# is relaunched, so its samplers keep running and are not "done" +CLEAN_LEAVE = frozenset({"STOPPED", "DONE", "FINISHED", "COMPLETED", "DRAINED"}) + + +def _clean_leave(left: dict | None) -> bool: + return bool(left) and str(left.get("reason") or "").upper() in CLEAN_LEAVE + + +def _store_launch(source: str) -> str | None: + """``