Skip to content

Commit

Permalink
Merge pull request #100 from trickest/feat/support-new-fleet-types
Browse files Browse the repository at this point in the history 
Feat/support new fleet types
  • Loading branch information
@mhmdiaa
mhmdiaa authored Nov 29, 2023
2 parents 91e2d97 + fc7da71 commit 8655a59
Show file tree
Hide file tree
Showing 10 changed files with 185 additions and 105 deletions.
4 changes: 3 additions & 1 deletion README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,9 @@ trickest execute --workflow <workflow_or_tool_name> --space <space_name> --confi
| --set-name | string | / | Sets the new workflow name and will copy the workflow to space and project supplied |
| --ci | boolean | false | Enable CI mode (in-progress executions will be stopped when the CLI is forcefully stopped - if not set, you will be asked for confirmation) |
| --create-project | boolean | false | If the project doesn't exist, create one using the project flag as its name (or workflow/tool name if project flag is not set) |
| --machines | string | / | Specify the number of machines (format: small-medium-large). Examples: 1-1-1, 0-0-3 |
| --machines | string | / | Specify the number of machines. Use one value for default/self-hosted machines (--machines 3) or three values for small-medium-large (--machines 1-1-1) |
| --fleet | string | / | The name of the fleet to use to execute the workflow


#### Provide parameters using **config.yaml** file

Expand Down
3 changes: 1 addition & 2 deletions cmd/execute/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -620,8 +620,7 @@ func readConfigMachines(config *map[string]interface{}, isTool bool, maximumMach
if maxMachines {
return maximumMachines
} else {
execMachines = maximumMachines
setMachinesToMinimum(execMachines)
*execMachines = setMachinesToMinimum(*maximumMachines)
}
}
}
Expand Down
66 changes: 42 additions & 24 deletions cmd/execute/execute.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ var (
outputNodesFlag string
ci bool
createProject bool
fleetName string
)

// ExecuteCmd represents the execute command
Expand All @@ -67,7 +68,7 @@ var ExecuteCmd = &cobra.Command{
}
}

