Merge pull request #38 from AllenNeuralDynamics:update-to-aind-behavi…

…or-service-09x

Update to aind-behavior-services 0.9 and add watchdog

examples/example_roi_trial_type.py

@@ -74,9 +74,7 @@ def mock_rig() -> AindForceForagingRig:
     )
     water_valve_calibration.output = WaterValveCalibrationOutput(slope=1, offset=0)  # For testing purposes
 
-    video_writer = rig.VideoWriterFfmpeg(
-        frame_rate=120, container_extension="mp4", output_arguments="-c:v h264_nvenc -vsync 0 -2pass "
-    )
+    video_writer = rig.VideoWriterFfmpeg(frame_rate=120, container_extension="mp4")
 
     load_cells_calibration = lcc.LoadCellsCalibration(
         output=lcc.LoadCellsCalibrationOutput(),

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "aind-behavior-force-foraging"
-description = "A library that defines AIND data schema for the Aind Behavior Force Foraing experiment."
+description = "A library that defines AIND data schema for the Aind Behavior Force Foraging experiment."
 authors = [
     { name = "Bruno Cruz", email = "[email protected]" },
 ]
@@ -19,7 +19,7 @@ readme = "README.md"
 dynamic = ["version"]
 
 dependencies = [
-  "aind_behavior_services>=0.8, <0.9",
+  "aind_behavior_services>=0.9, <0.10",
 ]
 
 [project.urls]
@@ -30,7 +30,7 @@ Changelog = "https://github.com/AllenNeuralDynamics/Aind.Behavior.ForceForaging/
 
 [project.optional-dependencies]
 
-launcher = ["aind_behavior_experiment_launcher[aind-services]>=0.2, <0.3"]
+launcher = ["aind_behavior_experiment_launcher[aind-services]>=0.3, <0.4"]
 
 dev = [
     "aind_behavior_force_foraging[launcher]",

src/DataSchemas/aind_behavior_force_foraging/launcher.py

@@ -1,29 +1,40 @@
+import datetime
+import os
+from functools import partial
 from pathlib import Path
+from typing import Callable, Optional
 
 import aind_behavior_experiment_launcher.launcher.behavior_launcher as behavior_launcher
-from aind_behavior_experiment_launcher.apps.app_service import BonsaiApp
-from aind_behavior_experiment_launcher.resource_monitor.resource_monitor_service import (
-    ResourceMonitor,
-    available_storage_constraint_factory,
-    remote_dir_exists_constraint_factory,
-)
+from aind_behavior_experiment_launcher import resource_monitor
+from aind_behavior_experiment_launcher.apps import BonsaiApp
+from aind_behavior_experiment_launcher.data_transfer import aind_watchdog
 from aind_behavior_services.session import AindBehaviorSessionModel
 
+from aind_behavior_force_foraging.data_mappers import AindDataMapperWrapper
 from aind_behavior_force_foraging.rig import AindForceForagingRig
 from aind_behavior_force_foraging.task_logic import AindForceForagingTaskLogic
 
 
 def make_launcher() -> behavior_launcher.BehaviorLauncher:
+    use_watchdog = False
     data_dir = r"C:/Data"
-    remote_dir = r"\\allen\aind\scratch\force-foraging\data"
+    remote_dir = Path(r"\\allen\aind\scratch\force-foraging\data")
     srv = behavior_launcher.BehaviorServicesFactoryManager()
-    srv.attach_bonsai_app(BonsaiApp(Path(r"./src/main.bonsai")))
-    srv.attach_data_transfer(behavior_launcher.robocopy_data_transfer_factory(Path(remote_dir)))
+    srv.attach_bonsai_app(BonsaiApp(r"./src/main.bonsai"))
+
+    if use_watchdog:
+        srv.attach_data_mapper(AindDataMapperWrapper.from_launcher)
+        srv.attach_data_transfer(
+            watchdog_data_transfer_factory(remote_dir, project_name="Cognitive flexibility in patch foraging")
+        )
+    else:
+        srv.attach_data_transfer(behavior_launcher.robocopy_data_transfer_factory(Path(remote_dir)))
+
     srv.attach_resource_monitor(
-        ResourceMonitor(
+        resource_monitor.ResourceMonitor(
             constrains=[
-                available_storage_constraint_factory(data_dir, 2e11),
-                remote_dir_exists_constraint_factory(Path(remote_dir)),
+                resource_monitor.available_storage_constraint_factory(data_dir, 2e11),
+                resource_monitor.remote_dir_exists_constraint_factory(Path(remote_dir)),
             ]
         )
     )
@@ -35,7 +46,6 @@ def make_launcher() -> behavior_launcher.BehaviorLauncher:
         data_dir=data_dir,
         config_library_dir=r"\\allen\aind\scratch\AindBehavior.db\AindForceForaging",
         temp_dir=r"./local/.temp",
-        repository_dir=None,
         allow_dirty=False,
         skip_hardware_validation=False,
         debug_mode=False,
@@ -45,6 +55,57 @@ def make_launcher() -> behavior_launcher.BehaviorLauncher:
     )
 
 
+def watchdog_data_transfer_factory(
+    destination: os.PathLike,
+    schedule_time: Optional[datetime.time] = datetime.time(hour=20),
+    project_name: Optional[str] = None,
+    **watchdog_kwargs,
+) -> Callable[[behavior_launcher.BehaviorLauncher], aind_watchdog.WatchdogDataTransferService]:
+    return partial(
+        _watchdog_data_transfer_factory,
+        destination=destination,
+        schedule_time=schedule_time,
+        project_name=project_name,
+        **watchdog_kwargs,
+    )
+
+
+def _watchdog_data_transfer_factory(
+    launcher: behavior_launcher.BehaviorLauncher,
+    destination: os.PathLike,
+    **watchdog_kwargs,
+) -> aind_watchdog.WatchdogDataTransferService:
+    if launcher.services_factory_manager.data_mapper is None:
+        raise ValueError("Data mapper service is not set. Cannot create watchdog.")
+    if not isinstance(launcher.services_factory_manager.data_mapper, AindDataMapperWrapper):
+        raise ValueError(
+            "Data mapper service is not of the correct type (AindDataMapperWrapper). Cannot create watchdog."
+        )
+    if not launcher.services_factory_manager.data_mapper.is_mapped():
+        raise ValueError("Data mapper has not mapped yet. Cannot create watchdog.")
+
+    if not isinstance(launcher.session_schema, AindBehaviorSessionModel):
+        raise ValueError(
+            "Session schema is not of the correct type (AindBehaviorSessionModel). Cannot create watchdog."
+        )
+
+    if not launcher.session_schema.session_name:
+        raise ValueError("Session name is not set. Cannot create watchdog.")
+
+    destination = Path(destination)
+    if launcher.group_by_subject_log:
+        destination = destination / launcher.session_schema.subject
+
+    watchdog = aind_watchdog.WatchdogDataTransferService(
+        source=launcher.session_directory,
+        aind_session_data_mapper=launcher.services_factory_manager.data_mapper._session_mapper,
+        session_name=launcher.session_schema.session_name,
+        destination=destination,
+        **watchdog_kwargs,
+    )
+    return watchdog
+
+
 def main():
     launcher = make_launcher()
     launcher.main()

src/DataSchemas/aind_behavior_session_model.json

@@ -1,7 +1,7 @@
 {
   "properties": {
     "aind_behavior_services_pkg_version": {
-      "default": "0.8.9",
+      "default": "0.9.0",
       "pattern": "^(0|[1-9]\\d*)\\.(0|[1-9]\\d*)\\.(0|[1-9]\\d*)(?:-((?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\\.(?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\\+([0-9a-zA-Z-]+(?:\\.[0-9a-zA-Z-]+)*))?$",
       "title": "aind_behavior_services package version",
       "type": "string"

src/DataSchemas/aind_force_foraging_rig.json

@@ -2257,7 +2257,7 @@
   },
   "properties": {
     "aind_behavior_services_pkg_version": {
-      "default": "0.8.9",
+      "default": "0.9.0",
       "pattern": "^(0|[1-9]\\d*)\\.(0|[1-9]\\d*)\\.(0|[1-9]\\d*)(?:-((?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\\.(?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\\+([0-9a-zA-Z-]+(?:\\.[0-9a-zA-Z-]+)*))?$",
       "title": "aind_behavior_services package version",
       "type": "string"

src/DataSchemas/aind_force_foraging_task_logic.json

@@ -40,7 +40,7 @@
           "title": "Rng Seed"
         },
         "aind_behavior_services_pkg_version": {
-          "default": "0.8.9",
+          "default": "0.9.0",
           "pattern": "^(0|[1-9]\\d*)\\.(0|[1-9]\\d*)\\.(0|[1-9]\\d*)(?:-((?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\\.(?:0|[1-9]\\d*|\\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\\+([0-9a-zA-Z-]+(?:\\.[0-9a-zA-Z-]+)*))?$",
           "title": "aind_behavior_services package version",
           "type": "string"

src/Extensions/AindBehaviorSessionModel.cs

@@ -15,7 +15,7 @@ namespace AindForceForagingDataSchema.Session
     public partial class AindBehaviorSessionModel
     {
 
-        private string _aindBehaviorServicesPkgVersion = "0.8.9";
+        private string _aindBehaviorServicesPkgVersion = "0.9.0";
 
         private string _version = "0.3.0";
 

src/Extensions/AindForceForagingRig.cs

@@ -5708,7 +5708,7 @@ public override string ToString()
     public partial class AindForceForagingRig
     {
 
-        private string _aindBehaviorServicesPkgVersion = "0.8.9";
+        private string _aindBehaviorServicesPkgVersion = "0.9.0";
 
         private string _version = "0.2.0";
 

src/Extensions/AindForceForagingTaskLogic.cs

@@ -127,7 +127,7 @@ public partial class AindForceForagingTaskParameters
 
         private double? _rngSeed;
 
-        private string _aindBehaviorServicesPkgVersion = "0.8.9";
+        private string _aindBehaviorServicesPkgVersion = "0.9.0";
 
         private Environment _environment = new Environment();
 

tests/__init__.py

@@ -1,11 +1,16 @@
 import glob
 import importlib.util
+import logging
 from pathlib import Path
 from types import ModuleType
 
 EXAMPLES_DIR = Path(__file__).parents[1] / "examples"
 JSON_ROOT = Path("./local").resolve()
 
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
+logging.disable(logging.CRITICAL)
+
 
 def build_example(script_path: str) -> ModuleType:
     module_name = Path(script_path).stem

tests/test_aind_data_mapper.py

@@ -0,0 +1,72 @@
+import sys
+import unittest
+from datetime import datetime
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from aind_behavior_force_foraging.data_mappers import (
+    AindRigDataMapper,
+    AindSessionDataMapper,
+)
+from aind_data_schema.core.rig import Rig
+from git import Repo
+
+sys.path.append(".")
+from examples.example_roi_trial_type import mock_rig, mock_session, mock_task_logic  # isort:skip # pylint: disable=wrong-import-position
+
+
+class TestAindSessionDataMapper(unittest.TestCase):
+    def setUp(self):
+        self.session_model = mock_session()
+        self.rig_model = mock_rig()
+        self.task_logic_model = mock_task_logic()
+        self.repository = Repo(Path("./"))
+        self.script_path = Path("./src/main.bonsai")
+        self.session_end_time = datetime.now()
+        self.session_directory = None
+
+        self.mapper = AindSessionDataMapper(
+            session_model=self.session_model,
+            rig_model=self.rig_model,
+            task_logic_model=self.task_logic_model,
+            repository=self.repository,
+            script_path=self.script_path,
+            session_end_time=self.session_end_time,
+        )
+
+    @patch("aind_behavior_force_foraging.data_mappers.AindSessionDataMapper._map")
+    def test_mock_map(self, mock_map):
+        mock_map.return_value = MagicMock()
+        result = self.mapper.map()
+        self.assertIsNotNone(result)
+        self.assertTrue(self.mapper.is_mapped())
+
+    def test_map(self):
+        mapped = self.mapper.map()
+        self.assertIsNotNone(mapped)
+
+
+class TestAindRigDataMapper(unittest.TestCase):
+    def setUp(self):
+        self.rig_schema_filename = "rig_schema.json"
+        self.db_root = MagicMock()
+        self.session_directory = MagicMock()
+        self.db_suffix = "test_suffix"
+        self.mapper = AindRigDataMapper(
+            rig_schema_filename=self.rig_schema_filename,
+            db_root=self.db_root,
+            db_suffix=self.db_suffix,
+        )
+
+    @patch("pathlib.Path.exists", return_value=True)
+    @patch("aind_behavior_force_foraging.data_mappers.model_from_json_file")
+    def test_mock_map(self, mock_model_from_json_file, mock_path_exists):
+        mock_model_from_json_file.return_value = MagicMock(spec=Rig)
+        result = self.mapper.map()
+        self.assertIsNotNone(result)
+        self.assertTrue(self.mapper.mapped)
+        self.assertIsInstance(result, Rig)
+
+
+if __name__ == "__main__":
+    unittest.main()