Memory leak from use of @task decorator when using asyncio and lxml

[andywirv]

Prefect Slack Discussion

N.B In the process of creating the minimal reproduction I have managed to resolve my immediate issue but I do think there is a problem here that others may face.

Fix/workaround : Stopping element objects being returned by casting to str and int where appropriate:

data[key] = extracted_value if key in numeric_fields else extracted_value

vs

data[key] = int(extracted_value) if key in numeric_fields else str(extracted_value)

This Prefect flow extracts file lists from a Google Cloud Storage (GCS) bucket and then processes them in batches. Each batch downloads XML files from GCS and then uses lxml to extract specific fields using XPath. Memory usage when executing as Prefect flow without the @task decorator is stable but with the decorator each iteration of the loop sees memory increase until all memory is exhausted.

This shows memory usage across 4 runs. (ignore line at 6Gb which is the memory limit)

1. @task decorator added to original code. Peaks over 2GB
2. @task decorator removed. Stable at 400Mb
3. @task decorator added. Code changed to cast xml element used in return. Stable at 300Mb
4. Code reverted and @task decorator added

Output from objgraph for one random element object (as profiling showed a steady increase in element objects over time)

Full project available here:
https://gitlab.com/vee-engineering-public/issues/prefect-xml-memory

import aiohttp
import asyncio
from asyncio import Semaphore
import aiohttp
from lxml import etree
from prefect import flow, task
from prefect.logging import get_run_logger
from gcloud.aio.storage import Storage, Bucket
from gcloud.aio.auth import Token
import time 

@flow
async def main(bucket_name, max_files_per_folder=1000, batch_size=500):
    file_list = await list_files(directory="", pattern='*.xml', bucket_name=bucket_name)
    file_list = file_list[:int(max_files_per_folder)]
    for i in range(0, len(file_list), batch_size):
        batch = file_list[i:i + batch_size]
        await process_batch(bucket_name=bucket_name, batch=batch)

@task(task_run_name=f"process_batch", persist_result=False)
async def process_batch(bucket_name, batch):
    logger = get_run_logger()
    data_batch = []
    semaphore = Semaphore(20)
    logger.info(f"Processing batch: {batch}")
        
    async with aiohttp.ClientSession() as session:
        async with await get_storage_client(session) as storage:
            for count, gcs_file_url in enumerate(batch):
                logger.debug(f"Processing file: {gcs_file_url} : {count}")
                xpath_fields = get_fields()
            data_batch = await asyncio.gather(
                *[getDataWithSemaphore(storage=storage, xpath_fields=xpath_fields, gcs_file_url=gcs_file_url, bucket_name=bucket_name, semaphore=semaphore) for gcs_file_url in batch]
            )
            # This is just to slow the flow down so not so much data is processed to get memory stats
            await slow_task()
            
    return data_batch

async def getDataWithSemaphore(storage, xpath_fields, gcs_file_url, bucket_name, semaphore):
    async with semaphore:
        return await getData(storage, xpath_fields, gcs_file_url, bucket_name)

async def getData(storage, xpathFields: dict, gcs_file_url: str, bucket_name: str):
    """
    Parses an XML file from a GCS URL to extract the data using XPath fields.
    param: gcs_file_url - The URL of the XML file in GCS.
    param: xpath_fields - The XPath fields to extract from the XML file.
    """
    logger = get_run_logger()
    logger.info(f"Parsing file: {gcs_file_url}")
    
    file = await getFile(storage, gcs_file_url, bucket_name)
    extracted_data, error = extract_data(file, xpathFields, gcs_file_url)
    if error:
        logger.error(error)
    logger.debug(f"Data extracted from {gcs_file_url}")
    return extracted_data

