Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit fe02cba

Browse files
AvasamAvasam
authored
Bump mypy-protobuf in sync_tensorflow script and improve generation scripts (#11740)
1 parent efad2fb commit fe02cba

94 files changed

Lines changed: 2346 additions & 2054 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎scripts/generate_proto_stubs.sh‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ echo "Working in $TMP_DIR"
3535
wget "$PROTOC_URL"
3636
mkdir protoc_install
3737
unzip "$PROTOC_FILENAME" -d protoc_install
38+
protoc_install/bin/protoc --version
3839

3940
# Fetch protoc-python (which contains all the .proto files)
4041
wget "$PYTHON_PROTOBUF_URL"
@@ -67,16 +68,22 @@ PROTO_FILES=$(grep "GenProto.*google" $PYTHON_PROTOBUF_DIR/python/setup.py | \
6768

6869
# And regenerate!
6970
# shellcheck disable=SC2086
70-
protoc_install/bin/protoc --proto_path="$PYTHON_PROTOBUF_DIR/src" --mypy_out="relax_strict_optional_primitives:$REPO_ROOT/stubs/protobuf" $PROTO_FILES
71+
protoc_install/bin/protoc \
72+
--proto_path="$PYTHON_PROTOBUF_DIR/src" \
73+
--mypy_out="relax_strict_optional_primitives:$REPO_ROOT/stubs/protobuf" \
74+
$PROTO_FILES
7175

7276
PYTHON_PROTOBUF_VERSION=$(jq -r '.[] | .languages.python' "$PYTHON_PROTOBUF_DIR/version.json")
7377

78+
# Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
79+
rm -rf "$TMP_DIR"
80+
# Must be in a git repository to run pre-commit
81+
cd "$REPO_ROOT"
82+
7483
sed --in-place="" \
7584
"s/extra_description = .*$/extra_description = \"Generated using [mypy-protobuf==$MYPY_PROTOBUF_VERSION](https:\/\/github.com\/nipunn1313\/mypy-protobuf\/tree\/v$MYPY_PROTOBUF_VERSION) on [protobuf v$PROTOBUF_VERSION](https:\/\/github.com\/protocolbuffers\/protobuf\/releases\/tag\/v$PROTOBUF_VERSION) (python protobuf==$PYTHON_PROTOBUF_VERSION)\"/" \
76-
"$REPO_ROOT/stubs/protobuf/METADATA.toml"
85+
stubs/protobuf/METADATA.toml
7786

78-
# Must be run in a git repository
79-
cd "$REPO_ROOT"
8087
# use `|| true` so the script still continues even if a pre-commit hook
8188
# applies autofixes (which will result in a nonzero exit code)
82-
pre-commit run --files $(git ls-files -- "$REPO_ROOT/stubs/protobuf/**_pb2.pyi") || true
89+
pre-commit run --files $(git ls-files -- "stubs/protobuf/**_pb2.pyi") || true

‎scripts/sync_tensorflow_protobuf_stubs.sh‎

Lines changed: 78 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -1,79 +1,97 @@
11
#!/bin/bash
2+
# Based on scripts/generate_proto_stubs.sh.
3+
# Generates the protobuf stubs for the given tensorflow version using mypy-protobuf.
4+
# Generally, new minor versions are a good time to update the stubs.
5+
26
set -euxo pipefail
37

4-
# Partly based on scripts/generate_proto_stubs.sh.
8+
# Need protoc >= 3.15 for explicit optional
9+
PROTOBUF_VERSION=25.3 # 4.25.3
10+
# Whenever you update TENSORFLOW_VERSION here, version should be updated
11+
# in stubs/tensorflow/METADATA.toml and vice-versa.
12+
TENSORFLOW_VERSION=2.12.1
13+
MYPY_PROTOBUF_VERSION=3.6.0
514

6-
# Generates the protobuf stubs for the given tensorflow version using mypy-protobuf.
7-
# Generally, new minor versions are a good time to update the stubs.
15+
if uname -a | grep Darwin; then
16+
# brew install coreutils wget
17+
PLAT=osx
18+
else
19+
PLAT=linux
20+
fi
821
REPO_ROOT="$(realpath "$(dirname "${BASH_SOURCE[0]}")"/..)"
22+
TMP_DIR="$(mktemp -d)"
23+
TENSORFLOW_FILENAME="v$TENSORFLOW_VERSION.zip"
24+
PROTOC_FILENAME="protoc-$PROTOBUF_VERSION-$PLAT-x86_64.zip"
25+
PROTOC_URL="https://github.com/protocolbuffers/protobuf/releases/download/v$PROTOBUF_VERSION/$PROTOC_FILENAME"
26+
TENSORFLOW_URL="https://github.com/tensorflow/tensorflow/archive/refs/tags/$TENSORFLOW_FILENAME"
927

10-
# This version should be consistent with the version in tensorflow's METADATA.toml.
11-
TENSORFLOW_VERSION=2.12.1
12-
# Latest mypy-protobuf has dependency on protobuf >4, which is incompatible at runtime
13-
# with tensorflow. However, the stubs produced do still work with tensorflow. So after
14-
# installing mypy-protobuf, before running stubtest on tensorflow you should downgrade
15-
# protobuf<4.
16-
MYPY_PROTOBUF_VERSION=3.5.0
28+
cd "$TMP_DIR"
29+
echo "Working in $TMP_DIR"
1730

18-
pip install pre-commit mypy-protobuf=="$MYPY_PROTOBUF_VERSION"
31+
# Install protoc
32+
wget "$PROTOC_URL"
33+
mkdir protoc_install
34+
unzip "$PROTOC_FILENAME" -d protoc_install
35+
protoc_install/bin/protoc --version
1936

20-
cd "$(dirname "$0")" > /dev/null
21-
cd ../stubs/tensorflow
22-
mkdir -p repository
23-
pushd repository &> /dev/null
24-
# If the script fails halfway, it's nice to be able to re-run it immediately
25-
if [ ! -d "tensorflow" ] ; then
26-
git clone --depth 1 --branch v"$TENSORFLOW_VERSION" https://github.com/tensorflow/tensorflow.git
27-
fi
28-
pushd tensorflow &> /dev/null
29-
# Folders here cover the more commonly used protobufs externally and
30-
# their dependencies. Tensorflow has more protobufs and can be added if requested.
31-
protoc --mypy_out "relax_strict_optional_primitives:$REPO_ROOT/stubs/tensorflow" \
32-
tensorflow/compiler/xla/*.proto \
33-
tensorflow/compiler/xla/service/*.proto \
34-
tensorflow/core/example/*.proto \
35-
tensorflow/core/framework/*.proto \
36-
tensorflow/core/protobuf/*.proto \
37-
tensorflow/core/protobuf/tpu/*.proto \
38-
tensorflow/core/util/*.proto \
39-
tensorflow/python/keras/protobuf/*.proto \
40-
tensorflow/tsl/protobuf/*.proto
41-
popd &> /dev/null
42-
popd &> /dev/null
37+
# Fetch tensorflow (which contains all the .proto files)
38+
wget "$TENSORFLOW_URL"
39+
unzip "$TENSORFLOW_FILENAME"
40+
TENSORFLOW_DIR="tensorflow-$TENSORFLOW_VERSION"
41+
42+
# Prepare virtualenv
43+
python3 -m venv .venv
44+
source .venv/bin/activate
45+
python3 -m pip install pre-commit mypy-protobuf=="$MYPY_PROTOBUF_VERSION"
46+
47+
# Remove existing pyi
48+
find "$REPO_ROOT/stubs/tensorflow/" -name "*_pb2.pyi" -delete
49+
50+
# Folders here cover the more commonly used protobufs externally and
51+
# their dependencies. Tensorflow has more protobufs and can be added if requested.
52+
protoc_install/bin/protoc \
53+
--proto_path="$TENSORFLOW_DIR" \
54+
--mypy_out "relax_strict_optional_primitives:$REPO_ROOT/stubs/tensorflow" \
55+
$TENSORFLOW_DIR/tensorflow/compiler/xla/*.proto \
56+
$TENSORFLOW_DIR/tensorflow/compiler/xla/service/*.proto \
57+
$TENSORFLOW_DIR/tensorflow/core/example/*.proto \
58+
$TENSORFLOW_DIR/tensorflow/core/framework/*.proto \
59+
$TENSORFLOW_DIR/tensorflow/core/protobuf/*.proto \
60+
$TENSORFLOW_DIR/tensorflow/core/protobuf/tpu/*.proto \
61+
$TENSORFLOW_DIR/tensorflow/core/util/*.proto \
62+
$TENSORFLOW_DIR/tensorflow/python/keras/protobuf/*.proto \
63+
$TENSORFLOW_DIR/tensorflow/tsl/protobuf/*.proto \
64+
65+
# Cleanup after ourselves, this is a temp dir, but it can still grow fast if run multiple times
66+
rm -rf "$TMP_DIR"
67+
# Must be in a git repository to run pre-commit
68+
cd "$REPO_ROOT"
4369

4470
# These protos exist in a folder with protos used in python, but are not
4571
# included in the python wheel. They are likely only used for other
4672
# language builds. stubtest was used to identify them by looking for
4773
# ModuleNotFoundError.
48-
rm tensorflow/compiler/xla/service/hlo_execution_profile_data_pb2.pyi \
49-
tensorflow/compiler/xla/service/hlo_profile_printer_data_pb2.pyi \
50-
tensorflow/compiler/xla/service/test_compilation_environment_pb2.pyi \
51-
tensorflow/compiler/xla/xla_pb2.pyi \
52-
tensorflow/core/protobuf/autotuning_pb2.pyi \
53-
tensorflow/core/protobuf/conv_autotuning_pb2.pyi \
54-
tensorflow/core/protobuf/critical_section_pb2.pyi \
55-
tensorflow/core/protobuf/eager_service_pb2.pyi \
56-
tensorflow/core/protobuf/master_pb2.pyi \
57-
tensorflow/core/protobuf/master_service_pb2.pyi \
58-
tensorflow/core/protobuf/replay_log_pb2.pyi \
59-
tensorflow/core/protobuf/tpu/compile_metadata_pb2.pyi \
60-
tensorflow/core/protobuf/worker_pb2.pyi \
61-
tensorflow/core/protobuf/worker_service_pb2.pyi \
62-
tensorflow/core/util/example_proto_fast_parsing_test_pb2.pyi
63-
74+
rm \
75+
stubs/tensorflow/tensorflow/compiler/xla/service/hlo_execution_profile_data_pb2.pyi \
76+
stubs/tensorflow/tensorflow/compiler/xla/service/hlo_profile_printer_data_pb2.pyi \
77+
stubs/tensorflow/tensorflow/compiler/xla/service/test_compilation_environment_pb2.pyi \
78+
stubs/tensorflow/tensorflow/compiler/xla/xla_pb2.pyi \
79+
stubs/tensorflow/tensorflow/core/protobuf/autotuning_pb2.pyi \
80+
stubs/tensorflow/tensorflow/core/protobuf/conv_autotuning_pb2.pyi \
81+
stubs/tensorflow/tensorflow/core/protobuf/critical_section_pb2.pyi \
82+
stubs/tensorflow/tensorflow/core/protobuf/eager_service_pb2.pyi \
83+
stubs/tensorflow/tensorflow/core/protobuf/master_pb2.pyi \
84+
stubs/tensorflow/tensorflow/core/protobuf/master_service_pb2.pyi \
85+
stubs/tensorflow/tensorflow/core/protobuf/replay_log_pb2.pyi \
86+
stubs/tensorflow/tensorflow/core/protobuf/tpu/compile_metadata_pb2.pyi \
87+
stubs/tensorflow/tensorflow/core/protobuf/worker_pb2.pyi \
88+
stubs/tensorflow/tensorflow/core/protobuf/worker_service_pb2.pyi \
89+
stubs/tensorflow/tensorflow/core/util/example_proto_fast_parsing_test_pb2.pyi \
6490

6591
sed --in-place="" \
6692
"s/extra_description = .*$/extra_description = \"Partially generated using [mypy-protobuf==$MYPY_PROTOBUF_VERSION](https:\/\/github.com\/nipunn1313\/mypy-protobuf\/tree\/v$MYPY_PROTOBUF_VERSION) on tensorflow==$TENSORFLOW_VERSION\"/" \
67-
"$REPO_ROOT/stubs/tensorflow/METADATA.toml"
68-
69-
# Cleanup last. If the script fails halfway, it's nice to be able to re-run it immediately
70-
rm -rf repository/
93+
stubs/tensorflow/METADATA.toml
7194

72-
# Must be run in a git repository
73-
cd $REPO_ROOT
7495
# use `|| true` so the script still continues even if a pre-commit hook
7596
# applies autofixes (which will result in a nonzero exit code)
76-
pre-commit run --files $(git ls-files -- "$REPO_ROOT/stubs/tensorflow/tensorflow") || true
77-
# Ruff takes two passes to fix everything, re-running all of pre-commit is *slow*
78-
# and we don't need --unsafe-fixes to remove imports
79-
ruff check "$REPO_ROOT/stubs/tensorflow/tensorflow" --fix --exit-zero
97+
pre-commit run --files $(git ls-files -- "stubs/tensorflow/**_pb2.pyi") || true

‎stubs/tensorflow/METADATA.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ version = "2.15.*"
22
upstream_repository = "https://github.com/tensorflow/tensorflow"
33
# requires a version of numpy with a `py.typed` file
44
requires = ["numpy>=1.20", "types-protobuf", "types-requests"]
5-
extra_description = "Partially generated using [mypy-protobuf==3.5.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.5.0) on tensorflow==2.12.1"
5+
extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) on tensorflow==2.12.1"
66
partial_stub = true
77

88
[tool.stubtest]

‎stubs/tensorflow/tensorflow/compiler/xla/autotune_results_pb2.pyi‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,10 @@ See the License for the specific language governing permissions and
1616
limitations under the License.
1717
==============================================================================
1818
"""
19+
1920
import builtins
2021
import collections.abc
21-
import typing as typing_extensions
22+
import typing
2223

2324
import google.protobuf.descriptor
2425
import google.protobuf.internal.containers
@@ -27,7 +28,7 @@ import tensorflow.tsl.protobuf.autotuning_pb2
2728

2829
DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
2930

30-
@typing_extensions.final
31+
@typing.final
3132
class AutotuneResults(google.protobuf.message.Message):
3233
"""A collection of algorithms for particular dot/convs. Usually this is "the
3334
best" algorithm for the particular dot/conv, although that's not strictly
@@ -43,7 +44,7 @@ class AutotuneResults(google.protobuf.message.Message):
4344

4445
DESCRIPTOR: google.protobuf.descriptor.Descriptor
4546

46-
@typing_extensions.final
47+
@typing.final
4748
class Entry(google.protobuf.message.Message):
4849
DESCRIPTOR: google.protobuf.descriptor.Descriptor
4950

@@ -60,15 +61,16 @@ class AutotuneResults(google.protobuf.message.Message):
6061
algorithms returned by cublasLt. Different version of cublasLt ->
6162
different list of algos -> different interpretation of results!
6263
"""
64+
6365
def __init__(
6466
self,
6567
*,
6668
device: builtins.str | None = ...,
6769
hlo: builtins.str | None = ...,
6870
result: tensorflow.tsl.protobuf.autotuning_pb2.AutotuneResult | None = ...,
6971
) -> None: ...
70-
def HasField(self, field_name: typing_extensions.Literal["result", b"result"]) -> builtins.bool: ...
71-
def ClearField(self, field_name: typing_extensions.Literal["device", b"device", "hlo", b"hlo", "result", b"result"]) -> None: ...
72+
def HasField(self, field_name: typing.Literal["result", b"result"]) -> builtins.bool: ...
73+
def ClearField(self, field_name: typing.Literal["device", b"device", "hlo", b"hlo", "result", b"result"]) -> None: ...
7274

7375
VERSION_FIELD_NUMBER: builtins.int
7476
DOTS_FIELD_NUMBER: builtins.int
@@ -85,6 +87,6 @@ class AutotuneResults(google.protobuf.message.Message):
8587
dots: collections.abc.Iterable[global___AutotuneResults.Entry] | None = ...,
8688
convs: collections.abc.Iterable[global___AutotuneResults.Entry] | None = ...,
8789
) -> None: ...
88-
def ClearField(self, field_name: typing_extensions.Literal["convs", b"convs", "dots", b"dots", "version", b"version"]) -> None: ...
90+
def ClearField(self, field_name: typing.Literal["convs", b"convs", "dots", b"dots", "version", b"version"]) -> None: ...
8991

9092
global___AutotuneResults = AutotuneResults

0 commit comments

Comments
 (0)