Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add CPU and memory options for task size #392

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion cmd/run.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ type runTask struct {
timeout int
timestampFormat string
platformVersion string
taskSizeCpu string
taskSizeMemory string
}

func runTaskCmd() *cobra.Command {
Expand All @@ -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
}
Expand All @@ -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)
}
Expand Down
15 changes: 13 additions & 2 deletions pkg/task/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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")
}
Expand Down Expand Up @@ -172,6 +175,8 @@ func NewTask(cluster, container, taskDefinitionName, command string, fargate boo
region: region,
timestampFormat: timestampFormat,
PlatformVersion: platformVersion,
taskSizeCpu: taskSizeCpu,
taskSizeMemory: taskSizeMemory,
}, nil
}

Expand All @@ -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{
Expand Down