Demo model service

openad/app/main.py

@@ -17,7 +17,7 @@
 from openad.app.main_lib import lang_parse, initialise, set_context, unset_context
 from openad.toolkit.toolkit_main import load_toolkit
 from openad.app import login_manager
-from openad.gui.gui_launcher import gui_init, GUI_SERVER, gui_shutdown
+from openad.gui.gui_launcher import gui_init, gui_shutdown
 from openad.gui.ws_server import ws_server  # Web socket server for gui - experimental
 from openad.helpers.output import output_table
 from openad.helpers.plugins import display_plugin_overview
@@ -38,6 +38,7 @@
 
 
 from openad.llm_assist.model_reference import SUPPORTED_TELL_ME_MODELS, SUPPORTED_TELL_ME_MODELS_SETTINGS
+from openad.openad_model_plugin.demo.launch_demo import terminate_model_service_demo
 
 # Helpers
 from openad.helpers.general import singular, confirm_prompt, get_case_insensitive_key
@@ -1207,9 +1208,17 @@ def cmd_line():
 
 
 def cleanup():
+    """
+    Cleanup function called on exit.
+    """
+    # Shut down GUI if it's running
     gui_shutdown(ignore_warning=True)
 
+    # Shut down model service demo if it's running
+    terminate_model_service_demo()
+
 
 atexit.register(cleanup)
+
 if __name__ == "__main__":
     cmd_line()

openad/app/main_lib.py

@@ -29,6 +29,7 @@
     detach_service_auth_group,
     list_auth_services,
     get_model_service_result,
+    model_service_demo,
 )
 
 # molecules
@@ -214,6 +215,8 @@ def lang_parse(cmd_pointer, parser):
         return detach_service_auth_group(cmd_pointer, parser)
     elif parser.getName() == "list_auth_services":
         return list_auth_services(cmd_pointer, parser)
+    elif parser.getName() == "model_service_demo":
+        return model_service_demo(cmd_pointer, parser)
 
     # @later -- move out all logic from here.
     # Language Model How To

openad/openad_model_plugin/auth_services.py

@@ -25,7 +25,7 @@ class LookupTable(TypedDict):
 
 def save_lookup_table(data: LookupTable):
     """save authentication lookup table to pickle file"""
-    logger.critical("saving auth lookup table")
+    logger.info("saving auth lookup table")
     with auth_lookup_lock:  # lock file to prevent concurrent access
         with open(auth_lookup_path, "wb") as file:
             pickle.dump(data, file)

openad/openad_model_plugin/catalog_model_services.py

@@ -27,6 +27,7 @@
 from openad.openad_model_plugin.config import DISPATCHER_SERVICE_PATH, SERVICE_MODEL_PATH, SERVICES_PATH
 from openad.openad_model_plugin.services import ModelService, UserProvidedConfig
 from openad.openad_model_plugin.utils import bcolors, get_logger
+from openad.openad_model_plugin.demo.launch_demo import launch_model_service_demo
 from pandas import DataFrame
 from tabulate import tabulate
 from tomlkit import parse
@@ -724,6 +725,16 @@ def get_model_service_result(cmd_pointer, parser):
     return result
 
 
+def model_service_demo(cmd_pointer, parser):
+    """
+    Spin up the model service demo in a subprocess.
+    """
+    restart = "restart" in parser.as_dict()
+    debug = "debug" in parser.as_dict()
+
+    return launch_model_service_demo(restart=restart, debug=debug)
+
+
 def service_catalog_grammar(statements: list, help: list):
     """This function creates the required grammar for managing cataloging services and model up or down"""
     logger.debug("catalog model service grammer")
@@ -1169,6 +1180,44 @@ def service_catalog_grammar(statements: list, help: list):
 Examples:
 - <cmd>get model service gen result 'xyz'</cmd>
 - <cmd>get model service 'my gen' result 'xyz'</cmd>
