[World logging] Fix world logging for dynamic batching

parlai/utils/world_logging.py

@@ -97,6 +97,14 @@ def _add_episode(self, episode):
         """
         self._logs.append(episode)
 
+    def _check_episode_done(self, parley) -> bool:
+        """
+        Check whether an episode is done for a given parley.
+        """
+        if parley[0]:
+            return parley[0].get('episode_done', False)
+        return False
+
     def _is_batch_world(self, world):
         return (
             isinstance(world, BatchWorld) or isinstance(world, DynamicBatchWorld)
@@ -111,7 +119,7 @@ def _log_batch(self, world):
             # in the buffer
             idx = parley[0]['dyn_batch_idx'] if 'dyn_batch_idx' in parley[0] else i
             self._add_msgs(parley, idx=idx)
-            if world.worlds[idx].episode_done():
+            if self._check_episode_done(parley):
                 self.reset_world(idx=idx)
 
     def log(self, world):

tests/test_dynamicbatching.py

@@ -124,6 +124,32 @@ def test_world_logging(self):
                         # we log the batch index in the teacher acts only
                         self.assertEquals(dyn_batch_idx, turn['dyn_batch_idx'])
 
+    def test_world_logging_buffersize(self):
+        """
+        Test world logging with dynamic batching.
+
+        Checks when the number of examples exceeds the buffersize.
+        """
+        with testing_utils.tempdir() as tmpdir:
+            save_report = os.path.join(tmpdir, 'report')
+            testing_utils.eval_model(
+                dict(
+                    model_file='zoo:unittest/transformer_generator2/model',
+                    task='integration_tests:RepeatTeacher:2000',
+                    world_logs=save_report + '.jsonl',
+                    report_filename=save_report,
+                    truncate=1024,
+                    dynamic_batching='full',
+                    batchsize=4,
+                ),
+                valid_datatype='train:evalmode',
+                skip_test=True,
+            )
+            convo_fle = str(save_report) + '.jsonl'
+            convos = Conversations(convo_fle)
+            # we expect there to be 2000 episodes logged in the convos
+            self.assertEquals(len(convos), 2000)
+
     def test_weird_batchsize(self):
         # intentionally a difficult number
         self._test_correct_processed(NUM_TEST, batchsize=7)