Add demo_remoteclass

sidi/test/demo_remoteclass.py

@@ -3,44 +3,50 @@ from threading import Thread
 from distributed import Client
 from sidi.scheduler import spawn_scheduler_process
 from sidi.worker import DaskWorker
-from sidi.remote import RemoteClass 
-from time import sleep  
+from sidi.remote import RemoteClass
+from time import sleep
+import atexit
 
-class Dummy(object): 
-	def __init__(self, a=1): 
-		self.a = a 
+class Dummy(object):
+    def __init__(self, a=1):
+        self.a = a
 
-	def do_work(self, a): 
-		print("[%d] Doing work" % os.getpid())
-		sleep(a) 
-		return a+1  
+    def do_work(self, a):
+        print("[%d] Doing work" % os.getpid())
+        sleep(a)
+        return a+1
 
 
 def mycallback(fut):
     print("callback time !")
 
 
+def cleanUp(P, W):
+    print("Stopping workers")
+    W.stop()
+    sleep(1)
+    print("Stopping scheduler")
+    P.terminate()
+    sleep(1)
+
+
 def main():
-	P = spawn_scheduler_process(addr="tcp://127.0.0.1", port=5455)
-	W = DaskWorker("tcp://127.0.0.1:5455", nprocs=1, nthreads=1)
-	T = Thread(target=W.start)
-	T.start()
-	cl = Client("tcp://127.0.0.1:5455")
-	
-	input("Instantiate remote class ?")
-	RC = RemoteClass(cl, Dummy, a=1)
-	print("...OK\n\n")
-	
-	sleep(2)
-	input("Ready to submit...")
-	futures = RC.submit_task("do_work", callback=mycallback, method_args=(1,))
-	print(futures)
-	return P, W
+    P = spawn_scheduler_process(addr="tcp://127.0.0.1", port=5455)
+    W = DaskWorker("tcp://127.0.0.1:5455", nprocs=1, nthreads=1)
+    T = Thread(target=W.start)
+    T.start()
+    cl = Client("tcp://127.0.0.1:5455")
+    atexit.register(cleanUp, P, W)
+
+    input("Instantiate remote class ?")
+    RC = RemoteClass(cl, Dummy, a=1)
+    print("...OK\n\n")
+
+    sleep(2)
+    input("Ready to submit...")
+    futures = RC.submit_task("do_work", callback=mycallback, method_args=(1,))
+    return futures
 
 
 if __name__ == "__main__":
-	try:
-		P, W = main()
-	finally:
-		W.stop()
-		P.terminate()
\ No newline at end of file
+    futures = main()