Merge pull request #436 from AllenNeuralDynamics/test-launcher

Upgrade aind-clabe to v0.8.0 and refactor launcher script

Author: bruno-f-cruz

.github/workflows/vr-foraging-cicd.yml

@@ -303,7 +303,9 @@ jobs:
         uses: actions/checkout@v5
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v3
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
 
       - name: Setup Graphviz
         uses: ts-graphviz/setup-graphviz@v2

docs/conf.py

@@ -90,5 +90,5 @@ def export_model_diagram(model: BaseModel, root: str = _static_path) -> None:
 
 # -- Dataset rendering
 
-with open(f"{_static_path}/dataset.txt", "w", encoding="utf-8") as f:
+with open(f"{_static_path}/dataset.html", "w", encoding="utf-8") as f:
     f.write(contract.render_dataset())

docs/dataset.rst

@@ -1,6 +1,10 @@
 Data Contract
 -------------
 
-.. literalinclude:: /_static/dataset.txt
-   :language: text
-   :caption: Dataset
+.. raw:: html
+
+   <iframe
+       src="_static/dataset.html"
+       style="width:100%; height:600px; border:none; border-radius:8px; box-shadow:0 2px 4px rgba(0,0,0,0.1);"
+       title="Data Contract visualization">
+   </iframe>

pyproject.toml

@@ -31,7 +31,7 @@ Changelog = "https://github.com/AllenNeuralDynamics/Aind.Behavior.VrForaging/rel
 data = ["contraqctor<0.6.0"]
 
 launcher = [
-    "aind-clabe[aind-services]>=0.7",
+    "aind-clabe[aind-services] >= 0.8.0 ,<0.9.0",
     "aind-data-schema>=2",
     "aind_behavior_vr_foraging[data]",
     ]

src/aind_behavior_vr_foraging/curricula/aind.Behavior.VrForaging.Curricula

@@ -41,7 +41,9 @@ def dataset(path: os.PathLike, version: str = __semver__) -> contraqctor.contrac
 
 def render_dataset(version: str = __semver__) -> str:
     """Renders the dataset as a tree-like structure for visualization."""
-    from contraqctor.contract.utils import print_data_stream_tree
+    from contraqctor.contract.utils import print_data_stream_tree_html
 
     dataset_constructor = _dataset_lookup_helper(version)
-    return print_data_stream_tree(dataset_constructor(Path("<RootPath>")), show_missing_indicator=False, show_type=True)
+    return print_data_stream_tree_html(
+        dataset_constructor(Path("<RootPath>")), show_missing_indicator=False, show_type=True
+    )

src/aind_behavior_vr_foraging/data_contract/__init__.py

