github.com/informationsea/shellflow@v0.1.3/execute_sge.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"os/exec"
     7  	"path"
     8  	"regexp"
     9  	"strconv"
    10  	"strings"
    11  )
    12  
    13  const sgeTaskIDFileName = "sge-taskid.txt"
    14  
    15  func FollowUpSge(jobLogRoot string) (bool, error) {
    16  	rc, err := os.Open(path.Join(jobLogRoot, "rc"))
    17  	if err == nil {
    18  		defer rc.Close()
    19  		return false, nil
    20  	} else if os.IsNotExist(err) {
    21  		sgeTaskIDFile, err := os.Open(path.Join(jobLogRoot, sgeTaskIDFileName))
    22  		if err == nil {
    23  			var sgeTaskID int
    24  			i, err := fmt.Fscanf(sgeTaskIDFile, "%d", &sgeTaskID)
    25  			if err != nil {
    26  				return false, err
    27  			}
    28  			if i != 1 {
    29  				return false, fmt.Errorf("Cannot read pid: %d", i)
    30  			}
    31  
    32  			cmd := exec.Command("qstat", "-j", fmt.Sprintf("%d", sgeTaskID))
    33  			err = cmd.Run()
    34  
    35  			if err != nil {
    36  				rc, err = os.OpenFile(path.Join(jobLogRoot, "rc"), os.O_CREATE|os.O_WRONLY, 0644)
    37  				if err != nil {
    38  					return false, err
    39  				}
    40  				defer rc.Close()
    41  				_, err = fmt.Fprintf(rc, "1000")
    42  				if err != nil {
    43  					return false, err
    44  				}
    45  			}
    46  			return true, nil
    47  		} else if os.IsNotExist(err) {
    48  			return false, nil
    49  		} else {
    50  			return false, err
    51  		}
    52  	} else {
    53  		return false, err
    54  	}
    55  }
    56  
    57  var jobNameReplace = regexp.MustCompile("[^\\w_\\-.]")
    58  
    59  func ExecuteInSge(ge *TaskScripts) error {
    60  	sgeTaskID := make(map[int]string)
    61  
    62  	jobNameBase := jobNameReplace.ReplaceAllString(ge.jobName, "_")
    63  
    64  	for _, v := range ge.builder.Tasks {
    65  		if v.ShouldSkip {
    66  			fmt.Printf("skipping: %s\n", v.ShellScript)
    67  			continue
    68  		}
    69  
    70  		if v.CommandConfiguration.RunImmediate {
    71  			err := ExecuteLocalSingleOneTask(ge, v)
    72  			if err != nil {
    73  				return err
    74  			}
    75  			continue
    76  		}
    77  
    78  		scriptInfo := ge.scripts[v.ID]
    79  		qsub := []string{"-wd", scriptInfo.JobRoot, "-terse", "-o", path.Join(scriptInfo.JobRoot, "run.stdout"), "-e", path.Join(scriptInfo.JobRoot, "run.stderr")}
    80  
    81  		var holdID strings.Builder
    82  		first := true
    83  		for _, d := range v.DependentTaskID {
    84  			if u, ok := sgeTaskID[d]; ok && strings.TrimSpace(u) != "" {
    85  				if first {
    86  					first = false
    87  				} else {
    88  					holdID.WriteString(",")
    89  				}
    90  				holdID.WriteString(strings.TrimSpace(u))
    91  			}
    92  		}
    93  
    94  		if holdID.Len() > 0 {
    95  			qsub = append(qsub, "-hold_jid", holdID.String())
    96  		}
    97  
    98  		qsub = append(qsub, "-N", jobNameBase+"__ID-"+strconv.Itoa(v.ID))
    99  
   100  		if len(v.CommandConfiguration.SGEOption) > 0 {
   101  			qsub = append(qsub, v.CommandConfiguration.SGEOption...)
   102  		}
   103  
   104  		qsub = append(qsub, scriptInfo.RunScriptPath)
   105  
   106  		// write SGE options
   107  		submitArgs, err := os.OpenFile(path.Join(scriptInfo.JobRoot, "sge-submit-args.txt"), os.O_CREATE|os.O_WRONLY, 0644)
   108  		if err != nil {
   109  			return fmt.Errorf("cannot write SGE option log: %s", err.Error())
   110  		}
   111  		defer submitArgs.Close()
   112  		for _, v := range qsub {
   113  			submitArgs.WriteString(v)
   114  			submitArgs.WriteString("\n")
   115  		}
   116  
   117  		// run SGE
   118  		cmd := exec.Command("qsub", qsub...)
   119  
   120  		out, err := cmd.Output()
   121  		if err != nil {
   122  			return fmt.Errorf("cannot run qsub successfully: %s", err.Error())
   123  		}
   124  		currentTaskID := strings.TrimSpace(string(out))
   125  		sgeTaskID[v.ID] = currentTaskID
   126  
   127  		taskid, err := os.OpenFile(path.Join(scriptInfo.JobRoot, "sge-taskid.txt"), os.O_CREATE|os.O_WRONLY, 0644)
   128  		if err != nil {
   129  			return fmt.Errorf("Cannot open SGE task ID log file: %s", err.Error())
   130  		}
   131  		defer taskid.Close()
   132  		fmt.Fprintf(taskid, "%s\n", currentTaskID)
   133  		fmt.Printf("Submit ID:%s  : %s\n", currentTaskID, v.ShellScript)
   134  	}
   135  
   136  	return nil
   137  }