aws-cdk-neuronx-patterns/lib/sagemaker-inference-toolkit-tnx/sagemaker-inference-toolkit-tnx-compiler.d.ts

> [!WARNING] > This library is experimental module.

import { Size } from "aws-cdk-lib";
import * as ec2 from "aws-cdk-lib/aws-ec2";
import { ContainerImage } from "aws-cdk-lib/aws-ecs";
import { IBucket } from "aws-cdk-lib/aws-s3";
import { Construct } from "constructs";
import { INeuronxImage, INeuronxInstanceType, Model, OptLevel, QuantDtype } from "../base/neuronx";
import { INeuronxContainerImage, NeuronxCompiledModel } from "../base/neuronx-compiler";
/**
 * Compile options.
 */
export interface SageMakerInferenceToolkitTnxCompileOptions {
    /**
     * Number of tensor parallel groups.
     * @default - calc from parameters and quantDtype
     */
    readonly tpDegree?: number;
    /**
     * Maximum number of sequences per iteration.
     * @default 1
     */
    readonly batchSize?: number;
    /**
     * @default - No quant
     */
    readonly quantDtype?: QuantDtype;
    /**
     * @default 4096
     */
    readonly nPositions?: number;
    /**
     * @default OptLevel.BEST_BALANCE
     */
    readonly optLevel?: OptLevel;
}
/**
 * The model compiled by Neuronx compiler.
 */
export interface SageMakerInferenceToolkitTnxCompiledModel extends NeuronxCompiledModel {
    readonly quantDtype?: QuantDtype;
    readonly optLevel?: OptLevel;
    readonly tpDegree: number;
    readonly nPositions: number;
}
/**
 * Compile image for SageMakerInferenceToolkitTnxCompile.
 * @example
 * new SageMakerInferenceToolkitTnxCompileImage(scope, 'CompileImage', PytorchTrainingNeuronxImage.LATEST);
 */
export declare class SageMakerInferenceToolkitTnxCompileImage implements INeuronxContainerImage {
    readonly image: ContainerImage;
    readonly neuronSdkVersion: string;
    constructor(scope: Construct, id: string, neruonxImage: INeuronxImage);
}
/**
 * Props of NeuronxCompile.
 */
export interface SageMakerInferenceToolkitTnxCompilerProps {
    /**
     * VPC in which this will launch compile worker instance.
     */
    readonly vpc: ec2.IVpc;
    /**
     * The bucket to upload compiled artifacts.
     */
    readonly bucket: IBucket;
    /**
     * The model to be compiled.
     */
    readonly model: Model;
    /**
     * The instance type of compile worker instance.
     */
    readonly neuronxInstanceType?: INeuronxInstanceType;
    /**
     * The root volume of worker instance.
     * @default - N bilion parameters * 5GiB EBS
     */
    readonly volumeSize?: Size;
    /**
     * Compile runtime image.
     * @default - latest image
     */
    readonly image?: SageMakerInferenceToolkitTnxCompileImage;
    /**
     * Neuronx compile options.
     * @default - Each properties are set default.
     */
    readonly compileOptions?: SageMakerInferenceToolkitTnxCompileOptions;
    /**
     * Whether or not to use spot instances. Spot instances are less expensive EC2 instances that can be reclaimed by EC2 at any time; your job will be given two minutes of notice before reclamation.
     *
     * @default false
     */
    readonly spot?: boolean;
    /**
     * The VPC Subnets this Compute Environment will launch instances in.
     *
     * @default - new subnets will be created
     */
    readonly vpcSubnets?: ec2.SubnetSelection;
}
/**
 * Neuronx compile construct. Compile the model to work with Inferentia2 and Trainium1 and upload it to an S3 bucket.
 */
export declare class SageMakerInferenceToolkitTnxCompiler extends Construct {
    readonly quantDtype?: QuantDtype;
    readonly optLevel?: OptLevel;
    readonly tpDegree: number;
    readonly nPositions: number;
    model: Model;
    private readonly compiler;
    constructor(scope: Construct, id: string, props: SageMakerInferenceToolkitTnxCompilerProps);
    compile(): SageMakerInferenceToolkitTnxCompiledModel;
}