Add heartbeat

libs/libcommon/src/libcommon/queue.py

@@ -37,7 +37,7 @@ def __get__(self, instance: object, cls: Type[U]) -> QuerySet[U]:
 # END monkey patching ### hack ###
 
 
-class Status(enum.Enum):
+class Status(str, enum.Enum):
     WAITING = "waiting"
     STARTED = "started"
     SUCCESS = "success"
@@ -46,7 +46,7 @@ class Status(enum.Enum):
     SKIPPED = "skipped"
 
 
-class Priority(enum.Enum):
+class Priority(str, enum.Enum):
     NORMAL = "normal"
     LOW = "low"
 
@@ -64,6 +64,7 @@ class JobDict(TypedDict):
     created_at: datetime
     started_at: Optional[datetime]
     finished_at: Optional[datetime]
+    last_heartbeat: Optional[datetime]
 
 
 class JobInfo(TypedDict):
@@ -121,6 +122,7 @@ class Job(Document):
         created_at (`datetime`): The creation date of the job.
         started_at (`datetime`, optional): When the job has started.
         finished_at (`datetime`, optional): When the job has finished.
+        last_heartbeat (`datetime`, optional): Last time the running job got a heartbeat from the worker.
     """
 
     meta = {
@@ -148,6 +150,7 @@ class Job(Document):
     created_at = DateTimeField(required=True)
     started_at = DateTimeField()
     finished_at = DateTimeField()
+    last_heartbeat = DateTimeField()
 
     def to_dict(self) -> JobDict:
         return {
@@ -163,6 +166,7 @@ def to_dict(self) -> JobDict:
             "created_at": self.created_at,
             "started_at": self.started_at,
             "finished_at": self.finished_at,
+            "last_heartbeat": self.last_heartbeat,
         }
 
     objects = QuerySetManager["Job"]()

services/worker/dev.Dockerfile

@@ -30,12 +30,14 @@ COPY services/worker/poetry.lock ./services/worker/poetry.lock
 COPY services/worker/pyproject.toml ./services/worker/pyproject.toml
 COPY libs/libcommon ./libs/libcommon
 WORKDIR /src/services/worker/
-RUN poetry install --no-cache
+RUN --mount=type=cache,target=/home/.cache/pypoetry/cache \
+    --mount=type=cache,target=/home/.cache/pypoetry/artifacts \
+    poetry install --no-root
 
 # FOR LOCAL DEVELOPMENT ENVIRONMENT
 # No need to copy the source code since we map a volume in docker-compose-base.yaml
 # Removed: COPY services/worker/src ./src
 # Removed: RUN poetry install --no-cache
 # However we need to install the package when the container starts
 # Added: poetry install
-ENTRYPOINT ["/bin/sh", "-c" , "poetry install && poetry run python src/worker/main.py"]
+ENTRYPOINT ["/bin/sh", "-c" , "poetry install --only-root && poetry run python src/worker/main.py"]

services/worker/pyproject.toml

@@ -40,6 +40,7 @@ transformers = "^4.26.1"
 trec-car-tools = { path = "vendors/trec-car-tools/python3" }
 typer = "^0.4.2"
 wget = "^3.2"
+mirakuru = "^2.4.2"
 
 [tool.poetry.group.dev.dependencies]
 bandit = "^1.7.4"

services/worker/src/worker/config.py

@@ -19,6 +19,7 @@
 WORKER_MAX_LOAD_PCT = 70
 WORKER_MAX_MEMORY_PCT = 80
 WORKER_SLEEP_SECONDS = 15
+WORKER_HEARTBEAT_TIME_INTERVAL_SECONDS = 60
 
 
 def get_empty_str_list() -> List[str]:
@@ -34,6 +35,8 @@ class WorkerConfig:
     only_job_types: list[str] = field(default_factory=get_empty_str_list)
     sleep_seconds: int = WORKER_SLEEP_SECONDS
     storage_paths: List[str] = field(default_factory=get_empty_str_list)
+    state_path: Optional[str] = None
+    heartbeat_time_interval_seconds: int = WORKER_HEARTBEAT_TIME_INTERVAL_SECONDS
 
     @classmethod
     def from_env(cls) -> "WorkerConfig":
@@ -47,6 +50,10 @@ def from_env(cls) -> "WorkerConfig":
                 sleep_seconds=env.int(name="SLEEP_SECONDS", default=WORKER_SLEEP_SECONDS),
                 only_job_types=env.list(name="ONLY_JOB_TYPES", default=get_empty_str_list()),
                 storage_paths=env.list(name="STORAGE_PATHS", default=get_empty_str_list()),
+                state_path=env.str(name="STATE_PATH", default=None),
+                heartbeat_time_interval_seconds=env.int(
+                    name="HEARTBEAT_TIME_INTERVAL_SECONDS", default=WORKER_HEARTBEAT_TIME_INTERVAL_SECONDS
+                ),
             )
 
 

services/worker/src/worker/loop.py

@@ -1,12 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 The HuggingFace Authors.
 
+import json
 import logging
 import random
 import time
 from dataclasses import dataclass, field
+from typing import Optional, TypedDict
 
-from libcommon.queue import EmptyQueueError, Queue
+from filelock import FileLock
+from libcommon.queue import EmptyQueueError, JobInfo, Queue
 from psutil import cpu_count, disk_usage, getloadavg, swap_memory, virtual_memory
 
 from worker.config import WorkerConfig
@@ -17,6 +20,10 @@ class UnknownJobTypeError(Exception):
     pass
 
 
+class WorkerState(TypedDict):
+    current_job_info: Optional[JobInfo]
+
+
 @dataclass
 class Loop:
     """
