285 lines
8.2 KiB
Python
285 lines
8.2 KiB
Python
import sys
|
|
import time
|
|
from collections import namedtuple
|
|
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
|
|
from functools import partial
|
|
from multiprocessing import Pool, cpu_count
|
|
from pathlib import Path
|
|
|
|
import rich
|
|
from rich.progress import (
|
|
BarColumn,
|
|
Progress,
|
|
TaskProgressColumn,
|
|
TextColumn,
|
|
TimeElapsedColumn,
|
|
TimeRemainingColumn,
|
|
)
|
|
|
|
from tatsu.util.unicode_characters import (
|
|
U_CHECK_MARK,
|
|
U_CROSSED_SWORDS,
|
|
)
|
|
|
|
from ..util import identity, memory_use, program_name, try_read
|
|
|
|
__all__ = ['parallel_proc', 'processing_loop']
|
|
|
|
|
|
ERROR_LOG_FILENAME = 'ERROR.log'
|
|
EOLCH = '\r' if sys.stderr.isatty() else '\n'
|
|
sys.setrecursionlimit(2**16)
|
|
|
|
__Task = namedtuple('__Task', 'payload args kwargs')
|
|
|
|
|
|
class ParprocResult:
|
|
def __init__(self, payload):
|
|
self.payload = payload
|
|
self.outcome = None
|
|
self.exception = None
|
|
self.linecount = 0
|
|
self.time = 0
|
|
self.memory = 0
|
|
|
|
@property
|
|
def success(self):
|
|
return self.exception is None
|
|
|
|
def __str__(self):
|
|
return str(self.__dict__)
|
|
|
|
|
|
def process_payload(process, task, pickable=identity, reraise=False):
|
|
start_time = time.process_time()
|
|
result = ParprocResult(task.payload)
|
|
try:
|
|
outcome = process(task.payload, *task.args, **task.kwargs)
|
|
result.memory = memory_use()
|
|
if hasattr(outcome, 'linecount'):
|
|
result.linecount = outcome.linecount
|
|
else:
|
|
result.linecount = len(try_read(task.payload).splitlines())
|
|
result.outcome = pickable(outcome)
|
|
except KeyboardInterrupt:
|
|
return None
|
|
except Exception as e:
|
|
result.exception = e
|
|
finally:
|
|
result.time = time.process_time() - start_time
|
|
|
|
return result
|
|
|
|
|
|
def _executor_pmap(executor, process, tasks):
|
|
nworkers = max(1, cpu_count())
|
|
n = nworkers * 8
|
|
chunks = [tasks[i:i + n] for i in range(0, len(tasks), n)]
|
|
for chunk in chunks:
|
|
with executor(max_workers=nworkers) as ex:
|
|
futures = [ex.submit(process, task) for task in chunk]
|
|
for future in as_completed(futures):
|
|
yield future.result()
|
|
|
|
|
|
def _thread_pmap(process, tasks):
|
|
yield from _executor_pmap(ThreadPoolExecutor, process, tasks)
|
|
|
|
|
|
def _process_pmap(process, tasks):
|
|
yield from _executor_pmap(ProcessPoolExecutor, process, tasks)
|
|
|
|
|
|
def _imap_pmap(process, tasks):
|
|
nworkers = max(1, cpu_count())
|
|
|
|
n = nworkers * 4
|
|
chunks = [tasks[i:i + n] for i in range(0, len(tasks), n)]
|
|
|
|
count = sum(len(c) for c in chunks)
|
|
if len(tasks) != count:
|
|
raise RuntimeError('number of chunked tasks different %d != %d' % (len(tasks), count))
|
|
for chunk in chunks:
|
|
with Pool(processes=nworkers) as pool:
|
|
yield from pool.imap_unordered(process, chunk)
|
|
|
|
|
|
_pmap = _imap_pmap
|
|
|
|
|
|
def parallel_proc(payloads, process, *args, **kwargs):
|
|
pickable = kwargs.pop('pickable', identity)
|
|
parallel = kwargs.pop('parallel', True)
|
|
reraise = kwargs.pop('reraise', False)
|
|
|
|
process = partial(process_payload, process, pickable=pickable, reraise=reraise)
|
|
tasks = [__Task(payload, args, kwargs) for payload in payloads]
|
|
|
|
try:
|
|
if len(tasks) == 1:
|
|
yield process(tasks[0])
|
|
else:
|
|
pmap = _pmap if parallel else map
|
|
yield from pmap(process, tasks)
|
|
except KeyboardInterrupt:
|
|
return
|
|
|
|
|
|
def _build_progressbar(total):
|
|
progress = Progress(
|
|
TextColumn(f"[progress.description]{program_name()}"),
|
|
BarColumn(),
|
|
# *Progress.get_default_columns(),
|
|
TaskProgressColumn(),
|
|
TimeElapsedColumn(),
|
|
TimeRemainingColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
refresh_per_second=1,
|
|
speed_estimate_period=30.0,
|
|
)
|
|
task = progress.add_task('', total=total)
|
|
return progress, task
|
|
|
|
|
|
def processing_loop(filenames, process, *args, reraise=False, **kwargs): # pylint: disable=too-many-locals
|
|
try:
|
|
total = len(filenames)
|
|
total_time = 0
|
|
run_time = 0
|
|
start_time = time.time()
|
|
results = parallel_proc(filenames, process, *args, **kwargs)
|
|
results = results or []
|
|
count = 0
|
|
success_count = 0
|
|
success_linecount = 0
|
|
progress, progress_task = _build_progressbar(total)
|
|
|
|
if total == 1:
|
|
log = sys.stderr
|
|
else:
|
|
log = Path(ERROR_LOG_FILENAME).open('w')
|
|
|
|
with progress:
|
|
for result in results:
|
|
if result is None:
|
|
continue
|
|
count += 1
|
|
|
|
total_time = time.time() - start_time
|
|
filename = Path(result.payload).name
|
|
if result.exception:
|
|
icon = f'[red]{U_CROSSED_SWORDS}'
|
|
else:
|
|
icon = f'[green]{U_CHECK_MARK}'
|
|
|
|
progress.update(
|
|
progress_task,
|
|
advance=1,
|
|
description=f'{icon} {filename}',
|
|
)
|
|
|
|
if result.exception:
|
|
print(file=log)
|
|
try:
|
|
print('ERROR:', result.payload, file=log)
|
|
print(result.exception, file=log)
|
|
except Exception:
|
|
# in case of errors while serializing the exception
|
|
print('EXCEPTION', type(result.exception).__name__, file=log)
|
|
if reraise:
|
|
raise result.exception
|
|
elif result.outcome is not None:
|
|
success_count += 1
|
|
success_linecount += result.linecount
|
|
run_time += result.time
|
|
yield result
|
|
|
|
log.flush()
|
|
|
|
progress.update(progress_task, advance=0, description='')
|
|
progress.stop()
|
|
file_process_summary(
|
|
filenames, total_time, run_time, success_count, success_linecount, log,
|
|
)
|
|
except KeyboardInterrupt:
|
|
return
|
|
|
|
|
|
def file_process_progress(latest_result, count, total, total_time):
|
|
filename = latest_result.payload
|
|
|
|
percent = count / total
|
|
mb_memory = (latest_result.memory + memory_use()) // (1024 * 1024)
|
|
|
|
eta = (total - count) * 0.8 * total_time / (0.2 * count)
|
|
bar = '[%-24s]' % ('#' * round(24 * percent))
|
|
|
|
print(
|
|
'%3d/%-3d' % (count, total),
|
|
bar,
|
|
'%3d%%' % (100 * percent),
|
|
# format_hours(total_time),
|
|
f'{format_hours(eta)}ETA',
|
|
format_minutes(latest_result),
|
|
'%3dMiB' % mb_memory if mb_memory else '',
|
|
(Path(filename).name + ' ' * 80)[:40],
|
|
file=sys.stderr,
|
|
end=EOLCH)
|
|
|
|
|
|
def format_minutes(result):
|
|
return f'{result.time / 60:3.0f}:{result.time % 60:04.1f}'
|
|
|
|
|
|
def format_hours(time):
|
|
return f'{time // 3600:2.0f}:{(time // 60) % 60:02.0f}:{time % 60:02.0f}'
|
|
|
|
|
|
def file_process_summary(
|
|
filenames, total_time, run_time, success_count, success_linecount, log,
|
|
):
|
|
filecount = 0
|
|
linecount = 0
|
|
for fname in filenames:
|
|
filecount += 1
|
|
|
|
nlines = len(try_read(fname).splitlines())
|
|
linecount += nlines
|
|
|
|
failures = filecount - success_count
|
|
|
|
summary_text = '''\
|
|
-----------------------------------------------------------------------
|
|
{:12,d} files input
|
|
{:12,d} source lines input
|
|
{:12,d} total lines processed
|
|
{:12,d} successes
|
|
{:12,d} failures
|
|
{:12.1f}% success rate
|
|
{:>13s} time
|
|
{:>13s} run time
|
|
'''
|
|
summary_text = '\n'.join(s.strip() for s in summary_text.splitlines())
|
|
|
|
summary = summary_text.format(
|
|
filecount,
|
|
linecount,
|
|
success_linecount,
|
|
success_count,
|
|
failures,
|
|
100 * success_count / filecount if filecount != 0 else 0,
|
|
format_hours(total_time),
|
|
format_hours(run_time),
|
|
)
|
|
print(EOLCH + 80 * ' ', file=log)
|
|
print(summary, file=log)
|
|
print(EOLCH + 80 * ' ', file=sys.stderr)
|
|
|
|
if failures:
|
|
rich.print(f'[red bold]FAILURES logged to [green]{log.name}!')
|
|
print(file=sys.stderr)
|
|
if log != sys.stderr:
|
|
print(summary, file=sys.stderr)
|
|
if failures:
|
|
sys.exit(1)
|