fix: Traverse symlinks for disk size calculation of models stored on PVC (#55)

Resolve symlinks (created by PVC puller) in order to calculate the
size of the target model instead of the size of the symlink itself.

Resolves #54

---------

Signed-off-by: Golan Levy <[email protected]>

Author: GolanLevy

model-serving-puller/puller/puller.go

@@ -209,9 +209,10 @@ func (s *Puller) ProcessLoadModelRequest(ctx context.Context, req *mmesh.LoadMod
 	}
 
 	// update the model key to add the disk size
-	if size, err1 := getModelDiskSize(modelFullPath); err1 != nil {
+	if size, err1 := s.getModelDiskSize(modelFullPath); err1 != nil {
 		s.Log.Error(err1, "Model disk size will not be included in the LoadModelRequest due to error", "model_key", modelKey)
 	} else {
+		s.Log.Info("Calculated disk size", "modelFullPath", modelFullPath, "disk_size", size)
 		modelKey.DiskSizeBytes = size
 	}
 
@@ -230,18 +231,36 @@ func (s *Puller) ProcessLoadModelRequest(ctx context.Context, req *mmesh.LoadMod
 	return req, nil
 }
 
-func getModelDiskSize(modelPath string) (int64, error) {
+func (s *Puller) getModelDiskSize(modelPath string) (int64, error) {
 	// This walks the local filesystem and accumulates the size of the model
 	// It would be more efficient to accumulate the size as the files are downloaded,
 	// but this would require refactoring because the s3 download iterator does not return a size.
 	var size int64
-	err := filepath.Walk(modelPath, func(_ string, info os.FileInfo, err error) error {
+	err := filepath.Walk(modelPath, func(path string, info os.FileInfo, err error) error {
 		if err != nil {
 			return err
 		}
-		if !info.IsDir() {
+
+		// Calculating the size of the resolved path (for pvc) instead of the symlink itself.
+		// We are not expecting to have infinite recursion since otherwise the serving runtime would not be able to load the model
+		if info.Mode()&os.ModeSymlink != 0 {
+			resolvedPath, resolveErr := os.Readlink(path)
+
+			if resolveErr != nil {
+				s.Log.Error(resolveErr, "Failed to resolve symlink path", "path", path)
+			} else {
+				symlinkTargetSize, symlinkTargetSizeErr := s.getModelDiskSize(resolvedPath)
+
+				if symlinkTargetSizeErr != nil {
+					return symlinkTargetSizeErr
+				}
+
+				size += symlinkTargetSize
+			}
+		} else if !info.IsDir() {
 			size += info.Size()
 		}
+
 		return nil
 	})
 	if err != nil {

model-serving-puller/puller/puller_test.go

@@ -607,12 +607,31 @@ func Test_getModelDiskSize(t *testing.T) {
 		{"testModelSize/2", 39375276},
 	}
 
+	p, _ := newPullerWithMock(t)
+
+	// creating symlinks on runtime since git is not managing symlinks by default
+	symlinkName := "symlink"
+	symlinkPath := filepath.Join(RootModelDir, symlinkName)
+
 	for _, tt := range diskSizeTests {
 		t.Run("", func(t *testing.T) {
 			fullPath := filepath.Join(RootModelDir, tt.modelPath)
-			diskSize, err := getModelDiskSize(fullPath)
+			validatePathDiskSize(t, p, fullPath, tt.expectedSize)
+
+			// testing symlink to the same path
+			err := os.Symlink(fullPath, symlinkPath)
+			assert.NoError(t, err)
+
+			validatePathDiskSize(t, p, symlinkPath, tt.expectedSize)
+
+			err = os.Remove(symlinkPath)
 			assert.NoError(t, err)
-			assert.EqualValues(t, tt.expectedSize, diskSize)
 		})
 	}
 }
+
+func validatePathDiskSize(t *testing.T, p *Puller, path string, expectedSize int64) {
+	diskSize, err := p.getModelDiskSize(path)
+	assert.NoError(t, err)
+	assert.EqualValues(t, expectedSize, diskSize)
+}