mlflow

package module
v0.1.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jun 28, 2026 License: MIT Imports: 13 Imported by: 0

README

mlflow

Go Reference Go Report Card

Standalone, idiomatic Go client for the MLflow REST API v2.0 (tracking + model registry). There is no official Go MLflow client; community efforts are server-focused and partial. This is a clean, dependency-free tracking + registry client so Go programs can report experiments, runs, metrics, and models to MLflow natively.

  • No dependencies. Standard library only. Requires Go 1.23+.
  • Complete tracking + registry. Experiments, runs, logging, model registry (registered models, versions, stages, aliases), and artifact list/download.
  • Idiomatic. context.Context first, explicit errors, typed structs, functional options, an injectable transport so methods are unit-testable without a server.

Install

go get github.com/guygrigsby/mlflow

Usage

package main

import (
	"context"
	"log"
	"time"

	"github.com/guygrigsby/mlflow"
)

func main() {
	// "" reads MLFLOW_TRACKING_URI from the environment.
	c, err := mlflow.NewClient("")
	if err != nil {
		log.Fatal(err)
	}
	ctx := context.Background()

	exp, err := c.GetOrCreateExperiment(ctx, "lm-100m-en")
	if err != nil {
		log.Fatal(err)
	}

	now := time.Now().UnixMilli()
	run, err := c.CreateRun(ctx, exp, mlflow.WithRunName("r1"), mlflow.WithStartTime(now))
	if err != nil {
		log.Fatal(err)
	}
	id := run.Info.RunID

	err = c.LogBatch(ctx, id,
		[]mlflow.Metric{{Key: "loss", Value: 1.5, Timestamp: now, Step: 0}},
		[]mlflow.Param{{Key: "lr", Value: "4e-4"}},
		[]mlflow.RunTag{{Key: "stage", Value: "train"}})
	if err != nil {
		log.Fatal(err)
	}

	if err := c.UpdateRun(ctx, id, mlflow.StatusFinished, time.Now().UnixMilli()); err != nil {
		log.Fatal(err)
	}
}

Configuration

NewClient takes the tracking URI (or "" to read MLFLOW_TRACKING_URI) plus options:

mlflow.NewClient(uri,
	mlflow.WithBearerToken(token),       // Authorization: Bearer
	mlflow.WithBasicAuth(user, pass),    // HTTP basic auth
	mlflow.WithHTTPClient(httpClient),   // default: 30s timeout
	mlflow.WithMaxRetries(3),            // 5xx + transport errors, exponential backoff
)

The client is synchronous and best-effort: it retries transient failures (5xx and transport errors) with exponential backoff, never retries 4xx, and never buffers. Callers decide what to do on error.

Errors

Non-2xx responses map to *mlflow.APIError (Code, Message, HTTPStatus). Use mlflow.IsNotFound(err) to test for RESOURCE_DOES_NOT_EXIST:

exp, err := c.GetExperimentByName(ctx, "missing")
if mlflow.IsNotFound(err) {
	// create it, etc.
}

Scope

Implements 100% of the MLflow REST API v2.0 for tracking and the model registry, plus artifact list and proxy get-artifact (download). Out of scope: artifact-store upload (the per-backend S3/GCS/Azure plumbing), the auth/permissions admin plugin, and MLflow 3.x GenAI surfaces (traces, logged-models, datasets).

License

MIT

Documentation

Overview

Package mlflow is a client for the MLflow REST API (tracking + model registry).

Construct a Client with NewClient (it reads MLFLOW_TRACKING_URI from the environment when given ""), then call typed methods. The client is synchronous and best-effort: it retries transient failures but never buffers, so callers decide what to do on error.

Example

Example shows a full run lifecycle: resolve an experiment, start a run, log a batch of metrics/params/tags, then mark the run finished. 

package main

import (
	"context"
	"log"
	"time"

	"github.com/guygrigsby/mlflow"
)

