python/go grpc interface for async uploads

CONTRIBUTING.md

@@ -18,12 +18,13 @@ But perhaps the most useful thing you can do is **use the tool**. Join the [Disc
 
 ## Project structure
 
-There are two main parts to the codebase:
+There are three main parts to the codebase:
 
-- `go/`: This contains the `replicate` command-line interface. It also provides a shared library that the Python library uses in `go/pkg/shared/`. This is called with subprocess and jsonrpc via stdout/in (it's like CGI RPC!).
+- `go/`: This contains the `replicate` command-line interface. It also provides a shared library that the Python library uses in `go/pkg/shared/`. The shared library runs in a standalone GRPC server in a Python subprocess.
 - `python/`: This is the `replicate` Python library. The Python package also includes the `replicate` Go command-line interface and a Go shared library.
+- `proto/`: This defines the interface between the Go server and the Python client.
 
-The main mechanism that is shared between these two parts is the storage mechanism – reading/saving files on Amazon S3 or Google Cloud Storage. By implementing this in Go, we don't have to add a bazillion dependencies to the Python project. All other abstractions are mostly duplicated across the two languages (repositories, experiments, checkpoints, etc), but this line might move over time.
+The Python library acts as a thin client on top of the Go GRPC server, and all the heavy lifting is done in Go.
 
 The other parts are:
 
@@ -100,6 +101,8 @@ This will build the CLI and the Python package:
 
 The built Python packages are in `python/dist/`. These contain both the CLI and the Python library.
 
+To generate the Protobuf implementations you need to install the required Protobuf tools. This is documented in `proto/Makefile`. Once they're installed, simply run `make build` from the `proto` folder.
+
 ## Release
 
 This will release both the CLI and Python package:

Makefile

@@ -81,3 +81,8 @@ verify-go-version:
 .PHONY: verify-python-version
 verify-python-version:
 	@./makefile-scripts/verify-python-version.sh
+
+.PHONY: fmt
+fmt:
+	cd go && $(MAKE) fmt
+	cd python && $(MAKE) fmt

go/Makefile

@@ -95,8 +95,8 @@ lint:
 	go run github.com/golangci/golangci-lint/cmd/golangci-lint run ./...
 
 .PHONY: fmt
-fmt: install-goimports
-	goimports --local replicate.ai -w -d .
+fmt:
+	go run golang.org/x/tools/cmd/goimports --local replicate.ai -w -d .
 
 .PHONY: mod-tidy
 mod-tidy:

go/cmd/replicate-shared/main.go

@@ -1,9 +1,13 @@
 package main
 
 import (
-	"github.com/replicate/replicate/go/pkg/shared"
+	"github.com/replicate/replicate/go/pkg/cli"
+	"github.com/replicate/replicate/go/pkg/console"
 )
 
 func main() {
-	shared.Serve()
+	cmd := cli.NewDaemonCommand()
+	if err := cmd.Execute(); err != nil {
+		console.Fatal("%s", err)
+	}
 }

go/go.mod

@@ -9,7 +9,8 @@ require (
 	github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 // indirect
 	github.com/ghodss/yaml v1.0.0
 	github.com/go-bindata/go-bindata v3.1.2+incompatible
-	github.com/golangci/golangci-lint v1.34.1
+	github.com/golang/protobuf v1.4.3
+	github.com/golangci/golangci-lint v1.32.2
 	github.com/hashicorp/go-uuid v1.0.2
 	github.com/kami-zh/go-capturer v0.0.0-20171211120116-e492ea43421d
 	github.com/logrusorgru/aurora v2.0.3+incompatible
@@ -29,7 +30,11 @@ require (
 	github.com/xtgo/uuid v0.0.0-20140804021211-a0b114877d4c // indirect
 	golang.org/x/oauth2 v0.0.0-20201109201403-9fd604954f58
 	golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9
-	golang.org/x/tools v0.0.0-20201201161351-ac6f37ff4c2a
-	google.golang.org/api v0.36.0
-	gotest.tools/gotestsum v0.6.0
+	golang.org/x/sys v0.0.0-20201119102817-f84b799fce68 // indirect
+	golang.org/x/tools v0.0.0-20201121010211-780cb80bd7fb
+	google.golang.org/api v0.32.0
+	google.golang.org/genproto v0.0.0-20200921151605-7abf4a1a14d5
+	google.golang.org/grpc v1.33.2
+	google.golang.org/protobuf v1.25.0
+	gotest.tools/gotestsum v0.5.2
 )

go/pkg/cli/checkout.go

@@ -3,7 +3,6 @@ package cli
 import (
 	"fmt"
 	"os"
-	"path"
 	"path/filepath"
 
 	"github.com/logrusorgru/aurora"
@@ -12,7 +11,6 @@ import (
 	"github.com/replicate/replicate/go/pkg/console"
 	"github.com/replicate/replicate/go/pkg/files"
 	"github.com/replicate/replicate/go/pkg/project"
-	"github.com/replicate/replicate/go/pkg/repository"
 )
 
 type checkoutOpts struct {
@@ -42,24 +40,8 @@ func newCheckoutCommand() *cobra.Command {
 	return cmd
 }
 
-// Returns the repository requested by opts.repositoryURL
-func getRepositoryFromOpts(opts checkoutOpts) (repository.Repository, error) {
-	repositoryURL, projectDir, err := getRepositoryURLFromStringOrConfig(opts.repositoryURL)
-	if err != nil {
-		return nil, err
-	}
-
-	repo, err := getRepository(repositoryURL, projectDir)
-	if err != nil {
-		return nil, err
-	}
-
-	return repo, nil
-}
-
 // Returns the experiment and the most appropriate checkpoint for that experiment.
-func getExperimentAndCheckpoint(prefix string, repo repository.Repository) (*project.Experiment, *project.Checkpoint, error) {
-	proj := project.NewProject(repo)
+func getExperimentAndCheckpoint(prefix string, proj *project.Project, projectDir string) (*project.Experiment, *project.Checkpoint, error) {
 	result, err := proj.CheckpointOrExperimentFromPrefix(prefix)
 	if err != nil {
 		return nil, nil, err
@@ -151,12 +133,17 @@ func overwriteDisplayPathPrompt(displayPath string, force bool) error {
 func checkoutCheckpoint(opts checkoutOpts, args []string) error {
 	prefix := args[0]
 
-	repo, err := getRepositoryFromOpts(opts)
+	repositoryURL, projectDir, err := getRepositoryURLFromStringOrConfig(opts.repositoryURL)
+	if err != nil {
+		return err
+	}
+	repo, err := getRepository(repositoryURL, projectDir)
 	if err != nil {
 		return err
 	}
 
-	experiment, checkpoint, err := getExperimentAndCheckpoint(prefix, repo)
+	proj := project.NewProject(repo, projectDir)
+	experiment, checkpoint, err := getExperimentAndCheckpoint(prefix, proj, projectDir)
 	if err != nil {
 		return err
 	}
@@ -191,112 +178,8 @@ func checkoutCheckpoint(opts checkoutOpts, args []string) error {
 
 	checkoutPath := opts.checkoutPath
 	if checkoutPath == "" {
-		return checkoutEntireCheckpoint(outputDir, repo, experiment, checkpoint)
-	} else {
-		return checkoutFileOrDir(outputDir, checkoutPath, repo, experiment, checkpoint)
-	}
-}
-
-// checkout all the files from an experiment or checkpoint
-func checkoutEntireCheckpoint(outputDir string, repo repository.Repository, experiment *project.Experiment, checkpoint *project.Checkpoint) error {
-	// Extract the tarfile
-	experimentFilesExist := true
-	checkpointFilesExist := true
-
-	if err := repo.GetPathTar(path.Join("experiments", experiment.ID+".tar.gz"), outputDir); err != nil {
-		// Ignore does not exist errors
-		if _, ok := err.(*repository.DoesNotExistError); ok {
-			console.Debug("No experiment data found")
-			experimentFilesExist = false
-		} else {
-			return err
-		}
+		return proj.CheckoutCheckpoint(checkpoint, experiment, outputDir, false)
 	} else {
-		console.Info("Copied the files from experiment %s to %q", experiment.ShortID(), filepath.Join(outputDir, experiment.Path))
-	}
-
-	// Overlay checkpoint on top of experiment
-	if checkpoint != nil {
-
-		if err := repo.GetPathTar(path.Join("checkpoints", checkpoint.ID+".tar.gz"), outputDir); err != nil {
-			if _, ok := err.(*repository.DoesNotExistError); ok {
-				console.Debug("No checkpoint data found")
-				checkpointFilesExist = false
-			} else {
-				return err
-
-			}
-		} else {
-			console.Info("Copied the files from checkpoint %s to %q", checkpoint.ShortID(), filepath.Join(outputDir, checkpoint.Path))
-		}
-
+		return proj.CheckoutFileOrDirectory(checkpoint, experiment, outputDir, checkoutPath)
 	}
-
-	if !experimentFilesExist && !checkpointFilesExist {
-		// Just an experiment, no checkpoints
-		if checkpoint == nil {
-			return fmt.Errorf("The experiment %s does not have any files associated with it. You need to pass the 'path' argument to 'init()' to check out files.", experiment.ShortID())
-		}
-		return fmt.Errorf("Neither the experiment %s nor the checkpoint %s has any files associated with it. You need to pass the 'path' argument to 'init()' or 'checkpoint()' to check out files.", experiment.ShortID(), checkpoint.ShortID())
-	}
-
-	console.Info(`If you want to run this experiment again, this is how it was run:
-
-  ` + experiment.Command + `
-`)
-
-	return nil
-
-}
-
-// checkout all the files from an experiment or checkpoint
-func checkoutFileOrDir(outputDir string, checkoutPath string, repo repository.Repository, experiment *project.Experiment, checkpoint *project.Checkpoint) error {
-	// Extract the tarfile
-	experimentFilesExist := true
-	checkpointFilesExist := true
-
-	if err := repo.GetPathItemTar(path.Join("experiments", experiment.ID+".tar.gz"), checkoutPath, outputDir); err != nil {
-		// Ignore does not exist errors
-		if _, ok := err.(*repository.DoesNotExistError); ok {
-			console.Debug("No experiment data found")
-			experimentFilesExist = false
-		} else {
-			return err
-		}
-	} else {
-		console.Info("Copied the path %s from experiment %s to %q", checkoutPath, experiment.ShortID(), filepath.Join(outputDir, experiment.Path))
-	}
-
-	// Overlay checkpoint on top of experiment
-	if checkpoint != nil {
-
-		if err := repo.GetPathItemTar(path.Join("checkpoints", checkpoint.ID+".tar.gz"), checkoutPath, outputDir); err != nil {
-			if _, ok := err.(*repository.DoesNotExistError); ok {
-				console.Debug("No checkpoint data found")
-				checkpointFilesExist = false
-			} else {
-				return err
-
-			}
-		} else {
-			console.Info("Copied the path %s from checkpoint %s to %q", checkoutPath, checkpoint.ShortID(), filepath.Join(outputDir, checkpoint.Path))
-		}
-
-	}
-
-	if !experimentFilesExist && !checkpointFilesExist {
-		// Just an experiment, no checkpoints
-		if checkpoint == nil {
-			return fmt.Errorf("The experiment %s does not have the path %s associated with it. You need to pass the 'path' argument to 'init()' to check out files.", experiment.ShortID(), checkoutPath)
-		}
-		return fmt.Errorf("Neither the experiment %s nor the checkpoint %s has the path %s associated with it. You need to pass the 'path' argument to 'init()' or 'checkpoint()' to check out files.", experiment.ShortID(), checkpoint.ShortID(), checkoutPath)
-	}
-
-	console.Info(`If you want to run this experiment again, this is how it was run:
-
-  ` + experiment.Command + `
-`)
-
-	return nil
-
 }

go/pkg/cli/common.go

@@ -28,20 +28,31 @@ func addRepositoryURLFlagVar(cmd *cobra.Command, opt *string) {
 }
 
 // getRepositoryURLFromStringOrConfig attempts to get it from passed string from --repository,
-// otherwise finds replicate.yaml recursively
+// otherwise finds replicate.yaml recursively.
+// The project directory is determined by the following logic:
+// * If an explicit directory is passed with -D, that is used
+// * Else, if repository URL isn't manually passed with -R, the directory of replicate.yaml is used
+// * Otherwise, the current working directory is used
+// Returns (repositoryURL, projectDir, error)
 func getRepositoryURLFromStringOrConfig(repositoryURL string) (string, string, error) {
+	projectDir := global.ProjectDirectory
 	if repositoryURL == "" {
-		conf, projectDir, err := config.FindConfigInWorkingDir(global.ProjectDirectory)
+		conf, confProjectDir, err := config.FindConfigInWorkingDir(global.ProjectDirectory)
 		if err != nil {
 			return "", "", err
 		}
-		return conf.Repository, projectDir, nil
+		if repositoryURL == "" {
+			repositoryURL = conf.Repository
+		}
+		if global.ProjectDirectory == "" {
+			projectDir = confProjectDir
+		} else {
+			projectDir = global.ProjectDirectory
+		}
 	}
 
-	// if global.ProjectDirectory == "", abs of that is cwd
-	// FIXME (bfirsh): this does not look up directories for replicate.yaml, so might be the wrong
-	// projectDir. It should probably use return value of FindConfigInWorkingDir.
-	projectDir, err := filepath.Abs(global.ProjectDirectory)
+	// abs of "" if cwd
+	projectDir, err := filepath.Abs(projectDir)
 	if err != nil {
 		return "", "", fmt.Errorf("Failed to determine absolute directory of '%s': %w", global.ProjectDirectory, err)
 	}
@@ -71,14 +82,14 @@ func getProjectDir() (string, error) {
 // getRepository returns the project's repository, with caching if needed
 // This is not in repository package so we can do user interface stuff around syncing
 func getRepository(repositoryURL, projectDir string) (repository.Repository, error) {
-	repo, err := repository.ForURL(repositoryURL)
+	repo, err := repository.ForURL(repositoryURL, projectDir)
 	if err != nil {
 		return nil, err
 	}
 	// projectDir might be "" if you use --repository option
 	if repository.NeedsCaching(repo) && projectDir != "" {
 		console.Info("Fetching new data from %q...", repo.RootURL())
-		repo, err = repository.NewCachedMetadataRepository(repo, projectDir)
+		repo, err = repository.NewCachedMetadataRepository(projectDir, repo)
 		if err != nil {
 			return nil, err
 		}

go/pkg/cli/daemon.go

@@ -0,0 +1,40 @@
+package cli
+
+import (
+	"github.com/spf13/cobra"
+
+	"github.com/replicate/replicate/go/pkg/project"
+	"github.com/replicate/replicate/go/pkg/shared"
+)
+
+func NewDaemonCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:  "replicate-daemon <socket-path>",
+		RunE: runDaemon,
+	}
+	setPersistentFlags(cmd)
+	addRepositoryURLFlag(cmd)
+	return cmd
+}
+
+func runDaemon(cmd *cobra.Command, args []string) error {
+	socketPath := args[0]
+
+	projectGetter := func() (proj *project.Project, err error) {
+		repositoryURL, projectDir, err := getRepositoryURLFromFlagOrConfig(cmd)
+		if err != nil {
+			return nil, err
+		}
+		repo, err := getRepository(repositoryURL, projectDir)
+		if err != nil {
+			return nil, err
+		}
+		proj = project.NewProject(repo, projectDir)
+		return proj, err
+	}
+
+	if err := shared.Serve(projectGetter, socketPath); err != nil {
+		return err
+	}
+	return nil
+}

go/pkg/cli/diff.go

@@ -50,7 +50,7 @@ func diffCheckpoints(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
-	proj := project.NewProject(repo)
+	proj := project.NewProject(repo, projectDir)
 	au := getAurora()
 	return printDiff(os.Stdout, au, proj, prefix1, prefix2)
 }
@@ -145,7 +145,7 @@ func printMapDiff(w *tabwriter.Writer, au aurora.Aurora, map1, map2 map[string]s
 // Returns a map of checkpoint things we want to show in diff
 func checkpointToMap(checkpoint *project.Checkpoint) map[string]string {
 	return map[string]string{
-		"Step":    strconv.Itoa(checkpoint.Step),
+		"Step":    strconv.FormatInt(checkpoint.Step, 10),
 		"Created": checkpoint.Created.In(timezone).Format(time.RFC1123),
 		"Path":    checkpoint.Path,
 	}

go/pkg/cli/diff_test.go

@@ -21,7 +21,7 @@ func TestDiffSameExperiment(t *testing.T) {
 
 	conf := &config.Config{}
 	repo := createShowTestData(t, workingDir, conf)
-	proj := project.NewProject(repo)
+	proj := project.NewProject(repo, workingDir)
 
 	au := aurora.NewAurora(false)
 	out := new(bytes.Buffer)
@@ -59,7 +59,7 @@ func TestDiffDifferentExperiment(t *testing.T) {
 
 	conf := &config.Config{}
 	repo := createShowTestData(t, workingDir, conf)
-	proj := project.NewProject(repo)
+	proj := project.NewProject(repo, workingDir)
 
 	au := aurora.NewAurora(false)
 	out := new(bytes.Buffer)