From afaa2218b4cd710044dd46cbfe55196ded11d29d Mon Sep 17 00:00:00 2001
From: Muammar El Khatib <muammarelkhatib@gmail.com>
Date: Sun, 26 Jan 2020 12:15:24 -0800
Subject: [PATCH] Quick hack to make %srun non-blocking.

This commit is related to #7.
---
 slurm_magic.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/slurm_magic.py b/slurm_magic.py
index 3c316b5..2c75240 100644
--- a/slurm_magic.py
+++ b/slurm_magic.py
@@ -160,13 +160,17 @@ def sview(self, line):
         pass
 
     def _execute(self, line, input=None, stderr=False):
+        non_blocking = ["srun"]
+
         name = inspect.stack()[1][3]
         process = Popen([name] + line.split(), stdin=PIPE, stdout=PIPE, stderr=PIPE)
-        stdout, stderr = process.communicate(input)
-        if stderr:
-            return stdout.decode("utf-8"), stderr.decode("utf-8")
-        else:
-            return stdout.decode("utf-8")
+
+        if name not in non_blocking:
+            stdout, stderr = process.communicate(input)
+            if stderr:
+                return stdout.decode("utf-8"), stderr.decode("utf-8")
+            else:
+                return stdout.decode("utf-8")
 
 
 def load_ipython_extension(ip):