tatara-engine 0.2.209

Drivers, nix eval, cluster logic, P2P, and scheduling engine for tatara
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

use anyhow::{bail, Context, Result};
use async_trait::async_trait;
use chrono::Utc;
use std::path::Path;
use std::process::Stdio;
use std::time::Duration;
use tokio::process::Command;
use tokio::sync::mpsc;
use tracing::debug;

use tatara_core::domain::allocation::TaskRunState;
use tatara_core::domain::job::{DriverType, Task, TaskConfig};

use super::exec::{is_process_alive, send_signal_and_wait};
use super::{Driver, LogEntry, TaskHandle};

pub struct NixDriver;

#[async_trait]
impl Driver for NixDriver {
    fn name(&self) -> &str {
        "nix"
    }

    async fn available(&self) -> bool {
        Command::new("nix")
            .arg("--version")
            .output()
            .await
            .map(|o| o.status.success())
            .unwrap_or(false)
    }

    async fn start(&self, task: &Task, alloc_dir: &Path) -> Result<TaskHandle> {
        let (flake_ref, args) = match &task.config {
            TaskConfig::Nix { flake_ref, args } => (flake_ref.clone(), args.clone()),
            _ => bail!("NixDriver received non-nix task config"),
        };

        let log_dir = alloc_dir.join(&task.name);
        tokio::fs::create_dir_all(&log_dir).await?;

        let stdout_file = std::fs::File::create(log_dir.join("stdout.log"))
            .context("Failed to create stdout log")?;
        let stderr_file = std::fs::File::create(log_dir.join("stderr.log"))
            .context("Failed to create stderr log")?;

        // `nix run <flake_ref> -- <args>`
        let mut cmd = Command::new("nix");
        cmd.arg("run")
            .arg(&flake_ref)
            .arg("--")
            .args(&args)
            .envs(&task.env)
            .stdout(Stdio::from(stdout_file))
            .stderr(Stdio::from(stderr_file))
            .kill_on_drop(false);

        let child = cmd.spawn().context("Failed to spawn nix run")?;
        let pid = child.id();

        debug!(
            task = %task.name,
            flake_ref = %flake_ref,
            pid = ?pid,
            "Nix driver started process"
        );

        std::mem::forget(child);

        Ok(TaskHandle {
            driver: DriverType::Nix,
            pid,
            container_id: None,
            started_at: Utc::now(),
        })
    }

    async fn stop(&self, handle: &TaskHandle, timeout: Duration) -> Result<()> {
        let pid = handle.pid.context("No PID in nix task handle")?;
        send_signal_and_wait(pid, timeout).await
    }

    async fn status(&self, handle: &TaskHandle) -> Result<TaskRunState> {
        let pid = handle.pid.context("No PID in nix task handle")?;
        if is_process_alive(pid) {
            Ok(TaskRunState::Running)
        } else {
            Ok(TaskRunState::Dead)
        }
    }

    async fn logs(&self, _handle: &TaskHandle) -> Result<mpsc::Receiver<LogEntry>> {
        let (tx, rx) = mpsc::channel(256);
        tokio::spawn(async move {
            let _ = tx;
        });
        Ok(rx)
    }
}