fleet = util.GetFleetInfo()
fleet = util.GetFleetInfo(fleetName)
if fleet == nil {
return
}
Expand All @@ -87,35 +88,51 @@ var ExecuteCmd = &cobra.Command{
allNodes, roots = CreateTrees(version, false)
if maxMachines {
executionMachines = version.MaxMachines
}
} else if machineConfiguration != "" {
machines := &types.Machines{}

if machineConfiguration != "" {
pattern := `^\d+-\d+-\d+$`
regex := regexp.MustCompile(pattern)
if regex.MatchString(machineConfiguration) {
parts := strings.Split(machineConfiguration, "-")
// Managed enterprise fleet, 3 types of machines, small-medium-large
if len(fleet.Machines) == 3 {
pattern := `^\d+-\d+-\d+$`
regex := regexp.MustCompile(pattern)

machines := &types.Machines{}
if regex.MatchString(machineConfiguration) {
parts := strings.Split(machineConfiguration, "-")

if small, err := strconv.Atoi(parts[0]); err == nil && small != 0 {
machines.Small = &small
}
if small, err := strconv.Atoi(parts[0]); err == nil && small != 0 {
machines.Small = &small
}

if medium, err := strconv.Atoi(parts[1]); err == nil && medium != 0 {
machines.Medium = &medium
}
if medium, err := strconv.Atoi(parts[1]); err == nil && medium != 0 {
machines.Medium = &medium
}

if large, err := strconv.Atoi(parts[2]); err == nil && large != 0 {
machines.Large = &large
if large, err := strconv.Atoi(parts[2]); err == nil && large != 0 {
machines.Large = &large
}
executionMachines = *machines
} else {
fmt.Printf("Invalid machine configuration \"%s\".\n", machineConfiguration)
fmt.Println("Please use the format: small-medium-large (e.g., 0-0-3)")
os.Exit(1)
}

executionMachines = *machines

} else {
fmt.Printf("Invalid machine configuration \"%s\".\n", machineConfiguration)
fmt.Println("Please use the format: small-medium-large (e.g., 0-0-3)")
os.Exit(0)
defaultOrSelfHosted, err := strconv.Atoi(machineConfiguration)
if err != nil {
fmt.Printf("Invalid machine configuration \"%s\".\n", machineConfiguration)
os.Exit(1)
}

if fleet.Type == "MANAGED" {
machines.Default = &defaultOrSelfHosted
} else if fleet.Type == "HOSTED" {
machines.SelfHosted = &defaultOrSelfHosted
}
}
executionMachines = *machines

} else {
executionMachines = setMachinesToMinimum(version.MaxMachines)
}

outputNodes := make([]string, 0)
Expand All @@ -130,7 +147,7 @@ var ExecuteCmd = &cobra.Command{
os.Exit(0)
}

createRun(version.ID, fleet.ID, watch, &executionMachines, outputNodes, outputsDirectory)
createRun(version.ID, fleet.ID, watch, outputNodes, outputsDirectory)
},
}

Expand All @@ -141,12 +158,13 @@ func init() {
ExecuteCmd.Flags().BoolVar(&showParams, "show-params", false, "Show parameters in the workflow tree")
// ExecuteCmd.Flags().StringVar(&workflowYAML, "file", "", "Workflow YAML file to execute")
ExecuteCmd.Flags().BoolVar(&maxMachines, "max", false, "Use maximum number of machines for workflow execution")
ExecuteCmd.Flags().StringVar(&machineConfiguration, "machines", "", "Specify the number of machines (format: small-medium-large). Examples: 1-1-1, 0-0-3")
ExecuteCmd.Flags().StringVar(&machineConfiguration, "machines", "", "Specify the number of machines. Use one value for default/self-hosted machines (--machines 3) or three values for small-medium-large (--machines 1-1-1)")
ExecuteCmd.Flags().BoolVar(&downloadAllNodes, "output-all", false, "Download all outputs when the execution is finished")
ExecuteCmd.Flags().StringVar(&outputNodesFlag, "output", "", "A comma separated list of nodes which outputs should be downloaded when the execution is finished")
ExecuteCmd.Flags().StringVar(&outputsDirectory, "output-dir", "", "Path to directory which should be used to store outputs")
ExecuteCmd.Flags().BoolVar(&ci, "ci", false, "Run in CI mode (in-progreess executions will be stopped when the CLI is forcefully stopped - if not set, you will be asked for confirmation)")
ExecuteCmd.Flags().BoolVar(&createProject, "create-project", false, "If the project doesn't exist, create it using the project flag as its name (or workflow name if not set)")
ExecuteCmd.Flags().StringVar(&fleetName, "fleet", "", "The name of the fleet to use to execute the workflow")
}

func readWorkflowYAMLandCreateVersion(fileName string, workflowName string, objectPath string) *types.WorkflowVersionDetailed {
Expand Down
125 changes: 70 additions & 55 deletions cmd/execute/helpers.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,11 +98,10 @@ func getScripts(pageSize int, search string, name string) []types.Script {
return scripts.Results
}

func createRun(versionID, fleetID uuid.UUID, watch bool, machines *types.Machines, outputNodes []string, outputsDir string) {
func createRun(versionID, fleetID uuid.UUID, watch bool, outputNodes []string, outputsDir string) {

run := types.CreateRun{
VersionID: versionID,
Vault: fleet.Vault,
Machines: executionMachines,
Fleet: &fleetID,
}
Expand All @@ -119,18 +118,6 @@ func createRun(versionID, fleetID uuid.UUID, watch bool, machines *types.Machine
os.Exit(0)
}

if resp.Status() != http.StatusCreated {
run.Fleet = nil

data, err := json.Marshal(run)
if err != nil {
fmt.Println("Error encoding create run request!")
os.Exit(0)
}

resp = request.Trickest.Post().Body(data).DoF("execution/")
}

if resp.Status() != http.StatusCreated {
request.ProcessUnexpectedResponse(resp)
}
Expand All @@ -151,9 +138,9 @@ func createRun(versionID, fleetID uuid.UUID, watch bool, machines *types.Machine
if watch {
WatchRun(createRunResp.ID, outputsDir, nodesToDownload, nil, false, &executionMachines, showParams)
} else {
availableMachines := GetAvailableMachines()
availableMachines := GetAvailableMachines(fleetName)
fmt.Println("Run successfully created! ID: " + createRunResp.ID.String())
fmt.Print("Machines:\n" + FormatMachines(*machines, false))
fmt.Print("Machines:\n" + FormatMachines(executionMachines, false))
fmt.Print("\nAvailable:\n" + FormatMachines(availableMachines, false))
}
}
Expand Down Expand Up @@ -193,7 +180,8 @@ func createNewVersion(version *types.WorkflowVersionDetailed) *types.WorkflowVer
return nil
}

newVersion := output.GetWorkflowVersionByID(newVersionInfo.ID)
fleet := util.GetFleetInfo(fleetName)
newVersion := output.GetWorkflowVersionByID(newVersionInfo.ID, fleet.ID)
return newVersion
}

Expand Down Expand Up @@ -396,8 +384,8 @@ func processInvalidInputType(newPNode, existingPNode types.PrimitiveNode) {
os.Exit(0)
}

func GetAvailableMachines() types.Machines {
hiveInfo := util.GetFleetInfo()
func GetAvailableMachines(fleetName string) types.Machines {
hiveInfo := util.GetFleetInfo(fleetName)
availableMachines := types.Machines{}
for _, machine := range hiveInfo.Machines {
if machine.Name == "small" {
Expand All @@ -412,6 +400,14 @@ func GetAvailableMachines() types.Machines {
available := machine.Total - machine.Running
availableMachines.Large = &available
}
if machine.Name == "default" {
available := machine.Total - machine.Running
availableMachines.Default = &available
}
if machine.Name == "self_hosted" {
available := machine.Total - machine.Running
availableMachines.SelfHosted = &available
}
}
return availableMachines
}
Expand Down Expand Up @@ -477,7 +473,7 @@ func stopRun(runID uuid.UUID) {
}
}

func setMachinesToMinimum(machines *types.Machines) {
func setMachinesToMinimum(machines types.Machines) types.Machines {
if machines.Small != nil {
*machines.Small = 1
}
Expand All @@ -487,50 +483,50 @@ func setMachinesToMinimum(machines *types.Machines) {
if machines.Large != nil {
*machines.Large = 1
}
if machines.Default != nil {
*machines.Default = 1
}
if machines.SelfHosted != nil {
*machines.SelfHosted = 1
}

return machines
}

func FormatMachines(machines types.Machines, inline bool) string {
var small, medium, large string
if machines.Small != nil {
small = "small: " + strconv.Itoa(*machines.Small)
}
if machines.Medium != nil {
medium = "medium: " + strconv.Itoa(*machines.Medium)
}
if machines.Large != nil {
large = "large: " + strconv.Itoa(*machines.Large)
}
smallMachines := formatSize("small", machines.Small)
mediumMachines := formatSize("medium", machines.Medium)
largeMachines := formatSize("large", machines.Large)
selfHostedMachines := formatSize("self hosted", machines.SelfHosted)
defaultMachines := formatSize("default", machines.Default)

out := ""
var out string
if inline {
if small != "" {
out = small
}
if medium != "" {
if small != "" {
out += ", "
}
out += medium
}
if large != "" {
if small != "" || medium != "" {
out += ", "
}
out += large
}
out = joinNonEmptyValues(", ", smallMachines, mediumMachines, largeMachines, selfHostedMachines, defaultMachines)
} else {
if small != "" {
out = " " + small + "\n "
}
if medium != "" {
out += medium + "\n "
}
if large != "" {
out += large + "\n"
out = joinNonEmptyValues("\n ", " "+smallMachines, mediumMachines, largeMachines, selfHostedMachines, defaultMachines)
}
return out
}

func formatSize(sizeName string, size *int) string {
if size != nil {
return sizeName + ": " + strconv.Itoa(*size)
}
return ""
}

func joinNonEmptyValues(separator string, values ...string) string {
var nonEmptyValues []string

for _, value := range values {
if value != "" {
nonEmptyValues = append(nonEmptyValues, value)
}
}

return out
result := strings.Join(nonEmptyValues, separator)
return result
}

func getNodeNameFromConnectionID(id string) string {
Expand Down Expand Up @@ -596,6 +592,25 @@ func uploadFilesIfNeeded(primitiveNodes map[string]*types.PrimitiveNode) {
}

func maxMachinesTypeCompatible(machines, maxMachines types.Machines) bool {
if machines.Default != nil && *machines.Default > *maxMachines.Default {
return false
}
if machines.SelfHosted != nil && *machines.SelfHosted > *maxMachines.SelfHosted {
return false
}

if machines.Small != nil && *machines.Small > *maxMachines.Small {
return false
}

if machines.Medium != nil && *machines.Medium > *maxMachines.Medium {
return false
}

if machines.Large != nil && *machines.Large > *maxMachines.Large {
return false
}

if (machines.Small != nil && maxMachines.Small == nil) ||
(machines.Medium != nil && maxMachines.Medium == nil) ||
(machines.Large != nil && maxMachines.Large == nil) {
Expand Down
2 changes: 1 addition & 1 deletion cmd/execute/watch.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ func WatchRun(runID uuid.UUID, downloadPath string, nodesToDownload map[string]o
out := ""
out += fmt.Sprintf(fmtStr, "Name:", run.WorkflowName)
out += fmt.Sprintf(fmtStr, "Status:", strings.ToLower(run.Status))
availableMachines := GetAvailableMachines()
availableMachines := GetAvailableMachines(fleetName)
out += fmt.Sprintf(fmtStr, "Machines:", FormatMachines(*machines, true)+
" (currently available: "+FormatMachines(availableMachines, true)+")")
out += fmt.Sprintf(fmtStr, "Created:", run.CreatedDate.In(time.Local).Format(time.RFC1123)+
Expand Down
3 changes: 0 additions & 3 deletions cmd/get/get.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,9 +85,6 @@ var GetCmd = &cobra.Command{
status = strings.ToLower(runs[0].Status)
}
out += fmt.Sprintf(fmtStr, "Status:", status)
availableBees := execute.GetAvailableMachines()
out += fmt.Sprintf(fmtStr, "Max machines:", execute.FormatMachines(version.MaxMachines, true)+
" (currently available: "+execute.FormatMachines(availableBees, true)+")")
out += fmt.Sprintf(fmtStr, "Created:", workflow.CreatedDate.In(time.Local).Format(time.RFC1123)+
" ("+util.FormatDuration(time.Since(workflow.CreatedDate))+" ago)")

Expand Down
Loading

0 comments on commit 8655a59

Please sign in to comment.