diff --git a/kolkhoz/capture.py b/kolkhoz/capture.py index 53dea56..56adef9 100644 --- a/kolkhoz/capture.py +++ b/kolkhoz/capture.py @@ -18,6 +18,7 @@ import os import fsspec +from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper from PIL import Image from pravda import Pravda, PravdaConfig, Snapshot @@ -46,12 +47,22 @@ def storage_filesystem(settings: PravdaSettings): Pravda resolves each snapshot's ``prefix`` against this same base path, so opening ``/`` on this filesystem locates the artifact for both local paths and remote (``gs://``/``s3://``) URLs. + + Synchronous backends (e.g. the local filesystem) are wrapped so reads use + the same async API as remote ones. This mirrors Pravda's own + ``Storage.from_url`` and, crucially, keeps every artifact read on the + running event loop: reading via the async API avoids fsspec's sync bridge, + which would otherwise drive the shared (async, e.g. ``gcsfs``) instance from + its background loop while Pravda has bound its session to this loop — + raising "got Future attached to a different loop". """ fs, _ = fsspec.core.url_to_fs(settings.storage_base_path) + if not fs.async_impl: + fs = AsyncFileSystemWrapper(fs) return fs -def read_artifact(fs, snapshot: Snapshot, filename: str | None) -> bytes: +async def read_artifact(fs, snapshot: Snapshot, filename: str | None) -> bytes: """Read a snapshot artifact blob from the shared storage backend. ``snapshot.prefix`` is the backend-resolved directory (base path plus the @@ -59,14 +70,16 @@ def read_artifact(fs, snapshot: Snapshot, filename: str | None) -> bytes: content-addressed name Pravda stored. Both are required for a stored artifact: a missing one is a malformed snapshot, so this fails loud rather than returning empty bytes. + + Reads through the async API on the current event loop; see + ``storage_filesystem`` for why the sync ``fs.open`` path is avoided. """ if snapshot.prefix is None: raise ValueError(f"snapshot {snapshot.id} has no storage prefix") if filename is None: raise ValueError(f"snapshot {snapshot.id} has no artifact filename") path = os.path.join(snapshot.prefix, filename) - with fs.open(path, "rb") as fh: - return fh.read() + return await fs._cat_file(path) def is_blank(blob: bytes) -> bool: diff --git a/kolkhoz/cli.py b/kolkhoz/cli.py index cefd048..a9343a1 100644 --- a/kolkhoz/cli.py +++ b/kolkhoz/cli.py @@ -43,10 +43,10 @@ async def extract_snapshot( client: OpenAI, ) -> list[dict]: """Extract flattened holder observations from one captured snapshot.""" - text = read_artifact(fs, snapshot, snapshot.plaintext).decode( + text = (await read_artifact(fs, snapshot, snapshot.plaintext)).decode( "utf-8", errors="replace" ) - html = read_artifact(fs, snapshot, snapshot.rendered_html).decode( + html = (await read_artifact(fs, snapshot, snapshot.rendered_html)).decode( "utf-8", errors="replace" ) @@ -56,7 +56,7 @@ async def extract_snapshot( if reason is not None: log.info(" → %s → including screenshot", reason) if snapshot.screenshot is not None: - blob = read_artifact(fs, snapshot, snapshot.screenshot) + blob = await read_artifact(fs, snapshot, snapshot.screenshot) if not is_blank(blob): screenshot_blob = blob metadata = metadata_from_html(snapshot.url, html) @@ -158,7 +158,7 @@ async def _run_pipeline( if holders: hits += 1 - write_outputs(groups, config.paths) + await write_outputs(groups, config.paths) log.info("extraction: %d hit, %d miss", hits, extracted - hits) diff --git a/kolkhoz/export.py b/kolkhoz/export.py index 3a792e2..b990950 100644 --- a/kolkhoz/export.py +++ b/kolkhoz/export.py @@ -6,6 +6,7 @@ from datetime import datetime import fsspec +from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper from pravda import Snapshot from kolkhoz.config import PathsConfig @@ -56,19 +57,23 @@ def holder_to_record( return {name: record[name] for name in EXPORT_FIELDS} -def write_outputs(groups: dict[str, list[dict]], paths: PathsConfig) -> None: +async def write_outputs(groups: dict[str, list[dict]], paths: PathsConfig) -> None: """Write exactly the records produced by the current run.""" fs, base = fsspec.core.url_to_fs(paths.output_base_path) + if not fs.async_impl: + # Consistently use async API on pipeline's event loop + fs = AsyncFileSystemWrapper(fs) date = datetime.now().strftime(EXPORT_DATE_FORMAT) total = 0 for dataset, records in groups.items(): out_dir = os.path.join(base, dataset) out_file = os.path.join(out_dir, f"{date}.jsonl") - fs.makedirs(out_dir, exist_ok=True) - with fs.open(out_file, "wb") as fh: - for record in records: - fh.write(json.dumps(record, ensure_ascii=False).encode("utf-8")) - fh.write(b"\n") + await fs._makedirs(out_dir, exist_ok=True) + payload = b"".join( + json.dumps(record, ensure_ascii=False).encode("utf-8") + b"\n" + for record in records + ) + await fs._pipe_file(out_file, payload) total += len(records) log.info("wrote %d record(s) → %s", len(records), out_file) log.info("wrote %d record(s) across %d dataset(s)", total, len(groups))