diff --git a/.github/scripts/dispatch.py b/.github/scripts/dispatch.py index 9083f45ed..b0b19778f 100644 --- a/.github/scripts/dispatch.py +++ b/.github/scripts/dispatch.py @@ -38,6 +38,8 @@ def sanitize(s: str) -> str: COVERALLS_TOKEN = os.environ['COVERALLS_REPO_TOKEN'] PR_NUMBER = os.environ.get('PR_NUMBER', '') JOB_TIMEOUT = int(os.environ.get('JOB_TIMEOUT_SECONDS', '7200')) +OPENALYX_PASSWORD = os.environ.get('OPENALYX_PASSWORD', '') +OPENALYX_USER = os.environ.get('OPENALYX_USER', 'intbrainlab') # Base image for the Lightning job. python:3.12 verified to ship git/bash/pip. # uv (installed at runtime) provides the actual test Python via PY_VERSION, so # this base Python version is only used to bootstrap `pip install uv`. @@ -83,6 +85,8 @@ def sanitize(s: str) -> str: 'PY_VERSION': PY_VERSION, 'REPO_URL': REPO_URL, 'INTEGRATION_DATA_DIR': INTEGRATION_DATA_DIR, + 'OPENALYX_PASSWORD': OPENALYX_PASSWORD, + 'OPENALYX_USER': OPENALYX_USER, # --- Coveralls auth + parallel grouping --- 'COVERALLS_REPO_TOKEN': COVERALLS_TOKEN, 'COVERALLS_PARALLEL': 'true', diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml index 91388593a..2fae3c69d 100644 --- a/.github/workflows/integration-tests.yml +++ b/.github/workflows/integration-tests.yml @@ -35,6 +35,7 @@ jobs: # --- consumed by dispatch.py / Job.run --- LIGHTNING_TEAMSPACE: ${{ vars.LIGHTNING_TEAMSPACE }} # "owner/teamspace" INTEGRATION_DATA_DIR: ${{ vars.INTEGRATION_DATA_DIR }} + OPENALYX_PASSWORD: ${{ secrets.ALYX_PWD }} REPO_URL: https://github.com/${{ github.repository }}.git COVERALLS_REPO_TOKEN: ${{ secrets.COVERALLS_REPO_TOKEN }} PY_VERSION: ${{ matrix.python_version }} diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index f3fb87650..e6d958cfc 100644 --- a/.github/workflows/unit-tests.yml +++ b/.github/workflows/unit-tests.yml @@ -39,6 +39,8 @@ jobs: python-version: 3.13 env: ONE_SAVE_ON_DELETE: false + OPENALYX_USER: intbrainlab + OPENALYX_PASSWORD: ${{ secrets.ALYX_PWD }} steps: - uses: actions/checkout@v7 - name: Set up Python ${{ matrix.python-version }} diff --git a/CHANGELOG.md b/CHANGELOG.md index 73f4ca1cf..5ced37cc7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,9 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [3.5.0] Unreleased +### Added +- `ibllib.pipes.spec.TaskSpec`: plain-data task specification for creating Alyx tasks without importing the task class +- `ibllib.pipes.routing`: route tasks to environments by their executable; `ENV_PATHS` maps environment labels to server venvs +- `ibllib.pipes.plan`: plan the tasks of other repositories (e.g. mpci) in the environment that runs them, via `python -m ibllib.pipes.plan` + ### Changed - SpikeSortingLoaders - merging clusters with channels doesn't require reading spikes if metrics are available +- `local_server.task_queue` and `list_queued_envs` determine a task's environment from its executable and only import task classes of the requested environments +- `dynamic_pipeline.make_pipeline` no longer imports mpci; mesoscope tasks are planned by `mpci.alyx.pipeline:plan` in the mpci env (requires mpci with the `plan` function) +- `local_server.job_creator` keeps the raw_session.flag file if an external task planner fails, so the missing tasks are created on the next run +- `Pipeline.create_alyx_tasks` creates tasks in dependency order and computes task levels from the parents + +### Fixed +- Task time outs are now stored on Alyx: tasks were posted with the key `time_out_sec` instead of the model field `time_out_secs`. `Pipeline.create_alyx_tasks` asserts that time outs don't exceed the Alyx field maximum (32767 s) before creating any tasks + +### Removed +- `dynamic_pipeline.get_mesoscope_tasks` ## [4.0.1] 2026-05-22 diff --git a/ibllib/ephys/sync_probes.py b/ibllib/ephys/sync_probes.py index d804bef81..02bf8285f 100644 --- a/ibllib/ephys/sync_probes.py +++ b/ibllib/ephys/sync_probes.py @@ -84,7 +84,7 @@ def get_sync_fronts(auxiliary_name): # exits if sync label not found for current probe if auxiliary_name not in sync_map: return - isync = np.in1d(sync['channels'], np.array([sync_map[auxiliary_name]])) + isync = np.isin(sync['channels'], np.array([sync_map[auxiliary_name]])) # only returns syncs if we get fronts for all probes if np.all(~isync): return diff --git a/ibllib/pipes/__init__.py b/ibllib/pipes/__init__.py index fe3b72d43..fe624814c 100644 --- a/ibllib/pipes/__init__.py +++ b/ibllib/pipes/__init__.py @@ -8,6 +8,10 @@ inferring the acquisition hardware from the task protocol. The new session's pipeline tasks are then registered for another process (or server) to query. +Tasks from repositories installed in other environments (e.g. mpci) are planned within their own +environment as plain-data :class:`spec.TaskSpec` objects (see :mod:`plan`), so the job creator +doesn't need to import them. + Another process calls :func:`local_server.task_queue` to get a list of queued tasks from Alyx, then :func:`local_server.tasks_runner` to loop through tasks. Each task is run by calling :func:`tasks.run_alyx_task` with a dictionary of task information, including the Task class and its diff --git a/ibllib/pipes/dynamic_pipeline.py b/ibllib/pipes/dynamic_pipeline.py index a4f521e01..1f3a60de8 100644 --- a/ibllib/pipes/dynamic_pipeline.py +++ b/ibllib/pipes/dynamic_pipeline.py @@ -34,6 +34,7 @@ import ibllib.io.raw_data_loaders as rawio import ibllib.io.session_params as sess_params import ibllib.pipes.tasks as mtasks +from ibllib.pipes.plan import get_external_tasks import ibllib.pipes.base_tasks as bstasks import ibllib.pipes.widefield_tasks as wtasks import ibllib.pipes.sync_tasks as stasks @@ -516,8 +517,10 @@ def get_audio_tasks(acquisition_description, **kwargs): devices = acquisition_description.get('devices', {}) audio_tasks = OrderedDict() if 'microphone' in devices: - ((microphone, micro_kwargs),) = devices['microphone'].items() - micro_kwargs['device_collection'] = micro_kwargs.pop('collection') + ((microphone, micro_info),) = devices['microphone'].items() + # Rename the collection key without modifying the acquisition description + micro_kwargs = {k: v for k, v in micro_info.items() if k != 'collection'} + micro_kwargs['device_collection'] = micro_info['collection'] if sync_kwargs['sync'] == 'bpod': audio_tasks['AudioRegisterRaw'] = type('AudioRegisterRaw', (atasks.AudioSync,), {})( **kwargs, **sync_kwargs, **micro_kwargs, collection=micro_kwargs['device_collection'] @@ -534,8 +537,10 @@ def get_wfield_tasks(acquisition_description, sync_tasks, **kwargs): wfield_tasks = OrderedDict() if 'widefield' in devices: - ((_, wfield_kwargs),) = devices['widefield'].items() - wfield_kwargs['device_collection'] = wfield_kwargs.pop('collection') + ((_, wfield_info),) = devices['widefield'].items() + # Rename the collection key without modifying the acquisition description + wfield_kwargs = {k: v for k, v in wfield_info.items() if k != 'collection'} + wfield_kwargs['device_collection'] = wfield_info['collection'] wfield_tasks['WideFieldRegisterRaw'] = type('WidefieldRegisterRaw', (wtasks.WidefieldRegisterRaw,), {})( **kwargs, **wfield_kwargs ) @@ -558,16 +563,6 @@ def get_wfield_tasks(acquisition_description, sync_tasks, **kwargs): return wfield_tasks -def get_mesoscope_tasks(acquisition_description, **kwargs): - if 'mesoscope' not in acquisition_description.get('devices', {}): - return OrderedDict() - - import mpci.alyx.pipeline - - pipe = mpci.alyx.pipeline.make_pipeline(acquisition_description, **kwargs) - return pipe.tasks - - def get_photometry_tasks(acquisition_description, **kwargs): devices = acquisition_description.get('devices', {}) photometry_tasks = OrderedDict() @@ -650,7 +645,10 @@ def make_pipeline(session_path, **pkwargs): Returns ------- ibllib.pipes.tasks.Pipeline - A task pipeline object. + A task pipeline object. Tasks from other repositories (see + :data:`ibllib.pipes.plan.PLANNERS`) are :class:`ibllib.pipes.spec.TaskSpec` objects, planned + in the environment that runs them. If any of these planners fail, the pipeline + `planner_errors` attribute maps the device to the error message. """ # NB: this pattern is a pattern for dynamic class creation # tasks['SyncPulses'] = type('SyncPulses', (epp.EphysPulses,), {})(session_path=session_path) @@ -694,17 +692,19 @@ def make_pipeline(session_path, **pkwargs): wfield_tasks = get_wfield_tasks(acquisition_description, sync_parent_tasks, **kwargs) tasks.update(wfield_tasks) - # Mesoscope tasks - mesoscope_tasks = get_mesoscope_tasks(acquisition_description, **kwargs) - tasks.update(mesoscope_tasks) - # photometry tasks # photometry_tasks = get_photometry_tasks(acquisition_description, **kwargs) # tasks.update(photometry_tasks) + # Tasks from other repositories (e.g. mesoscope), planned as specs in the env that runs them + context = {'tasks': [t.to_spec().to_dict() for t in tasks.values()]} + external_tasks, planner_errors = get_external_tasks(acquisition_description, session_path, context=context) + tasks.update(external_tasks) + # combine: make pipeline and add tasks p = mtasks.Pipeline(session_path=session_path, **pkwargs) p.tasks = tasks + p.planner_errors = planner_errors return p diff --git a/ibllib/pipes/local_server.py b/ibllib/pipes/local_server.py index d1d570101..9e5d7fcbd 100644 --- a/ibllib/pipes/local_server.py +++ b/ibllib/pipes/local_server.py @@ -22,6 +22,7 @@ from ibllib import __version__ as ibllib_version from ibllib.pipes import tasks +from ibllib.pipes.routing import task_env from ibllib.time import date2isostr from ibllib.oneibl.registration import IBLRegistrationClient from ibllib.oneibl.data_handlers import get_local_data_repository @@ -133,7 +134,15 @@ def job_creator(root_path, one=None, dry=False, rerun=False): else: rerun__status__in = ['Waiting'] pipe.create_alyx_tasks(rerun__status__in=rerun__status__in) - flag_file.unlink() + if pipe.planner_errors: + # Keep the flag file so that the missing tasks are created on the next run + _logger.error( + 'Failed to plan %s tasks for session %s; keeping flag file to retry', + ', '.join(pipe.planner_errors), + session_path.relative_to(root_path), + ) + else: + flag_file.unlink() if pipe is not None: pipes.append(pipe) except Exception: @@ -147,6 +156,9 @@ def list_available_envs(root=Path.home() / 'Documents/PYTHON/envs'): """ List all the envs within `root` dir. + NB: Environment labels don't necessarily match the venv directory names; use + :func:`ibllib.pipes.routing.installed_envs` to list the installed environment labels. + Parameters ---------- root : str, pathlib.Path @@ -164,27 +176,92 @@ def list_available_envs(root=Path.home() / 'Documents/PYTHON/envs'): return [None] -def list_queued_envs(one=None): +def _waiting_tasks(alyx, lab=None): + """ + Query the waiting tasks of a lab for the local data repository. + + Parameters + ---------- + alyx : one.webclient.AlyxClient + An Alyx instance. + lab : str + Lab name as per Alyx, otherwise try to infer from local Globus install. + + Returns + ------- + list of dict, None + A list of Alyx tasks with a 'Waiting' status, or None if the lab could not be determined. + """ + if lab is None: + _logger.debug('Trying to infer lab from globus installation') + lab = get_lab_from_endpoint_id(alyx=alyx) + if lab is None: + _logger.error('No lab provided or found') + return # if the lab is none, this will return empty tasks each time + data_repo = get_local_data_repository(alyx) + return alyx.rest( + 'tasks', 'list', status='Waiting', django=f'session__lab__name__in,{lab},data_repository__name,{data_repo}', no_cache=True + ) + + +def list_queued_envs(one=None, lab=None): """ The set of all envs in the list of waiting tasks. + The environment of each task is determined from its executable, without importing the task class. + + Parameters + ---------- + one : one.api.OneAlyx + An instance of ONE. + lab : str + Lab name as per Alyx, otherwise try to infer from local Globus install. + Returns ------- set All environments required to process waiting tasks. """ one = one or ONE(mode='remote', cache_rest=None) - waiting_tasks = task_queue(mode='large', alyx=one.alyx, env=list_available_envs()) - envs_in_queue = set() - for task_exe in map(lambda x: x['executable'], waiting_tasks): - envs_in_queue.add(tasks.str2class(task_exe).env) - return envs_in_queue + return {task_env(t['executable']) for t in _waiting_tasks(one.alyx, lab=lab) or []} + + +def is_job_size(task, mode): + """ + Check whether a task is of a given job size. + + NB: This imports the task class. + + Parameters + ---------- + task : dict + An Alyx task dictionary. + mode : {'all', 'small', 'large'} + The job size to check. + + Returns + ------- + bool + True if the task class job size matches the mode (or mode is 'all'). False if the task + class could not be imported. + """ + if mode == 'all': + return True + try: + return tasks.str2class(task['executable']).job_size == mode + except (ImportError, AttributeError): + _logger.error('Task %s not found in this env', task['executable']) + return False def task_queue(mode='all', lab=None, alyx=None, env=(None,)): """ Query waiting jobs from the specified Lab + The environment of each task is determined from its executable (see + :func:`ibllib.pipes.routing.task_env`), so only the task classes of the given environments are + imported (to determine the job size). + Parameters ---------- mode : {'all', 'small', 'large'} @@ -193,37 +270,23 @@ def task_queue(mode='all', lab=None, alyx=None, env=(None,)): Lab name as per Alyx, otherwise try to infer from local Globus install. alyx : one.webclient.AlyxClient An Alyx instance. - env : list - One or more environments to filter by. See :prop:`ibllib.pipes.tasks.Task.env`. + env : str, list + One or more environment labels to filter by, where None is the base environment. See + :data:`ibllib.pipes.routing.ROUTES`. Returns ------- list of dict A list of Alyx tasks associated with `lab` that have a 'Waiting' status. """ - - def predicate(task): - try: - classe = tasks.str2class(task['executable']) - return (mode == 'all' or classe.job_size == mode) and classe.env in env - except ModuleNotFoundError: - _logger.error('Task %s not found in this env', task['executable']) - return False - + env = (env,) if env is None or isinstance(env, str) else tuple(env) alyx = alyx or AlyxClient(cache_rest=None) - if lab is None: - _logger.debug('Trying to infer lab from globus installation') - lab = get_lab_from_endpoint_id(alyx=alyx) - if lab is None: - _logger.error('No lab provided or found') - return # if the lab is none, this will return empty tasks each time - data_repo = get_local_data_repository(alyx) - # Filter for tasks - waiting_tasks = alyx.rest( - 'tasks', 'list', status='Waiting', django=f'session__lab__name__in,{lab},data_repository__name,{data_repo}', no_cache=True - ) - # Filter tasks by size - filtered_tasks = filter(predicate, waiting_tasks) + waiting_tasks = _waiting_tasks(alyx, lab=lab) + if waiting_tasks is None: + return + # Filter tasks by environment, then by size + filtered_tasks = (t for t in waiting_tasks if task_env(t['executable']) in env) + filtered_tasks = filter(lambda t: is_job_size(t, mode), filtered_tasks) # Order tasks by priority sorted_tasks = sorted(filtered_tasks, key=lambda d: d['priority'], reverse=True) diff --git a/ibllib/pipes/plan.py b/ibllib/pipes/plan.py new file mode 100644 index 000000000..105d7d134 --- /dev/null +++ b/ibllib/pipes/plan.py @@ -0,0 +1,233 @@ +"""Plan pipeline tasks from other repositories in the environment that runs them. + +Some pipeline tasks are defined in repositories that are installed in their own environment (e.g. +the mesoscope tasks in mpci). Instead of importing these repositories when creating a pipeline, +their tasks are planned by running a *planner* function within the task environment, which returns +plain-data :class:`ibllib.pipes.spec.TaskSpec` objects. + +A planner is a function with the signature ``planner(session_path, context=None)`` that returns a +:class:`ibllib.pipes.tasks.Pipeline`, a map of task name to task, or a list of tasks, where each +task is a Task instance, TaskSpec or task dictionary. The context is a dictionary with the key +'tasks', a list of the other (core) pipeline task specs as dicts, whose names may be used as +parents. + +The PLANNERS map determines which planner is called for a given acquisition description device. + +Examples +-------- +Plan the mesoscope tasks of a session from the command line (within the mpci environment) + +>>> python -m ibllib.pipes.plan mpci.alyx.pipeline:plan /path/to/subject/2020-01-01/001 --output specs.json + +Plan the mesoscope tasks from the base environment + +>>> specs = plan_in_env('mpci.alyx.pipeline:plan', 'mpci', session_path) +""" + +import argparse +import importlib +import importlib.util +import json +import logging +import subprocess +import sys +import tempfile +from collections import OrderedDict +from pathlib import Path + +from ibllib.pipes.routing import env_python +from ibllib.pipes.spec import TaskSpec + +_logger = logging.getLogger(__name__) + +PLANNERS = { + 'mesoscope': ('mpci', 'mpci.alyx.pipeline:plan'), +} +"""dict of tuple: Map of acquisition description device to (environment label, planner target).""" + + +class PlannerError(Exception): + """Failed to plan the tasks of an external repository.""" + + +def load_planner(target): + """ + Import a planner function. + + Parameters + ---------- + target : str + The planner function as 'module:function', e.g. 'mpci.alyx.pipeline:plan'. + + Returns + ------- + function + The planner function. + """ + module, _, function = target.partition(':') + return getattr(importlib.import_module(module), function) + + +def to_specs(tasks): + """ + Convert the output of a planner function to a list of task specs. + + Parameters + ---------- + tasks : ibllib.pipes.tasks.Pipeline, dict, list + A pipeline, a map of task name to task, or a list of tasks, where each task is a Task + instance, TaskSpec or task dictionary. + + Returns + ------- + list of TaskSpec + The task specs. + """ + tasks = getattr(tasks, 'tasks', tasks) # Pipeline -> tasks map + tasks = tasks.values() if isinstance(tasks, dict) else tasks + specs = [] + for t in tasks: + if isinstance(t, TaskSpec): + specs.append(t) + elif isinstance(t, dict): + specs.append(TaskSpec.from_dict(t)) + else: + specs.append(t.to_spec()) + return specs + + +def plan(target, session_path, context=None): + """ + Plan tasks in the current environment. + + Parameters + ---------- + target : str + The planner function as 'module:function', e.g. 'mpci.alyx.pipeline:plan'. + session_path : str, pathlib.Path + The session path. + context : dict, optional + The planning context, see module docstring. + + Returns + ------- + list of TaskSpec + The task specs. + """ + return to_specs(load_planner(target)(Path(session_path), context=context)) + + +def plan_in_env(target, env, session_path, context=None, env_paths=None, timeout=900): + """ + Plan tasks in another environment by calling this module in a subprocess. + + Parameters + ---------- + target : str + The planner function as 'module:function', e.g. 'mpci.alyx.pipeline:plan'. + env : str + The environment label, e.g. 'mpci'. + session_path : str, pathlib.Path + The session path. + context : dict, optional + The planning context, see module docstring. + env_paths : dict, optional + A map of environment label to virtual environment location. Defaults to + ibllib.pipes.routing.ENV_PATHS. + timeout : float + The maximum time in seconds to wait for the subprocess. + + Returns + ------- + list of TaskSpec + The task specs. + + Raises + ------ + PlannerError + The environment is not installed or the planner failed. + """ + if not (python := env_python(env, env_paths)): + raise PlannerError(f'Environment "{env}" not installed') + with tempfile.TemporaryDirectory() as tmp: + context_file, output_file = Path(tmp, 'context.json'), Path(tmp, 'specs.json') + context_file.write_text(json.dumps(context or {})) + cmd = [str(python), '-m', 'ibllib.pipes.plan', target, str(session_path)] + cmd += ['--context', str(context_file), '--output', str(output_file)] + _logger.info('Planning %s tasks in "%s" env', target, env) + try: + process = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) + except subprocess.TimeoutExpired as ex: + raise PlannerError(f'{target} timed out after {timeout}s in "{env}" env') from ex + if process.returncode != 0 or not output_file.exists(): + raise PlannerError(f'{target} failed in "{env}" env:\n{process.stderr[-5000:]}') + return [TaskSpec.from_dict(d) for d in json.loads(output_file.read_text())] + + +def get_external_tasks(acquisition_description, session_path, context=None, planners=None, env_paths=None): + """ + Plan the tasks of external repositories for the devices in an acquisition description. + + For each device in PLANNERS, the planner is called in its environment if installed, otherwise + in the current environment if the planner module is importable (e.g. on a development machine). + + Parameters + ---------- + acquisition_description : dict + The acquisition description. + session_path : str, pathlib.Path + The session path. + context : dict, optional + The planning context, see module docstring. + planners : dict, optional + A map of device to (environment label, planner target). Defaults to PLANNERS. + env_paths : dict, optional + A map of environment label to virtual environment location. Defaults to + ibllib.pipes.routing.ENV_PATHS. + + Returns + ------- + collections.OrderedDict + A map of task name to TaskSpec. + dict + A map of device to error message for any planners that failed. + """ + planners = PLANNERS if planners is None else planners + devices = acquisition_description.get('devices', {}) + specs, errors = OrderedDict(), {} + for device, (env, target) in planners.items(): + if device not in devices: + continue + try: + if env_python(env, env_paths): + device_specs = plan_in_env(target, env, session_path, context=context, env_paths=env_paths) + elif importlib.util.find_spec(target.partition(':')[0].split('.', 1)[0]): # top-level package + _logger.debug('"%s" env not installed; planning %s tasks in current env', env, device) + device_specs = plan(target, session_path, context=context) + else: + raise PlannerError(f'Environment "{env}" not installed') + except Exception as ex: + _logger.error('Failed to plan %s tasks: %s', device, ex) + errors[device] = str(ex) + continue + specs.update((s.name, s) for s in device_specs) + return specs, errors + + +def main(argv=None): + parser = argparse.ArgumentParser(description='Plan pipeline tasks and save them as JSON task specs.') + parser.add_argument('target', help="The planner function as 'module:function', e.g. 'mpci.alyx.pipeline:plan'") + parser.add_argument('session_path', type=Path, help='The session path.') + parser.add_argument('--context', type=Path, help='A JSON file containing the planning context.') + parser.add_argument('--output', type=Path, help='The JSON file to save the task specs to (default: stdout).') + args = parser.parse_args(argv) + context = json.loads(args.context.read_text()) if args.context else None + specs = json.dumps([s.to_dict() for s in plan(args.target, args.session_path, context=context)]) + if args.output: + args.output.write_text(specs) + else: + sys.stdout.write(specs) + + +if __name__ == '__main__': + main() diff --git a/ibllib/pipes/routing.py b/ibllib/pipes/routing.py new file mode 100644 index 000000000..0b13f1317 --- /dev/null +++ b/ibllib/pipes/routing.py @@ -0,0 +1,128 @@ +"""Map pipeline tasks to the Python environments that run them. + +Tasks are routed by their Alyx ``executable`` string alone so that a task queue can be filtered +without importing the task classes. This means a server only ever imports the classes of the +environment it is running in. + +The top-level package of an executable usually determines the environment (e.g. all ``mpci`` +tasks run in the ``mpci`` env). The exceptions are ibllib tasks that have their own environment, +e.g. DLC and spike sorting. Any executable without a matching route is run in the base +environment (``None``). + +The :attr:`ibllib.pipes.tasks.Task.env` class attribute must match the route of the task's +executable; this is checked when tasks are created on Alyx. + +Examples +-------- +>>> task_env('mpci.suite2p.task.MesoscopePreprocess') +'mpci' +>>> task_env('ibllib.pipes.video_tasks.DLC') +'dlc' +>>> task_env('ibllib.pipes.video_tasks.VideoCompress') is None +True +""" + +from pathlib import Path + +ROUTES = { + 'ibllib.pipes.video_tasks.DLC': 'dlc', + 'ibllib.pipes.video_tasks.LightningPose': 'litpose', + 'ibllib.pipes.video_tasks.LightningAction': 'litaction', + 'ibllib.pipes.ephys_tasks.SpikeSorting': 'iblsorter', + 'mpci': 'mpci', + 'mpci.chronic.roicat': 'roicat', +} +"""dict of str: Map of dotted executable prefix to environment label; the longest match wins.""" + +_PYTHON = Path.home() / 'Documents' / 'PYTHON' +ENV_PATHS = { + 'dlc': _PYTHON / 'envs' / 'dlcenv', + 'litpose': _PYTHON / 'envs' / 'litpose', + 'litaction': _PYTHON / 'envs' / 'litaction', + 'iblsorter': _PYTHON / 'SPIKE_SORTING' / 'ibl-sorter' / '.venv', + 'mpci': _PYTHON / 'envs' / 'suite2p', + 'roicat': _PYTHON / 'envs' / 'roicat', +} +"""dict of str: Map of environment label to the location of its virtual environment on a server.""" + + +def task_env(executable, routes=None): + """ + Return the environment label of a task executable. + + Parameters + ---------- + executable : str + A task executable, e.g. 'mpci.suite2p.task.MesoscopePreprocess'. + routes : dict of str, optional + A map of dotted executable prefix to environment label. Defaults to ROUTES. + + Returns + ------- + str, None + The environment label, or None if the task is run in the base environment. + """ + routes = ROUTES if routes is None else routes + matches = [k for k in routes if executable == k or executable.startswith(k.rstrip('.') + '.')] + return routes[max(matches, key=len)] if matches else None + + +def env_path(env, env_paths=None): + """ + Return the virtual environment location of an environment label. + + Parameters + ---------- + env : str + An environment label, e.g. 'mpci'. + env_paths : dict of str, optional + A map of environment label to virtual environment location. Defaults to ENV_PATHS. + + Returns + ------- + pathlib.Path, None + The virtual environment location, or None if the label is unknown. + """ + env_paths = ENV_PATHS if env_paths is None else env_paths + return Path(env_paths[env]) if env in env_paths else None + + +def env_python(env, env_paths=None): + """ + Return the Python executable of an environment if it is installed. + + Parameters + ---------- + env : str + An environment label, e.g. 'mpci'. + env_paths : dict of str, optional + A map of environment label to virtual environment location. Defaults to ENV_PATHS. + + Returns + ------- + pathlib.Path, None + The Python executable, or None if the environment is not installed. + """ + if not (path := env_path(env, env_paths)): + return None + # POSIX venv, Windows venv, Windows conda env / base install + candidates = (path / 'bin' / 'python', path / 'Scripts' / 'python.exe', path / 'python.exe') + return next((python for python in candidates if python.exists()), None) + + +def installed_envs(env_paths=None): + """ + Return the environment labels installed on this machine. + + Parameters + ---------- + env_paths : dict of str, optional + A map of environment label to virtual environment location. Defaults to ENV_PATHS. + + Returns + ------- + list of str + The installed environment labels, including None (the base environment). + """ + env_paths = ENV_PATHS if env_paths is None else env_paths + return [None, *sorted(k for k in env_paths if env_python(k, env_paths))] diff --git a/ibllib/pipes/spec.py b/ibllib/pipes/spec.py new file mode 100644 index 000000000..350494dfa --- /dev/null +++ b/ibllib/pipes/spec.py @@ -0,0 +1,230 @@ +"""Plain-data pipeline task specifications. + +A :class:`TaskSpec` holds everything required to create a task on Alyx: its name, executable, +run arguments, parent names and resources. Unlike a :class:`ibllib.pipes.tasks.Task` instance, +a spec can be created and posted without importing the task class, so pipelines can combine tasks +from repositories that are installed in different environments. + +This module deliberately depends only on the standard library. + +Examples +-------- +Create a spec from a task instance + +>>> spec = task.to_spec() + +Specs are JSON serializable + +>>> specs = [TaskSpec.from_dict(d) for d in json.loads(json.dumps([s.to_dict() for s in specs]))] + +Sort specs so that parents come before their children and compute the task levels + +>>> specs = sort_specs(specs) +""" + +import importlib +from dataclasses import dataclass, field, asdict, fields, replace + + +def executable_name(task): + """ + Return the executable name of a task as it should be stored in Alyx. + + When the class is created dynamically using the type() built-in function, the base class is + returned so that the class can be re-instantiated from the Alyx record. + + Parameters + ---------- + task : ibllib.pipes.tasks.Task + A task instance. + + Returns + ------- + str + The full module plus class name. + """ + if task.__module__ == 'abc': + return f'{task.__class__.__base__.__module__}.{task.__class__.__base__.__name__}' + return f'{task.__module__}.{task.name}' + + +@dataclass +class TaskSpec: + """A plain-data specification of a pipeline task.""" + + name: str + """str: The task name, unique within a session (e.g. 'Trials_ChoiceWorldTrialsNidq_00').""" + executable: str + """str: The full module plus class name of the task (e.g. 'ibllib.pipes.video_tasks.DLC').""" + arguments: dict = field(default_factory=dict) + """dict: The keyword arguments passed to the task class constructor. Must be JSON serializable.""" + parents: list = field(default_factory=list) + """list of str: The names of the parent tasks.""" + level: int = 0 + """int: The level in the pipeline hierarchy. Computed from the parents by :func:`sort_specs`.""" + priority: int = 30 + io_charge: int = 5 + gpu: int = 0 + cpu: int = 1 + ram: int = 4 + time_out_secs: int = 3600 * 2 + env: str = None + """str: The environment label of the task class. Not stored on Alyx, see :mod:`ibllib.pipes.routing`.""" + job_size: str = 'small' + """str: The job size of the task class. Not stored on Alyx.""" + + @classmethod + def from_task(cls, task): + """ + Create a spec from a task instance. + + Parameters + ---------- + task : ibllib.pipes.tasks.Task + A task instance. + + Returns + ------- + TaskSpec + The task specification. + """ + return cls( + name=task.name, + executable=executable_name(task), + arguments=task.kwargs, + parents=[p.name for p in task.parents], + level=task.level, + priority=task.priority, + io_charge=task.io_charge, + gpu=task.gpu, + cpu=task.cpu, + ram=task.ram, + time_out_secs=task.time_out_secs, + env=task.env, + job_size=task.job_size, + ) + + @classmethod + def from_dict(cls, d): + """ + Create a spec from a dictionary. + + Accepts the output of :meth:`to_dict`, :meth:`to_alyx` and + :meth:`ibllib.pipes.tasks.Pipeline.create_tasks_list_from_pipeline`. Unknown keys are + ignored. + + Parameters + ---------- + d : dict + A task dictionary. The parents must be task names, not Alyx IDs. + + Returns + ------- + TaskSpec + The task specification. + """ + d = dict(d) + if 'time_out_sec' in d: # legacy task dict key + d.setdefault('time_out_secs', d.pop('time_out_sec')) + names = {f.name for f in fields(cls)} + spec = cls(**{k: v for k, v in d.items() if k in names}) + spec.arguments = spec.arguments or {} + spec.parents = list(spec.parents or []) + return spec + + def to_dict(self): + """dict: A JSON serializable dictionary of the spec.""" + return asdict(self) + + def to_alyx(self, parents=None, **kwargs): + """ + Return the dictionary used to create the task on Alyx. + + Parameters + ---------- + parents : list of str, optional + The parent task Alyx IDs. Defaults to the parent names. + kwargs + Other fields to set, e.g. session, graph, module, data_repository. + + Returns + ------- + dict + An Alyx task dictionary. + """ + return { + 'executable': self.executable, + 'priority': self.priority, + 'io_charge': self.io_charge, + 'gpu': self.gpu, + 'cpu': self.cpu, + 'ram': self.ram, + 'parents': self.parents if parents is None else parents, + 'level': self.level, + 'time_out_secs': self.time_out_secs, + 'status': 'Waiting', + 'log': None, + 'name': self.name, + 'arguments': self.arguments, + **kwargs, + } + + def instantiate(self, session_path, **kwargs): + """ + Instantiate the task class. This imports the task module. + + Parameters + ---------- + session_path : str, pathlib.Path + The session path. + kwargs + Extra keyword arguments passed to the task constructor, e.g. one, location. + + Returns + ------- + ibllib.pipes.tasks.Task + A task instance (without parents). + """ + module, name = self.executable.rsplit('.', 1) + task_class = getattr(importlib.import_module(module), name) + return task_class(session_path, **self.arguments, **kwargs) + + +def sort_specs(specs): + """ + Sort task specs so that parents come before their children, and compute each task's level. + + The sort is stable: specs that are already in a valid order are not moved. + + Parameters + ---------- + specs : iterable of TaskSpec + The task specs. Each parent name must be the name of another spec. + + Returns + ------- + list of TaskSpec + The sorted task specs (copies, the input specs are not modified). + + Raises + ------ + ValueError + Duplicate task names, missing parents or circular dependencies. + """ + specs = list(specs) + names = [s.name for s in specs] + if len(set(names)) != len(names): + raise ValueError(f'Duplicate task names: {sorted({n for n in names if names.count(n) > 1})}') + if missing := {p for s in specs for p in s.parents} - set(names): + raise ValueError(f'Parent tasks not found: {sorted(missing)}') + levels, out, pending = {}, [], specs + while pending: + # Take the first spec whose parents are all placed + i = next((i for i, s in enumerate(pending) if all(p in levels for p in s.parents)), None) + if i is None: + raise ValueError(f'Circular task dependencies: {sorted(s.name for s in pending)}') + spec = pending.pop(i) + spec = replace(spec, level=max((levels[p] for p in spec.parents), default=-1) + 1) + levels[spec.name] = spec.level + out.append(spec) + return out diff --git a/ibllib/pipes/tasks.py b/ibllib/pipes/tasks.py index 494e82b6b..e8e48c4d3 100644 --- a/ibllib/pipes/tasks.py +++ b/ibllib/pipes/tasks.py @@ -87,7 +87,9 @@ from ibllib.oneibl import data_handlers from ibllib.oneibl.data_handlers import get_local_data_repository from ibllib.oneibl.registration import get_lab -from iblutil.util import Bunch, flatten, ensure_list +from ibllib.pipes.spec import TaskSpec, executable_name, sort_specs +from ibllib.pipes.routing import task_env +from iblutil.util import flatten, ensure_list import one.params from one.api import ONE from one import webclient @@ -96,6 +98,8 @@ _logger = logging.getLogger(__name__) TASK_STATUS_SET = {'Waiting', 'Held', 'Started', 'Errored', 'Empty', 'Complete', 'Incomplete', 'Abandoned'} +MAX_TIME_OUT_SECS = 32767 +"""int: The maximum task time out, as stored in an Alyx SmallIntegerField.""" class Task(abc.ABC): @@ -197,6 +201,17 @@ def signature(self) -> Dict[str, List]: def name(self): return self.__class__.__name__ + def to_spec(self): + """ + Return a plain-data specification of the task for creating it on Alyx. + + Returns + ------- + ibllib.pipes.spec.TaskSpec + The task specification. + """ + return TaskSpec.from_task(self) + def path2eid(self): """ Fetch the experiment UUID from the Task session path, without using the REST cache. @@ -660,6 +675,7 @@ def __init__(self, session_path=None, one=None, eid=None, tasks=None): # eID for newer sessions may not be in cache so use remote query self.eid = one.path2eid(session_path, query_type='remote') if self.one else None self.label = self.__module__ + '.' + type(self).__name__ + self.planner_errors = {} # map of device to error message for failed external task planners self.tasks = tasks or {} if not isinstance(self.tasks, OrderedDict): self.tasks = OrderedDict(self.tasks) @@ -673,11 +689,43 @@ class to be able to re-instantiate the class from the alyx dictionary on the cli :param obj: :return: string containing the full module plus class name """ - if obj.__module__ == 'abc': - exec_name = f'{obj.__class__.__base__.__module__}.{obj.__class__.__base__.__name__}' - else: - exec_name = f'{obj.__module__}.{obj.name}' - return exec_name + return executable_name(obj) + + def task_specs(self, tasks_list=None): + """ + Return the pipeline tasks as specs, sorted so that parents come before their children. + + Parameters + ---------- + tasks_list : list of dict, TaskSpec or Task, optional + The tasks to convert. If None, uses self.tasks, whose values may be Task instances or + TaskSpec objects. Task dicts must have parent names, not Alyx IDs. + + Returns + ------- + list of ibllib.pipes.spec.TaskSpec + The sorted task specs with computed levels. + """ + specs = [] + for t in self.tasks.values() if tasks_list is None else tasks_list: + if isinstance(t, dict): + specs.append(TaskSpec.from_dict(t)) + elif isinstance(t, TaskSpec): + specs.append(t) + else: + specs.append(t.to_spec()) + for spec in specs: + if (env := task_env(spec.executable)) != spec.env: + _logger.warning( + '%s: task env "%s" does not match the env "%s" routed from its executable "%s"; ' + 'the task will be run in the "%s" env (see ibllib.pipes.routing)', + spec.name, + spec.env, + env, + spec.executable, + env, + ) + return sort_specs(specs) def make_graph(self, out_dir=None, show=True): if not out_dir: @@ -690,12 +738,11 @@ def make_graph(self, out_dir=None, show=True): e.node('root', label=self.label) e.attr('node', shape='ellipse') - for k in self.tasks: - j = self.tasks[k] + for j in self.task_specs(): if len(j.parents) == 0: e.edge('root', j.name) else: - [e.edge(p.name, j.name) for p in j.parents] + [e.edge(p, j.name) for p in j.parents] m.subgraph(e) m.attr(label=r'\n\Pre-processing\n') @@ -716,8 +763,9 @@ def create_alyx_tasks(self, rerun__status__in=None, tasks_list=None): rerun__status__in : list, str To re-run tasks if they already exist, specify one or more statuses strings to will be re-run, or '__all__' to re-run all tasks. - tasks_list : list - The list of tasks to create on Alyx. If None, uses self.tasks. + tasks_list : list of dict, TaskSpec or Task + The list of tasks to create on Alyx. If None, uses self.tasks. Task dicts must have + parent names, not Alyx IDs. Returns ------- @@ -733,56 +781,19 @@ def create_alyx_tasks(self, rerun__status__in=None, tasks_list=None): return tasks_alyx_pre = self.one.alyx.rest('tasks', 'list', session=self.eid, graph=self.name, no_cache=True) tasks_alyx = [] - # creates all the tasks by iterating through the ordered dict - - if tasks_list is not None: - task_items = tasks_list - # need to add in the session eid and the parents - else: - task_items = self.tasks.values() - - for t in task_items: + specs = self.task_specs(tasks_list) + # Check all tasks before creating any, as the Alyx field would reject these + for spec in specs: + assert spec.time_out_secs is None or spec.time_out_secs <= MAX_TIME_OUT_SECS, ( + f'{spec.name} time_out_secs of {spec.time_out_secs} exceeds maximum of {MAX_TIME_OUT_SECS}' + ) + # creates all the tasks in order, such that parents are created before their children + for spec in specs: # get the parents' alyx ids to reference in the database - if isinstance(t, dict): - t = Bunch(t) - executable = t.executable - arguments = t.arguments - t['time_out_secs'] = t['time_out_sec'] - if len(t.parents) > 0: - pnames = t.parents - else: - executable = self._get_exec_name(t) - arguments = t.kwargs - if len(t.parents): - pnames = [p.name for p in t.parents] - - if len(t.parents): - parents_ids = [ta['id'] for ta in tasks_alyx if ta['name'] in pnames] - else: - parents_ids = [] - - task_dict = { - 'executable': executable, - 'priority': t.priority, - 'io_charge': t.io_charge, - 'gpu': t.gpu, - 'cpu': t.cpu, - 'ram': t.ram, - 'module': self.label, - 'parents': parents_ids, - 'level': t.level, - 'time_out_sec': t.time_out_secs, - 'session': self.eid, - 'status': 'Waiting', - 'log': None, - 'name': t.name, - 'graph': self.name, - 'arguments': arguments, - } - if self.data_repo: - task_dict.update({'data_repository': self.data_repo}) + parents_ids = [ta['id'] for ta in tasks_alyx if ta['name'] in spec.parents] + task_dict = spec.to_alyx(parents=parents_ids, **self._alyx_fields()) # if the task already exists, patch it otherwise, create it - talyx = next(filter(lambda x: x['name'] == t.name, tasks_alyx_pre), []) + talyx = next(filter(lambda x: x['name'] == spec.name, tasks_alyx_pre), []) if len(talyx) == 0: talyx = self.one.alyx.rest('tasks', 'create', data=task_dict) elif talyx['status'] in rerun__status__in: @@ -790,44 +801,20 @@ def create_alyx_tasks(self, rerun__status__in=None, tasks_list=None): tasks_alyx.append(talyx) return tasks_alyx + def _alyx_fields(self): + """dict: The pipeline-level fields of each Alyx task dictionary.""" + fields = {'module': self.label, 'session': self.eid, 'graph': self.name} + if self.data_repo: + fields['data_repository'] = self.data_repo + return fields + def create_tasks_list_from_pipeline(self): """ From a pipeline with tasks, creates a list of dictionaries containing task description that can be used to upload to create alyx tasks :return: """ - tasks_list = [] - for k, t in self.tasks.items(): - # get the parents' alyx ids to reference in the database - if len(t.parents): - parent_names = [p.name for p in t.parents] - else: - parent_names = [] - - task_dict = { - 'executable': self._get_exec_name(t), - 'priority': t.priority, - 'io_charge': t.io_charge, - 'gpu': t.gpu, - 'cpu': t.cpu, - 'ram': t.ram, - 'module': self.label, - 'parents': parent_names, - 'level': t.level, - 'time_out_sec': t.time_out_secs, - 'session': self.eid, - 'status': 'Waiting', - 'log': None, - 'name': t.name, - 'graph': self.name, - 'arguments': t.kwargs, - } - if self.data_repo: - task_dict.update({'data_repository': self.data_repo}) - - tasks_list.append(task_dict) - - return tasks_list + return [spec.to_alyx(**self._alyx_fields()) for spec in self.task_specs()] def run(self, status__in=('Waiting',), machine=None, clobber=True, **kwargs): """ diff --git a/ibllib/tests/integration/tasks/test_dynamic_pipeline.py b/ibllib/tests/integration/tasks/test_dynamic_pipeline.py index 1c8a705b9..6ae65e286 100644 --- a/ibllib/tests/integration/tasks/test_dynamic_pipeline.py +++ b/ibllib/tests/integration/tasks/test_dynamic_pipeline.py @@ -1,18 +1,18 @@ import sys +import importlib.util import logging import shutil import tempfile from pathlib import Path -from collections import OrderedDict from one.registration import RegistrationClient from one.api import ONE from ibllib.pipes.local_server import job_creator, tasks_runner import ibllib.pipes.dynamic_pipeline as dynamic -from ibllib.pipes.tasks import Pipeline +from ibllib.pipes.plan import plan_in_env import ibllib.io.session_params as sess_params from ibllib.io.raw_data_loaders import patch_settings import unittest -from unittest.mock import patch, MagicMock +from unittest.mock import patch from ibllib.tests import base @@ -98,29 +98,35 @@ def test_photometry(self): shutil.copytree(src, self.session_path) self.check_pipeline() + @unittest.skipIf(importlib.util.find_spec('mpci') is None, 'mpci not installed') def test_mesoscope(self): - """Test that the mesoscope pipeline is created when the mesoscope device is present.""" - # get_mesoscope_tasks does a local `import mpci.alyx.pipeline`, so faking it requires - # sys.modules entries for every level of the dotted path (mpci, mpci.alyx, - # mpci.alyx.pipeline), with the parent -> child attributes wired up to match, since the - # real import machinery normally does that wiring for us. - pipe = Pipeline(session_path=self.session_path, tasks={'MesoscopeRegisterSnapshots': 'mocked_task'}) - pipeline_mock = MagicMock() - pipeline_mock.make_pipeline.return_value = pipe - alyx_mock = MagicMock(pipeline=pipeline_mock) - mpci_mock = MagicMock(alyx=alyx_mock) - fake_modules = {'mpci': mpci_mock, 'mpci.alyx': alyx_mock, 'mpci.alyx.pipeline': pipeline_mock} - with patch.dict(sys.modules, fake_modules): - experiment_description = {'devices': {'foo': {'bar': 'baz'}}} - # Without mesoscope device, the pipeline should not be created - ret = dynamic.get_mesoscope_tasks(experiment_description) - self.assertEqual(ret, OrderedDict()) - pipeline_mock.make_pipeline.assert_not_called() - # With mesoscope device, the make_pipeline should be called - experiment_description['devices']['mesoscope'] = {'collection': 'raw_imaging_data'} - ret = dynamic.get_mesoscope_tasks(experiment_description) - pipeline_mock.make_pipeline.assert_called_once_with(experiment_description) - self.assertEqual(ret, pipe.tasks) + """Test that the mesoscope tasks are planned by mpci when the mesoscope device is present.""" + src = self.data_path.joinpath('mesoscope', 'test', '2023-03-03', '002') + self.session_path.mkdir(parents=True) + shutil.copy(src / '_ibl_experiment.description.yaml', self.session_path) + for folder in src.glob('raw_task_data_*'): + shutil.copytree(folder, self.session_path / folder.name) + shutil.copy(self.folder_path / 'mesoscope' / 'pipeline_tasks.yaml', self.session_path) + # Plan the mpci tasks in the current environment + with patch('ibllib.pipes.plan.env_python', return_value=None): + self.check_pipeline() + # Plan the mpci tasks in a subprocess, using the current environment as the mpci env + with ( + patch.dict('ibllib.pipes.routing.ENV_PATHS', {'mpci': Path(sys.executable).parents[1]}), + patch('ibllib.pipes.plan.plan_in_env', wraps=plan_in_env) as plan_mock, + ): + self.check_pipeline() + plan_mock.assert_called_once() + # If the planner fails, the core tasks are still created and the error is recorded + with ( + patch('ibllib.pipes.plan.env_python', return_value=None), + patch('ibllib.pipes.plan.plan', side_effect=RuntimeError('planner failed')), + self.assertLogs('ibllib.pipes.plan', 'ERROR'), + ): + pipe = dynamic.make_pipeline(self.session_path) + self.assertEqual({'mesoscope': 'planner failed'}, pipe.planner_errors) + self.assertEqual(12, len(pipe.tasks)) + self.assertFalse(any(name.startswith('Mesoscope') for name in pipe.tasks)) def test_chained(self): """Test pipeline creation when there are multiple task protocols run within a session""" diff --git a/ibllib/tests/integration/test_brainbox_io.py b/ibllib/tests/integration/test_brainbox_io.py index b8c8677ef..8ce10f90d 100644 --- a/ibllib/tests/integration/test_brainbox_io.py +++ b/ibllib/tests/integration/test_brainbox_io.py @@ -1,5 +1,6 @@ import logging import hashlib +import os import unittest import pandas as pd @@ -26,7 +27,12 @@ def test_spike_interface(self): Those are the specifications for the spike interface tests to pass. :return: """ - one = ONE(base_url='https://openalyx.internationalbrainlab.org', silent=True, password='international') + one = ONE( + base_url='https://openalyx.internationalbrainlab.org', + silent=True, + username=os.getenv('OPENALYX_USER'), + password=os.getenv('OPENALYX_PASSWORD'), + ) pid = '80f6ffdd-f692-450f-ab19-cd6d45bfd73e' ssl = SpikeSortingLoader(pid=pid, one=one) @@ -42,7 +48,12 @@ def test_spike_interface(self): class TestReadChannels(unittest.TestCase): def test_read_channels(self): - one = ONE(base_url='https://openalyx.internationalbrainlab.org', silent=True, password='international') + one = ONE( + base_url='https://openalyx.internationalbrainlab.org', + silent=True, + username=os.getenv('OPENALYX_USER'), + password=os.getenv('OPENALYX_PASSWORD'), + ) pid = '511afaa5-fdc4-4166-b4c0-4629ec5e652e' ssl = SpikeSortingLoader(one=one, pid=pid) channels = ssl.load_channels(revision='2024-05-06') diff --git a/ibllib/tests/integration/test_open_alyx.py b/ibllib/tests/integration/test_open_alyx.py index f2e25ed0f..a9c23cd55 100644 --- a/ibllib/tests/integration/test_open_alyx.py +++ b/ibllib/tests/integration/test_open_alyx.py @@ -1,3 +1,4 @@ +import os import tempfile import unittest import logging @@ -11,13 +12,18 @@ _logger.setLevel(10) ba = AllenAtlas() -one = ONE(base_url='https://openalyx.internationalbrainlab.org', silent=True, password='international') class TestReadSpikeSorting(unittest.TestCase): def test_spike_sorting_loader(self): # insertions = one.alyx.rest('insertions', 'list') pid = 'da8dfec1-d265-44e8-84ce-6ae9c109b8bd' + one = ONE( + base_url='https://openalyx.internationalbrainlab.org', + silent=True, + username=os.getenv('OPENALYX_USER'), + password=os.getenv('OPENALYX_PASSWORD'), + ) ssl = SpikeSortingLoader(pid=pid, one=one, atlas=ba) spikes, clusters, channels = ssl.load_spike_sorting() SpikeSortingLoader.merge_clusters(spikes, clusters, channels) @@ -36,7 +42,11 @@ def test_streamer_object_with_spike_sorting_loader(self): t0 = 50 self.td = tempfile.TemporaryDirectory() tmp_one = ONE( - base_url='https://openalyx.internationalbrainlab.org', password='international', silent=True, cache_dir=self.td.name + base_url='https://openalyx.internationalbrainlab.org', + username=os.getenv('OPENALYX_USER'), + password=os.getenv('OPENALYX_PASSWORD'), + silent=True, + cache_dir=self.td.name, ) sl = SpikeSortingLoader(pid=pid, one=tmp_one) diff --git a/ibllib/tests/test_dynamic_pipeline.py b/ibllib/tests/test_dynamic_pipeline.py index 6ab73ee94..cf45b292c 100644 --- a/ibllib/tests/test_dynamic_pipeline.py +++ b/ibllib/tests/test_dynamic_pipeline.py @@ -155,3 +155,40 @@ def test_sync_label(self): self.assertEqual('timeline', dyn._sync_label('nidq', acquisition_software='timeline')) self.assertEqual('nidq', dyn._sync_label('nidq', acquisition_software='spikeglx')) self.assertEqual('tdms', dyn._sync_label('tdms')) + + def test_get_audio_tasks(self): + """Test pipes.dynamic_pipeline.get_audio_tasks function.""" + description = { + 'devices': {'microphone': {'microphone': {'collection': 'raw_task_data_00', 'sync_label': 'audio'}}}, + 'sync': {'bpod': {'collection': 'raw_task_data_00'}}, + } + expected = deepcopy(description) + session_path = Path('/subject/2020-01-01/001') + task = dyn.get_audio_tasks(description, session_path=session_path)['AudioRegisterRaw'] + self.assertEqual('AudioSync', type(task).__base__.__name__) + self.assertEqual('raw_task_data_00', task.kwargs['device_collection']) + self.assertEqual('raw_task_data_00', task.kwargs['collection']) + self.assertEqual(expected, description, 'acquisition description should not be modified') + description['sync'] = {'nidq': {'collection': 'raw_sync_data', 'acquisition_software': 'timeline'}} + expected = deepcopy(description) + task = dyn.get_audio_tasks(description, session_path=session_path)['AudioRegisterRaw'] + self.assertEqual('AudioCompress', type(task).__base__.__name__) + self.assertEqual({'sync_label': 'audio', 'device_collection': 'raw_task_data_00'}, task.kwargs) + self.assertEqual(expected, description, 'acquisition description should not be modified') + + def test_get_wfield_tasks(self): + """Test pipes.dynamic_pipeline.get_wfield_tasks function.""" + description = { + 'devices': {'widefield': {'widefield': {'collection': 'raw_widefield_data', 'sync_label': 'frame_trigger'}}}, + 'sync': {'nidq': {'collection': 'raw_sync_data', 'acquisition_software': 'spikeglx', 'extension': 'bin'}}, + } + expected = deepcopy(description) + session_path = Path('/subject/2020-01-01/001') + for _ in range(2): # previously raised a KeyError on the second call + tasks = dyn.get_wfield_tasks(description, [], session_path=session_path) + self.assertEqual(expected, description, 'acquisition description should not be modified') + expected_names = ['WideFieldRegisterRaw', 'WidefieldCompress', 'WidefieldPreprocess', 'WidefieldSync', 'WidefieldFOV'] + self.assertEqual(expected_names, list(tasks)) + kwargs = tasks['WidefieldSync'].kwargs + self.assertEqual(('raw_widefield_data', 'frame_trigger'), (kwargs['device_collection'], kwargs['sync_label'])) + self.assertNotIn('collection', kwargs) diff --git a/ibllib/tests/test_pipes.py b/ibllib/tests/test_pipes.py index 350db7dbf..33c72172f 100644 --- a/ibllib/tests/test_pipes.py +++ b/ibllib/tests/test_pipes.py @@ -46,6 +46,7 @@ def setUp(self): raw_behaviour_data.parent.joinpath('raw_session.flag').touch() fu.populate_task_settings(raw_behaviour_data, patch={'PYBPOD_PROTOCOL': 'ephys_optoChoiceWorld6.0.1'}) + @mock.patch.dict('ibllib.pipes.routing.ROUTES', {'ibllib.tests.test_pipes.EnvTask': 'suite2p'}) @mock.patch('ibllib.pipes.local_server.get_local_data_repository') def test_task_queue(self, lab_repo_mock): """Test ibllib.pipes.local_server.task_queue function.""" @@ -54,6 +55,7 @@ def test_task_queue(self, lab_repo_mock): {'executable': 'ibllib.tests.test_pipes.EnvTask', 'priority': EnvTask.priority}, # 80 {'executable': 'ibllib.pipes.ephys_tasks.SpikeSorting', 'priority': SpikeSorting.priority}, # 60 {'executable': 'ibllib.pipes.base_tasks.RegisterRawDataTask', 'priority': RegisterRawDataTask.priority}, # 100 + {'executable': 'mpci.not_installed.MesoscopeTask', 'priority': 100}, # not importable ] alyx = mock.Mock(spec=AlyxClient) alyx.rest.return_value = tasks @@ -71,8 +73,41 @@ def test_task_queue(self, lab_repo_mock): queue = local_server.task_queue(mode='small', lab='foolab', alyx=alyx, env=('suite2p',)) self.assertEqual([], queue) # Expect only register task as it's the only small job - queue = local_server.task_queue(mode='small', lab='foolab', alyx=alyx) + with mock.patch('ibllib.pipes.tasks.str2class', wraps=local_server.tasks.str2class) as str2class: + queue = local_server.task_queue(mode='small', lab='foolab', alyx=alyx) + # Only the classes of tasks in the base env should be imported + str2class.assert_called_once_with(tasks[2]['executable']) self.assertEqual([tasks[2]], queue) + # A single env label may be passed + queue = local_server.task_queue(lab='foolab', alyx=alyx, env='mpci') + self.assertEqual([tasks[3]], queue) + # An env's tasks that can't be imported are excluded when filtering by size + with self.assertLogs(local_server._logger, 'ERROR'): + queue = local_server.task_queue(mode='large', lab='foolab', alyx=alyx, env='mpci') + self.assertEqual([], queue) + # Check list_queued_envs, which returns the env labels without importing the task classes + one = mock.Mock(spec=ONE, alyx=alyx) + with mock.patch('ibllib.pipes.tasks.str2class') as str2class: + self.assertEqual({None, 'suite2p', 'iblsorter', 'mpci'}, local_server.list_queued_envs(one, lab='foolab')) + str2class.assert_not_called() + + @mock.patch('ibllib.pipes.local_server.IBLRegistrationClient') + @mock.patch('ibllib.pipes.local_server.make_pipeline') + def test_job_creator(self, make_pipeline_mock, _): + """Test the job_creator keeps the flag file if any external task planners failed.""" + make_pipeline_mock.return_value.planner_errors = {'mesoscope': 'Environment "mpci" not installed'} + flag_files = sorted(self.tmpdir.rglob('raw_session.flag')) + self.assertEqual(2, len(flag_files)) + with self.assertLogs(local_server._logger, 'ERROR') as log: + pipes, _ = local_server.job_creator(self.tmpdir, one=mock.Mock(spec=ONE)) + self.assertIn('keeping flag file', log.records[-1].getMessage()) + self.assertEqual(2, len(pipes)) + self.assertEqual(2, make_pipeline_mock.return_value.create_alyx_tasks.call_count) + self.assertTrue(all(f.exists() for f in flag_files)) + # Without errors, the flag files are removed + make_pipeline_mock.return_value.planner_errors = {} + local_server.job_creator(self.tmpdir, one=mock.Mock(spec=ONE)) + self.assertFalse(any(f.exists() for f in flag_files)) class TestPipesMisc(unittest.TestCase): diff --git a/ibllib/tests/test_pipes_spec.py b/ibllib/tests/test_pipes_spec.py new file mode 100644 index 000000000..3e186956e --- /dev/null +++ b/ibllib/tests/test_pipes_spec.py @@ -0,0 +1,286 @@ +"""Tests for the ibllib.pipes.routing, ibllib.pipes.spec and ibllib.pipes.plan modules.""" + +import importlib +import inspect +import json +import sys +import tempfile +import unittest +from collections import OrderedDict +from pathlib import Path +from unittest import mock + +import ibllib.pipes +from ibllib.pipes import routing, plan +from ibllib.pipes.spec import TaskSpec, sort_specs, executable_name +from ibllib.pipes.tasks import Task, Pipeline + + +class Task00(Task): + priority = 90 + job_size = 'large' + + def _run(self, **_): + pass + + +class Task01(Task): + env = 'foo' + + def _run(self, **_): + pass + + +def planner(session_path, context=None): + """A planner for testing ibllib.pipes.plan (must be importable in a subprocess).""" + if context and context.get('raise'): + raise RuntimeError('planner failed') + parents = [context['tasks'][0]['name']] if context and context.get('tasks') else [] + t0 = type('PlannedTask', (Task00,), {})(session_path, foo='bar', parents=[]) + return Pipeline( + session_path=session_path, + tasks={ + 'PlannedTask': t0, + 'PlannedSpec': TaskSpec('PlannedSpec', 'mpci.foo.Bar', parents=['PlannedTask', *parents], env='mpci'), + }, + ) + + +class TestRouting(unittest.TestCase): + """Tests for the ibllib.pipes.routing module.""" + + def test_task_env(self): + self.assertEqual('mpci', routing.task_env('mpci.suite2p.task.MesoscopePreprocess')) + self.assertEqual('roicat', routing.task_env('mpci.chronic.roicat.task.ROICaTTask')) # longest match + self.assertEqual('dlc', routing.task_env('ibllib.pipes.video_tasks.DLC')) + self.assertIsNone(routing.task_env('ibllib.pipes.video_tasks.DLCFoo')) # not a dotted prefix + self.assertIsNone(routing.task_env('mpcifoo.task.Task')) + self.assertIsNone(routing.task_env('ibllib.pipes.video_tasks.VideoCompress')) + self.assertEqual('bar', routing.task_env('foo.Task', routes={'foo.': 'bar'})) + + def test_installed_envs(self): + with tempfile.TemporaryDirectory() as tmp: + env_paths = {k: Path(tmp, k) for k in ('foo', 'bar', 'baz', 'win', 'conda')} + # POSIX venv, Windows venv and Windows conda env layouts + pythons = { + 'foo': ('bin', 'python'), + 'bar': ('bin', 'python'), + 'win': ('Scripts', 'python.exe'), + 'conda': ('python.exe',), + } + for env, parts in pythons.items(): + env_paths[env].joinpath(*parts).parent.mkdir(parents=True, exist_ok=True) + env_paths[env].joinpath(*parts).touch() + self.assertEqual([None, 'bar', 'conda', 'foo', 'win'], routing.installed_envs(env_paths)) + for env, parts in pythons.items(): + self.assertEqual(env_paths[env].joinpath(*parts), routing.env_python(env, env_paths)) + self.assertIsNone(routing.env_python('baz', env_paths)) + self.assertIsNone(routing.env_python('unknown', env_paths)) + + def test_ibllib_task_envs(self): + """Test each ibllib Task subclass env matches routing.task_env for its executable. + + If this fails, update ibllib.pipes.routing.ROUTES, otherwise the task will not be run. + """ + n = 0 + for file in sorted(Path(ibllib.pipes.__file__).parent.glob('*tasks*.py')): + module = importlib.import_module(f'ibllib.pipes.{file.stem}') + for name, cls in inspect.getmembers(module, inspect.isclass): + if issubclass(cls, Task) and cls.__module__ == module.__name__: + executable = f'{cls.__module__}.{name}' + with self.subTest(executable=executable): + self.assertEqual(cls.env, routing.task_env(executable)) + n += 1 + self.assertGreater(n, 50) + + +class TestTaskSpec(unittest.TestCase): + """Tests for the ibllib.pipes.spec module.""" + + def setUp(self): + self.session_path = Path('/subject/2020-01-01/001') + self.t0 = type('Task00_foo', (Task00,), {})(self.session_path, foo='bar') + self.t1 = Task01(self.session_path, parents=[self.t0]) + + def test_from_task(self): + spec = self.t0.to_spec() + self.assertEqual('Task00_foo', spec.name) + # The dynamic class base (NB: the module name depends on how the tests are run) + self.assertEqual(f'{Task00.__module__}.Task00', spec.executable) + self.assertEqual(executable_name(self.t0), spec.executable) + self.assertEqual({'foo': 'bar'}, spec.arguments) + self.assertEqual(([], 0, 90, 'large', None), (spec.parents, spec.level, spec.priority, spec.job_size, spec.env)) + spec = self.t1.to_spec() + self.assertEqual(f'{Task01.__module__}.Task01', spec.executable) + self.assertEqual((['Task00_foo'], 1, 'foo'), (spec.parents, spec.level, spec.env)) + + def test_dicts(self): + spec = self.t1.to_spec() + # Round trip through JSON + self.assertEqual(spec, TaskSpec.from_dict(json.loads(json.dumps(spec.to_dict())))) + # Alyx task dict + d = spec.to_alyx(parents=['uuid'], session='eid', graph='Pipeline') + self.assertEqual((['uuid'], 'eid', 'Pipeline', 'Waiting'), (d['parents'], d['session'], d['graph'], d['status'])) + self.assertEqual(spec.time_out_secs, d['time_out_secs']) + self.assertEqual(spec.parents, spec.to_alyx()['parents']) + # Alyx task dict (with parent names) to spec; env and job_size are not stored on Alyx + d = spec.to_alyx(session='eid') + d['time_out_secs'] = 10 + self.assertEqual(10, TaskSpec.from_dict(d).time_out_secs) + # Legacy task dicts (e.g. pipeline_tasks.yaml fixtures) use the key 'time_out_sec' + d['time_out_sec'] = d.pop('time_out_secs') + self.assertEqual(TaskSpec.from_dict(d), TaskSpec(**{**spec.to_dict(), 'env': None, 'time_out_secs': 10})) + + def test_sort_specs(self): + a, b, c, d = (TaskSpec(x, 'foo.Task') for x in 'abcd') + b.parents, c.parents, d.parents = ['c'], ['a'], ['a', 'b'] + specs = sort_specs([a, b, c, d]) + self.assertEqual(['a', 'c', 'b', 'd'], [s.name for s in specs]) + self.assertEqual([0, 1, 2, 3], [s.level for s in specs]) + self.assertEqual(0, b.level, 'input specs should not be modified') + # Already sorted specs are unchanged + self.assertEqual(specs, sort_specs(specs)) + with self.assertRaises(ValueError, msg='duplicate names'): + sort_specs([a, a]) + with self.assertRaises(ValueError, msg='missing parent'): + sort_specs([b]) + a.parents = ['d'] + with self.assertRaises(ValueError, msg='circular'): + sort_specs([a, b, c, d]) + + def test_instantiate(self): + task = self.t1.to_spec().instantiate(self.session_path, location='remote') + self.assertIsInstance(task, Task01) + self.assertEqual(('remote', self.session_path), (task.location, task.session_path)) + + def test_pipeline_task_specs(self): + """Test Pipeline methods with a mix of Task instances and specs.""" + spec = TaskSpec('Task02', 'mpci.foo.Task', parents=['Task01'], env='mpci') + pipe = Pipeline(session_path=self.session_path, tasks={'Task02': spec, 'Task00_foo': self.t0, 'Task01': self.t1}) + with self.assertLogs('ibllib.pipes.tasks', 'WARNING') as log: + specs = pipe.task_specs() + # The Task01 env doesn't match its route + self.assertEqual(1, len(log.records)) + self.assertIn('Task01', log.records[0].getMessage()) + self.assertEqual(['Task00_foo', 'Task01', 'Task02'], [s.name for s in specs]) + self.assertEqual([0, 1, 2], [s.level for s in specs]) + with self.assertLogs('ibllib.pipes.tasks', 'WARNING'): + task_list = pipe.create_tasks_list_from_pipeline() + self.assertEqual(['Task00_foo', 'Task01', 'Task02'], [t['name'] for t in task_list]) + self.assertEqual(['Task01'], task_list[-1]['parents']) + self.assertEqual('Pipeline', task_list[-1]['graph']) + # NB: env and job_size are not stored in Alyx task dicts + with self.assertLogs('ibllib.pipes.tasks', 'WARNING'): + from_dicts = pipe.task_specs(task_list) + stored = lambda s: {k: v for k, v in s.to_dict().items() if k not in ('env', 'job_size')} # noqa: E731 + self.assertEqual(list(map(stored, specs)), list(map(stored, from_dicts))) + with self.assertLogs('ibllib.pipes.tasks', 'WARNING'): + graph = pipe.make_graph(out_dir=tempfile.gettempdir(), show=False) + for edge in ('root -> Task00_foo', 'Task00_foo -> Task01', 'Task01 -> Task02'): + self.assertIn(edge, graph.source) + + def test_create_alyx_tasks(self): + """Test Pipeline.create_alyx_tasks creates parents first and patches existing tasks.""" + spec = TaskSpec('Task02', 'mpci.foo.Task', parents=['Task00_foo'], env='mpci') + one = mock.MagicMock() + one.alyx.cache_mode = None + pipe = Pipeline(session_path=self.session_path, one=one, eid='eid', tasks=OrderedDict(Task02=spec, Task00_foo=self.t0)) + existing = {'id': 'id_Task00_foo', 'name': 'Task00_foo', 'status': 'Waiting'} + one.alyx.rest.side_effect = lambda *args, **kwargs: { + 'list': [existing], + 'partial_update': existing, + 'create': {'id': 'id_' + kwargs.get('data', {}).get('name', ''), **kwargs.get('data', {})}, + }[args[1]] + tasks = pipe.create_alyx_tasks(rerun__status__in=['Waiting']) + self.assertEqual(['Task00_foo', 'Task02'], [t['name'] for t in tasks]) + self.assertEqual(['id_Task00_foo'], tasks[1]['parents']) + self.assertEqual(1, tasks[1]['level']) + calls = [c.args[1] for c in one.alyx.rest.call_args_list] + self.assertEqual(['list', 'partial_update', 'create'], calls) + self.assertEqual(Task00.time_out_secs, one.alyx.rest.call_args.kwargs['data']['time_out_secs']) + # Time outs greater than the Alyx field maximum raise before any tasks are created + one.alyx.rest.reset_mock() + pipe.tasks['Task02'] = TaskSpec('Task02', 'mpci.foo.Task', parents=['Task00_foo'], env='mpci', time_out_secs=32768) + with self.assertRaises(AssertionError): + pipe.create_alyx_tasks() + self.assertEqual(['list'], [c.args[1] for c in one.alyx.rest.call_args_list]) + + +class TestPlan(unittest.TestCase): + """Tests for the ibllib.pipes.plan module.""" + + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.session_path = Path(tmp.name, 'subject', '2020-01-01', '001') + self.target = f'{__name__}:planner' + # An env that points to the current Python environment, for testing subprocess calls + self.env_paths = {'test': Path(sys.prefix), 'missing': Path(tmp.name, 'missing')} + self.description = {'devices': {'foo': {}}} + + def test_to_specs(self): + pipe = planner(self.session_path) + specs = plan.to_specs(pipe) + self.assertEqual(['PlannedTask', 'PlannedSpec'], [s.name for s in specs]) + self.assertEqual(specs, plan.to_specs(pipe.tasks)) + self.assertEqual(specs, plan.to_specs(list(pipe.tasks.values()))) + self.assertEqual(specs, plan.to_specs([s.to_dict() for s in specs])) + + def test_plan(self): + context = {'tasks': [{'name': 'CoreTask'}]} + specs = plan.plan(self.target, self.session_path, context=context) + self.assertEqual(['PlannedTask', 'CoreTask'], specs[1].parents) + self.assertEqual({'foo': 'bar'}, specs[0].arguments) + + @unittest.skipIf(routing.env_python('test', {'test': Path(sys.prefix)}) is None, 'Python executable not found in sys.prefix') + def test_plan_in_env(self): + context = {'tasks': [{'name': 'CoreTask'}]} + specs = plan.plan_in_env(self.target, 'test', self.session_path, context=context, env_paths=self.env_paths) + self.assertEqual(plan.plan(self.target, self.session_path, context=context), specs) + with self.assertRaises(plan.PlannerError) as ex: + plan.plan_in_env(self.target, 'test', self.session_path, context={'raise': True}, env_paths=self.env_paths) + self.assertIn('planner failed', str(ex.exception)) + with self.assertRaises(plan.PlannerError): + plan.plan_in_env(self.target, 'missing', self.session_path, env_paths=self.env_paths) + + def test_get_external_tasks(self): + planners = {'foo': ('missing', self.target), 'bar': ('missing', 'notapkg.foo:plan')} + kwargs = {'planners': planners, 'env_paths': self.env_paths} + # Env not installed so the planner is called in the current env (the planner is importable) + with mock.patch('ibllib.pipes.plan.plan_in_env') as plan_in_env: + specs, errors = plan.get_external_tasks(self.description, self.session_path, **kwargs) + plan_in_env.assert_not_called() + self.assertEqual({}, errors) + self.assertEqual(['PlannedTask', 'PlannedSpec'], list(specs)) + # Env installed so the planner is called in a subprocess + planners['foo'] = ('test', self.target) + with mock.patch('ibllib.pipes.plan.plan_in_env', return_value=[TaskSpec('Foo', 'foo.Task')]) as plan_in_env: + specs, errors = plan.get_external_tasks(self.description, self.session_path, **kwargs) + plan_in_env.assert_called_once() + self.assertEqual(['Foo'], list(specs)) + # Neither the env nor the planner are installed + self.description['devices']['bar'] = {} + with self.assertLogs('ibllib.pipes.plan', 'ERROR'), mock.patch('ibllib.pipes.plan.plan_in_env', return_value=[]): + specs, errors = plan.get_external_tasks(self.description, self.session_path, **kwargs) + self.assertEqual({'bar': 'Environment "missing" not installed'}, errors) + # Planner errors are caught + with self.assertLogs('ibllib.pipes.plan', 'ERROR'): + specs, errors = plan.get_external_tasks( + self.description, self.session_path, context={'raise': True}, planners={'foo': ('missing', self.target)} + ) + self.assertEqual({'foo': 'planner failed'}, errors) + # Devices not in the description are ignored + specs, errors = plan.get_external_tasks({}, self.session_path, planners=planners) + self.assertEqual(({}, {}), (dict(specs), errors)) + + def test_main(self): + with tempfile.TemporaryDirectory() as tmp: + output, context = Path(tmp, 'specs.json'), Path(tmp, 'context.json') + context.write_text(json.dumps({'tasks': [{'name': 'CoreTask'}]})) + plan.main([self.target, str(self.session_path), '--context', str(context), '--output', str(output)]) + specs = [TaskSpec.from_dict(d) for d in json.loads(output.read_text())] + self.assertEqual(['PlannedTask', 'CoreTask'], specs[1].parents) + + +if __name__ == '__main__': + unittest.main()