Merge pull request #95 from AllenNeuralDynamics:feat-executable-service

Add App service and BonsaiApp concrete implementation

src/aind_behavior_services/launcher/_service.py

@@ -9,7 +9,5 @@ def __init__(self, logger: Optional[logging.Logger], *args, **kwargs) -> None: .
 
     def validate(self, *args, **kwargs) -> bool: ...
 
-    def register(self, *args, **kwargs) -> None: ...
-
     @property
     def logger(self) -> Optional[logging.Logger]: ...

src/aind_behavior_services/launcher/apps.py

@@ -0,0 +1,154 @@
+import glob
+import logging
+import os
+import subprocess
+from pathlib import Path
+from typing import Dict, Optional, Self
+
+from aind_behavior_services.launcher._service import IService
+from aind_behavior_services.launcher.ui_helper import UIHelper
+from aind_behavior_services.utils import run_bonsai_process
+
+
+class App(IService):
+    def __init__(self, logger: logging.Logger, *args, **kwargs) -> None:
+        self._logger = logger
+
+    def validate(self, *args, **kwargs) -> bool:
+        raise NotImplementedError
+
+    @property
+    def logger(self) -> logging.Logger:
+        return self._logger
+
+
+class BonsaiApp(App):
+    executable: os.PathLike
+    workflow: os.PathLike
+    is_editor_mode: bool
+    is_start_flag: bool
+    layout: Optional[os.PathLike | str]
+    layout_directory: Optional[os.PathLike]
+    additional_properties: Optional[Dict[str, str]]
+    cwd: Optional[os.PathLike]
+    timeout: Optional[float]
+    print_cmd: bool
+    _result: Optional[subprocess.CompletedProcess]
+    _ui_helper: UIHelper
+
+    def __init__(
+        self,
+        logger: logging.Logger,
+        workflow: os.PathLike,
+        executable: os.PathLike = Path("./bonsai/bonsai.exe"),
+        /,
+        is_editor_mode: bool = True,
+        is_start_flag: bool = True,
+        layout: Optional[os.PathLike] = None,
+        layout_dir: Optional[os.PathLike] = None,
+        additional_properties: Optional[Dict[str, str]] = None,
+        cwd: Optional[os.PathLike] = None,
+        timeout: Optional[float] = None,
+        ui_helper: Optional[UIHelper] = None,
+    ) -> None:
+        super().__init__(logger)
+        self.executable = Path(executable).resolve()
+        self.workflow = Path(workflow).resolve()
+        self.is_editor_mode = is_editor_mode
+        self.is_start_flag = is_start_flag if not self.is_editor_mode else True
+        self.layout = layout
+        self.layout_directory = layout_dir
+        self.additional_properties = additional_properties or {}
+        self.cwd = cwd
+        self.timeout = timeout
+        self._result = None
+        self._ui_helper = ui_helper or UIHelper(logger, print)
+
+    @property
+    def result(self) -> subprocess.CompletedProcess:
+        if self._result is None:
+            raise RuntimeError("The app has not been run yet.")
+        return self._result
+
+    def validate(self, *args, **kwargs) -> bool:
+        if not Path(self.executable).exists():
+            raise FileNotFoundError(f"Executable not found: {self.executable}")
+        if not Path(self.workflow).exists():
+            raise FileNotFoundError(f"Workflow file not found: {self.workflow}")
+        if self.layout and not Path(self.layout).exists():
+            raise FileNotFoundError(f"Layout file not found: {self.layout}")
+        if self.layout_directory and not Path(self.layout_directory).exists():
+            raise FileNotFoundError(f"Layout directory not found: {self.layout_directory}")
+        return True
+
+    def run(self) -> subprocess.CompletedProcess:
+        self.validate()
+
+        if self.is_editor_mode:
+            self.logger.warning("Bonsai is running in editor mode. Cannot assert successful completion.")
+        self.logger.info("Bonsai process running...")
+        proc = run_bonsai_process(
+            workflow_file=self.workflow,
+            bonsai_exe=self.executable,
+            is_editor_mode=self.is_editor_mode,
+            is_start_flag=self.is_start_flag,
+            layout=self.layout,
+            additional_properties=self.additional_properties,
+            cwd=self.cwd,
+            timeout=self.timeout,
+            print_cmd=self.print_cmd,
+        )
+        self._result = proc
+        self.logger.info("Bonsai process completed.")
+        return proc
+
+    def output_from_result(self, allow_stderr: Optional[bool]) -> Self:
+        proc = self.result
+        try:
+            proc.check_returncode()
+        except subprocess.CalledProcessError as e:
+            self._log_process_std_output("Bonsai", proc)
+            raise e
+        else:
+            self.logger.info("Result from bonsai process is valid.")
+            self._log_process_std_output("Bonsai", proc)
+
+            if len(proc.stdout) > 0:
+                self.logger.error("Bonsai process finished with errors.")
+                if allow_stderr is None:
+                    allow_stderr = self._ui_helper.prompt_yes_no_question("Would you like to see the error message?")
+                if allow_stderr is False:
+                    raise subprocess.CalledProcessError(1, proc.args)
+        return self
+
+    def prompt_visualizer_layout_input(
+        self,
+        directory: Optional[os.PathLike] = None,
+    ) -> Optional[str | os.PathLike]:
+        # This could use some refactoring. The bonsai CLI logic is:
+        # 1. If a layout is provided, use that.
+        # 2. If a layout is not provided, use the default layout
+        # 3. if the layout is passed as "" (empty string) no layout is used.
+        layout_schemas_path = directory if directory is not None else self.layout_directory
+        available_layouts = glob.glob(os.path.join(str(layout_schemas_path), "*.bonsai.layout"))
+        picked: Optional[str | os.PathLike] = None
+        has_pick = False
+        while has_pick is False:
+            try:
+                available_layouts.insert(0, "None")
+                picked = self._ui_helper.prompt_pick_file_from_list(
+                    available_layouts, prompt="Pick a visualizer layout:", override_zero=("Default", None)
+                )
+                if picked == "None":
+                    picked = ""
+            except ValueError as e:
+                self.logger.error("Invalid choice. Try again. %s", e)
+            has_pick = True
+        self.layout = picked
+        return self.layout
+
+    def _log_process_std_output(self, process_name: str, proc: subprocess.CompletedProcess) -> None:
+        if len(proc.stdout) > 0:
+            self.logger.info("%s full stdout dump: \n%s", process_name, proc.stdout)
+        if len(proc.stderr) > 0:
+            self.logger.error("%s full stderr dump: \n%s", process_name, proc.stderr)

