finish readme

Signed-off-by: Jan Fiedler <[email protected]>

plugins/flytekit-memray/README.md

@@ -1 +1,55 @@
-# Memray
+# Memray Profiling Plugin
+
+Memray tracks and reports memory allocations, both in python code and in compilled extension modules.
+This Memray Profiling plugin enables memory tracking on the Flyte task level and renders a memgraph profiling graph on Flyte Deck.
+
+To install the plugin, run the following command:
+
+```bash
+pip install flytekitplugins-memray
+```
+
+Example
+```python
+from flytekit import workflow, task, ImageSpec
+from flytekitplugins.memray import memray_profiling
+import time
+
+
+image = ImageSpec(
+    name="memray_demo",
+    packages=["flytekitplugins_memray"],
+    env={"PYTHONMALLOC": "malloc"},
+    registry="<your_cr_registry>",
+)
+
+
+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(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)
+```

plugins/flytekit-memray/flytekitplugins/memray/profiling.py

@@ -12,22 +12,34 @@ def __init__(
         self,
         task_function: Optional[Callable] = None,
         memray_html_reporter: str = "flamegraph",
-        memray_reporter_args: Optional[List[str]] = [],
+        memray_reporter_args: List[str] = [],
     ):
-        """Memray Profiling Plugin.
+        """Memray profiling plugin.
         Args:
+            task_function (function, optional): The user function to be decorated. Defaults to None.
+            memray_html_reporter (str): The name of the memray reporter which generates an html report.
+                Today there is only 'flamegraph' & 'table'.
+            memray_reporter_args (List[str], optional): A list of arguments to pass to the reporter commands.
+                See the [flamegraph](https://bloomberg.github.io/memray/flamegraph.html#reference)
+                and [table](https://bloomberg.github.io/memray/table.html#cli-reference) docs for details on supported arguments.
         """
+
         if memray_html_reporter not in ["flamegraph", "table"]:
             raise ValueError(
                 f"{memray_html_reporter} is not a supported html reporter."
             )
 
+        if not all(
+            isinstance(arg, str) and "--" in arg for arg in memray_reporter_args
+        ):
+            raise ValueError(
+                f"unrecognized arguments for {memray_html_reporter} reporter. Please check https://bloomberg.github.io/memray/{memray_html_reporter}.html"
+            )
+
         self.dir_name = "memray"
         self.memray_html_reporter = memray_html_reporter
         self.memray_reporter_args = memray_reporter_args
 
-        # All kwargs need to be passed up so that the function wrapping works for both
-        # `@wandb_init` and `@wandb_init(...)`
         super().__init__(
             task_function,
             memray_html_reporter=memray_html_reporter,

plugins/flytekit-memray/flytekitplugins/memray/tests/test_memray_profiling.py

@@ -0,0 +1,39 @@
+from unittest.mock import Mock, patch
+import pytest
+from flytekit import task, current_context
+from flytekitplugins.memray import memray_profiling
+
+
+@task(enable_deck=True)
+@memray_profiling
+def heavy_compute(i: int) -> int:
+    return i + 1
+
+
+def test_local_exec():
+    heavy_compute(i=7)
+    assert (
+        len(current_context().decks) == 6
+    )  # memray flamegraph, timeline, input, and output, source code, dependencies
+
+
+def test_errors():
+    reporter = "summary"
+    with pytest.raises(
+        ValueError, match=f"{reporter} is not a supported html reporter."
+    ):
+        memray_profiling(memray_html_reporter=reporter)
+
+    reporter = "flamegraph"
+    with pytest.raises(
+        ValueError,
+        match=f"unrecognized arguments for {reporter} reporter. Please check https://bloomberg.github.io/memray/{reporter}.html",
+    ):
+        memray_profiling(memray_reporter_args=["--leaks", "trash"])
+
+    reporter = "flamegraph"
+    with pytest.raises(
+        ValueError,
+        match=f"unrecognized arguments for {reporter} reporter. Please check https://bloomberg.github.io/memray/{reporter}.html",
+    ):
+        memray_profiling(memray_reporter_args=[0, 1, 2])