We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
The workers fail with `mongoengine.errors.FieldDoesNotExist: The fields "{'force'}" do not exist on the document "Job"`
```
INFO: 2023-05-13 10:50:18,007 - root - Worker loop started
INFO: 2023-05-13 10:50:18,023 - root - Starting heartbeat.
ERROR: 2023-05-13 10:50:18,115 - asyncio - Task exception was never retrieved
future: <Task finished name='Task-2' coro=<every() done, defined at /src/services/worker/src/worker/executor.py:26> exception=FieldDoesNotExist('The fields "{\'force\'}" do not exist on the document "Job"')>
Traceback (most recent call last):
  File "/src/services/worker/src/worker/executor.py", line 30, in every
    out = func(*args, **kwargs)
  File "/src/services/worker/src/worker/executor.py", line 117, in kill_zombies
    zombies = queue.get_zombies(max_seconds_without_heartbeat=self.max_seconds_without_heartbeat_for_zombies)
  File "/src/libs/libcommon/src/libcommon/queue.py", line 622, in get_zombies
    zombies = [
  File "/src/libs/libcommon/src/libcommon/queue.py", line 622, in <listcomp>
    zombies = [
  File "/src/services/worker/.venv/lib/python3.9/site-packages/mongoengine/queryset/queryset.py", line 110, in _iter_results
    self._populate_cache()
  File "/src/services/worker/.venv/lib/python3.9/site-packages/mongoengine/queryset/queryset.py", line 129, in _populate_cache
    self._result_cache.append(next(self))
  File "/src/services/worker/.venv/lib/python3.9/site-packages/mongoengine/queryset/base.py", line 1599, in __next__
    doc = self._document._from_son(
  File "/src/services/worker/.venv/lib/python3.9/site-packages/mongoengine/base/document.py", line 836, in _from_son
    obj = cls(__auto_convert=False, _created=created, **data)
  File "/src/services/worker/.venv/lib/python3.9/site-packages/mongoengine/base/document.py", line 99, in __init__
    raise FieldDoesNotExist(msg)
mongoengine.errors.FieldDoesNotExist: The fields "{'force'}" do not exist on the document "Job"
Traceback (most recent call last):
  File "/src/services/worker/src/worker/main.py", line 61, in <module>
    worker_executor.start()
  File "/src/services/worker/src/worker/executor.py", line 89, in start
    loop.run_until_complete(
  File "/usr/local/lib/python3.9/asyncio/base_events.py", line 645, in run_until_complete
    raise RuntimeError('Event loop stopped before Future completed.')
RuntimeError: Event loop stopped before Future completed.
```

