diff --git a/openad/app/main.py b/openad/app/main.py
index 946d58d8..f0e30bb3 100644
--- a/openad/app/main.py
+++ b/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()
diff --git a/openad/app/main_lib.py b/openad/app/main_lib.py
index 52ce19c4..77879c5c 100644
--- a/openad/app/main_lib.py
+++ b/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
diff --git a/openad/openad_model_plugin/auth_services.py b/openad/openad_model_plugin/auth_services.py
index 30bbd06c..9daeb4ff 100644
--- a/openad/openad_model_plugin/auth_services.py
+++ b/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)
diff --git a/openad/openad_model_plugin/catalog_model_services.py b/openad/openad_model_plugin/catalog_model_services.py
index b3739772..b35f2654 100644
--- a/openad/openad_model_plugin/catalog_model_services.py
+++ b/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:
- get model service gen result 'xyz'
- get model service 'my gen' result 'xyz'
+""",
+ )
+ )
+
+ # ---
+ # 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:
+pip install git+https://github.com/acceleratedscience/openad_service_utils.git@0.3.1
+
+Further instructions are provided once the service is launched.
+It will shut down automatically when OpenAD is terminated.
+
+Optional clauses:
+restart
+ Reboot the service
+debug
+ Display the logs from the subprocess
+
+Examples:
+- model service demo
+- model service demo restart
+- model service demo debug
+
""",
)
)
diff --git a/openad/openad_model_plugin/demo/launch_demo.py b/openad/openad_model_plugin/demo/launch_demo.py
new file mode 100644
index 00000000..0d726b19
--- /dev/null
+++ b/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"
+ "pip install git+https://github.com/acceleratedscience/openad_service_utils.git@0.3.1"
+ )
+ output_warning(msg, return_val=False)
+ return False
+
+
+def _print_success(new=True):
+ """
+ Success message & instructions.
+ """
+ main_msg = (
+ (
+ "Demo model service started at http://localhost:8034\n"
+ f"PID: {DEMO_PROCESS.pid}"
+ )
+ if new
+ else (
+ "Demo model service already running at http://localhost:8034\n"
+ f"PID: {DEMO_PROCESS.pid} / To restart the demo service, run model service demo restart"
+ )
+ )
+
+ msg = [
+ main_msg,
+ "",
+ "Next up, run:",
+ "catalog model service from remote 'http://localhost:8034' as demo_service",
+ "",
+ "To test the service:",
+ "demo_service ?",
+ "demo_service get molecule property num_atoms for CC",
+ "demo_service get molecule property num_atoms for NCCc1c[nH]c2ccc(O)cc12",
+ ]
+ 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
diff --git a/openad/openad_model_plugin/demo/model_service_demo.py b/openad/openad_model_plugin/demo/model_service_demo.py
new file mode 100644
index 00000000..204792f2
--- /dev/null
+++ b/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)