Testing1

WORKSPACE

@@ -41,8 +41,8 @@ tf_configure(name = "local_config_tf")
 # 3. Request the new archive to be mirrored on mirror.bazel.build for more
 #    reliable downloads.
 
-_TENSORFLOW_GIT_COMMIT = "5bc9d26649cca274750ad3625bd93422617eed4b"  # tf 2.16.1
-_TENSORFLOW_ARCHIVE_SHA256 = "fe592915c85d1a89c20f3dd89db0772ee22a0fbda78e39aa46a778d638a96abc"
+_TENSORFLOW_GIT_COMMIT = "3c92ac03cab816044f7b18a86eb86aa01a294d95"  # tf 2.17.1
+_TENSORFLOW_ARCHIVE_SHA256 = "317dd95c4830a408b14f3e802698eb68d70d81c7c7cfcd3d28b0ba023fe84a68"
 
 http_archive(
     name = "org_tensorflow",
@@ -51,6 +51,8 @@ http_archive(
         "https://github.com/tensorflow/tensorflow/archive/%s.tar.gz" % _TENSORFLOW_GIT_COMMIT,
     ],
     strip_prefix = "tensorflow-%s" % _TENSORFLOW_GIT_COMMIT,
+    patches = ["//third_party:tensorflow.patch"],
+    patch_args = ["-p1"],
 )
 
 load("//third_party:python_configure.bzl", "local_python_configure")
@@ -62,6 +64,21 @@ local_python_configure(name = "local_execution_config_python")
 load("//struct2tensor:workspace.bzl", "struct2tensor_workspace")
 struct2tensor_workspace()
 
+# ===== Protobuf 4.25.6 dependency =====
+# Must be declared BEFORE TensorFlow's workspaces to override the version they pull
+http_archive(
+    name = "com_google_protobuf",
+    sha256 = "4e6727bc5d23177edefa3ad86fd2f5a92cd324151636212fd1f7f13aef3fd2b7",
+    strip_prefix = "protobuf-4.25.6",
+    urls = [
+        "https://github.com/protocolbuffers/protobuf/archive/v4.25.6.tar.gz",
+    ],
+)
+
+# Load Protobuf dependencies
+load("@com_google_protobuf//:protobuf_deps.bzl", "protobuf_deps")
+protobuf_deps()
+
 # Initialize TensorFlow's external dependencies.
 load("@org_tensorflow//tensorflow:workspace3.bzl", "tf_workspace3")
 tf_workspace3()

struct2tensor/BUILD

@@ -84,7 +84,7 @@ s2t_pytype_library(
         "path.py",
     ],
     deps = [
-        "@com_github_tensorflow_metadata//tensorflow_metadata/proto/v0:py_metadata_v0_proto_py",
+        "@com_github_tensorflow_metadata//tensorflow_metadata/proto/v0:metadata_v0_proto_py_pb2",
         "@com_google_protobuf//:protobuf_python",
     ],
 )

struct2tensor/proto/BUILD

