No need to use different implementations for pull models between win/unix for now
This commit is contained in:
parent
ab0738b920
commit
99af02a4f9
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
Loading…
Reference in New Issue