Merge pull request #44 from willgittoes-dd/willgittoes-dd/headless-op…

…eration-by-default Fix crash on machines without a display
astrofrog · Sep 6, 2018 · e604aec · e604aec
2 parents 27f629a + 69bfef9
commit e604aec
Showing 1 changed file with 14 additions and 12 deletions.
diff --git a/psrecord/main.py b/psrecord/main.py
@@ -209,24 +209,26 @@ def monitor(pid, logfile=None, plot=None, duration=None, interval=None,
 
     if plot:
 
+        # Use non-interactive backend, to enable operation on headless machines
         import matplotlib.pyplot as plt
+        with plt.rc_context({'backend': 'Agg'}):
 
-        fig = plt.figure()
-        ax = fig.add_subplot(1, 1, 1)
+            fig = plt.figure()
+            ax = fig.add_subplot(1, 1, 1)
 
-        ax.plot(log['times'], log['cpu'], '-', lw=1, color='r')
+            ax.plot(log['times'], log['cpu'], '-', lw=1, color='r')
 
-        ax.set_ylabel('CPU (%)', color='r')
-        ax.set_xlabel('time (s)')
-        ax.set_ylim(0., max(log['cpu']) * 1.2)
+            ax.set_ylabel('CPU (%)', color='r')
+            ax.set_xlabel('time (s)')
+            ax.set_ylim(0., max(log['cpu']) * 1.2)
 
-        ax2 = ax.twinx()
+            ax2 = ax.twinx()
 
-        ax2.plot(log['times'], log['mem_real'], '-', lw=1, color='b')
-        ax2.set_ylim(0., max(log['mem_real']) * 1.2)
+            ax2.plot(log['times'], log['mem_real'], '-', lw=1, color='b')
+            ax2.set_ylim(0., max(log['mem_real']) * 1.2)
 
-        ax2.set_ylabel('Real Memory (MB)', color='b')
+            ax2.set_ylabel('Real Memory (MB)', color='b')
 
-        ax.grid()
+            ax.grid()
 
-        fig.savefig(plot)
+            fig.savefig(plot)