132 lines
2.7 KiB
Go
132 lines
2.7 KiB
Go
package runpod
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
var Endpoint string
|
|
var Authorization string
|
|
|
|
type Request struct {
|
|
Id string
|
|
Status string
|
|
}
|
|
|
|
type StatusResp struct {
|
|
DelayTime int `json:"delayTime"`
|
|
ExecutionTime int `json:"executionTime"`
|
|
Id string `json:"id"`
|
|
Status string `json:"status"`
|
|
Output []map[string]string `json:"output"`
|
|
}
|
|
|
|
func postReq(data []byte, path string) (resp *http.Response, err error) {
|
|
req, err := http.NewRequest("POST", Endpoint+path, bytes.NewBuffer(data))
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
req.Header.Add("Content-Type", "application/json")
|
|
req.Header.Add("Authorization", Authorization)
|
|
|
|
client := &http.Client{}
|
|
resp, err = client.Do(req)
|
|
|
|
statusOK := resp.StatusCode >= 200 && resp.StatusCode < 300
|
|
if !statusOK {
|
|
return resp, fmt.Errorf("runpod '%s' resp HTTP code '%d'", path, resp.StatusCode)
|
|
}
|
|
return
|
|
}
|
|
|
|
func getReq(path string) (resp *http.Response, err error) {
|
|
req, err := http.NewRequest("GET", Endpoint+path, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
req.Header.Add("Content-Type", "application/json")
|
|
req.Header.Add("Authorization", Authorization)
|
|
|
|
client := &http.Client{}
|
|
resp, err = client.Do(req)
|
|
|
|
statusOK := resp.StatusCode >= 200 && resp.StatusCode < 300
|
|
if !statusOK {
|
|
return resp, fmt.Errorf("runpod '%s' resp HTTP code '%d'", path, resp.StatusCode)
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// Returns the status of the serverless endpoint, can be used to check if it's operational
|
|
// returns a map because I'm too lazy to make the proper struct
|
|
func Health() (res map[string]interface{}, err error) {
|
|
resp, err := getReq("health")
|
|
if err != nil {
|
|
return
|
|
}
|
|
err = json.NewDecoder(resp.Body).Decode(&res)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
return
|
|
}
|
|
|
|
// Return the Status of a job
|
|
func Status(id string) (status StatusResp, err error) {
|
|
resp, err := getReq("status/" + id)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
err = json.NewDecoder(resp.Body).Decode(&status)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
return
|
|
}
|
|
|
|
// Run a job
|
|
func Run(data []byte) (req Request, err error) {
|
|
resp, err := postReq(data, "run")
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
err = json.NewDecoder(resp.Body).Decode(&req)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
return
|
|
}
|
|
|
|
// Poll job until it succeeded or failed and return it
|
|
func Poll(id string, delay time.Duration) (status StatusResp, err error) {
|
|
for {
|
|
status, err = Status(id)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
switch status.Status {
|
|
case "COMPLETED":
|
|
return
|
|
case "IN_QUEUE", "IN_PROGRESS":
|
|
time.Sleep(delay)
|
|
continue
|
|
default:
|
|
return status, fmt.Errorf("job failed with status '%s'", status.Status)
|
|
}
|
|
}
|
|
}
|