-
Notifications
You must be signed in to change notification settings - Fork 62
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Rework build process to support both directories and single files
Rework kit build to handle both layers specified as directories and as regular files. In particular, we need to be able to package and extract both model: path: ./model (directory) and model: path: ./model.onnx (file) This requires significant reworking of the build process as previously we assumed we could always 1) create a directory and 2) extract files into it. In addition, this commit changes the build process to stream the build to a temporary file rather than an in-memory buffer to avoid requiring a lot of memory.
- Loading branch information
Showing
6 changed files
with
197 additions
and
114 deletions.
There are no files selected for viewing
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,7 +1,28 @@ | ||
package artifact | ||
|
||
import "kitops/pkg/lib/constants" | ||
|
||
type Model struct { | ||
Repository string | ||
Layers []ModelLayer | ||
Config *KitFile | ||
} | ||
|
||
type ModelLayer struct { | ||
Path string | ||
MediaType string | ||
} | ||
|
||
func (l *ModelLayer) Type() string { | ||
switch l.MediaType { | ||
case constants.CodeLayerMediaType: | ||
return "code" | ||
case constants.DataSetLayerMediaType: | ||
return "dataset" | ||
case constants.ModelConfigMediaType: | ||
return "config" | ||
case constants.ModelLayerMediaType: | ||
return "model" | ||
} | ||
return "<unknown>" | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,146 @@ | ||
package storage | ||
|
||
import ( | ||
"archive/tar" | ||
"compress/gzip" | ||
"fmt" | ||
"io" | ||
"kitops/pkg/artifact" | ||
"kitops/pkg/output" | ||
"os" | ||
"path/filepath" | ||
"strings" | ||
|
||
"github.com/opencontainers/go-digest" | ||
ocispec "github.com/opencontainers/image-spec/specs-go/v1" | ||
) | ||
|
||
func compressLayer(layer *artifact.ModelLayer) (string, ocispec.Descriptor, error) { | ||
pathInfo, err := os.Stat(layer.Path) | ||
if err != nil { | ||
return "", ocispec.DescriptorEmptyJSON, err | ||
} | ||
tempFile, err := os.CreateTemp("", "kitops_layer_*") | ||
if err != nil { | ||
return "", ocispec.DescriptorEmptyJSON, fmt.Errorf("failed to create temporary file: %w", err) | ||
} | ||
tempFileName := tempFile.Name() | ||
output.Debugf("Compressing layer to temporary file %s", tempFileName) | ||
|
||
digester := digest.Canonical.Digester() | ||
mw := io.MultiWriter(tempFile, digester.Hash()) | ||
|
||
// Note: we have to close gzip writer before reading digest from digester as closing is what writes the GZIP footer | ||
gzw := gzip.NewWriter(mw) | ||
tw := tar.NewWriter(gzw) | ||
|
||
// Wrapper function for closing writers before returning an error | ||
handleErr := func(err error) (string, ocispec.Descriptor, error) { | ||
// Don't care about these errors since we'll be deleting the file anyways | ||
_ = tw.Close() | ||
_ = gzw.Close() | ||
_ = tempFile.Close() | ||
removeTempFile(tempFileName) | ||
return "", ocispec.DescriptorEmptyJSON, err | ||
} | ||
|
||
if pathInfo.Mode().IsRegular() { | ||
if err := writeHeaderToTar(pathInfo.Name(), pathInfo, tw); err != nil { | ||
return handleErr(err) | ||
} | ||
if err := writeFileToTar(layer.Path, pathInfo, tw); err != nil { | ||
return handleErr(err) | ||
} | ||
} else if pathInfo.IsDir() { | ||
if err := writeDirToTar(layer.Path, tw); err != nil { | ||
return handleErr(err) | ||
} | ||
} else { | ||
return handleErr(fmt.Errorf("path %s is neither a file nor a directory", layer.Path)) | ||
} | ||
|
||
callAndPrintError(tw.Close, "Failed to close tar writer: %s") | ||
callAndPrintError(gzw.Close, "Failed to close gzip writer: %s") | ||
|
||
tempFileInfo, err := tempFile.Stat() | ||
if err != nil { | ||
removeTempFile(tempFileName) | ||
return "", ocispec.DescriptorEmptyJSON, fmt.Errorf("failed to stat temporary file: %w", err) | ||
} | ||
callAndPrintError(tempFile.Close, "Failed to close temporary file: %s") | ||
|
||
desc := ocispec.Descriptor{ | ||
MediaType: layer.MediaType, | ||
Digest: digester.Digest(), | ||
Size: tempFileInfo.Size(), | ||
} | ||
return tempFileName, desc, nil | ||
} | ||
|
||
func writeDirToTar(basePath string, tw *tar.Writer) error { | ||
// We'll want paths in the tarball to be relative to the *parent* of basePath since we want | ||
// to compress the directory pointed at by basePath | ||
trimPath := filepath.Dir(basePath) | ||
return filepath.Walk(basePath, func(file string, fi os.FileInfo, err error) error { | ||
if err != nil { | ||
return err | ||
} | ||
// Skip anything that's not a regular file or directory | ||
if !fi.Mode().IsRegular() && !fi.Mode().IsDir() { | ||
return nil | ||
} | ||
|
||
relPath := strings.TrimPrefix(strings.Replace(file, trimPath, "", -1), string(filepath.Separator)) | ||
if relPath == "" { | ||
relPath = filepath.Base(basePath) | ||
} | ||
if err := writeHeaderToTar(relPath, fi, tw); err != nil { | ||
return err | ||
} | ||
if fi.IsDir() { | ||
return nil | ||
} | ||
return writeFileToTar(file, fi, tw) | ||
}) | ||
} | ||
|
||
func writeHeaderToTar(name string, fi os.FileInfo, tw *tar.Writer) error { | ||
header, err := tar.FileInfoHeader(fi, "") | ||
if err != nil { | ||
return fmt.Errorf("failed to generate header for %s: %w", name, err) | ||
} | ||
header.Name = name | ||
if err := tw.WriteHeader(header); err != nil { | ||
return fmt.Errorf("failed to write header: %w", err) | ||
} | ||
output.Debugf("Wrote header %s to tar file", header.Name) | ||
return nil | ||
} | ||
|
||
func writeFileToTar(file string, fi os.FileInfo, tw *tar.Writer) error { | ||
f, err := os.Open(file) | ||
if err != nil { | ||
return fmt.Errorf("failed to open file for archiving: %w", err) | ||
} | ||
defer f.Close() | ||
|
||
if written, err := io.Copy(tw, f); err != nil { | ||
return fmt.Errorf("failed to add file to archive: %w", err) | ||
} else if written != fi.Size() { | ||
return fmt.Errorf("error writing file: %w", err) | ||
} | ||
output.Debugf("Wrote file %s to tar file", file) | ||
return nil | ||
} | ||
|
||
func callAndPrintError(f func() error, msg string) { | ||
if err := f(); err != nil { | ||
output.Errorf(msg, err) | ||
} | ||
} | ||
|
||
func removeTempFile(filepath string) { | ||
if err := os.Remove(filepath); err != nil && !os.IsNotExist(err) { | ||
output.Errorf("Failed to clean up temporary file %s: %s", filepath, err) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters