Where to put xray_recorder.configure(...AsyncContext()) in AWS Lambda functions?

[xjia1]

If running inside AWS Lambda, and the function is using asyncio so I need

xray_recorder.configure(service='...', context=AsyncContext())

However if I put it in the global, or in my lambda handler, I always get errors like

cannot find the current segment/subsegment, please make sure you have a segment open
No segment found, cannot begin subsegment

I think it's because Lambda has its own segment created, presumably it's the LambdaContext defined in aws_xray_sdk/core/lambda_launcher.py?

So how should I configure it properly so that @xray_recorder.capture_async can work?

[willarmiros]

However if I put it in the global, or in my lambda handler, I always get errors like

Are you saying that merely adding the line to configure the xray_recorder causes those errors? Or is @xray_recorder.capture_async causing the errors and that configuration is just failing to fix them?

What function are you decorating with @xray_recorder.capture_async? If you are decorating the Lambda function handler, that might be causing those errors since a segment is not generated until the beginning of the handler function.

In response to your question, the line

xray_recorder.configure(service='...', context=AsyncContext())

should be able to be included globally or in the lambda handler, I don't necessarily think it's causing the segment not found errors.

[xjia1]

Please take a look at https://github.com/stfairy/repro-xray-issue to reproduce this issue.

[willarmiros]

Hi @stfairy,
Thank you for the repro code. After doing a little deep dive, I realized the root cause of the error is because setting the context to async_context actually overrides the default lambda_context that must be present in Lambda environments for X-Ray to work. So actually just removing the context=AsyncContext() configuration parameter and keeping everything else the same will record a subsegment for the download of each file.

That being said, this is clearly a bug to not be able to support AsyncContext and LambdaContext at once, since they are not mutually exclusive. Also, even though removing the AsyncContext works in this simple example, it will likely cause issues with more complex applications of aiohttp, we should prioritize a fix soon.

[heitorlessa]

Do you actually need that to use AsyncContext Lambda? I'm trying to learn more about Async and X-Ray operations so this might be a dummy question

In my tests, I simply got aws_xray_sdk.core.xray_recorder and all of its operations worked just fine - aioboto, aiohttp, etc.

Code exerpt run in Python3.7 runtime

import asyncio
import json

import aws_xray_sdk.core
from aws_xray_sdk.ext.aiohttp.client import aws_xray_trace_config

import aiohttp

recorder = aws_xray_sdk.core.xray_recorder
...


async def aiohttp_task():
    async with aiohttp.ClientSession(trace_configs=[aws_xray_trace_config()]) as session:
        async with session.get("https://httpbin.org/json") as resp:
            resp = await resp.json()
            print(resp)
            return resp

def lambda_handler(evt, ctx):
    ret = asyncio.run(aiohttp_task())
    return {
        "statusCode": 200,
        "body": json.dumps(ret)
    }

[srprash]

Hi @heitorlessa
The AsyncContext provides overriding methods for TaskFactory and some custom logic for TaskLocalStorage. Not having the AsyncContext in an async program using X-Ray may result in unintended behavior while dealing with multiple tasks and reference management. We are still investigating this issue.

[wojciech-dabrowski]

@willarmiros Any update on this? It prevents from using X-Ray with async code, which sounds ridiculous :)

[NathanielRN]

Hi @stfairy !

Why the error occurs

In your reproduction code you first globally configure the recorder using xray_recorder.configure(service='repro_xray_issue', context=AsynContext()). However, because there is no event loop available on initialization, AsyncContext() will get a new event loop from the asyncio library:

aws-xray-sdk-python/aws_xray_sdk/core/async_context.py

Lines 17 to 22 in b5b99b7

	def __init__(self, *args, loop=None, use_task_factory=True, **kwargs):
	super(AsyncContext, self).__init__(*args, **kwargs)
	self._loop = loop
	if loop is None:
	self._loop = asyncio.get_event_loop()

The issue is that the loop is not set as the running loop even though it is created. Later, when you use asyncio.run(my_async_func()) to start your running your async code, asyncio will create its own loop and set it as the running loop.

In asyncio:

def run(main, *, debug=None):
# ...
    loop = events.new_event_loop()
        try:
            events.set_event_loop(loop)
# ...

Finally, when the @xray_recorder.capture_async('## name') decorator tries to create the subsegment, it will use the AsyncContext._loop variable which does not equal ≠ the loop that asyncio created! This is important becauseAsyncContext will use its _loop variable to add subsegment information.

aws-xray-sdk-python/aws_xray_sdk/core/async_context.py

Lines 53 to 58 in b5b99b7

	if _GTE_PY37:
	task = asyncio.current_task(loop=self._loop)
	else:
	task = asyncio.Task.current_task(loop=self._loop)
	if task is None:
	return None

AsyncContext expects that its _loop has a task with the current segment. But because the loops aren't the same, no such segment exists on _loop, and you get a No segment found, cannot begin subsegment error.

Solutions

Using AsyncContext

