Skip to content

Commit 0dfe129

Browse files
author
Ramkumar Chinchani
committed
Initial commit for OCI dist spec v1.1.0 agent support
Partially addresses kubeflow/community#682 Signed-off-by: Ramkumar Chinchani <[email protected]>
1 parent 11e0ab2 commit 0dfe129

File tree

6 files changed

+119
-15
lines changed

6 files changed

+119
-15
lines changed

Makefile

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,9 @@ deploy-ci: manifests
126126
deploy-helm: manifests
127127
helm install kserve-crd charts/kserve-crd/ --wait --timeout 180s
128128
helm install kserve charts/kserve-resources/ --wait --timeout 180s
129+
# deploy a OCI dist spec v1.1.0 registry
130+
helm repo add project-zot http://zotregistry.dev/helm-charts
131+
helm install --set service.port=5000 zot project-zot/zot
129132

130133
undeploy:
131134
kubectl delete -k config/default

pkg/agent/storage/oci.go

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
/*
2+
Copyright 2021 The KServe 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 storage
18+
19+
import (
20+
"fmt"
21+
"io"
22+
"net/http"
23+
"net/url"
24+
"path/filepath"
25+
"strings"
26+
)
27+
28+
type OCIProvider struct {
29+
Client *http.Client
30+
}
31+
32+
func (m *OCIProvider) DownloadModel(modelDir string, modelName string, storageUri string) error {
33+
log.Info("Download model ", "modelName", modelName, "storageUri", storageUri, "modelDir", modelDir)
34+
uri, err := url.Parse(storageUri)
35+
if err != nil {
36+
return fmt.Errorf("unable to parse storage uri: %w", err)
37+
}
38+
OCIDownloader := &OCIDownloader{
39+
StorageUri: storageUri,
40+
ModelDir: modelDir,
41+
ModelName: modelName,
42+
Uri: uri,
43+
}
44+
if err := OCIDownloader.Download(*m.Client); err != nil {
45+
return err
46+
}
47+
return nil
48+
}
49+
50+
type OCIDownloader struct {
51+
StorageUri string
52+
ModelDir string
53+
ModelName string
54+
Uri *url.URL
55+
}
56+
57+
func (h *OCIDownloader) Download(client http.Client) error {
58+
// Create request
59+
req, err := http.NewRequest("GET", h.StorageUri, nil)
60+
if err != nil {
61+
return err
62+
}
63+
64+
// Query request
65+
resp, err := client.Do(req)
66+
if err != nil {
67+
return fmt.Errorf("failed to make a request: %w", err)
68+
}
69+
70+
defer func(Body io.ReadCloser) {
71+
closeErr := Body.Close()
72+
if closeErr != nil {
73+
log.Error(closeErr, "failed to close body")
74+
}
75+
}(resp.Body)
76+
if resp.StatusCode != 200 {
77+
return fmt.Errorf("URI: %s returned a %d response code", h.StorageUri, resp.StatusCode)
78+
}
79+
80+
// Write content into file(s)
81+
contentType := resp.Header.Get("Content-type")
82+
fileDirectory := filepath.Join(h.ModelDir, h.ModelName)
83+
84+
if strings.Contains(contentType, "application/zip") {
85+
if err := extractZipFiles(resp.Body, fileDirectory); err != nil {
86+
return err
87+
}
88+
} else if strings.Contains(contentType, "application/x-tar") || strings.Contains(contentType, "application/x-gtar") ||
89+
strings.Contains(contentType, "application/x-gzip") || strings.Contains(contentType, "application/gzip") {
90+
if err := extractTarFiles(resp.Body, fileDirectory); err != nil {
91+
return err
92+
}
93+
} else {
94+
paths := strings.Split(h.Uri.Path, "/")
95+
fileName := paths[len(paths)-1]
96+
fileFullName := filepath.Join(fileDirectory, fileName)
97+
file, err := createNewFile(fileFullName)
98+
if err != nil {
99+
return err
100+
}
101+
if _, err = io.Copy(file, resp.Body); err != nil {
102+
return fmt.Errorf("unable to copy file content: %w", err)
103+
}
104+
}
105+
106+
return nil
107+
}

pkg/agent/storage/provider.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,11 @@ const (
2929
//File Protocol = "file://"
3030
HTTPS Protocol = "https://"
3131
HTTP Protocol = "http://"
32+
// OCI dist spec v1.1.0
33+
OCI Protocol = "oci://"
3234
)
3335

34-
var SupportedProtocols = []Protocol{S3, GCS, HTTPS, HTTP}
36+
var SupportedProtocols = []Protocol{S3, GCS, HTTPS, HTTP, OCI}
3537

3638
func GetAllProtocol() (protocols []string) {
3739
for _, protocol := range SupportedProtocols {

pkg/agent/storage/utils.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,11 @@ func GetProvider(providers map[Protocol]Provider, protocol Protocol) (Provider,
173173
providers[HTTP] = &HTTPSProvider{
174174
Client: httpsClient,
175175
}
176+
case OCI:
177+
httpsClient := &http.Client{}
178+
providers[OCI] = &OCIProvider{
179+
Client: httpsClient,
180+
}
176181
}
177182

178183
return providers[protocol], nil

pkg/apis/serving/v1beta1/openapi_generated.go

Lines changed: 0 additions & 14 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pkg/utils/utils_test.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -520,6 +520,7 @@ func TestIsPrefixSupported(t *testing.T) {
520520
"GCS://",
521521
"HTTP://",
522522
"HTTPS://",
523+
"OCI://",
523524
}
524525
scenarios := map[string]struct {
525526
input string

0 commit comments

Comments
 (0)