apputils - DBTaskManager (#57)

Author: mfranciszkiewicz

python/golem_task_api/apputils/__init__.py

@@ -0,0 +1,21 @@
+from pathlib import Path
+
+import peewee
+from typing import Iterable, Type
+
+
+database = peewee.DatabaseProxy()
+
+
+def initialize_database(
+    db: peewee.Database,
+    db_path: Path,
+    models: Iterable[Type[peewee.Model]],
+) -> None:
+    """ Initialize and bind the corresponding database instance
+        to the database proxy. Create tables for the given models """
+    db_path.parent.mkdir(parents=True, exist_ok=True)
+    db.init(str(db_path))
+
+    database.initialize(db)
+    database.create_tables(models, safe=True)

python/golem_task_api/apputils/database/__init__.py

@@ -0,0 +1,33 @@
+import enum
+
+from typing import Set, Type
+
+import peewee
+
+
+class EnumField(peewee.CharField):
+
+    def __init__(
+            self,
+            enum_type: Type[enum.Enum],
+            *args,
+            **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+
+        self._type: Type[enum.Enum] = enum_type
+        self._values: Set[enum.Enum] = set([e.value for e in self._type])
+
+    def db_value(self, value):
+        if isinstance(value, self._type):
+            return value.value
+
+        value = self.coerce(value)
+        if value not in self._values:
+            raise TypeError(
+                f"Expected {self._type.__name__} "
+                f"type or one of {self._values}")
+        return value
+
+    def python_value(self, value):
+        return self._type(value)

python/golem_task_api/apputils/database/fields.py

@@ -0,0 +1,126 @@
+import enum
+
+from abc import ABC, abstractmethod
+from typing import Optional, Dict, List, Tuple
+
+
+class SubtaskStatus(enum.Enum):
+    WAITING = None
+    COMPUTING = 'computing'
+    VERIFYING = 'verifying'
+    SUCCESS = 'success'
+    FAILURE = 'failure'
+    ABORTED = 'aborted'
+
+    @classmethod
+    def default(cls) -> 'SubtaskStatus':
+        return cls.WAITING
+
+    def is_computable(self) -> bool:
+        return self in (self.WAITING, self.FAILURE, self.ABORTED)
+
+
+class TaskManager(ABC):
+    """ Responsible for managing the internal task state.
+
+        Based on a concept of a:
+
+        - part
+
+          An outcome of splitting a task into separate units of work. There
+          usually exists a constant number of parts.
+
+        - subtask
+
+          A clone of a chosen task part that will be assigned to a
+          computing node.
+
+          Each subtask is given a unique identifier in order to distinguish
+          computation attempts of the same part, which may fail due to
+          unexpected errors or simply time out.
+
+          A successful subtask computation concludes the computation of the
+          corresponding part.
+
+        This class is responsible for correlating subtasks with task parts and
+        managing their status.
+    """
+
+    @abstractmethod
+    def create_task(
+            self,
+            part_count: int
+    ) -> None:
+        """ Persist a "part_count" number of task parts and other necessary
+            task information """
+        raise NotImplementedError
+
+    @abstractmethod
+    def abort_task(
+            self,
+    ) -> None:
+        """ Change the statuses of currently assigned subtasks to ABORTED.
+            Subtasks must have a computable or COMPUTING status """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_part(
+            self,
+            part_num: int
+    ) -> Optional[object]:
+        """ Return a task part object with the given part_num, if exists """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_part_num(
+            self,
+            subtask_id: str
+    ) -> Optional[int]:
+        """ Return a task part object's number for the given subtask_id,
+            if exists """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_subtask(
+            self,
+            subtask_id: str
+    ) -> Optional[object]:
+        """ Return a subtask object with the given subtask_id, if exists """
+        raise NotImplementedError
+
+    @abstractmethod
+    def start_subtask(
+            self,
+            part_num: int,
+            subtask_id: str
+    ) -> None:
+        """ Persist a new subtask with the given subtask_id and a COMPUTING
+            status. Assign that subtask to a task part object with the given
+            part_num """
+        raise NotImplementedError
+
+    @abstractmethod
+    def update_subtask_status(
+        self,
+        subtask_id: str,
+        status: SubtaskStatus,
+    ) -> None:
+        """ Update the status of a subtask with the given subtask_id """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_subtasks_statuses(
+            self,
+            part_nums: List[int],
+    ) -> Dict[int, Tuple[SubtaskStatus, Optional[str]]]:
+        """ Return (subtask status, subtask id) tuples mapped to part numbers
+            with assigned subtasks or (WAITING, None) tuples otherwise.
+            Task parts are chosen from the "part_nums" pool """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_next_computable_part_num(
+            self,
+    ) -> Optional[int]:
+        """ Return the next computable task part's number, if exists """
+        raise NotImplementedError

python/golem_task_api/apputils/task/__init__.py

@@ -0,0 +1,213 @@
+from pathlib import Path
+from typing import Optional, List, Dict, Tuple, Type
+
+import peewee
+from golem_task_api.apputils.database import (
+    database,
+    initialize_database,
+)
+from golem_task_api.apputils.database.fields import EnumField
+from golem_task_api.apputils.task import SubtaskStatus, TaskManager
+
+
+class Part(peewee.Model):
+
+    class Meta:
+        database = database
+
+    num = peewee.IntegerField(
+        primary_key=True,
+        unique=True)
+    subtask = peewee.DeferredForeignKey(
+        'Subtask',
+        null=True,
+        index=True,
+        backref='parts')
+
+
+class Subtask(peewee.Model):
+
+    class Meta:
+        database = database
+
+    id = peewee.CharField(
+        primary_key=True,
+        unique=True)
+    part = peewee.ForeignKeyField(
+        Part,
+        null=True,
+        index=True,
+        backref='subtasks')
+    status = EnumField(
+        SubtaskStatus,
+        default=SubtaskStatus.default)
+
+
+class DBTaskManager(TaskManager):
+    """ TaskManager subclass with a SQLite database backend for persistence """
+
+    def __init__(
+            self,
+            work_dir: Path,
+            part_model: Type[Part] = Part,
+            subtask_model: Type[Subtask] = Subtask,
+    ) -> None:
+
+        self._part = part_model
+        self._subtask = subtask_model
+        models = [self._part, self._subtask]
+
+        initialize_database(
+            db=self._database,
+            db_path=work_dir / 'task.db',
+            models=models)
+
+    @property
+    def _database(self) -> peewee.Database:
+        """ Return an instance of a chosen `peewee.Database` subclass """
+        return peewee.SqliteDatabase(
+            database=None,
+            thread_safe=True,
+            pragmas=(
+                ('foreign_keys', True),
+                ('busy_timeout', 1000),
+                ('journal_mode', 'WAL'),
+            ),
+        )
+
+    def create_task(
+            self,
+            part_count: int,
+    ) -> None:
+        with database.atomic():
+            for num in range(part_count):
+                self._part.create(num=num, subtask=None)
+
+    def abort_task(
+            self,
+    ) -> None:
+
+        self._subtask.update(
+            status=SubtaskStatus.ABORTED,
+        ).where(
+            self._subtask.status != SubtaskStatus.SUCCESS,
+            self._subtask.id.in_(
+                self._part.select(
+                    self._part.subtask.id
+                )
+            )
+        ).execute()
+
+    def get_part(
+            self,
+            part_num: int,
+    ) -> Optional[Part]:
+        try:
+            return self._part.get(num=part_num)
+        except peewee.DoesNotExist:
+            return None
+
+    def get_part_num(
+            self,
+            subtask_id: str,
+    ) -> Optional[int]:
+
+        result = self._part.select(
+            self._part.num
+        ).join(
+            self._subtask
+        ).where(
+            self._subtask.id == subtask_id
+        ).execute()
+
+        return result[0].num if result else None
+
+    def get_subtask(
+            self,
+            subtask_id: str,
+    ) -> Optional[Subtask]:
+        try:
+            return self._subtask.get(id=subtask_id)
+        except peewee.DoesNotExist:
+            return None
+
+    def start_subtask(
+            self,
+            part_num: int,
+            subtask_id: str,
+    ) -> None:
+
+        status = self.get_subtasks_statuses([part_num])
+        if status and part_num in status:
+            if not status[part_num][0].is_computable():
+                raise RuntimeError(f"Subtask {part_num} already started")
+
+        with database.atomic():
+            part = self._part.get(self._part.num == part_num)
+            part.subtask = subtask_id
+            part.save()
+
+            self._subtask.create(
+                id=subtask_id,
+                part=part,
+                status=SubtaskStatus.COMPUTING)
+
+    def update_subtask_status(
+        self,
+        subtask_id: str,
+        status: SubtaskStatus
+    ) -> None:
+
+        self._subtask.update(
+            status=status,
+        ).where(
+            self._subtask.id == subtask_id,
+        ).execute()
+
+    def get_subtasks_statuses(
+            self,
+            part_nums: List[int],
+    ) -> Dict[int, Tuple[SubtaskStatus, Optional[str]]]:
+
+        parts = self._part.select(
+            self._part.num,
+            self._subtask.status,
+            self._subtask.id,
+        ).join(
+            self._subtask,
+            join_type=peewee.JOIN.LEFT_OUTER,
+        ).where(
+            self._part.num.in_(part_nums),
+        ).execute()
+
+        return {
+            p.num: (
+                SubtaskStatus(p.subtask.status if p.subtask else None),
+                p.subtask.id if p.subtask else None
+            )
+            for p in parts
+        }
+
+    def get_next_computable_part_num(
+            self,
+    ) -> Optional[int]:
+
+        results = self._part.select(
+            self._part.num
+        ).where(
+            ~peewee.fn.EXISTS(
+                self._subtask.select(
+                    self._subtask.id,
+                ).where(
+                    self._subtask.id == self._part.subtask,
+                    self._subtask.status.not_in((
+                        SubtaskStatus.ABORTED,
+                        SubtaskStatus.FAILURE
+                    ))
+                )
+            )
+        ).order_by(
+            self._part.num.asc()
+        ).limit(1).execute()
+
+        return results[0].num if results else None