mlflow-client-go/model_types.go

package mlflow

import (
	"encoding/json"
	"fmt"
	"reflect"
	"sort"
)

type ModelVersionStatus string

const ModelVersionStatusPendingRegistration = "PENDING_REGISTRATION"
const ModelVersionStatusFailedRegistration = "FAILED_REGISTRATION"
const ModelVersionStatusReady = "READY"

const ModelVersionStageNone = "None"
const ModelVersionStageStaging = "Staging"
const ModelVersionStageProduction = "Production"
const ModelVersionStageArchived = "Archived"

type RegisteredModelTag struct {
	Key   string `json:"key,omitempty"`
	Value string `json:"value,omitempty"`
}

type ModelVersionTag struct {
	Key   string `json:"key,omitempty"`
	Value string `json:"value,omitempty"`
}

type RegisteredModelAlias struct {
	Alias   string `json:"alias,omitempty"`
	Version string `json:"version,omitempty"`
}

type CreateRegisteredModelRequest struct {
	Name        string               `json:"name,omitempty"`
	Tags        []RegisteredModelTag `json:"tags,omitempty"`
	Description string               `json:"description,omitempty"`
}

func (r *CreateRegisteredModelRequest) serialize() []byte {
	rv, _ := json.Marshal(r)
	return rv
}

type CreateRegisteredModelResponse struct {
	RegisteredModel RegisteredModel `json:"registered_model,omitempty"`
}

type ModelVersion struct {
	Name                string             `json:"name,omitempty"`
	Version             string             `json:"version,omitempty"`
	CreationTimestamp   TimestampMs        `json:"creation_timestamp,omitempty"`
	LastUpdateTimestamp TimestampMs        `json:"last_update_timestamp,omitempty"`
	UserId              string             `json:"user_id,omitempty"`
	CurrentStage        string             `json:"current_stage,omitempty"`
	Description         string             `json:"description,omitempty"`
	Source              string             `json:"source,omitempty"`
	RunId               string             `json:"run_id,omitempty"`
	Status              ModelVersionStatus `json:"status,omitempty"`
	StatusMessage       string             `json:"status_message,omitempty"`
	Tags                []ModelVersionTag  `json:"tags,omitempty"`
	RunLink             string             `json:"run_link,omitempty"`
	Aliases             []string           `json:"aliases,omitempty"`
}

func (ref *ModelVersion) GetTag(key string) *ModelVersionTag {
	for _, tag := range ref.Tags {
		if tag.Key == key {
			return &tag
		}
	}
	return nil
}

func (ref *ModelVersion) IsEqual(target *ModelVersion) bool {
	vRef := reflect.ValueOf(*ref)
	vTarget := reflect.ValueOf(*target)

	typeOfV := vRef.Type()
	values := make([]interface{}, vRef.NumField())
	for i := 0; i < vRef.NumField(); i++ {
		values[i] = vRef.Field(i).Interface()
		refName := typeOfV.Field(i).Name
		refType := typeOfV.Field(i).Type.String()
		refVal := vRef.Field(i)

		targetVal := reflect.Indirect(vTarget).FieldByName(refName)

		switch refType {
		case "string", "mlflow.ModelVersionStatus":
			if refVal.String() != targetVal.String() {
				fmt.Println("NE STRING: ", refType, refName, refVal, targetVal, (refVal.String() == targetVal.String()))
				return false
			}
		case "mlflow.TimestampMs":
			if refVal.Int() != targetVal.Int() {
				fmt.Println("NE TimestampMs: ", refType, refName, refVal, targetVal, (refVal.Int() == targetVal.Int()))
				return false
			}
		case "[]mlflow.ModelVersionTag":
			sliceRef, _ := refVal.Interface().([]ModelVersionTag)
			sliceTarget, _ := targetVal.Interface().([]ModelVersionTag)

			if len(sliceRef) != len(sliceTarget) {
				//fmt.Println("Not equal - len mismatch")
				return false
			}

			for _, val := range sliceRef {
				if tag2 := target.GetTag(val.Key); tag2 == nil || tag2.Value != val.Value {
					return false
				}
			}
		case "[]string":
			sliceRef, _ := refVal.Interface().([]string)
			sliceTarget, _ := targetVal.Interface().([]string)
			sort.Strings(sliceRef)
			sort.Strings(sliceTarget)
			if len(sliceRef) != len(sliceTarget) {
				//fmt.Println("Not equal - len mismatch")
				return false
			}

			for xI, xRef := range sliceRef {
				if xRef != sliceTarget[xI] {
					return false
				}
			}
		default:
			fmt.Println("Unknown: ", refType, refName, refVal, targetVal)
		}

	}

	return true
}

type RegisteredModel struct {
	Name                string                 `json:"name,omitempty"`
	CreationTimestamp   Timestamp              `json:"creation_timestamp,omitempty"`
	LastUpdateTimestamp Timestamp              `json:"last_update_timestamp,omitempty"`
	UserId              string                 `json:"user_id,omitempty"`
	Description         string                 `json:"description,omitempty"`
	LatestVersions      []ModelVersion         `json:"latest_versions,omitempty"`
	Tags                []RegisteredModelTag   `json:"tags,omitempty"`
	Aliases             []RegisteredModelAlias `json:"aliases,omitempty"`
}

type GetRegisteredModelResponse CreateRegisteredModelResponse

type RenameRegisteredModelRequrest struct {
	Name    string `json:"name,omitempty"`
	NewName string `json:"new_name,omitempty"`
}

func (r *RenameRegisteredModelRequrest) serialize() []byte {
	rv, _ := json.Marshal(r)
	return rv
}

type RenameRegisteredModelResponse CreateRegisteredModelResponse

