From 9bc4c858c6e3a09227846964bd0950ba361eded9 Mon Sep 17 00:00:00 2001 From: Michael Reneer Date: Thu, 4 Jan 2024 16:24:01 -0800 Subject: [PATCH] Deptecate the `type_signature` parameter for the `tff.program.ReleaseManager.release` method. PiperOrigin-RevId: 595831028 --- .../python/program/file_release_manager.py | 14 +++++---- .../python/program/logging_release_manager.py | 6 ++-- .../python/program/memory_release_manager.py | 7 +++-- .../python/program/release_manager.py | 30 ++++++++++++------- .../program/tensorboard_release_manager.py | 8 +++-- 5 files changed, 43 insertions(+), 22 deletions(-) diff --git a/tensorflow_federated/python/program/file_release_manager.py b/tensorflow_federated/python/program/file_release_manager.py index 4a5be3fae2..9600d36880 100644 --- a/tensorflow_federated/python/program/file_release_manager.py +++ b/tensorflow_federated/python/program/file_release_manager.py @@ -28,7 +28,7 @@ import os import os.path import random -from typing import Union +from typing import Optional, Union import numpy as np import tensorflow as tf @@ -275,11 +275,13 @@ async def _remove_values_greater_than_key(self, key: int) -> None: ) self._latest_key = key + # TODO: b/305743962 - Deprecate `type_signature` and temporarily give `key` a + # default value. async def release( self, value: release_manager.ReleasableStructure, - type_signature: computation_types.Type, - key: int, + type_signature: Optional[computation_types.Type] = None, + key: int = 0, ) -> None: """Releases `value` from a federated program. @@ -380,11 +382,13 @@ def _get_path_for_key(self, key: release_manager.Key) -> str: basename = f'{self._prefix}{str(key)}' return os.path.join(self._root_dir, basename) + # TODO: b/305743962 - Deprecate `type_signature` and temporarily give `key` a + # default value. async def release( self, value: release_manager.ReleasableStructure, - type_signature: computation_types.Type, - key: release_manager.Key, + type_signature: Optional[computation_types.Type] = None, + key: release_manager.Key = None, ) -> None: """Releases `value` from a federated program. diff --git a/tensorflow_federated/python/program/logging_release_manager.py b/tensorflow_federated/python/program/logging_release_manager.py index a25ecb9b32..2a8c144b51 100644 --- a/tensorflow_federated/python/program/logging_release_manager.py +++ b/tensorflow_federated/python/program/logging_release_manager.py @@ -39,11 +39,13 @@ class LoggingReleaseManager( containing value references, each value reference is materialized. """ + # TODO: b/305743962 - Deprecate `type_signature` and temporarily give `key` a + # default value. async def release( self, value: release_manager.ReleasableStructure, - type_signature: computation_types.Type, - key: Optional[release_manager.Key], + type_signature: Optional[computation_types.Type] = None, + key: Optional[release_manager.Key] = None, ) -> None: """Releases `value` from a federated program. diff --git a/tensorflow_federated/python/program/memory_release_manager.py b/tensorflow_federated/python/program/memory_release_manager.py index c0e62fe9cc..ddb6c57d43 100644 --- a/tensorflow_federated/python/program/memory_release_manager.py +++ b/tensorflow_federated/python/program/memory_release_manager.py @@ -15,6 +15,7 @@ import collections from collections.abc import Hashable +from typing import Optional from tensorflow_federated.python.common_libs import py_typecheck from tensorflow_federated.python.core.impl.types import computation_types @@ -42,11 +43,13 @@ def __init__(self): """Returns an initialized `tff.program.MemoryReleaseManager`.""" self._values = collections.OrderedDict() + # TODO: b/305743962 - Deprecate `type_signature` and temporarily give `key` a + # default value. async def release( self, value: release_manager.ReleasableStructure, - type_signature: computation_types.Type, - key: Hashable, + type_signature: Optional[computation_types.Type] = None, + key: Hashable = None, ) -> None: """Releases `value` from a federated program. diff --git a/tensorflow_federated/python/program/release_manager.py b/tensorflow_federated/python/program/release_manager.py index 32e7f78a90..eb6cd8467b 100644 --- a/tensorflow_federated/python/program/release_manager.py +++ b/tensorflow_federated/python/program/release_manager.py @@ -57,12 +57,14 @@ class ReleaseManager(abc.ABC, Generic[ReleasableStructure, Key]): to customer storage in a federated program. """ + # TODO: b/305743962 - Deprecate `type_signature` and temporarily give `key` a + # default value. @abc.abstractmethod async def release( self, value: ReleasableStructure, - type_signature: computation_types.Type, - key: Key, + type_signature: Optional[computation_types.Type] = None, + key: Key = None, ) -> None: """Releases `value` from a federated program. @@ -163,11 +165,13 @@ def __init__( self._release_manager = release_manager self._filter_fn = filter_fn + # TODO: b/305743962 - Deprecate `type_signature` and temporarily give `key` a + # default value. async def release( self, value: ReleasableStructure, - type_signature: computation_types.Type, - key: Key, + type_signature: Optional[computation_types.Type] = None, + key: Key = None, ) -> None: """Releases `value` from a federated program. @@ -376,11 +380,13 @@ def __init__( self._release_managers = release_managers + # TODO: b/305743962 - Deprecate `type_signature` and temporarily give `key` a + # default value. async def release( self, value: ReleasableStructure, - type_signature: computation_types.Type, - key: Key, + type_signature: Optional[computation_types.Type] = None, + key: Key = None, ) -> None: """Releases `value` from a federated program. @@ -450,11 +456,13 @@ def __init__( f'Unexpected `periodicity` found: {type(periodicity)}.' ) + # TODO: b/305743962 - Deprecate `type_signature` and temporarily give `key` a + # default value. async def release( self, value: ReleasableStructure, - type_signature: computation_types.Type, - key: Key, + type_signature: Optional[computation_types.Type] = None, + key: Key = None, ) -> None: """Releases `value` from a federated program. @@ -523,11 +531,13 @@ def __init__( self._count = 0 self._delay = delay + # TODO: b/305743962 - Deprecate `type_signature` and temporarily give `key` a + # default value. async def release( self, value: ReleasableStructure, - type_signature: computation_types.Type, - key: Key, + type_signature: Optional[computation_types.Type] = None, + key: Key = None, ) -> None: """Releases `value` from a federated program. diff --git a/tensorflow_federated/python/program/tensorboard_release_manager.py b/tensorflow_federated/python/program/tensorboard_release_manager.py index 1b3068cb7f..59b6d8edca 100644 --- a/tensorflow_federated/python/program/tensorboard_release_manager.py +++ b/tensorflow_federated/python/program/tensorboard_release_manager.py @@ -14,7 +14,7 @@ """Utilities for releasing values from a federated program to TensorBoard.""" import os -from typing import Union +from typing import Optional, Union import numpy as np import tensorflow as tf @@ -69,11 +69,13 @@ def __init__(self, summary_dir: Union[str, os.PathLike[str]]): summary_dir = os.fspath(summary_dir) self._summary_writer = tf.summary.create_file_writer(summary_dir) + # TODO: b/305743962 - Deprecate `type_signature` and temporarily give `key` a + # default value. async def release( self, value: release_manager.ReleasableStructure, - type_signature: computation_types.Type, - key: int, + type_signature: Optional[computation_types.Type] = None, + key: int = 0, ) -> None: """Releases `value` from a federated program.