fixup: check safetensors files by file extension (#13)

fixup: check weight files by file extension

We noticed that during modctl builds, some non-weight files
might be mistakenly identified as weight files (application/vnd.cnai.model.doc.v1.raw),
such as the `tiktoken.model` file in kimi k2. This could
cause these files to be excluded when `excludeWeights = true`,
affecting inference engine startup.

This PR fixes it by checking file extensions to accurately
identify weight files, preventing the above problem.

Currently, this solution applies only to .safetensors files,
if additional weight file formats are introduced in the future,
the detection logic will need to be further extended.

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

Author: imeoer

pkg/service/model.go

@@ -1,15 +1,14 @@
 package service
 
 import (
+	"path/filepath"
 	"sync"
 	"time"
 
-	legacymodelspec "github.com/dragonflyoss/model-spec/specs-go/v1"
 	"github.com/modelpack/modctl/pkg/backend"
 	modctlConfig "github.com/modelpack/modctl/pkg/config"
 	"github.com/modelpack/model-csi-driver/pkg/logger"
 	"github.com/modelpack/model-csi-driver/pkg/utils"
-	modelspec "github.com/modelpack/model-spec/specs-go/v1"
 	"github.com/pkg/errors"
 	"golang.org/x/net/context"
 )
@@ -24,26 +23,17 @@ type ModelArtifact struct {
 	artifact *backend.InspectedModelArtifact
 }
 
-func isSafetensorIndexFile(layer backend.InspectedModelArtifactLayer) bool {
-	return layer.Filepath == safetensorIndexFilePath
-}
-
 func isWeightLayer(layer backend.InspectedModelArtifactLayer) bool {
-	if layer.MediaType == modelspec.MediaTypeModelWeightRaw ||
-		layer.MediaType == modelspec.MediaTypeModelWeight ||
-		layer.MediaType == modelspec.MediaTypeModelWeightGzip ||
-		layer.MediaType == modelspec.MediaTypeModelWeightZstd {
-		return true
-	}
-	if layer.MediaType == legacymodelspec.MediaTypeModelWeightRaw ||
-		layer.MediaType == legacymodelspec.MediaTypeModelWeight ||
-		layer.MediaType == legacymodelspec.MediaTypeModelWeightGzip ||
-		layer.MediaType == legacymodelspec.MediaTypeModelWeightZstd {
+	// For *.safetensors files
+	if filepath.Ext(layer.Filepath) == ".safetensors" {
 		return true
 	}
-	if isSafetensorIndexFile(layer) {
+
+	// For safetensors index file
+	if layer.Filepath == "model.safetensors.index.json" {
 		return true
 	}
+
 	return false
 }
 

pkg/service/model_test.go

@@ -41,7 +41,7 @@ func TestModelArtifact(t *testing.T) {
 						MediaType: modelspec.MediaTypeModelWeightRaw,
 						Digest:    "sha256:layer1",
 						Size:      3 * 1024 * 1024,
-						Filepath:  "zoo.safetensors",
+						Filepath:  "bar.zoo.safetensors",
 					},
 				},
 			}, nil
@@ -60,7 +60,7 @@ func TestModelArtifact(t *testing.T) {
 
 	paths, err := modelArtifact.GetPatterns(ctx, false)
 	require.NoError(t, err)
-	require.Equal(t, []string{"foo.safetensors", "README.md", "zoo.safetensors"}, paths)
+	require.Equal(t, []string{"foo.safetensors", "README.md", "bar.zoo.safetensors"}, paths)
 
 	paths, err = modelArtifact.GetPatterns(ctx, true)
 	require.NoError(t, err)

pkg/service/puller.go

@@ -16,10 +16,6 @@ import (
 	"github.com/pkg/errors"
 )
 
-const (
-	safetensorIndexFilePath = "model.safetensors.index.json"
-)
-
 type PullHook interface {
 	BeforePullLayer(desc ocispec.Descriptor, manifest ocispec.Manifest)
 	AfterPullLayer(desc ocispec.Descriptor, err error)