Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

try more complete lifecycle #2744

Open
wants to merge 5 commits into
base: async-loop-issue
Choose a base branch
from
Open

try more complete lifecycle #2744

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a31b713
try more complete lifecycle
wild-endeavor Sep 11, 2024
0d2a7ad
lint
wild-endeavor Sep 11, 2024
c16b4d7
remove test
wild-endeavor Sep 11, 2024
1f13722
lint
wild-endeavor Sep 11, 2024
7f7d374
wrap the loop around the whole dispatch execute
wild-endeavor Sep 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 7 additions & 21 deletions flytekit/bin/entrypoint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
import sys
import tempfile
import traceback
import warnings
from sys import exit
from typing import Callable, List, Optional

Expand Down Expand Up @@ -46,6 +45,7 @@
from flytekit.models.core import identifier as _identifier
from flytekit.tools.fast_registration import download_distribution as _download_distribution
from flytekit.tools.module_loader import load_object_from_module
from flytekit.utils.loop_handling import use_event_loop


def get_version_message():
Expand All @@ -69,23 +69,6 @@ def _compute_array_job_index():
return offset


def _get_working_loop():
"""Returns a running event loop."""
try:
return asyncio.get_running_loop()
except RuntimeError:
with warnings.catch_warnings():
warnings.simplefilter("error", DeprecationWarning)
try:
return asyncio.get_event_loop_policy().get_event_loop()
# Since version 3.12, DeprecationWarning is emitted if there is no
# current event loop.
except DeprecationWarning:
loop = asyncio.get_event_loop_policy().new_event_loop()
asyncio.set_event_loop(loop)
return loop


