Update tar decompress (#200)

* Update stream decompressor

Make the stream decompressor a bit better and cleaner. This removes some
additional code that was no longer needed once we moved to bufio.Reader
for tar peeking

* Emit log  on compressed tars

Compressed tar files for model weights can significantly impact
performance. The log will help identify when that happens so that the
user can opt to use uncompressed tar files where relevant.

pkg/consumer/tar_extractor.go

@@ -1,6 +1,7 @@
 package consumer
 
 import (
+	"bufio"
 	"fmt"
 	"io"
 
@@ -14,7 +15,7 @@ type TarExtractor struct {
 var _ Consumer = &TarExtractor{}
 
 func (f *TarExtractor) Consume(reader io.Reader, destPath string) error {
-	err := extract.TarFile(reader, destPath, f.Overwrite)
+	err := extract.TarFile(bufio.NewReader(reader), destPath, f.Overwrite)
 	if err != nil {
 		return fmt.Errorf("error extracting file: %w", err)
 	}

pkg/extract/compression.go

@@ -30,7 +30,6 @@ var _ decompressor = bzip2Decompressor{}
 var _ decompressor = xzDecompressor{}
 var _ decompressor = lzwDecompressor{}
 var _ decompressor = lz4Decompressor{}
-var _ decompressor = noOpDecompressor{}
 
 // decompressor represents different compression formats.
 type decompressor interface {
@@ -43,20 +42,13 @@ func detectFormat(input []byte) decompressor {
 	inputSize := len(input)
 
 	if inputSize < 2 {
-		return noOpDecompressor{}
+		return nil
 	}
 	// pad to 8 bytes
 	if inputSize < 8 {
 		input = append(input, make([]byte, peekSize-inputSize)...)
 	}
 
-	//  magic16 := binary.BigEndian.Uint16(input)
-	//  magic32 := binary.BigEndian.Uint32(input)
-	//  // We need to pre-pend the padding since we're reading into something bigendian and exceeding the
-	//  // 48bits size of the magic number bytes. The 16 and 32 bit magic numbers are complete bytes and
-	//  // therefore do not need any padding.
-	//  magic48 := binary.BigEndian.Uint64(append(make([]byte, 2), input[0:6]...))
-
 	switch true {
 	case bytes.HasPrefix(input, gzipMagic):
 		log.Debug().
@@ -95,8 +87,9 @@ func detectFormat(input []byte) decompressor {
 		log.Debug().
 			Str("type", "none").
 			Msg("Compression Format")
-		return noOpDecompressor{}
+		return nil
 	}
+
 }
 
 type gzipDecompressor struct{}
@@ -131,9 +124,3 @@ type lz4Decompressor struct{}
 func (d lz4Decompressor) decompress(r io.Reader) (io.Reader, error) {
 	return lz4.NewReader(r), nil
 }
-
-type noOpDecompressor struct{}
-
-func (d noOpDecompressor) decompress(r io.Reader) (io.Reader, error) {
-	return r, nil
-}

pkg/extract/compression_test.go

@@ -31,12 +31,12 @@ func TestDetectFormat(t *testing.T) {
 		{
 			name:       "Less than 2 bytes",
 			input:      []byte{0x1f},
-			expectType: "extract.noOpDecompressor",
+			expectType: "",
 		},
 		{
 			name:       "UNKNOWN",
 			input:      []byte{0xde, 0xad},
-			expectType: "extract.noOpDecompressor",
+			expectType: "",
 		},
 	}
 

pkg/extract/tar.go

@@ -23,19 +23,25 @@ type link struct {
 	newName  string
 }
 
-func TarFile(r io.Reader, destDir string, overwrite bool) error {
+func TarFile(r *bufio.Reader, destDir string, overwrite bool) error {
 	var links []*link
+	var reader io.Reader = r
+
+	log := logging.GetLogger()
 
 	startTime := time.Now()
-	peekableReader := bufio.NewReader(r)
-	peekData, err := peekableReader.Peek(peekSize)
+	peekData, err := r.Peek(peekSize)
 	if err != nil {
 		return fmt.Errorf("error reading peek data: %w", err)
 	}
-	decompressor := detectFormat(peekData)
-	reader, err := decompressor.decompress(peekableReader)
-	if err != nil {
-		return fmt.Errorf("error creating decompressed stream: %w", err)
+	if decompressor := detectFormat(peekData); decompressor != nil {
+		reader, err = decompressor.decompress(reader)
+		if err != nil {
+			return fmt.Errorf("error creating decompressed stream: %w", err)
+		}
+		log.Info().
+			Str("decompressor", fmt.Sprintf("%T", decompressor)).
+			Msg("Tar Compression Detected: Compression can significantly slowdown pget (e.g. for model weights)")
 	}
 	tarReader := tar.NewReader(reader)
 	logger := logging.GetLogger()