Source code for composer.utils.eval_client.mosaicml_lambda_eval_client

# Copyright 2022 MosaicML Composer authors
# SPDX-License-Identifier: Apache-2.0

"""MCLI compatible eval client."""
import logging
import os
import time
from http import HTTPStatus
from typing import Dict, List

import mcli
import numpy as np

from composer.utils.eval_client.eval_client import EvalClient

__all__ = ['MosaicMLLambdaEvalClient']
log = logging.getLogger(__name__)


[docs]class MosaicMLLambdaEvalClient(EvalClient): """Utility for creating a client for and invoking an AWS Lambda through MCLI.""" def __init__(self, backoff: int = 3, num_retries: int = 5) -> None: """Checks that the requisite environment variables are in the EvalClient. `MOSAICML_ACCESS_TOKEN_ENV_VAR` environment variable must be set to access the platform. """ from composer.loggers.mosaicml_logger import \ MOSAICML_ACCESS_TOKEN_ENV_VAR # in-line import to avoid circular import if MOSAICML_ACCESS_TOKEN_ENV_VAR not in os.environ: raise RuntimeError('Cannot use MosaicML Lambda Client Eval without setting MOSAICML_ACCESS_TOKEN_ENV_VAR.') log.debug('Running code eval through MosaicMLLambdaEvalClient.') self.backoff = backoff self.num_retries = num_retries
[docs] def invoke(self, payload: List[List[List[Dict[str, str]]]]) -> List[List[List[bool]]]: """Invoke a batch of provided payloads for code evaluations.""" num_beams = len(payload[0]) num_tests = [len(generation_payload[0]) for generation_payload in payload] cum_tests = (np.cumsum([0] + num_tests[:-1]) * num_beams).tolist() test_cases = [ test_case for generation_payload in payload for beam_payload in generation_payload for test_case in beam_payload ] ret_helper = [False] * len(test_cases) for i in range(self.num_retries): try: ret_helper = mcli.get_code_eval_output(test_cases).data # pyright: ignore[reportGeneralTypeIssues] break except mcli.MAPIException as e: if e.status >= 500: if i == self.num_retries - 1: log.error(f'Failed to get code eval output after {self.num_retries} retries. Error: {e}') log.warning(f'Failed to get code eval output, retrying in {self.backoff**i} seconds.') time.sleep(self.backoff**i) elif e.status == HTTPStatus.UNAUTHORIZED: raise RuntimeError('Failed to get code eval output due to UNAUTHORIZED error. ' 'Please ensure you have access to MosaicMLLambdaEvalClient.') from e else: log.error(f'Failed to get code eval output with unexpected MAPIException. Error: {e}') break except TimeoutError as e: if i == self.num_retries - 1: log.error(f'Failed to get code eval output after {self.num_retries} retries. Error: {e}') log.warning(f'Failed to get code eval output, retrying in {self.backoff**i} seconds.') time.sleep(self.backoff**i) except Exception as e: log.error(f'Failed to get code eval output with unexpected error. Error: {e}') break ret = [[[ret_helper[cum_tests[i] + j * num_tests[i] + k] for k in range(num_tests[i])] for j in range(num_beams)] for i in range(len(payload))] return ret