func main() {
	c, err := mlflow.NewClient("") // reads MLFLOW_TRACKING_URI
	if err != nil {
		log.Fatal(err)
	}
	ctx := context.Background()

	exp, err := c.GetOrCreateExperiment(ctx, "lm-100m-en")
	if err != nil {
		log.Fatal(err)
	}

	now := time.Now().UnixMilli()
	run, err := c.CreateRun(ctx, exp, mlflow.WithRunName("r1"), mlflow.WithStartTime(now))
	if err != nil {
		log.Fatal(err)
	}

	err = c.LogBatch(ctx, run.Info.RunID,
		[]mlflow.Metric{{Key: "loss", Value: 1.5, Timestamp: now, Step: 0}},
		[]mlflow.Param{{Key: "lr", Value: "4e-4"}},
		[]mlflow.RunTag{{Key: "stage", Value: "train"}})
	if err != nil {
		log.Fatal(err)
	}

	if err := c.UpdateRun(ctx, run.Info.RunID, mlflow.StatusFinished, time.Now().UnixMilli()); err != nil {
		log.Fatal(err)
	}
}

Index

Examples

Constants

This section is empty.

Variables

This section is empty.

Functions

func IsNotFound 

func IsNotFound(err error) bool

IsNotFound reports whether err is an APIError with RESOURCE_DOES_NOT_EXIST.

Example

ExampleIsNotFound shows branching on a missing resource. 

package main

import (
	"context"
	"fmt"
	"log"

	"github.com/guygrigsby/mlflow"
)

func main() {
	c, err := mlflow.NewClient("")
	if err != nil {
		log.Fatal(err)
	}
	ctx := context.Background()

	if _, err := c.GetExperimentByName(ctx, "does-not-exist"); mlflow.IsNotFound(err) {
		fmt.Println("not found, creating")
		_, _ = c.CreateExperiment(ctx, "does-not-exist")
	}
}

Types

type APIError 

type APIError struct {
	Code       string // MLflow error_code, e.g. "RESOURCE_DOES_NOT_EXIST"
	Message    string
	HTTPStatus int
}

APIError is a non-2xx response from MLflow, mapped from its {"error_code","message"} envelope.

func (*APIError) Error 

func (e *APIError) Error() string

type Client 

type Client struct {
	// contains filtered or unexported fields
}

func NewClient 

func NewClient(trackingURI string, opts ...Option) (*Client, error)

NewClient builds a Client. trackingURI == "" reads MLFLOW_TRACKING_URI from the environment; empty in both is an error. A trailing slash is trimmed and "/api/2.0/mlflow/<path>" is appended by the transport.

func (*Client) CreateExperiment 

func (c *Client) CreateExperiment(ctx context.Context, name string, tags ...ExperimentTag) (string, error)

CreateExperiment creates a new experiment with the given name and optional tags. Returns the new experiment's ID.

func (*Client) CreateModelVersion 

func (c *Client) CreateModelVersion(ctx context.Context, name, source string, runID string, tags ...ModelVersionTag) (*ModelVersion, error)

CreateModelVersion registers a new model version. runID may be empty.

func (*Client) CreateRegisteredModel 

func (c *Client) CreateRegisteredModel(ctx context.Context, name string, tags ...RegisteredModelTag) (*RegisteredModel, error)

CreateRegisteredModel creates a new registered model with the given name and optional tags.

func (*Client) CreateRun 

func (c *Client) CreateRun(ctx context.Context, experimentID string, opts ...RunOption) (*Run, error)

CreateRun creates a new run under experimentID.

func (*Client) DeleteExperiment 

func (c *Client) DeleteExperiment(ctx context.Context, id string) error

DeleteExperiment marks an experiment as deleted.

func (*Client) DeleteModelVersion 

func (c *Client) DeleteModelVersion(ctx context.Context, name, version string) error

DeleteModelVersion deletes a specific model version.

