Source code for schrodinger.job.driver_decorator

"""
Main function decorator designed for backend jobs. For example::

    @driver_decorator.main_wrapper("A super nifty backend")
    def main():
        do_some_cool_work()

Copyright Schrodinger, LLC. All rights reserved.
"""

import datetime

from decorator import decorator

from schrodinger.job import jobcontrol
from schrodinger.utils import cmdline
from schrodinger.utils import log

logger = log.get_output_logger(__file__)


@decorator
def main_wrapper(func, driver_name="Job", *args, **kwargs):
    """
    The backend job's application header string is logged only if job is
    running under jobcontrol.

    Running time is logged after job succeeds, or when an exception is caught.
    Running time won't be logged if sys.exit() in called inside main(), or
    job is terminated via keyboard interruption.
    """

    def _driver_wrapper(*args, **kwargs):

        backend = jobcontrol.get_backend()
        if backend:
            logger.info(backend.getJob().getApplicationHeaderString())

        t1 = datetime.datetime.now()
        status = "successfully completed."
        try:
            result = func(*args, **kwargs)
        except KeyboardInterrupt:
            status = "interrupted."
            raise
        except SystemExit as e:
            status = f"exited.\n\nSystemExit: {e}"
            raise
        except Exception:
            status = "failed."
            raise
        finally:
            t2 = datetime.datetime.now()
            dt = str(t2 - t1)[:-4]  # Hack to only print only hundredths of sec
            logger.info("")
            logger.info(f"Current time: {t2.ctime()}")
            logger.info(f"Elapsed time: {dt}")
            logger.info(f"{driver_name} {status}")
            logger.info("")

        return result

    return cmdline.main_wrapper(_driver_wrapper, *args, **kwargs)