captain-workflow-manager 0.1.0

//! This library helps you run pipelines of jobs that depend on each other.
//! Its modularity allows to run jobs either locally, on a cluster or other remote computing resources.
//! 
//! Consider a set of jobs as illustrated below, where two kind of source files are used to generate
//! two kind of intermediate files which are in turn used to generate a final result.
//!
//! ```text
//!                          ┌─────────────────┐
//! ┌─────────────────┐      │  Source File B: │
//! │  Source File A: │      │    - Param1     │
//! │    - Param1     │      │    - Param2     │
//! └─────────┬───────┘      └───────┬─────────┘
//!           │                      │
//!           │                      │
//!           │           ┌──────────┴───────────────┐
//!           │           │                          │
//!           │           │                          │
//!           │           │                          │
//!           ▼           ▼                          ▼
//!       ┌──────────────────────┐       ┌──────────────────────┐
//!       │ Intermediate File A: │       │ Intermediate File B: │
//!       │   - Param1           │       │   - Param1           │
//!       │   - Param2           │       │   - Param2           │
//!       └────────────────┬─────┘       └──────┬───────────────┘
//!                        │                    │
//!                        └─────────┬──────────┘
//!                                  │
//!                                  │
//!                                  ▼
//!                          ┌───────────────┐
//!                          │ Final Result: │
//!                          │   - Param1    │
//!                          │   - Param2    │
//!                          └───────────────┘
//! ```
//! 
//! Each kind of file has one or more "parameters".
//! 
//! This dependency could be represented by the following `Job` enum, assuming that `Param1` is an integer and `Param2` a string:
//! 
//! ```rust
//! enum Job {
//!     SourceFileA {param1: u16},
//!     SourceFileB {param1: u16, param2: &'static str},
//!     IntermediateFileA {param1: u16, param2: &'static str},
//!     IntermediateFileB {param1: u16, param2: &'static str},
//!     FinalResult {param1: u16, param2: &'static str},
//! }
//! ```
//! 
//! To manage this set of jobs using *captain*, one would first implement the [`JobUnit`] trait on it. And then run it using the job [`Scheduler`].
//! The back-end to run it on is chosen by selecting an [`ExecutorBuilder`].
use std::{any::Any, collections::{HashMap, HashSet}, fmt::{self, Debug}, hash::Hash, io::{self, Write}, path::PathBuf};

use executor::ExecutorBuilder;
use log::trace;
use thiserror::Error;

use crossbeam_channel::SendError;
use petgraph::prelude::*;

use crate::executor::{Executor, ExecutorResult};

pub mod executor;

/// A trait that must be implemented by struct representing job units.
/// 
/// A single struct is used to represent all the possible job units and their dependencies,
/// so most of the checks are performed at run-time.
pub trait JobUnit: Debug + Any + Sized + Copy + Hash + Eq + Send {
    /// Generate a vector containing all dependencies of the `self` 
    fn deps(&self) -> Vec<Self>;
    /// Returns the command to run to generate the output file of the job
    /// 
    /// If no command is specified (ie. this function returns [`None`]),
    /// the output file must be present on disk.
    fn cmd(&self) -> Option<String>;
    /// Path of the output file
    fn out_file(&self) -> PathBuf;
    /// Path of the log file of the job, if any
    fn log_file(&self) -> Option<PathBuf>;
    /// Name of the job
    fn name(&self) -> &'static str;

    /// A nice-ness score akin to Unix niceness: the higher the highest priority.
    fn nice(&self) -> u16 {
        0
    }
}

struct ExpandedDescr<'a, J: JobUnit>(&'a J);

impl<'a, J: JobUnit> Debug for ExpandedDescr<'a, J> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let inner = self.0;
        let mut dbg_struct = f.debug_struct(inner.name());
        dbg_struct
        .field("dependencies", &inner.deps())
        .field("output file",&inner.out_file());
        match &inner.cmd() {
            Some(cmd) =>  dbg_struct.field("command",cmd),
            x@None => dbg_struct.field("command",&x),
        };
        match &inner.log_file() {
            Some(log) => dbg_struct.field("log",log),
            x@None => dbg_struct.field("log",&x),
        };
        dbg_struct.finish()
    }
}