def _dispatch_execute(
ctx: FlyteContext,
load_task: Callable[[], PythonTask],
Expand Down Expand Up @@ -125,7 +108,8 @@ def _dispatch_execute(
if inspect.iscoroutine(outputs):
# Handle eager-mode (async) tasks
logger.info("Output is a coroutine")
outputs = _get_working_loop().run_until_complete(outputs)
loop = asyncio.get_running_loop()
outputs = loop.run_until_complete(outputs)

# Step3a
if isinstance(outputs, VoidPromise):
Expand Down Expand Up @@ -418,7 +402,8 @@ def load_task():
f"Test detected, returning. Args were {inputs} {output_prefix} {raw_output_data_prefix} {resolver} {resolver_args}"
)
return
_dispatch_execute(ctx, load_task, inputs, output_prefix)
with use_event_loop():
_dispatch_execute(ctx, load_task, inputs, output_prefix)


def _execute_map_task(
Expand Down Expand Up @@ -479,7 +464,8 @@ def load_task():
)
return

_dispatch_execute(ctx, load_task, inputs, output_prefix)
with use_event_loop():
_dispatch_execute(ctx, load_task, inputs, output_prefix)


def normalize_inputs(
Expand Down
Empty file added flytekit/utils/__init__.py
View file
Open in desktop
Empty file.
34 changes: 34 additions & 0 deletions flytekit/utils/loop_handling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
import asyncio
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from signal import SIGINT, SIGTERM

from flytekit.loggers import logger


def handler(loop, s: int):
loop.stop()
logger.debug(f"Shutting down loop at {id(loop)} via {s!s}")
loop.remove_signal_handler(SIGTERM)
loop.add_signal_handler(SIGINT, lambda: None)


@contextmanager
def use_event_loop():
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you leave a comment here that explains why we need this functionality?

loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
Comment on lines +18 to +19
Copy link
Member

@thomasjpfan thomasjpfan Sep 11, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If there is an existing event loop running, then I do not think a library should be modifying the global event loop.

If there is another event loop running with tasks from another library, then switching out the event loop from underneath them can cause problems. For example, if another library is actively scheduling work, they will have task end up in two different loops.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can add a try block around this then, to look for an existing event loop, and use that if found. But I think that is not a good idea. Still learning, but I feel libraries should not be creating event loops on import. As of today, this is true of flytekit's dependencies (except for the unionai library).

The lack of an event loop is why the error came about in the first place right? asyncio.run() creates and then destroys/cleans up the event loop. But if there's one that already exists (which is the case in the main issue with these unionfses), it will basically ignore it, and create a new one. After asyncio.run() completes, it does not restore the existing one, which is why when the data persistence layer goes to look for one, it can't find one and errors.

My issue with adding a try block... and then using the event loop if we find one.

  • asyncio.run doesn't do this.
  • In our current case, we end up using an event loop in the union library. If the union library is not installed, then we use one created by flytekit. Difference in behavior is not good I feel. If we want to differentiate, I feel the way to do it is: if user runs pyflyte-execute, then flytekit creates and manages its own event loop, if the user runs union-execute (which i know doesn't exist today), then unionai creates and manages its own event loop. Changing where the event loop is managed based on what library is installed feels bad. Also if a user installs some other library that we don't know about that creates its own event loop, and that is somehow triggered before this code is (yes i know this is unlikely since we load all flytekit code before loading user code), then flytekit ends up using some random event loop.
  • If we had had code in flytekit to use an existing loop, and the unionai library is as it stands, then flytekit would've ended up using the event loop created as part of a grpc.aio channel. That channel may get cleaned up or deleted right?
  • When we get to more advanced async stuff, we'll may very well want to explicitly handle cleanup in case of termination. Relying on unionai's loop to clean up things in flytekit feels incorrect also.
  • I'm not suggesting we add this code every where in flytekit/union... this is only on the main entrypoint of an executable. I think in this scenario, it's okay to take ownership. In other parts of the code, yes, definitely use the loop available, and maybe even crash if one's not found (rather than grpcio's behavior).

Event loops are thread singletons, there's one per thread. I feel that if your library needs an event loop on import (not on calling an executable like pyflyte or pyflyte-execute), then you should run it in a different thread.

Alternatively we can also check to see if there's an event loop, save it to a variable if so, and then restore it later, basically what async run does, but with one extra step, but honestly I'd rather see it fail. Seeing the failure allowed us to find an issue in the union library.

Copy link
Member

@thomasjpfan thomasjpfan Sep 12, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

unionfs does not create the event loop during import. The event loop is created when union:// fsspec protocol is used which initializes union's fsspec implementation.

if the user runs union-execute (which i know doesn't exist today), then unionai creates and manages its own event loop.

In principle, I am okay with this, but async python libraries do different things:

  1. https://github.com/grpc/grpc/blob/1f70b34fb15ffb409553c5f7055cfd5cca61e98e/src/python/grpcio/grpc/_cython/_cygrpc/aio/common.pyx.pxi#L177-L194 : If there is no global event loop, then create one
  2. https://github.com/fsspec/filesystem_spec/blob/76ca4a68885d572880ac6800f079738df562f02c/fsspec/asyn.py#L318-L321: If there is no loop passed in, then create a event loop on another thread
  3. https://github.com/Tinche/aiofiles/blob/7582fda077e0e5b58a4f7fa3f11d0f53ef36eed5/src/aiofiles/threadpool/__init__.py#L79-L80: Errors if there is not a running event loop

My preferred solution is 2, if there was a good way to pass in the new event loop into a library that does 1.

Copy link
Contributor Author

@wild-endeavor wild-endeavor Sep 16, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thinking about this more, yeah shouldn't we do 2 always? We can't prevent user code from arbitrarily running asyncio.run in their code and as we've seen with our own usage, that will destroy the then-current event loop. So basically the rule is:
Anytime an event loop is needed by flytekit/union, including when it needed by a downstream library like in the case of async grpc (and we'll just have to know), always create a loop on another thread and

There will probably be bugs related to usage of id in artifacts and maybe need more contextvars in general but we can address those.

pass in the new event loop into a library that does 1.

in the case of grpc aio, doesn't just having the loop set count as passing it in?

executor = ThreadPoolExecutor()
loop.set_default_executor(executor)
for sig in (SIGTERM, SIGINT):
loop.add_signal_handler(sig, handler, loop, sig)
try:
yield loop
finally:
tasks = asyncio.all_tasks(loop=loop)
for t in tasks:
logger.debug(f"canceling {t.get_name()}")
t.cancel()
group = asyncio.gather(*tasks, return_exceptions=True)
loop.run_until_complete(group)
executor.shutdown(wait=True)
loop.close()
25 changes: 1 addition & 24 deletions tests/flytekit/unit/experimental/test_eager_workflows.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@

from flytekit import dynamic, task, workflow

from flytekit.bin.entrypoint import _get_working_loop, _dispatch_execute
from flytekit.bin.entrypoint import _dispatch_execute
from flytekit.core import context_manager
from flytekit.core.promise import VoidPromise
from flytekit.exceptions.user import FlyteValidationException
Expand Down Expand Up @@ -281,26 +281,3 @@ async def eager_wf_flyte_directory() -> str:

result = asyncio.run(eager_wf_flyte_directory())
assert result == "some data"


@mock.patch("flytekit.core.utils.load_proto_from_file")
@mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data")
@mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data")
@mock.patch("flytekit.core.utils.write_proto_to_file")
def test_eager_workflow_dispatch(mock_write_to_file, mock_put_data, mock_get_data, mock_load_proto, event_loop):
"""Test that event loop is preserved after executing eager workflow via dispatch."""

@eager
async def eager_wf():
await asyncio.sleep(0.1)
return

ctx = context_manager.FlyteContext.current_context()
with context_manager.FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.execution_state.with_params(mode=context_manager.ExecutionState.Mode.TASK_EXECUTION)
)
) as ctx:
_dispatch_execute(ctx, lambda: eager_wf, "inputs path", "outputs prefix")
loop_after_execute = asyncio.get_event_loop_policy().get_event_loop()
assert event_loop == loop_after_execute
Loading