#
# Copyright 2015 by Justin MacCallum, Alberto Perez, Ken Dill
# All rights reserved
#
"""
A module for launching replica exchange runs
"""
import logging
import os
import socket
from typing import Union
from openmm import version as mm_version # type: ignore
import meld
from meld import runner, util, vault
Handler = Union[logging.StreamHandler, logging.FileHandler]
logger = logging.getLogger(__name__)
[docs]def log_versions() -> None:
"""Record version numbers to log"""
logger.info("Meld version is %s", meld.__version__)
logger.info("OpenMM_Meld version is %s", mm_version.full_version)
[docs]def launch(
platform: str,
console_handler: Handler,
debug: bool = False,
console_log: bool = False,
) -> None:
"""
Launch a replica exchange run
Args:
platform: platform to run on [Reference, CPU, CUDA]
console_handler: log handler for console logging
debug: log debugging information
console_log: display logging on console
"""
logger.info("loading data store")
store = vault.DataStore.load_data_store()
logger.info("initializing communicator")
communicator = store.load_communicator()
communicator.initialize()
#
# setup logging
#
hostname = socket.gethostname()
hostid = f"{hostname}:{communicator.rank:03d}"
meld_logger = logging.getLogger("meld")
# this filter adds the hostid to each logging record so
# it gets printed out as part of the logging output
hostid_log_filter = util.HostNameContextFilter(hostid)
# remove the console handler, so that
# we can add a new handler below without
# duplicate logging messages
meld_logger.removeHandler(console_handler)
if not console_log:
# setup file
log_path = os.path.join(store.log_dir, f"remd_{communicator.rank:03d}.log")
handler: Handler = logging.FileHandler(
filename=log_path,
mode="a",
)
else:
fmt = "%(hostid)s %(asctime)s %(levelname)s %(name)s: %(message)s"
fmt = fmt.format(hostid)
datefmt = "%Y-%m-%d %H:%M:%S"
formatter = logging.Formatter(fmt=fmt, datefmt=datefmt)
handler = logging.StreamHandler()
handler.setFormatter(formatter)
meld_logger.addHandler(handler)
handler.addFilter(hostid_log_filter)
level = logging.DEBUG if debug else logging.INFO
handler.setLevel(level)
meld_logger.setLevel(level)
meld_logger.propagate = False
if communicator.is_leader():
logger.info("Launching replica exchange on leader")
else:
logger.info("Launching replica exchange on worker")
log_versions()
logger.info("Loading system")
system = store.load_system()
logger.info("Loading run options")
options = store.load_run_options()
system_runner = runner.get_runner(
system, options, comm=communicator, platform=platform
)
if communicator.is_leader():
store.initialize(mode="a")
remd_runner = store.load_remd_runner()
remd_runner.run(communicator, system_runner, store)
else:
remd_runner = store.load_remd_runner().to_worker()
remd_runner.run(communicator, system_runner)