spikeforest 0.12.9__py3-none-any.whl

spikeforest/__init__.py

@@ -0,0 +1,9 @@
+from .load_extractors.load_recording_extractor import load_recording_extractor
+from .load_extractors.load_sorting_extractor import load_sorting_extractor
+from .load_spikeforest_recordings.load_spikeforest_recordings import load_spikeforest_recordings
+from .load_spikeforest_recordings.load_spikeforest_recording import load_spikeforest_recording
+from .load_spikeforest_sorting_outputs.load_spikeforest_sorting_outputs import load_spikeforest_sorting_outputs
+from .load_spikeforest_sorting_outputs.load_spikeforest_sorting_output import load_spikeforest_sorting_output
+from .load_extractors import load_recording_extractor, load_sorting_extractor
+
+from .version import __version__

spikeforest/_common/calling_framework.py

@@ -0,0 +1,173 @@
+#!/usr/bin/python
+from argparse import ArgumentParser, Namespace
+
+import datetime
+import os
+from typing import Any, TypedDict, NamedTuple, Union
+import hither2 as hi
+
+
+# NamedTuple is probably cleaner, but keeping a dict is more convenient for screen output.
+class StandardArgs(TypedDict):
+    test: int
+    timeout_min: int
+    outfile: str
+    workercount: int
+    job_cache: Union[str, None]
+    use_container: bool
+    use_slurm: bool
+    slurm_max_jobs_per_alloc: int
+    slurm_max_simultaneous_allocs: int
+    slurm_command: str
+
+class HitherConfiguration(TypedDict):
+    job_handler: Any
+    job_cache: Any
+    use_container: bool
+    log: Any
+
+
+RECORDING_URI_KEY = 'recordingUri'
+GROUND_TRUTH_URI_KEY = 'sortingTrueUri'
+SORTING_FIRINGS_URI_KEY = 'firings'
+
+def add_standard_args(parser: ArgumentParser) -> ArgumentParser:
+    """Adds standard command-line arguments for interacting with hither/slurm calling conventions.
+    Included arguments are --verbose (-v|vv|vvv...), --test (-t), --outfile (-o), --workercount (-w),
+    --job-cache, --no-job-cache, --use-container, --no-container, --use-slurm, --slurm-partition,
+    --slurm-accept-shared-nodes, --slurm-jobs-per-allocation, --slurm-max-simultaneous-allocations,
+    --slurm-gpus-per-node, --timeout-min, and --check-config.
+
+    Args:
+        parser (argparse.ArgumentParser): An initialized argparse ArgumentParser to extend.
+
+    Returns:
+        (argparse.Namespace): The result of calling parser.parse_args().
+    """
+    parser.add_argument('--verbose', '-v', action='count', default=0,
+        help="Set verbosity level. Add vs for more verbosity.")
+    # Note: Whatever 'number of iterations' means for your application should be locally defined.
+    parser.add_argument('--test', '-t', action='store', type=int, default=0,
+        help="If non-zero, this will set a maximum number of iterations before quitting, " +
+        "to give a usable sample without processing the entire data set.")
+    parser.add_argument('--timeout-min', '-T', action='store', type=int, default=0,
+        help="If non-zero, this will set a maximum duration for any job before it is cancelled.")
+    parser.add_argument('--outfile', '-o', action='store', default=None,
+        help='If set, output (but not warnings/messages) will be written to this file (instead of to STDOUT). ' +
+             'Any existing file will NOT be overwritten; the program will abort instead.')
+    parser.add_argument('--workercount', '-w', action='store', type=int, default=4,
+        help="If set, determines the number of worker threads for a parallel job handler. Ignored if using slurm.")
+    parser.add_argument('--job-cache', action='store', type=str, default='default-job-cache',
+        help="If set, indicates the feed name for the job cache feed.")
+    parser.add_argument('--no-job-cache', action='store_true', default=False,
+        help="If set, job cache will not be used, and any value for --job-cache will be ignored.")
+    parser.add_argument('--use-container', '-C', action='store_true', default=False,
+        help='If set, hither calls will use containerization. If unset, containerization may still be used if ' +
+        'environment variable HITHER_USE_CONTAINER is set to "1" or "TRUE".')
+    parser.add_argument('--no-container', action='store_true', default=False,
+        help='Override HITHER_USE_CONTAINER environment variable to suppress container use. Ignored if --use-container is set.')
+    parser.add_argument('--use-slurm', action='store_true', default=False,
+        help='If set, this script will use a SlurmJobHandler and attempt to run jobs on the configured cluster. The exact ' +
+        'call used by the slurm job handler to acquire resources can be customized with command-line arguments.')
+    parser.add_argument('--slurm-partition', action='store', type=str, default="ccm",
+        help='If set, slurm will use the specified text as a partition name to request. Note that slurm must be explicitly ' +
+        'requested with the --use-slurm flag; if it is not, this value is ignored.')
+    parser.add_argument('--slurm-accept-shared-nodes', action='store_true', default=False,
+        help='If set, slurm calls will be made without --exclusive. Note that slurm must still be explicitly ' +
+        'requested with the --use-slurm flag; if it is not, this value is ignored.')
+    parser.add_argument('--slurm-jobs-per-allocation', action='store', type=int, default=6,
+        help='Controls the max length of job processing queues for slurm nodes. Default 6.')
+    parser.add_argument('--slurm-max-simultaneous-allocations', action='store', type=int, default=5,
+        help='The maximum number of job processing queues/slurm nodes to be requested. Default 5.')
+    parser.add_argument('--slurm-gpus-per-node', action='store', type=int, default=0,
+        help='If set, slurm commands will require this many GPUs per allocated node.')
+    parser.add_argument('--check-config', action='store_true', default=False,
+        help='Debugging tool. If set, program will simply quit with a description of the parsed configuration.')
+    return parser
+
+def parse_shared_configuration(parsed: Namespace) -> StandardArgs:
+    """Generates a dictionary of configuration options for running jobs against hither, regardless of back-end.
+
+    Args:
+        parsed (Any): parsed ArgParse parser, with at least the keys defined in add_standard_args.
+
+    Raises:
+        Exception: Fatal error if attempting to use an existing file as an output.
+
+    Returns:
+        StandardArgs: A dictionary of compiled run values, to be passed to future calls from this file.
+    """    
+
+    # example srun_command: srun --exclusive -n 1 -p <partition>
+    slurm_command = f"srun -n 1 -p {parsed.slurm_partition} {'--exclusive' if not parsed.slurm_accept_shared_nodes else ''}"
+    if parsed.slurm_gpus_per_node > 0:
+        slurm_command += f" --gpus-per-node={parsed.slurm_gpus_per_node}"
+    if parsed.outfile is not None and parsed.outfile != '' and os.path.exists(parsed.outfile) and parsed.outfile != "/dev/null":
+        raise Exception('Error: Requested to write to an existing output file. Aborting to avoid overwriting file.')
+    # configure verbosity for the run
+    print_per_verbose.__dict__['verbosity_level'] = parsed.verbose or 0
+
+    # We would very much like to avoid this manual copying, unfortunately the argsparse module and the typing module
+    # don't play at all nicely with each other. Maybe fix this later (keyword TAP/typed arg parser)
+    return StandardArgs(
+        test                          = parsed.test,
+        timeout_min                   = parsed.timeout_min,
+        outfile                       = parsed.outfile,
+        workercount                   = max(parsed.workercount, 1),
+        # As a reminder, argparse converts internal -es to _s to keep the identifiers valid
+        job_cache                     = None if parsed.no_job_cache else parsed.job_cache,
+        use_container                 = parsed.use_container or \
+                                            ((not parsed.no_container) \
+                                             and os.getenv('HITHER_USE_CONTAINER') in ['TRUE', '1']),
+        use_slurm                     = parsed.use_slurm,
+        slurm_max_jobs_per_alloc      = parsed.slurm_jobs_per_allocation,
+        slurm_max_simultaneous_allocs = parsed.slurm_max_simultaneous_allocations,
+        slurm_command                 = slurm_command
+    )
+
+# TODO: print_per_verbose, _fmt_time belong in a different file?
+def print_per_verbose(lvl: int, msg: str) -> None:
+    # verbosity_level is a static value, initialized from command-line argument at setup time in init_configuration().
+    # This does not play nicely with containerization, but global arguments variables don't either; rather than passing
+    # needless verbosity parameters around, we're just going to stay silent if the verbosity_level property isn't set.
+    if ('verbosity_level' not in print_per_verbose.__dict__): return
+    if (print_per_verbose.verbosity_level < lvl): return
+    tabs = max(0, lvl - 1)
+    print("\t" * tabs + msg)
+
+def _fmt_time(t: Union[float, None]) -> str:
+    if not t: return 'TIME NOT SPECIFIED'
+    return datetime.datetime.fromtimestamp(t).isoformat()
+
+def extract_hither_config(args: StandardArgs) -> HitherConfiguration:
+    use_container = args['use_container']
+    if args['test'] != 0: print(f"\tRunning in TEST MODE--Execution will stop after processing {args['test']} sortings!\n")
+
+    if use_container:
+        print_per_verbose(1, f"Using {'Singularity' if os.getenv('HITHER_USE_SINGULARITY') else 'Docker'} containers.")
+    else:
+        print_per_verbose(1, "Running without containers.")
+    # Define job cache and job handler
+    jc = None if args['job_cache'] == None else hi.JobCache(feed_name=args['job_cache'])
+    if args['use_slurm']:
+        jh = hi.SlurmJobHandler(
+            num_jobs_per_allocation=args['slurm_max_jobs_per_alloc'],
+            max_simultaneous_allocations=args['slurm_max_simultaneous_allocs'],
+            srun_command=args['slurm_command']
+        )
+    else:
+        jh = hi.ParallelJobHandler(num_workers=args['workercount'])
+    log = hi.Log()
+    timeout_sec = None if args['timeout_min'] == 0 else 60 * args['timeout_min']
+    return HitherConfiguration(
+        job_cache=jc,
+        job_handler=jh,
+        use_container=use_container,
+        job_timeout_sec=timeout_sec,
+        show_console=False,
+        log=log
+    )
+    
+def call_cleanup(config: HitherConfiguration) -> None:
+    if config['job_handler'] is not None:
+        config['job_handler'].cleanup()