diff --git a/core/cli/run.go b/core/cli/run.go index de6621ba9..197f63850 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -42,6 +42,7 @@ type RunCMD struct { BackendsPath string `env:"LOCALAI_BACKENDS_PATH,BACKENDS_PATH" type:"path" default:"${basepath}/backends" help:"Path containing backends used for inferencing" group:"backends"` BackendsSystemPath string `env:"LOCALAI_BACKENDS_SYSTEM_PATH,BACKEND_SYSTEM_PATH" type:"path" default:"/var/lib/local-ai/backends" help:"Path containing system backends used for inferencing" group:"backends"` ModelsPath string `env:"LOCALAI_MODELS_PATH,MODELS_PATH" type:"path" default:"${basepath}/models" help:"Path containing models used for inferencing" group:"storage"` + DownloadStagingPath string `env:"LOCALAI_DOWNLOAD_STAGING_PATH,DOWNLOAD_STAGING_PATH" type:"path" default:"${basepath}/downloading" help:"Path where in-flight downloads are staged before extraction" group:"storage"` ArtifactDownloadConcurrency int `env:"LOCALAI_ARTIFACT_DOWNLOAD_CONCURRENCY" help:"How many files of a model artifact to download at once. 1 (the default) downloads sequentially. Raising it helps artifacts split into many files on a fast link, at the cost of more concurrent load on the models volume" group:"storage" default:"1"` GeneratedContentPath string `env:"LOCALAI_GENERATED_CONTENT_PATH,GENERATED_CONTENT_PATH" type:"path" default:"${generatedcontentpath}" help:"Location for generated content (e.g. images, audio, videos)" group:"storage"` UploadPath string `env:"LOCALAI_UPLOAD_PATH,UPLOAD_PATH" type:"path" default:"${uploadpath}" help:"Path to store uploads from files api" group:"storage"` @@ -280,6 +281,7 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { system.WithBackendSystemPath(r.BackendsSystemPath), system.WithModelPath(r.ModelsPath), system.WithBackendPath(r.BackendsPath), + system.WithStagingPath(r.DownloadStagingPath), system.WithBackendImagesReleaseTag(r.BackendImagesReleaseTag), system.WithBackendImagesBranchTag(r.BackendImagesBranchTag), system.WithBackendDevSuffix(r.BackendDevSuffix), diff --git a/core/gallery/backends.go b/core/gallery/backends.go index 21159eb21..b2a16f997 100644 --- a/core/gallery/backends.go +++ b/core/gallery/backends.go @@ -339,6 +339,7 @@ func InstallBackend(ctx context.Context, systemState *system.SystemState, modelL if optsErr != nil { return fmt.Errorf("backend %q: %w", config.Name, optsErr) } + downloadOpts = append(downloadOpts, downloader.WithStagingDir(systemState.StagingPath)) // PreferDevelopmentBackends installs the development image as the primary URI, // keeping the released image reachable as the first fallback — instead of only diff --git a/core/gallery/upgrade.go b/core/gallery/upgrade.go index 3891c60e9..162ef533f 100644 --- a/core/gallery/upgrade.go +++ b/core/gallery/upgrade.go @@ -283,6 +283,7 @@ func UpgradeBackend(ctx context.Context, systemState *system.SystemState, modelL if err != nil { return fmt.Errorf("upgrade %q: %w", backendName, err) } + downloadOpts = append(downloadOpts, downloader.WithStagingDir(systemState.StagingPath)) backendPath := filepath.Join(systemState.Backend.BackendsPath, backendName) tmpPath := backendPath + ".upgrade-tmp" diff --git a/docs/content/reference/cli-reference.md b/docs/content/reference/cli-reference.md index 73fa07722..31bb79013 100644 --- a/docs/content/reference/cli-reference.md +++ b/docs/content/reference/cli-reference.md @@ -28,6 +28,7 @@ model, see [Model Configuration]({{% relref "advanced/model-configuration" %}}). | Parameter | Default | Description | Environment Variable | |-----------|---------|-------------|----------------------| | `--models-path` | `BASEPATH/models` | Path containing models used for inferencing | `$LOCALAI_MODELS_PATH`, `$MODELS_PATH` | +| `--download-staging-path` | `BASEPATH/downloading` | Path where in-flight downloads are staged before extraction | `$LOCALAI_DOWNLOAD_STAGING_PATH`, `$DOWNLOAD_STAGING_PATH` | | `--data-path` | `BASEPATH/data` | Path for persistent data (collectiondb, agent state, tasks, jobs). Separates mutable data from configuration | `$LOCALAI_DATA_PATH` | | `--generated-content-path` | `TMPDIR/localai-UID/generated/content` | Location for assets generated by backends (e.g. stablediffusion, images, audio, videos). Defaults under the OS temp dir (`$TMPDIR`, falling back to `/tmp`), scoped to the current user's UID so accounts sharing a host never collide. | `$LOCALAI_GENERATED_CONTENT_PATH`, `$GENERATED_CONTENT_PATH` | | `--upload-path` | `TMPDIR/localai-UID/upload` | Path to store uploads from files API. Defaults under the OS temp dir (`$TMPDIR`, falling back to `/tmp`), scoped to the current user's UID. | `$LOCALAI_UPLOAD_PATH`, `$UPLOAD_PATH` | diff --git a/pkg/downloader/uri.go b/pkg/downloader/uri.go index 33010c4a2..64ef160a8 100644 --- a/pkg/downloader/uri.go +++ b/pkg/downloader/uri.go @@ -71,6 +71,7 @@ type downloadOptions struct { verifier ImageVerifier bearerToken string transferProgress TransferProgressSink + stagingDir string } // DownloadOption configures DownloadFileWithContext / DownloadFile. @@ -93,6 +94,10 @@ func WithBearerToken(token string) DownloadOption { return func(o *downloadOptions) { o.bearerToken = token } } +func WithStagingDir(dir string) DownloadOption { + return func(o *downloadOptions) { o.stagingDir = dir } +} + // WithTransferProgress attaches a sink for raw HTTP download byte progress. func WithTransferProgress(sink TransferProgressSink) DownloadOption { return func(o *downloadOptions) { o.transferProgress = sink } @@ -627,7 +632,7 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string return fmt.Errorf("failed to open tarball: %s", err.Error()) } - return oci.ExtractOCIImage(ctx, img, url, filePath, downloadStatus) + return oci.ExtractOCIImage(ctx, img, url, filePath, dopts.stagingDir, downloadStatus) } url = URI(url).OCIReference() @@ -653,7 +658,7 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string xlog.Info("Image signature verified", "ref", pinned) } - return oci.ExtractOCIImage(ctx, img, url, filePath, downloadStatus) + return oci.ExtractOCIImage(ctx, img, url, filePath, dopts.stagingDir, downloadStatus) } // Check for cancellation before starting diff --git a/pkg/oci/extract_internal_test.go b/pkg/oci/extract_internal_test.go index b9b99ac11..afa357740 100644 --- a/pkg/oci/extract_internal_test.go +++ b/pkg/oci/extract_internal_test.go @@ -117,6 +117,32 @@ func buildChainedLinkTar() []byte { return buf.Bytes() } +var _ = Describe("OCI download staging", func() { + It("uses the configured directory without OS temp space and cleans up", func() { + root := GinkgoT().TempDir() + staging := filepath.Join(root, "staging") + destination := filepath.Join(root, "destination") + Expect(os.Mkdir(destination, 0755)).To(Succeed()) + layer := buildLayer(tar.Header{ + Name: "payload", Mode: 0644, + PAXRecords: map[string]string{"content": "downloaded bytes"}, + }) + image, err := mutate.AppendLayers(empty.Image, layer) + Expect(err).NotTo(HaveOccurred()) + GinkgoT().Setenv("TMPDIR", filepath.Join(root, "missing-temp")) + GinkgoT().Setenv("TMP", filepath.Join(root, "missing-temp")) + GinkgoT().Setenv("TEMP", filepath.Join(root, "missing-temp")) + + Expect(ExtractOCIImage(context.Background(), image, "test/image", destination, staging, nil)).To(Succeed()) + content, err := os.ReadFile(filepath.Join(destination, "payload")) + Expect(err).NotTo(HaveOccurred()) + Expect(string(content)).To(Equal("downloaded bytes")) + entries, err := os.ReadDir(staging) + Expect(err).NotTo(HaveOccurred()) + Expect(entries).To(BeEmpty()) + }) +}) + var _ = Describe("Tar extraction fallback for link-less filesystems", func() { It("downloads a layered image once and preserves whiteouts before copying links", func() { base := buildLayer( diff --git a/pkg/oci/image.go b/pkg/oci/image.go index 56f9c9b41..7adeeeb59 100644 --- a/pkg/oci/image.go +++ b/pkg/oci/image.go @@ -28,6 +28,7 @@ import ( "github.com/google/go-containerregistry/pkg/v1/tarball" "github.com/mudler/LocalAI/pkg/credentials" "github.com/mudler/LocalAI/pkg/xio" + "github.com/mudler/xlog" ) // ref: https://github.com/mudler/luet/blob/master/pkg/helpers/docker/docker.go#L117 @@ -323,10 +324,34 @@ func (pw *progressWriter) Write(p []byte) (int, error) { return n, nil } -// ExtractOCIImage will extract a given targetImage into a given targetDestination -func ExtractOCIImage(ctx context.Context, img v1.Image, imageRef string, targetDestination string, downloadStatus func(string, string, string, float64)) error { - // Create a temporary tar file - tmpTarFile, err := os.CreateTemp("", "localai-oci-*.tar") +// ExtractOCIImage will extract a given targetImage into a given +// targetDestination. A non-empty stagingDir holds the in-flight image +// tar and layers (the caller picks a directory whose filesystem can +// hold them); empty falls back to the OS temp directory, which on many +// systems is a RAM-backed tmpfs far smaller than a backend image. +func ExtractOCIImage(ctx context.Context, img v1.Image, imageRef string, targetDestination, stagingDir string, downloadStatus func(string, string, string, float64)) error { + // Every download works in its own directory: the image tar and the + // layer staging below both live there, nothing is shared between + // concurrent downloads, and one removal cleans a download up. + var downloadDir string + var err error + if stagingDir != "" { + if err := os.MkdirAll(stagingDir, 0755); err != nil { + return fmt.Errorf("failed to create staging directory: %v", err) + } + downloadDir, err = os.MkdirTemp(stagingDir, "localai-oci-*") + } else { + downloadDir, err = os.MkdirTemp("", "localai-oci-*") + } + if err != nil { + return fmt.Errorf("failed to create download directory: %v", err) + } + defer func() { + if err := os.RemoveAll(downloadDir); err != nil { + xlog.Warn("Failed to remove OCI download directory", "path", downloadDir, "error", err) + } + }() + tmpTarFile, err := os.CreateTemp(downloadDir, "image-*.tar") if err != nil { return fmt.Errorf("failed to create temporary tar file: %v", err) } @@ -504,7 +529,7 @@ func DownloadOCIImageTar(ctx context.Context, img v1.Image, imageRef string, tar } // Create a temporary directory to store the compressed layers - tmpDir, err := os.MkdirTemp("", "localai-oci-layers-*") + tmpDir, err := os.MkdirTemp(filepath.Dir(tarFilePath), "localai-oci-layers-*") if err != nil { return fmt.Errorf("failed to create temporary directory: %v", err) } diff --git a/pkg/oci/image_test.go b/pkg/oci/image_test.go index 447bc90f6..c94cdfecc 100644 --- a/pkg/oci/image_test.go +++ b/pkg/oci/image_test.go @@ -32,7 +32,7 @@ var _ = Describe("OCI", func() { Expect(err).NotTo(HaveOccurred()) defer os.RemoveAll(dir) - err = ExtractOCIImage(context.TODO(), img, imageName, dir, nil) + err = ExtractOCIImage(context.TODO(), img, imageName, dir, "", nil) Expect(err).NotTo(HaveOccurred()) }) }) diff --git a/pkg/system/state.go b/pkg/system/state.go index e87df528a..e631c5c33 100644 --- a/pkg/system/state.go +++ b/pkg/system/state.go @@ -19,6 +19,9 @@ type SystemState struct { Backend Backend Model Model VRAM uint64 + // StagingPath holds in-flight downloads before extraction; empty + // falls back to the OS temp directory. + StagingPath string systemCapabilities string @@ -53,6 +56,12 @@ func WithBackendSystemPath(path string) SystemStateOptions { } } +func WithStagingPath(path string) SystemStateOptions { + return func(s *SystemState) { + s.StagingPath = path + } +} + func WithModelPath(path string) SystemStateOptions { return func(s *SystemState) { s.Model.ModelsPath = path