Skip to content

Commit

Permalink
Merge pull request #1133 from NNPDF/keyboard_interrupt
Browse files Browse the repository at this point in the history 
Keyboard interrupt handling
  • Loading branch information
@Zaharid
Zaharid authored Apr 15, 2021
2 parents b09549e + 523068e commit 332f141
Show file tree
Hide file tree
Showing 3 changed files with 167 additions and 89 deletions.
35 changes: 18 additions & 17 deletions validphys2/src/validphys/loader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
import pathlib
import functools
import logging
import numbers
import re
import tempfile
import shutil
Expand All @@ -30,6 +29,7 @@
PositivitySetSpec, DataSetSpec, PDF, Cuts, DataGroupSpec,
peek_commondata_metadata, CutsPolicy,
InternalCutsWrapper)
from validphys.utils import tempfile_cleaner
from validphys import lhaindex
import NNPDF as nnpath

Expand Down Expand Up @@ -616,8 +616,6 @@ def download_and_extract(url, local_path):
os.unlink(archive_dest.name)




def _key_or_loader_error(f):
@functools.wraps(f)
def f_(*args, **kwargs):
Expand Down Expand Up @@ -744,21 +742,24 @@ def download_fit(self, fitname):
if not fitname in self.remote_fits:
raise FitNotFound("Could not find fit '{}' in remote index {}".format(fitname, self.fit_index))

tempdir = pathlib.Path(tempfile.mkdtemp(prefix='fit_download_deleteme_',
dir=self.resultspath))
download_and_extract(self.remote_fits[fitname], tempdir)
#Handle old-style fits compressed with 'results' as root.
old_style_res = tempdir/'results'
if old_style_res.is_dir():
move_target = old_style_res / fitname
else:
move_target = tempdir/fitname
if not move_target.is_dir():
raise RemoteLoaderError(f"Unknown format for fit in {tempdir}. Expecting a folder {move_target}")
with tempfile_cleaner(
root=self.resultspath,
exit_func=shutil.rmtree,
exc=KeyboardInterrupt,
prefix='fit_download_deleteme_',
) as tempdir:
download_and_extract(self.remote_fits[fitname], tempdir)
#Handle old-style fits compressed with 'results' as root.
old_style_res = tempdir/'results'
if old_style_res.is_dir():
move_target = old_style_res / fitname
else:
move_target = tempdir/fitname
if not move_target.is_dir():
raise RemoteLoaderError(f"Unknown format for fit in {tempdir}. Expecting a folder {move_target}")

fitpath = self.resultspath / fitname
shutil.move(move_target, fitpath)
shutil.rmtree(tempdir)
fitpath = self.resultspath / fitname
shutil.move(move_target, fitpath)


if lhaindex.isinstalled(fitname):
Expand Down
149 changes: 77 additions & 72 deletions validphys2/src/validphys/scripts/postfit.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
import itertools
from glob import glob
import logging
import tempfile

import lhapdf

Expand All @@ -31,6 +30,7 @@
from validphys import fitveto
from validphys.core import PDF
from validphys.fitveto import NSIGMA_DISCARD_ARCLENGTH, NSIGMA_DISCARD_CHI2, INTEG_THRESHOLD
from validphys.utils import tempfile_cleaner

