No need to use different implementations for pull models between win/unix for now

This commit is contained in:
Ignasi 2025-03-06 12:10:37 +01:00 committed by Jacob Howard
parent ab0738b920
commit 99af02a4f9
No known key found for this signature in database
GPG Key ID: 3E8B8F7FEB46FC66
3 changed files with 23 additions and 46 deletions

View File

@ -1,6 +1,7 @@
package models
import (
"context"
"encoding/json"
"errors"
"fmt"
@ -235,3 +236,25 @@ func (m *Manager) GetModelPath(ref string) (string, error) {
}
return fmt.Sprintf("%s/blobs/%s/%s", m.distributionClient.GetStorePath(), parts[0], parts[1]), nil
}
// PullModel pulls a model to local storage. Any error it returns is suitable
// for writing back to the client.
func (m *Manager) PullModel(ctx context.Context, model string) error {
// Restrict model pull concurrency.
select {
case <-m.pullTokens:
case <-ctx.Done():
return context.Canceled
}
defer func() {
m.pullTokens <- struct{}{}
}()
// Pull the model using the Docker model distribution client
m.log.Infoln("Pulling model:", model)
if _, err := m.distributionClient.PullModel(ctx, model); err != nil {
return fmt.Errorf("error while pulling model: %w", err)
}
return nil
}

View File

@ -1,30 +0,0 @@
//go:build !windows
package models
import (
"context"
"fmt"
)
// PullModel pulls a model to local storage. Any error it returns is suitable
// for writing back to the client.
func (m *Manager) PullModel(ctx context.Context, model string) error {
// Restrict model pull concurrency.
select {
case <-m.pullTokens:
case <-ctx.Done():
return context.Canceled
}
defer func() {
m.pullTokens <- struct{}{}
}()
// Pull the model using the Docker model distribution client
m.log.Infoln("Pulling model:", model)
if _, err := m.distributionClient.PullModel(ctx, model); err != nil {
return fmt.Errorf("error while pulling model: %w", err)
}
return nil
}

View File

@ -1,16 +0,0 @@
package models
import (
"context"
"errors"
)
// PullModel is unimplemented on Windows.
func (m *Manager) PullModel(ctx context.Context, model string) error {
// TODO: The Hugging Face Hub package that we're using doesn't build on
// Windows due to non-portable syscall package usage. With
// docker/model-distribution just a few days away, it's not worth patching
// or reimplementing. Once we switch to docker/model-distribution, delete
// the _unix.go / _windows.go variants and move PullModel into manager.go.
return errors.New("model pulls not yet supported on Windows")
}