Repository: huggingface/dataset-viewer

Issue:
"2 or more items returned, instead of 1" returned by /splits-next
https://datasets-server.huggingface.co/splits-next?dataset=GEM/turku_paraphrase_corpus

```
{"error":"2 or more items returned, instead of 1"}
```

It comes from the MultipleObjectsReturned error in mongoengine.


Discussion:
See https://observablehq.com/@huggingface/internal-datasets-server-cache-contents to get the list of erroneous datasets
"Only" 15 datasets affected until now, see https://grafana.huggingface.tech/d/SaHl2KX7z/datasets-server-admin-cache-and-queue?orgId=1&from=now-30m&to=now&refresh=30s ("Errors in /splits-next") 

<img width="717" alt="Capture d’écran 2022-08-24 à 15 49 38" src="https://user-images.githubusercontent.com/1676121/186510018-9dafe17b-dad3-4541-b79f-8a6fcdbc9cc4.png">


[GEM/turku_paraphrase_corpus](https://datasets-server.huggingface.co/splits-next?dataset=GEM/turku_paraphrase_corpus)
[HarrisDePerceptron/sv_corpora_parliament_processed](https://datasets-server.huggingface.co/splits-next?dataset=HarrisDePerceptron/sv_corpora_parliament_processed)
[KETI-AIR/klue](https://datasets-server.huggingface.co/splits-next?dataset=KETI-AIR/klue)
[allegro/summarization-polish-summaries-corpus](https://datasets-server.huggingface.co/splits-next?dataset=allegro/summarization-polish-summaries-corpus)
[castorini/msmarco_v1_doc_doc2query-t5_expansions](https://datasets-server.huggingface.co/splits-next?dataset=castorini/msmarco_v1_doc_doc2query-t5_expansions)
[chitra/contradictionNLI](https://datasets-server.huggingface.co/splits-next?dataset=chitra/contradictionNLI)
[huggan/selfie2anime](https://datasets-server.huggingface.co/splits-next?dataset=huggan/selfie2anime)
[fever/feverous](https://datasets-server.huggingface.co/splits-next?dataset=fever/feverous)
[Skelebor/book_titles_and_descriptions_en_clean](https://datasets-server.huggingface.co/splits-next?dataset=Skelebor/book_titles_and_descriptions_en_clean)
[ThierryZhou/test](https://datasets-server.huggingface.co/splits-next?dataset=ThierryZhou/test)
[poleval2019_cyberbullying](https://datasets-server.huggingface.co/splits-next?dataset=poleval2019_cyberbullying)
[qangaroo](https://datasets-server.huggingface.co/splits-next?dataset=qangaroo)
[scicite](https://datasets-server.huggingface.co/splits-next?dataset=scicite)
[sede](https://datasets-server.huggingface.co/splits-next?dataset=sede)
[sem_eval_2020_task_11](https://datasets-server.huggingface.co/splits-next?dataset=sem_eval_2020_task_11)
[simple_questions_v2](https://datasets-server.huggingface.co/splits-next?dataset=simple_questions_v2)
[tapaco](https://datasets-server.huggingface.co/splits-next?dataset=tapaco)
[wnut_17](https://datasets-server.huggingface.co/splits-next?dataset=wnut_17)
[ASCCCCCCCC/amazon_zh_simple](https://datasets-server.huggingface.co/splits-next?dataset=ASCCCCCCCC/amazon_zh_simple)

Raised by

https://github.com/huggingface/datasets-server/blob/63ecc62e05d566466deaafc6b97eb667ad9ffa25/services/worker/src/worker/refresh.py#L42

and more specifically by:

https://github.com/huggingface/datasets-server/blob/63ecc62e05d566466deaafc6b97eb667ad9ffa25/libs/libqueue/src/libqueue/queue.py#L285-L296

and

https://github.com/huggingface/datasets-server/blob/63ecc62e05d566466deaafc6b97eb667ad9ffa25/libs/libqueue/src/libqueue/queue.py#L248-L254

As stated:

> raises MultipleObjectsReturned if more than one entry -> should never occur, we let it raises

but it occurs from time to time.

It means that two `/first-rows` jobs exist for the same tuple `(dataset, config, split)` while the logic of the application is only to have one such job among the `WAITING`and `STARTED` jobs.
The solution could be to handle the `MultipleObjectsReturned` exception, and in that case, cancel all of them, ie. move them all to the `CANCELLED` status and create a fresh new one.

Resolving patch:
diff --git a/libs/libqueue/src/libqueue/queue.py b/libs/libqueue/src/libqueue/queue.py
--- a/libs/libqueue/src/libqueue/queue.py
+++ b/libs/libqueue/src/libqueue/queue.py
@@ -5,6 +5,7 @@
 from typing import Generic, List, Optional, Tuple, Type, TypedDict, TypeVar
 
 from mongoengine import Document, DoesNotExist, connect
+from mongoengine.errors import MultipleObjectsReturned
 from mongoengine.fields import DateTimeField, EnumField, IntField, StringField
 from mongoengine.queryset.queryset import QuerySet
 
@@ -245,13 +246,20 @@ def get_datetime() -> datetime:
     return datetime.now(timezone.utc)
 
 
-def add_job(existing_jobs: QuerySet[AnyJob], new_job: AnyJob):
+def add_job(existing_jobs: QuerySet[AnyJob], new_job: AnyJob) -> AnyJob:
+    pending_jobs = existing_jobs.filter(status__in=[Status.WAITING, Status.STARTED])
     try:
-        # Check if a non-finished job already exists
-        existing_jobs.filter(status__in=[Status.WAITING, Status.STARTED]).get()
+        # If one non-finished job exists, return it
+        return pending_jobs.get()
     except DoesNotExist:
-        new_job.save()
-    # raises MultipleObjectsReturned if more than one entry -> should never occur, we let it raise
+        # None exist, create one
+        return new_job.save()
+    except MultipleObjectsReturned:
+        # should not happen, but it's not enforced in the database
+        # (we could have one in WAITING status and another one in STARTED status)
+        # it it happens, we "cancel" all of them, and re-run the same function
+        pending_jobs.update(finished_at=get_datetime(), status=Status.CANCELLED)
+        return add_job(existing_jobs, new_job)
 
 
 def add_dataset_job(dataset_name: str, retries: Optional[int] = 0) -> None:


Test patch:
diff --git a/libs/libqueue/tests/test_queue.py b/libs/libqueue/tests/test_queue.py
--- a/libs/libqueue/tests/test_queue.py
+++ b/libs/libqueue/tests/test_queue.py
@@ -2,12 +2,15 @@
 
 from libqueue.queue import (
     EmptyQueue,
+    FirstRowsJob,
+    Status,
     add_first_rows_job,
     add_splits_job,
     clean_database,
     connect_to_queue,
     finish_first_rows_job,
     finish_splits_job,
+    get_datetime,
     get_first_rows_job,
     get_first_rows_jobs_count_by_status,
     get_splits_job,
@@ -68,6 +71,41 @@ def test_add_job() -> None:
     finish_splits_job(job_id, success=True)
 
 
+def test_add_job_with_broken_collection() -> None:
+    dataset_name = "dataset_broken"
+    config_name = "config_broken"
+    split_name = "split_broken"
+    # ensure the jobs are cancelled with more than one exist in a "pending" status
+    # we "manually" create two jobs in a "pending" status for the same split
+    # (we normally cannot do that with the exposed methods)
+    job_1 = FirstRowsJob(
+        dataset_name=dataset_name,
+        config_name=config_name,
+        split_name=split_name,
+        created_at=get_datetime(),
+        status=Status.WAITING,
+    ).save()
+    job_2 = FirstRowsJob(
+        dataset_name=dataset_name,
+        config_name=config_name,
+        split_name=split_name,
+        created_at=get_datetime(),
+        started_at=get_datetime(),
+        status=Status.STARTED,
+    ).save()
+    # then we add a job: it should create a new job in the "WAITING" status
+    # and the two other jobs should be cancelled
+    add_first_rows_job(dataset_name=dataset_name, config_name=config_name, split_name=split_name)
+    assert (
+        FirstRowsJob.objects(
+            dataset_name=dataset_name, config_name=config_name, split_name=split_name, status__in=[Status.WAITING]
+        ).count()
+        == 1
+    )
+    assert FirstRowsJob.objects(pk=job_1.pk).get().status == Status.CANCELLED
+    assert FirstRowsJob.objects(pk=job_2.pk).get().status == Status.CANCELLED
+
+
 def test_max_jobs_per_dataset() -> None:
     add_first_rows_job("dataset", "config", "split1")
     assert is_first_rows_response_in_process("dataset", "config", "split1") is True
diff --git a/services/worker/tests/responses/test_first_rows.py b/services/worker/tests/responses/test_first_rows.py
--- a/services/worker/tests/responses/test_first_rows.py
+++ b/services/worker/tests/responses/test_first_rows.py
@@ -7,7 +7,6 @@
 from ..utils import ASSETS_BASE_URL, HF_ENDPOINT, HF_TOKEN, get_default_config_split
 
 
-@pytest.mark.wip
 @pytest.mark.parametrize(
     "name,use_token,error_code,cause",
     [