Repository: huggingface/dataset-viewer

Issue:
Run the backfill script in parallel?
The backfilll script is very slow (one dataset at a time) and also it breaks due to OOM (https://cd.internal.huggingface.tech/applications/argo-cd/datasets-server-prod?view=pods&node=%2FPod%2Fdatasets-server%2Fprod-datasets-server-job-backfill-28412480-2n9ps%2F0)

Maybe we could run the loop iterations in parallel, and also allocate more memory to the pod


Discussion:
See https://github.com/huggingface/datasets-server/pull/2375#issuecomment-1924086210
With the increasing number of datasets, the backfill can take a lot too long (more than 24h!) so that we even overlap between daily cron jobs. We should improve this

<img width="305" alt="Capture d’écran 2024-03-29 à 11 04 22" src="https://github.com/huggingface/datasets-server/assets/1676121/267596a7-65bf-4f5a-9dff-15f005cffd6e">

Current state:


https://github.com/huggingface/datasets-server/assets/1676121/ee7affb3-0133-4f8e-b3a8-fa754c1f5fb1





<img width="321" alt="Capture d’écran 2024-04-04 à 12 03 21" src="https://github.com/huggingface/datasets-server/assets/1676121/155b1194-06c7-4f4b-a7e6-0780863ba3b7">


Resolving patch:
diff --git a/jobs/cache_maintenance/src/cache_maintenance/backfill.py b/jobs/cache_maintenance/src/cache_maintenance/backfill.py
--- a/jobs/cache_maintenance/src/cache_maintenance/backfill.py
+++ b/jobs/cache_maintenance/src/cache_maintenance/backfill.py
@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 The HuggingFace Authors.
 
+import concurrent.futures
 import logging
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional
 
 from libcommon.dtos import Priority
@@ -10,13 +11,16 @@
 from libcommon.simple_cache import get_all_datasets, get_datasets_with_retryable_errors
 from libcommon.storage_client import StorageClient
 
+MAX_BACKFILL_WORKERS = 8
+LOG_BATCH = 100
+
 
 @dataclass
 class BackfillStatistics:
     num_total_datasets: int = 0
     num_analyzed_datasets: int = 0
     num_error_datasets: int = 0
-    operations: OperationsStatistics = OperationsStatistics()
+    operations: OperationsStatistics = field(default_factory=OperationsStatistics)
 
     def add(self, other: "BackfillStatistics") -> None:
         self.num_total_datasets += other.num_total_datasets
@@ -69,39 +73,67 @@ def backfill_retryable_errors(
     )
 
 
+def try_backfill_dataset(
+    dataset: str,
+    hf_endpoint: str,
+    blocked_datasets: list[str],
+    hf_token: Optional[str] = None,
+    storage_clients: Optional[list[StorageClient]] = None,
+) -> BackfillStatistics:
+    try:
+        return BackfillStatistics(
+            num_analyzed_datasets=1,
+            operations=backfill_dataset(
+                dataset=dataset,
+                hf_endpoint=hf_endpoint,
+                blocked_datasets=blocked_datasets,
+                hf_token=hf_token,
+                priority=Priority.LOW,
+                hf_timeout_seconds=None,
+                storage_clients=storage_clients,
+            ),
+        )
+
+    except Exception as e:
+        logging.warning(f"failed to update_dataset {dataset}: {e}")
+        return BackfillStatistics(num_analyzed_datasets=1, num_error_datasets=1)
+
+
 def backfill_datasets(
     dataset_names: set[str],
     hf_endpoint: str,
     blocked_datasets: list[str],
     hf_token: Optional[str] = None,
     storage_clients: Optional[list[StorageClient]] = None,
-) -> None:
+) -> BackfillStatistics:
     logging.info(f"analyzing {len(dataset_names)} datasets in the database")
+
     statistics = BackfillStatistics(num_total_datasets=len(dataset_names))
-    log_batch = 100
-
-    for dataset in dataset_names:
-        try:
-            statistics.add(
-                BackfillStatistics(
-                    num_analyzed_datasets=1,
-                    operations=backfill_dataset(
-                        dataset=dataset,
-                        hf_endpoint=hf_endpoint,
-                        blocked_datasets=blocked_datasets,
-                        hf_token=hf_token,
-                        priority=Priority.LOW,
-                        hf_timeout_seconds=None,
-                        storage_clients=storage_clients,
-                    ),
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_BACKFILL_WORKERS) as executor:
+        futures_to_dataset = {
+            executor.submit(
+                try_backfill_dataset, dataset, hf_endpoint, blocked_datasets, hf_token, storage_clients
+            ): dataset
+            for dataset in dataset_names
+        }
+        # Start the load operations and gives stats on the progress
+        for future in concurrent.futures.as_completed(futures_to_dataset):
+            dataset = futures_to_dataset[future]
+            try:
+                dataset_statistics = future.result()
+            except Exception as e:
+                logging.warning(f"{dataset} generated an exception: {e}")
+                dataset_statistics = BackfillStatistics(
+                    num_total_datasets=1, num_analyzed_datasets=1, num_error_datasets=1
                 )
-            )
-        except Exception as e:
-            logging.warning(f"failed to update_dataset {dataset}: {e}")
-            statistics.add(BackfillStatistics(num_analyzed_datasets=1, num_error_datasets=1))
-        logging.debug(statistics.get_log())
-        if statistics.num_analyzed_datasets % log_batch == 0:
-            logging.info(statistics.get_log())
+            finally:
+                statistics.add(dataset_statistics)
+                logging.debug(statistics.get_log())
+                if statistics.num_analyzed_datasets % LOG_BATCH == 0:
+                    logging.info(statistics.get_log())
 
     logging.info(statistics.get_log())
     logging.info("backfill completed")
