Windows AMI: Add possibility to use different owners specified in the filter param

.github/scripts/validate_scale_config.py

@@ -37,7 +37,11 @@
     "additionalProperties": False,
     "properties": {
         "ami_experiment": {"type": "object"},
-        "ami": {"type": "string"},
+        "ami": {
+            "type": "string",
+            "pattern": "^[A-Za-z0-9 -_. ]+(\|[0-9]+)?$",
+            "description": "AMI Name|AWS Account (optional).",
+        },
         "disk_size": {"type": "number"},
         "instance_type": {"type": "string"},
         "is_ephemeral": {"type": "boolean"},

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.test.ts

@@ -418,6 +418,25 @@ describe('findAmiID', () => {
     await findAmiID(metrics, 'REGION', 'FILTER2');
     expect(mockEC2.describeImages).toBeCalledTimes(3);
   });
+
+  it('handles filter with separator and custom account ID', async () => {
+    const result = await findAmiID(metrics, 'REGION', 'my-image|123456789012');
+    expect(mockEC2.describeImages).toBeCalledTimes(1);
+    expect(mockEC2.describeImages).toBeCalledWith({
+      Owners: ['123456789012'],
+      Filters: [
+        {
+          Name: 'name',
+          Values: ['my-image'],
+        },
+        {
+          Name: 'state',
+          Values: ['available'],
+        },
+      ],
+    });
+    expect(result).toBe('ami-AGDGADU113');
+  });
 });
 
 describe('tryReuseRunner', () => {

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.ts

@@ -77,19 +77,42 @@ export function resetRunnersCaches() {
 
 export async function findAmiID(metrics: Metrics, region: string, filter: string, owners = 'amazon'): Promise<string> {
   const ec2 = new EC2({ region: region });
+  // Check if filter contains separator '|' and extract image name and account ID
+  let imageName = filter;
+  let actualOwners = owners;
+
+  if (filter.includes('|')) {
+    const parts = filter.split('|');
+    if (parts.length === 2) {
+      imageName = parts[0].trim();
+      const extractedOwner = parts[1].trim();
+
+      // Check if the extracted owner is only numbers (AWS account ID format)
+      if (/^\d+$/.test(extractedOwner)) {
+        actualOwners = extractedOwner;
+      } else {
+        console.error(
+          `Invalid account ID format: '${extractedOwner}'. Account ID must` +
+            ` contain only numbers. Using default value '${actualOwners}'`,
+        );
+      }
+    }
+  }
+
   const filters = [
-    { Name: 'name', Values: [filter] },
+    { Name: 'name', Values: [imageName] },
     { Name: 'state', Values: ['available'] },
   ];
-  return redisCached('awsEC2', `findAmiID-${region}-${filter}-${owners}`, 10 * 60, 0.5, () => {
+
+  return redisCached('awsEC2', `findAmiID-${region}-${imageName}-${actualOwners}`, 10 * 60, 0.5, () => {
     return expBackOff(() => {
       return metrics.trackRequestRegion(
         region,
         metrics.ec2DescribeImagesSuccess,
         metrics.ec2DescribeImagesFailure,
         () => {
           return ec2
-            .describeImages({ Owners: [owners], Filters: filters })
+            .describeImages({ Owners: [actualOwners], Filters: filters })
             .promise()
             .then((data: EC2.DescribeImagesResult) => {
               /* istanbul ignore next */