Source code for schrodinger.application.phase.license

"""
Module for token checkouts in phase related backends.
NOTE: Tampering with licensing is a violation of the license agreement.

Copyright Schrodinger LLC, All Rights Reserved.
"""

import sys

from schrodinger import gpgpu
from schrodinger.job import jobcontrol
from schrodinger.utils import license
from schrodinger.utils import log

# Logging
logger = log.get_output_logger(__file__)

SHAPE_SCREEN_LICENSE_NAME = license.SHAPE_SCREEN
PHASE_DBSEARCH_LICENSE_ID = license.PHASE_DBSEARCH


def num_of_shape_screen_tokens(use_gpu: bool = False) -> int:
    """
    :param use_gpu: whether to checkout tokens for GPU usage or not
    :return: number of tokens to check out
    """

    CPU_SHAPE_SCREEN_TOKENS = 1
    GPU_SHAPE_SCREEN_TOKENS = 8

    if not use_gpu:
        return CPU_SHAPE_SCREEN_TOKENS

    backend = jobcontrol.get_backend()
    host = backend.getJob().HostEntry if backend else "localhost"
    return gpgpu.get_scaled_token_count(host, GPU_SHAPE_SCREEN_TOKENS)


def checkout_phase_dbsearch_tokens(num_tokens: int = 1,
                                   want_exit: bool = True) -> license.License:
    """
    Checks out a PHASE_DBSEARCH license consuming the requested number
    of tokens.

    :param num_tokens: Number of tokens to check out
    :param want_exit: Whether to call exit on failure or raise RuntimeError
    :return: License consuming num_tokens PHASE_DBSEARCH tokens
    """

    tokens = license.License(PHASE_DBSEARCH_LICENSE_ID, num_tokens)
    if not tokens.isValid():
        msg = (f'{num_tokens}-token PHASE_DBSEARCH license is unavailable. '
               'Contact help@schrodinger.com to request an appropriate '
               'license.')
        if want_exit:
            logger.error(msg)
            sys.exit(1)
        else:
            raise RuntimeError(msg)

    return tokens


def checkout_shape_screen_tokens(use_gpu: bool = False) -> license.License:
    """
    Checks out the appropriate number of tokens for a shape screening job.

    :param use_gpu: whether to checkout tokens for GPU usage or not
    :return: license token
    """

    tokens = license.License(SHAPE_SCREEN_LICENSE_NAME,
                             num_of_shape_screen_tokens(use_gpu))
    if not tokens.isValid():
        logger.error("SHAPE_SCREEN license token is unavailable. Contact "
                     "help@schrodinger.com to request additional licenses.")
        sys.exit(1)
    return tokens