aws-cdk-neuronx-patterns
Version:
> [!WARNING] > This library is experimental module.
83 lines • 12.6 kB
JavaScript
;
var _a;
Object.defineProperty(exports, "__esModule", { value: true });
exports.NeuronxCrossCompiler = void 0;
const JSII_RTTI_SYMBOL_1 = Symbol.for("jsii.rtti");
const aws_cdk_lib_1 = require("aws-cdk-lib");
const batch = require("aws-cdk-lib/aws-batch");
const ec2 = require("aws-cdk-lib/aws-ec2");
const neuronx_1 = require("../neuronx");
const neuronx_compiler_base_1 = require("./neuronx-compiler-base");
/**
* Neuronx cross-compiler construct.
* Compile the model on a non-Neuron instance and upload the artifacts to an S3 bucket.
* This avoids the need for expensive Neuron instances during the compilation phase.
*
* The compilation uses `vllm serve` which performs model tracing and neuronx-cc compilation
* entirely on CPU. The resulting artifacts are compatible with Neuron instances for inference.
*/
class NeuronxCrossCompiler 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 compileInstanceType = props.compileInstanceType ??
ec2.InstanceType.of(ec2.InstanceClass.C7I, ec2.InstanceSize.XLARGE4);
const computeEnvironment = new batch.ManagedEc2EcsComputeEnvironment(this, "ComputeEnvironment", {
vpc: props.vpc,
vpcSubnets: props.vpcSubnets,
instanceTypes: [compileInstanceType],
useOptimalInstanceClasses: false,
launchTemplate,
spot: props.spot,
});
if (launchTemplate instanceof ec2.LaunchTemplate) {
computeEnvironment.node.defaultChild.addPropertyOverride("ComputeResources.LaunchTemplate.Version", launchTemplate.latestVersionNumber);
}
aws_cdk_lib_1.Tags.of(computeEnvironment).add("Name", "neuronx-cross-compile-worker");
return {
computeEnvironment,
instanceRole: computeEnvironment.instanceRole,
};
}
createJobDefinition(props) {
const neuronxInstanceType = props.neuronxInstanceType;
const targetPlatform = neuronxInstanceType.instanceType
.toString()
.split(".")[0];
const jobDefinition = new batch.EcsJobDefinition(this, "JobDefinition", {
container: new batch.EcsEc2ContainerDefinition(this, "ContainerDefinition", {
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")}`,
NEURON_PLATFORM_TARGET_OVERRIDE: targetPlatform,
...props.environment,
},
command: props.command,
secrets: props.secrets,
}),
});
return jobDefinition;
}
}
exports.NeuronxCrossCompiler = NeuronxCrossCompiler;
_a = JSII_RTTI_SYMBOL_1;
NeuronxCrossCompiler[_a] = { fqn: "aws-cdk-neuronx-patterns.NeuronxCrossCompiler", version: "0.3.0" };
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoibmV1cm9ueC1jcm9zcy1jb21waWxlci5qcyIsInNvdXJjZVJvb3QiOiIiLCJzb3VyY2VzIjpbIi4uLy4uLy4uL3NyYy9iYXNlL25ldXJvbngtY29tcGlsZXIvbmV1cm9ueC1jcm9zcy1jb21waWxlci50cyJdLCJuYW1lcyI6W10sIm1hcHBpbmdzIjoiOzs7OztBQUFBLDZDQUF5QztBQUN6QywrQ0FBK0M7QUFDL0MsMkNBQTJDO0FBRTNDLHdDQUdvQjtBQUNwQixtRUFJaUM7QUFnQmpDOzs7Ozs7O0dBT0c7QUFDSCxNQUFhLG9CQUFxQixTQUFRLDJDQUFtQjtJQUMzRCxZQUFZLEtBQWdCLEVBQUUsRUFBVSxFQUFFLEtBQWdDO1FBQ3hFLEtBQUssQ0FBQyxLQUFLLEVBQUUsRUFBRSxFQUFFLEtBQUssQ0FBQyxDQUFDO0lBQzFCLENBQUM7SUFFUyx3QkFBd0IsQ0FDaEMsS0FBK0I7UUFFL0IsTUFBTSxVQUFVLEdBQ2QsS0FBSyxDQUFDLFVBQVUsRUFBRSxXQUFXLEVBQUU7WUFDL0IsSUFBSSxDQUFDLElBQUksQ0FDUCxJQUFJLENBQUMsVUFBVSxDQUFDLFdBQVcsRUFBRTtnQkFDM0IscUNBQTJCLENBQUMsSUFBSSxDQUFDLFdBQVcsRUFBRTtnQkFDOUMscUNBQTJCLENBQUMsSUFBSSxDQUFDLFdBQVcsRUFBRSxDQUNqRCxDQUFDO1FBQ0osTUFBTSxjQUFjLEdBQUcsSUFBSSxHQUFHLENBQUMsY0FBYyxDQUFDLElBQUksRUFBRSxnQkFBZ0IsRUFBRTtZQUNwRSxZQUFZLEVBQUU7Z0JBQ1o7b0JBQ0UsVUFBVSxFQUFFLFdBQVc7b0JBQ3ZCLE1BQU0sRUFBRSxHQUFHLENBQUMsaUJBQWlCLENBQUMsR0FBRyxDQUFDLFVBQVUsRUFBRTt3QkFDNUMsVUFBVSxFQUFFLEdBQUcsQ0FBQyxtQkFBbUIsQ0FBQyxHQUFHO3dCQUN2QyxTQUFTLEVBQUUsSUFBSTtxQkFDaEIsQ0FBQztpQkFDSDthQUNGO1NBQ0YsQ0FBQyxDQUFDO1FBRUgsTUFBTSxtQkFBbUIsR0FDdEIsS0FBbUMsQ0FBQyxtQkFBbUI7WUFDeEQsR0FBRyxDQUFDLFlBQVksQ0FBQyxFQUFFLENBQUMsR0FBRyxDQUFDLGFBQWEsQ0FBQyxHQUFHLEVBQUUsR0FBRyxDQUFDLFlBQVksQ0FBQyxPQUFPLENBQUMsQ0FBQztRQUV2RSxNQUFNLGtCQUFrQixHQUFHLElBQUksS0FBSyxDQUFDLCtCQUErQixDQUNsRSxJQUFJLEVBQ0osb0JBQW9CLEVBQ3BCO1lBQ0UsR0FBRyxFQUFFLEtBQUssQ0FBQyxHQUFHO1lBQ2QsVUFBVSxFQUFFLEtBQUssQ0FBQyxVQUFVO1lBQzVCLGFBQWEsRUFBRSxDQUFDLG1CQUFtQixDQUFDO1lBQ3BDLHlCQUF5QixFQUFFLEtBQUs7WUFDaEMsY0FBYztZQUNkLElBQUksRUFBRSxLQUFLLENBQUMsSUFBSTtTQUNqQixDQUNGLENBQUM7UUFDRixJQUFJLGNBQWMsWUFBWSxHQUFHLENBQUMsY0FBYyxFQUFFLENBQUM7WUFFL0Msa0JBQWtCLENBQUMsSUFBSSxDQUFDLFlBQ3pCLENBQUMsbUJBQW1CLENBQ25CLHlDQUF5QyxFQUN6QyxjQUFjLENBQUMsbUJBQW1CLENBQ25DLENBQUM7UUFDSixDQUFDO1FBRUQsa0JBQUksQ0FBQyxFQUFFLENBQUMsa0JBQWtCLENBQUMsQ0FBQyxHQUFHLENBQUMsTUFBTSxFQUFFLDhCQUE4QixDQUFDLENBQUM7UUFFeEUsT0FBTztZQUNMLGtCQUFrQjtZQUNsQixZQUFZLEVBQUUsa0JBQWtCLENBQUMsWUFBYTtTQUMvQyxDQUFDO0lBQ0osQ0FBQztJQUVTLG1CQUFtQixDQUMzQixLQUErQjtRQUUvQixNQUFNLG1CQUFtQixHQUFHLEtBQUssQ0FBQyxtQkFBbUIsQ0FBQztRQUN0RCxNQUFNLGNBQWMsR0FBRyxtQkFBbUIsQ0FBQyxZQUFZO2FBQ3BELFFBQVEsRUFBRTthQUNWLEtBQUssQ0FBQyxHQUFHLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQztRQUVqQixNQUFNLGFBQWEsR0FBRyxJQUFJLEtBQUssQ0FBQyxnQkFBZ0IsQ0FBQyxJQUFJLEVBQUUsZUFBZSxFQUFFO1lBQ3RFLFNBQVMsRUFBRSxJQUFJLEtBQUssQ0FBQyx5QkFBeUIsQ0FDNUMsSUFBSSxFQUNKLHFCQUFxQixFQUNyQjtnQkFDRSxLQUFLLEVBQUUsS0FBSyxDQUFDLEtBQUssQ0FBQyxLQUFLO2dCQUN4QixNQUFNLEVBQUUsa0JBQUksQ0FBQyxTQUFTLENBQ3BCLElBQUksQ0FBQyxJQUFJLENBQUMsbUJBQW1CLENBQUMsTUFBTSxDQUFDLFdBQVcsRUFBRSxHQUFHLElBQUksQ0FBQyxDQUMzRDtnQkFDRCxHQUFHLEVBQUUsbUJBQW1CLENBQUMsSUFBSTtnQkFDN0IsV0FBVyxFQUFFO29CQUNYLHdCQUF3QixFQUFFLEdBQUcsS0FBSyxDQUFDLE1BQU0sQ0FBQyxjQUFjLENBQUMsc0JBQXNCLENBQUMsRUFBRTtvQkFDbEYsK0JBQStCLEVBQUUsY0FBYztvQkFDL0MsR0FBRyxLQUFLLENBQUMsV0FBVztpQkFDckI7Z0JBQ0QsT0FBTyxFQUFFLEtBQUssQ0FBQyxPQUFPO2dCQUN0QixPQUFPLEVBQUUsS0FBSyxDQUFDLE9BQU87YUFDdkIsQ0FDRjtTQUNGLENBQUMsQ0FBQztRQUVILE9BQU8sYUFBYSxDQUFDO0lBQ3ZCLENBQUM7O0FBMUZILG9EQTJGQyIsInNvdXJjZXNDb250ZW50IjpbImltcG9ydCB7IFNpemUsIFRhZ3MgfSBmcm9tIFwiYXdzLWNkay1saWJcIjtcbmltcG9ydCAqIGFzIGJhdGNoIGZyb20gXCJhd3MtY2RrLWxpYi9hd3MtYmF0Y2hcIjtcbmltcG9ydCAqIGFzIGVjMiBmcm9tIFwiYXdzLWNkay1saWIvYXdzLWVjMlwiO1xuaW1wb3J0IHsgQ29uc3RydWN0IH0gZnJvbSBcImNvbnN0cnVjdHNcIjtcbmltcG9ydCB7XG4gIE5ldXJvbk9wdGltaXplZE1hY2hpbmVJbWFnZSxcbiAgUHl0b3JjaFRyYWluaW5nTmV1cm9ueEltYWdlLFxufSBmcm9tIFwiLi4vbmV1cm9ueFwiO1xuaW1wb3J0IHtcbiAgTmV1cm9ueENvbXBpbGVyQmFzZSxcbiAgTmV1cm9ueENvbXBpbGVyQmFzZVByb3BzLFxuICBDb21wdXRlRW52aXJvbm1lbnRSZXN1bHQsXG59IGZyb20gXCIuL25ldXJvbngtY29tcGlsZXItYmFzZVwiO1xuXG4vKipcbiAqIFByb3BzIG9mIE5ldXJvbnhDcm9zc0NvbXBpbGVyLlxuICovXG5leHBvcnQgaW50ZXJmYWNlIE5ldXJvbnhDcm9zc0NvbXBpbGVyUHJvcHMgZXh0ZW5kcyBOZXVyb254Q29tcGlsZXJCYXNlUHJvcHMge1xuICAvKipcbiAgICogVGhlIEVDMiBpbnN0YW5jZSB0eXBlIHRvIHVzZSBmb3IgY3Jvc3MtY29tcGlsYXRpb24uXG4gICAqIFRoaXMgc2hvdWxkIGJlIGEgbm9uLU5ldXJvbiBpbnN0YW5jZSB0eXBlIHdpdGggc3VmZmljaWVudCBtZW1vcnkgYW5kIENQVVxuICAgKiBmb3IgbW9kZWwgY29tcGlsYXRpb24uXG4gICAqXG4gICAqIEBkZWZhdWx0IGVjMi5JbnN0YW5jZVR5cGUub2YoZWMyLkluc3RhbmNlQ2xhc3MuQzdJLCBlYzIuSW5zdGFuY2VTaXplLlhMQVJHRTQpXG4gICAqL1xuICByZWFkb25seSBjb21waWxlSW5zdGFuY2VUeXBlPzogZWMyLkluc3RhbmNlVHlwZTtcbn1cblxuLyoqXG4gKiBOZXVyb254IGNyb3NzLWNvbXBpbGVyIGNvbnN0cnVjdC5cbiAqIENvbXBpbGUgdGhlIG1vZGVsIG9uIGEgbm9uLU5ldXJvbiBpbnN0YW5jZSBhbmQgdXBsb2FkIHRoZSBhcnRpZmFjdHMgdG8gYW4gUzMgYnVja2V0LlxuICogVGhpcyBhdm9pZHMgdGhlIG5lZWQgZm9yIGV4cGVuc2l2ZSBOZXVyb24gaW5zdGFuY2VzIGR1cmluZyB0aGUgY29tcGlsYXRpb24gcGhhc2UuXG4gKlxuICogVGhlIGNvbXBpbGF0aW9uIHVzZXMgYHZsbG0gc2VydmVgIHdoaWNoIHBlcmZvcm1zIG1vZGVsIHRyYWNpbmcgYW5kIG5ldXJvbngtY2MgY29tcGlsYXRpb25cbiAqIGVudGlyZWx5IG9uIENQVS4gVGhlIHJlc3VsdGluZyBhcnRpZmFjdHMgYXJlIGNvbXBhdGlibGUgd2l0aCBOZXVyb24gaW5zdGFuY2VzIGZvciBpbmZlcmVuY2UuXG4gKi9cbmV4cG9ydCBjbGFzcyBOZXVyb254Q3Jvc3NDb21waWxlciBleHRlbmRzIE5ldXJvbnhDb21waWxlckJhc2Uge1xuICBjb25zdHJ1Y3RvcihzY29wZTogQ29uc3RydWN0LCBpZDogc3RyaW5nLCBwcm9wczogTmV1cm9ueENyb3NzQ29tcGlsZXJQcm9wcykge1xuICAgIHN1cGVyKHNjb3BlLCBpZCwgcHJvcHMpO1xuICB9XG5cbiAgcHJvdGVjdGVkIGNyZWF0ZUNvbXB1dGVFbnZpcm9ubWVudChcbiAgICBwcm9wczogTmV1cm9ueENvbXBpbGVyQmFzZVByb3BzLFxuICApOiBDb21wdXRlRW52aXJvbm1lbnRSZXN1bHQge1xuICAgIGNvbnN0IHZvbHVtZVNpemUgPVxuICAgICAgcHJvcHMudm9sdW1lU2l6ZT8udG9HaWJpYnl0ZXMoKSA/P1xuICAgICAgTWF0aC5jZWlsKFxuICAgICAgICB0aGlzLndlaWdodFNpemUudG9HaWJpYnl0ZXMoKSArXG4gICAgICAgICAgUHl0b3JjaFRyYWluaW5nTmV1cm9ueEltYWdlLnNpemUudG9HaWJpYnl0ZXMoKSArXG4gICAgICAgICAgTmV1cm9uT3B0aW1pemVkTWFjaGluZUltYWdlLnNpemUudG9HaWJpYnl0ZXMoKSxcbiAgICAgICk7XG4gICAgY29uc3QgbGF1bmNoVGVtcGxhdGUgPSBuZXcgZWMyLkxhdW5jaFRlbXBsYXRlKHRoaXMsIFwiTGF1bmNoVGVtcGxhdGVcIiwge1xuICAgICAgYmxvY2tEZXZpY2VzOiBbXG4gICAgICAgIHtcbiAgICAgICAgICBkZXZpY2VOYW1lOiBcIi9kZXYveHZkYVwiLFxuICAgICAgICAgIHZvbHVtZTogZWMyLkJsb2NrRGV2aWNlVm9sdW1lLmVicyh2b2x1bWVTaXplLCB7XG4gICAgICAgICAgICB2b2x1bWVUeXBlOiBlYzIuRWJzRGV2aWNlVm9sdW1lVHlwZS5HUDMsXG4gICAgICAgICAgICBlbmNyeXB0ZWQ6IHRydWUsXG4gICAgICAgICAgfSksXG4gICAgICAgIH0sXG4gICAgICBdLFxuICAgIH0pO1xuXG4gICAgY29uc3QgY29tcGlsZUluc3RhbmNlVHlwZSA9XG4gICAgICAocHJvcHMgYXMgTmV1cm9ueENyb3NzQ29tcGlsZXJQcm9wcykuY29tcGlsZUluc3RhbmNlVHlwZSA/P1xuICAgICAgZWMyLkluc3RhbmNlVHlwZS5vZihlYzIuSW5zdGFuY2VDbGFzcy5DN0ksIGVjMi5JbnN0YW5jZVNpemUuWExBUkdFNCk7XG5cbiAgICBjb25zdCBjb21wdXRlRW52aXJvbm1lbnQgPSBuZXcgYmF0Y2guTWFuYWdlZEVjMkVjc0NvbXB1dGVFbnZpcm9ubWVudChcbiAgICAgIHRoaXMsXG4gICAgICBcIkNvbXB1dGVFbnZpcm9ubWVudFwiLFxuICAgICAge1xuICAgICAgICB2cGM6IHByb3BzLnZwYyxcbiAgICAgICAgdnBjU3VibmV0czogcHJvcHMudnBjU3VibmV0cyxcbiAgICAgICAgaW5zdGFuY2VUeXBlczogW2NvbXBpbGVJbnN0YW5jZVR5cGVdLFxuICAgICAgICB1c2VPcHRpbWFsSW5zdGFuY2VDbGFzc2VzOiBmYWxzZSxcbiAgICAgICAgbGF1bmNoVGVtcGxhdGUsXG4gICAgICAgIHNwb3Q6IHByb3BzLnNwb3QsXG4gICAgICB9LFxuICAgICk7XG4gICAgaWYgKGxhdW5jaFRlbXBsYXRlIGluc3RhbmNlb2YgZWMyLkxhdW5jaFRlbXBsYXRlKSB7XG4gICAgICAoXG4gICAgICAgIGNvbXB1dGVFbnZpcm9ubWVudC5ub2RlLmRlZmF1bHRDaGlsZCBhcyBiYXRjaC5DZm5Db21wdXRlRW52aXJvbm1lbnRcbiAgICAgICkuYWRkUHJvcGVydHlPdmVycmlkZShcbiAgICAgICAgXCJDb21wdXRlUmVzb3VyY2VzLkxhdW5jaFRlbXBsYXRlLlZlcnNpb25cIixcbiAgICAgICAgbGF1bmNoVGVtcGxhdGUubGF0ZXN0VmVyc2lvbk51bWJlcixcbiAgICAgICk7XG4gICAgfVxuXG4gICAgVGFncy5vZihjb21wdXRlRW52aXJvbm1lbnQpLmFkZChcIk5hbWVcIiwgXCJuZXVyb254LWNyb3NzLWNvbXBpbGUtd29ya2VyXCIpO1xuXG4gICAgcmV0dXJuIHtcbiAgICAgIGNvbXB1dGVFbnZpcm9ubWVudCxcbiAgICAgIGluc3RhbmNlUm9sZTogY29tcHV0ZUVudmlyb25tZW50Lmluc3RhbmNlUm9sZSEsXG4gICAgfTtcbiAgfVxuXG4gIHByb3RlY3RlZCBjcmVhdGVKb2JEZWZpbml0aW9uKFxuICAgIHByb3BzOiBOZXVyb254Q29tcGlsZXJCYXNlUHJvcHMsXG4gICk6IGJhdGNoLklKb2JEZWZpbml0aW9uIHtcbiAgICBjb25zdCBuZXVyb254SW5zdGFuY2VUeXBlID0gcHJvcHMubmV1cm9ueEluc3RhbmNlVHlwZTtcbiAgICBjb25zdCB0YXJnZXRQbGF0Zm9ybSA9IG5ldXJvbnhJbnN0YW5jZVR5cGUuaW5zdGFuY2VUeXBlXG4gICAgICAudG9TdHJpbmcoKVxuICAgICAgLnNwbGl0KFwiLlwiKVswXTtcblxuICAgIGNvbnN0IGpvYkRlZmluaXRpb24gPSBuZXcgYmF0Y2guRWNzSm9iRGVmaW5pdGlvbih0aGlzLCBcIkpvYkRlZmluaXRpb25cIiwge1xuICAgICAgY29udGFpbmVyOiBuZXcgYmF0Y2guRWNzRWMyQ29udGFpbmVyRGVmaW5pdGlvbihcbiAgICAgICAgdGhpcyxcbiAgICAgICAgXCJDb250YWluZXJEZWZpbml0aW9uXCIsXG4gICAgICAgIHtcbiAgICAgICAgICBpbWFnZTogcHJvcHMuaW1hZ2UuaW1hZ2UsXG4gICAgICAgICAgbWVtb3J5OiBTaXplLm1lYmlieXRlcyhcbiAgICAgICAgICAgIE1hdGguY2VpbChuZXVyb254SW5zdGFuY2VUeXBlLm1lbW9yeS50b01lYmlieXRlcygpICogMC45NSksXG4gICAgICAgICAgKSxcbiAgICAgICAgICBjcHU6IG5ldXJvbnhJbnN0YW5jZVR5cGUudkNwdSxcbiAgICAgICAgICBlbnZpcm9ubWVudDoge1xuICAgICAgICAgICAgTkVVUk9OX0NPTVBJTEVfQ0FDSEVfVVJMOiBgJHtwcm9wcy5idWNrZXQuczNVcmxGb3JPYmplY3QoXCJuZXVyb24tY29tcGlsZS1jYWNoZVwiKX1gLFxuICAgICAgICAgICAgTkVVUk9OX1BMQVRGT1JNX1RBUkdFVF9PVkVSUklERTogdGFyZ2V0UGxhdGZvcm0sXG4gICAgICAgICAgICAuLi5wcm9wcy5lbnZpcm9ubWVudCxcbiAgICAgICAgICB9LFxuICAgICAgICAgIGNvbW1hbmQ6IHByb3BzLmNvbW1hbmQsXG4gICAgICAgICAgc2VjcmV0czogcHJvcHMuc2VjcmV0cyxcbiAgICAgICAgfSxcbiAgICAgICksXG4gICAgfSk7XG5cbiAgICByZXR1cm4gam9iRGVmaW5pdGlvbjtcbiAgfVxufVxuIl19