ch-frb-make-acq-inventory

#!/usr/bin/env python

import sys
import argparse


####################################################################################################
#
# Argument parsing, checking


class MyParser(argparse.ArgumentParser):
    def print_usage(self, file=None):
        if file is None:
            file = sys.stdout

        print >>file, "Usage: ch-frb-make-acq-inventory [-D] [--nchunks-per-part=NN] [--nchunks-overlap=NN] <device>"
        print >>file, "The <device> argument should be either 'ssd' or 'nfs'."
        print >>file, "    --nchunks-per-part: long acqs will be split into parts, with specified number of 1s chunks (default 3600)"
        print >>file, "    --nchunks-overlap: if acq is split into parts, they overlap by specified number of 1s chunks (default 600)"
        print >>file, "    -D: do not back up old 'acq_json' directory (by default, a backup 'acq_json_deleteme_xxx' is created"
        # print >>file, '    -h: show longer help message and exit'


    def error(self, message=None):
        self.print_usage(sys.stderr)

        if message is not None:
            print >>sys.stderr, '\nError:', message

        sys.exit(2)

    
    def print_help(self, file=None):
        self.print_usage(file)

        # print >>file
        # print >>file, "HELP TEXT HERE"


parser = MyParser()

parser.add_argument('device')
parser.add_argument('-D', action='store_true', help='runs the pipeline with no output directory')
parser.add_argument('--nchunks-per-part', dest='nchunks_per_part', type=int, default=3600)
parser.add_argument('--nchunks-overlap', dest='nchunks_overlap', type=int, default=600)

args = parser.parse_args()

# Globals.
# Each acq is now divided into multiple parts.
# We take the parts to be 60 minutes long, overlapping by 10 minutes.
nchunks_per_part = args.nchunks_per_part
nchunks_overlap = args.nchunks_overlap
toplevel_dir = None
nfs_hack = False

if args.device == 'ssd':
    toplevel_dir = '/local'
elif args.device == 'nfs':
    toplevel_dir = '/frb-archiver-1'
    nfs_hack = True
else:
    parser.error("Fatal: <device> argument must be either 'ssd' or 'nfs'.")

if nchunks_overlap < 0:
    parser.error("Fatal: 'nchunks_overlap' must be >= 0")
if nchunks_per_part <= nchunks_overlap:
    parser.error("Fatal: 'nchunks_per_part' must be greater than 'nchunks_overlap'")


####################################################################################################


import os
import re
import time
import json
import errno
import rf_pipelines


def xmakedirs(name):
    """A version of os.makedirs() which does not raise an exception if the directory already exists."""

    # Ensures correct handling of
    #   makedirs(os.path.dirname(f))
    # in the case where f is just an ordinary filename with no slashes

    if name == '':
        return

    try:
        os.makedirs(name)
    except OSError, exc:
        if (exc.errno != errno.EEXIST) or not os.path.isdir(name):
            raise


def write_runlist(json_filename, run_list):
    if len(run_list) == 0:
        return

    f = open(json_filename, 'w')
    json.dump(run_list, f, indent=4)
    print >>f, ''  # extra newline

    f.close()
    print 'wrote %s' % json_filename


def replace_beginning_of_string(s, old_prefix, new_prefix):
    assert s.startswith(old_prefix)
    return new_prefix + s[len(old_prefix):]


def delete_json_dir(dirname):
    """Returns True if deleted, False if directory could not be deleted due to stray files."""

    ret = True

    for b in os.listdir(dirname):
        if b in [ '.', '..' ]:
            continue
        
        a = os.path.join(dirname, b)

        if os.path.isdir(a):
            if not delete_json_dir(a):
                ret = False
        elif a.endswith('.json'):
            os.remove(a)
        else:
            print 'Stray file %s, directory will not be deleted' % a
            ret = False

    if ret:
        os.rmdir(dirname)

    return ret


####################################################################################################


