Retrieval task uses too much disk space

[nqbao]

I was running evaluation for MS-MARCO-VN dataset (which is huge) on a container with 30gb disk space. After 5hr the task failed due to out of disk

2025-11-06 04:16:52,087 - mteb.deprecated_evaluator - ERROR - Error while evaluating MSMARCO-VN: [Errno 122] Disk quota exceeded
Traceback (most recent call last):
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 3689, in _map_single
    writer.write(example)
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 623, in write
    self.write_examples_on_file()
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 581, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 701, in write_batch
    self.write_table(pa_table, writer_batch_size)
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 719, in write_table
    self.pa_writer.write_table(pa_table, writer_batch_size)
  File "pyarrow/ipc.pxi", line 616, in pyarrow.lib._CRecordBatchWriter.write_table
  File "pyarrow/error.pxi", line 89, in pyarrow.lib.check_status
  File "/usr/local/lib/python3.12/dist-packages/fsspec/implementations/local.py", line 422, in write
    return self.f.write(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: [Errno 122] Disk quota exceeded

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 3724, in _map_single
    writer.finalize()
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 728, in finalize
    self.write_examples_on_file()
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 581, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 701, in write_batch
    self.write_table(pa_table, writer_batch_size)
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 719, in write_table
    self.pa_writer.write_table(pa_table, writer_batch_size)
  File "pyarrow/ipc.pxi", line 616, in pyarrow.lib._CRecordBatchWriter.write_table
  File "pyarrow/error.pxi", line 89, in pyarrow.lib.check_status
  File "/usr/local/lib/python3.12/dist-packages/fsspec/implementations/local.py", line 422, in write
    return self.f.write(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: [Errno 122] Disk quota exceeded

During handling of the above exception, another exception occurred:

OSError: [Errno 122] Disk quota exceeded

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/workspace/evals/mteb_eval.py", line 68, in <module>
    score = eval_dataset(task_name, model_name, batch_size=32)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/evals/mteb_eval.py", line 61, in eval_dataset
    return evaluation.run(model, encode_kwargs=encode_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mteb/deprecated_evaluator.py", line 544, in run
    raise e
  File "/usr/local/lib/python3.12/dist-packages/mteb/deprecated_evaluator.py", line 500, in run
    results, tick, tock = self._run_eval(
                          ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mteb/deprecated_evaluator.py", line 186, in _run_eval
    results = task.evaluate(
              ^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mteb/abstasks/retrieval.py", line 310, in evaluate
    return super().evaluate(
           ^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mteb/abstasks/abstask.py", line 183, in evaluate
    scores[hf_subset] = self._evaluate_subset(
                        ^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mteb/abstasks/retrieval.py", line 372, in _evaluate_subset
    results = retriever(
              ^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mteb/_evaluators/retrieval_evaluator.py", line 62, in __call__
    return search_model.search(
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mteb/models/search_wrappers.py", line 120, in search
    result_heaps = self._full_corpus_search(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mteb/models/search_wrappers.py", line 164, in _full_corpus_search
    create_dataloader(
  File "/usr/local/lib/python3.12/dist-packages/mteb/_create_dataloaders.py", line 414, in create_dataloader
    return _create_document_dataloader(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mteb/_create_dataloaders.py", line 369, in _create_document_dataloader
    return _create_dataloader_for_retrieval_corpus(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mteb/_create_dataloaders.py", line 74, in _create_dataloader_for_retrieval_corpus
    new_ds = dataset.map(_corpus_to_dict, desc="Converting corpus dict")
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 562, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 3341, in map
    for rank, done, content in Dataset._map_single(**unprocessed_kwargs):
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 3658, in _map_single
    with contextlib.ExitStack() as stack:
  File "/usr/lib/python3.12/contextlib.py", line 610, in __exit__
    raise exc_details[1]
  File "/usr/lib/python3.12/contextlib.py", line 595, in __exit__
    if cb(*exc_details):
       ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 486, in __exit__
    self.close()
  File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_writer.py", line 496, in close
    self.stream.close()  # This also closes self.pa_writer if it is opened
    ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/fsspec/implementations/local.py", line 440, in close
    return self.f.close()

I know i can always provision more disk space, but is this expected to use this much? Most of this comes from datasets caching. Should i disable caching or is there a way to clean up cache for each batch?

[KennethEnevoldsen]

Hmm yea that seems like too much. @nqbao can I ask you to see how much is used by ds = datasets.load_dataset("GreenNode/msmarco-vn")? This kinda gives us the lower bound of what are aiming for.

The dataset itself is 3.58gb on huggingface

[Samoed]

Doing

ds_corpus = datasets.load_dataset("GreenNode/msmarco-vn", "corpus")

Took 4 GB loading corpus (before generating corpus split, just to load 14 parquet files, but kaggle might can report with delay). After generating corpus split 13 GB.

Doing

ds_corpus = ds_corpus.cast_column("_id", datasets.Value("string")).rename_column("_id", "id")

From 13 -> 20 GB.

Overall, task.load_data() took 20 GB of disk space, with dependencies it was 28 GB.

Each "batch" processing (50_000 sentences) adds 100 MB.

I can reupload with a renamed column name, but it still will be ~13-15 GB.

[nqbao]

Here is some quick numbers on my end

The dataset itself is ~3.5gb like you mention
All the cache-xxx data are ~19gb

Look like only corpus has this issue

3.1M    ./queries/0.0.0/85d1ad4cc9070b8d019d65f5af1631a2ab91e294
5.0M    ./queries/0.0.0
6.9M    ./queries
23G     ./corpus/0.0.0/85d1ad4cc9070b8d019d65f5af1631a2ab91e294
23G     ./corpus/0.0.0
23G     ./corpus
1.4M    ./default/0.0.0/85d1ad4cc9070b8d019d65f5af1631a2ab91e294
2.3M    ./default/0.0.0
3.3M    ./default
23G     .
23G     total

Below are some examples, each batch seems to add 60mb disk space. There are 177 batches total so extra 10gb to go through the batch.

root@07e1bb073721:/workspace/hf_home/datasets/GreenNode___msmarco-vn/corpus/0.0.0/85d1ad4cc9070b8d019d65f5af1631a2ab91e294# ll -ah | grep cache
-rw-rw-rw- 1 root root  58M Nov  6 06:55 cache-008e983c7c5aab53.arrow
-rw-rw-rw- 1 root root  57M Nov  6 13:22 cache-02ae9aba407d2920.arrow
-rw-rw-rw- 1 root root  58M Nov  6 05:52 cache-0572bf244de496c4.arrow
-rw-rw-rw- 1 root root  56M Nov  6 13:02 cache-05b07db99dbe8880.arrow
-rw-rw-rw- 1 root root  57M Nov  6 14:14 cache-05b0b1792eb703de.arrow
-rw-rw-rw- 1 root root  58M Nov  6 06:59 cache-06a4ca6726549e0b.arrow
-rw-rw-rw- 1 root root 6.5G Nov  6 04:44 cache-0ab267415f836065.arrow
-rw-rw-rw- 1 root root  58M Nov  6 06:47 cache-0ebf1b2a258a3275.arrow
-rw-rw-rw- 1 root root  58M Nov  6 08:06 cache-0ee6658c15b19092.arrow
-rw-rw-rw- 1 root root  58M Nov  6 06:31 cache-145f4eee7821147b.arrow
-rw-rw-rw- 1 root root  63M Nov  6 10:49 cache-1795dcfdff7a590d.arrow
-rw-rw-rw- 1 root root  59M Nov  6 10:15 cache-182d4b6c939ad75a.arrow
-rw-rw-rw- 1 root root  55M Nov  6 11:58 cache-1941c81f10175a03.arrow
-rw-rw-rw- 1 root root  59M Nov  6 10:19 cache-1b0db89b244936ac.arrow
-rw-rw-rw- 1 root root  58M Nov  6 14:46 cache-1c108659da0ab8b1.arrow
-rw-rw-rw- 1 root root  58M Nov  6 09:34 cache-1cea1a64341c9f97.arrow
-rw-rw-rw- 1 root root  58M Nov  6 14:50 cache-1d49d1ca6de070c6.arrow

[Samoed]

I think we can change map to batched map or just direct looping over rows

[nqbao]

Can we disable cache? https://huggingface.co/docs/datasets/en/cache#enable-or-disable-caching

[Samoed]

You can disable cache during your run. I don't think we want to disable globally

[nqbao]

Just for my knowledge, this is expected right? This is by design of datasets as it will materialize transform into disk.

[Samoed]

this is expected right

I'm not sure what behavior do you expect.

This is by design of datasets as it will materialize transform into disk.

Yes

[nqbao]

I'm not sure what behavior do you expect.

I meant we dont use this much of disk space during evaluation.