The cpu to query can now be specified

it defaults to zero, which is what was used for MSRs before. for CPUID before tasksel or similar would be needed. It also honors this option
ScaleComputing · Sep 12, 2024 · c472517 · c472517
1 parent 4f615df
commit c472517
Show file tree

Hide file tree

Showing 5 changed files with 73 additions and 10 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -13,6 +13,7 @@ serde = { version = "1.0", features = ["derive"] }
 serde_yaml = "0.8"
 enum_dispatch = "0.3.8"
 serde_json = "1.0.117"
+core_affinity = "^0.8.1"
 
 [target.'cfg(target_os = "linux")'.dependencies]
 kvm-ioctls = { version = "0.17", optional = true }

diff --git a/src/kvm.rs b/src/kvm.rs
@@ -86,6 +86,6 @@ impl MsrStore for KvmMsrInfo {
                     None
                 }
             })
-            .ok_or_else(|| msr::Error::NotAvailible {})
+            .ok_or_else(|| msr::Error::NotAvailible("/dev/kvm".to_string()))
     }
 }
diff --git a/src/main.rs b/src/main.rs
@@ -3,6 +3,7 @@
 // use std::io::BufWriter;
 
 use clap::{self, Args, Parser, Subcommand, ValueEnum};
+use core_affinity::CoreId;
 use cpuinfo::facts::{FactSet, Facter, GenericFact};
 use cpuinfo::layout::LeafDesc;
 use cpuinfo::msr::MsrStore;
@@ -36,6 +37,8 @@ enum CommandOpts {
 
 #[derive(Clone, Args)]
 struct Disp {
+    #[arg(short, long, default_value = "0")]
+    cpu: usize,
     #[arg(short, long)]
     raw: bool,
     #[arg(long)]
@@ -50,6 +53,9 @@ struct Disp {
 
 impl Command for Disp {
     fn run(&self, config: &Definition) -> Result<(), Box<dyn std::error::Error>> {
+        if !core_affinity::set_for_current(CoreId { id: self.cpu }) {
+            panic!("Unable to pin to core {}", self.cpu);
+        }
         if self.raw {
             display_raw()
         } else {
@@ -86,7 +92,7 @@ impl Command for Disp {
             if !self.skip_msr {
                 #[cfg(target_os = "linux")]
                 {
-                    match msr::linux::LinuxMsrStore::new() {
+                    match msr::linux::LinuxMsrStore::new(self.cpu) {
                         Ok(linux_store) => {
                             println!("MSRS:");
                             for msr in &config.msrs {
@@ -132,6 +138,8 @@ enum FactsOutput {
 
 #[derive(Clone, Args)]
 struct Facts {
+    #[arg(short, long, default_value = "0")]
+    cpu: usize,
     #[cfg(all(target_os = "linux", feature = "kvm"))]
     #[arg(short, long)]
     use_kvm: bool,
@@ -172,6 +180,9 @@ fn collect_facts(
 
 impl Command for Facts {
     fn run(&self, config: &Definition) -> Result<(), Box<dyn std::error::Error>> {
+        if !core_affinity::set_for_current(CoreId { id: self.cpu }) {
+            panic!("Unable to pin to core {}", self.cpu);
+        }
         let (cpuid_source, msr_source): (_, Box<dyn MsrStore>) = {
             #[cfg(all(target_os = "linux", feature = "kvm"))]
             {
@@ -188,7 +199,7 @@ impl Command for Facts {
                     let msr = {
                         #[cfg(feature = "use_msr")]
                         {
-                            match msr::linux::LinuxMsrStore::new() {
+                            match msr::linux::LinuxMsrStore::new(self.cpu) {
                                 Ok(store) => Box::new(store) as Box<dyn MsrStore>,
                                 Err(e) => {
                                     eprintln!("Error accessing MSRs: {}", e);

diff --git a/src/msr.rs b/src/msr.rs
@@ -9,7 +9,7 @@ use std::{convert, error, fmt, io};
 
 #[derive(Debug)]
 pub enum Error {
-    NotAvailible,
+    NotAvailible(String),
     IOError(io::Error),
 }
 
@@ -18,7 +18,7 @@ pub type Result<V> = std::result::Result<V, Error>;
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
-            Error::NotAvailible => write!(f, "MSR Feature not availible"),
+            Error::NotAvailible(name) => write!(f, "MSR Feature not availible file: {}", name),
             Error::IOError(e) => write!(f, "IOError: {}", e),
         }
     }
@@ -48,7 +48,7 @@ impl MsrStore for EmptyMSR {
         true
     }
     fn get_value<'a>(&self, _desc: &'a MSRDesc) -> std::result::Result<MSRValue<'a>, Error> {
-        Err(Error::NotAvailible)
+        Err(Error::NotAvailible("".to_string()))
     }
 }
 
@@ -63,14 +63,15 @@ pub mod linux {
     }
 
     impl LinuxMsrStore {
-        pub fn new() -> Result<LinuxMsrStore> {
+        pub fn new(cpu: usize) -> Result<LinuxMsrStore> {
+            let file_name = format!("/dev/cpu/{}/msr", cpu);
             Ok(LinuxMsrStore {
                 msr_device: fs::OpenOptions::new()
                     .read(true)
-                    .open("/dev/cpu/0/msr")
+                    .open(file_name.clone())
                     .map_err(|e| match e.kind() {
-                        io::ErrorKind::NotFound => Error::NotAvailible,
-                        io::ErrorKind::PermissionDenied => Error::NotAvailible,
+                        io::ErrorKind::NotFound => Error::NotAvailible(file_name),
+                        io::ErrorKind::PermissionDenied => Error::NotAvailible(file_name),
                         _ => Error::IOError(e),
                     })?,
             })