Upgrade scale up/down lambdas to aws sdk v3 by clee2000 · Pull Request #7061 · pytorch/test-infra

clee2000 · 2025-08-28T04:56:42Z

Upgrade to aws sdk v3

Main change is getting rid of the promise calls since I think they just directly return promises instead of requests

This changes a lot of mocks in the testing so I'm not sure how good running just yarn test is

Testing:
Mangled scaleDown to only run listInstances and listSSMParameters
In terraform-aws-github-runner/modules/runners/lambdas/runners:

yarn build; cd dist;node -e 'require("./index").scaleDown({}, {}, {});' > t.log

I also tried to terminate a runner and it worked

Deployed to pytorch-canary and ran some jobs, seems ok

vercel · 2025-08-28T04:56:48Z

The latest updates on your projects. Learn more about Vercel for GitHub.

1 Skipped Deployment

Project	Deployment	Preview	Updated (UTC)
torchci	Ignored	Preview	Sep 2, 2025 9:06pm

clee2000 · 2025-09-02T16:06:00Z

    "@types/aws-lambda": "^8.10.72",
    "@types/express": "^4.17.11",
-    "@types/jest": "^26.0.20",
+    "@types/jest": "29",


had to update jest because I kept getting some error with mocks, seems v3 isn't compatible with older jest?

Shall we upgrade to 30 then, the latest version:

