Przeglądaj źródła

feat(serve): add --force flag to ignore occupied nodes

Skips the occupancy check during cluster selection — occupied nodes are kept in the available set (reachable is still required — can't SSH to an unreachable node). Implemented by clearing occupied flags on the probe state in prepare_cluster when force_occupied is set, so both select_cluster (which checks free_count) and the downstream node filter see reachable nodes as free.

Use to run a lightweight test recipe alongside a live serve, or to force onto a busy node. Dangerous for GPU recipes (OOM, port conflicts); the operator is responsible for coexistence.
Alexander Korolev 2 miesięcy temu
rodzic
commit
892c28d1a4
6 zmienionych plików z 97 dodań i 1 usunięć
  1. 6 0
      AGENTS.md
  2. 8 0
      src/cli.rs
  3. 3 0
      src/docker/mod.rs
  4. 3 0
      src/main.rs
  5. 16 1
      src/orchestrate/mod.rs
  6. 61 0
      src/orchestrate/tests.rs

+ 6 - 0
AGENTS.md

@@ -28,6 +28,12 @@ all clusters and runs one single-node container there (no `--nnodes`/`--tp`/
 forces multi-node even if the recipe sets `solo = true`. A recipe can also
 set `solo = true` as the default.
 
+`serve --force` skips the occupancy check during cluster selection —
+occupied nodes are kept in the available set (reachable is still required).
+Use to run a lightweight test recipe alongside a live serve, or to force onto
+a busy node. Dangerous for GPU recipes (OOM, port conflicts); the operator is
+responsible for coexistence.
+
 > ⚠️ **Cluster may be serving live traffic.** A model (often the one serving this
 > very agent) may be running on the cluster at any time. Before running `stop`,
 > `power off`, or `power switch`, ALWAYS first run `health` and confirm with the

+ 8 - 0
src/cli.rs

@@ -59,6 +59,14 @@ pub enum Commands {
         #[arg(long)]
         force_build: bool,
 
+        /// Run even if nodes are occupied (another container running).
+        /// Skips the occupancy check during cluster selection — use to run
+        /// a lightweight test recipe alongside a live serve, or to force
+        /// onto a busy node. Dangerous for GPU recipes (OOM, port
+        /// conflicts); the operator is responsible for coexistence.
+        #[arg(long)]
+        force: bool,
+
         /// Serve single-node on the first free node found (overrides recipe
         /// `solo`). Use `--no-solo` to force multi-node even if the recipe
         /// sets `solo = true`.

+ 3 - 0
src/docker/mod.rs

@@ -27,6 +27,8 @@ pub struct ServeOpts {
     pub sync_image_source: Option<String>,
     pub build_node: Option<String>,
     pub force_build: bool,
+    /// Ignore occupied nodes during cluster selection (`--force`).
+    pub force: bool,
     /// CLI override for recipe `solo`: `Some(true)` (--solo), `Some(false)`
     /// (--no-solo), `None` (unset → use recipe `solo` field).
     pub solo: Option<bool>,
@@ -96,6 +98,7 @@ pub async fn serve(
                 pinger: &SshNodePinger,
                 auto_power_on: opts.auto_power_on,
                 boot_timing: BootTiming::default(),
+                force_occupied: opts.force,
             },
         )
         .await?

+ 3 - 0
src/main.rs

@@ -95,6 +95,7 @@ async fn main() -> Result<()> {
             sync_image_source,
             build_node,
             force_build,
+            force,
             solo,
         } => {
             handle_serve(
@@ -111,6 +112,7 @@ async fn main() -> Result<()> {
                     sync_image_source,
                     build_node,
                     force_build,
+                    force,
                     solo,
                 },
                 follow,
@@ -233,6 +235,7 @@ async fn handle_build_cmd(
         sync_image_source: None,
         build_node: Some(build_node.to_string()),
         force_build,
+        force: false,
         solo: None,
     };
     docker::build::handle_build(config, &cluster, recipe, &recipe_obj, &opts, build_node).await?;

+ 16 - 1
src/orchestrate/mod.rs

@@ -50,6 +50,10 @@ pub struct PrepareRequest<'a, P: NodeProber, C: PowerController, N: NodePinger>
     pub pinger: &'a N,
     pub auto_power_on: bool,
     pub boot_timing: BootTiming,
+    /// `--force`: skip the occupancy filter during cluster selection.
+    /// Nodes with another container running are kept in the available set
+    /// (reachable is still required — can't SSH to an unreachable node).
+    pub force_occupied: bool,
 }
 
 pub async fn prepare_cluster<P, C, N>(
@@ -83,6 +87,17 @@ where
             state
         };
 
+        // `--force`: clear occupied flags so select_cluster counts reachable
+        // nodes (ignoring occupancy) and the downstream filter keeps them.
+        let mut state = state;
+        if req.force_occupied {
+            for ns in state.nodes.values_mut() {
+                if ns.reachable {
+                    ns.occupied = false;
+                }
+            }
+        }
+
         let (c, n) = select_cluster(config, cluster_name, recipe.required_nodes, &state)?;
         let avail: Vec<Node> = c
             .nodes
@@ -91,7 +106,7 @@ where
                 state
                     .nodes
                     .get(&n.name)
-                    .is_some_and(|ns| ns.reachable && !ns.occupied)
+                    .is_some_and(|ns| ns.reachable && (req.force_occupied || !ns.occupied))
             })
             .collect();
         let c = Cluster { nodes: avail, ..c };

+ 61 - 0
src/orchestrate/tests.rs

@@ -189,6 +189,7 @@ async fn cluster_off_auto_power_on_powers_and_selects() {
                 poll_interval_secs: 0,
                 max_poll_secs: 10,
             },
+            force_occupied: false,
         },
     )
     .await
