Merge pull request #29 from guardian/update-scale-in-protection

Add/remove scale in protection for worker instances
guardian · Feb 13, 2024 · 5f7f748 · 5f7f748
2 parents 6aead63 + 72ff7e4
commit 5f7f748
Show file tree

Hide file tree

Showing 8 changed files with 555 additions and 12 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/cdk/lib/__snapshots__/transcription-service.test.ts.snap b/packages/cdk/lib/__snapshots__/transcription-service.test.ts.snap
@@ -16,6 +16,7 @@ exports[`The TranscriptionService stack matches the snapshot 1`] = `
       "GuAllowPolicy",
       "GuAllowPolicy",
       "GuAllowPolicy",
+      "GuAllowPolicy",
       "GuInstanceRole",
       "GuDescribeEC2Policy",
       "GuLogShippingPolicy",
@@ -420,6 +421,27 @@ exports[`The TranscriptionService stack matches the snapshot 1`] = `
       },
       "Type": "AWS::IAM::Policy",
     },
+    "SetInstanceProtection5DD5610F": {
+      "Properties": {
+        "PolicyDocument": {
+          "Statement": [
+            {
+              "Action": "autoscaling:SetInstanceProtection",
+              "Effect": "Allow",
+              "Resource": "arn:aws:autoscaling:test-region:745349931791:autoScalingGroup:*:autoScalingGroupName/transcription-service-workers-TEST",
+            },
+          ],
+          "Version": "2012-10-17",
+        },
+        "PolicyName": "SetInstanceProtection5DD5610F",
+        "Roles": [
+          {
+            "Ref": "InstanceRoleTranscriptionserviceworkerA8AC6615",
+          },
+        ],
+      },
+      "Type": "AWS::IAM::Policy",
+    },
     "TranscriptDestinationTopic831E3125": {
       "Properties": {
         "Tags": [
@@ -539,7 +561,7 @@ exports[`The TranscriptionService stack matches the snapshot 1`] = `
             ],
           },
         },
-        "NewInstancesProtectedFromScaleIn": false,
+        "NewInstancesProtectedFromScaleIn": true,
         "Tags": [
           {
             "Key": "gu:cdk:version",

diff --git a/packages/cdk/lib/transcription-service.ts b/packages/cdk/lib/transcription-service.ts
@@ -52,6 +52,7 @@ export class TranscriptionService extends GuStack {
 		const APP_NAME = 'transcription-service';
 		const apiId = `${APP_NAME}-${props.stage}`;
 		const isProd = props.stage === 'PROD';
+		const autoScalingGroupName = `transcription-service-workers-${this.stage}`;
 		if (!props.env?.region) throw new Error('region not provided in props');
 
 		const workerAmi = new GuAmiParameter(this, {
@@ -239,6 +240,12 @@ export class TranscriptionService extends GuStack {
 					],
 					resources: [loggingStreamArn],
 				}),
+				new GuAllowPolicy(this, 'SetInstanceProtection', {
+					actions: ['autoscaling:SetInstanceProtection'],
+					resources: [
+						`arn:aws:autoscaling:${props.env.region}:${GuardianAwsAccounts.Investigations}:autoScalingGroup:*:autoScalingGroupName/${autoScalingGroupName}`,
+					],
+				}),
 			],
 		});
 
@@ -286,7 +293,7 @@ export class TranscriptionService extends GuStack {
 			{
 				minCapacity: 0,
 				maxCapacity: isProd ? 20 : 4,
-				autoScalingGroupName: `transcription-service-workers-${this.stage}`,
+				autoScalingGroupName,
 				vpc: GuVpc.fromIdParameter(this, 'InvestigationsInternetEnabledVpc', {
 					availabilityZones: ['eu-west-1a', 'eu-west-1b', 'eu-west-1c'],
 				}),
@@ -298,7 +305,7 @@ export class TranscriptionService extends GuStack {
 				},
 				// we might want to set this to true once we are actually doing transcriptions to protect the instance from
 				// being terminated before it has a chance to complete a transcription job.
-				newInstancesProtectedFromScaleIn: false,
+				newInstancesProtectedFromScaleIn: true,
 				mixedInstancesPolicy: {
 					launchTemplate,
 					instancesDistribution: {

diff --git a/packages/worker/package.json b/packages/worker/package.json
@@ -12,6 +12,7 @@
 	"author": "",
 	"license": "ISC",
 	"dependencies": {
+		"@aws-sdk/client-auto-scaling": "^3.496.0",
 		"@aws-sdk/client-sns": "^3.496.0",
 		"@guardian/transcription-service-backend-common": "1.0.0"
 	},

diff --git a/packages/worker/src/asg.ts b/packages/worker/src/asg.ts
@@ -0,0 +1,35 @@
+import {
+	AutoScalingClient,
+	SetInstanceProtectionCommand,
+} from '@aws-sdk/client-auto-scaling';
+import { readFile } from './util';
+
+export const updateScaleInProtection = async (
+	region: string,
+	stage: string,
+	value: boolean,
+) => {
+	const clientConfig = {
+		region,
+	};
+	try {
+		if (stage !== 'DEV') {
+			const instanceId = readFile('/var/lib/cloud/data/instance-id');
+			console.log(`instanceId: ${instanceId}`);
+			const autoScalingClient = new AutoScalingClient(clientConfig);
+			const input = {
+				InstanceIds: [instanceId.trim()],
+				AutoScalingGroupName: `transcription-service-workers-${stage}`,
+				ProtectedFromScaleIn: value,
+			};
+			const command = new SetInstanceProtectionCommand(input);
+			await autoScalingClient.send(command);
+			console.log(
+				`Updated scale-in protection to value ${value} for instance ${instanceId}`,
+			);
+		}
+	} catch (error) {
+		console.log(`Could not remove scale-in protection`, error);
+		throw error;
+	}
+};
diff --git a/packages/worker/src/index.ts b/packages/worker/src/index.ts
@@ -18,22 +18,27 @@ import {
 	getOrCreateContainer,
 } from './transcribe';
 import path from 'path';
+import { updateScaleInProtection } from './asg';
 
 const main = async () => {
 	const config = await getConfig();
+	const stage = config.app.stage;
+	const region = config.aws.region;
 
 	const numberOfThreads = config.app.stage === 'PROD' ? 16 : 2;
 
-	const client = getSQSClient(config.aws.region, config.aws.localstackEndpoint);
+	const client = getSQSClient(region, config.aws.localstackEndpoint);
 	const message = await getNextMessage(client, config.app.taskQueueUrl);
 
 	if (isFailure(message)) {
 		console.error(`Failed to fetch message due to ${message.errorMsg}`);
+		await updateScaleInProtection(region, stage, false);
 		return;
 	}
 
 	if (!message.message) {
 		console.log('No messages available');
+		await updateScaleInProtection(region, stage, false);
 		return;
 	}
 
@@ -57,8 +62,6 @@ const main = async () => {
 
 		const fileToTranscribe = await getFileFromS3(config, job.s3Key);
 
-		console.log('file is here');
-
 		// docker container to run ffmpeg and whisper on file
 		const containerId = await getOrCreateContainer(
 			path.parse(fileToTranscribe).dir,
@@ -122,6 +125,8 @@ const main = async () => {
 				0,
 			);
 		}
+	} finally {
+		await updateScaleInProtection(region, stage, false);
 	}
 };
 

diff --git a/packages/worker/src/transcribe.ts b/packages/worker/src/transcribe.ts
@@ -1,6 +1,6 @@
 import { spawn } from 'child_process';
 import path from 'path';
-import * as fs from 'fs';
+import { readFile } from './util';
 
 interface ProcessResult {
 	code?: number;
@@ -139,11 +139,6 @@ const getDuration = (ffmpegOutput: string) => {
 	return duration;
 };
 
-const readFile = (filePath: string): string => {
-	const file = fs.readFileSync(filePath, 'utf8');
-	return file;
-};
-
 export const getTranscriptionText = async (
 	containerId: string,
 	wavPath: string,

diff --git a/packages/worker/src/util.ts b/packages/worker/src/util.ts
@@ -0,0 +1,6 @@
+import * as fs from 'fs';
+
+export const readFile = (filePath: string): string => {
+	const file = fs.readFileSync(filePath, 'utf8');
+	return file;
+};