inferlab_runtime/
container.rs

1//! One authority for docker container lifecycle mechanics
2//! ([[RFC-0003:C-RUNTIME-WORKFLOWS]]): bounded client invocation with
3//! concurrently drained pipes, and confirmed container removal on the
4//! launch machine. Callers own their argv assembly, evidence shapes, and
5//! operator-facing messages — this module owns only the mechanics that
6//! have actually failed in the field: the pipe-capacity deadlock, the
7//! deadline kill, and the container that outlives its docker client.
8
9use serde::{Deserialize, Serialize};
10use std::io::{Read, Write};
11use std::os::unix::process::CommandExt;
12use std::path::Path;
13use std::process::{Command, ExitStatus, Stdio};
14use std::time::Duration;
15use wait_timeout::ChildExt;
16
17use crate::operation_bound::{OperationBound, Remaining, duration_millis};
18
19/// How long a removal may take before it is reported as unconfirmed: a
20/// wedged daemon is a plausible cause of the original deadline miss, and a
21/// cleanup path must not hang on its own cleanup.
22pub const REMOVAL_TIMEOUT: Duration = Duration::from_secs(30);
23const INTERRUPT_POLL_INTERVAL: Duration = Duration::from_millis(50);
24const COMMAND_REAP_GRACE: Duration = Duration::from_secs(5);
25const COMMAND_IO_DRAIN_GRACE: Duration = Duration::from_secs(5);
26
27/// The `--gpus` argv pair for a device spec (a single index or a
28/// comma-joined list). The literal quotes are part of the value: docker's
29/// `--gpus` parser splits an unquoted `device=0,1` on the comma and exposes
30/// only the first device (verified on real hardware in v1.5).
31pub fn docker_device_args(spec: &str) -> [String; 2] {
32    ["--gpus".to_owned(), format!("\"device={spec}\"")]
33}
34
35/// A bounded invocation that ran to an observable end.
36pub enum BoundedWait {
37    Exited {
38        status: ExitStatus,
39        stdout: Vec<u8>,
40        stderr: Vec<u8>,
41    },
42    /// The effective wait bound expired: the client was killed and reaped.
43    /// The caller remains responsible for deciding whether the owning
44    /// operation expired or a shorter subordinate attempt cap fired.
45    Expired {
46        kill: std::io::Result<()>,
47        operation_elapsed_ms: u64,
48        cleanup: Option<CommandCleanupEvidence>,
49    },
50    /// A new operator interruption arrived after this command started. The
51    /// owned process group was terminated and reaped before returning.
52    Interrupted {
53        kill: std::io::Result<()>,
54        operation_elapsed_ms: u64,
55        cleanup: CommandCleanupEvidence,
56    },
57}
58
59#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
60#[serde(deny_unknown_fields)]
61pub struct CommandCleanupEvidence {
62    pub trigger: CommandCleanupTrigger,
63    pub elapsed_ms: u64,
64    pub reap_grace_ms: u64,
65    pub io_drain_grace_ms: u64,
66    pub kill_attempted: bool,
67    pub verified: bool,
68    pub error: Option<String>,
69}
70
71#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
72#[serde(rename_all = "snake_case")]
73pub enum CommandCleanupTrigger {
74    Deadline,
75    Interruption,
76    WaitFailure,
77}
78
79/// Where a bounded invocation failed before producing a wait outcome.
80pub enum BoundedError {
81    Launch(std::io::Error),
82    Stdin(std::io::Error),
83    Wait(std::io::Error),
84    WaitCleanup {
85        source: std::io::Error,
86        operation_elapsed_ms: u64,
87        cleanup: CommandCleanupEvidence,
88    },
89}
90
91/// Spawn `argv`, optionally write `stdin_payload`, and wait under the
92/// operation's remaining budget and optional attempt cap while draining
93/// stdout and stderr concurrently: a child whose
94/// output exceeds the pipe capacity would otherwise block writing and
95/// never exit, turning a large response or traceback into a timeout.
96pub fn run_with_bound<S: AsRef<std::ffi::OsStr>>(
97    argv: &[S],
98    cwd: Option<&Path>,
99    stdin_payload: Option<&[u8]>,
100    bound: &OperationBound,
101    attempt_cap: Option<Duration>,
102) -> Result<BoundedWait, BoundedError> {
103    run_with_bound_mode(
104        argv,
105        cwd,
106        stdin_payload,
107        bound,
108        attempt_cap,
109        true,
110        crate::interrupt::received,
111    )
112}
113
114pub fn run_cleanup_with_bound<S: AsRef<std::ffi::OsStr>>(
115    argv: &[S],
116    cwd: Option<&Path>,
117    stdin_payload: Option<&[u8]>,
118    bound: &OperationBound,
119    attempt_cap: Option<Duration>,
120) -> Result<BoundedWait, BoundedError> {
121    run_with_bound_mode(
122        argv,
123        cwd,
124        stdin_payload,
125        bound,
126        attempt_cap,
127        false,
128        crate::interrupt::received,
129    )
130}
131
132fn run_with_bound_mode<S: AsRef<std::ffi::OsStr>, F: FnMut() -> bool>(
133    argv: &[S],
134    cwd: Option<&Path>,
135    stdin_payload: Option<&[u8]>,
136    bound: &OperationBound,
137    attempt_cap: Option<Duration>,
138    interruptible: bool,
139    mut interrupted: F,
140) -> Result<BoundedWait, BoundedError> {
141    let attempt = bound.attempt(attempt_cap);
142    if matches!(attempt.remaining(), Remaining::Expired) {
143        return Ok(BoundedWait::Expired {
144            kill: Ok(()),
145            operation_elapsed_ms: bound.elapsed_ms(),
146            cleanup: None,
147        });
148    }
149    let Some(program) = argv.first() else {
150        return Err(BoundedError::Launch(std::io::Error::new(
151            std::io::ErrorKind::InvalidInput,
152            "external command argv is empty",
153        )));
154    };
155    let mut command = Command::new(program);
156    command
157        .args(&argv[1..])
158        .stdin(if stdin_payload.is_some() {
159            Stdio::piped()
160        } else {
161            Stdio::null()
162        })
163        .stdout(Stdio::piped())
164        .stderr(Stdio::piped())
165        .process_group(0);
166    if let Some(cwd) = cwd {
167        command.current_dir(cwd);
168    }
169    let mut child = command.spawn().map_err(BoundedError::Launch)?;
170    let stdin_write = if let Some(payload) = stdin_payload {
171        let mut stdin = child
172            .stdin
173            .take()
174            .ok_or_else(|| BoundedError::Stdin(std::io::Error::other("stdin was not piped")))?;
175        let payload = payload.to_owned();
176        Some(std::thread::spawn(move || {
177            stdin.write_all(&payload).and_then(|()| stdin.flush())
178        }))
179    } else {
180        None
181    };
182    let mut stdout_pipe = child.stdout.take();
183    let stdout_drain = std::thread::spawn(move || {
184        let mut buffer = Vec::new();
185        if let Some(pipe) = stdout_pipe.as_mut() {
186            let _ = pipe.read_to_end(&mut buffer);
187        }
188        buffer
189    });
190    let mut stderr_pipe = child.stderr.take();
191    let stderr_drain = std::thread::spawn(move || {
192        let mut buffer = Vec::new();
193        if let Some(pipe) = stderr_pipe.as_mut() {
194            let _ = pipe.read_to_end(&mut buffer);
195        }
196        buffer
197    });
198    let status = loop {
199        if interruptible && interrupted() {
200            let operation_elapsed_ms = bound.elapsed_ms();
201            let (kill, cleanup) = terminate_group_and_finish_io(
202                &mut child,
203                stdin_write,
204                stdout_drain,
205                stderr_drain,
206                CommandCleanupTrigger::Interruption,
207            );
208            return Ok(BoundedWait::Interrupted {
209                kill,
210                operation_elapsed_ms,
211                cleanup,
212            });
213        }
214        let Some(wait) = wait_slice(&attempt) else {
215            let operation_elapsed_ms = bound.elapsed_ms();
216            let (kill, cleanup) = terminate_group_and_finish_io(
217                &mut child,
218                stdin_write,
219                stdout_drain,
220                stderr_drain,
221                CommandCleanupTrigger::Deadline,
222            );
223            return Ok(BoundedWait::Expired {
224                kill,
225                operation_elapsed_ms,
226                cleanup: Some(cleanup),
227            });
228        };
229        match child.wait_timeout(wait) {
230            Ok(Some(status)) => break status,
231            Ok(None) => {}
232            Err(error) => {
233                let operation_elapsed_ms = bound.elapsed_ms();
234                let (_, cleanup) = terminate_group_and_finish_io(
235                    &mut child,
236                    stdin_write,
237                    stdout_drain,
238                    stderr_drain,
239                    CommandCleanupTrigger::WaitFailure,
240                );
241                return Err(BoundedError::WaitCleanup {
242                    source: error,
243                    operation_elapsed_ms,
244                    cleanup,
245                });
246            }
247        }
248    };
249    while !io_finished(&stdin_write, &stdout_drain, &stderr_drain) {
250        if interruptible && interrupted() {
251            let operation_elapsed_ms = bound.elapsed_ms();
252            let (kill, cleanup) = terminate_group_and_finish_io(
253                &mut child,
254                stdin_write,
255                stdout_drain,
256                stderr_drain,
257                CommandCleanupTrigger::Interruption,
258            );
259            return Ok(BoundedWait::Interrupted {
260                kill,
261                operation_elapsed_ms,
262                cleanup,
263            });
264        }
265        let Some(wait) = wait_slice(&attempt) else {
266            let operation_elapsed_ms = bound.elapsed_ms();
267            let (kill, cleanup) = terminate_group_and_finish_io(
268                &mut child,
269                stdin_write,
270                stdout_drain,
271                stderr_drain,
272                CommandCleanupTrigger::Deadline,
273            );
274            return Ok(BoundedWait::Expired {
275                kill,
276                operation_elapsed_ms,
277                cleanup: Some(cleanup),
278            });
279        };
280        std::thread::sleep(wait);
281    }
282    let stdin_result = join_writer_result(stdin_write);
283    let stdout = join_drain(stdout_drain);
284    let stderr = join_drain(stderr_drain);
285    stdin_result.map_err(BoundedError::Stdin)?;
286    let stdout = stdout?;
287    let stderr = stderr?;
288    Ok(BoundedWait::Exited {
289        status,
290        stdout,
291        stderr,
292    })
293}
294
295fn wait_slice(attempt: &crate::operation_bound::AttemptBound) -> Option<Duration> {
296    match attempt.remaining() {
297        Remaining::Finite(remaining) => Some(remaining.min(INTERRUPT_POLL_INTERVAL)),
298        Remaining::Expired => None,
299        Remaining::Unbounded => Some(INTERRUPT_POLL_INTERVAL),
300    }
301}
302
303fn io_finished(
304    writer: &Option<std::thread::JoinHandle<std::io::Result<()>>>,
305    stdout: &std::thread::JoinHandle<Vec<u8>>,
306    stderr: &std::thread::JoinHandle<Vec<u8>>,
307) -> bool {
308    writer.as_ref().is_none_or(|writer| writer.is_finished())
309        && stdout.is_finished()
310        && stderr.is_finished()
311}
312
313fn terminate_group_and_finish_io(
314    child: &mut std::process::Child,
315    writer: Option<std::thread::JoinHandle<std::io::Result<()>>>,
316    stdout: std::thread::JoinHandle<Vec<u8>>,
317    stderr: std::thread::JoinHandle<Vec<u8>>,
318    trigger: CommandCleanupTrigger,
319) -> (std::io::Result<()>, CommandCleanupEvidence) {
320    let started = std::time::Instant::now();
321    let kill = terminate_process_group(child);
322    let result = match kill {
323        Ok(()) => finish_io_with_grace(writer, stdout, stderr),
324        Err(error) => Err(error),
325    };
326    let evidence = CommandCleanupEvidence {
327        trigger,
328        elapsed_ms: duration_millis(started.elapsed()),
329        reap_grace_ms: duration_millis(COMMAND_REAP_GRACE),
330        io_drain_grace_ms: duration_millis(COMMAND_IO_DRAIN_GRACE),
331        kill_attempted: true,
332        verified: result.is_ok(),
333        error: result.as_ref().err().map(ToString::to_string),
334    };
335    (result, evidence)
336}
337
338fn terminate_process_group(child: &mut std::process::Child) -> std::io::Result<()> {
339    let pid = rustix::process::Pid::from_raw(child.id() as i32)
340        .ok_or_else(|| std::io::Error::other("child process id is zero"))?;
341    let signal = rustix::process::kill_process_group(pid, rustix::process::Signal::KILL);
342    if signal.is_err() && child.try_wait()?.is_none() {
343        child.kill()?;
344    }
345    match child.wait_timeout(COMMAND_REAP_GRACE)? {
346        Some(_) => Ok(()),
347        None => Err(std::io::Error::new(
348            std::io::ErrorKind::TimedOut,
349            format!(
350                "command child did not reap within {} seconds",
351                COMMAND_REAP_GRACE.as_secs()
352            ),
353        )),
354    }
355}
356
357fn finish_io_with_grace(
358    writer: Option<std::thread::JoinHandle<std::io::Result<()>>>,
359    stdout: std::thread::JoinHandle<Vec<u8>>,
360    stderr: std::thread::JoinHandle<Vec<u8>>,
361) -> std::io::Result<()> {
362    let started = std::time::Instant::now();
363    while !io_finished(&writer, &stdout, &stderr) {
364        if started.elapsed() >= COMMAND_IO_DRAIN_GRACE {
365            return Err(std::io::Error::new(
366                std::io::ErrorKind::TimedOut,
367                format!(
368                    "command I/O did not drain within {} seconds after termination",
369                    COMMAND_IO_DRAIN_GRACE.as_secs()
370                ),
371            ));
372        }
373        std::thread::sleep(INTERRUPT_POLL_INTERVAL);
374    }
375    // A writer blocked when the operation expires normally observes a broken
376    // pipe after the process group is killed. Cleanup needs to prove that the
377    // writer thread finished; its write result still belongs to the terminal
378    // operation and must not reclassify verified termination as cleanup
379    // failure.
380    join_writer_completion(writer)?;
381    let _ = join_drain(stdout).map_err(bounded_error_into_io)?;
382    let _ = join_drain(stderr).map_err(bounded_error_into_io)?;
383    Ok(())
384}
385
386fn join_writer_completion(
387    writer: Option<std::thread::JoinHandle<std::io::Result<()>>>,
388) -> std::io::Result<()> {
389    let Some(writer) = writer else {
390        return Ok(());
391    };
392    writer
393        .join()
394        .map(|_| ())
395        .map_err(|_| std::io::Error::other("stdin writer thread panicked"))
396}
397
398fn bounded_error_into_io(error: BoundedError) -> std::io::Error {
399    match error {
400        BoundedError::Launch(error) | BoundedError::Stdin(error) | BoundedError::Wait(error) => {
401            error
402        }
403        BoundedError::WaitCleanup { source, .. } => source,
404    }
405}
406
407fn join_writer_result(
408    writer: Option<std::thread::JoinHandle<std::io::Result<()>>>,
409) -> std::io::Result<()> {
410    let Some(writer) = writer else {
411        return Ok(());
412    };
413    match writer.join() {
414        Ok(result) => result,
415        Err(_) => Err(std::io::Error::other("stdin writer thread panicked")),
416    }
417}
418
419fn join_drain(drain: std::thread::JoinHandle<Vec<u8>>) -> Result<Vec<u8>, BoundedError> {
420    drain
421        .join()
422        .map_err(|_| BoundedError::Wait(std::io::Error::other("output drain thread panicked")))
423}
424
425/// The outcome of a confirmed-removal attempt.
426pub enum Removal {
427    Confirmed { already_absent: bool },
428    Unconfirmed(RemovalFailure),
429}
430
431/// Why a removal stayed unconfirmed; callers format their own messages.
432pub enum RemovalFailure {
433    Launch(std::io::Error),
434    Wait(std::io::Error),
435    WaitCleanup {
436        source: std::io::Error,
437        operation_elapsed_ms: u64,
438        client_cleanup: CommandCleanupEvidence,
439    },
440    Deadline {
441        operation_elapsed_ms: u64,
442        client_cleanup: Option<CommandCleanupEvidence>,
443    },
444    Exit {
445        status: ExitStatus,
446        stderr: String,
447    },
448    Ssh(String),
449}
450
451/// `docker rm -f` on the container's launch machine, always under
452/// [`REMOVAL_TIMEOUT`]: the SSH client itself runs through the bounded
453/// invocation, because the transport's connect and keepalive options
454/// detect dead connections but cannot bound a live connection to a wedged
455/// remote daemon. A container the daemon no longer knows is the
456/// confirmation, not a failure ([[RFC-0003:C-RUNTIME-WORKFLOWS]]).
457pub fn remove_container(target: Option<&str>, container: &str) -> Removal {
458    let bound = OperationBound::finite(REMOVAL_TIMEOUT);
459    let argv = match target {
460        Some(target) => crate::ssh::ssh_argv(
461            target,
462            &format!("docker rm -f {}", crate::shell::shell_quote(container)),
463        ),
464        None => vec![
465            "docker".to_owned(),
466            "rm".to_owned(),
467            "-f".to_owned(),
468            container.to_owned(),
469        ],
470    };
471    let (status, stderr) = match run_cleanup_with_bound(&argv, None, None, &bound, None) {
472        Ok(BoundedWait::Exited { status, stderr, .. }) => {
473            (status, String::from_utf8_lossy(&stderr).into_owned())
474        }
475        Ok(BoundedWait::Expired {
476            operation_elapsed_ms,
477            cleanup,
478            ..
479        }) => {
480            return Removal::Unconfirmed(RemovalFailure::Deadline {
481                operation_elapsed_ms,
482                client_cleanup: cleanup,
483            });
484        }
485        Ok(BoundedWait::Interrupted { .. }) => {
486            return Removal::Unconfirmed(RemovalFailure::Wait(std::io::Error::new(
487                std::io::ErrorKind::Interrupted,
488                "container cleanup was interrupted",
489            )));
490        }
491        Err(BoundedError::Launch(error)) => {
492            return Removal::Unconfirmed(match target {
493                Some(target) => {
494                    RemovalFailure::Ssh(format!("failed to launch SSH for {target:?}: {error}"))
495                }
496                None => RemovalFailure::Launch(error),
497            });
498        }
499        Err(BoundedError::Stdin(error)) | Err(BoundedError::Wait(error)) => {
500            return Removal::Unconfirmed(RemovalFailure::Wait(error));
501        }
502        Err(BoundedError::WaitCleanup {
503            source,
504            operation_elapsed_ms,
505            cleanup,
506        }) => {
507            return Removal::Unconfirmed(RemovalFailure::WaitCleanup {
508                source,
509                operation_elapsed_ms,
510                client_cleanup: cleanup,
511            });
512        }
513    };
514    if status.success() {
515        return Removal::Confirmed {
516            already_absent: false,
517        };
518    }
519    if stderr.contains("No such container") {
520        return Removal::Confirmed {
521            already_absent: true,
522        };
523    }
524    // A `--rm` container whose exit races this removal: the daemon owns an
525    // in-flight removal it refuses to double-remove, so confirmation is
526    // observing the container disappear within the same removal deadline.
527    if stderr.contains("is already in progress")
528        && confirm_container_absent(target, container, &bound)
529    {
530        return Removal::Confirmed {
531            already_absent: true,
532        };
533    }
534    Removal::Unconfirmed(RemovalFailure::Exit { status, stderr })
535}
536
537/// Poll the daemon on the container's launch machine until it no longer
538/// knows the container, bounded by [`REMOVAL_TIMEOUT`]. Any answer other
539/// than a definitive not-found keeps polling; the deadline decides.
540fn confirm_container_absent(target: Option<&str>, container: &str, bound: &OperationBound) -> bool {
541    let argv = match target {
542        Some(target) => crate::ssh::ssh_argv(
543            target,
544            &format!(
545                "docker container inspect --format {{{{.Id}}}} {}",
546                crate::shell::shell_quote(container)
547            ),
548        ),
549        None => vec![
550            "docker".to_owned(),
551            "container".to_owned(),
552            "inspect".to_owned(),
553            "--format".to_owned(),
554            "{{.Id}}".to_owned(),
555            container.to_owned(),
556        ],
557    };
558    loop {
559        if let Ok(BoundedWait::Exited { status, stderr, .. }) =
560            run_cleanup_with_bound(&argv, None, None, bound, None)
561            && !status.success()
562            && String::from_utf8_lossy(&stderr).contains("No such container")
563        {
564            return true;
565        }
566        if bound.is_expired() {
567            return false;
568        }
569        let sleep = match bound.attempt(Some(Duration::from_millis(250))).remaining() {
570            Remaining::Finite(duration) => duration,
571            Remaining::Expired => return false,
572            Remaining::Unbounded => Duration::from_millis(250),
573        };
574        std::thread::sleep(sleep);
575    }
576}
577
578#[cfg(test)]
579mod tests {
580    use super::{BoundedWait, CommandCleanupTrigger, run_with_bound, run_with_bound_mode};
581    use crate::operation_bound::OperationBound;
582    use std::time::Duration;
583
584    #[test]
585    fn unbounded_command_wait_terminates_when_the_process_exits() {
586        let argv = vec!["sh".to_owned(), "-c".to_owned(), "exit 0".to_owned()];
587        let outcome = run_with_bound(&argv, None, None, &OperationBound::unbounded(), None);
588
589        assert!(matches!(
590            outcome,
591            Ok(BoundedWait::Exited { status, .. }) if status.success()
592        ));
593    }
594
595    #[test]
596    fn pending_interruption_terminates_an_unbounded_command() {
597        let argv = vec!["sh".to_owned(), "-c".to_owned(), "sleep 60".to_owned()];
598        let outcome = run_with_bound_mode(
599            &argv,
600            None,
601            None,
602            &OperationBound::unbounded(),
603            None,
604            true,
605            || true,
606        );
607
608        let (operation_elapsed_ms, cleanup) = match outcome {
609            Ok(BoundedWait::Interrupted {
610                kill: Ok(()),
611                operation_elapsed_ms,
612                cleanup,
613            }) => (operation_elapsed_ms, cleanup),
614            other => {
615                assert!(
616                    matches!(other, Ok(BoundedWait::Interrupted { .. })),
617                    "interrupted command did not return cleanup evidence"
618                );
619                return;
620            }
621        };
622        assert!(operation_elapsed_ms < 1_000);
623        assert!(cleanup.verified, "{cleanup:?}");
624        assert_eq!(cleanup.trigger, CommandCleanupTrigger::Interruption);
625        assert!(cleanup.kill_attempted);
626        assert_eq!(cleanup.reap_grace_ms, 5_000);
627        assert_eq!(cleanup.io_drain_grace_ms, 5_000);
628    }
629
630    #[test]
631    fn blocked_stdin_write_cannot_outlive_the_owner_bound() {
632        let argv = vec![
633            "sh".to_owned(),
634            "-c".to_owned(),
635            "while :; do :; done".to_owned(),
636        ];
637        let payload = vec![b'x'; 1024 * 1024];
638        let bound = OperationBound::finite(Duration::from_millis(50));
639
640        let outcome = run_with_bound(&argv, None, Some(&payload), &bound, None);
641
642        assert!(matches!(
643            outcome,
644            Ok(BoundedWait::Expired { kill: Ok(()), .. })
645        ));
646    }
647
648    #[test]
649    fn inherited_output_pipe_cannot_outlive_the_owner_bound() {
650        let argv = vec![
651            "sh".to_owned(),
652            "-c".to_owned(),
653            "sleep 60 & exit 0".to_owned(),
654        ];
655        let bound = OperationBound::finite(Duration::from_millis(50));
656
657        let outcome = run_with_bound(&argv, None, None, &bound, None);
658
659        assert!(matches!(
660            outcome,
661            Ok(BoundedWait::Expired { kill: Ok(()), .. })
662        ));
663    }
664}
inferlab_runtime/container.rs

inferlab_runtime/
container.rs
