Update jmm models to list all tagged models

pkg/cmd/build/build.go

@@ -117,12 +117,13 @@ func (options *BuildOptions) RunBuild() error {
 	model.Config = jozufile
 
 	modelStorePath := options.storageHome
+	repo := ""
 	tag := ""
 	if options.modelRef != nil {
-		modelStorePath = path.Join(options.storageHome, options.modelRef.Registry, options.modelRef.Repository)
+		repo = path.Join(options.modelRef.Registry, options.modelRef.Repository)
 		tag = options.modelRef.Reference
 	}
-	store := storage.NewLocalStore(modelStorePath)
+	store := storage.NewLocalStore(modelStorePath, repo)
 	manifestDesc, err := store.SaveModel(model, tag)
 	if err != nil {
 		fmt.Println(err)

pkg/cmd/models/cmd.go

@@ -2,7 +2,12 @@ package models
 
 import (
 	"fmt"
+	"io/fs"
 	"jmm/pkg/lib/storage"
+	"os"
+	"path"
+	"path/filepath"
+	"text/tabwriter"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -18,11 +23,13 @@ var (
 )
 
 type ModelsOptions struct {
-	configHome string
+	configHome  string
+	storageHome string
 }
 
 func (opts *ModelsOptions) complete() {
 	opts.configHome = viper.GetString("config")
+	opts.storageHome = path.Join(opts.configHome, "storage")
 }
 
 func (opts *ModelsOptions) validate() error {
@@ -51,14 +58,62 @@ func RunCommand(options *ModelsOptions) func(*cobra.Command, []string) {
 			fmt.Println(err)
 			return
 		}
-		store := storage.NewLocalStore(opts.configHome)
 
-		summary, err := listModels(store)
+		storeDirs, err := findRepos(opts.storageHome)
 		if err != nil {
 			fmt.Println(err)
-			return
 		}
 
-		fmt.Println(summary)
+		var allInfoLines []string
+		for _, storeDir := range storeDirs {
+			store := storage.NewLocalStore(opts.storageHome, storeDir)
+
+			infolines, err := listModels(store)
+			if err != nil {
+				fmt.Println(err)
+				return
+			}
+			allInfoLines = append(allInfoLines, infolines...)
+		}
+
+		printSummary(allInfoLines)
+
+	}
+}
+
+func findRepos(storePath string) ([]string, error) {
+	var indexPaths []string
+	err := filepath.WalkDir(storePath, func(file string, info fs.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+		if info.Name() == "index.json" && !info.IsDir() {
+			dir := filepath.Dir(file)
+			relDir, err := filepath.Rel(storePath, dir)
+			if err != nil {
+				return err
+			}
+			if relDir == "." {
+				relDir = ""
+			}
+			indexPaths = append(indexPaths, relDir)
+		}
+		return nil
+	})
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil, nil
+		}
+		return nil, fmt.Errorf("failed to read local storage: %w", err)
+	}
+	return indexPaths, nil
+}
+
+func printSummary(lines []string) {
+	tw := tabwriter.NewWriter(os.Stdout, 0, 2, 3, ' ', 0)
+	fmt.Fprintln(tw, ModelsTableHeader)
+	for _, line := range lines {
+		fmt.Fprintln(tw, line)
 	}
+	tw.Flush()
 }

pkg/cmd/models/models.go

@@ -4,60 +4,62 @@ Copyright © 2024 Jozu.com
 package models
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
 	"jmm/pkg/artifact"
 	"jmm/pkg/lib/storage"
 	"math"
-	"text/tabwriter"
+	"strings"
 
-	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
 const (
-	ModelsTableHeader = "DIGEST\tMAINTAINER\tMODEL FORMAT\tSIZE"
-	ModelsTableFmt    = "%s\t%s\t%s\t%s\t"
+	ModelsTableHeader = "REPOSITORY\tTAG\tMAINTAINER\tMODEL FORMAT\tSIZE\tDIGEST"
+	ModelsTableFmt    = "%s\t%s\t%s\t%s\t%s\t%s\t"
 )
 
