Remove tff.program.X, use federated_language.program instead.

For each: * FederatedDataSource * FederatedDataSourceIterator * check_in_federated_context * ComputationArg * contains_only_server_placed_data * FederatedContext * LoggingReleaseManager * MemoryReleaseManager * ProgramStateExistsError * ProgramStateManager * ProgramStateNotFoundError * ProgramStateStructure * ProgramStateValue * DelayedReleaseManager * FilteringReleaseManager * GroupingReleaseManager * NotFilterableError * PeriodicReleaseManager * ReleasableStructure * ReleasableValue * ReleaseManager * MaterializableStructure * MaterializableTypeSignature * MaterializableValue * MaterializableValueReference * materialize_value * MaterializedStructure * MaterializedValue See: * google-parfait/federated-language@2ec477b * 0ae85a7 for more information. PiperOrigin-RevId: 705561178
google-parfait · Dec 12, 2024 · 8fa237b · 8fa237b
1 parent 913415c
commit 8fa237b
Show file tree

Hide file tree

Showing 24 changed files with 248 additions and 202 deletions.
diff --git a/RELEASE.md b/RELEASE.md
@@ -8,12 +8,6 @@ and this project adheres to
 
 ## Unreleased
 
-*   Updated `MeasuredProcessOutput` to be a `NamedTuple`.
-*   Added number of round retries to training metrics in
-    `tff.learning.programs.train_model`.
-*   Fix buffer overrun in `AggVectorIterator` when passing in an empty
-    `TensorData`.
-
 ### Added
 
 *   `tff.StructType.items()`, this API makes it easier to iterate over
@@ -55,6 +49,12 @@ and this project adheres to
     of the state of a `DPTensorAggregatorBundle`.
 *   `DPTensorAggregatorBundle::TakeOutputs` calls the inner aggregator's
     `ReportWithEpsilonAndDelta` methods and stitches the outputs together.
+*   Number of round retries to training metrics in
+    `tff.learning.programs.train_model`.
+
+### Fixed
+
+*   Buffer overrun in `AggVectorIterator` when passing in an empty `TensorData`.
 
 ### Changed
 
@@ -68,10 +68,40 @@ and this project adheres to
     check that the input is valid before passing to the aggregators it contains.
 *   Moved the tests of compatibility from `DPQuantileAggregator::MergeWith` to
     `DPQuantileAggregator::IsCompatible`.
+*   Updated `MeasuredProcessOutput` to be a `NamedTuple`.
 
 ### Removed
 
 *   `tff.types.tensorflow_to_type`, this function is no longer used.
+*   `tff.program.X`, use `federated_language.program` instead, for each:
+    *   `FederatedDataSource`
+    *   `FederatedDataSourceIterator`
+    *   `check_in_federated_context`
+    *   `ComputationArg`
+    *   `contains_only_server_placed_data`
+    *   `FederatedContext`
+    *   `LoggingReleaseManager`
+    *   `MemoryReleaseManager`
+    *   `ProgramStateExistsError`
+    *   `ProgramStateManager`
+    *   `ProgramStateNotFoundError`
+    *   `ProgramStateStructure`
+    *   `ProgramStateValue`
+    *   `DelayedReleaseManager`
+    *   `FilteringReleaseManager`
+    *   `GroupingReleaseManager`
+    *   `NotFilterableError`
+    *   `PeriodicReleaseManager`
+    *   `ReleasableStructure`
+    *   `ReleasableValue`
+    *   `ReleaseManager`
+    *   `MaterializableStructure`
+    *   `MaterializableTypeSignature`
+    *   `MaterializableValue`
+    *   `MaterializableValueReference`
+    *   `materialize_value`
+    *   `MaterializedStructure`
+    *   `MaterializedValue`
 
 ## Release 0.88.0
 

diff --git a/docs/program/federated_program.md b/docs/program/federated_program.md
@@ -85,8 +85,8 @@ def main() -> None:
 
   # Construct the platform-agnostic components.
   summary_dir = os.path.join(FLAGS.output_dir, 'summary')
