Repository: pfnet/pytorch-pfn-extras

Issue:
Trainer: copy config in Handler
Handler should not modify config dict given by user.


Discussion:
to do this, we should separate logic, and handler options in different config options too :(

Resolving patch:
diff --git a/pytorch_pfn_extras/engine.py b/pytorch_pfn_extras/engine.py
--- a/pytorch_pfn_extras/engine.py
+++ b/pytorch_pfn_extras/engine.py
@@ -107,11 +107,12 @@ def create_trainer(
         writer: Optional['writing.Writer'] = None,
         evaluator: Optional['_Evaluator'] = None,
         device: 'DeviceLike' = 'cpu',
-        options: Optional[Dict[str, Any]] = None,
         logic: Optional[handler_module.Logic] = None,
         transform_model: Callable[
             [str, torch.nn.Module], torch.nn.Module] = default_transform_model,
         handler_class: Optional[Type[handler_module.BaseHandler]] = None,
+        options: Optional[Dict[str, Any]] = None,
+        runtime_options: Optional[Dict[str, Any]] = None,
 ) -> '_Trainer':
     """Creates a trainer object.
 
@@ -140,9 +141,6 @@ def create_trainer(
             :func:`pytorch_pfn_extras.engine.create_evaluator`.
         device (str or torch.device):
             Device name used for selecting a corresponding runtime class.
-        options:
-            Option that is set to a logic object. When using the default logic
-            class, See the documentation of `ppe.handler.Logic` for details.
         logic:
             A logic object. If `None` is given, an logic object is instantiated
             from the default logic class.
@@ -152,24 +150,30 @@ def create_trainer(
         handler_class:
             A handler class that instantiates a handler object. If `None` is
             given, `ppe.handler.Handler` is used as a default handler class.
+        options:
+            Options that are set to the handler and logic object.
+            See the documentation of `ppe.handler.Handler` and
+            `ppe.handler.Logic` for details.
+        runtime_options:
+            Options that are set to the runtime object. See the documentation
+            of `ppe.runtime.PyTorchRuntime` for details.
     """
 
-    if options is None:
-        options = {}
-    else:
-        options = options.copy()
-    if logic is None:
-        logic = handler_module.Logic()
-    logic.set_options(options)
-    if handler_class is None:
-        handler_class = handler_module.Handler
+    options = options.copy() if options else {}
+    runtime_options = runtime_options.copy() if runtime_options else {}
+    logic = handler_module.Logic() if logic is None else logic
+    handler_class = handler_class if handler_class else handler_module.Handler
 
     entry_runtime_cls = runtime_registry.get_runtime_class_for_device_spec(
         device)
-    entry_runtime = entry_runtime_cls(device, options.pop('runtime', {}))
-    handler = handler_class(logic, entry_runtime, options)
+    entry_runtime = entry_runtime_cls(device, runtime_options)
+    handler = handler_class(logic, entry_runtime, {})
+
+    # Handler options are popped first, then pass the remainings to Logic.
+    handler.consume_options(options)
+    logic.consume_options(options)
     if len(options) > 0:
-        raise ValueError('Unknown configuration options:', options)
+        raise ValueError('Unknown options: ', options)
 
     from pytorch_pfn_extras.training._trainer import _Trainer
     return _Trainer(
@@ -187,9 +191,10 @@ def create_evaluator(
         progress_bar: bool = False,
         device: 'DeviceLike' = 'cpu',
         metrics: Optional[List['MetricType']] = None,
-        options: Optional[Dict[str, Any]] = None,
         logic: Optional[handler_module.Logic] = None,
         handler_class: Optional[Type[handler_module.BaseHandler]] = None,
+        options: Optional[Dict[str, Any]] = None,
+        runtime_options: Optional[Dict[str, Any]] = None,
 ) -> '_Evaluator':
     """Creates an evaluator object. The return value of this function is
     expected to be fed to `ppe.engine.create_trainer` as an argument.
@@ -206,36 +211,37 @@ def create_evaluator(
         metrics (list of metrics):
             List of metrics, which computes various quantities and update
             output for the reporting.
-        options (dict):
-            Option that is set to a logic object. When using the default logic
-            class, See the documentation of `ppe.handler.Logic` for details.
         logic:
             A logic object. If `None` is given, an logic object is instantiated
             from the default logic class.
         handler_class:
             A handler class that instantiates a handler object. If `None` is
             given, `ppe.handler.Handler` is used as a default handler class.
+        options:
+            Options that are set to the handler and logic object.
+            See the documentation of `ppe.handler.Handler` and
+            `ppe.handler.Logic` for details.
+        runtime_options:
+            Options that are set to the runtime object. See the documentation
+            of `ppe.handler.Handler` for details.
     """
 
-    if metrics is None:
-        metrics = []
-    if options is None:
-        options = {}
-    else:
-        options = options.copy()
-    if logic is None:
-        logic = handler_module.Logic()
-    logic.set_options(options)
-    if handler_class is None:
-        handler_class = handler_module.Handler
+    metrics = metrics if metrics else []
+    options = options.copy() if options else {}
+    runtime_options = runtime_options.copy() if runtime_options else {}
+    logic = handler_module.Logic() if logic is None else logic
+    handler_class = handler_class if handler_class else handler_module.Handler
 
     entry_runtime_cls = runtime_registry.get_runtime_class_for_device_spec(
         device)
-    entry_runtime = entry_runtime_cls(device, options.pop('runtime', {}))
+    entry_runtime = entry_runtime_cls(device, runtime_options)
     handler = handler_class(logic, entry_runtime, options)
 
+    # Handler options are popped first, then pass the remainings to Logic.
+    handler.consume_options(options)
+    logic.consume_options(options)
     if len(options) > 0:
-        raise ValueError('Unknown configuration options ', options)
+        raise ValueError('Unknown options: ', options)
 
     from pytorch_pfn_extras.training._evaluator import _Evaluator
     return _Evaluator(
diff --git a/pytorch_pfn_extras/handler.py b/pytorch_pfn_extras/handler.py
--- a/pytorch_pfn_extras/handler.py
+++ b/pytorch_pfn_extras/handler.py
@@ -29,7 +29,7 @@ def torch_autocast(enabled=True):
 
 class BaseHandler:
 
-    def __init__(self, logic, *args, **kwargs):
+    def __init__(self, logic, options, *args, **kwargs):
         """Base class of Handler.
 
         .. seealso:
@@ -39,7 +39,22 @@ def __init__(self, logic, *args, **kwargs):
             logic (Logic): A logic.
         """
         super().__init__()
+        options = options.copy() if options else {}
         self._logic = logic
+        self.consume_options(options)
+
+    def consume_options(self, options):
+        """A method to update options of Handler.
+
+        Note that the given dict will be modified.
+
+        Args:
+            options (dict): Option key-values to be set.
+
+        .. seealso:
+           :meth:`pytorch_pfn_extras.handler.Handler.consume_options`
+        """
+        pass
 
     def train_setup(self, trainer, loader):
         """A method called only once when starting a training run.
@@ -169,13 +184,13 @@ def eval_post_step(self, evaluator, batch_idx, batch, outputs):
 
 class Handler(BaseHandler):
 
-    def __init__(self, logic, entry_runtime, config):
+    def __init__(self, logic, entry_runtime, options):
         """A set of callback functions to perform device-specific operations.
 
         Args:
             logic (Logic): A logic.
             entry_runtime (BaseRuntime): A runtime object.
-            config (dict): A config.
+            options (dict): The configuration options.
 
                 * ``'autocast'`` (bool):
                     If ``True``, torch.cuda.amp.autocast is enabled.
@@ -191,15 +206,19 @@ def __init__(self, logic, entry_runtime, config):
                 * ``'async'`` (bool):
                     If ``True``, async mode is enabled. Default is ``False``.
         """
-        super().__init__(logic)
+        super().__init__(logic, options)
+        self.pending_iters = defaultdict(list)
+
         # This is used to send the batch to the appropiate device
         self._entry_runtime = entry_runtime
         self._ppe_modules = []
-        self._autocast = config.pop('autocast', False)
-        self._eval_report_keys = config.pop('eval_report_keys', [])
-        self._train_report_keys = config.pop('train_report_keys', [])
-        self._async = config.pop('async', False)
-        self.pending_iters = defaultdict(list)
+
+    def consume_options(self, options):
+        super().consume_options(options)
+        self._autocast = options.pop('autocast', False)
+        self._eval_report_keys = options.pop('eval_report_keys', [])
+        self._train_report_keys = options.pop('train_report_keys', [])
+        self._async = options.pop('async', False)
         if self._autocast and not _amp_enabled:
             raise RuntimeError('Requested AMP features but torch.cuda.amp'
                                ' is not enabled')
@@ -460,6 +479,21 @@ def train_post_step(self, trainer, batch_idx, batch, outputs):
 
 
 class BaseLogic:
+    def __init__(self, options=None):
+        super().__init__()
+        options = options.copy() if options else {}
+        self.consume_options(options)
+
+    def consume_options(self, options):
+        """A method to update options of Logic.
+
+        Note that the given dict will be modified.
+
+        Args:
+            options (dict): Option key-values to be set.
+        """
+        pass
+
     def train_epoch_begin(self, models, epoch, loader):
         """A method called when starting a new epoch of training.
 
@@ -552,10 +586,15 @@ def __init__(self, model_name='main', options=None):
             * ``'grad_scaler'`` (torch.cuda.amp.GradScaler):
                 A gradient scaler that outputs are applied to.
         """
-        if options is None:
-            options = {}
-        self.set_options(options)
+        super().__init__(options)
         self.model_name = model_name
+
+    def consume_options(self, options):
+        super().consume_options(options)
+
+        self.backward_outputs = options.pop('backward_outputs', None)
+        self._grad_scaler = options.pop('grad_scaler', None)
+
         if self._grad_scaler is not None:
             if not _amp_enabled:
                 raise RuntimeError('Requested AMP features but torch.cuda.amp'
@@ -564,10 +603,6 @@ def __init__(self, model_name='main', options=None):
                 raise RuntimeError('grad_scaler should be a '
                                    'torch.cuda.amp.GradScaler object')
 
-    def set_options(self, options):
-        self.backward_outputs = options.pop('backward_outputs', None)
-        self._grad_scaler = options.pop('grad_scaler', None)
-
     def _forward(self, model, batch):
         if isinstance(batch, tuple) and hasattr(batch, '_fields'):
             return model(batch)
diff --git a/pytorch_pfn_extras/runtime/_runtime.py b/pytorch_pfn_extras/runtime/_runtime.py
--- a/pytorch_pfn_extras/runtime/_runtime.py
+++ b/pytorch_pfn_extras/runtime/_runtime.py
@@ -25,17 +25,17 @@ class BaseRuntime:
     Args:
         device_spec (torch.device or str):
             The device that modules and tensors are transferred to.
-        config (dict):
+        options (dict):
             A configuration dictionary that can be used from runtime method.
     """
 
     def __init__(
             self,
             device_spec: DeviceLike,
-            config: Optional[Dict[str, Any]] = None,
+            options: Optional[Dict[str, Any]] = None,
     ) -> None:
         self.device_spec = device_spec
-        self.config = config
+        self.options = options
 
     def convert_batch(self, args: Any) -> Any:
         """Transfers the given batch to the specific device.


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/test_handler.py b/tests/pytorch_pfn_extras_tests/test_handler.py
--- a/tests/pytorch_pfn_extras_tests/test_handler.py
+++ b/tests/pytorch_pfn_extras_tests/test_handler.py
@@ -7,8 +7,8 @@
 
 
 class MockRuntime(ppe.runtime.BaseRuntime):
-    def __init__(self, device, config):
-        super().__init__(device, config)
+    def __init__(self, device, options):
+        super().__init__(device, options)
         self._initialize_called = False
         self._train_epoch_begin_called = False
         self._called_module = None
@@ -104,14 +104,14 @@ def eval_step(self, models, batch_idx, batch):
 
 
 class HandlerTester:
-    def _get_handler(self, config=None):
-        if config is None:
-            config = {}
+    def _get_handler(self, options=None):
+        if options is None:
+            options = {}
         ppe.runtime.runtime_registry.register('test_rt', MockRuntime)
         trainer = MockTrainer()
         logic = MockLogic()
         handler = ppe.handler.Handler(
-            logic, MockRuntime('test_rt', {}), config
+            logic, MockRuntime('test_rt', {}), options
         )
         return handler, trainer, logic
 
@@ -132,6 +132,7 @@ def _assert_called(self, module, to_move, function):
 
 
 class TestHandlerTrainSync(HandlerTester):
+
     @pytest.mark.parametrize(
         'to_move', [('self',), ('sm1',), ('sm2',), ('sm1', 'sm2')]
     )
@@ -180,8 +181,8 @@ def test_train_step(self, to_move):
         'to_move', [('self',), ('sm1',), ('sm2',), ('sm1', 'sm2')]
     )
     def test_train_post_step(self, to_move):
-        config = {'train_report_keys': ['output']}
-        handler, trainer, _ = self._get_handler(config)
+        options = {'train_report_keys': ['output']}
+        handler, trainer, _ = self._get_handler(options)
         module = trainer.models['main']
         self._move_modules(module, to_move)
         reporter = ppe.reporting.Reporter()
@@ -192,8 +193,8 @@ def test_train_post_step(self, to_move):
 
 
 class TestHandlerValidationSync(HandlerTester):
-    def _get_handler(self, config=None):
-        handler, _, logic = super()._get_handler(config)
+    def _get_handler(self, options=None):
+        handler, _, logic = super()._get_handler(options)
         evaluator = MockEvaluator()
         return handler, evaluator, logic
 
@@ -236,8 +237,8 @@ def test_eval_step(self, to_move):
         'to_move', [('self',), ('sm1',), ('sm2',), ('sm1', 'sm2')]
     )
     def test_train_post_step(self, to_move):
-        config = {'eval_report_keys': ['output']}
-        handler, evaluator, _ = self._get_handler(config)
+        options = {'eval_report_keys': ['output']}
+        handler, evaluator, _ = self._get_handler(options)
         module = evaluator.models['main']
         self._move_modules(module, to_move)
         reporter = ppe.reporting.Reporter()
@@ -248,8 +249,8 @@ def test_train_post_step(self, to_move):
 
 
 class AsyncRuntime(MockRuntime):
-    def __init__(self, device, config):
-        super().__init__(device, config)
+    def __init__(self, device, options):
+        super().__init__(device, options)
         # Returns a result once every 10 items
         self._period = 10
         self._count = 0
@@ -262,18 +263,18 @@ def get_pending_result(self, module, blocking):
 
 
 class TestAsyncHandler:
-    def _get_handler(self, config):
+    def _get_handler(self, options):
         ppe.runtime.runtime_registry.register('test_rt', AsyncRuntime)
         logic = MockLogic()
         handler = ppe.handler.Handler(
-            logic, AsyncRuntime('test_rt', {}), config
+            logic, AsyncRuntime('test_rt', {}), options
         )
         return handler
 
     def test_train_step_async(self):
-        config = {'eval_report_keys': ['output'], 'async': True}
+        options = {'eval_report_keys': ['output'], 'async': True}
         trainer = MockTrainer()
-        handler = self._get_handler(config)
+        handler = self._get_handler(options)
         ppe.to(trainer.models['main'], 'test_rt')
         prev_batch_idx = 0
 
@@ -294,8 +295,8 @@ def callback(batch_idx, outs, is_deferred):
         assert len(handler.pending_iters['main']) == 0
 
     def test_eval_step_async(self):
-        config = {'eval_report_keys': ['output'], 'async': True}
-        handler = self._get_handler(config)
+        options = {'eval_report_keys': ['output'], 'async': True}
+        handler = self._get_handler(options)
         evaluator = MockEvaluator()
         ppe.to(evaluator.models['main'], 'test_rt')
         prev_batch_idx = 0
@@ -317,9 +318,9 @@ def callback(batch_idx, outs, is_deferred):
         assert len(handler.pending_iters['main']) == 0
 
     def test_setup_multi_device_split_invalid(self):
-        config = {'eval_report_keys': ['output'], 'async': True}
+        options = {'eval_report_keys': ['output'], 'async': True}
         trainer = MockTrainer()
-        handler = self._get_handler(config)
+        handler = self._get_handler(options)
         ppe.to(trainer.models['main'].sm1, 'test_rt')
         ppe.to(trainer.models['main'].sm2, 'cpu')
         with pytest.raises(RuntimeError, match='models splitted'):
@@ -379,6 +380,7 @@ def test_autocast_not_enabled(self):
 
 
 class TestLogic:
+
     def test_train_epoch_begin(self):
         # Check that the DataLoader has the sampler updated
         class _MockedDL:
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py b/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py
--- a/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py
@@ -86,6 +86,23 @@ def test_trainer(device, path):
     trainer.run(data, data)
 
 
+def test_trainer_invalid_options(path):
+    device = 'cpu'
+    model = MyModel()
+    ppe.to(model, device)
+    model_with_loss = MyModelWithLossFn(model)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+    extensions = _make_extensions()
+    options = {'UNKNOWN_OPTIONS': True}
+    with pytest.raises(ValueError, match="UNKNOWN_OPTIONS"):
+        engine.create_trainer(
+            model_with_loss, optimizer, 20,
+            device=device, extensions=extensions,
+            out_dir=path,
+            options=options,
+        )
+
+
 @pytest.mark.parametrize('device', ['cpu', 'cuda'])
 @pytest.mark.parametrize('progress_bar', [True, False])
 def test_train_with_evaluator(device, progress_bar, path):