Source code for schrodinger.application.desmond.license

"""
Module for license related code in Desmond workflows.

Copyright Schrodinger LLC, All Rights Reserved.
"""
import getpass
import json
import os
from pathlib import Path
from typing import Dict
from typing import List

import requests

from schrodinger.application.desmond import bld_def
from schrodinger.application.desmond import bld_ver
from schrodinger.infra import licensing
from schrodinger.job import jobcontrol
from schrodinger import gpgpu
from schrodinger.application.desmond.constants import PRODUCT

MODEL2_ACCESS_FILE = Path(os.environ["SCHRODINGER"], "config",
                          "per_compound.json")
MODEL2_ENV_NAME = 'SCHRODINGER_PER_COMPOUND_LICENSING'
DESMOND_GPU_TOKEN_MULTIPLIER = licensing.DESMOND_GPU_TOKEN_MULTIPLIER


class Model2CheckoutError(Exception):
    pass


def add_fep_lic(args: List[str]) -> List[str]:
    """
    Add the command line arguments to properly handle the FEP GPU license.

    NOTE: If the same license is present with a different
    number of tokens, this will override the current value.

    :param args: List of command line arguments. This function will look into
        `args` for the -HOST and -cpu options to determine the number of processors.
        If these are not found, this function will assume one processor.
    :param use_custom_charges: Set to True when using custom charges (F16).

    :return: The updated command.
    """
    return _deduplicate_lics(args + ["-lic", fep_lic(get_host(args))])


def add_md_lic(args: List[str]) -> List[str]:
    """
    Add the command line arguments to properly handle the
    MD GPU license.

    NOTE: If the same license is present with a different
    number of tokens, this will override the current value.

    :param args: List of command line arguments. This function will look into
        `args` for the -HOST and -cpu options to determine the number of processors.
        If these are not found, this function will assume one processor.

    :return: The updated command.
    """
    return _deduplicate_lics(args + ["-lic", md_lic(get_host(args))])


def add_watermap_lic(args: List[str]) -> List[str]:
    """
    Add the command line arguments to properly handle the
    Watermap GPU license.

    NOTE: If the same license is present with a different
    number of tokens, this will override the current value.

    :param args: List of command line arguments. This function will look into
        `args` for the -HOST and -cpu options to determine the number of processors.
        If these are not found, this function will assume one processor.

    :return: The updated command.
    """
    return _deduplicate_lics(args + ["-lic", watermap_lic(get_host(args))])


def protein_mutation_lics() -> List[str]:
    """
    Return the license strings needed for Plop/Macromodel, used in the
    protein mutation generator stage.
    """
    return [
        f"{licensing.getFeatureName(licensing.PSP_PLOP)}:8",
        f'{licensing.getFeatureName(licensing.MMOD_MACROMODEL)}:2'
    ]


def epik_lic() -> str:
    """
    Return the license string needed for Epik, used in constant pH.
    """
    return f'{licensing.getFeatureName(licensing.EPIK_MAIN)}:1'


def fep_lic(host: str) -> str:
    """
    Return the license string for FEP GPU, given the -HOST value.
    """
    license_string = licensing.getFeatureName(licensing.FEP_GPGPU)
    multiplier = gpgpu.get_scaled_token_count(_get_hostname(host),
                                              DESMOND_GPU_TOKEN_MULTIPLIER)
    return f"{license_string}:{_get_num_procs(host) * multiplier}"


def md_lic(host: str) -> str:
    """
    Return the license string for MD GPU, given the -HOST value.
    """
    license_string = licensing.getFeatureName(
        licensing.DESMOND_ACADEMIC) if _is_academic(
        ) else licensing.getFeatureName(licensing.DESMOND_GPGPU)
    multiplier = gpgpu.get_scaled_token_count(_get_hostname(host),
                                              DESMOND_GPU_TOKEN_MULTIPLIER)
    return f"{license_string}:{_get_num_procs(host) * multiplier}"


