launch_node.rs 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  1. //! Per-node launch body for [`start_cluster`]: resolves path references and
  2. //! NCCL GID per node, builds the composed model dir if configured, then
  3. //! assembles + runs the `docker run` command (or prints it in dry-run).
  4. //!
  5. //! Kept as a submodule so `src/docker/mod.rs` stays under the ≤500-line gate;
  6. //! the per-node work needs SSH (per-node HF snapshots and GID indexes can
  7. //! differ across nodes) so it must run inside the cluster spawn loop, not
  8. //! per-cluster.
  9. use crate::config::Node;
  10. use crate::docker::{command, compose};
  11. /// Per-node inputs for [`launch_node`], grouped so the function stays under
  12. /// the argument-count lint without scattering fields.
  13. pub(crate) struct LaunchNodeCtx<'a> {
  14. pub node: &'a Node,
  15. pub rn: &'a str,
  16. pub dry: bool,
  17. pub engine: &'a str,
  18. pub is_head: bool,
  19. pub cleanup: String,
  20. pub dm_ctx: &'a command::DockerCommandCtx,
  21. pub hf_cache: &'a std::path::Path,
  22. pub path_refs: &'a [command::PathRef],
  23. }
  24. /// Run one node's launch: returns `(node.name, rank, engine, command result)`.
  25. pub(crate) async fn launch_node(
  26. config: &crate::config::Config,
  27. ctx: LaunchNodeCtx<'_>,
  28. ) -> (String, u32, String, anyhow::Result<String>) {
  29. let LaunchNodeCtx {
  30. node,
  31. rn,
  32. dry,
  33. engine,
  34. is_head,
  35. cleanup,
  36. dm_ctx,
  37. hf_cache,
  38. path_refs,
  39. } = ctx;
  40. let fail = |e: anyhow::Error| (node.name.clone(), node.rank, engine.to_string(), Err(e));
  41. if !dry && let Err(e) = crate::ssh::run_ssh(config, node, cleanup, dry).await {
  42. // Ignore cleanup errors — a stale container that won't die must not
  43. // abort the launch (the `docker run --name` will surface any real
  44. // conflict).
  45. tracing::warn!("cleanup on {}: {e}", node.name);
  46. }
  47. // Resolve model path references per node from that node's HF cache (the
  48. // snapshot hash can differ across nodes). Dry-run skips SSH and emits a
  49. // descriptive placeholder path instead. The resolved paths are substituted
  50. // into the recipe args/env via `{{<name>}}` before building the docker
  51. // command — no bind-mount is created for path refs (the HF cache is
  52. // already mounted at /root/.cache/huggingface, so symlinks resolve
  53. // correctly).
  54. let mut resolved_refs: Vec<(String, String)> = Vec::new();
  55. for spec in path_refs {
  56. let resolved_path = if dry {
  57. command::path_ref_placeholder(spec)
  58. } else {
  59. match crate::docker::path_ref::resolve_path_ref_on_node(config, node, spec, hf_cache)
  60. .await
  61. {
  62. Ok(p) => p,
  63. Err(e) => return fail(e),
  64. }
  65. };
  66. resolved_refs.push((spec.name.clone(), resolved_path));
  67. }
  68. let mut recipe = command::substitute_path_refs(&dm_ctx.recipe, &resolved_refs);
  69. let mut extra_mounts = dm_ctx.extra_mounts.clone();
  70. // Resolve NCCL_IB_GID_INDEX per node when the cluster config does not pin
  71. // `ib_gid_index` (auto-detect: scan this node's sysfs IB GID table for the
  72. // RoCE-v2 IPv4 GID — it re-numbers across reboots, mirroring the reference
  73. // serving stack). Dry-run skips SSH and renders a placeholder.
  74. let nccl =
  75. match crate::detect::resolve_nccl_for_node(config, node, dm_ctx.nccl.as_ref(), dry).await {
  76. Ok(n) => n,
  77. Err(e) => return fail(e),
  78. };
  79. // Build composed model dir if configured: symlinks into the HF cache +
  80. // generated files, bind-mounted at `mount`. `model` is rewritten to the
  81. // mount path so vLLM loads from the composed dir.
  82. if let Some(ref cm) = recipe.composed_model {
  83. match compose::build_composed_dir_on_node(config, node, rn, node.rank, cm, hf_cache, dry)
  84. .await
  85. {
  86. Ok((hp, cp)) => {
  87. extra_mounts.push(command::ResolvedMount {
  88. host: hp,
  89. container: cp.clone(),
  90. read_only: false,
  91. env: None,
  92. });
  93. recipe.model = cp;
  94. }
  95. Err(e) => return fail(e),
  96. }
  97. }
  98. let ctx = command::DockerCommandCtx {
  99. recipe,
  100. extra_mounts,
  101. nccl,
  102. ..dm_ctx.clone()
  103. };
  104. let cmd = match command::build_docker_command(&ctx, node, is_head) {
  105. Ok(c) => c.join(" "),
  106. Err(e) => return fail(e),
  107. };
  108. let result = crate::ssh::run_ssh_with_retry(config, node, cmd, dry, 2).await;
  109. (node.name.clone(), node.rank, engine.to_string(), result)
  110. }