From 9a3aa1be1a2473824a8caa14a0309030e49eb9f8 Mon Sep 17 00:00:00 2001 From: keidai Date: Mon, 18 Nov 2024 17:43:10 +0900 Subject: [PATCH] Add CPU and memory options for task size --- cmd/run.go | 6 +++++- pkg/task/task.go | 15 +++++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/cmd/run.go b/cmd/run.go index 85db038..f156136 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -19,6 +19,8 @@ type runTask struct { timeout int timestampFormat string platformVersion string + taskSizeCpu string + taskSizeMemory string } func runTaskCmd() *cobra.Command { @@ -40,6 +42,8 @@ func runTaskCmd() *cobra.Command { flags.IntVarP(&r.timeout, "timeout", "t", 0, "Timeout seconds") flags.StringVarP(&r.timestampFormat, "timestamp-format", "", "[2006-01-02 15:04:05.999999999 -0700 MST]", "Format of timestamp for outputs. You should follow the style of Time.Format (see https://golang.org/pkg/time/#pkg-constants)") flags.StringVarP(&r.platformVersion, "platform-version", "p", "", "The platform version that your tasks in the service are running on. A platform version is specified only for tasks using the Fargate launch type. If one isn't specified, the LATEST platform version is used by default.") + flags.StringVar(&r.taskSizeCpu, "task-size-cpu", "", "The hard limit of CPU units to present for the task. If both task-size-cpu and task-size-memory are set, overwrite task definition.") + flags.StringVar(&r.taskSizeMemory, "task-size-memory", "", "The hard limit of memory to present to the task. If both task-size-cpu and task-size-memory are set, overwrite task definition.") return cmd } @@ -49,7 +53,7 @@ func (r *runTask) run(cmd *cobra.Command, args []string) { if !verbose { log.SetLevel(log.WarnLevel) } - t, err := task.NewTask(r.cluster, r.container, r.taskDefinition, r.command, r.fargate, r.subnets, r.securityGroups, r.platformVersion, (time.Duration(r.timeout) * time.Second), r.timestampFormat, profile, region) + t, err := task.NewTask(r.cluster, r.container, r.taskDefinition, r.command, r.fargate, r.subnets, r.securityGroups, r.platformVersion, (time.Duration(r.timeout) * time.Second), r.timestampFormat, profile, region, r.taskSizeCpu, r.taskSizeMemory) if err != nil { log.Fatal(err) } diff --git a/pkg/task/task.go b/pkg/task/task.go index 065bdd9..1acc063 100644 --- a/pkg/task/task.go +++ b/pkg/task/task.go @@ -13,7 +13,7 @@ At first, you have to get a task definition. The task definition is used to run For example: - t, err := task.NewTask("cluster-name", "container-name", "task-definition-arn or family", "commands", false, "", 300 * time.Second, "profile", "region") + t, err := task.NewTask("cluster-name", "container-name", "task-definition-arn or family", "commands", false, "", 300 * time.Second, "profile", "region", "task-size-cpu", "task-size-memory") // At first you have to get a task definition. taskDef, err := t.taskDefinition.DescribeTaskDefinition(t.TaskDefinitionName) @@ -105,12 +105,15 @@ type Task struct { profile string region string timestampFormat string + // If you wat to override CPU and Memory, please set these values. + taskSizeCpu string + taskSizeMemory string } // NewTask returns a new Task struct, and initialize aws ecs API client. // If you want to run the task as Fargate, please provide fargate flag to true, and your subnet IDs for awsvpc. // If you don't want to run the task as Fargate, please provide empty string for subnetIDs. -func NewTask(cluster, container, taskDefinitionName, command string, fargate bool, subnetIDs, securityGroupIDs, platformVersion string, timeout time.Duration, timestampFormat, profile, region string) (*Task, error) { +func NewTask(cluster, container, taskDefinitionName, command string, fargate bool, subnetIDs, securityGroupIDs, platformVersion string, timeout time.Duration, timestampFormat, profile, region, taskSizeCpu, taskSizeMemory string) (*Task, error) { if cluster == "" { return nil, errors.New("Cluster name is required") } @@ -172,6 +175,8 @@ func NewTask(cluster, container, taskDefinitionName, command string, fargate boo region: region, timestampFormat: timestampFormat, PlatformVersion: platformVersion, + taskSizeCpu: taskSizeCpu, + taskSizeMemory: taskSizeMemory, }, nil } @@ -187,6 +192,12 @@ func (t *Task) RunTask(ctx context.Context, taskDefinition *ecstypes.TaskDefinit containerOverride, }, } + + if len(t.taskSizeCpu) > 0 && len(t.taskSizeMemory) > 0 { + override.Cpu = aws.String(t.taskSizeCpu) + override.Memory = aws.String(t.taskSizeMemory) + } + var params *ecs.RunTaskInput if len(t.Subnets) > 0 { vpcConfiguration := &ecstypes.AwsVpcConfiguration{