Import cycle in tensorflow protobufs

Question

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?

Answer 1

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:

• tensorflow/compiler/xla/xla.proto

  • import "tensorflow/compiler/xla/service/hlo.proto"

• tensorflow/compiler/xla/service/hlo.proto

  • import "tensorflow/compiler/xla/xla_data.proto"

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.