nerdd-link 0.1.0__py3-none-any.whl

nerdd_link/__init__.py

@@ -0,0 +1,4 @@
+from .actions import *
+from .channels 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 .write_output_action import *

nerdd_link/actions/action.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from stringcase import spinalcase
+
+from ..channels import Channel, Topic
+from ..types import Message
+
+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):
+            await self._process_message(message)
+
+    @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,59 @@
+import logging
+import os
+
+from nerdd_module import Model
+
+from ..channels import Channel
+from ..delegates import ReadCheckpointModel
+from ..types import CheckpointMessage
+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: Model, 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 input file to the job is stored in the file data_dir/job_id/input/
+        checkpoints_file = f"{self.data_dir}/jobs/{job_id}/input/checkpoint_{checkpoint_id}.pickle"
+        checkpoint_results_file = (
+            f"{self.data_dir}/jobs/{job_id}/results/checkpoint_{checkpoint_id}.pickle"
+        )
+
+        # create the results directory
+        os.makedirs(f"{self.data_dir}/jobs/{job_id}/results", exist_ok=True)
+
+        # create a model that reads the checkpoint file
+        model = ReadCheckpointModel(
+            base_model=self.model,
+            job_id=job_id,
+            checkpoint_id=checkpoint_id,
+            channel=self.channel,
+            checkpoints_file=checkpoints_file,
+            results_file=checkpoint_results_file,
+        )
+
+        # predict the checkpoint
+        model.predict(
+            input=None,
+            **params,
+        )
+
+    def _get_group_name(self) -> str:
+        model_name = self.model.__class__.__name__
+        return model_name

nerdd_link/actions/process_jobs_action.py

@@ -0,0 +1,138 @@
+import logging
+import os
+from pickle import dump
+
+from nerdd_module.input import DepthFirstExplorer
+from nerdd_module.model import ReadInputStep
+
+from ..channels import Channel
+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.data_dir = 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 the directory data_dir/sources/
+        # (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)
+        sources_dir = os.path.join(self.data_dir, "sources")
+
+        # 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=sources_dir,
+        )
+
+        read_input_step = ReadInputStep(explorer, message.source_id)
+
+        # create a directory for the job
+        os.makedirs(f"{self.data_dir}/jobs/{job_id}/input", exist_ok=True)
+
+        # 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
+        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 open(f"{self.data_dir}/jobs/{job_id}/input/checkpoint_{i}.pickle", "wb") as f:
+                dump(batch[:num_store], 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
+
+            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",
+                size=num_entries,
+            )
+        )

nerdd_link/actions/register_module_action.py

@@ -0,0 +1,30 @@
+import logging
+
+from nerdd_module import Model
+from stringcase import spinalcase
+
+from ..channels import Channel
+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):
+        super().__init__(channel.system_topic())
+        # TODO: do this differently
+        assert hasattr(model, "get_config")
+        self._model = model
+
+    async def _process_message(self, message: SystemMessage) -> None:
+        # send the initialization message
+        config = self._model.get_config()
+        logger.info(f"Send registration message for module {config.name}")
+        await self.channel.modules_topic().send(ModuleMessage(**config.model_dump()))
+
+    def _get_group_name(self) -> str:
+        model_name = spinalcase(self._model.__class__.__name__)
+        return model_name

nerdd_link/actions/write_output_action.py

