Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion examples/_mocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,4 +97,4 @@ def create_fake_rig():
computer_name = os.getenv("COMPUTERNAME")
os.makedirs(_dir := f"{LIB_CONFIG}/Rig/{computer_name}", exist_ok=True)
with open(f"{_dir}/rig1.json", "w", encoding="utf-8") as f:
f.write(RigModel().model_dump_json(indent=2))
f.write(RigModel(data_directory=r"./local/data").model_dump_json(indent=2))
86 changes: 86 additions & 0 deletions examples/client_behavior_launcher.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
import logging
from pathlib import Path

from _mocks import (
LIB_CONFIG,
AindBehaviorSessionModel,
RigModel,
TaskLogicModel,
create_fake_rig,
create_fake_subjects,
)
from pydantic_settings import CliApp

from clabe import resource_monitor
from clabe.apps import BonsaiApp
from clabe.launcher import Launcher, LauncherCliArgs, experiment
from clabe.pickers import DefaultBehaviorPicker, DefaultBehaviorPickerSettings
from clabe.xml_rpc import XmlRpcClient, XmlRpcClientSettings

logger = logging.getLogger(__name__)


@experiment()
async def client_experiment(launcher: Launcher) -> None:
"""Demo experiment showcasing CLABE functionality."""
picker = DefaultBehaviorPicker(
launcher=launcher,
settings=DefaultBehaviorPickerSettings(config_library_dir=LIB_CONFIG),
experimenter_validator=lambda _: True,
)

session = picker.pick_session(AindBehaviorSessionModel)
rig = picker.pick_rig(RigModel)
launcher.register_session(session, rig.data_directory)
trainer_state, task_logic = picker.pick_trainer_state(TaskLogicModel)

resource_monitor.ResourceMonitor(
constrains=[
resource_monitor.available_storage_constraint_factory_from_rig(rig, 2e11),
]
).run()

xml_rpc_client = XmlRpcClient(settings=XmlRpcClientSettings(server_url="http://localhost:8000", token="42"))

bonsai_root = Path(r"C:\git\AllenNeuralDynamics\Aind.Behavior.VrForaging")
session_response = xml_rpc_client.upload_model(session, "session.json")
rig_response = xml_rpc_client.upload_model(rig, "rig.json")
task_logic_response = xml_rpc_client.upload_model(task_logic, "task_logic.json")
assert rig_response.path is not None
assert session_response.path is not None
assert task_logic_response.path is not None

bonsai_app_result = await xml_rpc_client.run_async(
BonsaiApp(
workflow=bonsai_root / "src/test_deserialization.bonsai",
executable=bonsai_root / "bonsai/bonsai.exe",
additional_externalized_properties={
"RigPath": rig_response.path,
"SessionPath": session_response.path,
"TaskLogicPath": task_logic_response.path,
},
).command
)
print(bonsai_app_result)
return


def main():
create_fake_subjects()
create_fake_rig()
behavior_cli_args = CliApp.run(
LauncherCliArgs,
cli_args=[
"--debug-mode",
"--allow-dirty",
"--skip-hardware-validation",
],
)

launcher = Launcher(settings=behavior_cli_args)
launcher.run_experiment(client_experiment)
return None


