tensorflowgoprotocol-bufferstensorflow-servingbuf

Import cycle in tensorflow protobufs


I am trying to write client code to talk to a tensorflow server. I need the golang compiled protobufs for tensorflow and tensorflow_serving. These are not easy to come by, I managed to do so through this. Basically, using buf to generate them. Here is the buf yaml:

version: v1
managed:
  enabled: true
  optimize_for: CODE_SIZE

  # Go
  go_package_prefix:
    default: "some/path"

plugins:
  - plugin: buf.build/protocolbuffers/go
    out: gen/proto/go

This runs successfully, but running the application logs:

 package command-line-arguments
     imports my-package/internal/infer
     imports my-package/internal/infer/tensorflow_serving/apis
     imports my-package/internal/infer/tensorflow/core/protobuf
     imports my-package/internal/infer/tensorflow/compiler/xla/stream_executor
     imports my-package/internal/infer/tensorflow/compiler/xla
     imports my-package/internal/infer/tensorflow/compiler/xla/service
     imports my-package/internal/infer/tensorflow/compiler/xla: import cycle not allowed

Note that everything under tensorflow and tensorflow_serving is directly compiled from the original repositories.

It is surprising to me that something as widely used as tensorflow should have an import cycle, but maybe it does. How can I resolve this?


Solution

  • TL;DR

    The root cause is that the repositry https://github.com/tensorflow/tensorflow does not organize the proto files correctly (or at least does not make it friendly for Go).

    The following two files lead to the import cycle (xla->xla/service->xla) in Go:

    Since xla_data.proto does not import any other files, we can move it into its own package to break the import cycle. We can utilize the override feature of buf to do this. Here is the final buf.gen.yaml file:

    version: v1
    managed:
      enabled: true
      go_package_prefix:
        default: example.com/mymodule/internal
      override:
        GO_PACKAGE:
          # move the generated xla_data.pb.go file into package xla/data to break the import cycle.
          tensorflow/compiler/xla/xla_data.proto: 'example.com/mymodule/internal/tensorflow/compiler/xla/data'
    plugins:
      - name: go
        out: internal
        opt:
          - module=example.com/mymodule/internal
    
      - name: go-grpc
        out: internal
        opt:
          - module=example.com/mymodule/internal
    

    The Full Setup to Compile Tensorflow proto Files with buf

    Here is the final directory structure:

    ├── buf.gen.yaml
    ├── buf.work.yaml
    ├── buf.yaml
    ├── go.mod
    ├── go.sum
    ├── internal
    │   ├── tensorflow
    │   └── tensorflow_serving
    └── testdata
        ├── serving
        └── tensorflow
    

    buf.gen.yaml: see the "TL;DR" section.

    buf.work.yaml:

    version: v1
    directories:
      - testdata/serving
      - testdata/tensorflow
    

    buf.yaml:

    version: v1
    breaking:
      use:
        - FILE
    lint:
      use:
        - DEFAULT
    

    Here is my environment:

    $ go version
    go version go1.20.3 linux/amd64
    $ buf --version
    1.17.0
    $ protoc --version
    libprotoc 3.12.4
    $ protoc-gen-go --version
    protoc-gen-go v1.30.0
    $ protoc-gen-go-grpc --version
    protoc-gen-go-grpc 1.3.0
    $ git version
    git version 2.37.2
    

    Now execute the following commands in the root of this directory:

    $ go mod init example.com/mymodule
    $ go get google.golang.org/grpc
    $ git clone https://github.com/tensorflow/tensorflow.git testdata/tensorflow
    $ git clone https://github.com/tensorflow/serving.git testdata/serving
    $ buf generate
    $ go build ./...
    

    Notes:

    1. The tensorflow repositories are cloned into the testdata directory so that go build will ignore them.
    2. The setup generates the files into the internal directory. You can modify the buf.gen.yaml file to place them anywhere you want.
    3. go build ./... does not report any error. But I'm not sure whether the generated files are valid or not.