diff --git a/nemo_curator/stages/text/download/common_crawl/download.py b/nemo_curator/stages/text/download/common_crawl/download.py index a323c599cf..2c3630fc94 100644 --- a/nemo_curator/stages/text/download/common_crawl/download.py +++ b/nemo_curator/stages/text/download/common_crawl/download.py @@ -22,8 +22,8 @@ import pandas as pd import requests +from fastwarc.warc import ArchiveIterator, WarcRecordType from loguru import logger -from warcio.archiveiterator import ArchiveIterator from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.text.download import DocumentDownloader @@ -285,10 +285,11 @@ def _read_warc_record_s3(self, row: pd.Series) -> bytes | None: try: stream = io.BytesIO(decompressed) - archive_iterator = ArchiveIterator(stream) + archive_iterator = ArchiveIterator( + stream, record_types=WarcRecordType.response, auto_decode="all", strict_mode=False + ) for record in archive_iterator: - if record.rec_type == "response": - return record.content_stream().read() + return record.reader.read() except Exception as e: # noqa: BLE001 logger.debug(f"Failed to parse WARC record {filename}: {e}, returning decompressed bytes") return decompressed @@ -302,13 +303,13 @@ def _read_warc_record_s3(self, row: pd.Series) -> bytes | None: logger.warning(f"S3 fetch failed for {filename}: {e}") return None - def _read_warc_record(self, row: pd.Series) -> bytes | None: # noqa: C901, PLR0911 + def _read_warc_record(self, row: pd.Series) -> bytes | None: # noqa: PLR0911 """Fetch a single WARC record using HTTPS range request. This method: 1. Fetches gzip-compressed WARC record bytes via HTTP range request 2. Decompresses the gzip content - 3. Parses the WARC record format using warcio + 3. Parses the WARC record format using fastwarc 4. Extracts and returns the HTTP response body (the actual content) """ filename = None @@ -349,14 +350,15 @@ def _read_warc_record(self, row: pd.Series) -> bytes | None: # noqa: C901, PLR0 # Content might not be gzip-compressed, use as-is decompressed = raw_bytes - # Parse the WARC record using warcio to extract HTTP response body + # Parse the WARC record using fastwarc to extract HTTP response body try: stream = io.BytesIO(decompressed) - archive_iterator = ArchiveIterator(stream) + archive_iterator = ArchiveIterator( + stream, record_types=WarcRecordType.response, auto_decode="all", strict_mode=False + ) for record in archive_iterator: - if record.rec_type == "response": - # Return the HTTP response body (content after HTTP headers) - return record.content_stream().read() + # Return the HTTP response body (content after HTTP headers) + return record.reader.read() except Exception as e: # noqa: BLE001 logger.debug(f"Failed to parse WARC record {filename}: {e}, returning decompressed bytes") return decompressed diff --git a/nemo_curator/stages/text/download/common_crawl/warc_iterator.py b/nemo_curator/stages/text/download/common_crawl/warc_iterator.py index fa430820c7..44cdeed9f0 100644 --- a/nemo_curator/stages/text/download/common_crawl/warc_iterator.py +++ b/nemo_curator/stages/text/download/common_crawl/warc_iterator.py @@ -16,12 +16,10 @@ from pathlib import Path from typing import Any +from fastwarc.warc import ArchiveIterator, WarcRecordType from fsspec.core import url_to_fs from loguru import logger -# TODO: Consider using fastwarc https://github.com/NVIDIA-NeMo/Curator/issues/778 -from warcio.archiveiterator import ArchiveIterator - from nemo_curator.stages.text.download import DocumentIterator @@ -46,19 +44,44 @@ def iterate(self, file_path: str) -> Iterator[dict[str, Any]]: num_records = 0 fs, fs_path = url_to_fs(file_path_str, **self.storage_options) with fs.open(fs_path, "rb") as file_pointer: - archive_iterator = ArchiveIterator(file_pointer, arc2warc=True) + # fastwarc wraps any file-like object and sniffs gzip itself, so the fsspec + # handle can be passed straight through. Non-response records are discarded + # in C++ before their headers reach Python, and auto_decode="all" keeps the + # HTTP body decoded from its Content-Encoding, as this iterator did before. + # strict_mode=False resynchronizes past a record with an unparseable WARC + # header instead of silently ending the file there, which is the default. + # + # Gaps that come with fastwarc 0.x, none of them reachable from Common Crawl's + # own files: a record resynchronized past is dropped silently -- the parser + # exposes no skip counter or callback and does not surface the record even + # with record_types=any_type, so there is nothing this loop can log, and a + # short record count is the only symptom; chunked transfer-encoding is not + # decoded even with auto_decode="all"; and the HTTP preamble is only stripped + # from records that declare Content-Type: application/http, which Common Crawl + # emits. ARC input, which the previous arc2warc=True accepted, is not supported. + archive_iterator = ArchiveIterator( + file_pointer, record_types=WarcRecordType.response, auto_decode="all", strict_mode=False + ) while True: try: rec = next(archive_iterator) - if rec.rec_type == "response": - content = rec.content_stream().read() - warc_id = rec.rec_headers.get_header("WARC-Record-ID")[10:-1] - url = rec.rec_headers.get_header("WARC-Target-URI") - yield {"url": url, "warc_id": warc_id, "source_id": filename, "content": content} - num_records += 1 except StopIteration: # End of file reached normally break + except Exception as e: # noqa: BLE001 + # next() has its own try because a stream the parser cannot open at all + # fails here (fastwarc raises StreamError) rather than while reading a + # record, and leaves nothing to resynchronize to. Report it once and stop + # instead of calling next() again on a stream that is already dead. + logger.error(f"Error processing record {num_records} in {filename}: {e!s}") + break + + try: + content = rec.reader.read() + warc_id = rec.headers.get("WARC-Record-ID")[10:-1] + url = rec.headers.get("WARC-Target-URI") + yield {"url": url, "warc_id": warc_id, "source_id": filename, "content": content} + num_records += 1 except Exception as e: # noqa: BLE001 # Handle corruption or other errors logger.error(f"Error processing record {num_records} in {filename}: {e!s}") diff --git a/pyproject.toml b/pyproject.toml index 34d97e159f..c9b19cef44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -197,13 +197,16 @@ translation_all = [ text_cpu = [ # Download / Extract "beautifulsoup4", + # fastwarc and resiliparse pin each other exactly, so this cap also holds resiliparse + # below 1.x; both are used well beyond WARC parsing (see html_extractors/resiliparse.py). + # fastwarc 1.x also drops strict_mode's resync past a corrupt WARC header (warc_iterator.py). + "fastwarc<1", "justext", "lxml", "pycld2", "resiliparse", "s5cmd", "trafilatura==2.0.0", - "warcio", # Filters "fasttext==0.9.3", "sentencepiece", diff --git a/tests/stages/text/download/common_crawl/test_warc_iterator.py b/tests/stages/text/download/common_crawl/test_warc_iterator.py index 82519fda16..23fda2b963 100644 --- a/tests/stages/text/download/common_crawl/test_warc_iterator.py +++ b/tests/stages/text/download/common_crawl/test_warc_iterator.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gzip +from collections.abc import Callable from pathlib import Path from unittest import mock @@ -21,6 +23,24 @@ from nemo_curator.stages.text.download.common_crawl.warc_iterator import CommonCrawlWarcIterator +_OK_BODY = b"ok" +_OK_HTTP = b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n" + _OK_BODY + + +def _response_record(record_id: str, http_payload: bytes, version: str = "WARC/1.0") -> bytes: + """Build a single WARC response record wrapping an already-serialized HTTP response.""" + header = ( + f"{version}\r\n" + f"WARC-Type: response\r\n" + f"WARC-Record-ID: \r\n" + f"WARC-Date: 2022-01-01T00:00:00Z\r\n" + f"WARC-Target-URI: http://example.com/{record_id}\r\n" + f"Content-Type: application/http;msgtype=response\r\n" + f"Content-Length: {len(http_payload)}\r\n" + f"\r\n" + ).encode() + return header + http_payload + b"\r\n\r\n" + class TestCommonCrawlWarcIterator: """Test suite for CommonCrawlWarcIterator - focused on core logic correctness.""" @@ -63,7 +83,7 @@ def test_error_processing_record_continues(self, tmp_path: Path) -> None: raw_warc_path = tmp_path / "test.warc" # Create a WARC file with a response record that has no WARC-Record-ID header - # This will cause the get_header to return None, leading to the subscriptable error + # This makes headers.get return None, leading to the subscriptable error http_response = "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\nTest\r\n" http_response_bytes = http_response.encode("utf-8") content_length = len(http_response_bytes) @@ -119,6 +139,8 @@ def test_mixed_record_types_response_only_with_correct_values(self, tmp_path: Pa "content": "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n

