Skip to content

Commit 6b95da4

Browse files
chlinschlins
authored
feat: add the inspect command (#24)
Signed-off-by: chlins <[email protected]>
1 parent a5bdde0 commit 6b95da4

File tree

5 files changed

+336
-0
lines changed

5 files changed

+336
-0
lines changed

‎cmd/inspect.go‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
/*
2+
* Copyright 2024 The CNAI Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package cmd
18+
19+
import (
20+
"context"
21+
"encoding/json"
22+
"fmt"
23+
24+
"github.com/CloudNativeAI/modctl/pkg/backend"
25+
26+
"github.com/spf13/cobra"
27+
"github.com/spf13/viper"
28+
)
29+
30+
// inspectCmd represents the modctl command for inspect.
31+
var inspectCmd = &cobra.Command{
32+
Use: "inspect [flags] <target>",
33+
Short: "A command line tool for modctl inspect",
34+
Args: cobra.ExactArgs(1),
35+
DisableAutoGenTag: true,
36+
SilenceUsage: true,
37+
FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
38+
RunE: func(cmd *cobra.Command, args []string) error {
39+
return runInspect(context.Background(), args[0])
40+
},
41+
}
42+
43+
// init initializes inspect command.
44+
func init() {
45+
flags := rmCmd.Flags()
46+
47+
if err := viper.BindPFlags(flags); err != nil {
48+
panic(fmt.Errorf("bind cache inspect flags to viper: %w", err))
49+
}
50+
}
51+
52+
// runInspect runs the inspect modctl.
53+
func runInspect(ctx context.Context, target string) error {
54+
b, err := backend.New()
55+
if err != nil {
56+
return err
57+
}
58+
59+
if target == "" {
60+
return fmt.Errorf("target is required")
61+
}
62+
63+
inspected, err := b.Inspect(ctx, target)
64+
if err != nil {
65+
return err
66+
}
67+
68+
data, err := json.MarshalIndent(inspected, "", " ")
69+
if err != nil {
70+
return err
71+
}
72+
73+
fmt.Println(string(data))
74+
return nil
75+
}

‎cmd/root.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,4 +65,5 @@ func init() {
6565
rootCmd.AddCommand(pushCmd)
6666
rootCmd.AddCommand(rmCmd)
6767
rootCmd.AddCommand(pruneCmd)
68+
rootCmd.AddCommand(inspectCmd)
6869
}

‎pkg/backend/backend.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,9 @@ type Backend interface {
4747

4848
// Prune prunes the unused blobs and clean up the storage.
4949
Prune(ctx context.Context) ([]string, error)
50+
51+
// Inspect inspects the model artifact.
52+
Inspect(ctx context.Context, target string) (*InspectedModelArtifact, error)
5053
}
5154

5255
// backend is the implementation of Backend.

‎pkg/backend/inspect.go‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/*
2+
* Copyright 2024 The CNAI Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package backend
18+
19+
import (
20+
"context"
21+
"encoding/json"
22+
"fmt"
23+
24+
modelspec "github.com/CloudNativeAI/modctl/pkg/spec"
25+
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
26+
)
27+
28+
// InspectedModelArtifact is the data structure for model artifact that has been inspected.
29+
type InspectedModelArtifact struct {
30+
// ID is the image id of the model artifact.
31+
ID string `json:"Id"`
32+
// Digest is the digest of the model artifact.
33+
Digest string `json:"Digest"`
34+
// Architecture is the architecture of the model.
35+
Architecture string `json:"Architecture"`
36+
// Created is the creation time of the model artifact.
37+
Created string `json:"Created"`
38+
// Family is the family of the model.
39+
Family string `json:"Family"`
40+
// Format is the format of the model.
41+
Format string `json:"Format"`
42+
// Name is the name of the model.
43+
Name string `json:"Name"`
44+
// ParamSize is the param size of the model.
45+
ParamSize string `json:"ParamSize"`
46+
// Precision is the precision of the model.
47+
Precision string `json:"Precision"`
48+
// Quantization is the quantization of the model.
49+
Quantization string `json:"Quantization"`
50+
// Layers is the layers of the model artifact.
51+
Layers []InspectedModelArtifactLayer `json:"Layers"`
52+
}
53+
54+
// InspectedModelArtifactLayer is the data structure for model artifact layer that has been inspected.
55+
type InspectedModelArtifactLayer struct {
56+
// Digest is the digest of the model artifact layer.
57+
Digest string `json:"Digest"`
58+
// Size is the size of the model artifact layer.
59+
Size int64 `json:"Size"`
60+
// Filepath is the filepath of the model artifact layer.
61+
Filepath string `json:"Filepath"`
62+
}
63+
64+
// Inspect inspects the target from the storage.
65+
func (b *backend) Inspect(ctx context.Context, target string) (*InspectedModelArtifact, error) {
66+
ref, err := ParseReference(target)
67+
if err != nil {
68+
return nil, fmt.Errorf("failed to parse target: %w", err)
69+
}
70+
71+
repo, tag := ref.Repository(), ref.Tag()
72+
manifestRaw, digest, err := b.store.PullManifest(ctx, repo, tag)
73+
if err != nil {
74+
return nil, fmt.Errorf("failed to get manifest: %w", err)
75+
}
76+
77+
var manifest ocispec.Manifest
78+
if err := json.Unmarshal(manifestRaw, &manifest); err != nil {
79+
return nil, fmt.Errorf("failed to unmarshal manifest: %w", err)
80+
}
81+
82+
inspectedModelArtifact := &InspectedModelArtifact{
83+
ID: manifest.Config.Digest.String(),
84+
Digest: digest,
85+
Architecture: manifest.Annotations[modelspec.AnnotationArchitecture],
86+
Created: manifest.Annotations[modelspec.AnnotationCreated],
87+
Family: manifest.Annotations[modelspec.AnnotationFamily],
88+
Format: manifest.Annotations[modelspec.AnnotationFormat],
89+
Name: manifest.Annotations[modelspec.AnnotationName],
90+
ParamSize: manifest.Annotations[modelspec.AnnotationParamSize],
91+
Precision: manifest.Annotations[modelspec.AnnotationPrecision],
92+
Quantization: manifest.Annotations[modelspec.AnnotationQuantization],
93+
}
94+
95+
for _, layer := range manifest.Layers {
96+
inspectedModelArtifact.Layers = append(inspectedModelArtifact.Layers, InspectedModelArtifactLayer{
97+
Digest: layer.Digest.String(),
98+
Size: layer.Size,
99+
Filepath: layer.Annotations[modelspec.AnnotationFilepath],
100+
})
101+
}
102+
103+
return inspectedModelArtifact, nil
104+
}

‎pkg/backend/inspect_test.go‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
/*
2+
* Copyright 2024 The CNAI Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package backend
18+
19+
import (
20+
"context"
21+
"testing"
22+
23+
"github.com/CloudNativeAI/modctl/test/mocks/storage"
24+
"github.com/stretchr/testify/assert"
25+
)
26+
27+
func TestInspect(t *testing.T) {
28+
mockStore := &storage.Storage{}
29+
b := &backend{store: mockStore}
30+
ctx := context.Background()
31+
target := "example.com/repo:tag"
32+
manifest := `{
33+
"schemaVersion": 2,
34+
"mediaType": "application/vnd.oci.image.manifest.v1+json",
35+
"artifactType": "application/vnd.cnai.model.manifest.v1+json",
36+
"config": {
37+
"mediaType": "application/vnd.oci.image.config.v1+json",
38+
"digest": "sha256:144ac462bafbbc7cc6c9e6b325049a0aca1b6ffa2f6cfb0a80ec64bc690bec04",
39+
"size": 46
40+
},
41+
"layers": [
42+
{
43+
"mediaType": "application/vnd.oci.image.layer.v1.tar",
44+
"digest": "sha256:5a96686deb327903f4310e9181ef2ee0bc7261e5181bd23ccdce6c575b6120a2",
45+
"size": 13312,
46+
"annotations": {
47+
"org.cnai.model.filepath": "LICENSE",
48+
"org.cnai.model.license": "true"
49+
},
50+
"artifactType": "application/vnd.cnai.model.layer.v1.tar"
51+
},
52+
{
53+
"mediaType": "application/vnd.oci.image.layer.v1.tar",
54+
"digest": "sha256:44a6e989cc7084ef35aedf1dd7090204ccc928829c51ce79d7d59c346a228333",
55+
"size": 5632,
56+
"annotations": {
57+
"org.cnai.model.filepath": "README.md",
58+
"org.cnai.model.readme": "true"
59+
},
60+
"artifactType": "application/vnd.cnai.model.layer.v1.tar"
61+
},
62+
{
63+
"mediaType": "application/vnd.oci.image.layer.v1.tar",
64+
"digest": "sha256:a4e7c313c8addcc5f8ac3d87d48a9af7eb89bf8819c869c9fa0cad1026397b0c",
65+
"size": 2560,
66+
"annotations": {
67+
"org.cnai.model.config": "true",
68+
"org.cnai.model.filepath": "config.json"
69+
},
70+
"artifactType": "application/vnd.cnai.model.layer.v1.tar"
71+
},
72+
{
73+
"mediaType": "application/vnd.oci.image.layer.v1.tar",
74+
"digest": "sha256:567f11b7338855adbaf58c8e195455860400ef148fc7f02ebc446efdb8b7c515",
75+
"size": 1536,
76+
"annotations": {
77+
"org.cnai.model.config": "true",
78+
"org.cnai.model.filepath": "foo/bar.json"
79+
},
80+
"artifactType": "application/vnd.cnai.model.layer.v1.tar"
81+
},
82+
{
83+
"mediaType": "application/vnd.oci.image.layer.v1.tar",
84+
"digest": "sha256:628ce381719b65598622e3f71844192f84e135d937c7b5a8116582edbe3b1f5d",
85+
"size": 2048,
86+
"annotations": {
87+
"org.cnai.model.config": "true",
88+
"org.cnai.model.filepath": "generation_config.json"
89+
},
90+
"artifactType": "application/vnd.cnai.model.layer.v1.tar"
91+
},
92+
{
93+
"mediaType": "application/vnd.oci.image.layer.v1.tar",
94+
"digest": "sha256:0480097912f4dd530382c69f00d41409bc51f62ea146a04d70c0254791f4ac32",
95+
"size": 7033344,
96+
"annotations": {
97+
"org.cnai.model.config": "true",
98+
"org.cnai.model.filepath": "tokenizer.json"
99+
},
100+
"artifactType": "application/vnd.cnai.model.layer.v1.tar"
101+
},
102+
{
103+
"mediaType": "application/vnd.oci.image.layer.v1.tar",
104+
"digest": "sha256:ebea935e6c2de57780addfc0262c30c2f83afb1457a124fd9b22370e6cb5bc34",
105+
"size": 9216,
106+
"annotations": {
107+
"org.cnai.model.config": "true",
108+
"org.cnai.model.filepath": "tokenizer_config.json"
109+
},
110+
"artifactType": "application/vnd.cnai.model.layer.v1.tar"
111+
},
112+
{
113+
"mediaType": "application/vnd.oci.image.layer.v1.tar",
114+
"digest": "sha256:3a2844a891e19d1d183ac12918a497116309ba9abe0523cdcf1874cf8aebe8e0",
115+
"size": 2778624,
116+
"annotations": {
117+
"org.cnai.model.config": "true",
118+
"org.cnai.model.filepath": "vocab.json"
119+
},
120+
"artifactType": "application/vnd.cnai.model.layer.v1.tar"
121+
}
122+
],
123+
"annotations": {
124+
"org.cnai.model.architecture": "transformer",
125+
"org.cnai.model.created": "2024-11-11T21:16:41+08:00",
126+
"org.cnai.model.family": "qwen2",
127+
"org.cnai.model.format": "tensorflow",
128+
"org.cnai.model.name": "Qwen2.5-0.5B",
129+
"org.cnai.model.param.size": "0.49B",
130+
"org.cnai.model.precision": "int8",
131+
"org.cnai.model.quantization": "gptq"
132+
}
133+
}`
134+
135+
mockStore.On("PullManifest", ctx, "example.com/repo", "tag").Return([]byte(manifest), "sha256:2bc8836f5910ec63a01109e20db67c2ad7706cb19bef5a303bc86fa5572ec9a2", nil)
136+
137+
inspected, err := b.Inspect(ctx, target)
138+
assert.NoError(t, err)
139+
assert.Equal(t, "sha256:144ac462bafbbc7cc6c9e6b325049a0aca1b6ffa2f6cfb0a80ec64bc690bec04", inspected.ID)
140+
assert.Equal(t, "sha256:2bc8836f5910ec63a01109e20db67c2ad7706cb19bef5a303bc86fa5572ec9a2", inspected.Digest)
141+
assert.Equal(t, "transformer", inspected.Architecture)
142+
assert.Equal(t, "2024-11-11T21:16:41+08:00", inspected.Created)
143+
assert.Equal(t, "qwen2", inspected.Family)
144+
assert.Equal(t, "tensorflow", inspected.Format)
145+
assert.Equal(t, "Qwen2.5-0.5B", inspected.Name)
146+
assert.Equal(t, "0.49B", inspected.ParamSize)
147+
assert.Equal(t, "int8", inspected.Precision)
148+
assert.Equal(t, "gptq", inspected.Quantization)
149+
assert.Len(t, inspected.Layers, 8)
150+
assert.Equal(t, "sha256:5a96686deb327903f4310e9181ef2ee0bc7261e5181bd23ccdce6c575b6120a2", inspected.Layers[0].Digest)
151+
assert.Equal(t, "LICENSE", inspected.Layers[0].Filepath)
152+
assert.Equal(t, int64(13312), inspected.Layers[0].Size)
153+
}

0 commit comments

Comments
 (0)