-  metrics_manager = tff.program.GroupingReleaseManager([
-      tff.program.LoggingReleaseManager(),
+  metrics_manager = federated_language.program.GroupingReleaseManager([
+      federated_language.program.LoggingReleaseManager(),
       tff.program.TensorBoardReleaseManager(summary_dir),
   ])
   program_state_dir = os.path.join(..., 'program_state')
@@ -177,11 +177,11 @@ For example:
 async def program_logic(
     initialize: tff.Computation,
     train: tff.Computation,
-    data_source: tff.program.FederatedDataSource,
+    data_source: federated_language.program.FederatedDataSource,
     total_rounds: int,
     num_clients: int,
-    metrics_manager: tff.program.ReleaseManager[
-        tff.program.ReleasableStructure, int
+    metrics_manager: federated_language.program.ReleaseManager[
+        federated_language.program.ReleasableStructure, int
     ],
 ) -> None:
   state = initialize()

diff --git a/docs/program/guide.md b/docs/program/guide.md
@@ -27,7 +27,7 @@ program logic that has a type signature.
 ```python {.good}
 async def program_logic(
     train: tff.Computation,
-    data_source: tff.program.FederatedDataSource,
+    data_source: federated_language.program.FederatedDataSource,
     ...
 ) -> None:
   """Trains a federated model for some number of rounds.
@@ -53,14 +53,14 @@ the program logic that has a type signature.
 ```python {.good}
 def _check_program_logic_type_signatures(
     train: tff.Computation,
-    data_source: tff.program.FederatedDataSource,
+    data_source: federated_language.program.FederatedDataSource,
     ...
 ) -> None:
   ...
 
 async def program_logic(
     train: tff.Computation,
-    data_source: tff.program.FederatedDataSource,
+    data_source: federated_language.program.FederatedDataSource,
     ...
 ) -> None:
   _check_program_logic_type_signatures(
@@ -73,13 +73,13 @@ async def program_logic(
 ### Type Annotations
 
 **Do** provide a well defined Python type for each
-[`tff.program.ReleaseManager`](https://www.tensorflow.org/federated/api_docs/python/tff/program/ReleaseManager)
+[`federated_language.program.ReleaseManager`](https://www.tensorflow.org/federated/api_docs/python/tff/program/ReleaseManager)
 parameter supplied to the program logic.
 
 ```python {.good}
 async def program_logic(
     metrics_manager: Optional[
-        tff.program.ReleaseManager[tff.program.ReleasableStructure, int]
+        federated_language.program.ReleaseManager[federated_language.program.ReleasableStructure, int]
     ] = None,
     ...
 ) -> None:
@@ -98,7 +98,7 @@ async def program_logic(
 
 ```python {.bad}
 async def program_logic(
-    metrics_manager: Optional[tff.program.ReleaseManager] = None,
+    metrics_manager: Optional[federated_language.program.ReleaseManager] = None,
     ...
 ) -> None:
   ...
@@ -148,7 +148,7 @@ async def program_loic(...) -> None:
 
 ```python {.good}
 async def program_logic(
-    metrics_manager: Optional[tff.program.ReleaseManager] = None,
+    metrics_manager: Optional[federated_language.program.ReleaseManager] = None,
     ...
 ) -> None:
   """Trains a federated model for some number of rounds.

diff --git a/examples/learning/federated_program/vizier/BUILD b/examples/learning/federated_program/vizier/BUILD
@@ -3,7 +3,10 @@ load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test")
 py_library(
     name = "data_sources",
     srcs = ["data_sources.py"],
-    deps = ["//tensorflow_federated"],
+    deps = [
+        "//tensorflow_federated",
+        "@federated_language//federated_language",
+    ],
 )
 
 py_library(
@@ -23,6 +26,7 @@ py_binary(
         ":learning_process",
         ":vizier_service",
         "//tensorflow_federated",
+        "@federated_language//federated_language",
     ],
 )
 

diff --git a/examples/learning/federated_program/vizier/data_sources.py b/examples/learning/federated_program/vizier/data_sources.py
@@ -11,20 +11,19 @@
 # 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.
-"""The `tff.program.FederatedDataSource`s for this federated program."""
+"""The `federated_language.program.FederatedDataSource`s for this federated program."""
 
+import federated_language
 import tensorflow as tf
 import tensorflow_federated as tff
 
 
-def create_data_sources() -> (
-    tuple[
-        tff.program.FederatedDataSource,
-        tff.program.FederatedDataSource,
-        tf.TensorSpec,
-    ]
-):
-  """Creates the `tff.program.FederatedDataSource`s for this program.
+def create_data_sources() -> tuple[
+    federated_language.program.FederatedDataSource,
+    federated_language.program.FederatedDataSource,
+    tf.TensorSpec,
+]:
+  """Creates the `federated_language.program.FederatedDataSource`s for this program.
 
   Returns:
     A `tuple` containing the train data source, evaluation data source, and the

diff --git a/examples/learning/federated_program/vizier/program.py b/examples/learning/federated_program/vizier/program.py
@@ -32,6 +32,7 @@
 from absl import app
 from absl import flags
 from absl import logging
+import federated_language
 import tensorflow_federated as tff
 from vizier import pyvizier
 from vizier.client import client_abc
@@ -125,27 +126,33 @@ def _train_process_factory(
 
   def _model_output_manager_factory(
       trial: client_abc.TrialInterface,
-  ) -> tff.program.ReleaseManager[tff.program.ReleasableStructure, str]:
+  ) -> federated_language.program.ReleaseManager[
+      federated_language.program.ReleasableStructure, str
+  ]:
     del trial  # Unused.
-    return tff.program.LoggingReleaseManager()
+    return federated_language.program.LoggingReleaseManager()
 
   def _metrics_manager_factory(
       trial: client_abc.TrialInterface,
-  ) -> tff.program.ReleaseManager[tff.program.ReleasableStructure, int]:
+  ) -> federated_language.program.ReleaseManager[
+      federated_language.program.ReleasableStructure, int
+  ]:
     del trial  # Unused.
-    return tff.program.LoggingReleaseManager()
+    return federated_language.program.LoggingReleaseManager()
 
   def _program_state_manager_factory(
       trial: client_abc.TrialInterface,
-  ) -> tff.program.ProgramStateManager:
+  ) -> federated_language.program.ProgramStateManager:
     trial_name = f'trial_{trial.id}'
     root_dir = os.path.join(_OUTPUT_DIR.value, experiment_name, trial_name)
     return tff.program.FileProgramStateManager(root_dir)
 
   def _evaluation_manager_factory(
       trial: client_abc.TrialInterface,
   ) -> tff.learning.programs.EvaluationManager:
-    aggregated_metrics_manager = tff.program.LoggingReleaseManager()
+    aggregated_metrics_manager = (
+        federated_language.program.LoggingReleaseManager()
+    )
 
     def _create_state_manager_fn(
         name: str,
@@ -164,11 +171,13 @@ def _create_process_fn(
     ) -> tuple[
         tff.learning.templates.LearningProcess,
         Optional[
-            tff.program.ReleaseManager[tff.program.ReleasableStructure, int]
+            federated_language.program.ReleaseManager[
+                federated_language.program.ReleasableStructure, int
+            ]
         ],
     ]:
       del name  # Unused.
-      release_manager = tff.program.LoggingReleaseManager()
+      release_manager = federated_language.program.LoggingReleaseManager()
       return (evaluation_process, release_manager)
 
     return tff.learning.programs.EvaluationManager(

diff --git a/examples/program/BUILD b/examples/program/BUILD
@@ -16,7 +16,10 @@ py_library(
 py_library(
     name = "program_logic",
     srcs = ["program_logic.py"],
-    deps = ["//tensorflow_federated"],
+    deps = [
+        "//tensorflow_federated",
+        "@federated_language//federated_language",
+    ],
 )
 
 py_binary(
@@ -26,6 +29,7 @@ py_binary(
         ":computations",
         ":program_logic",
         "//tensorflow_federated",
+        "@federated_language//federated_language",
     ],
 )
 
@@ -36,5 +40,6 @@ py_test(
         ":computations",
         ":program_logic",
         "//tensorflow_federated",
+        "@federated_language//federated_language",
     ],
 )
diff --git a/examples/program/program.py b/examples/program/program.py
@@ -52,6 +52,7 @@
 
 from absl import app
 from absl import flags
+import federated_language
 import tensorflow as tf
 import tensorflow_federated as tff
 
@@ -97,9 +98,11 @@ def main(argv: Sequence[str]) -> None:
   # Configure the platform-agnostic components.
 
   # Create release managers with access to customer storage.
-  train_metrics_managers = [tff.program.LoggingReleaseManager()]
-  evaluation_metrics_managers = [tff.program.LoggingReleaseManager()]
-  model_output_manager = tff.program.LoggingReleaseManager()
+  train_metrics_managers = [federated_language.program.LoggingReleaseManager()]
+  evaluation_metrics_managers = [
+      federated_language.program.LoggingReleaseManager()
+  ]
+  model_output_manager = federated_language.program.LoggingReleaseManager()
 
   if _OUTPUT_DIR.value is not None:
     summary_dir = os.path.join(_OUTPUT_DIR.value, 'summary')
@@ -113,20 +116,24 @@ def main(argv: Sequence[str]) -> None:
   # Group the metrics release managers; program logic may accept a single
   # release manager to make the implementation of the program logic simpler and
   # easier to maintain, the program can use a
-  # `tff.program.GroupingReleaseManager` to release values to multiple
-  # destinations.
+  # `federated_language.program.GroupingReleaseManager` to release values to
+  # multiple destinations.
   #
   # Filter the metrics before they are released; the program can use a
-  # `tff.program.FilteringReleaseManager` to limit the values that are
-  # released by the program logic. If a formal privacy guarantee is not
+  # `federated_language.program.FilteringReleaseManager` to limit the values
+  # that are released by the program logic. If a formal privacy guarantee is not
   # required, it may be ok to release all the metrics.
-  train_metrics_manager = tff.program.FilteringReleaseManager(
-      tff.program.GroupingReleaseManager(train_metrics_managers),
+  train_metrics_manager = federated_language.program.FilteringReleaseManager(
+      federated_language.program.GroupingReleaseManager(train_metrics_managers),
       _filter_metrics,
   )
-  evaluation_metrics_manager = tff.program.FilteringReleaseManager(
-      tff.program.GroupingReleaseManager(evaluation_metrics_managers),
-      _filter_metrics,
+  evaluation_metrics_manager = (
+      federated_language.program.FilteringReleaseManager(
+          federated_language.program.GroupingReleaseManager(
+              evaluation_metrics_managers
+          ),
+          _filter_metrics,
+      )
   )
 
   # Create a program state manager with access to platform storage.