Skip to main content

entrenar/prune/pipeline/
stage.rs

1//! Pruning pipeline stage enum
2//!
3//! Defines the stages of the pruning pipeline workflow.
4
5use serde::{Deserialize, Serialize};
6
7/// Current stage of the pruning pipeline.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
9pub enum PruningStage {
10    /// Not started.
11    #[default]
12    Idle,
13    /// Collecting calibration data.
14    Calibrating,
15    /// Computing importance scores.
16    ComputingImportance,
17    /// Applying pruning masks.
18    Pruning,
19    /// Fine-tuning after pruning.
20    FineTuning,
21    /// Evaluating pruned model.
22    Evaluating,
23    /// Exporting pruned model.
24    Exporting,
25    /// Pipeline complete.
26    Complete,
27    /// Pipeline failed.
28    Failed,
29}
30
31impl PruningStage {
32    /// Check if the pipeline is in an active (non-terminal) state.
33    pub fn is_active(&self) -> bool {
34        matches!(
35            self,
36            PruningStage::Calibrating
37                | PruningStage::ComputingImportance
38                | PruningStage::Pruning
39                | PruningStage::FineTuning
40                | PruningStage::Evaluating
41                | PruningStage::Exporting
42        )
43    }
44
45    /// Check if the pipeline is complete (success or failure).
46    pub fn is_terminal(&self) -> bool {
47        matches!(self, PruningStage::Complete | PruningStage::Failed)
48    }
49
50    /// Get display name for the stage.
51    pub fn display_name(&self) -> &'static str {
52        match self {
53            PruningStage::Idle => "Idle",
54            PruningStage::Calibrating => "Calibrating",
55            PruningStage::ComputingImportance => "Computing Importance",
56            PruningStage::Pruning => "Pruning",
57            PruningStage::FineTuning => "Fine-Tuning",
58            PruningStage::Evaluating => "Evaluating",
59            PruningStage::Exporting => "Exporting",
60            PruningStage::Complete => "Complete",
61            PruningStage::Failed => "Failed",
62        }
63    }
64}