@@ -96,7 +103,7 @@ def sleep(self) -> None:
         time.sleep(duration)
 
     def run(self) -> None:
-        logging.info("Worker started")
+        logging.info("Worker loop started")
         try:
             while True:
                 if self.has_resources() and self.process_next_job():
@@ -119,13 +126,23 @@ def process_next_job(self) -> bool:
                     f" ${', '.join(self.worker_config.only_job_types)}). The queue should not have provided this"
                     " job. It is in an inconsistent state. Please report this issue to the datasets team."
                 )
+            self.set_worker_state(current_job_info=job_info)
             logging.debug(f"job assigned: {job_info}")
         except EmptyQueueError:
+            self.set_worker_state(current_job_info=None)
             logging.debug("no job in the queue")
             return False
 
         job_runner = self.job_runner_factory.create_job_runner(job_info)
         finished_status = job_runner.run()
         self.queue.finish_job(job_id=job_runner.job_id, finished_status=finished_status)
+        self.set_worker_state(current_job_info=None)
         logging.debug(f"job finished with {finished_status.value}: {job_runner}")
         return True
+
+    def set_worker_state(self, current_job_info: Optional[JobInfo]) -> None:
+        worker_state: WorkerState = {"current_job_info": current_job_info}
+        if self.worker_config.state_path:
+            with FileLock(self.worker_config.state_path + ".lock"):
+                with open(self.worker_config.state_path, "w") as worker_state_f:
+                    json.dump(worker_state, worker_state_f)

services/worker/src/worker/main.py

@@ -1,53 +1,90 @@
-# SPDX-License-Identifier: Apache-2.0
-# Copyright 2022 The HuggingFace Authors.
+import json
+import logging
+import os
+import sys
+import tempfile
+import time
+from typing import Optional
 
+from filelock import FileLock
 from libcommon.log import init_logging
-from libcommon.processing_graph import ProcessingGraph
-from libcommon.resources import CacheMongoResource, QueueMongoResource
-from libcommon.storage import init_assets_dir
+from libcommon.queue import Job, Status, get_datetime
+from libcommon.resources import QueueMongoResource
+from mirakuru import OutputExecutor
 
+from worker import start_worker_loop
 from worker.config import AppConfig