+""",
+        )
+    )
+
+    # ---
+    # Model service demo
+    statements.append(
+        py.Forward(
+            model
+            + service
+            + py.CaselessKeyword("demo")
+            + py.Optional(py.CaselessKeyword("restart")("restart") | py.CaselessKeyword("debug")("debug"))
+        )("model_service_demo")
+    )
+    help.append(
+        help_dict_create(
+            name="model service demo",
+            category="Model",
+            command="model service demo",
+            description="""Launch a demo service to learn about the OpenAD model service.
+
+Before you can run the demo service, you'll need to install the service tools:
+<cmd>pip install git+https://github.com/acceleratedscience/openad_service_utils.git@0.3.1</cmd>
+
+Further instructions are provided once the service is launched.
+It will shut down automatically when OpenAD is terminated.
+
+Optional clauses:
+<cmd>restart</cmd>
+    Reboot the service
+<cmd>debug</cmd>
+    Display the logs from the subprocess
+
+Examples:
+- <cmd>model service demo</cmd>
+- <cmd>model service demo restart</cmd>
+- <cmd>model service demo debug</cmd>
+
 """,
         )
     )

openad/openad_model_plugin/demo/launch_demo.py

@@ -0,0 +1,145 @@
+"""
+Launch and shutdown of the model service demo.
+See model_service_demo.py for more info and the actual demo service.
+"""
+
+import os
+import sys
+import threading
+import subprocess
+from openad.helpers.general import confirm_prompt
+from openad.helpers.output import output_error, output_text, output_success, output_warning
+
+DEMO_PROCESS = None
+
+
+def launch_model_service_demo(restart=False, debug=False):
+    """
+    Spin up the model service demo in a subprocess.
+    """
+
+    global DEMO_PROCESS
+
+    # Process already running
+    if DEMO_PROCESS:
+        # Try restart
+        if restart or debug:
+            success = terminate_model_service_demo()
+            if not success:
+                return
+
+        # Remind instructions
+        else:
+            return _print_success(new=False)
+
+    # Make sure openad_service_utils are installed
+    utils_installed = _verify_utils_installed()
+    if not utils_installed:
+        return
+
+    service_path = os.path.join(os.path.dirname(__file__), "model_service_demo.py")
+    command = [sys.executable, service_path]
+
+    try:
+        DEMO_PROCESS = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,  # Redirect stderr to stdout for combined logging
+            text=True,  # Decode output as text (Python 3.6+)
+            bufsize=1,  # Line-buffered output
+        )
+
+        # Log the subprocess' stdout
+        if debug:
+
+            def log_output():
+                for line in iter(DEMO_PROCESS.stdout.readline, ""):
+                    print(f"DEMO SERVICE: {line.strip()}")
+                DEMO_PROCESS.stdout.close()
+
+            # Start the logging thread
+            log_thread = threading.Thread(target=log_output, daemon=True)
+            log_thread.start()
+
+        # Success message
+        return _print_success()
+    except Exception as e:  # pylint: disable=broad-except
+        return output_error(f"Failed to start model service demo: {e}")
+
+
+def _verify_utils_installed():
+    """
+    Make sure openad_service_utils are installed.
+    """
+    try:
+        from openad_service_utils import start_server
+
+        return True
+    except ImportError:
+        msg = (
+            "Install openad_service_utils to use the demo model service:\n"
+            "<cmd>pip install git+https://github.com/acceleratedscience/openad_service_utils.git@0.3.1</cmd>"
+        )
+        output_warning(msg, return_val=False)
+        return False
+
+
+def _print_success(new=True):
+    """
+    Success message & instructions.
+    """
+    main_msg = (
+        (
+            "<success>Demo model service started at <yellow>http://localhost:8034</yellow></success>\n"
+            f"<soft>PID: {DEMO_PROCESS.pid}</soft>"
+        )
+        if new
+        else (
+            "<yellow>Demo model service already running at <reset>http://localhost:8034</reset></yellow>\n"
+            f"<soft>PID: {DEMO_PROCESS.pid} / To restart the demo service, run <cmd>model service demo restart</cmd></soft>"
+        )
+    )
+
+    msg = [
+        main_msg,
+        "",
+        "Next up, run:",
+        "<cmd>catalog model service from remote 'http://localhost:8034' as demo_service</cmd>",
+        "",
+        "To test the service:",
+        "<cmd>demo_service ?</cmd>",
+        "<cmd>demo_service get molecule property num_atoms for CC</cmd>",
+        "<cmd>demo_service get molecule property num_atoms for NCCc1c[nH]c2ccc(O)cc12</cmd>",
+    ]
+    return output_text("\n".join(msg), edge=True, pad=1)
+
+
+def terminate_model_service_demo():
+    """
+    Terminate the model service demo.
+    """
+    global DEMO_PROCESS
+    if DEMO_PROCESS is None:
+        return True
+
+    if DEMO_PROCESS:
+        try:
+            DEMO_PROCESS.terminate()
+            DEMO_PROCESS.wait(timeout=1)
+            output_success(f"Demo model service terminated - PID: {DEMO_PROCESS.pid}", return_val=False)
+            DEMO_PROCESS = None
+            return True
+        except Exception as err1:  # pylint: disable=broad-except
+            try:
+                # Force kill if terminate fails
+                DEMO_PROCESS.kill()
+                DEMO_PROCESS.wait(timeout=5)
+                output_success(f"Demo model service killed - PID: {DEMO_PROCESS.pid}", return_val=False)
+                DEMO_PROCESS = None
+                return True
+            except Exception as err2:  # pylint: disable=broad-except
+                output_error(
+                    [f"Failed to terminate model service demo with PID: {DEMO_PROCESS.pid}", err1, err2],
+                    return_val=False,
+                )
+                return False

openad/openad_model_plugin/demo/model_service_demo.py

@@ -0,0 +1,73 @@
+"""
+Model service demo used for tutorials and examples.
+
+To launch:
+    model service demo
+
+Model repo:
+    https://github.com/acceleratedscience/openad-service-demo
+"""
+
+import os
+from typing import Any
+from pydantic.v1 import Field
+from openad_service_utils import (
+    start_server,
+    SimplePredictor,
+    PredictorTypes,
+    DomainSubmodule,
+)
+
+
+# Model imports
+from rdkit import Chem
+
+
+class DemoPredictor(SimplePredictor):
+    """
+    Return the number of atoms in a molecule.
+    """
+
+    # fmt:off
+    domain: DomainSubmodule = DomainSubmodule("molecules")   # <-- edit here
+    algorithm_name: str = "rdkit"                            # <-- edit here
+    algorithm_application: str = "num_atoms"                 # <-- edit here
+    algorithm_version: str = "v0"
+    property_type: PredictorTypes = PredictorTypes.MOLECULE  # <-- edit here
+    # fmt:on
+
+    # User provided params for api / model inference
+    batch_size: int = Field(description="Prediction batch size", default=128)
+    workers: int = Field(description="Number of data loading workers", default=8)
+    device: str = Field(description="Device to be used for inference", default="cpu")
+
+    def setup(self):
+        """Model setup. Loads the model and tokenizer, if any. Runs once.
+
+        To wrap a model, copy and modify the standalone model setup and load
+        code here. Remember to change variables to instance variables, so they
+        can be used in the `predict` method.
+        """
+        self.model = None
+        self.tokenizer = []
+        self.model_path = os.path.join(self.get_model_location(), "model.ckpt")  # load model
+
+    def predict(self, sample: Any):
+        """
+        Run predictions.
+        """
+        # -----------------------User Code goes in here------------------------
+        smiles = sample
+        mol = Chem.MolFromSmiles(smiles)  # pylint: disable=no-member
+        num_atoms = mol.GetNumAtoms()
+        result = num_atoms
+        # ---------------------------------------------------------------------
+        return result
+
+
+# Register the class in global scope
+DemoPredictor.register(no_model=True)
+
+if __name__ == "__main__":
+    # Start the server
+    start_server(port=8034)