class beamdir:
    """
    Represents one beam directory of the form ACQDIR/beam_NNNN.
    We "inventory" all data files in the directory, and divide them into overlapping parts.
    """

    def __init__(self, parent_abspath, beamdir_basename):
        m = re.match(r'beam_(\d+)$', beamdir_basename)
        if not m:
            raise RuntimeError("internal error: bad beamdir_basename '%s'" % beamdir_basename)

        self.beam_id = int(m.group(1))
        self.parent_abspath = parent_abspath
        self.parent_basename = os.path.basename(parent_abspath)  # = acqname
        self.beamdir_basename = beamdir_basename
        self.beamdir_abspath = os.path.join(parent_abspath, beamdir_basename)
        self.data_file_basenames = [ ]
        self.in_progress = False
        self.retry_time = 0.0
        self.nfiles = 0
        
        # System time before the readdir() call
        t0 = time.time()

        # Hash ichunk -> basename
        h = { }

        for f in os.listdir(self.beamdir_abspath):
            m = re.match(r'chunk_(\d+).msg$', f)

            if not m:
                f = os.path.join(self.beamdir_abspath, f)
                print >>sys.stderr, 'warning: stray file %s in acquisition directory' % f
                continue

            n = int(m.group(1))

            if h.has_key(n):
                f1 = os.path.join(self.beamdir_abspath, f)
                f2 = os.path.join(self.beamdir_abspath, h[n])
                raise RuntimeError('ambiguous file pair (%s, %s)' % (f1,f2))

            h[n] = f

        if len(h) == 0:
            return

        n0 = min(h.keys())
        n1 = max(h.keys())+1

        if sorted(h.keys()) != range(n0, n1):
            raise RuntimeError('non-contiguous time indices in acquisition directory %s' % self.beamdir_abspath)

        self.data_file_basenames = [ h[n] for n in xrange(n0,n1) ]
        self.nfiles = len(self.data_file_basenames)

        # Time when last file was changed
        last_file = os.path.join(self.beamdir_abspath, self.data_file_basenames[-1])
        t1 = os.stat(last_file).st_ctime

        if t1 >= t0 - 60.0:
            self.in_progress = True
            self.retry_time = t1 - t0 + 60.0


    def __repr__(self):
        return "beamdir('%s', files=[%s .. %s])" % (self.beamdir_abspath, self.data_file_basenames[0], self.data_file_basenames[-1])

    
    def write_json_files(self, json_dir):
        """Returns runlist, i.e. list of (run_name, json_basename) pairs."""

        runlist = [ ]

        if self.in_progress:
            nparts = self.nfiles // nchunks_per_part
        else:
            nparts = self.nfiles / float(nchunks_per_part)
            nparts = max(int(nparts+0.5), 1)

        for ipart in xrange(nparts):
            i0 = (ipart) * nchunks_per_part
            i1 = (ipart+1) * nchunks_per_part
            
            if (ipart > 0):
                i0 -= nchunks_overlap
            if (ipart == (nparts-1)) and not self.in_progress:
                i1 = self.nfiles

            s = ('_part%d' % ipart) if (nparts > 1) else ''
            run_name = '%s_beam%d%s' % (self.parent_basename, self.beam_id, s)
            json_basename = run_name + '.json'
            json_filename = os.path.join(json_dir, json_basename)
            data_abspath_list = [ ]

            if not nfs_hack:
                # Normal case
                data_abspath_list = [ os.path.join(self.beamdir_abspath, self.data_file_basenames[i]) for i in xrange(i0,i1) ]
            else:
                # Hack for NFS load-balancing
                for i in xrange(i0,i1):
                    d = replace_beginning_of_string(self.beamdir_abspath, '/frb-archiver-1/', ('/frb-archiver-%d/' % (i%4+1)))
                    f = os.path.join(d, self.data_file_basenames[i])
                    data_abspath_list.append(f)
        
            # rf_pipelines stream object
            stream = rf_pipelines.chime_frb_stream_from_filename_list(data_abspath_list)
            rf_pipelines.utils.json_write(json_filename, stream, clobber=True, verbose=True)

            runlist.append((run_name, json_basename))

        return runlist


    def is_nonempty(self):
        return (self.nfiles >= nchunks_per_part) if self.in_progress else (self.nfiles > 0)


class acqdir:
    def __init__(self, parent_abspath, acqdir_basename):
        self.parent_abspath = parent_abspath
        self.acqdir_basename = acqdir_basename
        self.acqdir_abspath = os.path.join(parent_abspath, acqdir_basename)
        self.beamdir_hash = { }   # hash (integer beam_id) -> (beamdir object)
        self.in_progress = False
        self.retry_time = 0.0

        for beamdir_basename in os.listdir(self.acqdir_abspath):
            beamdir_abspath = os.path.join(self.acqdir_abspath, beamdir_basename)

            m = re.match(r'beam_(\d+)$', beamdir_basename)
            if not m:
                continue
            if not os.path.isdir(beamdir_abspath):
                continue

            beam_id = int(m.group(1))
            if (beam_id < 0) or (beam_id >= 10000):
                raise RuntimeError('bad beam id %d (appearing in filename %s)' % (beam_id, beamdir_abspath))

            if self.beamdir_hash.has_key(beam_id):
                d1 = self.beamdir_hash[beam_id].beamdir_abspath
                d2 = beamdir_abspath
                raise RuntimeError('ambiguous directory pair (%s, %s)' % (d1,d2))

            bd = beamdir(self.acqdir_abspath, beamdir_basename)

            if bd.nfiles == 0:
                continue

            self.beamdir_hash[beam_id] = bd

            if bd.in_progress:
                self.in_progress = True
                self.retry_time = max(self.retry_time, bd.retry_time)

        if self.in_progress:
            for bd in self.beamdir_hash.values():
                bd.in_progress = True
                bd.retry_time = self.retry_time


    def __repr__(self):
        return "acqdir('%s', %s)" % (self.parent_abspath, self.beamdir_hash)


    def write_json_files(self, json_dir):
        """Returns runlist, i.e. list of (run_name, json_relpath) pairs."""

        xmakedirs(json_dir)
        runlist = [ ]

        for (beam_id, bd) in sorted(self.beamdir_hash.iteritems()):
            runlist2 = bd.write_json_files(json_dir)
            runlist.extend(runlist2)

        runlist_filename = os.path.join(json_dir, ('%s_runlist.json' % self.acqdir_basename))
        write_runlist(runlist_filename, runlist)

        return runlist


    def is_nonempty(self):
        for bd in self.beamdir_hash.values():
            if bd.is_nonempty():
                return True

        return False


