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


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_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