func (*Client) DeleteModelVersionTag 

func (c *Client) DeleteModelVersionTag(ctx context.Context, name, version, key string) error

DeleteModelVersionTag deletes a tag from a model version.

func (*Client) DeleteRegisteredModel 

func (c *Client) DeleteRegisteredModel(ctx context.Context, name string) error

DeleteRegisteredModel deletes a registered model by name.

func (*Client) DeleteRegisteredModelAlias 

func (c *Client) DeleteRegisteredModelAlias(ctx context.Context, name, alias string) error

DeleteRegisteredModelAlias removes an alias from a registered model.

func (*Client) DeleteRegisteredModelTag 

func (c *Client) DeleteRegisteredModelTag(ctx context.Context, name, key string) error

DeleteRegisteredModelTag deletes a tag from a registered model.

func (*Client) DeleteRun 

func (c *Client) DeleteRun(ctx context.Context, runID string) error

DeleteRun soft-deletes a run.

func (*Client) DeleteRunTag 

func (c *Client) DeleteRunTag(ctx context.Context, runID, key string) error

DeleteRunTag removes a tag from a run.

func (*Client) DownloadArtifact 

func (c *Client) DownloadArtifact(ctx context.Context, runID, path string) ([]byte, error)

DownloadArtifact fetches the raw bytes of a single artifact file.

func (*Client) GetExperimentByName 

func (c *Client) GetExperimentByName(ctx context.Context, name string) (*Experiment, error)

GetExperimentByName fetches an experiment by name.

func (*Client) GetLatestVersions 

func (c *Client) GetLatestVersions(ctx context.Context, name string, stages ...string) ([]ModelVersion, error)

GetLatestVersions returns the latest model versions for a registered model, optionally filtered to the given stages.

func (*Client) GetMetricHistory 

func (c *Client) GetMetricHistory(ctx context.Context, runID, metricKey string) ([]Metric, error)

GetMetricHistory returns the full history of a metric from metrics/get-history.

func (*Client) GetModelVersion 

func (c *Client) GetModelVersion(ctx context.Context, name, version string) (*ModelVersion, error)

GetModelVersion retrieves a specific model version by name and version number.

func (*Client) GetModelVersionByAlias 

func (c *Client) GetModelVersionByAlias(ctx context.Context, name, alias string) (*ModelVersion, error)

GetModelVersionByAlias retrieves the model version that the given alias points to.

func (*Client) GetModelVersionDownloadURI 

func (c *Client) GetModelVersionDownloadURI(ctx context.Context, name, version string) (string, error)

GetModelVersionDownloadURI returns the URI for downloading a model version's artifacts.

func (*Client) GetOrCreateExperiment 

func (c *Client) GetOrCreateExperiment(ctx context.Context, name string) (string, error)

GetOrCreateExperiment returns the ID of an existing experiment by name, creating it if it does not exist.

func (*Client) GetRegisteredModel 

func (c *Client) GetRegisteredModel(ctx context.Context, name string) (*RegisteredModel, error)

GetRegisteredModel retrieves a registered model by name.

func (*Client) GetRun 

func (c *Client) GetRun(ctx context.Context, runID string) (*Run, error)

GetRun fetches a run by ID.

func (*Client) ListArtifacts 

func (c *Client) ListArtifacts(ctx context.Context, runID, path string) (files []FileInfo, nextPageToken string, err error)

ListArtifacts returns the artifact listing for runID under path. Pass path=="" to list the root. nextPageToken is non-empty when more pages remain.

func (*Client) LogBatch 

func (c *Client) LogBatch(ctx context.Context, runID string, metrics []Metric, params []Param, tags []RunTag) error

LogBatch posts metrics, params, and tags in chunks that satisfy the MLflow API limits per call: ≤1000 metrics, ≤100 params, ≤100 tags, ≤1000 total. Chunks are sent sequentially; the first error aborts and is returned.

func (*Client) LogInputs 