log = logging.getLogger()
log.setLevel(logging.DEBUG)
Expand Down Expand Up @@ -124,79 +124,84 @@ def _postfit(results: str, nrep: int, chi2_threshold: float, arclength_threshold

# Paths
nnfit_path = result_path / 'nnfit' # Path of nnfit replica output
final_postfit_path = result_path / 'postfit'
# Create a temporary path to store work in progress and move it to
# the final location in the end,
postfit_path = pathlib.Path(tempfile.mkdtemp(prefix='postfit_work_deleteme_',
dir=result_path))
final_postfit_path = result_path / 'postfit'
LHAPDF_path = postfit_path/fitname # Path for LHAPDF grid output

if not fitdata.check_nnfit_results_path(result_path):
raise PostfitError('Postfit cannot find a valid results path')
if not fitdata.check_lhapdf_info(result_path, fitname):
raise PostfitError('Postfit cannot find a valid LHAPDF info file')

nrep = int(nrep)
if at_least_nrep:
log.warning(f"Postfit aiming for at least {nrep} replicas")
else:
log.warning(f"Postfit aiming for {nrep} replicas")

# Generate postfit and LHAPDF directory
if final_postfit_path.is_dir():
log.warning(f"Removing existing postfit directory: {postfit_path}")
shutil.rmtree(final_postfit_path)
os.mkdir(LHAPDF_path)

# Setup postfit log
postfitlog = logging.FileHandler(postfit_path/'postfit.log', mode='w')
log.addHandler(postfitlog)

# Perform postfit selection
passing_paths = filter_replicas(postfit_path, nnfit_path, fitname, chi2_threshold, arclength_threshold, integ_threshold)
if len(passing_paths) < nrep:
raise PostfitError("Number of requested replicas is too large")
# Select the first nrep passing replicas
if at_least_nrep:
selected_paths = passing_paths
else:
selected_paths = passing_paths[:nrep]


# Copy info file
info_source_path = nnfit_path.joinpath(f'{fitname}.info')
info_target_path = LHAPDF_path.joinpath(f'{fitname}.info')
shutil.copy2(info_source_path, info_target_path)
set_lhapdf_info(info_target_path, len(selected_paths))

# Generate symlinks
for drep, source_path in enumerate(selected_paths, 1):
# Symlink results to postfit directory
source_dir = pathlib.Path(source_path).resolve()
target_dir = postfit_path.joinpath('replica_%d' % drep)
relative_symlink(source_dir, target_dir)
# Symlink results to pdfset directory
source_grid = source_dir.joinpath(fitname+'.dat')
target_file = f'{fitname}_{drep:04d}.dat'
target_grid = LHAPDF_path.joinpath(target_file)
relative_symlink(source_grid, target_grid)

log.info(f"{len(selected_paths)} replicas written to the postfit folder")

# Generate final PDF with replica 0
log.info("Beginning construction of replica 0")
# It's important that this is prepended, so that any existing instance of
# `fitname` is not read from some other path
lhapdf.pathsPrepend(str(postfit_path))
generatingPDF = PDF(fitname)
lhio.generate_replica0(generatingPDF)

# Test replica 0
try:
lhapdf.mkPDF(fitname, 0)
except RuntimeError as e:
raise PostfitError("CRITICAL ERROR: Failure in reading replica zero") from e
postfit_path.rename(final_postfit_path)
with tempfile_cleaner(
root=result_path,
exit_func=shutil.move,
exc=(KeyboardInterrupt, PostfitError),
prefix="postfit_work_deleteme_",
dst=final_postfit_path,
) as postfit_path:

LHAPDF_path = postfit_path/fitname # Path for LHAPDF grid output

if not fitdata.check_nnfit_results_path(result_path):
raise PostfitError('Postfit cannot find a valid results path')
if not fitdata.check_lhapdf_info(result_path, fitname):
raise PostfitError('Postfit cannot find a valid LHAPDF info file')

nrep = int(nrep)
if at_least_nrep:
log.warning(f"Postfit aiming for at least {nrep} replicas")
else:
log.warning(f"Postfit aiming for {nrep} replicas")

# Generate postfit and LHAPDF directory
if final_postfit_path.is_dir():
log.warning(f"Removing existing postfit directory: {final_postfit_path}")
shutil.rmtree(final_postfit_path)
os.mkdir(LHAPDF_path)

# Setup postfit log
postfitlog = logging.FileHandler(postfit_path/'postfit.log', mode='w')
log.addHandler(postfitlog)

# Perform postfit selection
passing_paths = filter_replicas(postfit_path, nnfit_path, fitname, chi2_threshold, arclength_threshold, integ_threshold)
if len(passing_paths) < nrep:
raise PostfitError("Number of requested replicas is too large")
# Select the first nrep passing replicas
if at_least_nrep:
selected_paths = passing_paths
else:
selected_paths = passing_paths[:nrep]


# Copy info file
info_source_path = nnfit_path.joinpath(f'{fitname}.info')
info_target_path = LHAPDF_path.joinpath(f'{fitname}.info')
shutil.copy2(info_source_path, info_target_path)
set_lhapdf_info(info_target_path, len(selected_paths))

# Generate symlinks
for drep, source_path in enumerate(selected_paths, 1):
# Symlink results to postfit directory
source_dir = pathlib.Path(source_path).resolve()
target_dir = postfit_path.joinpath('replica_%d' % drep)
relative_symlink(source_dir, target_dir)
# Symlink results to pdfset directory
source_grid = source_dir.joinpath(fitname+'.dat')
target_file = f'{fitname}_{drep:04d}.dat'
target_grid = LHAPDF_path.joinpath(target_file)
relative_symlink(source_grid, target_grid)

log.info(f"{len(selected_paths)} replicas written to the postfit folder")

# Generate final PDF with replica 0
log.info("Beginning construction of replica 0")
# It's important that this is prepended, so that any existing instance of
# `fitname` is not read from some other path
lhapdf.pathsPrepend(str(postfit_path))
generatingPDF = PDF(fitname)
lhio.generate_replica0(generatingPDF)

# Test replica 0
try:
lhapdf.mkPDF(fitname, 0)
except RuntimeError as e:
raise PostfitError("CRITICAL ERROR: Failure in reading replica zero") from e
log.info("\n\n*****************************************************************\n")
log.info("Postfit complete")
log.info("Please upload your results with:")
Expand Down
72 changes: 72 additions & 0 deletions validphys2/src/validphys/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,80 @@
@author: Zahari Kassabov
"""
import contextlib
import shutil
import pathlib
import tempfile

import numpy as np


@contextlib.contextmanager
def tempfile_cleaner(root, exit_func, exc, prefix=None, **kwargs):
"""A context manager to handle temporary directory creation and
clean-up upon raising an expected exception.
Parameters
----------
root: str
The root directory to create the temporary directory in.
exit_func: Callable
The exit function to call upon exiting the context manager.
Usually one of ``shutil.move`` or ``shutil.rmtree``. Use the former
if the temporary directory will be the final result directory and the
latter if the temporary directory will contain the result directory, for
example when downloading a resource.
exc: Exception
The exception to catch within the ``with`` block.
prefix: optional[str]
A prefix to prepend to the temporary directory.
**kwargs: dict
Keyword arguments to provide to ``exit_func``.
Returns
-------
tempdir: pathlib.Path
The path to the temporary directory.
Example
-------
The following example creates a temporary directory prepended with
``tutorial_`` in the ``/tmp`` directory. The context manager will listen
for a ``KeyboardInterrupt`` and will clean up if this exception is
raised. Upon completion of the ``with`` block, it will rename the
temporary to ``completed`` as the ``dst``, using ``shutil.move``. The
final directory will contain an empty file called ``new_file``, which
we created within the ``with`` block.
.. code-block:: python
:linenos:
import shutil
from validphys.utils import tempfile_cleaner
with tempfile_cleaner(
root="/tmp",
exit_func=shutil.move,
exc=KeyboardInterrupt,
prefix="tutorial_",
dst="completed",
) as tempdir:
new_file = tempdir / "new_file"
input("Press enter to continue or Ctrl-C to interrupt:\\n")
new_file.touch()
"""
try:
tempdir = pathlib.Path(tempfile.mkdtemp(prefix=prefix, dir=root))
yield tempdir
except exc:
shutil.rmtree(tempdir)
raise
else:
# e.g shutil.rmtree, shutil.move etc
exit_func(tempdir, **kwargs)

def split_by(it, crit):
"""Split ``it`` in two lists, the first is such that ``crit`` evaluates to
True and the second such it doesn't. Crit can be either a function or an
Expand Down

0 comments on commit 332f141

Please sign in to comment.