xref: /aosp_15_r20/external/federated-compute/fcp/artifact_building/BUILD (revision 14675a029014e728ec732f129a32e299b2da0601) 
1load("@rules_python//python:defs.bzl", "py_library", "py_test")
2
3default_visibility = ["//fcp:internal"]
4
5py_library(
6    name = "artifact_constants",
7    srcs = ["artifact_constants.py"],
8    srcs_version = "PY3",
9    visibility = default_visibility + [
10    ],
11)
12
13py_library(
14    name = "checkpoint_type",
15    srcs = ["checkpoint_type.py"],
16    srcs_version = "PY3",
17    visibility = default_visibility + [
18    ],
19)
20
21py_library(
22    name = "checkpoint_utils",
23    srcs = ["checkpoint_utils.py"],
24    srcs_version = "PY3",
25    visibility = default_visibility + [
26    ],
27    deps = [
28        ":artifact_constants",
29        ":tensor_utils",
30        ":type_checks",
31        ":variable_helpers",
32        "//fcp/protos:plan_py_pb2",
33    ],
34)
35
36py_test(
37    name = "checkpoint_utils_test",
38    srcs = [
39        "checkpoint_utils_test.py",
40    ],
41    python_version = "PY3",
42    srcs_version = "PY3",
43    deps = [
44        ":checkpoint_utils",
45        "//fcp/protos:plan_py_pb2",
46        "@com_google_protobuf//:protobuf_python",
47    ],
48)
49
50py_library(
51    name = "data_spec",
52    srcs = ["data_spec.py"],
53    srcs_version = "PY3",
54    visibility = default_visibility + [
55    ],
56    deps = [
57        ":type_checks",
58        "//fcp/protos:plan_py_pb2",
59    ],
60)
61
62py_test(
63    name = "data_spec_test",
64    srcs = [
65        "data_spec_test.py",
66    ],
67    python_version = "PY3",
68    srcs_version = "PY3",
69    deps = [
70        ":data_spec",
71        "//fcp/protos:plan_py_pb2",
72    ],
73)
74
75py_library(
76    name = "federated_compute_plan_builder",
77    srcs = ["federated_compute_plan_builder.py"],
78    srcs_version = "PY3",
79    visibility = default_visibility + [
80    ],
81    deps = [
82        ":artifact_constants",
83        ":checkpoint_type",
84        ":checkpoint_utils",
85        ":data_spec",
86        ":graph_helpers",
87        ":proto_helpers",
88        ":tensor_utils",
89        ":type_checks",
90        ":variable_helpers",
91        "//fcp/protos:plan_py_pb2",
92        "//fcp/tensorflow:append_slices_py",
93        "//fcp/tensorflow:delete_file_py",
94    ],
95)
96
97py_library(
98    name = "graph_helpers",
99    srcs = ["graph_helpers.py"],
100    srcs_version = "PY3",
101    visibility = default_visibility + [
102    ],
103    deps = [
104        ":data_spec",
105        ":tensor_utils",
106        ":type_checks",
107        "//fcp/tensorflow:external_dataset_py",
108    ],
109)
110
111py_test(
112    name = "graph_helpers_test",
113    size = "small",
114    srcs = ["graph_helpers_test.py"],
115    python_version = "PY3",
116    srcs_version = "PY3",
117    deps = [
118        ":data_spec",
119        ":graph_helpers",
120        ":variable_helpers",
121        "//fcp/protos:plan_py_pb2",
122    ],
123)
124
125py_library(
126    name = "plan_utils",
127    srcs = ["plan_utils.py"],
128    srcs_version = "PY3",
129    visibility = default_visibility + [
130    ],
131    deps = [
132        ":tensor_utils",
133        "//fcp/protos:plan_py_pb2",
134    ],
135)
136
137py_test(
138    name = "plan_utils_test",
139    srcs = [
140        "plan_utils_test.py",
141    ],
142    python_version = "PY3",
143    srcs_version = "PY3",
144    deps = [
145        ":checkpoint_utils",
146        ":plan_utils",
147        ":test_utils",
148        "//fcp/protos:plan_py_pb2",
149    ],
150)
151
152py_library(
153    name = "proto_helpers",
154    srcs = ["proto_helpers.py"],
155    srcs_version = "PY3",
156    visibility = default_visibility + [
157    ],
158    deps = [
159        ":tensor_utils",
160        ":type_checks",
161        "//fcp/protos:plan_py_pb2",
162    ],
163)
164
165py_test(
166    name = "proto_helpers_test",
167    size = "small",
168    srcs = ["proto_helpers_test.py"],
169    python_version = "PY3",
170    srcs_version = "PY3",
171    deps = [
172        ":proto_helpers",
173        ":variable_helpers",
174    ],
175)
176
177py_library(
178    name = "tensor_utils",
179    srcs = ["tensor_utils.py"],
180    srcs_version = "PY3",
181    visibility = default_visibility + [
182    ],
183)
184
185py_test(
186    name = "tensor_utils_test",
187    srcs = [
188        "tensor_utils_test.py",
189    ],
190    python_version = "PY3",
191    srcs_version = "PY3",
192    deps = [
193        ":tensor_utils",
194        "@com_google_protobuf//:protobuf_python",
195    ],
196)
197
198py_library(
199    name = "test_utils",
200    srcs = ["test_utils.py"],
201    srcs_version = "PY3",
202    visibility = default_visibility + [
203    ],
204    deps = ["//fcp/protos:plan_py_pb2"],
205)
206
207py_test(
208    name = "test_utils_test",
209    srcs = [
210        "test_utils_test.py",
211    ],
212    python_version = "PY3",
213    srcs_version = "PY3",
214    deps = [
215        ":checkpoint_utils",
216        ":test_utils",
217        "//fcp/protos:plan_py_pb2",
218    ],
219)
220
221py_library(
222    name = "type_checks",
223    srcs = ["type_checks.py"],
224    srcs_version = "PY3",
225    visibility = default_visibility + [
226    ],
227)
228
229py_test(
230    name = "type_checks_test",
231    srcs = [
232        "type_checks_test.py",
233    ],
234    python_version = "PY3",
235    srcs_version = "PY3",
236    deps = [":type_checks"],
237)
238
239py_library(
240    name = "variable_helpers",
241    srcs = ["variable_helpers.py"],
242    srcs_version = "PY3",
243    visibility = default_visibility + [
244    ],
245    deps = [
246        ":tensor_utils",
247        ":type_checks",
248    ],
249)
250
251py_test(
252    name = "variable_helpers_test",
253    size = "small",
254    srcs = ["variable_helpers_test.py"],
255    python_version = "PY3",
256    srcs_version = "PY3",
257    deps = [
258        ":artifact_constants",
259        ":variable_helpers",
260    ],
261)
262