func (c *Client) LogInputs(ctx context.Context, runID string, datasets []DatasetInput) error

LogInputs posts dataset inputs to runs/log-inputs.

func (*Client) LogMetric 

func (c *Client) LogMetric(ctx context.Context, runID, key string, value float64, timestampMs, step int64) error

LogMetric posts a single metric value to runs/log-metric.

Example

ExampleClient_LogMetric streams a metric across steps during a training loop. 

package main

import (
	"context"
	"log"
	"time"

	"github.com/guygrigsby/mlflow"
)

func main() {
	c, err := mlflow.NewClient("")
	if err != nil {
		log.Fatal(err)
	}
	ctx := context.Background()

	runID := "abc123"
	losses := []float64{2.1, 1.7, 1.4}
	for step, loss := range losses {
		ts := time.Now().UnixMilli()
		if err := c.LogMetric(ctx, runID, "loss", loss, ts, int64(step)); err != nil {
			log.Fatal(err)
		}
	}
}

func (*Client) LogParam 

func (c *Client) LogParam(ctx context.Context, runID, key, value string) error

LogParam posts a single parameter to runs/log-parameter.

func (*Client) RenameRegisteredModel 

func (c *Client) RenameRegisteredModel(ctx context.Context, name, newName string) (*RegisteredModel, error)

RenameRegisteredModel renames a registered model.

func (*Client) RestoreExperiment 

func (c *Client) RestoreExperiment(ctx context.Context, id string) error

RestoreExperiment restores a deleted experiment.

func (*Client) RestoreRun 

func (c *Client) RestoreRun(ctx context.Context, runID string) error

RestoreRun undeletes a previously deleted run.

func (*Client) SearchExperiments 

func (c *Client) SearchExperiments(ctx context.Context, req SearchExperimentsRequest) ([]Experiment, string, error)

SearchExperiments searches experiments with the given criteria and returns matching experiments and the next-page token (empty string when no more pages).

func (*Client) SearchModelVersions 

func (c *Client) SearchModelVersions(ctx context.Context, filter string, maxResults int64, orderBy []string, pageToken string) ([]ModelVersion, string, error)

SearchModelVersions searches model versions with optional filter, ordering, and pagination. Returns the matching versions and the next page token.

func (*Client) SearchRegisteredModels 

func (c *Client) SearchRegisteredModels(ctx context.Context, filter string, maxResults int64, orderBy []string, pageToken string) ([]RegisteredModel, string, error)

SearchRegisteredModels searches registered models with optional filter, ordering, and pagination. Returns the matching models and the next page token (empty when there are no more pages).

func (*Client) SearchRuns 

func (c *Client) SearchRuns(ctx context.Context, req SearchRunsRequest) ([]Run, string, error)

SearchRuns searches runs matching req. Returns the matching runs and an opaque next-page token (empty when no further pages exist).

func (*Client) SetExperimentTag 

func (c *Client) SetExperimentTag(ctx context.Context, id, key, value string) error

SetExperimentTag sets a tag on an experiment.

func (*Client) SetModelVersionTag 

func (c *Client) SetModelVersionTag(ctx context.Context, name, version, key, value string) error

SetModelVersionTag sets a tag on a model version.

func (*Client) SetRegisteredModelAlias 

func (c *Client) SetRegisteredModelAlias(ctx context.Context, name, alias, version string) error

SetRegisteredModelAlias creates or updates an alias pointing to a specific model version.

func (*Client) SetRegisteredModelTag 

func (c *Client) SetRegisteredModelTag(ctx context.Context, name, key, value string) error

SetRegisteredModelTag sets a tag on a registered model.

func (*Client) SetRunTag 

func (c *Client) SetRunTag(ctx context.Context, runID, key, value string) error

SetRunTag sets a tag on a run.

func (*Client) SetTag 

func (c *Client) SetTag(ctx context.Context, runID, key, value string) error

