Allow attaching tags to models; fix errors when pushing existing content to storage

pkg/artifact/model.go

@@ -2,7 +2,6 @@ package artifact
 
 type Model struct {
 	Repository string
-	Tag        string
-	Layers     []*ModelLayer
+	Layers     []ModelLayer
 	Config     *JozuFile
 }

pkg/cmd/build/build.go

@@ -6,16 +6,20 @@ package build
 import (
 	"fmt"
 	"os"
+	"path"
+	"strings"
 
 	"jmm/pkg/artifact"
 	"jmm/pkg/lib/storage"
 
-	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
+	"oras.land/oras-go/v2/registry"
 )
 
-const DEFAULT_MODEL_FILE = "Jozufile"
+const (
+	DEFAULT_MODEL_FILE = "Jozufile"
+)
 
 var (
 	shortDesc = `Build a model`
@@ -28,13 +32,17 @@ var (
 )
 
 type BuildFlags struct {
-	ModelFile string
+	ModelFile  string
+	FullTagRef string
 }
 
 type BuildOptions struct {
-	ModelFile  string
-	ContextDir string
-	JozuHome   string
+	ModelFile   string
+	ContextDir  string
+	configHome  string
+	storageHome string
+	modelRef    *registry.Reference
+	extraRefs   []string
 }
 
 func NewCmdBuild() *cobra.Command {
@@ -76,8 +84,9 @@ func (options *BuildOptions) Complete(cmd *cobra.Command, argsIn []string) error
 	if options.ModelFile == "" {
 		options.ModelFile = options.ContextDir + "/" + DEFAULT_MODEL_FILE
 	}
-	fmt.Println("config: ", viper.GetString("config"))
-	options.JozuHome = viper.GetString("config")
+	options.configHome = viper.GetString("config")
+	fmt.Println("config: ", options.configHome)
+	options.storageHome = path.Join(options.configHome, "storage")
 	return nil
 }
 
@@ -104,17 +113,30 @@ func (options *BuildOptions) RunBuild() error {
 	// 3. Tar the build context and push to local registry
 	layer := artifact.NewLayer(options.ContextDir)
 	model := &artifact.Model{}
-	model.Layers = append(model.Layers, layer)
+	model.Layers = append(model.Layers, *layer)
 	model.Config = jozufile
 
-	store := storage.NewLocalStore(options.JozuHome)
-	var manifest *v1.Manifest
-	manifest, err = store.SaveModel(model)
+	modelStorePath := options.storageHome
+	tag := ""
+	if options.modelRef != nil {
+		modelStorePath = path.Join(options.storageHome, options.modelRef.Registry, options.modelRef.Repository)
+		tag = options.modelRef.Reference
+	}
+	store := storage.NewLocalStore(modelStorePath)
+	manifestDesc, err := store.SaveModel(model, tag)
 	if err != nil {
 		fmt.Println(err)
 		return err
 	}
-	fmt.Println("Model saved: ", manifest.Config.Digest)
+
+	for _, tag := range options.extraRefs {
+		if err := store.TagModel(*manifestDesc, tag); err != nil {
+			return err
+		}
+	}
+
+	fmt.Println("Model saved: ", manifestDesc.Digest)
+
 	return nil
 }
 
@@ -123,13 +145,28 @@ func (o *BuildFlags) ToOptions() (*BuildOptions, error) {
 	if o.ModelFile != "" {
 		options.ModelFile = o.ModelFile
 	}
+	if o.FullTagRef != "" {
+		// References _must_ contain host; use localhost to mark local-only
+		if !strings.Contains(o.FullTagRef, "/") {
+			o.FullTagRef = fmt.Sprintf("localhost/%s", o.FullTagRef)
+		}
+
+		// Handle multiple tags specified with commas, e.g. <registry>/<repo>:tag1,tag2
+		refs := strings.Split(o.FullTagRef, ",")
+		modelRef, err := registry.ParseReference(refs[0])
+		if err != nil {
+			return nil, err
+		}
+		options.modelRef = &modelRef
+		options.extraRefs = refs[1:]
+	}
 	return options, nil
 }
 
 func (flags *BuildFlags) AddFlags(cmd *cobra.Command) {
 	cmd.Flags().StringVarP(&flags.ModelFile, "file", "f", "", "Path to the model file")
+	cmd.Flags().StringVarP(&flags.FullTagRef, "tag", "t", "", "Tag for the model. Example: -t registry/repository:tag1,tag2")
 	cmd.Args = cobra.ExactArgs(1)
-
 }
 
 func NewBuildFlags() *BuildFlags {

pkg/lib/storage/common.go

@@ -0,0 +1,17 @@
+package storage
+
+import (
+	"fmt"
+	"regexp"
+)
+
+var (
+	validTagRegex = regexp.MustCompile(`^[a-zA-Z0-9_][a-zA-Z0-9._-]{0,127}$`)
+)
+
+func validateTag(tag string) error {
+	if !validTagRegex.MatchString(tag) {
+		return fmt.Errorf("invalid tag")
+	}
+	return nil
+}

pkg/lib/storage/local.go

@@ -18,15 +18,14 @@ import (
 )
 
 type LocalStore struct {
-	Storage   *oci.Store
+	storage   *oci.Store
 	indexPath string
 }
 
 // Assert LocalStore implements the Store interface.
 var _ Store = (*LocalStore)(nil)
 
-func NewLocalStore(jozuhome string) *LocalStore {
-	storeHome := filepath.Join(jozuhome, ".jozuStore")
+func NewLocalStore(storeHome string) Store {
 	indexPath := filepath.Join(storeHome, "index.json")
 
 	store, err := oci.New(storeHome)
@@ -35,32 +34,44 @@ func NewLocalStore(jozuhome string) *LocalStore {
 	}
 
 	return &LocalStore{
-		Storage:   store,
+		storage:   store,
 		indexPath: indexPath,
 	}
 }
 
-func (store *LocalStore) SaveModel(model *artifact.Model) (*ocispec.Manifest, error) {
+func (store *LocalStore) SaveModel(model *artifact.Model, tag string) (*ocispec.Descriptor, error) {
 	config, err := store.saveConfigFile(model.Config)
 	if err != nil {
 		return nil, err
 	}
 	for _, layer := range model.Layers {
-		_, err = store.saveContentLayer(layer)
+		_, err = store.saveContentLayer(&layer)
 		if err != nil {
 			return nil, err
 		}
 	}
 
-	manifest, err := store.saveModelManifest(model, config)
+	manifestDesc, err := store.saveModelManifest(model, config, tag)
 	if err != nil {
 		return nil, err
 	}
-	return manifest, nil
+	return manifestDesc, nil
+}
+
+func (store *LocalStore) TagModel(manifestDesc ocispec.Descriptor, tag string) error {
+	if err := validateTag(tag); err != nil {
+		return err
+	}
+
+	if err := store.storage.Tag(context.Background(), manifestDesc, tag); err != nil {
+		return fmt.Errorf("failed to tag manifest: %w", err)
+	}
+
+	return nil
 }
 
 func (store *LocalStore) Fetch(ctx context.Context, desc ocispec.Descriptor) ([]byte, error) {
-	bytes, err := content.FetchAll(ctx, store.Storage, desc)
+	bytes, err := content.FetchAll(ctx, store.storage, desc)
 	return bytes, err
 }
 
@@ -93,39 +104,59 @@ func (store *LocalStore) saveContentLayer(layer *artifact.ModelLayer) (*ocispec.
 		Digest:    digest.FromBytes(buf.Bytes()),
 		Size:      int64(buf.Len()),
 	}
-	err = store.Storage.Push(ctx, desc, buf)
-	layer.Descriptor = desc
+
+	exists, err := store.storage.Exists(ctx, desc)
 	if err != nil {
 		return nil, err
 	}
-	fmt.Println("Saved model layer: ", desc.Digest)
+	if exists {
+		fmt.Println("Model layer already saved: ", desc.Digest)
+	} else {
+		// Does not exist in storage, need to push
+		err = store.storage.Push(ctx, desc, buf)
+		if err != nil {
+			return nil, err
+		}
+		fmt.Println("Saved model layer: ", desc.Digest)
+	}
+
 	return &desc, nil
 }
 
 func (store *LocalStore) saveConfigFile(model *artifact.JozuFile) (*ocispec.Descriptor, error) {
 	ctx := context.Background()
-	buf, err := model.MarshalToJSON()
+	modelBytes, err := model.MarshalToJSON()
 	if err != nil {
 		return nil, err
 	}
 	desc := ocispec.Descriptor{
 		MediaType: constants.ModelConfigMediaType,
-		Digest:    digest.FromBytes(buf),
-		Size:      int64(len(buf)),
+		Digest:    digest.FromBytes(modelBytes),
+		Size:      int64(len(modelBytes)),
 	}
-	err = store.Storage.Push(ctx, desc, bytes.NewReader(buf))
+
+	exists, err := store.storage.Exists(ctx, desc)
 	if err != nil {
 		return nil, err
 	}
+	if !exists {
+		// Does not exist in storage, need to push
+		err = store.storage.Push(ctx, desc, bytes.NewReader(modelBytes))
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	return &desc, nil
 }
 
-func (store *LocalStore) saveModelManifest(model *artifact.Model, config *ocispec.Descriptor) (*ocispec.Manifest, error) {
+func (store *LocalStore) saveModelManifest(model *artifact.Model, config *ocispec.Descriptor, tag string) (*ocispec.Descriptor, error) {
 	ctx := context.Background()
 	manifest := ocispec.Manifest{
-		Versioned: specs.Versioned{SchemaVersion: 2},
-		Config:    *config,
-		Layers:    []ocispec.Descriptor{},
+		Versioned:   specs.Versioned{SchemaVersion: 2},
+		Config:      *config,
+		Layers:      []ocispec.Descriptor{},
+		Annotations: map[string]string{},
 	}
 	// Add the layers to the manifest
 	for _, layer := range model.Layers {
@@ -142,9 +173,25 @@ func (store *LocalStore) saveModelManifest(model *artifact.Model, config *ocispe
 		Digest:    digest.FromBytes(manifestBytes),
 		Size:      int64(len(manifestBytes)),
 	}
-	err = store.Storage.Push(ctx, desc, bytes.NewReader(manifestBytes))
-	if err != nil {
+
+	if exists, err := store.storage.Exists(ctx, desc); err != nil {
 		return nil, err
+	} else if !exists {
+		// Does not exist in storage, need to push
+		err = store.storage.Push(ctx, desc, bytes.NewReader(manifestBytes))
+		if err != nil {
+			return nil, err
+		}
 	}
-	return &manifest, nil
+
+	if tag != "" {
+		if err := validateTag(tag); err != nil {
+			return nil, err
+		}
+		if err := store.storage.Tag(context.Background(), desc, tag); err != nil {
+			return nil, fmt.Errorf("failed to tag manifest: %w", err)
+		}
+	}
+
+	return &desc, nil
 }

pkg/lib/storage/store.go

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

pkg/lib/testing/testing.go

@@ -69,7 +69,11 @@ func (s *TestStore) ParseIndexJson() (*ocispec.Index, error) {
 	return nil, TestingNotFoundError
 }
 
+func (*TestStore) TagModel(ocispec.Descriptor, string) error {
+	return fmt.Errorf("tag model is not implemented for testing")
+}
+
 // SaveModel is not yet implemented!
-func (*TestStore) SaveModel(*artifact.Model) (*ocispec.Manifest, error) {
+func (*TestStore) SaveModel(*artifact.Model, string) (*ocispec.Descriptor, error) {
 	return nil, fmt.Errorf("save model is not implemented for testing")
 }