myxce/src/frontend/schedulerclient/write.go

package schedulerclient

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
)

// ── Write types (Phase 1) ────────────────────────────────────────────────────

// JobSpec is a unified job submission specification. Canonical units are
// WalltimeSeconds (seconds) and MemoryBytes (bytes); the HTTP layer converts
// them to the apiserver's wire format ("HH:MM" and "<N> MB").
type JobSpec struct {
	Name             string   `json:"name,omitempty"`             // job_name; array syntax "myjob[1-10]" is passed through
	Command          string   `json:"command"`                    // required
	Queue            string   `json:"queue,omitempty"`            // empty → apiserver defaults to "normal"
	Project          string   `json:"project,omitempty"`          // project_name
	CWD              string   `json:"cwd,omitempty"`              // working directory
	OutFile          string   `json:"out_file,omitempty"`         // standard output file
	ErrFile          string   `json:"err_file,omitempty"`         // standard error file
	Env              []string `json:"env,omitempty"`              // "KEY=VALUE"
	NumProcessors    int      `json:"num_processors,omitempty"`   // number of processors
	MaxNumProcessors int      `json:"maxnum_processors,omitempty"`
	GpuReq           string   `json:"gpureq,omitempty"`           // gpureq resource string
	WalltimeSeconds  int      `json:"walltime_seconds,omitempty"` // run_limit → "HH:MM"
	MemoryBytes      int64    `json:"memory_bytes,omitempty"`     // mem_limit → "<MB> MB"
	Exclusive        bool     `json:"exclusive,omitempty"`
	Rerunable        bool     `json:"rerunable,omitempty"`
	DependCond       string   `json:"depend_cond,omitempty"` // dependency condition string
	JobDescription   string   `json:"description,omitempty"`
}

// SubmitJobResult is the apiserver's response to a successful job submission.
type SubmitJobResult struct {
	JobID      int64  `json:"jobid"` // NB: lowercase field name in the apiserver response
	JobIDHuman string `json:"job_id_human"` // cluster-apiserver-timestamp-seq (usable in xce_test/xcecli)
	QueueName  string `json:"queue_name"`
	Message    string `json:"message"`
}

// JobControlRequest is the body sent to POST /xce/v1/jobs/:jobids/control.
type JobControlRequest struct {
	Action    string `json:"action"` // kill | stop | resume | checkpoint
	Signal    string `json:"signal,omitempty"`
	Comment   string `json:"comment,omitempty"`
	ForceKill bool   `json:"force_kill,omitempty"`
	Method    string `json:"method,omitempty"`
}

// ControlResult is the apiserver's per-item control response.
type ControlResult struct {
	Status   string        `json:"status"` // accepted | failed | partially_accepted | no_items_processed
	Accepted []ControlItem `json:"accepted"`
	Failed   []ControlItem `json:"failed"`
}

// ControlItem reports the outcome for one job (or array element).
type ControlItem struct {
	UniqueJobID int64  `json:"unique_job_id"`
	JobID       int64  `json:"job_id"`
	ArrayIdx    int    `json:"array_idx"`
	Msg         string `json:"msg"`
}

// apiJobSubmitRequest mirrors the subset of apiserver models.JobSubmitRequest
// JSON tags we send. Zero-valued fields are omitted (omitempty) so the
// apiserver applies its own defaults (e.g. queue_name → "normal").
type apiJobSubmitRequest struct {
	JobName          string   `json:"job_name,omitempty"`
	Command          string   `json:"command"`
	QueueName        string   `json:"queue_name,omitempty"`
	ProjectName      string   `json:"project_name,omitempty"`
	CWD              string   `json:"cwd,omitempty"`
	OutFile          string   `json:"out_file,omitempty"`
	ErrFile          string   `json:"err_file,omitempty"`
	Env              []string `json:"env,omitempty"`
	NumProcessors    int      `json:"num_processors,omitempty"`
	MaxNumProcessors int      `json:"maxnum_processors,omitempty"`
	Gpureq           string   `json:"gpureq,omitempty"`
	Exclusive        bool     `json:"exclusive,omitempty"`
	Rerunable        bool     `json:"rerunable,omitempty"`
	DependCond       string   `json:"depend_cond,omitempty"`
	JobDescription   string   `json:"job_description,omitempty"`
	RunLimit         string   `json:"run_limit,omitempty"` // "HH:MM" (bare ints are MINUTES — never sent)
	MemLimit         string   `json:"mem_limit,omitempty"` // "<N> MB"
}

