control-tower 6.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (446) hide show
  1. control_tower/__init__.py +3 -0
  2. control_tower/backends/__init__.py +31 -0
  3. control_tower/backends/fleet/__init__.py +1 -0
  4. control_tower/backends/fleet/ami.py +45 -0
  5. control_tower/backends/fleet/aws_identity.py +9 -0
  6. control_tower/backends/fleet/bootstrap.sh +349 -0
  7. control_tower/backends/fleet/config.py +40 -0
  8. control_tower/backends/fleet/controller.py +2058 -0
  9. control_tower/backends/fleet/job.py +325 -0
  10. control_tower/backends/fleet/run.py +399 -0
  11. control_tower/backends/fleet/run_job.py +270 -0
  12. control_tower/backends/fleet/worker.py +1045 -0
  13. control_tower/backends/fleet/worker_cli.py +68 -0
  14. control_tower/backends/hawk/__init__.py +7 -0
  15. control_tower/backends/hawk/config.py +114 -0
  16. control_tower/backends/hawk/download.py +93 -0
  17. control_tower/backends/hawk/middleman.py +64 -0
  18. control_tower/backends/hawk/model_closure.py +39 -0
  19. control_tower/backends/hawk/preflight.py +142 -0
  20. control_tower/backends/hawk/pull_envs.py +160 -0
  21. control_tower/backends/hawk/runner.py +1050 -0
  22. control_tower/backends/hawk/sandbox.py +561 -0
  23. control_tower/backends/hawk/yaml_generator.py +206 -0
  24. control_tower/backends/local.py +223 -0
  25. control_tower/cli/__init__.py +3 -0
  26. control_tower/cli/base.py +207 -0
  27. control_tower/cli/cloud.py +132 -0
  28. control_tower/cli/completions/completion_callbacks.py +125 -0
  29. control_tower/cli/completions/completions_build.py +59 -0
  30. control_tower/cli/completions/completions_hooks.py +209 -0
  31. control_tower/cli/extras.py +89 -0
  32. control_tower/cli/login.py +77 -0
  33. control_tower/cli/main.py +51 -0
  34. control_tower/cli/mcp_meta.py +30 -0
  35. control_tower/cli/protocols.py +140 -0
  36. control_tower/cli/run.py +893 -0
  37. control_tower/cli/task_sets.py +123 -0
  38. control_tower/datasets/__init__.py +0 -0
  39. control_tower/datasets/check.py +88 -0
  40. control_tower/datasets/cli.py +599 -0
  41. control_tower/datasets/external.py +675 -0
  42. control_tower/datasets/importers/__init__.py +47 -0
  43. control_tower/datasets/importers/eval_log.py +596 -0
  44. control_tower/datasets/importers/monitoring_bench.py +766 -0
  45. control_tower/datasets/importers/mythos_5_incident.py +241 -0
  46. control_tower/datasets/importers/sessions.py +159 -0
  47. control_tower/datasets/importers/shade_arena.py +588 -0
  48. control_tower/datasets/importers/sleight_bench.py +487 -0
  49. control_tower/datasets/model.py +45 -0
  50. control_tower/environments/build.py +191 -0
  51. control_tower/environments/env_args.py +98 -0
  52. control_tower/environments/environment.py +92 -0
  53. control_tower/environments/internet_simulator/.dockerignore +3 -0
  54. control_tower/environments/internet_simulator/Dockerfile +89 -0
  55. control_tower/environments/internet_simulator/__init__.py +0 -0
  56. control_tower/environments/internet_simulator/__main__.py +4 -0
  57. control_tower/environments/internet_simulator/compose.yml +31 -0
  58. control_tower/environments/internet_simulator/constants.py +53 -0
  59. control_tower/environments/internet_simulator/core/assets.py +304 -0
  60. control_tower/environments/internet_simulator/core/exposed_services.py +71 -0
  61. control_tower/environments/internet_simulator/core/generate_ca.py +68 -0
  62. control_tower/environments/internet_simulator/core/logs.py +212 -0
  63. control_tower/environments/internet_simulator/core/network_mapping.py +63 -0
  64. control_tower/environments/internet_simulator/core/tls.py +124 -0
  65. control_tower/environments/internet_simulator/core/trajectory.py +17 -0
  66. control_tower/environments/internet_simulator/core/webhook.py +195 -0
  67. control_tower/environments/internet_simulator/malware/malware.c +143 -0
  68. control_tower/environments/internet_simulator/modules/forward_proxy.py +431 -0
  69. control_tower/environments/internet_simulator/modules/provider_policy.py +390 -0
  70. control_tower/environments/internet_simulator/modules/response_rules.py +193 -0
  71. control_tower/environments/internet_simulator/modules/shell_sessions.py +206 -0
  72. control_tower/environments/internet_simulator/modules/upstream_proxy.py +289 -0
  73. control_tower/environments/internet_simulator/modules/upstream_specs.py +221 -0
  74. control_tower/environments/internet_simulator/personas.py +150 -0
  75. control_tower/environments/internet_simulator/requirements.txt +3 -0
  76. control_tower/environments/internet_simulator/resolve_override.c +75 -0
  77. control_tower/environments/internet_simulator/runtimes/__init__.py +0 -0
  78. control_tower/environments/internet_simulator/runtimes/docker.py +630 -0
  79. control_tower/environments/internet_simulator/runtimes/hawk.py +193 -0
  80. control_tower/environments/internet_simulator/runtimes/provisioning.py +39 -0
  81. control_tower/environments/internet_simulator/server.py +660 -0
  82. control_tower/environments/internet_simulator/source_hash.py +78 -0
  83. control_tower/environments/list_utils.py +482 -0
  84. control_tower/environments/registry.py +275 -0
  85. control_tower/environments/remote/__init__.py +0 -0
  86. control_tower/environments/remote/process.py +248 -0
  87. control_tower/environments/remote/remote.py +407 -0
  88. control_tower/environments/remote/sample.py +244 -0
  89. control_tower/environments/setting_runtime.py +176 -0
  90. control_tower/environments/settings.py +1170 -0
  91. control_tower/environments/settings.yml +26 -0
  92. control_tower/environments/settings_cli.py +610 -0
  93. control_tower/environments/validate.py +116 -0
  94. control_tower/eval/README.md +38 -0
  95. control_tower/eval/__init__.py +1 -0
  96. control_tower/eval/game/__init__.py +11 -0
  97. control_tower/eval/game/blue.py +570 -0
  98. control_tower/eval/game/harness.py +146 -0
  99. control_tower/eval/game/history.py +221 -0
  100. control_tower/eval/game/messages.py +75 -0
  101. control_tower/eval/game/red.py +346 -0
  102. control_tower/eval/game/registry.py +29 -0
  103. control_tower/eval/game/solver.py +214 -0
  104. control_tower/eval/game/view.py +114 -0
  105. control_tower/eval/game/white.py +37 -0
  106. control_tower/eval/launch/__init__.py +1 -0
  107. control_tower/eval/launch/assemble.py +416 -0
  108. control_tower/eval/launch/dataset.py +611 -0
  109. control_tower/eval/launch/dispatch.py +142 -0
  110. control_tower/eval/launch/inspect_tasks.py +248 -0
  111. control_tower/eval/launch/metadata.py +39 -0
  112. control_tower/eval/launch/naming.py +37 -0
  113. control_tower/eval/launch/options.py +40 -0
  114. control_tower/eval/launch/preflight.py +90 -0
  115. control_tower/eval/launch/selection.py +151 -0
  116. control_tower/eval/log/README.md +76 -0
  117. control_tower/eval/log/__init__.py +0 -0
  118. control_tower/eval/log/action.py +72 -0
  119. control_tower/eval/log/blue_calls.py +93 -0
  120. control_tower/eval/log/convert.py +1773 -0
  121. control_tower/eval/log/eval_trace_metadata.py +15 -0
  122. control_tower/eval/log/extended_thinking.py +56 -0
  123. control_tower/eval/log/format_history.py +35 -0
  124. control_tower/eval/log/from_eval_log.py +141 -0
  125. control_tower/eval/log/merge.py +87 -0
  126. control_tower/eval/log/monitor_scores.py +112 -0
  127. control_tower/eval/log/sample.py +403 -0
  128. control_tower/eval/log/task_args.py +23 -0
  129. control_tower/eval/log/task_scores.py +110 -0
  130. control_tower/eval/log/tool_calls.py +40 -0
  131. control_tower/eval/log/translation.py +342 -0
  132. control_tower/eval/log/utils.py +85 -0
  133. control_tower/eval/log/workspace_diff.py +16 -0
  134. control_tower/eval/monitoring/__init__.py +1 -0
  135. control_tower/eval/monitoring/posthoc.py +339 -0
  136. control_tower/eval/monitoring/scorer.py +409 -0
  137. control_tower/extras/__init__.py +0 -0
  138. control_tower/extras/claude_host.py +26 -0
  139. control_tower/extras/dashboard.py +269 -0
  140. control_tower/extras/dashboard_cli.py +19 -0
  141. control_tower/extras/env_description/__init__.py +0 -0
  142. control_tower/extras/env_description/audit/README.md +80 -0
  143. control_tower/extras/env_description/audit/__init__.py +0 -0
  144. control_tower/extras/env_description/audit/cli.py +81 -0
  145. control_tower/extras/env_description/audit/models.py +63 -0
  146. control_tower/extras/env_description/audit/prompts/system.md +68 -0
  147. control_tower/extras/env_description/audit/prompts/user.md +30 -0
  148. control_tower/extras/env_description/audit/task.py +278 -0
  149. control_tower/extras/env_description/generate/README.md +75 -0
  150. control_tower/extras/env_description/generate/cli.py +146 -0
  151. control_tower/extras/env_description/generate/prompts/system.md +62 -0
  152. control_tower/extras/env_description/generate/prompts/user.md +21 -0
  153. control_tower/extras/env_description/generate/task.py +359 -0
  154. control_tower/extras/env_description/shared/concurrency.py +72 -0
  155. control_tower/extras/env_description/shared/spec.md +79 -0
  156. control_tower/extras/env_description/shared/spec_loader.py +24 -0
  157. control_tower/extras/env_source.py +230 -0
  158. control_tower/extras/integration_test_quality/README.md +59 -0
  159. control_tower/extras/integration_test_quality/cli.py +115 -0
  160. control_tower/extras/integration_test_quality/prompts/analyze.md +44 -0
  161. control_tower/extras/integration_test_quality/prompts/check_tasks.md +29 -0
  162. control_tower/extras/integration_test_quality/prompts/coverage.md +27 -0
  163. control_tower/extras/integration_test_quality/prompts/fix.md +14 -0
  164. control_tower/extras/integration_test_quality/task.py +179 -0
  165. control_tower/extras/pairing_analysis/README.md +30 -0
  166. control_tower/extras/pairing_analysis/cli.py +82 -0
  167. control_tower/extras/pairing_analysis/prompts/pairing.md +43 -0
  168. control_tower/extras/pairing_analysis/task.py +367 -0
  169. control_tower/extras/prompts.py +28 -0
  170. control_tower/extras/scorer_idempotency/base.py +216 -0
  171. control_tower/extras/scorer_idempotency/checks/_json.py +31 -0
  172. control_tower/extras/scorer_idempotency/checks/chroma.py +285 -0
  173. control_tower/extras/scorer_idempotency/checks/docker.py +366 -0
  174. control_tower/extras/scorer_idempotency/checks/elasticsearch.py +293 -0
  175. control_tower/extras/scorer_idempotency/checks/filesystem.py +155 -0
  176. control_tower/extras/scorer_idempotency/checks/generated_manifest.json +13 -0
  177. control_tower/extras/scorer_idempotency/checks/influxdb.py +549 -0
  178. control_tower/extras/scorer_idempotency/checks/listener.py +81 -0
  179. control_tower/extras/scorer_idempotency/checks/minio.py +360 -0
  180. control_tower/extras/scorer_idempotency/checks/mongo.py +338 -0
  181. control_tower/extras/scorer_idempotency/checks/openldap.py +241 -0
  182. control_tower/extras/scorer_idempotency/checks/postgres.py +317 -0
  183. control_tower/extras/scorer_idempotency/checks/process.py +134 -0
  184. control_tower/extras/scorer_idempotency/checks/rabbitmq.py +369 -0
  185. control_tower/extras/scorer_idempotency/checks/redis.py +291 -0
  186. control_tower/extras/scorer_idempotency/checks/sqlite.py +309 -0
  187. control_tower/extras/scorer_idempotency/checks/sshd.py +270 -0
  188. control_tower/extras/scorer_idempotency/checks/vault.py +376 -0
  189. control_tower/extras/scorer_idempotency/cli.py +575 -0
  190. control_tower/extras/scorer_idempotency/compose_analyzer.py +182 -0
  191. control_tower/extras/scorer_idempotency/detection.py +236 -0
  192. control_tower/extras/scorer_idempotency/dockerfile_analyzer.py +377 -0
  193. control_tower/extras/scorer_idempotency/llm_generator.py +448 -0
  194. control_tower/extras/scorer_idempotency/runner.py +449 -0
  195. control_tower/extras/scorer_idempotency/runtime_prober.py +378 -0
  196. control_tower/extras/side_task_folder.py +115 -0
  197. control_tower/extras/storage.py +74 -0
  198. control_tower/extras/task_analysis/compare_side_tasks/README.md +25 -0
  199. control_tower/extras/task_analysis/compare_side_tasks/cli.py +25 -0
  200. control_tower/extras/task_analysis/compare_side_tasks/task.py +127 -0
  201. control_tower/extras/task_analysis/doc_audit/README.md +77 -0
  202. control_tower/extras/task_analysis/doc_audit/audits/base.py +129 -0
  203. control_tower/extras/task_analysis/doc_audit/audits/goal.py +92 -0
  204. control_tower/extras/task_analysis/doc_audit/audits/sd.py +98 -0
  205. control_tower/extras/task_analysis/doc_audit/cli.py +185 -0
  206. control_tower/extras/task_analysis/doc_audit/heuristics/base.py +53 -0
  207. control_tower/extras/task_analysis/doc_audit/heuristics/goal.py +215 -0
  208. control_tower/extras/task_analysis/doc_audit/heuristics/sd.py +222 -0
  209. control_tower/extras/task_analysis/doc_audit/prompts/goal/system.md +109 -0
  210. control_tower/extras/task_analysis/doc_audit/prompts/goal/user.md +21 -0
  211. control_tower/extras/task_analysis/doc_audit/prompts/sd/system.md +311 -0
  212. control_tower/extras/task_analysis/doc_audit/prompts/sd/user.md +26 -0
  213. control_tower/extras/task_analysis/doc_audit/sd_audit_labels.json +98 -0
  214. control_tower/extras/task_analysis/doc_audit/task.py +200 -0
  215. control_tower/extras/task_analysis/failure_type_audit/cli.py +125 -0
  216. control_tower/extras/task_analysis/failure_type_audit/prompts/system.md +9 -0
  217. control_tower/extras/task_analysis/failure_type_audit/prompts/user.md +23 -0
  218. control_tower/extras/task_analysis/failure_type_audit/task.py +190 -0
  219. control_tower/extras/task_analysis/goal_completeness/README.md +56 -0
  220. control_tower/extras/task_analysis/goal_completeness/cli.py +125 -0
  221. control_tower/extras/task_analysis/goal_completeness/prompts/system.md +44 -0
  222. control_tower/extras/task_analysis/goal_completeness/prompts/user.md +50 -0
  223. control_tower/extras/task_analysis/goal_completeness/task.py +320 -0
  224. control_tower/extras/task_analysis/goal_completeness/viewer.py +142 -0
  225. control_tower/extras/task_analysis/reward_hacking/README.md +46 -0
  226. control_tower/extras/task_analysis/reward_hacking/cli.py +115 -0
  227. control_tower/extras/task_analysis/reward_hacking/prompts/system.md +71 -0
  228. control_tower/extras/task_analysis/reward_hacking/prompts/user.md +62 -0
  229. control_tower/extras/task_analysis/reward_hacking/task.py +360 -0
  230. control_tower/extras/task_analysis/reward_hacking/viewer.py +122 -0
  231. control_tower/extras/task_analysis/rt_alignment/README.md +45 -0
  232. control_tower/extras/task_analysis/rt_alignment/cli.py +119 -0
  233. control_tower/extras/task_analysis/rt_alignment/prompts/system.md +50 -0
  234. control_tower/extras/task_analysis/rt_alignment/prompts/user.md +57 -0
  235. control_tower/extras/task_analysis/rt_alignment/task.py +357 -0
  236. control_tower/extras/task_analysis/rt_alignment/viewer.py +124 -0
  237. control_tower/extras/task_analysis/side_task_review/README.md +43 -0
  238. control_tower/extras/task_analysis/side_task_review/cli.py +87 -0
  239. control_tower/extras/task_analysis/side_task_review/prompts/review.md +30 -0
  240. control_tower/extras/task_analysis/side_task_review/task.py +98 -0
  241. control_tower/extras/task_analysis/validate_cover/README.md +32 -0
  242. control_tower/extras/task_analysis/validate_cover/cli.py +97 -0
  243. control_tower/extras/task_analysis/validate_cover/container_script.py +230 -0
  244. control_tower/extras/task_analysis/validate_cover/full_env_context.py +66 -0
  245. control_tower/extras/task_analysis/validate_cover/task.py +97 -0
  246. control_tower/extras/task_analysis/validate_cover/validate_cover_prompt.md +35 -0
  247. control_tower/extras/timing_report/__init__.py +17 -0
  248. control_tower/extras/timing_report/cli.py +71 -0
  249. control_tower/extras/timing_report/report.py +324 -0
  250. control_tower/extras/traj_analysis/submit_guessability/README.md +49 -0
  251. control_tower/extras/traj_analysis/submit_guessability/cli.py +125 -0
  252. control_tower/extras/traj_analysis/submit_guessability/prompts/system.md +29 -0
  253. control_tower/extras/traj_analysis/submit_guessability/prompts/user.md +70 -0
  254. control_tower/extras/traj_analysis/submit_guessability/task.py +351 -0
  255. control_tower/extras/traj_analysis/submit_guessability/viewer.py +159 -0
  256. control_tower/extras/traj_analysis/submit_guessability_verified/README.md +52 -0
  257. control_tower/extras/traj_analysis/submit_guessability_verified/cli.py +143 -0
  258. control_tower/extras/traj_analysis/submit_guessability_verified/prompts/system.md +36 -0
  259. control_tower/extras/traj_analysis/submit_guessability_verified/prompts/user.md +61 -0
  260. control_tower/extras/traj_analysis/submit_guessability_verified/task.py +673 -0
  261. control_tower/extras/traj_analysis/submit_guessability_verified/viewer.py +186 -0
  262. control_tower/live/__init__.py +0 -0
  263. control_tower/live/agents.py +117 -0
  264. control_tower/live/assets/dev_container_extensions.json +3 -0
  265. control_tower/live/assets/live.gitignore +53 -0
  266. control_tower/live/assets/pyrightconfig.json +4 -0
  267. control_tower/live/assets/settings.json +3 -0
  268. control_tower/live/cli.py +1133 -0
  269. control_tower/live/compose.py +117 -0
  270. control_tower/live/context/print_context.py +98 -0
  271. control_tower/live/context/task_getters.py +65 -0
  272. control_tower/live/context/traj_dir.py +334 -0
  273. control_tower/live/integration_tests.py +128 -0
  274. control_tower/live/sandbox_ops.py +134 -0
  275. control_tower/live/score.py +85 -0
  276. control_tower/live/session.py +513 -0
  277. control_tower/live/setup.py +294 -0
  278. control_tower/live/task_state.py +145 -0
  279. control_tower/live/traj/__init__.py +0 -0
  280. control_tower/live/traj/gen_traj.py +319 -0
  281. control_tower/live/traj/git_status.py +33 -0
  282. control_tower/live/traj/git_traj.py +735 -0
  283. control_tower/live/traj/git_traj_helpers/FORMAT.md +25 -0
  284. control_tower/live/traj/git_traj_helpers/call_insert.py +15 -0
  285. control_tower/live/traj/git_traj_helpers/call_str_replace.py +20 -0
  286. control_tower/live/traj/macros.py +100 -0
  287. control_tower/live/traj/migrate_traj.py +80 -0
  288. control_tower/live/traj/patch.py +62 -0
  289. control_tower/live/traj/solution.py +242 -0
  290. control_tower/live/traj/templ/add.py +14 -0
  291. control_tower/live/traj/templ/add_tool_call.py +85 -0
  292. control_tower/live/traj/templ/bash.py +127 -0
  293. control_tower/live/traj/templ/git_utils.py +93 -0
  294. control_tower/live/traj/templ/input_utils.py +145 -0
  295. control_tower/live/traj/templ/template_from.py +118 -0
  296. control_tower/live/traj/templ/tool_call_builder.py +119 -0
  297. control_tower/live/traj/tool_call_utils.py +283 -0
  298. control_tower/live/traj/traj.py +294 -0
  299. control_tower/live/tui/__init__.py +0 -0
  300. control_tower/live/tui/popup.py +293 -0
  301. control_tower/live/tui/tui.py +1237 -0
  302. control_tower/live/tui/tui.tcss +32 -0
  303. control_tower/metrics/__init__.py +0 -0
  304. control_tower/metrics/audit.py +357 -0
  305. control_tower/metrics/control.py +120 -0
  306. control_tower/models/__init__.py +3 -0
  307. control_tower/models/aliases.py +59 -0
  308. control_tower/models/cache_scopes.py +26 -0
  309. control_tower/models/inspect_rates.py +32 -0
  310. control_tower/models/openrouter_provider.py +65 -0
  311. control_tower/models/pricing.py +183 -0
  312. control_tower/models/rates.py +240 -0
  313. control_tower/models/role_model.py +179 -0
  314. control_tower/models/usage_hooks.py +14 -0
  315. control_tower/monitors/__init__.py +40 -0
  316. control_tower/monitors/_monitor.py +722 -0
  317. control_tower/monitors/checklist.py +95 -0
  318. control_tower/monitors/jev.py +217 -0
  319. control_tower/monitors/jev_questions.yml +32 -0
  320. control_tower/monitors/presets.yml +29 -0
  321. control_tower/monitors/prompts/action-goals.md +78 -0
  322. control_tower/monitors/prompts/action.md +70 -0
  323. control_tower/monitors/prompts/goal-analysis.md +72 -0
  324. control_tower/monitors/prompts/goals.md +76 -0
  325. control_tower/monitors/prompts/no-history.md +37 -0
  326. control_tower/monitors/prompts/off-task.md +86 -0
  327. control_tower/monitors/prompts/quick-goals.md +65 -0
  328. control_tower/monitors/prompts/quick.md +57 -0
  329. control_tower/monitors/prompts/simple.md +19 -0
  330. control_tower/monitors/prompts/strong-goals.md +87 -0
  331. control_tower/monitors/prompts/strong.md +79 -0
  332. control_tower/monitors/prompts/trajectory.md +16 -0
  333. control_tower/monitors/registry.py +18 -0
  334. control_tower/monitors/single_goal.py +67 -0
  335. control_tower/monitors/steps/checklist.md +38 -0
  336. control_tower/monitors/steps/single_goal.md +75 -0
  337. control_tower/monitors/steps/single_goal_analysis.md +71 -0
  338. control_tower/monitors/steps/three_scores.md +107 -0
  339. control_tower/monitors/three_scores.py +68 -0
  340. control_tower/policies/__init__.py +1 -0
  341. control_tower/policies/attack.py +185 -0
  342. control_tower/policies/honest.py +23 -0
  343. control_tower/policies/obvious_attack.py +67 -0
  344. control_tower/policies/prompts/attack.md +7 -0
  345. control_tower/policies/prompts/attack_planning.md +22 -0
  346. control_tower/policies/replay.py +201 -0
  347. control_tower/policies/scripted.py +115 -0
  348. control_tower/protocols/__init__.py +1 -0
  349. control_tower/protocols/defer_to_trusted.py +141 -0
  350. control_tower/protocols/null_blue.py +27 -0
  351. control_tower/protocols/prompts/trusted_editing.md +9 -0
  352. control_tower/protocols/resample.py +191 -0
  353. control_tower/protocols/trusted_editing.py +143 -0
  354. control_tower/protocols/trusted_monitoring.py +90 -0
  355. control_tower/protocols/watcher.py +352 -0
  356. control_tower/py.typed +0 -0
  357. control_tower/runs/README.md +93 -0
  358. control_tower/runs/__init__.py +0 -0
  359. control_tower/runs/convert.py +100 -0
  360. control_tower/runs/convert_eval_log_to_run.py +92 -0
  361. control_tower/runs/eval_log_io.py +200 -0
  362. control_tower/runs/io.py +190 -0
  363. control_tower/runs/models.py +64 -0
  364. control_tower/runs/mongo_docs.py +359 -0
  365. control_tower/runs/runs_cli.py +735 -0
  366. control_tower/runs/s3_upload.py +167 -0
  367. control_tower/runs/source.py +24 -0
  368. control_tower/runs/timing_utils.py +149 -0
  369. control_tower/runs/viewer/__init__.py +0 -0
  370. control_tower/runs/viewer/api.py +464 -0
  371. control_tower/runs/viewer/auth.py +50 -0
  372. control_tower/runs/viewer/io.py +204 -0
  373. control_tower/runs/viewer/models.py +25 -0
  374. control_tower/runs/viewer/upload.py +216 -0
  375. control_tower/runs/viewer/urls.py +501 -0
  376. control_tower/sandbox/__init__.py +6 -0
  377. control_tower/sandbox/_alpine_rootfs.py +230 -0
  378. control_tower/sandbox/_ssh.py +471 -0
  379. control_tower/sandbox/cli.py +236 -0
  380. control_tower/sandbox/constants.py +21 -0
  381. control_tower/sandbox/docker.py +685 -0
  382. control_tower/sandbox/docker_in_vm.py +527 -0
  383. control_tower/sandbox/firecracker.py +363 -0
  384. control_tower/sandbox/inspect_tools.py +46 -0
  385. control_tower/sandbox/qemu.py +251 -0
  386. control_tower/sandbox/replay/__init__.py +23 -0
  387. control_tower/sandbox/replay/environment.py +435 -0
  388. control_tower/sandbox/replay/score.py +37 -0
  389. control_tower/sandbox/simulated/__init__.py +25 -0
  390. control_tower/sandbox/simulated/actions.py +90 -0
  391. control_tower/sandbox/simulated/compose.py +192 -0
  392. control_tower/sandbox/simulated/environment.py +190 -0
  393. control_tower/sandbox/simulated/judge.py +158 -0
  394. control_tower/sandbox/simulated/prompts/judge.md +31 -0
  395. control_tower/sandbox/simulated/prompts/simulator.md +74 -0
  396. control_tower/sandbox/simulated/prompts/simulator_no_exec.md +22 -0
  397. control_tower/sandbox/simulated/scratch/Dockerfile +11 -0
  398. control_tower/sandbox/simulated/scratch/compose.yaml +42 -0
  399. control_tower/sandbox/simulated/session.py +155 -0
  400. control_tower/scanners.py +130 -0
  401. control_tower/task/__init__.py +0 -0
  402. control_tower/task/final_diff.py +88 -0
  403. control_tower/task/sample_metadata.py +46 -0
  404. control_tower/task/scorers.py +327 -0
  405. control_tower/task/select_utils.py +106 -0
  406. control_tower/task/selectors.py +49 -0
  407. control_tower/task/setup.py +115 -0
  408. control_tower/task/space.py +45 -0
  409. control_tower/task/task.py +230 -0
  410. control_tower/task/task_sets/__init__.py +203 -0
  411. control_tower/task/traj_dir.py +39 -0
  412. control_tower/task/workspace_baseline.py +581 -0
  413. control_tower/tools/__init__.py +0 -0
  414. control_tower/tools/answer.py +45 -0
  415. control_tower/tools/bash_tools.py +37 -0
  416. control_tower/tools/declarations.py +23 -0
  417. control_tower/tools/diff_display_text_editor.py +66 -0
  418. control_tower/tools/registry.py +34 -0
  419. control_tower/tools/split_text_editor_tools.py +125 -0
  420. control_tower/tools/task_tools.py +204 -0
  421. control_tower/tools/text_editor_output.py +53 -0
  422. control_tower/tools/tool_calls.py +44 -0
  423. control_tower/tools/tool_registry.py +363 -0
  424. control_tower/utils/__init__.py +0 -0
  425. control_tower/utils/cli_args.py +20 -0
  426. control_tower/utils/direct_api_generator.py +124 -0
  427. control_tower/utils/docs_utils.py +18 -0
  428. control_tower/utils/env.py +191 -0
  429. control_tower/utils/file_patterns.py +89 -0
  430. control_tower/utils/fs.py +137 -0
  431. control_tower/utils/function_registry/__init__.py +0 -0
  432. control_tower/utils/function_registry/base_model.py +30 -0
  433. control_tower/utils/function_registry/registry.py +535 -0
  434. control_tower/utils/function_registry/selection.py +68 -0
  435. control_tower/utils/paths.py +91 -0
  436. control_tower/utils/sanitize.py +29 -0
  437. control_tower/utils/suggestions.py +25 -0
  438. control_tower/utils/time.py +8 -0
  439. control_tower/utils/transcript.py +20 -0
  440. control_tower/utils/versioning.py +327 -0
  441. control_tower-6.0.0.dist-info/METADATA +91 -0
  442. control_tower-6.0.0.dist-info/RECORD +446 -0
  443. control_tower-6.0.0.dist-info/WHEEL +5 -0
  444. control_tower-6.0.0.dist-info/entry_points.txt +5 -0
  445. control_tower-6.0.0.dist-info/licenses/LICENSE +21 -0
  446. control_tower-6.0.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,2058 @@
