Add memray plugin

[fiedlerNr9]

Why are the changes needed?

• Enables memray profiling on Flyte task level
• renders memray report into Flytedeck

What changes were proposed in this pull request?

• Adding memray flytekit plugin

How was this patch tested?

• unit tests
• tested local & remote run

Setup process

from flytekit import workflow, task, ImageSpec
from flytekitplugins.memray import memray_profiling
import time


flytekit_hash = "82d5ac739f5f02998edb9538c58cf93c8f6e501b"
flytekitplugins_memray = f"git+https://github.com/flyteorg/flytekit.git@{flytekit_hash}#subdirectory=plugins/flytekit-memray"

image = ImageSpec(
    name="memray_demo",
    python_version="3.11.10",
    apt_packages=["git"],
    packages=[flytekitplugins_memray],
    registry="ghcr.io/fiedlernr9",
)


def generate_data(n: int):
    leak_list = []
    for _ in range(n):  # Arbitrary large number for demonstration
        large_data = " " * 10**6  # 1 MB string
        leak_list.append(large_data)  # Keeps appending without releasing
        time.sleep(0.1)  # Slow down the loop to observe memory changes


@task(container_image=image, enable_deck=True)
@memray_profiling(memray_html_reporter="table")
def memory_usage(n: int) -> str:
    generate_data(n=n)

    return "Well"


@task(container_image=image, enable_deck=True)
@memray_profiling(trace_python_allocators=True, memray_reporter_args=["--leaks"])
def memory_leakage(n: int) -> str:
    generate_data(n=n)

    return "Well"


@workflow
def wf(n: int = 500):
    memory_usage(n=n)
    memory_leakage(n=n)

Screenshots

Flamegraph

Table

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 43.89%. Comparing base (3fc51af) to head (3541f5e).
Report is 31 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2875      +/-   ##
==========================================
- Coverage   45.53%   43.89%   -1.65%     
==========================================
  Files         196      199       +3     
  Lines       20418    20820     +402     
  Branches     2647     2676      +29     
==========================================
- Hits         9298     9138     -160     
- Misses      10658    11448     +790     
+ Partials      462      234     -228     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[eapolinario]

This is wicked cool!

Can you add flytekit-memray to https://github.com/flyteorg/flytekit/blob/master/.github/workflows/pythonbuild.yml#L319-L364 ?

[fiedlerNr9]

This is wicked cool!

Can you add flytekit-memray to https://github.com/flyteorg/flytekit/blob/master/.github/workflows/pythonbuild.yml#L319-L364 ?

✅

[thomasjpfan]

Instead of hard coding this into the environment, can we now trace_python_allocators=True?

[fiedlerNr9]

I tested that and its not throwing any warnings but the results look different though:
Thats the task I tested without having the env variable set:

@task(container_image=image, enable_deck=True)
@memray_profiling(trace_python_allocators=True, memray_reporter_args=["--leaks"])
def memory_leakage(n: int) -> str:
    generate_data(n=n)

    return "Well"

Thats the result:

Not sure if this is expected

[thomasjpfan]

For completeness, what do you see when you set trace_python_allocators=False?

[fiedlerNr9]

This, which is expected i guess

[thomasjpfan]

It's weird that it gives two different flamegraphs. The new flamegraph makes more sense to me because the tracker is wrapping the user code and you can clearly see the generete_data.

I can not really see where the generate_data is on your original flamegraph.

[fiedlerNr9]

Oh man, I mixed up my demos. Its looking exactly the same with trace_python_allocaters=False and having the env variable set. Sorry for the confusion, I will update the Readme shortly

[thomasjpfan]

Not actionable for this PR I wish there was a way to ensure that enable_deck=True when using memray_profiling. Otherwise, we just add overhead without any reports.

@eapolinario @pingsutw What do you think of making deck_fields=None and set enable_decks=True?

https://github.com/flyteorg/flytekit/blob/master/flytekit/core/task.py#L203-L210

[eapolinario]

What if we added a way to carry over task kwargs from successive flytekit decorators? Something like #2911.

[eapolinario]

replace mentions to wandb

[fiedlerNr9]

What exactly is supposed to be in this comment section of __init__.py?
I planned to just get rid of it actually?

[eapolinario]

This is exposed as the landing page for the plugin in the docs, e.g. https://docs.flyte.org/en/latest/api/flytekit/plugins/wandb.html

[eapolinario]

In fact, in order to show this new plugin in https://docs.flyte.org/en/latest/api/flytekit/plugins/index.html we need to update the index and add the corresponding file under like the one for wandb.

[fiedlerNr9]

Good catch - updated!
edit: uargh, looks like changes in flytesnacks and flyte are necessary to get the mondocs build to succeed- will follow up

[fiedlerNr9]

Removed adding docs from this PR, since we need a published version of thi plugin on pypi to continue with the process of adding docs to flytesnacks & flyte. Will follow up with docs as soon this is merged.
Ok with you @eapolinario & can you please have a last review if times allows