def watermap_lic(host: str) -> str:
    """
    Return the license string for Watermap GPU, given the -HOST value.
    """
    license_string = licensing.getFeatureName(licensing.DESMOND_WATERMAP_GPGPU)
    multiplier = gpgpu.get_scaled_token_count(_get_hostname(host),
                                              DESMOND_GPU_TOKEN_MULTIPLIER)
    return f"{license_string}:{_get_num_procs(host) * multiplier}"


def _get_model2_filename():
    """
    Returns the model2 filename from the environment variable,
    or None if it's not defined.
    """
    return os.getenv(MODEL2_ENV_NAME)


def _post_model2(endpoint: str, additional_data: Dict = None):
    access_fname = _get_model2_filename()
    if not access_fname:
        raise Model2CheckoutError(
            'Could not find FEP+ Model 2 configuration file.')
    params = json.loads(Path(access_fname).read_text())
    url = params['server']
    params.update(additional_data or {})
    return requests.post(f"{url}/{endpoint}", json=params)


def has_model2_file() -> bool:
    access_fname = _get_model2_filename()
    if not access_fname:
        return False
    return Path(access_fname).exists()


def is_model2_server_available() -> bool:
    try:
        return _post_model2('ping').ok
    except Exception:
        # Any exception here means we couldn't connect
        return False


def checkout_model2_compounds(num_compounds: int, product: PRODUCT):
    """
    Checkout the specified number of model2 compounds.

    :param product: Name of the product to use.
    """
    if product == PRODUCT.IFD_MD:
        for _ in range(num_compounds):
            licensing.checkout_ifd_md_completed_compound()
    else:
        licensing.checkout_fep_completed_compound(num_compounds)
    access_fname = _get_model2_filename()
    if not access_fname:
        return
    if not has_model2_file():
        return
    if not is_model2_server_available():
        raise Model2CheckoutError(
            'Need to make sure per compound server is accessible to run jobs. Please contact support.'
        )

    data = {}
    try:
        data['username'] = getpass.getuser()
    except Exception:
        # A number of different exceptions could be raised
        # so just catch them all.
        data['username'] = ''
    data['num_compounds'] = num_compounds
    data['product'] = product
    try:
        _post_model2('checkout', data)
    except Exception as e:
        raise Model2CheckoutError(
            f'Could not checkout compounds: {e}. Please contact support.')


def get_host(args: List[str]) -> str:
    """
    Return the value for the -HOST argument if found,
    otherwise default to 'localhost:1'.
    """
    if '-HOST' in args:
        return args[args.index('-HOST') + 1]
    return 'localhost:1'


def _get_num_procs(host: str) -> int:
    """
    Return the number of processors (GPUs or CPUs) from command line
    arguments.
    """
    return jobcontrol.host_str_to_list(host)[0][-1] or 1


def _get_hostname(host: str) -> str:
    """
    Return just the host name from a string 'host:number'.
    """
    return jobcontrol.host_str_to_list(host)[0][0]


def _is_academic() -> bool:
    """
    Return True if this is an academic build. False otherwise.
    """
    return bld_ver.desmond_build_version() == bld_def.bld_types[
        bld_def.DESMOND_ACADEMIC]


def _deduplicate_lics(args: List[str]) -> List[str]:
    """
    Go through the command line arguments `args`
    and return the arguments without duplicate licenses.
    NOTE: If the same license is present with a different
    number of tokens, the one that appears last in the list
    will be used.
    """
    lic_map = {}
    updated_args = []
    i = 0
    while i < len(args):
        if args[i] == '-lic':
            lic_value = args[i + 1]
            lic_map[lic_value.split(':')[0]] = ['-lic', lic_value]
            # Skip -lic value
            i += 1
        else:
            updated_args.append(args[i])
        i += 1
    return sum(sorted(lic_map.values()), updated_args)