Skip to main content

inferlab_runtime/server/
observation.rs

1use super::{
2    HostProcessHandle, LogSyncError, ProcessCommandError, ProcessHandle, ProcessObserver,
3    ProcessStatus, REMOTE_LOG_SYNC_DEADLINE, SshProcessHandle, SystemProcessRuntime,
4};
5use crate::operation_bound::OperationBound;
6use crate::process_group::process_start_time;
7use crate::shell::shell_quote_path;
8use crate::ssh::{ssh_argv, ssh_output};
9use std::fs;
10use std::path::Path;
11use std::process::{Command, Output};
12
13pub(super) fn verified_local_status(handle: &HostProcessHandle) -> ProcessStatus {
14    verified_local_status_under(handle, None, false)
15}
16
17pub(super) fn verified_local_status_with_bound(
18    handle: &HostProcessHandle,
19    bound: &OperationBound,
20) -> ProcessStatus {
21    verified_local_status_under(handle, Some(bound), false)
22}
23
24pub(super) fn verified_local_status_under(
25    handle: &HostProcessHandle,
26    bound: Option<&OperationBound>,
27    cleanup: bool,
28) -> ProcessStatus {
29    if let Err(error) = handle.validate() {
30        return status_error(error);
31    }
32    match process_start_time(handle.leader_pid) {
33        Ok(Some(actual)) if actual != handle.leader_start_time_ticks => ProcessStatus {
34            queried: true,
35            alive: false,
36            error: Some(format!(
37                "managed process {} exited and its pid was reused: recorded start time {}, observed {}",
38                handle.leader_pid, handle.leader_start_time_ticks, actual
39            )),
40        },
41        Ok(Some(_)) => {
42            match process_group_has_live_members_under(handle.process_group, bound, cleanup) {
43                Ok(alive) => ProcessStatus {
44                    queried: true,
45                    alive,
46                    error: None,
47                },
48                Err(error) => status_error(error.to_string()),
49            }
50        }
51        Ok(None) => {
52            match process_group_has_live_members_under(handle.process_group, bound, cleanup) {
53                Ok(false) => ProcessStatus {
54                    queried: true,
55                    alive: false,
56                    error: None,
57                },
58                Ok(true) => status_error(format!(
59                    "process-group {} still has members but recorded leader {} no longer exists; ownership cannot be verified",
60                    handle.process_group, handle.leader_pid
61                )),
62                Err(error) => status_error(error.to_string()),
63            }
64        }
65        Err(error) => status_error(error.to_string()),
66    }
67}
68
69pub(super) fn verified_ssh_status(handle: &SshProcessHandle) -> ProcessStatus {
70    verified_ssh_status_under(handle, None)
71}
72
73pub(super) fn verified_ssh_status_with_bound(
74    handle: &SshProcessHandle,
75    bound: &OperationBound,
76) -> ProcessStatus {
77    verified_ssh_status_under(handle, Some(bound))
78}
79
80pub(super) fn verified_ssh_status_under(
81    handle: &SshProcessHandle,
82    bound: Option<&OperationBound>,
83) -> ProcessStatus {
84    if let Err(error) = handle.validate() {
85        return status_error(error);
86    }
87    let script = format!(
88        "set -eu; pid={}; expected={}; if [ -r /proc/$pid/stat ]; then actual=$(awk '{{print $22}}' /proc/$pid/stat); if [ \"$actual\" != \"$expected\" ]; then printf 'stale %s\\n' \"$actual\"; exit 4; fi; elif {}; then printf 'unknown leader-missing\\n'; exit 5; else printf 'dead\\n'; exit 3; fi; if {}; then printf 'alive\\n'; exit 0; fi; printf 'dead\\n'; exit 3",
89        handle.leader_pid,
90        handle.leader_start_time_ticks,
91        remote_group_alive_script(&handle.process_group.to_string()),
92        remote_group_alive_script(&handle.process_group.to_string()),
93    );
94    let output = match bound {
95        Some(bound) => run_status_command(&ssh_argv(&handle.target, &script), bound),
96        None => ssh_output(&handle.target, &script).map_err(|source| ProcessCommandError::Ssh {
97            operation: "process status command".to_owned(),
98            source,
99        }),
100    };
101    match output {
102        Ok(output) if output.status.success() => ProcessStatus {
103            queried: true,
104            alive: true,
105            error: None,
106        },
107        Ok(output) if output.status.code() == Some(3) => ProcessStatus {
108            queried: true,
109            alive: false,
110            error: None,
111        },
112        Ok(output) if output.status.code() == Some(4) => ProcessStatus {
113            queried: true,
114            alive: false,
115            error: Some(format!(
116                "managed SSH process {} exited and its pid was reused: {}",
117                handle.leader_pid,
118                String::from_utf8_lossy(&output.stdout).trim()
119            )),
120        },
121        Ok(output) if output.status.code() == Some(5) => status_error(format!(
122            "SSH process-group {} ownership could not be verified: {}",
123            handle.process_group,
124            String::from_utf8_lossy(&output.stdout).trim()
125        )),
126        Ok(output) => status_error(format!(
127            "SSH status exited with {}: {}",
128            output.status,
129            String::from_utf8_lossy(&output.stderr).trim()
130        )),
131        Err(error) => status_error(error.to_string()),
132    }
133}
134
135fn status_error(error: String) -> ProcessStatus {
136    ProcessStatus {
137        queried: false,
138        alive: false,
139        error: Some(error),
140    }
141}
142
143pub(super) fn remote_group_alive_script(group: &str) -> String {
144    format!(
145        "ps -eo pgid=,stat= | awk -v pgid={group} '$1 == pgid && $2 !~ /^Z/ {{ found=1 }} END {{ exit !found }}'"
146    )
147}
148
149pub(super) fn fetch_remote_file(
150    target: &str,
151    remote: &Path,
152    local: &Path,
153    bound: &OperationBound,
154    cleanup: bool,
155) -> Result<(), LogSyncError> {
156    let argv = ssh_argv(target, &format!("cat -- {}", shell_quote_path(remote)));
157    let output = if cleanup {
158        run_cleanup_command(&argv, bound, "remote log synchronization")
159    } else {
160        run_status_command(&argv, bound)
161    }
162    .map_err(|source| LogSyncError::ReadRemote {
163        path: remote.to_path_buf(),
164        source,
165    })?;
166    if !output.status.success() {
167        return Err(LogSyncError::RemoteExit {
168            path: remote.to_path_buf(),
169            status: output.status,
170            stderr: String::from_utf8_lossy(&output.stderr).trim().to_owned(),
171        });
172    }
173    fs::write(local, output.stdout).map_err(|source| LogSyncError::WriteLocal {
174        path: local.to_path_buf(),
175        source,
176    })
177}
178
179fn process_group_has_live_members_under(
180    process_group: u32,
181    bound: Option<&OperationBound>,
182    cleanup: bool,
183) -> Result<bool, ProcessCommandError> {
184    let argv = ["ps", "-eo", "pid=,pgid=,stat="];
185    let output = match bound {
186        Some(bound) if cleanup => run_cleanup_command(&argv, bound, "process cleanup status"),
187        Some(bound) => run_status_command(&argv, bound),
188        None => Command::new(argv[0])
189            .args(&argv[1..])
190            .output()
191            .map_err(|source| ProcessCommandError::Launch {
192                operation: "process-group query".to_owned(),
193                source,
194            }),
195    }?;
196    if !output.status.success() {
197        return Err(ProcessCommandError::Exit {
198            operation: "process-group query".to_owned(),
199            status: output.status,
200            stderr: String::from_utf8_lossy(&output.stderr).trim().to_owned(),
201        });
202    }
203    let process_group = process_group.to_string();
204    Ok(String::from_utf8_lossy(&output.stdout)
205        .lines()
206        .filter_map(|line| {
207            let mut fields = line.split_whitespace();
208            let _pid = fields.next()?;
209            let group = fields.next()?;
210            let state = fields.next()?;
211            Some((group, state))
212        })
213        .any(|(group, state)| group == process_group && !state.starts_with('Z')))
214}
215
216pub(super) fn run_status_command<S: AsRef<std::ffi::OsStr>>(
217    argv: &[S],
218    bound: &OperationBound,
219) -> Result<Output, ProcessCommandError> {
220    let operation = "process status command";
221    match crate::container::run_with_bound(argv, None, None, bound, None) {
222        Ok(crate::container::BoundedWait::Exited {
223            status,
224            stdout,
225            stderr,
226        }) => Ok(Output {
227            status,
228            stdout,
229            stderr,
230        }),
231        Ok(crate::container::BoundedWait::Expired { kill, .. }) => {
232            kill.map_err(|source| ProcessCommandError::Io {
233                operation: "process status cleanup".to_owned(),
234                source,
235            })?;
236            Err(ProcessCommandError::Deadline {
237                operation: "process status attempt".to_owned(),
238            })
239        }
240        Ok(crate::container::BoundedWait::Interrupted { kill, .. }) => {
241            kill.map_err(|source| ProcessCommandError::Io {
242                operation: "process status cleanup".to_owned(),
243                source,
244            })?;
245            Err(ProcessCommandError::Interrupted {
246                operation: "process status attempt".to_owned(),
247            })
248        }
249        Err(crate::container::BoundedError::Launch(source)) => Err(ProcessCommandError::Launch {
250            operation: operation.to_owned(),
251            source,
252        }),
253        Err(
254            crate::container::BoundedError::Stdin(error)
255            | crate::container::BoundedError::Wait(error),
256        ) => Err(ProcessCommandError::Io {
257            operation: operation.to_owned(),
258            source: error,
259        }),
260        Err(crate::container::BoundedError::WaitCleanup {
261            source, cleanup, ..
262        }) => Err(ProcessCommandError::WaitCleanup {
263            operation: operation.to_owned(),
264            source,
265            cleanup: cleanup.error.unwrap_or_else(|| {
266                if cleanup.verified {
267                    "verified"
268                } else {
269                    "unverified"
270                }
271                .to_owned()
272            }),
273        }),
274    }
275}
276
277pub(super) fn run_cleanup_command<S: AsRef<std::ffi::OsStr>>(
278    argv: &[S],
279    bound: &OperationBound,
280    operation: &str,
281) -> Result<Output, ProcessCommandError> {
282    match crate::container::run_cleanup_with_bound(argv, None, None, bound, None) {
283        Ok(crate::container::BoundedWait::Exited {
284            status,
285            stdout,
286            stderr,
287        }) => Ok(Output {
288            status,
289            stdout,
290            stderr,
291        }),
292        Ok(crate::container::BoundedWait::Expired { kill, .. }) => {
293            kill.map_err(|source| ProcessCommandError::Io {
294                operation: format!("{operation} child cleanup"),
295                source,
296            })?;
297            Err(ProcessCommandError::Deadline {
298                operation: operation.to_owned(),
299            })
300        }
301        Ok(crate::container::BoundedWait::Interrupted { kill, .. }) => {
302            kill.map_err(|source| ProcessCommandError::Io {
303                operation: format!("{operation} child cleanup"),
304                source,
305            })?;
306            Err(ProcessCommandError::Interrupted {
307                operation: operation.to_owned(),
308            })
309        }
310        Err(crate::container::BoundedError::Launch(source)) => Err(ProcessCommandError::Launch {
311            operation: operation.to_owned(),
312            source,
313        }),
314        Err(
315            crate::container::BoundedError::Stdin(error)
316            | crate::container::BoundedError::Wait(error),
317        ) => Err(ProcessCommandError::Io {
318            operation: operation.to_owned(),
319            source: error,
320        }),
321        Err(crate::container::BoundedError::WaitCleanup {
322            source, cleanup, ..
323        }) => Err(ProcessCommandError::WaitCleanup {
324            operation: operation.to_owned(),
325            source,
326            cleanup: cleanup.error.unwrap_or_else(|| {
327                if cleanup.verified {
328                    "verified"
329                } else {
330                    "unverified"
331                }
332                .to_owned()
333            }),
334        }),
335    }
336}
337
338impl ProcessObserver for SystemProcessRuntime {
339    fn status(&self, handle: &ProcessHandle) -> ProcessStatus {
340        match handle {
341            ProcessHandle::Local(handle) => verified_local_status(handle),
342            ProcessHandle::Ssh(handle) => verified_ssh_status(handle),
343        }
344    }
345
346    fn status_with_bound(&self, handle: &ProcessHandle, bound: &OperationBound) -> ProcessStatus {
347        match handle {
348            ProcessHandle::Local(handle) => verified_local_status_with_bound(handle, bound),
349            ProcessHandle::Ssh(handle) => verified_ssh_status_with_bound(handle, bound),
350        }
351    }
352
353    fn sync_logs(
354        &self,
355        handle: &ProcessHandle,
356        stdout: &Path,
357        stderr: &Path,
358        cleanup: bool,
359    ) -> Result<(), LogSyncError> {
360        match handle {
361            ProcessHandle::Local(_) => Ok(()),
362            ProcessHandle::Ssh(handle) => {
363                let bound = OperationBound::finite(REMOTE_LOG_SYNC_DEADLINE);
364                fetch_remote_file(&handle.target, &handle.stdout, stdout, &bound, cleanup)?;
365                fetch_remote_file(&handle.target, &handle.stderr, stderr, &bound, cleanup)
366            }
367        }
368    }
369}