feat(slurmdb): add JobsService
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
dailz
218
internal/slurm/slurmdb_jobs.go
Normal file
218
internal/slurm/slurmdb_jobs.go
Normal file
@@ -0,0 +1,218 @@
|
||||
package slurm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/url"
|
||||
)
|
||||
|
||||
// GetSlurmdbJobsOptions specifies optional query parameters for SlurmdbJobsService.GetJobs.
|
||||
// All 41 parameters correspond to the GET /slurmdb/v0.0.40/jobs endpoint.
|
||||
type GetSlurmdbJobsOptions struct {
|
||||
Account *string `url:"account,omitempty"`
|
||||
Association *string `url:"association,omitempty"`
|
||||
Cluster *string `url:"cluster,omitempty"`
|
||||
Constraints *string `url:"constraints,omitempty"`
|
||||
CpusMax *string `url:"cpus_max,omitempty"`
|
||||
CpusMin *string `url:"cpus_min,omitempty"`
|
||||
SchedulerUnset *string `url:"scheduler_unset,omitempty"`
|
||||
ScheduledOnSubmit *string `url:"scheduled_on_submit,omitempty"`
|
||||
ScheduledByMain *string `url:"scheduled_by_main,omitempty"`
|
||||
ScheduledByBackfill *string `url:"scheduled_by_backfill,omitempty"`
|
||||
JobStarted *string `url:"job_started,omitempty"`
|
||||
ExitCode *string `url:"exit_code,omitempty"`
|
||||
ShowDuplicates *string `url:"show_duplicates,omitempty"`
|
||||
SkipSteps *string `url:"skip_steps,omitempty"`
|
||||
DisableTruncateUsageTime *string `url:"disable_truncate_usage_time,omitempty"`
|
||||
WholeHetjob *string `url:"whole_hetjob,omitempty"`
|
||||
DisableWholeHetjob *string `url:"disable_whole_hetjob,omitempty"`
|
||||
DisableWaitForResult *string `url:"disable_wait_for_result,omitempty"`
|
||||
UsageTimeAsSubmitTime *string `url:"usage_time_as_submit_time,omitempty"`
|
||||
ShowBatchScript *string `url:"show_batch_script,omitempty"`
|
||||
ShowJobEnvironment *string `url:"show_job_environment,omitempty"`
|
||||
Format *string `url:"format,omitempty"`
|
||||
Groups *string `url:"groups,omitempty"`
|
||||
JobName *string `url:"job_name,omitempty"`
|
||||
NodesMax *string `url:"nodes_max,omitempty"`
|
||||
NodesMin *string `url:"nodes_min,omitempty"`
|
||||
Partition *string `url:"partition,omitempty"`
|
||||
Qos *string `url:"qos,omitempty"`
|
||||
Reason *string `url:"reason,omitempty"`
|
||||
Reservation *string `url:"reservation,omitempty"`
|
||||
ReservationID *string `url:"reservation_id,omitempty"`
|
||||
State *string `url:"state,omitempty"`
|
||||
Step *string `url:"step,omitempty"`
|
||||
TimelimitMax *string `url:"timelimit_max,omitempty"`
|
||||
TimelimitMin *string `url:"timelimit_min,omitempty"`
|
||||
EndTime *string `url:"end_time,omitempty"`
|
||||
StartTime *string `url:"start_time,omitempty"`
|
||||
SubmitTime *string `url:"submit_time,omitempty"`
|
||||
Node *string `url:"node,omitempty"`
|
||||
Users *string `url:"users,omitempty"`
|
||||
Wckey *string `url:"wckey,omitempty"`
|
||||
}
|
||||
|
||||
// GetJobs queries the SlurmDBD for jobs matching the given options.
|
||||
func (s *SlurmdbJobsService) GetJobs(ctx context.Context, opts *GetSlurmdbJobsOptions) (*OpenapiSlurmdbdJobsResp, *Response, error) {
|
||||
path := "slurmdb/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.Account != nil {
|
||||
q.Set("account", *opts.Account)
|
||||
}
|
||||
if opts.Association != nil {
|
||||
q.Set("association", *opts.Association)
|
||||
}
|
||||
if opts.Cluster != nil {
|
||||
q.Set("cluster", *opts.Cluster)
|
||||
}
|
||||
if opts.Constraints != nil {
|
||||
q.Set("constraints", *opts.Constraints)
|
||||
}
|
||||
if opts.CpusMax != nil {
|
||||
q.Set("cpus_max", *opts.CpusMax)
|
||||
}
|
||||
if opts.CpusMin != nil {
|
||||
q.Set("cpus_min", *opts.CpusMin)
|
||||
}
|
||||
if opts.SchedulerUnset != nil {
|
||||
q.Set("scheduler_unset", *opts.SchedulerUnset)
|
||||
}
|
||||
if opts.ScheduledOnSubmit != nil {
|
||||
q.Set("scheduled_on_submit", *opts.ScheduledOnSubmit)
|
||||
}
|
||||
if opts.ScheduledByMain != nil {
|
||||
q.Set("scheduled_by_main", *opts.ScheduledByMain)
|
||||
}
|
||||
if opts.ScheduledByBackfill != nil {
|
||||
q.Set("scheduled_by_backfill", *opts.ScheduledByBackfill)
|
||||
}
|
||||
if opts.JobStarted != nil {
|
||||
q.Set("job_started", *opts.JobStarted)
|
||||
}
|
||||
if opts.ExitCode != nil {
|
||||
q.Set("exit_code", *opts.ExitCode)
|
||||
}
|
||||
if opts.ShowDuplicates != nil {
|
||||
q.Set("show_duplicates", *opts.ShowDuplicates)
|
||||
}
|
||||
if opts.SkipSteps != nil {
|
||||
q.Set("skip_steps", *opts.SkipSteps)
|
||||
}
|
||||
if opts.DisableTruncateUsageTime != nil {
|
||||
q.Set("disable_truncate_usage_time", *opts.DisableTruncateUsageTime)
|
||||
}
|
||||
if opts.WholeHetjob != nil {
|
||||
q.Set("whole_hetjob", *opts.WholeHetjob)
|
||||
}
|
||||
if opts.DisableWholeHetjob != nil {
|
||||
q.Set("disable_whole_hetjob", *opts.DisableWholeHetjob)
|
||||
}
|
||||
if opts.DisableWaitForResult != nil {
|
||||
q.Set("disable_wait_for_result", *opts.DisableWaitForResult)
|
||||
}
|
||||
if opts.UsageTimeAsSubmitTime != nil {
|
||||
q.Set("usage_time_as_submit_time", *opts.UsageTimeAsSubmitTime)
|
||||
}
|
||||
if opts.ShowBatchScript != nil {
|
||||
q.Set("show_batch_script", *opts.ShowBatchScript)
|
||||
}
|
||||
if opts.ShowJobEnvironment != nil {
|
||||
q.Set("show_job_environment", *opts.ShowJobEnvironment)
|
||||
}
|
||||
if opts.Format != nil {
|
||||
q.Set("format", *opts.Format)
|
||||
}
|
||||
if opts.Groups != nil {
|
||||
q.Set("groups", *opts.Groups)
|
||||
}
|
||||
if opts.JobName != nil {
|
||||
q.Set("job_name", *opts.JobName)
|
||||
}
|
||||
if opts.NodesMax != nil {
|
||||
q.Set("nodes_max", *opts.NodesMax)
|
||||
}
|
||||
if opts.NodesMin != nil {
|
||||
q.Set("nodes_min", *opts.NodesMin)
|
||||
}
|
||||
if opts.Partition != nil {
|
||||
q.Set("partition", *opts.Partition)
|
||||
}
|
||||
if opts.Qos != nil {
|
||||
q.Set("qos", *opts.Qos)
|
||||
}
|
||||
if opts.Reason != nil {
|
||||
q.Set("reason", *opts.Reason)
|
||||
}
|
||||
if opts.Reservation != nil {
|
||||
q.Set("reservation", *opts.Reservation)
|
||||
}
|
||||
if opts.ReservationID != nil {
|
||||
q.Set("reservation_id", *opts.ReservationID)
|
||||
}
|
||||
if opts.State != nil {
|
||||
q.Set("state", *opts.State)
|
||||
}
|
||||
if opts.Step != nil {
|
||||
q.Set("step", *opts.Step)
|
||||
}
|
||||
if opts.TimelimitMax != nil {
|
||||
q.Set("timelimit_max", *opts.TimelimitMax)
|
||||
}
|
||||
if opts.TimelimitMin != nil {
|
||||
q.Set("timelimit_min", *opts.TimelimitMin)
|
||||
}
|
||||
if opts.EndTime != nil {
|
||||
q.Set("end_time", *opts.EndTime)
|
||||
}
|
||||
if opts.StartTime != nil {
|
||||
q.Set("start_time", *opts.StartTime)
|
||||
}
|
||||
if opts.SubmitTime != nil {
|
||||
q.Set("submit_time", *opts.SubmitTime)
|
||||
}
|
||||
if opts.Node != nil {
|
||||
q.Set("node", *opts.Node)
|
||||
}
|
||||
if opts.Users != nil {
|
||||
q.Set("users", *opts.Users)
|
||||
}
|
||||
if opts.Wckey != nil {
|
||||
q.Set("wckey", *opts.Wckey)
|
||||
}
|
||||
u.RawQuery = q.Encode()
|
||||
req.URL = u
|
||||
}
|
||||
|
||||
var result OpenapiSlurmdbdJobsResp
|
||||
resp, err := s.client.Do(ctx, req, &result)
|
||||
if err != nil {
|
||||
return nil, resp, err
|
||||
}
|
||||
return &result, resp, nil
|
||||
}
|
||||
|
||||
// GetJob retrieves info for a specific job from SlurmDBD.
|
||||
func (s *SlurmdbJobsService) GetJob(ctx context.Context, jobID string) (*OpenapiSlurmdbdJobsResp, *Response, error) {
|
||||
path := fmt.Sprintf("slurmdb/v0.0.40/job/%s", jobID)
|
||||
req, err := s.client.NewRequest("GET", path, nil)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
var result OpenapiSlurmdbdJobsResp
|
||||
resp, err := s.client.Do(ctx, req, &result)
|
||||
if err != nil {
|
||||
return nil, resp, err
|
||||
}
|
||||
return &result, resp, nil
|
||||
}
|
||||
246
internal/slurm/slurmdb_jobs_test.go
Normal file
246
internal/slurm/slurmdb_jobs_test.go
Normal file
@@ -0,0 +1,246 @@
|
||||
package slurm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSlurmdbJobsService_GetJobs(t *testing.T) {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/slurmdb/v0.0.40/jobs", func(w http.ResponseWriter, r *http.Request) {
|
||||
testMethod(t, r, "GET")
|
||||
fmt.Fprint(w, `{"jobs": []}`)
|
||||
})
|
||||
server := httptest.NewServer(mux)
|
||||
defer server.Close()
|
||||
|
||||
client, _ := NewClient(server.URL, nil)
|
||||
resp, _, err := client.SlurmdbJobs.GetJobs(context.Background(), nil)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if resp == nil {
|
||||
t.Fatal("expected non-nil response")
|
||||
}
|
||||
if len(resp.Jobs) != 0 {
|
||||
t.Errorf("expected empty jobs, got %d", len(resp.Jobs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlurmdbJobsService_GetJobsWithOptions(t *testing.T) {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/slurmdb/v0.0.40/jobs", func(w http.ResponseWriter, r *http.Request) {
|
||||
testMethod(t, r, "GET")
|
||||
q := r.URL.Query()
|
||||
if q.Get("users") != "alice,bob" {
|
||||
t.Errorf("expected users=alice,bob, got %s", q.Get("users"))
|
||||
}
|
||||
if q.Get("account") != "science" {
|
||||
t.Errorf("expected account=science, got %s", q.Get("account"))
|
||||
}
|
||||
if q.Get("cluster") != "cluster1" {
|
||||
t.Errorf("expected cluster=cluster1, got %s", q.Get("cluster"))
|
||||
}
|
||||
if q.Get("state") != "RUNNING" {
|
||||
t.Errorf("expected state=RUNNING, got %s", q.Get("state"))
|
||||
}
|
||||
if q.Get("start_time") != "1700000000" {
|
||||
t.Errorf("expected start_time=1700000000, got %s", q.Get("start_time"))
|
||||
}
|
||||
if q.Get("partition") != "gpu" {
|
||||
t.Errorf("expected partition=gpu, got %s", q.Get("partition"))
|
||||
}
|
||||
fmt.Fprint(w, `{"jobs": []}`)
|
||||
})
|
||||
server := httptest.NewServer(mux)
|
||||
defer server.Close()
|
||||
|
||||
client, _ := NewClient(server.URL, nil)
|
||||
opts := &GetSlurmdbJobsOptions{
|
||||
Users: Ptr("alice,bob"),
|
||||
Account: Ptr("science"),
|
||||
Cluster: Ptr("cluster1"),
|
||||
State: Ptr("RUNNING"),
|
||||
StartTime: Ptr("1700000000"),
|
||||
Partition: Ptr("gpu"),
|
||||
}
|
||||
resp, _, err := client.SlurmdbJobs.GetJobs(context.Background(), opts)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if resp == nil {
|
||||
t.Fatal("expected non-nil response")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlurmdbJobsService_GetJob(t *testing.T) {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/slurmdb/v0.0.40/job/12345", func(w http.ResponseWriter, r *http.Request) {
|
||||
testMethod(t, r, "GET")
|
||||
fmt.Fprint(w, `{"jobs": [{"job_id": 12345, "name": "my-job"}]}`)
|
||||
})
|
||||
server := httptest.NewServer(mux)
|
||||
defer server.Close()
|
||||
|
||||
client, _ := NewClient(server.URL, nil)
|
||||
resp, _, err := client.SlurmdbJobs.GetJob(context.Background(), "12345")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if resp == nil {
|
||||
t.Fatal("expected non-nil response")
|
||||
}
|
||||
if len(resp.Jobs) != 1 {
|
||||
t.Fatalf("expected 1 job, got %d", len(resp.Jobs))
|
||||
}
|
||||
if resp.Jobs[0].JobID == nil || *resp.Jobs[0].JobID != 12345 {
|
||||
t.Errorf("expected job_id=12345, got %v", resp.Jobs[0].JobID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlurmdbJobsService_GetJobs_ServerError(t *testing.T) {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/slurmdb/v0.0.40/jobs", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
fmt.Fprint(w, `{"errors": [{"error": "internal error"}]}`)
|
||||
})
|
||||
server := httptest.NewServer(mux)
|
||||
defer server.Close()
|
||||
|
||||
client, _ := NewClient(server.URL, nil)
|
||||
_, _, err := client.SlurmdbJobs.GetJobs(context.Background(), nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlurmdbJobsService_GetJob_ServerError(t *testing.T) {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/slurmdb/v0.0.40/job/99999", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
fmt.Fprint(w, `{"errors": [{"error": "job not found"}]}`)
|
||||
})
|
||||
server := httptest.NewServer(mux)
|
||||
defer server.Close()
|
||||
|
||||
client, _ := NewClient(server.URL, nil)
|
||||
_, _, err := client.SlurmdbJobs.GetJob(context.Background(), "99999")
|
||||
if err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlurmdbJobsService_GetJobs_AllQueryParams(t *testing.T) {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/slurmdb/v0.0.40/jobs", func(w http.ResponseWriter, r *http.Request) {
|
||||
testMethod(t, r, "GET")
|
||||
q := r.URL.Query()
|
||||
expectedParams := map[string]string{
|
||||
"account": "acct1",
|
||||
"association": "assoc1",
|
||||
"cluster": "cl1",
|
||||
"constraints": "c1",
|
||||
"cpus_max": "64",
|
||||
"cpus_min": "1",
|
||||
"scheduler_unset": "1",
|
||||
"scheduled_on_submit": "1",
|
||||
"scheduled_by_main": "1",
|
||||
"scheduled_by_backfill": "1",
|
||||
"job_started": "1",
|
||||
"exit_code": "0",
|
||||
"show_duplicates": "1",
|
||||
"skip_steps": "1",
|
||||
"disable_truncate_usage_time": "1",
|
||||
"whole_hetjob": "1",
|
||||
"disable_whole_hetjob": "0",
|
||||
"disable_wait_for_result": "1",
|
||||
"usage_time_as_submit_time": "1",
|
||||
"show_batch_script": "1",
|
||||
"show_job_environment": "1",
|
||||
"format": "json",
|
||||
"groups": "grp1",
|
||||
"job_name": "jn1",
|
||||
"nodes_max": "10",
|
||||
"nodes_min": "1",
|
||||
"partition": "part1",
|
||||
"qos": "normal",
|
||||
"reason": "r1",
|
||||
"reservation": "res1",
|
||||
"reservation_id": "100",
|
||||
"state": "RUNNING",
|
||||
"step": "0",
|
||||
"timelimit_max": "3600",
|
||||
"timelimit_min": "60",
|
||||
"end_time": "1700001000",
|
||||
"start_time": "1700000000",
|
||||
"submit_time": "1699999000",
|
||||
"node": "n1",
|
||||
"users": "u1",
|
||||
"wckey": "w1",
|
||||
}
|
||||
for k, v := range expectedParams {
|
||||
got := q.Get(k)
|
||||
if got != v {
|
||||
t.Errorf("expected %s=%s, got %s", k, v, got)
|
||||
}
|
||||
}
|
||||
fmt.Fprint(w, `{"jobs": []}`)
|
||||
})
|
||||
server := httptest.NewServer(mux)
|
||||
defer server.Close()
|
||||
|
||||
client, _ := NewClient(server.URL, nil)
|
||||
opts := &GetSlurmdbJobsOptions{
|
||||
Account: Ptr("acct1"),
|
||||
Association: Ptr("assoc1"),
|
||||
Cluster: Ptr("cl1"),
|
||||
Constraints: Ptr("c1"),
|
||||
CpusMax: Ptr("64"),
|
||||
CpusMin: Ptr("1"),
|
||||
SchedulerUnset: Ptr("1"),
|
||||
ScheduledOnSubmit: Ptr("1"),
|
||||
ScheduledByMain: Ptr("1"),
|
||||
ScheduledByBackfill: Ptr("1"),
|
||||
JobStarted: Ptr("1"),
|
||||
ExitCode: Ptr("0"),
|
||||
ShowDuplicates: Ptr("1"),
|
||||
SkipSteps: Ptr("1"),
|
||||
DisableTruncateUsageTime: Ptr("1"),
|
||||
WholeHetjob: Ptr("1"),
|
||||
DisableWholeHetjob: Ptr("0"),
|
||||
DisableWaitForResult: Ptr("1"),
|
||||
UsageTimeAsSubmitTime: Ptr("1"),
|
||||
ShowBatchScript: Ptr("1"),
|
||||
ShowJobEnvironment: Ptr("1"),
|
||||
Format: Ptr("json"),
|
||||
Groups: Ptr("grp1"),
|
||||
JobName: Ptr("jn1"),
|
||||
NodesMax: Ptr("10"),
|
||||
NodesMin: Ptr("1"),
|
||||
Partition: Ptr("part1"),
|
||||
Qos: Ptr("normal"),
|
||||
Reason: Ptr("r1"),
|
||||
Reservation: Ptr("res1"),
|
||||
ReservationID: Ptr("100"),
|
||||
State: Ptr("RUNNING"),
|
||||
Step: Ptr("0"),
|
||||
TimelimitMax: Ptr("3600"),
|
||||
TimelimitMin: Ptr("60"),
|
||||
EndTime: Ptr("1700001000"),
|
||||
StartTime: Ptr("1700000000"),
|
||||
SubmitTime: Ptr("1699999000"),
|
||||
Node: Ptr("n1"),
|
||||
Users: Ptr("u1"),
|
||||
Wckey: Ptr("w1"),
|
||||
}
|
||||
resp, _, err := client.SlurmdbJobs.GetJobs(context.Background(), opts)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if resp == nil {
|
||||
t.Fatal("expected non-nil response")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user