cloud/aws/node/generate-setup-script.py (231 lines of code) (raw):

#!/usr/bin/env python3 # # This script automates the generation of `setup.ps1` in the `scripts` subdirectory # import json, re, subprocess, yaml, sys from pathlib import Path class Utility: @staticmethod def log(message): """ Logs a message to stderr """ print('[generate-setup-script.py]: {}'.format(message), flush=True, file=sys.stderr) @staticmethod def capture(command, **kwargs): """ Executes the specified command and captures its output """ # Log the command being executed Utility.log(command) # Attempt to execute the specified command result = subprocess.run( command, check = True, capture_output = True, universal_newlines = True, **kwargs ) # Return the contents of stdout return result.stdout.strip() @staticmethod def writeFile(filename, data): """ Writes data to the specified file """ return Path(filename).write_bytes(data.encode('utf-8')) @staticmethod def commentForStep(name): """ Returns a descriptive comment for the build step with the specified name """ return { 'ConfigureDirectories': '# Create each of our directories', 'DownloadKubernetes': '# Download the Kubernetes components', 'DownloadEKSArtifacts': '# Download the EKS artifacts archive', 'ExtractEKSArtifacts': '# Extract the EKS artifacts archive', 'MoveEKSArtifacts': '# Move the EKS files into place', 'ExecuteBuildScripts': '# Perform EKS worker node setup', 'RemoveEKSArtifactDownloadDirectory': '# Perform cleanup', 'InstallContainers': '\n'.join([ '# Install the Windows Containers feature', '# (Note: this is actually a no-op here, since we install the feature beforehand in startup.ps1)' ]) }.get(name, None) @staticmethod def parseConstants(constants): """ Parses an EC2 ImageBuilder component's constants list """ parsed = {} for entry in constants: for key, values in entry.items(): parsed[key] = values['value'] return parsed @staticmethod def replaceConstants(string, constants): """ Converts EC2 ImageBuilder constant references to PowerShell variable references """ # If the value of a constant is used as a magic value rather than a reference, # replace it with a reference to the variable representing the constant instead transformed = string for key, value in constants.items(): transformed = transformed.replace(value, '${}'.format(key)) # Convert `{{ variable }}` syntax to PowerShell `$variable` syntax # (Note that we don't bother to wrap the variable names in curly braces, since we know that none # of the variable names contain special characters, and they're only ever interpolated as either # part of a filesystem path surrounded by separators, or as a parameter surrounded by whitespace) return re.sub('{{ (.+?) }}', '$\\1', transformed) @staticmethod def replaceSystemPaths(path): """ Replaces hard-coded system paths with the equivalent environment variables """ replaced = path replaced = replaced.replace('C:\\Program Files', '$env:ProgramFiles') replaced = replaced.replace('C:\\ProgramData', '$env:ProgramData') return replaced @staticmethod def s3UriToHttpsUrl(s3Uri): """ Converts an `s3://` URI to an HTTPS URL """ url = s3Uri.replace('s3://', '') components = url.split('/', 1) return 'https://{}.s3.amazonaws.com/{}'.format(components[0], components[1]) # Retrieve the contents of the "Amazon EKS Optimized Windows AMI" EC2 ImageBuilder component componentData = json.loads(Utility.capture([ 'aws', 'imagebuilder', 'get-component', '--region=us-east-1', '--component-build-version-arn', 'arn:aws:imagebuilder:us-east-1:aws:component/eks-optimized-ami-windows/1.24.0' ])) # Parse the pipeline YAML data and extract the list of constants pipelineData = yaml.load(componentData['component']['data'], Loader=yaml.Loader) constants = Utility.parseConstants(pipelineData['constants']) # Extract the steps for the "build" phase buildSteps = [p['steps'] for p in pipelineData['phases'] if p['name'] == 'build'][0] print('CONSTANTS:') print(json.dumps(constants, indent=4)) print() print('BUILD STEPS:') print(json.dumps(buildSteps, indent=4)) # Prepend our header to the generated PowerShell code generated = '''<# THIS FILE IS AUTOMATICALLY GENERATED, DO NOT EDIT! This script is based on the logic from the "Amazon EKS Optimized Windows AMI" EC2 ImageBuilder component, with modifications to use containerd 1.7.0. The original ImageBuilder component logic is Copyright Amazon.com, Inc. or its affiliates, and is licensed under the MIT License. #> # Halt execution if we encounter an error $ErrorActionPreference = 'Stop' # Applies in-place patches to a file function PatchFile { Param ( $File, $Patches ) $patched = Get-Content -Path $File -Raw $Patches.GetEnumerator() | ForEach-Object { $patched = $patched.Replace($_.Key, $_.Value) } Set-Content -Path $File -Value $patched -NoNewline } ''' # Inject an additional constant for the parent of the temp directory, immediately before the child directory tempPath = {k:v for k,v in constants.items() if k == 'TempPath'} otherConstants = {k:v for k,v in constants.items() if k != 'TempPath'} constants = {**otherConstants, 'TempRoot': 'C:\\TempEKSArtifactDir', **tempPath} # Define variables for each of our constants generated += '# Constants\n' existingConstants = {} for key, value in constants.items(): transformed = Utility.replaceConstants(value, existingConstants) transformed = Utility.replaceSystemPaths(transformed) generated += '${} = "{}"\n'.format(key, transformed) existingConstants[key] = value # Process each build step in turn for step in buildSteps: # Determine whether we have custom preprocessing logic for the step name = step['name'] if name == 'ConfigureDirectories': # Add the temp directory to the list of directories to be created step['loop']['forEach'] += [constants['TempRoot']] elif name == 'DownloadKubernetes': # Inject the driver installation step immediately prior to the Kubernetes download step generated += '\n'.join([ '', '# Install the NVIDIA GPU drivers', "$driverBucket = 'ec2-windows-nvidia-drivers'", "$driver = Get-S3Object -BucketName $driverBucket -KeyPrefix 'latest' -Region 'us-east-1' | Where-Object {$_.Key.Contains('server2022')}", 'Copy-S3Object -BucketName $driverBucket -Key $driver.Key -LocalFile "$TempRoot\driver.exe" -Region \'us-east-1\'', "Start-Process -FilePath \"$TempRoot\driver.exe\" -ArgumentList @('-s', '-noreboot') -NoNewWindow -Wait", '' ]) elif name == 'ExtractEKSArtifacts': # Remove the redundant directory creation command step['inputs']['commands'] = [ c for c in step['inputs']['commands'] if not c.startswith('New-Item') ] # Use absolute file and directory paths rather than relative paths step['inputs']['commands'] = [ c.replace('EKS-Artifacts.zip', '"C:\\EKS-Artifacts.zip"').replace('TempEKSArtifactDir', 'C:\\TempEKSArtifactDir') for c in step['inputs']['commands'] ] elif name == 'InstallContainerRuntimes': # Inject the containerd 1.7.0 download step, along with our configuration patching steps, immediately prior to the containerd installation step generated += '\n'.join([ '', '# -------', '', '# TEMPORARY UNTIL EKS ADDS SUPPORT FOR CONTAINERD v1.7.0:', '# Download and extract the containerd 1.7.0 release build', '$containerdTarball = "$TempPath\\containerd-1.7.0.tar.gz"', '$containerdFiles = "$TempPath\\containerd-1.7.0"', '$webClient.DownloadFile(\'https://github.com/containerd/containerd/releases/download/v1.7.0/containerd-1.7.0-windows-amd64.tar.gz\', $containerdTarball)', 'New-Item -Path "$containerdFiles" -ItemType Directory -Force | Out-Null', 'tar.exe -xvzf "$containerdTarball" -C "$containerdFiles"', '', '# Move the containerd files into place', 'Move-Item -Path "$containerdFiles\\bin\\containerd.exe" -Destination "$ContainerdPath\\containerd.exe" -Force', 'Move-Item -Path "$containerdFiles\\bin\\containerd-shim-runhcs-v1.exe" -Destination "$ContainerdPath\\containerd-shim-runhcs-v1.exe" -Force', 'Move-Item -Path "$containerdFiles\\bin\\ctr.exe" -Destination "$ContainerdPath\\ctr.exe" -Force', '', '# Clean up the containerd intermediate files', 'Remove-Item -Path "$containerdFiles" -Recurse -Force', 'Remove-Item -Path "$containerdTarball" -Force', '', '# -------', '', '# Patch the containerd setup script to configure a log file (rather than just discarding log output) and to use the upstream pause', '# container image rather than the EKS version, since the latter appears to cause errors when attempting to create Windows Pods', 'PatchFile -File "$TempPath\Add-ContainerdRuntime.ps1" -Patches @{', ' "containerd --register-service" = "containerd --register-service --log-file \'C:\\ProgramData\\containerd\\root\\output.log\'";', ' "amazonaws.com/eks/pause-windows:latest" = "registry.k8s.io/pause:3.9"', '}', '', '# Add the full Windows Server 2022 base image and the pause image to the list of images to pre-pull', '$baseLayersFile = "$TempPath\eks.baselayers.config"', '$baseLayers = Get-Content -Path $baseLayersFile -Raw | ConvertFrom-Json', '$baseLayers.2022 += "mcr.microsoft.com/windows/server:ltsc2022"', '$baseLayers.2022 += "registry.k8s.io/pause:3.9"', '$patchedJson = ConvertTo-Json -Depth 100 -InputObject $baseLayers', 'Set-Content -Path $baseLayersFile -Value $patchedJson -NoNewline', '', ]) # Simplify the containerd installation command step['inputs']['commands'] = [ '', '# Register containerd as the EKS container runtime', 'Push-Location $TempPath', '& .\Add-ContainerdRuntime.ps1 -Path "$ContainerdPath"', 'Pop-Location' ] elif name == 'ExecuteBuildScripts': # Prefix each script invocation with the call operator step['loop']['forEach'] = [ '& {}'.format(command) for command in step['loop']['forEach'] ] # Strip away the boilerplate code surrounding each script invocation step['inputs']['commands'] = ['Push-Location $TempPath'] + step['loop']['forEach'] + ['Pop-Location'] # ------- # If we have a descriptive comment for the step then include it above its generated code comment = Utility.commentForStep(name) if comment != None: generated += '\n{}\n'.format(comment) # ------- # Generate code for the step based on its action type action = step['action'] if action == 'CreateFolder': directories = [Utility.replaceConstants(d, constants) for d in step['loop']['forEach']] generated += '\n'.join([ 'foreach ($dir in @({})) {{'.format(', '.join(directories)), '\tNew-Item -Path $dir -ItemType Directory -Force | Out-Null', '}' ]) elif action == 'DeleteFolder': generated += '\n'.join([ 'Remove-Item -Path "{}" -Recurse -Force'.format(Utility.replaceConstants(input['path'], constants)) for input in step['inputs'] ]) elif action == 'MoveFile': generated += '\n'.join([ 'Move-Item -Path "{}" -Destination "{}" -Force'.format( Utility.replaceConstants(input['source'], constants), Utility.replaceConstants(input['destination'], constants) ) for input in step['inputs'] ]) elif action == 'S3Download': generated += '\n'.join([ '$webClient.DownloadFile("{}", "{}")'.format( Utility.s3UriToHttpsUrl(input['source']), Utility.replaceConstants(input['destination'], constants) ) for input in step['inputs'] ]) elif action == 'ExecutePowerShell': generated += '\n'.join([ Utility.replaceConstants(c, constants).replace("'", '"') for c in step['inputs']['commands'] if not c.startswith('$ErrorActionPreference') ]) elif action == 'Reboot': Utility.log('Ignoring reboot step.') continue else: raise RuntimeError('Unknown build step action: {}'.format(action)) # ------- # Add a trailing newline after each non-ignored step generated += '\n' # Write the generated code to the output script file outfile = Path(__file__).parent / 'scripts' / 'setup.ps1' Utility.writeFile(outfile, generated) Utility.log('Wrote generated code to {}'.format(outfile))