labdata 0.1.22__tar.gz → 0.1.24__tar.gz

{labdata-0.1.22 → labdata-0.1.24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: labdata
-Version: 0.1.22
+Version: 0.1.24
 Summary: Package to manage data in experimental neuroscience labs
 Author-email: Joao Couto <jpcouto@gmail.com>
 Project-URL: Homepage, https://github.com/jcouto/labdata

{labdata-0.1.22 → labdata-0.1.24}/labdata/__init__.py

@@ -1,4 +1,5 @@
-VERSION = "0.1.22"
+VERSION = "0.1.24"
+__version__ = VERSION
 
 from .utils import *
 from .copy import copy_to_upload_server

{labdata-0.1.22 → labdata-0.1.24}/labdata/accounts.py

@@ -78,13 +78,22 @@ def database_grants(name, projects, database_name='lab_data', sandbox=False):
     return [g.replace('%', '%%') for g in grants]
 
 
+def split_host(host, port=3306):
+    '''Return ``(address, port)`` of a database host written ``address`` or ``address:port``;
+    ``port`` is used when the host has none.'''
+    address, _, number = str(host).strip().rpartition(':')
+    if address and number.isdigit():
+        return address, int(number)
+    return str(host).strip(), int(port)
+
+
 def ask_database_admin(host, default_user=None):
     '''Ask for the database admin login and return a pymysql connection.
 
     Parameters
     ----------
     host : str
-        Database address.
+        Database address, ``address`` or ``address:port``.
     default_user : str, optional
         User offered as default.
 
@@ -96,7 +105,8 @@ def ask_database_admin(host, default_user=None):
     print(f'Database admin login for {host} (used for this command only, not saved).')
     user = input(f'  Admin user [{default_user or ""}]: ').strip() or default_user
     password = getpass('  Admin password: ')
-    return pymysql.connect(host=host, user=user, password=password, autocommit=True)
+    address, port = split_host(host)
+    return pymysql.connect(host=address, port=port, user=user, password=password, autocommit=True)
 
 
 def database_user_exists(conn, name):
@@ -147,6 +157,49 @@ def database_user_grants(conn, name):
     return [r[0] for r in _rows(conn, f"SHOW GRANTS FOR '{name}'@'%%'", ())]
 
 
+def _grant_pattern(schema):
+    '''Return the regular expression of a schema name of a GRANT, where ``_`` and ``%`` are
+    wildcards unless escaped with ``\\``.'''
+    out, i = '', 0
+    while i < len(schema):
+        c = schema[i]
+        if c == '\\' and i + 1 < len(schema):
+            out += re.escape(schema[i + 1])
+            i += 2
+            continue
+        out += '.*' if c == '%' else '.' if c == '_' else re.escape(c)
+        i += 1
+    return re.compile(out)
+
+
+def database_access(conn, databases):
+    '''Return ``{user: [databases]}`` for the database users with privileges other than
+    REFERENCES on ``databases``; ``'*'`` stands for a grant on every database. Users whose
+    grants can not be read are left out.'''
+    access = dict()
+    for user in list_database_users(conn):
+        try:
+            grants = database_user_grants(conn, user)
+        except Exception:
+            continue
+        for grant in grants:
+            m = re.match(r'GRANT (.+?) ON (\S+?)\.\S+ TO ', grant)
+            if not m:
+                continue
+            privileges = set(p.strip().upper() for p in m.group(1).split(','))
+            if privileges <= {'USAGE', 'REFERENCES'}:
+                continue
+            schema = m.group(2).strip('`')
+            if schema == '*':
+                access.setdefault(user, set()).add('*')
+                continue
+            pattern = _grant_pattern(schema)
+            for database in databases:
+                if pattern.fullmatch(database):
+                    access.setdefault(user, set()).add(database)
+    return {user: sorted(names) for user, names in sorted(access.items())}
+
+
 def create_s3_user(session, name, groups=(), policies=(), like=None):
     '''Create the IAM user ``name`` with one access key, and the permissions of ``groups``, ``policies`` or user ``like``.
 
@@ -836,7 +889,8 @@ def grant_database_project(conn, name, project, database_name='lab_data'):
 def connect_database_admin(host, user, password):
     '''Return a pymysql connection for ``user``, after checking it may create users and grant permissions.'''
     import pymysql
-    conn = pymysql.connect(host=host, user=user, password=password, autocommit=True)
+    address, port = split_host(host)
+    conn = pymysql.connect(host=address, port=port, user=user, password=password, autocommit=True)
     grants = ' '.join(r[0] for r in _rows(conn, 'SHOW GRANTS', ())).upper()
     if 'CREATE USER' not in grants and 'ALL PRIVILEGES ON *.*' not in grants:
         conn.close()

{labdata-0.1.22 → labdata-0.1.24}/labdata/cli.py

@@ -41,6 +41,10 @@ class CLI_parser():
         Create and manage the database and S3 accounts of lab members (admin use).
     access
         Manage what the members of a project can read in S3 (admin use).
+    share
+        Copy tables, with their parents and part tables, from one project to another (admin use).
+    export
+        Copy a selection to its own database, a dump folder or a Docker image (admin use).
     '''
     def __init__(self):
         parser = argparse.ArgumentParser(
@@ -74,6 +78,8 @@ Administration commands (don't run on experimental computers):
             {tcolor["r"]("cluster")} create|status <name>                  Creates or checks the AWS ParallelCluster
             {tcolor["r"]("account")} create|show|list|grant|revoke <name>  Creates and manages database and S3 accounts
             {tcolor["r"]("access")} show|sync|share <project>             Sets what a project's members can read in S3
+            {tcolor["r"]("share")} <table> --from <project> --to <project>   Shares tables with another project of the lab
+            {tcolor["r"]("export")} <table> --docker <folder>              Exports data to a database, dump or Docker image
 
             ''')
         
