feat(cli): add command to create custom OCI images from directories

Signed-off-by: Ettore Di Giacinto <[email protected]>

Author: mudler

core/cli/util.go

@@ -4,18 +4,24 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
 
+	"github.com/mholt/archiver/v3"
 	"github.com/rs/zerolog/log"
 
 	gguf "github.com/gpustack/gguf-parser-go"
 	cliContext "github.com/mudler/LocalAI/core/cli/context"
 	"github.com/mudler/LocalAI/core/config"
 	"github.com/mudler/LocalAI/core/gallery"
 	"github.com/mudler/LocalAI/pkg/downloader"
+	"github.com/mudler/LocalAI/pkg/oci"
 )
 
 type UtilCMD struct {
 	GGUFInfo         GGUFInfoCMD         `cmd:"" name:"gguf-info" help:"Get information about a GGUF file"`
+	CreateOCIImage   CreateOCIImageCMD   `cmd:"" name:"create-oci-image" help:"Create an OCI image from a file or a directory"`
 	HFScan           HFScanCMD           `cmd:"" name:"hf-scan" help:"Checks installed models for known security issues. WARNING: this is a best-effort feature and may not catch everything!"`
 	UsecaseHeuristic UsecaseHeuristicCMD `cmd:"" name:"usecase-heuristic" help:"Checks a specific model config and prints what usecase LocalAI will offer for it."`
 }
@@ -36,6 +42,35 @@ type UsecaseHeuristicCMD struct {
 	ModelsPath string `env:"LOCALAI_MODELS_PATH,MODELS_PATH" type:"path" default:"${basepath}/models" help:"Path containing models used for inferencing" group:"storage"`
 }
 
+type CreateOCIImageCMD struct {
+	Input     []string `arg:"" help:"Input file or directory to create an OCI image from"`
+	Output    string   `default:"image.tar" help:"Output OCI image name"`
+	ImageName string   `default:"localai" help:"Image name"`
+	Platform  string   `default:"linux/amd64" help:"Platform of the image"`
+}
+
+func (u *CreateOCIImageCMD) Run(ctx *cliContext.Context) error {
+	log.Info().Msg("Creating OCI image from input")
+
+	dir, err := os.MkdirTemp("", "localai")
+	if err != nil {
+		return err
+	}
+	defer os.RemoveAll(dir)
+	err = archiver.Archive(u.Input, filepath.Join(dir, "archive.tar"))
+	if err != nil {
+		return err
+	}
+	log.Info().Msgf("Creating '%s' as '%s' from %v", u.Output, u.Input, u.Input)
+
+	platform := strings.Split(u.Platform, "/")
+	if len(platform) != 2 {
+		return fmt.Errorf("invalid platform: %s", u.Platform)
+	}
+
+	return oci.CreateTar(filepath.Join(dir, "archive.tar"), u.Output, u.ImageName, platform[1], platform[0])
+}
+
 func (u *GGUFInfoCMD) Run(ctx *cliContext.Context) error {
 	if u.Args == nil || len(u.Args) == 0 {
 		return fmt.Errorf("no GGUF file provided")

pkg/oci/tarball.go

@@ -0,0 +1,82 @@
+package oci
+
+import (
+	"io"
+	"os"
+
+	containerdCompression "github.com/containerd/containerd/archive/compression"
+	"github.com/google/go-containerregistry/pkg/name"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/empty"
+	"github.com/google/go-containerregistry/pkg/v1/mutate"
+	"github.com/google/go-containerregistry/pkg/v1/tarball"
+	"github.com/pkg/errors"
+)
+
+func imageFromTar(imagename, architecture, OS string, opener func() (io.ReadCloser, error)) (name.Reference, v1.Image, error) {
+	newRef, err := name.ParseReference(imagename)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	layer, err := tarball.LayerFromOpener(opener)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	baseImage := empty.Image
+	cfg, err := baseImage.ConfigFile()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	cfg.Architecture = architecture
+	cfg.OS = OS
+
+	baseImage, err = mutate.ConfigFile(baseImage, cfg)
+	if err != nil {
+		return nil, nil, err
+	}
+	img, err := mutate.Append(baseImage, mutate.Addendum{
+		Layer: layer,
+		History: v1.History{
+			CreatedBy: "localai",
+			Comment:   "Custom image",
+		},
+	})
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return newRef, img, nil
+}
+
+// CreateTar a imagetarball from a standard tarball
+func CreateTar(srctar, dstimageTar, imagename, architecture, OS string) error {
+
+	dstFile, err := os.Create(dstimageTar)
+	if err != nil {
+		return errors.Wrap(err, "Cannot create "+dstimageTar)
+	}
+	defer dstFile.Close()
+
+	newRef, img, err := imageFromTar(imagename, architecture, OS, func() (io.ReadCloser, error) {
+		f, err := os.Open(srctar)
+		if err != nil {
+			return nil, errors.Wrap(err, "Cannot open "+srctar)
+		}
+		decompressed, err := containerdCompression.DecompressStream(f)
+		if err != nil {
+			return nil, errors.Wrap(err, "Cannot open "+srctar)
+		}
+
+		return decompressed, nil
+	})
+	if err != nil {
+		return err
+	}
+
+	// NOTE: We might also stream that back to the daemon with daemon.Write(tag, img)
+	return tarball.Write(newRef, img, dstFile)
+
+}