diff --git a/pkg/artifact/model.go b/pkg/artifact/model.go index cb33e47a..c2faaff6 100644 --- a/pkg/artifact/model.go +++ b/pkg/artifact/model.go @@ -2,7 +2,6 @@ package artifact type Model struct { Repository string - Tag string - Layers []*ModelLayer + Layers []ModelLayer Config *JozuFile } diff --git a/pkg/cmd/build/build.go b/pkg/cmd/build/build.go index 00534010..4c5529d0 100644 --- a/pkg/cmd/build/build.go +++ b/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. /: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 { diff --git a/pkg/lib/storage/common.go b/pkg/lib/storage/common.go new file mode 100644 index 00000000..814161ea --- /dev/null +++ b/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 +} diff --git a/pkg/lib/storage/local.go b/pkg/lib/storage/local.go index 65b644c1..7e11e115 100644 --- a/pkg/lib/storage/local.go +++ b/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 } diff --git a/pkg/lib/storage/store.go b/pkg/lib/storage/store.go index 5d68ff8d..f33655f7 100644 --- a/pkg/lib/storage/store.go +++ b/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) } diff --git a/pkg/lib/testing/testing.go b/pkg/lib/testing/testing.go index 84e3618d..36a2c5e8 100644 --- a/pkg/lib/testing/testing.go +++ b/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") }