├── README.md ├── hpc.go.save └── hpc.go /README.md: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /hpc.go.save: -------------------------------------------------------------------------------- 1 | package hpc 2 | 3 | import ( 4 | "bufio" 5 | "bytes" 6 | "errors" 7 | "fmt" 8 | "io" 9 | "io/ioutil" 10 | "os" 11 | "os/exec" 12 | "strconv" 13 | "strings" 14 | "syscall" 15 | "time" 16 | ) 17 | 18 | type Job struct { 19 | ScriptContents string 20 | NativeSpecs []string 21 | UID int 22 | GID int 23 | OutputScriptPth string 24 | } 25 | 26 | //This function exists to help the native spec system 27 | func Contains(illegalArgs []string, elementToTest string) bool { 28 | for _, illegalArg := range illegalArgs { 29 | if elementToTest == illegalArg { 30 | return true 31 | } 32 | if strings.Contains(elementToTest, illegalArg) { 33 | return true 34 | } 35 | } 36 | return false 37 | } 38 | 39 | func RemoveIllegalParams(input []string, illegalParams []string) []string { 40 | skip := false 41 | var output []string 42 | for i, parameter := range input { 43 | if skip { 44 | skip = false 45 | continue 46 | } 47 | 48 | if Contains(illegalParams, parameter) { 49 | if !(i+1 > len(input)-1) { 50 | if strings.HasPrefix(input[i+1], "-") { 51 | skip = false 52 | continue 53 | } 54 | } 55 | 56 | skip = true 57 | continue 58 | } 59 | output = append(output, parameter) 60 | 61 | } 62 | return output 63 | } 64 | 65 | func BuildScript(cmd, filenameSuffix string, myUid, myGid int, pth string) (err error, scriptPath string) { 66 | uniqueID := strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10) 67 | time.Sleep(50 * time.Millisecond) 68 | 69 | os.MkdirAll(fmt.Sprint(pth,"/scripts"), 0740) 70 | os.Chown(fmt.Sprint(pth,"/scripts"), myUid, myGid) 71 | 72 | batchScriptFull := pth + "/scripts/" + filenameSuffix + uniqueID + ".bash" 73 | batchScript, scriptErr := os.Create(batchScriptFull) 74 | if scriptErr != nil { 75 | return fmt.Errorf("Unable to create Script file: %v", scriptErr), "" 76 | } 77 | fmt.Fprintf(batchScript, "#!/bin/bash\n") 78 | fmt.Fprintf(batchScript, cmd) 79 | 80 | chmodErr := os.Chmod(batchScriptFull, 0750) 81 | chownErr := os.Chown(batchScriptFull, myUid, myGid) 82 | 83 | if chmodErr != nil || chownErr != nil { 84 | return fmt.Errorf("Could not set the script to be executable: %v", chmodErr), "" 85 | } 86 | batchScript.Close() 87 | 88 | return nil, batchScriptFull 89 | 90 | } 91 | 92 | func DetectBatchSystem() (num int) { 93 | osEnv := os.Environ() 94 | counter := 0 95 | for _, entry := range osEnv { 96 | if strings.Contains(entry, "LSF_BINDIR") { 97 | counter = counter + 1 98 | } 99 | } 100 | 101 | _, err := exec.LookPath("sbatch") 102 | if err == nil { 103 | counter = counter + 3 104 | } 105 | 106 | //if COBALT 107 | // counter = counter + 6 108 | 109 | return counter 110 | //1 = LSF 111 | //3 = Slurm 112 | //6 = Cobalt 113 | //4 = LSF + Slurm 114 | //10 = LSF + Slurm + Cobalt 115 | //9 = Slurm + Cobalt 116 | 117 | } 118 | 119 | func (j *Job) Run() (err error, out string) { 120 | switch DetectBatchSystem() { 121 | case 1: 122 | return RunLSF(j) 123 | case 3: 124 | return RunSlurm(j) 125 | default: 126 | return errors.New("Batch System Detection Error"), "" 127 | } 128 | return nil, "" 129 | } 130 | 131 | func RunLSF(j *Job) (err error, out string) { 132 | err, Script := BuildScript(j.ScriptContents, "batch_script", j.UID, j.GID, j.OutputScriptPth) 133 | 134 | if err != nil { 135 | return err, "" 136 | } 137 | 138 | file, err := ioutil.ReadFile(Script) 139 | if err != nil { 140 | return err, "" 141 | } 142 | 143 | fileText := string(file) 144 | 145 | var cmd *exec.Cmd 146 | 147 | if strings.Contains(fileText, "clone") { 148 | cmd = exec.Command("/bin/bash", Script) 149 | } else { 150 | outputScriptPath := fmt.Sprint(j.OutputScriptPth, "/lsf_out.log") 151 | errorScriptPath := fmt.Sprint(j.OutputScriptPth, "/lsf_err.log") 152 | 153 | var Specs []string 154 | if len(j.NativeSpecs) != 0 { 155 | //Defines an array of illegal arguments which will not be passed in as native specifications 156 | illegalArguments := []string{"-e","-o","-eo"} 157 | Specs = RemoveIllegalParams(j.NativeSpecs, illegalArguments) 158 | } 159 | 160 | cmd = exec.Command("bsub", "-o", outputScriptPath, "-e", errorScriptPath, strings.Join(Specs, " "), Script) 161 | } 162 | cmd.SysProcAttr = &syscall.SysProcAttr{} 163 | cmd.SysProcAttr.Credential = &syscall.Credential{Uid: uint32(j.UID), Gid: uint32(j.GID)} 164 | cmd.Env = append(os.Environ()) 165 | 166 | var stdBuffer bytes.Buffer 167 | mw := io.MultiWriter(os.Stdout, &stdBuffer) 168 | cmd.Stdout = mw 169 | cmd.Stderr = mw 170 | if err := cmd.Run(); err != nil { 171 | return err, "" 172 | } 173 | 174 | var commandOut string 175 | 176 | if !strings.Contains(fileText, "clone") { 177 | err, commandOut = GetOutput(Script,fmt.Sprint(j.OutputScriptPth,"/lsf_out.log")) 178 | if err != nil { 179 | return err, "" 180 | } 181 | } 182 | 183 | return nil, commandOut 184 | } 185 | 186 | func RunSlurm(j *Job) (err error, out string) { 187 | 188 | err, Script := BuildScript(j.ScriptContents, "batch_script", j.UID, j.GID, j.OutputScriptPth) 189 | if err != nil { 190 | return err, "" 191 | } 192 | 193 | 194 | file, err := ioutil.ReadFile(Script) 195 | if err != nil { 196 | return err, "" 197 | } 198 | 199 | fileText := string(file) 200 | 201 | var cmd *exec.Cmd 202 | 203 | if strings.Contains(fileText, "clone") { 204 | cmd = exec.Command("/bin/bash", Script) 205 | } else { 206 | outputScriptPath := fmt.Sprint(j.OutputScriptPth, "/slurm_out.log") 207 | 208 | var Specs []string 209 | if len(j.NativeSpecs) != 0 { 210 | //Defines an array of illegal arguments which will not be passed in as native specifications 211 | illegalArguments := []string{"-o"} 212 | Specs = RemoveIllegalParams(j.NativeSpecs, illegalArguments) 213 | } 214 | if len(Specs) != 0{ 215 | cmd = exec.Command("sbatch","-o",outputScriptPath, Script) 216 | }else { 217 | cmd = exec.Command("sbatch","-o",outputScriptPath,strings.Join(Specs, " "), Script) 218 | } 219 | } 220 | cmd.SysProcAttr = &syscall.SysProcAttr{} 221 | cmd.SysProcAttr.Credential = &syscall.Credential{Uid: uint32(j.UID), Gid: uint32(j.GID)} 222 | cmd.Env = append(os.Environ()) 223 | 224 | var stdBuffer bytes.Buffer 225 | mw := io.MultiWriter(os.Stdout, &stdBuffer) 226 | cmd.Stdout = mw 227 | cmd.Stderr = mw 228 | 229 | os.Remove(fmt.Sprint(j.OutputScriptPth, "/slurm_out.log")) 230 | 231 | if err := cmd.Run(); err != nil { 232 | return err, "" 233 | } 234 | 235 | var commandOut string 236 | 237 | if !strings.Contains(fileText, "clone") { 238 | fmt.Println(Script) 239 | err, commandOut = GetOutputSlurm(fmt.Sprint(j.OutputScriptPth, "/slurm_out.log")) 240 | if err != nil {return err, ""} 241 | } 242 | return nil, commandOut 243 | } 244 | 245 | func GetOutputSlurm(outputFile string) (err error, output string){ 246 | retry := true 247 | var file []byte 248 | 249 | for retry { 250 | _, err := os.Open(outputFile) 251 | if err != nil{ 252 | continue 253 | } else { 254 | file, err = ioutil.ReadFile(outputFile) 255 | if err != nil { 256 | return err, "" 257 | } 258 | fmt.Println(string(file)) 259 | os.Remove(outputFile) 260 | return nil, string(file) 261 | } 262 | } 263 | } 264 | 265 | func GetOutput(scriptName string, outputFile string) (err error, output string) { 266 | retry := true 267 | 268 | for retry { 269 | 270 | file, err := os.Open(outputFile) 271 | if err != nil { 272 | continue 273 | } 274 | scanner := bufio.NewScanner(file) 275 | for scanner.Scan() { 276 | line := scanner.Text() 277 | if strings.Contains(line, scriptName) { 278 | retry = false 279 | } 280 | 281 | } 282 | file.Close() 283 | file = nil 284 | } 285 | 286 | file, err := os.Open(outputFile) 287 | if err != nil { 288 | return err, "" 289 | } 290 | 291 | var lineArray []string 292 | var subLineArray []string 293 | var startingLine int = 0 294 | var endingLine int 295 | 296 | scanner := bufio.NewScanner(file) 297 | for scanner.Scan() { 298 | lineArray = append(lineArray, scanner.Text()) 299 | } 300 | if err := scanner.Err(); err != nil { 301 | return err, "" 302 | } 303 | 304 | for i, line := range lineArray { 305 | if strings.Contains(line, scriptName) { 306 | startingLine = i 307 | break 308 | } 309 | } 310 | 311 | for i, line := range lineArray[startingLine:] { 312 | if strings.Contains(line, "PS:") { 313 | endingLine = startingLine + i 314 | break 315 | } 316 | } 317 | 318 | for _, line := range lineArray[startingLine+34 : endingLine-1] { 319 | subLineArray = append(subLineArray, line) 320 | } 321 | 322 | file.Close() 323 | file = nil 324 | return nil, strings.Join(subLineArray, "\n") 325 | 326 | } 327 | -------------------------------------------------------------------------------- /hpc.go: -------------------------------------------------------------------------------- 1 | package hpc 2 | 3 | import ( 4 | "bufio" 5 | "bytes" 6 | "errors" 7 | "fmt" 8 | "io" 9 | "io/ioutil" 10 | "os" 11 | "os/exec" 12 | "strconv" 13 | "strings" 14 | "syscall" 15 | "time" 16 | ) 17 | 18 | type Job struct { 19 | ScriptContents string 20 | NativeSpecs []string 21 | Bank string 22 | UID int 23 | GID int 24 | OutputScriptPth string 25 | BatchExecution bool 26 | } 27 | 28 | //This function exists to help the native spec system 29 | func Contains(illegalArgs []string, elementToTest string) bool { 30 | for _, illegalArg := range illegalArgs { 31 | if elementToTest == illegalArg { 32 | return true 33 | } 34 | if strings.Contains(elementToTest, illegalArg) { 35 | return true 36 | } 37 | } 38 | return false 39 | } 40 | 41 | //Removes specified arguments 42 | func RemoveIllegalParams(input []string, illegalParams []string) []string { 43 | skip := false 44 | var output []string 45 | for i, parameter := range input { 46 | if skip { 47 | skip = false 48 | continue 49 | } 50 | 51 | if Contains(illegalParams, parameter) { 52 | if !(i+1 > len(input)-1) { 53 | if strings.HasPrefix(input[i+1], "-") { 54 | skip = false 55 | continue 56 | } 57 | } 58 | 59 | skip = true 60 | continue 61 | } 62 | output = append(output, parameter) 63 | 64 | } 65 | return output 66 | } 67 | 68 | //Creates a script returns the absolute path 69 | func BuildScript(cmd, filenameSuffix string, myUid, myGid int, pth string) (err error, scriptPath string) { 70 | uniqueID := strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10) 71 | time.Sleep(50 * time.Millisecond) 72 | 73 | os.MkdirAll(fmt.Sprint(pth, "/scripts"), 0740) 74 | os.Chown(fmt.Sprint(pth, "/scripts"), myUid, myGid) 75 | 76 | batchScriptFull := pth + "/scripts/" + filenameSuffix + uniqueID + ".bash" 77 | batchScript, scriptErr := os.Create(batchScriptFull) 78 | if scriptErr != nil { 79 | return fmt.Errorf("Unable to create Script file: %v", scriptErr), "" 80 | } 81 | fmt.Fprintf(batchScript, "#!/bin/bash\n") 82 | fmt.Fprintf(batchScript, cmd) 83 | 84 | chmodErr := os.Chmod(batchScriptFull, 0750) 85 | chownErr := os.Chown(batchScriptFull, myUid, myGid) 86 | 87 | if chmodErr != nil || chownErr != nil { 88 | return fmt.Errorf("Could not set the script to be executable: %v", chmodErr), "" 89 | } 90 | batchScript.Close() 91 | 92 | return nil, batchScriptFull 93 | 94 | } 95 | 96 | //Finds which batch system is installed on the system returns a specific number 97 | func DetectBatchSystem() (num int) { 98 | osEnv := os.Environ() 99 | counter := 0 100 | for _, entry := range osEnv { 101 | if strings.Contains(entry, "LSF_BINDIR") { 102 | counter = counter + 1 103 | } 104 | } 105 | 106 | _, err := exec.LookPath("sbatch") 107 | if err == nil { 108 | counter = counter + 3 109 | } 110 | 111 | //if COBALT 112 | // counter = counter + 6 113 | 114 | return counter 115 | //1 = LSF 116 | //3 = Slurm 117 | //6 = Cobalt 118 | //4 = LSF + Slurm 119 | //10 = LSF + Slurm + Cobalt 120 | //9 = Slurm + Cobalt 121 | 122 | } 123 | 124 | //Initial func run. Gets batch system and calls the corresponding run 125 | func (j *Job) Run() (err error, out string) { 126 | switch DetectBatchSystem() { 127 | case 1: 128 | return RunLSF(j) 129 | case 3: 130 | return RunSlurm(j) 131 | default: 132 | return errors.New("Batch System Detection Error"), "" 133 | } 134 | return nil, "" 135 | } 136 | 137 | func RunLSF(j *Job) (err error, out string) { 138 | //Create Script Check for Errors 139 | err, Script := BuildScript(j.ScriptContents, "batch_script", j.UID, j.GID, j.OutputScriptPth) 140 | if err != nil { 141 | return err, "" 142 | } 143 | 144 | //Create empty command var 145 | var cmd *exec.Cmd 146 | 147 | //Determin if script to be run should be done locally or through the batch system 148 | if j.BatchExecution == false { 149 | cmd = exec.Command("/bin/bash", Script) 150 | } else { 151 | //Get output script paths 152 | outputScriptPath := fmt.Sprint(j.OutputScriptPth, "/lsf_out.log") 153 | errorScriptPath := fmt.Sprint(j.OutputScriptPth, "/lsf_err.log") 154 | //Hancle Native Specs 155 | var Specs []string 156 | if len(j.NativeSpecs) != 0 { 157 | //Defines an array of illegal arguments which will not be passed in as native specifications 158 | illegalArguments := []string{"-e", "-o", "-eo"} 159 | Specs = RemoveIllegalParams(j.NativeSpecs, illegalArguments) 160 | } 161 | 162 | //Assemle bash command 163 | if j.Bank == "" { 164 | cmd = exec.Command("bsub", "-o", outputScriptPath, "-e", errorScriptPath, strings.Join(Specs, " "), Script) 165 | } else { 166 | cmd = exec.Command("bsub","-G", j.Bank , "-o", outputScriptPath, "-e", errorScriptPath, strings.Join(Specs, " "), Script) 167 | } 168 | } 169 | //Assign setUID information and env. vars 170 | cmd.SysProcAttr = &syscall.SysProcAttr{} 171 | cmd.SysProcAttr.Credential = &syscall.Credential{Uid: uint32(j.UID), Gid: uint32(j.GID)} 172 | cmd.Env = append(os.Environ()) 173 | 174 | //Handle Std out and Std err 175 | var stdBuffer bytes.Buffer 176 | mw := io.MultiWriter(os.Stdout, &stdBuffer) 177 | cmd.Stdout = mw 178 | cmd.Stderr = mw 179 | //Run the command, check for errors 180 | if err := cmd.Run(); err != nil { 181 | return err, "" 182 | } 183 | 184 | //Create empty output var 185 | var commandOut string 186 | 187 | //If command was run with Batch system get output 188 | if j.BatchExecution == true { 189 | err, commandOut = GetOutput(Script, fmt.Sprint(j.OutputScriptPth, "/lsf_out.log")) 190 | if err != nil { 191 | return err, "" 192 | } 193 | } 194 | //Return output 195 | return nil, commandOut 196 | } 197 | 198 | func RunSlurm(j *Job) (err error, out string) { 199 | //Create Script Check for Errors 200 | err, Script := BuildScript(j.ScriptContents, "batch_script", j.UID, j.GID, j.OutputScriptPth) 201 | if err != nil { 202 | return err, "" 203 | } 204 | 205 | //Open script and get its contents 206 | file, err := ioutil.ReadFile(Script) 207 | if err != nil { 208 | return err, "" 209 | } 210 | fileText := string(file) 211 | 212 | //Create empty command var 213 | var cmd *exec.Cmd 214 | 215 | //Determin if script to be run should be done locally or through the batch system 216 | if strings.Contains(fileText, "clone") { 217 | cmd = exec.Command("/bin/bash", Script) 218 | } else { 219 | //Get output script paths 220 | outputScriptPath := fmt.Sprint(j.OutputScriptPth, "/slurm_out.log") 221 | //Hancle Native Specs 222 | var Specs []string 223 | if len(j.NativeSpecs) != 0 { 224 | //Defines an array of illegal arguments which will not be passed in as native specifications 225 | illegalArguments := []string{"-o"} 226 | Specs = RemoveIllegalParams(j.NativeSpecs, illegalArguments) 227 | } 228 | //If native specs were defined attach them to the end. Assemble bash command 229 | if j.Bank == ""{ 230 | if len(Specs) != 0 { 231 | cmd = exec.Command("sbatch", "-o", outputScriptPath, Script) 232 | } else { 233 | cmd = exec.Command("sbatch", "-o", outputScriptPath, strings.Join(Specs, " "), Script) 234 | } 235 | } else { 236 | if len(Specs) != 0 { 237 | cmd = exec.Command("sbatch","-A",j.Bank, "-o", outputScriptPath, Script) 238 | } else { 239 | cmd = exec.Command("sbatch","-A",j.Bank, "-o", outputScriptPath, strings.Join(Specs, " "), Script) 240 | } 241 | 242 | } 243 | } 244 | //Assign setUID information and env. vars 245 | cmd.SysProcAttr = &syscall.SysProcAttr{} 246 | cmd.SysProcAttr.Credential = &syscall.Credential{Uid: uint32(j.UID), Gid: uint32(j.GID)} 247 | cmd.Env = append(os.Environ()) 248 | 249 | //Handle Std out and Std err 250 | var stdBuffer bytes.Buffer 251 | mw := io.MultiWriter(os.Stdout, &stdBuffer) 252 | cmd.Stdout = mw 253 | cmd.Stderr = mw 254 | //Run the command, check for errors 255 | if err := cmd.Run(); err != nil { 256 | return err, "" 257 | } 258 | 259 | //Create empty output var 260 | var commandOut string 261 | 262 | //If command was run with Batch system get output 263 | if !strings.Contains(fileText, "clone") { 264 | err, commandOut = GetOutputSlurm(fmt.Sprint(j.OutputScriptPth, "/slurm_out.log")) 265 | if err != nil { 266 | return err, "" 267 | } 268 | } 269 | //Return output 270 | return nil, commandOut 271 | } 272 | 273 | //Gets contents of a slurm output file and returns it when availible 274 | func GetOutputSlurm(outputFile string) (err error, output string) { 275 | retry := true 276 | for retry { 277 | if _, err := os.Stat(outputFile); os.IsNotExist(err) { 278 | time.Sleep(10 * time.Millisecond) 279 | continue 280 | } else { 281 | time.Sleep(50 * time.Millisecond) 282 | file, err := ioutil.ReadFile(outputFile) 283 | if err != nil { 284 | return err, "" 285 | } 286 | os.Remove(outputFile) 287 | return nil, string(file) 288 | } 289 | } 290 | return nil, "" 291 | } 292 | 293 | //Parses lsf output file, finds the script that was just run, returns output if any when availible 294 | func GetOutput(scriptName string, outputFile string) (err error, output string) { 295 | retry := true 296 | 297 | for retry { 298 | file, err := os.Open(outputFile) 299 | if err != nil { 300 | continue 301 | } 302 | scanner := bufio.NewScanner(file) 303 | for scanner.Scan() { 304 | line := scanner.Text() 305 | if strings.Contains(line, scriptName) { 306 | retry = false 307 | } 308 | } 309 | file.Close() 310 | file = nil 311 | } 312 | 313 | file, err := os.Open(outputFile) 314 | if err != nil { 315 | return err, "" 316 | } 317 | 318 | var lineArray []string 319 | var subLineArray []string 320 | var startingLine int = 0 321 | var endingLine int 322 | 323 | scanner := bufio.NewScanner(file) 324 | for scanner.Scan() { 325 | lineArray = append(lineArray, scanner.Text()) 326 | } 327 | if err := scanner.Err(); err != nil { 328 | return err, "" 329 | } 330 | 331 | for i, line := range lineArray { 332 | if strings.Contains(line, scriptName) { 333 | startingLine = i 334 | break 335 | } 336 | } 337 | 338 | for i, line := range lineArray[startingLine:] { 339 | if strings.Contains(line, "PS:") { 340 | endingLine = startingLine + i 341 | break 342 | } 343 | } 344 | 345 | for _, line := range lineArray[startingLine+34 : endingLine-1] { 346 | subLineArray = append(subLineArray, line) 347 | } 348 | 349 | file.Close() 350 | file = nil 351 | return nil, strings.Join(subLineArray, "\n") 352 | 353 | } 354 | --------------------------------------------------------------------------------