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

[WIP] Fixing instrumentation for workflows #16290

Merged
merged 5 commits into from
Oct 3, 2024
Merged

[WIP] Fixing instrumentation for workflows #16290

Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
34d958a
Fix instrumentation for workflows
logan-markewich Sep 29, 2024
6966e2c
linting
logan-markewich Sep 30, 2024
ab455db
fix tests
logan-markewich Sep 30, 2024
a287665
bump dep
logan-markewich Oct 2, 2024
497fc3c
inherit the context to use in future
logan-markewich Oct 3, 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
55 changes: 48 additions & 7 deletions llama-index-core/llama_index/core/instrumentation/dispatcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import asyncio
from functools import partial
from contextlib import contextmanager
from contextvars import ContextVar, Token
from typing import Any, Callable, Generator, List, Optional, Dict, Protocol
Expand Down Expand Up @@ -251,6 +253,7 @@ def wrapper(func: Callable, instance: Any, args: list, kwargs: dict) -> Any:
bound_args = inspect.signature(func).bind(*args, **kwargs)
id_ = f"{func.__qualname__}-{uuid.uuid4()}"
tags = active_instrument_tags.get()
result = None

token = active_span_id.set(id_)
parent_id = None if token.old_value is Token.MISSING else token.old_value
Expand All @@ -261,20 +264,58 @@ def wrapper(func: Callable, instance: Any, args: list, kwargs: dict) -> Any:
parent_id=parent_id,
tags=tags,
)

def handle_future_result(
future: asyncio.Future,
span_id: str,
bound_args: inspect.BoundArguments,
instance: Any,
) -> None:
try:
result = future.result()
self.span_exit(
id_=span_id,
bound_args=bound_args,
instance=instance,
result=result,
)
return result
except BaseException as e:
self.event(SpanDropEvent(span_id=span_id, err_str=str(e)))
self.span_drop(
id_=span_id, bound_args=bound_args, instance=instance, err=e
)
raise
finally:
active_span_id.reset(token)

try:
result = func(*args, **kwargs)
if isinstance(result, asyncio.Future):
# If the result is a Future, wrap it
new_future = asyncio.ensure_future(result)
new_future.add_done_callback(
Copy link
Contributor

Choose a reason for hiding this comment

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

This makes sense to me! Wonder where that Context is being spun up in the arize error you shared

Copy link
Contributor

Choose a reason for hiding this comment

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

Thank you for the quick fix @logan-markewich! Just wanted to callout that the same error occurs with MLflow tracing (while the trace itself is logged correctly).

From my observation, I think the error is because of trying to reset ContextVar by a token that is created in the main thread, while resetting it in the async task - Arize and MLflow works fine nevertheless because neither relies on active_span_id. I think setting active_span_id inside handle_future_result should fix it (and it makes more sense that the async task manages its own active span).

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

@B-Step62 I don't think I can move the active_span_id into handle_future_result, since its needed for span_enter, and handle_future_result only runs once the span ends 🤔

Copy link
Contributor

@B-Step62 B-Step62 Oct 2, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

@logan-markewich I see. I think the problem here is we are trying to put both sync and async execution into a single active_span. In reality, they are executed in separate time.

  1. First we enter the wrapper.
  2. Then we create a future object and exit the wrapper immediately.
  3. After a while, the future is scheduled for execution.
  4. Finally, the task is completed.

Currently, we are trying to create a single span that stays "active" all the way from 1 to 4. However, in practice, there is some time gap between 2 and 3, and another sync operation should be able to take over the "active" span during that gap.

Therefore, I think we should manage active span like this

  1. First we enter the wrapper. -> Create a span and set to active
  2. Then we create a future object and exit the wrapper immediately. -> Reset active span
  3. After a while, the future is scheduled for execution. -> Set span to active. This lives in the separate context so should not conflict with the active span in the main context, while any spans created within the async task will recognize this as the parent span.
  4. Finally, the task is completed. -> Reset active span

How does this sound? This should be correct in theory, but not 100% sure if this works.

partial(
handle_future_result,
span_id=id_,
bound_args=bound_args,
instance=instance,
)
)
return new_future
else:
# For non-Future results, proceed as before
self.span_exit(
id_=id_, bound_args=bound_args, instance=instance, result=result
)
return result
except BaseException as e:
self.event(SpanDropEvent(span_id=id_, err_str=str(e)))
self.span_drop(id_=id_, bound_args=bound_args, instance=instance, err=e)
raise
else:
self.span_exit(
id_=id_, bound_args=bound_args, instance=instance, result=result
)
return result
finally:
# clean up
active_span_id.reset(token)
if not isinstance(result, asyncio.Future):
logan-markewich marked this conversation as resolved.
Show resolved Hide resolved
active_span_id.reset(token)

@wrapt.decorator
async def async_wrapper(
Expand Down
Loading