labdata 0.1.24__tar.gz → 0.1.25__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: labdata
-Version: 0.1.24
+Version: 0.1.25
 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.24 → labdata-0.1.25}/labdata/__init__.py

@@ -1,4 +1,4 @@
-VERSION = "0.1.24"
+VERSION = "0.1.25"
 __version__ = VERSION
 
 from .utils import *

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

@@ -32,7 +32,7 @@ class CLI_parser():
     logpipe
         Pipe stdout into a ``ComputeTask.task_log`` entry.
     run_container
-        Launch a local Apptainer container interactively.
+        Open a shell or Jupyter Lab in a container, here or on a node of a remote cluster.
     build_container
         Build Apptainer ``.sif`` images from definition files.
     cluster
@@ -56,7 +56,7 @@ Command to start the dashboard webpage:
 
 Data manipulation commands are:
 
-            {tcolor["y"]("subjects")}                            List subjects
+            {tcolor["y"]("subjects")}                            List subjects (--add <subject> adds one)
             {tcolor["y"]("sessions")} -a <subject>               List sessions 
             {tcolor["y"]("get")} -a <subject> -s <session>       Download data from one session if not already there
             {tcolor["y"]("put")} -a <subject> -s <session>       Copies a dataset to the server to be used
@@ -71,7 +71,7 @@ Data analysis commands:
 Other
             logpipe <compute_task_number>                Appends stdout log to a ComputeTask
             build-container                              Builds and uploads singularity/apptainer containers to S3            
