Source code for schrodinger.application.desmond.stage.jobs

import os
from dataclasses import dataclass
from typing import Optional

import numpy as np

from schrodinger.application.desmond import cmj
from schrodinger.application.desmond.picklejar import Picklable


@dataclass
class SnapcoreCoordinates:

    fh_ref_coord: np.array
    fh_mut_coord: np.array
    sh_ref_coord: np.array
    sh_mut_coord: np.array


class DesmondJob(cmj.Job):
    """

    """

    @staticmethod
    def err_handler(job):
        """

        """
        cmj.JobErrorHandler.default(job)
        if (job.status == cmj.JobStatus.KILLED):
            # Inhibits retrials. Instead, we revive the job from its .cpt file.
            job.status.set(cmj.JobStatus.NON_RETRIABLE_FAILURE)
        else:
            log_fname = job.output.log_file()
            if not log_fname:  # No log file is associated with this job.
                return
            if not os.path.isfile(log_fname):
                job._print("quiet", f"Log file not found: {log_fname}")
            else:
                with open(log_fname, "r") as f:
                    log_content = f.read()
                if (-1 != log_content.find(
                        "out-of-range result; the simulation likely is")):
                    job._print(
                        "quiet",
                        "DIAGNOSIS: The error message indicates that your system is unstable and \"blowing up\". "
                        "Re-check your input files and system setup for atom overlaps or other instabilities."
                    )
                elif (-1 != log_content.find(
                        "Not all processes managed to acquire a GPU.")):
                    job._print(
                        "quiet",
                        "DIAGNOSIS: Desmond cannot find a GPU to run on. Your hardware, queuing system and/or "
                        " schrodinger.hosts file may be incorrectly configured."
                    )
                else:
                    import re
                    if (re.compile("An allocation for .+ bytes failed!").search(
                            log_content)):
                        job._print(
                            "quiet",
                            "DIAGNOSIS: The GPU ran out of memory. It may be that your system is too large to run "
                            "on this GPU model or that more than one simulation is running on the GPU due to an "
                            "incorrect configuration.")

    def __init__(self,
                 task="",
                 gid=None,
                 permanent_restrain=None,
                 permanent_group=None,
                 *arg,
                 **kwarg):
        """

        """
        kwarg.setdefault("err_handler", DesmondJob.err_handler)

        cmj.Job.__init__(self, *arg, **kwarg)

        self.task = task
        self.gid = gid
        self.rid = None  # Replica ID
        self.permanent_restrain = permanent_restrain
        self.permanent_group = permanent_group
        self.need_host = True

        parent = self.parent
        if (self.gid is None and hasattr(parent, "gid")):
            self.gid = parent.gid
        if (hasattr(parent, "cross_link")):
            self.cross_link = parent.cross_link
        # __init__

    def describe(self):
        """

        """
        cmj.Job.describe(self)
        self._print("verbose", "  Task       : %s" % self.task)


class DesmondBackendJob(DesmondJob):
    USE_GPU = True


class FepJob(DesmondBackendJob, Picklable):
    """

    """

    id = 0, Picklable

    def __init__(self,
                 fepid=None,
                 n_win=None,
                 i_win=None,
                 fepout=None,
                 egout=None,
                 fragname=None,
                 snapcore: Optional[SnapcoreCoordinates] = None,
                 *arg,
                 **kwarg):
        """

        """
        DesmondJob.__init__(self, *arg, **kwarg)

        self.fepid = fepid
        self.n_win = n_win
        self.i_win = i_win
        self.fragname = fragname
        self.snapcore = snapcore

        self.output.add(fepout, tag="dE")
        self.output.add(egout, tag="eg")

        parent = self.parent
        if (self.fepid is None and hasattr(parent, "fepid")):
            self.fepid = parent.fepid
        if (self.n_win is None and hasattr(parent, "n_win")):
            self.n_win = parent.n_win
        if (self.i_win is None and hasattr(parent, "i_win")):
            self.i_win = parent.i_win
        if (self.fragname is None and hasattr(parent, "fragname")):
            self.fragname = parent.fragname

        if (self.fepid is None):
            self.fepid = FepJob.id
            FepJob.id += 1
        # __init__

    def describe(self):
        """

        """
        DesmondJob.describe(self)
        self._print(
            "verbose", "  Fep ID     : %d\n"
            "  Mutation   : %s\n"
            "  Lambda     : %d (%d windows)" % (
                self.fepid,
                self.fragname,
                self.i_win,
                self.n_win,
            ))