Created
August 27, 2019 15:21
-
-
Save mxk/b86769306037c9dc49b44d52764cbbdc to your computer and use it in GitHub Desktop.
TensorFlow Serving gRPC interface generator for Go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
//+build ignore | |
// TensorFlow Serving gRPC interface generator. | |
// | |
// This script works around a bunch of issues (as of 2019-08-25) between Go's | |
// protobuf compiler plugin, Go modules, and definitions of TensorFlow and | |
// TensorFlow Serving proto files. It assumes that protoc and protoc-gen-go are | |
// on your PATH. | |
// | |
// git clone -b r1.15 https://github.com/tensorflow/tensorflow.git | |
// git clone -b r1.14 https://github.com/tensorflow/serving.git | |
// go run protoc.go | |
// go mod edit -replace=github.com/tensorflow/tensorflow/tensorflow/go/core=./proto/tensorflow/core | |
// cd proto/tensorflow/core && go mod init github.com/tensorflow/tensorflow/tensorflow/go/core && cd - | |
// go build ./proto/tensorflow/serving | |
package main | |
import ( | |
"fmt" | |
"io/ioutil" | |
"os" | |
"os/exec" | |
"path/filepath" | |
) | |
const protoDir = "proto/" | |
var opts = []string{"-Iserving", "-Itensorflow"} | |
var cmds = []ProtocCmd{{ | |
PkgDir: protoDir + "tensorflow/core/example", | |
Inputs: []string{"tensorflow/tensorflow/core/example/*.proto"}, | |
}, { | |
PkgDir: protoDir + "tensorflow/core/framework", | |
Inputs: []string{"tensorflow/tensorflow/core/framework/*.proto"}, | |
}, { | |
PkgDir: protoDir + "tensorflow/core/lib/core", | |
Inputs: []string{"tensorflow/tensorflow/core/lib/core/*.proto"}, | |
}, { | |
GoOpts: "import_path=protobuf", | |
PkgDir: protoDir + "tensorflow/core/protobuf", | |
Inputs: []string{ | |
"tensorflow/tensorflow/core/protobuf/*.proto", | |
"tensorflow/tensorflow/stream_executor/*.proto", | |
}, | |
}, { | |
GoOpts: "plugins=grpc,import_path=serving", | |
PkgDir: protoDir + "tensorflow/serving", | |
Inputs: []string{ | |
"serving/tensorflow_serving/apis/*.proto", | |
"serving/tensorflow_serving/config/*.proto", | |
"serving/tensorflow_serving/core/*.proto", | |
"serving/tensorflow_serving/sources/storage_path/*.proto", | |
"serving/tensorflow_serving/util/*.proto", | |
}, | |
}} | |
func main() { | |
for _, cmd := range cmds { | |
fmt.Fprintln(os.Stderr, "==>", cmd.PkgDir) | |
if err := cmd.Run(); err != nil { | |
if e, ok := err.(*exec.ExitError); ok { | |
os.Exit(e.ExitCode()) | |
} | |
fmt.Fprintln(os.Stderr, err) | |
os.Exit(1) | |
} | |
} | |
} | |
// ProtocCmd executes protoc to generate sources for a single Go package. | |
type ProtocCmd struct { | |
GoOpts string // --go_out options | |
PkgDir string // Final output directory | |
Inputs []string // Input files or glob patterns | |
} | |
func (pc *ProtocCmd) Run() error { | |
// Use a temporary protoc output directory | |
root := filepath.Dir(pc.PkgDir) | |
os.MkdirAll(root, 0777) | |
tmp, err := ioutil.TempDir(root, filepath.Base(pc.PkgDir)+".") | |
if err != nil { | |
return err | |
} | |
defer os.RemoveAll(tmp) | |
// Run protoc | |
cmd := exec.Command("protoc", opts...) | |
cmd.Args = append(cmd.Args, "--go_out="+pc.GoOpts+":"+tmp) | |
for _, in := range pc.Inputs { | |
files, err := filepath.Glob(in) | |
if err != nil { | |
return err | |
} | |
cmd.Args = append(cmd.Args, files...) | |
} | |
cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr | |
err = cmd.Run() | |
// Move generated files to PkgDir | |
os.RemoveAll(pc.PkgDir) | |
if err := os.MkdirAll(pc.PkgDir, 0777); err != nil { | |
return err | |
} | |
walkErr := filepath.Walk(tmp, func(path string, fi os.FileInfo, err error) error { | |
if err == nil && fi.Mode().IsRegular() { | |
err = os.Rename(path, filepath.Join(pc.PkgDir, fi.Name())) | |
} | |
return err | |
}) | |
if err == nil { | |
err = walkErr | |
} | |
return err | |
} |
Are you using modules or GOPATH? This step needs to be done from a directory (or subdirectory) with your main go.mod file. If you're using GOPATH, then you'll need to actually copy the contents of ./proto/tensorflow/core
to $GOPATH/src/github.com/tensorflow/tensorflow/tensorflow/go/core
, though I haven't tested this.
Ok!!! I'm using modules but what I did not understand was that the protoc.go file had to belong to my project. Once I run the protoc.go file with the other command lines, everything is working fine!
Thanks a lot for your answer.
does it also work for Tensorflow serving 2.1 and Tensorflow 2.1?
what is go mod edit -replace=github.com/tensorflow/tensorflow/tensorflow/go/core=./proto/tensorflow/core
used for?
what does that?
cd proto/tensorflow/core && go mod init github.com/tensorflow/tensorflow/tensorflow/go/core && cd -
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Hello,
I have an issue to run your process. While doing:
go mod edit -replace=github.com/tensorflow/tensorflow/tensorflow/go/core=./proto/tensorflow/core
I get the following error message:
go: cannot find main module; see 'go help modules'