SetTag posts a single run tag to runs/set-tag. It is an alias for SetRunTag.

func (*Client) TransitionModelVersionStage 

func (c *Client) TransitionModelVersionStage(ctx context.Context, name, version, stage string, archiveExisting bool) (*ModelVersion, error)

TransitionModelVersionStage transitions a model version to a new stage. When archiveExisting is true, all other versions in the target stage are archived.

func (*Client) UpdateExperiment 

func (c *Client) UpdateExperiment(ctx context.Context, id, newName string) error

UpdateExperiment renames an experiment.

func (*Client) UpdateModelVersion 

func (c *Client) UpdateModelVersion(ctx context.Context, name, version, description string) (*ModelVersion, error)

UpdateModelVersion updates the description of a model version.

func (*Client) UpdateRegisteredModel 

func (c *Client) UpdateRegisteredModel(ctx context.Context, name, description string) (*RegisteredModel, error)

UpdateRegisteredModel updates the description of a registered model.

func (*Client) UpdateRun 

func (c *Client) UpdateRun(ctx context.Context, runID string, status RunStatus, endTime int64) error

UpdateRun sets the status and optionally the end time of a run. endTime is included only when nonzero.

type Dataset 

type Dataset struct {
	Name       string `json:"name"`
	Digest     string `json:"digest,omitempty"`
	SourceType string `json:"source_type,omitempty"`
	Source     string `json:"source,omitempty"`
	Schema     string `json:"schema,omitempty"`
	Profile    string `json:"profile,omitempty"`
}

Dataset describes a logged dataset.

type DatasetInput 

type DatasetInput struct {
	Tags    []RunTag `json:"tags,omitempty"`
	Dataset Dataset  `json:"dataset"`
}

DatasetInput is a dataset associated with a run via LogInputs.

type Experiment 

type Experiment struct {
	ExperimentID     string          `json:"experiment_id"`
	Name             string          `json:"name"`
	ArtifactLocation string          `json:"artifact_location,omitempty"`
	LifecycleStage   string          `json:"lifecycle_stage,omitempty"`
	Tags             []ExperimentTag `json:"tags,omitempty"`
}

type ExperimentTag 

type ExperimentTag struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

type FileInfo 

type FileInfo struct {
	Path     string `json:"path"`
	IsDir    bool   `json:"is_dir"`
	FileSize int64  `json:"file_size,omitempty"`
}

FileInfo describes a single artifact entry returned by ListArtifacts.

type Metric 

type Metric struct {
	Key       string  `json:"key"`
	Value     float64 `json:"value"`
	Timestamp int64   `json:"timestamp"` // unix ms
	Step      int64   `json:"step"`
}

type ModelVersion 

type ModelVersion struct {
	Name                 string            `json:"name"`
	Version              string            `json:"version"`
	CreationTimestamp    int64             `json:"creation_timestamp,omitempty"`
	LastUpdatedTimestamp int64             `json:"last_updated_timestamp,omitempty"`
	CurrentStage         string            `json:"current_stage,omitempty"`
	Description          string            `json:"description,omitempty"`
	Source               string            `json:"source,omitempty"`
	RunID                string            `json:"run_id,omitempty"`
	Status               string            `json:"status,omitempty"`
	Tags                 []ModelVersionTag `json:"tags,omitempty"`
	Aliases              []string          `json:"aliases,omitempty"`
}

ModelVersion describes a specific version of a registered model.

type ModelVersionTag 

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

ModelVersionTag is a key-value tag on a model version.

type Option 

type Option func(*config)

func WithBasicAuth 

func WithBasicAuth(u, p string) Option

WithBasicAuth sets HTTP basic auth on every request.

func WithBearerToken 

func WithBearerToken(tok string) Option

WithBearerToken sets an Authorization: Bearer <tok> header on every request.

func WithHTTPClient 

func WithHTTPClient(h *http.Client) Option

