348 lines
13 KiB
Python
348 lines
13 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
Report the disk usage of the raw/ and res/ directories of every pgroup in
|
|
/sf/cristallina/data/ that the current user has access to, plus the size and
|
|
last modification date of the pgroup's work directory on Ra.
|
|
|
|
Access is decided by membership of the unix group with the pgroup's name
|
|
(e.g. p21528); use --all to instead try every pgroup that is readable.
|
|
|
|
Also reports the archiving state of each pgroup, using the same logic as
|
|
/sf/daq/bin/sf-daq_writing_state (new / open / closed / archived).
|
|
|
|
/das/work is not mounted here, so the work directories are checked by running
|
|
find over ssh on Ra. This uses your existing Kerberos ticket (run kinit if it
|
|
has expired) - no passwords or keys are stored in this script.
|
|
|
|
Sizes are in TB (1 TB = 1024**4 bytes) and pgroups are printed without the "p".
|
|
|
|
While running, a status line listing the directories still being scanned is
|
|
printed every few minutes (--status-interval). Ctrl-C stops the scans (also
|
|
those on Ra) and still prints/writes the results collected so far, with the
|
|
unfinished entries marked "not done".
|
|
|
|
usage:
|
|
python pgroup_disk_usage.py # all accessible pgroups
|
|
python pgroup_disk_usage.py 22547 p23610 # only these pgroups ("p" optional)
|
|
python pgroup_disk_usage.py -j 16 --csv sizes.csv
|
|
python pgroup_disk_usage.py -t 600 # give up on a directory after 10 min
|
|
python pgroup_disk_usage.py --no-work # skip the Ra work directories
|
|
"""
|
|
|
|
import argparse
|
|
import csv
|
|
import grp
|
|
import os
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from datetime import datetime
|
|
|
|
DATA_DIR = "/sf/cristallina/data"
|
|
SUBDIRS = ("raw", "res")
|
|
RA_HOST = "ra.psi.ch"
|
|
NOT_DONE = "not done"
|
|
|
|
# run on Ra: prints "START <pgroup>" when a work directory scan begins, then
|
|
# "<pgroup> <bytes> <newest mtime>" or "<pgroup> <error>" when it ends.
|
|
# The watchdog writes a line every 20 s; if that fails the ssh connection is
|
|
# gone (e.g. Ctrl-C here), so it kills all the scans instead of leaving them
|
|
# running on Ra. timeout --foreground keeps find in the same process group.
|
|
RA_SCRIPT = r"""
|
|
check() {
|
|
p=$1; d=/das/work/${p:0:3}/$p
|
|
if [ ! -e "$d" ]; then echo "$p missing"; return; fi
|
|
if [ ! -r "$d/" ] || [ ! -x "$d/" ]; then echo "$p no_access"; return; fi
|
|
echo "START $p"
|
|
out=$(timeout --foreground $TIMEOUT find "$d/" -printf '%s %T@\n' 2>/dev/null \
|
|
| awk '{s += $1; if ($2 > m) m = $2} END {printf "%d %d", s, m}'
|
|
exit ${PIPESTATUS[0]})
|
|
if [ $? -eq 124 ]; then echo "$p timeout"; return; fi
|
|
echo "$p ${out:-error}"
|
|
}
|
|
export -f check
|
|
export TIMEOUT
|
|
(
|
|
trap '' PIPE
|
|
while sleep 20 >/dev/null 2>&1; do echo ALIVE || break; done
|
|
kill -TERM 0
|
|
) &
|
|
watchdog=$!
|
|
echo "$PGROUPS" | tr ' ' '\n' | xargs -P "$JOBS" -I{} bash -c 'check {}'
|
|
kill $watchdog 2>/dev/null
|
|
"""
|
|
|
|
print_lock = threading.Lock()
|
|
running = {} # (pgroup, subdir) -> start time of scans in progress
|
|
running_lock = threading.Lock()
|
|
|
|
|
|
def log(msg):
|
|
with print_lock:
|
|
print(msg, flush=True)
|
|
|
|
|
|
def mark_started(pgroup, subdir):
|
|
with running_lock:
|
|
running[(pgroup, subdir)] = time.time()
|
|
|
|
|
|
def mark_finished(pgroup, subdir):
|
|
"""Remove from the running list and return how long it took (s)."""
|
|
with running_lock:
|
|
started = running.pop((pgroup, subdir), None)
|
|
return time.time() - started if started else 0
|
|
|
|
|
|
def minutes(seconds):
|
|
return f"{seconds / 60:.1f} min"
|
|
|
|
|
|
def my_groups():
|
|
"""Names of all unix groups the current user belongs to."""
|
|
names = set()
|
|
for gid in os.getgroups():
|
|
try:
|
|
names.add(grp.getgrgid(gid).gr_name)
|
|
except KeyError:
|
|
pass
|
|
return names
|
|
|
|
|
|
def find_pgroups(requested=None, check_all=False):
|
|
pgroups = sorted(
|
|
d for d in os.listdir(DATA_DIR)
|
|
if d.startswith("p") and os.path.isdir(os.path.join(DATA_DIR, d))
|
|
)
|
|
if requested:
|
|
requested = {"p" + r.lstrip("p") for r in requested}
|
|
return [p for p in pgroups if p in requested]
|
|
if check_all:
|
|
return pgroups
|
|
groups = my_groups()
|
|
return [p for p in pgroups if p in groups]
|
|
|
|
|
|
def archive_state(pgroup):
|
|
"""Same logic as /sf/daq/bin/sf-daq_writing_state."""
|
|
run_info = os.path.join(DATA_DIR, pgroup, "raw", "run_info")
|
|
if not os.path.exists(run_info):
|
|
return "new"
|
|
if os.path.exists(os.path.join(run_info, "ARCHIVED")):
|
|
return "archived"
|
|
if os.path.exists(os.path.join(run_info, "CLOSED")):
|
|
return "closed"
|
|
return "open"
|
|
|
|
|
|
def dir_size(path, timeout=None):
|
|
"""Size of path in bytes using du, or an error string."""
|
|
if not os.path.exists(path):
|
|
return "missing"
|
|
if not os.access(path, os.R_OK | os.X_OK):
|
|
return "no access"
|
|
# trailing slash so du follows the raw/res symlinks to gpfs
|
|
try:
|
|
out = subprocess.run(
|
|
["du", "-sb", path + "/"],
|
|
capture_output=True, text=True, timeout=timeout,
|
|
)
|
|
except subprocess.TimeoutExpired:
|
|
return "timeout"
|
|
if not out.stdout:
|
|
return "error"
|
|
# du still prints a total if some subdirectories are unreadable
|
|
return int(out.stdout.split()[0])
|
|
|
|
|
|
def local_size(pgroup, subdir, timeout):
|
|
mark_started(pgroup, subdir)
|
|
return dir_size(os.path.join(DATA_DIR, pgroup, subdir), timeout)
|
|
|
|
|
|
def work_usage(pgroups, jobs, timeout, host, work, progress, procs):
|
|
"""
|
|
Size and newest mtime of each /das/work pgroup directory, checked over ssh
|
|
on Ra. Fills work[pgroup] = (size, modified) as results arrive; on failure
|
|
size is an error string.
|
|
"""
|
|
env = f"PGROUPS='{' '.join(pgroups)}' JOBS={jobs} TIMEOUT={int(timeout) if timeout else 0}"
|
|
# BatchMode: fail instead of asking for a password if kerberos is not set up
|
|
cmd = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=15", host, f"{env} bash -s"]
|
|
try:
|
|
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE, text=True)
|
|
except OSError as e:
|
|
log(f" could not run ssh: {e}")
|
|
work.update({p: ("ssh error", "") for p in pgroups})
|
|
return
|
|
procs.append(proc)
|
|
proc.stdin.write(RA_SCRIPT)
|
|
proc.stdin.close()
|
|
for line in proc.stdout:
|
|
parts = line.split()
|
|
if len(parts) == 2 and parts[0] == "START" and parts[1] in pgroups:
|
|
mark_started(parts[1], "work")
|
|
continue
|
|
if not parts or parts[0] not in pgroups:
|
|
continue
|
|
p = parts[0]
|
|
if len(parts) == 3:
|
|
size, mtime = int(parts[1]), int(parts[2])
|
|
modified = datetime.fromtimestamp(mtime).strftime("%Y-%m-%d") if mtime else "empty"
|
|
work[p] = (size, modified)
|
|
else:
|
|
work[p] = (parts[1].replace("_", " "), "")
|
|
progress(p, "work", work[p][0])
|
|
proc.wait()
|
|
err = proc.stderr.read().strip()
|
|
# returncode < 0 means ssh was stopped by us (Ctrl-C), not a failure
|
|
if proc.returncode > 0 and not work:
|
|
log(f" ssh to {host} failed (is your kerberos ticket valid? try kinit):\n {err}")
|
|
work.update({p: ("ssh error", "") for p in pgroups})
|
|
|
|
|
|
def print_running(prefix):
|
|
"""Print the scans still in progress, longest running first."""
|
|
now = time.time()
|
|
with running_lock:
|
|
items = sorted(running.items(), key=lambda item: item[1])
|
|
if not items:
|
|
return
|
|
log(f"{prefix} {len(items)} directories still being scanned:")
|
|
for (pgroup, subdir), started in items:
|
|
log(f" {pgroup[1:]}/{subdir:<4} running for {minutes(now - started)}")
|
|
|
|
|
|
def tb(size):
|
|
"""Bytes to TB (1024**4), or pass through an error string."""
|
|
if not isinstance(size, int):
|
|
return size
|
|
return f"{size / 1024**4:.3f}"
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
parser.add_argument("pgroups", nargs="*", help="only check these pgroups")
|
|
parser.add_argument("--all", action="store_true", help="check every pgroup, not only those in your unix groups")
|
|
parser.add_argument("-j", "--jobs", type=int, default=8, help="number of parallel du/find processes (default 8)")
|
|
parser.add_argument("-t", "--timeout", type=float, default=None, help="timeout in seconds per directory")
|
|
parser.add_argument("-s", "--status-interval", type=float, default=5,
|
|
help="minutes between 'still being scanned' status lines (default 5, 0 = off)")
|
|
parser.add_argument("--no-work", action="store_true", help="do not check the work directories on Ra")
|
|
parser.add_argument("--ra-host", default=RA_HOST, help=f"host to check work directories on (default {RA_HOST})")
|
|
parser.add_argument("--csv", help="also write results (in TB) to this csv file")
|
|
args = parser.parse_args()
|
|
|
|
pgroups = find_pgroups(args.pgroups, args.all)
|
|
if not pgroups:
|
|
print("no pgroups found")
|
|
return
|
|
print(f"checking {len(pgroups)} pgroups in {DATA_DIR} with {args.jobs} parallel jobs:")
|
|
print(" " + " ".join(p[1:] for p in pgroups))
|
|
if not args.no_work:
|
|
print(f"work directories are checked on {args.ra_host} at the same time")
|
|
print("press Ctrl-C to stop early and still get the results so far")
|
|
print(flush=True)
|
|
|
|
jobs = [(p, s) for p in pgroups for s in SUBDIRS]
|
|
n_total = len(jobs) + (0 if args.no_work else len(pgroups))
|
|
done = [0]
|
|
start = time.time()
|
|
|
|
def progress(pgroup, subdir, size):
|
|
took = mark_finished(pgroup, subdir)
|
|
with print_lock:
|
|
done[0] += 1
|
|
print(f" [{done[0]:>3}/{n_total}] {pgroup[1:]}/{subdir:<4} {tb(size):>10} TB"
|
|
f" took {minutes(took):>9} ({minutes(time.time() - start)} elapsed)", flush=True)
|
|
|
|
# periodically show what is still running, so slow directories stand out
|
|
stop_status = threading.Event()
|
|
|
|
def status_loop():
|
|
while not stop_status.wait(args.status_interval * 60):
|
|
print_running(f" -- {minutes(time.time() - start)} elapsed,")
|
|
|
|
if args.status_interval > 0:
|
|
threading.Thread(target=status_loop, daemon=True).start()
|
|
|
|
results = {}
|
|
work = {}
|
|
ssh_procs = []
|
|
interrupted = False
|
|
pool = ThreadPoolExecutor(max_workers=args.jobs + 1)
|
|
try:
|
|
if not args.no_work:
|
|
work_future = pool.submit(work_usage, pgroups, args.jobs, args.timeout,
|
|
args.ra_host, work, progress, ssh_procs)
|
|
futures = {pool.submit(local_size, p, s, args.timeout): (p, s) for p, s in jobs}
|
|
for future in as_completed(futures):
|
|
job = futures[future]
|
|
results[job] = future.result()
|
|
progress(job[0], job[1], results[job])
|
|
if not args.no_work:
|
|
work_future.result()
|
|
pool.shutdown()
|
|
except KeyboardInterrupt:
|
|
interrupted = True
|
|
print(flush=True)
|
|
print_running("interrupted -")
|
|
# the running du/ssh also get the Ctrl-C; make sure ssh is gone so the
|
|
# watchdog on Ra stops the remote scans
|
|
for proc in ssh_procs:
|
|
proc.terminate()
|
|
pool.shutdown(wait=False, cancel_futures=True)
|
|
finally:
|
|
stop_status.set()
|
|
|
|
word = "stopped after" if interrupted else "done in"
|
|
print(f"\n{word} {minutes(time.time() - start)}\n")
|
|
|
|
header = f"{'pgroup':<10}{'raw (TB)':>12}{'res (TB)':>12} {'state':<10}"
|
|
if not args.no_work:
|
|
header += f"{'work (TB)':>12} {'work modified':<14}"
|
|
print(header)
|
|
print("-" * len(header))
|
|
totals = {"raw": 0, "res": 0, "work": 0}
|
|
rows = []
|
|
for p in pgroups:
|
|
raw = results.get((p, "raw"), NOT_DONE)
|
|
res = results.get((p, "res"), NOT_DONE)
|
|
state = archive_state(p)
|
|
line = f"{p[1:]:<10}{tb(raw):>12}{tb(res):>12} {state:<10}"
|
|
row = [p[1:], tb(raw), tb(res), state]
|
|
sizes = {"raw": raw, "res": res}
|
|
if not args.no_work:
|
|
work_size, work_modified = work.get(p, (NOT_DONE, ""))
|
|
line += f"{tb(work_size):>12} {work_modified:<14}"
|
|
row += [tb(work_size), work_modified]
|
|
sizes["work"] = work_size
|
|
for s, x in sizes.items():
|
|
if isinstance(x, int):
|
|
totals[s] += x
|
|
rows.append(row)
|
|
print(line)
|
|
print("-" * len(header))
|
|
line = f"{'TOTAL':<10}{tb(totals['raw']):>12}{tb(totals['res']):>12} {'':<10}"
|
|
if not args.no_work:
|
|
line += f"{tb(totals['work']):>12}"
|
|
print(line)
|
|
if interrupted:
|
|
print("(interrupted - totals only include finished directories)")
|
|
|
|
if args.csv:
|
|
with open(args.csv, "w", newline="") as f:
|
|
writer = csv.writer(f)
|
|
columns = ["pgroup", "raw_TB", "res_TB", "state"]
|
|
if not args.no_work:
|
|
columns += ["work_TB", "work_modified"]
|
|
writer.writerow(columns)
|
|
writer.writerows(rows)
|
|
print(f"\nwritten to {args.csv}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|