Source code for schrodinger.application.livedesign.job_utils

"""
LiveDesign protocol job launching utils

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

import os
import shlex

from distutils import util as dutil

from schrodinger.job import queue
from schrodinger.job import jobcontrol

from schrodinger.utils import fileutils


def get_procs(backend):
    """
    Determine the HOST and number of CPUs to be used

    :param backend: Job backend object
    :type backend: 'schrodinger.job.jobcontrol._Backend', None

    :return: String of the hostname:processors
    :rtype: str
    """

    if backend:
        host_list = jobcontrol.get_backend_host_list()
    else:
        host_list = jobcontrol.get_command_line_host_list()

    return jobcontrol.host_list_to_str(host_list)


def use_smart_dist():
    """
    Enable/disable Smart Distribution when running on a queuing system

    Default behavior (True) is to use Smart Distribution when the env var is not
    set

    :return: True/False based on the env var setting
    :rtype: bool
    """

    smart_dist_env = os.getenv('SCHRODINGER_LD_SMARTDIST')

    if smart_dist_env is None:
        return True
    else:
        return bool(dutil.strtobool(smart_dist_env))


class LigPrepJob(queue.JobControlJob):
    """
    LigPrep Job object
    """

    base_command = ['ligprep']

    def __init__(self,
                 infile,
                 outfile,
                 host_str,
                 ligprep_inp=None,
                 ligprep_args=''):
        """
        Create a LigPrep job object using in the input, output and job settings

        :param infile: Input SD file
        :type infile: str

        :param outfile: Output Maestro file
        :type outfile: str

        :param host_str: Str containing hostname:cpus
        :type host_str: str

        :param ligprep_inp: File containing LigPrep keyword/value arguments
        :type ligprep_inp: str, None

        :param ligprep_args: LigPrep options
        :type ligprep_args: str, None
        """

        self.infile = infile
        self.outfile = outfile
        self.ligprep_inp = ligprep_inp
        self.ligprep_args = ligprep_args
        self.host_str = host_str

        outfile_basename = fileutils.get_basename(self.outfile)

        self.logfile = f'{outfile_basename}.log'
        self.dropped_file = f'{outfile_basename}-dropped.sdf'

        command = ['ligprep', '-isd', self.infile, '-omae', self.outfile]

        if self.ligprep_inp:
            command.extend(['-inp', self.ligprep_inp])

        if self.ligprep_args:
            command.extend(shlex.split(self.ligprep_args))

        # When this object is added to the queue, it is submitted as a single
        # job even if multiple CPUs are specified. The -HOST must be added here
        # in order for LigPrep to split across the number of specified CPUs
        command.extend(['-HOST', self.host_str])

        super().__init__(command)