nerdd-link 0.2.11__py3-none-any.whl

nerdd_link/__init__.py

@@ -0,0 +1,6 @@
+from .actions import *
+from .channels import *
+from .converters import *
+from .files import *
+from .input import *
+from .types import *

nerdd_link/actions/__init__.py

@@ -0,0 +1,5 @@
+from .action import *
+from .predict_checkpoints_action import *
+from .process_jobs_action import *
+from .register_module_action import *
+from .serialize_job_action import *

nerdd_link/actions/action.py

@@ -0,0 +1,41 @@
+import logging
+from abc import ABC, abstractmethod
+from asyncio import CancelledError
+from typing import Generic, TypeVar
+
+from stringcase import spinalcase
+
+from ..channels import Channel, Topic
+from ..types import Message
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=Message)
+
+
+class Action(ABC, Generic[T]):
+    def __init__(self, input_topic: Topic[T]):
+        self._input_topic = input_topic
+
+    async def run(self) -> None:
+        consumer_group = spinalcase(self._get_group_name())
+        async for message in self._input_topic.receive(consumer_group):
+            try:
+                await self._process_message(message)
+            except CancelledError:
+                # the consumer was cancelled, stop processing messages
+                break
+            except Exception:
+                # log the error and continue processing the next message
+                logger.error("Error processing message", exc_info=True)
+
+    @abstractmethod
+    async def _process_message(self, message: T) -> None:
+        pass
+
+    @property
+    def channel(self) -> Channel:
+        return self._input_topic.channel
+
+    def _get_group_name(self) -> str:
+        return self.__class__.__name__

nerdd_link/actions/predict_checkpoints_action.py

@@ -0,0 +1,76 @@
+import logging
+from multiprocessing import Queue
+
+from nerdd_module import SimpleModel
+
+from ..channels import Channel
+from ..delegates import ReadCheckpointModel
+from ..files import FileSystem
+from ..types import CheckpointMessage, ResultCheckpointMessage, ResultMessage
+from .action import Action
+
+__all__ = ["PredictCheckpointsAction"]
+
+logger = logging.getLogger(__name__)
+
+
+class PredictCheckpointsAction(Action[CheckpointMessage]):
+    # Accept a batch of input molecules on the "<job-type>-checkpoints" topic
+    # (generated in the previous step) and process them. Results are written to
+    # the "results" topic.
+
+    def __init__(self, channel: Channel, model: SimpleModel, data_dir: str) -> None:
+        super().__init__(channel.checkpoints_topic(model))
+        self._model = model
+        self._data_dir = data_dir
+
+    async def _process_message(self, message: CheckpointMessage) -> None:
+        job_id = message.job_id
+        checkpoint_id = message.checkpoint_id
+        params = message.params
+        logger.info(f"Predict checkpoint {checkpoint_id} of job {job_id}")
+
+        # The Kafka consumers and producers run in the current asyncio event loop and (by
+        # observation) it seems that calling the produce method of a Kafka producer in a different
+        # event loop / thread / process doesn't seem to work (hangs indefinitely). Therefore, we
+        # create a queue in this event loop / thread and other tasks send messages to the queue
+        # instead of directly to the Kafka producer. This event loop will wait for new messages in
+        # this queue and forward them to the Kafka producer.
+        queue: Queue = Queue()
+
+        file_system = FileSystem(self._data_dir)
+
+        # create a wrapper model that
+        # * reads the checkpoint file instead of normal input
+        # * does preprocessing, prediction, and postprocessing like the encapsulated model
+        # * does not write to the specified results file, but to the checkpoints file instead
+        # * sends the results to the results topic
+        model = ReadCheckpointModel(
+            base_model=self._model,
+            job_id=job_id,
+            file_system=file_system,
+            checkpoint_id=checkpoint_id,
+            queue=queue,
+        )
+
+        # predict the checkpoint
+        # assign input=None, because the checkpoint file is already provided in ReadCheckpointModel
+        model.predict(
+            input=None,
+            **params,
+        )
+
+        # Wait for the prediction to finish and the results to be sent.
+        while True:
+            record = queue.get()
+            if record is not None:
+                await self.channel.results_topic().send(ResultMessage(job_id=job_id, **record))
+            else:
+                await self.channel.result_checkpoints_topic().send(
+                    ResultCheckpointMessage(job_id=job_id, checkpoint_id=checkpoint_id)
+                )
+                break
+
+    def _get_group_name(self) -> str:
+        model_id = self._model.get_config().id
+        return f"predict-checkpoints-{model_id}"

