Skip to content

Commit 55d652b

Browse files
tbrandtbrandaws
andauthored
Job to copy video async (#996)
Co-authored-by: Taichiro Suzuki <taichirs@amazon.co.jp>
1 parent a0efc89 commit 55d652b

10 files changed

Lines changed: 1204 additions & 859 deletions

File tree

package-lock.json

Lines changed: 799 additions & 693 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
import {
2+
S3Client,
3+
GetObjectCommand,
4+
PutObjectCommand,
5+
DeleteObjectsCommand,
6+
ListObjectsV2Command,
7+
} from '@aws-sdk/client-s3';
8+
import { Readable } from 'stream';
9+
import { VideoJob } from 'generative-ai-use-cases';
10+
import { updateJobStatus } from './repositoryVideoJob';
11+
12+
export interface CopyVideoJobParams {
13+
job: VideoJob;
14+
}
15+
16+
const BUCKET_NAME: string = process.env.BUCKET_NAME!;
17+
const videoBucketRegionMap = JSON.parse(
18+
process.env.VIDEO_BUCKET_REGION_MAP ?? '{}'
19+
);
20+
21+
const copyAndDeleteObject = async (
22+
jobId: string,
23+
srcBucket: string,
24+
srcRegion: string,
25+
dstBucket: string,
26+
dstRegion: string
27+
) => {
28+
const srcS3 = new S3Client({ region: srcRegion });
29+
const dstS3 = new S3Client({ region: dstRegion });
30+
31+
const { Body, ContentType, ContentLength } = await srcS3.send(
32+
new GetObjectCommand({
33+
Bucket: srcBucket,
34+
Key: `${jobId}/output.mp4`,
35+
})
36+
);
37+
38+
const chunks = [];
39+
for await (const chunk of Body as Readable) {
40+
chunks.push(chunk);
41+
}
42+
const fileBuffer = Buffer.concat(chunks);
43+
44+
await dstS3.send(
45+
new PutObjectCommand({
46+
Bucket: dstBucket,
47+
Key: `${jobId}/output.mp4`,
48+
Body: fileBuffer,
49+
ContentType,
50+
ContentLength,
51+
})
52+
);
53+
54+
const listRes = await srcS3.send(
55+
new ListObjectsV2Command({
56+
Bucket: srcBucket,
57+
Prefix: jobId,
58+
})
59+
);
60+
61+
const objects = listRes.Contents?.map((object) => ({
62+
Key: object.Key,
63+
}));
64+
65+
await srcS3.send(
66+
new DeleteObjectsCommand({
67+
Bucket: srcBucket,
68+
Delete: {
69+
Objects: objects,
70+
},
71+
})
72+
);
73+
};
74+
75+
export const handler = async (event: CopyVideoJobParams): Promise<void> => {
76+
const job = event.job;
77+
const jobId = job.jobId;
78+
const dstRegion = process.env.AWS_DEFAULT_REGION!;
79+
const dstBucket = BUCKET_NAME;
80+
const srcRegion = job.region;
81+
const srcBucket = videoBucketRegionMap[srcRegion];
82+
83+
try {
84+
await copyAndDeleteObject(
85+
jobId,
86+
srcBucket,
87+
srcRegion,
88+
dstBucket,
89+
dstRegion
90+
);
91+
92+
await updateJobStatus(job, 'Completed');
93+
} catch (error) {
94+
console.error(error);
95+
await updateJobStatus(job, 'Failed');
96+
}
97+
};

packages/cdk/lambda/repositoryVideoJob.ts

Lines changed: 48 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -15,20 +15,20 @@ import {
1515
GetAsyncInvokeCommand,
1616
ValidationException,
1717
} from '@aws-sdk/client-bedrock-runtime';
18-
import {
19-
S3Client,
20-
GetObjectCommand,
21-
PutObjectCommand,
22-
DeleteObjectsCommand,
23-
ListObjectsV2Command,
24-
} from '@aws-sdk/client-s3';
2518
import { initBedrockClient } from './utils/bedrockApi';
26-
import { Readable } from 'stream';
19+
import { CopyVideoJobParams } from './copyVideoJob';
20+
import {
21+
LambdaClient,
22+
InvokeCommand,
23+
InvocationType,
24+
} from '@aws-sdk/client-lambda';
2725

2826
const BUCKET_NAME: string = process.env.BUCKET_NAME!;
2927
const TABLE_NAME: string = process.env.TABLE_NAME!;
28+
const COPY_VIDEO_JOB_FUNCTION_ARN = process.env.COPY_VIDEO_JOB_FUNCTION_ARN!;
3029
const dynamoDb = new DynamoDBClient({});
3130
const dynamoDbDocument = DynamoDBDocumentClient.from(dynamoDb);
31+
const lambda = new LambdaClient({});
3232

3333
export const createJob = async (
3434
_userId: string,
@@ -67,63 +67,28 @@ export const createJob = async (
6767
return item;
6868
};
6969

70-
const copyAndDeleteObject = async (
71-
jobId: string,
72-
srcBucket: string,
73-
srcRegion: string,
74-
dstBucket: string,
75-
dstRegion: string
76-
) => {
77-
const srcS3 = new S3Client({ region: srcRegion });
78-
const dstS3 = new S3Client({ region: dstRegion });
79-
80-
const { Body, ContentType, ContentLength } = await srcS3.send(
81-
new GetObjectCommand({
82-
Bucket: srcBucket,
83-
Key: `${jobId}/output.mp4`,
84-
})
85-
);
86-
87-
const chunks = [];
88-
for await (const chunk of Body as Readable) {
89-
chunks.push(chunk);
90-
}
91-
const fileBuffer = Buffer.concat(chunks);
92-
93-
await dstS3.send(
94-
new PutObjectCommand({
95-
Bucket: dstBucket,
96-
Key: `${jobId}/output.mp4`,
97-
Body: fileBuffer,
98-
ContentType,
99-
ContentLength,
100-
})
101-
);
102-
103-
const listRes = await srcS3.send(
104-
new ListObjectsV2Command({
105-
Bucket: srcBucket,
106-
Prefix: jobId,
107-
})
108-
);
109-
110-
const objects = listRes.Contents?.map((object) => ({
111-
Key: object.Key,
112-
}));
113-
114-
await srcS3.send(
115-
new DeleteObjectsCommand({
116-
Bucket: srcBucket,
117-
Delete: {
118-
Objects: objects,
119-
},
120-
})
121-
);
70+
export const updateJobStatus = async (job: VideoJob, status: string) => {
71+
const updateCommand = new UpdateCommand({
72+
TableName: TABLE_NAME,
73+
Key: {
74+
id: job.id,
75+
createdDate: job.createdDate,
76+
},
77+
UpdateExpression: 'set #status = :status',
78+
ExpressionAttributeNames: {
79+
'#status': 'status',
80+
},
81+
ExpressionAttributeValues: {
82+
':status': status,
83+
},
84+
});
85+
86+
await dynamoDbDocument.send(updateCommand);
12287
};
12388

12489
const checkAndUpdateJob = async (
12590
job: VideoJob
126-
): Promise<'InProgress' | 'Completed' | 'Failed'> => {
91+
): Promise<'InProgress' | 'Completed' | 'Failed' | 'Finalizing'> => {
12792
try {
12893
const client = await initBedrockClient(job.region);
12994
const command = new GetAsyncInvokeCommand({
@@ -145,45 +110,29 @@ const checkAndUpdateJob = async (
145110
}
146111
}
147112

148-
if (res.status !== 'InProgress') {
149-
if (res.status === 'Completed') {
150-
const jobId = job.jobId;
151-
const dstBucket = BUCKET_NAME;
152-
const dstRegion = process.env.AWS_DEFAULT_REGION!;
153-
const srcRegion = job.region;
154-
const videoBucketRegionMap = JSON.parse(
155-
process.env.VIDEO_BUCKET_REGION_MAP ?? '{}'
156-
);
157-
const srcBucket = videoBucketRegionMap[srcRegion];
158-
159-
await copyAndDeleteObject(
160-
jobId,
161-
srcBucket,
162-
srcRegion,
163-
dstBucket,
164-
dstRegion
165-
);
166-
}
167-
168-
const updateCommand = new UpdateCommand({
169-
TableName: TABLE_NAME,
170-
Key: {
171-
id: job.id,
172-
createdDate: job.createdDate,
173-
},
174-
UpdateExpression: 'set #status = :status',
175-
ExpressionAttributeNames: {
176-
'#status': 'status',
177-
},
178-
ExpressionAttributeValues: {
179-
':status': res.status,
180-
},
181-
});
182-
183-
await dynamoDbDocument.send(updateCommand);
113+
// Video generation is complete, but the video copying is not finished.
114+
// We will run the copy job to set the status to "Finalizing".
115+
if (res.status === 'Completed') {
116+
const params: CopyVideoJobParams = { job };
117+
118+
await lambda.send(
119+
new InvokeCommand({
120+
FunctionName: COPY_VIDEO_JOB_FUNCTION_ARN,
121+
InvocationType: InvocationType.Event,
122+
Payload: JSON.stringify(params),
123+
})
124+
);
125+
126+
await updateJobStatus(job, 'Finalizing');
127+
return 'Finalizing';
128+
} else if (res.status === 'Failed') {
129+
// Since video generation has failed, we will not copy the video and will terminate with a Failed status.
130+
await updateJobStatus(job, 'Failed');
131+
return 'Failed';
132+
} else {
133+
// This res.status will be InProgress only.
134+
return res.status!;
184135
}
185-
186-
return res.status!;
187136
} catch (e) {
188137
console.error(e);
189138
return job.status;

packages/cdk/lib/construct/api.ts

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -284,10 +284,11 @@ export class Api extends Construct {
284284
}
285285
table.grantWriteData(generateVideoFunction);
286286

287-
const listVideoJobs = new NodejsFunction(this, 'ListVideoJobs', {
287+
const copyVideoJob = new NodejsFunction(this, 'CopyVideoJob', {
288288
runtime: Runtime.NODEJS_LATEST,
289-
entry: './lambda/listVideoJobs.ts',
289+
entry: './lambda/copyVideoJob.ts',
290290
timeout: Duration.minutes(15),
291+
memorySize: 512,
291292
environment: {
292293
MODEL_REGION: modelRegion,
293294
MODEL_IDS: JSON.stringify(modelIds),
@@ -304,7 +305,7 @@ export class Api extends Construct {
304305
});
305306
for (const region of Object.keys(props.videoBucketRegionMap)) {
306307
const bucketName = props.videoBucketRegionMap[region];
307-
listVideoJobs.role?.addToPrincipalPolicy(
308+
copyVideoJob.role?.addToPrincipalPolicy(
308309
new PolicyStatement({
309310
effect: Effect.ALLOW,
310311
actions: ['s3:GetObject', 's3:DeleteObject', 's3:ListBucket'],
@@ -315,8 +316,30 @@ export class Api extends Construct {
315316
})
316317
);
317318
}
318-
fileBucket.grantWrite(listVideoJobs);
319+
fileBucket.grantWrite(copyVideoJob);
320+
table.grantWriteData(copyVideoJob);
321+
322+
const listVideoJobs = new NodejsFunction(this, 'ListVideoJobs', {
323+
runtime: Runtime.NODEJS_LATEST,
324+
entry: './lambda/listVideoJobs.ts',
325+
timeout: Duration.minutes(15),
326+
environment: {
327+
MODEL_REGION: modelRegion,
328+
MODEL_IDS: JSON.stringify(modelIds),
329+
IMAGE_GENERATION_MODEL_IDS: JSON.stringify(imageGenerationModelIds),
330+
VIDEO_GENERATION_MODEL_IDS: JSON.stringify(videoGenerationModelIds),
331+
VIDEO_BUCKET_REGION_MAP: JSON.stringify(props.videoBucketRegionMap),
332+
CROSS_ACCOUNT_BEDROCK_ROLE_ARN: crossAccountBedrockRoleArn ?? '',
333+
BUCKET_NAME: fileBucket.bucketName,
334+
TABLE_NAME: table.tableName,
335+
COPY_VIDEO_JOB_FUNCTION_ARN: copyVideoJob.functionArn,
336+
},
337+
bundling: {
338+
nodeModules: ['@aws-sdk/client-bedrock-runtime'],
339+
},
340+
});
319341
table.grantReadWriteData(listVideoJobs);
342+
copyVideoJob.grantInvoke(listVideoJobs);
320343

321344
const deleteVideoJob = new NodejsFunction(this, 'DeleteVideoJob', {
322345
runtime: Runtime.NODEJS_LATEST,

packages/cdk/package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
"@aws-sdk/client-bedrock-runtime": "^3.755.0",
3333
"@aws-sdk/client-dynamodb": "^3.755.0",
3434
"@aws-sdk/client-kendra": "^3.755.0",
35+
"@aws-sdk/client-lambda": "^3.755.0",
3536
"@aws-sdk/client-s3": "^3.755.0",
3637
"@aws-sdk/client-sagemaker-runtime": "^3.755.0",
3738
"@aws-sdk/client-transcribe": "^3.755.0",

0 commit comments

Comments
 (0)