Test Page

\r\n", "id": "response123", "target_uri": "http://example.com/page", + # Marks the payload as an HTTP response, as Common Crawl's own WARC files do. + "content_type": "application/http;msgtype=response", }, { "type": "metadata", @@ -147,6 +169,9 @@ def test_mixed_record_types_response_only_with_correct_values(self, tmp_path: Pa if config["target_uri"]: header_parts.append(f"WARC-Target-URI: {config['target_uri']}\r\n") + if config.get("content_type"): + header_parts.append(f"Content-Type: {config['content_type']}\r\n") + header_parts.append(f"Content-Length: {content_length}\r\n\r\n") warc_record = "".join(header_parts).encode() + content_bytes + b"\r\n\r\n" @@ -166,12 +191,78 @@ def test_mixed_record_types_response_only_with_correct_values(self, tmp_path: Pa assert record["url"] == "http://example.com/page" assert record["warc_id"] == "response123" # Stripped assert record["source_id"] == "mixed_types.warc" - # The content should be just the HTML body (warcio extracts body from HTTP response) + # The content should be just the HTML body, with the HTTP headers stripped assert record["content"] == html_content # Verify the content contains expected HTML assert b"

Test Page

" in record["content"] + @pytest.mark.parametrize( + ("content_encoding", "encode"), + [ + (None, lambda body: body), + ("gzip", gzip.compress), + ], + ) + def test_http_body_is_returned_decoded( + self, + tmp_path: Path, + content_encoding: str | None, + encode: Callable[[bytes], bytes], + ) -> None: + """The yielded content is the decoded HTTP body, whatever Content-Encoding the server used.""" + body = b"decoded" + encoding_header = f"Content-Encoding: {content_encoding}\r\n" if content_encoding else "" + http_payload = (f"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n{encoding_header}\r\n").encode() + encode(body) + + raw_warc_path = tmp_path / "encoded.warc" + raw_warc_path.write_bytes(_response_record("encoded123", http_payload)) + + records = list(CommonCrawlWarcIterator().iterate(str(raw_warc_path))) + + assert len(records) == 1 + assert records[0]["content"] == body + + # expected_errors is 0 for both corrupt-record cases on purpose: fastwarc 0.x + # resynchronizes past the bad record and offers no way to observe that it did, + # so the drop cannot be logged (see the comment in CommonCrawlWarcIterator.iterate). + @pytest.mark.parametrize( + ("warc_bytes", "expected_ids", "expected_errors"), + [ + ( + _response_record("corrupt", _OK_HTTP, version="WARC/XX") + _response_record("good", _OK_HTTP), + ["good"], + 0, + ), + ( + _response_record("first", _OK_HTTP) + + _response_record("corrupt", _OK_HTTP, version="WARC/XX") + + _response_record("last", _OK_HTTP), + ["first", "last"], + 0, + ), + (b"\x00\x01not-a-warc\r\n\r\n" + _response_record("good", _OK_HTTP), [], 1), + ], + ids=["corrupt-first", "corrupt-mid", "unreadable-stream"], + ) + def test_corrupt_record_does_not_abandon_the_rest_of_the_file( + self, + tmp_path: Path, + warc_bytes: bytes, + expected_ids: list[str], + expected_errors: int, + ) -> None: + """A corrupt WARC header is resynchronized past; a stream the parser cannot open is logged once.""" + raw_warc_path = tmp_path / "corrupt.warc" + raw_warc_path.write_bytes(warc_bytes) + + with mock.patch.object(logger, "error") as mock_logger: + yielded = list(CommonCrawlWarcIterator().iterate(str(raw_warc_path))) + + assert [record["warc_id"] for record in yielded] == expected_ids + assert [record["content"] for record in yielded] == [_OK_BODY] * len(expected_ids) + assert mock_logger.call_count == expected_errors + def test_output_columns(self) -> None: """Test that output_columns returns the expected column names.""" iterator = CommonCrawlWarcIterator() diff --git a/uv.lock b/uv.lock index 479c9079dc..e0acf95eb3 100644 --- a/uv.lock +++ b/uv.lock @@ -5567,6 +5567,7 @@ all = [ { name = "easydict" }, { name = "einops" }, { name = "fasttext" }, + { name = "fastwarc" }, { name = "ftfy" }, { name = "google-cloud-translate" }, { name = "gpustat" }, @@ -5625,7 +5626,6 @@ all = [ { name = "transformers" }, { name = "vllm", version = "0.22.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform != 'darwin' and sys_platform != 'linux'" }, { name = "vllm", version = "0.22.0+cu129", source = { registry = "https://wheels.vllm.ai/0.22.0/cu129" }, extra = ["otel", "runai"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "warcio" }, { name = "whisperx", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, ] audio-common = [ @@ -5759,6 +5759,7 @@ math-cpu = [ { name = "beautifulsoup4" }, { name = "boto3" }, { name = "fasttext" }, + { name = "fastwarc" }, { name = "ftfy" }, { name = "justext" }, { name = "lxml" }, @@ -5771,7 +5772,6 @@ math-cpu = [ { name = "sentence-transformers" }, { name = "sentencepiece" }, { name = "trafilatura" }, - { name = "warcio" }, ] math-cuda12 = [ { name = "beautifulsoup4" }, @@ -5780,6 +5780,7 @@ math-cuda12 = [ { name = "cudf-streaming-cu12" }, { name = "cuml-cu12" }, { name = "fasttext" }, + { name = "fastwarc" }, { name = "ftfy" }, { name = "gpustat" }, { name = "justext" }, @@ -5800,7 +5801,6 @@ math-cuda12 = [ { name = "trafilatura" }, { name = "vllm", version = "0.22.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform != 'darwin' and sys_platform != 'linux'" }, { name = "vllm", version = "0.22.0+cu129", source = { registry = "https://wheels.vllm.ai/0.22.0/cu129" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "warcio" }, ] sdg-cpu = [ { name = "data-designer" }, @@ -5817,6 +5817,7 @@ sdg-cuda12 = [ text-cpu = [ { name = "beautifulsoup4" }, { name = "fasttext" }, + { name = "fastwarc" }, { name = "ftfy" }, { name = "justext" }, { name = "lxml" }, @@ -5828,7 +5829,6 @@ text-cpu = [ { name = "sentence-transformers" }, { name = "sentencepiece" }, { name = "trafilatura" }, - { name = "warcio" }, ] text-cuda12 = [ { name = "beautifulsoup4" }, @@ -5836,6 +5836,7 @@ text-cuda12 = [ { name = "cudf-streaming-cu12" }, { name = "cuml-cu12" }, { name = "fasttext" }, + { name = "fastwarc" }, { name = "ftfy" }, { name = "gpustat" }, { name = "justext" }, @@ -5854,7 +5855,6 @@ text-cuda12 = [ { name = "sentencepiece" }, { name = "trafilatura" }, { name = "vllm", version = "0.22.0+cu129", source = { registry = "https://wheels.vllm.ai/0.22.0/cu129" }, extra = ["otel", "runai"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "warcio" }, ] translation-all = [ { name = "aiohttp" }, @@ -5979,6 +5979,7 @@ requires-dist = [ { name = "easydict", marker = "extra == 'video-cpu'" }, { name = "einops", marker = "extra == 'video-cpu'" }, { name = "fasttext", marker = "extra == 'text-cpu'", specifier = "==0.9.3" }, + { name = "fastwarc", marker = "extra == 'text-cpu'", specifier = "<1" }, { name = "fsspec" }, { name = "ftfy", marker = "extra == 'text-cpu'", specifier = ">=6.3.1" }, { name = "google-cloud-translate", marker = "extra == 'translation-google'" }, @@ -6114,7 +6115,6 @@ requires-dist = [ { name = "vllm", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'math-cuda12'", specifier = ">=0.13" }, { name = "vllm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'math-cuda12'", specifier = ">=0.13", index = "https://wheels.vllm.ai/0.22.0/cu129" }, { name = "vllm", extras = ["flashinfer", "otel", "runai"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'vllm') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'vllm')", specifier = "==0.22.0+cu129", index = "https://wheels.vllm.ai/0.22.0/cu129" }, - { name = "warcio", marker = "extra == 'text-cpu'" }, { name = "whisperx", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin' and extra == 'audio-common'", specifier = ">=3.8.4" }, ] provides-extras = ["cuda12", "cv2", "vllm", "inference-server", "deduplication-cuda12", "audio-common", "audio-cpu", "audio-cuda12", "image-cpu", "image-cuda12", "translation-common", "translation-metrics", "translation-segmentation", "translation-aws", "translation-google", "translation-nmt", "translation-all", "text-cpu", "lance", "text-cuda12", "video-cpu", "video-cuda12", "video-media", "math-cpu", "math-cuda12", "interleaved-cpu", "interleaved-cuda12", "sdg-cpu", "sdg-cuda12", "all"] @@ -12016,18 +12016,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c6/c4/c7bed5e981679c74e9fbb22c03ff31c42e95f266199d03d8d325f4d0e6df/wandb-0.28.0-py3-none-win_amd64.whl", hash = "sha256:ac1f82292e2da4f98297b78c3a46726b3a6c5734ecb75fc39b8db2c8a4989159", size = 24525214, upload-time = "2026-06-23T00:38:44.549Z" }, ] -[[package]] -name = "warcio" -version = "1.7.5" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "six" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f2/2b/d825506924cb4508c90cd950dbda2a4dbfa9f5609e2ae76b53deaba656db/warcio-1.7.5.tar.gz", hash = "sha256:7247b57e68074cfd9433cb6dc226f8567d6777052abec2d3c78346cffa4d19b9", size = 61691, upload-time = "2024-12-10T21:31:48.939Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5c/f0/3f19085980f8a4485f4265cc9dba1099b2fa35ef7552390a8446e149c293/warcio-1.7.5-py2.py3-none-any.whl", hash = "sha256:ca96130bde7747e49da714097d144c6ff939458d4f93e1beb1e42455db4326d4", size = 40568, upload-time = "2024-12-10T21:31:46.291Z" }, -] - [[package]] name = "wasabi" version = "1.1.3"