diff --git a/mosaic/runtime/warehouse.py b/mosaic/runtime/warehouse.py index 2849eb9..a456683 100644 --- a/mosaic/runtime/warehouse.py +++ b/mosaic/runtime/warehouse.py @@ -83,15 +83,15 @@ def set_profiler(self): super().set_profiler() async def put(self, obj, uid=None, publish=False, **kwargs): - return await self.put_remote(self.uid, obj, uid=uid) + return await self.put_remote(self.uid, obj, uid=uid, **kwargs) async def get(self, uid, **kwargs): - return await self.get_remote(self.uid, uid=uid) + return await self.get_remote(self.uid, uid=uid, **kwargs) async def drop(self, uid, **kwargs): - return await self.drop_remote(self.uid, uid=uid) + return await self.drop_remote(self.uid, uid=uid, **kwargs) - async def put_remote(self, sender_id, obj, uid=None, publish=False): + async def put_remote(self, sender_id, obj, uid=None, publish=False, **kwargs): """ Put an object into the warehouse. @@ -117,7 +117,7 @@ async def put_remote(self, sender_id, obj, uid=None, publish=False): if publish: await self.publish(sender_id, uid) - async def get_remote(self, sender_id, uid, warehouse_id=None, node_id=None): + async def get_remote(self, sender_id, uid, warehouse_id=None, node_id=None, **kwargs): """ Retrieve an object from the warehouse. @@ -167,7 +167,7 @@ async def get_remote(self, sender_id, uid, warehouse_id=None, node_id=None): return obj - async def exec_remote(self, sender_id, uid, func, func_args=None, func_kwargs=None): + async def exec_remote(self, sender_id, uid, func, func_args=None, func_kwargs=None, **kwargs): """ Retrieve an object from the warehouse and execute function on it. @@ -202,7 +202,7 @@ async def exec_remote(self, sender_id, uid, func, func_args=None, func_kwargs=No warehouse_obj = WarehouseObject(obj, uid=obj_id) return warehouse_obj - async def drop_remote(self, sender_id, uid, propagate=False): + async def drop_remote(self, sender_id, uid, propagate=False, **kwargs): """ Delete an object from the warehouse.