Add arguments to CreateJobHeader (#101)

* Add SlurmParams struct to Input

* Update PrepareJobFile and Post functions to accept SlurmParams

* Refactor CreateJobHeader function to accept parameters for job configuration

* Update job posting parameters

* Update max_concurrent value and add slurm configuration

* Update markdownlint and trufflehog versions

* Remove UseSlurm field and update related code

.trunk/trunk.yaml

@@ -27,12 +27,12 @@ lint:
     - [email protected]
     - [email protected]
     - [email protected]
-    - markdownlint@0.38.0
+    - markdownlint@0.39.0
     - [email protected]
     - [email protected]
     - [email protected]
     - [email protected]
-    - trufflehog@3.64.0
+    - trufflehog@3.66.1
     - [email protected]
 actions:
   enabled:

example/example_haddock30.yml

@@ -1,12 +1,15 @@
 general:
   executable: /workspaces/haddock-runner/example/haddock3.sh
-  max_concurrent: 1
+  max_concurrent: 4
   haddock_dir: /opt/haddock3
   receptor_suffix: _r_u
   ligand_suffix: _l_u
   input_list: /workspaces/haddock-runner/example/input_list.txt
   work_dir: /workspaces/haddock-runner/bm-goes-here
-  use_slurm: true
+
+slurm:
+  # partition: mypartition
+  cpus_per_task: 1
 
 scenarios:
   - name: true-interface

input/input.go

@@ -22,6 +22,7 @@ import (
 // Input is the input structure
 type Input struct {
 	General   GeneralStruct
+	Slurm     SlurmParams
 	Scenarios []Scenario
 }
 
@@ -34,7 +35,16 @@ type GeneralStruct struct {
 	InputList         string `yaml:"input_list"`
 	WorkDir           string `yaml:"work_dir"`
 	MaxConcurrent     int    `yaml:"max_concurrent"`
-	UseSlurm          bool   `yaml:"use_slurm"`
+}
+
+type SlurmParams struct {
+	Partition       string `yaml:"partition"`
+	Cpus_per_task   int    `yaml:"cpus_per_task"`
+	Ntasks_per_node int    `yaml:"ntasks_per_node"`
+	Nodes           int    `yaml:"nodes"`
+	Time            string `yaml:"time"`
+	Account         string `yaml:"account"`
+	Mail_user       string `yaml:"mail_user"`
 }
 
 // Scenario is the scenario structure
@@ -216,7 +226,7 @@ func ValidateRunCNSParams(known map[string]interface{}, params map[string]interf
 // ValidateExecutionModes checks if the execution modes are valid
 func (inp *Input) ValidateExecutionModes() error {
 
-	if inp.General.UseSlurm {
+	if inp.Slurm == (SlurmParams{}) {
 		// Check if the executable is HADDOCK3
 		if utils.IsHaddock24(inp.General.HaddockDir) {
 			err := errors.New("cannot use `use_slurm` with HADDOCK2")

input/input_test.go

@@ -579,7 +579,6 @@ func TestValidateExecutionModes(t *testing.T) {
 			name: "valid",
 			fields: fields{
 				General: GeneralStruct{
-					UseSlurm:   true,
 					HaddockDir: haddock3Dir,
 				},
 				Scenarios: []Scenario{
@@ -599,7 +598,6 @@ func TestValidateExecutionModes(t *testing.T) {
 			name: "invalid-haddock2",
 			fields: fields{
 				General: GeneralStruct{
-					UseSlurm:   true,
 					HaddockDir: haddock2Dir,
 				},
 				Scenarios: []Scenario{},
@@ -610,7 +608,6 @@ func TestValidateExecutionModes(t *testing.T) {
 			name: "invalid-haddock3",
 			fields: fields{
 				General: GeneralStruct{
-					UseSlurm:   true,
 					HaddockDir: haddock3Dir,
 				},
 				Scenarios: []Scenario{

main.go

@@ -208,7 +208,7 @@ func main() {
 				if err != nil {
 					glog.Exit("Failed to clean job: " + err.Error())
 				}
-				err = j.Post(haddockVersion, inp.General.HaddockExecutable, inp.General.UseSlurm)
+				err = j.Post(haddockVersion, inp.General.HaddockExecutable, inp.Slurm)
 				if err != nil {
 					glog.Exit("Failed to post job: " + err.Error())
 				}
@@ -225,7 +225,7 @@ func main() {
 				}
 
 			default:
-				err := j.Post(haddockVersion, inp.General.HaddockExecutable, inp.General.UseSlurm)
+				err := j.Post(haddockVersion, inp.General.HaddockExecutable, inp.Slurm)
 				if err != nil {
 					glog.Exit("Failed to post job: " + err.Error())
 				}
@@ -241,10 +241,6 @@ func main() {
 	wg.Wait()
 
 	glog.Info("############################################")
-	if inp.General.UseSlurm {
-		glog.Info("haddock-runner finished successfully (things might still be running on the cluster)")
-	} else {
-		glog.Info("haddock-runner finished successfully")
-	}
+	glog.Info("haddock-runner finished successfully")
 
 }

runner/jobs.go

@@ -219,13 +219,21 @@ func (j Job) Run(version int, cmd string) (string, error) {
 }
 
 // PrepareJobFile prepares the job file, returns the path to the job file
-func (j *Job) PrepareJobFile(executable string) error {
+func (j *Job) PrepareJobFile(executable string, slurm input.SlurmParams) error {
 
 	var header string
 	var body string
 
 	// Create the JobFile
-	header = utils.CreateJobHeader()
+	header = utils.CreateJobHeader(
+		slurm.Partition,
+		slurm.Account,
+		slurm.Mail_user,
+		slurm.Time,
+		slurm.Cpus_per_task,
+		slurm.Nodes,
+		slurm.Ntasks_per_node,
+	)
 
 	// Create the JobBody
 	body = utils.CreateJobBody(executable, j.Path)
@@ -369,10 +377,10 @@ func (j Job) Clean() error {
 
 }
 
-func (j Job) Post(haddockVersion int, executable string, slurm bool) error {
+func (j Job) Post(haddockVersion int, executable string, slurm input.SlurmParams) error {
 
-	if slurm {
-		err := j.PrepareJobFile(executable)
+	if slurm != (input.SlurmParams{}) {
+		err := j.PrepareJobFile(executable, slurm)
 		if err != nil {
 			glog.Error("Failed to prepare job file: " + err.Error())
 			return err

runner/jobs_test.go

@@ -617,6 +617,7 @@ func TestJobPrepareJobFile(t *testing.T) {
 	}
 	type args struct {
 		executable string
+		slurm      input.SlurmParams
 	}
 
 	tests := []struct {
@@ -629,6 +630,7 @@ func TestJobPrepareJobFile(t *testing.T) {
 			name: "pass by creating a job file",
 			args: args{
 				executable: "echo test",
+				slurm:      input.SlurmParams{},
 			},
 			fields: fields{
 				j: Job{
@@ -640,6 +642,7 @@ func TestJobPrepareJobFile(t *testing.T) {
 			name: "fail by passing a non-existing path",
 			args: args{
 				executable: "echo test",
+				slurm:      input.SlurmParams{},
 			},
 			fields: fields{
 				j: Job{
@@ -652,7 +655,7 @@ func TestJobPrepareJobFile(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			err := tt.fields.j.PrepareJobFile(tt.args.executable)
+			err := tt.fields.j.PrepareJobFile(tt.args.executable, tt.args.slurm)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("PrepareJobFile() error = %v, wantErr %v", err, tt.wantErr)
 				return
@@ -801,7 +804,7 @@ func TestJobPost(t *testing.T) {
 	type args struct {
 		haddockVersion int
 		executable     string
-		slurm          bool
+		slurm          input.SlurmParams
 	}
 
 	tests := []struct {
@@ -820,7 +823,9 @@ func TestJobPost(t *testing.T) {
 			args: args{
 				haddockVersion: 2,
 				executable:     "echo test",
-				slurm:          true,
+				slurm: input.SlurmParams{
+					Partition: "test",
+				},
 			},
 			wantErr: false,
 		},
@@ -834,7 +839,9 @@ func TestJobPost(t *testing.T) {
 			args: args{
 				haddockVersion: 3,
 				executable:     "echo test",
-				slurm:          false,
+				slurm: input.SlurmParams{
+					Partition: "test",
+				},
 			},
 			wantErr: true,
 		},
@@ -848,7 +855,7 @@ func TestJobPost(t *testing.T) {
 			args: args{
 				haddockVersion: 3,
 				executable:     "echo test",
-				slurm:          true,
+				slurm:          input.SlurmParams{},
 			},
 			wantErr: true,
 		},

utils/utils.go

@@ -249,16 +249,34 @@ func SearchInLog(filePath, searchString string) (bool, error) {
 	return false, nil
 }
 
-func CreateJobHeader() string {
+func CreateJobHeader(partition, account, mail_user, runtime string, cpus_per_task, nodes, ntasks_per_node int) string {
 
 	header := "#!/bin/bash\n"
 	header += "#SBATCH --job-name=haddock\n"
 	header += "#SBATCH --output=haddock-%j.out\n"
 	header += "#SBATCH --error=haddock-%j.err\n"
-	header += "#SBATCH --time=7-00:00:00\n"
-	header += "#SBATCH --nodes=1\n"
-	header += "#SBATCH --ntasks-per-node=1\n"
-	header += "#SBATCH --cpus-per-task=1\n"
+	if partition != "" {
+		header += "#SBATCH --partition=" + partition + "\n"
+	}
+	if account != "" {
+		header += "#SBATCH --account=" + account + "\n"
+	}
+	if mail_user != "" {
+		header += "#SBATCH --mail-user=" + mail_user + "\n"
+		header += "#SBATCH --mail-type=ALL\n"
+	}
+	if cpus_per_task != 0 {
+		header += "#SBATCH --cpus-per-task=" + strconv.Itoa(cpus_per_task) + "\n"
+	}
+	if nodes != 0 {
+		header += "#SBATCH --nodes=" + strconv.Itoa(nodes) + "\n"
+	}
+	if ntasks_per_node != 0 {
+		header += "#SBATCH --ntasks-per-node=" + strconv.Itoa(ntasks_per_node) + "\n"
+	}
+	if runtime != "" {
+		header += "#SBATCH --time=" + runtime + "\n"
+	}
 	header += "\n"
 
 	return header

utils/utils_test.go

@@ -406,7 +406,15 @@ func TestSearchInLog(t *testing.T) {
 func TestCreateJobHeader(t *testing.T) {
 	// TODO: Improve this test, now its just checking if the function runs without errors
 	//  it should check if the output is correct
-	output := CreateJobHeader()
+	output := CreateJobHeader(
+		"partition",
+		"account",
+		"mail_user",
+		"runtime",
+		1,
+		1,
+		1,
+	)
 
 	if output == "" {
 		t.Errorf("Failed to create job header")