1
+ """Fleet controller: orchestrates distributed EC2 eval.
2
+
3
+ Runs on the user's machine. Creates SQS queues, enqueues jobs,
4
+ launches EC2 workers, monitors progress, and returns S3 result locations.
5
+ """
6
+
7
+ import contextlib
8
+ import gzip
9
+ import hashlib
10
+ import json
11
+ import logging
12
+ import math
13
+ import os
14
+ import shutil
15
+ import tempfile
16
+ import time
17
+ from collections import defaultdict
18
+ from collections.abc import Collection
19
+ from concurrent.futures import ThreadPoolExecutor, as_completed
20
+ from dataclasses import dataclass
21
+ from datetime import UTC, datetime
22
+ from importlib.metadata import distribution
23
+ from pathlib import Path
24
+ from typing import TYPE_CHECKING, Any
25
+
26
+ import boto3
27
+
28
+ if TYPE_CHECKING:
29
+ from botocore.client import BaseClient
30
+ import re
31
+ import subprocess
32
+
33
+ import botocore.config
34
+ import botocore.exceptions
35
+ from inspect_ai.log import EvalConfig, EvalLog, read_eval_log
36
+
37
+ from control_tower.backends.fleet.ami import get_worker_ami_id
38
+ from control_tower.backends.fleet.aws_identity import caller_identity_name
39
+ from control_tower.backends.fleet.config import Ec2Config
40
+ from control_tower.backends.fleet.job import (
41
+ FLEET_RUNNER_SCHEMA_VERSION,
42
+ SQS_MAX_VISIBILITY_S,
43
+ SQS_VISIBILITY_TIMEOUT_S,
44
+ FleetManifest,
45
+ InspectFleetJobWire,
46
+ JobDoneMessage,
47
+ JobPayload,
48
+ LifecycleMessage,
49
+ SampleDoneMessage,
50
+ compute_fleet_id,
51
+ enqueue_jobs,
52
+ list_eval_keys,
53
+ )
54
+ from control_tower.eval.log.merge import merge_eval_logs
55
+ from control_tower.models.rates import default_rate_table
56
+ from control_tower.sandbox.constants import FLEET_PLANNING_CONCURRENCY
57
+ from control_tower.utils.env import EnvVar
58
+ from control_tower.utils.paths import cache_dir, find_repo_root
59
+ from control_tower.utils.versioning import git_changed_paths, setting_repo_dirs
60
+ from ct_setting_sdk.errors import OperationError, ValidationError
61
+
62
+ logger: logging.Logger = logging.getLogger(__name__)
63
+
64
+
65
+ INSTANCE_PROFILE_NAME = "linuxbench-worker"
66
+ FLEET_SSM_PREFIX = "/linuxbench/fleet"
67
+ WORKER_SECRET_ENV_VARS = (
68
+ "ANTHROPIC_API_KEY",
69
+ "OPENAI_API_KEY",
70
+ "OPENROUTER_API_KEY",
71
+ "WATCHER_API_KEY",
72
+ "WATCHER_API_URL",
73
+ "GH_TOKEN",
74
+ "DOCKERHUB_USERNAME",
75
+ "DOCKERHUB_TOKEN",
76
+ )
77
+
78
+ # Default fleet settings – imported from shared constants module
79
+ FAILURE_RATE_THRESHOLD = 0.20 # pause if >20% failed
80
+ STALE_WARN_S = 600 # warn after 10 min of no progress
81
+ STALE_ABORT_S = 10800 # abort after 3 hours of no progress
82
+ STALE_HARD_ABORT_S = 21600 # 6hr absolute ceiling even with lively workers
83
+ ROOT_VOLUME_SIZE_GB = 300
84
+ SQS_QUEUE_DELETE_COOLDOWN_S = 60
85
+ SQS_PURGE_WINDOW_S = 60
86
+ DLQ_RECONCILE_INTERVAL_S = 60
87
+ # Extra receives per monitor tick when the done queue returns a full batch.
88
+ DONE_QUEUE_RECEIVERS = 8
89
+
90
+
91
+ @dataclass(frozen=True)
92
+ class FleetRunResult:
93
+ """Generic result produced by the fleet runner after jobs are in S3."""
94
+
95
+ fleet_id: str
96
+ s3_bucket: str
97
+ s3_prefix: str
98
+ total_expected: int
99
+ total_completed: int
100
+ region: str
101
+ eval_log: EvalLog | None
102
+ setting_repos: dict[str, str]
103
+ shipped_dirty: dict[Path, bool]
104
+
105
+
106
+ class FleetStreamingLog:
107
+ """Append fleet events to <log_dir>/events.jsonl, one JSON object per line (per-sample results arrive via run_job.FleetSampleStreamer)."""
108
+
109
+ def __init__(self, *, log_dir: Path, resume: bool = False):
110
+ log_dir.mkdir(parents=True, exist_ok=True)
111
+ self.path: Path = log_dir / "events.jsonl"
112
+ mode = "a" if resume else "w"
113
+ self._f = self.path.open(mode, buffering=1)
114
+
115
+ def event(self, event_type: str, **fields: Any) -> None:
116
+ record = {
117
+ "ts": datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
118
+ "type": event_type,
119
+ **fields,
120
+ }
121
+ self._f.write(json.dumps(record, default=str) + "\n")
122
+ self._f.flush()
123
+
124
+ def load_seen_results(self) -> set[tuple[str, str]]:
125
+ """Load (job_id, status) tuples from existing result events."""
126
+ seen: set[tuple[str, str]] = set()
127
+ if not self.path.exists():
128
+ return seen
129
+ with self.path.open("r") as f:
130
+ for line in f:
131
+ try:
132
+ event = json.loads(line)
133
+ if event.get("type") == "result":
134
+ job_id = event.get("job_id")
135
+ status = event.get("status")
136
+ if job_id and status:
137
+ seen.add((job_id, status))
138
+ except (json.JSONDecodeError, KeyError):
139
+ continue
140
+ return seen
141
+
142
+ def close(self) -> None:
143
+ self._f.close()
144
+
145
+
146
+ class ProgressMonitor:
147
+ """Tracks fleet progress from SQS done-queue messages.
148
+
149
+ Encapsulates all monitor state: job statuses, durations, failure tracking,
150
+ instance liveness, and stale-progress detection. The controller drives the
151
+ loop; this class owns the bookkeeping.
152
+ """
153
+
154
+ def __init__(
155
+ self,
156
+ *,
157
+ sqs_client: "BaseClient",
158
+ done_queue_url: str,
159
+ work_queue_url: str,
160
+ dlq_url: str,
161
+ instance_ids: list[str],
162
+ concurrency_per_worker: int,
163
+ price_per_hour: float,
164
+ instance_type: str,
165
+ stale_abort_s: int,
166
+ stream_log: FleetStreamingLog,
167
+ pending_job_ids: set[str],
168
+ initial_seen_results: set[tuple[str, str]],
169
+ ):
170
+ self._sqs = sqs_client
171
+ self._done_queue_url = done_queue_url
172
+ self._work_queue_url = work_queue_url
173
+ self._dlq_url = dlq_url
174
+ self._total_expected = len(pending_job_ids)
175
+ self._pending_job_ids = pending_job_ids
176
+ self._expected_instances = set(instance_ids)
177
+ self._concurrency_per_worker = concurrency_per_worker
178
+ self._num_workers = len(instance_ids)
179
+ self._price_per_hour = price_per_hour
180
+ self._instance_type = instance_type
181
+ self._stale_abort_s = stale_abort_s
182
+ self._stream = stream_log
183
+
184
+ # Job tracking
185
+ self.seen_jobs: dict[str, str] = {} # job_id -> last known status
186
+ # (job_id, status) keys of already-logged sample results. Keying on the
187
+ # status as well as the job_id lets a genuine retry that changes the
188
+ # outcome (e.g. failed -> completed) surface as a fresh result event,
189
+ # while still collapsing exact SQS redeliveries of the same message.
190
+ self._seen_sample_results = initial_seen_results
191
+ self.error_summaries: list[str] = []
192
+ self.durations: list[float] = []
193
+ # (eval_duration_s, upload_s) for completed jobs that reported upload_s.
194
+ self.upload_samples: list[tuple[float, float]] = []
195
+ self.total_retries: int = 0
196
+ self.retried_jobs: int = 0
197
+ self._repeated_failures_by_job: defaultdict[str, int] = defaultdict(int)
198
+
199
+ # Instance tracking
200
+ self.instance_last_seen: dict[str, float] = {}
201
+ self.failed_bootstraps: set[str] = set()
202
+ self._first_worker_time: float | None = None
203
+ self._first_completion_time: float | None = None
204
+
205
+ # Timing
206
+ self._start = time.monotonic()
207
+ self._last_progress = time.monotonic()
208
+ self._last_completion = time.monotonic() # for hard abort (completions only)
209
+ self._stale_warned = False
210
+ self._lively_note_time: float = 0.0
211
+
212
+ # DLQ / work queue stats
213
+ self._last_dlq_check = 0.0
214
+ self._last_dlq_reconcile = 0.0
215
+ self._last_queue_check = 0.0
216
+ self.dlq_count: int = 0
217
+
218
+ @property
219
+ def completed(self) -> int:
220
+ return sum(1 for s in self.seen_jobs.values() if s == "completed")
221
+
222
+ @property
223
+ def failed(self) -> int:
224
+ return sum(1 for s in self.seen_jobs.values() if s == "failed")
225
+
226
+ @property
227
+ def repeated_failure_count(self) -> int:
228
+ return sum(self._repeated_failures_by_job.values())
229
+
230
+ @property
231
+ def permanent_failure_count(self) -> int:
232
+ return sum(1 for s in self.seen_jobs.values() if s == "permanent_failure")
233
+
234
+ @property
235
+ def terminal(self) -> int:
236
+ # DLQ jobs are counted via permanent_failure after explicit reconciliation.
237
+ # dlq_count (approximate) is display-only.
238
+ return self.completed + self.permanent_failure_count
239
+
240
+ @property
241
+ def in_flight(self) -> int:
242
+ return self._total_expected - self.terminal
243
+
244
+ @property
245
+ def is_done(self) -> bool:
246
+ return self.terminal >= self._total_expected
247
+
248
+ @property
249
+ def elapsed_s(self) -> float:
250
+ return time.monotonic() - self._start
251
+
252
+ @property
253
+ def stale_s(self) -> float:
254
+ return time.monotonic() - self._last_progress
255
+
256
+ def poll_and_process(self) -> bool:
257
+ """Poll done queue once, process messages. Returns True if progress was made."""
258
+ self._check_queue_stats()
259
+ self._reconcile_dlq()
260
+
261
+ try:
262
+ first = self._sqs.receive_message(
263
+ QueueUrl=self._done_queue_url,
264
+ MaxNumberOfMessages=10,
265
+ WaitTimeSeconds=20,
266
+ ).get("Messages", [])
267
+ batches = [first]
268
+ if len(first) == 10:
269
+ with ThreadPoolExecutor(max_workers=DONE_QUEUE_RECEIVERS) as pool:
270
+ batches += pool.map(
271
+ lambda _: self._sqs.receive_message(
272
+ QueueUrl=self._done_queue_url,
273
+ MaxNumberOfMessages=10,
274
+ WaitTimeSeconds=0,
275
+ ).get("Messages", []),
276
+ range(DONE_QUEUE_RECEIVERS),
277
+ )
278
+ except botocore.exceptions.ClientError as e:
279
+ error_code = e.response.get("Error", {}).get("Code", "")
280
+ if "NonExistentQueue" in error_code or "QueueDoesNotExist" in str(e):
281
+ raise RuntimeError(
282
+ f"SQS queue no longer exists: {self._done_queue_url}. "
283
+ "This usually means another eval with the same run-name "
284
+ "already completed and cleaned up the queue."
285
+ ) from e
286
+ logger.warning("Failed to poll done queue", exc_info=True)
287
+ time.sleep(5)
288
+ return False
289
+ except Exception:
290
+ logger.warning("Failed to poll done queue", exc_info=True)
291
+ time.sleep(5)
292
+ return False
293
+
294
+ progress = False
295
+ completed_before = self.completed
296
+
297
+ for msg in (msg for batch in batches for msg in batch):
298
+ try:
299
+ made_progress = self._handle_message(msg)
300
+ if made_progress:
301
+ progress = True
302
+ except Exception:
303
+ logger.warning(
304
+ "Bad done message, deleting to prevent redelivery", exc_info=True
305
+ )
306
+ with ThreadPoolExecutor(max_workers=len(batches)) as pool:
307
+ list(pool.map(self._delete_done_messages, [b for b in batches if b]))
308
+
309
+ if progress:
310
+ self._last_progress = time.monotonic()
311
+ self._stale_warned = False
312
+ if self.completed > completed_before:
313
+ self._last_completion = time.monotonic()
314
+
315
+ return progress
316
+
317
+ def _has_lively_workers(self) -> bool:
318
+ """Check if any workers have sent a message recently (within 10 min)."""
319
+ if not self.instance_last_seen:
320
+ return False
321
+ now = time.monotonic()
322
+ for t in self.instance_last_seen.values():
323
+ if now - t < 600: # heard from within last 10min
324
+ return True
325
+ return False
326
+
327
+ def check_stale(self, instance_ids: list[str]) -> str | None:
328
+ """Check for stale progress. Returns an action string or None.
329
+
330
+ Possible returns: "abort", "warn", or None.
331
+ """
332
+ # Hard abort: 6hr (or the stale-abort timeout, if longer) with no completions, regardless of liveness.
333
+ # Uses _last_completion (not _last_progress) so a stream of failures
334
+ # can't prevent the abort from firing.
335
+ no_completion_s = time.monotonic() - self._last_completion
336
+ if no_completion_s > max(STALE_HARD_ABORT_S, self._stale_abort_s):
337
+ print(f"\n\n No completions for {no_completion_s / 60:.0f} minutes.")
338
+ return "abort"
339
+
340
+ if instance_ids and self.failed_bootstraps.issuperset(instance_ids):
341
+ print(
342
+ f"\n Every worker failed to bootstrap ({len(instance_ids)} of {len(instance_ids)})."
343
+ )
344
+ return "abort"
345
+
346
+ lively = self._has_lively_workers()
347
+
348
+ if lively:
349
+ # Workers are alive but no completions — likely long-running evals
350
+ if self.stale_s > STALE_WARN_S:
351
+ now = time.monotonic()
352
+ if now - self._lively_note_time > 300: # rate-limit to every 5min
353
+ self._lively_note_time = now
354
+ active = sum(
355
+ 1 for t in self.instance_last_seen.values() if now - t < 600
356
+ )
357
+ print(
358
+ f"\n No completions for {self.stale_s / 60:.0f}m but "
359
+ f"{active} worker(s) active — long-running evals likely in progress"
360
+ )
361
+ return None
362
+
363
+ # Workers are silent — use original escalation
364
+ if self.stale_s > self._stale_abort_s:
365
+ print(f"\n\n No progress for {self.stale_s / 60:.0f} minutes.")
366
+ return "abort"
367
+ if not self._stale_warned and self.stale_s > STALE_WARN_S:
368
+ self._stale_warned = True
369
+ quiet = self._quiet_workers()
370
+ print(
371
+ "\n Warning: no progress for 10+ minutes. Workers may have failed to bootstrap."
372
+ "\n Check /var/log/lb-worker-bootstrap.log on worker instances via SSM or SSH."
373
+ f"\n Instance IDs: {instance_ids}"
374
+ )
375
+ if quiet:
376
+ print(f" Quiet workers (>5m): {', '.join(i[-8:] for i in quiet)}")
377
+ return "warn"
378
+ return None
379
+
380
+ def check_failure_rate(self) -> bool:
381
+ """Returns True if failure rate warrants aborting."""
382
+ grace_period = 180 # 3 min for workers to bootstrap
383
+ total_seen = len(self.seen_jobs)
384
+ if self.elapsed_s <= grace_period or total_seen < 5:
385
+ return False
386
+
387
+ if self.failed == total_seen:
388
+ return True
389
+ if (
390
+ self.repeated_failure_count >= 5
391
+ and self.repeated_failure_count / total_seen > FAILURE_RATE_THRESHOLD
392
+ ):
393
+ return True
394
+ return False
395
+
396
+ def _cost_estimate(self) -> float:
397
+ """Running EC2 cost estimate."""
398
+ return self._num_workers * self.elapsed_s * self._price_per_hour / 3600
399
+
400
+ def print_progress(self) -> None:
401
+ """Print the current progress line."""
402
+ pf = self.permanent_failure_count
403
+ dlq_str = f", {self.dlq_count} DLQ" if self.dlq_count else ""
404
+ pf_str = f", {pf} non-retryable" if pf else ""
405
+ eta_str = self._eta_string()
406
+ elapsed_min = self.elapsed_s / 60
407
+ print(
408
+ f"\r [{self.terminal}/{self._total_expected}] "
409
+ f"{self.completed} done, {self.failed} retrying{dlq_str}{pf_str}, "
410
+ f"~{self.in_flight} in-flight{eta_str} [{elapsed_min:.0f}m] | ~${self._cost_estimate():.2f}",
411
+ end="",
412
+ flush=True,
413
+ )
414
+
415
+ def print_summary(self) -> None:
416
+ """Print final summary."""
417
+ pf = self.permanent_failure_count
418
+ dlq_str = (
419
+ f", {self.dlq_count} permanently failed (DLQ)" if self.dlq_count else ""
420
+ )
421
+ pf_str = f", {pf} non-retryable" if pf else ""
422
+ elapsed_min = self.elapsed_s / 60
423
+ print(
424
+ f"\n\nFleet complete: {self.completed} succeeded, {self.failed} failed"
425
+ f"{dlq_str}{pf_str}"
426
+ f" out of {self._total_expected} ({elapsed_min:.1f}m total)"
427
+ )
428
+ self._stream.event(
429
+ "fleet_summary",
430
+ completed=self.completed,
431
+ failed=self.failed,
432
+ permanent_failure=pf,
433
+ dlq=self.dlq_count,
434
+ total_expected=self._total_expected,
435
+ elapsed_s=round(self.elapsed_s, 1),
436
+ )
437
+ cost = self._cost_estimate()
438
+ print(
439
+ f" EC2 cost: ~${cost:.2f} "
440
+ f"(~{cost / self._price_per_hour:.1f} {self._instance_type} worker-hours "
441
+ f"@ ${self._price_per_hour:.3f}/hr)"
442
+ )
443
+ if self._first_worker_time is not None:
444
+ print(
445
+ f" Time to first worker: {self._first_worker_time - self._start:.0f}s"
446
+ )
447
+ if self._first_completion_time is not None:
448
+ print(
449
+ f" Time to first completion: {self._first_completion_time - self._start:.0f}s"
450
+ )
451
+ if self.durations:
452
+ avg = sum(self.durations) / len(self.durations)
453
+ mn, mx = min(self.durations), max(self.durations)
454
+ print(f" Task duration: avg={avg:.0f}s, min={mn:.0f}s, max={mx:.0f}s")
455
+ if self.upload_samples:
456
+ total_eval = sum(eval_s for eval_s, _ in self.upload_samples)
457
+ total_upload = sum(upload_s for _, upload_s in self.upload_samples)
458
+ avg_upload = total_upload / len(self.upload_samples)
459
+ split = ""
460
+ if total_eval + total_upload > 0:
461
+ upload_pct = total_upload / (total_eval + total_upload) * 100
462
+ split = f", {upload_pct:.1f}% of worker time on upload"
463
+ print(
464
+ f" Result upload: avg={avg_upload:.1f}s/job{split} "
465
+ "(rest is eval exec: LLM + docker)"
466
+ )
467
+ if self.total_retries:
468
+ print(
469
+ f" API retries: {self.total_retries} across "
470
+ f"{self.retried_jobs} jobs (rate-limit + transient)"
471
+ )
472
+
473
+ def print_failure_report(self) -> None:
474
+ """Print top errors for abort decision."""
475
+ total_seen = len(self.seen_jobs)
476
+ print(
477
+ f"\n\nHigh failure rate: {self.failed}/{total_seen} "
478
+ f"({self.failed / max(total_seen, 1):.0%})"
479
+ f", {self.repeated_failure_count} persistent"
480
+ )
481
+ print("Top errors:")
482
+ for err in self.error_summaries[-5:]:
483
+ print(f" - {err[:200]}")
484
+
485
+ # -- Internal ----------------------------------------------------------
486
+
487
+ def _handle_message(self, msg: dict) -> bool:
488
+ """Process a single SQS message. Returns True if it's real progress."""
489
+ data = json.loads(msg["Body"])
490
+
491
+ # Track instance activity for all message types
492
+ if iid := data.get("instance_id"):
493
+ self.instance_last_seen[iid] = time.monotonic()
494
+
495
+ # Lifecycle messages: informational, never count as progress
496
+ if "type" in data:
497
+ if data.get("type") == "sample_done":
498
+ sample_done = SampleDoneMessage.model_validate(data)
499
+ # Skip stale sample_done messages for jobs not in the current pending set
500
+ if sample_done.job_id not in self._pending_job_ids:
501
+ return False
502
+ result_key = (sample_done.job_id, sample_done.status)
503
+ if result_key not in self._seen_sample_results:
504
+ self._stream.event(
505
+ "result",
506
+ job_id=sample_done.job_id,
507
+ status=sample_done.status,
508
+ task=sample_done.task,
509
+ model=sample_done.model,
510
+ scores=sample_done.scores,
511
+ error=sample_done.error,
512
+ duration_s=sample_done.duration_s,
513
+ instance_id=sample_done.instance_id,
514
+ )
515
+ self._seen_sample_results.add(result_key)
516
+ return False
517
+
518
+ lifecycle = LifecycleMessage.model_validate(data)
519
+ if lifecycle.type == "worker_online":
520
+ if self._first_worker_time is None:
521
+ self._first_worker_time = time.monotonic()
522
+ bootstrap_s = time.monotonic() - self._start
523
+ print(
524
+ f"\n Worker online: {lifecycle.instance_id} (+{bootstrap_s:.0f}s)"
525
+ )
526
+ self._stream.event(
527
+ "worker_online",
528
+ instance_id=lifecycle.instance_id,
529
+ bootstrap_s=round(bootstrap_s, 1),
530
+ )
531
+ elif lifecycle.type == "bootstrap_failed":
532
+ self.failed_bootstraps.add(lifecycle.instance_id)
533
+ self.instance_last_seen.pop(lifecycle.instance_id, None)
534
+ print(
535
+ f"\n Worker {lifecycle.instance_id} failed to bootstrap:\n"
536
+ + "\n".join(
537
+ f" {line}" for line in (lifecycle.detail or "").splitlines()
538
+ )
539
+ )
540
+ self._stream.event(
541
+ "bootstrap_failed",
542
+ instance_id=lifecycle.instance_id,
543
+ detail=lifecycle.detail,
544
+ )
545
+ return False
546
+
547
+ done = JobDoneMessage.model_validate(data)
548
+ job_id = done.job_id
549
+ status = done.status
550
+
551
+ # Skip stale messages for jobs not in the current pending set
552
+ if job_id not in self._pending_job_ids:
553
+ return False
554
+
555
+ # Track first completion
556
+ if status == "completed" and self._first_completion_time is None:
557
+ self._first_completion_time = time.monotonic()
558
+
559
+ prev = self.seen_jobs.get(job_id)
560
+
561
+ progress = True
562
+
563
+ if status == "permanent_failure":
564
+ if done.error_summary:
565
+ self.error_summaries.append(done.error_summary)
566
+ cat = done.failure_category or "unknown"
567
+ iid_short = done.instance_id[-8:]
568
+ logger.warning(
569
+ "Job %s permanently failed [%s] on %s: %s",
570
+ job_id,
571
+ cat,
572
+ iid_short,
573
+ done.error_summary[:200],
574
+ )
575
+ if prev is None or prev == "failed":
576
+ self.seen_jobs[job_id] = "permanent_failure"
577
+ elif prev in ("permanent_failure", "completed"):
578
+ progress = False
579
+ elif prev is None:
580
+ self.seen_jobs[job_id] = status
581
+ if status == "failed":
582
+ if done.error_summary:
583
+ self.error_summaries.append(done.error_summary)
584
+ cat = done.failure_category or "unknown"
585
+ iid_short = done.instance_id[-8:]
586
+ logger.warning(
587
+ "Job %s failed [%s] on %s: %s",
588
+ job_id,
589
+ cat,
590
+ iid_short,
591
+ done.error_summary[:200],
592
+ )
593
+ if (done.receive_count or 1) >= 3:
594
+ self._repeated_failures_by_job[job_id] += 1
595
+ if status == "completed":
596
+ if done.retries:
597
+ self.total_retries += done.retries
598
+ self.retried_jobs += 1
599
+ if done.duration_s:
600
+ self.durations.append(done.duration_s)
601
+ if done.upload_s is not None:
602
+ self.upload_samples.append((done.duration_s, done.upload_s))
603
+ elif prev in ("failed", "permanent_failure") and status == "completed":
604
+ # Late completion overrides failure/DLQ status — the job actually succeeded
605
+ self.seen_jobs[job_id] = "completed"
606
+ self._repeated_failures_by_job.pop(job_id, None)
607
+ if done.retries:
608
+ self.total_retries += done.retries
609
+ self.retried_jobs += 1
610
+ if done.duration_s:
611
+ self.durations.append(done.duration_s)
612
+ if done.upload_s is not None:
613
+ self.upload_samples.append((done.duration_s, done.upload_s))
614
+ elif prev == "failed" and status == "failed":
615
+ if (done.receive_count or 1) >= 3:
616
+ self._repeated_failures_by_job[job_id] += 1
617
+ if done.error_summary:
618
+ self.error_summaries.append(done.error_summary)
619
+ else:
620
+ progress = False
621
+
622
+ # Log job_done event only if this represents real progress (deduplication)
623
+ if progress:
624
+ self._stream.event(
625
+ "job_done",
626
+ job_id=job_id,
627
+ status=status,
628
+ instance_id=done.instance_id,
629
+ duration_s=done.duration_s,
630
+ failure_category=done.failure_category,
631
+ receive_count=done.receive_count,
632
+ error_summary=done.error_summary,
633
+ )
634
+
635
+ return progress
636
+
637
+ def _check_queue_stats(self) -> None:
638
+ """Periodically check DLQ and work queue stats (every 60s)."""
639
+ now = time.monotonic()
640
+
641
+ if now - self._last_dlq_check > 60:
642
+ self._last_dlq_check = now
643
+ try:
644
+ attrs = self._sqs.get_queue_attributes(
645
+ QueueUrl=self._dlq_url,
646
+ AttributeNames=["ApproximateNumberOfMessages"],
647
+ )
648
+ self.dlq_count = int(
649
+ attrs["Attributes"].get("ApproximateNumberOfMessages", 0)
650
+ )
651
+ except Exception:
652
+ logger.warning(
653
+ "Failed to read DLQ depth; dashboard count is stale",
654
+ exc_info=True,
655
+ )
656
+
657
+ if now - self._last_queue_check > 60:
658
+ self._last_queue_check = now
659
+ try:
660
+ attrs = self._sqs.get_queue_attributes(
661
+ QueueUrl=self._work_queue_url,
662
+ AttributeNames=[
663
+ "ApproximateNumberOfMessages",
664
+ "ApproximateNumberOfMessagesNotVisible",
665
+ ],
666
+ )
667
+ q_available = attrs["Attributes"].get(
668
+ "ApproximateNumberOfMessages", "?"
669
+ )
670
+ q_inflight = attrs["Attributes"].get(
671
+ "ApproximateNumberOfMessagesNotVisible", "?"
672
+ )
673
+ print(
674
+ f"\n [queue] available={q_available}, in-flight={q_inflight}",
675
+ end="",
676
+ flush=True,
677
+ )
678
+ except Exception:
679
+ logger.warning(
680
+ "Failed to read work-queue depth; dashboard count is stale",
681
+ exc_info=True,
682
+ )
683
+
684
+ quiet = self._quiet_workers()
685
+ if quiet:
686
+ print(
687
+ f"\n [quiet] {len(quiet)} worker(s) silent >5m: "
688
+ + ", ".join(i[-8:] for i in quiet),
689
+ end="",
690
+ flush=True,
691
+ )
692
+
693
+ def _reconcile_dlq(self) -> None:
694
+ """Read DLQ message bodies and mark those job_ids as permanently failed."""
695
+ now = time.monotonic()
696
+ if now - self._last_dlq_reconcile < DLQ_RECONCILE_INTERVAL_S:
697
+ return
698
+ self._last_dlq_reconcile = now
699
+ while True:
700
+ try:
701
+ resp = self._sqs.receive_message(
702
+ QueueUrl=self._dlq_url,
703
+ MaxNumberOfMessages=10,
704
+ WaitTimeSeconds=0,
705
+ )
706
+ except Exception:
707
+ logger.debug("Failed to read DLQ", exc_info=True)
708
+ return
709
+ msgs = resp.get("Messages", [])
710
+ if not msgs:
711
+ return
712
+ for msg in msgs:
713
+ try:
714
+ data = json.loads(msg["Body"])
715
+ job_id = data.get("job_id", "unknown")
716
+ prev = self.seen_jobs.get(job_id)
717
+ if prev != "completed" and prev != "permanent_failure":
718
+ self.seen_jobs[job_id] = "permanent_failure"
719
+ logger.warning(
720
+ "Job %s found in DLQ — marking permanent failure", job_id
721
+ )
722
+ except Exception:
723
+ # Can't extract job_id — use receipt handle as a synthetic key
724
+ # to ensure this counts toward terminalization
725
+ synthetic_id = f"dlq-unknown-{msg['ReceiptHandle'][:16]}"
726
+ if synthetic_id not in self.seen_jobs:
727
+ self.seen_jobs[synthetic_id] = "permanent_failure"
728
+ logger.warning(
729
+ "Unparseable DLQ message %s, marking as permanent failure",
730
+ synthetic_id,
731
+ )
732
+ try:
733
+ self._sqs.delete_message(
734
+ QueueUrl=self._dlq_url, ReceiptHandle=msg["ReceiptHandle"]
735
+ )
736
+ except Exception:
737
+ logger.debug("Failed to delete DLQ message", exc_info=True)
738
+
739
+ def _quiet_workers(self) -> list[str]:
740
+ """Return instance IDs that haven't sent a message in 5+ minutes."""
741
+ if not self.instance_last_seen:
742
+ return []
743
+ now = time.monotonic()
744
+ return [iid for iid, t in self.instance_last_seen.items() if now - t > 300]
745
+
746
+ def _delete_done_messages(self, messages: list[dict]) -> None:
747
+ """Delete one received batch from the done queue (best-effort, idempotent processing)."""
748
+ try:
749
+ self._sqs.delete_message_batch(
750
+ QueueUrl=self._done_queue_url,
751
+ Entries=[
752
+ {"Id": str(i), "ReceiptHandle": msg["ReceiptHandle"]}
753
+ for i, msg in enumerate(messages)
754
+ ],
755
+ )
756
+ except Exception:
757
+ logger.debug(
758
+ "Failed to delete done messages (will be re-delivered)", exc_info=True
759
+ )
760
+
761
+ def _active_worker_count(self) -> int:
762
+ """Count workers heard from in the last 5 minutes."""
763
+ if not self.instance_last_seen:
764
+ return len(self._expected_instances)
765
+ now = time.monotonic()
766
+ active = sum(1 for t in self.instance_last_seen.values() if now - t < 300)
767
+ return active or len(self._expected_instances)
768
+
769
+ def _eta_string(self) -> str:
770
+ """Compute ETA string from recent durations (needs 3+ completions)."""
771
+ if len(self.durations) < 3:
772
+ return ""
773
+ recent = self.durations[-20:]
774
+ avg_dur = sum(recent) / len(recent)
775
+ active_slots = self._active_worker_count() * self._concurrency_per_worker
776
+ eta_min = self.in_flight * avg_dur / 60 / max(active_slots, 1)
777
+ return f" | est ~{eta_min:.0f}m left"
778
+
779
+
780
+ class FleetController:
781
+ """Orchestrates a distributed eval run across EC2 workers."""
782
+
783
+ def __init__(self, config: Ec2Config):
784
+ self.config = config
785
+ self._worker_setup_s3_key = ""
786
+ env = EnvVar()
787
+ self.s3_bucket: str = env.fleet_s3_bucket
788
+
789
+ if env.aws_profile and env.aws_access_key_id:
790
+ raise ValidationError(
791
+ "Both AWS_PROFILE and AWS_ACCESS_KEY_ID are set — the static key would silently "
792
+ "take precedence over the profile. Unset one (usually: remove the "
793
+ "AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY lines from .env when using SSO)."
794
+ )
795
+ self._ec2 = boto3.client("ec2", region_name=config.region)
796
+ self._pricing = boto3.client("pricing", region_name="us-east-1")
797
+ self._sqs = boto3.client(
798
+ "sqs",
799
+ region_name=config.region,
800
+ config=botocore.config.Config(max_pool_connections=16),
801
+ )
802
+ self._s3 = boto3.client(
803
+ "s3",
804
+ region_name=config.region,
805
+ config=botocore.config.Config(max_pool_connections=64),
806
+ )
807
+ # Each instance costs 1 + one-per-setting-package PutParameter calls; standard mode backs off on throttling.
808
+ self._ssm = boto3.client(
809
+ "ssm",
810
+ region_name=config.region,
811
+ config=botocore.config.Config(
812
+ retries={"mode": "standard", "max_attempts": 10}
813
+ ),
814
+ )
815
+
816
+ def run(
817
+ self,
818
+ *,
819
+ inspect_jobs: list[InspectFleetJobWire],
820
+ eval_args: dict[str, Any],
821
+ setting_dirs: dict[str, Path],
822
+ env_dirs: list[Path],
823
+ shipped_files: Collection[Path],
824
+ extra_env_var_names: set[str],
825
+ setting_package_secrets: dict[str, str],
826
+ metadata: dict[str, Any] | None = None,
827
+ ) -> FleetRunResult | None:
828
+ """Run generic Inspect jobs on workers given the local checkouts of ``setting_dirs`` and of the repos holding ``env_dirs``, which may leave untracked only the ``shipped_files`` that reach workers another way; return the S3 result location."""
829
+ run_start = time.monotonic()
830
+ for package, values in setting_package_secrets.items():
831
+ if not re.fullmatch(r"[A-Za-z0-9_.-]+", package):
832
+ raise ValidationError(f"Invalid setting package name: {package!r}")
833
+ if len(values.encode()) > 4096:
834
+ raise ValidationError(
835
+ f"Setting {package}'s {package}.env exceeds the 4 KB SSM parameter limit."
836
+ )
837
+
838
+ epoch_by_eval_name = {
839
+ job.output_key: epoch_index
840
+ for job in inspect_jobs
841
+ if isinstance(epoch_index := job.metadata.get("epoch_index"), int)
842
+ }
843
+
844
+ eval_args_wire = EvalConfig.model_validate(eval_args)
845
+ longest_job_s = SQS_MAX_VISIBILITY_S - SQS_VISIBILITY_TIMEOUT_S
846
+ if (eval_args_wire.time_limit or 0) > longest_job_s:
847
+ raise ValidationError(
848
+ f"`time_limit` is {eval_args_wire.time_limit}s, but a fleet job can run for at most {longest_job_s}s."
849
+ )
850
+ eval_args_dict = eval_args_wire.model_dump(mode="json", exclude_none=True)
851
+
852
+ # 1. Compute fleet identity
853
+
854
+ code_root = _code_root()
855
+ git_sha, _ = _git_snapshot(code_root)
856
+ setting_paths = setting_repo_dirs(setting_dirs, env_dirs)
857
+ snapshots = {dest: _git_snapshot(repo) for dest, repo in setting_paths.items()}
858
+ setting_repos = {
859
+ dest: f"{sha}+dirty" if dirty else sha
860
+ for dest, (sha, dirty) in snapshots.items()
861
+ }
862
+ shipped_dirty = {
863
+ setting_paths[dest]: dirty for dest, (_, dirty) in snapshots.items()
864
+ }
865
+ trees = {
866
+ label: subprocess.run(
867
+ ["git", "rev-parse", f"{sha}^{{tree}}"],
868
+ cwd=path,
869
+ capture_output=True,
870
+ text=True,
871
+ check=True,
872
+ ).stdout.strip()
873
+ for label, path, sha in [
874
+ ("control-tower", code_root, git_sha),
875
+ *((dest, setting_paths[dest], s) for dest, (s, _) in snapshots.items()),
876
+ ]
877
+ }
878
+ for label, path in [("control-tower", code_root), *setting_paths.items()]:
879
+ changed = {
880
+ file: tracked
881
+ for file, tracked in (git_changed_paths(path) or {}).items()
882
+ if not any(
883
+ file.is_relative_to(repo)
884
+ for repo in setting_paths.values()
885
+ if path in repo.parents
886
+ )
887
+ }
888
+ modified = [file for file, tracked in changed.items() if tracked]
889
+ unshipped = [
890
+ file
891
+ for file, tracked in changed.items()
892
+ if not tracked and file.resolve() not in shipped_files
893
+ ]
894
+ if modified:
895
+ print(
896
+ f" WARNING: {label}: shipping {len(modified)} uncommitted file(s) from local worktree"
897
+ )
898
+ if unshipped:
899
+ raise ValidationError(
900
+ f"{label} has {len(unshipped)} untracked file(s), which workers would not get: "
901
+ "`git add` them to ship them, or gitignore or remove them. "
902
+ f"First: {unshipped[0].relative_to(path)}"
903
+ )
904
+
905
+ worker_setup_script = b""
906
+ if self.config.worker_setup:
907
+ try:
908
+ worker_setup_script = Path(self.config.worker_setup).read_bytes()
909
+ except OSError as exc:
910
+ raise ValidationError(
911
+ f"--ec2-arg worker_setup={self.config.worker_setup!r} is not a readable file: {exc}"
912
+ ) from exc
913
+ worker_setup_hash = (
914
+ hashlib.sha256(worker_setup_script).hexdigest()[:16]
915
+ if worker_setup_script
916
+ else ""
917
+ )
918
+ worker_bootstrap: dict[str, Any] = {
919
+ "settings": {dest: trees[dest] for dest in snapshots},
920
+ "instance_type": self.config.instance_type,
921
+ **({"worker_setup": worker_setup_hash} if worker_setup_hash else {}),
922
+ **({"ami": self.config.ami} if self.config.ami else {}),
923
+ **(
924
+ {"worker_user": self.config.worker_user}
925
+ if self.config.worker_user != "root"
926
+ else {}
927
+ ),
928
+ }
929
+ fleet_id = compute_fleet_id(
930
+ trees["control-tower"],
931
+ inspect_jobs,
932
+ eval_args_dict,
933
+ new_fleet=self.config.new_fleet,
934
+ worker_bootstrap=worker_bootstrap,
935
+ )
936
+ s3_prefix = f"{fleet_id}/"
937
+
938
+ # 2. Build job list
939
+ jobs = [
940
+ JobPayload(
941
+ job_id=inspect_job.job_id, fleet_id=fleet_id, inspect_job=inspect_job
942
+ )
943
+ for inspect_job in inspect_jobs
944
+ ]
945
+ total_jobs = len(jobs) # Immutable: total across all invocations
946
+
947
+ if total_jobs == 0:
948
+ print("No jobs to run.")
949
+ return None
950
+
951
+ jobs = self._with_output_prefix(jobs, s3_prefix)
952
+ inspect_jobs = [job.inspect_job for job in jobs]
953
+
954
+ # 3. Check for existing fleet (resume)
955
+ existing_manifest = self._get_manifest(s3_prefix)
956
+ is_resume = existing_manifest is not None and not self.config.new_fleet
957
+ done_keys = set(list_eval_keys(self._s3, self.s3_bucket, s3_prefix))
958
+ pending_jobs = [j for j in jobs if j.inspect_job.output_key not in done_keys]
959
+ if existing_manifest and not self.config.new_fleet:
960
+ print(
961
+ f"Resuming existing fleet {fleet_id} ({existing_manifest.total_jobs} jobs)"
962
+ )
963
+ print(
964
+ f" {total_jobs - len(pending_jobs)} done, {len(pending_jobs)} remaining"
965
+ )
966
+ elif existing_manifest:
967
+ logger.info(
968
+ "Existing manifest at s3://%s/%smanifest.json found but "
969
+ "--new-fleet was set; starting fresh fleet %s.",
970
+ self.s3_bucket,
971
+ s3_prefix,
972
+ fleet_id,
973
+ )
974
+ else:
975
+ logger.info(
976
+ "No existing manifest at s3://%s/%smanifest.json; "
977
+ "starting fresh fleet %s.",
978
+ self.s3_bucket,
979
+ s3_prefix,
980
+ fleet_id,
981
+ )
982
+ if not pending_jobs:
983
+ print(f"All {total_jobs} jobs already complete.")
984
+ try:
985
+ return self._collect_fleet_result(
986
+ fleet_id,
987
+ s3_prefix,
988
+ total_jobs,
989
+ epoch_by_eval_name=epoch_by_eval_name,
990
+ setting_repos=setting_repos,
991
+ shipped_dirty=shipped_dirty,
992
+ )
993
+ finally:
994
+ self._cleanup_fleet_resources(fleet_id, wait_for_workers=False)
995
+
996
+ instance_types = list(
997
+ dict.fromkeys(
998
+ [self.config.instance_type, *self.config.instance_type_fallbacks]
999
+ )
1000
+ )
1001
+ gpu_types = [
1002
+ info["InstanceType"]
1003
+ for info in self._ec2.describe_instance_types(InstanceTypes=instance_types)[
1004
+ "InstanceTypes"
1005
+ ]
1006
+ if "GpuInfo" in info
1007
+ ]
1008
+ if gpu_types and self.config.concurrency_per_worker != 1:
1009
+ raise ValidationError(
1010
+ f"A {', '.join(gpu_types)} worker has a GPU and runs one sample at a time: "
1011
+ "pass `--ec2-arg concurrency_per_worker=1`."
1012
+ )
1013
+ prices: dict[str, float] = {}
1014
+ for instance_type in instance_types:
1015
+ [product] = self._pricing.get_products(
1016
+ ServiceCode="AmazonEC2",
1017
+ Filters=[
1018
+ {"Type": "TERM_MATCH", "Field": field, "Value": value}
1019
+ for field, value in {
1020
+ "instanceType": instance_type,
1021
+ "regionCode": self.config.region,
1022
+ "operatingSystem": "Linux",
1023
+ "tenancy": "Shared",
1024
+ "preInstalledSw": "NA",
1025
+ "capacitystatus": "Used",
1026
+ "licenseModel": "No License required",
1027
+ }.items()
1028
+ ],
1029
+ )["PriceList"]
1030
+ [on_demand] = json.loads(product)["terms"]["OnDemand"].values()
1031
+ [dimension] = on_demand["priceDimensions"].values()
1032
+ prices[instance_type] = float(dimension["pricePerUnit"]["USD"])
1033
+
1034
+ # 4. Size the fleet so every pending job starts at once, up to max_workers
1035
+ planning_cpw = min(
1036
+ self.config.concurrency_per_worker, FLEET_PLANNING_CONCURRENCY
1037
+ )
1038
+ num_workers = max(
1039
+ 1,
1040
+ min(self.config.max_workers, math.ceil(len(pending_jobs) / planning_cpw)),
1041
+ )
1042
+
1043
+ # 5. Print cost estimate (for pending work only)
1044
+ fleet_cost_per_hour = num_workers * max(prices.values())
1045
+ self._print_estimate(len(pending_jobs), num_workers, fleet_cost_per_hour)
1046
+
1047
+ if self.config.estimate:
1048
+ return None
1049
+
1050
+ # 6. Wait for confirmation
1051
+ if not self.config.auto_confirm:
1052
+ self._wait_for_confirm()
1053
+
1054
+ stream_log: FleetStreamingLog | None = None
1055
+ wait_for_workers = is_resume
1056
+ try:
1057
+ # 7. Create/reuse queues
1058
+ t0 = time.monotonic()
1059
+ work_queue_url, work_created = self._ensure_queue(
1060
+ f"lb-fleet-{fleet_id}-work",
1061
+ VisibilityTimeout=str(SQS_VISIBILITY_TIMEOUT_S),
1062
+ ReceiveMessageWaitTimeSeconds="20",
1063
+ )
1064
+ done_queue_url, done_created = self._ensure_queue(
1065
+ f"lb-fleet-{fleet_id}-done",
1066
+ ReceiveMessageWaitTimeSeconds="20",
1067
+ )
1068
+ dlq_url, dlq_created = self._ensure_queue(f"lb-fleet-{fleet_id}-dlq")
1069
+ self._configure_dlq(work_queue_url, dlq_url)
1070
+ print(f" Queues ready ({time.monotonic() - t0:.1f}s)")
1071
+
1072
+ # 8. Write manifest
1073
+ if not existing_manifest:
1074
+ manifest = FleetManifest(
1075
+ fleet_id=fleet_id,
1076
+ total_jobs=len(jobs),
1077
+ runner_schema_version=FLEET_RUNNER_SCHEMA_VERSION,
1078
+ inspect_jobs=inspect_jobs,
1079
+ eval_args=eval_args_wire,
1080
+ metadata=metadata or {},
1081
+ git_sha=git_sha,
1082
+ setting_repos=setting_repos,
1083
+ created_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
1084
+ )
1085
+ self._put_manifest(manifest, s3_prefix)
1086
+
1087
+ # 9. Enqueue the pending jobs
1088
+ purges: list[float] = []
1089
+ if not work_created:
1090
+ purges.append(self._purge_queue(work_queue_url, "work"))
1091
+ if not done_created and not is_resume:
1092
+ purges.append(self._purge_queue(done_queue_url, "done"))
1093
+ if not dlq_created:
1094
+ purges.append(self._purge_queue(dlq_url, "DLQ"))
1095
+
1096
+ last_purge = max(purges, default=None)
1097
+ if last_purge is not None:
1098
+ remaining = SQS_PURGE_WINDOW_S - (time.monotonic() - last_purge)
1099
+ if remaining > 0:
1100
+ print(f"Waiting {remaining:.0f}s for the queue purges to settle...")
1101
+ time.sleep(remaining)
1102
+
1103
+ t0 = time.monotonic()
1104
+ print(f"Enqueueing {len(pending_jobs)} jobs...")
1105
+ enqueue_jobs(self._sqs, work_queue_url, pending_jobs)
1106
+ print(f" Enqueued {len(pending_jobs)} ({time.monotonic() - t0:.1f}s)")
1107
+
1108
+ # 10. Upload code artifact to S3
1109
+ t0 = time.monotonic()
1110
+ print("Uploading code artifact to S3...")
1111
+ code_s3_key = self._upload_git_artifact(code_root, git_sha)
1112
+ print(f" Code artifact ready ({time.monotonic() - t0:.1f}s)")
1113
+
1114
+ setting_archives: list[tuple[str, str]] = []
1115
+ if setting_repos:
1116
+ t0 = time.monotonic()
1117
+ print(f"Uploading {len(setting_repos)} setting repo(s) to S3...")
1118
+ setting_archives = [
1119
+ (self._upload_git_artifact(setting_paths[dest], sha), dest)
1120
+ for dest, (sha, _) in snapshots.items()
1121
+ ]
1122
+ print(f" Setting artifacts ready ({time.monotonic() - t0:.1f}s)")
1123
+
1124
+ rates_s3_key = (
1125
+ "" if EnvVar().pricing == "off" else self._upload_rate_table()
1126
+ )
1127
+ if worker_setup_hash:
1128
+ self._worker_setup_s3_key = f"worker-setup/{worker_setup_hash}.sh"
1129
+ self._s3.put_object(
1130
+ Bucket=self.s3_bucket,
1131
+ Key=self._worker_setup_s3_key,
1132
+ Body=worker_setup_script,
1133
+ )
1134
+
1135
+ # 11. Launch workers
1136
+ t0 = time.monotonic()
1137
+ instance_ids, fleet_price_per_hour = self._launch_workers(
1138
+ num_workers,
1139
+ prices=prices,
1140
+ fleet_id=fleet_id,
1141
+ user_data=self._build_user_data(
1142
+ code_s3_key=code_s3_key,
1143
+ rates_s3_key=rates_s3_key,
1144
+ s3_prefix=s3_prefix,
1145
+ work_queue_url=work_queue_url,
1146
+ done_queue_url=done_queue_url,
1147
+ setting_archives=setting_archives,
1148
+ env_secret_packages=sorted(setting_package_secrets),
1149
+ fleet_secrets_prefix=f"{FLEET_SSM_PREFIX}/{fleet_id}",
1150
+ ),
1151
+ secrets=self._worker_secrets(extra_env_var_names),
1152
+ setting_package_secrets=setting_package_secrets,
1153
+ )
1154
+ wait_for_workers = True
1155
+ print(
1156
+ f"Launched {len(instance_ids)} workers ({time.monotonic() - t0:.1f}s): "
1157
+ f"{instance_ids}"
1158
+ )
1159
+
1160
+ setup_s = time.monotonic() - run_start
1161
+ print(f"\n Setup complete in {setup_s:.1f}s")
1162
+
1163
+ # 12. Monitor progress
1164
+ # Scale the stale-abort timeout based on time_limit so that
1165
+ # individual sample timeouts have a chance to fire before the fleet
1166
+ # controller gives up. Default: max(STALE_ABORT_S, time_limit * 1.5).
1167
+ sample_tl = eval_args_wire.time_limit
1168
+ if sample_tl is not None and sample_tl > 0:
1169
+ stale_abort_s = max(STALE_ABORT_S, int(sample_tl * 1.5))
1170
+ else:
1171
+ stale_abort_s = STALE_ABORT_S
1172
+ if stale_abort_s != STALE_ABORT_S:
1173
+ logger.info(
1174
+ "Stale-abort timeout scaled to %ds (time_limit=%ds)",
1175
+ stale_abort_s,
1176
+ sample_tl,
1177
+ )
1178
+ stream_log_dir = Path("logs/fleet") / fleet_id
1179
+ stream_log = FleetStreamingLog(log_dir=stream_log_dir, resume=is_resume)
1180
+ monitor = ProgressMonitor(
1181
+ sqs_client=self._sqs,
1182
+ done_queue_url=done_queue_url,
1183
+ work_queue_url=work_queue_url,
1184
+ dlq_url=dlq_url,
1185
+ instance_ids=instance_ids,
1186
+ concurrency_per_worker=planning_cpw,
1187
+ price_per_hour=fleet_price_per_hour / len(instance_ids),
1188
+ instance_type="/".join(instance_types),
1189
+ stale_abort_s=stale_abort_s,
1190
+ stream_log=stream_log,
1191
+ pending_job_ids={j.job_id for j in pending_jobs},
1192
+ initial_seen_results=stream_log.load_seen_results(),
1193
+ )
1194
+ print(f"\nMonitoring fleet {fleet_id}: 0/{len(pending_jobs)}")
1195
+ print(f" Streaming: tail -f {stream_log.path}")
1196
+ stream_log.event(
1197
+ "fleet_start",
1198
+ fleet_id=fleet_id,
1199
+ expected=len(pending_jobs),
1200
+ workers=len(instance_ids),
1201
+ instance_type=self.config.instance_type,
1202
+ )
1203
+
1204
+ self._run_monitor_loop(monitor, fleet_id, instance_ids)
1205
+ monitor.print_summary()
1206
+ return self._collect_fleet_result(
1207
+ fleet_id,
1208
+ s3_prefix,
1209
+ total_jobs,
1210
+ epoch_by_eval_name=epoch_by_eval_name,
1211
+ setting_repos=setting_repos,
1212
+ shipped_dirty=shipped_dirty,
1213
+ )
1214
+ finally:
1215
+ if stream_log is not None:
1216
+ stream_log.close()
1217
+ # 13. Tear down fleet resources. S3 results remain for the caller to consume.
1218
+ self._cleanup_fleet_resources(fleet_id, wait_for_workers=wait_for_workers)
1219
+
1220
+ def _collect_fleet_result(
1221
+ self,
1222
+ fleet_id: str,
1223
+ s3_prefix: str,
1224
+ total_expected: int,
1225
+ *,
1226
+ epoch_by_eval_name: dict[str, int] | None = None,
1227
+ setting_repos: dict[str, str],
1228
+ shipped_dirty: dict[Path, bool],
1229
+ ) -> FleetRunResult:
1230
+ job_order = {name: index for index, name in enumerate(epoch_by_eval_name or {})}
1231
+ eval_keys = sorted(
1232
+ list_eval_keys(self._s3, self.s3_bucket, s3_prefix),
1233
+ key=lambda key: (job_order.get(Path(key).name, len(job_order)), key),
1234
+ )
1235
+ print(f"\nResults: found {len(eval_keys)}/{total_expected} .eval files in S3")
1236
+ eval_log = (
1237
+ self._download_and_merge_eval_logs(eval_keys, epoch_by_eval_name or {})
1238
+ if eval_keys
1239
+ else None
1240
+ )
1241
+ completed = len(eval_log.samples or []) if eval_log is not None else 0
1242
+ if completed < total_expected:
1243
+ print(
1244
+ f" Warning: {total_expected - completed}/{total_expected} jobs "
1245
+ "produced no sample (may have failed)"
1246
+ )
1247
+
1248
+ return FleetRunResult(
1249
+ fleet_id=fleet_id,
1250
+ s3_bucket=self.s3_bucket,
1251
+ s3_prefix=s3_prefix,
1252
+ total_expected=total_expected,
1253
+ total_completed=completed,
1254
+ region=self.config.region,
1255
+ eval_log=eval_log,
1256
+ setting_repos=setting_repos,
1257
+ shipped_dirty=shipped_dirty,
1258
+ )
1259
+
1260
+ def _download_and_merge_eval_logs(
1261
+ self, eval_keys: list[str], epoch_by_eval_name: dict[str, int]
1262
+ ) -> EvalLog | None:
1263
+ with tempfile.TemporaryDirectory(prefix="lb-fleet-merge-") as tmpdir:
1264
+ tmppath = Path(tmpdir)
1265
+ print(f"Downloading {len(eval_keys)} .eval files...")
1266
+ paths_by_key: dict[str, Path] = {}
1267
+
1268
+ max_workers = min(64, len(eval_keys))
1269
+ with ThreadPoolExecutor(max_workers=max_workers) as pool:
1270
+ futures = {}
1271
+ for idx, key in enumerate(eval_keys):
1272
+ local_path = tmppath / f"{idx}-{Path(key).name}"
1273
+ future = pool.submit(
1274
+ self._s3.download_file,
1275
+ self.s3_bucket,
1276
+ key,
1277
+ str(local_path),
1278
+ )
1279
+ futures[future] = (key, local_path)
1280
+ for future in as_completed(futures):
1281
+ key, path = futures[future]
1282
+ try:
1283
+ future.result()
1284
+ paths_by_key[key] = path
1285
+ except Exception as e:
1286
+ logger.error("Failed to download %s: %s", key, e)
1287
+
1288
+ present_keys = [key for key in eval_keys if key in paths_by_key]
1289
+ eval_paths = [paths_by_key[key] for key in present_keys]
1290
+ print(f" Downloaded {len(eval_paths)} files")
1291
+ eval_logs = list(pool.map(read_eval_log, map(str, eval_paths)))
1292
+
1293
+ # Fleet epoch jobs all emit epoch 1 with the same deterministic sample
1294
+ # id; stamp the job's true epoch so merged (id, epoch) keys are stable
1295
+ # across re-merges.
1296
+ for key, log in zip(present_keys, eval_logs, strict=True):
1297
+ epoch_index = epoch_by_eval_name.get(Path(key).name)
1298
+ if epoch_index is not None:
1299
+ for sample in log.samples or []:
1300
+ sample.epoch = epoch_index + 1
1301
+
1302
+ combined = merge_eval_logs(eval_logs)
1303
+ if combined is not None:
1304
+ print(f" Merged {len(combined.samples or [])} eval samples")
1305
+ return combined
1306
+
1307
+ def _run_monitor_loop(
1308
+ self,
1309
+ monitor: ProgressMonitor,
1310
+ fleet_id: str,
1311
+ instance_ids: list[str],
1312
+ ) -> None:
1313
+ """Drive the progress monitor until completion or abort."""
1314
+ last_liveness_check = time.monotonic()
1315
+ try:
1316
+ while not monitor.is_done:
1317
+ action = monitor.check_stale(instance_ids)
1318
+ if action is None and time.monotonic() - last_liveness_check > 60:
1319
+ last_liveness_check = time.monotonic()
1320
+ resp = self._ec2.describe_instances(
1321
+ Filters=[
1322
+ {"Name": "tag:FleetId", "Values": [fleet_id]},
1323
+ {
1324
+ "Name": "instance-state-name",
1325
+ "Values": ["pending", "running"],
1326
+ },
1327
+ ]
1328
+ )
1329
+ if not any(res["Instances"] for res in resp["Reservations"]):
1330
+ print(
1331
+ f"\n\n Every worker is gone: none of the {len(instance_ids)} "
1332
+ "launched instances is running. Bootstrap logs, if any: "
1333
+ f"s3://{self.s3_bucket}/{fleet_id}/logs/"
1334
+ )
1335
+ action = "abort"
1336
+ if action == "abort":
1337
+ print(" Terminating fleet and aborting.")
1338
+ self._terminate_fleet_instances(fleet_id)
1339
+ raise OperationError(
1340
+ f"Fleet {fleet_id} aborted due to lack of progress. "
1341
+ "S3 results preserved. Fix the issue and re-run to resume."
1342
+ )
1343
+
1344
+ # Poll and process
1345
+ monitor.poll_and_process()
1346
+ monitor.print_progress()
1347
+
1348
+ # Check failure rate
1349
+ if monitor.check_failure_rate():
1350
+ monitor.print_failure_report()
1351
+ print("\nTerminating workers to stop the bleeding...")
1352
+ self._terminate_fleet_instances(fleet_id)
1353
+ raise OperationError(
1354
+ f"Fleet {fleet_id} stopped on its failure rate. "
1355
+ "S3 results preserved. Fix the issue and re-run to resume."
1356
+ )
1357
+ except KeyboardInterrupt:
1358
+ # Best effort: stop workers from pulling new jobs so they can drain
1359
+ # in-flight work and self-terminate on idle.
1360
+ purge_failure: botocore.exceptions.ClientError | None = None
1361
+ try:
1362
+ self._purge_queue(monitor._work_queue_url, "work")
1363
+ except self._sqs.exceptions.PurgeQueueInProgress:
1364
+ logger.info("Work queue purge already in progress")
1365
+ except botocore.exceptions.ClientError as exc:
1366
+ purge_failure = exc
1367
+ monitor.print_summary()
1368
+ queue_state = (
1369
+ "Work queue purged; workers will self-terminate after idle timeout."
1370
+ if purge_failure is None
1371
+ else f"Work queue purge FAILED ({purge_failure}); workers keep pulling jobs."
1372
+ )
1373
+ print(
1374
+ f"\nInterrupted. {queue_state}"
1375
+ f"\nRe-run the same command to resume, or terminate instances:"
1376
+ f"\n aws ec2 describe-instances --filters Name=tag:FleetId,Values={fleet_id} "
1377
+ f"Name=instance-state-name,Values=running --query Reservations[].Instances[].InstanceId"
1378
+ )
1379
+ raise SystemExit(1)
1380
+
1381
+ def _with_output_prefix(
1382
+ self, jobs: list[JobPayload], s3_prefix: str
1383
+ ) -> list[JobPayload]:
1384
+ """Fill each generic job's output key after the fleet ID is known."""
1385
+ updated: list[JobPayload] = []
1386
+ for job in jobs:
1387
+ base_output_key = job.inspect_job.output_key
1388
+ inspect_job = job.inspect_job.model_copy(
1389
+ update={"output_key": f"{s3_prefix}{base_output_key.lstrip('/')}"}
1390
+ )
1391
+ updated.append(job.model_copy(update={"inspect_job": inspect_job}))
1392
+ return updated
1393
+
1394
+ def _ensure_queue(self, name: str, **extra_attrs: str) -> tuple[str, bool]:
1395
+ """Create the SQS queue if it does not exist.
1396
+
1397
+ Returns the URL and whether this call created it: the only case a caller
1398
+ can skip purging on, and one AWS's message counts cannot answer.
1399
+ """
1400
+ try:
1401
+ resp = self._sqs.get_queue_url(QueueName=name)
1402
+ return resp["QueueUrl"], False
1403
+ except self._sqs.exceptions.QueueDoesNotExist:
1404
+ pass
1405
+
1406
+ attrs = {"MessageRetentionPeriod": "86400"}
1407
+ attrs.update(extra_attrs)
1408
+
1409
+ deadline = time.monotonic() + SQS_QUEUE_DELETE_COOLDOWN_S + 10
1410
+ while True:
1411
+ try:
1412
+ resp = self._sqs.create_queue(QueueName=name, Attributes=attrs)
1413
+ break
1414
+ except self._sqs.exceptions.QueueDeletedRecently:
1415
+ if time.monotonic() >= deadline:
1416
+ raise
1417
+ logger.info("Queue %s deleted recently; retrying after cooldown", name)
1418
+ time.sleep(5)
1419
+ logger.info("Created queue: %s", name)
1420
+ return resp["QueueUrl"], True
1421
+
1422
+ def _configure_dlq(self, work_queue_url: str, dlq_url: str) -> None:
1423
+ """Attach a dead-letter queue to the work queue."""
1424
+ dlq_arn = self._sqs.get_queue_attributes(
1425
+ QueueUrl=dlq_url, AttributeNames=["QueueArn"]
1426
+ )["Attributes"]["QueueArn"]
1427
+
1428
+ self._sqs.set_queue_attributes(
1429
+ QueueUrl=work_queue_url,
1430
+ Attributes={
1431
+ "RedrivePolicy": json.dumps(
1432
+ {"deadLetterTargetArn": dlq_arn, "maxReceiveCount": 3}
1433
+ )
1434
+ },
1435
+ )
1436
+
1437
+ def _delete_queues(self, fleet_id: str) -> None:
1438
+ """Clean up queues for a fleet."""
1439
+ for suffix in ("work", "done", "dlq"):
1440
+ name = f"lb-fleet-{fleet_id}-{suffix}"
1441
+ try:
1442
+ url = self._sqs.get_queue_url(QueueName=name)["QueueUrl"]
1443
+ self._sqs.delete_queue(QueueUrl=url)
1444
+ logger.info("Deleted queue: %s", name)
1445
+ except Exception:
1446
+ logger.warning(
1447
+ "Failed to delete queue %s; it may be orphaned", name, exc_info=True
1448
+ )
1449
+
1450
+ def _purge_queue(self, queue_url: str, label: str = "") -> float:
1451
+ """Purge an SQS queue.
1452
+
1453
+ Returns when the purge was accepted so callers can wait out the window
1454
+ in which SQS may still delete newly sent messages.
1455
+ """
1456
+ self._sqs.purge_queue(QueueUrl=queue_url)
1457
+ if label:
1458
+ logger.info("Purged %s queue", label)
1459
+ return time.monotonic()
1460
+
1461
+ def _put_manifest(self, manifest: FleetManifest, s3_prefix: str) -> None:
1462
+ self._s3.put_object(
1463
+ Bucket=self.s3_bucket,
1464
+ Key=f"{s3_prefix}manifest.json",
1465
+ Body=manifest.model_dump_json(indent=2),
1466
+ ContentType="application/json",
1467
+ )
1468
+
1469
+ def _get_manifest(self, s3_prefix: str) -> FleetManifest | None:
1470
+ key = f"{s3_prefix}manifest.json"
1471
+ try:
1472
+ resp = self._s3.get_object(Bucket=self.s3_bucket, Key=key)
1473
+ data = json.loads(resp["Body"].read())
1474
+ return FleetManifest.model_validate(data)
1475
+ except self._s3.exceptions.NoSuchKey:
1476
+ return None
1477
+
1478
+ def _ensure_fleet_security_group(self) -> str:
1479
+ """Create or find a no-inbound security group for fleet workers."""
1480
+ sg_name = "linuxbench-fleet-worker-sg"
1481
+
1482
+ vpcs = self._ec2.describe_vpcs(
1483
+ Filters=[{"Name": "is-default", "Values": ["true"]}]
1484
+ )["Vpcs"]
1485
+ vpc_id = vpcs[0]["VpcId"] if vpcs else None
1486
+
1487
+ filters = [{"Name": "group-name", "Values": [sg_name]}]
1488
+ if vpc_id:
1489
+ filters.append({"Name": "vpc-id", "Values": [vpc_id]})
1490
+ groups = self._ec2.describe_security_groups(Filters=filters)["SecurityGroups"]
1491
+ if groups:
1492
+ return groups[0]["GroupId"]
1493
+
1494
+ create_kwargs: dict = dict(
1495
+ GroupName=sg_name,
1496
+ Description="LinuxBench fleet workers - outbound only, no inbound",
1497
+ )
1498
+ if vpc_id:
1499
+ create_kwargs["VpcId"] = vpc_id
1500
+ resp = self._ec2.create_security_group(**create_kwargs)
1501
+ sg_id = resp["GroupId"]
1502
+ logger.info(
1503
+ "Created fleet security group: %s (%s) in VPC %s", sg_name, sg_id, vpc_id
1504
+ )
1505
+ return sg_id
1506
+
1507
+ def _upload_git_artifact(self, repo: Path, git_sha: str) -> str:
1508
+ """Upload a tarball of ``git_sha`` from ``repo``, keyed by that SHA."""
1509
+
1510
+ s3_key = f"artifacts/{git_sha}.tar.gz"
1511
+
1512
+ # Check if already uploaded (shared across fleets with same SHA)
1513
+ try:
1514
+ self._s3.head_object(Bucket=self.s3_bucket, Key=s3_key)
1515
+ except botocore.exceptions.ClientError as e:
1516
+ if e.response["Error"]["Code"] != "404":
1517
+ raise
1518
+ else:
1519
+ logger.info("Artifact already exists at s3://%s/%s", self.s3_bucket, s3_key)
1520
+ return s3_key
1521
+
1522
+ common_dir = subprocess.run(
1523
+ ["git", "rev-parse", "--path-format=absolute", "--git-common-dir"],
1524
+ cwd=repo,
1525
+ check=True,
1526
+ capture_output=True,
1527
+ text=True,
1528
+ ).stdout.strip()
1529
+ with (
1530
+ tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=True) as tmp,
1531
+ tempfile.TemporaryDirectory(prefix="ct-archive-") as git_dir,
1532
+ ):
1533
+ # $GIT_DIR/info/attributes outranks in-tree .gitattributes, so this disables export-ignore.
1534
+ subprocess.run(
1535
+ ["git", "init", "--bare", "-q", git_dir],
1536
+ check=True,
1537
+ capture_output=True,
1538
+ )
1539
+ (Path(git_dir) / "objects" / "info" / "alternates").write_text(
1540
+ f"{common_dir}/objects\n"
1541
+ )
1542
+ (Path(git_dir) / "info" / "attributes").write_text(
1543
+ "* -export-ignore -export-subst\n"
1544
+ )
1545
+ subprocess.run(
1546
+ ["git", "--git-dir", git_dir, "archive", "--format=tar.gz"]
1547
+ + ["-o", tmp.name, git_sha],
1548
+ cwd=git_dir,
1549
+ check=True,
1550
+ capture_output=True,
1551
+ )
1552
+
1553
+ self._s3.upload_file(tmp.name, self.s3_bucket, s3_key)
1554
+
1555
+ logger.info("Uploaded %s to s3://%s/%s", repo, self.s3_bucket, s3_key)
1556
+ return s3_key
1557
+
1558
+ def _upload_rate_table(self) -> str:
1559
+ """Upload the driver's rate table so workers price with the same rates."""
1560
+ table = default_rate_table().model_dump_json()
1561
+ s3_key = f"rates/{hashlib.sha256(table.encode()).hexdigest()[:16]}.json"
1562
+ self._s3.put_object(Bucket=self.s3_bucket, Key=s3_key, Body=table.encode())
1563
+ logger.info("Uploaded rate table to s3://%s/%s", self.s3_bucket, s3_key)
1564
+ return s3_key
1565
+
1566
+ def upload_recording(self, path: Path) -> str:
1567
+ """Upload a local recording for workers to replay; return the ``s3://`` URL they load it from."""
1568
+ if path.is_dir():
1569
+ raise ValidationError(
1570
+ f"--ec2 ships a recording as one file, but {path} is a directory; pass the .eval file in it."
1571
+ )
1572
+ with path.open("rb") as file:
1573
+ digest = hashlib.file_digest(file, "sha256").hexdigest()
1574
+ s3_key = f"recordings/{digest[:16]}{path.suffix}"
1575
+ if not self.config.estimate:
1576
+ self._s3.upload_file(str(path), self.s3_bucket, s3_key)
1577
+ return f"s3://{self.s3_bucket}/{s3_key}"
1578
+
1579
+ def _launch_workers(
1580
+ self,
1581
+ num_workers: int,
1582
+ *,
1583
+ prices: dict[str, float],
1584
+ fleet_id: str,
1585
+ user_data: str,
1586
+ secrets: dict[str, str],
1587
+ setting_package_secrets: dict[str, str],
1588
+ ) -> tuple[list[str], float]:
1589
+ """Launch EC2 worker instances."""
1590
+ ami_id = get_worker_ami_id(self.config.region, self.config.ami)
1591
+ sg_id = self._ensure_fleet_security_group()
1592
+
1593
+ try:
1594
+ launched_by = caller_identity_name(self.config.region)
1595
+ except Exception as e:
1596
+ raise RuntimeError(
1597
+ f"Cannot determine AWS identity: {e}\n"
1598
+ "The scoped fleet policy requires a LaunchedBy tag matching your identity on RunInstances. "
1599
+ "Check your AWS credentials: AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY in .env, or "
1600
+ "AWS_PROFILE pointing at a logged-in SSO profile (aws sso login --profile <name>)."
1601
+ ) from e
1602
+
1603
+ fleet_tags = [
1604
+ {"Key": "ManagedBy", "Value": "linuxbench"},
1605
+ {"Key": "FleetId", "Value": fleet_id},
1606
+ {"Key": "LaunchedBy", "Value": launched_by},
1607
+ ]
1608
+
1609
+ # Resolve instance profile ARN
1610
+ iam = boto3.client("iam", region_name=self.config.region)
1611
+ try:
1612
+ profile_resp = iam.get_instance_profile(
1613
+ InstanceProfileName=INSTANCE_PROFILE_NAME
1614
+ )
1615
+ profile_arn = profile_resp["InstanceProfile"]["Arn"]
1616
+ except Exception:
1617
+ raise SystemExit(
1618
+ f"IAM instance profile '{INSTANCE_PROFILE_NAME}' not found.\n"
1619
+ "Create it with SQS, S3, and SSM permissions for fleet workers."
1620
+ )
1621
+
1622
+ launch_kwargs: dict = dict(
1623
+ ImageId=ami_id,
1624
+ InstanceType=self.config.instance_type,
1625
+ MinCount=1,
1626
+ MaxCount=num_workers,
1627
+ UserData=gzip.compress(user_data.encode()),
1628
+ IamInstanceProfile={"Arn": profile_arn},
1629
+ InstanceInitiatedShutdownBehavior="terminate",
1630
+ MetadataOptions={"HttpTokens": "required", "HttpPutResponseHopLimit": 1},
1631
+ BlockDeviceMappings=[
1632
+ {
1633
+ "DeviceName": "/dev/sda1",
1634
+ "Ebs": {
1635
+ "VolumeSize": ROOT_VOLUME_SIZE_GB,
1636
+ "VolumeType": "gp3",
1637
+ "DeleteOnTermination": True,
1638
+ },
1639
+ }
1640
+ ],
1641
+ TagSpecifications=[
1642
+ {
1643
+ "ResourceType": "instance",
1644
+ "Tags": [
1645
+ {"Key": "Name", "Value": f"lb-fleet-{fleet_id[:8]}"},
1646
+ *fleet_tags,
1647
+ ],
1648
+ },
1649
+ ],
1650
+ NetworkInterfaces=[
1651
+ {
1652
+ "DeviceIndex": 0,
1653
+ "AssociatePublicIpAddress": True,
1654
+ "Groups": [sg_id],
1655
+ }
1656
+ ],
1657
+ )
1658
+ if self.config.spot:
1659
+ launch_kwargs["InstanceMarketOptions"] = {
1660
+ "MarketType": "spot",
1661
+ "SpotOptions": {
1662
+ "SpotInstanceType": "one-time",
1663
+ "InstanceInterruptionBehavior": "terminate",
1664
+ },
1665
+ }
1666
+
1667
+ instance_ids: list[str] = []
1668
+ fleet_price_per_hour = 0.0
1669
+ for instance_type, price_per_hour in prices.items():
1670
+ if (shortfall := num_workers - len(instance_ids)) == 0:
1671
+ break
1672
+ launched = self._launch_with_fallback(
1673
+ {**launch_kwargs, "InstanceType": instance_type, "MaxCount": shortfall},
1674
+ quota_fatal=not instance_ids,
1675
+ )
1676
+ instance_ids += launched
1677
+ fleet_price_per_hour += len(launched) * price_per_hour
1678
+ if not instance_ids:
1679
+ raise SystemExit(
1680
+ f"No EC2 capacity for any of {', '.join(prices)} in {self.config.region}. "
1681
+ "Retry later or add `--ec2-arg instance_type_fallbacks=...`."
1682
+ )
1683
+ try:
1684
+ self._put_instance_secrets(
1685
+ fleet_id,
1686
+ instance_ids,
1687
+ secrets,
1688
+ setting_package_secrets,
1689
+ fleet_tags,
1690
+ )
1691
+ except Exception:
1692
+ self._ec2.terminate_instances(InstanceIds=instance_ids)
1693
+ raise
1694
+
1695
+ if len(instance_ids) < num_workers:
1696
+ logger.warning(
1697
+ "Requested %d workers but only %d launched (AWS capacity). "
1698
+ "Fleet will run slower but still complete.",
1699
+ num_workers,
1700
+ len(instance_ids),
1701
+ )
1702
+ print(
1703
+ f" Warning: only {len(instance_ids)}/{num_workers} workers launched "
1704
+ f"(insufficient capacity). Fleet will run slower."
1705
+ )
1706
+
1707
+ # Print console link for debugging
1708
+ console_url = (
1709
+ f"https://{self.config.region}.console.aws.amazon.com/ec2/home"
1710
+ f"?region={self.config.region}#Instances:tag:FleetId={fleet_id};instanceState=running,pending"
1711
+ )
1712
+ print(f" EC2 console: {console_url}")
1713
+
1714
+ return instance_ids, fleet_price_per_hour
1715
+
1716
+ def _launch_with_fallback(
1717
+ self, launch_kwargs: dict, *, quota_fatal: bool
1718
+ ) -> list[str]:
1719
+ """Launch instances, falling back from spot→on-demand."""
1720
+ try:
1721
+ response = self._ec2.run_instances(**launch_kwargs)
1722
+ instance_ids = [inst["InstanceId"] for inst in response["Instances"]]
1723
+ if self.config.spot and instance_ids:
1724
+ print(f" Launched {len(instance_ids)} spot instances")
1725
+ return instance_ids
1726
+ except botocore.exceptions.ClientError as e:
1727
+ error_code = e.response["Error"]["Code"]
1728
+
1729
+ # Spot capacity unavailable — fall back to on-demand
1730
+ if self.config.spot and error_code in (
1731
+ "InsufficientInstanceCapacity",
1732
+ "SpotMaxPriceTooLow",
1733
+ "MaxSpotInstanceCountExceeded",
1734
+ ):
1735
+ print(
1736
+ f" Spot capacity unavailable ({error_code}), falling back to on-demand..."
1737
+ )
1738
+ fallback_kwargs = {
1739
+ k: v
1740
+ for k, v in launch_kwargs.items()
1741
+ if k != "InstanceMarketOptions"
1742
+ }
1743
+ try:
1744
+ response = self._ec2.run_instances(**fallback_kwargs)
1745
+ instance_ids = [
1746
+ inst["InstanceId"] for inst in response["Instances"]
1747
+ ]
1748
+ print(f" Launched {len(instance_ids)} on-demand instances")
1749
+ return instance_ids
1750
+ except botocore.exceptions.ClientError as e2:
1751
+ error_code = e2.response["Error"]["Code"]
1752
+ # Fall through to common error handling below
1753
+ e = e2
1754
+
1755
+ if error_code == "InsufficientInstanceCapacity":
1756
+ print(f" No {launch_kwargs['InstanceType']} capacity")
1757
+ return []
1758
+ if (
1759
+ "VcpuLimitExceeded" in error_code
1760
+ or "InstanceLimitExceeded" in error_code
1761
+ ):
1762
+ if not quota_fatal:
1763
+ print(
1764
+ f" vCPU quota reached before {launch_kwargs['InstanceType']}"
1765
+ )
1766
+ return []
1767
+ raise SystemExit(
1768
+ f"AWS vCPU quota exceeded for {launch_kwargs['InstanceType']} in {self.config.region}.\n"
1769
+ f"Request a quota increase at: https://console.aws.amazon.com/servicequotas/\n"
1770
+ f"Or reduce --ec2-arg max_workers=N (currently {self.config.max_workers})"
1771
+ ) from e
1772
+ if "PendingVerification" in error_code:
1773
+ raise SystemExit(
1774
+ f"AWS account is pending verification in {self.config.region}.\n"
1775
+ "You cannot launch EC2 instances until verification completes.\n"
1776
+ "This usually resolves within minutes but can take up to 4 hours.\n"
1777
+ "Fleet state in S3 is preserved. Re-run the same command once "
1778
+ "verification completes; queues will be recreated."
1779
+ ) from e
1780
+ raise
1781
+
1782
+ def _worker_secrets(self, extra_env_var_names: set[str]) -> dict[str, str]:
1783
+ names = sorted(set(WORKER_SECRET_ENV_VARS) | extra_env_var_names)
1784
+ for name in names:
1785
+ if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name) or name.startswith(
1786
+ "AWS_"
1787
+ ):
1788
+ raise ValueError(f"Invalid worker env var name: {name!r}")
1789
+ secrets = {name: value for name in names if (value := os.environ.get(name))}
1790
+ if len(json.dumps(secrets)) > 4096:
1791
+ largest = sorted(secrets, key=lambda n: -len(secrets[n]))[:3]
1792
+ raise ValueError(
1793
+ "Worker secrets exceed the 4 KB SSM parameter limit; largest: "
1794
+ + ", ".join(f"{n} ({len(secrets[n])} chars)" for n in largest)
1795
+ )
1796
+ return secrets
1797
+
1798
+ def _put_instance_secrets(
1799
+ self,
1800
+ fleet_id: str,
1801
+ instance_ids: list[str],
1802
+ secrets: dict[str, str],
1803
+ setting_package_secrets: dict[str, str],
1804
+ tags: list[dict[str, str]],
1805
+ ) -> None:
1806
+ """Store one SecureString per instance and one per instance and setting package, each tagged with the ARN that may read it.
1807
+
1808
+ Bootstrap polls for the instance's own parameter, so the setting packages' are written first.
1809
+ """
1810
+ account_id = boto3.client(
1811
+ "sts", region_name=self.config.region
1812
+ ).get_caller_identity()["Account"]
1813
+ for instance_id in instance_ids:
1814
+ instance_arn = (
1815
+ f"arn:aws:ec2:{self.config.region}:{account_id}:instance/{instance_id}"
1816
+ )
1817
+ name = f"{FLEET_SSM_PREFIX}/{fleet_id}/{instance_id}"
1818
+ values = {
1819
+ **{
1820
+ f"{name}.env.{package}": package_values
1821
+ for package, package_values in setting_package_secrets.items()
1822
+ },
1823
+ name: json.dumps(secrets),
1824
+ }
1825
+ for parameter, value in values.items():
1826
+ self._ssm.put_parameter(
1827
+ Name=parameter,
1828
+ Value=value,
1829
+ Type="SecureString",
1830
+ Tags=[*tags, {"Key": "InstanceArn", "Value": instance_arn}],
1831
+ )
1832
+
1833
+ def _delete_fleet_secrets(self, fleet_id: str) -> None:
1834
+ pages = self._ssm.get_paginator("get_parameters_by_path").paginate(
1835
+ Path=f"{FLEET_SSM_PREFIX}/{fleet_id}/"
1836
+ )
1837
+ names = [p["Name"] for page in pages for p in page["Parameters"]]
1838
+ for i in range(0, len(names), 10):
1839
+ self._ssm.delete_parameters(Names=names[i : i + 10])
1840
+
1841
+ def _build_user_data(
1842
+ self,
1843
+ *,
1844
+ code_s3_key: str,
1845
+ rates_s3_key: str,
1846
+ s3_prefix: str,
1847
+ work_queue_url: str,
1848
+ done_queue_url: str,
1849
+ setting_archives: list[tuple[str, str]],
1850
+ env_secret_packages: list[str],
1851
+ fleet_secrets_prefix: str,
1852
+ ) -> str:
1853
+ """Fill the bootstrap.sh template into EC2 user-data for a worker.
1854
+
1855
+ Secrets are never embedded: the script carries the SSM prefix under
1856
+ which the instance role reads its own instance's parameter.
1857
+
1858
+ Args:
1859
+ setting_archives: Each archive's S3 key and the worker directory it is extracted into.
1860
+ """
1861
+
1862
+ template_path = Path(__file__).parent / "bootstrap.sh"
1863
+ template = template_path.read_text()
1864
+
1865
+ def _escape_bash_sq(s: str) -> str:
1866
+ """Escape a string for embedding in bash single quotes."""
1867
+ return s.replace("'", "'\\''")
1868
+
1869
+ return template.format(
1870
+ s3_bucket=self.s3_bucket,
1871
+ s3_prefix=s3_prefix,
1872
+ code_s3_key=code_s3_key,
1873
+ rates_s3_key=_escape_bash_sq(rates_s3_key),
1874
+ worker_setup_s3_key=_escape_bash_sq(self._worker_setup_s3_key),
1875
+ region=self.config.region,
1876
+ work_queue_url=work_queue_url,
1877
+ done_queue_url=done_queue_url,
1878
+ concurrency=self.config.concurrency_per_worker,
1879
+ idle_timeout=300,
1880
+ fleet_secrets_prefix=_escape_bash_sq(fleet_secrets_prefix),
1881
+ setting_archives=" ".join(
1882
+ f"{key}={dest}" for key, dest in setting_archives
1883
+ ),
1884
+ env_secret_packages=_escape_bash_sq(" ".join(env_secret_packages)),
1885
+ worker_user=_escape_bash_sq(self.config.worker_user),
1886
+ inspect_log_level=_escape_bash_sq(
1887
+ os.environ.get("INSPECT_LOG_LEVEL", "warning")
1888
+ ),
1889
+ docker_org_name_ssm_param=_escape_bash_sq(
1890
+ EnvVar().dockerhub_username_ssm_param
1891
+ ),
1892
+ docker_org_pat_ssm_param=_escape_bash_sq(
1893
+ EnvVar().dockerhub_token_ssm_param
1894
+ ),
1895
+ )
1896
+
1897
+ def _terminate_fleet_instances(self, fleet_id: str) -> None:
1898
+ """Terminate EC2 instances for this fleet owned by the current user.
1899
+
1900
+ Always filters by LaunchedBy tag to prevent cross-user termination.
1901
+ If we can't determine the caller's AWS identity, we skip termination
1902
+ entirely rather than risk terminating another user's instances.
1903
+ """
1904
+ try:
1905
+ username = caller_identity_name(self.config.region)
1906
+ except Exception:
1907
+ logger.warning(
1908
+ "Cannot determine AWS identity — skipping fleet termination "
1909
+ "to avoid cross-user impact. Instances will self-terminate on idle."
1910
+ )
1911
+ return
1912
+
1913
+ filters = [
1914
+ {"Name": "tag:FleetId", "Values": [fleet_id]},
1915
+ {"Name": "tag:LaunchedBy", "Values": [username]},
1916
+ {"Name": "instance-state-name", "Values": ["running", "pending"]},
1917
+ ]
1918
+
1919
+ resp = self._ec2.describe_instances(Filters=filters)
1920
+ instance_ids = [
1921
+ inst["InstanceId"]
1922
+ for res in resp["Reservations"]
1923
+ for inst in res["Instances"]
1924
+ ]
1925
+ if instance_ids:
1926
+ self._ec2.terminate_instances(InstanceIds=instance_ids)
1927
+ print(f"Terminated {len(instance_ids)} instances (owner: {username})")
1928
+
1929
+ def _cleanup_fleet_resources(
1930
+ self, fleet_id: str, *, wait_for_workers: bool = True
1931
+ ) -> None:
1932
+ """Best-effort cleanup for queues and tagged workers after terminal results."""
1933
+ if wait_for_workers:
1934
+ logger.info("Waiting 30s for workers to drain before cleanup...")
1935
+ time.sleep(30)
1936
+ try:
1937
+ self._delete_queues(fleet_id)
1938
+ except Exception as e:
1939
+ logger.warning("Failed to delete fleet queues for %s: %s", fleet_id, e)
1940
+ try:
1941
+ self._terminate_fleet_instances(fleet_id)
1942
+ except Exception as e:
1943
+ logger.warning(
1944
+ "Failed to terminate fleet instances for %s: %s", fleet_id, e
1945
+ )
1946
+ try:
1947
+ self._delete_fleet_secrets(fleet_id)
1948
+ except Exception as e:
1949
+ logger.warning("Failed to delete fleet secrets for %s: %s", fleet_id, e)
1950
+
1951
+ def _print_estimate(
1952
+ self, total_jobs: int, num_workers: int, fleet_cost_per_hour: float
1953
+ ) -> None:
1954
+ print(f"\nDistributed EC2 eval: {total_jobs} samples")
1955
+ spot_label = " (spot)" if self.config.spot else ""
1956
+ print(
1957
+ f"Launching {num_workers}x {self.config.instance_type}{spot_label} ({min(self.config.concurrency_per_worker, FLEET_PLANNING_CONCURRENCY)} per worker) in {self.config.region}"
1958
+ )
1959
+ price_note = (
1960
+ "" if self.config.region == "us-east-1" else " (prices assume us-east-1)"
1961
+ )
1962
+ if self.config.spot:
1963
+ price_note += (
1964
+ " (spot pricing varies; estimate uses on-demand rates as ceiling)"
1965
+ )
1966
+ print(
1967
+ f"Est compute: ~${fleet_cost_per_hour:.2f}/hr | 1h: ~${fleet_cost_per_hour:.0f} | 3h: ~${fleet_cost_per_hour * 3:.0f} | 6h: ~${fleet_cost_per_hour * 6:.0f}{price_note}"
1968
+ )
1969
+
1970
+ def _wait_for_confirm(self) -> None:
1971
+ """Wait 5 seconds, abort on Ctrl+C."""
1972
+ print("Ctrl+C in 5s to cancel...")
1973
+ try:
1974
+ for i in range(5, 0, -1):
1975
+ print(f" {i}...", end=" ", flush=True)
1976
+ time.sleep(1)
1977
+ print()
1978
+ except KeyboardInterrupt:
1979
+ print("\nAborted.")
1980
+ raise SystemExit(0)
1981
+
1982
+
1983
+ def _code_root() -> Path:
1984
+ """The control-tower checkout fleet workers run: for a git or package-index install, its commit or release tag fetched into the cache; otherwise the checkout this runs from."""
1985
+ dist = distribution("control-tower")
1986
+ direct_url = json.loads(dist.read_text("direct_url.json") or "{}")
1987
+ vcs_info = direct_url.get("vcs_info") or {}
1988
+ if vcs_info.get("vcs") == "git":
1989
+ url, ref = direct_url["url"], vcs_info["commit_id"]
1990
+ elif not direct_url:
1991
+ url, ref = "https://github.com/linuxarena/control-tower", f"v{dist.version}"
1992
+ else:
1993
+ return find_repo_root(Path(__file__).resolve().parent)
1994
+ checkout = cache_dir("fleet-checkouts") / ref
1995
+ if not checkout.is_dir():
1996
+ with tempfile.TemporaryDirectory(dir=checkout.parent) as scratch:
1997
+ for args in (
1998
+ ["init", "-q"],
1999
+ ["fetch", "-q", "--depth", "1", url, ref],
2000
+ ["checkout", "-q", "FETCH_HEAD"],
2001
+ ):
2002
+ try:
2003
+ subprocess.run(
2004
+ ["git", *args],
2005
+ cwd=scratch,
2006
+ check=True,
2007
+ capture_output=True,
2008
+ text=True,
2009
+ )
2010
+ except subprocess.CalledProcessError as exc:
2011
+ raise OperationError(
2012
+ f"Cannot fetch {ref} from {url} for fleet workers ({exc.stderr.strip()})."
2013
+ ) from exc
2014
+ # Another driver fetching the same ref may have put its checkout there first.
2015
+ with contextlib.suppress(OSError):
2016
+ os.rename(scratch, checkout)
2017
+ return checkout
2018
+
2019
+
2020
+ def _git_snapshot(path: Path) -> tuple[str, bool]:
2021
+ """The git object to archive for ``path``: HEAD if clean, else the dirty worktree's tree, narrowed to ``path`` inside its repo; and whether it differs from HEAD's."""
2022
+
2023
+ def git(*args: str, env: dict[str, str] | None = None) -> str:
2024
+ return subprocess.run(
2025
+ ["git", *args],
2026
+ cwd=path,
2027
+ capture_output=True,
2028
+ text=True,
2029
+ check=True,
2030
+ env=env,
2031
+ ).stdout.strip()
2032
+
2033
+ try:
2034
+ prefix = git("rev-parse", "--show-prefix")
2035
+ # A private copy of the index, so a concurrent git command holding index.lock can't fail this.
2036
+ with tempfile.TemporaryDirectory(prefix="ct-index-") as tmp:
2037
+ index = Path(tmp) / "index"
2038
+ real_index = Path(path, git("rev-parse", "--git-path", "index"))
2039
+ if real_index.exists():
2040
+ shutil.copy2(real_index, index)
2041
+ env = {**os.environ, "GIT_INDEX_FILE": str(index)}
2042
+ git("add", "--update", "--", ":/", env=env)
2043
+ tree = git("write-tree", env=env)
2044
+ if tree == git("rev-parse", "HEAD^{tree}"):
2045
+ return git("rev-parse", f"HEAD:{prefix}" if prefix else "HEAD"), False
2046
+ tree = git("rev-parse", f"{tree}:{prefix}")
2047
+ head_tree = subprocess.run(
2048
+ ["git", "rev-parse", "--verify", "--quiet", f"HEAD:{prefix}"],
2049
+ cwd=path,
2050
+ capture_output=True,
2051
+ text=True,
2052
+ ).stdout.strip()
2053
+ return tree, tree != head_tree
2054
+ except subprocess.CalledProcessError as exc:
2055
+ raise ValidationError(
2056
+ f"Cannot ship {path} to fleet workers: git cannot archive it "
2057
+ f"({exc.stderr.strip()}). Workers get only tracked files from a git checkout."
2058
+ ) from exc