Skip to content

Commit 834beb7

Browse files
bruno-f-cruzbruno-f-cruz
authored
Merge pull request #190 from AllenNeuralDynamics/development
Merge development branch
2 parents c975bdb + 47e531f commit 834beb7

File tree

7 files changed

+299
-353
lines changed

7 files changed

+299
-353
lines changed

‎examples/_mocks.py‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
import datetime
2+
import logging
3+
import os
4+
from typing import Any, Dict, Literal, Optional, Union
5+
6+
import git
7+
from aind_behavior_curriculum import Stage, TrainerState
8+
from aind_behavior_services.rig import AindBehaviorRigModel
9+
from aind_behavior_services.session import AindBehaviorSessionModel
10+
from aind_behavior_services.task_logic import AindBehaviorTaskLogicModel
11+
from pydantic import Field
12+
13+
from clabe.data_mapper import DataMapper
14+
15+
logger = logging.getLogger(__name__)
16+
17+
TASK_NAME = "RandomTask"
18+
LIB_CONFIG = rf"local\AindBehavior.db\{TASK_NAME}"
19+
20+
21+
### Task-specific definitions
22+
class RigModel(AindBehaviorRigModel):
23+
rig_name: str = Field(default="TestRig", description="Rig name")
24+
version: Literal["0.0.0"] = "0.0.0"
25+
26+
27+
class TaskLogicModel(AindBehaviorTaskLogicModel):
28+
version: Literal["0.0.0"] = "0.0.0"
29+
name: Literal[TASK_NAME] = TASK_NAME
30+
31+
32+
mock_trainer_state = TrainerState[Any](
33+
curriculum=None,
34+
is_on_curriculum=False,
35+
stage=Stage(name="TestStage", task=TaskLogicModel(name=TASK_NAME, task_parameters={"foo": "bar"})),
36+
)
37+
38+
39+
class MockAindDataSchemaSession:
40+
def __init__(
41+
self,
42+
computer_name: Optional[str] = None,
43+
repository: Optional[Union[os.PathLike, git.Repo]] = None,
44+
task_name: Optional[str] = None,
45+
):
46+
self.computer_name = computer_name
47+
self.repository = repository
48+
self.task_name = task_name
49+
50+
def __str__(self) -> str:
51+
return f"MockAindDataSchemaSession(computer_name={self.computer_name}, repository={self.repository}, task_name={self.task_name})"
52+
53+
54+
class DemoAindDataSchemaSessionDataMapper(DataMapper[MockAindDataSchemaSession]):
55+
def __init__(
56+
self,
57+
rig_model: RigModel,
58+
session_model: AindBehaviorSessionModel,
59+
task_logic_model: TaskLogicModel,
60+
repository: Union[os.PathLike, git.Repo],
61+
script_path: os.PathLike,
62+
session_end_time: Optional[datetime.datetime] = None,
63+
output_parameters: Optional[Dict] = None,
64+
):
65+
super().__init__()
66+
self.session_model = session_model
67+
self.rig_model = rig_model
68+
self.task_logic_model = task_logic_model
69+
self.repository = repository
70+
self.script_path = script_path
71+
self.session_end_time = session_end_time
72+
self.output_parameters = output_parameters
73+
self._mapped: Optional[MockAindDataSchemaSession] = None
74+
75+
def map(self) -> MockAindDataSchemaSession:
76+
self._mapped = MockAindDataSchemaSession(
77+
computer_name=self.rig_model.computer_name, repository=self.repository, task_name=self.task_logic_model.name
78+
)
79+
print("#" * 50)
80+
print("THIS IS MAPPED DATA!")
81+
print("#" * 50)
82+
print(self._mapped)
83+
return self._mapped
84+
85+
86+
def create_fake_subjects():
87+
subjects = ["00000", "123456"]
88+
for subject in subjects:
89+
os.makedirs(f"{LIB_CONFIG}/Subjects/{subject}", exist_ok=True)
90+
with open(f"{LIB_CONFIG}/Subjects/{subject}/task_logic.json", "w", encoding="utf-8") as f:
91+
f.write(TaskLogicModel(task_parameters={"subject": subject}).model_dump_json(indent=2))
92+
with open(f"{LIB_CONFIG}/Subjects/{subject}/trainer_state.json", "w", encoding="utf-8") as f:
93+
f.write(mock_trainer_state.model_dump_json(indent=2))
94+
95+
96+
def create_fake_rig():
97+
computer_name = os.getenv("COMPUTERNAME")
98+
os.makedirs(_dir := f"{LIB_CONFIG}/Rig/{computer_name}", exist_ok=True)
99+
with open(f"{_dir}/rig1.json", "w", encoding="utf-8") as f:
100+
f.write(RigModel().model_dump_json(indent=2))