<img width="622" alt="Capture d’écran 2023-05-13 à 12 52 06" src="https://github.com/huggingface/datasets-server/assets/1676121/3e4d7e1a-dd02-476d-b671-57ce03d28d7b">


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `services/worker/src/worker/job_manager.py`
Content:
```
1 # SPDX-License-Identifier: Apache-2.0
2 # Copyright 2022 The HuggingFace Authors.
3 
4 import logging
5 from http import HTTPStatus
6 from typing import Literal, Optional
7 
8 from libcommon.config import CommonConfig
9 from libcommon.dataset import DatasetNotFoundError, get_dataset_git_revision
10 from libcommon.exceptions import CustomError
11 from libcommon.processing_graph import ProcessingGraph, ProcessingStep
12 from libcommon.simple_cache import (
13     DoesNotExist,
14     get_response_without_content_params,
15     upsert_response_params,
16 )
17 from libcommon.state import DatasetState
18 from libcommon.utils import JobInfo, JobParams, Priority, Status, orjson_dumps
19 
20 from worker.common_exceptions import (
21     JobManagerCrashedError,
22     JobManagerExceededMaximumDurationError,
23     NoGitRevisionError,
24     ResponseAlreadyComputedError,
25     TooBigContentError,
26     UnexpectedError,
27 )
28 from worker.config import AppConfig, WorkerConfig
29 from worker.job_runner import JobRunner
30 
31 # List of error codes that should trigger a retry.
32 ERROR_CODES_TO_RETRY: list[str] = ["ClientConnectionError"]
33 
34 
35 class JobManager:
36     """
37     A job manager is a class that handles a job runner compute, for a specific processing step.
38 
39     Args:
40         job_info (:obj:`JobInfo`):
41             The job to process. It contains the job_id, the job type, the dataset, the config, the split
42             and the priority level.
43         common_config (:obj:`CommonConfig`):
44             The common config.
45         processing_step (:obj:`ProcessingStep`):
46             The processing step to process.
47     """
48 
49     job_id: str
50     job_params: JobParams
51     priority: Priority
52     worker_config: WorkerConfig
53     common_config: CommonConfig
54     processing_step: ProcessingStep
55     processing_graph: ProcessingGraph
56     _dataset_git_revision: Optional[str] = None
57     job_runner: JobRunner
58 
59     def __init__(
60         self,
61         job_info: JobInfo,
62         app_config: AppConfig,
63         job_runner: JobRunner,
64         processing_graph: ProcessingGraph,
65     ) -> None:
66         self.job_info = job_info
67         self.job_type = job_info["type"]
68         self.job_id = job_info["job_id"]
69         self.priority = job_info["priority"]
70         self.job_params = job_info["params"]
71         self.common_config = app_config.common
72         self.worker_config = app_config.worker
73         self.job_runner = job_runner
74         self.processing_graph = processing_graph
75         self.processing_step = self.job_runner.processing_step
76         self.setup()
77 
78     def setup(self) -> None:
79         job_type = self.job_runner.get_job_type()
80         if self.processing_step.job_type != job_type:
81             raise ValueError(
82                 f"The processing step's job type is {self.processing_step.job_type}, but the job manager only"
83                 f" processes {job_type}"
84             )
85         if self.job_type != job_type:
86             raise ValueError(
87                 f"The submitted job type is {self.job_type}, but the job manager only processes {job_type}"
88             )
89 
90     def __str__(self) -> str:
91         return f"JobManager(job_id={self.job_id} dataset={self.job_params['dataset']} job_info={self.job_info}"
92 
93     def log(self, level: int, msg: str) -> None:
94         logging.log(level=level, msg=f"[{self.processing_step.job_type}] {msg}")
95 
96     def debug(self, msg: str) -> None:
97         self.log(level=logging.DEBUG, msg=msg)
98 
99     def info(self, msg: str) -> None:
100         self.log(level=logging.INFO, msg=msg)
101 
102     def warning(self, msg: str) -> None:
103         self.log(level=logging.WARNING, msg=msg)
104 
105     def exception(self, msg: str) -> None:
106         self.log(level=logging.ERROR, msg=msg)
107 
108     def critical(self, msg: str) -> None:
109         self.log(level=logging.CRITICAL, msg=msg)
110 
111     def run(self) -> Literal[Status.SUCCESS, Status.ERROR]:
112         try:
113             self.info(f"compute {self}")
114             result: Literal[Status.SUCCESS, Status.ERROR] = Status.SUCCESS if self.process() else Status.ERROR
115         except Exception:
116             self.exception(f"error while computing {self}")
117             result = Status.ERROR
118         self.backfill()
119         return result
120 
121     def get_dataset_git_revision(self) -> Optional[str]:
122         """Get the git revision of the dataset repository."""
123         if self._dataset_git_revision is None:
124             self._dataset_git_revision = get_dataset_git_revision(
125                 dataset=self.job_params["dataset"],
126                 hf_endpoint=self.common_config.hf_endpoint,
127                 hf_token=self.common_config.hf_token,
128             )
129         return self._dataset_git_revision
130 
131     def raise_if_parallel_response_exists(self, parallel_cache_kind: str, parallel_job_version: int) -> None:
132         try:
133             existing_response = get_response_without_content_params(
134                 kind=parallel_cache_kind,
135                 job_params=self.job_params,
136             )
137 
138             dataset_git_revision = self.get_dataset_git_revision()
139             if (
140                 existing_response["http_status"] == HTTPStatus.OK
141                 and existing_response["job_runner_version"] == parallel_job_version
142                 and existing_response["progress"] == 1.0  # completed response
143                 and dataset_git_revision is not None
144                 and existing_response["dataset_git_revision"] == dataset_git_revision
145             ):
146                 raise ResponseAlreadyComputedError(
147                     f"Response has already been computed and stored in cache kind: {parallel_cache_kind}. Compute will"
148                     " be skipped."
149                 )
150         except DoesNotExist:
151             logging.debug(f"no cache found for {parallel_cache_kind}.")
152 
153     def process(
154         self,
155     ) -> bool:
156         dataset_git_revision = None
157         try:
158             dataset_git_revision = self.get_dataset_git_revision()
159             if dataset_git_revision is None:
160                 self.debug(f"the dataset={self.job_params['dataset']} has no git revision, don't update the cache")
161                 raise NoGitRevisionError(f"Could not get git revision for dataset {self.job_params['dataset']}")
162             try:
163                 self.job_runner.pre_compute()
164                 parallel_job_runner = self.job_runner.get_parallel_job_runner()
165                 if parallel_job_runner:
166                     self.raise_if_parallel_response_exists(
167                         parallel_cache_kind=parallel_job_runner["job_type"],
168                         parallel_job_version=parallel_job_runner["job_runner_version"],
169                     )
170 
171                 job_result = self.job_runner.compute()
172                 content = job_result.content
173 
174                 # Validate content size
175                 if len(orjson_dumps(content)) > self.worker_config.content_max_bytes:
176                     raise TooBigContentError(
177                         "The computed response content exceeds the supported size in bytes"
178                         f" ({self.worker_config.content_max_bytes})."
179                     )
180             finally:
181                 # ensure the post_compute hook is called even if the compute raises an exception
182                 self.job_runner.post_compute()
183             upsert_response_params(
184                 kind=self.processing_step.cache_kind,
185                 job_params=self.job_params,
186                 content=content,
187                 http_status=HTTPStatus.OK,
188                 job_runner_version=self.job_runner.get_job_runner_version(),
189                 dataset_git_revision=dataset_git_revision,
190                 progress=job_result.progress,
191             )
192             self.debug(f"dataset={self.job_params['dataset']} job_info={self.job_info} is valid, cache updated")
193             return True
194         except DatasetNotFoundError:
195             # To avoid filling the cache, we don't save this error. Otherwise, DoS is possible.
196             self.debug(f"the dataset={self.job_params['dataset']} could not be found, don't update the cache")
197             return False
198         except Exception as err:
199             e = err if isinstance(err, CustomError) else UnexpectedError(str(err), err)
200             upsert_response_params(
201                 kind=self.processing_step.cache_kind,
202                 job_params=self.job_params,
203                 content=dict(e.as_response()),
204                 http_status=e.status_code,
205                 error_code=e.code,
206                 details=dict(e.as_response_with_cause()),
207                 job_runner_version=self.job_runner.get_job_runner_version(),
208                 dataset_git_revision=dataset_git_revision,
209             )
210             self.debug(f"response for job_info={self.job_info} had an error, cache updated")
211             return False
212 
213     def backfill(self) -> None:
214         """Evaluate the state of the dataset and backfill the cache if necessary."""
215         DatasetState(
216             dataset=self.job_params["dataset"],
217             processing_graph=self.processing_graph,
218             revision=self.get_dataset_git_revision(),
219             error_codes_to_retry=ERROR_CODES_TO_RETRY,
220             priority=self.priority,
221         ).backfill()
222 
223     def set_crashed(self, message: str, cause: Optional[BaseException] = None) -> None:
224         error = JobManagerCrashedError(message=message, cause=cause)
225         upsert_response_params(
226             kind=self.processing_step.cache_kind,
227             job_params=self.job_params,
228             content=dict(error.as_response()),
229             http_status=error.status_code,
230             error_code=error.code,
231             details=dict(error.as_response_with_cause()),
232             job_runner_version=self.job_runner.get_job_runner_version(),
233             dataset_git_revision=self.get_dataset_git_revision(),
234         )
235         logging.debug(
236             f"response for dataset={self.job_params['dataset']} job_info={self.job_info} had an error (crashed), cache"
237             " updated"
238         )
239 
240     def set_exceeded_maximum_duration(self, message: str, cause: Optional[BaseException] = None) -> None:
241         error = JobManagerExceededMaximumDurationError(message=message, cause=cause)
242         upsert_response_params(
243             kind=self.processing_step.cache_kind,
244             job_params=self.job_params,
245             content=dict(error.as_response()),
246             http_status=error.status_code,
247             error_code=error.code,
248             details=dict(error.as_response_with_cause()),
249             job_runner_version=self.job_runner.get_job_runner_version(),
250             dataset_git_revision=self.get_dataset_git_revision(),
251         )
252         logging.debug(
253             f"response for dataset={self.job_params['dataset']} job_info={self.job_info} had an error (exceeded"
254             " maximum duration), cache updated"
255         )
256
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/services/worker/src/worker/job_manager.py b/services/worker/src/worker/job_manager.py
--- a/services/worker/src/worker/job_manager.py
+++ b/services/worker/src/worker/job_manager.py
@@ -115,17 +115,25 @@
         except Exception:
             self.exception(f"error while computing {self}")
             result = Status.ERROR
-        self.backfill()
+        revision = self.get_dataset_git_revision(allow_raise=False)
+        if revision is not None:
+            self.backfill(revision=revision)
         return result
 
-    def get_dataset_git_revision(self) -> Optional[str]:
+    def get_dataset_git_revision(self, allow_raise: bool = True) -> Optional[str]:
         """Get the git revision of the dataset repository."""
         if self._dataset_git_revision is None:
-            self._dataset_git_revision = get_dataset_git_revision(
-                dataset=self.job_params["dataset"],
-                hf_endpoint=self.common_config.hf_endpoint,
-                hf_token=self.common_config.hf_token,
-            )
+            try:
+                self._dataset_git_revision = get_dataset_git_revision(
+                    dataset=self.job_params["dataset"],
+                    hf_endpoint=self.common_config.hf_endpoint,
+                    hf_token=self.common_config.hf_token,
+                )
+            except Exception as e:
+                if allow_raise:
+                    raise e
+                else:
+                    return None
         return self._dataset_git_revision
 
     def raise_if_parallel_response_exists(self, parallel_cache_kind: str, parallel_job_version: int) -> None:
@@ -210,12 +218,12 @@
             self.debug(f"response for job_info={self.job_info} had an error, cache updated")
             return False
 
-    def backfill(self) -> None:
+    def backfill(self, revision: str) -> None:
         """Evaluate the state of the dataset and backfill the cache if necessary."""
         DatasetState(
             dataset=self.job_params["dataset"],
             processing_graph=self.processing_graph,
-            revision=self.get_dataset_git_revision(),
+            revision=revision,
             error_codes_to_retry=ERROR_CODES_TO_RETRY,
             priority=self.priority,
         ).backfill()
@@ -230,7 +238,7 @@
             error_code=error.code,
             details=dict(error.as_response_with_cause()),
             job_runner_version=self.job_runner.get_job_runner_version(),
-            dataset_git_revision=self.get_dataset_git_revision(),
+            dataset_git_revision=self.get_dataset_git_revision(allow_raise=False),
         )
         logging.debug(
             f"response for dataset={self.job_params['dataset']} job_info={self.job_info} had an error (crashed), cache"
@@ -247,7 +255,7 @@
             error_code=error.code,
             details=dict(error.as_response_with_cause()),
             job_runner_version=self.job_runner.get_job_runner_version(),
-            dataset_git_revision=self.get_dataset_git_revision(),
+            dataset_git_revision=self.get_dataset_git_revision(allow_raise=False),
         )
         logging.debug(
             f"response for dataset={self.job_params['dataset']} job_info={self.job_info} had an error (exceeded"

{"golden_diff": "diff --git a/services/worker/src/worker/job_manager.py b/services/worker/src/worker/job_manager.py\n--- a/services/worker/src/worker/job_manager.py\n+++ b/services/worker/src/worker/job_manager.py\n@@ -115,17 +115,25 @@\n         except Exception:\n             self.exception(f\"error while computing {self}\")\n             result = Status.ERROR\n-        self.backfill()\n+        revision = self.get_dataset_git_revision(allow_raise=False)\n+        if revision is not None:\n+            self.backfill(revision=revision)\n         return result\n \n-    def get_dataset_git_revision(self) -> Optional[str]:\n+    def get_dataset_git_revision(self, allow_raise: bool = True) -> Optional[str]:\n         \"\"\"Get the git revision of the dataset repository.\"\"\"\n         if self._dataset_git_revision is None:\n-            self._dataset_git_revision = get_dataset_git_revision(\n-                dataset=self.job_params[\"dataset\"],\n-                hf_endpoint=self.common_config.hf_endpoint,\n-                hf_token=self.common_config.hf_token,\n-            )\n+            try:\n+                self._dataset_git_revision = get_dataset_git_revision(\n+                    dataset=self.job_params[\"dataset\"],\n+                    hf_endpoint=self.common_config.hf_endpoint,\n+                    hf_token=self.common_config.hf_token,\n+                )\n+            except Exception as e:\n+                if allow_raise:\n+                    raise e\n+                else:\n+                    return None\n         return self._dataset_git_revision\n \n     def raise_if_parallel_response_exists(self, parallel_cache_kind: str, parallel_job_version: int) -> None:\n@@ -210,12 +218,12 @@\n             self.debug(f\"response for job_info={self.job_info} had an error, cache updated\")\n             return False\n \n-    def backfill(self) -> None:\n+    def backfill(self, revision: str) -> None:\n         \"\"\"Evaluate the state of the dataset and backfill the cache if necessary.\"\"\"\n         DatasetState(\n             dataset=self.job_params[\"dataset\"],\n             processing_graph=self.processing_graph,\n-            revision=self.get_dataset_git_revision(),\n+            revision=revision,\n             error_codes_to_retry=ERROR_CODES_TO_RETRY,\n             priority=self.priority,\n         ).backfill()\n@@ -230,7 +238,7 @@\n             error_code=error.code,\n             details=dict(error.as_response_with_cause()),\n             job_runner_version=self.job_runner.get_job_runner_version(),\n-            dataset_git_revision=self.get_dataset_git_revision(),\n+            dataset_git_revision=self.get_dataset_git_revision(allow_raise=False),\n         )\n         logging.debug(\n             f\"response for dataset={self.job_params['dataset']} job_info={self.job_info} had an error (crashed), cache\"\n@@ -247,7 +255,7 @@\n             error_code=error.code,\n             details=dict(error.as_response_with_cause()),\n             job_runner_version=self.job_runner.get_job_runner_version(),\n-            dataset_git_revision=self.get_dataset_git_revision(),\n+            dataset_git_revision=self.get_dataset_git_revision(allow_raise=False),\n         )\n         logging.debug(\n             f\"response for dataset={self.job_params['dataset']} job_info={self.job_info} had an error (exceeded\"\n", "issue": "The workers fail with `mongoengine.errors.FieldDoesNotExist: The fields \"{'force'}\" do not exist on the document \"Job\"`\n```\r\nINFO: 2023-05-13 10:50:18,007 - root - Worker loop started\r\nINFO: 2023-05-13 10:50:18,023 - root - Starting heartbeat.\r\nERROR: 2023-05-13 10:50:18,115 - asyncio - Task exception was never retrieved\r\nfuture: <Task finished name='Task-2' coro=<every() done, defined at /src/services/worker/src/worker/executor.py:26> exception=FieldDoesNotExist('The fields \"{\\'force\\'}\" do not exist on the document \"Job\"')>\r\nTraceback (most recent call last):\r\n  File \"/src/services/worker/src/worker/executor.py\", line 30, in every\r\n    out = func(*args, **kwargs)\r\n  File \"/src/services/worker/src/worker/executor.py\", line 117, in kill_zombies\r\n    zombies = queue.get_zombies(max_seconds_without_heartbeat=self.max_seconds_without_heartbeat_for_zombies)\r\n  File \"/src/libs/libcommon/src/libcommon/queue.py\", line 622, in get_zombies\r\n    zombies = [\r\n  File \"/src/libs/libcommon/src/libcommon/queue.py\", line 622, in <listcomp>\r\n    zombies = [\r\n  File \"/src/services/worker/.venv/lib/python3.9/site-packages/mongoengine/queryset/queryset.py\", line 110, in _iter_results\r\n    self._populate_cache()\r\n  File \"/src/services/worker/.venv/lib/python3.9/site-packages/mongoengine/queryset/queryset.py\", line 129, in _populate_cache\r\n    self._result_cache.append(next(self))\r\n  File \"/src/services/worker/.venv/lib/python3.9/site-packages/mongoengine/queryset/base.py\", line 1599, in __next__\r\n    doc = self._document._from_son(\r\n  File \"/src/services/worker/.venv/lib/python3.9/site-packages/mongoengine/base/document.py\", line 836, in _from_son\r\n    obj = cls(__auto_convert=False, _created=created, **data)\r\n  File \"/src/services/worker/.venv/lib/python3.9/site-packages/mongoengine/base/document.py\", line 99, in __init__\r\n    raise FieldDoesNotExist(msg)\r\nmongoengine.errors.FieldDoesNotExist: The fields \"{'force'}\" do not exist on the document \"Job\"\r\nTraceback (most recent call last):\r\n  File \"/src/services/worker/src/worker/main.py\", line 61, in <module>\r\n    worker_executor.start()\r\n  File \"/src/services/worker/src/worker/executor.py\", line 89, in start\r\n    loop.run_until_complete(\r\n  File \"/usr/local/lib/python3.9/asyncio/base_events.py\", line 645, in run_until_complete\r\n    raise RuntimeError('Event loop stopped before Future completed.')\r\nRuntimeError: Event loop stopped before Future completed.\r\n```\r\n\r\n<img width=\"622\" alt=\"Capture d\u2019e\u0301cran 2023-05-13 a\u0300 12 52 06\" src=\"https://github.com/huggingface/datasets-server/assets/1676121/3e4d7e1a-dd02-476d-b671-57ce03d28d7b\">\r\n\n", "before_files": [{"content": "# SPDX-License-Identifier: Apache-2.0\n# Copyright 2022 The HuggingFace Authors.\n\nimport logging\nfrom http import HTTPStatus\nfrom typing import Literal, Optional\n\nfrom libcommon.config import CommonConfig\nfrom libcommon.dataset import DatasetNotFoundError, get_dataset_git_revision\nfrom libcommon.exceptions import CustomError\nfrom libcommon.processing_graph import ProcessingGraph, ProcessingStep\nfrom libcommon.simple_cache import (\n    DoesNotExist,\n    get_response_without_content_params,\n    upsert_response_params,\n)\nfrom libcommon.state import DatasetState\nfrom libcommon.utils import JobInfo, JobParams, Priority, Status, orjson_dumps\n\nfrom worker.common_exceptions import (\n    JobManagerCrashedError,\n    JobManagerExceededMaximumDurationError,\n    NoGitRevisionError,\n    ResponseAlreadyComputedError,\n    TooBigContentError,\n    UnexpectedError,\n)\nfrom worker.config import AppConfig, WorkerConfig\nfrom worker.job_runner import JobRunner\n\n# List of error codes that should trigger a retry.\nERROR_CODES_TO_RETRY: list[str] = [\"ClientConnectionError\"]\n\n\nclass JobManager:\n    \"\"\"\n    A job manager is a class that handles a job runner compute, for a specific processing step.\n\n    Args:\n        job_info (:obj:`JobInfo`):\n            The job to process. It contains the job_id, the job type, the dataset, the config, the split\n            and the priority level.\n        common_config (:obj:`CommonConfig`):\n            The common config.\n        processing_step (:obj:`ProcessingStep`):\n            The processing step to process.\n    \"\"\"\n\n    job_id: str\n    job_params: JobParams\n    priority: Priority\n    worker_config: WorkerConfig\n    common_config: CommonConfig\n    processing_step: ProcessingStep\n    processing_graph: ProcessingGraph\n    _dataset_git_revision: Optional[str] = None\n    job_runner: JobRunner\n\n    def __init__(\n        self,\n        job_info: JobInfo,\n        app_config: AppConfig,\n        job_runner: JobRunner,\n        processing_graph: ProcessingGraph,\n    ) -> None:\n        self.job_info = job_info\n        self.job_type = job_info[\"type\"]\n        self.job_id = job_info[\"job_id\"]\n        self.priority = job_info[\"priority\"]\n        self.job_params = job_info[\"params\"]\n        self.common_config = app_config.common\n        self.worker_config = app_config.worker\n        self.job_runner = job_runner\n        self.processing_graph = processing_graph\n        self.processing_step = self.job_runner.processing_step\n        self.setup()\n\n    def setup(self) -> None:\n        job_type = self.job_runner.get_job_type()\n        if self.processing_step.job_type != job_type:\n            raise ValueError(\n                f\"The processing step's job type is {self.processing_step.job_type}, but the job manager only\"\n                f\" processes {job_type}\"\n            )\n        if self.job_type != job_type:\n            raise ValueError(\n                f\"The submitted job type is {self.job_type}, but the job manager only processes {job_type}\"\n            )\n\n    def __str__(self) -> str:\n        return f\"JobManager(job_id={self.job_id} dataset={self.job_params['dataset']} job_info={self.job_info}\"\n\n    def log(self, level: int, msg: str) -> None:\n        logging.log(level=level, msg=f\"[{self.processing_step.job_type}] {msg}\")\n\n    def debug(self, msg: str) -> None:\n        self.log(level=logging.DEBUG, msg=msg)\n\n    def info(self, msg: str) -> None:\n        self.log(level=logging.INFO, msg=msg)\n\n    def warning(self, msg: str) -> None:\n        self.log(level=logging.WARNING, msg=msg)\n\n    def exception(self, msg: str) -> None:\n        self.log(level=logging.ERROR, msg=msg)\n\n    def critical(self, msg: str) -> None:\n        self.log(level=logging.CRITICAL, msg=msg)\n\n    def run(self) -> Literal[Status.SUCCESS, Status.ERROR]:\n        try:\n            self.info(f\"compute {self}\")\n            result: Literal[Status.SUCCESS, Status.ERROR] = Status.SUCCESS if self.process() else Status.ERROR\n        except Exception:\n            self.exception(f\"error while computing {self}\")\n            result = Status.ERROR\n        self.backfill()\n        return result\n\n    def get_dataset_git_revision(self) -> Optional[str]:\n        \"\"\"Get the git revision of the dataset repository.\"\"\"\n        if self._dataset_git_revision is None:\n            self._dataset_git_revision = get_dataset_git_revision(\n                dataset=self.job_params[\"dataset\"],\n                hf_endpoint=self.common_config.hf_endpoint,\n                hf_token=self.common_config.hf_token,\n            )\n        return self._dataset_git_revision\n\n    def raise_if_parallel_response_exists(self, parallel_cache_kind: str, parallel_job_version: int) -> None:\n        try:\n            existing_response = get_response_without_content_params(\n                kind=parallel_cache_kind,\n                job_params=self.job_params,\n            )\n\n            dataset_git_revision = self.get_dataset_git_revision()\n            if (\n                existing_response[\"http_status\"] == HTTPStatus.OK\n                and existing_response[\"job_runner_version\"] == parallel_job_version\n                and existing_response[\"progress\"] == 1.0  # completed response\n                and dataset_git_revision is not None\n                and existing_response[\"dataset_git_revision\"] == dataset_git_revision\n            ):\n                raise ResponseAlreadyComputedError(\n                    f\"Response has already been computed and stored in cache kind: {parallel_cache_kind}. Compute will\"\n                    \" be skipped.\"\n                )\n        except DoesNotExist:\n            logging.debug(f\"no cache found for {parallel_cache_kind}.\")\n\n    def process(\n        self,\n    ) -> bool:\n        dataset_git_revision = None\n        try:\n            dataset_git_revision = self.get_dataset_git_revision()\n            if dataset_git_revision is None:\n                self.debug(f\"the dataset={self.job_params['dataset']} has no git revision, don't update the cache\")\n                raise NoGitRevisionError(f\"Could not get git revision for dataset {self.job_params['dataset']}\")\n            try:\n                self.job_runner.pre_compute()\n                parallel_job_runner = self.job_runner.get_parallel_job_runner()\n                if parallel_job_runner:\n                    self.raise_if_parallel_response_exists(\n                        parallel_cache_kind=parallel_job_runner[\"job_type\"],\n                        parallel_job_version=parallel_job_runner[\"job_runner_version\"],\n                    )\n\n                job_result = self.job_runner.compute()\n                content = job_result.content\n\n                # Validate content size\n                if len(orjson_dumps(content)) > self.worker_config.content_max_bytes:\n                    raise TooBigContentError(\n                        \"The computed response content exceeds the supported size in bytes\"\n                        f\" ({self.worker_config.content_max_bytes}).\"\n                    )\n            finally:\n                # ensure the post_compute hook is called even if the compute raises an exception\n                self.job_runner.post_compute()\n            upsert_response_params(\n                kind=self.processing_step.cache_kind,\n                job_params=self.job_params,\n                content=content,\n                http_status=HTTPStatus.OK,\n                job_runner_version=self.job_runner.get_job_runner_version(),\n                dataset_git_revision=dataset_git_revision,\n                progress=job_result.progress,\n            )\n            self.debug(f\"dataset={self.job_params['dataset']} job_info={self.job_info} is valid, cache updated\")\n            return True\n        except DatasetNotFoundError:\n            # To avoid filling the cache, we don't save this error. Otherwise, DoS is possible.\n            self.debug(f\"the dataset={self.job_params['dataset']} could not be found, don't update the cache\")\n            return False\n        except Exception as err:\n            e = err if isinstance(err, CustomError) else UnexpectedError(str(err), err)\n            upsert_response_params(\n                kind=self.processing_step.cache_kind,\n                job_params=self.job_params,\n                content=dict(e.as_response()),\n                http_status=e.status_code,\n                error_code=e.code,\n                details=dict(e.as_response_with_cause()),\n                job_runner_version=self.job_runner.get_job_runner_version(),\n                dataset_git_revision=dataset_git_revision,\n            )\n            self.debug(f\"response for job_info={self.job_info} had an error, cache updated\")\n            return False\n\n    def backfill(self) -> None:\n        \"\"\"Evaluate the state of the dataset and backfill the cache if necessary.\"\"\"\n        DatasetState(\n            dataset=self.job_params[\"dataset\"],\n            processing_graph=self.processing_graph,\n            revision=self.get_dataset_git_revision(),\n            error_codes_to_retry=ERROR_CODES_TO_RETRY,\n            priority=self.priority,\n        ).backfill()\n\n    def set_crashed(self, message: str, cause: Optional[BaseException] = None) -> None:\n        error = JobManagerCrashedError(message=message, cause=cause)\n        upsert_response_params(\n            kind=self.processing_step.cache_kind,\n            job_params=self.job_params,\n            content=dict(error.as_response()),\n            http_status=error.status_code,\n            error_code=error.code,\n            details=dict(error.as_response_with_cause()),\n            job_runner_version=self.job_runner.get_job_runner_version(),\n            dataset_git_revision=self.get_dataset_git_revision(),\n        )\n        logging.debug(\n            f\"response for dataset={self.job_params['dataset']} job_info={self.job_info} had an error (crashed), cache\"\n            \" updated\"\n        )\n\n    def set_exceeded_maximum_duration(self, message: str, cause: Optional[BaseException] = None) -> None:\n        error = JobManagerExceededMaximumDurationError(message=message, cause=cause)\n        upsert_response_params(\n            kind=self.processing_step.cache_kind,\n            job_params=self.job_params,\n            content=dict(error.as_response()),\n            http_status=error.status_code,\n            error_code=error.code,\n            details=dict(error.as_response_with_cause()),\n            job_runner_version=self.job_runner.get_job_runner_version(),\n            dataset_git_revision=self.get_dataset_git_revision(),\n        )\n        logging.debug(\n            f\"response for dataset={self.job_params['dataset']} job_info={self.job_info} had an error (exceeded\"\n            \" maximum duration), cache updated\"\n        )\n", "path": "services/worker/src/worker/job_manager.py"}], "after_files": [{"content": "# SPDX-License-Identifier: Apache-2.0\n# Copyright 2022 The HuggingFace Authors.\n\nimport logging\nfrom http import HTTPStatus\nfrom typing import Literal, Optional\n\nfrom libcommon.config import CommonConfig\nfrom libcommon.dataset import DatasetNotFoundError, get_dataset_git_revision\nfrom libcommon.exceptions import CustomError\nfrom libcommon.processing_graph import ProcessingGraph, ProcessingStep\nfrom libcommon.simple_cache import (\n    DoesNotExist,\n    get_response_without_content_params,\n    upsert_response_params,\n)\nfrom libcommon.state import DatasetState\nfrom libcommon.utils import JobInfo, JobParams, Priority, Status, orjson_dumps\n\nfrom worker.common_exceptions import (\n    JobManagerCrashedError,\n    JobManagerExceededMaximumDurationError,\n    NoGitRevisionError,\n    ResponseAlreadyComputedError,\n    TooBigContentError,\n    UnexpectedError,\n)\nfrom worker.config import AppConfig, WorkerConfig\nfrom worker.job_runner import JobRunner\n\n# List of error codes that should trigger a retry.\nERROR_CODES_TO_RETRY: list[str] = [\"ClientConnectionError\"]\n\n\nclass JobManager:\n    \"\"\"\n    A job manager is a class that handles a job runner compute, for a specific processing step.\n\n    Args:\n        job_info (:obj:`JobInfo`):\n            The job to process. It contains the job_id, the job type, the dataset, the config, the split\n            and the priority level.\n        common_config (:obj:`CommonConfig`):\n            The common config.\n        processing_step (:obj:`ProcessingStep`):\n            The processing step to process.\n    \"\"\"\n\n    job_id: str\n    job_params: JobParams\n    priority: Priority\n    worker_config: WorkerConfig\n    common_config: CommonConfig\n    processing_step: ProcessingStep\n    processing_graph: ProcessingGraph\n    _dataset_git_revision: Optional[str] = None\n    job_runner: JobRunner\n\n    def __init__(\n        self,\n        job_info: JobInfo,\n        app_config: AppConfig,\n        job_runner: JobRunner,\n        processing_graph: ProcessingGraph,\n    ) -> None:\n        self.job_info = job_info\n        self.job_type = job_info[\"type\"]\n        self.job_id = job_info[\"job_id\"]\n        self.priority = job_info[\"priority\"]\n        self.job_params = job_info[\"params\"]\n        self.common_config = app_config.common\n        self.worker_config = app_config.worker\n        self.job_runner = job_runner\n        self.processing_graph = processing_graph\n        self.processing_step = self.job_runner.processing_step\n        self.setup()\n\n    def setup(self) -> None:\n        job_type = self.job_runner.get_job_type()\n        if self.processing_step.job_type != job_type:\n            raise ValueError(\n                f\"The processing step's job type is {self.processing_step.job_type}, but the job manager only\"\n                f\" processes {job_type}\"\n            )\n        if self.job_type != job_type:\n            raise ValueError(\n                f\"The submitted job type is {self.job_type}, but the job manager only processes {job_type}\"\n            )\n\n    def __str__(self) -> str:\n        return f\"JobManager(job_id={self.job_id} dataset={self.job_params['dataset']} job_info={self.job_info}\"\n\n    def log(self, level: int, msg: str) -> None:\n        logging.log(level=level, msg=f\"[{self.processing_step.job_type}] {msg}\")\n\n    def debug(self, msg: str) -> None:\n        self.log(level=logging.DEBUG, msg=msg)\n\n    def info(self, msg: str) -> None:\n        self.log(level=logging.INFO, msg=msg)\n\n    def warning(self, msg: str) -> None:\n        self.log(level=logging.WARNING, msg=msg)\n\n    def exception(self, msg: str) -> None:\n        self.log(level=logging.ERROR, msg=msg)\n\n    def critical(self, msg: str) -> None:\n        self.log(level=logging.CRITICAL, msg=msg)\n\n    def run(self) -> Literal[Status.SUCCESS, Status.ERROR]:\n        try:\n            self.info(f\"compute {self}\")\n            result: Literal[Status.SUCCESS, Status.ERROR] = Status.SUCCESS if self.process() else Status.ERROR\n        except Exception:\n            self.exception(f\"error while computing {self}\")\n            result = Status.ERROR\n        revision = self.get_dataset_git_revision(allow_raise=False)\n        if revision is not None:\n            self.backfill(revision=revision)\n        return result\n\n    def get_dataset_git_revision(self, allow_raise: bool = True) -> Optional[str]:\n        \"\"\"Get the git revision of the dataset repository.\"\"\"\n        if self._dataset_git_revision is None:\n            try:\n                self._dataset_git_revision = get_dataset_git_revision(\n                    dataset=self.job_params[\"dataset\"],\n                    hf_endpoint=self.common_config.hf_endpoint,\n                    hf_token=self.common_config.hf_token,\n                )\n            except Exception as e:\n                if allow_raise:\n                    raise e\n                else:\n                    return None\n        return self._dataset_git_revision\n\n    def raise_if_parallel_response_exists(self, parallel_cache_kind: str, parallel_job_version: int) -> None:\n        try:\n            existing_response = get_response_without_content_params(\n                kind=parallel_cache_kind,\n                job_params=self.job_params,\n            )\n\n            dataset_git_revision = self.get_dataset_git_revision()\n            if (\n                existing_response[\"http_status\"] == HTTPStatus.OK\n                and existing_response[\"job_runner_version\"] == parallel_job_version\n                and existing_response[\"progress\"] == 1.0  # completed response\n                and dataset_git_revision is not None\n                and existing_response[\"dataset_git_revision\"] == dataset_git_revision\n            ):\n                raise ResponseAlreadyComputedError(\n                    f\"Response has already been computed and stored in cache kind: {parallel_cache_kind}. Compute will\"\n                    \" be skipped.\"\n                )\n        except DoesNotExist:\n            logging.debug(f\"no cache found for {parallel_cache_kind}.\")\n\n    def process(\n        self,\n    ) -> bool:\n        dataset_git_revision = None\n        try:\n            dataset_git_revision = self.get_dataset_git_revision()\n            if dataset_git_revision is None:\n                self.debug(f\"the dataset={self.job_params['dataset']} has no git revision, don't update the cache\")\n                raise NoGitRevisionError(f\"Could not get git revision for dataset {self.job_params['dataset']}\")\n            try:\n                self.job_runner.pre_compute()\n                parallel_job_runner = self.job_runner.get_parallel_job_runner()\n                if parallel_job_runner:\n                    self.raise_if_parallel_response_exists(\n                        parallel_cache_kind=parallel_job_runner[\"job_type\"],\n                        parallel_job_version=parallel_job_runner[\"job_runner_version\"],\n                    )\n\n                job_result = self.job_runner.compute()\n                content = job_result.content\n\n                # Validate content size\n                if len(orjson_dumps(content)) > self.worker_config.content_max_bytes:\n                    raise TooBigContentError(\n                        \"The computed response content exceeds the supported size in bytes\"\n                        f\" ({self.worker_config.content_max_bytes}).\"\n                    )\n            finally:\n                # ensure the post_compute hook is called even if the compute raises an exception\n                self.job_runner.post_compute()\n            upsert_response_params(\n                kind=self.processing_step.cache_kind,\n                job_params=self.job_params,\n                content=content,\n                http_status=HTTPStatus.OK,\n                job_runner_version=self.job_runner.get_job_runner_version(),\n                dataset_git_revision=dataset_git_revision,\n                progress=job_result.progress,\n            )\n            self.debug(f\"dataset={self.job_params['dataset']} job_info={self.job_info} is valid, cache updated\")\n            return True\n        except DatasetNotFoundError:\n            # To avoid filling the cache, we don't save this error. Otherwise, DoS is possible.\n            self.debug(f\"the dataset={self.job_params['dataset']} could not be found, don't update the cache\")\n            return False\n        except Exception as err:\n            e = err if isinstance(err, CustomError) else UnexpectedError(str(err), err)\n            upsert_response_params(\n                kind=self.processing_step.cache_kind,\n                job_params=self.job_params,\n                content=dict(e.as_response()),\n                http_status=e.status_code,\n                error_code=e.code,\n                details=dict(e.as_response_with_cause()),\n                job_runner_version=self.job_runner.get_job_runner_version(),\n                dataset_git_revision=dataset_git_revision,\n            )\n            self.debug(f\"response for job_info={self.job_info} had an error, cache updated\")\n            return False\n\n    def backfill(self, revision: str) -> None:\n        \"\"\"Evaluate the state of the dataset and backfill the cache if necessary.\"\"\"\n        DatasetState(\n            dataset=self.job_params[\"dataset\"],\n            processing_graph=self.processing_graph,\n            revision=revision,\n            error_codes_to_retry=ERROR_CODES_TO_RETRY,\n            priority=self.priority,\n        ).backfill()\n\n    def set_crashed(self, message: str, cause: Optional[BaseException] = None) -> None:\n        error = JobManagerCrashedError(message=message, cause=cause)\n        upsert_response_params(\n            kind=self.processing_step.cache_kind,\n            job_params=self.job_params,\n            content=dict(error.as_response()),\n            http_status=error.status_code,\n            error_code=error.code,\n            details=dict(error.as_response_with_cause()),\n            job_runner_version=self.job_runner.get_job_runner_version(),\n            dataset_git_revision=self.get_dataset_git_revision(allow_raise=False),\n        )\n        logging.debug(\n            f\"response for dataset={self.job_params['dataset']} job_info={self.job_info} had an error (crashed), cache\"\n            \" updated\"\n        )\n\n    def set_exceeded_maximum_duration(self, message: str, cause: Optional[BaseException] = None) -> None:\n        error = JobManagerExceededMaximumDurationError(message=message, cause=cause)\n        upsert_response_params(\n            kind=self.processing_step.cache_kind,\n            job_params=self.job_params,\n            content=dict(error.as_response()),\n            http_status=error.status_code,\n            error_code=error.code,\n            details=dict(error.as_response_with_cause()),\n            job_runner_version=self.job_runner.get_job_runner_version(),\n            dataset_git_revision=self.get_dataset_git_revision(allow_raise=False),\n        )\n        logging.debug(\n            f\"response for dataset={self.job_params['dataset']} job_info={self.job_info} had an error (exceeded\"\n            \" maximum duration), cache updated\"\n        )\n", "path": "services/worker/src/worker/job_manager.py"}]}