type GetLatestModelVersionsRequest struct {
	Name   string   `json:"name,omitempty"`
	Stages []string `json:"stages,omitempty"`
}

func (r *GetLatestModelVersionsRequest) serialize() []byte {
	rv, _ := json.Marshal(r)
	return rv
}

type GetLatestModelVersionsReponse struct {
	ModelVersions []ModelVersion `json:"model_versions,omitempty"`
}

type GetModelVersionRequest struct {
	Name    string `json:"name,omitempty"`
	Version string `json:"version,omitempty"`
}

type GetModelVersionResponse struct {
	ModelVersion ModelVersion `json:"model_version,omitempty"`
}

type DeleteModelVersionRequest GetModelVersionRequest

func (r *DeleteModelVersionRequest) serialize() []byte {
	rv, _ := json.Marshal(r)
	return rv
}

type CreateModelVersionRequest struct {
	Name        string            `json:"name,omitempty"`
	Source      string            `json:"source,omitempty"`
	RunId       string            `json:"run_id,omitempty"`
	Tags        []ModelVersionTag `json:"tags,omitempty"`
	RunLink     string            `json:"run_link,omitempty"`
	Description string            `json:"description,omitempty"`
}

type CreateModelVersionResponse GetModelVersionResponse

func (r *CreateModelVersionRequest) serialize() []byte {
	rv, _ := json.Marshal(r)
	return rv
}

type GetDownloadUriForModelVersionArtifactsRequest GetModelVersionRequest

type GetDownloadUriForModelVersionArtifactsResponse struct {
	ArtifactUri string `json:"artifact_uri,omitempty"`
}

type SearchModelVersionsRequest struct {
	// Valid keys are '{'version_number', 'source_path', 'run_id', 'name'}'
	Filter    string `json:"filter,omitempty"`
	PageSize  int64  `json:"max_results,omitempty"`
	OrderBy   string `json:"order_by,omitempty"`
	PageToken string `json:"page_token,omitempty"`
}

func (e *SearchModelVersionsRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

func (e *SearchModelVersionsRequest) __init() *SearchModelVersionsRequest {
	if e.PageSize == 0 {
		e.PageSize = 1000
	}

	return e
}

type SearchModelVersionsResponse struct {
	ModelVersions []ModelVersion `json:"model_versions,omitempty"`
	NextPageToken string         `json:"next_page_token,omitempty"`
}

type UpdateModelVersionRequest struct {
	Name        string `json:"name,omitempty"`
	Version     string `json:"version,omitempty"`
	Description string `json:"description,omitempty"`
}

type UpdateModelVersionResponse GetModelVersionResponse

func (e *UpdateModelVersionRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type UpdateRegisteredModelRequest struct {
	Name        string `json:"name,omitempty"`
	Description string `json:"description,omitempty"`
}

type UpdateRegisteredModelResponse GetRegisteredModelResponse

func (e *UpdateRegisteredModelRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type TransitionModelVersionStageRequest struct {
	Name                    string `json:"name,omitempty"`
	Version                 string `json:"version,omitempty"`
	Stage                   string `json:"stage,omitempty"`
	ArchiveExistingVersions bool   `json:"archive_existing_versions,omitempty"`
}

func (e *TransitionModelVersionStageRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type TransitionModelVersionStageResponse GetModelVersionResponse

type SearchRegisteredModelsRequest struct {
	// Valid keys are '{'name'}'
	Filter     string `json:"filter,omitempty"`
	MaxResults int64  `json:"max_results,omitempty"`
	OrderBy    string `json:"order_by,omitempty"`
	PageToken  string `json:"page_token,omitempty"`
}

func (e *SearchRegisteredModelsRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

func (e *SearchRegisteredModelsRequest) __init() *SearchRegisteredModelsRequest {
	if e.MaxResults == 0 {
		e.MaxResults = 1000
	}

	return e
}

type SearchRegisteredModelsResponse struct {
	RegisteredModels []RegisteredModel `json:"registered_models,omitempty"`
	NextPageToken    string            `json:"next_page_token,omitempty"`
}

type SetRegisteredModelTagRequest struct {
	RegisteredModelTag
	Name string `json:"name,omitempty"`
}

func (e *SetRegisteredModelTagRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type SetModelVersionTagRequest struct {
	ModelVersionTag
	Name    string `json:"name,omitempty"`
	Version string `json:"version,omitempty"`
}

func (e *SetModelVersionTagRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type DeleteRegisteredModelTagRequest struct {
	Key  string `json:"key,omitempty"`
	Name string `json:"name,omitempty"`
}

func (e *DeleteRegisteredModelTagRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type DeleteModelVersionTagRequest struct {
	Key     string `json:"key,omitempty"`
	Name    string `json:"name,omitempty"`
	Version string `json:"version,omitempty"`
}

func (e *DeleteModelVersionTagRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type SetRegisteredModelAliasRequest struct {
	Name    string `json:"name,omitempty"`
	Alias   string `json:"alias,omitempty"`
	Version string `json:"version,omitempty"`
}

func (e *SetRegisteredModelAliasRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type GetModelVersionByAliasRequest struct {
	Name  string `json:"name,omitempty"`
	Alias string `json:"alias,omitempty"`
}

func (e *GetModelVersionByAliasRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type DeleteRegisteredModelAliasRequest struct {
	Name  string `json:"name,omitempty"`
	Alias string `json:"alias,omitempty"`
}

func (e *DeleteRegisteredModelAliasRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type GetModelVersionByAliasResponse GetModelVersionResponse

type SearchModelVersionsRequestAdvanced struct {
	SearchModelVersionsRequest
	Stage string
	Limit int64
}