@@ -615,7 +621,9 @@ Administration commands (don't run on experimental computers):
             whose output goes to ``task_log`` and ``~/.labdata/logs``,
             instead of in the terminal.
         --queue
-            Print the current job queue for all targets and exit.
+            Print the scheduler queue of each target, then the ``ComputeTask``
+            table by status with the last line of each task log, and exit.
+            With an analysis name, only the tasks of that analysis.
         --delete-computes
             Delete all ``ComputeTask`` entries for this analysis (careful).
         --
@@ -666,6 +674,8 @@ Available analysis are
             targets = np.unique([prefs['compute']['default_target']]+
                                 list(prefs['compute']['remotes'].keys()))
             for q in targets:
+                if q == 'local':  # no scheduler queue, the tasks are in the table below
+                    continue
                 queue = check_queue(target = q)
                 if queue is None:
                     continue
@@ -678,7 +688,14 @@ Available analysis are
                         cc = tcolor['g']
                     else:
                         cc = tcolor['r']
-                    print(f"\t {j['jobid']} {j['name']} {cc('running')} time: {tcolor['y'](j['time'])}\n\t\tpartition: {j['partition']} node: {j['nodelist(reason)']} ")
+                    print(f"\t {j['jobid']} {j['name']} {cc(j['state'].lower())} time: {tcolor['y'](j['time'])}\n\t\tpartition: {j['partition']} node: {j['nodelist(reason)']} ")
+            # the ComputeTask table, optionally for one analysis: labdata run <analysis> --queue
+            qparser = argparse.ArgumentParser(add_help = False)
+            qparser.add_argument('analysis', nargs = '?', default = None)
+            qparser.add_argument('-p','--project', default = _get_project())
+            qargs, _ = qparser.parse_known_args(argum)
+            from .utils import load_project_schema
+            self._print_compute_tasks(load_project_schema(qargs.project), analysis = qargs.analysis)
             return
         args = parser.parse_args(argum)
         if '--delete-computes' in argum:
@@ -1220,6 +1237,9 @@ examples:
         --s3-like, --s3-group, --s3-policy
             Extra IAM permissions: copy the groups and policies of a user, or
             add groups and managed policies (names or ARNs).
+        --database-name
+            Lab database (``database.name``) of the projects, for a copy made with
+            ``labdata share --database-name``. Default: the one of the preference file.
         --profile, --region
             AWS admin profile and region (asked when omitted).
         --dry-run
@@ -1246,6 +1266,7 @@ examples:
   labdata account create
   labdata account create <name> -p <project> --restore --no-sandbox
   labdata account grant <name> -p <other project>
+  labdata account grant <name> -p other_project --database-name newlab    # a copy in another lab database
   labdata account show <name>
 ''')
         parser.add_argument('action', action = 'store',
@@ -1276,13 +1297,17 @@ examples:
                             help = 'AWS profile with rights to manage IAM (asked when omitted).')
         parser.add_argument('--region', action = 'store', default = None, type = str,
                             help = 'AWS region (asked when omitted).')
+        parser.add_argument('--database-name', action = 'store', default = None, type = str,
+                            help = 'Lab database (database.name) of the projects, e.g. a copy made with '
+                                   '"labdata share --database-name" (default: the one of the preference file).')
         parser.add_argument('--dry-run', action = 'store_true', default = False,
                             help = 'Print the plan and change nothing.')
         args = parser.parse_args(sys.argv[2:])
 
         from . import accounts
         host = prefs['database'].get('database.host')
-        database_name = prefs['database'].get('database.name', 'lab_data')
+        database_name = args.database_name or prefs['database'].get('database.name', 'lab_data')
+        self._account_database_name = database_name
         admin_user = prefs['database'].get('database.user')
 
         if args.action == 'list':
@@ -1756,9 +1781,13 @@ examples:
     def _write_account_file(self, name, projects, password, keys):
         '''Write the preference file of a user with their new credentials.'''
         from . import accounts
+        base = dict(prefs)
+        database_name = getattr(self, '_account_database_name', None)
+        if database_name:
+            base['database'] = dict(prefs['database'], **{'database.name': database_name})
         path = accounts.write_private_json(
             self._account_file(name),
-            accounts.new_user_preferences(prefs, name, projects, database_password = password, s3_keys = keys))
+            accounts.new_user_preferences(base, name, projects, database_password = password, s3_keys = keys))
         print(f'  ✓ wrote {path} (mode 600)')
         print(f'    Give it to {name} as <their labdata folder>/user_preferences.json, then delete your copy.')
         if password is None:
@@ -1920,6 +1949,361 @@ examples:
             write_log(final = True)   # the last lines, often the summary
         print('Pipe closed.')
 
+    def _export_restriction(self, args):
+        '''Return the restriction of -a/-s/-d as a list of dicts, or None.'''
+        fields = [('subject_name', args.subject), ('session_name', args.session), ('dataset_name', args.datatype)]
+        restriction = [dict()]
+        for field, values in fields:
+            if values:
+                restriction = [dict(r, **{field: v}) for r in restriction for v in values]
+        return None if restriction == [dict()] else restriction
+
+    def _export_common_arguments(self, parser):
+        self._add_default_arguments(parser, level = 3)
+        parser.add_argument('--plugin', nargs = '+', default = [],
+                            help = 'Plugins whose tables are seeds (restricted by -a/-s/-d where they have those fields).')
+        parser.add_argument('--exclude', nargs = '+', default = None,
+                            help = 'Tables left out with their part tables (class names, e.g. ComputeTask Dataset.DataFiles).')
+        parser.add_argument('--schema-map', nargs = '+', default = [],
+                            help = 'source=target names for databases that do not start with the lab database name.')
+        parser.add_argument('-n', '--n-jobs', type = int, default = DEFAULT_N_JOBS,
+                            help = f'Copy threads (default {DEFAULT_N_JOBS}).')
+        parser.add_argument('--dry-run', action = 'store_true', default = False,
+                            help = 'Print the tables, rows and estimated size; copy nothing.')
+        parser.add_argument('--existing', default = 'ask', choices = ['ask', 'keep', 'update', 'replace'],
+                            help = 'Rows of the selection already in the target: ask (default), keep them, '
+                                   'update their values, or replace them (delete the exported rows and copy them again).')
+        parser.add_argument('--host', default = None, help = 'Another MySQL server to copy to: address or address:port.')
+        parser.add_argument('--port', default = 3306, type = int, help = 'Port of --host when it has none (default 3306).')
+        parser.add_argument('--db-user', default = None,
+                            help = 'Database user: the login for --host, or the account created in the --docker image (asked when not given).')
+        parser.add_argument('--db-password', default = None,
+                            help = 'Its password (asked when not given; better asked than typed in the shell history).')
+
+    def _target_login(self, args):
+        '''Return the --host target with its login from --db-user/--db-password, asked when not given.'''
+        from getpass import getpass
+        from .accounts import split_host
+        print(f'Database login for {args.host} (used for this command only, not saved).')
+        user = args.db_user or input('  User: ').strip()
+        password = args.db_password if args.db_password is not None else getpass(f'  Password for {user}: ')
+        host, port = split_host(args.host, args.port)
+        return dict(host = host, port = port, user = user, password = password)
+
+    def _lab_server_login(self):
+        '''Ask for the login that writes the copy on the lab server, unless the environment has one.
+
+        The login is checked, then handed to this process in the environment (password encrypted)
+        before the schema is loaded; it is not saved and stays out of the shell history.
+        '''
+        if os.environ.get('LABDATA_DATABASE_USER') and os.environ.get('LABDATA_DATABASE_PASS'):
+            return
+        import pymysql
+        from getpass import getpass
+        host = os.environ.get('LABDATA_DATABASE_HOST') or prefs['database'].get('database.host')
+        default = prefs['database'].get('database.user') or ''
+        print(f'Database login for {host} that writes the copy (used for this command only, not saved).')
+        user = input(f'  User [{default}]: ').strip() or default
+        password = getpass(f'  Password for {user}: ')
+        from .accounts import split_host
+        address, port = split_host(host)
+        try:
+            pymysql.connect(host = address, port = port, user = user, password = password).close()
+        except pymysql.err.OperationalError as err:
+            print(tcolor['r'](f'Could not log in as {user}: {err.args[-1]}'))
+            exit(1)
+        os.environ['LABDATA_DATABASE_USER'] = user
+        os.environ['LABDATA_DATABASE_PASS'] = 'encrypted:' + encrypt_string(password)
+
+    def _image_login(self, args):
+        '''Return the database account of a --docker image from --db-user/--db-password, asked when not given.'''
+        from getpass import getpass
+        user = args.db_user or input('Database account of the image [labdata]: ').strip() or 'labdata'
+        password = args.db_password
+        if password is None:
+            password = getpass(f'  Password for {user} [leave empty to generate one]: ')
+        return user, password or None
+
+    def _print_export_access(self, report, target = None):
+        '''Print who can read the databases of an export and how to give an account access.'''
+        from . import accounts
+        databases, project = report.attrs.get('databases', []), report.attrs.get('project')
+        database_name = report.attrs.get('database_name')
+        server = None
+        if target:
+            server = target['host'] if int(target.get('port', 3306)) == 3306 else f"{target['host']}:{target['port']}"
+        where = f'on {server}' if target else 'on this server'
+        try:
+            if target:
+                from .export.selection import _open_connection
+                conn = _open_connection(target)
+            else:
+                import datajoint as dj
+                conn = dj.conn()
+            access = accounts.database_access(conn, databases)
+        except Exception as err:
+            print(f'Could not read the database grants {where}: {err}')
+            access = None
+        if access is not None:
+            everything = [u for u, names in access.items() if '*' in names]
+            some = {u: names for u, names in access.items() if '*' not in names}
+            print(f'\nAccess to {", ".join(databases)} {where} (database grants):')
+            for user, names in some.items():
+                print(f'    {user}: {", ".join(names)}')
+            if not some:
+                print('    no account has grants on these databases.')
+            if everything:
+                print(f'    and the accounts with grants on every database: {", ".join(everything)}')
+        option = '' if database_name == prefs['database'].get('database.name') else f' --database-name {database_name}'
+        print('\nTo give an account access to the project (database grants, Project.User and, '
+              'with an IAM user, the S3 group of the project):')
+        if target:
+            print(f'    LABDATA_PREFERENCES=<preference file for {server}> '
+                  f'labdata account grant <name> -p {project}{option}')
+            print(f'    (labdata account uses the server of the preference file; '
+                  f'"labdata account create <name> -p {project}{option}" for a new account)')
+        else:
+            print(f'    labdata account grant <name> -p {project}{option}')
+            print(f'    labdata account create <name> -p {project}{option}    # a new account')
+
+    def _print_export_report(self, report):
+        copied = report[report['rows'].fillna(0) > 0]
+        mode = report.attrs.get('existing', 'keep')
+        for row in copied.itertuples():
+            existing = int(row.existing or 0) if 'existing' in report else 0
+            note = f', {existing} were there ({mode})' if existing else ''
+            print(f'\t{row.table}  {int(row.rows)} rows{note}')
+        print(tcolor['g'](f"Done: {int(copied['rows'].sum())} rows in {len(copied)} tables."))
+
+    def share(self):
+        '''Copy tables from one project to another, with their parents and part tables.
+
+        The lab-wide tables (File, LabMember, Setup, ...) are shared by the projects and are not
+        copied. The copy runs on the database server. Needs read access to the source project and
+        write access to the target databases.
+
+        Usage
+        -----
+        ::
+
+            labdata share EphysRecording --from couto_arc --to other_project -a JC1013
+            labdata share EphysRecording --from couto_arc --to other_project --descendants --dry-run
+            labdata share --plugin ratetask --from couto_arc --to other_project
+            labdata share EphysRecording --from couto_arc --host 10.0.0.5 --db-user admin -a JC1013
+            labdata share Dataset.DataFiles --from couto_arc --to other_project --database-name newlab -a JC1013
+
+        With --host the project goes to another MySQL server (same project name unless --to), with
+        the lab-wide rows it uses (File, LabMember, ...); the password is asked. With a new
+        --database-name on this server it goes to the databases ``newlab`` and
+        ``newlab_<project>...``, with the lab-wide rows it uses. The preference entries that open
+        the copy are written to a JSON file (--access-file).
+        '''
+        parser = argparse.ArgumentParser(
+            description = 'Share tables, with their parents and part tables, with another project.',
+            usage = 'labdata share <table ...> --from <project> --to <project> [-a subject] [-s session] [options]',
+            formatter_class = argparse.RawDescriptionHelpFormatter,
+            epilog = '''examples:
+  labdata share EphysRecording --from couto_arc --to other_project -a JC1013 --dry-run
+  labdata share EphysRecording --from couto_arc --to other_project --descendants      # with spike sorting
+  labdata share --plugin ratetask --from couto_arc --to other_project -a JC1013
+  labdata share EphysRecording --from couto_arc --host 10.0.0.5 --db-user admin -a JC1013
+
+  # the File rows of the datasets of JC1013 (with their Session and Dataset rows), on this
+  # server, to the lab-wide database newlab and the project databases newlab_other_project...
+  labdata share Dataset.DataFiles --from couto_arc --to other_project --database-name newlab -a JC1013
+
+The lab-wide database (File, LabMember, Setup, ...) is shared by the projects of a server and is
+not copied, unless --host or a new --database-name is given; then only the rows the selection
+uses are copied. Use a table with the -a/-s/-d fields as seed (File has none, so -a/-s/-d would
+not restrict it).
+
+After the copy, the entries that open the target are written to a JSON file (--access-file,
+default labdata_access_<database name>_<project>.json): copy them into the "database" section of
+the preference file, or set them with LABDATA_DATABASE_PROJECT for the project.
+''')
+        parser.add_argument('table', nargs = '*', default = [],
+                            help = 'Seed tables (class names, e.g. EphysRecording, SpikeSorting).')
+        parser.add_argument('--from', dest = 'source', required = True, help = 'Source project.')
+        parser.add_argument('--to', dest = 'target', default = None,
+                            help = 'Target project (default with --host: the source project).')
+        parser.add_argument('--database-name', default = None,
+                            help = 'database.name of the lab in the target (default: the same as here). On this server, '
+                                   'a new name copies the lab-wide rows the selection uses to that database.')
+        parser.add_argument('--access-file', default = None,
+                            help = 'JSON file with the preference entries that open the target '
+                                   '(default: labdata_access_<database name>_<project>.json).')
+        parser.add_argument('--descendants', action = 'store_true', default = False,
+                            help = 'Also copy the tables derived from the seeds (spike sorting, ...).')
+        self._export_common_arguments(parser)
+        args = parser.parse_args(sys.argv[2:])
+        if not args.table and not args.plugin:
+            parser.error('give a table or --plugin.')
+        if args.target is None and args.host is None and args.database_name is None:
+            parser.error('give --to, --host or --database-name.')
+        target = None
+        if args.host is not None and not args.dry_run:
+            target = self._target_login(args)
+        elif args.host is not None:
+            from .accounts import split_host
+            address, port = split_host(args.host, args.port)
+            target = dict(host = address, port = port)
+        elif not args.dry_run:
+            self._lab_server_login()
+        from .export import share_project, DEFAULT_EXCLUDE, ExportCancelled
+        try:
+            report = share_project(args.table or None, args.source, args.target,
+                                   target = target, database_name = args.database_name, existing = args.existing,
+                                   restriction = self._export_restriction(args), plugins = args.plugin,
+                                   descendants = args.descendants,
+                                   exclude = DEFAULT_EXCLUDE if args.exclude is None else args.exclude,
+                                   schema_map = dict(m.split('=', 1) for m in args.schema_map),
+                                   n_jobs = args.n_jobs, dry_run = args.dry_run)
+        except ExportCancelled as err:
+            print(err)
+            return
+        except ValueError as err:
+            print(tcolor['r'](str(err)))
+            exit(1)
+        if not args.dry_run:
+            self._print_export_report(report)
+            from .export import project_access
+            host = None
+            if target is not None:
+                host = target['host'] if target['port'] == 3306 else f"{target['host']}:{target['port']}"
+            access = project_access(args.target or args.source, args.database_name, host)
+            name = access['database']['database.name']
+            path = Path(args.access_file or f'labdata_access_{name}_{access["database"]["database.project"]}.json')
+            with open(path, 'w') as fd:
+                json.dump(access, fd, indent = 4)
+            print(f'Preference entries to open {name} / {access["database"]["database.project"]} in {path}; '
+                  'copy them into the "database" section of the preference file.')
+            self._print_export_access(report, target)
+
+    def export(self):
+        '''Copy a selection, with the lab-wide rows it uses, to its own database, a dump or a Docker image.
+
+        Derived tables are included unless --no-descendants. Raw files stay where they are unless
+        --files bucket (copied to another storage) or --files download (to a folder).
+
+        Usage
+        -----
+        ::
+
+            labdata export EphysRecording -p couto_arc -a JC1013 --dry-run
+            labdata export EphysRecording -p couto_arc -a JC1013 --database-name release_2026
+            labdata export EphysRecording -p couto_arc --host 10.0.0.5 --db-user admin
+            labdata export EphysRecording -p couto_arc --dump release_dump
+            labdata export EphysRecording -p couto_arc --docker release --tag lab/release:1 \\
+                    --pip git+https://github.com/org/package --blank-personal --build
+            labdata export EphysRecording -p couto_arc --docker release --link-data --dump release_data \\
+                    --tag lab/release:1 --build
+
+        Run the image with ``docker run -p 8888:8888 <tag>`` (Jupyter) or
+        ``docker run -p 8501:8501 <tag> dashboard`` (dashboard). With --link-data the rows are not
+        in the image: mount the data folder with ``-v <data folder>:/dump:ro`` and keep the loaded
+        database in a volume with ``-v <volume>:/var/lib/mysql``.
+        '''
+        parser = argparse.ArgumentParser(
+            description = 'Export a selection to a database, a dump folder or a Docker image.',
+            usage = 'labdata export <table ...> [-p project] [-a subject] [--database-name name] '
+                    '[--host h --user u | --dump folder | --docker folder --tag tag [--link-data]] [options]')
+        parser.add_argument('table', nargs = '*', default = [],
+                            help = 'Seed tables (class names, e.g. EphysRecording).')
+        project = _get_project()
+        parser.add_argument('-p', '--project', default = project, type = str, help = f'Source project: {project}')
+        parser.add_argument('--database-name', default = None,
+                            help = 'New database.name of the exported copy (required on the lab server).')
+        parser.add_argument('--dump', default = None,
+                            help = 'Write gzipped SQL files to this folder (with --docker --link-data: the data folder).')
+        parser.add_argument('--docker', default = None, help = 'Write a Docker build context to this folder.')
+        parser.add_argument('--link-data', action = 'store_true', default = False,
+                            help = 'With --docker: leave the rows out of the image; they go to a separate data folder '
+                                   '(--dump, default <docker folder>_data) mounted at /dump when the container starts.')
+        parser.add_argument('--tag', default = None, help = 'Image tag (with --docker).')
+        parser.add_argument('--build', action = 'store_true', default = False, help = 'Run docker build (with --docker).')
+        parser.add_argument('--platform', default = None, help = 'docker build --platform, e.g. linux/amd64.')
+        parser.add_argument('--pip', nargs = '+', default = [],
+                            help = 'Packages to install in the image (PyPI names or git+https:// URLs).')
+        parser.add_argument('--requirements', default = None, help = 'Requirements file to install in the image.')
+        parser.add_argument('--labdata-source', default = None,
+                            help = 'pip source of labdata for the image (default: the code running now).')
+        parser.add_argument('--no-dashboard', action = 'store_true', default = False,
+                            help = 'Leave the dashboard dependencies out of the image.')
+        parser.add_argument('--pref', nargs = '+', default = [],
+                            help = 'key=value entries for the preference file of the image (value read as JSON when it parses).')
+        parser.add_argument('--overwrite', action = 'store_true', default = False,
+                            help = 'Replace an existing --docker folder (and its data folder with --link-data).')
+        parser.add_argument('--files', default = 'none', choices = ['none', 'bucket', 'download'],
+                            help = 'Raw files: leave them (none), copy them to --files-storage (bucket) or to --files-path (download).')
+        parser.add_argument('--files-storage', default = None, help = 'Storage name of the preference file for --files bucket.')
+        parser.add_argument('--files-path', default = None, help = 'Folder for --files download.')
+        parser.add_argument('--blank-personal', action = 'store_true', default = False,
+                            help = 'Export LabMember rows with email, first and last name set to NULL.')
+        parser.add_argument('--strip', nargs = '+', default = [],
+                            help = 'Further nullable columns to export as NULL: Table.column (e.g. Subject.subject_dob).')
+        parser.add_argument('--no-descendants', action = 'store_true', default = False,
+                            help = 'Leave out the tables derived from the seeds.')
+        self._export_common_arguments(parser)
+        args = parser.parse_args(sys.argv[2:])
+        if not args.table and not args.plugin:
+            parser.error('give a table or --plugin.')
+        if args.link_data and args.docker is None:
+            parser.error('--link-data is used with --docker.')
+        if args.host is not None and (args.dump is not None or args.docker is not None):
+            parser.error('use one of --host, --dump or --docker.')
+        if args.dump is not None and args.docker is not None and not args.link_data:
+            parser.error('use one of --dump or --docker, or both with --link-data (the dump is the data folder).')
+        from .export import export, EXPORT_EXCLUDE, ExportCancelled
+        strip = dict()
+        for item in args.strip:
+            table, column = item.rsplit('.', 1)
+            strip.setdefault(table, []).append(column)
+        prefs_extra = dict()
+        for item in args.pref:
+            key, value = item.split('=', 1)
+            try:
+                prefs_extra[key] = json.loads(value)
+            except ValueError:
+                prefs_extra[key] = value
+        options = dict(project = args.project, restriction = self._export_restriction(args), plugins = args.plugin,
+                       database_name = args.database_name, descendants = not args.no_descendants,
+                       exclude = EXPORT_EXCLUDE if args.exclude is None else args.exclude,
+                       schema_map = dict(m.split('=', 1) for m in args.schema_map),
+                       blank_personal = args.blank_personal, strip_columns = strip,
+                       files = args.files, files_storage = args.files_storage, files_path = args.files_path,
+                       existing = args.existing, n_jobs = args.n_jobs, dry_run = args.dry_run)
+        if args.docker is not None:
+            from .export import build_export_image
+            db_user, db_password = ('labdata', None) if args.dry_run else self._image_login(args)
+            options.pop('existing')   # an image is built from a new dump folder
+            result = build_export_image(args.table or None, args.docker, tag = args.tag, pip = args.pip,
+                                        db_user = db_user, db_password = db_password,
+                                        requirements = args.requirements, labdata_source = args.labdata_source,
+                                        dashboard = not args.no_dashboard, platform = args.platform,
+                                        build = args.build, prefs_extra = prefs_extra,
+                                        link_data = args.link_data, data_dir = args.dump,
+                                        overwrite = args.overwrite, **options)
+        else:
+            target = None
+            if args.host is not None and not args.dry_run:
+                target = self._target_login(args)
+            elif args.dump is not None:
+                target = args.dump
+            elif not args.dry_run:
+                self._lab_server_login()
+            try:
+                result = export(args.table or None, target, **options)
+            except ExportCancelled as err:
+                print(err)
+                return
+            except ValueError as err:
+                print(tcolor['r'](str(err)))
+                exit(1)
+        if not args.dry_run:
+            self._print_export_report(result['report'])
+            if args.docker is not None and not args.build:
+                print(f"Build the image with: docker build -t {args.tag or '<tag>'} {args.docker}")
+
     def run_container(self):
         '''
         Runs a container
@@ -2051,7 +2435,54 @@ examples:
                                 type = str,
                                 help= f'Select project: {project}')
         return parser