@@ -245,6 +246,7 @@ async fn cluster_off_no_auto_power_fails() {
                 poll_interval_secs: 0,
                 max_poll_secs: 10,
             },
+            force_occupied: false,
         },
     )
     .await;
@@ -300,6 +302,7 @@ async fn cluster_already_on_selects_directly() {
                 poll_interval_secs: 0,
                 max_poll_secs: 10,
             },
+            force_occupied: false,
         },
     )
     .await
@@ -354,6 +357,7 @@ async fn occupied_nodes_excluded() {
                 poll_interval_secs: 0,
                 max_poll_secs: 10,
             },
+            force_occupied: false,
         },
     )
     .await
@@ -363,6 +367,63 @@ async fn occupied_nodes_excluded() {
     assert_eq!(cluster.nodes[0].name, "n2");
 }
 
+#[tokio::test]
+async fn force_occupied_keeps_occupied_nodes() {
+    // `--force` skips the occupancy filter: an occupied node is kept in the
+    // available set (reachable is still required). Mirror of
+    // `occupied_nodes_excluded` with force_occupied: true.
+    let c = test_cluster("c", &["n1", "n2"]);
+    let config = test_config(vec![c.clone()]);
+    let recipe = test_recipe(Some(2));
+
+    let mut state = ClusterState::default();
+    state.nodes.insert(
+        "n1".into(),
+        NodeState {
+            reachable: true,
+            occupied: true,
+        },
+    );
+    state.nodes.insert(
+        "n2".into(),
+        NodeState {
+            reachable: true,
+            occupied: false,
+        },
+    );
+
+    let prober = MockProber { state };
+    let pinger = MockPinger { reachable: true };
+    let power = MockPower {
+        results: HashMap::new(),
+        turned_on: std::sync::Mutex::new(Vec::new()),
+    };
+
+    let (cluster, _nccl) = prepare_cluster(
+        &config,
+        &recipe,
+        "auto",
+        None,
+        &PrepareRequest {
+            prober: &prober,
+            power: Some(&power),
+            pinger: &pinger,
+            auto_power_on: false,
+            boot_timing: BootTiming {
+                boot_wait_secs: 0,
+                poll_interval_secs: 0,
+                max_poll_secs: 10,
+            },
+            force_occupied: true,
+        },
+    )
+    .await
+    .unwrap();
+
+    // Both nodes kept (n1 occupied but force_occupied=true).
+    assert_eq!(cluster.nodes.len(), 2);
+}
+
 #[test]
 fn dry_run_resolves_without_probing_when_all_nodes_unreachable() {
     // Dry-run must resolve a cluster even when no node is reachable — the