-from worker.job_runner_factory import JobRunnerFactory
-from worker.loop import Loop
-from worker.resources import LibrariesResource
+from worker.loop import WorkerState
+
+WORKER_STATE_FILE_NAME = "worker_state.json"
+START_WORKER_LOOP_PATH = start_worker_loop.__file__
+
+
+class WorkerExecutor:
+    def __init__(self, app_config: AppConfig) -> None:
+        self.app_config = app_config
+
+    def _create_worker_loop_executor(self) -> OutputExecutor:
+        banner = self.app_config.worker.state_path
+        if not banner:
+            raise ValueError("Failed to create the executor because WORKER_STATE_PATH is missing.")
+        start_worker_loop_command = [
+            sys.executable,
+            START_WORKER_LOOP_PATH,
+            "--print-worker-state-path",
+        ]
+        return OutputExecutor(start_worker_loop_command, banner, timeout=10)
+
+    def start(self) -> None:
+        worker_loop_executor = self._create_worker_loop_executor()
+        worker_loop_executor.start()  # blocking until the banner is printed
+        logging.info("Starting heartbeat.")
+        while worker_loop_executor.running():
+            self.heartbeat()
+            time.sleep(self.app_config.worker.heartbeat_time_interval_seconds)
+        worker_loop_executor.stop()
+
+    def get_state(self) -> WorkerState:
+        worker_state_path = self.app_config.worker.state_path
+        if not worker_state_path:
+            raise ValueError("Failed to get worker state because WORKER_STATE_PATH is missing.")
+        if os.path.exists(worker_state_path):
+            with FileLock(worker_state_path + ".lock"):
+                try:
+                    with open(worker_state_path, "r") as worker_state_f:
+                        worker_state = json.load(worker_state_f)
+                        return WorkerState(current_job_info=worker_state.get("current_job_info"))
+                except json.JSONDecodeError:
+                    return WorkerState(current_job_info=None)
+        else:
+            return WorkerState(current_job_info=None)
+
+    def get_current_job(self) -> Optional[Job]:
+        worker_state = self.get_state()
+        if worker_state["current_job_info"]:
+            job = Job.objects.with_id(worker_state["current_job_info"]["job_id"])  # type: ignore
+            if job and isinstance(job, Job) and job.status == Status.STARTED:
+                return job
+        return None
+
+    def heartbeat(self) -> None:
+        current_job = self.get_current_job()
+        if current_job:
+            current_job.update(last_heartbeat=get_datetime())
+
 
 if __name__ == "__main__":
-    app_config = AppConfig.from_env()
-
-    init_logging(log_level=app_config.common.log_level)
-    # ^ set first to have logs as soon as possible
-    assets_directory = init_assets_dir(directory=app_config.assets.storage_directory)
-
-    processing_graph = ProcessingGraph(app_config.processing_graph.specification)
-
-    with (
-        LibrariesResource(
-            hf_endpoint=app_config.common.hf_endpoint,
-            init_hf_datasets_cache=app_config.datasets_based.hf_datasets_cache,
-            numba_path=app_config.numba.path,
-        ) as libraries_resource,
-        CacheMongoResource(
-            database=app_config.cache.mongo_database, host=app_config.cache.mongo_url
-        ) as cache_resource,
-        QueueMongoResource(
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        if "WORKER_STATE_PATH" not in os.environ:
+            os.environ["WORKER_STATE_PATH"] = os.path.join(tmp_dir, WORKER_STATE_FILE_NAME)
+
+        app_config = AppConfig.from_env()
+        init_logging(log_level=app_config.common.log_level)
+
+        with QueueMongoResource(
             database=app_config.queue.mongo_database, host=app_config.queue.mongo_url
-        ) as queue_resource,
-    ):
-        if not cache_resource.is_available():
-            raise RuntimeError("The connection to the cache database could not be established. Exiting.")
-        if not queue_resource.is_available():
-            raise RuntimeError("The connection to the queue database could not be established. Exiting.")
-
-        job_runner_factory = JobRunnerFactory(
-            app_config=app_config,
-            processing_graph=processing_graph,
-            hf_datasets_cache=libraries_resource.hf_datasets_cache,
-            assets_directory=assets_directory,
-        )
-        loop = Loop(
-            library_cache_paths=libraries_resource.storage_paths,
-            job_runner_factory=job_runner_factory,
-            max_jobs_per_namespace=app_config.queue.max_jobs_per_namespace,
-            worker_config=app_config.worker,
-        )
-        loop.run()
+        ) as queue_resource:
+            if not queue_resource.is_available():
+                raise RuntimeError("The connection to the queue database could not be established. Exiting.")
+            worker_executor = WorkerExecutor(app_config)
+            worker_executor.start()