// jobSpecToAPI converts a JobSpec into the wire request. run_limit is rendered
// as "H:MM" and mem_limit as "<MB> MB" using the package's decimal
// mbToBytes convention (1 MB = 1_000_000 bytes), so "4 GB" → 4e9 bytes →
// "4000 MB" → CustomLimit parses cleanly.
func jobSpecToAPI(s JobSpec) apiJobSubmitRequest {
	req := apiJobSubmitRequest{
		JobName:          s.Name,
		Command:          s.Command,
		QueueName:        s.Queue,
		ProjectName:      s.Project,
		CWD:              s.CWD,
		OutFile:          s.OutFile,
		ErrFile:          s.ErrFile,
		Env:              s.Env,
		NumProcessors:    s.NumProcessors,
		MaxNumProcessors: s.MaxNumProcessors,
		Gpureq:           s.GpuReq,
		Exclusive:        s.Exclusive,
		Rerunable:        s.Rerunable,
		DependCond:       s.DependCond,
		JobDescription:   s.JobDescription,
	}
	if s.WalltimeSeconds > 0 {
		req.RunLimit = formatWalltime(s.WalltimeSeconds)
	}
	if s.MemoryBytes > 0 {
		req.MemLimit = formatMemory(s.MemoryBytes)
	}
	return req
}

// formatWalltime renders seconds as "H:MM", matching the apiserver's
// CustomInt.String() output. Sub-minute seconds are truncated (walltime
// granularity is minutes).
func formatWalltime(secs int) string {
	return fmt.Sprintf("%d:%02d", secs/3600, (secs%3600)/60)
}

// formatMemory renders bytes as "<N> MB" using the package's decimal
// mbToBytes convention. Rounds up so sub-MB amounts don't become "0 MB".
func formatMemory(bytes int64) string {
	mb := bytes / mbToBytes
	if bytes%mbToBytes != 0 {
		mb++
	}
	return strconv.Itoa(int(mb)) + " MB"
}

// ── shared POST helper ───────────────────────────────────────────────────────

// doPost performs an authenticated POST with a JSON body and decodes the
// envelope. expectedCode is the envelope code on success (201 create, 202
// control). NB: ReturnSuccess always uses HTTP 200; the logical status is the
// envelope's code field. ReturnError sets both HTTP status and envelope code
// to the error code.
func (hc *HTTPClient) doPost(ctx context.Context, url string, body any, expectedCode int) (*apiEnvelope, error) {
	req := hc.rc.R().SetContext(ctx).
		SetHeader("Authorization", hc.token).
		SetHeader("Content-Type", "application/json")

	var env apiEnvelope
	resp, err := req.SetBody(body).SetResult(&env).Post(url)
	if err != nil {
		return nil, fmt.Errorf("post %s: %w", url, err)
	}
	if resp.StatusCode() != http.StatusOK {
		// resty does not decode non-2xx bodies into SetResult, so read the raw
		// body (the apiserver's error envelope) for a useful message.
		return nil, fmt.Errorf("post %s: HTTP %d: %s", url, resp.StatusCode(), resp.String())
	}
	if env.Code != expectedCode {
		return nil, fmt.Errorf("post %s: api error %d: %s", url, env.Code, env.Msg)
	}
	return &env, nil
}

// ── write method implementations ─────────────────────────────────────────────

func (hc *HTTPClient) SubmitJob(ctx context.Context, spec JobSpec) (*SubmitJobResult, error) {
	if strings.TrimSpace(spec.Command) == "" {
		return nil, fmt.Errorf("submit job: command is required")
	}

	url := hc.baseURL + "/xce/v1/jobs/create"
	env, err := hc.doPost(ctx, url, jobSpecToAPI(spec), http.StatusCreated)
	if err != nil {
		return nil, fmt.Errorf("submit job: %w", err)
	}

	var res SubmitJobResult
	if err := json.Unmarshal(env.Data, &res); err != nil {
		return nil, fmt.Errorf("submit job: parse result: %w", err)
	}
	return &res, nil
}

// controlJob sends a job-control action and returns a descriptive error unless
// every requested item was accepted.
func (hc *HTTPClient) controlJob(ctx context.Context, id, action string) error {
	if strings.TrimSpace(id) == "" {
		return fmt.Errorf("%s job: job id is required", action)
	}

	url := hc.baseURL + "/xce/v1/jobs/" + id + "/control"
	env, err := hc.doPost(ctx, url, JobControlRequest{Action: action}, http.StatusAccepted)
	if err != nil {
		return fmt.Errorf("%s job %s: %w", action, id, err)
	}

	var res ControlResult
	if err := json.Unmarshal(env.Data, &res); err != nil {
		return fmt.Errorf("%s job %s: parse result: %w", action, id, err)
	}
	switch res.Status {
	case "accepted":
		return nil
	case "no_items_processed":
		return fmt.Errorf("%s job %s: no matching job found", action, id)
	default: // "failed" or "partially_accepted"
		return fmt.Errorf("%s job %s: status=%s failed=%+v", action, id, res.Status, res.Failed)
	}
}

func (hc *HTTPClient) CancelJob(ctx context.Context, id string) error {
	return hc.controlJob(ctx, id, "kill")
}

func (hc *HTTPClient) HoldJob(ctx context.Context, id string) error {
	return hc.controlJob(ctx, id, "stop")
}

func (hc *HTTPClient) ReleaseJob(ctx context.Context, id string) error {
	return hc.controlJob(ctx, id, "resume")
}