nerdd_link/actions/process_jobs_action.py

@@ -0,0 +1,157 @@
+import logging
+from pickle import dump
+from typing import Any
+
+from nerdd_module.input import DepthFirstExplorer
+from nerdd_module.model import ReadInputStep
+from rdkit.Chem import Mol
+from rdkit.Chem.PropertyMol import PropertyMol
+
+from ..channels import Channel
+from ..files import FileSystem
+from ..types import CheckpointMessage, JobMessage, LogMessage
+from ..utils import batched
+from .action import Action
+
+__all__ = ["ProcessJobsAction"]
+
+logger = logging.getLogger(__name__)
+
+
+class ProcessJobsAction(Action[JobMessage]):
+    # Accept new jobs (on the "<job_type>-jobs" topic). For each job, the program
+    # iterates through all molecules in the input (files), writes them as batches
+    # into checkpoint files and sends checkpoint messages (for each batch) to the
+    # "<job_type>-checkpoints" topic. Also, the number of molecules read is
+    # reported to the topic "job-sizes".
+
+    def __init__(
+        self,
+        channel: Channel,
+        checkpoint_size: int,
+        max_num_molecules: int,
+        num_test_entries: int,
+        ratio_valid_entries: float,
+        maximum_depth: int,
+        max_num_lines_mol_block: int,
+        data_dir: str,
+    ) -> None:
+        super().__init__(channel.jobs_topic())
+        # relevant for chunking
+        self._checkpoint_size = checkpoint_size
+        self._max_num_molecules = max_num_molecules
+        # parameters of DepthFirstExplorer
+        self._num_test_entries = num_test_entries
+        self._ratio_valid_entries = ratio_valid_entries
+        self._maximum_depth = maximum_depth
+        # used as kwargs in DepthFirstExplorer
+        self._max_num_lines_mol_block = max_num_lines_mol_block
+        self._file_system = FileSystem(data_dir)
+
+    async def _process_message(self, message: JobMessage) -> None:
+        job_id = message.id
+        job_type = message.job_type
+        logger.info(f"Received a new job {job_id} of type {job_type}")
+
+        # the input file to the job is stored in a designated sources directory
+        # (the file is allowed to reference other files, but setting the data_dir
+        # to the sources directory ensures that we never read files outside of the
+        # sources directory)
+        data_dir = self._file_system.get_sources_dir()
+
+        # create a reader (explorer) for the input file
+        explorer = DepthFirstExplorer(
+            num_test_entries=self._num_test_entries,
+            threshold=self._ratio_valid_entries,
+            maximum_depth=self._maximum_depth,
+            # extra args
+            max_num_lines_mol_block=self._max_num_lines_mol_block,
+            data_dir=data_dir,
+        )
+
+        read_input_step = ReadInputStep(explorer, message.source_id)
+
+        # read the input file
+        entries = read_input_step()
+
+        # iterate through the entries
+        # create batches of size checkpoint_size
+        # limit the number of molecules to max_num_molecules
+        batches = batched(entries, self._checkpoint_size)
+        num_entries = 0
+        num_checkpoints = 0
+        for i, batch in enumerate(batches):
+            # max_num_molecules might be reached within the batch
+            num_store = min(len(batch), self._max_num_molecules - num_entries)
+
+            # store batch in data_dir
+            with self._file_system.get_checkpoint_file_handle(job_id, i, "wb") as f:
+                results = list(batch[:num_store])
+
+                # TODO: use a model for storing the batches
+
+                # check all items for mol values and use PropertyMol for those
+                # in order to keep molecular properties (thanks, RDKit! :/ )
+                def _check_value(value: Any) -> Any:
+                    if isinstance(value, Mol):
+                        return PropertyMol(value)
+                    return value
+
+                def _check_item(item: dict) -> dict:
+                    return {key: _check_value(value) for key, value in item.items()}
+
+                results = [_check_item(item) for item in results]
+
+                dump(results, f)
+
+            # send a tuple to topic cypstrate-checkpoints
+            await self.channel.checkpoints_topic(job_type).send(
+                CheckpointMessage(
+                    job_id=job_id,
+                    checkpoint_id=i,
+                    params=message.params,
+                )
+            )
+
+            num_entries += num_store
+            num_checkpoints += 1
+
+            if num_entries >= self._max_num_molecules:
+                break
+
+        logger.info(f"Wrote {i+1} checkpoints containing {num_entries} entries for job {job_id}")
+
+        # send a warning message if there were more molecules in the job than allowed
+        too_many_molecules = num_store < len(batch)
+        try:
+            # try to get another entry
+            next(entries)
+
+            # if we get here, there was another entry and we need to send a warning
+            too_many_molecules = True
+        except StopIteration:
+            pass
+
+        if too_many_molecules:
+            await self.channel.logs_topic().send(
+                LogMessage(
+                    job_id=job_id,
+                    message_type="warning",
+                    message=(
+                        f"The provided job contains more than "
+                        f"{self._max_num_molecules} input structures. Only the "
+                        f"first {self._max_num_molecules} will be processed."
+                    ),
+                )
+            )
+
+        # at the end, send a tuple to topic job-sizes with the overall size
+        # of the job
+        await self.channel.logs_topic().send(
+            LogMessage(
+                job_id=job_id,
+                message_type="report_job_size",
+                num_entries=num_entries,
+                num_checkpoints=num_checkpoints,
+            )
+        )

