Release v0.2.0

GitOrigin-RevId: eb73a08756cf3c4ffcfabd079af92ad0b60304e2

.gitignore

@@ -27,8 +27,12 @@
 .idea/
 **/*.iml
 
-# python venv
+# python
+__pycache__/
 venv/
 
 # code coverage output
 genhtml/
+
+# python cache
+**/__pycache__/**

BUILDING.md

@@ -16,12 +16,13 @@ Images can be built using the provided script
 
 The expected output for images built within the provided [Dockerfile](Dockerfile):
 ```
-bazel-bin/shuffler/services/aggregator/aggregator_image/index.json:      "digest": "sha256:375a7525eb34a9cff17c0f21f572485a4f19318ede5dc85dd15fa8390800f203"
-bazel-bin/shuffler/services/collector/collector_image/index.json:      "digest": "sha256:aa51fa96b69612eda957d2931cd04a570cf6deb9dbdff0cf87ee673b43bc6cbe"
-bazel-bin/shuffler/services/modelupdater/model_updater_image/index.json:      "digest": "sha256:0c57e62a78cc112ecfa91fbb2eb6a16a484264cb65d2add500bcccbabae8e51c"
-bazel-bin/shuffler/services/taskassignment/task_assignment_image/index.json:      "digest": "sha256:bcc328634272c591a5a58f3a8406d8849ecb393fe8cf0cb62091338a99e3d772"
-bazel-bin/shuffler/services/taskmanagement/task_management_image/index.json:      "digest": "sha256:001060ec8a7fe40f761aa125474389021cd8404930739c8f8340240bfd21c221"
-bazel-bin/shuffler/services/taskscheduler/task_scheduler_image/index.json:      "digest": "sha256:3a8a7bd66646aac0e55fbc510291f32a1e4b9b03f3e64ea72fb9ceace985750e"
+bazel-bin/shuffler/services/aggregator/aggregator_image/index.json:      "digest": "sha256:c04b35b41fc2990e47f4c6c1e575f4dc06eadfef5add6498959977b012c008d4"
+bazel-bin/shuffler/services/collector/collector_image/index.json:      "digest": "sha256:4f6bef12c33bac516448d44e423ffc9a7a6d273fe9c2de327c45ec0bad74295a"
+bazel-bin/shuffler/services/modelupdater/model_updater_image/index.json:      "digest": "sha256:05f73cbec7dea193ae81e51c99a3c02de6cd0e62c27e267b9a43e31e6c8d530a"
+bazel-bin/shuffler/services/taskassignment/task_assignment_image/index.json:      "digest": "sha256:80439bfca81d86d69010540c26bb5e7371f7fe98d5538a1ac0813f6f8c69f35a"
+bazel-bin/shuffler/services/taskbuilder/task_builder_image/index.json:      "digest": "sha256:3959610c3a6896e0f77358d94a85900b088290951f0cad53557c6298a4d86bf2"
+bazel-bin/shuffler/services/taskmanagement/task_management_image/index.json:      "digest": "sha256:4f9239def6a6ee004aec5a55dbed47b255e83705250d749923d96b6d1e27b2b0"
+bazel-bin/shuffler/services/taskscheduler/task_scheduler_image/index.json:      "digest": "sha256:a15ef0b318a133c7e5d6112ae458fca0f344e2646bb2094323d6d27a8c721049"
 ```
 
 ## Publishing

WORKSPACE

@@ -84,6 +84,7 @@ maven_install(
         "com.google.crypto.tink:tink:1.12.0",
         "com.google.inject:guice:7.0.0",
         "org.slf4j:slf4j-api:2.0.7",
+        "com.google.cloud:google-cloud-logging-logback:0.131.8-alpha",
         # spring
         "org.springframework.boot:spring-boot-autoconfigure:3.1.1",
         "org.springframework.boot:spring-boot-starter-actuator:3.1.1",
@@ -145,10 +146,10 @@ http_archive(
 http_archive(
     name = "coordinator-services-and-shared-libraries",
     build_file = "coordinator.BUILD",
-    sha256 = "90e147c05801bbcd9b25491bea0239c17404f61cdff236edc3da6a522b59ecd0",
-    strip_prefix = "coordinator-services-and-shared-libraries-1.5.0",
+    sha256 = "7879ff1c7bf0f96a6214725fe87bf39ef5297db4aab8d39d7b6fef49c2643c55",
+    strip_prefix = "coordinator-services-and-shared-libraries-1.7.0",
     urls = [
-        "https://github.com/privacysandbox/coordinator-services-and-shared-libraries/archive/refs/tags/v1.5.0.tar.gz",
+        "https://github.com/privacysandbox/coordinator-services-and-shared-libraries/archive/refs/tags/v1.7.0.tar.gz",
     ],
 )
 
@@ -169,6 +170,29 @@ http_archive(
     url = "https://github.com/project-oak/oak/archive/0acf3f6dc0af2035d40884fe1258b1e0e7db5488.tar.gz",
 )
 
+# The following enables the use of the library functions in the differential-
+# privacy github repo
+http_archive(
+    name = "com_google_cc_differential_privacy",
+    sha256 = "6e6e1cd7a819695caae408f4fa938129ab7a86e83fe2410137c85e50131abbe0",
+    strip_prefix = "differential-privacy-3.0.0/cc",
+    url = "https://github.com/google/differential-privacy/archive/refs/tags/v3.0.0.tar.gz",
+)
+
+http_archive(
+    name = "com_google_differential_privacy",
+    sha256 = "6e6e1cd7a819695caae408f4fa938129ab7a86e83fe2410137c85e50131abbe0",
+    strip_prefix = "differential-privacy-3.0.0",
+    url = "https://github.com/google/differential-privacy/archive/refs/tags/v3.0.0.tar.gz",
+)
+
+http_archive(
+    name = "org_tensorflow_federated",
+    sha256 = "2bb1e641a84f05bc1776eefc98eb108454807071ae0bed45ee5e8bee896511be",
+    strip_prefix = "tensorflow-federated-9c5a51af41fccc8e720a63aa754465767113877d",
+    url = "https://github.com/tensorflow/federated/archive/9c5a51af41fccc8e720a63aa754465767113877d.tar.gz",
+)
+
 # Tensorflow v2.14.0
 http_archive(
     name = "org_tensorflow",
@@ -208,12 +232,10 @@ http_archive(
     url = "https://github.com/bazelbuild/rules_python/releases/download/0.29.0/rules_python-0.29.0.tar.gz",
 )
 
-load("@rules_python//python:repositories.bzl", "py_repositories")
+load("@rules_python//python:repositories.bzl", "py_repositories", "python_register_toolchains")
 
 py_repositories()
 
-load("@rules_python//python:repositories.bzl", "python_register_toolchains")
-
 python_register_toolchains(
     name = "python",
     ignore_root_user_error = True,
@@ -243,20 +265,20 @@ load("@rules_proto//proto:repositories.bzl", "rules_proto_dependencies", "rules_
 
 rules_proto_toolchains()
 
+load("@federatedcompute//fcp/tensorflow/pip_tf:defs.bzl", "TF_ADDITIVE_BUILD_CONTENT")
 load("@python//:defs.bzl", "interpreter")
 load("@rules_python//python:pip.bzl", "package_annotation", "pip_parse")
-load("@federatedcompute//fcp/tensorflow/pip_tf:defs.bzl", "TF_ADDITIVE_BUILD_CONTENT")
 
 http_archive(
     name = "differential_privacy",
+    repo_mapping = {
+        "@accounting_py_pip_deps": "@pypi",
+    },
     sha256 = "6e6e1cd7a819695caae408f4fa938129ab7a86e83fe2410137c85e50131abbe0",
     strip_prefix = "differential-privacy-3.0.0/python/dp_accounting/",
     urls = [
         "https://github.com/google/differential-privacy/archive/refs/tags/v3.0.0.tar.gz",
     ],
-    repo_mapping = {
-        "@accounting_py_pip_deps": "@pypi"
-    },
 )
 
 pip_parse(
@@ -458,3 +480,23 @@ llvm_toolchain(
 load("@llvm_toolchain//:toolchains.bzl", "llvm_register_toolchains")
 
 llvm_register_toolchains()
+
+
+################
+# pybind11
+################
+
+http_archive(
+  name = "pybind11_bazel",
+  strip_prefix = "pybind11_bazel-2.12.0",
+  sha256 = "a58c25c5fe063a70057fa20cb8e15f3bda19b1030305bcb533af1e45f36a4a55",
+  urls = ["https://github.com/pybind/pybind11_bazel/archive/pybind11_bazel-2.12.0.zip"],
+)
+# We still require the pybind library.
+http_archive(
+  name = "pybind11",
+  build_file = "@pybind11_bazel//:pybind11-BUILD.bazel",
+  strip_prefix = "pybind11-2.12.0",
+  sha256 = "411f77380c43798506b39ec594fc7f2b532a13c4db674fcf2b1ca344efaefb68",
+  urls = ["https://github.com/pybind/pybind11/archive/pybind11-2.12.0.zip"],
+)

java/src/it/java/com/google/ondevicepersonalization/federatedcompute/endtoendtests/BUILD

@@ -33,6 +33,7 @@ java_binary(
         ":end_to_end_args",
         "//java/src/it/java/com/google/ondevicepersonalization/federatedcompute/endtoendtests/clients:device",
         "//java/src/it/java/com/google/ondevicepersonalization/federatedcompute/endtoendtests/clients:partner",
+        "//java/src/main/java/com/google/ondevicepersonalization/federatedcompute/shuffler/common:compression_utils",
         "//shuffler/proto:task_java_proto",
         "@maven//:com_beust_jcommander",
         "@maven//:com_google_guava_guava",

java/src/it/java/com/google/ondevicepersonalization/federatedcompute/endtoendtests/EndToEndArgs.java

@@ -1,3 +1,19 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * 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.
+ */
+
 package com.google.ondevicepersonalization.federatedcompute.endtoendtests;
 
 import com.beust.jcommander.Parameter;
