This CL updates the Snapshotter class to allow None as a valid va…

…lue for `snapshot_ttl_seconds`. This TTL value is passed to `paths.process_path` which already supports a `None` TTL.

PiperOrigin-RevId: 675680440
Change-Id: If44d8bcb01f53306660bdcd19a23b570b8f4a20e

acme/tf/savers.py

@@ -286,15 +286,16 @@ def __init__(
       *,
       directory: str = '~/acme/',
       time_delta_minutes: float = 30.0,
-      snapshot_ttl_seconds: int = _DEFAULT_SNAPSHOT_TTL,
+      snapshot_ttl_seconds: int | None = _DEFAULT_SNAPSHOT_TTL,
   ):
     """Builds the saver object.
 
     Args:
       objects_to_save: Mapping specifying what to snapshot.
       directory: Which directory to put the snapshot in.
       time_delta_minutes: How often to save the snapshot, in minutes.
-      snapshot_ttl_seconds: TTL (time to leave) in seconds for snapshots.
+      snapshot_ttl_seconds: TTL (time to live) in seconds for snapshots. If
+        `None`, then snapshots will be created in `directory` without a TTL.
     """
     objects_to_save = objects_to_save or {}
 

acme/tf/savers_test.py

@@ -201,6 +201,38 @@ def test_snapshot(self):
     assert np.allclose(outputs1, outputs2)
     assert all(tree.map_structure(np.allclose, list(grads1), list(grads2)))
 
+  def test_snapshot_no_ttl(self):
+    """Test that snapshotter correctly calls saves/restores snapshots w/o a TTL."""
+    # Create a test network.
+    net1 = networks.LayerNormMLP([10, 10])
+    spec = specs.Array([10], dtype=np.float32)
+    tf2_utils.create_variables(net1, [spec])
+
+    # Save the test network.
+    directory = self.get_tempdir()
+    objects_to_save = {'net': net1}
+    snapshotter = tf2_savers.Snapshotter(
+        objects_to_save, directory=directory, snapshot_ttl_seconds=None
+    )
+    snapshotter.save()
+
+    # Reload the test network.
+    net2 = tf.saved_model.load(os.path.join(snapshotter.directory, 'net'))
+    inputs = tf2_utils.add_batch_dim(tf2_utils.zeros_like(spec))
+
+    with tf.GradientTape() as tape:
+      outputs1 = net1(inputs)
+      loss1 = tf.math.reduce_sum(outputs1)
+      grads1 = tape.gradient(loss1, net1.trainable_variables)
+
+    with tf.GradientTape() as tape:
+      outputs2 = net2(inputs)
+      loss2 = tf.math.reduce_sum(outputs2)
+      grads2 = tape.gradient(loss2, net2.trainable_variables)
+
+    assert np.allclose(outputs1, outputs2)
+    assert all(tree.map_structure(np.allclose, list(grads1), list(grads2)))
+
   def test_snapshot_distribution(self):
     """Test that snapshotter correctly calls saves/restores snapshots."""
     # Create a test network.