feat: support package code to the artifact (#74)

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

Author: chlins

docs/getting-started.md

@@ -54,6 +54,10 @@ CONFIG generation_config.json
 
 # Model weight, support glob path pattern.
 MODEL *.safetensors
+
+# Specify code, support glob path pattern.
+CODE *.py
+
 ```
 
 Then run the following command to build the model artifact:

pkg/backend/build.go

@@ -88,6 +88,10 @@ func (b *backend) getProcessors(modelfile modelfile.Modelfile) []processor.Proce
 		processors = append(processors, processor.NewModelProcessor(b.store, modelspec.MediaTypeModelWeight, models))
 	}
 
+	if codes := modelfile.GetCodes(); len(codes) > 0 {
+		processors = append(processors, processor.NewCodeProcessor(b.store, modelspec.MediaTypeModelCode, codes))
+	}
+
 	return processors
 }
 

pkg/backend/build_test.go

@@ -28,13 +28,15 @@ func TestGetProcessors(t *testing.T) {
 	modelfile := &modelfile.Modelfile{}
 	modelfile.On("GetConfigs").Return([]string{"config1", "config2"})
 	modelfile.On("GetModels").Return([]string{"model1", "model2"})
+	modelfile.On("GetCodes").Return([]string{"1.py", "2.py"})
 
 	b := &backend{}
 	processors := b.getProcessors(modelfile)
 
-	assert.Len(t, processors, 4)
+	assert.Len(t, processors, 5)
 	assert.Equal(t, "readme", processors[0].Name())
 	assert.Equal(t, "license", processors[1].Name())
 	assert.Equal(t, "model_config", processors[2].Name())
 	assert.Equal(t, "model", processors[3].Name())
+	assert.Equal(t, "code", processors[4].Name())
 }

pkg/backend/processor/code.go

@@ -0,0 +1,49 @@
+/*
+ *     Copyright 2025 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package processor
+
+import (
+	"context"
+
+	"github.com/CloudNativeAI/modctl/pkg/storage"
+
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+// NewCodeProcessor creates a new code processor.
+func NewCodeProcessor(store storage.Storage, mediaType string, patterns []string) Processor {
+	return &codeProcessor{
+		base: &base{
+			store:     store,
+			mediaType: mediaType,
+			patterns:  patterns,
+		},
+	}
+}
+
+// codeProcessor is the processor to process the code file.
+type codeProcessor struct {
+	base *base
+}
+
+func (p *codeProcessor) Name() string {
+	return "code"
+}
+
+func (p *codeProcessor) Process(ctx context.Context, workDir, repo string) ([]ocispec.Descriptor, error) {
+	return p.base.Process(ctx, workDir, repo)
+}

pkg/backend/processor/code_test.go

@@ -0,0 +1,69 @@
+/*
+ *     Copyright 2025 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package processor
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/CloudNativeAI/modctl/test/mocks/storage"
+	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/suite"
+)
+
+type codeProcessorSuite struct {
+	suite.Suite
+	mockStore *storage.Storage
+	processor Processor
+	workDir   string
+}
+
+func (s *codeProcessorSuite) SetupTest() {
+	s.mockStore = &storage.Storage{}
+	s.processor = NewCodeProcessor(s.mockStore, modelspec.MediaTypeModelCode, []string{"*.py"})
+	// generate test files for prorcess.
+	s.workDir = s.Suite.T().TempDir()
+	if err := os.WriteFile(filepath.Join(s.workDir, "test.py"), []byte(""), 0644); err != nil {
+		s.Suite.T().Fatal(err)
+	}
+}
+
+func (s *codeProcessorSuite) TestName() {
+	assert.Equal(s.Suite.T(), "code", s.processor.Name())
+}
+
+func (s *codeProcessorSuite) TestProcess() {
+	ctx := context.Background()
+	repo := "test-repo"
+	s.mockStore.On("PushBlob", ctx, repo, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
+
+	desc, err := s.processor.Process(ctx, s.workDir, repo)
+	assert.NoError(s.Suite.T(), err)
+	assert.NotNil(s.Suite.T(), desc)
+	assert.Equal(s.Suite.T(), "sha256:1234567890abcdef", desc[0].Digest.String())
+	assert.Equal(s.Suite.T(), int64(1024), desc[0].Size)
+	assert.Equal(s.Suite.T(), "test.py", desc[0].Annotations[modelspec.AnnotationFilepath])
+}
+
+func TestCodeProcessorSuite(t *testing.T) {
+	suite.Run(t, new(codeProcessorSuite))
+}