388 lines
11 KiB
Go
388 lines
11 KiB
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
|
|
}
|