Merge branch 'threading' into 'master'

Threading

Closes #5 and #18

See merge request pythonspeed/memory-profiler!4

Author: itamarst

Cargo.toml

@@ -2,4 +2,9 @@
 
 members = [
         "memapi",
-]
+]
+
+[profile.release]
+# Needed for profiling, disable for real release:
+#debug = true
+lto = true

Makefile

@@ -20,7 +20,8 @@ filprofiler/libpymemprofile_api.so: Cargo.lock memapi/Cargo.toml memapi/src/*.rs
 	cp -f target/release/libpymemprofile_api.so filprofiler/
 
 test: build
-	fil-profile example.py
+	env RUST_BACKTRACE=1 cargo test
+	env RUST_BACKTRACE=1 py.test
 
 .PHONY: clean
 clean:

filprofiler/_script.py

@@ -38,14 +38,14 @@ def stage_1():
 
 def stage_2():
     """Main CLI interface. Presumes LD_PRELOAD etc. has been set by stage_1()."""
-    usage = "fil-profile [-o /path/to/output.svg] [-m module | /path/to/script.py ] [arg] ..."
+    usage = "fil-profile [-o /path/to/output-dir/] [-m module | /path/to/script.py ] [arg] ..."
     parser = ArgumentParser(usage=usage)
     parser.add_argument(
         "-o",
-        dest="svg_output_path",
+        dest="output_path",
         action="store",
-        default="memory-graph.svg",
-        help="Path where the SVG with the profile flamegraph is written.",
+        default="fil-result",
+        help="Directory where the profiling results written.",
     )
     parser.add_argument(
         "-m",
@@ -74,7 +74,7 @@ def stage_2():
             "__package__": None,
             "__cached__": None,
         }
-    trace(code, globals_, arguments.svg_output_path)
+    trace(code, globals_, arguments.output_path)
 
 
 if __name__ == "__main__":

filprofiler/_tracer.py

@@ -1,6 +1,9 @@
 """Trace code, so that libpymemprofile_api.so know's where we are."""
 
+import atexit
+import os
 import sys
+import threading
 from ctypes import CDLL, RTLD_GLOBAL
 
 from ._utils import library_path
@@ -14,25 +17,42 @@
 
 def start_tracing():
     preload.fil_reset()
+    threading.settrace(_start_thread_trace)
     _profiler.start_tracing()
 
 
-def stop_tracing(svg_output_path: str):
+def _start_thread_trace(frame, event, arg):
+    """Trace function that can be passed to sys.settrace.
+
+    All this does is register the underlying C trace function, using the
+    mechanism described in
+    https://github.com/nedbat/coveragepy/blob/master/coverage/ctracer/tracer.c's
+    CTracer_call.
+    """
+    if event == "call":
+        _profiler.start_tracing()
+    return _start_thread_trace
+
+
+def stop_tracing(output_path: str):
     sys.settrace(None)
-    path = svg_output_path.encode("utf-8")
+    path = output_path.encode("utf-8")
     preload.fil_dump_peak_to_flamegraph(path)
-    with open(path) as f:
+    svg_path = os.path.join(output_path, "peak-memory.svg")
+    with open(svg_path) as f:
         data = f.read().replace(
             "SUBTITLE-HERE",
             """Made with the Fil memory profiler. <a href="https://pythonspeed.com/products/filmemoryprofiler/" style="text-decoration: underline;" target="_parent">Try it on your code!</a>""",
         )
-    with open(path, "w") as f:
+    with open(svg_path, "w") as f:
         f.write(data)
 
 
-def trace(code, globals_, svg_output_path: str):
+def trace(code, globals_, output_path: str):
+    """
+    Given code (Python or code object), run it under the tracer until the
+    program exits.
+    """
+    atexit.register(stop_tracing, output_path)
     start_tracing()
-    try:
-        exec(code, globals_, None)
-    finally:
-        stop_tracing(svg_output_path)
+    exec(code, globals_, None)

memapi/Cargo.toml

@@ -24,10 +24,6 @@ features = ["ffi"]
 [dev-dependencies]
 proptest = "0.9.5"
 
-[profile.release]
-# Needed for profiling, disable for real release:
-#debug = true
-
 [lib]
 name = "pymemprofile_api"
 crate_type = ["cdylib"]

memapi/src/lib.rs

@@ -1,4 +1,3 @@
-use smallstr::SmallString;
 use std::ffi::CStr;
 use std::os::raw::c_char;
 use std::str;
@@ -25,16 +24,9 @@ pub unsafe extern "C" fn pymemprofile_start_call(
     file_name: *const c_char,
     func_name: *const c_char,
 ) {
-    let function_name = SmallString::from_str(str::from_utf8_unchecked(
-        CStr::from_ptr(func_name).to_bytes(),
-    ));
-    let module_name = SmallString::from_str(str::from_utf8_unchecked(
-        CStr::from_ptr(file_name).to_bytes(),
-    ));
-    let call_site = memorytracking::CallSite {
-        function_name,
-        module_name,
-    };
+    let function_name = str::from_utf8_unchecked(CStr::from_ptr(func_name).to_bytes());
+    let module_name = str::from_utf8_unchecked(CStr::from_ptr(file_name).to_bytes());
+    let call_site = memorytracking::CallSite::new(module_name, function_name);
     memorytracking::start_call(call_site);
 }
 

memapi/src/memorytracking.rs

@@ -7,8 +7,13 @@ use smallstr::SmallString;
 use std::cell::RefCell;
 use std::collections;
 use std::fmt;
+use std::fs;
+use std::io::Write;
+use std::path::Path;
 use std::sync::Mutex;
 
+/// The current Python callstack. We use u32 IDs instead of CallSite objects for
+/// performance reasons.
 #[derive(Clone, Debug, PartialEq)]
 struct Callstack {
     calls: Vec<u32>,
@@ -41,18 +46,28 @@ impl Callstack {
 
 thread_local!(static THREAD_CALLSTACK: RefCell<Callstack> = RefCell::new(Callstack::new()));
 
-/// A particular place where a call happened:
+/// A particular place where a call happened.
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct CallSite {
-    pub module_name: SmallString<[u8; 24]>,
-    pub function_name: SmallString<[u8; 24]>,
+    module_name: SmallString<[u8; 24]>,
+    function_name: SmallString<[u8; 24]>,
+}
+
+impl CallSite {
+    pub fn new(module_name: &str, function_name: &str) -> CallSite {
+        CallSite {
+            module_name: SmallString::from_str(module_name),
+            function_name: SmallString::from_str(function_name),
+        }
+    }
 }
 
 impl fmt::Display for CallSite {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "{}:{}", self.module_name, self.function_name)
     }
 }
+
 /// Maps CallSites to integer identifiers used in CallStacks.
 struct CallSites {
     max_id: u32,
@@ -67,6 +82,7 @@ impl CallSites {
         }
     }
 
+    /// Add a (possibly) new CallSite, returning its ID.
     fn get_or_insert_id(&mut self, call_site: CallSite) -> u32 {
         let max_id = &mut self.max_id;
         let result = self.callsite_to_id.entry(call_site).or_insert_with(|| {
@@ -77,6 +93,7 @@ impl CallSites {
         *result
     }
 
+    /// Get map from IDs to CallSites.
     fn get_reverse_map(&self) -> HashMap<u32, CallSite> {
         let mut result = HashMap::default();
         for (call_site, csid) in &(self.callsite_to_id) {
@@ -86,12 +103,14 @@ impl CallSites {
     }
 }
 
+/// A specific call to malloc()/calloc().
 #[derive(Clone, Debug, PartialEq)]
 struct Allocation {
     callstack: Callstack,
     size: libc::size_t,
 }
 
+/// The main data structure tracking everything.
 struct AllocationTracker {
     current_allocations: imhashmap::HashMap<usize, Allocation>,
     peak_allocations: imhashmap::HashMap<usize, Allocation>,
@@ -149,22 +168,46 @@ impl<'a> AllocationTracker {
         by_call
     }
 
-    /// Dump all callstacks in peak memory usage to format used by flamegraph.
+    /// Dump all callstacks in peak memory usage to various files describing the
+    /// memory usage.
     fn dump_peak_to_flamegraph(&self, path: &str) {
+        let directory_path = Path::new(path);
+        if !directory_path.exists() {
+            fs::create_dir(directory_path).expect("Couldn't create the output directory.");
+        } else if !directory_path.is_dir() {
+            panic!("Output path must be a directory.");
+        }
+
         let by_call = self.combine_callstacks();
         let lines: Vec<String> = by_call
             .iter()
+            // Filter out callstacks with less than 1 KiB RAM usage.
+            // TODO maybe make this number configurable someday.
+            .filter(|(_, size)| **size >= 1024)
             .map(|(callstack, size)| {
                 format!("{} {:.0}", callstack, (*size as f64 / 1024.0).round())
             })
             .collect();
+        let raw_path = directory_path
+            .join("peak-memory.prof")
+            .to_str()
+            .unwrap()
+            .to_string();
+        if let Err(e) = write_lines(&lines, &raw_path) {
+            eprintln!("Error writing raw profiling data: {}", e);
+        }
+        let svg_path = directory_path
+            .join("peak-memory.svg")
+            .to_str()
+            .unwrap()
+            .to_string();
         match write_flamegraph(
             lines.iter().map(|s| s.as_ref()),
-            path,
+            &svg_path,
             self.peak_allocated_bytes,
         ) {
             Ok(_) => {
-                eprintln!("Wrote memory usage flamegraph to {}", path);
+                eprintln!("Wrote memory usage flamegraph to {}", svg_path);
             }
             Err(e) => {
                 eprintln!("Error writing SVG: {}", e);
@@ -218,6 +261,17 @@ pub fn dump_peak_to_flamegraph(path: &str) {
     allocations.dump_peak_to_flamegraph(path);
 }
 
+/// Write strings to disk, one line per string.
+fn write_lines(lines: &Vec<String>, path: &str) -> std::io::Result<()> {
+    let mut file = fs::File::create(path)?;
+    for line in lines.iter() {
+        file.write_all(line.as_bytes())?;
+        file.write_all(b"\n")?;
+    }
+    Ok(())
+}
+
+/// Write a flamegraph SVG to disk, given lines in summarized format.
 fn write_flamegraph<'a, I: IntoIterator<Item = &'a str>>(
     lines: I,
     path: &str,
@@ -250,7 +304,7 @@ fn write_flamegraph<'a, I: IntoIterator<Item = &'a str>>(
 
 #[cfg(test)]
 mod tests {
-    use super::{AllocationTracker, Callstack};
+    use super::{AllocationTracker, CallSite, CallSites, Callstack};
     use itertools::Itertools;
     use proptest::prelude::*;
     use std::collections;
@@ -308,22 +362,51 @@ mod tests {
         assert_eq!(tracker.peak_allocated_bytes, 2123);
     }
 
+    #[test]
+    fn callsites_notices_duplicate_callsites() {
+        let callsite1 = CallSite::new("a", "af");
+        let callsite2 = CallSite::new("b", "af");
+        let callsite3 = CallSite::new("a", "bf");
+        let mut callsites = CallSites::new();
+        let id1 = callsites.get_or_insert_id(callsite1.clone());
+        let id1b = callsites.get_or_insert_id(callsite1);
+        let id2 = callsites.get_or_insert_id(callsite2);
+        let id3 = callsites.get_or_insert_id(callsite3.clone());
+        let id3b = callsites.get_or_insert_id(callsite3.clone());
+        assert_eq!(id1, id1b);
+        assert_ne!(id1, id2);
+        assert_ne!(id1, id3);
+        assert_ne!(id2, id3);
+        assert_eq!(id3, id3b);
+    }
+
     #[test]
     fn combine_callstacks_and_sum_allocations() {
         let mut tracker = AllocationTracker::new();
+        let id1 = tracker
+            .call_sites
+            .get_or_insert_id(CallSite::new("a", "af"));
+
+        let id2 = tracker
+            .call_sites
+            .get_or_insert_id(CallSite::new("b", "bf"));
+
+        let id3 = tracker
+            .call_sites
+            .get_or_insert_id(CallSite::new("c", "cf"));
         let mut cs1 = Callstack::new();
-        cs1.start_call(1);
-        cs1.start_call(2);
+        cs1.start_call(id1);
+        cs1.start_call(id2);
         let mut cs2 = Callstack::new();
-        cs2.start_call(3);
+        cs2.start_call(id3);
 
         tracker.add_allocation(1, 1000, cs1.clone());
         tracker.add_allocation(2, 234, cs2.clone());
         tracker.add_allocation(3, 50000, cs1.clone());
 
         let mut expected: collections::HashMap<String, usize> = collections::HashMap::new();
-        expected.insert("a;b".to_string(), 51000);
-        expected.insert("c".to_string(), 234);
+        expected.insert("a:af;b:bf".to_string(), 51000);
+        expected.insert("c:cf".to_string(), 234);
         assert_eq!(expected, tracker.combine_callstacks());
     }
 }

python-benchmarks/threaded.py

@@ -4,31 +4,21 @@
 import numpy
 
 def h(i):
-    return numpy.ones((1024, 1024, 20), dtype=numpy.uint8)
+    return numpy.ones((1024, 1024, i), dtype=numpy.uint8)
 
 def child1():
     return h(30)
 
 def thread1():
     data = h(20)
-    sleep(1)
+    sleep(0.1)
     data2 = child1()
-    sleep(1)
-
-def child2():
-    return h(30)
-
-def thread2():
-    data = h(50)
-    sleep(1)
-    data2 = child2()
-    sleep(1)
+    sleep(0.1)
 
 threading.Thread(target=thread1).start()
-threading.Thread(target=thread2).start()
 
 def main():
     data = h(50)
-    sleep(5)
+    sleep(0.5)
 
 main()