elements 3.18.0.dev1__tar.gz → 3.18.2__tar.gz

{elements-3.18.0.dev1/elements.egg-info → elements-3.18.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: elements
-Version: 3.18.0.dev1
+Version: 3.18.2
 Summary: Building blocks for productive research.
 Home-page: http://github.com/danijar/elements
 Classifier: Intended Audience :: Science/Research

{elements-3.18.0.dev1 → elements-3.18.2}/elements/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '3.18.0.dev1'
+__version__ = '3.18.2'
 
 from .agg import Agg
 from .checkpoint import Checkpoint, Saveable
@@ -17,6 +17,7 @@ from .usage import Usage
 from .utils import timestamp
 from .uuid import UUID
 
+from . import checkpoint
 from . import logger
 from . import plotting
 from . import timer

{elements-3.18.0.dev1 → elements-3.18.2}/elements/checkpoint.py

@@ -92,10 +92,10 @@ class Checkpoint:
   def save(self, path=None, keys=None):
     assert self._directory or path
     if keys is None:
-      saveables = self._saveables
+      savefns = {k: v.save for k, v in self._saveables.items()}
     else:
       assert all([not k.startswith('_') for k in keys]), keys
-      saveables = {k: self._saveables[k] for k in keys}
+      savefns = {k: self._saveables[k].save for k in keys}
     if path:
       folder = None
     else:
@@ -104,8 +104,8 @@ class Checkpoint:
         folder += f'-{int(self._step):012d}'
       path = self._directory / folder
     printing.print_(f'Saving checkpoint: {path}')
-    save(path, saveables, self._write)
-    if folder:
+    save(path, savefns, self._write)
+    if folder and self._write:
       (self._directory / 'latest').write_text(folder)
       self._cleanup()
     print('Saved checkpoint.')
@@ -114,15 +114,15 @@ class Checkpoint:
   def load(self, path=None, keys=None):
     assert self._directory or path
     if keys is None:
-      saveables = self._saveables
+      loadfns = {k: v.load for k, v in self._saveables.items()}
     else:
       assert all([not k.startswith('_') for k in keys]), keys
-      saveables = {k: self._saveables[k] for k in keys}
+      loadfns = {k: self._saveables[k].load for k in keys}
     if not path:
       path = self.latest()
       assert path
     printing.print_(f'Loading checkpoint: {path}')
-    load(path, saveables)
+    load(path, loadfns)
     print('Loaded checkpoint.')
 
   def load_or_save(self):
@@ -152,40 +152,44 @@ def exists(path):
   return (path / 'done').exists()
 
 
-def save(path, saveables, write=True):
+def save(path, savefns, write=True):
   path = pathlib.Path(path)
   assert not exists(path), path
-  path.mkdir(parents=True)
-  for name, saveable in saveables.items():
+  write and path.mkdir(parents=True)
+  for name, savefn in savefns.items():
     try:
-      data = saveable.save()
+      data = savefn()
       if inspect.isgenerator(data):
         for i, shard in enumerate(data):
-          assert i < 1e7, i
+          assert i < 1e5, i
           if write:  # Iterate even if we're not writing.
-            buffer = pickle.dumps(shard)
-            (path / f'{name}-{i:06d}.pkl').write_bytes(buffer)
+            with timer.section('checkpoint_pickle'):
+              buffer = pickle.dumps(shard)
+            with timer.section('checkpoint_write'):
+              (path / f'{name}-{i:04d}.pkl').write_bytes(buffer)
       else:
         if write:
-          buffer = pickle.dumps(data)
-          (path / f'{name}.pkl').write_bytes(buffer)
+          with timer.section('checkpoint_pickle'):
+            buffer = pickle.dumps(data)
+          with timer.section('checkpoint_write'):
+            (path / f'{name}.pkl').write_bytes(buffer)
     except Exception:
       print(f"Error save '{name}' to checkpoint.")
       raise
-  (path / 'done').write_bytes(b'')
+  write and (path / 'done').write_bytes(b'')
 
 
-def load(path, saveables):
+def load(path, loadfns):
   path = pathlib.Path(path)
   assert exists(path), path
   filenames = set(path.glob('*'))
-  for name, saveable in saveables.items():
+  for name, loadfn in loadfns.items():
     try:
       if (path / f'{name}.pkl') in filenames:
         buffer = (path / f'{name}.pkl').read_bytes()
         data = pickle.loads(buffer)
-        saveable.load(data)
-      elif (path / f'{name}-000000.pkl') in filenames:
+        loadfn(data)
+      elif (path / f'{name}-0000.pkl') in filenames:
         shards = [x for x in filenames if x.name.startswith(f'{name}-')]
         shards = sorted(shards)
         def generator():
@@ -193,7 +197,7 @@ def load(path, saveables):
             buffer = filename.read_bytes()
             data = pickle.loads(buffer)
             yield data
-        saveable.load(generator())
+        loadfn(generator())
       else:
         raise KeyError(name)
     except Exception:

{elements-3.18.0.dev1 → elements-3.18.2/elements.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: elements
-Version: 3.18.0.dev1
+Version: 3.18.2
 Summary: Building blocks for productive research.
 Home-page: http://github.com/danijar/elements
 Classifier: Intended Audience :: Science/Research

{elements-3.18.0.dev1 → elements-3.18.2}/tests/test_checkpoint.py

@@ -6,7 +6,7 @@ class TestCheckpoint:
 
   def test_basic(self, tmpdir):
     path = elements.Path(tmpdir)
-    foo = Foo(42)
+    foo = SaveableMock(42)
     cp = elements.Checkpoint(path)
     cp.foo = foo
     foo.value = 12
@@ -15,7 +15,7 @@ class TestCheckpoint:
     filenames = set(x.name for x in cp.latest().glob('*'))
     assert filenames == {'foo.pkl', 'done'}
     del cp
-    foo = Foo(42)
+    foo = SaveableMock(42)
     cp = elements.Checkpoint(tmpdir)
     cp.foo = foo
     cp.load()
@@ -24,7 +24,7 @@ class TestCheckpoint:
   def test_load_or_save(self, tmpdir):
     path = elements.Path(tmpdir)
     for restart in range(3):
-      foo = Foo(42)
+      foo = SaveableMock(42)
       cp = elements.Checkpoint(path, keep=3)
       cp.foo = foo
       cp.load_or_save()
@@ -35,7 +35,7 @@ class TestCheckpoint:
   def test_keep(self, tmpdir, keep=3):
     path = elements.Path(tmpdir)
     cp = elements.Checkpoint(path, keep=keep)
-    cp.foo = Foo(0)
+    cp.foo = SaveableMock(0)
     for i in range(1, 6):
       cp.foo.value = i
       cp.save()
@@ -49,7 +49,7 @@ class TestCheckpoint:
     path = elements.Path(tmpdir)
     step = elements.Counter(0)
     cp = elements.Checkpoint(path, step=step, keep=3)
-    cp.foo = Foo(0)
+    cp.foo = SaveableMock(0)
     for _ in range(5):
       cp.foo.value = int(step)
       cp.save()
@@ -60,33 +60,22 @@ class TestCheckpoint:
     assert steps == {2, 3, 4}
 
   def test_generator(self, tmpdir):
-
-    class Bar:
-      def __init__(self, values):
-        self.values = values
-      def save(self):
-        for value in self.values:
-          yield {'value': value}
-      def load(self, data):
-        for i, shard in enumerate(data):
-          self.values[i] = shard['value']
-
     path = elements.Path(tmpdir)
     cp = elements.Checkpoint(path)
-    cp.bar = Bar([42, 12, 26])
+    cp.foo = GeneratorMock([42, 12, 26])
     cp.save()
     filenames = set(x.name for x in cp.latest().glob('*'))
     assert filenames == {
-        'bar-000000.pkl',
-        'bar-000001.pkl',
-        'bar-000002.pkl',
+        'foo-0000.pkl',
+        'foo-0001.pkl',
+        'foo-0002.pkl',
         'done',
     }
     del cp
     cp = elements.Checkpoint(path)
-    cp.bar = Bar([0, 0, 0])
+    cp.foo = GeneratorMock([0, 0, 0])
     cp.load()
-    assert cp.bar.values == [42, 12, 26]
+    assert cp.foo.values == [42, 12, 26]
 
   def test_saveable_inline(self, tmpdir):
     path = elements.Path(tmpdir)
@@ -101,25 +90,27 @@ class TestCheckpoint:
     assert foo == [42]
 
   def test_saveable_inherit(self, tmpdir):
-
-    class Bar(elements.Saveable):
-      def __init__(self, value):
-        super().__init__(['value'])
-        self.value = value
-
     path = elements.Path(tmpdir)
-    bar = Bar(42)
+    foo = SubclassMock(42)
     cp = elements.Checkpoint(path)
-    cp.bar = bar
+    cp.foo = foo
     cp.save()
-    bar.value = 12
+    foo.value = 12
     cp.load()
-    assert bar.value == 42
+    assert foo.value == 42
+
+  def test_write(self, tmpdir):
+    path = elements.Path(tmpdir)
+    cp = elements.Checkpoint(path, write=False)
+    cp.foo = SaveableMock(42)
+    cp.bar = GeneratorMock([1, 2, 3])
+    cp.save()
+    assert list(path.glob('*')) == []
 
   def test_path(self, tmpdir):
     path = elements.Path(tmpdir)
     cp = elements.Checkpoint()
-    cp.foo = Foo(42)
+    cp.foo = SaveableMock(42)
     cp.save(path / 'inner')
     assert set(path.glob('*')) == {path / 'inner'}
     cp.foo.value = 12
@@ -129,8 +120,8 @@ class TestCheckpoint:
   def test_keys(self, tmpdir):
     path = elements.Path(tmpdir)
     cp = elements.Checkpoint(path)
-    cp.foo = Foo(42)
-    cp.bar = Foo(12)
+    cp.foo = SaveableMock(42)
+    cp.bar = SaveableMock(12)
     cp.save(keys=['bar'])
     filenames = set(x.name for x in cp.latest().glob('*'))
     assert filenames == {'bar.pkl', 'done'}
@@ -143,7 +134,7 @@ class TestCheckpoint:
       cp.load()
 
 
-class Foo:
+class SaveableMock:
 
   def __init__(self, value):
     self.value = value
@@ -153,3 +144,25 @@ class Foo:
 
   def load(self, data):
     self.value = data['value']
+
+
+class SubclassMock(elements.Saveable):
+
+  def __init__(self, value):
+    super().__init__(['value'])
+    self.value = value
+
+
+class GeneratorMock:
+
+  def __init__(self, values):
+    self.values = values
+
+  def save(self):
+    for value in self.values:
+      shard = {'value': value}
+      yield shard
+
+  def load(self, data):
+    for i, shard in enumerate(data):
+      self.values[i] = shard['value']