sched: SchedAssert: add getLastCpu()

Similar to getFirstCpu(), it's sometimes useful to know which CPU run
a given PID for the last time in the trace.

3 files changed, 41 insertions, 0 deletions

diff --git a/bart/sched/SchedAssert.py b/bart/sched/SchedAssert.py
index dd0b427..5ecfec9 100755
--- a/bart/sched/SchedAssert.py
+++ b/bart/sched/SchedAssert.py
@@ -612,6 +612,19 @@ class SchedAssert(object):
         cpus = Utils.listify(cpus)
         return first_cpu in cpus
 
+    def getLastCpu(self, window=None):
+        """Return the last CPU the task ran on"""
+
+        agg = self._aggregator(sched_funcs.last_cpu)
+        result = agg.aggregate(level="cpu", window=window)
+        result = list(itertools.chain.from_iterable(result))
+
+        end_time = max(result)
+        if not end_time:
+            return -1
+
+        return result.index(end_time)
+
     def generate_events(self, level, start_id=0, window=None):
         """Generate events for the trace plot
 
diff --git a/bart/sched/functions.py b/bart/sched/functions.py
index cb3336e..d1b17d4 100644
--- a/bart/sched/functions.py
+++ b/bart/sched/functions.py
@@ -216,6 +216,27 @@ def first_cpu(series, window=None):
     else:
         return [float("inf")]
 
+def last_cpu(series, window=None):
+    """:func:`aggfunc` to calculate the time of
+    the last switch out event in the series
+    This is returned as a vector of unit length
+    so that it can be aggregated and reduced across
+    nodes to find the last cpu of a task
+
+    :param series: Input Time Series data
+    :type series: :mod:`pandas.Series`
+
+    :param window: A tuple indicating a time window
+    :type window: tuple
+    """
+    series = select_window(series, window)
+    series = series[series == SCHED_SWITCH_OUT]
+
+    if len(series):
+        return [series.index.values[-1]]
+    else:
+        return [0]
+
 def select_window(series, window):
     """Helper Function to select a portion of
     pandas time series
diff --git a/tests/test_sched_assert.py b/tests/test_sched_assert.py
index 029dde6..4f8c28b 100644
--- a/tests/test_sched_assert.py
+++ b/tests/test_sched_assert.py
@@ -70,6 +70,13 @@ class TestSchedAssert(utils_tests.SetupDirectory):
         self.assertAlmostEqual(s.getRuntime(window=window), expected_time,
                                places=9)
 
+    def test_get_last_cpu(self):
+        """SchedAssert.getLastCpu() gives you the last cpu in which a task ran"""
+        expected_last_cpu = 5
+
+        sa = SchedAssert("trace.dat", self.topology, execname="ls")
+        self.assertEqual(sa.getLastCpu(), expected_last_cpu)
+
 class TestSchedMultiAssert(utils_tests.SetupDirectory):
     def __init__(self, *args, **kwargs):
         self.big = [1,2]