-func listModels(store storage.Store) (string, error) {
+func listModels(store storage.Store) ([]string, error) {
 	index, err := store.ParseIndexJson()
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
-	manifests, err := manifestsFromIndex(index, store)
-	if err != nil {
-		return "", err
-	}
-
-	summary, err := printManifestsSummary(manifests, store)
-	if err != nil {
-		return "", err
-	}
-
-	return summary, nil
-}
-
-func manifestsFromIndex(index *ocispec.Index, store storage.Store) (map[digest.Digest]ocispec.Manifest, error) {
-	manifests := map[digest.Digest]ocispec.Manifest{}
+	var infolines []string
 	for _, manifestDesc := range index.Manifests {
-		manifestBytes, err := store.Fetch(context.Background(), manifestDesc)
+		manifest, err := getManifest(store, manifestDesc)
 		if err != nil {
-			return nil, fmt.Errorf("failed to read manifest %s: %w", manifestDesc.Digest, err)
+			return nil, err
 		}
-		manifest := ocispec.Manifest{}
-		if err := json.Unmarshal(manifestBytes, &manifest); err != nil {
-			return nil, fmt.Errorf("failed to parse manifest %s: %w", manifestDesc.Digest, err)
+		manifestConf, err := readManifestConfig(store, manifest)
+		if err != nil {
+			return nil, err
+		}
+		// TODO: filter list for our manifests only, ignore other artifacts
+		infoline, err := getManifestInfoLine(store.GetRepository(), manifestDesc, manifest, manifestConf)
+		if err != nil {
+			return nil, err
 		}
-		manifests[manifestDesc.Digest] = manifest
+		infolines = append(infolines, infoline)
 	}
-	return manifests, nil
+
+	return infolines, nil
+}
+
+func getManifest(store storage.Store, manifestDesc ocispec.Descriptor) (*ocispec.Manifest, error) {
+	manifestBytes, err := store.Fetch(context.Background(), manifestDesc)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read manifest %s: %w", manifestDesc.Digest, err)
+	}
+	manifest := &ocispec.Manifest{}
+	if err := json.Unmarshal(manifestBytes, &manifest); err != nil {
+		return nil, fmt.Errorf("failed to parse manifest %s: %w", manifestDesc.Digest, err)
+	}
+	return manifest, nil
 }
 
-func readManifestConfig(manifest *ocispec.Manifest, store storage.Store) (*artifact.JozuFile, error) {
+func readManifestConfig(store storage.Store, manifest *ocispec.Manifest) (*artifact.JozuFile, error) {
 	configBytes, err := store.Fetch(context.Background(), manifest.Config)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read config: %w", err)
@@ -69,34 +71,25 @@ func readManifestConfig(manifest *ocispec.Manifest, store storage.Store) (*artif
 	return config, nil
 }
 
-func printManifestsSummary(manifests map[digest.Digest]ocispec.Manifest, store storage.Store) (string, error) {
-	buf := bytes.Buffer{}
-	tw := tabwriter.NewWriter(&buf, 0, 2, 3, ' ', 0)
-	fmt.Fprintln(tw, ModelsTableHeader)
-	for digest, manifest := range manifests {
-		// TODO: filter this list for manifests we're interested in (build needs to set a manifest mediaType/artifactType)
-		line, err := getManifestInfoLine(digest, &manifest, store)
-		if err != nil {
-			return "", err
-		}
-		fmt.Fprintln(tw, line)
+func getManifestInfoLine(repo string, desc ocispec.Descriptor, manifest *ocispec.Manifest, config *artifact.JozuFile) (string, error) {
+	ref := desc.Annotations[ocispec.AnnotationRefName]
+	if ref == "" {
+		ref = "<none>"
 	}
-	tw.Flush()
-	return buf.String(), nil
-}
 
-func getManifestInfoLine(digest digest.Digest, manifest *ocispec.Manifest, store storage.Store) (string, error) {
-	config, err := readManifestConfig(manifest, store)
-	if err != nil {
-		return "", err
+	// Strip localhost from repo if present, since we added it
+	repo = strings.TrimPrefix(repo, "localhost/")
+	if repo == "" {
+		repo = "<none>"
 	}
+
 	var size int64
 	for _, layer := range manifest.Layers {
 		size += layer.Size
 	}
 	sizeStr := formatBytes(size)
 
-	info := fmt.Sprintf(ModelsTableFmt, digest, config.Maintainer, config.ModelFormat, sizeStr)
+	info := fmt.Sprintf(ModelsTableFmt, repo, ref, config.Maintainer, config.ModelFormat, sizeStr, desc.Digest)
 	return info, nil
 }
 

pkg/lib/storage/local.go

@@ -20,12 +20,14 @@ import (
 type LocalStore struct {
 	storage   *oci.Store
 	indexPath string
+	repo      string
 }
 
 // Assert LocalStore implements the Store interface.
 var _ Store = (*LocalStore)(nil)
 
-func NewLocalStore(storeHome string) Store {
+func NewLocalStore(storeRoot, repo string) Store {
+	storeHome := filepath.Join(storeRoot, repo)
 	indexPath := filepath.Join(storeHome, "index.json")
 
 	store, err := oci.New(storeHome)
@@ -36,6 +38,7 @@ func NewLocalStore(storeHome string) Store {
 	return &LocalStore{
 		storage:   store,
 		indexPath: indexPath,
+		repo:      repo,
 	}
 }
 
@@ -44,11 +47,12 @@ func (store *LocalStore) SaveModel(model *artifact.Model, tag string) (*ocispec.
 	if err != nil {
 		return nil, err
 	}
-	for _, layer := range model.Layers {
-		_, err = store.saveContentLayer(&layer)
+	for idx, layer := range model.Layers {
+		layerDesc, err := store.saveContentLayer(&layer)
 		if err != nil {
 			return nil, err
 		}
+		model.Layers[idx].Descriptor = *layerDesc
 	}
 
 	manifestDesc, err := store.saveModelManifest(model, config, tag)
@@ -89,6 +93,10 @@ func (store *LocalStore) ParseIndexJson() (*ocispec.Index, error) {
 	return index, nil
 }
 
+func (store *LocalStore) GetRepository() string {
+	return store.repo
+}
+
 func (store *LocalStore) saveContentLayer(layer *artifact.ModelLayer) (*ocispec.Descriptor, error) {
 	ctx := context.Background()
 

pkg/lib/storage/store.go

@@ -13,6 +13,7 @@ import (
 type Store interface {
 	SaveModel(model *artifact.Model, tag string) (*ocispec.Descriptor, error)
 	TagModel(manifestDesc ocispec.Descriptor, tag string) error
+	GetRepository() string
 	ParseIndexJson() (*ocispec.Index, error)
 	Fetch(context.Context, ocispec.Descriptor) ([]byte, error)
 }

pkg/lib/testing/testing.go

@@ -19,6 +19,8 @@ var TestingNotFoundError = errors.New("artifact not found")
 var TestingInvalidSizeError = errors.New("invalid size")
 
 type TestStore struct {
+	// Repository for this store
+	Repo string
 	// Map of digest to Manifest, to simulate retrieval from e.g. disk
 	Manifests map[digest.Digest]ocispec.Manifest
 	// Map of digest to Config, to simulate retrieval from e.g. disk
@@ -77,3 +79,7 @@ func (*TestStore) TagModel(ocispec.Descriptor, string) error {
 func (*TestStore) SaveModel(*artifact.Model, string) (*ocispec.Descriptor, error) {
 	return nil, fmt.Errorf("save model is not implemented for testing")
 }
+
+func (t *TestStore) GetRepository() string {
+	return t.Repo
+}