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)