gflow 0.4.15

A lightweight, single-node job scheduler written in Rust.
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

use super::super::state::ServerState;
use axum::{
    extract::{Path, State},
    http::StatusCode,
    response::IntoResponse,
    Json,
};
use gflow::core::job::JobState;
use gflow::debug;
use std::collections::HashMap;

// Debug endpoints
#[axum::debug_handler]
pub(in crate::multicall::gflowd::server) async fn debug_state(
    State(server_state): State<ServerState>,
) -> impl IntoResponse {
    let state = server_state.scheduler.read().await;

    // Get GPU info from the info() method
    let info = state.info();
    let gpu_slots: Vec<debug::DebugGpuSlot> = info
        .gpus
        .iter()
        .map(|gpu_info| debug::DebugGpuSlot {
            uuid: gpu_info.uuid.clone(),
            index: gpu_info.index,
            available: gpu_info.available,
        })
        .collect();

    let debug_state = debug::DebugState {
        jobs: state.jobs(),
        next_job_id: state.next_job_id(),
        total_memory_mb: state.total_memory_mb(),
        available_memory_mb: state.available_memory_mb(),
        gpu_slots,
        allowed_gpu_indices: info.allowed_gpu_indices,
    };

    (StatusCode::OK, Json(debug_state))
}

#[axum::debug_handler]
pub(in crate::multicall::gflowd::server) async fn debug_job(
    State(server_state): State<ServerState>,
    Path(id): Path<u32>,
) -> impl IntoResponse {
    let state = server_state.scheduler.read().await;

    state
        .get_job(id)
        .map(debug::DebugJobInfo::from_job)
        .map(Json)
        .ok_or(StatusCode::NOT_FOUND)
}

#[axum::debug_handler]
pub(in crate::multicall::gflowd::server) async fn debug_metrics(
    State(server_state): State<ServerState>,
) -> impl IntoResponse {
    let state = server_state.scheduler.read().await;
    let jobs = state.jobs();

    let jobs_by_state: HashMap<JobState, usize> =
        jobs.iter().fold(HashMap::new(), |mut acc, job| {
            *acc.entry(job.state).or_insert(0) += 1;
            acc
        });

    let jobs_by_user: HashMap<String, debug::UserJobStats> =
        jobs.iter().fold(HashMap::new(), |mut acc, job| {
            let stats = acc
                .entry(job.submitted_by.to_string())
                .or_insert(debug::UserJobStats {
                    submitted: 0,
                    running: 0,
                    finished: 0,
                    failed: 0,
                });
            stats.submitted += 1;
            match job.state {
                JobState::Running => stats.running += 1,
                JobState::Finished => stats.finished += 1,
                JobState::Failed => stats.failed += 1,
                _ => {}
            }
            acc
        });

    let debug_metrics = debug::DebugMetrics {
        jobs_by_state,
        jobs_by_user,
    };

    (StatusCode::OK, Json(debug_metrics))
}