@@ -9,7 +9,7 @@ licenses(["notice"])
 s2t_proto_library(
     name = "query_metadata_proto",
     srcs = ["query_metadata.proto"],
-    deps = ["@com_github_tensorflow_metadata//tensorflow_metadata/proto/v0:cc_metadata_v0_proto_cc"],
+    deps = ["@com_github_tensorflow_metadata//tensorflow_metadata/proto/v0:metadata_v0_proto"],
 )
 
 s2t_proto_library_cc(
@@ -19,10 +19,9 @@ s2t_proto_library_cc(
 
 s2t_proto_library_py(
     name = "query_metadata_py_pb2",
-    srcs = ["query_metadata.proto"],
     api_version = 2,
     oss_deps = [
-        "@com_github_tensorflow_metadata//tensorflow_metadata/proto/v0:py_metadata_v0_proto_py",
+        "@com_github_tensorflow_metadata//tensorflow_metadata/proto/v0:metadata_v0_proto_py_pb2",
     ],
     proto_library = "query_metadata_proto",
 )

struct2tensor/struct2tensor.bzl

@@ -14,17 +14,224 @@
 
 """Bazel macros used in OSS."""
 
-load("@com_google_protobuf//:protobuf.bzl", "cc_proto_library", "py_proto_library")
-load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
-load("@rules_cc//cc:cc_library.bzl", "cc_library")
+def _py_proto_library_impl(ctx):
+    """Implementation of py_proto_library rule."""
+    proto_deps = ctx.attr.deps
+
+    # Separate proto and Python dependencies
+    all_sources = []
+    py_infos = []
+
+    for dep in proto_deps:
+        if ProtoInfo in dep:
+            # It's a proto_library - collect proto sources
+            all_sources.extend(dep[ProtoInfo].direct_sources)
+        elif PyInfo in dep:
+            # It's already a py_library - collect its PyInfo for passthrough
+            py_infos.append(dep[PyInfo])
+
+    # Filter to only include sources from the workspace (not external packages)
+    # We can only declare outputs in our own package
+    workspace_sources = []
+    for src in all_sources:
+        # Filter out external sources (they start with external/ or ..)
+        if not src.short_path.startswith("external/") and not src.short_path.startswith("../"):
+            workspace_sources.append(src)
+
+    # Generate Python output files from proto sources
+    py_outputs = []
+    for proto_src in workspace_sources:
+        # Use just the basename to avoid path issues
+        basename = proto_src.basename[:-6]  # Remove .proto
+        py_file = ctx.actions.declare_file(basename + "_pb2.py")
+        py_outputs.append(py_file)
+
+    if py_outputs:
+        # Build proto_path arguments for protoc
+        # We need to include paths for workspace root and external dependencies
+        proto_path_args = []
+
+        # Add current directory to find workspace proto files
+        proto_path_args.append("--proto_path=.")
+
+        # Collect proto_path entries from all transitive dependencies
+        # Use dictionary as a set (Starlark doesn't have set type)
+        proto_paths = {".": True}
+
+        # Also add directories of workspace sources so imports like "any.proto"
+        # (in the same folder) resolve correctly.
+        for ws in workspace_sources:
+            ws_dir = "/".join(ws.short_path.split("/")[:-1])
+            if ws_dir and ws_dir not in proto_paths:
+                proto_paths[ws_dir] = True
+                proto_path_args.append("--proto_path=" + ws_dir)
+
+        for dep in proto_deps:
+            if ProtoInfo in dep:
+                # Add proto_source_root if available
+                if hasattr(dep[ProtoInfo], 'proto_source_root'):
+                    root = dep[ProtoInfo].proto_source_root
+                    if root and root not in proto_paths:
+                        proto_paths[root] = True
+                        proto_path_args.append("--proto_path=" + root)
+
+                # Also derive from file paths for more coverage
+                for src in dep[ProtoInfo].transitive_sources.to_list():
+                    # Use the directory containing the proto file's import root
+                    # For external/com_google_protobuf/src/google/protobuf/any.proto,
+                    # we want external/com_google_protobuf/src
+                    if src.path.startswith("external/com_google_protobuf/"):
+                        proto_path = "external/com_google_protobuf/src"
+                        if proto_path not in proto_paths:
+                            proto_paths[proto_path] = True
+                            proto_path_args.append("--proto_path=" + proto_path)
+                    elif src.path.startswith("external/"):
+                        # For other external repos like tensorflow_metadata
+                        # Extract external/repo_name
+                        parts = src.path.split("/")
+                        if len(parts) >= 2:
+                            proto_path = "/".join(parts[:2])
+                            if proto_path not in proto_paths:
+                                proto_paths[proto_path] = True
+                                proto_path_args.append("--proto_path=" + proto_path)
+
+                    # Also add Bazel root paths
+                    if src.root.path and src.root.path not in proto_paths:
+                        proto_paths[src.root.path] = True
+                        proto_path_args.append("--proto_path=" + src.root.path)
+
+# Build list of proto file paths - only include workspace sources
+        proto_file_args = []
+        for src in workspace_sources:
+            proto_file_args.append(src.short_path)
+
+        # Run protoc to generate Python files
+        # Use ctx.bin_dir.path as the output directory root
+        output_root = ctx.bin_dir.path
+
+        ctx.actions.run(
+            # Include workspace sources plus all transitive dependencies for imports
+            inputs = depset(direct = workspace_sources, transitive = [
+                dep[ProtoInfo].transitive_sources for dep in proto_deps if ProtoInfo in dep
+            ]),
+            outputs = py_outputs,
+            executable = ctx.executable._protoc,
+            arguments = [
+                "--python_out=" + output_root,
+            ] + proto_path_args + proto_file_args,
+            mnemonic = "ProtocPython",
+        )
+
+    # Collect transitive sources from both generated files and Python deps
+    all_transitive_sources = [depset(py_outputs)]
+    all_imports = [depset([ctx.bin_dir.path])] if py_outputs else []
+
+    for py_info in py_infos:
+        all_transitive_sources.append(py_info.transitive_sources)
+        if hasattr(py_info, 'imports'):
+            all_imports.append(py_info.imports)
+
+    # Return PyInfo provider so this can be used as a py_library dependency
+    # Merge proto-generated files with passthrough Python dependencies
+    return [
+        DefaultInfo(files = depset(py_outputs)),
+        PyInfo(
+            transitive_sources = depset(transitive = all_transitive_sources),
+            imports = depset(transitive = all_imports),
+            has_py2_only_sources = False,
+            has_py3_only_sources = True,
+        ),
+    ]
+
+_py_proto_library_rule = rule(
+    implementation = _py_proto_library_impl,
+    attrs = {
+        "deps": attr.label_list(
+            providers = [[ProtoInfo], [PyInfo]],  # Accept either ProtoInfo OR PyInfo
+            doc = "Proto library or Python library dependencies",
+        ),
+        "_protoc": attr.label(
+            default = "@com_google_protobuf//:protoc",
+            executable = True,
+            cfg = "exec",
+        ),
+    },
+    provides = [PyInfo],
+)
+
+# Wrapper for cc_proto_library to maintain compatibility with old Protobuf 3.x API
+def cc_proto_library(
+        name,
+        srcs = [],
+        deps = [],
+        cc_libs = [],
+        protoc = None,
+        default_runtime = None,
+        use_grpc_plugin = None,
+        testonly = 0,
+        visibility = None,
+        **kwargs):
+    """Wrapper for cc_proto_library that works with Protobuf 4.x."""
+    _ignore = [cc_libs, protoc, default_runtime, use_grpc_plugin, kwargs]
+
+    # Create proto_library first
+    native.proto_library(
+        name = name + "_proto",
+        srcs = srcs,
+        deps = [d + "_proto" if not d.startswith("@") else d for d in deps],
+        testonly = testonly,
+        visibility = visibility,
+    )
+
+    # Create cc_proto_library that depends on proto_library
+    native.cc_proto_library(
+        name = name,
+        deps = [":" + name + "_proto"],
+        testonly = testonly,
+        visibility = visibility,
+    )
 
 def s2t_pytype_library(
         name,
         srcs = [],
         deps = [],
         srcs_version = "PY3ONLY",
         testonly = False):
-    native.py_library(name = name, srcs = srcs, deps = deps, testonly = testonly)
+    """Python library that automatically wraps proto_library deps with PyInfo.
+
+    This wrapper wraps all dependencies with our custom py_proto_library_rule.
+    Dependencies that don't provide ProtoInfo will fail with a clear error.
+    Dependencies that do provide ProtoInfo (proto_library targets) will get PyInfo.
+    """
+    # Process dependencies to wrap them all with our custom rule
+    processed_deps = []
+    for dep in deps:
+        # Skip protobuf_python - it's already a proper Python library
+        if dep == "@com_google_protobuf//:protobuf_python":
+            processed_deps.append(dep)
+            continue
+
+        # Create a safe wrapper name for this dependency
+        safe_dep_name = dep.replace(":", "_").replace("//", "").replace("/", "_").replace("@", "").replace("-", "_").replace(".", "_")
+        wrapper_name = name + "_proto_wrapper_" + safe_dep_name
+
+        # Wrap all dependencies with our custom py_proto_library rule
+        # If the dep provides ProtoInfo, this will work and provide PyInfo
+        # If it doesn't provide ProtoInfo, it will fail with a clear error
+        _py_proto_library_rule(
+            name = wrapper_name,
+            deps = [dep],
+            testonly = testonly,
+        )
+        processed_deps.append(":" + wrapper_name)
+
+    native.py_library(
+        name = name,
+        srcs = srcs,
+        deps = processed_deps,
+        testonly = testonly,
+    )
+
 
 def s2t_proto_library(
         name,
@@ -52,18 +259,22 @@ def s2t_proto_library(
         testonly = testonly,
     )
 
-    use_grpc_plugin = None
-    if cc_grpc_version:
-        use_grpc_plugin = True
+    # Create a native proto_library for Python generation
+    # This is needed by s2t_proto_library_py
+    proto_lib_deps = [d + "_proto" if not d.startswith("@") else d for d in deps]
+    native.proto_library(
+        name = name + "_proto",
+        srcs = srcs,
+        deps = proto_lib_deps,
+        visibility = visibility,
+        testonly = testonly,
+    )
 
-    # TODO(martinz): replace with proto_library, when that works.
-    cc_proto_library(
+    # Create cc_proto_library that depends on the proto_library we just created
+    # Don't use our cc_proto_library wrapper to avoid duplicate proto_library creation
+    native.cc_proto_library(
         name = name,
-        srcs = srcs,
-        deps = deps,
-        cc_libs = ["@com_google_protobuf//:protobuf"],
-        protoc = "@com_google_protobuf//:protoc",
-        default_runtime = "@com_google_protobuf//:protobuf",
+        deps = [":" + name + "_proto"],
         testonly = testonly,
         visibility = visibility,
     )
@@ -76,7 +287,7 @@ DYNAMIC_DEPS = ["@local_config_tf//:libtensorflow_framework", "@local_config_tf/
 
 def s2t_dynamic_binary(name, deps):
     """Creates a .so file intended for linking with tensorflow_framework.so."""
-    cc_binary(
+    native.cc_binary(
         name = name,
         copts = DYNAMIC_COPTS,
         linkshared = 1,
@@ -89,7 +300,7 @@ def s2t_dynamic_library(
         deps = None):
     """Creates a static library intended for linking with tensorflow_framework.so."""
     true_deps = [] if deps == None else deps
-    cc_library(
+    native.cc_library(
         name = name,
         srcs = srcs,
         alwayslink = 1,
@@ -169,15 +380,25 @@ def s2t_proto_library_cc(
     )
 
 def s2t_proto_library_py(name, proto_library, srcs = [], deps = [], oss_deps = [], visibility = None, testonly = 0, api_version = None):
-    """Opensource py_proto_library."""
-    _ignore = [proto_library, api_version]
-    py_proto_library(
+    """Opensource py_proto_library.
+
+    Uses a custom rule implementation that properly generates Python from proto_library
+    and provides PyInfo for Python library dependencies.
+
+    Note: s2t_proto_library creates {name}_proto for the proto_library, so we append _proto.
+    """
+    _ignore = [api_version, srcs, deps]
+
+    if not proto_library:
+        fail("proto_library parameter is required for s2t_proto_library_py")
+
+    # s2t_proto_library creates a proto_library named {name}_proto
+    # So we need to reference it correctly
+    actual_proto_library = ":" + proto_library + "_proto"
+
+    # Use our custom py_proto_library rule
+    _py_proto_library_rule(
         name = name,
-        srcs = srcs,
-        srcs_version = "PY3ONLY",
-        deps = ["@com_google_protobuf//:well_known_types_py_pb2"] + oss_deps,
-        default_runtime = "@com_google_protobuf//:protobuf_python",
-        protoc = "@com_google_protobuf//:protoc",
+        deps = [actual_proto_library] + oss_deps,
         visibility = visibility,
-        testonly = testonly,
     )