from array import array
import os
import shutil
import logging
logger = logging.getLogger('communicator')
from time import sleep
from subprocess import Popen, PIPE
from io import StringIO
import pickle as pickle
import re
import numpy as np
from pathlib import Path
from eon.config import ConfigClass # Typing
[docs]
def bundled_job_name(bundle_dirname, slot):
"""Job id for slot *slot* inside a bundle named after the first job.
Bundles live in a directory named ``{state}_{first_wuid}``. Jobs in
the chunk are consecutive wuids, so slot *n* is ``first_wuid + n``.
"""
parts = str(bundle_dirname).split("_", 1)
if len(parts) != 2:
return str(bundle_dirname)
state, first = parts
try:
return f"{state}_{int(first) + int(slot)}"
except ValueError:
return str(bundle_dirname)
[docs]
def tryint(s):
try:
return int(s)
except:
return s
[docs]
def alphanum_key(s):
""" Turn a string into a list of string and number chunks.
"z23a" -> ["z", 23, "a"]
"""
return [ tryint(c) for c in re.split('([0-9]+)', s) ]
[docs]
def sort_nicely(l):
""" Sort the given list in the way that humans expect.
"""
l.sort(key=alphanum_key)
# MPI may only be constructed once per process; non-MPI communicators are
# cached per ConfigClass instance so dual-config / tests do not cross-wire.
_mpi_communicator = None
_comm_by_config_id = {}
[docs]
def reset_communicators():
"""Clear non-MPI communicator cache (tests / multi-config isolation)."""
global _mpi_communicator
_comm_by_config_id.clear()
# MPI is not reset: mpi4py world may only be initialized once.
[docs]
def get_communicator(config: ConfigClass):
"""Return a communicator for *config*.
*config* is required (no process-wide default). Non-MPI communicators are
memoized by ``id(config)``. MPI uses a single process-wide instance.
"""
global _mpi_communicator
if config is None:
raise TypeError("get_communicator() requires a ConfigClass instance")
if config.comm_type == 'mpi':
if _mpi_communicator is not None:
return _mpi_communicator
comm = MPI(config.path_scratch, config.comm_job_bundle_size, config=config)
_mpi_communicator = comm
return comm
key = id(config)
if key in _comm_by_config_id:
return _comm_by_config_id[key]
if config.comm_type == 'cluster':
comm = Script(config.path_scratch, config.comm_job_bundle_size,
config.comm_script_name_prefix,
config.comm_script_path,
config.comm_script_queued_jobs_cmd,
config.comm_script_cancel_job_cmd,
config.comm_script_submit_job_cmd,
config=config)
elif config.comm_type in ('local_lib', 'inprocess', 'local_inprocess'):
from eon.communicator_inprocess import LocalInProcess
bundle = getattr(config, 'comm_job_bundle_size', 1)
comm = LocalInProcess(config.path_scratch, bundle, config=config)
elif config.comm_type == 'local':
comm = Local(config.path_scratch, config.comm_local_client,
config.comm_local_ncpus, config.comm_job_bundle_size,
config=config)
else:
logger.error(str(config.comm_type)+" is an unknown communicator.")
raise ValueError()
_comm_by_config_id[key] = comm
return comm
[docs]
class CommunicatorError(Exception):
pass
[docs]
class EONClientError(Exception):
"""An eOn client finished without outputting results, it probably crashed."""
pass
RETURN_FILES_MANIFEST = "return_files.dat"
# Legacy harvest (no manifest) keeps these suffixes. A manifest overrides
# the filter so declared .json / .log files arrive and undeclared files do not.
_HARVEST_EXTENSIONS = (".con", ".dat", ".json", ".log")
[docs]
def read_return_files_manifest(jobpath):
"""Return the filenames listed in return_files.dat, or None if absent."""
manifest = Path(jobpath) / RETURN_FILES_MANIFEST
if not manifest.is_file():
return None
names = []
for line in manifest.read_text().splitlines():
name = line.strip()
if not name or name.startswith("#"):
continue
names.append(Path(name).name)
return names
def _read_result_file(path):
return StringIO(Path(path).read_text())
[docs]
def harvest_job_files(jobpath):
"""Harvest one non-bundled job directory into {filename: StringIO}.
When return_files.dat is present, only those names (plus results.dat)
are kept, any extension. A listed file that was never written is
logged and omitted. Without a manifest, keep .con/.dat/.json/.log.
"""
listed = read_return_files_manifest(jobpath)
files = {}
root = Path(jobpath)
if listed is None:
for filename in root.glob("*.*"):
if filename.suffix in _HARVEST_EXTENSIONS:
files[filename.name] = _read_result_file(filename)
return files
wanted = list(listed)
if "results.dat" not in wanted:
wanted.append("results.dat")
seen = set()
for fname in wanted:
if fname in seen:
continue
seen.add(fname)
path = root / fname
if not path.is_file():
logger.warning(
"returnFiles lists %s but it is missing under %s",
fname,
jobpath,
)
continue
files[fname] = _read_result_file(path)
return files
[docs]
class Communicator:
def __init__(self, scratchpath, bundle_size=1, config: ConfigClass = None):
if config is None:
raise TypeError("Communicator requires a ConfigClass instance")
self.config = config
scratch = Path(scratchpath)
if not scratch.is_dir():
scratch.mkdir(parents=True)
self.scratchpath = str(scratch)
self.bundle_size = bundle_size
[docs]
def submit_jobs(self, data, invariants):
'''Throws CommunicatorError if fails.'''
raise NotImplementedError()
[docs]
def get_results(self, results):
'''Returns a list of dictionaries containing the results.'''
raise NotImplementedError()
[docs]
def get_queue_size(self):
'''Returns the number of items waiting to run in the queue.'''
raise NotImplementedError()
[docs]
def queued_search_count(self):
"""Searches waiting, using the bundle size this communicator was built with.
``get_queue_size`` counts bundles. Multiplying by the live config
``comm_job_bundle_size`` is wrong if that value changed after
construction.
"""
return self.get_queue_size() * self.bundle_size
[docs]
def in_progress_search_count(self):
"""Searches running, using this communicator's bundle size."""
nprog = getattr(self, "get_number_in_progress", lambda: 0)()
return nprog * self.bundle_size
[docs]
def cancel_state(self, statenumber):
'''Returns the number of workunits that were canceled.'''
raise NotImplementedError()
[docs]
def get_bundle_size(self, job_path):
if not isinstance(job_path, list):
# List files in job_path.
fnames = os.listdir(job_path)
else:
# job_path is already a list of filenames.
fnames = job_path
# Count results*.dat files.
pattern = re.compile(r'results(?:_\d+)?.dat$')
size = sum(1
for fname in fnames
if pattern.match(fname))
is_bundle = not (size == 1 and "results.dat" in fnames)
return size, is_bundle
[docs]
def unbundle(self, resultpath, keep_result):
'''This method unbundles multiple jobs into multiple single
jobs so the akmc script can process them.
If the job did not return results (probably because it
crashed or was canceled), this method will raise
EONClientError.
'''
# These are the files in the result directory that we keep.
result_root = Path(resultpath)
jobpaths = [str(p) for p in result_root.iterdir() if p.is_dir()]
regex = re.compile(r"(\w+)_(\d+)(\.\w+)")
for jobpath in jobpaths:
dirname = Path(jobpath).name
if not keep_result(dirname):
continue
# Need to figure out how many jobs were bundled together
# and then create the new job directories with the split files.
bundle_size, is_bundle = self.get_bundle_size(jobpath)
if bundle_size == 0:
logger.error("Client running in %s returned no results. "
"Check its output for errors." % jobpath)
# GH: just log the error and continue instead of quitting
#raise EONClientError("Client running in %s returned no results. "
# "Check its output for errors." % jobpath)
continue
results = [
{'name': bundled_job_name(dirname, i)}
for i in range(bundle_size)
]
if not is_bundle:
# Only a single task inside this job, no need to unbundle.
harvested = harvest_job_files(jobpath)
results[0].update(harvested)
results[0]['number'] = 0
else:
# Several tasks bundled inside this job, we need to unbundle.
filenames = list(Path(jobpath).glob("*_[0-9]*.*"))
for filename in filenames:
if filename.suffix not in _HARVEST_EXTENSIONS:
continue
# parse filename
fname = filename.name
match = regex.match(fname)
if not match:
continue
parts = match.groups()
index = int(parts[1])
key = parts[0]+parts[2]
# Load data into stringIO object (should we just return filehandles?)
try:
filedata = StringIO(filename.read_text())
except (IOError, OSError):
logger.exception("Failed to read file %s" % filename)
continue
# add result to results
results[index][key] = filedata
results[index]['number'] = index
kept = []
for slot in results:
if any(str(k).startswith('results') for k in slot):
kept.append(slot)
else:
logger.warning("Failed to find a results.dat file for %s slot %s",
slot.get('name'), slot.get('number'))
yield kept
[docs]
def make_bundles(self, data, invariants):
'''This method is a generator that bundles together multiple jobs into a single job.
Example usage:
for jobpath in self.make_bundles(data, invariants):
do_stuff()'''
from eon import fileio as io
pot_invariants = io.load_potfiles(getattr(self.config, "path_pot", "") or "")
merged = dict(pot_invariants)
if invariants:
merged.update(invariants)
invariants = merged
# Split jobpaths in to lists of size self.bundle_size.
chunks = [ data[i:i+self.bundle_size] for i in range(0, len(data), self.bundle_size) ]
for chunk in chunks:
job_path = Path(self.scratchpath) / chunk[0]["id"]
job_path.mkdir()
for filename, (file_contents, file_permissions) in invariants.items():
dest = job_path / filename
dest.write_text(file_contents.getvalue())
dest.chmod(file_permissions)
# Concatenate all of the displacement and modes together.
n = 0
for job in chunk:
for basename in list(job.keys()):
splitname = basename.rsplit(".", 1)
if len(splitname)!=2:
continue
if self.bundle_size == 1:
filename = basename
else:
filename = "%s_%d.%s" % (splitname[0], n, splitname[1])
(job_path / filename).write_text(job[basename].getvalue())
n += 1
# Returns the jobpath to the new bigger workunit.
yield str(job_path)
[docs]
class MPI(Communicator):
def __init__(self, scratchpath, bundle_size, config: ConfigClass = None):
Communicator.__init__(self, scratchpath, bundle_size, config = config)
from mpi4py.MPI import COMM_WORLD
self.comm = COMM_WORLD
self.client_ranks = [ int(r) for r in os.environ['EON_CLIENT_RANKS'].split(os.pathsep) ]
self.config.comm_job_buffer_size = len(self.client_ranks)
scratch = Path(self.scratchpath)
self.resume_jobs = [p.name for p in scratch.iterdir() if p.is_dir()] if scratch.is_dir() else []
logger.info("Found %i jobs to resume in %s", len(self.resume_jobs), self.scratchpath)
[docs]
def submit_jobs(self, data, invariants):
ready_ranks = self.get_ready_ranks()
for jobpath in self.make_bundles(data, invariants):
rank = ready_ranks.pop()
tmp = np.empty(1, dtype='i')
self.comm.Recv(tmp, source=rank, tag=1)
buf = array('b')
bufval = jobpath+'\0'
buf.frombytes(bufval.encode())
self.comm.Send(buf, rank)
[docs]
def run_resume_jobs(self):
if len(self.resume_jobs) == 0: return
ready_ranks = self.get_ready_ranks()
while True:
if len(self.resume_jobs) == 0: break
if len(ready_ranks) == 0: break
jobdir = self.resume_jobs.pop()
rank = ready_ranks.pop()
jobpath = str(Path(self.scratchpath) / jobdir)
tmp = np.empty(1, dtype='i')
self.comm.Recv(tmp, source=rank, tag=1)
buf = array('b')
bufval = jobpath+'\0'
buf.frombytes(bufval.encode())
self.comm.Send(buf, rank)
[docs]
def get_ready_ranks(self):
ready_ranks = []
for rank in self.client_ranks:
ready = self.comm.Iprobe(rank, tag=1)
if ready:
ready_ranks.append(rank)
return ready_ranks
[docs]
def get_queue_size(self):
self.run_resume_jobs()
nready = len(self.get_ready_ranks())
nclients = len(self.client_ranks)
qs = nclients - nready
return qs
[docs]
def get_results(self, resultspath, keep_result):
'''Moves work from scratchpath to results path.'''
import mpi4py.MPI as MPI
status = MPI.Status()
while self.comm.Iprobe(source=MPI.ANY_SOURCE, tag=0, status=status):
buf = np.array(['\0']*1024, dtype="S1")
self.comm.Recv([buf, MPI.CHARACTER], source=status.source, tag=0)
strterm = np.where(buf == b'\0')
strindex = strterm[0][0]
jobdir = buf[:strindex].tobytes()
jobdir = Path(jobdir.decode()).name
if not (Path(self.scratchpath) / jobdir).is_dir():
logger.error("Client returned %s, which is not in %s; skipping",
jobdir, self.scratchpath)
continue
if self.config.debug_keep_all_results:
shutil.copytree(
Path(self.scratchpath) / jobdir,
Path(self.config.path_root) / self.config.debug_results_path / jobdir,
)
dest_dir = Path(resultspath) / jobdir
shutil.move(str(Path(self.scratchpath) / jobdir), str(dest_dir))
for bundle in self.unbundle(resultspath, keep_result):
for result in bundle:
yield result
[docs]
def get_number_in_progress(self):
return int(os.environ['EON_NUMBER_OF_CLIENTS'])
[docs]
def stop_clients(self):
"""Send STOPCAR to every ready client rank (ClientEON exits on that path)."""
n = 0
for rank in self.client_ranks:
if not self.comm.Iprobe(rank, tag=1):
continue
tmp = np.empty(1, dtype='i')
self.comm.Recv(tmp, source=rank, tag=1)
buf = array('b')
buf.frombytes(b'STOPCAR\0')
self.comm.Send(buf, rank)
n += 1
return n
[docs]
def cancel_state(self, state):
# Work already handed to a client cannot be recalled over this
# protocol; its result carries the old state number and the explorer's
# keep_result filter discards it. Stopping the clients here would
# leave the run without workers for the next state.
return 0
[docs]
def client_environment(base=None):
"""Environment for an eonclient the server starts outside an MPI launcher.
An eonclient built with MPI (rgpot calculator groups) loads libmpi, and
with it UCX, in every run. UCX patches the memory functions as the
library loads, which on conda's OpenMPI costs about 0.4 s per process
against 0.02 s for the job itself. A client the Local communicator
starts is one process with no MPI peers, so the memory events are of no
use to it. An explicit UCX_MEM_EVENTS from the user wins.
"""
env = dict(os.environ if base is None else base)
env.setdefault("UCX_MEM_EVENTS", "no")
return env
[docs]
class Local(Communicator):
def __init__(self, scratchpath, client, ncpus, bundle_size, config: ConfigClass = None):
Communicator.__init__(self, scratchpath, bundle_size, config = config)
# number of cpus to use
self.ncpus = ncpus
# path to the client
client_path = Path(client)
if os.sep in client or "/" in client:
self.client = str(client_path.resolve())
if not Path(self.client).is_file():
logger.error("Can't find client: %s", client)
raise CommunicatorError("Can't find client binary: %s"%client)
else:
import shutil
if client_path.is_file():
self.client = str(client_path.resolve())
elif shutil.which(client) is not None:
self.client = shutil.which(client)
else:
logger.error("Can't find client: %s", client)
raise CommunicatorError("Can't find client binary: %s"%client)
self.joblist = []
import atexit
# don't let clients hang around if the script dies
atexit.register(self.cleanup)
[docs]
def cleanup(self):
'''Kills the running eonclients and everything they started.'''
import signal
for job in self.joblist:
p = job[0]
if p.poll() is not None:
continue
try:
if os.name == 'posix':
os.killpg(p.pid, signal.SIGKILL)
else:
p.kill()
except OSError:
pass
[docs]
def get_results(self, resultspath, keep_result):
'''Moves work from scratchpath to results path.'''
scratch = Path(self.scratchpath)
jobdirs = [p.name for p in scratch.iterdir() if p.is_dir()]
for jobdir in jobdirs:
if self.config.debug_keep_all_results:
shutil.copytree(
scratch / jobdir,
Path(self.config.path_root) / self.config.debug_results_path / jobdir,
)
dest_dir = Path(resultspath) / jobdir
shutil.move(str(scratch / jobdir), str(dest_dir))
for bundle in self.unbundle(resultspath, keep_result):
for result in bundle:
yield result
# Clean out scratch directory
for path_name in scratch.iterdir():
if path_name.is_dir():
shutil.rmtree(path_name)
[docs]
def check_job(self, job):
p, jobpath = job[0], job[1]
# Close stdout/stderr file handles so the files can be read for diagnostics.
for handle in job[2:]:
if handle and not handle.closed:
handle.close()
if p.returncode == 0:
logger.info('Job finished: %s' % jobpath)
return True
else:
rc = p.returncode
# stderr is a file (stderr.dat), not a pipe: never call
# communicate() for it. Read both logs from disk for the failure
# path so a dead/hung client cannot leave the parent blocked on a
# full undrained pipe either.
errmsg = "job failed: %s (return code %s)" % (jobpath, rc)
saw_output = False
for name in ("stderr.dat", "stdout.dat"):
path = Path(jobpath) / name
try:
text = path.read_text().strip()
except OSError:
text = ""
if text:
saw_output = True
lines = text.splitlines()
tail = "\n ".join(lines[-10:])
errmsg += "\n %s (last lines): %s" % (name, tail)
if not saw_output:
errmsg += "\n no output captured; client may have failed to start"
if os.name == 'nt' and (rc < 0 or rc > 0x80000000):
errmsg += " (Windows error code 0x%X may indicate a missing DLL)" % (rc & 0xFFFFFFFF)
logger.warning(errmsg)
[docs]
def submit_jobs(self, data, invariants):
'''Run up to ncpu number of clients to process the work in jobpaths.
The job directories are moved to the scratch path before the calculation
is run. This method doesn't return anything.'''
for jobpath in self.make_bundles(data, invariants):
# move the job directory to the scratch directory
# update jobpath to be in the scratch directory
# Redirect both streams to files. stderr=PIPE with an undrained
# pipe deadlocks once the client writes past one OS pipe buffer
# (~64 KiB on Linux): the child blocks in anon_pipe_write while
# the parent only polls p.poll() and never reads the pipe until
# after exit. FPE floods and long client diagnostics both hit this.
fstdout = open(Path(jobpath) / "stdout.dat", 'w')
fstderr = open(Path(jobpath) / "stderr.dat", 'w')
# A session of its own lets cleanup() reach every process the
# client starts (ExtPot wrappers, srun, mpirun), not just the client.
# That session also forces subprocess down fork or vfork. fork runs
# atfork handlers; eon._forksafe parks the OpenBLAS pool first so
# server_lock is free when this process next factorizes.
p = Popen(self.client, cwd=jobpath, stdout=fstdout, stderr=fstderr,
env=client_environment(),
start_new_session=(os.name == 'posix'))
self.joblist.append((p, jobpath, fstdout, fstderr))
while len(self.joblist) == self.ncpus:
for i in range(len(self.joblist)):
p = self.joblist[i][0]
retval = p.poll()
if retval is None:
continue
else:
self.check_job(self.joblist[i])
self.joblist.pop(i)
break
sleep(0.1)
# wait for everything to finish
for job in self.joblist:
p = job[0]
p.wait()
self.check_job(job)
[docs]
def cancel_state(self, state):
return 0
[docs]
def get_queue_size(self):
return 0
[docs]
def get_number_in_progress(self):
return 0
[docs]
class Script(Communicator):
def __init__(self, scratch_path, bundle_size, name_prefix, scripts_path,
queued_jobs_cmd, cancel_job_cmd, submit_job_cmd, config: ConfigClass = None):
Communicator.__init__(self, scratch_path, bundle_size, config=config)
scripts = Path(scripts_path)
self.queued_jobs_cmd = str(scripts / queued_jobs_cmd)
self.cancel_job_cmd = str(scripts / cancel_job_cmd)
self.submit_job_cmd = str(scripts / submit_job_cmd)
self.job_id_path = str(Path(scratch_path) / "script_job_ids")
self.name_prefix = name_prefix
# read in job ids
try:
f = open(self.job_id_path, "rb")
self.jobids = pickle.load(f)
f.close()
except IOError:
self.jobids = {}
pass
[docs]
def save_jobids(self):
f = open(self.job_id_path, "wb")
pickle.dump(self.jobids, f)
f.close()
[docs]
def get_results(self, resultspath, keep_result):
'''Moves work from scratchpath to results path.'''
# queued_jobs.sh jobid1 jobid2 jobid 3
# the inverse of the jobids returned is
# job dirs needs to map
queued_jobs = self.get_queued_jobs()
finished_jobids = set(self.jobids.keys()) - set(queued_jobs)
finished_eonids = []
for jobid in finished_jobids:
finished_eonids.append(int(self.jobids.pop(jobid)))
scratch = Path(self.scratchpath)
jobdirs = [
p.name
for p in scratch.iterdir()
if p.is_dir() and int(p.name.rsplit("_", 1)[-1]) in finished_eonids
]
#try to return jobs in order
sort_nicely(jobdirs)
for jobdir in jobdirs:
if self.config.debug_keep_all_results:
shutil.copytree(
scratch / jobdir,
Path(self.config.path_root) / self.config.debug_results_path / jobdir,
)
dest_dir = Path(resultspath) / jobdir
shutil.move(str(scratch / jobdir), str(dest_dir))
for bundle in self.unbundle(resultspath, keep_result):
for result in bundle:
yield result
[docs]
def check_command(self, status, output, cmdname):
if status != 0:
logger.error(output)
raise CommunicatorError("'%s' returned a non-zero exit status"%cmdname)
[docs]
def submit_jobs(self, data, invariants):
for jobpath in self.make_bundles(data, invariants):
# submit_job.sh jobname jobpath
# should return a jobid
# need to associate this jobid with our jobid
jobpath = str(Path(jobpath).resolve())
jobname = "%s_%s" % (self.name_prefix, Path(jobpath).name)
eon_jobid = jobname.rsplit('_',1)[-1]
cmd = "%s %s %s" % (self.submit_job_cmd, jobname, jobpath)
p = Popen([self.submit_job_cmd,jobname,jobpath], stdout=PIPE, stderr=PIPE)
output, error = p.communicate()
output = output.decode()
error = error.decode()
status = p.returncode
self.check_command(status, output, cmd)
jobid = int(output.strip())
self.jobids[jobid] = eon_jobid
# Persist after each submit so a crash does not lose job ids.
self.save_jobids()
[docs]
def cancel_state(self, state):
'''Cancel the jobs still in the queue and drop their directories.
Jobs that already left the queue keep their ids and directories, so
the next get_results() harvests them; the explorer's keep_result
filter then decides whether they still count. A cancel that fails
(the job finished in between) is logged, not fatal. Returns the
number of jobs cancelled.
'''
queued = self.get_queued_jobs()
cancelled = 0
scratch = Path(self.scratchpath)
for job_id in queued:
p = Popen([self.cancel_job_cmd, str(job_id)], stdout=PIPE, stderr=PIPE)
output, error = p.communicate()
if p.returncode != 0:
logger.warning("cancel_job for %s failed (exit %i): %s", job_id,
p.returncode, (output + error).decode().strip())
continue
eon_jobid = self.jobids.pop(job_id)
for path in scratch.iterdir():
if path.is_dir() and path.name.rsplit("_", 1)[-1] == str(eon_jobid):
shutil.rmtree(path)
cancelled += 1
self.save_jobids()
return cancelled
[docs]
def get_queued_jobs(self):
p = Popen([self.queued_jobs_cmd,''], stdout=PIPE, stderr=PIPE)
output, error = p.communicate()
output = output.decode()
error = error.decode()
status = p.returncode
self.check_command(status, output, self.queued_jobs_cmd)
queued_job_ids = []
for line in output.split("\n"):
try:
queued_job_ids.append(int(line))
except ValueError:
continue
return list(set(self.jobids).intersection(queued_job_ids))
[docs]
def get_number_in_progress(self):
return 0
[docs]
def get_queue_size(self):
return len(self.get_queued_jobs())
eOn