clee2000 · 2025-09-02T16:06:32Z

  /* istanbul ignore next */
  if (!kms) {
-    AWS.config.update({
+    kms = new KMS({


Something about global configs not being supported anymore and to just do it in the initialization of the client?

clee2000 · 2025-09-02T16:08:18Z

  const accountId = splitARN[4];
  const queueName = splitARN[5];
-  return sqs.endpoint.href + accountId + '/' + queueName;
+  return `https://sqs.${region}.amazonaws.com/${accountId}/${queueName}`;


endpoint.href no longer seems to be a thing, so I construct it manually, hopefully it's correct. There's also a getQueueUrl function but it returns string | undefined

clee2000 · 2025-09-02T16:08:42Z

    });
  });
-  if (response.Failed.length || response.Successful.length < events.length) {
+  const failedCount = response.Failed?.length ?? 0;


return value became something | undefined

Copilot

Pull Request Overview

This PR upgrades the Lambda functions from AWS SDK v2 to v3, removing the .promise() calls since v3 directly returns promises. The upgrade includes updating dependencies, imports, type definitions, and comprehensive test mock updates.

Migrates from aws-sdk v2 to individual AWS SDK v3 client packages
Updates all service instantiations and removes .promise() calls
Refactors test mocks to align with v3's direct promise returns

Reviewed Changes

Copilot reviewed 20 out of 21 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/sqs.ts	Updates SQS service calls and queue URL construction for v3
terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.ts	Migrates EC2 and SSM service calls to v3 with updated type imports
terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/metrics.ts	Updates CloudWatch client and metric type definitions
terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/kms/index.ts	Migrates KMS client initialization and decrypt calls
terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/gh-auth.ts	Updates Secrets Manager client import and service calls
terraform-aws-github-runner/modules/runners/lambdas/runners/package.json	Replaces aws-sdk v2 with individual v3 client packages and updates Jest
Multiple test files	Updates all mocks to remove `.promise()` returns and use direct promise resolution

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

Copilot · 2025-09-03T19:10:13Z

                    Version: launchTemplateVersion,
                  },
-                  InstanceType: runnerParameters.runnerType.instance_type,
+                  InstanceType: runnerParameters.runnerType.instance_type as RunInstancesCommandInput['InstanceType'],


The type assertion here suggests a potential mismatch between the defined instance_type and what AWS SDK v3 expects. Consider updating the RunnerType interface's instance_type field to use the proper SDK v3 type instead of relying on type assertions.

Suggested change

InstanceType: runnerParameters.runnerType.instance_type as RunInstancesCommandInput['InstanceType'],

InstanceType: runnerParameters.runnerType.instance_type,

Copilot · 2025-09-03T19:10:13Z

-  protected getMetricType(metric: string): string {
-    if (Metrics.baseMetricTypes.has(metric)) return Metrics.baseMetricTypes.get(metric) as string;
+  protected getMetricType(metric: string): StandardUnit {
+    if (Metrics.baseMetricTypes.has(metric)) return Metrics.baseMetricTypes.get(metric)!;


Using the non-null assertion operator (!) after a .has() check is redundant. The has() check already guarantees the value exists, so you can safely cast or use as-assertion if needed for typing.

Suggested change

if (Metrics.baseMetricTypes.has(metric)) return Metrics.baseMetricTypes.get(metric)!;

if (Metrics.baseMetricTypes.has(metric)) return Metrics.baseMetricTypes.get(metric) as StandardUnit;

Copilot · 2025-09-03T19:10:14Z

  nock.disableNetConnect();

-  mocked(createClient).mockImplementation(produceMockedRedis);
+  (mocked(createClient) as any).mockImplementation(produceMockedRedis);


Using 'as any' defeats the purpose of TypeScript's type safety. Consider properly typing the mocked function or using a more specific type assertion that maintains type safety.

Suggested change

(mocked(createClient) as any).mockImplementation(produceMockedRedis);

(mocked(createClient) as jest.MockedFunction<typeof createClient>).mockImplementation(produceMockedRedis);

ZainRizvi

Looks reasonable, but please deploy these changes to canary and verify scale up/down works over there before merging

ZainRizvi · 2025-09-03T19:12:18Z

    "@types/aws-lambda": "^8.10.72",
    "@types/express": "^4.17.11",
-    "@types/jest": "^26.0.20",
+    "@types/jest": "29",


Shall we upgrade to 30 then, the latest version:

Similar to #7061 Mostly just getting rid of `promise()` Testing: just `yarn test` but idk how helpful that is since it mocks everything Deployed to pytorch-canary and it seems ok?

Catherine Lee added 27 commits August 27, 2025 12:51

update package versions

c1cd433

converter on runners.ts

8e6b701

copilot runners.ts

5becd31

converter on gh-auth.ts

bdd2590

converter on sqs.ts

b652ee2

copilot sqs.ts

d12e117

converter on metrics.ts

357972f

standardunit in metrics.ts

9cb8de2

converter on scale-down.ts

fa52de1

converter on kms/index.ts

9994a8f

copilot kms/index.ts

61af37b

runners.test.ts result -> command output

7b93cfc

ksm/index.test.ts

f976f25

upgrad jest

0b710d2

cache.test.ts mocking cast

b847d2d

gh-auth.test.ts idk mock stuff

2ac9677

gh-issues.test.ts change import for mocked

05b28a9

gh-runners.test.ts change sdk mock idk how this worked

a3a110f

metrics.test.ts awsk sdk mock

f28abcb

runnres.test.ts and a ton of stuff in runners.ts idk whats going on

6952669

scale-up-chron.test.ts and some axios thing

32a53df

scale-down.test.ts just mocked import

b581f7e

scale-up.test.ts just mocked import

4071731

sqs test, swap sqs to sqs client

8253274

lambda test, all tests now passing

81d6a3c

appease lint

336856c

remove unused mocks?

4ac6cb3

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 28, 2025

clee2000 force-pushed the csl/aws_sdk branch from dc1e2c9 to 4ac6cb3 Compare September 2, 2025 04:53

Catherine Lee added 3 commits September 1, 2025 22:34

swap back to ec2 in runners

4cab597

swap back to sqs from sqsclient

e2ad420

appease lint

0e59825

clee2000 commented Sep 2, 2025

View reviewed changes

clee2000 changed the title ~~[WIP] Upgrade scale up/down lambdas to aws sdk v3~~ Upgrade scale up/down lambdas to aws sdk v3 Sep 2, 2025

clee2000 marked this pull request as ready for review September 2, 2025 16:17

minor clean up

b82d997

clee2000 mentioned this pull request Sep 2, 2025

Upgrade webhook lambda for scale up/down to aws sdk v3 #7077

Merged

clee2000 requested review from a team, ZainRizvi, jeanschmidt and seemethere and removed request for ZainRizvi September 2, 2025 16:47

Catherine Lee added 2 commits September 2, 2025 14:01

fix some cloudwatch mocks that I don't think were being hit anyways

2c9bfd9

yarn format isn't opinionated enough

e6f2342

ZainRizvi requested a review from Copilot September 3, 2025 19:09

Copilot AI reviewed Sep 3, 2025

View reviewed changes

ZainRizvi approved these changes Sep 3, 2025

View reviewed changes

clee2000 merged commit 12626a2 into main Sep 5, 2025
5 checks passed

clee2000 deleted the csl/aws_sdk branch September 5, 2025 15:33

clee2000 mentioned this pull request Sep 10, 2025

Upgrade AWS lambda functions from version 2.x to 3.x of the AWS SDK for JavaScript pytorch/pytorch#137228

Closed

Navigation Menu

Uh oh!

Upgrade scale up/down lambdas to aws sdk v3 - #7061

Upgrade scale up/down lambdas to aws sdk v3#7061
clee2000 merged 33 commits into
mainfrom
csl/aws_sdk

clee2000 commented Aug 28, 2025 •

edited

Loading

Uh oh!

vercel Bot commented Aug 28, 2025 •

edited

Loading

Uh oh!

clee2000 Sep 2, 2025

Uh oh!

ZainRizvi Sep 3, 2025

Uh oh!

clee2000 Sep 2, 2025 •

edited

Loading

Uh oh!

clee2000 Sep 2, 2025

Uh oh!

clee2000 Sep 2, 2025

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Sep 3, 2025

Uh oh!

Copilot AI Sep 3, 2025

Uh oh!

Copilot AI Sep 3, 2025

Uh oh!

ZainRizvi left a comment

Uh oh!

ZainRizvi Sep 3, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	InstanceType: runnerParameters.runnerType.instance_type as RunInstancesCommandInput['InstanceType'],
	InstanceType: runnerParameters.runnerType.instance_type,

	if (Metrics.baseMetricTypes.has(metric)) return Metrics.baseMetricTypes.get(metric)!;
	if (Metrics.baseMetricTypes.has(metric)) return Metrics.baseMetricTypes.get(metric) as StandardUnit;

	(mocked(createClient) as any).mockImplementation(produceMockedRedis);
	(mocked(createClient) as jest.MockedFunction<typeof createClient>).mockImplementation(produceMockedRedis);

Uh oh!

Conversation

clee2000 commented Aug 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vercel Bot commented Aug 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

clee2000 Sep 2, 2025

Choose a reason for hiding this comment

Uh oh!

ZainRizvi Sep 3, 2025

Choose a reason for hiding this comment

Uh oh!

clee2000 Sep 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

clee2000 Sep 2, 2025

Choose a reason for hiding this comment

Uh oh!

clee2000 Sep 2, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Sep 3, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Sep 3, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Sep 3, 2025

Choose a reason for hiding this comment

Uh oh!

ZainRizvi left a comment

Choose a reason for hiding this comment

Uh oh!

ZainRizvi Sep 3, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

clee2000 commented Aug 28, 2025 •

edited

Loading

vercel Bot commented Aug 28, 2025 •

edited

Loading

clee2000 Sep 2, 2025 •

edited

Loading