[WIP] Fixing instrumentation for workflows

[logan-markewich]

Since workflows return a Future directly now, instrumentation will log the future instead of the actual result

So, to remedy this, we need to handle returning futures in our span decorators

This fix seems to work, but seeing the "error" below when using arize (note that this seems benign? since things are still logged to arize properly?)

ERROR [opentelemetry.context] Failed to detach context
Traceback (most recent call last):
  File "/Users/loganmarkewich/Library/Caches/pypoetry/virtualenvs/llama-index-caVs7DDe-py3.11/lib/python3.11/site-packages/opentelemetry/context/__init__.py", line 154, in detach
    _RUNTIME_CONTEXT.detach(token)
  File "/Users/loganmarkewich/Library/Caches/pypoetry/virtualenvs/llama-index-caVs7DDe-py3.11/lib/python3.11/site-packages/opentelemetry/context/contextvars_context.py", line 50, in detach
    self._current_context.reset(token)  # type: ignore
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: <Token var=<ContextVar name='current_context' default={} at 0x2acb049f0> at 0x2cb890500> was created in a different Context

I bit more validation is needed

Fixes #16283

[nerdai]

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

[B-Step62]

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).

[logan-markewich]

@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 🤔

[B-Step62]

@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.

[enrico-stauss]

I ran into the error log mentioned in the original post and found this issue. Now I don't know the internals of the LlamaIndex instrumentation well enough but maybe it helps in some way.