def get_fields():
    fields = {
        'ORG_EIN': {
            'xpath': '../irs:Return/irs:ReturnHeader/irs:Filer/irs:EIN/text()',
            'mode': 'REQUIRED',
        },
        'TAX_YEAR': {
            'xpath': '../irs:Return/irs:ReturnHeader/irs:TaxYr/text()',
            'mode': 'REQUIRED',
        },
        'RETURN_TYPE': {
            'xpath': '../irs:Return/irs:ReturnHeader/irs:ReturnTypeCd/text()',
            'mode': 'REQUIRED',
        },
        'BUSINESS_NAME_L1': {
            'xpath': '../irs:Return/irs:ReturnHeader/irs:Filer/irs:BusinessName/irs:BusinessNameLine1Txt/text()',
            'mode': 'REQUIRED',
        },
        'BUSINESS_NAME_L2': {
            'xpath': '../irs:Return/irs:ReturnHeader/irs:Filer/irs:BusinessName/irs:BusinessNameLine2Txt/text()',
            'mode': 'NULLABLE',
        },
        'BUSINESS_ADDR_L1': {
            'xpath': '../irs:Return/irs:ReturnHeader/irs:Filer/irs:USAddress/irs:AddressLine1Txt/text()',
            'mode': 'NULLABLE',
        },
        'BUSINESS_CITY': {
            'xpath': '../irs:Return/irs:ReturnHeader/irs:Filer/irs:USAddress/irs:CityNm/text()',
            'mode': 'NULLABLE',
        },
        'BUSINESS_STATE': {
            'xpath': '../irs:Return/irs:ReturnHeader/irs:Filer/irs:USAddress/irs:StateAbbreviationCd/text()',
            'mode': 'NULLABLE',
        },
        'ORG_WEBSITE': {
            'xpath': '../irs:Return/irs:ReturnData/*/irs:WebsiteAddressTxt/text()',
            'mode': 'NULLABLE',
        },
        'SYS_SCHEMA_VERSION': {
            'xpath': '../irs:Return/@returnVersion',
            'mode': 'REQUIRED',
        },
        'GROSS_RECEIPTS': {
            'xpath': '../irs:Return/irs:ReturnData/*/irs:GrossReceiptsAmt/text()',
            'mode': 'NULLABLE',
        },
        'INITIAL_RETURN': {
            'xpath': 'boolean(../irs:Return/irs:ReturnData/*/irs:InitialReturnInd/text()="X")',
            'mode': 'NULLABLE',
        },
    }
    return fields

async def list_files(directory: str, pattern: str, bucket_name: str):
    async with aiohttp.ClientSession() as session:
        client = Storage(session=session)
        bucket = Bucket(client, bucket_name)
        fileList = await bucket.list_blobs(prefix=directory, match_glob=pattern)
        return fileList
from lxml import etree
def extract_data(xml_data, xpath_fields, gcsFileUrl):
    """
    Extracts data from XML data using XPath expressions.
    
    :param xml_data: XML data as a string
    :param xpath_fields: Dictionary of field names and their XPath expressions
    """

    namespaces = {
        'irs': 'http://www.irs.gov/efile',
        'xsi': 'http://www.w3.org/2001/XMLSchema-instance'
    }

    numeric_fields = [
        "FRGN_ORG_GRANT_AMT",
        "FRGN_ORG_GRANT_NC_AMT", 
        "GRANT_US_ORG_AMT_NONCASH", 
        "GRANT_US_ORG_AMT_CASH",
        "ORG_EIN", "TAX_YEAR"
    ]

    root = etree.fromstring(xml_data)
    
    data = {}
    error = []
    for key, item in xpath_fields.items():
        try:
            value = root.xpath(item['xpath'], namespaces=namespaces)
            if isinstance(value, bool):
                data[key] = bool(value)
            else:
                if value:
                    extracted_value = value[0]
                    data[key] = extracted_value if key in numeric_fields else str(extracted_value)
                else:
                    data[key] = None
        except Exception as e:
            error.append(f"Error extracting {key}: {e} on {gcsFileUrl}")
            data[key] = None
    root = None
    return data, error

async def get_storage_client(session) -> Storage:
    logger = get_run_logger()
    try:
        token = Token(session=session,)
        storage_client = Storage(token=token)
        return storage_client
    except Exception as e:
        logger.error(f"Failed to create storage client: {e}")
        return None
    
async def getFile(storage: Storage, fileUrl: str, bucket_name: str):
    file = await storage.download(bucket=bucket_name, object_name=fileUrl)
    return file

async def slow_task():
    await asyncio.sleep(20)
    return
           
if __name__ == "__main__":
    parameters = {
        "bucket_name": "prefect-task-memory-issue",
        "max_files_per_folder": 1000,
        "batch_size": 500
    }
    asyncio.run(main(**parameters))