src/aind_behavior_services/launcher/resource_manager_service.py → src/aind_behavior_services/launcher/resource_monitor_service.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
 import logging
+import os
 import shutil
 from dataclasses import dataclass, field
 from typing import Callable, List, Optional
-import os
 
 from aind_behavior_services.launcher._service import IService
 
 
-class ResourceManager(IService):
+class ResourceMonitor(IService):
     def __init__(
         self,
         logger: logging.Logger,
@@ -25,9 +25,6 @@ def logger(self) -> logging.Logger:
     def validate(self, *args, **kwargs) -> bool:
         return True
 
-    def register(self, *args, **kwargs) -> None:
-        pass
-
     def add_constraint(self, constraint: Constraint) -> None:
         self.constraints.append(constraint)
 
@@ -78,4 +75,3 @@ def remote_dir_exists_constraint_factory(dir_path: os.PathLike) -> Constraint:
         kwargs={"dir_path": dir_path},
         fail_msg_handler=lambda dir_path: f"Directory {dir_path} does not exist.",
     )
-

src/aind_behavior_services/launcher/services.py

@@ -1,22 +1,31 @@
 import logging
 from typing import Optional, Self, TypeVar, Union
 
-import aind_behavior_services.launcher.resource_manager_service as resource_manager_service
+import aind_behavior_services.launcher.apps as apps
+import aind_behavior_services.launcher.resource_monitor_service as resource_monitor_service
 import aind_behavior_services.launcher.watchdog_service as watchdog_service
 
-SupportedServices = Union[watchdog_service.WatchdogService, resource_manager_service.ResourceManager]
+SupportedServices = Union[watchdog_service.WatchdogService, resource_monitor_service.ResourceMonitor, apps.BonsaiApp]
 TService = TypeVar("TService", bound=SupportedServices)
 
 
 class Services:
     _watchdog: Optional[watchdog_service.WatchdogService]
     _logger: Optional[logging.Logger]
-    _resource_manager: Optional[resource_manager_service.ResourceManager]
+    _resource_monitor: Optional[resource_monitor_service.ResourceMonitor]
+    _app: Optional[apps.BonsaiApp]
 
-    def __init__(self, logger: Optional[logging.Logger] = None):
-        self._watchdog = None
+    def __init__(
+        self,
+        logger: Optional[logging.Logger] = None,
+        watchdog: Optional[watchdog_service.WatchdogService] = None,
+        resource_monitor: Optional[resource_monitor_service.ResourceMonitor] = None,
+        app: Optional[apps.BonsaiApp] = None,
+    ) -> None:
         self._logger = logger
