hpc/internal/slurm/slurm_jobs.go

package slurm

import (
	"context"
	"fmt"
	"net/url"
)

// Job query flags for GetJobsOptions.Flags and GetNodeOptions.Flags.
const (
	JobFlagAll        = "ALL"
	JobFlagDetail     = "DETAIL"
	JobFlagMixed      = "MIXED"
	JobFlagLocal      = "LOCAL"
	JobFlagSibling    = "SIBLING"
	JobFlagFederation = "FEDERATION"
	JobFlagFuture     = "FUTURE"
)

// GetJobsOptions specifies optional parameters for GetJobs.
type GetJobsOptions struct {
	UpdateTime *string `url:"update_time,omitempty"`
	Flags      *string `url:"flags,omitempty"` // Use JobFlag* constants (e.g. JobFlagDetail)
}

// Job signal flags for DeleteJobOptions.Flags.
const (
	JobSignalFlagBatchJob          = "BATCH_JOB"
	JobSignalFlagArrayTask         = "ARRAY_TASK"
	JobSignalFlagFullStepsOnly     = "FULL_STEPS_ONLY"
	JobSignalFlagFullJob           = "FULL_JOB"
	JobSignalFlagFederationRequeue = "FEDERATION_REQUEUE"
	JobSignalFlagHurry             = "HURRY"
	JobSignalFlagOutOfMemory       = "OUT_OF_MEMORY"
	JobSignalFlagNoSiblingJobs     = "NO_SIBLING_JOBS"
	JobSignalFlagReservationJob    = "RESERVATION_JOB"
	JobSignalFlagWarningSent       = "WARNING_SENT"
)

// DeleteJobOptions specifies optional parameters for DeleteJob.
type DeleteJobOptions struct {
	Signal *string `url:"signal,omitempty"`
	Flags  *string `url:"flags,omitempty"` // Use JobSignalFlag* constants (e.g. JobSignalFlagHurry)
}

// GetJobs lists all jobs.
func (s *JobsService) GetJobs(ctx context.Context, opts *GetJobsOptions) (*OpenapiJobInfoResp, *Response, error) {
	path := "slurm/v0.0.40/jobs"
	req, err := s.client.NewRequest("GET", path, nil)
	if err != nil {
		return nil, nil, err
	}

	if opts != nil {
		u, parseErr := url.Parse(req.URL.String())
		if parseErr != nil {
			return nil, nil, parseErr
		}
		q := u.Query()
		if opts.UpdateTime != nil {
			q.Set("update_time", *opts.UpdateTime)
		}
		if opts.Flags != nil {
			q.Set("flags", *opts.Flags)
		}
		u.RawQuery = q.Encode()
		req.URL = u
	}

	var result OpenapiJobInfoResp
	resp, err := s.client.Do(ctx, req, &result)
	if err != nil {
		return nil, resp, err
	}
	return &result, resp, nil
}

// GetJob gets a single job by ID.
func (s *JobsService) GetJob(ctx context.Context, jobID string, opts *GetJobsOptions) (*OpenapiJobInfoResp, *Response, error) {
	path := fmt.Sprintf("slurm/v0.0.40/job/%s", jobID)
	req, err := s.client.NewRequest("GET", path, nil)
	if err != nil {
		return nil, nil, err
	}

	if opts != nil {
		u, parseErr := url.Parse(req.URL.String())
		if parseErr != nil {
			return nil, nil, parseErr
		}
		q := u.Query()
		if opts.UpdateTime != nil {
			q.Set("update_time", *opts.UpdateTime)
		}
		if opts.Flags != nil {
			q.Set("flags", *opts.Flags)
		}
		u.RawQuery = q.Encode()
		req.URL = u
	}

	var result OpenapiJobInfoResp
	resp, err := s.client.Do(ctx, req, &result)
	if err != nil {
		return nil, resp, err
	}
	return &result, resp, nil
}

// PostJob updates a job.
func (s *JobsService) PostJob(ctx context.Context, jobID string, desc *JobDescMsg) (*OpenapiJobPostResponse, *Response, error) {
	path := fmt.Sprintf("slurm/v0.0.40/job/%s", jobID)
	req, err := s.client.NewRequest("POST", path, desc)
	if err != nil {
		return nil, nil, err
	}

	var result OpenapiJobPostResponse
	resp, err := s.client.Do(ctx, req, &result)
	if err != nil {
		return nil, resp, err
	}
	return &result, resp, nil
}

// DeleteJob cancels (signals) a job.
func (s *JobsService) DeleteJob(ctx context.Context, jobID string, opts *DeleteJobOptions) (*OpenapiResp, *Response, error) {
	path := fmt.Sprintf("slurm/v0.0.40/job/%s", jobID)
	req, err := s.client.NewRequest("DELETE", path, nil)
	if err != nil {
		return nil, nil, err
	}

	if opts != nil {
		u, parseErr := url.Parse(req.URL.String())
		if parseErr != nil {
			return nil, nil, parseErr
		}
		q := u.Query()
		if opts.Signal != nil {
			q.Set("signal", *opts.Signal)
		}
		if opts.Flags != nil {
			q.Set("flags", *opts.Flags)
		}
		u.RawQuery = q.Encode()
		req.URL = u
	}

	var result OpenapiResp
	resp, err := s.client.Do(ctx, req, &result)
	if err != nil {
		return nil, resp, err
	}
	return &result, resp, nil
}

// SubmitJob submits a new job.
func (s *JobsService) SubmitJob(ctx context.Context, reqBody *JobSubmitReq) (*OpenapiJobSubmitResponse, *Response, error) {
	path := "slurm/v0.0.40/job/submit"
	req, err := s.client.NewRequest("POST", path, reqBody)
	if err != nil {
		return nil, nil, err
	}

	var result OpenapiJobSubmitResponse
	resp, err := s.client.Do(ctx, req, &result)
	if err != nil {
		return nil, resp, err
	}
	return &result, resp, nil
}