WithHTTPClient overrides the default &http.Client{Timeout: 30s}.

func WithMaxRetries 

func WithMaxRetries(n int) Option

WithMaxRetries sets the retry budget for 5xx + transport errors (default 3). Retries are not idempotent-aware: a write whose response is lost in transit (transport error or 5xx after the server applied it) may be re-sent, so a retried CreateRun/LogMetric/LogBatch can double-write. This is acceptable under the best-effort, accept-gaps contract; set 0 to disable.

type Param 

type Param struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

type RegisteredModel 

type RegisteredModel struct {
	Name                 string                 `json:"name"`
	CreationTimestamp    int64                  `json:"creation_timestamp,omitempty"`
	LastUpdatedTimestamp int64                  `json:"last_updated_timestamp,omitempty"`
	Description          string                 `json:"description,omitempty"`
	LatestVersions       []ModelVersion         `json:"latest_versions,omitempty"`
	Tags                 []RegisteredModelTag   `json:"tags,omitempty"`
	Aliases              []RegisteredModelAlias `json:"aliases,omitempty"`
}

RegisteredModel describes a model registered in the MLflow Model Registry.

type RegisteredModelAlias 

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

RegisteredModelAlias maps an alias name to a model version number.

type RegisteredModelTag 

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

RegisteredModelTag is a key-value tag on a registered model.

type Run 

type Run struct {
	Info RunInfo `json:"info"`
	Data RunData `json:"data"`
}

type RunData 

type RunData struct {
	Metrics []Metric `json:"metrics,omitempty"`
	Params  []Param  `json:"params,omitempty"`
	Tags    []RunTag `json:"tags,omitempty"`
}

type RunInfo 

type RunInfo struct {
	RunID          string    `json:"run_id"`
	ExperimentID   string    `json:"experiment_id"`
	RunName        string    `json:"run_name,omitempty"`
	Status         RunStatus `json:"status,omitempty"`
	StartTime      int64     `json:"start_time,omitempty"` // unix ms
	EndTime        int64     `json:"end_time,omitempty"`   // unix ms
	ArtifactURI    string    `json:"artifact_uri,omitempty"`
	LifecycleStage string    `json:"lifecycle_stage,omitempty"`
}

type RunOption 

type RunOption func(*runCreate)

RunOption configures a CreateRun request.

func WithRunName 

func WithRunName(name string) RunOption

WithRunName sets the run name.

func WithRunTags 

func WithRunTags(tags ...RunTag) RunOption

WithRunTags appends tags to the run.

func WithStartTime 

func WithStartTime(unixMs int64) RunOption

WithStartTime sets the run start time as unix milliseconds.

type RunStatus 

type RunStatus string
const (
	StatusRunning   RunStatus = "RUNNING"
	StatusScheduled RunStatus = "SCHEDULED"
	StatusFinished  RunStatus = "FINISHED"
	StatusFailed    RunStatus = "FAILED"
	StatusKilled    RunStatus = "KILLED"
)

type RunTag 

type RunTag struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

type SearchExperimentsRequest 

type SearchExperimentsRequest struct {
	Filter     string   `json:"filter,omitempty"`
	MaxResults int64    `json:"max_results,omitempty"`
	OrderBy    []string `json:"order_by,omitempty"`
	PageToken  string   `json:"page_token,omitempty"`
	ViewType   string   `json:"view_type,omitempty"`
}

SearchExperimentsRequest is the request body for SearchExperiments.

type SearchRunsRequest 

type SearchRunsRequest struct {
	ExperimentIDs []string `json:"experiment_ids"`
	Filter        string   `json:"filter,omitempty"`
	MaxResults    int32    `json:"max_results,omitempty"`
	OrderBy       []string `json:"order_by,omitempty"`
	PageToken     string   `json:"page_token,omitempty"`
	RunViewType   string   `json:"run_view_type,omitempty"`
}

SearchRunsRequest is the request body for SearchRuns.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.