Skip to content

Commit

Permalink
Merge pull request #63 from BlueBrain/hoc-workflow
Browse files Browse the repository at this point in the history
integrate ExportHoc into luigi workflow
  • Loading branch information
AurelienJaquier authored Oct 10, 2023
2 parents 7de83c1 + 897c7aa commit a503205
Show file tree
Hide file tree
Showing 6 changed files with 221 additions and 12 deletions.
10 changes: 10 additions & 0 deletions bluepyemodel/access_point/local.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@
from bluepyemodel.evaluation.evaluator import LEGACY_PRE_PROTOCOLS
from bluepyemodel.evaluation.evaluator import PRE_PROTOCOLS
from bluepyemodel.evaluation.fitness_calculator_configuration import FitnessCalculatorConfiguration
from bluepyemodel.export_emodel.utils import get_hoc_file_path
from bluepyemodel.export_emodel.utils import get_output_path_from_metadata
from bluepyemodel.model.mechanism_configuration import MechanismConfiguration
from bluepyemodel.model.neuron_model_configuration import NeuronModelConfiguration
from bluepyemodel.tools.mechanisms import get_mechanism_currents
Expand Down Expand Up @@ -826,3 +828,11 @@ def store_emodels_sonata(
):
"""Not Implemented"""
raise NotImplementedError

def sonata_exists(self, seed):
"""Returns True if the sonata hoc file has been exported"""
output_path = get_output_path_from_metadata(
"export_emodels_sonata", self.emodel_metadata, seed
)
hoc_file_path = get_hoc_file_path(output_path)
return Path(hoc_file_path).is_file()
5 changes: 4 additions & 1 deletion bluepyemodel/evaluation/evaluation.py
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,10 @@ def get_evaluator_from_access_point(

mechanisms_directory = access_point.get_mechanisms_directory()
if isinstance(access_point, LocalAccessPoint):
if Path.cwd() != access_point.emodel_dir and access_point.emodel_metadata.iteration:
if (
Path.cwd() != access_point.emodel_dir.resolve()
and access_point.emodel_metadata.iteration
):
delete_compiled_mechanisms()
if not (access_point.emodel_dir / "x86_64" / "special").is_file():
compile_mechs_in_emodel_dir(mechanisms_directory)
Expand Down
15 changes: 14 additions & 1 deletion bluepyemodel/export_emodel/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,17 @@
logger = logging.getLogger(__name__)


def get_output_path_from_metadata(output_base_dir, emodel_metadata, seed):
"""Get the output path from the emodel_metadata.
Args:
output_base_dir (str): output base directory
emodel_metadata (EModelMetadata): emodel metadata
seed (int): seed
"""
return f"./{output_base_dir}/{emodel_metadata.as_string(seed=seed)}/"


def get_output_path(emodel, output_dir=None, output_base_dir="export_emodels_hoc"):
"""Get the output path.
Expand All @@ -32,7 +43,9 @@ def get_output_path(emodel, output_dir=None, output_base_dir="export_emodels_hoc
using also emodel metadata in the path
"""
if output_dir is None:
output_dir = f"./{output_base_dir}/{emodel.emodel_metadata.as_string(seed=emodel.seed)}/"
output_dir = get_output_path_from_metadata(
output_base_dir, emodel.emodel_metadata, emodel.seed
)
output_path = pathlib.Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)

Expand Down
4 changes: 0 additions & 4 deletions bluepyemodel/tasks/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@ class EmodelAPIConfig(luigi.Config):
recipes_path = luigi.OptionalParameter(default=None)
final_path = luigi.OptionalParameter(default="./final.json")
legacy_dir_structure = luigi.BoolParameter(default=False)
extract_config = luigi.OptionalParameter(default=None)

# nexus parameters
forge_path = luigi.OptionalParameter(default=None)
Expand All @@ -45,15 +44,12 @@ def __init__(self, *args, **kwargs):
if self.api == "local":
if self.recipes_path is None:
raise ValueError("recipes_path cannot be None when api is set to 'local'")
if self.extract_config is None:
raise ValueError("extract_config cannot be None when api is set to 'local'")

self.api_args = {
"emodel_dir": self.emodel_dir,
"recipes_path": self.recipes_path,
"final_path": self.final_path,
"legacy_dir_structure": self.legacy_dir_structure,
"extract_config": self.extract_config,
}

if self.api == "nexus":
Expand Down
197 changes: 192 additions & 5 deletions bluepyemodel/tasks/emodel_creation/optimisation.py
Original file line number Diff line number Diff line change
Expand Up @@ -295,8 +295,7 @@ def remote_script(self):
'{"emodel_dir": None',
'"recipes_path": None',
'"final_path": None',
'"legacy_dir_structure": None',
'"extract_config": None}',
'"legacy_dir_structure": None}',
)
),
type=json.loads,
Expand Down Expand Up @@ -611,8 +610,7 @@ def remote_script(self):
'{"emodel_dir": None',
'"recipes_path": None',
'"final_path": None',
'"legacy_dir_structure": None',
'"extract_config": None}',
'"legacy_dir_structure": None}',
)
),
type=json.loads,
Expand Down Expand Up @@ -658,6 +656,194 @@ def output(self):
)


class StoreHocTarget(WorkflowTarget):
"""Check if the hoc files have been stored."""

def __init__(self, emodel, etype, ttype, mtype, species, brain_region, iteration_tag, seed):
"""Constructor.
Args:
emodel (str): name of the emodel. Has to match the name of the emodel
under which the configuration data are stored.
ttype (str): name of the ttype.
seed (int): seed used in the optimisation.
"""
super().__init__(
emodel=emodel,
etype=etype,
ttype=ttype,
mtype=mtype,
species=species,
brain_region=brain_region,
iteration_tag=iteration_tag,
)

