docker
/
model-runner
mirror of https://github.com/docker/model-runner.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Initialize model distribution in model-runner (#8) 

* Initialize model distribution in model-runner. This allows to remove model-distribution dependency in Pinata

* Fix main.go

* bump model-distribution

* gofumpt -l -extra -w .
This commit is contained in:
Ignasi 2025-04-10 17:34:56 +02:00 committed by GitHub
parent dc011af9e2
commit 055b8e451b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 34 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
go.mod
View File

@ -5,7 +5,7 @@ go 1.23.7
require (
github.com/containerd/containerd/v2 v2.0.4
github.com/containerd/platforms v1.0.0-rc.1
github.com/docker/model-distribution v0.0.0-20250410082334-7cc4acf00e0e
github.com/docker/model-distribution v0.0.0-20250410151231-bf9b59b512f7
github.com/opencontainers/go-digest v1.0.0
github.com/opencontainers/image-spec v1.1.1
github.com/sirupsen/logrus v1.9.3

4
go.sum
View File

@ -49,8 +49,8 @@ github.com/docker/go-connections v0.5.0 h1:USnMq7hx7gwdVZq1L49hLXaFtUdTADjXGp+uj
github.com/docker/go-connections v0.5.0/go.mod h1:ov60Kzw0kKElRwhNs9UlUHAE/F9Fe6GLaXnqyDdmEXc=
github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4=
github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
github.com/docker/model-distribution v0.0.0-20250410082334-7cc4acf00e0e h1:Buqp/Bv4M0AdVQlcGAE2Ua6rHf+3KtQuctZOX29Aexc=
github.com/docker/model-distribution v0.0.0-20250410082334-7cc4acf00e0e/go.mod h1:/JWSwYc3pihCpHqFzDUyoiRKegA1srfYESxRh/vJE10=
github.com/docker/model-distribution v0.0.0-20250410151231-bf9b59b512f7 h1:li7LReF/bddqOXbX7IfAnzRmYYAuMjG4C7xCMQLZPlI=
github.com/docker/model-distribution v0.0.0-20250410151231-bf9b59b512f7/go.mod h1:/JWSwYc3pihCpHqFzDUyoiRKegA1srfYESxRh/vJE10=
github.com/ebitengine/purego v0.8.2 h1:jPPGWs2sZ1UgOSgD2bClL0MJIqu58nOmIcBuXr62z1I=
github.com/ebitengine/purego v0.8.2/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ=
github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=

12
main.go
View File

@ -9,7 +9,6 @@ import (
"path/filepath"
"syscall"
"github.com/docker/model-distribution/pkg/distribution"
"github.com/docker/model-runner/pkg/inference"
"github.com/docker/model-runner/pkg/inference/backends/llamacpp"
"github.com/docker/model-runner/pkg/inference/models"
@ -32,14 +31,11 @@ func main() {
if err != nil {
log.Fatalf("Failed to get user home directory: %v", err)
}
distributionClient, err := distribution.NewClient(
distribution.WithStoreRootPath(filepath.Join(userHomeDir, ".docker", "models")),
)
if err != nil {
log.Fatalf("Failed to create distribution client: %v", err)
}
modelManager := models.NewManager(log, distributionClient)
modelManager := models.NewManager(log, models.ClientConfig{
StoreRootPath: filepath.Join(userHomeDir, ".docker", "models"),
Logger: log.WithFields(logrus.Fields{"component": "model-manager"}),
})
llamaCppBackend, err := llamacpp.New(
log,

29
pkg/inference/models/manager.go
View File

@ -12,6 +12,7 @@ import (
"github.com/docker/model-distribution/pkg/types"
"github.com/docker/model-runner/pkg/inference"
"github.com/docker/model-runner/pkg/logging"
"github.com/sirupsen/logrus"
)
const (
@ -33,14 +34,38 @@ type Manager struct {
distributionClient *distribution.Client
}
type ClientConfig struct {
// StoreRootPath is the root path for the model store.
StoreRootPath string
// Logger is the logger to use.
Logger *logrus.Entry
// Transport is the HTTP transport to use.
Transport http.RoundTripper
// UserAgent is the user agent to use.
UserAgent string
}
// NewManager creates a new model's manager.
func NewManager(log logging.Logger, client *distribution.Client) *Manager {
func NewManager(log logging.Logger, c ClientConfig) *Manager {
// Create the model distribution client.
distributionClient, err := distribution.NewClient(
distribution.WithStoreRootPath(c.StoreRootPath),
distribution.WithLogger(c.Logger),
distribution.WithTransport(c.Transport),
distribution.WithUserAgent(c.UserAgent),
)
if err != nil {
log.Errorf("Failed to create distribution client: %v", err)
// Continue without distribution client. The model manager will still
// respond to requests, but may return errors if the client is required.
}
// Create the manager.
m := &Manager{
log: log,
pullTokens: make(chan struct{}, maximumConcurrentModelPulls),
router: http.NewServeMux(),
distributionClient: client,
distributionClient: distributionClient,
}
// Register routes.