nerdd_link/actions/register_module_action.py

@@ -0,0 +1,37 @@
+import json
+import logging
+
+from nerdd_module import Model
+
+from ..channels import Channel
+from ..files import FileSystem
+from ..types import ModuleMessage, SystemMessage
+from .action import Action
+
+__all__ = ["RegisterModuleAction"]
+
+logger = logging.getLogger(__name__)
+
+
+class RegisterModuleAction(Action[SystemMessage]):
+    def __init__(self, channel: Channel, model: Model, data_dir: str):
+        super().__init__(channel.system_topic())
+        # TODO: do this differently
+        assert hasattr(model, "get_config")
+        self._model = model
+        self._file_system = FileSystem(data_dir)
+
+    async def _process_message(self, message: SystemMessage) -> None:
+        config = self._model.get_config()
+        logger.info(f"Registering module with id {config.id}")
+
+        # save module as json to file
+        module_file = self._file_system.get_module_file_path(config.id)
+        json.dump(config.model_dump(), open(module_file, "w"))
+
+        # send the initialization message
+        await self.channel.modules_topic().send(ModuleMessage(id=config.id))
+
+    def _get_group_name(self) -> str:
+        model_id = self._model.get_config().id
+        return f"register-module-{model_id}"

nerdd_link/actions/serialize_job_action.py

@@ -0,0 +1,64 @@
+import json
+import logging
+
+from nerdd_module import OutputStep
+
+from ..channels import Channel
+from ..delegates import ReadPickleStep, SerializeJobModel
+from ..files import FileSystem
+from ..types import SerializationRequestMessage, SerializationResultMessage
+from .action import Action
+
+__all__ = ["SerializeJobAction"]
+
+
+logger = logging.getLogger(__name__)
+
+
+class SerializeJobAction(Action[SerializationRequestMessage]):
+    def __init__(self, channel: Channel, data_dir: str) -> None:
+        super().__init__(channel.serialization_requests_topic())
+        self._file_system = FileSystem(data_dir)
+
+    async def _process_message(self, message: SerializationRequestMessage) -> None:
+        job_id = message.job_id
+        job_type = message.job_type
+        params = message.params
+        output_format = message.output_format
+        logger.info(f"Write output for job {job_id} in format {output_format}")
+
+        # remove specific parameter keys that could induce vulnerabilities
+        params.pop("output_file", None)
+        params.pop("output_format", None)
+
+        # obtain output file
+        output_file = self._file_system.get_output_file(job_id, output_format)
+
+        # get the configuration for the job_type
+        config_file = self._file_system.get_module_file_path(job_type)
+        config = json.load(open(config_file, "r"))
+
+        # create a fake model instance to get the postprocessing steps
+        model = SerializeJobModel(config)
+
+        steps = [
+            # read the result checkpoint files in the correct order
+            ReadPickleStep(self._file_system.iter_results_file_handles(job_id)),
+            # don't preprocess, don't do prediction, only post-process
+            *model._get_postprocessing_steps(output_format, output_file=output_file, **params),
+        ]
+
+        output_step = steps[-1]
+        assert isinstance(output_step, OutputStep), "The last step must be an OutputStep."
+
+        # build the pipeline from the list of steps
+        pipeline = None
+        for t in steps:
+            pipeline = t(pipeline)
+
+        # run the pipeline by calling the get_result method of the last step
+        output_step.get_result()
+
+        await self.channel.serialization_results_topic().send(
+            SerializationResultMessage(job_id=job_id, output_format=output_format)
+        )