@@ -54,6 +70,23 @@ public class EndToEndArgs {
   @Parameter(names = "--population_name", description = "Population name to use for the task")
   private String populationName;
 
+  @Parameter(
+      names = "--eval_training_population_name",
+      description = "Training population name to use by the eval task")
+  private String evalTrainingPopulationName;
+
+  @Parameter(
+      names = "--eval_training_task_id",
+      description = "Training task id used by the eval task",
+      validateWith = PositiveInteger.class)
+  private int evalTrainingTaskId = 1;
+
+  @Parameter(
+      names = "--eval_evey_k_iteration",
+      description = "K value for EveryKIteration selector",
+      validateWith = PositiveInteger.class)
+  private int evalEveryKIteration = 1;
+
   @Parameter(
       names = "--min_client_version",
       description = "Minimum client version to use for the task")
@@ -118,6 +151,7 @@ enum Operation {
     CREATE_AND_COMPLETE_TRAINING_TASK,
     CREATE_AND_COMPLETE_EVALUATION_TASK,
     CREATE_TASK,
+    CREATE_EVALUATION_TASK,
     RUN_TASK,
     GET_TASK
   }

java/src/it/java/com/google/ondevicepersonalization/federatedcompute/endtoendtests/EndToEndTest.java

@@ -31,6 +31,7 @@
 import com.google.ondevicepersonalization.federatedcompute.proto.TaskInfo;
 import com.google.ondevicepersonalization.federatedcompute.proto.TaskStatus;
 import com.google.ondevicepersonalization.federatedcompute.proto.TrainingInfo;
+import com.google.ondevicepersonalization.federatedcompute.shuffler.common.CompressionUtils;
 import java.time.InstantSource;
 import java.util.ArrayList;
 import java.util.List;
@@ -127,26 +128,43 @@ public static void main(String[] args) {
           endToEndTest.completeTask(endToEndTest.createTrainingTask());
           break;
         case CREATE_AND_COMPLETE_EVALUATION_TASK:
-          // create training and evaluation task
-          Task trainingTask = endToEndTest.createTrainingTask();
-          EvaluationInfo evaluationInfo =
-              EvaluationInfo.newBuilder()
-                  .setTrainingPopulationName(trainingTask.getPopulationName())
-                  .setTrainingTaskId(trainingTask.getTaskId())
-                  .setCheckPointSelector(
-                      CheckPointSelector.newBuilder()
-                          .setIterationSelector(
-                              EveryKIterationsCheckpointSelector.newBuilder().setSize(1)))
-                  .build();
-          // create an evaluation task under same population with training task
-          Task evaluationTask =
-              endToEndTest.createEvaluationTask(trainingTask.getPopulationName(), evaluationInfo);
-          endToEndTest.completeTasksInSamePopulation(
-              ImmutableList.of(trainingTask, evaluationTask));
-          break;
+          {
+            // create training and evaluation task
+            Task trainingTask = endToEndTest.createTrainingTask();
+            EvaluationInfo evaluationInfo =
+                EvaluationInfo.newBuilder()
+                    .setTrainingPopulationName(trainingTask.getPopulationName())
+                    .setTrainingTaskId(trainingTask.getTaskId())
+                    .setCheckPointSelector(
+                        CheckPointSelector.newBuilder()
+                            .setIterationSelector(
+                                EveryKIterationsCheckpointSelector.newBuilder().setSize(1)))
+                    .build();
+            // create an evaluation task under same population with training task
+            Task evaluationTask =
+                endToEndTest.createEvaluationTask(trainingTask.getPopulationName(), evaluationInfo);
+            endToEndTest.completeTasksInSamePopulation(
+                ImmutableList.of(trainingTask, evaluationTask));
+            break;
+          }
         case CREATE_TASK:
           endToEndTest.createTrainingTask();
           break;
+        case CREATE_EVALUATION_TASK:
+          {
+            EvaluationInfo evaluationInfo =
+                EvaluationInfo.newBuilder()
+                    .setTrainingPopulationName(endToEndArgs.getEvalTrainingPopulationName())
+                    .setTrainingTaskId(endToEndArgs.getEvalTrainingTaskId())
+                    .setCheckPointSelector(
+                        CheckPointSelector.newBuilder()
+                            .setIterationSelector(
+                                EveryKIterationsCheckpointSelector.newBuilder()
+                                    .setSize(endToEndArgs.getEvalEveryKIteration())))
+                    .build();
+            endToEndTest.createEvaluationTask(endToEndArgs.getPopulationName(), evaluationInfo);
+            break;
+          }
         case RUN_TASK:
           endToEndTest.runTask(endToEndArgs.getRunCount());
           break;
@@ -184,6 +202,13 @@ private static void validateInput(EndToEndArgs args) {
           Objects.requireNonNull(args.getInitCheckpoint());
         }
         break;
+      case CREATE_EVALUATION_TASK:
+        if (!args.isUseLocalResources()) {
+          Objects.requireNonNull(args.getClientPlan());
+          Objects.requireNonNull(args.getServerPlan());
+        }
+        Objects.requireNonNull(args.getEvalTrainingPopulationName());
+        break;
       case RUN_TASK:
         if (args.isEncrypt()) {
           Objects.requireNonNull(args.getPublicKeyUrl());
@@ -441,6 +466,6 @@ private void processTaskAssignment() throws Exception {
     if (encrypt) {
       downloadedGradient = device.encryptAndCompressGradient(downloadedGradient, publicKeyUrl);
     }
-    device.submitResult(downloadedGradient);
+    device.submitResult(CompressionUtils.compressWithGzip(downloadedGradient));
   }
 }

java/src/it/java/com/google/ondevicepersonalization/federatedcompute/endtoendtests/README.md

@@ -22,6 +22,12 @@ Creates a task
 bazel run //java/src/it/java/com/google/ondevicepersonalization/federatedcompute/endtoendtests:end_to_end_test -- --task_management_server <tm_server_url> --operation CREATE_TASK
 ```
 
+### CREATE_EVALUATION_TASK
+Creates a task
+```
+bazel run //java/src/it/java/com/google/ondevicepersonalization/federatedcompute/endtoendtests:end_to_end_test -- --task_management_server <tm_server_url> --operation CREATE_EVALUATION_TASK --eval_training_population_name <training_population_name> --eval_training_task_id <training_task_id>
+```
+
 ### RUN_TASK
 Runs `run_count` contributions against a given task.
 ```

java/src/it/java/com/google/ondevicepersonalization/federatedcompute/endtoendtests/clients/Device.java

@@ -168,6 +168,7 @@ public void submitResult(byte[] gradients) throws Exception {
     }
 
     String assignmentUri = taskAssignment.getSelfUri();
+    System.out.println("submitResult: " + assignmentUri);
     upload(assignmentUri, content);
   }
 
@@ -201,6 +202,7 @@ private void upload(String url, byte[] content) throws Exception {
             .PUT(HttpRequest.BodyPublishers.ofByteArray(content))
             .timeout(Duration.ofSeconds(10))
             .header("Content-Type", "application/octet-stream")
+            .header("content-encoding", "gzip")
             .build();
 
     HttpResponse<byte[]> response = client.send(request, BodyHandlers.ofByteArray());

java/src/main/java/com/google/ondevicepersonalization/federatedcompute/shuffler/aggregator/BUILD

@@ -26,8 +26,9 @@ springboot(
 java_library(
     name = "aggregator",
     srcs = ["AggregatorApplication.java"],
-    resources = ["//java/src/main/resources:aggregator_properties"],
+    resources = ["//java/src/main/resources:aggregator_properties", "//java/src/main/resources:logback_spring"],
     deps = [
+        "//java/src/main/java/com/google/ondevicepersonalization/federatedcompute/shuffler/aggregator/controllers:aggregator_controller",
         "//java/src/main/java/com/google/ondevicepersonalization/federatedcompute/shuffler/aggregator/core:aggregator_core_impl",
         "//java/src/main/java/com/google/ondevicepersonalization/federatedcompute/shuffler/aggregator/core/gcp:aggregator_config",
         "//java/src/main/java/com/google/ondevicepersonalization/federatedcompute/shuffler/aggregator/core/gcp:pub_sub_message_receiver",
@@ -39,6 +40,7 @@ java_library(
         "//java/src/main/java/com/google/ondevicepersonalization/federatedcompute/shuffler/common/dao/gcp:gcs_blob_dao",
         "//java/src/main/java/com/google/ondevicepersonalization/federatedcompute/shuffler/common/messaging/gcp:pub_sub_subscriber_config",
         "@maven//:io_micrometer_micrometer_registry_prometheus",
+        "@maven//:com_google_cloud_google_cloud_logging_logback",
         "@maven//:org_springframework_boot_spring_boot",
         "@maven//:org_springframework_boot_spring_boot_autoconfigure",
         "@maven//:org_springframework_boot_spring_boot_loader",

java/src/main/java/com/google/ondevicepersonalization/federatedcompute/shuffler/aggregator/controllers/AggregatorController.java

@@ -0,0 +1,36 @@
+// Copyright 2024 Google LLC
+//
+// 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.
+//
+////////////////////////////////////////////////////////////////////////////////
+
+package com.google.ondevicepersonalization.federatedcompute.shuffler.aggregator.controllers;
+
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RestController;
+
+/** Aggregator controller. */
+@RestController
+public final class AggregatorController {
+  @GetMapping("/ready")
+  public String ready() {
+    // TODO(291594777): Implement readiness check.
+    return "Greetings from Aggregator Spring Boot! Ready check. \n";
+  }
+
+  @GetMapping("/healthz")
+  public String healthz() {
+    // TODO(291594777): Implement health check.
+    return "Greetings from Aggregator Spring Boot! Health check. \n";
+  }
+}