+
+    return statistics
diff --git a/libs/libcommon/src/libcommon/operations.py b/libs/libcommon/src/libcommon/operations.py
--- a/libs/libcommon/src/libcommon/operations.py
+++ b/libs/libcommon/src/libcommon/operations.py
@@ -2,7 +2,7 @@
 # Copyright 2022 The HuggingFace Authors.
 
 import logging
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional, Union
 
 from huggingface_hub.hf_api import DatasetInfo, HfApi
@@ -174,7 +174,7 @@ class OperationsStatistics:
     num_backfilled_datasets: int = 0
     num_deleted_datasets: int = 0
     num_untouched_datasets: int = 0
-    tasks: TasksStatistics = TasksStatistics()
+    tasks: TasksStatistics = field(default_factory=TasksStatistics)
 
     def add(self, other: "OperationsStatistics") -> None:
         self.num_backfilled_datasets += other.num_backfilled_datasets


Test patch:
diff --git a/jobs/cache_maintenance/tests/test_backfill.py b/jobs/cache_maintenance/tests/test_backfill.py
new file mode 100644
--- /dev/null
+++ b/jobs/cache_maintenance/tests/test_backfill.py
@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 The HuggingFace Authors.
+
+from unittest.mock import patch
+
+import pytest
+from libcommon.operations import OperationsStatistics
+
+from cache_maintenance.backfill import backfill_datasets, try_backfill_dataset
+
+from .constants import CI_APP_TOKEN, CI_HUB_ENDPOINT
+
+
+@pytest.mark.parametrize(
+    "num_backfilled_datasets,num_untouched_datasets,num_deleted_datasets",
+    [
+        (1, 0, 0),
+        (0, 1, 0),
+        (0, 0, 1),
+    ],
+)
+def test_try_backfill_dataset(
+    num_backfilled_datasets: int, num_untouched_datasets: int, num_deleted_datasets: int
+) -> None:
+    with patch(
+        "cache_maintenance.backfill.backfill_dataset",
+        return_value=OperationsStatistics(
+            num_backfilled_datasets=num_backfilled_datasets,
+            num_untouched_datasets=num_untouched_datasets,
+            num_deleted_datasets=num_deleted_datasets,
+        ),
+    ):
+        statistics = try_backfill_dataset(
+            dataset="dataset",
+            hf_endpoint=CI_HUB_ENDPOINT,
+            blocked_datasets=[],
+            hf_token=CI_APP_TOKEN,
+            storage_clients=[],
+        )
+    assert statistics.num_total_datasets == 0
+    assert statistics.num_analyzed_datasets == 1
+    assert statistics.num_error_datasets == 0
+    assert statistics.operations.num_backfilled_datasets == num_backfilled_datasets
+    assert statistics.operations.num_untouched_datasets == num_untouched_datasets
+    assert statistics.operations.num_deleted_datasets == num_deleted_datasets
+
+
+def test_try_backfill_dataset_exception() -> None:
+    with patch("cache_maintenance.backfill.backfill_dataset", side_effect=Exception("error")):
+        statistics = try_backfill_dataset(
+            dataset="dataset",
+            hf_endpoint=CI_HUB_ENDPOINT,
+            blocked_datasets=[],
+            hf_token=CI_APP_TOKEN,
+            storage_clients=[],
+        )
+    assert statistics.num_analyzed_datasets == 1
+    assert statistics.num_error_datasets == 1
+
+
+def test_backfill_datasets() -> None:
+    num_datasets = 200
+    with patch(
+        "cache_maintenance.backfill.backfill_dataset",
+        return_value=OperationsStatistics(
+            num_backfilled_datasets=1,
+            num_untouched_datasets=0,
+            num_deleted_datasets=0,
+        ),
+    ):
+        statistics = backfill_datasets(
+            dataset_names={f"a{i}" for i in range(0, num_datasets)},
+            hf_endpoint=CI_HUB_ENDPOINT,
+            blocked_datasets=[],
+            hf_token=CI_APP_TOKEN,
+            storage_clients=[],
+        )
+    assert statistics.num_total_datasets == num_datasets
+    assert statistics.num_analyzed_datasets == num_datasets
+    assert statistics.num_error_datasets == 0
+    assert statistics.operations.num_backfilled_datasets == num_datasets
+    assert statistics.operations.num_untouched_datasets == 0
+    assert statistics.operations.num_deleted_datasets == 0