-        self._resource_manager = None
+        self._watchdog = watchdog
+        self._resource_monitor = resource_monitor
+        self._app = app
 
     @property
     def logger(self) -> logging.Logger:
@@ -37,27 +46,28 @@ def watchdog(self) -> Optional[watchdog_service.WatchdogService]:
     def register_watchdog(self, watchdog: watchdog_service.WatchdogService) -> Self:
         if self._watchdog is not None:
             raise ValueError("Watchdog already registered")
-        watchdog.register()
         self._watchdog = watchdog
         return self
 
     @property
-    def resource_manager(self) -> Optional[resource_manager_service.ResourceManager]:
-        return self._resource_manager
+    def resource_monitor(self) -> Optional[resource_monitor_service.ResourceMonitor]:
+        return self._resource_monitor
 
-    def register_resource_manager(self, resource_manager: resource_manager_service.ResourceManager) -> Self:
-        if self._resource_manager is not None:
+    def register_resource_monitor(self, resource_monitor: resource_monitor_service.ResourceMonitor) -> Self:
+        if self._resource_monitor is not None:
             raise ValueError("Resource manager already registered")
-        self._resource_manager = resource_manager
+        self._resource_monitor = resource_monitor
         return self
 
-    def register(self, service: TService) -> Self:
-        if isinstance(service, watchdog_service.WatchdogService):
-            return self.register_watchdog(service)
-        elif isinstance(service, resource_manager_service.ResourceManager):
-            return self.register_resource_manager(service)
-        else:
-            raise ValueError(f"Unsupported service: {service}")
+    @property
+    def app(self) -> Optional[apps.BonsaiApp]:
+        return self._app
+
+    def register_app(self, app: apps.BonsaiApp) -> Self:
+        if self._app is not None:
+            raise ValueError("App already registered")
+        self._app = app
+        return self
 
     def validate_service(self, obj: Optional[TService]) -> bool:
         if obj is None:

src/aind_behavior_services/launcher/ui_helper.py

@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import Callable, Optional, Tuple
+from typing import Any, Callable, Optional, Tuple, TypeVar
 
 from aind_behavior_services.db_utils import SubjectDataBase
 from aind_behavior_services.rig import AindBehaviorRigModel
@@ -16,28 +16,33 @@ def __init__(self, logger: logging.Logger, print_func: Callable[[str], None] = p
         self._print = print_func
         self._logger = logger
 
+    T = TypeVar("T", bound=Any)
+
     def prompt_pick_file_from_list(
         self,
         available_files: list[str],
         prompt: str = "Choose a file:",
-        override_zero: Tuple[Optional[str], Optional[str]] = ("Enter manually", None),
-    ) -> str:
+        zero_label: Optional[str] = None,
+        zero_value: Optional[T] = None,
+        zero_as_input: bool = True,
+        override_zero: Tuple[Optional[str], Optional[Any]] = ("Enter manually", None),
+    ) -> Optional[str | T]:
         self._print(prompt)
-        if override_zero[0] is not None:
-            self._print(f"0: {override_zero[0]}")
+        if zero_label is not None:
+            self._print(f"0: {zero_label}")
         for i, file in enumerate(available_files):
             self._print(f"{i+1}: {os.path.split(file)[1]}")
         choice = int(input("Choice: "))
         if choice < 0 or choice >= len(available_files) + 1:
             raise ValueError
         if choice == 0:
-            if override_zero[0] is None:
+            if zero_label is None:
                 raise ValueError
-            if override_zero[1] is not None:
-                return override_zero[1]
             else:
-                path = str(input(override_zero[0]))
-            return path
+                if zero_as_input:
+                    return str(input(override_zero[0]))
+                else:
+                    return zero_value
         else:
             return available_files[choice - 1]
 
@@ -56,8 +61,10 @@ def choose_subject(self, subject_list: SubjectDataBase) -> str:
         while subject is None:
             try:
                 subject = self.prompt_pick_file_from_list(
-                    list(subject_list.subjects.keys()), prompt="Choose a subject:", override_zero=(None, None)
+                    list(subject_list.subjects.keys()), prompt="Choose a subject:", zero_label=None
                 )
+                if not isinstance(subject, str):
+                    raise ValueError("Return value is not a string type.")
             except ValueError as e:
                 self._logger.error("Invalid choice. Try again. %s", e)
         return subject