class acq_inventory:
    def __init__(self, toplevel_abspath):
        self.toplevel_abspath = toplevel_abspath
        self.acqdir_hash = { }   # hash (acq_name) -> (acqdir object)

        for acqdir_basename in os.listdir(toplevel_abspath):
            acqdir_abspath = os.path.join(toplevel_abspath, acqdir_basename)
            if not os.path.isdir(acqdir_abspath):
                continue

            ad = acqdir(toplevel_abspath, acqdir_basename)
            if len(ad.beamdir_hash) > 0:
                self.acqdir_hash[acqdir_basename] = ad


    def __repr__(self):
        return "acq_inventory('%s', %s)" % (self.toplevel_abspath, self.acqdir_hash)

    
    def write_json_files(self, json_dir):
        master_run_list = [ ]

        for (acq_name, ad) in sorted(self.acqdir_hash.iteritems()):
            run_list = ad.write_json_files(os.path.join(json_dir, acq_name))
            run_list = [ (run_name, os.path.join(ad.acqdir_basename,f)) for (run_name,f) in run_list ]
            master_run_list.extend(run_list)

        runlist_filename = os.path.join(json_dir, 'master_runlist.json')
        write_runlist(runlist_filename, master_run_list)

        
    def is_nonempty(self):
        for ad in self.acqdir_hash.values():
            if ad.is_nonempty():
                return True

        return False


    def announce_acquisitions_in_progress(self):
        for (acq_name, ad) in sorted(self.acqdir_hash.iteritems()):
            if ad.in_progress:
                print '%s: classified as acquisition in progress, the most recent ~hour of data may not be indexed' % acq_name
                print '  If this classification is incorrect, try rerunning ch-frb-make-acq-inventory in %s seconds' % ad.retry_time


####################################################################################################


# Reminder: global variable 'toplevel_dir' is defined near the beginning of this file.
user_dir = os.path.join(toplevel_dir, os.environ['USER'])
input_acqdir = os.path.join(toplevel_dir, 'acq_data')
output_json_dir = os.path.join(user_dir, 'acq_json')
backup_json_dir = os.path.join(user_dir, ('acq_json_deleteme_%s' % time.strftime('%y-%m-%d-%X')))

if not os.path.exists(input_acqdir):
    raise RuntimeError("acqdir '%s' does not exist, maybe filesystem is not mounted?")

print >>sys.stdout, 'Inspecting %s, this may take a minute...' % input_acqdir
sys.stdout.flush()

a = acq_inventory(input_acqdir)

if not a.is_nonempty():
    print 'No acqusitions could be indexed, acq_json directory will not be created.'
    a.announce_acquisitions_in_progress()
    sys.exit(0)


backup_flag = False

if os.path.exists(output_json_dir):
    assert not os.path.exists(backup_json_dir)
    print 'Backing up %s -> %s' % (output_json_dir, backup_json_dir)
    os.rename(output_json_dir, backup_json_dir)

    assert not os.path.exists(output_json_dir)
    backup_flag = True

try:
    a.write_json_files(output_json_dir)
except:
    if backup_flag:
        print 'Restoring backup %s -> %s' % (backup_json_dir, output_json_dir)
        os.rename(backup_json_dir, output_json_dir)
    raise

a.announce_acquisitions_in_progress()

if args.D and backup_flag:
    print 'Deleting backup %s' % backup_json_dir
    delete_json_dir(backup_json_dir)
else:
    print "Note: backup directory '%s' will need to be deleted by hand (to do this automatically, use the -D flag)" % backup_json_dir