forked from flux-framework/flux-core
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Problem: we currently cannot integrate flux with Python asyncio Description: This will allow us to run tasks using asyncio. It adds a few extra functions to Flux futures to support being run alongside asyncio tasks, and a flux.asyncio module that provides a FluxEventLoop and a custom selector to handle integration of the flux native watchers with Python native asyncio. Signed-off-by: vsoch <vsoch@users.noreply.github.com>
- Loading branch information
Showing
15 changed files
with
440 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
| from flux.asyncio.events import FluxEventLoop, loop, submit |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| ############################################################### | ||
| # Copyright 2022 Lawrence Livermore National Security, LLC | ||
| # (c.f. AUTHORS, NOTICE.LLNS, COPYING) | ||
| # | ||
| # This file is part of the Flux resource manager framework. | ||
| # For details, see https://github.com/flux-framework. | ||
| # | ||
| # SPDX-License-Identifier: LGPL-3.0 | ||
| ############################################################### | ||
|
|
||
| import asyncio | ||
|
|
||
| import flux | ||
| import flux.core.watchers | ||
| from flux.asyncio.selector import FluxSelector | ||
|
|
||
| # The loop *must* have the same handle as the submit or else the loop will | ||
| # run forever. | ||
| HANDLE = flux.Flux() | ||
|
|
||
|
|
||
| async def submit(jobspec, flux_handle=None): | ||
| """Submit a Flux jobspec and return a job ID. | ||
| Example usage | ||
| ============= | ||
| # ensure the script is running in an active flux instance | ||
| import asyncio | ||
| from flux.asyncio import loop | ||
| import flux.job | ||
| fluxsleep = flux.job.JobspecV1.from_command(['sleep', '2']) | ||
| fluxecho = flux.job.JobspecV1.from_command(['echo', 'pancakes']) | ||
| tasks = [ | ||
| loop.create_task(asyncio.sleep(5)), | ||
| loop.create_task(flux.asyncio.submit(fluxecho)), | ||
| loop.create_task(flux.asyncio.submit(fluxsleep)), | ||
| ] | ||
| asyncio.set_event_loop(loop) | ||
| results = loop.run_until_complete(asyncio.gather(*tasks)) | ||
| # [JobID(456004999315456), JobID(456004999315457)] | ||
| """ | ||
| handle = flux_handle or HANDLE | ||
| uid = await flux.job.submit_async(handle, jobspec) | ||
| return uid | ||
|
|
||
|
|
||
| class FluxEventLoop(asyncio.SelectorEventLoop): | ||
| """An asyncio loop that handles running Flux.""" | ||
|
|
||
| def __init__(self, flux_handle=None): | ||
| # This can be provided by a user that knows what they are doing. | ||
| # The handle to the submit job and loop must be the same. | ||
| if not flux_handle: | ||
| flux_handle = HANDLE | ||
| selector = FluxSelector(flux_handle) | ||
|
|
||
| # Reverse reference is needed for watchers | ||
| selector.loop = self | ||
| super().__init__(selector) | ||
|
|
||
| @property | ||
| def selector(self): | ||
| return self._selector | ||
|
|
||
|
|
||
| # This loop needs to be accessible from all places! | ||
| loop = FluxEventLoop() # pylint: disable=invalid-name |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,151 @@ | ||
| ############################################################### | ||
| # Copyright 2022 Lawrence Livermore National Security, LLC | ||
| # (c.f. AUTHORS, NOTICE.LLNS, COPYING) | ||
| # | ||
| # This file is part of the Flux resource manager framework. | ||
| # For details, see https://github.com/flux-framework. | ||
| # | ||
| # SPDX-License-Identifier: LGPL-3.0 | ||
| ############################################################### | ||
|
|
||
| import selectors | ||
| import signal | ||
|
|
||
| import flux.constants | ||
| from flux.core.watchers import FDWatcher | ||
|
|
||
|
|
||
| def stop_callback(handle, _watcher, _fd_int, events, args=None): | ||
| """ | ||
| This is called by the watcher, and stops the reactor add adds ready jobs. | ||
| """ | ||
| handle.reactor_stop() | ||
|
|
||
| # TimerWatcher won't know a file descriptor, just needs to stop the reactor | ||
| if not args: | ||
| return | ||
| selector = args["select"] | ||
| selector_event = get_selector_event(events) | ||
| selector.ready.add((args["key"], selector_event & args["key"].events)) | ||
|
|
||
|
|
||
| def get_selector_event(events): | ||
| """ | ||
| Given an int, return the corresponding flux identifier. | ||
| """ | ||
| event = 0 | ||
| if events & flux.constants.FLUX_POLLIN: | ||
| event |= selectors.EVENT_READ | ||
| if events & flux.constants.FLUX_POLLOUT: | ||
| event |= selectors.EVENT_WRITE | ||
| return event | ||
|
|
||
|
|
||
| def get_flux_event(events): | ||
| """ | ||
| Given an int, return the corresponding selector identifier | ||
| """ | ||
| event = 0 | ||
| if events & selectors.EVENT_READ: | ||
| event |= flux.constants.FLUX_POLLIN | ||
| if events & selectors.EVENT_WRITE: | ||
| event |= flux.constants.FLUX_POLLOUT | ||
| return event | ||
|
|
||
|
|
||
| class FluxSelector( | ||
| selectors._BaseSelectorImpl | ||
| ): # pylint: disable=protected-access # type: ignore | ||
| """ | ||
| A Flux selector supports registering file objects to be monitored for | ||
| specific I/O events (for Flux). | ||
| """ | ||
|
|
||
| def __init__(self, handle): | ||
| super().__init__() | ||
| self.handle = handle | ||
| self.ready = set() | ||
| self._watchers = {} | ||
|
|
||
| def register(self, fileobj, events, data=None): | ||
| """ | ||
| Register a new file descriptor event. | ||
| """ | ||
| key = super().register(fileobj, events, data) | ||
| watcher = FDWatcher( | ||
| self.handle, | ||
| fileobj, | ||
| get_flux_event(events), | ||
| stop_callback, | ||
| args={"key": key, "select": self}, | ||
| ) | ||
| watcher.start() | ||
| self._watchers[fileobj] = watcher | ||
| return key | ||
|
|
||
| def unregister(self, fileobj): | ||
| """ | ||
| Remove the key and the watcher. | ||
| """ | ||
| try: | ||
| key = self._fd_to_key.pop(self._fileobj_lookup(fileobj)) | ||
| self._watchers[key.fileobj].stop() | ||
| except KeyError: | ||
| raise KeyError("{!r} is not registered".format(fileobj)) from None | ||
| return key | ||
|
|
||
| def select(self, timeout=None): | ||
| """ | ||
| Perform the actual selection, until some monitored file objects are | ||
| ready or a timeout expires. | ||
| Parameters: | ||
| timeout -- if timeout > 0, this specifies the maximum wait time, in | ||
| seconds, waited for by a flux TimerWatcher | ||
| if timeout <= 0, the select() call won't block, and will | ||
| report the currently ready file objects | ||
| if timeout is None, select() will block until a monitored | ||
| file object becomes ready | ||
| Returns: | ||
| list of (key, events) for ready file objects | ||
| `events` is a bitwise mask of EVENT_READ|EVENT_WRITE | ||
| """ | ||
| reactor = self.handle.get_reactor() | ||
| reactor_interrupted = False | ||
|
|
||
| def reactor_interrupt(handle, *_args): | ||
| # ensure reactor_interrupted from enclosing scope: | ||
| nonlocal reactor_interrupted | ||
| reactor_interrupted = True | ||
| handle.reactor_stop(reactor) | ||
|
|
||
| with self.handle.signal_watcher_create(signal.SIGINT, reactor_interrupt): | ||
| with self.handle.in_reactor(): | ||
|
|
||
| # Ensure previous events are cleared | ||
| self.ready.clear() | ||
|
|
||
| # 0 == "run until I tell you to stop" | ||
| if timeout is not None: | ||
| if timeout > 0: | ||
|
|
||
| # Block for a specified timeout | ||
| with self.handle.timer_watcher_create(timeout, stop_callback): | ||
| watcher_count = self.handle.flux_reactor_run(reactor, 0) | ||
|
|
||
| # If timeout <= 0, select won't block | ||
| else: | ||
| watcher_count = self.handle.flux_reactor_run( | ||
| reactor, flux.constants.FLUX_REACTOR_NOWAIT | ||
| ) | ||
|
|
||
| # If timeout is None, block until a monitored object ready | ||
| else: | ||
| watcher_count = self.handle.flux_reactor_run(reactor, 0) | ||
|
|
||
| if reactor_interrupted: | ||
| raise KeyboardInterrupt | ||
|
|
||
| if watcher_count < 0: | ||
| self.handle.raise_if_exception() | ||
|
|
||
| return list(self.ready) |
Oops, something went wrong.