remove tmp directories created by previous servers (#559)

* remove tmp directories created by previous servers

* clean up on server stop

* Update routes.go

* Update server/routes.go

Co-authored-by: Jeffrey Morgan <jmorganca@gmail.com>

* create top-level temp ollama dir

* check file exists before creating

---------

Co-authored-by: Jeffrey Morgan <jmorganca@gmail.com>
Co-authored-by: Michael Yang <mxyng@pm.me>
This commit is contained in:
Bruce MacDonald
2023-09-21 20:38:49 +01:00
committed by GitHub
parent 8c83701e9f
commit 4cba75efc5
6 changed files with 56 additions and 60 deletions

View File

@@ -267,7 +267,7 @@ func filenameWithPath(path, f string) (string, error) {
return f, nil
}
func CreateModel(ctx context.Context, name string, path string, fn func(resp api.ProgressResponse)) error {
func CreateModel(ctx context.Context, workDir, name string, path string, fn func(resp api.ProgressResponse)) error {
mp := ParseModelPath(name)
var manifest *ManifestV2
@@ -524,7 +524,7 @@ func CreateModel(ctx context.Context, name string, path string, fn func(resp api
}
// generate the embedding layers
embeddingLayers, err := embeddingLayers(embed)
embeddingLayers, err := embeddingLayers(workDir, embed)
if err != nil {
return err
}
@@ -581,7 +581,7 @@ type EmbeddingParams struct {
}
// embeddingLayers loads the associated LLM and generates the embeddings to be stored from an input file
func embeddingLayers(e EmbeddingParams) ([]*LayerReader, error) {
func embeddingLayers(workDir string, e EmbeddingParams) ([]*LayerReader, error) {
layers := []*LayerReader{}
if len(e.files) > 0 {
// check if the model is a file path or a model name
@@ -594,7 +594,7 @@ func embeddingLayers(e EmbeddingParams) ([]*LayerReader, error) {
model = &Model{ModelPath: e.model}
}
if err := load(context.Background(), model, e.opts, defaultSessionDuration); err != nil {
if err := load(context.Background(), workDir, model, e.opts, defaultSessionDuration); err != nil {
return nil, fmt.Errorf("load model to generate embeddings: %v", err)
}

View File

@@ -58,7 +58,7 @@ var loaded struct {
var defaultSessionDuration = 5 * time.Minute
// load a model into memory if it is not already loaded, it is up to the caller to lock loaded.mu before calling this function
func load(ctx context.Context, model *Model, reqOpts map[string]interface{}, sessionDuration time.Duration) error {
func load(ctx context.Context, workDir string, model *Model, reqOpts map[string]interface{}, sessionDuration time.Duration) error {
opts := api.DefaultOptions()
if err := opts.FromMap(model.Options); err != nil {
log.Printf("could not load model options: %v", err)
@@ -94,7 +94,7 @@ func load(ctx context.Context, model *Model, reqOpts map[string]interface{}, ses
loaded.Embeddings = model.Embeddings
}
llmModel, err := llm.New(model.ModelPath, model.AdapterPaths, opts)
llmModel, err := llm.New(workDir, model.ModelPath, model.AdapterPaths, opts)
if err != nil {
return err
}
@@ -130,6 +130,7 @@ func load(ctx context.Context, model *Model, reqOpts map[string]interface{}, ses
llmModel.SetOptions(opts)
}
}
loaded.expireAt = time.Now().Add(sessionDuration)
if loaded.expireTimer == nil {
@@ -150,6 +151,7 @@ func load(ctx context.Context, model *Model, reqOpts map[string]interface{}, ses
loaded.digest = ""
})
}
loaded.expireTimer.Reset(sessionDuration)
return nil
}
@@ -172,8 +174,11 @@ func GenerateHandler(c *gin.Context) {
return
}
sessionDuration := defaultSessionDuration // TODO: set this duration from the request if specified
if err := load(c.Request.Context(), model, req.Options, sessionDuration); err != nil {
workDir := c.GetString("workDir")
// TODO: set this duration from the request if specified
sessionDuration := defaultSessionDuration
if err := load(c.Request.Context(), workDir, model, req.Options, sessionDuration); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
@@ -245,7 +250,9 @@ func EmbeddingHandler(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if err := load(c.Request.Context(), model, req.Options, 5*time.Minute); err != nil {
workDir := c.GetString("workDir")
if err := load(c.Request.Context(), workDir, model, req.Options, 5*time.Minute); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
@@ -335,6 +342,8 @@ func CreateModelHandler(c *gin.Context) {
return
}
workDir := c.GetString("workDir")
ch := make(chan any)
go func() {
defer close(ch)
@@ -345,7 +354,7 @@ func CreateModelHandler(c *gin.Context) {
ctx, cancel := context.WithCancel(c.Request.Context())
defer cancel()
if err := CreateModel(ctx, req.Name, req.Path, fn); err != nil {
if err := CreateModel(ctx, workDir, req.Name, req.Path, fn); err != nil {
ch <- gin.H{"error": err.Error()}
}
}()
@@ -519,8 +528,20 @@ func Serve(ln net.Listener, allowOrigins []string) error {
)
}
workDir, err := os.MkdirTemp("", "ollama")
if err != nil {
return err
}
defer os.RemoveAll(workDir)
r := gin.Default()
r.Use(cors.New(config))
r.Use(
cors.New(config),
func(c *gin.Context) {
c.Set("workDir", workDir)
c.Next()
},
)
r.GET("/", func(c *gin.Context) {
c.String(http.StatusOK, "Ollama is running")
@@ -546,12 +567,10 @@ func Serve(ln net.Listener, allowOrigins []string) error {
// listen for a ctrl+c and stop any loaded llm
signals := make(chan os.Signal, 1)
signal.Notify(signals, syscall.SIGINT)
signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
go func() {
<-signals
if loaded.llm != nil {
loaded.llm.Close()
}
os.RemoveAll(workDir)
os.Exit(0)
}()