/// A job to be run by an Executor
#[derive(Debug)]
pub struct ToRun<J: JobUnit> {
    job_idx: NodeIndex,
    job_unit: J,
}

impl<J: JobUnit> ToRun<J> {
    /// Get the job number of the job to run.
    pub fn job_idx(&self) -> NodeIndex {
        self.job_idx
    }

    /// The job to run description
    pub fn job_unit(&self) -> &J {
        &self.job_unit
    }
}

/// The entry point of the crate, implementing the job scheduling logic
/// 
/// Example
/// ------
/// 
/// ```no_run
/// use captain_workflow_manager::{executor, Scheduler};
/// # use std::path::PathBuf;
/// # #[derive(Debug,Clone, Copy, PartialEq, Eq, Hash)]
/// # struct J {}
/// # impl captain_workflow_manager::JobUnit for J {
/// # fn deps(&self) -> Vec<Self> { vec![]}
/// # fn cmd(&self) -> Option<String> {None}
/// # fn out_file(&self) -> PathBuf { PathBuf::new()}
/// # fn log_file(&self) -> Option<PathBuf> {None}
/// # fn name(&self) -> &'static str {&""}
/// # }
/// # let some_jobs = unimplemented!();
/// let executor = executor::SlurmExecutorBuilder { max_jobs: 10_000 };
/// let scheduler = Scheduler::new(executor);
/// scheduler.run(some_jobs)?;
/// # Ok::<(), captain_workflow_manager::Error<J>>(())
/// ```
pub struct Scheduler<ExB: ExecutorBuilder> {
    ex_builder: ExB,
}

