From dd956d508c95d4c61196f44bdba5ce04bd667470 Mon Sep 17 00:00:00 2001
From: Jan Fiedler <jan@union.ai>
Date: Fri, 8 Nov 2024 12:09:01 -0800
Subject: [PATCH] add memray example

---
 examples/memray_plugin/Dockerfile             | 27 ++++++++
 examples/memray_plugin/README.md              | 20 ++++++
 .../memray_plugin/memray_plugin/__init__.py   |  0
 .../memray_plugin/memray_example.py           | 66 +++++++++++++++++++
 examples/memray_plugin/requirements.in        |  1 +
 5 files changed, 114 insertions(+)
 create mode 100644 examples/memray_plugin/Dockerfile
 create mode 100644 examples/memray_plugin/README.md
 create mode 100644 examples/memray_plugin/memray_plugin/__init__.py
 create mode 100644 examples/memray_plugin/memray_plugin/memray_example.py
 create mode 100644 examples/memray_plugin/requirements.in

diff --git a/examples/memray_plugin/Dockerfile b/examples/memray_plugin/Dockerfile
new file mode 100644
index 000000000..65da827a3
--- /dev/null
+++ b/examples/memray_plugin/Dockerfile
@@ -0,0 +1,27 @@
+FROM python:3.11-slim-bookworm
+LABEL org.opencontainers.image.source=https://github.com/flyteorg/flytesnacks
+
+WORKDIR /root
+ENV VENV /opt/venv
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8
+ENV PYTHONPATH /root
+
+WORKDIR /root
+
+ENV VENV /opt/venv
+# Virtual environment
+RUN python3 -m venv ${VENV}
+ENV PATH="${VENV}/bin:$PATH"
+
+# Install Python dependencies
+COPY requirements.in /root
+RUN pip install -r /root/requirements.in
+
+# Copy the actual code
+COPY . /root
+
+# This tag is supplied by the build script and will be used to determine the version
+# when registering tasks, workflows, and launch plans
+ARG tag
+ENV FLYTE_INTERNAL_IMAGE $tag
diff --git a/examples/memray_plugin/README.md b/examples/memray_plugin/README.md
new file mode 100644
index 000000000..4a0c71549
--- /dev/null
+++ b/examples/memray_plugin/README.md
@@ -0,0 +1,20 @@
+(memray_plugin)=
+
+# Memray Profiling
+
+```{eval-rst}
+.. tags:: Integration, Profiling, Observability
+```
+
+Memray tracks and reports memory allocations, both in python code and in compiled extension modules.
+This Memray Profiling plugin enables memory tracking on the Flyte task level and renders a memgraph profiling graph on Flyte Deck.
+
+First, install the Memray plugin:
+
+```bash
+pip install flytekitplugins-memray
+```
+
+```{auto-examples-toc}
+memray_example
+```
diff --git a/examples/memray_plugin/memray_plugin/__init__.py b/examples/memray_plugin/memray_plugin/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/examples/memray_plugin/memray_plugin/memray_example.py b/examples/memray_plugin/memray_plugin/memray_example.py
new file mode 100644
index 000000000..69351a2f0
--- /dev/null
+++ b/examples/memray_plugin/memray_plugin/memray_example.py
@@ -0,0 +1,66 @@
+# %% [markdown]
+# (memray_example)=
+#
+# # Memray Profiling Example
+# Memray tracks and reports memory allocations, both in python code and in compiled extension modules.
+# This Memray Profiling plugin enables memory tracking on the Flyte task level and renders a memgraph profiling graph on Flyte Deck.
+# %%
+from flytekit import workflow, task, ImageSpec
+from flytekitplugins.memray import memray_profiling
+import time
+
+# %% [markdown]
+# First, we use `ImageSpec` to construct a container that contains the dependencies for the
+# tasks, we want to profile:
+# %%
+image = ImageSpec(
+    name="memray_demo",
+    packages=["flytekitplugins_memray"],
+    registry="<your_cr_registry>",
+)
+
+
+# %% [markdown]
+# Next, we define a dummy function that generates data in memory without releasing:
+# %%
+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
+
+
+# %% [markdown]
+# Example of profiling the memory usage of `generate_data()` via the memray `table` html reporter
+# %%
+@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"
+
+
+# %% [markdown]
+# Example of profiling the memory leackage of `generate_data()` via the memray `flamegraph` html reporter
+# %%
+
+
+@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"
+
+
+# %% [markdown]
+# Put everything together in a workflow.
+# %%
+
+
+@workflow
+def wf(n: int = 500):
+    memory_usage(n=n)
+    memory_leakage(n=n)
diff --git a/examples/memray_plugin/requirements.in b/examples/memray_plugin/requirements.in
new file mode 100644
index 000000000..f4663640b
--- /dev/null
+++ b/examples/memray_plugin/requirements.in
@@ -0,0 +1 @@
+flytekitplugins-memray
\ No newline at end of file