‎examples/behavior_launcher.py‎

Lines changed: 10 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,20 @@
11
import asyncio
2-
import datetime
32
import logging
4-
import os
53
from pathlib import Path
6-
from typing import Any, Dict, Literal, Optional, Union
7-
8-
import git
9-
from aind_behavior_curriculum import Stage, TrainerState
10-
from aind_behavior_services.rig import AindBehaviorRigModel
11-
from aind_behavior_services.session import AindBehaviorSessionModel
12-
from aind_behavior_services.task_logic import AindBehaviorTaskLogicModel
13-
from pydantic import Field
4+
5+
from _mocks import (
6+
LIB_CONFIG,
7+
AindBehaviorSessionModel,
8+
DemoAindDataSchemaSessionDataMapper,
9+
RigModel,
10+
TaskLogicModel,
11+
create_fake_rig,
12+
create_fake_subjects,
13+
)
1414
from pydantic_settings import CliApp
1515

1616
from clabe import resource_monitor
1717
from clabe.apps import CurriculumApp, CurriculumSettings, PythonScriptApp
18-
from clabe.data_mapper import DataMapper
1918
from clabe.launcher import (
2019
Launcher,
2120
LauncherCliArgs,
@@ -24,74 +23,6 @@
2423

2524
logger = logging.getLogger(__name__)
2625

27-
TASK_NAME = "RandomTask"
28-
LIB_CONFIG = rf"local\AindBehavior.db\{TASK_NAME}"
29-
30-
31-
### Task-specific definitions
32-
class RigModel(AindBehaviorRigModel):
33-
rig_name: str = Field(default="TestRig", description="Rig name")
34-
version: Literal["0.0.0"] = "0.0.0"
35-
36-
37-
class TaskLogicModel(AindBehaviorTaskLogicModel):
38-
version: Literal["0.0.0"] = "0.0.0"
39-
name: Literal[TASK_NAME] = TASK_NAME
40-
41-
42-
mock_trainer_state = TrainerState[Any](
43-
curriculum=None,
44-
is_on_curriculum=False,
45-
stage=Stage(name="TestStage", task=TaskLogicModel(name=TASK_NAME, task_parameters={"foo": "bar"})),
46-
)
47-
48-
49-
class MockAindDataSchemaSession:
50-
def __init__(
51-
self,
52-
computer_name: Optional[str] = None,
53-
repository: Optional[Union[os.PathLike, git.Repo]] = None,
54-
task_name: Optional[str] = None,
55-
):
56-
self.computer_name = computer_name
57-
self.repository = repository
58-
self.task_name = task_name
59-
60-
def __str__(self) -> str:
61-
return f"MockAindDataSchemaSession(computer_name={self.computer_name}, repository={self.repository}, task_name={self.task_name})"
62-
63-
64-
class DemoAindDataSchemaSessionDataMapper(DataMapper[MockAindDataSchemaSession]):
65-
def __init__(
66-
self,
67-
rig_model: RigModel,
68-
session_model: AindBehaviorSessionModel,
69-
task_logic_model: TaskLogicModel,
70-
repository: Union[os.PathLike, git.Repo],
71-
script_path: os.PathLike,
72-
session_end_time: Optional[datetime.datetime] = None,
73-
output_parameters: Optional[Dict] = None,
74-
):
75-
super().__init__()
76-
self.session_model = session_model
77-
self.rig_model = rig_model
78-
self.task_logic_model = task_logic_model
79-
self.repository = repository
80-
self.script_path = script_path
81-
self.session_end_time = session_end_time
82-
self.output_parameters = output_parameters
83-
self._mapped: Optional[MockAindDataSchemaSession] = None
84-
85-
def map(self) -> MockAindDataSchemaSession:
86-
self._mapped = MockAindDataSchemaSession(
87-
computer_name=self.rig_model.computer_name, repository=self.repository, task_name=self.task_logic_model.name
88-
)
89-
print("#" * 50)
90-
print("THIS IS MAPPED DATA!")
91-
print("#" * 50)
92-
print(self._mapped)
93-
return self._mapped
94-
9526

9627
async def experiment(launcher: Launcher) -> None:
9728
monitor = resource_monitor.ResourceMonitor(
@@ -145,23 +76,6 @@ def fmt(value: str) -> str:
14576
return
14677

14778

148-
def create_fake_subjects():
149-
subjects = ["00000", "123456"]
150-
for subject in subjects:
151-
os.makedirs(f"{LIB_CONFIG}/Subjects/{subject}", exist_ok=True)
152-
with open(f"{LIB_CONFIG}/Subjects/{subject}/task_logic.json", "w", encoding="utf-8") as f:
153-
f.write(TaskLogicModel(task_parameters={"subject": subject}).model_dump_json(indent=2))
154-
with open(f"{LIB_CONFIG}/Subjects/{subject}/trainer_state.json", "w", encoding="utf-8") as f:
155-
f.write(mock_trainer_state.model_dump_json(indent=2))
156-
157-
158-
def create_fake_rig():
159-
computer_name = os.getenv("COMPUTERNAME")
160-
os.makedirs(_dir := f"{LIB_CONFIG}/Rig/{computer_name}", exist_ok=True)
161-
with open(f"{_dir}/rig1.json", "w", encoding="utf-8") as f:
162-
f.write(RigModel().model_dump_json(indent=2))
163-
164-
16579
def main():
16680
create_fake_subjects()
16781
create_fake_rig()

‎src/clabe/apps/_bonsai.py‎

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
1+
import hashlib
12
import logging
23
import os
4+
import random
35
from os import PathLike
46
from pathlib import Path
57
from typing import Dict, Optional
68

9+
import pydantic
710
from aind_behavior_services import AindBehaviorRigModel, AindBehaviorSessionModel, AindBehaviorTaskLogicModel
811

9-
from clabe.launcher._base import Launcher
10-
1112
from ._base import Command, CommandResult, ExecutableApp, identity_parser
1213
from ._executors import _DefaultExecutorMixin
1314

@@ -177,7 +178,7 @@ def __init__(
177178
self,
178179
workflow: os.PathLike,
179180
*,
180-
launcher: Launcher,
181+
temp_directory: os.PathLike,
181182
rig: Optional[AindBehaviorRigModel] = None,
182183
session: Optional[AindBehaviorSessionModel] = None,
183184
task_logic: Optional[AindBehaviorTaskLogicModel] = None,
@@ -232,14 +233,35 @@ def __init__(
232233
"""
233234
additional_externalized_properties = kwargs.pop("additional_externalized_properties", {}) or {}
234235
if rig:
235-
additional_externalized_properties["RigPath"] = os.path.abspath(launcher.save_temp_model(model=rig))
236+
additional_externalized_properties["RigPath"] = os.path.abspath(self._save_temp_model(model=rig))
236237
if session:
237-
additional_externalized_properties["SessionPath"] = os.path.abspath(launcher.save_temp_model(model=session))
238+
additional_externalized_properties["SessionPath"] = os.path.abspath(self._save_temp_model(model=session))
238239
if task_logic:
239240
additional_externalized_properties["TaskLogicPath"] = os.path.abspath(
240-
launcher.save_temp_model(model=task_logic)
241+
self._save_temp_model(model=task_logic)
241242
)
242243
super().__init__(
243244
workflow=workflow, additional_externalized_properties=additional_externalized_properties, **kwargs
244245
)
245-
self._launcher = launcher
246+
self._temp_directory = Path(temp_directory)
247+
248+
def _save_temp_model(self, model: pydantic.BaseModel) -> Path:
249+
"""
250+
Saves a temporary JSON representation of a pydantic model.
251+
252+
Args:
253+
model: The pydantic model to save
254+
directory: The directory to save the file in.
255+
256+
Returns:
257+
Path: The path to the saved file
258+
"""
259+
self._temp_directory.mkdir(parents=True, exist_ok=True)
260+
261+
random_data = str(random.random()).encode("utf-8")
262+
sha_hash = hashlib.sha256(random_data).hexdigest()[:8]
263+
264+
fpath = self._temp_directory / f"{model.__class__.__name__}_{sha_hash}.json"
265+
with open(fpath, "w+", encoding="utf-8") as f:
266+
f.write(model.model_dump_json(indent=2))
267+
return Path(fpath)

0 commit comments

Comments
 (0)