1 files changed, 27 insertions, 0 deletions
diff --git a/rumba/utils.py b/rumba/utils.py
new file mode 100644
index 0000000..5348b0a
--- /dev/null
+++ b/rumba/utils.py
@@ -0,0 +1,27 @@
+import time
+
+import rumba.log as log
+import rumba.model as model
+
+logger = log.get_logger(__name__)
+
+
+class ExperimentManager(object):
+
+    def __init__(self, experiment, do_swap_out=True):
+        assert isinstance(experiment, model.Experiment), \
+            'An experiment instance is required.'
+        self.experiment = experiment
+        self.do_swap_out = do_swap_out
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.do_swap_out:
+            self.experiment.swap_out()
+        if exc_val is not None:
+            logger.error('Something went wrong. Got %s: %s',
+                         type(exc_val).__name__, str(exc_val))
+        time.sleep(0.1)  # Give the queue logger enough time to flush.
+        return True  # Suppress the exception we logged: no traceback.