mlflow-client-go/run.go

package mlflow

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
)

type SearchRunReply struct {
	Runs          []Run  `json:"runs,omitempty"`
	NextPageToken string `json:"next_page_token,omitempty"`
}

type UpdateRunRequest struct {
	RunId   string      `json:"run_id,omitempty"`
	Status  RunStatus   `json:"status,omitempty"`
	EndTime TimestampMs `json:"end_time,omitempty"`
	RunName string      `json:"run_name,omitempty"`
}

type UpdateRunReply struct {
	RunInfo RunInfo `json:"run_info,omitempty"`
}

func CreateRun(conf Config, req CreateRunRequest) (*Run, error) {
	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/create", "application/json", bytes.NewReader([]byte(req.serialize())))
	r := CreateRunReply{}
	err = apiReadReply(resp, err, &r)
	if err != nil {
		return nil, err
	}

	return &r.Run, nil
}

func DeleteRun(conf Config, id string) error {
	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/delete", "application/json", bytes.NewReader([]byte(`{"run_id":"`+id+`"}`)))
	err = apiReadReply(resp, err, nil)
	if err != nil {
		return err
	}

	return nil
}

func RestoreRun(conf Config, id string) (*Run, error) {
	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/restore", "application/json", bytes.NewReader([]byte(`{"run_id":"`+id+`"}`)))
	err = apiReadReply(resp, err, nil)
	if err != nil {
		return nil, err
	}

	return GetRun(conf, id)
}

func GetRun(conf Config, id string) (*Run, error) {
	resp, err := http.Get(conf.ApiURI + "/api/2.0/mlflow/runs/get" + "?run_id=" + id)
	run := GetRunReply{}
	err = apiReadReply(resp, err, &run)
	if err != nil {
		return nil, err
	}

	return &run.Run, nil
}

func SearchRuns(conf Config, opts SearchRunRequest) (*SearchRunReply, error) {
	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/search", "application/json", opts.getReader())
	rv := SearchRunReply{}
	err = apiReadReply(resp, err, &rv)
	if err != nil {
		return nil, err
	}

	return &rv, nil
}

func UpdateRun(conf Config, opts UpdateRunRequest) (*RunInfo, error) {
	args, err := json.Marshal(opts)
	if err != nil {
		return nil, err
	}
	rv := UpdateRunReply{}
	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/update", "application/json", bytes.NewReader([]byte(args)))

	err = apiReadReply(resp, err, &rv)
	if err != nil {
		return nil, err
	}

	return &rv.RunInfo, nil
}

func LogMetric(conf Config, opts LogMetricRequest) error {
	args, err := json.Marshal(opts.__init())
	if err != nil {
		return err
	}

	fmt.Println(opts)
	fmt.Println(string(args))

	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/log-metric", "application/json", bytes.NewReader([]byte(args)))
	err = apiReadReply(resp, err, nil)
	if err != nil {
		return err
	}

	return nil
}

func LogBatch(conf Config, opts LogBatchRequest) error {
	for i, metric := range opts.Metrics {
		opts.Metrics[i] = *metric.__init()
	}

	args, err := json.Marshal(opts)
	if err != nil {
		return err
	}
	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/log-batch", "application/json", bytes.NewReader([]byte(args)))
	err = apiReadReply(resp, err, nil)
	if err != nil {
		return err
	}

	return nil
}

func LogInputs(conf Config, opts LogInputsRequest) error {
	args, err := json.Marshal(opts)
	if err != nil {
		return err
	}

	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/log-inputs", "application/json", bytes.NewReader([]byte(args)))
	err = apiReadReply(resp, err, nil)
	if err != nil {
		return err
	}

	return nil
}

func SetTag(conf Config, opts SetTagRequest) error {
	args, err := json.Marshal(opts)
	if err != nil {
		return err
	}

	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/set-tag", "application/json", bytes.NewReader([]byte(args)))
	err = apiReadReply(resp, err, nil)
	if err != nil {
		return err
	}

	return nil
}

func DeleteTag(conf Config, opts DeleteTagRequest) error {
	args, err := json.Marshal(opts)
	if err != nil {
		return err
	}

	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/delete-tag", "application/json", bytes.NewReader([]byte(args)))
	err = apiReadReply(resp, err, nil)
	if err != nil {
		return err
	}

	return nil
}

func LogParam(conf Config, opts LogParamRequest) error {
	args, err := json.Marshal(opts)
	if err != nil {
		return err
	}

	resp, err := http.Post(conf.ApiURI+"/api/2.0/mlflow/runs/log-parameter", "application/json", bytes.NewReader([]byte(args)))
	err = apiReadReply(resp, err, nil)
	if err != nil {
		return err
	}

	return nil
}
func GetMetricHistory(conf Config, opts GetMetricHistoryRequest) (*GetMetricHistoryResponse, error) {

	rv := GetMetricHistoryResponse{}
	resp, err := http.Get(conf.ApiURI + "/api/2.0/mlflow/metrics/get-history" + UrlEncode(*opts.__init()))

	err = apiReadReply(resp, err, &rv)
	if err != nil {
		return nil, err
	}

	return &rv, nil
}