nerdd_link/channels/__init__.py

@@ -0,0 +1,3 @@
+from .channel import *
+from .kafka_channel import *
+from .memory_channel import *

nerdd_link/channels/channel.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import AsyncIterable, Generic, TypeVar, Union, cast
+
+from nerdd_module import Model
+from stringcase import spinalcase
+
+from ..types import (
+    CheckpointMessage,
+    JobMessage,
+    LogMessage,
+    Message,
+    ModuleMessage,
+    ResultCheckpointMessage,
+    ResultMessage,
+    SerializationRequestMessage,
+    SerializationResultMessage,
+    SystemMessage,
+)
+
+__all__ = ["Channel", "Topic"]
+
+T = TypeVar("T", bound=Message)
+
+
+def get_job_type(job_type_or_model: Union[str, Model]) -> str:
+    if isinstance(job_type_or_model, Model):
+        model = job_type_or_model
+
+        # create topic name from model name by
+        # * converting to spinal case, (e.g. "MyModel" -> "my-model")
+        # * converting to lowercase (just to be sure) and
+        # * removing all characters except dash and alphanumeric characters
+        # TODO: move to Module Id
+        topic_name = spinalcase(model.name)
+        topic_name = topic_name.lower()
+        topic_name = "".join([c for c in topic_name if str.isalnum(c) or c == "-"])
+        return topic_name
+    else:
+        return spinalcase(job_type_or_model)
+
+
+class Topic(Generic[T]):
+    def __init__(self, channel: Channel, name: str):
+        self._channel = channel
+        self._name = name
+
+    async def receive(self, consumer_group: str) -> AsyncIterable[T]:
+        async for msg in self.channel.iter_messages(self._name, consumer_group):
+            yield cast(T, msg)
+
+    async def send(self, message: T) -> None:
+        await self.channel.send(self._name, message)
+
+    @property
+    def channel(self) -> Channel:
+        return self._channel
+
+    def __repr__(self) -> str:
+        return f"Topic({self._name})"
+
+
+class Channel(ABC):
+    def __init__(self) -> None:
+        self._is_running = False
+
+    async def start(self) -> None:
+        self._is_running = True
+        await self._start()
+
+    async def _start(self) -> None:  # noqa: B027
+        pass
+
+    async def stop(self) -> None:
+        await self._stop()
+        self._is_running = False
+
+    async def _stop(self) -> None:  # noqa: B027
+        pass
+
+    async def __aenter__(self) -> Channel:
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type: type, exc_value: Exception, traceback: object) -> None:
+        await self.stop()
+
+    #
+    # RECEIVE
+    #
+    async def iter_messages(self, topic: str, consumer_group: str) -> AsyncIterable[Message]:
+        if not self._is_running:
+            raise RuntimeError("Channel is not running. Call start() first.")
+        async for message in self._iter_messages(topic, consumer_group):
+            yield message
+
+    # Insane glitch: we need to use "def _iter_messages" instead of "async def _iter_messages"
+    # here, because the method doesn't use "yield" and so the type checker will assume that the
+    # actual type is Coroutine[AsyncIterable[Message], None, None].
+    @abstractmethod
+    def _iter_messages(self, topic: str, consumer_group: str) -> AsyncIterable[Message]:
+        pass
+
+    #
+    # SEND
+    #
+    async def send(self, topic: str, message: Message) -> None:
+        if not self._is_running:
+            raise RuntimeError("Channel is not running. Call start() first.")
+        await self._send(topic, message)
+
+    @abstractmethod
+    async def _send(self, topic: str, message: Message) -> None:
+        pass
+
+    #
+    # TOPICS
+    #
+    def modules_topic(self) -> Topic[ModuleMessage]:
+        return Topic[ModuleMessage](self, "modules")
+
+    def jobs_topic(self) -> Topic[JobMessage]:
+        return Topic[JobMessage](self, "jobs")
+
+    def checkpoints_topic(self, job_type_or_model: Union[str, Model]) -> Topic[CheckpointMessage]:
+        job_type = get_job_type(job_type_or_model)
+        topic_name = f"{job_type}-checkpoints"
+        return Topic[CheckpointMessage](self, topic_name)
+
+    def results_topic(self) -> Topic[ResultMessage]:
+        return Topic[ResultMessage](self, "results")
+
+    def result_checkpoints_topic(self) -> Topic[ResultCheckpointMessage]:
+        return Topic[ResultCheckpointMessage](self, "result-checkpoints")
+
+    def serialization_requests_topic(self) -> Topic[SerializationRequestMessage]:
+        return Topic[SerializationRequestMessage](self, "serialization-requests")
+
+    def serialization_results_topic(self) -> Topic[SerializationResultMessage]:
+        return Topic[SerializationResultMessage](self, "serialization-results")
+
+    def logs_topic(self) -> Topic[LogMessage]:
+        return Topic[LogMessage](self, "logs")
+
+    def system_topic(self) -> Topic[SystemMessage]:
+        return Topic[SystemMessage](self, "system")

