1 files changed, 20 insertions, 5 deletions

diff --git a/src/python/m5/__init__.py b/src/python/m5/__init__.py
index 579785a46..f4f5be2d1 100644
--- a/src/python/m5/__init__.py
+++ b/src/python/m5/__init__.py
@@ -34,7 +34,7 @@ import cc_main
 # import a few SWIG-wrapped items (those that are likely to be used
 # directly by user scripts) completely into this module for
 # convenience
-from cc_main import simulate, SimLoopExitEvent, setCheckpointDir
+from cc_main import simulate, SimLoopExitEvent
 
 # import the m5 compile options
 import defines
@@ -213,29 +213,44 @@ atexit.register(cc_main.doExitCleanup)
 # matter since most scripts will probably 'from m5.objects import *'.
 import objects
 
+# This loops until all objects have been fully drained.
 def doDrain(root):
+    all_drained = drain(root)
+    while (not all_drained):
+        all_drained = drain(root)
+
+# Tries to drain all objects.  Draining might not be completed unless
+# all objects return that they are drained on the first call.  This is
+# because as objects drain they may cause other objects to no longer
+# be drained.
+def drain(root):
+    all_drained = False
     drain_event = cc_main.createCountedDrain()
     unready_objects = root.startDrain(drain_event, True)
     # If we've got some objects that can't drain immediately, then simulate
     if unready_objects > 0:
         drain_event.setCount(unready_objects)
         simulate()
+    else:
+        all_drained = True
     cc_main.cleanupCountedDrain(drain_event)
+    return all_drained
 
 def resume(root):
     root.resume()
 
-def checkpoint(root):
+def checkpoint(root, dir):
     if not isinstance(root, objects.Root):
         raise TypeError, "Object is not a root object. Checkpoint must be called on a root object."
     doDrain(root)
     print "Writing checkpoint"
-    cc_main.serializeAll()
+    cc_main.serializeAll(dir)
     resume(root)
 
-def restoreCheckpoint(root):
+def restoreCheckpoint(root, dir):
     print "Restoring from checkpoint"
-    cc_main.unserializeAll()
+    cc_main.unserializeAll(dir)
+    resume(root)
 
 def changeToAtomic(system):
     if not isinstance(system, objects.Root) and not isinstance(system, System):