/// The Error type of this crate
#[derive(Debug, Error)]
pub enum Error<J: JobUnit> {
    #[error("error sending job to worker")]
    ErrorSendingJob(#[from] SendError<J>),
    #[error("system I/O error")]
    IoError(#[from] io::Error),
    #[error("Error in subcommand: {msg}")]
    SubCommandError { msg: String },
    #[error("Command-less job unit file is absent: {file}")]
    MissingCommandLessOutputFile { file: PathBuf },
}

impl<ExB: ExecutorBuilder> Scheduler<ExB> {
    /// Creates a new scheduler objects, with the provided [`ExecutorBuilder`]
    /// 
    /// Executor builders can be found in the [`executor`] module, or custom ones can be implemented.
    pub fn new(ex_builder: ExB) -> Self {
        Scheduler { ex_builder }
    }

    /// Run each job in `final_job_units` and its transitive dependencies.
    /// 
    /// A job won't be run if its output file is already present.
    /// (Currently no finer dependency freshness management is attempted.)
    /// Moreover, dependencies of such jobs won't be run either.
    pub fn run<J: JobUnit>(self, final_job_units: Vec<J>) -> Result<(), Error<J>> {
        let mut graph = petgraph::stable_graph::StableDiGraph::new();
        let mut stack = final_job_units;
        let mut already_processed = HashSet::new();
        let mut indexes = HashMap::new();
        while let Some(job_unit) = stack.pop() {
            let newly_seen = already_processed.insert(job_unit);
            if !newly_seen {
                continue;
            }
            if job_unit.out_file().exists() {
                continue;
            }
            let &mut ju_index = indexes
                .entry(job_unit)
                .or_insert_with_key(|j| graph.add_node(Some(*j)));
            let deps = job_unit.deps();
            for dep in deps {
                if dep.out_file().exists() {
                    continue;
                }
                if !already_processed.contains(&dep) {
                    stack.push(dep);
                }
                let &mut dep_index = indexes
                    .entry(dep)
                    .or_insert_with_key(|j| graph.add_node(Some(*j)));
                debug_assert!(!graph.contains_edge(ju_index, dep_index));
                graph.add_edge(ju_index, dep_index, ());
            }
        }
        let (executor, to_run_channel_sndr, done_channel_sndr, done_channel_rcvr) =
            self.ex_builder.init();

        let no_deps_nodes: Vec<_> = graph
            .node_indices()
            .filter(|&idx| graph.neighbors_directed(idx, Outgoing).next().is_none())
            .collect();
        let start_idx = graph.add_node(None);
        for job_idx in no_deps_nodes {
            graph.add_edge(job_idx, start_idx, ());
        }
        let start_count: usize = graph.node_count();
        let mut done_count: usize = 0;
        let mut failed_count: usize = 0;
        done_channel_sndr
            .send(ExecutorResult {
                job_idx: start_idx,
                result: Ok(()),
            })
            .unwrap();

        let done_iter = done_channel_rcvr.into_iter();
        trace!("Entering done_iter main loop.");
        for done in done_iter {
            trace!("Iterating done_iter main loop.");
            let j_idx = done.job_idx;
            if let Err(msg) = done.result {
                trace!("Job errored.");
                let job_unit = graph.node_weight(j_idx).unwrap().unwrap();
                {
                    let stderr = io::stderr();
                    let mut lock = stderr.lock();
                    writeln!(
                        lock,
                        "There was an error running the following job: {:#?} interpreted as {:#?}\nMessage was: {}",
                        job_unit, ExpandedDescr(&job_unit), msg
                    )
                    .unwrap();

                    let mut stack = vec![j_idx];
                    while let Some(idx) = stack.pop() {
                        stack.extend(graph.neighbors_directed(idx, Incoming));
                        writeln!(
                            lock,
                            "Because a dependency failed, deleting job: {:?}\n",
                            graph.node_weight(idx).unwrap()
                        )
                        .unwrap();
                        graph.remove_node(idx);
                        failed_count += 1;
                    }
                }
                continue;
                // return Err(Error::SubCommandError { msg });
            }
            trace!("Job was succesful.");

            let was_required_by: Vec<_> = graph.neighbors_directed(j_idx, Incoming).collect();

            graph.remove_node(j_idx);
            done_count += 1;
            assert_eq!(done_count + failed_count + graph.node_count(), start_count);
            // eprintln!("DEBUG:foo");
            trace!("Writting progress.");
            writeln!(
                io::stdout().lock(),
                "{} done ({:5.1}%), {} failed ({:5.1}%), {} remaining ({:5.1}%)\n",
                done_count,
                100. * (done_count as f64) / (start_count as f64),
                failed_count,
                100. * (failed_count as f64) / (start_count as f64),
                graph.node_count(),
                100. * (graph.node_count() as f64) / (start_count as f64)
            )
            .unwrap();
            if graph.node_count() == 0 {
                break;
            }
            // eprintln!("DEBUG:bar");
            trace!("Submitting dependent jobs.");
            for job_idx in was_required_by {
                if graph.neighbors_directed(job_idx, Outgoing).count() == 0 {
                    let j_u = graph.node_weight(job_idx).unwrap().unwrap();
                    let out = j_u.out_file();
                    if out.exists() {
                        writeln!(
                            io::stdout().lock(),
                            "Output already present, skipping job: {:?}\n",
                            j_u
                        )
                        .unwrap();
                        done_channel_sndr
                            .send(ExecutorResult {
                                result: Ok(()),
                                job_idx,
                            })
                            .unwrap();
                    } else {
                        if j_u.cmd().is_none() {
                            return Err(Error::MissingCommandLessOutputFile {
                                file: j_u.out_file(),
                            });
                        }
                        // eprintln!("DEBUG:zob");
                        to_run_channel_sndr
                            .send(ToRun {
                                job_idx,
                                job_unit: j_u,
                            })
                            .unwrap();
                        // eprintln!("DEBUG:kola");
                    }
                }
            }
        }
        drop(to_run_channel_sndr);
        executor.join();
        Ok(())
    }
}