Add clabe launcher infrastructure

README.md

@@ -66,6 +66,25 @@ uv run dynamic-foraging -h
 
 You may need to install optional dependencies depending on the sub-commands you run.
 
+
+## 🎮 Experiment launcher (temporarily CLABE)
+
+To manage experiments and input files, this repository contains a launcher script that can be used to run the Dynamic Foraging task. A default script is located at `./scripts/aind-launcher.py`. It can be run from the command line as follows:
+
+```powershell
+uv run clabe run ./scripts/aind-launcher.py
+# or uv run ./scripts/main.py
+```
+
+Additional arguments can be passed to the script as needed. For instance to allow the script to run with uncommitted changes in the repository, the `--allow-dirty` flag can be used:
+
+```powershell
+uv run clabe run ./scripts/aind-launcher.py --allow-dirty
+```
+
+or via a `./local/clabe.yml` file. Additional custom launcher scripts can be created and used as needed. See documentation in the [`clabe` repository](https://allenneuraldynamics.github.io/clabe/) for more details.
+
+
 ## 🔍 Primary data quality-control
 
 Once an experiment is collected, the primary data quality-control script can be run to check the data for issues. This script can be launcher using:

pyproject.toml

@@ -39,6 +39,7 @@ dev = [
     "ruff",
     "codespell",
     "aind-behavior-dynamic-foraging[data]",
+    "aind-clabe[aind-services]>=0.10"
 ]
 
 docs = [

scripts/aind-launcher.py

@@ -0,0 +1,162 @@
+import logging
+import typing as ty
+from pathlib import Path
+
+import clabe.resource_monitor
+import pandas as pd
+from aind_behavior_services.rig.aind_manipulator import ManipulatorPosition
+from aind_behavior_services.session import Session
+from clabe.apps import (
+    AindBehaviorServicesBonsaiApp,
+)
+from clabe.data_transfer.robocopy import RobocopyService, RobocopySettings
+from clabe.launcher import Launcher, LauncherCliArgs, experiment
+from clabe.pickers import ByAnimalModifier, DefaultBehaviorPicker, DefaultBehaviorPickerSettings
+from pydantic_settings import CliApp
+
+from aind_behavior_dynamic_foraging import data_contract
+from aind_behavior_dynamic_foraging.rig import AindDynamicForagingRig
+from aind_behavior_dynamic_foraging.task_logic import AindDynamicForagingTaskLogic
+
+logger = logging.getLogger(__name__)
+
+
+@experiment()
+async def iso_force_experiment(launcher: Launcher) -> None:
+    picker = DefaultBehaviorPicker(
+        launcher=launcher,
+        settings=DefaultBehaviorPickerSettings(
+            config_library_dir=r"\\allen\aind\scratch\AindBehavior.db\AindDynamicForaging"
+        ),
+    )
+
+    session = picker.pick_session(Session)
+    task_logic = picker.pick_task(AindDynamicForagingTaskLogic)
+    rig = picker.pick_rig(AindDynamicForagingRig)
+    ensure_rig_and_computer_name(rig)
+
+    # Post-fetching modifications
+    manipulator_modifier = ByAnimalManipulatorModifier(
+        subject_db_path=picker.subject_dir / session.subject,
+        model_path="manipulator.calibration.initial_position",
+        model_name="manipulator_init.json",
+        launcher=launcher,
+    )
+    manipulator_modifier.inject(rig)
+
+    launcher.register_session(session, rig.data_directory)
+
+    clabe.resource_monitor.ResourceMonitor(
+        constrains=[
+            clabe.resource_monitor.available_storage_constraint_factory(rig.data_directory, 2e11),
+        ]
+    ).run()
+
+    bonsai_app = AindBehaviorServicesBonsaiApp(
+        workflow=Path(r"./src/main.bonsai"),
+        temp_directory=launcher.temp_dir,
+        rig=rig,
+        session=session,
+        task=task_logic,
+    )
+    await bonsai_app.run_async()
+
+    # Update manipulator initial position for next session
+    try:
+        manipulator_modifier.dump()
+    except Exception as e:
+        logger.error("Failed to update manipulator initial position: %s", e)
+
+    # Run data qc
+    if picker.ui_helper.prompt_yes_no_question("Would you like to generate a qc report?"):
+        try:
+            import webbrowser
+
+            from contraqctor.qc.reporters import HtmlReporter
+
+            from aind_behavior_dynamic_foraging.data_qc.suite import make_qc_runner
+
+            vr_dataset = data_contract.dataset(launcher.session_directory)
+            runner = make_qc_runner(vr_dataset)
+            qc_path = launcher.session_directory / "Behavior" / "Logs" / "qc_report.html"
+            reporter = HtmlReporter(output_path=qc_path)
+            runner.run_all_with_progress(reporter=reporter)
+            webbrowser.open(qc_path.as_uri(), new=2)
+        except Exception as e:
+            logger.error("Failed to run data QC: %s", e)
+
+    # Transfer data
+    is_transfer = picker.ui_helper.prompt_yes_no_question("Would you like to transfer data?")
+    if not is_transfer:
+        logger.info("Data transfer skipped by user.")
+        return
+
+    launcher.copy_logs()
+    settings = RobocopySettings(destination=r"\\allen\aind\scratch\AindDynamicForaging\data")
+    assert launcher.session.session_name is not None, "Session name is None"
+    settings.destination = Path(settings.destination) / launcher.session.session_name
+    RobocopyService(source=launcher.session_directory, settings=settings).transfer()
+    return
+
+
+def ensure_rig_and_computer_name(rig: AindDynamicForagingRig) -> None:
+    """Ensures rig and computer name are set from environment variables if available, otherwise defaults to rig configuration values."""
+
+    import os
+
+    rig_name = os.environ.get("aibs_comp_id", None)
+    computer_name = os.environ.get("hostname", None)
+
+    if rig_name is None:
+        logger.warning(
+            "'aibs_comp_id' environment variable not set. Defaulting to rig name from configuration. %s", rig.rig_name
+        )
+        rig_name = rig.rig_name
+    if computer_name is None:
+        computer_name = rig.computer_name
+        logger.warning(
+            "'hostname' environment variable not set. Defaulting to computer name from configuration. %s",
+            rig.computer_name,
+        )
+
+    if rig_name != rig.rig_name or computer_name != rig.computer_name:
+        logger.warning(
+            "Rig name or computer name from environment variables do not match the rig configuration. "
+            "Forcing rig name: %s and computer name: %s from environment variables.",
+            rig_name,
+            computer_name,
+        )
+        rig.rig_name = rig_name
+        rig.computer_name = computer_name
+
+
+class ByAnimalManipulatorModifier(ByAnimalModifier[AindDynamicForagingRig]):
+    """Modifier to set and update manipulator initial position based on animal-specific data."""
+
+    def __init__(
+        self, subject_db_path: Path, model_path: str, model_name: str, *, launcher: Launcher, **kwargs
+    ) -> None:
+        super().__init__(subject_db_path, model_path, model_name, **kwargs)
+        self._launcher = launcher
+
+    def _process_before_dump(self) -> ManipulatorPosition:
+        _dataset = data_contract.dataset(self._launcher.session_directory)
+        manipulator_parking_position: pd.DataFrame = ty.cast(
+            pd.DataFrame, _dataset["Behavior"]["SoftwareEvents"]["InitialManipulatorPosition"].read()
+        )
+        return ManipulatorPosition.model_validate(manipulator_parking_position.iloc[0]["data"])
+
+
+class ClabeCli(LauncherCliArgs):
+    def cli_cmd(self):
+        launcher = Launcher(settings=self)
+        launcher.run_experiment(iso_force_experiment)
+        return None
+
+
+def main() -> None:
+    CliApp().run(ClabeCli)
+
+
+if __name__ == "__main__":
+    main()

src/Extensions/Hardware.bonsai

@@ -290,7 +290,7 @@
             </Expression>
             <Expression xsi:type="beh:Format">
               <harp:MessageType>Write</harp:MessageType>
-              <harp:Register xsi:type="beh:Camera0Frequency" />
+              <harp:Register xsi:type="beh:Camera1Frequency" />
             </Expression>
             <Expression xsi:type="MulticastSubject">
               <Name>HarpBehaviorCommands</Name>
@@ -305,7 +305,7 @@
               </Combinator>
             </Expression>
             <Expression xsi:type="beh:Parse">
-              <harp:Register xsi:type="beh:TimestampedCamera0Frame" />
+              <harp:Register xsi:type="beh:TimestampedCamera1Frame" />
             </Expression>
             <Expression xsi:type="MemberSelector">
               <Selector>Value</Selector>

src/Extensions/Logging.bonsai

@@ -815,8 +815,29 @@
                         <Expression xsi:type="GroupWorkflow">
                           <Name>AdditionalSoftwareEvents</Name>
                           <Workflow>
-                            <Nodes />
-                            <Edges />
+                            <Nodes>
+                              <Expression xsi:type="SubscribeSubject">
+                                <Name>ManipulatorPosition</Name>
+                              </Expression>
+                              <Expression xsi:type="Combinator">
+                                <Combinator xsi:type="rx:Take">
+                                  <rx:Count>1</rx:Count>
+                                </Combinator>
+                              </Expression>
+                              <Expression xsi:type="Combinator">
+                                <Combinator xsi:type="p1:CreateSoftwareEvent">
+                                  <p1:EventName>InitialManipulatorPosition</p1:EventName>
+                                </Combinator>
+                              </Expression>
+                              <Expression xsi:type="MulticastSubject">
+                                <Name>SoftwareEvent</Name>
+                              </Expression>
+                            </Nodes>
+                            <Edges>
+                              <Edge From="0" To="1" Label="Source1" />
+                              <Edge From="1" To="2" Label="Source1" />
+                              <Edge From="2" To="3" Label="Source1" />
+                            </Edges>
                           </Workflow>
                         </Expression>
                       </Nodes>

src/Extensions/OperationControl.bonsai

@@ -640,6 +640,16 @@
             <Expression xsi:type="rx:BehaviorSubject" TypeArguments="p2:Unit">
               <rx:Name>ResetQuickRetract</rx:Name>
             </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="DoubleProperty">
+                <Value>0</Value>
+              </Combinator>
+            </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="rx:Take">
+                <rx:Count>1</rx:Count>
+              </Combinator>
+            </Expression>
             <Expression xsi:type="rx:BehaviorSubject" TypeArguments="sys:Double">
               <rx:Name>ManualSpoutDelta</rx:Name>
             </Expression>
@@ -713,6 +723,9 @@
                 </Edges>
               </Workflow>
             </Expression>
+            <Expression xsi:type="Combinator">
+              <Combinator xsi:type="rx:Merge" />
+            </Expression>
             <Expression xsi:type="rx:BehaviorSubject">
               <Name>ManipulatorBiasTracker</Name>
             </Expression>
@@ -746,19 +759,22 @@
           </Nodes>
           <Edges>
             <Edge From="2" To="3" Label="Source1" />
-            <Edge From="3" To="5" Label="Source1" />
-            <Edge From="4" To="5" Label="Source2" />
-            <Edge From="5" To="6" Label="Source1" />
-            <Edge From="6" To="10" Label="Source1" />
+            <Edge From="3" To="13" Label="Source1" />
+            <Edge From="4" To="5" Label="Source1" />
+            <Edge From="5" To="7" Label="Source1" />
+            <Edge From="6" To="7" Label="Source2" />
             <Edge From="7" To="8" Label="Source1" />
-            <Edge From="8" To="9" Label="Source1" />
-            <Edge From="9" To="10" Label="Source2" />
+            <Edge From="8" To="12" Label="Source1" />
+            <Edge From="9" To="10" Label="Source1" />
             <Edge From="10" To="11" Label="Source1" />
-            <Edge From="11" To="12" Label="Source1" />
-            <Edge From="12" To="13" Label="Source1" />
+            <Edge From="11" To="12" Label="Source2" />
+            <Edge From="12" To="13" Label="Source2" />
             <Edge From="13" To="14" Label="Source1" />
+            <Edge From="14" To="15" Label="Source1" />
             <Edge From="15" To="16" Label="Source1" />
             <Edge From="16" To="17" Label="Source1" />
+            <Edge From="18" To="19" Label="Source1" />
+            <Edge From="19" To="20" Label="Source1" />
           </Edges>
         </Workflow>
       </Expression>
@@ -825,7 +841,7 @@
             <Expression xsi:type="beh:CreateMessage">
               <harp:MessageType>Write</harp:MessageType>
               <harp:Payload xsi:type="beh:CreateStartCamerasPayload">
-                <beh:StartCameras>CameraOutput0</beh:StartCameras>
+                <beh:StartCameras>CameraOutput1</beh:StartCameras>
               </harp:Payload>
             </Expression>
             <Expression xsi:type="BitwiseNot" />
@@ -845,7 +861,7 @@
             <Expression xsi:type="beh:CreateMessage">
               <harp:MessageType>Write</harp:MessageType>
               <harp:Payload xsi:type="beh:CreateStopCamerasPayload">
-                <beh:StopCameras>CameraOutput0</beh:StopCameras>
+                <beh:StopCameras>CameraOutput1</beh:StopCameras>
               </harp:Payload>
             </Expression>
             <Expression xsi:type="Combinator">

src/aind_behavior_dynamic_foraging/data_contract/_dataset.py

@@ -250,6 +250,13 @@ def make_dataset(
                                     root_path / "behavior/SoftwareEvents/RngSeed.json"
                                 ),
                             ),
+                            SoftwareEvents(
+                                name="InitialManipulatorPosition",
+                                description="An event emitted at the start of the experiment to indicate the initial manipulator position.",
+                                reader_params=SoftwareEvents.make_params(
+                                    root_path / "behavior/SoftwareEvents/InitialManipulatorPosition.json"
+                                ),
+                            ),
                             SoftwareEvents(
                                 name="EndExperiment",
                                 description="An event emitted when the experiment ends.",

src/aind_behavior_dynamic_foraging/data_qc/__init__.py

@@ -20,7 +20,7 @@ class DataQcCli(BaseSettings, cli_kebab_case=True):
     def cli_cmd(self):
         """Run data quality checks on the dataset located at the specified path."""
         from ..data_contract import dataset
-        from ._suite import make_qc_runner
+        from .suite import make_qc_runner
 
         this_dataset = dataset(Path(self.data_path), self.version)
         runner = make_qc_runner(this_dataset)

src/main.bonsai

@@ -11,7 +11,7 @@
   <Workflow>
     <Nodes>
       <Expression xsi:type="ExternalizedMapping">
-        <Property Name="TaskLogicPath" />
+        <Property Name="TaskLogicPath" DisplayName="TaskPath" />
         <Property Name="RigPath" />
         <Property Name="SessionPath" />
       </Expression>