if __name__ == "__main__":
main()
4 changes: 2 additions & 2 deletions src/clabe/apps/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
CommandResult,
ExecutableApp,
Executor,
OutputParser,
StdCommand,
_OutputParser,
identity_parser,
)
from ._bonsai import AindBehaviorServicesBonsaiApp, BonsaiApp
Expand All @@ -26,7 +26,7 @@
"AsyncExecutor",
"Executor",
"identity_parser",
"OutputParser",
"_OutputParser",
"PythonScriptApp",
"ExecutableApp",
"StdCommand",
Expand Down
64 changes: 47 additions & 17 deletions src/clabe/apps/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ class ExecutableApp(Protocol):
class MyApp(ExecutableApp):
@property
def command(self) -> Command:
return Command(cmd="echo hello", output_parser=identity_parser)
return Command(cmd=["echo", "hello"], output_parser=identity_parser)
```
"""

Expand Down Expand Up @@ -170,7 +170,7 @@ async def run_async(self, command: "Command") -> CommandResult:

TOutput = TypeVar("TOutput")

OutputParser: TypeAlias = Callable[[CommandResult], TOutput]
_OutputParser: TypeAlias = Callable[[CommandResult], TOutput]


class Command(Generic[TOutput]):
Expand All @@ -181,14 +181,20 @@ class Command(Generic[TOutput]):
Supports both synchronous and asynchronous execution patterns with type-safe
output parsing.

Commands are provided as a list of strings, which is consistent with subprocess
and executed directly without shell interpretation. This approach:
- Avoids shell injection vulnerabilities
- Handles arguments with spaces correctly without manual quoting
- Is more portable across platforms

Attributes:
cmd: The command string to execute
cmd: The command to execute as a list of strings
result: The result of command execution (available after execution)

Example:
```python
# Create a simple command
cmd = Command(cmd="echo hello", output_parser=identity_parser)
# Create a command
cmd = Command(cmd=["python", "-c", "print('hello')"], output_parser=identity_parser)

# Execute with a synchronous executor
executor = LocalExecutor()
Expand All @@ -198,24 +204,25 @@ class Command(Generic[TOutput]):
def parse_json(result: CommandResult) -> dict:
return json.loads(result.stdout)

