oc-schedulerd/argo_builder.go

277 lines
7.6 KiB
Go
Raw Normal View History

2024-05-02 09:52:28 +02:00
// A class that translates the informations held in the graph object
// via its lists of components into an argo file, using the a list of
// link ID to build the dag
package main
import (
"fmt"
2024-05-03 10:58:06 +02:00
"os"
2024-05-02 09:52:28 +02:00
"slices"
"strings"
2024-07-02 12:24:23 +02:00
"time"
2024-05-02 09:52:28 +02:00
"github.com/beego/beego/v2/core/logs"
2024-05-07 19:45:20 +02:00
"github.com/nwtgck/go-fakelish"
2024-05-02 09:52:28 +02:00
"gopkg.in/yaml.v3"
)
type ArgoBuilder struct {
2024-05-03 10:58:06 +02:00
graph Graph
2024-05-02 09:52:28 +02:00
branches [][]string
2024-05-03 10:58:06 +02:00
Workflow Workflow
2024-05-02 09:52:28 +02:00
}
2024-05-03 10:58:06 +02:00
type Workflow struct {
2024-05-07 19:45:20 +02:00
ApiVersion string `yaml:"apiVersion"`
Kind string `yaml:"kind"`
Metadata struct {
GenerateName string `yaml:"generateName"`
} `yaml:"metadata"`
Spec Spec `yaml:"spec,omitempty"`
2024-05-02 09:52:28 +02:00
}
2024-05-07 19:45:20 +02:00
type Spec struct {
Entrypoint string `yaml:"entrypoint"`
Arguments []Parameter `yaml:"arguments,omitempty"`
Volumes []VolumeClaimTemplate `yaml:"volumeClaimTemplates,omitempty"`
Templates []Template `yaml:"templates"`
}
2024-05-02 09:52:28 +02:00
func (b *ArgoBuilder) CreateDAG() bool {
fmt.Println("list of branches : ", b.branches)
b.createTemplates()
2024-05-02 09:52:28 +02:00
b.createDAGstep()
2024-05-07 19:45:20 +02:00
b.createVolumes()
b.Workflow.Spec.Entrypoint = "dag"
b.Workflow.ApiVersion = "argoproj.io/v1alpha1"
b.Workflow.Kind = "Workflow"
2024-07-02 12:24:23 +02:00
random_name := generateWfName()
b.Workflow.Metadata.GenerateName = "oc-test-" + random_name
2024-05-07 19:45:20 +02:00
2024-05-03 10:58:06 +02:00
yamlified, err := yaml.Marshal(b.Workflow)
2024-05-02 09:52:28 +02:00
if err != nil {
2024-05-03 10:58:06 +02:00
logs.Error("Could not transform object to yaml file")
return false
2024-05-02 09:52:28 +02:00
}
2024-07-02 12:24:23 +02:00
// Give a unique name to each argo file with its timestamp DD:MM:YYYY_hhmmss
current_timestamp := time.Now().Format("02_01_2006_150405")
file_name := random_name + "_" + current_timestamp + ".yml"
workflows_dir := "workflows_argo/"
err = os.WriteFile(workflows_dir + file_name , []byte(yamlified), 0660)
2024-05-07 19:45:20 +02:00
if err != nil {
2024-05-03 10:58:06 +02:00
logs.Error("Could not write the yaml file")
return false
2024-05-03 10:58:06 +02:00
}
2024-07-02 12:24:23 +02:00
fmt.Println("Created " + file_name)
2024-05-02 09:52:28 +02:00
return true
}
func (b *ArgoBuilder) createTemplates() {
2024-05-02 09:52:28 +02:00
for _, comp := range b.graph.Computings{
image_name := strings.Split(comp.Command," ")[0] // TODO : decide where to store the image name, GUI or models.computing.Image
temp_container := Container{Image: image_name} // TODO : decide where to store the image name, GUI or models.computing.Image
temp_container.Command = getComputingCommands(comp.Command)
temp_container.Args = getComputingArgs(comp.Arguments,comp.Command)
input_names := getComputingEnvironmentName(comp.Environment)
var inputs_container []Parameter
for _, name := range input_names {
inputs_container = append(inputs_container, Parameter{Name: name})
}
argo_name := getArgoName(comp.Name,comp.ID)
new_temp := Template{Name: argo_name, Container: temp_container}
new_temp.Inputs.Parameters = inputs_container
2024-05-07 19:45:20 +02:00
new_temp.Container.VolumeMounts = append(new_temp.Container.VolumeMounts, VolumeMount{Name: "workdir",MountPath: "/mnt/vol"}) // TODO : replace this with a search of the storage / data source name
b.Workflow.Spec.Templates = append(b.Workflow.Spec.Templates, new_temp)
}
}
2024-05-02 09:52:28 +02:00
func (b *ArgoBuilder) createDAGstep() {
2024-05-03 10:58:06 +02:00
new_dag := Dag{}
2024-05-02 09:52:28 +02:00
for _, comp := range b.graph.Computings{
unique_name := getArgoName(comp.Name,comp.ID)
2024-05-03 10:58:06 +02:00
step := Task{Name: unique_name, Template: unique_name}
comp_envs := getComputingEnvironment(comp.Environment)
for name, value := range comp_envs {
step.Arguments.Parameters = append(step.Arguments.Parameters, Parameter{Name: name, Value: value})
}
2024-05-03 10:58:06 +02:00
// retrieves the name (computing.name-computing.ID)
step.Dependencies = b.getDependency(comp.ID)
2024-05-03 10:58:06 +02:00
new_dag.Tasks = append(new_dag.Tasks, step)
2024-05-02 09:52:28 +02:00
}
2024-05-03 10:58:06 +02:00
2024-05-07 19:45:20 +02:00
b.Workflow.Spec.Templates = append (b.Workflow.Spec.Templates, Template{Name: "dag", Dag: new_dag})
}
2024-05-03 10:58:06 +02:00
2024-05-07 19:45:20 +02:00
func (b *ArgoBuilder) createVolumes() {
// For testing purposes we only declare one volume, mounted in each computing
new_volume := VolumeClaimTemplate{}
new_volume.Metadata.Name = "workdir"
new_volume.Spec.AccessModes = []string{"ReadWriteOnce"}
new_volume.Spec.Resources.Requests.Storage = "1Gi"
b.Workflow.Spec.Volumes = append(b.Workflow.Spec.Volumes, new_volume)
2024-05-02 09:52:28 +02:00
}
func (b *ArgoBuilder) getDependency(current_computing_id string) (dependencies []string) {
var dependencies_id []string
2024-05-02 09:52:28 +02:00
for _, link := range b.graph.Links {
if current_computing_id == link.Destination && b.graph.getComponentType(link.Source) == "computing" && !slices.Contains(dependencies_id,link.Source) {
dependencies_id = append(dependencies_id, link.Source)
2024-05-02 09:52:28 +02:00
}
}
for _, dependency := range dependencies_id {
dependency_name := getArgoName(b.graph.getComponentName(dependency),dependency)
dependencies = append(dependencies, dependency_name)
}
2024-05-02 09:52:28 +02:00
return
2024-05-02 09:52:28 +02:00
}
// func (b *ArgoBuilder) componentInBranch(component_id string, branch []string) bool {
// for _, link := range branch {
// if b.graph.Links[link].Source == component_id || b.graph.Links[link].Destination == component_id {
// return true
// }
// }
// return false
// }
2024-05-02 09:52:28 +02:00
// func (b *ArgoBuilder) findPreviousComputing(computing_id string, branch []string, index int) string {
// for i := index; i >= 0 ; i-- {
// previousLink := b.graph.Links[branch[i]]
// if previousLink.Source != computing_id && b.graph.getComponentType(previousLink.Source) == "computing"{
// name := getArgoName(b.graph.getComponentName(previousLink.Source),previousLink.Source)
// return name
// }
// if previousLink.Destination != computing_id && b.graph.getComponentType(previousLink.Destination) == "computing"{
// name := getArgoName(b.graph.getComponentName(previousLink.Destination),previousLink.Destination)
// return name
// }
// }
// return ""
// }
func getComputingCommands(user_input string) (list_command []string) {
user_input = removeImageName(user_input)
if len(user_input) == 0 {
return
}
list_command = strings.Split(user_input, " ")
for i := range list_command {
list_command[i] = list_command[i]
}
return
}
func getComputingArgs(user_input []string, command string) (list_args []string) {
if len(user_input) == 0 {
return
}
// quickfix that might need improvement
if(strings.Contains(command,"sh -c")){
list_args = append(list_args, strings.Join(user_input," "))
return
}
for _, arg := range user_input{
list_args = append(list_args, arg)
}
return
}
// Currently implements code to overcome problems in data structure
func getComputingEnvironment(user_input []string) (map_env map[string]string) {
if len(user_input) == 0 {
return
}
if(len(user_input) == 1){
user_input = strings.Split(user_input[0],",")
}
map_env = make(map[string]string,0)
for _, str := range user_input {
new_pair := strings.Split(str,"=")
if(len(new_pair) != 2) {
logs.Error("Error extracting the environment variable from ", str)
panic(0)
}
map_env[new_pair[0]] = new_pair[1]
}
return
}
func getComputingEnvironmentName(user_input []string) (list_names []string){
env_map := getComputingEnvironment(user_input)
for name := range env_map {
list_names = append(list_names, name)
}
return
2024-05-07 19:45:20 +02:00
}
func generateWfName() (Name string){
2024-05-07 19:45:20 +02:00
Name = fakelish.GenerateFakeWord(5, 8) + "-" + fakelish.GenerateFakeWord(5, 8)
return
}
func getArgoName(raw_name string, component_id string) (formatedName string){
formatedName = strings.ReplaceAll(raw_name," ","-")
formatedName += "-" + component_id
formatedName = strings.ToLower(formatedName)
return
2024-05-07 19:45:20 +02:00
}
func printYAML(data interface{}) {
yamlData, err := yaml.Marshal(data)
if err != nil {
fmt.Printf("Error marshalling YAML: %v\n", err)
return
}
fmt.Println(string(yamlData))
}
func removeImageName(user_input string) string {
// First command is the name of the container for now
if len(strings.Split(user_input, " ")) == 1 {
return ""
}
slice_input := strings.Split(user_input, " ")
new_slice := slice_input[1:]
user_input = strings.Join(new_slice," ")
return user_input
2024-05-02 09:52:28 +02:00
}