Add orbax/checkpoint/_src/logging/BUILD and updated logging dependency to the updated build target.

PiperOrigin-RevId: 715154137

Author: liangyaning33

checkpoint/orbax/checkpoint/_src/BUILD

@@ -0,0 +1,26 @@
+package(
+    default_applicable_licenses = ["//:package_license"],
+    default_visibility = ["//visibility:public"],
+)
+
+py_library(
+    name = "asyncio_utils",
+    srcs = ["asyncio_utils.py"],
+)
+
+py_test(
+    name = "asyncio_utils_test",
+    srcs = ["asyncio_utils_test.py"],
+    deps = [":asyncio_utils"],
+)
+
+py_library(
+    name = "composite",
+    srcs = ["composite.py"],
+)
+
+py_test(
+    name = "composite_test",
+    srcs = ["composite_test.py"],
+    deps = [":composite"],
+)

checkpoint/orbax/checkpoint/_src/logging/BUILD

@@ -0,0 +1,55 @@
+package(
+    default_applicable_licenses = ["//:package_license"],
+    default_visibility = ["//visibility:public"],
+)
+
+py_library(
+    name = "abstract_logger",
+    srcs = ["abstract_logger.py"],
+)
+
+py_library(
+    name = "step_statistics",
+    srcs = ["step_statistics.py"],
+)
+
+py_library(
+    name = "cloud_logger",
+    srcs = ["cloud_logger.py"],
+    deps = [":abstract_logger"],
+)
+
+py_test(
+    name = "cloud_logger_test",
+    srcs = ["cloud_logger_test.py"],
+    deps = [":cloud_logger"],
+)
+
+py_library(
+    name = "standard_logger",
+    srcs = ["standard_logger.py"],
+    deps = [":abstract_logger"],
+)
+
+py_test(
+    name = "standard_logger_test",
+    srcs = ["standard_logger_test.py"],
+    deps = [":standard_logger"],
+)
+
+py_library(
+    name = "composite_logger",
+    srcs = ["composite_logger.py"],
+    deps = [":abstract_logger"],
+)
+
+py_test(
+    name = "composite_logger_test",
+    srcs = ["composite_logger_test.py"],
+    deps = [
+        ":abstract_logger",
+        ":cloud_logger",
+        ":composite_logger",
+        ":standard_logger",
+    ],
+)

checkpoint/orbax/checkpoint/_src/metadata/BUILD

@@ -62,6 +62,7 @@ py_test(
 py_library(
     name = "checkpoint",
     srcs = ["checkpoint.py"],
+    deps = ["//checkpoint/orbax/checkpoint/_src/logging:step_statistics"],
 )
 
 py_test(
@@ -71,6 +72,7 @@ py_test(
         ":checkpoint",
         ":root_metadata_serialization",
         ":step_metadata_serialization",
+        "//checkpoint/orbax/checkpoint/_src/logging:step_statistics",
     ],
 )
 
@@ -89,6 +91,7 @@ py_library(
     deps = [
         ":checkpoint",
         ":metadata_serialization_utils",
+        "//checkpoint/orbax/checkpoint/_src/logging:step_statistics",
     ],
 )
 

requirements.txt

@@ -20,6 +20,7 @@ etils[epath]
 simplejson
 chex
 optax
+mock
 nest_asyncio
 tensorstore
 humanize