nerdd_link/channels/kafka_channel.py

@@ -0,0 +1,109 @@
+import json
+import logging
+from typing import AsyncIterable, Dict, Tuple
+
+from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
+from aiokafka.errors import CommitFailedError
+
+from ..types import Message
+from .channel import Channel
+
+__all__ = ["KafkaChannel"]
+
+logger = logging.getLogger(__name__)
+
+
+class KafkaChannel(Channel):
+    def __init__(self, broker_url: str) -> None:
+        super().__init__()
+        self._broker_url = broker_url
+        self._consumers: Dict[Tuple[str, str], AIOKafkaConsumer] = {}
+
+    async def _start(self) -> None:
+        self._producer = AIOKafkaProducer(
+            bootstrap_servers=[self._broker_url],
+            value_serializer=lambda v: json.dumps(v.model_dump()).encode("utf-8"),
+        )
+        logger.info(f"Connecting to Kafka broker {self._broker_url} and starting a producer...")
+        await self._producer.start()
+
+        for consumer in self._consumers.values():
+            await consumer.start()
+
+    async def _stop(self) -> None:
+        await self._producer.stop()
+        for consumer in self._consumers.values():
+            await consumer.stop()
+
+    async def _iter_messages(self, topic: str, consumer_group: str) -> AsyncIterable[Message]:
+        key = (topic, consumer_group)
+
+        if key not in self._consumers:
+            # create consumer
+            consumer = AIOKafkaConsumer(
+                topic,
+                bootstrap_servers=[self._broker_url],
+                auto_offset_reset="earliest",
+                group_id=consumer_group,
+                enable_auto_commit=False,
+                # consume only one message at a time
+                max_poll_records=1,
+                # one message should be consumed within 10 minutes
+                max_poll_interval_ms=10 * 60 * 1000,
+                session_timeout_ms=10 * 60 * 1000,
+                # send heartbeat every minute
+                heartbeat_interval_ms=1 * 60 * 1000,
+            )
+            await consumer.start()
+            self._consumers[key] = consumer
+            logger.info(
+                f"Connecting to Kafka broker {self._broker_url} and starting a consumer on "
+                f"topic {topic}."
+            )
+
+        consumer = self._consumers[key]
+
+        try:
+            async for message in consumer:
+                message_obj = json.loads(message.value)
+                yield Message(**message_obj)
+                try:
+                    await consumer.commit()
+                except CommitFailedError as e:
+                    logger.error(f"Commit failed: {e}... trying again.")
+        finally:
+            await consumer.stop()
+
+        # try:
+        #     while True:
+        #         # we use polling (instead of iterating through the consumer messages)
+        #         # to be able to cancel the consumer
+        #         messages = await self.kafka_consumer.getmany(timeout_ms=1000)
+
+        #         if messages:
+        #             for _, message_list in messages.items():
+        #                 for message in message_list:
+        #                     result = json.loads(message.value)
+        #                     logger.info(f"Received message on {message.topic}")
+
+        #                     try:
+        #                         for consumer in self.consumers:
+        #                             await consumer.consume(result)
+
+        #                         logger.info("Committing message")
+        #                         await self.kafka_consumer.commit()
+        #                     except Exception:
+        #                         logger.info("Rolling back message")
+        #                         logger.error(traceback.format_exc())
+        # except asyncio.CancelledError:
+        #     logger.info("Stopping ConsumeKafkaTopicLifespan")
+        #     await self.kafka_consumer.stop()
+        # except Exception as e:
+        #     logger.error(e)
+        #     logger.error(traceback.format_exc())
+
+    async def _send(self, topic: str, message: Message) -> None:
+        await self._producer.send_and_wait(
+            topic,
+            message,
+        )