cmd = Command(cmd="get-data --json", output_parser=parse_json)
cmd = Command(cmd=["get-data", "--json"], output_parser=parse_json)
data = cmd.execute(executor)
```
"""

def __init__(self, cmd: str, output_parser: OutputParser[TOutput]) -> None:
def __init__(self, cmd: list[str], output_parser: _OutputParser[TOutput]) -> None:
"""Initialize the Command instance.

Args:
cmd: The command string to execute
cmd: The command to execute as a list of strings. The first element
is the program to run, followed by its arguments.
output_parser: Function to parse the command result into desired output type

Example:
```python
# Create a simple command
cmd = Command(cmd="echo hello", output_parser=identity_parser)
cmd = Command(cmd=["echo", "hello"], output_parser=identity_parser)
```
"""
self._cmd = cmd
self._cmd: list[str] = cmd
self._output_parser = output_parser
self._result: Optional[CommandResult] = None

Expand All @@ -227,16 +234,30 @@ def result(self) -> CommandResult:
return self._result

@property
def cmd(self) -> str:
"""Get the command string."""
def cmd(self) -> list[str]:
"""Get the command as a list of strings."""
return self._cmd

def append_arg(self, args: str | list[str]) -> Self:
"""Append an argument to the command."""
"""Append arguments to the command.

Args:
args: Argument(s) to append. Can be a single string or list of strings.
Empty strings are filtered out.

Returns:
Self for method chaining.

Example:
```python
cmd = Command(cmd=["python"], output_parser=identity_parser)
cmd.append_arg(["-m", "pytest"]) # Results in ["python", "-m", "pytest"]
```
"""
if isinstance(args, str):
args = [args]
args = [arg for arg in args if arg]
self._cmd = (self.cmd + f" {' '.join(args)}").strip()
self._cmd = self._cmd + args
return self

def execute(self, executor: Executor) -> TOutput:
Expand Down Expand Up @@ -267,9 +288,18 @@ def _parse_output(self, result: CommandResult) -> TOutput:


class StdCommand(Command[CommandResult]):
"""Standard command that returns the raw CommandResult."""
"""Standard command that returns the raw CommandResult.

A convenience class that creates a Command with the identity_parser,
returning the raw CommandResult without transformation.

Example:
```python
cmd = StdCommand(["echo", "hello"])
```
"""

def __init__(self, cmd: str) -> None:
def __init__(self, cmd: list[str]) -> None:
super().__init__(cmd, identity_parser)


Expand Down
30 changes: 18 additions & 12 deletions src/clabe/apps/_bonsai.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import random
from os import PathLike
from pathlib import Path
from typing import Dict, Optional
from typing import Dict, List, Optional

import pydantic
from aind_behavior_services import AindBehaviorRigModel, AindBehaviorSessionModel, AindBehaviorTaskLogicModel
Expand Down Expand Up @@ -117,13 +117,18 @@ def _build_bonsai_process_command(
is_editor_mode: bool = True,
is_start_flag: bool = True,
additional_properties: Optional[Dict[str, str]] = None,
) -> str:
) -> List[str]:
"""
Builds a shell command for running a Bonsai workflow via subprocess.
Builds a command list for running a Bonsai workflow via subprocess.

Constructs the complete command string with all necessary flags and properties
for executing a Bonsai workflow. Handles editor mode, start flag, and
externalized properties.
Constructs the complete command as a list of arguments with all necessary
flags and properties for executing a Bonsai workflow. Handles editor mode,
start flag, and externalized properties.

Using list format is preferred over string format as it:
- Avoids shell injection vulnerabilities
- Handles paths with spaces correctly without manual quoting
- Is more portable across platforms

Args:
workflow_file: Path to the Bonsai workflow file
Expand All @@ -133,7 +138,7 @@ def _build_bonsai_process_command(
additional_properties: Dictionary of externalized properties to pass. Defaults to None

Returns:
str: The complete command string
List[str]: The complete command as a list of arguments

Example:
```python
Expand All @@ -142,19 +147,20 @@ def _build_bonsai_process_command(
is_editor_mode=False,
additional_properties={"SubjectName": "Mouse123"}
)
# Returns: '"bonsai.exe" "workflow.bonsai" --no-editor -p:"SubjectName"="Mouse123"'
# Returns: ["bonsai.exe", "workflow.bonsai", "--no-editor", "-p:SubjectName=Mouse123"]
```
"""
output_cmd: str = f'"{bonsai_exe}" "{workflow_file}"'
output_cmd: List[str] = [str(bonsai_exe), str(workflow_file)]

if is_editor_mode:
if is_start_flag:
output_cmd += " --start"
output_cmd.append("--start")
else:
output_cmd += " --no-editor"
output_cmd.append("--no-editor")

if additional_properties:
for param, value in additional_properties.items():
output_cmd += f' -p:"{param}"="{value}"'
output_cmd.append(f"-p:{param}={value}")

return output_cmd

Expand Down
10 changes: 5 additions & 5 deletions src/clabe/apps/_curriculum.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ class CurriculumSettings(ServiceSettings):

__yml_section__: t.ClassVar[t.Optional[str]] = "curriculum"

script: str = "curriculum run"
script: list[str] = ["curriculum", "run"]
project_directory: os.PathLike = Path(".")
input_trainer_state: t.Optional[os.PathLike] = None
data_directory: t.Optional[os.PathLike] = None
Expand Down Expand Up @@ -145,18 +145,18 @@ def __init__(
raise ValueError("Data directory is not set.")

kwargs: dict[str, t.Any] = { # Must use kebab casing
"data-directory": f'"{self._settings.data_directory}"',
"input-trainer-state": f'"{self._settings.input_trainer_state}"',
"data-directory": str(self._settings.data_directory),
"input-trainer-state": str(self._settings.input_trainer_state),
}
if self._settings.curriculum is not None:
kwargs["curriculum"] = f'"{self._settings.curriculum}"'
kwargs["curriculum"] = str(self._settings.curriculum)

python_script_app_kwargs = python_script_app_kwargs or {}
self._python_script_app = PythonScriptApp(
script=settings.script,
project_directory=settings.project_directory,
extra_uv_arguments="-q",
additional_arguments=" ".join(f"--{key} {value}" for key, value in kwargs.items()),
additional_arguments=[arg for kv in kwargs.items() for arg in ("--" + kv[0], str(kv[1]))],
**python_script_app_kwargs,
)

Expand Down
Loading