@@ -46,8 +46,8 @@ def cli_cmd(self):
 
         repo = Repo(self.repo_path)
         session_mapped = AindSessionDataMapper(
-            session_model=session,
-            rig_model=rig,
+            session=session,
+            rig=rig,
             task_logic_model=task_logic,
             repository=repo,
             script_path=Path("./src/main.bonsai"),

src/aind_behavior_vr_foraging/data_mappers/__init__.py

@@ -4,7 +4,7 @@
 import platform
 from decimal import Decimal
 from pathlib import Path
-from typing import Any, Callable, Optional, cast
+from typing import Optional, cast
 
 import aind_behavior_services.rig as AbsRig
 from aind_behavior_services import calibration as AbsCalibration
@@ -14,7 +14,6 @@
 from aind_data_schema_models import coordinates as aind_schema_model_coordinates
 from aind_data_schema_models import modalities, units
 from clabe.data_mapper import aind_data_schema as ads
-from clabe.launcher import Launcher
 
 from aind_behavior_vr_foraging.rig import AindVrForagingRig
 
@@ -45,10 +44,10 @@ def get_spawned_device(self, name: str) -> devices.Device:
 class AindRigDataMapper(ads.AindDataSchemaRigDataMapper):
     def __init__(
         self,
-        rig_model: AindVrForagingRig,
+        rig: AindVrForagingRig,
     ):
         super().__init__()
-        self.rig_model = rig_model
+        self.rig_model = rig
         self._mapped: Optional[instrument.Instrument] = None
 
     def rig_schema(self):
@@ -75,17 +74,6 @@ def mapped(self) -> instrument.Instrument:
     def is_mapped(self) -> bool:
         return self.mapped is not None
 
-    @classmethod
-    def build_runner(cls) -> Callable[[Launcher[AindVrForagingRig, Any, Any]], "AindRigDataMapper"]:
-        def _new(
-            launcher: Launcher[AindVrForagingRig, Any, Any],
-        ) -> "AindRigDataMapper":
-            new = cls(rig_model=launcher.get_rig(strict=True))
-            new.map()
-            return new
-
-        return _new
-
     ## From here on, private methods only!
     ## Lasciate ogne speranza, voi ch'entrate
 

src/aind_behavior_vr_foraging/data_mappers/_rig.py

@@ -3,7 +3,7 @@
 import os
 import sys
 from pathlib import Path
-from typing import Callable, Dict, List, Optional, Union
+from typing import List, Optional, Union
 
 import aind_behavior_services.rig as AbsRig
 import git
@@ -17,7 +17,6 @@
 from clabe.apps import CurriculumSuggestion
 from clabe.data_mapper import aind_data_schema as ads
 from clabe.data_mapper import helpers as data_mapper_helpers
-from clabe.launcher import Launcher
 
 from aind_behavior_vr_foraging.rig import AindVrForagingRig
 from aind_behavior_vr_foraging.task_logic import AindVrForagingTaskLogic
@@ -30,24 +29,22 @@
 class AindSessionDataMapper(ads.AindDataSchemaSessionDataMapper):
     def __init__(
         self,
-        session_model: AindBehaviorSessionModel,
-        rig_model: AindVrForagingRig,
-        task_logic_model: AindVrForagingTaskLogic,
+        session: AindBehaviorSessionModel,
+        rig: AindVrForagingRig,
+        task_logic: AindVrForagingTaskLogic,
         repository: Union[os.PathLike, git.Repo] = Path("."),
         script_path: os.PathLike = Path("./src/main.bonsai"),
         session_end_time: Optional[datetime.datetime] = None,
         curriculum_suggestion: Optional[CurriculumSuggestion] = None,
-        output_parameters: Optional[Dict] = None,
     ):
-        self.session_model = session_model
-        self.rig_model = rig_model
-        self.task_logic_model = task_logic_model
+        self.session_model = session
+        self.rig_model = rig
+        self.task_logic_model = task_logic
         self.repository = repository
         if isinstance(self.repository, os.PathLike | str):
             self.repository = git.Repo(Path(self.repository))
         self.script_path = script_path
         self._session_end_time = session_end_time
-        self.output_parameters = output_parameters
         self._mapped: Optional[acquisition.Acquisition] = None
         self.curriculum = curriculum_suggestion
 
@@ -60,28 +57,6 @@ def session_end_time(self) -> datetime.datetime:
     def session_schema(self):
         return self.mapped
 
-    @classmethod
-    def build_runner(
-        cls,
-        curriculum_suggestion: Optional[Callable[[], CurriculumSuggestion | None]] = None,
-    ) -> Callable[
-        [Launcher[AindVrForagingRig, AindBehaviorSessionModel, AindVrForagingTaskLogic]], "AindSessionDataMapper"
-    ]:
-        def _new(
-            launcher: Launcher[AindVrForagingRig, AindBehaviorSessionModel, AindVrForagingTaskLogic],
-        ) -> "AindSessionDataMapper":
-            new = cls(
-                session_model=launcher.get_session(strict=True),
-                rig_model=launcher.get_rig(strict=True),
-                task_logic_model=launcher.get_task_logic(strict=True),
-                repository=launcher.repository,
-                curriculum_suggestion=curriculum_suggestion() if curriculum_suggestion is not None else None,
-            )
-            new.map()
-            return new
-
-        return _new
-
     @property
     def session_name(self) -> str:
         if self.session_model.session_name is None:
@@ -97,7 +72,7 @@ def mapped(self) -> acquisition.Acquisition:
     def is_mapped(self) -> bool:
         return self.mapped is not None
 
-    def map(self) -> Optional[acquisition.Acquisition]:
+    def map(self) -> acquisition.Acquisition:
         logger.info("Mapping aind-data-schema Session.")
         try:
             self._mapped = self._map()
@@ -120,7 +95,7 @@ def _map(self) -> acquisition.Acquisition:
             acquisition_end_time=utcnow(),
             acquisition_start_time=self.session_model.date,
             experimenters=self.session_model.experimenter,
-            acquisition_type=self.session_model.experiment,
+            acquisition_type=self.session_model.experiment or self.task_logic_model.name,
             coordinate_system=_make_origin_coordinate_system(),
             data_streams=self._get_data_streams(),
             calibrations=self._get_calibrations(),
@@ -245,7 +220,7 @@ def _get_stimulus_epochs(self) -> List[acquisition.StimulusEpoch]:
                 stimulus_start_time=self.session_model.date,
                 stimulus_end_time=self.session_end_time,
                 configurations=stimulus_epoch_configurations,
-                stimulus_name=self.session_model.experiment,
+                stimulus_name=self.session_model.experiment or self.task_logic_model.name,
                 stimulus_modalities=stimulus_modalities,
                 performance_metrics=performance_metrics,
                 curriculum_status=curriculum_status,

src/aind_behavior_vr_foraging/data_mappers/_session.py

@@ -1,24 +1,16 @@
 import enum
 import logging
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Type, TypeVar, Union
+from typing import List, Optional, Type, TypeVar, Union
 
 import aind_behavior_services.calibration as AbsCalibration
 import pydantic
 from aind_behavior_services.utils import get_fields_of_type, utcnow
 from aind_data_schema.components import coordinates, measurements
 from aind_data_schema.core import acquisition
 from aind_data_schema_models import units
-from clabe.launcher import Launcher, Promise
 
 from aind_behavior_vr_foraging.rig import AindVrForagingRig
 
-if TYPE_CHECKING:
-    from ._rig import AindRigDataMapper
-    from ._session import AindSessionDataMapper
-else:
-    AindRigDataMapper = AindSessionDataMapper = Any
-
 TTo = TypeVar("TTo", bound=pydantic.BaseModel)
 
 T = TypeVar("T")
@@ -38,22 +30,6 @@ def coerce_to_aind_data_schema(value: Union[pydantic.BaseModel, dict], target_ty
     return target_type(**_normalized_input)
 
 
-def write_ads_mappers(
-    session_mapper: Promise[[Launcher], AindSessionDataMapper], rig_mapper: Promise[[Launcher], AindRigDataMapper]
-) -> Callable[[Launcher], None]:
-    def _run(launcher: Launcher) -> None:
-        session_directory = launcher.session_directory
-        _session = session_mapper.result.mapped
-        _rig = rig_mapper.result.mapped
-        _session.instrument_id = _rig.instrument_id
-        logger.info("Writing session.json to %s", session_directory)
-        _session.write_standard_file(Path(session_directory))
-        logger.info("Writing rig.json to %s", session_directory)
-        _rig.write_standard_file(Path(session_directory))
-
-    return _run
-
-
 def _get_water_calibration(rig_model: AindVrForagingRig) -> List[measurements.VolumeCalibration]:
     def _mapper(
         device_name: Optional[str], water_calibration: AbsCalibration.water_valve.WaterValveCalibration