Skip to main content

joule_profiler_cli/
lib.rs

1use std::{collections::HashSet, path::PathBuf};
2
3use clap::{ArgAction, Parser, ValueEnum};
4
5use anyhow::{Result, bail};
6pub use commands::ProfilerCommand;
7#[cfg(feature = "rapl")]
8use serde::Deserialize;
9
10use crate::{
11    config::{overrides::ConfigOverride, table::ConfigTable},
12    output::{
13        displayer::Displayer,
14        formats::{OutputFormat, csv::CsvOutput, json::JsonOutput, terminal::TerminalOutput},
15    },
16};
17
18mod commands;
19pub mod config;
20mod logging;
21mod output;
22
23/// joule-profiler: measure program energy consumption
24#[allow(clippy::struct_excessive_bools)]
25#[derive(Parser, Debug)]
26#[command(name = "joule-profiler")]
27#[command(
28    version,
29    about = "Measure program metrics from various sources like RAPL, perf_event or NVML"
30)]
31pub struct CliArgs {
32    /// Verbosity (-v, -vv, -vvv)
33    #[arg(short = 'v', long = "verbose", action = ArgAction::Count)]
34    pub verbose: u8,
35
36    /// Output format to export the results in. (e.g., terminal, json, csv)
37    #[arg(long = "output-format")]
38    pub output_format: Option<OutputFormat>,
39
40    /// Output file for CSV/JSON. (else `data<TIMESTAMP>`.csv/json)
41    #[arg(short = 'o', long = "output-file")]
42    pub output_file: Option<String>,
43
44    /// Sources activation list. All sources must be separated with a comma (e.g., "rapl,nvml").
45    #[cfg_attr(
46        feature = "rapl",
47        arg(long, value_delimiter = ',', default_value = "rapl")
48    )]
49    #[cfg_attr(not(feature = "rapl"), arg(long, value_delimiter = ','))]
50    pub sources: Vec<Source>,
51
52    #[allow(clippy::doc_markdown)]
53    /// Override a configuration key, repeatable. (e.g., -D profiler.rapl_backend=powercap)
54    ///
55    /// KEY is the dotted path of the key in the configuration file, so any
56    /// key a configuration file can set is settable here:
57    ///   -D profiler.output_format=json
58    ///   -D sources.rapl.sockets_spec=[0,1]
59    ///   -D sources.cgroup.create_cgroup=false
60    ///
61    /// VALUE is read as TOML, falling back to a plain string, so durations,
62    /// regexes and paths need no quoting. Overrides are applied on top of the
63    /// --config file, and configuring a source enables it.
64    #[arg(
65        short = 'D',
66        long = "define",
67        value_name = "KEY=VALUE",
68        verbatim_doc_comment
69    )]
70    pub overrides: Vec<ConfigOverride>,
71
72    /// TOML configuration file. Every key it sets can also be set with -D.
73    #[arg(long = "config")]
74    pub config_file: Option<PathBuf>,
75
76    /// The command to execute.
77    #[command(subcommand)]
78    pub command: ProfilerCommand,
79}
80
81impl CliArgs {
82    pub fn from_args() -> Self {
83        Self::parse()
84    }
85
86    pub fn validate(&self) -> Result<()> {
87        let mut seen = HashSet::new();
88
89        for source in &self.sources {
90            if !seen.insert(source) {
91                bail!("Duplicate source specified: {source}");
92            }
93        }
94
95        Ok(())
96    }
97}
98
99#[cfg(not(any(
100    feature = "rapl",
101    feature = "perf_event",
102    feature = "nvml",
103    feature = "amdsmi",
104    feature = "procfs",
105    feature = "cgroup",
106)))]
107compile_error!("At least one source feature must be enabled");
108
109#[derive(Clone, Debug, PartialEq, Eq, Hash, ValueEnum)]
110pub enum Source {
111    #[cfg(feature = "rapl")]
112    Rapl,
113
114    #[cfg(feature = "perf_event")]
115    #[value(alias = "perf_event")]
116    Perf,
117
118    #[cfg(feature = "nvml")]
119    Nvml,
120
121    #[cfg(feature = "amdsmi")]
122    #[value(name = "amdsmi")]
123    AmdSmi,
124
125    #[cfg(feature = "procfs")]
126    Procfs,
127
128    #[cfg(feature = "cgroup")]
129    Cgroup,
130}
131
132impl std::fmt::Display for Source {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        let s = match self {
135            #[cfg(feature = "rapl")]
136            Source::Rapl => "rapl",
137
138            #[cfg(feature = "perf_event")]
139            Source::Perf => "perf",
140
141            #[cfg(feature = "nvml")]
142            Source::Nvml => "nvml",
143
144            #[cfg(feature = "amdsmi")]
145            Source::AmdSmi => "amdsmi",
146
147            #[cfg(feature = "procfs")]
148            Source::Procfs => "procfs",
149
150            #[cfg(feature = "cgroup")]
151            Source::Cgroup => "cgroup",
152        };
153        write!(f, "{s}")
154    }
155}
156
157#[cfg(all(
158    feature = "rapl",
159    not(any(feature = "rapl-backend-perf", feature = "rapl-backend-powercap"))
160))]
161compile_error!("Enable at least one backend: `rapl-backend-perf` or `rapl-backend-powercap`.");
162
163/// RAPL backend, selected with `profiler.rapl_backend` in the configuration.
164///
165/// The RAPL source always comes with a backend, so this enum always has at
166/// least one variant.
167#[cfg(feature = "rapl")]
168#[derive(Debug, Default, Clone, Deserialize)]
169pub enum RaplBackend {
170    #[cfg(feature = "rapl-backend-perf")]
171    #[default]
172    #[serde(rename = "perf")]
173    Perf,
174
175    // Only the default when it is the sole backend built in: two `#[default]`
176    // variants do not compile.
177    #[cfg(feature = "rapl-backend-powercap")]
178    #[cfg_attr(not(feature = "rapl-backend-perf"), default)]
179    #[serde(rename = "powercap")]
180    Powercap,
181}
182
183/// Builds the displayer for the configured output format.
184///
185/// Reads only from `config_table`: `apply_cli` must have already resolved
186/// any CLI override into `profiler_config` beforehand.
187pub fn config_table_to_displayer(config_table: &ConfigTable) -> Result<Box<dyn Displayer>> {
188    let output_file = config_table.profiler_config.output_file.clone();
189
190    let displayer = match config_table.profiler_config.output_format {
191        OutputFormat::Terminal => TerminalOutput.into(),
192        OutputFormat::Json => JsonOutput::new(output_file)?.into(),
193        OutputFormat::Csv => CsvOutput::try_new(output_file)?.into(),
194    };
195
196    Ok(displayer)
197}
198
199pub fn init_logging(verbose: u8) {
200    logging::init_logging(verbose);
201}
202
203#[cfg(test)]
204mod tests {
205    use clap::ValueEnum;
206    use joule_profiler_core::source::MetricReader;
207
208    use super::*;
209
210    fn cli_args_with_sources(sources: Vec<Source>) -> CliArgs {
211        CliArgs {
212            verbose: 0,
213            output_format: None,
214            output_file: None,
215            sources,
216            overrides: Vec::new(),
217            config_file: None,
218            command: ProfilerCommand::ListSensors,
219        }
220    }
221
222    #[test]
223    fn source_display_matches_every_metric_reader_get_id() {
224        #[cfg(feature = "rapl-backend-perf")]
225        assert_eq!(
226            Source::Rapl.to_string(),
227            joule_profiler_source_rapl::perf::Rapl::get_id()
228        );
229        #[cfg(feature = "rapl-backend-powercap")]
230        assert_eq!(
231            Source::Rapl.to_string(),
232            joule_profiler_source_rapl::powercap::Rapl::get_id()
233        );
234        #[cfg(feature = "perf_event")]
235        {
236            type DefaultPerfEvent = joule_profiler_source_perf_event::PerfEvent;
237            assert_eq!(Source::Perf.to_string(), DefaultPerfEvent::get_id());
238        }
239        #[cfg(feature = "cgroup")]
240        {
241            type DefaultCgroup = joule_profiler_source_cgroup::Cgroup;
242            assert_eq!(Source::Cgroup.to_string(), DefaultCgroup::get_id());
243        }
244        #[cfg(feature = "procfs")]
245        {
246            type DefaultProcfs = joule_profiler_source_procfs::Procfs;
247            assert_eq!(Source::Procfs.to_string(), DefaultProcfs::get_id());
248        }
249        #[cfg(feature = "nvml")]
250        {
251            type DefaultNvml = joule_profiler_source_nvml::Nvml;
252            assert_eq!(Source::Nvml.to_string(), DefaultNvml::get_id());
253        }
254        #[cfg(feature = "amdsmi")]
255        {
256            type DefaultAmdSmi = joule_profiler_source_amdsmi::AmdSmi;
257            assert_eq!(Source::AmdSmi.to_string(), DefaultAmdSmi::get_id());
258        }
259    }
260
261    #[cfg(feature = "perf_event")]
262    #[test]
263    fn source_value_enum_accepts_perf_event_alias() {
264        let parsed = Source::from_str("perf_event", false).unwrap();
265        assert_eq!(parsed, Source::Perf);
266    }
267
268    /// Every variant this build has must parse back from the name it displays.
269    #[test]
270    fn source_value_enum_accepts_canonical_names() {
271        for source in Source::value_variants() {
272            let parsed = Source::from_str(&source.to_string(), false).unwrap();
273            assert_eq!(&parsed, source);
274        }
275    }
276
277    #[test]
278    fn validate_rejects_duplicate_sources() {
279        let source = Source::value_variants()[0].clone();
280        let cli = cli_args_with_sources(vec![source.clone(), source]);
281        assert!(cli.validate().is_err());
282    }
283
284    #[test]
285    fn validate_accepts_distinct_sources() {
286        let cli = cli_args_with_sources(Source::value_variants().to_vec());
287        assert!(cli.validate().is_ok());
288    }
289}