self.seed = seed

def exists(self):
"""Check if the hoc is stored for all given seeds."""
batch_size = self.access_point.pipeline_settings.optimisation_batch_size
checked_for_all_seeds = [
self.access_point.sonata_exists(seed=seed)
for seed in range(self.seed, self.seed + batch_size)
]
return all(checked_for_all_seeds)


class ExportHoc(WorkflowTaskRequiringMechanisms, IPyParallelTask):
"""Luigi wrapper for export_emodels_sonata.
Parameters:
seed (int): seed used in the optimisation.
graceful_killer (multiprocessing.Event): event triggered when USR1 signal is received.
Has to use multiprocessing event for communicating between processes
when there is more than 1 luigi worker. Skip task if set.
"""

seed = luigi.IntParameter(default=1)
graceful_killer = multiprocessing.Event()

def __init__(self, *args, **kwargs):
""" """
super().__init__(*args, **kwargs)

# set self.batch_size here to easily handle it in the argparse argument passing
self.batch_size = self.access_point.pipeline_settings.optimisation_batch_size

def requires(self):
""" """

compile_mechanisms = self.access_point.pipeline_settings.compile_mechanisms

to_run = [
StoreBestModels(
emodel=self.emodel,
etype=self.etype,
ttype=self.ttype,
mtype=self.mtype,
species=self.species,
brain_region=self.brain_region,
iteration_tag=self.iteration_tag,
seed=self.seed,
)
]
if compile_mechanisms:
to_run.append(
CompileMechanisms(
emodel=self.emodel,
etype=self.etype,
ttype=self.ttype,
mtype=self.mtype,
species=self.species,
brain_region=self.brain_region,
iteration_tag=self.iteration_tag,
)
)

return to_run

@WorkflowTask.check_mettypes
def run(self):
"""Prepare self.args, then call bbp-workflow's IPyParallelTask's run()."""
attrs = [
"backend",
"emodel",
"etype",
"ttype",
"mtype",
"species",
"brain_region",
"iteration_tag",
"seed",
"batch_size",
]

self.prepare_args_for_remote_script(attrs)

super().run()

def remote_script(self):
"""Catch arguments from parsing, and run validation."""
# -- imports -- #
import argparse
import json

from bluepyemodel import access_point
from bluepyemodel.export_emodel.export_emodel import export_emodels_sonata
from bluepyemodel.tools.multiprocessing import get_mapper

# -- parsing -- #
parser = argparse.ArgumentParser()
parser.add_argument("--backend", default=None, type=str)
parser.add_argument("--api_from_config", default="local", type=str)
parser.add_argument(
"--api_args_from_config",
default=", ".join(
(
'{"emodel_dir": None',
'"recipes_path": None',
'"final_path": None',
'"legacy_dir_structure": None}',
)
),
type=json.loads,
)
parser.add_argument("--emodel", default=None, type=str)
parser.add_argument("--etype", default=None, type=str)
parser.add_argument("--ttype", default=None, type=str)
parser.add_argument("--mtype", default=None, type=str)
parser.add_argument("--species", default=None, type=str)
parser.add_argument("--brain_region", default=None, type=str)
parser.add_argument("--iteration_tag", default=None, type=str)
parser.add_argument("--ipyparallel_profile", default=None, type=str)
parser.add_argument("--seed", default=None, type=int)
parser.add_argument("--batch_size", default=None, type=int)

args = parser.parse_args()

# -- run validation -- #
mapper = get_mapper(args.backend, ipyparallel_profile=args.ipyparallel_profile)
access_pt = access_point.get_access_point(
access_point=args.api_from_config,
emodel=args.emodel,
etype=args.etype,
ttype=args.ttype,
mtype=args.mtype,
species=args.species,
brain_region=args.brain_region,
iteration_tag=args.iteration_tag,
**args.api_args_from_config,
)

export_emodels_sonata(
access_pt,
only_validated=False,
only_best=False,
seeds=list(range(args.seed, args.seed + args.batch_size)),
map_function=mapper,
)
if args.api_from_config == "nexus":
access_pt.store_emodels_sonata(
only_validated=False,
only_best=False,
seeds=list(range(args.seed, args.seed + args.batch_size)),
map_function=mapper,
)

def output(self):
""" """
return StoreHocTarget(
emodel=self.emodel,
etype=self.etype,
ttype=self.ttype,
mtype=self.mtype,
species=self.species,
brain_region=self.brain_region,
iteration_tag=self.iteration_tag,
seed=self.seed,
)


class EModelCreationTarget(WorkflowTarget):
"""Check if the the model is validated for any seed."""

Expand Down Expand Up @@ -717,10 +903,11 @@ class EModelCreation(WorkflowTask):
when there is more than 1 luigi worker. Exit loop if set.
"""

# pylint: disable=no-self-argument, not-callable

seed = luigi.IntParameter(default=1)
graceful_killer = multiprocessing.Event()

@staticmethod
def check_mettypes(func):
"""Decorator to check mtype, etype and ttype presence on nexus"""

Expand Down
2 changes: 1 addition & 1 deletion bluepyemodel/tools/mechanisms.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ def compile_mechs_in_emodel_dir(mechanisms_directory):
mechanisms_dir (Path): path to the directory containing the
mod files to compile.
"""
# pylint: disable=broad-exception-caught
# pylint: disable=broad-except
cwd = os.getcwd()

try:
Expand Down

0 comments on commit a503205

Please sign in to comment.