diff --git a/pkg/cli/build.go b/pkg/cli/build.go index 45b530e954..ea17bee44f 100644 --- a/pkg/cli/build.go +++ b/pkg/cli/build.go @@ -8,7 +8,10 @@ import ( "github.com/spf13/cobra" "github.com/spf13/pflag" + "github.com/replicate/cog/pkg/coglog" "github.com/replicate/cog/pkg/config" + "github.com/replicate/cog/pkg/docker" + "github.com/replicate/cog/pkg/http" "github.com/replicate/cog/pkg/image" "github.com/replicate/cog/pkg/util/console" ) @@ -57,8 +60,17 @@ func newBuildCommand() *cobra.Command { func buildCommand(cmd *cobra.Command, args []string) error { ctx := cmd.Context() + command := docker.NewDockerCommand() + client, err := http.ProvideHTTPClient(ctx, command) + if err != nil { + return err + } + logClient := coglog.NewClient(client) + logCtx := logClient.StartBuild(buildFast, buildLocalImage) + cfg, projectDir, err := config.GetConfig(projectDirFlag) if err != nil { + logClient.EndBuild(ctx, err, logCtx) return err } if cfg.Build.Fast { @@ -75,14 +87,17 @@ func buildCommand(cmd *cobra.Command, args []string) error { err = config.ValidateModelPythonVersion(cfg) if err != nil { + logClient.EndBuild(ctx, err, logCtx) return err } - if err := image.Build(ctx, cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, nil, buildLocalImage); err != nil { + if err := image.Build(ctx, cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, nil, buildLocalImage, command); err != nil { + logClient.EndBuild(ctx, err, logCtx) return err } console.Infof("\nImage built as %s", imageName) + logClient.EndBuild(ctx, nil, logCtx) return nil } diff --git a/pkg/cli/push.go b/pkg/cli/push.go index c2aba5afea..3b8e07ac0a 100644 --- a/pkg/cli/push.go +++ b/pkg/cli/push.go @@ -9,9 +9,11 @@ import ( "github.com/replicate/go/uuid" + "github.com/replicate/cog/pkg/coglog" "github.com/replicate/cog/pkg/config" "github.com/replicate/cog/pkg/docker" "github.com/replicate/cog/pkg/global" + "github.com/replicate/cog/pkg/http" "github.com/replicate/cog/pkg/image" "github.com/replicate/cog/pkg/util/console" ) @@ -44,8 +46,17 @@ func newPushCommand() *cobra.Command { func push(cmd *cobra.Command, args []string) error { ctx := cmd.Context() + command := docker.NewDockerCommand() + client, err := http.ProvideHTTPClient(ctx, command) + if err != nil { + return err + } + logClient := coglog.NewClient(client) + logCtx := logClient.StartPush(buildFast, buildLocalImage) + cfg, projectDir, err := config.GetConfig(projectDirFlag) if err != nil { + logClient.EndPush(ctx, err, logCtx) return err } if cfg.Build.Fast { @@ -58,17 +69,23 @@ func push(cmd *cobra.Command, args []string) error { } if imageName == "" { - return fmt.Errorf("To push images, you must either set the 'image' option in cog.yaml or pass an image name as an argument. For example, 'cog push r8.im/your-username/hotdog-detector'") + err = fmt.Errorf("To push images, you must either set the 'image' option in cog.yaml or pass an image name as an argument. For example, 'cog push r8.im/your-username/hotdog-detector'") + logClient.EndPush(ctx, err, logCtx) + return err } replicatePrefix := fmt.Sprintf("%s/", global.ReplicateRegistryHost) if strings.HasPrefix(imageName, replicatePrefix) { if err := docker.ManifestInspect(ctx, imageName); err != nil && strings.Contains(err.Error(), `"code":"NAME_UNKNOWN"`) { - return fmt.Errorf("Unable to find Replicate existing model for %s. Go to replicate.com and create a new model before pushing.", imageName) + err = fmt.Errorf("Unable to find Replicate existing model for %s. Go to replicate.com and create a new model before pushing.", imageName) + logClient.EndPush(ctx, err, logCtx) + return err } } else { if buildLocalImage { - return fmt.Errorf("Unable to push a local image model to a non replicate host, please disable the local image flag before pushing to this host.") + err = fmt.Errorf("Unable to push a local image model to a non replicate host, please disable the local image flag before pushing to this host.") + logClient.EndPush(ctx, err, logCtx) + return err } } @@ -83,7 +100,7 @@ func push(cmd *cobra.Command, args []string) error { startBuildTime := time.Now() - if err := image.Build(ctx, cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, annotations, buildLocalImage); err != nil { + if err := image.Build(ctx, cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, annotations, buildLocalImage, command); err != nil { return err } @@ -94,14 +111,13 @@ func push(cmd *cobra.Command, args []string) error { console.Info("Fast push enabled.") } - command := docker.NewDockerCommand() err = docker.Push(ctx, imageName, buildFast, projectDir, command, docker.BuildInfo{ BuildTime: buildDuration, BuildID: buildID.String(), - }) + }, client) if err != nil { if strings.Contains(err.Error(), "404") { - return fmt.Errorf("Unable to find existing Replicate model for %s. "+ + err = fmt.Errorf("Unable to find existing Replicate model for %s. "+ "Go to replicate.com and create a new model before pushing."+ "\n\n"+ "If the model already exists, you may be getting this error "+ @@ -110,8 +126,12 @@ func push(cmd *cobra.Command, args []string) error { "or `sudo cog push` instead of `cog push`, "+ "which causes Docker to use the wrong Docker credentials.", imageName) + logClient.EndPush(ctx, err, logCtx) + return err } - return fmt.Errorf("Failed to push image: %w", err) + err = fmt.Errorf("Failed to push image: %w", err) + logClient.EndPush(ctx, err, logCtx) + return err } console.Infof("Image '%s' pushed", imageName) @@ -119,6 +139,7 @@ func push(cmd *cobra.Command, args []string) error { replicatePage := fmt.Sprintf("https://%s", strings.Replace(imageName, global.ReplicateRegistryHost, global.ReplicateWebsiteHost, 1)) console.Infof("\nRun your model on Replicate:\n %s", replicatePage) } + logClient.EndPush(ctx, nil, logCtx) return nil } diff --git a/pkg/coglog/client.go b/pkg/coglog/client.go new file mode 100644 index 0000000000..1d13593e91 --- /dev/null +++ b/pkg/coglog/client.go @@ -0,0 +1,163 @@ +package coglog + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/url" + "strconv" + "strings" + "time" + + "github.com/replicate/cog/pkg/env" + "github.com/replicate/cog/pkg/util/console" +) + +type Client struct { + client *http.Client +} + +type BuildLogContext struct { + started time.Time + fast bool + localImage bool +} + +type PushLogContext struct { + started time.Time + fast bool + localImage bool +} + +type buildLog struct { + DurationMs float32 `json:"length_ms"` + BuildError *string `json:"error"` + Fast bool `json:"fast"` + LocalImage bool `json:"local_image"` +} + +type pushLog struct { + DurationMs float32 `json:"length_ms"` + BuildError *string `json:"error"` + Fast bool `json:"fast"` + LocalImage bool `json:"local_image"` +} + +func NewClient(client *http.Client) *Client { + return &Client{ + client: client, + } +} + +func (c *Client) StartBuild(fast bool, localImage bool) BuildLogContext { + logContext := BuildLogContext{ + started: time.Now(), + fast: fast, + localImage: localImage, + } + return logContext +} + +func (c *Client) EndBuild(ctx context.Context, err error, logContext BuildLogContext) bool { + var errorStr *string = nil + if err != nil { + errStr := err.Error() + errorStr = &errStr + } + buildLog := buildLog{ + DurationMs: float32(time.Now().Sub(logContext.started).Milliseconds()), + BuildError: errorStr, + Fast: logContext.fast, + LocalImage: logContext.localImage, + } + + jsonData, err := json.Marshal(buildLog) + if err != nil { + console.Warn("Failed to marshal JSON for build log: " + err.Error()) + return false + } + + err = c.postLog(ctx, jsonData) + if err != nil { + console.Warn(err.Error()) + return false + } + + return true +} + +func (c *Client) StartPush(fast bool, localImage bool) PushLogContext { + logContext := PushLogContext{ + started: time.Now(), + fast: fast, + localImage: localImage, + } + return logContext +} + +func (c *Client) EndPush(ctx context.Context, err error, logContext PushLogContext) bool { + var errorStr *string = nil + if err != nil { + errStr := err.Error() + errorStr = &errStr + } + pushLog := pushLog{ + DurationMs: float32(time.Now().Sub(logContext.started).Milliseconds()), + BuildError: errorStr, + Fast: logContext.fast, + LocalImage: logContext.localImage, + } + + jsonData, err := json.Marshal(pushLog) + if err != nil { + console.Warn("Failed to marshal JSON for build log: " + err.Error()) + return false + } + + err = c.postLog(ctx, jsonData) + if err != nil { + console.Warn(err.Error()) + return false + } + + return true +} + +func (c *Client) postLog(ctx context.Context, jsonData []byte) error { + disabled, err := DisableFromEnvironment() + if err != nil { + return err + } + if disabled { + return errors.New("Cog logging disabled") + } + + url := buildURL() + req, err := http.NewRequestWithContext(ctx, http.MethodPut, url.String(), bytes.NewReader(jsonData)) + if err != nil { + return err + } + resp, err := c.client.Do(req) + if err != nil { + return err + } + if resp.StatusCode != http.StatusOK { + return errors.New("Bad response from build log: " + strconv.Itoa(resp.StatusCode)) + } + return nil +} + +func baseURL() url.URL { + return url.URL{ + Scheme: env.SchemeFromEnvironment(), + Host: HostFromEnvironment(), + } +} + +func buildURL() url.URL { + url := baseURL() + url.Path = strings.Join([]string{"", "v1", "build"}, "/") + return url +} diff --git a/pkg/coglog/client_test.go b/pkg/coglog/client_test.go new file mode 100644 index 0000000000..a6438d6501 --- /dev/null +++ b/pkg/coglog/client_test.go @@ -0,0 +1,54 @@ +package coglog + +import ( + "net/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/replicate/cog/pkg/env" +) + +func TestLogBuild(t *testing.T) { + // Setup mock http server + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + url, err := url.Parse(server.URL) + require.NoError(t, err) + t.Setenv(env.SchemeEnvVarName, url.Scheme) + t.Setenv(CoglogHostEnvVarName, url.Host) + + client := NewClient(http.DefaultClient) + logContext := client.StartBuild(false, false) + success := client.EndBuild(t.Context(), nil, logContext) + require.True(t, success) +} + +func TestLogBuildDisabled(t *testing.T) { + t.Setenv(CoglogDisableEnvVarName, "true") + client := NewClient(http.DefaultClient) + logContext := client.StartBuild(false, false) + success := client.EndBuild(t.Context(), nil, logContext) + require.False(t, success) +} + +func TestLogPush(t *testing.T) { + // Setup mock http server + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + url, err := url.Parse(server.URL) + require.NoError(t, err) + t.Setenv(env.SchemeEnvVarName, url.Scheme) + t.Setenv(CoglogHostEnvVarName, url.Host) + + client := NewClient(http.DefaultClient) + logContext := client.StartPush(false, false) + success := client.EndPush(t.Context(), nil, logContext) + require.True(t, success) +} diff --git a/pkg/coglog/env.go b/pkg/coglog/env.go new file mode 100644 index 0000000000..b12191fe3e --- /dev/null +++ b/pkg/coglog/env.go @@ -0,0 +1,25 @@ +package coglog + +import ( + "os" + "strconv" +) + +const CoglogHostEnvVarName = "R8_COGLOG_HOST" +const CoglogDisableEnvVarName = "R8_COGLOG_DISABLE" + +func HostFromEnvironment() string { + host := os.Getenv(CoglogHostEnvVarName) + if host == "" { + host = "coglog.replicate.delivery" + } + return host +} + +func DisableFromEnvironment() (bool, error) { + disable := os.Getenv(CoglogDisableEnvVarName) + if disable == "" { + disable = "false" + } + return strconv.ParseBool(disable) +} diff --git a/pkg/coglog/env_test.go b/pkg/coglog/env_test.go new file mode 100644 index 0000000000..0ba69e1127 --- /dev/null +++ b/pkg/coglog/env_test.go @@ -0,0 +1,20 @@ +package coglog + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestHostFromEnvironment(t *testing.T) { + const testHost = "coglog" + t.Setenv(CoglogHostEnvVarName, testHost) + require.Equal(t, HostFromEnvironment(), testHost) +} + +func TestDisabledFromEnvironment(t *testing.T) { + t.Setenv(CoglogDisableEnvVarName, "true") + disabled, err := DisableFromEnvironment() + require.NoError(t, err) + require.True(t, disabled) +} diff --git a/pkg/docker/push.go b/pkg/docker/push.go index 0024a989bc..0cfee0bc86 100644 --- a/pkg/docker/push.go +++ b/pkg/docker/push.go @@ -2,10 +2,10 @@ package docker import ( "context" + "net/http" "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" @@ -16,11 +16,7 @@ type BuildInfo struct { BuildID string } -func Push(ctx context.Context, image string, fast bool, projectDir string, command command.Command, buildInfo BuildInfo) error { - client, err := http.ProvideHTTPClient(ctx, command) - if err != nil { - return err - } +func Push(ctx context.Context, image string, fast bool, projectDir string, command command.Command, buildInfo BuildInfo, client *http.Client) error { webClient := web.NewClient(command, client) if err := webClient.PostPushStart(ctx, buildInfo.BuildID, buildInfo.BuildTime); err != nil { diff --git a/pkg/docker/push_test.go b/pkg/docker/push_test.go index d9dfbbb385..5c1bdca110 100644 --- a/pkg/docker/push_test.go +++ b/pkg/docker/push_test.go @@ -14,6 +14,7 @@ import ( "github.com/replicate/cog/pkg/docker/dockertest" "github.com/replicate/cog/pkg/dockercontext" "github.com/replicate/cog/pkg/env" + cogHttp "github.com/replicate/cog/pkg/http" "github.com/replicate/cog/pkg/monobeam" "github.com/replicate/cog/pkg/web" "github.com/replicate/cog/pkg/weights" @@ -69,9 +70,11 @@ func TestPush(t *testing.T) { // Setup mock docker command command := dockertest.NewMockCommand() + client, err := cogHttp.ProvideHTTPClient(t.Context(), command) + require.NoError(t, err) // Run fast push - err = Push(t.Context(), "r8.im/username/modelname", true, dir, command, BuildInfo{}) + err = Push(t.Context(), "r8.im/username/modelname", true, dir, command, BuildInfo{}, client) require.NoError(t, err) } @@ -138,8 +141,10 @@ func TestPushWithWeight(t *testing.T) { // Setup mock docker command command := dockertest.NewMockCommand() + client, err := cogHttp.ProvideHTTPClient(t.Context(), command) + require.NoError(t, err) // Run fast push - err = Push(t.Context(), "r8.im/username/modelname", true, dir, command, BuildInfo{}) + err = Push(t.Context(), "r8.im/username/modelname", true, dir, command, BuildInfo{}, client) require.NoError(t, err) } diff --git a/pkg/http/user_agent.go b/pkg/http/user_agent.go index b54fb315ee..558aca28bb 100644 --- a/pkg/http/user_agent.go +++ b/pkg/http/user_agent.go @@ -2,10 +2,23 @@ package http import ( "fmt" + "runtime" "github.com/replicate/cog/pkg/global" ) func UserAgent() string { - return fmt.Sprintf("Cog/%s", global.Version) + var platform string + switch runtime.GOOS { + case "linux": + platform = "Linux" + case "windows": + platform = "Windows" + case "darwin": + platform = "macOS" + default: + platform = runtime.GOOS + } + + return fmt.Sprintf("Cog/%s (%s)", global.Version, platform) } diff --git a/pkg/image/build.go b/pkg/image/build.go index cd82386f30..e845957c56 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -36,7 +36,7 @@ var errGit = errors.New("git error") // Build a Cog model from a config // // This is separated out from docker.Build(), so that can be as close as possible to the behavior of 'docker build'. -func Build(ctx context.Context, cfg *config.Config, dir, imageName string, secrets []string, noCache, separateWeights bool, useCudaBaseImage string, progressOutput string, schemaFile string, dockerfileFile string, useCogBaseImage *bool, strip bool, precompile bool, fastFlag bool, annotations map[string]string, localImage bool) error { +func Build(ctx context.Context, cfg *config.Config, dir, imageName string, secrets []string, noCache, separateWeights bool, useCudaBaseImage string, progressOutput string, schemaFile string, dockerfileFile string, useCogBaseImage *bool, strip bool, precompile bool, fastFlag bool, annotations map[string]string, localImage bool, dockerCommand command.Command) error { console.Infof("Building Docker image from environment in cog.yaml as %s...", imageName) if fastFlag { console.Info("Fast build enabled.") @@ -61,8 +61,7 @@ func Build(ctx context.Context, cfg *config.Config, dir, imageName string, secre return fmt.Errorf("Failed to build Docker image: %w", err) } } else { - command := docker.NewDockerCommand() - generator, err := dockerfile.NewGenerator(cfg, dir, fastFlag, command, localImage) + generator, err := dockerfile.NewGenerator(cfg, dir, fastFlag, dockerCommand, localImage) if err != nil { return fmt.Errorf("Error creating Dockerfile generator: %w", err) }