aws_manager/ec2/plugins/
mod.rs

1pub mod scripts;
2
3use std::{
4    collections::HashSet,
5    io::{self, Error, ErrorKind},
6    str::FromStr,
7};
8
9use crate::ec2;
10use serde::{Deserialize, Serialize};
11
12/// Defines the Arch type.
13#[derive(
14    Deserialize, Serialize, std::clone::Clone, std::cmp::Eq, std::fmt::Debug, std::hash::Hash,
15)]
16pub enum Plugin {
17    #[serde(rename = "imds")]
18    Imds,
19    #[serde(rename = "provider-id")]
20    ProviderId,
21    #[serde(rename = "vercmp")]
22    Vercmp,
23    #[serde(rename = "setup-local-disks")]
24    SetupLocalDisks,
25    #[serde(rename = "mount-bpf-fs")]
26    MountBpfFs,
27
28    #[serde(rename = "time-sync")]
29    TimeSync,
30    #[serde(rename = "system-limit-bump")]
31    SystemLimitBump,
32    #[serde(rename = "aws-cli")]
33    AwsCli,
34    #[serde(rename = "ssm-agent")]
35    SsmAgent,
36    #[serde(rename = "cloudwatch-agent")]
37    CloudwatchAgent,
38
39    #[serde(rename = "static-volume-provisioner")]
40    StaticVolumeProvisioner,
41    #[serde(rename = "static-ip-provisioner")]
42    StaticIpProvisioner,
43
44    #[serde(rename = "anaconda")]
45    Anaconda,
46    #[serde(rename = "python")]
47    Python,
48
49    #[serde(rename = "rust")]
50    Rust,
51    #[serde(rename = "go")]
52    Go,
53
54    #[serde(rename = "docker")]
55    Docker,
56    #[serde(rename = "containerd")]
57    Containerd,
58    #[serde(rename = "runc")]
59    Runc,
60    #[serde(rename = "cni-plugins")]
61    CniPlugins,
62
63    #[serde(rename = "aws-cfn-helper")]
64    AwsCfnHelper,
65    #[serde(rename = "saml2aws")]
66    Saml2Aws,
67    #[serde(rename = "aws-iam-authenticator")]
68    AwsIamAuthenticator,
69    #[serde(rename = "ecr-credential-helper")]
70    EcrCredentialHelper,
71    #[serde(rename = "ecr-credential-provider")]
72    EcrCredentialProvider,
73
74    #[serde(rename = "kubelet")]
75    Kubelet,
76    #[serde(rename = "kubectl")]
77    Kubectl,
78    #[serde(rename = "helm")]
79    Helm,
80    #[serde(rename = "terraform")]
81    Terraform,
82
83    #[serde(rename = "ssh-key-with-email")]
84    SshKeyWithEmail,
85
86    /// ref. <https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/enhanced-networking-ena.html#ena-requirements>
87    /// ref. <https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/enhanced-networking-ena.html>
88    /// ref. <https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/enhanced-networking-ena.html#enhanced-networking-ena-ubuntu>
89    #[serde(rename = "ena")]
90    Ena,
91
92    #[serde(rename = "nvidia-driver")]
93    NvidiaDriver,
94    #[serde(rename = "nvidia-cuda-toolkit")]
95    NvidiaCudaToolkit,
96    #[serde(rename = "nvidia-container-toolkit")]
97    NvidiaContainerToolkit,
98
99    #[serde(rename = "amd-radeon-gpu-driver")]
100    AmdRadeonGpuDriver,
101
102    #[serde(rename = "protobuf-compiler")]
103    ProtobufCompiler,
104    #[serde(rename = "cmake")]
105    Cmake,
106    #[serde(rename = "gcc7")]
107    Gcc7,
108
109    #[serde(rename = "dev-bark")]
110    DevBark,
111    #[serde(rename = "dev-faiss-gpu")]
112    DevFaissGpu,
113
114    #[serde(rename = "eks-worker-node-ami-scratch")]
115    EksWorkerNodeAmiScratch,
116    #[serde(rename = "eks-worker-node-ami-reuse")]
117    EksWorkerNodeAmiReuse,
118
119    #[serde(rename = "ami-info")]
120    AmiInfo,
121    #[serde(rename = "cluster-info")]
122    ClusterInfo,
123
124    #[serde(rename = "post-init-script")]
125    PostInitScript,
126
127    #[serde(rename = "cleanup-image-packages")]
128    CleanupImagePackages,
129    #[serde(rename = "cleanup-image-tmp-dir")]
130    CleanupImageTmpDir,
131    #[serde(rename = "cleanup-image-aws-credentials")]
132    CleanupImageAwsCredentials,
133    #[serde(rename = "cleanup-image-ssh-keys")]
134    CleanupImageSshKeys,
135
136    Unknown(String),
137}
138
139impl std::convert::From<&str> for Plugin {
140    fn from(s: &str) -> Self {
141        match s {
142            "imds" => Plugin::Imds,
143            "provider-id" => Plugin::ProviderId,
144            "vercmp" => Plugin::Vercmp,
145            "setup-local-disks" => Plugin::SetupLocalDisks,
146            "mount-bpf-fs" => Plugin::MountBpfFs,
147            "time-sync" => Plugin::TimeSync,
148            "system-limit-bump" => Plugin::SystemLimitBump,
149            "aws-cli" => Plugin::AwsCli,
150            "ssm-agent" => Plugin::SsmAgent,
151            "cloudwatch-agent" => Plugin::CloudwatchAgent,
152            "static-volume-provisioner" => Plugin::StaticVolumeProvisioner,
153            "static-ip-provisioner" => Plugin::StaticIpProvisioner,
154            "anaconda" => Plugin::Anaconda,
155            "python" => Plugin::Python,
156            "rust" => Plugin::Rust,
157            "go" => Plugin::Go,
158            "docker" => Plugin::Docker,
159            "containerd" => Plugin::Containerd,
160            "runc" => Plugin::Runc,
161            "cni-plugins" => Plugin::CniPlugins,
162            "aws-cfn-helper" => Plugin::AwsCfnHelper,
163            "saml2aws" => Plugin::Saml2Aws,
164            "aws-iam-authenticator" => Plugin::AwsIamAuthenticator,
165            "ecr-credential-helper" => Plugin::EcrCredentialHelper,
166            "ecr-credential-provider" => Plugin::EcrCredentialProvider,
167            "kubelet" => Plugin::Kubelet,
168            "kubectl" => Plugin::Kubectl,
169            "helm" => Plugin::Helm,
170            "terraform" => Plugin::Terraform,
171            "ssh-key-with-email" => Plugin::SshKeyWithEmail,
172            "ena" => Plugin::Ena,
173            "nvidia-driver" => Plugin::NvidiaDriver,
174            "nvidia-cuda-toolkit" => Plugin::NvidiaCudaToolkit,
175            "nvidia-container-toolkit" => Plugin::NvidiaContainerToolkit,
176            "amd-radeon-gpu-driver" => Plugin::AmdRadeonGpuDriver,
177            "protobuf-compiler" => Plugin::ProtobufCompiler,
178            "cmake" => Plugin::Cmake,
179            "gcc7" => Plugin::Gcc7,
180            "dev-bark" => Plugin::DevBark,
181            "dev-faiss-gpu" => Plugin::DevFaissGpu,
182            "eks-worker-node-ami-scratch" => Plugin::EksWorkerNodeAmiScratch,
183            "eks-worker-node-ami-reuse" => Plugin::EksWorkerNodeAmiReuse,
184            "ami-info" => Plugin::AmiInfo,
185            "cluster-info" => Plugin::ClusterInfo,
186            "post-init-script" => Plugin::PostInitScript,
187            "cleanup-image-packages" => Plugin::CleanupImagePackages,
188            "cleanup-image-tmp-dir" => Plugin::CleanupImageTmpDir,
189            "cleanup-image-aws-credentials" => Plugin::CleanupImageAwsCredentials,
190            "cleanup-image-ssh-keys" => Plugin::CleanupImageSshKeys,
191            other => Plugin::Unknown(other.to_owned()),
192        }
193    }
194}
195
196impl std::str::FromStr for Plugin {
197    type Err = std::convert::Infallible;
198
199    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
200        Ok(Plugin::from(s))
201    }
202}
203
204impl Plugin {
205    /// Returns the `&str` value of the enum member.
206    pub fn as_str(&self) -> &str {
207        match self {
208            Plugin::Imds => "imds",
209            Plugin::ProviderId => "provider-id",
210            Plugin::Vercmp => "vercmp",
211            Plugin::SetupLocalDisks => "setup-local-disks",
212            Plugin::MountBpfFs => "mount-bpf-fs",
213            Plugin::TimeSync => "time-sync",
214            Plugin::SystemLimitBump => "system-limit-bump",
215            Plugin::AwsCli => "aws-cli",
216            Plugin::SsmAgent => "ssm-agent",
217            Plugin::CloudwatchAgent => "cloudwatch-agent",
218            Plugin::StaticVolumeProvisioner => "static-volume-provisioner",
219            Plugin::StaticIpProvisioner => "static-ip-provisioner",
220            Plugin::Anaconda => "anaconda",
221            Plugin::Python => "python",
222            Plugin::Rust => "rust",
223            Plugin::Go => "go",
224            Plugin::Docker => "docker",
225            Plugin::Containerd => "containerd",
226            Plugin::Runc => "runc",
227            Plugin::CniPlugins => "cni-plugins",
228            Plugin::AwsCfnHelper => "aws-cfn-helper",
229            Plugin::Saml2Aws => "saml2aws",
230            Plugin::AwsIamAuthenticator => "aws-iam-authenticator",
231            Plugin::EcrCredentialHelper => "ecr-credential-helper",
232            Plugin::EcrCredentialProvider => "ecr-credential-provider",
233            Plugin::Kubelet => "kubelet",
234            Plugin::Kubectl => "kubectl",
235            Plugin::Helm => "helm",
236            Plugin::Terraform => "terraform",
237            Plugin::SshKeyWithEmail => "ssh-key-with-email",
238            Plugin::Ena => "ena",
239            Plugin::NvidiaDriver => "nvidia-driver",
240            Plugin::NvidiaCudaToolkit => "nvidia-cuda-toolkit",
241            Plugin::NvidiaContainerToolkit => "nvidia-container-toolkit",
242            Plugin::AmdRadeonGpuDriver => "amd-radeon-gpu-driver",
243            Plugin::ProtobufCompiler => "protobuf-compiler",
244            Plugin::Cmake => "cmake",
245            Plugin::Gcc7 => "gcc7",
246            Plugin::DevBark => "dev-bark",
247            Plugin::DevFaissGpu => "dev-faiss-gpu",
248            Plugin::EksWorkerNodeAmiScratch => "eks-worker-node-ami-scratch",
249            Plugin::EksWorkerNodeAmiReuse => "eks-worker-node-ami-reuse",
250            Plugin::AmiInfo => "ami-info",
251            Plugin::ClusterInfo => "cluster-info",
252            Plugin::PostInitScript => "post-init-script",
253            Plugin::CleanupImagePackages => "cleanup-image-packages",
254            Plugin::CleanupImageTmpDir => "cleanup-image-tmp-dir",
255            Plugin::CleanupImageAwsCredentials => "cleanup-image-aws-credentials",
256            Plugin::CleanupImageSshKeys => "cleanup-image-ssh-keys",
257            Plugin::Unknown(s) => s.as_ref(),
258        }
259    }
260
261    /// Returns the ranking in the sort.
262    /// Must in order; volume mount first + ip provision
263    /// useful when we create a base instance for AMI.
264    pub fn rank(&self) -> u32 {
265        match self {
266            Plugin::Imds => 0,
267            Plugin::ProviderId => 1,
268            Plugin::Vercmp => 2,
269            Plugin::SetupLocalDisks => 3,
270            Plugin::MountBpfFs => 4,
271            Plugin::TimeSync => 5,
272            Plugin::SystemLimitBump => 6,
273            Plugin::AwsCli => 7,
274            Plugin::SsmAgent => 8,
275            Plugin::CloudwatchAgent => 9,
276
277            Plugin::StaticVolumeProvisioner => 20,
278            Plugin::StaticIpProvisioner => 21,
279
280            Plugin::Anaconda => 25,
281            Plugin::Python => 25,
282
283            Plugin::Rust => 26,
284            Plugin::Go => 27,
285            Plugin::Docker => 28,
286            Plugin::Containerd => 29,
287            Plugin::Runc => 30,
288            Plugin::CniPlugins => 31,
289
290            Plugin::AwsCfnHelper => 32,
291            Plugin::Saml2Aws => 33,
292
293            Plugin::AwsIamAuthenticator => 34,
294            Plugin::EcrCredentialHelper => 35,
295            Plugin::EcrCredentialProvider => 36,
296
297            Plugin::Kubelet => 37,
298            Plugin::Kubectl => 38,
299            Plugin::Helm => 50,
300            Plugin::Terraform => 51,
301
302            Plugin::SshKeyWithEmail => 68,
303
304            Plugin::Ena => 100,
305
306            Plugin::NvidiaDriver => 200,
307            Plugin::NvidiaCudaToolkit => 201,
308            Plugin::NvidiaContainerToolkit => 202,
309
310            Plugin::AmdRadeonGpuDriver => 300,
311
312            Plugin::ProtobufCompiler => 60000,
313            Plugin::Cmake => 60001,
314            Plugin::Gcc7 => 60002,
315
316            Plugin::DevBark => 80000,
317            Plugin::DevFaissGpu => 80001,
318
319            Plugin::EksWorkerNodeAmiScratch => 99990,
320            Plugin::EksWorkerNodeAmiReuse => 99991,
321
322            Plugin::AmiInfo => u32::MAX - 2000,
323            Plugin::ClusterInfo => u32::MAX - 1999,
324
325            Plugin::PostInitScript => u32::MAX - 1000,
326
327            Plugin::CleanupImagePackages => u32::MAX - 10,
328            Plugin::CleanupImageTmpDir => u32::MAX - 9,
329            Plugin::CleanupImageAwsCredentials => u32::MAX - 8,
330            Plugin::CleanupImageSshKeys => u32::MAX - 5,
331
332            Plugin::Unknown(_) => u32::MAX,
333        }
334    }
335
336    /// Returns all the `&str` values of the enum members.
337    pub fn values() -> &'static [&'static str] {
338        &[
339            "imds",                          //
340            "provider-id",                   //
341            "vercmp",                        //
342            "setup-local-disks",             //
343            "mount-bpf-fs",                  //
344            "time-sync",                     //
345            "system-limit-bump",             //
346            "aws-cli",                       //
347            "ssm-agent",                     //
348            "cloudwatch-agent",              //
349            "static-volume-provisioner",     //
350            "static-ip-provisioner",         //
351            "anaconda",                      //
352            "python",                        //
353            "rust",                          //
354            "go",                            //
355            "docker",                        //
356            "containerd",                    //
357            "runc",                          //
358            "cni-plugins",                   //
359            "protobuf-compiler",             //
360            "aws-cfn-helper",                //
361            "saml2aws",                      //
362            "aws-iam-authenticator",         //
363            "ecr-credential-helper",         //
364            "ecr-credential-provider",       //
365            "kubelet",                       //
366            "kubectl",                       //
367            "helm",                          //
368            "terraform",                     //
369            "ssh-key-with-email",            //
370            "ena",                           //
371            "nvidia-driver",                 //
372            "nvidia-cuda-toolkit",           //
373            "nvidia-container-toolkit",      //
374            "amd-radeon-gpu-driver",         //
375            "cmake",                         //
376            "gcc7",                          //
377            "dev-bark",                      //
378            "dev-faiss-gpu",                 //
379            "eks-worker-node-ami-scratch",   //
380            "eks-worker-node-ami-reuse",     //
381            "ami-info",                      //
382            "cluster-info",                  //
383            "post-init-script",              //
384            "cleanup-image-packages",        //
385            "cleanup-image-ssh-keys",        //
386            "cleanup-image-aws-credentials", //
387        ]
388    }
389
390    pub fn all() -> Vec<String> {
391        vec![
392            Plugin::Imds.as_str().to_string(),
393            Plugin::ProviderId.as_str().to_string(),
394            Plugin::Vercmp.as_str().to_string(),
395            Plugin::SetupLocalDisks.as_str().to_string(),
396            Plugin::MountBpfFs.as_str().to_string(),
397            Plugin::TimeSync.as_str().to_string(),
398            Plugin::SystemLimitBump.as_str().to_string(),
399            Plugin::AwsCli.as_str().to_string(),
400            Plugin::SsmAgent.as_str().to_string(),
401            Plugin::CloudwatchAgent.as_str().to_string(),
402            Plugin::StaticVolumeProvisioner.as_str().to_string(),
403            Plugin::StaticIpProvisioner.as_str().to_string(),
404            Plugin::Anaconda.as_str().to_string(),
405            Plugin::Python.as_str().to_string(),
406            Plugin::Rust.as_str().to_string(),
407            Plugin::Go.as_str().to_string(),
408            Plugin::Docker.as_str().to_string(),
409            Plugin::Containerd.as_str().to_string(),
410            Plugin::Runc.as_str().to_string(),
411            Plugin::CniPlugins.as_str().to_string(),
412            Plugin::AwsCfnHelper.as_str().to_string(),
413            Plugin::Saml2Aws.as_str().to_string(),
414            Plugin::AwsIamAuthenticator.as_str().to_string(),
415            Plugin::EcrCredentialHelper.as_str().to_string(),
416            Plugin::EcrCredentialProvider.as_str().to_string(),
417            Plugin::Kubelet.as_str().to_string(),
418            Plugin::Kubectl.as_str().to_string(),
419            Plugin::Helm.as_str().to_string(),
420            Plugin::Terraform.as_str().to_string(),
421            Plugin::SshKeyWithEmail.as_str().to_string(),
422            Plugin::Ena.as_str().to_string(),
423            Plugin::NvidiaDriver.as_str().to_string(),
424            Plugin::NvidiaCudaToolkit.as_str().to_string(),
425            Plugin::NvidiaContainerToolkit.as_str().to_string(),
426            Plugin::AmdRadeonGpuDriver.as_str().to_string(),
427            Plugin::ProtobufCompiler.as_str().to_string(),
428            Plugin::Cmake.as_str().to_string(),
429            Plugin::Gcc7.as_str().to_string(),
430            Plugin::DevBark.as_str().to_string(),
431            Plugin::DevFaissGpu.as_str().to_string(),
432            Plugin::EksWorkerNodeAmiScratch.as_str().to_string(),
433            Plugin::EksWorkerNodeAmiReuse.as_str().to_string(),
434            Plugin::AmiInfo.as_str().to_string(),
435            Plugin::ClusterInfo.as_str().to_string(),
436            Plugin::PostInitScript.as_str().to_string(),
437            Plugin::CleanupImagePackages.as_str().to_string(),
438            Plugin::CleanupImageTmpDir.as_str().to_string(),
439            Plugin::CleanupImageAwsCredentials.as_str().to_string(),
440            Plugin::CleanupImageSshKeys.as_str().to_string(),
441        ]
442    }
443
444    pub fn default() -> Vec<String> {
445        vec![
446            Plugin::Imds.as_str().to_string(),
447            Plugin::ProviderId.as_str().to_string(),
448            Plugin::Vercmp.as_str().to_string(),
449            Plugin::SetupLocalDisks.as_str().to_string(),
450            Plugin::MountBpfFs.as_str().to_string(),
451            Plugin::TimeSync.as_str().to_string(),
452            Plugin::SystemLimitBump.as_str().to_string(),
453            Plugin::SsmAgent.as_str().to_string(),
454            Plugin::CloudwatchAgent.as_str().to_string(),
455            Plugin::Anaconda.as_str().to_string(),
456            Plugin::AwsCfnHelper.as_str().to_string(),
457        ]
458    }
459}
460
461impl AsRef<str> for Plugin {
462    fn as_ref(&self) -> &str {
463        self.as_str()
464    }
465}
466
467/// To be printed/echoed before clean up process (e.g., before cleaning up SSH keys).
468pub const INIT_SCRIPT_COMPLETE_MSG: &str = "INIT SCRIPT COMPLETE";
469
470pub fn create(
471    arch_type: ec2::ArchType,
472    os_type: ec2::OsType,
473    plugins_str: Vec<String>,
474    require_static_ip_provisioner: bool,
475    s3_bucket: &str,
476    id: &str,
477    region: &str,
478    volume_type: &str,
479    volume_size: u32,
480    volume_iops: u32,
481    volume_throughput: u32,
482    ssh_key_email: Option<String>,
483    aws_secret_key_id: Option<String>,
484    aws_secret_access_key: Option<String>,
485    post_init_script: Option<String>,
486) -> io::Result<(Vec<Plugin>, String)> {
487    let mut plugins_set = HashSet::new();
488    for p in plugins_str.iter() {
489        let plugin = Plugin::from_str(p).map_err(|e| {
490            Error::new(
491                ErrorKind::InvalidInput,
492                format!("failed to convert '{p}' to plugin {}", e),
493            )
494        })?;
495        plugins_set.insert(plugin);
496    }
497
498    if plugins_set.contains(&Plugin::StaticVolumeProvisioner) {
499        if require_static_ip_provisioner {
500            plugins_set.insert(Plugin::StaticIpProvisioner);
501        }
502    }
503    // pick either anaconda or python
504    if plugins_set.contains(&Plugin::Anaconda) {
505        if plugins_set.contains(&Plugin::Python) {
506            log::info!("anaconda specifies thus overriding python plugin");
507            plugins_set.remove(&Plugin::Python);
508        }
509    } else if plugins_set.contains(&Plugin::Python) {
510        log::info!("only python plugin, without anaconda");
511    }
512    if arch_type.is_nvidia() {
513        plugins_set.insert(Plugin::NvidiaDriver);
514    }
515
516    if plugins_set.contains(&Plugin::AwsCfnHelper) {
517        if !plugins_set.contains(&Plugin::Anaconda) && !plugins_set.contains(&Plugin::Python) {
518            return Err(Error::new(
519                ErrorKind::InvalidInput,
520                format!(
521                    "specified '{}' requires pip/python install plugin",
522                    Plugin::AwsCfnHelper.as_str()
523                ),
524            ));
525        }
526    }
527    if plugins_set.contains(&Plugin::EcrCredentialProvider) && !plugins_set.contains(&Plugin::Go) {
528        return Err(Error::new(
529            ErrorKind::InvalidInput,
530            format!(
531                "'{}' requires '{}' plugin",
532                Plugin::EcrCredentialProvider.as_str(),
533                Plugin::Go.as_str(),
534            ),
535        ));
536    }
537    if plugins_set.contains(&Plugin::TimeSync) && !plugins_set.contains(&Plugin::Imds) {
538        return Err(Error::new(
539            ErrorKind::InvalidInput,
540            format!(
541                "'{}' requires '{}' plugin",
542                Plugin::TimeSync.as_str(),
543                Plugin::Imds.as_str(),
544            ),
545        ));
546    }
547    if plugins_set.contains(&Plugin::Ena) && !plugins_set.contains(&Plugin::Imds) {
548        return Err(Error::new(
549            ErrorKind::InvalidInput,
550            format!(
551                "'{}' requires '{}' plugin",
552                Plugin::Ena.as_str(),
553                Plugin::Imds.as_str(),
554            ),
555        ));
556    }
557    if plugins_set.contains(&Plugin::NvidiaCudaToolkit) {
558        if !arch_type.is_nvidia() {
559            return Err(Error::new(
560                ErrorKind::InvalidInput,
561                format!(
562                    "specified '{}' but arch type is '{}' (not nvidia)",
563                    Plugin::NvidiaCudaToolkit.as_str(),
564                    arch_type.as_str()
565                ),
566            ));
567        }
568        if !plugins_set.contains(&Plugin::NvidiaDriver) {
569            return Err(Error::new(
570                ErrorKind::InvalidInput,
571                format!(
572                    "specified '{}' but no '{}'",
573                    Plugin::NvidiaCudaToolkit.as_str(),
574                    Plugin::NvidiaDriver.as_str()
575                ),
576            ));
577        }
578    }
579    if plugins_set.contains(&Plugin::NvidiaCudaToolkit) {
580        if !arch_type.is_nvidia() {
581            return Err(Error::new(
582                ErrorKind::InvalidInput,
583                format!(
584                    "specified '{}' but arch type is '{}' (not nvidia)",
585                    Plugin::NvidiaCudaToolkit.as_str(),
586                    arch_type.as_str()
587                ),
588            ));
589        }
590        if !plugins_set.contains(&Plugin::NvidiaDriver) {
591            return Err(Error::new(
592                ErrorKind::InvalidInput,
593                format!(
594                    "specified '{}' but no '{}'",
595                    Plugin::NvidiaContainerToolkit.as_str(),
596                    Plugin::NvidiaDriver.as_str()
597                ),
598            ));
599        }
600    }
601
602    if plugins_set.contains(&Plugin::DevBark) {
603        if !plugins_set.contains(&Plugin::StaticVolumeProvisioner) {
604            return Err(Error::new(
605                ErrorKind::InvalidInput,
606                format!(
607                    "specified '{}' requires '{}'",
608                    Plugin::DevBark.as_str(),
609                    Plugin::StaticVolumeProvisioner.as_str()
610                ),
611            ));
612        }
613    }
614    if plugins_set.contains(&Plugin::DevFaissGpu) {
615        if !plugins_set.contains(&Plugin::StaticVolumeProvisioner) {
616            return Err(Error::new(
617                ErrorKind::InvalidInput,
618                format!(
619                    "specified '{}' requires '{}'",
620                    Plugin::DevFaissGpu.as_str(),
621                    Plugin::StaticVolumeProvisioner.as_str()
622                ),
623            ));
624        }
625    }
626
627    if plugins_set.contains(&Plugin::EksWorkerNodeAmiScratch)
628        && plugins_set.contains(&Plugin::EksWorkerNodeAmiReuse)
629    {
630        return Err(Error::new(
631            ErrorKind::InvalidInput,
632            format!(
633                "'{}' conflicts with '{}'",
634                Plugin::EksWorkerNodeAmiScratch.as_str(),
635                Plugin::EksWorkerNodeAmiReuse.as_str()
636            ),
637        ));
638    }
639
640    if post_init_script.is_some() {
641        if !plugins_set.contains(&Plugin::PostInitScript) {
642            log::warn!(
643                "'{}' but no post init script specified -- adding to set",
644                Plugin::PostInitScript.as_str()
645            );
646        }
647        plugins_set.insert(Plugin::PostInitScript);
648    }
649
650    let mut plugins = Vec::new();
651    for p in plugins_set.iter() {
652        plugins.push(p.clone());
653    }
654    plugins.sort();
655
656    // TODO: make this configurable?
657    let provisioner_initial_wait_random_seconds = 10;
658
659    let mut contents = scripts::start(os_type.clone())?;
660    let mut updated_bash_profile = false;
661    for p in plugins.iter() {
662        match p {
663            Plugin::Imds => {
664                let d = scripts::imds(os_type.clone())?;
665                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
666                contents.push_str(&d);
667            }
668            Plugin::ProviderId => {
669                let d = scripts::provider_id(os_type.clone())?;
670                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
671                contents.push_str(&d);
672            }
673            Plugin::Vercmp => {
674                let d = scripts::vercmp(os_type.clone())?;
675                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
676                contents.push_str(&d);
677            }
678            Plugin::SetupLocalDisks => {
679                let d = scripts::setup_local_disks(os_type.clone())?;
680                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
681                contents.push_str(&d);
682            }
683            Plugin::MountBpfFs => {
684                let d = scripts::mount_bpf_fs(os_type.clone())?;
685                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
686                contents.push_str(&d);
687            }
688            Plugin::TimeSync => {
689                let d = scripts::time_sync(os_type.clone())?;
690                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
691                contents.push_str(&d);
692            }
693            Plugin::SystemLimitBump => {
694                let d = scripts::system_limit_bump(os_type.clone())?;
695                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
696                contents.push_str(&d);
697
698                if !updated_bash_profile {
699                    updated_bash_profile = true;
700
701                    let d = scripts::update_bash_profile(
702                        os_type.clone(),
703                        plugins_set.contains(&Plugin::Anaconda),
704                        plugins_set.contains(&Plugin::Python),
705                        plugins_set.contains(&Plugin::Rust),
706                        plugins_set.contains(&Plugin::NvidiaCudaToolkit),
707                        plugins_set.contains(&Plugin::Go),
708                        plugins_set.contains(&Plugin::Kubectl),
709                        plugins_set.contains(&Plugin::Helm),
710                        plugins_set.contains(&Plugin::StaticVolumeProvisioner),
711                    )?;
712                    contents.push_str(
713                        "###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n",
714                    );
715                    contents.push_str(&d);
716                }
717            }
718            Plugin::AwsCli => {
719                let d = scripts::aws_cli(arch_type.clone(), os_type.clone())?;
720                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
721                contents.push_str(&d);
722            }
723            Plugin::SsmAgent => {
724                let d = scripts::ssm_agent(os_type.clone())?;
725                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
726                contents.push_str(&d);
727            }
728            Plugin::CloudwatchAgent => {
729                let d = scripts::cloudwatch_agent(os_type.clone())?;
730                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
731                contents.push_str(&d);
732            }
733
734            Plugin::StaticVolumeProvisioner => {
735                // https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-ec2-launchtemplate-blockdevicemapping-ebs.html#cfn-ec2-launchtemplate-blockdevicemapping-ebs-volumetype
736                // https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-ec2-launchtemplate-blockdevicemapping-ebs.html#cfn-ec2-launchtemplate-blockdevicemapping-ebs-volumesize
737                // https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-ec2-launchtemplate-blockdevicemapping-ebs.html#cfn-ec2-launchtemplate-blockdevicemapping-ebs-iops
738                //
739                // only for gp3
740                // https://aws.amazon.com/ebs/volume-types/
741                // https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-ec2-launchtemplate-blockdevicemapping-ebs.html#cfn-ec2-launchtemplate-blockdevicemapping-ebs-throughput
742                // "1000" does not work -- "InvalidParameterValue - Throughput (MiBps) to iops ratio of 0.333333 is too high; maximum is 0.250000 MiBps per iops."
743                let d = scripts::static_volume_provisioner(
744                    os_type.clone(),
745                    id,
746                    region,
747                    volume_type,
748                    volume_size,
749                    volume_iops,
750                    volume_throughput,
751                    "/dev/xvdb",
752                    provisioner_initial_wait_random_seconds,
753                )?;
754                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
755                contents.push_str(&d);
756
757                if !updated_bash_profile {
758                    updated_bash_profile = true;
759
760                    let d = scripts::update_bash_profile(
761                        os_type.clone(),
762                        plugins_set.contains(&Plugin::Anaconda),
763                        plugins_set.contains(&Plugin::Python),
764                        plugins_set.contains(&Plugin::Rust),
765                        plugins_set.contains(&Plugin::NvidiaCudaToolkit),
766                        plugins_set.contains(&Plugin::Go),
767                        plugins_set.contains(&Plugin::Kubectl),
768                        plugins_set.contains(&Plugin::Helm),
769                        plugins_set.contains(&Plugin::StaticVolumeProvisioner),
770                    )?;
771                    contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
772                    contents.push_str(&d);
773                }
774            }
775            Plugin::StaticIpProvisioner => {
776                let d = scripts::static_ip_provisioner(
777                    os_type.clone(),
778                    id,
779                    region,
780                    provisioner_initial_wait_random_seconds,
781                )?;
782                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
783                contents.push_str(&d);
784            }
785
786            Plugin::Anaconda => {
787                let d = scripts::anaconda(os_type.clone())?;
788                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
789                contents.push_str(&d);
790            }
791            Plugin::Python => {
792                let d = scripts::python(os_type.clone())?;
793                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
794                contents.push_str(&d);
795            }
796
797            Plugin::Rust => {
798                let d = scripts::rust(os_type.clone())?;
799                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
800                contents.push_str(&d);
801            }
802            Plugin::Go => {
803                let d = scripts::go(os_type.clone())?;
804                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
805                contents.push_str(&d);
806            }
807
808            Plugin::Docker => {
809                let d = scripts::docker(os_type.clone())?;
810                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
811                contents.push_str(&d);
812            }
813            Plugin::Containerd => {
814                let d = scripts::containerd(os_type.clone())?;
815                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
816                contents.push_str(&d);
817            }
818            Plugin::Runc => {
819                let d = scripts::runc(os_type.clone())?;
820                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
821                contents.push_str(&d);
822            }
823            Plugin::CniPlugins => {
824                let d = scripts::cni_plugins(os_type.clone())?;
825                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
826                contents.push_str(&d);
827            }
828
829            Plugin::AwsCfnHelper => {
830                let d = scripts::aws_cfn_helper(
831                    os_type.clone(),
832                    if plugins_set.contains(&Plugin::Anaconda) {
833                        "/home/ubuntu/anaconda3/bin"
834                    } else {
835                        ""
836                    },
837                )?;
838                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
839                contents.push_str(&d);
840            }
841            Plugin::Saml2Aws => {
842                let d = scripts::saml2aws(os_type.clone())?;
843                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
844                contents.push_str(&d);
845            }
846            Plugin::AwsIamAuthenticator => {
847                let d = scripts::aws_iam_authenticator(os_type.clone())?;
848                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
849                contents.push_str(&d);
850            }
851            Plugin::EcrCredentialHelper => {
852                let d = scripts::ecr_credential_helper(os_type.clone())?;
853                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
854                contents.push_str(&d);
855            }
856            Plugin::EcrCredentialProvider => {
857                let d = scripts::ecr_credential_provider(os_type.clone())?;
858                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
859                contents.push_str(&d);
860            }
861
862            Plugin::Kubelet => {
863                let d = scripts::kubelet(os_type.clone())?;
864                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
865                contents.push_str(&d);
866            }
867            Plugin::Kubectl => {
868                let d = scripts::kubectl(os_type.clone())?;
869                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
870                contents.push_str(&d);
871            }
872            Plugin::Helm => {
873                let d = scripts::helm(os_type.clone())?;
874                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
875                contents.push_str(&d);
876            }
877            Plugin::Terraform => {
878                let d = scripts::terraform(os_type.clone())?;
879                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
880                contents.push_str(&d);
881            }
882
883            Plugin::SshKeyWithEmail => {
884                if ssh_key_email.is_none() {
885                    return Err(Error::new(
886                        ErrorKind::Other,
887                        format!(
888                            "plugin {} specified but empty email",
889                            Plugin::SshKeyWithEmail.as_str()
890                        ),
891                    ));
892                }
893                let d = scripts::ssh_key_with_email(
894                    os_type.clone(),
895                    ssh_key_email.clone().unwrap().as_str(),
896                )?;
897                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
898                contents.push_str(&d);
899            }
900
901            Plugin::Ena => {
902                let d = scripts::ena(os_type.clone())?;
903                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
904                contents.push_str(&d);
905            }
906
907            Plugin::NvidiaDriver => {
908                let d = scripts::nvidia_driver(arch_type.clone(), os_type.clone())?;
909                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
910                contents.push_str(&d);
911            }
912            Plugin::NvidiaCudaToolkit => {
913                let d = scripts::nvidia_cuda_toolkit(os_type.clone())?;
914                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
915                contents.push_str(&d);
916            }
917            Plugin::NvidiaContainerToolkit => {
918                let d = scripts::nvidia_container_toolkit(os_type.clone())?;
919                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
920                contents.push_str(&d);
921            }
922
923            Plugin::AmdRadeonGpuDriver => {
924                let d = scripts::amd_radeon_gpu_driver(arch_type.clone(), os_type.clone())?;
925                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
926                contents.push_str(&d);
927            }
928
929            Plugin::ProtobufCompiler => {
930                let d = scripts::protobuf_compiler(os_type.clone())?;
931                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
932                contents.push_str(&d);
933            }
934            Plugin::Cmake => {
935                let d = scripts::cmake(
936                    os_type.clone(),
937                    if plugins_set.contains(&Plugin::Anaconda) {
938                        "/home/ubuntu/anaconda3/bin"
939                    } else {
940                        ""
941                    },
942                )?;
943                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
944                contents.push_str(&d);
945            }
946            Plugin::Gcc7 => {
947                let d = scripts::gcc7(os_type.clone())?;
948                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
949                contents.push_str(&d);
950            }
951
952            Plugin::DevBark => {
953                let d = scripts::dev_bark(
954                    os_type.clone(),
955                    if plugins_set.contains(&Plugin::Anaconda) {
956                        "/home/ubuntu/anaconda3/bin"
957                    } else {
958                        ""
959                    },
960                    plugins_set.contains(&Plugin::StaticVolumeProvisioner),
961                )?;
962                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
963                contents.push_str(&d);
964            }
965            Plugin::DevFaissGpu => {
966                let d = scripts::dev_faiss_gpu(
967                    os_type.clone(),
968                    plugins_set.contains(&Plugin::StaticVolumeProvisioner),
969                )?;
970                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
971                contents.push_str(&d);
972            }
973
974            Plugin::EksWorkerNodeAmiScratch => {
975                let d = scripts::eks_worker_node_ami_scratch(os_type.clone())?;
976                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
977                contents.push_str(&d);
978            }
979            Plugin::EksWorkerNodeAmiReuse => {
980                let d = scripts::eks_worker_node_ami_reuse(os_type.clone())?;
981                contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
982                contents.push_str(&d);
983            }
984
985            Plugin::AmiInfo
986            | Plugin::ClusterInfo
987            | Plugin::PostInitScript
988            | Plugin::CleanupImagePackages
989            | Plugin::CleanupImageTmpDir
990            | Plugin::CleanupImageAwsCredentials
991            | Plugin::CleanupImageSshKeys => {
992                log::info!(
993                    "skipping {}, saving it to write after bash profile at the end",
994                    p.as_str()
995                )
996            }
997
998            _ => {
999                return Err(Error::new(
1000                    ErrorKind::Other,
1001                    format!("unknown plugin {}", p.as_str()),
1002                ))
1003            }
1004        }
1005    }
1006
1007    if !updated_bash_profile {
1008        let d = scripts::update_bash_profile(
1009            os_type.clone(),
1010            plugins_set.contains(&Plugin::Anaconda),
1011            plugins_set.contains(&Plugin::Python),
1012            plugins_set.contains(&Plugin::Rust),
1013            plugins_set.contains(&Plugin::NvidiaCudaToolkit),
1014            plugins_set.contains(&Plugin::Go),
1015            plugins_set.contains(&Plugin::Kubectl),
1016            plugins_set.contains(&Plugin::Helm),
1017            plugins_set.contains(&Plugin::StaticVolumeProvisioner),
1018        )?;
1019
1020        contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1021        contents.push_str(&d);
1022    }
1023
1024    if let Some(secret_key_id) = &aws_secret_key_id {
1025        let access_key = aws_secret_access_key.clone().unwrap();
1026        let d = scripts::aws_key(os_type.clone(), region, secret_key_id.as_str(), &access_key)?;
1027
1028        contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1029        contents.push_str(&d);
1030    }
1031
1032    if plugins_set.contains(&Plugin::AmiInfo) {
1033        let d = scripts::ami_info(os_type.clone())?;
1034
1035        contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1036        contents.push_str(&d);
1037    }
1038    if plugins_set.contains(&Plugin::ClusterInfo) {
1039        let d = scripts::cluster_info(
1040            s3_bucket,
1041            id,
1042            plugins_set.contains(&Plugin::StaticVolumeProvisioner),
1043        );
1044
1045        contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1046        contents.push_str(&d);
1047    }
1048
1049    if plugins_set.contains(&Plugin::PostInitScript) {
1050        contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1051        contents.push_str("###########################\n# USER-DEFINED POST INIT SCRIPT\n");
1052        contents.push_str(&post_init_script.unwrap());
1053    }
1054
1055    if plugins_set.contains(&Plugin::CleanupImagePackages) {
1056        let d = scripts::cleanup_image_packages(os_type.clone())?;
1057
1058        contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1059        contents.push_str(&d);
1060    }
1061    if plugins_set.contains(&Plugin::CleanupImageTmpDir) {
1062        let d = scripts::cleanup_image_tmp_dir(os_type.clone())?;
1063
1064        contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1065        contents.push_str(&d);
1066    }
1067
1068    if plugins_set.contains(&Plugin::CleanupImageAwsCredentials) {
1069        let d = scripts::cleanup_image_aws_credentials(os_type.clone())?;
1070
1071        contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1072        contents.push_str(&d);
1073    }
1074
1075    // do before the SSH key clean-ups
1076    contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1077    contents
1078        .push_str(format!("###########################\n# {INIT_SCRIPT_COMPLETE_MSG}\n").as_str());
1079    contents.push_str(format!("echo \"{INIT_SCRIPT_COMPLETE_MSG}\"\n\n").as_str());
1080
1081    if plugins_set.contains(&Plugin::CleanupImageSshKeys) {
1082        let d = scripts::cleanup_image_ssh_keys(os_type.clone())?;
1083
1084        contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1085        contents.push_str(&d);
1086    }
1087
1088    contents.push_str("###########################\nset +x\necho \"\"\necho \"\"\necho \"\"\necho \"\"\necho \"\"\nset -x\n\n\n\n\n");
1089    Ok((plugins, contents))
1090}
1091
1092pub fn to_strings(plugins: Vec<Plugin>) -> Vec<String> {
1093    let mut ss = Vec::new();
1094    for s in plugins.iter() {
1095        ss.push(s.as_str().to_string());
1096    }
1097    ss
1098}
1099
1100impl Ord for Plugin {
1101    fn cmp(&self, plugin: &Plugin) -> std::cmp::Ordering {
1102        self.rank().cmp(&(plugin.rank()))
1103    }
1104}
1105
1106impl PartialOrd for Plugin {
1107    fn partial_cmp(&self, plugin: &Plugin) -> Option<std::cmp::Ordering> {
1108        Some(self.cmp(plugin))
1109    }
1110}
1111
1112impl PartialEq for Plugin {
1113    fn eq(&self, plugin: &Plugin) -> bool {
1114        self.cmp(plugin) == std::cmp::Ordering::Equal
1115    }
1116}
1117
1118/// RUST_LOG=debug cargo test --package aws-manager --lib -- ec2::plugins::test_sort --exact --show-output
1119#[test]
1120fn test_sort() {
1121    let expected: Vec<Plugin> = vec![
1122        Plugin::Imds,
1123        Plugin::ProviderId,
1124        Plugin::Vercmp,
1125        Plugin::SetupLocalDisks,
1126        Plugin::MountBpfFs,
1127        Plugin::TimeSync,
1128        Plugin::SystemLimitBump,
1129        Plugin::AwsCli,
1130        Plugin::SsmAgent,
1131        Plugin::CloudwatchAgent,
1132        Plugin::Anaconda,
1133        Plugin::Go,
1134        Plugin::Docker,
1135        Plugin::Containerd,
1136        Plugin::Runc,
1137        Plugin::Ena,
1138        Plugin::NvidiaDriver,
1139        Plugin::DevBark,
1140        Plugin::EksWorkerNodeAmiReuse,
1141        Plugin::AmiInfo,
1142        Plugin::PostInitScript,
1143        Plugin::CleanupImagePackages,
1144        Plugin::CleanupImageTmpDir,
1145        Plugin::CleanupImageAwsCredentials,
1146        Plugin::CleanupImageSshKeys,
1147    ];
1148
1149    let mut unsorted: Vec<Plugin> = vec![
1150        Plugin::NvidiaDriver,
1151        Plugin::EksWorkerNodeAmiReuse,
1152        Plugin::CloudwatchAgent,
1153        Plugin::CleanupImageSshKeys,
1154        Plugin::Runc,
1155        Plugin::CleanupImagePackages,
1156        Plugin::CleanupImageTmpDir,
1157        Plugin::PostInitScript,
1158        Plugin::CleanupImageAwsCredentials,
1159        Plugin::SsmAgent,
1160        Plugin::MountBpfFs,
1161        Plugin::Imds,
1162        Plugin::SystemLimitBump,
1163        Plugin::Containerd,
1164        Plugin::Vercmp,
1165        Plugin::DevBark,
1166        Plugin::ProviderId,
1167        Plugin::TimeSync,
1168        Plugin::Docker,
1169        Plugin::SetupLocalDisks,
1170        Plugin::Anaconda,
1171        Plugin::AmiInfo,
1172        Plugin::Go,
1173        Plugin::AwsCli,
1174        Plugin::Ena,
1175    ];
1176    unsorted.sort();
1177
1178    assert_eq!(expected, unsorted);
1179    assert_eq!(unsorted[0], Plugin::Imds);
1180}
aws_manager/ec2/plugins/mod.rs

aws_manager/ec2/plugins/
mod.rs
