Fix fuzzy dedup producing zero duplicates under Ray

When _load_fuzzy_dupe_map_shard and _compute_fuzzy_dedup_stats created a
ZephyrContext with client=LocalClient(), the coordinator job called
current_client() which auto-detected Ray and spawned workers as separate
processes. The add_to_dup_map closure then modified serialized copies of
shard_dup_map, leaving the original dict empty. Wrap both call sites with
set_current_client so the LocalClient propagates via context variables.

Author: ravwojdyla

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

@@ -7,6 +7,7 @@
 import logging
 from marin.utils import rebase_file_path
 import pyarrow as pa
+from fray.v2.client import set_current_client
 from fray.v2.local_backend import LocalClient
 from marin.processing.classification.deduplication.dedup_commons import (
     DEFAULT_FILETYPES,
@@ -31,23 +32,25 @@
 
 def _compute_fuzzy_dedup_stats(shards: list[str] | Sequence[str], method: str, level: str) -> DupCounters:
     with log_time(f"Compute fuzzy deduplication stats from {len(shards)} shards"):
-        ctx = ZephyrContext(client=LocalClient(), name="fuzzy-dup-counts")
-        result: DupCounters = ctx.execute(  # type: ignore[bad-assignment]
-            Dataset.from_list(shards)
-            .load_parquet(columns=["component_id"])
-            # Compute the per-component statistics and then roll them up into a single counter group
-            .group_by(
-                key=lambda r: r["component_id"],
-                reducer=lambda _, items: DupCounters(
-                    method=method,
-                    level=level,
-                    total=(total := sum(1 for _ in items)),
-                    dups=total if total > 1 else 0,
-                    unique=1,
-                ),
-            )
-            .reduce(partial(sum, start=DupCounters(method=method, level=level))),
-        )[0]
+        client = LocalClient()
+        with set_current_client(client):
+            ctx = ZephyrContext(client=client, name="fuzzy-dup-counts")
+            result: DupCounters = ctx.execute(  # type: ignore[bad-assignment]
+                Dataset.from_list(shards)
+                .load_parquet(columns=["component_id"])
+                # Compute the per-component statistics and then roll them up into a single counter group
+                .group_by(
+                    key=lambda r: r["component_id"],
+                    reducer=lambda _, items: DupCounters(
+                        method=method,
+                        level=level,
+                        total=(total := sum(1 for _ in items)),
+                        dups=total if total > 1 else 0,
+                        unique=1,
+                    ),
+                )
+                .reduce(partial(sum, start=DupCounters(method=method, level=level))),
+            )[0]
     return result
 
 
@@ -63,10 +66,12 @@ def add_to_dup_map(record: dict):
         shard_dup_map[record["id"]] = record["fuzzy_duplicate"]
 
     with log_time(f"Load fuzzy duplicate map from {len(shards)} shards"):
-        ctx = ZephyrContext(client=LocalClient(), name="fuzzy-dup-map")
-        ctx.execute(
-            Dataset.from_list(shards).load_parquet().map(add_to_dup_map),
-        )
+        client = LocalClient()
+        with set_current_client(client):
+            ctx = ZephyrContext(client=client, name="fuzzy-dup-map")
+            ctx.execute(
+                Dataset.from_list(shards).load_parquet().map(add_to_dup_map),
+            )
 
     return shard_dup_map