diff --git a/charm4py/ray/__init__.py b/charm4py/ray/__init__.py index 61abc17..c88982f 100644 --- a/charm4py/ray/__init__.py +++ b/charm4py/ray/__init__.py @@ -1 +1 @@ -from .api import get_object_store, init, remote, get, wait, put +from .api import get_object_store, init, remote, get, wait, put, is_initialized, shutdown diff --git a/charm4py/ray/api.py b/charm4py/ray/api.py index c49f625..24d1cd8 100644 --- a/charm4py/ray/api.py +++ b/charm4py/ray/api.py @@ -10,6 +10,13 @@ def init(): charm.thisProxy.updateGlobals({'object_store' : object_store,}, awaitable=True, module_name='charm4py.ray.api').get() +def is_initialized(): + # if the object store exists, ray is initialized + if 'object_store' in globals(): + return True + else: + return False + def get_object_store(): global object_store @@ -103,4 +110,8 @@ def put(obj): from ..charm import charm fut = charm.threadMgr.createFuture(store=True) fut.create_object(obj) - return fut \ No newline at end of file + return fut + +def shutdown(): + global object_store + del object_store \ No newline at end of file diff --git a/examples/ray/simple.py b/examples/ray/simple.py index 800f6b9..25d2cb7 100644 --- a/examples/ray/simple.py +++ b/examples/ray/simple.py @@ -25,6 +25,7 @@ def add(self, a, b): def main(args): ray.init() + assert ray.is_initialized() # create 3 instances of MyChare, distributed among cores by the runtime arr = [Compute.remote(i) for i in range(4)] @@ -43,6 +44,9 @@ def main(args): while len(not_ready) > 0: ready, not_ready = ray.wait(not_ready) print("Fetched value: ", ray.get(ready)) + + ray.shutdown() + assert not ray.is_initialized() exit()