def get_optimum_neuron_pipeline()

in src/sagemaker_huggingface_inference_toolkit/optimum_utils.py [0:0]

• 16 lines of code
• 6 McCabe index (conditional complexity)

def get_optimum_neuron_pipeline(task, model_dir):
    """Method to get optimum neuron pipeline for a given task. Method checks if task is supported by optimum neuron and if required environment variables are set, in case model is not converted. If all checks pass, optimum neuron pipeline is returned. If checks fail, an error is raised."""
    from optimum.neuron.pipelines.transformers.base import NEURONX_SUPPORTED_TASKS, pipeline
    from optimum.neuron.utils import NEURON_FILE_NAME

    # check task support
    if task not in NEURONX_SUPPORTED_TASKS:
        raise ValueError(
            f"Task {task} is not supported by optimum neuron and inf2. Supported tasks are: {list(NEURONX_SUPPORTED_TASKS.keys())}"
        )

    # check if model is already converted and has input shapes available
    export = True
    if NEURON_FILE_NAME in os.listdir(model_dir):
        export = False
    if export:
        logger.info("Model is not converted. Checking if required environment variables are set and converting model.")

    # get static input shapes to run inference
    input_shapes = get_input_shapes(model_dir)
    # set NEURON_RT_NUM_CORES to 1 to avoid conflicts with multiple HTTP workers
    os.environ["NEURON_RT_NUM_CORES"] = "1"
    # get optimum neuron pipeline
    neuron_pipe = pipeline(task, model=model_dir, export=export, input_shapes=input_shapes)

    return neuron_pipe