aws-cdk-neuronx-patterns
Version:
> [!WARNING] > This library is experimental module.
69 lines • 9.88 kB
JavaScript
;
var _a;
Object.defineProperty(exports, "__esModule", { value: true });
exports.NeuronxNativeCompiler = void 0;
const JSII_RTTI_SYMBOL_1 = Symbol.for("jsii.rtti");
const aws_cdk_lib_1 = require("aws-cdk-lib");
const ec2 = require("aws-cdk-lib/aws-ec2");
const aws_batch_1 = require("../aws-batch");
const neuronx_1 = require("../neuronx");
const neuronx_compiler_base_1 = require("./neuronx-compiler-base");
/**
* Neuronx compiler construct.
* Compile the model to work with Inferentia2 and Trainium1 and upload it to an S3 bucket.
*/
class NeuronxNativeCompiler extends neuronx_compiler_base_1.NeuronxCompilerBase {
constructor(scope, id, props) {
super(scope, id, props);
}
createComputeEnvironment(props) {
const volumeSize = props.volumeSize?.toGibibytes() ??
Math.ceil(this.weightSize.toGibibytes() +
neuronx_1.PytorchTrainingNeuronxImage.size.toGibibytes() +
neuronx_1.NeuronOptimizedMachineImage.size.toGibibytes());
const launchTemplate = new ec2.LaunchTemplate(this, "LaunchTemplate", {
blockDevices: [
{
deviceName: "/dev/xvda",
volume: ec2.BlockDeviceVolume.ebs(volumeSize, {
volumeType: ec2.EbsDeviceVolumeType.GP3,
encrypted: true,
}),
},
],
});
const neuronxInstanceType = props.neuronxInstanceType ?? neuronx_1.NeuronxInstanceType.INF2_48XLARGE;
const computeEnvironment = new aws_batch_1.NeuronxBatchComputeEnvironment(this, "ComputeEnvironment", {
vpc: props.vpc,
vpcSubnets: props.vpcSubnets,
instanceTypes: [neuronxInstanceType.instanceType],
useOptimalInstanceClasses: false,
launchTemplate,
spot: props.spot,
});
aws_cdk_lib_1.Tags.of(computeEnvironment).add("Name", "neuronx-compile-worker");
return {
computeEnvironment,
instanceRole: computeEnvironment.instanceRole,
};
}
createJobDefinition(props) {
const neuronxInstanceType = props.neuronxInstanceType ?? neuronx_1.NeuronxInstanceType.INF2_48XLARGE;
return new aws_batch_1.NeuronxBatchEcsJobDefinition(this, "JobDefinition", {
neuronxInstanceType,
image: props.image.image,
memory: aws_cdk_lib_1.Size.mebibytes(Math.ceil(neuronxInstanceType.memory.toMebibytes() * 0.95)),
cpu: neuronxInstanceType.vCpu,
environment: {
NEURON_COMPILE_CACHE_URL: `${props.bucket.s3UrlForObject("neuron-compile-cache")}`,
...props.environment,
},
command: props.command,
secrets: props.secrets,
});
}
}
exports.NeuronxNativeCompiler = NeuronxNativeCompiler;
_a = JSII_RTTI_SYMBOL_1;
NeuronxNativeCompiler[_a] = { fqn: "aws-cdk-neuronx-patterns.NeuronxNativeCompiler", version: "0.3.0" };
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoibmV1cm9ueC1uYXRpdmUtY29tcGlsZXIuanMiLCJzb3VyY2VSb290IjoiIiwic291cmNlcyI6WyIuLi8uLi8uLi9zcmMvYmFzZS9uZXVyb254LWNvbXBpbGVyL25ldXJvbngtbmF0aXZlLWNvbXBpbGVyLnRzIl0sIm5hbWVzIjpbXSwibWFwcGluZ3MiOiI7Ozs7O0FBQUEsNkNBQXlDO0FBRXpDLDJDQUEyQztBQUUzQyw0Q0FHc0I7QUFDdEIsd0NBSW9CO0FBQ3BCLG1FQU1pQztBQU9qQzs7O0dBR0c7QUFDSCxNQUFhLHFCQUFzQixTQUFRLDJDQUFtQjtJQUM1RCxZQUFZLEtBQWdCLEVBQUUsRUFBVSxFQUFFLEtBQWlDO1FBQ3pFLEtBQUssQ0FBQyxLQUFLLEVBQUUsRUFBRSxFQUFFLEtBQUssQ0FBQyxDQUFDO0lBQzFCLENBQUM7SUFFUyx3QkFBd0IsQ0FDaEMsS0FBK0I7UUFFL0IsTUFBTSxVQUFVLEdBQ2QsS0FBSyxDQUFDLFVBQVUsRUFBRSxXQUFXLEVBQUU7WUFDL0IsSUFBSSxDQUFDLElBQUksQ0FDUCxJQUFJLENBQUMsVUFBVSxDQUFDLFdBQVcsRUFBRTtnQkFDM0IscUNBQTJCLENBQUMsSUFBSSxDQUFDLFdBQVcsRUFBRTtnQkFDOUMscUNBQTJCLENBQUMsSUFBSSxDQUFDLFdBQVcsRUFBRSxDQUNqRCxDQUFDO1FBQ0osTUFBTSxjQUFjLEdBQUcsSUFBSSxHQUFHLENBQUMsY0FBYyxDQUFDLElBQUksRUFBRSxnQkFBZ0IsRUFBRTtZQUNwRSxZQUFZLEVBQUU7Z0JBQ1o7b0JBQ0UsVUFBVSxFQUFFLFdBQVc7b0JBQ3ZCLE1BQU0sRUFBRSxHQUFHLENBQUMsaUJBQWlCLENBQUMsR0FBRyxDQUFDLFVBQVUsRUFBRTt3QkFDNUMsVUFBVSxFQUFFLEdBQUcsQ0FBQyxtQkFBbUIsQ0FBQyxHQUFHO3dCQUN2QyxTQUFTLEVBQUUsSUFBSTtxQkFDaEIsQ0FBQztpQkFDSDthQUNGO1NBQ0YsQ0FBQyxDQUFDO1FBRUgsTUFBTSxtQkFBbUIsR0FDdkIsS0FBSyxDQUFDLG1CQUFtQixJQUFJLDZCQUFtQixDQUFDLGFBQWEsQ0FBQztRQUNqRSxNQUFNLGtCQUFrQixHQUFHLElBQUksMENBQThCLENBQzNELElBQUksRUFDSixvQkFBb0IsRUFDcEI7WUFDRSxHQUFHLEVBQUUsS0FBSyxDQUFDLEdBQUc7WUFDZCxVQUFVLEVBQUUsS0FBSyxDQUFDLFVBQVU7WUFDNUIsYUFBYSxFQUFFLENBQUMsbUJBQW1CLENBQUMsWUFBWSxDQUFDO1lBQ2pELHlCQUF5QixFQUFFLEtBQUs7WUFDaEMsY0FBYztZQUNkLElBQUksRUFBRSxLQUFLLENBQUMsSUFBSTtTQUNqQixDQUNGLENBQUM7UUFDRixrQkFBSSxDQUFDLEVBQUUsQ0FBQyxrQkFBa0IsQ0FBQyxDQUFDLEdBQUcsQ0FBQyxNQUFNLEVBQUUsd0JBQXdCLENBQUMsQ0FBQztRQUVsRSxPQUFPO1lBQ0wsa0JBQWtCO1lBQ2xCLFlBQVksRUFBRSxrQkFBa0IsQ0FBQyxZQUFZO1NBQzlDLENBQUM7SUFDSixDQUFDO0lBRVMsbUJBQW1CLENBQzNCLEtBQStCO1FBRS9CLE1BQU0sbUJBQW1CLEdBQ3ZCLEtBQUssQ0FBQyxtQkFBbUIsSUFBSSw2QkFBbUIsQ0FBQyxhQUFhLENBQUM7UUFDakUsT0FBTyxJQUFJLHdDQUE0QixDQUFDLElBQUksRUFBRSxlQUFlLEVBQUU7WUFDN0QsbUJBQW1CO1lBQ25CLEtBQUssRUFBRSxLQUFLLENBQUMsS0FBSyxDQUFDLEtBQUs7WUFDeEIsTUFBTSxFQUFFLGtCQUFJLENBQUMsU0FBUyxDQUNwQixJQUFJLENBQUMsSUFBSSxDQUFDLG1CQUFtQixDQUFDLE1BQU0sQ0FBQyxXQUFXLEVBQUUsR0FBRyxJQUFJLENBQUMsQ0FDM0Q7WUFDRCxHQUFHLEVBQUUsbUJBQW1CLENBQUMsSUFBSTtZQUM3QixXQUFXLEVBQUU7Z0JBQ1gsd0JBQXdCLEVBQUUsR0FBRyxLQUFLLENBQUMsTUFBTSxDQUFDLGNBQWMsQ0FBQyxzQkFBc0IsQ0FBQyxFQUFFO2dCQUNsRixHQUFHLEtBQUssQ0FBQyxXQUFXO2FBQ3JCO1lBQ0QsT0FBTyxFQUFFLEtBQUssQ0FBQyxPQUFPO1lBQ3RCLE9BQU8sRUFBRSxLQUFLLENBQUMsT0FBTztTQUN2QixDQUFDLENBQUM7SUFDTCxDQUFDOztBQXBFSCxzREFxRUMiLCJzb3VyY2VzQ29udGVudCI6WyJpbXBvcnQgeyBTaXplLCBUYWdzIH0gZnJvbSBcImF3cy1jZGstbGliXCI7XG5pbXBvcnQgKiBhcyBiYXRjaCBmcm9tIFwiYXdzLWNkay1saWIvYXdzLWJhdGNoXCI7XG5pbXBvcnQgKiBhcyBlYzIgZnJvbSBcImF3cy1jZGstbGliL2F3cy1lYzJcIjtcbmltcG9ydCB7IENvbnN0cnVjdCB9IGZyb20gXCJjb25zdHJ1Y3RzXCI7XG5pbXBvcnQge1xuICBOZXVyb254QmF0Y2hDb21wdXRlRW52aXJvbm1lbnQsXG4gIE5ldXJvbnhCYXRjaEVjc0pvYkRlZmluaXRpb24sXG59IGZyb20gXCIuLi9hd3MtYmF0Y2hcIjtcbmltcG9ydCB7XG4gIE5ldXJvbk9wdGltaXplZE1hY2hpbmVJbWFnZSxcbiAgTmV1cm9ueEluc3RhbmNlVHlwZSxcbiAgUHl0b3JjaFRyYWluaW5nTmV1cm9ueEltYWdlLFxufSBmcm9tIFwiLi4vbmV1cm9ueFwiO1xuaW1wb3J0IHtcbiAgSU5ldXJvbnhDb250YWluZXJJbWFnZSxcbiAgTmV1cm9ueENvbXBpbGVkTW9kZWwsXG4gIE5ldXJvbnhDb21waWxlckJhc2UsXG4gIE5ldXJvbnhDb21waWxlckJhc2VQcm9wcyxcbiAgQ29tcHV0ZUVudmlyb25tZW50UmVzdWx0LFxufSBmcm9tIFwiLi9uZXVyb254LWNvbXBpbGVyLWJhc2VcIjtcblxuLyoqXG4gKiBQcm9wcyBvZiBOZXVyb254TmF0aXZlQ29tcGlsZXIuXG4gKi9cbmV4cG9ydCBpbnRlcmZhY2UgTmV1cm9ueE5hdGl2ZUNvbXBpbGVyUHJvcHMgZXh0ZW5kcyBOZXVyb254Q29tcGlsZXJCYXNlUHJvcHMge31cblxuLyoqXG4gKiBOZXVyb254IGNvbXBpbGVyIGNvbnN0cnVjdC5cbiAqIENvbXBpbGUgdGhlIG1vZGVsIHRvIHdvcmsgd2l0aCBJbmZlcmVudGlhMiBhbmQgVHJhaW5pdW0xIGFuZCB1cGxvYWQgaXQgdG8gYW4gUzMgYnVja2V0LlxuICovXG5leHBvcnQgY2xhc3MgTmV1cm9ueE5hdGl2ZUNvbXBpbGVyIGV4dGVuZHMgTmV1cm9ueENvbXBpbGVyQmFzZSB7XG4gIGNvbnN0cnVjdG9yKHNjb3BlOiBDb25zdHJ1Y3QsIGlkOiBzdHJpbmcsIHByb3BzOiBOZXVyb254TmF0aXZlQ29tcGlsZXJQcm9wcykge1xuICAgIHN1cGVyKHNjb3BlLCBpZCwgcHJvcHMpO1xuICB9XG5cbiAgcHJvdGVjdGVkIGNyZWF0ZUNvbXB1dGVFbnZpcm9ubWVudChcbiAgICBwcm9wczogTmV1cm9ueENvbXBpbGVyQmFzZVByb3BzLFxuICApOiBDb21wdXRlRW52aXJvbm1lbnRSZXN1bHQge1xuICAgIGNvbnN0IHZvbHVtZVNpemUgPVxuICAgICAgcHJvcHMudm9sdW1lU2l6ZT8udG9HaWJpYnl0ZXMoKSA/P1xuICAgICAgTWF0aC5jZWlsKFxuICAgICAgICB0aGlzLndlaWdodFNpemUudG9HaWJpYnl0ZXMoKSArXG4gICAgICAgICAgUHl0b3JjaFRyYWluaW5nTmV1cm9ueEltYWdlLnNpemUudG9HaWJpYnl0ZXMoKSArXG4gICAgICAgICAgTmV1cm9uT3B0aW1pemVkTWFjaGluZUltYWdlLnNpemUudG9HaWJpYnl0ZXMoKSxcbiAgICAgICk7XG4gICAgY29uc3QgbGF1bmNoVGVtcGxhdGUgPSBuZXcgZWMyLkxhdW5jaFRlbXBsYXRlKHRoaXMsIFwiTGF1bmNoVGVtcGxhdGVcIiwge1xuICAgICAgYmxvY2tEZXZpY2VzOiBbXG4gICAgICAgIHtcbiAgICAgICAgICBkZXZpY2VOYW1lOiBcIi9kZXYveHZkYVwiLFxuICAgICAgICAgIHZvbHVtZTogZWMyLkJsb2NrRGV2aWNlVm9sdW1lLmVicyh2b2x1bWVTaXplLCB7XG4gICAgICAgICAgICB2b2x1bWVUeXBlOiBlYzIuRWJzRGV2aWNlVm9sdW1lVHlwZS5HUDMsXG4gICAgICAgICAgICBlbmNyeXB0ZWQ6IHRydWUsXG4gICAgICAgICAgfSksXG4gICAgICAgIH0sXG4gICAgICBdLFxuICAgIH0pO1xuXG4gICAgY29uc3QgbmV1cm9ueEluc3RhbmNlVHlwZSA9XG4gICAgICBwcm9wcy5uZXVyb254SW5zdGFuY2VUeXBlID8/IE5ldXJvbnhJbnN0YW5jZVR5cGUuSU5GMl80OFhMQVJHRTtcbiAgICBjb25zdCBjb21wdXRlRW52aXJvbm1lbnQgPSBuZXcgTmV1cm9ueEJhdGNoQ29tcHV0ZUVudmlyb25tZW50KFxuICAgICAgdGhpcyxcbiAgICAgIFwiQ29tcHV0ZUVudmlyb25tZW50XCIsXG4gICAgICB7XG4gICAgICAgIHZwYzogcHJvcHMudnBjLFxuICAgICAgICB2cGNTdWJuZXRzOiBwcm9wcy52cGNTdWJuZXRzLFxuICAgICAgICBpbnN0YW5jZVR5cGVzOiBbbmV1cm9ueEluc3RhbmNlVHlwZS5pbnN0YW5jZVR5cGVdLFxuICAgICAgICB1c2VPcHRpbWFsSW5zdGFuY2VDbGFzc2VzOiBmYWxzZSxcbiAgICAgICAgbGF1bmNoVGVtcGxhdGUsXG4gICAgICAgIHNwb3Q6IHByb3BzLnNwb3QsXG4gICAgICB9LFxuICAgICk7XG4gICAgVGFncy5vZihjb21wdXRlRW52aXJvbm1lbnQpLmFkZChcIk5hbWVcIiwgXCJuZXVyb254LWNvbXBpbGUtd29ya2VyXCIpO1xuXG4gICAgcmV0dXJuIHtcbiAgICAgIGNvbXB1dGVFbnZpcm9ubWVudCxcbiAgICAgIGluc3RhbmNlUm9sZTogY29tcHV0ZUVudmlyb25tZW50Lmluc3RhbmNlUm9sZSxcbiAgICB9O1xuICB9XG5cbiAgcHJvdGVjdGVkIGNyZWF0ZUpvYkRlZmluaXRpb24oXG4gICAgcHJvcHM6IE5ldXJvbnhDb21waWxlckJhc2VQcm9wcyxcbiAgKTogYmF0Y2guSUpvYkRlZmluaXRpb24ge1xuICAgIGNvbnN0IG5ldXJvbnhJbnN0YW5jZVR5cGUgPVxuICAgICAgcHJvcHMubmV1cm9ueEluc3RhbmNlVHlwZSA/PyBOZXVyb254SW5zdGFuY2VUeXBlLklORjJfNDhYTEFSR0U7XG4gICAgcmV0dXJuIG5ldyBOZXVyb254QmF0Y2hFY3NKb2JEZWZpbml0aW9uKHRoaXMsIFwiSm9iRGVmaW5pdGlvblwiLCB7XG4gICAgICBuZXVyb254SW5zdGFuY2VUeXBlLFxuICAgICAgaW1hZ2U6IHByb3BzLmltYWdlLmltYWdlLFxuICAgICAgbWVtb3J5OiBTaXplLm1lYmlieXRlcyhcbiAgICAgICAgTWF0aC5jZWlsKG5ldXJvbnhJbnN0YW5jZVR5cGUubWVtb3J5LnRvTWViaWJ5dGVzKCkgKiAwLjk1KSxcbiAgICAgICksXG4gICAgICBjcHU6IG5ldXJvbnhJbnN0YW5jZVR5cGUudkNwdSxcbiAgICAgIGVudmlyb25tZW50OiB7XG4gICAgICAgIE5FVVJPTl9DT01QSUxFX0NBQ0hFX1VSTDogYCR7cHJvcHMuYnVja2V0LnMzVXJsRm9yT2JqZWN0KFwibmV1cm9uLWNvbXBpbGUtY2FjaGVcIil9YCxcbiAgICAgICAgLi4ucHJvcHMuZW52aXJvbm1lbnQsXG4gICAgICB9LFxuICAgICAgY29tbWFuZDogcHJvcHMuY29tbWFuZCxcbiAgICAgIHNlY3JldHM6IHByb3BzLnNlY3JldHMsXG4gICAgfSk7XG4gIH1cbn1cblxuZXhwb3J0IHsgSU5ldXJvbnhDb250YWluZXJJbWFnZSwgTmV1cm9ueENvbXBpbGVkTW9kZWwgfTtcbiJdfQ==