-        
+
+    def _print_compute_tasks(self, schema, analysis = None, limit = 10):
+        '''Print the ComputeTask table by status, with the last line of each task log.
+
+        Shows the most recent ``limit`` tasks of each status and the count of the rest.
+        '''
+        tasks = schema.ComputeTask()
+        if not analysis is None:
+            tasks = tasks & dict(task_name = analysis)
+        title = 'ComputeTask table' + (f' [{analysis}]' if not analysis is None else '')
+        print(tcolor['m'](f"++++++++++++++++++++ {title} ++++++++++++++++++++"))
+        if not len(tasks):
+            print('\t no compute tasks.')
+            return
+        statuses = [('WORKING', tcolor['c']), ('WAITING', tcolor['y']), ('ON SERVER', tcolor['b']),
+                    ('FAILED', tcolor['r']), ('CANCELED', tcolor['o']), ('COMPLETED', tcolor['g'])]
+        def _brief(log, width = 100):
+            if log is None:
+                return ''
+            if isinstance(log, bytes):
+                log = log.decode('utf-8', errors = 'replace')
+            lines = [l.strip() for l in str(log).splitlines() if len(l.strip())]
+            line = ' '.join(lines[-1].split()) if len(lines) else ''
+            return line if len(line) <= width else line[:width-3] + '...'
+        def _time(t):
+            return t.strftime('%Y-%m-%d %H:%M') if not t is None else ''
+        for status, cc in statuses:
+            sel = tasks & f'task_status LIKE "%{status}%"'
+            n = len(sel)
+            if not n:
+                continue
+            print(cc(f"{status} ({n})"))
+            rows = sel.fetch('job_id', 'task_name', 'task_target', 'task_host', 'subject_name',
+                             'session_name', 'task_starttime', 'task_endtime', 'task_log',
+                             as_dict = True, order_by = 'job_id DESC', limit = limit)
+            for r in rows:
+                where = '/'.join([str(r[k]) for k in ('subject_name', 'session_name') if not r[k] is None])
+                when = _time(r['task_endtime'] if status in ('COMPLETED', 'FAILED', 'CANCELED') and
+                             not r['task_endtime'] is None else r['task_starttime'])
+                print(f"\t {r['job_id']} {r['task_name']} {where} "
+                      f"target: {r['task_target']} host: {r['task_host']}" +
+                      (f" {tcolor['y'](when)}" if len(when) else ''))
+                log = _brief(r['task_log'])
+                if len(log):
+                    print(f"\t\t{log}")
+            if n > limit:
+                print(f"\t ... and {n - limit} more")
+
     def _get_default_arg(self,argument,cli_arg = 'submit', default = None):
         # checks if there is a default in the options
         if not f'{cli_arg}_defaults' in labdata_preferences.keys():

{labdata-0.1.22 → labdata-0.1.24}/labdata/compute/__init__.py

@@ -5,6 +5,7 @@ Tasks can be scheduled and run on dedicated containers using job schedulers.
 
 Available compute tasks:
 - SpksCompute: Spike sorting using Kilosort/Phy via SPKS
+- SpikeInterfaceCompute: Spike sorting with the sorters of SpikeInterface
 - DeeplabcutCompute: Animal pose estimation using DeepLabCut
 - LightningPoseCompute: Animal pose estimation using Lightning Pose
 
@@ -32,6 +33,7 @@ The compute tasks handle:
 from .utils import *
 from .singularity import build_singularity_container
 from .ephys import SpksCompute
+from .spikeinterface import SpikeInterfaceCompute
 from .pose import DeeplabcutCompute, LightningPoseCompute
 from .caiman import CaimanCompute
 from .suite2p import Suite2pCompute