allow building ml_metadata_store_server image on ARM64

.bazelversion

@@ -0,0 +1 @@
+6.1.0

ml_metadata/.bazelversion

@@ -1,14 +1 @@
-# Copyright 2022 Google LLC. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
 6.1.0

ml_metadata/postgresql.BUILD

@@ -6,7 +6,7 @@ exports_files(["COPYRIGHT"])
 
 # This is OSS version of PostgreSQL bazel build target.
 # The build target is consistent with Tensorflow's dependency on PostgreSQL.
-# Reference: https://github.com/tensorflow/io/commit/6a2b9b3e94fd80f9a7f0b982089e8a69751bd059
+# Reference: https://github.com/tensorflow/io/commit/a1171cdd20e658ef3f1a8b3bf66dd4e228ceae30
 cc_library(
     name = "postgresql",
     srcs = glob([
@@ -43,8 +43,8 @@ cc_library(
         "src/common/unicode_norm.c",
         "src/common/username.c",
         "src/common/wait_error.c",
-        "src/interfaces/libpq/fe-auth.c",
         "src/interfaces/libpq/fe-auth-scram.c",
+        "src/interfaces/libpq/fe-auth.c",
         "src/interfaces/libpq/fe-connect.c",
         "src/interfaces/libpq/fe-exec.c",
         "src/interfaces/libpq/fe-lobj.c",
@@ -62,8 +62,6 @@ cc_library(
         "src/port/path.c",
         "src/port/pg_bitutils.c",
         "src/port/pg_crc32c_sb8.c",
-        # Comment this line out to force usage of sb8 algorithm of crc32c
-        # "src/port/pg_crc32c_sse42_choose.c",
         "src/port/pg_strong_random.c",
         "src/port/pgcheckdir.c",
         "src/port/pgmkdirp.c",
@@ -80,12 +78,28 @@ cc_library(
         "src/port/tar.c",
         "src/port/thread.c",
     ] + select({
-        "@//ml_metadata:macos": [],
+        "@bazel_tools//src/conditions:darwin": [],
+        "@bazel_tools//src/conditions:windows": [
+            "src/interfaces/libpq/pthread-win32.c",
+            "src/interfaces/libpq/win32.c",
+            "src/port/dirmod.c",
+            "src/port/getaddrinfo.c",
+            "src/port/inet_aton.c",
+            "src/port/open.c",
+            "src/port/strlcpy.c",
+            "src/port/win32error.c",
+            "src/port/win32setlocale.c",
+            "src/port/pthread-win32.h",
+        ],
         "//conditions:default": [
             "src/port/getpeereid.c",
-            "src/port/strlcat.c",
             "src/port/strlcpy.c",
         ],
+    }) + select({
+        "@platforms//cpu:x86_64": [
+            "src/port/pg_crc32c_sse42_choose.c",
+        ],
+        "//conditions:default": [],
     }),
     hdrs = [
         "config/pg_config.h",
@@ -98,7 +112,14 @@ cc_library(
     defines = [
         "FRONTEND",
     ] + select({
-        "@//ml_metadata:macos": [
+        "@bazel_tools//src/conditions:windows": [
+            "BLCKSZ=8192",
+            "XLOG_BLCKSZ=8192",
+            'PG_MAJORVERSION=\\"12\\"',
+            "HAVE_LIBZ=1",
+            "WIN32",
+        ],
+        "@bazel_tools//src/conditions:darwin": [
             "HAVE_DECL_STRLCPY=1",
             "HAVE_STRLCPY=1",
             "HAVE_STRUCT_SOCKADDR_STORAGE_SS_LEN=1",
@@ -114,9 +135,18 @@ cc_library(
         "src/include",
         "src/interfaces/libpq",
     ] + select({
+        "@bazel_tools//src/conditions:windows": [
+            "src/include/port/win32",
+            "src/include/port/win32_msvc",
+            "src/port",
+        ],
         "//conditions:default": [],
     }),
     linkopts = select({
+        "@bazel_tools//src/conditions:windows": [
+            "-DEFAULTLIB:ws2_32.lib",
+            "-DEFAULTLIB:shell32.lib",
+        ],
         "//conditions:default": [],
     }),
     deps = [],
@@ -125,7 +155,10 @@ cc_library(
 genrule(
     name = "pg_config_os_h",
     srcs = select({
-        "@//ml_metadata:macos": [
+        "@bazel_tools//src/conditions:windows": [
+            "src/include/port/win32.h",
+        ],
+        "@bazel_tools//src/conditions:darwin": [
             "src/include/port/darwin.h",
         ],
         "//conditions:default": [
@@ -135,18 +168,22 @@ genrule(
     outs = [
         "config/pg_config_os.h",
     ],
-    cmd = "cp $< $@",
+    cmd = ("cp $< $@"),
 )
 
 genrule(
     name = "pg_config_ext_h",
     srcs = select({
+        "@bazel_tools//src/conditions:windows": [
+            "src/include/pg_config_ext.h.win32",
+        ],
         "//conditions:default": [
             "src/include/pg_config_ext.h.in",
         ],
     }),
     outs = ["config/pg_config_ext.h"],
     cmd = select({
+        "@bazel_tools//src/conditions:windows": ("cp $< $@"),
         "//conditions:default": (
             "sed " +
             "-e 's/undef PG_INT64_TYPE/define PG_INT64_TYPE long int/g' " +
@@ -357,6 +394,9 @@ genrule(
 genrule(
     name = "pg_config_h",
     srcs = select({
+        "@bazel_tools//src/conditions:windows": [
+            "src/include/pg_config.h.win32",
+        ],
         "//conditions:default": [
             "src/include/pg_config.h.in",
         ],
@@ -365,7 +405,8 @@ genrule(
         "config/pg_config.h",
     ],
     cmd = select({
-        "//conditions:default": "\n".join([
+        "@bazel_tools//src/conditions:windows": ("cp $< $@"),
+        "//conditions:default": ("\n".join([
             "cat <<'EOF' >$@",
             "/* src/include/pg_config.h.  Generated from pg_config.h.in by configure.  */",
             "/* src/include/pg_config.h.in.  Generated from configure.in by autoheader.  */",
@@ -1145,9 +1186,11 @@ genrule(
             "/* Define to 1 if you have __cpuid. */",
             "/* #undef HAVE__CPUID */",
             "",
+            "#if defined __x86_64__",
             "/* Define to 1 if you have __get_cpuid. */",
             "#define HAVE__GET_CPUID 1",
             "",
+            "#endif",
             "/* Define to 1 if your compiler understands _Static_assert. */",
             "#define HAVE__STATIC_ASSERT 1",
             "",
@@ -1312,15 +1355,13 @@ genrule(
             "/* #undef USE_PAM */",
             "",
             "/* Define to 1 to use software CRC-32C implementation (slicing-by-8). */",
-            # Force usage of sb8 algorithm of crc32c
-            "#define USE_SLICING_BY_8_CRC32C 1",
+            "/* #undef USE_SLICING_BY_8_CRC32C */",
             "",
             "/* Define to 1 use Intel SSE 4.2 CRC instructions. */",
             "/* #undef USE_SSE42_CRC32C */",
             "",
             "/* Define to 1 to use Intel SSE 4.2 CRC instructions with a runtime check. */",
-            # Do not check at runtime but force usage of sb8 algorithm of crc32c
-            "/* #undef USE_SSE42_CRC32C_WITH_RUNTIME_CHECK */",
+            "#define USE_SSE42_CRC32C_WITH_RUNTIME_CHECK 1",
             "",
             "/* Define to build with systemd support. (--with-systemd) */",
             "/* #undef USE_SYSTEMD */",
@@ -1334,6 +1375,15 @@ genrule(
             "/* Define to select unnamed POSIX semaphores. */",
             "/* #undef USE_UNNAMED_POSIX_SEMAPHORES */",
             "",
+            "/* Define to use native Windows API for random number generation */",
+            "/* #undef USE_WIN32_RANDOM */",
+            "",
+            "/* Define to select Win32-style semaphores. */",
+            "/* #undef USE_WIN32_SEMAPHORES */",
+            "",
+            "/* Define to select Win32-style shared memory. */",
+            "/* #undef USE_WIN32_SHARED_MEMORY */",
+            "",
             "/* Define to 1 if `wcstombs_l' requires <xlocale.h>. */",
             "#define WCSTOMBS_L_IN_XLOCALE 1",
             "",
@@ -1405,6 +1455,6 @@ genrule(
             "   pointer, if such a type exists, and if the system does not define it. */",
             "/* #undef uintptr_t */",
             "EOF",
-        ]),
+        ])),
     }),
-)
+)

ml_metadata/tools/docker_server/Dockerfile

@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO(b/195701120) Introduces l.gcr.io/google/bazel:5.3.0 when it is available
-# and makes sure that it uses ubuntu 20.04 as base image. Currently the lastest
-# version only supports bazel 6.1.0.
 FROM ubuntu:20.04 as builder
 
+ARG TARGETARCH
+
 RUN export DEBIAN_FRONTEND=noninteractive && \
   apt-get update -y --option Acquire::Retries=3 && \
   apt-get install --no-install-recommends -y -q --option Acquire::Retries=3 \
@@ -35,17 +34,16 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
   python-is-python3 \
   python3-dev
 
-# Set up Bazel 6.1.0
-ENV BAZEL_VERSION 6.1.0
+# Install Bazelisk
+ENV BAZELISK_VERSION 1.19.0
 WORKDIR /
-RUN mkdir /bazel && \
-    cd /bazel && \
-    curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \
-    curl -H "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36" -fSsL -o /bazel/LICENSE.txt https://raw.githubusercontent.com/bazelbuild/bazel/master/LICENSE && \
-    chmod +x bazel-*.sh && \
-    ./bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \
-    cd / && \
-    rm -f /bazel/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh
+RUN case "${TARGETARCH}" in \
+      amd64) BAZELISK_ARCH="amd64" ;; \
+      arm64) BAZELISK_ARCH="arm64" ;; \
+      *) echo "Unsupported architecture: ${TARGETARCH}"; exit 1 ;; \
+    esac && \
+    curl -fSsL https://github.com/bazelbuild/bazelisk/releases/download/v${BAZELISK_VERSION}/bazelisk-linux-${BAZELISK_ARCH} -o /usr/local/bin/bazel && \
+    chmod +x /usr/local/bin/bazel
 
 ADD . /mlmd-src
 WORKDIR /mlmd-src

ml_metadata/tools/docker_server/README.md

@@ -1,6 +1,7 @@
-Files for building the [Docker](http://www.docker.com) image for running the
-MLMD gRPC server.
+Files for building the container image that runs the MLMD gRPC server (`gcr.io/tfx-oss-public/ml_metadata_store_server`).
 
-To build a docker image, run
-`./ml_metadata/tools/docker_server/build_docker_image.sh` under root directory
-of github checkout.
+To build the image, run the following command:
+
+```shell
+./ml_metadata/tools/docker_server/build_docker_image.sh
+```

ml_metadata/tools/docker_server/build_docker_image.sh

@@ -1,3 +1,4 @@
+#!/usr/bin/env bash
 # Copyright 2019 Google LLC. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,11 +15,18 @@
 
 # Convenience script to build docker image for MLMD gRPC server.
 
-set -u -x
+set -euo pipefail
 
 DOCKER_IMAGE_REPO=${DOCKER_IMAGE_REPO:-"gcr.io/tfx-oss-public/ml_metadata_store_server"}
 DOCKER_IMAGE_TAG=${DOCKER_IMAGE_TAG:-"latest"}
 DOCKER_FILE=${DOCKER_FILE:-"Dockerfile"}
 
+# Change to repository root.
+REPOSITORY_ROOT_PATH=$(git rev-parse --show-toplevel)
+cd "${REPOSITORY_ROOT_PATH}"
+
 # Run docker build command.
-docker build -t ${DOCKER_IMAGE_REPO}:${DOCKER_IMAGE_TAG} -f ml_metadata/tools/docker_server/${DOCKER_FILE} .
+docker build \
+  -t ${DOCKER_IMAGE_REPO}:${DOCKER_IMAGE_TAG} \
+  -f ml_metadata/tools/docker_server/${DOCKER_FILE} \
+  .