One solution is to put xray_recorder.configure(...AsyncContext() after asyncio.run() but before your async function continues:

async def download_files():
    xray_recorder.configure(service='repro_xray_issue', context=AsyncContext())
    async with xray_recorder.in_segment_async('my_segment_name') as segment:
        # ...

asyncio.run(download_files())

This will let AsyncContext use the currently running loop set by asyncio as its own _loop and add subsegments successfully.

Notice that you must still create your own my_segment_name. This will work in Lambda. However, notice that this will replace the Lambda context resulting in 2 separate traces because there are 2 segments.

This solution is important when you use async code to have multiple segments. The AsyncContext will make sure that new custom subsegments make it to the segment associated with the current task (as explained above) which could be one of many concurrent tasks.

If you are not running on Lambda, this is most likely the solution you need.

Without AsyncContext

As was pointed out before, if you simply remove the context=AsyncContext() configuration you will see the following traces in X-Ray:

Here there is only 1 segment (the one created by LambdaContext in Lambda environments) and the async calls become subsegments of that call.

In this case, any custom subsegments will go under the 1 segment. Since a LambdaContext does not allow you to create new segments, all your subsegments would go under that segment.

Conclusion

I hope that helps and please feel free to ask anymore questions! What helped me discover this behavior was understanding how the event loops are used by asyncio and aws-xray-sdk and looking at the starter code we have on the README.md

@wojciech-dabrowski Please let me know if this explanation helped you! 🙂

[abivolmv]

Hi ,
these solutions will not work if we use asyncio.gather() or asyncio.wait(). There is a longer discussion about it here

[NathanielRN]

Hi @abivolmv,

Sorry to hear it's not working, would you be able to provide simple reproduction code to better describe your issue?

For example, I noticed that in the original description of that issue you posted, the user is trying to configure xray_recorder globally before they call asyncio.gather.

I'm guessing that this would cause a similar problem to the one I pointed out above, because asyncio.gather, like asyncio.run will create and begin its own loop which the AsynContext() method is not aware of because it also created its own loop which differs from asyncio.gather's loop.

In /usr/local/Cellar/python@3.9/3.9.5/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/tasks.py

def gather(*coros_or_futures, loop=None, return_exceptions=False):
    ...
    if not coros_or_futures:
        if loop is None:
            # (NathanielRN): No running loop, so it makes its own which is
            # not seen from previously instantiated `AsyncContext`!
            loop = events.get_event_loop()
        else:
            warnings.warn("The loop argument is deprecated since Python 3.8, "
                          "and scheduled for removal in Python 3.10.",
                          DeprecationWarning, stacklevel=2)

Hopefully there's a way you can can call the asyncio.wait() method first and then configure the AsyncContext so that the problem can be solved. Let me know if this helps! 🙂

[abivolmv]

Hi @NathanielRN , thanks for the quick reply.
I actually use .run and then .wait. So I tried to put in-between the xray_recorder.configure - is that right ?

I tried both solutions from above and either I see no segments and have the issue mentioned in the current description or it fails with the error mentioned in the linked issue.

Some sample code would be as follows (a minimalistic state of my code):

import asyncio
from io import BytesIO

import aioboto3
from aws_xray_sdk.core import xray_recorder, patch_all
from aws_xray_sdk.core.async_context import AsyncContext

# here I tried adding your solution (not added when solution to use AsyncContext() is used below)
# xray_recorder.configure(service='repro_xray_issue')  # uncomment this for second solution trial

patch_all()


def lambda_handler():
    asyncio.run(main())


async def main():
    # here I tried adding your solution (not added when solution to remove AsyncContext() is used above)
    xray_recorder.configure(service='repro_xray_issue', context=AsyncContext())    # comment this for second solution trial
    async with xray_recorder.in_segment_async('my_segment_name') as segment:  # comment this for second solution trial
        filelike1 = BytesIO()
        filelike2 = BytesIO()
        res1, res2 = await asyncio.gather(s3_get(filelike1), s3_get(filelike2))  # .wait() doesn't work either


@xray_recorder.capture_async('s3_get')  # tried with this and without also
async def s3_get(filelike):
    async with aioboto3.Session().client('s3') as s3:
        return await s3.download_fileobj('s3-bucket', 'test.txt', filelike)

if __name__ == '__main__':
    lambda_handler()

First solution will fail with error and second will not display the S3 segments in xray
If I try with :

await s3_get(filelike1)
await s3_get(filelike2)

then it works and I also see the segments in the trace. But I would like to run these calls concurrently.

[NathanielRN]

@abivolmv Thanks for providing that repro code! Actually, you understood my post exactly and this is the code I had in mind 😄

I copied your code exactly and ran it as python3 repro_code.py. However, I actually do see the traces as I expect in X-Ray:

Notice that unfortunately they do render as subsegments of subsegments instead of subsegments of the segment, but that's what we expect. You can still see them running concurrently in this image.

I am running the demon locally as explained in the docs if perhaps that can be compared with how you are running it?

Finally I am using:

• Python 3.9
• aws-xray-sdk 2.8

If you see anything you are doing differently please call it out! Hopefully this helps 🙂