feat: padding the diff_ids in the model config (#121)

Signed-off-by: chlins <[email protected]>

Author: chlins

pkg/backend/build.go

@@ -78,7 +78,7 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri
 
 	layers = append(layers, layerDescs...)
 	// build the image config.
-	configDesc, err := builder.BuildConfig(ctx, hooks.NewHooks(
+	configDesc, err := builder.BuildConfig(ctx, layers, hooks.NewHooks(
 		hooks.WithOnStart(func(name string, size int64, reader io.Reader) io.Reader {
 			return pb.Add(internalpb.NormalizePrompt("Building config"), name, size, reader)
 		}),

pkg/backend/build/builder.go

@@ -33,6 +33,7 @@ import (
 
 	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
 	sha256 "github.com/minio/sha256-simd"
+	godigest "github.com/opencontainers/go-digest"
 	spec "github.com/opencontainers/image-spec/specs-go"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
@@ -58,7 +59,7 @@ type Builder interface {
 	BuildLayer(ctx context.Context, mediaType, workDir, path string, hooks hooks.Hooks) (ocispec.Descriptor, error)
 
 	// BuildConfig builds the config blob of the artifact.
-	BuildConfig(ctx context.Context, hooks hooks.Hooks) (ocispec.Descriptor, error)
+	BuildConfig(ctx context.Context, layers []ocispec.Descriptor, hooks hooks.Hooks) (ocispec.Descriptor, error)
 
 	// BuildManifest builds the manifest blob of the artifact.
 	BuildManifest(ctx context.Context, layers []ocispec.Descriptor, config ocispec.Descriptor, annotations map[string]string, hooks hooks.Hooks) (ocispec.Descriptor, error)
@@ -148,8 +149,8 @@ func (ab *abstractBuilder) BuildLayer(ctx context.Context, mediaType, workDir, p
 	return ab.strategy.OutputLayer(ctx, mediaType, workDir, relPath, info.Size()+tarHeaderSize, reader, hooks)
 }
 
-func (ab *abstractBuilder) BuildConfig(ctx context.Context, hooks hooks.Hooks) (ocispec.Descriptor, error) {
-	config, err := buildModelConfig(ab.modelfile)
+func (ab *abstractBuilder) BuildConfig(ctx context.Context, layers []ocispec.Descriptor, hooks hooks.Hooks) (ocispec.Descriptor, error) {
+	config, err := buildModelConfig(ab.modelfile, layers)
 	if err != nil {
 		return ocispec.Descriptor{}, fmt.Errorf("failed to build model config: %w", err)
 	}
@@ -189,7 +190,7 @@ func (ab *abstractBuilder) BuildManifest(ctx context.Context, layers []ocispec.D
 }
 
 // buildModelConfig builds the model config.
-func buildModelConfig(modelfile modelfile.Modelfile) (*modelspec.Model, error) {
+func buildModelConfig(modelfile modelfile.Modelfile, layers []ocispec.Descriptor) (*modelspec.Model, error) {
 	config := modelspec.ModelConfig{
 		Architecture: modelfile.GetArch(),
 		Format:       modelfile.GetFormat(),
@@ -205,8 +206,14 @@ func buildModelConfig(modelfile modelfile.Modelfile) (*modelspec.Model, error) {
 		Name:      modelfile.GetName(),
 	}
 
+	diffIDs := make([]godigest.Digest, 0, len(layers))
+	for _, layer := range layers {
+		diffIDs = append(diffIDs, layer.Digest)
+	}
+
 	fs := modelspec.ModelFS{
-		Type: "layers",
+		Type:    "layers",
+		DiffIDs: diffIDs,
 	}
 
 	return &modelspec.Model{

pkg/backend/build/builder_test.go

@@ -31,6 +31,7 @@ import (
 	storagemock "github.com/CloudNativeAI/modctl/test/mocks/storage"
 
 	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
+	godigest "github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
@@ -159,7 +160,7 @@ func (s *BuilderTestSuite) TestBuildConfig() {
 		s.mockOutputStrategy.On("OutputConfig", mock.Anything, modelspec.MediaTypeModelConfig, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
 			Return(expectedDesc, nil).Once()
 
-		desc, err := s.builder.BuildConfig(context.Background(), hooks.NewHooks())
+		desc, err := s.builder.BuildConfig(context.Background(), []ocispec.Descriptor{}, hooks.NewHooks())
 		s.NoError(err)
 		s.Equal(expectedDesc, desc)
 
@@ -179,7 +180,7 @@ func (s *BuilderTestSuite) TestBuildConfig() {
 		s.mockOutputStrategy.On("OutputConfig", mock.Anything, modelspec.MediaTypeModelConfig, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
 			Return(ocispec.Descriptor{}, errors.New("output error")).Once()
 
-		_, err := s.builder.BuildConfig(context.Background(), hooks.NewHooks())
+		_, err := s.builder.BuildConfig(context.Background(), []ocispec.Descriptor{}, hooks.NewHooks())
 		s.Error(err)
 		s.True(strings.Contains(err.Error(), "output error"))
 	})
@@ -239,7 +240,10 @@ func (s *BuilderTestSuite) TestBuildModelConfig() {
 	mockModelfile.On("GetFamily").Return("llama")
 	mockModelfile.On("GetName").Return("llama-2")
 
-	model, err := buildModelConfig(mockModelfile)
+	model, err := buildModelConfig(mockModelfile, []ocispec.Descriptor{
+		{Digest: godigest.Digest("sha256:layer-1")},
+		{Digest: godigest.Digest("sha256:layer-2")},
+	})
 	s.NoError(err)
 	s.NotNil(model)
 
@@ -255,6 +259,9 @@ func (s *BuilderTestSuite) TestBuildModelConfig() {
 	s.WithinDuration(time.Now(), *model.Descriptor.CreatedAt, 5*time.Second)
 
 	s.Equal("layers", model.ModelFS.Type)
+	s.Len(model.ModelFS.DiffIDs, 2)
+	s.Equal("sha256:layer-1", model.ModelFS.DiffIDs[0].String())
+	s.Equal("sha256:layer-2", model.ModelFS.DiffIDs[1].String())
 }
 
 func TestBuilderSuite(t *testing.T) {