This commit is contained in:
Blake Mizerany 2024-04-03 20:52:27 -07:00
parent f7cfe946dc
commit 76a202c04e
5 changed files with 99 additions and 97 deletions

View File

@ -99,7 +99,7 @@ func (s *Server) handlePush(_ http.ResponseWriter, r *http.Request) error {
// commit the manifest to the registry // commit the manifest to the registry
requirements, err = c.Push(r.Context(), params.Name, man, &registry.PushParams{ requirements, err = c.Push(r.Context(), params.Name, man, &registry.PushParams{
Uploaded: uploads, CompleteParts: uploads,
}) })
if err != nil { if err != nil {
return err return err

View File

@ -23,7 +23,7 @@ type PushRequest struct {
// Parts is a list of upload parts that the client upload in the previous // Parts is a list of upload parts that the client upload in the previous
// push. // push.
Uploaded []CompletePart `json:"part_uploads"` CompleteParts []CompletePart `json:"part_uploads"`
} }
type Requirement struct { type Requirement struct {

View File

@ -24,7 +24,7 @@ func (c *Client) oclient() *ollama.Client {
} }
type PushParams struct { type PushParams struct {
Uploaded []apitype.CompletePart CompleteParts []apitype.CompletePart
} }
// Push pushes a manifest to the server. // Push pushes a manifest to the server.
@ -34,7 +34,7 @@ func (c *Client) Push(ctx context.Context, ref string, manifest []byte, p *PushP
v, err := ollama.Do[apitype.PushResponse](ctx, c.oclient(), "POST", "/v1/push", &apitype.PushRequest{ v, err := ollama.Do[apitype.PushResponse](ctx, c.oclient(), "POST", "/v1/push", &apitype.PushRequest{
Ref: ref, Ref: ref,
Manifest: manifest, Manifest: manifest,
Uploaded: p.Uploaded, CompleteParts: p.CompleteParts,
}) })
if err != nil { if err != nil {
return nil, err return nil, err

View File

@ -101,9 +101,9 @@ func (s *Server) handlePush(w http.ResponseWriter, r *http.Request) error {
} }
completePartsByUploadID := make(map[string]completeParts) completePartsByUploadID := make(map[string]completeParts)
for _, pu := range pr.Uploaded { for _, mcp := range pr.CompleteParts {
// parse the URL // parse the URL
u, err := url.Parse(pu.URL) u, err := url.Parse(mcp.URL)
if err != nil { if err != nil {
return err return err
} }
@ -117,8 +117,7 @@ func (s *Server) handlePush(w http.ResponseWriter, r *http.Request) error {
if err != nil { if err != nil {
return oweb.Mistake("invalid", "url", "invalid or missing PartNumber") return oweb.Mistake("invalid", "url", "invalid or missing PartNumber")
} }
etag := pu.ETag if mcp.ETag == "" {
if etag == "" {
return oweb.Mistake("invalid", "etag", "missing") return oweb.Mistake("invalid", "etag", "missing")
} }
cp, ok := completePartsByUploadID[uploadID] cp, ok := completePartsByUploadID[uploadID]
@ -128,7 +127,7 @@ func (s *Server) handlePush(w http.ResponseWriter, r *http.Request) error {
} }
cp.parts = append(cp.parts, minio.CompletePart{ cp.parts = append(cp.parts, minio.CompletePart{
PartNumber: partNumber, PartNumber: partNumber,
ETag: etag, ETag: mcp.ETag,
}) })
completePartsByUploadID[uploadID] = cp completePartsByUploadID[uploadID] = cp
} }

View File

@ -28,11 +28,18 @@ import (
"kr.dev/diff" "kr.dev/diff"
) )
func testPush(t *testing.T, chunkSize int64) { func TestPushBasic(t *testing.T) {
t.Run(fmt.Sprintf("chunkSize=%d", chunkSize), func(t *testing.T) { const MB = 1024 * 1024
mc := startMinio(t, true) mc := startMinio(t, true)
const MB = 1024 * 1024 defer func() {
mcc := &minio.Core{Client: mc}
// fail if there are any incomplete uploads
for x := range mcc.ListIncompleteUploads(context.Background(), "test", "theKey", true) {
t.Errorf("incomplete: %v", x)
}
}()
// Upload two small layers and one large layer that will // Upload two small layers and one large layer that will
// trigger a multipart upload. // trigger a multipart upload.
@ -59,7 +66,7 @@ func testPush(t *testing.T, chunkSize int64) {
} }
if len(requirements) < 3 { if len(requirements) < 3 {
t.Fatalf("expected at least 3 requirements; got %d", len(requirements)) t.Errorf("expected at least 3 requirements; got %d", len(requirements))
t.Logf("requirements: %v", requirements) t.Logf("requirements: %v", requirements)
} }
@ -75,13 +82,13 @@ func testPush(t *testing.T, chunkSize int64) {
} }
requirements, err = c.Push(context.Background(), ref, manifest, &PushParams{ requirements, err = c.Push(context.Background(), ref, manifest, &PushParams{
Uploaded: uploaded, CompleteParts: uploaded,
}) })
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
if len(requirements) != 0 { if len(requirements) != 0 {
t.Fatalf("unexpected requirements: %v", requirements) t.Errorf("unexpected requirements: %v", requirements)
} }
var paths []string var paths []string
@ -138,12 +145,6 @@ func testPush(t *testing.T, chunkSize int64) {
t.Errorf("[%d] %s", i, msg) t.Errorf("[%d] %s", i, msg)
} }
} }
})
}
func TestPush(t *testing.T) {
testPush(t, 0)
testPush(t, 1)
} }
// TestBasicPresignS3MultipartReferenceDoNotDelete tests the basic flow of // TestBasicPresignS3MultipartReferenceDoNotDelete tests the basic flow of
@ -318,9 +319,11 @@ func startMinio(t *testing.T, trace bool) *minio.Client {
if err != nil { if err != nil {
t.Fatalf("startMinio: %v", err) t.Fatalf("startMinio: %v", err)
} }
if mc.IsOnline() { // try list buckets to see if server is up
if _, err := mc.ListBuckets(ctx); err == nil {
break break
} }
t.Logf("startMinio: server is offline; retrying")
} }
if trace { if trace {