-
Notifications
You must be signed in to change notification settings - Fork 590
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Send push timings to the server #2152
Changes from 8 commits
1055e20
2e03200
eebf0af
2a1b6ea
953ac32
4db7ec0
6010117
4b311e5
db772df
711b479
75068ae
936ef84
224a33f
d47593f
12c69ee
fa79bb1
fc16456
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Large diffs are not rendered by default.
Original file line number | Diff line number | Diff line change | ||||||
---|---|---|---|---|---|---|---|---|
|
@@ -2,22 +2,58 @@ package docker | |||||||
|
||||||||
import ( | ||||||||
"context" | ||||||||
"fmt" | ||||||||
"math/rand/v2" | ||||||||
"strings" | ||||||||
"time" | ||||||||
|
||||||||
"github.com/replicate/cog/pkg/docker/command" | ||||||||
"github.com/replicate/cog/pkg/http" | ||||||||
"github.com/replicate/cog/pkg/monobeam" | ||||||||
"github.com/replicate/cog/pkg/util/console" | ||||||||
"github.com/replicate/cog/pkg/web" | ||||||||
) | ||||||||
|
||||||||
func Push(image string, fast bool, projectDir string, command command.Command) error { | ||||||||
func Push(image string, fast bool, projectDir string, command command.Command, buildTime time.Duration) error { | ||||||||
ctx := context.Background() | ||||||||
client, err := http.ProvideHTTPClient(command) | ||||||||
if err != nil { | ||||||||
return err | ||||||||
} | ||||||||
webClient := web.NewClient(command, client) | ||||||||
|
||||||||
// For the timing flow, on error we will just log and continue since | ||||||||
// this is just a loss of push timing information | ||||||||
imageID := "" | ||||||||
if fast { | ||||||||
client, err := http.ProvideHTTPClient(command) | ||||||||
// Generate 256 bit random hash (4x64 bits) to use as an upload ID | ||||||||
for i := 0; i < 4; i++ { | ||||||||
// Ignoring the linter warning about math/rand/v2 not being cryptographically secure | ||||||||
// because this just needs to be a "unique enough" ID for a cache between when the | ||||||||
// push starts and ends, which should only be ~a week max. | ||||||||
imageID = fmt.Sprintf("%s%x", imageID, rand.Int64()) //nolint:gosec | ||||||||
} | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||||||||
} else { | ||||||||
imageMeta, err := command.Inspect(image) | ||||||||
if err != nil { | ||||||||
return err | ||||||||
console.Warnf("Failed to inspect image: %v", err) | ||||||||
} | ||||||||
webClient := web.NewClient(command, client) | ||||||||
parts := strings.Split(imageMeta.ID, ":") | ||||||||
if len(parts) != 2 { | ||||||||
console.Warn("Image ID was not of the form sha:hash") | ||||||||
} else { | ||||||||
imageID = parts[1] | ||||||||
} | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: Ooh! Use strings.Cut instead! |
||||||||
} | ||||||||
|
||||||||
err = webClient.PostBuildStart(ctx, imageID, buildTime) | ||||||||
if err != nil { | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: I don't see
Suggested change
|
||||||||
console.Warnf("Failed to send build timings to server: %v", err) | ||||||||
} | ||||||||
|
||||||||
if fast { | ||||||||
monobeamClient := monobeam.NewClient(client) | ||||||||
return FastPush(context.Background(), image, projectDir, command, webClient, monobeamClient) | ||||||||
return FastPush(ctx, image, projectDir, command, webClient, monobeamClient, imageID) | ||||||||
} | ||||||||
return StandardPush(image, command) | ||||||||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -10,6 +10,9 @@ import ( | |
"os" | ||
"strconv" | ||
"strings" | ||
"time" | ||
|
||
"github.com/replicate/go/types" | ||
|
||
"github.com/replicate/cog/pkg/config" | ||
"github.com/replicate/cog/pkg/docker/command" | ||
|
@@ -56,6 +59,7 @@ type Version struct { | |
OpenAPISchema map[string]any `json:"openapi_schema"` | ||
RuntimeConfig RuntimeConfig `json:"runtime_config"` | ||
Virtual bool `json:"virtual"` | ||
UploadID string `json:"upload_id"` | ||
} | ||
|
||
func NewClient(dockerCommand command.Command, client *http.Client) *Client { | ||
|
@@ -65,12 +69,47 @@ func NewClient(dockerCommand command.Command, client *http.Client) *Client { | |
} | ||
} | ||
|
||
func (c *Client) PostNewVersion(ctx context.Context, image string, weights []File, files []File) error { | ||
func (c *Client) PostBuildStart(ctx context.Context, imageHash string, buildTime time.Duration) error { | ||
jsonBody := map[string]any{ | ||
"image_hash": imageHash, | ||
"build_time": types.Duration(buildTime).String(), | ||
"push_start_time": time.Now().UTC(), | ||
} | ||
|
||
jsonData, err := json.Marshal(jsonBody) | ||
if err != nil { | ||
return util.WrapError(err, "failed to marshal JSON for build start") | ||
} | ||
|
||
url := webBaseURL() | ||
url.Path = strings.Join([]string{"", "api", "models", "build-start"}, "/") | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: Given there's no variable data in this string, I think it should be a |
||
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url.String(), bytes.NewReader(jsonData)) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
resp, err := c.client.Do(req) | ||
if err != nil { | ||
return err | ||
} | ||
defer resp.Body.Close() | ||
|
||
if resp.StatusCode != http.StatusCreated { | ||
return errors.New("Bad response from new version endpoint: " + strconv.Itoa(resp.StatusCode)) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: it's not a huge deal in this context since the process is not long-lived, but |
||
} | ||
|
||
return nil | ||
} | ||
|
||
func (c *Client) PostNewVersion(ctx context.Context, image string, weights []File, files []File, uploadID string) error { | ||
version, err := c.versionFromManifest(image, weights, files) | ||
if err != nil { | ||
return util.WrapError(err, "failed to build new version from manifest") | ||
} | ||
|
||
version.UploadID = uploadID | ||
|
||
jsonData, err := json.Marshal(version) | ||
if err != nil { | ||
return util.WrapError(err, "failed to marshal JSON for new version") | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks like maybe your
asdf
installation isn't happy 🤔 Thetoolchain
should stay in lockstep with what is in.tool-versions
.