-            run-container                                Launches a container in execution mode
+            run-container <container> [-t <remote>]      Opens a shell or Jupyter Lab in a container, here or on a cluster node
 
 Administration commands (don't run on experimental computers):
             {tcolor["r"]("upload")}                                       Sends pending data to S3 (applies upload rules)
@@ -207,6 +207,8 @@ Administration commands (don't run on experimental computers):
             labdata subjects -f <name_substring>
             labdata subjects -s M
             labdata subjects --include-size
+            labdata subjects --add <subject> --dob <YYYY-MM-DD> --sex <M|F|U> -u <user>
+            labdata subjects --add <subject> --dob <YYYY-MM-DD> --sex <M|F|U> --strain <strain> -u <user>
 
         Options
         -------
@@ -220,6 +222,12 @@ Administration commands (don't run on experimental computers):
             Project name.
         --include-size
             Show total raw data size (TB) and file count per subject.
+        --add
+            Add a subject instead of listing, with ``--dob``, ``--sex``, ``--strain``
+            and ``-u`` (the lab member). At a terminal, without ``--strain`` or with
+            one that is not on record, the strain is picked from a numbered list of
+            the Strain table. A missing option prints which are missing, with the
+            strains and users on record.
         '''
         parser = argparse.ArgumentParser(
             description='List sessions and datatypes',
@@ -246,7 +254,15 @@ Administration commands (don't run on experimental computers):
         parser.add_argument('--include-size', action = 'store_true',
                             default = False,
                             help= 'Display the size of the raw data for each subject.')
-        
+        parser.add_argument('--add', default = None, type = str, metavar = 'SUBJECT',
+                            help = 'Add a subject; needs --dob, --sex and -u, and --strain or a pick from the list.')
+        parser.add_argument('--dob', default = None, type = str,
+                            help = 'Date of birth of the subject to add (YYYY-MM-DD).')
+        parser.add_argument('--sex', default = None, type = str,
+                            help = 'Sex of the subject to add (M, F or U).')
+        parser.add_argument('--strain', default = None, type = str,
+                            help = 'Strain of the subject to add, from the Strain table; left out, pick it from a list.')
+
         args = parser.parse_args(sys.argv[2:])
         from .utils import load_project_schema
         from datajoint.errors import AccessError
@@ -256,6 +272,9 @@ Administration commands (don't run on experimental computers):
             txt = f'No permission to user {prefs["database"]["database.user"]} on schema {args.project}.'
             print(f'{tcolor["r"](txt)}')
             sys.exit(-1)
+        if not args.add is None:
+            self._add_subject(schema, args)
+            return
         query = schema.Subject()
         if not args.user in [None,'none']:
             query = query & f'user_name = "{args.user}"'
@@ -296,6 +315,74 @@ Administration commands (don't run on experimental computers):
                 print(tt,flush = True)
             print('')
 
+    def _add_subject(self, schema, args):
+        '''Insert the subject of ``labdata subjects --add`` after checking each field.'''
+        def stop(txt):
+            print(tcolor['r'](txt))
+            sys.exit(1)
+        name = args.add
+        if len(schema.Subject() & dict(subject_name = name)):
+            stop(f'Subject {name} is already in the database.')
+        # at a terminal the strain is chosen from a list; elsewhere --strain is required
+        interactive = can_prompt()
+        options = dict(dob = args.dob, sex = args.sex, u = args.user)
+        if not interactive:
+            options['strain'] = args.strain
+        missing = [('-' if k == 'u' else '--') + k for k, v in options.items() if v is None]
+        strains = list(schema.Strain().fetch('strain_name'))
+        users = schema.LabMember().fetch('user_name')
+        if len(missing):
+            stop(f'Adding {name} needs {", ".join(missing)}.\n'
+                 f'Strains: {", ".join(strains)}\nUsers: {", ".join(users)}')
+        try:
+            dob = datetime.strptime(args.dob, '%Y-%m-%d').date()
+        except ValueError:
+            stop(f'Date of birth {args.dob} is expected as YYYY-MM-DD.')
+        sex = args.sex.upper()
+        if not sex in ['M', 'F', 'U']:
+            stop(f'Sex {args.sex} is expected as M, F or U.')
+        if not args.user in users:
+            stop(f'User {args.user} is not in the LabMember table. Users: {", ".join(users)}')
+        strain = args.strain
+        if not strain in strains:
+            if not interactive:
+                stop(f'Strain {strain} is not in the Strain table. Strains: {", ".join(strains)}')
+            if not strain is None:
+                print(f'Strain {strain} is not in the Strain table.')
+            strain = self._choose_strain(schema)
+            if strain is None:
+                stop(f'Canceled adding subject {name}.')
+        subject = dict(subject_name = name, subject_dob = dob, subject_sex = sex,
+                       strain_name = strain, user_name = args.user)
+        try:
+            schema.Subject().insert1(subject)
+        except Exception as err:
+            stop(f'Could not add subject {name}: {err}')
+        print(f"{tcolor['g']('Added')} subject {name}\t{sex}\t{dob}\t{strain}\t{args.user}")
+
+    def _choose_strain(self, schema):
+        '''Return the strain picked at the prompt from a numbered list of the Strain table, or None.
+
+        Each strain shows the number of subjects of the project that have it.
+        '''
+        counts = schema.Strain.aggr(schema.Subject, n = 'count(subject_name)', keep_all_rows = True)
+        strains, n = counts.fetch('strain_name', 'n', order_by = 'strain_name')
+        print('Strains:')
+        for i, (strain, count) in enumerate(zip(strains, n)):
+            print(f'\t{tcolor["c"](str(i+1) + ".")} {strain}  ({count} subjects)')
+        while True:
+            try:
+                answer = input(f'Strain number, 1-{len(strains)} (enter to cancel): ').strip()
+            except EOFError:
+                return None
+            if answer == '':
+                return None
+            if answer.isdigit() and 1 <= int(answer) <= len(strains):
+                return strains[int(answer) - 1]
+            if answer in strains:
+                return answer
+            print(f'{answer} is not on the list; type a number from 1 to {len(strains)}.')
+
     def sessions(self):
         '''List sessions and datasets for one or more subjects.
 
@@ -444,6 +531,9 @@ Administration commands (don't run on experimental computers):
         database for the requested upload rule.  **This does not upload to the
         cloud** — that happens when ``labdata upload`` is later run on the server.
 
+        When Qt cannot be imported, or on Linux without a display, the files are
+        copied from the terminal with ``labdata.copy.copy_to_upload_server``.
+
         Usage
         -----
         ::
@@ -509,13 +599,26 @@ Administration commands (don't run on experimental computers):
             w = LABDATA_PUT()
             sys.exit(app.exec_())
         else:
-            from .widgets import QApplication, ServerCopyWidget,QMessageBox
-            app = QApplication(sys.argv)
             filepaths = args.filepaths
             if len(filepaths) == 1:
                 if Path(filepaths[0]).is_dir(): # then select all files inside the folder
                     filepaths = list(Path(filepaths[0]).rglob('**/*'))
                     filepaths = list(filter(lambda f: f.is_file(),filepaths))
+            # without Qt or a display, copy from the terminal with copy_to_upload_server
+            reason = None
+            if sys.platform.startswith('linux') and not (os.environ.get('DISPLAY') or
+                                                         os.environ.get('WAYLAND_DISPLAY')):
+                reason = 'There is no display for the copy window'
+            else:
+                try:
+                    from .widgets import QApplication, ServerCopyWidget,QMessageBox
+                except ImportError as err:
+                    reason = f'Qt is not available ({err})'
+            if not reason is None:
+                print(f'{reason}; copying from the terminal.', flush = True)
+                self._put_from_terminal(filepaths, args)
+                return
+            app = QApplication(sys.argv)
             try:
                 w = ServerCopyWidget(src_filepaths = filepaths,
                                      upload_rule = args.rule,
@@ -531,6 +634,44 @@ Administration commands (don't run on experimental computers):
                 raise(OSError(f"{tcolor['r']('Upload failed')} {err}."))
             app.exit()
             #sys.exit(app.exec_())
+
+    def _put_from_terminal(self, filepaths, args):
+        '''Copy files to the upload server with ``copy_to_upload_server``, without the Qt window.
+
+        ``labdata put`` uses this when Qt cannot be imported or there is no display.
+        '''
+        from .copy import copy_to_upload_server
+        filepaths = [Path(f).resolve() for f in filepaths]
+        if not len(filepaths):
+            print(tcolor['r']('No files to copy.'))
+            return
+        missing = [str(f) for f in filepaths if not f.is_file()]
+        if len(missing):
+            raise OSError(f"{tcolor['r']('Upload failed')}: these are not files {missing}.")
+        names = dict(subject_name = args.subject,
+                     session_name = args.session,
+                     dataset_name = args.datatype,
+                     dataset_type = args.datatype_name)
+        names = {k: v[0] for k, v in names.items() if not v is None}
+        rule = args.rule[0] if not args.rule is None else None
+        nbytes = sum(f.stat().st_size for f in filepaths)
+        txt = f"{len(filepaths)} files ({nbytes/1e9:.2f} GB) to {prefs.get('upload_path')}"
+        if args.ask:
+            if input(f'Copy {txt}? [y/N] ').strip().lower() not in ('y', 'yes'):
+                print(tcolor['r']('Canceled the copy to the upload server.'))
+                return
+        print(f'Copying {txt}.', flush = True)
+        try:
+            res = copy_to_upload_server(filepaths,
+                                        job_rule = rule,
+                                        overwrite = args.overwrite,
+                                        project = args.project,
+                                        **names)
+        except Exception as err:
+            raise OSError(f"{tcolor['r']('Upload failed')} {err}.") from err
+        print(f"{tcolor['g']('Copied')} {len(res)} files: upload job {res[0]['job_id']} "
+              f"waits for labdata upload on {prefs.get('upload_host')}.", flush = True)
+
     def clean(self):
         '''Delete local copies of files that are confirmed uploaded to S3.
 
@@ -2305,26 +2446,112 @@ the preference file, or set them with LABDATA_DATABASE_PROJECT for the project.
                 print(f"Build the image with: docker build -t {args.tag or '<tag>'} {args.docker}")
 
     def run_container(self):
+        '''Open a shell or Jupyter Lab in a container, here or on a node of a remote cluster.
+
+        Usage
+        -----
+        ::
+
+            labdata run-container <container>
+            labdata run-container <container> --jupyter
+            labdata run-container <container> -t <remote> [--partition gpu] [--walltime 2:00:00]
+            labdata run-container <container> -t <remote> --jupyter
+            labdata run-container -t <remote>
+            labdata run-container -t <remote> --attach <slurm job id>
+
+        With ``-t`` naming a remote of the preference file, the command submits
+        a SLURM job there, waits until it starts and opens a shell in the
+        container on its node, or with ``--jupyter`` starts Jupyter Lab there
+        and tunnels it to this computer. The session has the environment of a
+        labdata job on that remote (project, credentials, ``pre_cmds``). Ending
+        the session ends the job; a lost connection leaves it running, to
+        reconnect with ``--attach``. Without a container name on a remote the
+        shell opens on the node itself.
+
+        Options
+        -------
+        container
+            Container name: the .sif file in the containers folder, without extension.
+        -t, --target
+            ``local`` (default), or a remote of the preference file.
+        --jupyter
+            Start Jupyter Lab instead of a shell.
+        --writable, --cuda
+            Local only: open the container writable, give it the GPUs. On a
+            remote the container gets the GPUs when the node has them.
+        --partition, --ncpus, --gpus, --memory, --walltime, --exclusive, --no-exclusive
+            Resources of the job on a remote. The defaults come from the
+            ``interactive`` entry of the remote's ``analysis_options``; without
+            it the job takes a whole node for 4 hours.
+        --attach
+            SLURM job id of a running job to open the session in.
+        --keep
+            Leave the job running when the session ends.
+        -p, --project
+            Project of the session on a remote.
         '''
-        Runs a container
-        '''
+        remotes = list((prefs['compute'].get('remotes') or {}).keys())
         parser = argparse.ArgumentParser(
-            description = 'Run a container',
-            usage = '''labdata run-container <container> ''')
+            description = 'Open a shell or Jupyter Lab in a container, here or on a node of a remote',
+            usage = '''labdata run-container <container> [-t <remote>] [--jupyter]''')
         parser.add_argument('container_name',
                             action = 'store',
-                            type = str,
-                            help = 'Container name (the .sif file in the local containers folder, without extension).')
+                            type = str, nargs = '?', default = None,
+                            help = 'Container name (the .sif file in the containers folder, without extension).')
         parser.add_argument('-t','--target',action = 'store', default = 'local', type = str,
-                            help = 'Where to run it: local or slurm.')
+                            help = f'Where to run it: local, or a remote [{tcolor["y"](", ".join(remotes))}].')
         parser.add_argument('--jupyter', action = 'store_true', default = False,
                             help = 'Start Jupyter Lab in the container instead of a shell.')
         parser.add_argument('--writable', action = 'store_true', default = False,
-                            help = 'Open the container writable (for sandbox containers).')
+                            help = 'Open the container writable (for sandbox containers, local only).')
         parser.add_argument('--cuda',action = 'store_true', default = False,
-                            help = 'Give the container access to the GPUs.')
+                            help = 'Give the container access to the GPUs (local only; on a remote when the node has them).')
+        parser.add_argument('--partition', action = 'store', default = None, type = str,
+                            help = 'SLURM partition of the job on a remote.')
+        parser.add_argument('--ncpus', action = 'store', default = None, type = int,
+                            help = 'CPUs of the job on a remote.')
+        parser.add_argument('--gpus', action = 'store', default = None, type = str,
+                            help = 'GPUs of the job on a remote (SLURM --gpus).')
+        parser.add_argument('--memory', action = 'store', default = None, type = str,
+                            help = 'Memory of the job on a remote, e.g. 32G.')
+        parser.add_argument('--walltime', action = 'store', default = None, type = str,
+                            help = 'Time limit of the job on a remote (default 4:00:00).')
+        parser.add_argument('--exclusive', action = 'store_const', const = True, dest = 'exclusive', default = None,
+                            help = 'Take the whole node (the default on a remote).')
+        parser.add_argument('--no-exclusive', action = 'store_const', const = False, dest = 'exclusive',
+                            help = 'Share the node with other jobs.')
+        parser.add_argument('--attach', action = 'store', default = None, type = int,
+                            help = 'SLURM job id (not a ComputeTask id) of a running job to open the session in.')
+        parser.add_argument('--keep', action = 'store_true', default = False,
+                            help = 'Leave the job running when the session ends.')
+        parser = self._add_default_arguments(parser, level = 0, include_project = True)
 
         args = parser.parse_args(sys.argv[2:])
+        if args.target in remotes:
+            from .compute.schedulers import run_container_remote
+            try:
+                run_container_remote(args.target,
+                                     container = args.container_name,
+                                     jupyter = args.jupyter,
+                                     partition = args.partition,
+                                     ncpus = args.ncpus,
+                                     gpus = args.gpus,
+                                     memory = args.memory,
+                                     walltime = args.walltime,
+                                     exclusive = args.exclusive,
+                                     attach = args.attach,
+                                     keep = args.keep,
+                                     project = args.project)
+            except (RuntimeError, ValueError) as err:
+                print(tcolor['r'](str(err)))
+                sys.exit(1)
+            return
+        if not args.target in ['local','slurm']:
+            print(f"{tcolor['r'](args.target)} is not a target: use local or one of the remotes "
+                  f"[{', '.join(remotes)}].")
+            sys.exit(1)
+        if args.container_name is None:
+            parser.error('the container name is needed here; it is optional only on a remote.')
         container = args.container_name
         container_file = None
         cmd = None
@@ -2334,11 +2561,10 @@ the preference file, or set them with LABDATA_DATABASE_PROJECT for the project.
             container_file = (container_store/container).with_suffix('.sif')
             if not container_file.exists():
                 print( f"Container {tcolor['r'](container)} not found in {tcolor['r'](container_store)}")
-                return 
-        
+                return
+
         if args.jupyter:
             cmd = 'jupyter lab --ip="*"'
-        print(args)
         if not cmd is None:
             from .compute.singularity import run_on_apptainer
             launchcmd = run_on_apptainer(container_file,

{labdata-0.1.24 → labdata-0.1.25}/labdata/compute/ephys.py

@@ -622,9 +622,10 @@ class SpksCompute(BaseCompute):
                                             **channel_statistics(segment.astype(np.float32)*gains)))
             del dat
         self.set_job_status(job_log = f'Completed {base_key}')
-        from labdata.schema import UnitMetrics
-        # limit number of jobs because of memory constraints
-        self.schema.UnitMetrics.populate(base_key, display_progress = True, processes = int(max(1,np.ceil(n_jobs/2))))
+        # every dataset of the sorting; limit number of jobs because of memory constraints
+        self.schema.UnitMetrics.populate(base_keys, display_progress = True, processes = int(max(1,np.ceil(n_jobs/2))))
+        # label and count the units under every computed criteria of the project
+        self.schema.UnitCount.populate(base_keys, display_progress = True)
     
     def extract_waveforms(self, udict, nchannels, results_folder, n_pre_samples, n_jobs, offsets, gains):
         '''Collect the waveforms of the units of one dataset from the filtered binary.