Hi, we’re hitting an issue where we see excessive memory usage from a task before the Pod gets OOMKilled. This happens even when we reduce the task to do basically nothing:

@task
def train_model(
    some: pd.DataFrame,
    pandas: pd.DataFrame,
    data: pd.DataFrame,
    frames: pd.DataFrame,
) -> train_model_output:
    print("foo")
    return train_model_output(
        model=None
    )

In this case, the biggest file in S3 is about 2.1GB large (Dataframe is about 12GB), so nothing crazy. One other is ~200MB, and others in the order of KBs. But the memory usage is silly high. In the logs it’s not even hitting the print statement, so Im guessing something is up with deserialization? Im seeing boto download logs, and after they stop the memory starts spiking (hence Im thinking perhaps deserialization). Any thoughts on this? Note: We run the exact same workflow and it runs fine for a smaller dataset (~500MB).

This seems like how the dataframes are built

Same behaviour when Im doing:

df_dataset_cols = kwtypes(
    some_field=int,
    some_other_field=str,
    and_another_field=List[int],
    ...
)

@task
def train_model(
    df: Annotated[StructuredDataset, df_dataset_cols],
    ...
) -> train_model_output:
    df = df.open(pd.DataFrame).all()
    ...

interesting

Yeah I understand, Ill try to repro with a minimal example with some random data

thank you

I think I have a relatively minimal example, when I add the column

c

with type

list[int]

I see the same stepwise behaviour resulting in eventual OOMKilled. If I run the example without that column, it finishes successfully:

from typing import List, NamedTuple

try:
    from typing import Annotated
except ImportError:
    from typing_extensions import Annotated

import pandas as pd
import numpy as np
import uuid

from flytekit import (
    task,
    workflow,
    kwtypes,
    Resources,
)
from flytekit.types.structured.structured_dataset import (
    StructuredDataset,
)


df_cols = kwtypes(
    a=int,
    b=int,
    c=List[str],
    d=str,
    e=str,
    f=str,
)


foo_output = NamedTuple(
    "FooOutput",
    [
        ("df_output", Annotated[pd.DataFrame, df_cols]),
    ],
)


@task(
    requests=Resources(cpu="1000m", mem="64Gi"),
    limits=Resources(cpu="2000m", mem="128Gi"),
)
def foo() -> foo_output:
    print("foo")

    num_rows = 50000000

    df1 = pd.DataFrame(
        np.random.randint(0, 100, size=(num_rows, 2)), columns=list("ab")
    )
    df2 = pd.DataFrame(
        np.array([str(uuid.uuid4()) for i in range(3 * num_rows)]).reshape(-1, 3),
        columns=list("def"),
    )
    df2["c"] = df2["d"].map(lambda x: x.split("-"))

    df_output = pd.concat([df1, df2], ignore_index=True, sort=False)
    return foo_output(df_output)


@task(
    requests=Resources(cpu="1000m", mem="64Gi"),
    limits=Resources(cpu="2000m", mem="128Gi"),
)
def bar(df_input: Annotated[StructuredDataset, df_cols]):
    print("bar")
    df_input = df_input.open(pd.DataFrame).all()
    print(df_input)


@workflow
def base_workflow():
    task_foo = foo()
    task_bar = bar(
        df_input=task_foo.df_output,
    )

cc @Pryce

@Geert I'm guessing that in your minimal example you're not even getting to

print('bar')

?

@Pryce you could try "Union" TLM to simply observe this in real time

Sounds good will try to repro on our end

Actually in the minimal example I do see the print, so guess here it's the open/all call? I can run more tests tomorrow if that helps! Thanks for checking on this

Maybe! Will let you know what my testing reveals when I get to it..

In our original workflow, it runs if I remove the

list[..]

columns, and instead df.explode() them, and output/input them as separate

pd.DataFrame

in the tasks, and merge the columns again in the next task into the original dataframe. _I_t just takes a lot of extra processing time to do the explode and merge again in the subsequent task, so it’s not really a fix at the moment (checking if I can speed it up). But no memory issues.

Hey @Geert thanks for your patience. I've been able to reproduce this internally with a smaller dataset. I'm still chasing down why those datastructures are ballooning so much though.