nerdd-link 0.2.2__tar.gz → 0.2.4__tar.gz

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nerdd-link
-Version: 0.2.2
+Version: 0.2.4
 Summary: Run a NERDD module as a service
 Author-email: Steffen Hirte <steffen.hirte@univie.ac.at>
 Maintainer-email: Steffen Hirte <steffen.hirte@univie.ac.at>

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/actions/predict_checkpoints_action.py

@@ -49,5 +49,5 @@ class PredictCheckpointsAction(Action[CheckpointMessage]):
         )
 
     def _get_group_name(self) -> str:
-        model_name = self._model.__class__.__name__
-        return model_name
+        model_id = self._model.get_config().id
+        return f"predict-checkpoints-{model_id}"

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/actions/register_module_action.py

@@ -1,9 +1,10 @@
+import json
 import logging
 
 from nerdd_module import Model
-from stringcase import spinalcase
 
 from ..channels import Channel
+from ..files import FileSystem
 from ..types import ModuleMessage, SystemMessage
 from .action import Action
 
@@ -13,18 +14,24 @@ logger = logging.getLogger(__name__)
 
 
 class RegisterModuleAction(Action[SystemMessage]):
-    def __init__(self, channel: Channel, model: Model):
+    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:
-        # 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()))
+        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_name = spinalcase(self._model.__class__.__name__)
-        return model_name
+        model_id = self._model.get_config().id
+        return f"register-module-{model_id}"

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/actions/serialize_job_action.py

@@ -1,10 +1,10 @@
+import json
 import logging
 
-from nerdd_module import Model, OutputStep
-from stringcase import spinalcase
+from nerdd_module import OutputStep
 
 from ..channels import Channel
-from ..delegates import ReadPickleStep
+from ..delegates import ReadPickleStep, SerializeJobModel
 from ..files import FileSystem
 from ..types import SerializationRequestMessage, SerializationResultMessage
 from .action import Action
@@ -16,13 +16,13 @@ logger = logging.getLogger(__name__)
 
 
 class SerializeJobAction(Action[SerializationRequestMessage]):
-    def __init__(self, channel: Channel, model: Model, data_dir: str) -> None:
-        super().__init__(channel.serialization_requests_topic(model))
-        self._model = model
+    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}")
@@ -31,15 +31,23 @@ class SerializeJobAction(Action[SerializationRequestMessage]):
         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)
 
-        # TODO: don't write the file if it exists
+        # 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),
+        ]
 
-        read_pickle_step = ReadPickleStep(self._file_system.iter_results_file_handles(job_id))
-        post_processing_steps = self._model._get_postprocessing_steps(
-            output_format, output_file=output_file, **params
-        )
-        steps = [read_pickle_step, *post_processing_steps]
         output_step = steps[-1]
         assert isinstance(output_step, OutputStep), "The last step must be an OutputStep."
 
@@ -54,7 +62,3 @@ class SerializeJobAction(Action[SerializationRequestMessage]):
         await self.channel.serialization_results_topic().send(
             SerializationResultMessage(job_id=job_id, output_format=output_format)
         )
-
-    def _get_group_name(self) -> str:
-        model_name = spinalcase(self._model.__class__.__name__)
-        return model_name

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/channels/channel.py

@@ -32,6 +32,7 @@ def get_job_type(job_type_or_model: Union[str, Model]) -> str:
         # * 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 == "-"])
@@ -133,12 +134,8 @@ class Channel(ABC):
     def result_checkpoints_topic(self) -> Topic[ResultCheckpointMessage]:
         return Topic[ResultCheckpointMessage](self, "result-checkpoints")
 
-    def serialization_requests_topic(
-        self, job_type_or_model: Union[str, Model]
-    ) -> Topic[SerializationRequestMessage]:
-        job_type = get_job_type(job_type_or_model)
-        topic_name = f"{job_type}-serialization-requests"
-        return Topic[SerializationRequestMessage](self, topic_name)
+    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")

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/cli/initialize_system.py

