Source code for ibllib.pipes.local_server

"""Lab server pipeline construction and task runner.

This is the module called by the job services on the lab servers.  See
iblscripts/deploy/serverpc/crons for the service scripts that employ this module.
"""
import logging
import time
from datetime import datetime
from pathlib import Path
import re
import subprocess
import sys
import traceback
import importlib
import importlib.metadata

from one.api import ONE
from one.webclient import AlyxClient
from one.remote.globus import get_lab_from_endpoint_id, get_local_endpoint_id
from one.alf.spec import is_session_path
from one.alf.path import session_path_parts

from ibllib import __version__ as ibllib_version
from ibllib.pipes import tasks
from ibllib.time import date2isostr
from ibllib.oneibl.registration import IBLRegistrationClient
from ibllib.oneibl.data_handlers import get_local_data_repository
from ibllib.io.session_params import read_params
from ibllib.pipes.dynamic_pipeline import make_pipeline, acquisition_description_legacy_session

_logger = logging.getLogger(__name__)
LARGE_TASKS = [
    'EphysVideoCompress', 'TrainingVideoCompress', 'SpikeSorting', 'EphysDLC', 'MesoscopePreprocess'
]


def _run_command(cmd):
    process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
    info, error = process.communicate()
    if process.returncode != 0:
        return None
    else:
        return info.decode('utf-8').strip()


def _get_volume_usage(vol, label=''):
    cmd = f'df {vol}'
    res = _run_command(cmd)
    # size_list = ['/dev/sdc1', '1921802500', '1427128132', '494657984', '75%', '/datadisk']
    size_list = re.split(' +', res.split('\n')[-1])
    fac = 1024 ** 2
    d = {'total': int(size_list[1]) / fac,
         'used': int(size_list[2]) / fac,
         'available': int(size_list[3]) / fac,
         'volume': size_list[5]}
    return {f"{label}_{k}": d[k] for k in d}




def report_health(alyx):
    """
    Get a few indicators and label the json field of the corresponding lab with them.
    """
    status = {'python_version': sys.version,
              'ibllib_version': ibllib_version,
              'phylib_version': importlib.metadata.version('phylib'),
              'local_time': date2isostr(datetime.now())}
    status.update(_get_volume_usage('/mnt/s0/Data', 'raid'))
    status.update(_get_volume_usage('/', 'system'))

    data_repos = alyx.rest('data-repository', 'list', globus_endpoint_id=get_local_endpoint_id())

    for dr in data_repos:
        alyx.json_field_update(endpoint='data-repository', uuid=dr['name'], field_name='json', data=status)





def job_creator(root_path, one=None, dry=False, rerun=False):
    """
    Create new sessions and pipelines.

    Server function that will look for 'raw_session.flag' files and for each:
     1) create the session on Alyx
     2) create the tasks to be run on Alyx

    For legacy sessions the raw data are registered separately, instead of within a pipeline task.

    Parameters
    ----------
    root_path : str, pathlib.Path
        Main path containing sessions or a session path.
    one : one.api.OneAlyx
        An ONE instance for registering the session(s).
    dry : bool
        If true, simply log the session_path(s) found, without registering anything.
    rerun : bool
        If true and session pipeline tasks already exist, set them all to waiting.

    Returns
    -------
    list of ibllib.pipes.tasks.Pipeline
        The pipelines created.
    list of dicts
        A list of any datasets registered (only for legacy sessions)
    """
    _logger.info('Start looking for new sessions...')
    if not one:
        one = ONE(cache_rest=None)
    rc = IBLRegistrationClient(one=one)
    flag_files = Path(root_path).glob('*/????-??-??/*/raw_session.flag')
    flag_files = filter(lambda x: is_session_path(x.parent), flag_files)
    pipes = []
    all_datasets = []
    for flag_file in flag_files:
        session_path = flag_file.parent
        if session_path_parts(session_path)[1] in ('test', 'test_subject'):
            _logger.debug('skipping test session %s', session_path)
            continue
        _logger.info(f'creating session for {session_path}')
        if dry:
            continue
        try:
            # if the subject doesn't exist in the database, skip
            rc.register_session(session_path, file_list=False)

            # NB: all sessions now extracted using dynamic pipeline
            if read_params(session_path) is None:
                # Create legacy experiment description file
                acquisition_description_legacy_session(session_path, save=True)
            pipe = make_pipeline(session_path, one=one)
            if rerun:
                rerun__status__in = '__all__'
            else:
                rerun__status__in = ['Waiting']
            pipe.create_alyx_tasks(rerun__status__in=rerun__status__in)
            flag_file.unlink()
            if pipe is not None:
                pipes.append(pipe)
        except Exception:
            _logger.error('Failed to register session %s:\n%s', session_path.relative_to(root_path), traceback.format_exc())
            continue

    return pipes, all_datasets





def task_queue(mode='all', lab=None, alyx=None, env=(None,)):
    """
    Query waiting jobs from the specified Lab

    Parameters
    ----------
    mode : {'all', 'small', 'large'}
        Whether to return all waiting tasks, or only small or large (specified in LARGE_TASKS) jobs.
    lab : str
        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`.

    Returns
    -------
    list of dict
        A list of Alyx tasks associated with `lab` that have a 'Waiting' status.
    """
    def predicate(task):
        classe = tasks.str2class(task['executable'])
        return (mode == 'all' or classe.job_size == mode) and classe.env in 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)
    # Order tasks by priority
    sorted_tasks = sorted(filtered_tasks, key=lambda d: d['priority'], reverse=True)

    return sorted_tasks





def tasks_runner(subjects_path, tasks_dict, one=None, dry=False, count=5, time_out=None, **kwargs):
    """
    Function to run a list of tasks (task dictionary from Alyx query) on a local server

    Parameters
    ----------
    subjects_path : str, pathlib.Path
        The location of the subject session folders, e.g. '/mnt/s0/Data/Subjects'.
    tasks_dict : list of dict
        A list of tasks to run. Typically the output of `task_queue`.
    one : one.api.OneAlyx
        An instance of ONE.
    dry : bool, default=False
        If true, simply prints the full session paths and task names without running the tasks.
    count : int, default=5
        The maximum number of tasks to run from the tasks_dict list.
    time_out : float, optional
        The time in seconds to run tasks before exiting. If set this will run tasks until the
        timeout has elapsed.  NB: Only checks between tasks and will not interrupt a running task.
    kwargs
        See ibllib.pipes.tasks.run_alyx_task.

    Returns
    -------
    list of pathlib.Path
        A list of datasets registered to Alyx.
    """
    if one is None:
        one = ONE(cache_rest=None)
    tstart = time.time()
    c = 0
    last_session = None
    all_datasets = []
    for tdict in tasks_dict:
        # if the count is reached or if the time_out has been elapsed, break the loop and return
        if c >= count or (time_out and time.time() - tstart > time_out):
            break
        # reconstruct the session local path. As many jobs belong to the same session
        # cache the result
        if last_session != tdict['session']:
            ses = one.alyx.rest('sessions', 'list', django=f"pk,{tdict['session']}")[0]
            session_path = Path(subjects_path).joinpath(
                Path(ses['subject'], ses['start_time'][:10], str(ses['number']).zfill(3)))
            last_session = tdict['session']
        if dry:
            print(session_path, tdict['name'])
        else:
            task, dsets = tasks.run_alyx_task(tdict=tdict, session_path=session_path, one=one, **kwargs)
            if dsets:
                all_datasets.extend(dsets)
                c += 1
    return all_datasets