@@ -0,0 +1,21 @@
+from nerdd_module import Model
+from stringcase import spinalcase
+
+from ..channels import Channel
+from ..types import SystemMessage
+from .action import Action
+
+__all__ = ["WriteOutputAction"]
+
+
+class WriteOutputAction(Action[SystemMessage]):
+    def __init__(self, channel: Channel, model: Model):
+        super().__init__(channel.system_topic())
+        self._model = model
+
+    async def _process_message(self, message: SystemMessage) -> None:
+        pass
+
+    def _get_group_name(self) -> str:
+        model_name = spinalcase(self._model.__class__.__name__)
+        return model_name

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,114 @@
+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  # type: ignore
+
+from ..types import (
+    CheckpointMessage,
+    JobMessage,
+    LogMessage,
+    Message,
+    ModuleMessage,
+    ResultCheckpointMessage,
+    ResultMessage,
+    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
+        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):
+    #
+    # RECEIVE
+    #
+    async def iter_messages(self, topic: str, consumer_group: str) -> AsyncIterable[Message]:
+        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:
+        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, job_type_or_model: Union[str, Model]
+    ) -> Topic[ResultCheckpointMessage]:
+        job_type = get_job_type(job_type_or_model)
+        topic_name = f"{job_type}-result-checkpoints"
+        return Topic[ResultCheckpointMessage](self, topic_name)
+
+    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,97 @@
+import asyncio
+import json
+import logging
+from typing import AsyncIterable, Dict, Tuple
+
+from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
+
+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] = {}
+
+        self._producer = AIOKafkaProducer(
+            bootstrap_servers=[self._broker_url],
+        )
+        # TODO: check value_serializer
+        # producer = AIOKafkaProducer(
+        #     bootstrap_servers=KAFKA_BROKER_URL,
+        #     value_serializer=lambda v: json.dumps(v).encode("utf-8"),
+        # )
+        asyncio.create_task(self._producer.start())
+        logger.info(f"Connecting to Kafka broker {self._broker_url} and starting a producer.")
+
+    async def _iter_messages(self, topic: str, consumer_group: str) -> AsyncIterable[Message]:
+        if consumer_group is not None:
+            consumer_group = f"{consumer_group}-consumer-group"
+
+        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,
+            )
+            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)
+                await consumer.commit()
+        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,
+            json.dumps(message.model_dump()).encode("utf-8"),
+        )

nerdd_link/channels/memory_channel.py

@@ -0,0 +1,33 @@
+import logging
+from typing import AsyncIterable, List, Tuple
+
+from ..types import Message
+from ..utils import ObservableList
+from .channel import Channel
+
+__all__ = ["MemoryChannel"]
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryChannel(Channel):
+    def __init__(self) -> None:
+        super().__init__()
+        self._messages = ObservableList[Tuple[str, Message]]()
+
+    def get_produced_messages(self) -> List[Tuple[str, Message]]:
+        return self._messages.get_items()
+
+    async def _iter_messages(self, topic: str, consumer_group: str) -> AsyncIterable[Message]:
+        async for _, new in self._messages.changes():
+            assert new is not None
+            (t, message) = new
+            if topic == t:
+                yield message
+
+    async def _send(self, topic: str, message: Message) -> None:
+        logger.info(f"Send message to topic {topic}")
+        self._messages.append((topic, message))
+
+    async def stop(self) -> None:
+        await self._messages.stop()

nerdd_link/cli/__init__.py

@@ -0,0 +1,3 @@
+from .initialize_system import *
+from .run_job_server import *
+from .run_prediction_server import *

nerdd_link/cli/initialize_system.py

@@ -0,0 +1,45 @@
+import logging
+
+import rich_click as click
+
+from ..channels import KafkaChannel
+from ..types import SystemMessage
+from ..utils import async_to_sync
+
+__all__ = ["initialize_system"]
+
+logger = logging.getLogger(__name__)
+
+
+@click.command(context_settings={"show_default": True})
+@click.option(
+    "--channel",
+    type=click.Choice(["kafka"], case_sensitive=False),
+    default="kafka",
+    help="Channel to use for communication with the model.",
+)
+@click.option("--broker-url", default="localhost:9092", help="Kafka broker to connect to.")
+@click.option(
+    "--log-level",
+    default="info",
+    type=click.Choice(["debug", "info", "warning", "error", "critical"], case_sensitive=False),
+    help="The logging level.",
+)
+@async_to_sync
+async def initialize_system(
+    # communication options
+    channel: str,
+    broker_url: str,
+    # log level
+    log_level: str,
+) -> None:
+    logging.basicConfig(level=log_level.upper())
+
+    channel_instance = None
+    if channel == "kafka":
+        channel_instance = KafkaChannel(broker_url)
+    else:
+        raise ValueError(f"Channel {channel} not supported.")
+
+    logging.info("Sending the system initialization message...")
+    await channel_instance.system_topic().send(SystemMessage())