@@ -41,5 +41,9 @@ async def initialize_system(
     else:
         raise ValueError(f"Channel {channel} not supported.")
 
+    await channel_instance.start()
+
     logging.info("Sending the system initialization message...")
     await channel_instance.system_topic().send(SystemMessage())
+
+    await channel_instance.stop()

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/cli/run_job_server.py

@@ -19,7 +19,11 @@ logger = logging.getLogger(__name__)
     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(
+    "--broker-url",
+    default="localhost:9092",
+    help="Broker url to connect to.",
+)
 @click.option(
     "--max-num-molecules",
     default=10_000,
@@ -86,6 +90,8 @@ async def run_job_server(
     else:
         raise ValueError(f"Channel {channel} not supported.")
 
+    await channel_instance.start()
+
     action = ProcessJobsAction(
         channel_instance,
         checkpoint_size,
@@ -105,3 +111,5 @@ async def run_job_server(
         logger.info("Shutting down server")
         task.cancel()
         await task
+
+        await channel_instance.stop()

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/cli/run_prediction_server.py

@@ -53,6 +53,8 @@ async def run_prediction_server(
     else:
         raise ValueError(f"Channel {channel} not supported.")
 
+    await channel_instance.start()
+
     # import the model class
     package_name, class_name = model_name.rsplit(".", 1)
     package = import_module(package_name)
@@ -79,3 +81,5 @@ async def run_prediction_server(
         for task in tasks:
             task.cancel()
         await asyncio.gather(*tasks, return_exceptions=True)
+
+        await channel_instance.stop()

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/cli/run_serialization_server.py

@@ -1,6 +1,5 @@
 import asyncio
 import logging
-from importlib import import_module
 from typing import List
 
 import rich_click as click
@@ -15,7 +14,6 @@ logger = logging.getLogger(__name__)
 
 
 @click.command(context_settings={"show_default": True})
-@click.argument("model-name")
 @click.option(
     "--channel",
     type=click.Choice(["kafka"], case_sensitive=False),
@@ -40,7 +38,6 @@ async def run_serialization_server(
     channel: str,
     broker_url: str,
     # options
-    model_name: str,
     data_dir: str,
     # log level
     log_level: str,
@@ -53,15 +50,10 @@ async def run_serialization_server(
     else:
         raise ValueError(f"Channel {channel} not supported.")
 
-    # import the model class
-    package_name, class_name = model_name.rsplit(".", 1)
-    package = import_module(package_name)
-    Model = getattr(package, class_name)
-    model = Model()
+    await channel_instance.start()
 
     serialize_job = SerializeJobAction(
         channel=channel_instance,
-        model=model,
         data_dir=data_dir,
     )
 
@@ -77,3 +69,5 @@ async def run_serialization_server(
         for task in tasks:
             task.cancel()
         await asyncio.gather(*tasks, return_exceptions=True)
+
+        await channel_instance.stop()

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/converters/__init__.py

@@ -1,2 +1,3 @@
 from .mol_pickle_converter import *
+from .mol_to_image_converter import *
 from .pickle_converter import *

nerdd_link-0.2.4/nerdd_link/converters/mol_to_image_converter.py

@@ -0,0 +1,59 @@
+from typing import Any
+from xml.dom import minidom
+
+from nerdd_module import Converter, ConverterConfig
+from rdkit.Chem import Mol
+from rdkit.Chem.Draw import MolDraw2DSVG
+
+__all__ = ["MolToImageConverter"]
+
+default_width = 300
+default_height = 180
+
+
+class MolToImageConverter(Converter):
+    def _convert(self, input: Any, context: dict) -> Any:
+        width = self.result_property.image_width
+        height = self.result_property.image_height
+
+        if width is None:
+            width = default_width
+        if height is None:
+            height = default_height
+
+        mol = input
+        if mol is None:
+            return None
+
+        assert isinstance(mol, Mol), f"Expected RDKit Mol object, but got {type(mol)}"
+
+        svg = MolDraw2DSVG(width, height)
+
+        # remove background
+        opts = svg.drawOptions()
+        opts.clearBackground = False
+
+        # add highlight circles around atoms during drawing
+        # (we will remove them later in post processing)
+        atoms = range(mol.GetNumAtoms())
+        colors = [[(0.8, 1, 1)]] * mol.GetNumAtoms()
+        radii = [0.5] * mol.GetNumAtoms()
+        atom_highlight = dict(zip(atoms, colors))
+        atom_radii = dict(zip(atoms, radii))
+        svg.DrawMoleculeWithHighlights(mol, "", atom_highlight, {}, atom_radii, [])
+        svg.FinishDrawing()
+
+        # post process SVG
+        xml = svg.GetDrawingText()
+        tree = minidom.parseString(xml)
+        root = tree.getElementsByTagName("svg")[0]
+
+        # make highlight circles invisible
+        for ellipse in root.getElementsByTagName("ellipse"):
+            ellipse.setAttribute("style", "opacity:0")
+
+        xml = tree.toxml()
+
+        return xml
+
+    config = ConverterConfig(data_types="mol", output_formats="json")

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/delegates/__init__.py

@@ -1,4 +1,5 @@
 from .pickle_writer import *
 from .read_checkpoint_model import *
 from .read_pickle_step import *
+from .serialize_job_model import *
 from .topic_writer import *

nerdd_link-0.2.4/nerdd_link/delegates/serialize_job_model.py

@@ -0,0 +1,19 @@
+from typing import List
+
+from nerdd_module import SimpleModel
+from nerdd_module.config import Configuration, DictConfiguration
+from rdkit.Chem import Mol
+
+__all__ = ["SerializeJobModel"]
+
+
+class SerializeJobModel(SimpleModel):
+    def __init__(self, config: dict) -> None:
+        super().__init__()
+        self._config = config
+
+    def _get_config(self) -> Configuration:
+        return DictConfiguration(self._config)
+
+    def _predict_mols(self, mols: List[Mol]) -> List[dict]:
+        return []

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/files/file_system.py

@@ -16,6 +16,11 @@ class FileSystem:
     #
     # DIRECTORIES
     #
+    def get_modules_dir(self) -> str:
+        result = os.path.join(self.root_path, "modules")
+        os.makedirs(result, exist_ok=True)
+        return result
+
     def get_sources_dir(self) -> str:
         result = os.path.join(self.root_path, "sources")
         os.makedirs(result, exist_ok=True)
@@ -49,6 +54,9 @@ class FileSystem:
     #
     # FILES
     #
+    def get_module_file_path(self, module_id: str) -> str:
+        return os.path.join(self.get_modules_dir(), module_id)
+
     def get_source_file_path(self, source_id: str) -> str:
         return os.path.join(self.get_sources_dir(), source_id)
 

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link/types/__init__.py

@@ -20,7 +20,7 @@ class Message(BaseModel):
 
 
 class ModuleMessage(Message):
-    name: str
+    id: str
 
 
 class CheckpointMessage(Message):
@@ -43,6 +43,7 @@ class JobMessage(Message):
 
 class SerializationRequestMessage(Message):
     job_id: str
+    job_type: str
     params: Dict[str, Any]
     output_format: str
 
@@ -55,6 +56,8 @@ class SerializationResultMessage(Message):
 class ResultMessage(Message):
     job_id: str
 
+    model_config = ConfigDict(extra="allow")
+
 
 class LogMessage(Message):
     job_id: str

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nerdd-link
-Version: 0.2.2
+Version: 0.2.4
 Summary: Run a NERDD module as a service
 Author-email: Steffen Hirte <steffen.hirte@univie.ac.at>
 Maintainer-email: Steffen Hirte <steffen.hirte@univie.ac.at>

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/nerdd_link.egg-info/SOURCES.txt

@@ -27,11 +27,13 @@ nerdd_link/cli/run_prediction_server.py
 nerdd_link/cli/run_serialization_server.py
 nerdd_link/converters/__init__.py
 nerdd_link/converters/mol_pickle_converter.py
+nerdd_link/converters/mol_to_image_converter.py
 nerdd_link/converters/pickle_converter.py
 nerdd_link/delegates/__init__.py
 nerdd_link/delegates/pickle_writer.py
 nerdd_link/delegates/read_checkpoint_model.py
 nerdd_link/delegates/read_pickle_step.py
+nerdd_link/delegates/serialize_job_model.py
 nerdd_link/delegates/split_and_merge_step.py
 nerdd_link/delegates/topic_writer.py
 nerdd_link/files/__init__.py

{nerdd_link-0.2.2 → nerdd_link-0.2.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "nerdd-link"
-version = "0.2.2"
+version = "0.2.4"
 description = "Run a NERDD module as a service"
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -86,7 +86,7 @@ nerdd_link = ["py.typed"]
 
 [tool.pytest.ini_options]
 log_cli = 1
-log_cli_level = "INFO"
+log_cli_level = "ERROR"
 addopts = "-x --cov-report term --cov=nerdd_link"
 
 [tool.pytest-watcher]