Create convenience decorator as_step_fn for building StepSpecs

This allows for decorating a function that you want to use in a `StepSpec` and avoiding the `lambda op: ...` boilerplate.

Author: wmoss

experiments/dedup/fineweb_10bt_exact.py

@@ -33,9 +33,8 @@ def build_steps() -> list[StepSpec]:
         name="exact_dedup_fineweb_10bt",
         output_path_prefix=f"{marin_prefix()}/tmp/{OUTPUT_PREFIX}",
         deps=[download],
-        fn=lambda op: dedup_exact_paragraph(
+        fn=dedup_exact_paragraph(
             input_paths=os.path.join(download.output_path, "sample/10BT"),
-            output_path=op,
             max_parallelism=128,
         ),
     )

lib/marin/src/marin/execution/step_spec.py

@@ -8,13 +8,27 @@
 import json
 from collections.abc import Callable
 from dataclasses import dataclass
-from functools import cached_property
+from functools import cached_property, wraps
 from typing import Any
 from urllib.parse import urlparse
 
 from rigging.filesystem import marin_prefix
 
 
+def as_step_fn(fn: Callable[..., Any]) -> Callable[..., Any]:
+    """Decorator that allows a function to either be called normally with output_path,
+    or curried by omitting output_path, returning a Callable[[str], Any] suitable for StepSpec.fn."""
+
+    @wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        def inner(output_path: str) -> Any:
+            return fn(*args, output_path=output_path, **kwargs)
+
+        return inner
+
+    return wrapper
+
+
 def _is_relative_path(url_or_path: str) -> bool:
     """Return True if the path is relative (not a URL and doesn't start with /)."""
     if urlparse(url_or_path).scheme:

lib/marin/src/marin/processing/classification/deduplication/exact.py

@@ -23,6 +23,7 @@
 from fray import ResourceConfig
 from zephyr import ZephyrContext, counters, write_parquet_file
 from zephyr.dataset import Dataset
+from marin.execution.step_spec import as_step_fn
 
 logger = logging.getLogger(__name__)
 
@@ -46,6 +47,7 @@ def _iter_has_more_than_one(records: Iterator[T]) -> tuple[bool, T, Iterator[T]]
     return has_more_than_one, first, itertools.chain([first], rest)
 
 
+@as_step_fn
 def dedup_exact_paragraph(
     *,
     input_paths: str | list[str],