class_docstr: Deep LR scheduler tests to check how we handle preemption.
class_docstr_tok: ["Deep","LR","scheduler","tests","to","check","how","we","handle","preemption","."]
code: def _test_scheduler(self, **kwargs):
      from parlai.scripts.train_model import TrainModel, TrainLoop
  
      # shallow copy to prevent overwrites
      kwargs = kwargs.copy()
      with testing_utils.tempdir() as tmpdir:
          kwargs["model"] = "test_agents/unigram"
          kwargs["task"] = "integration_tests"
          kwargs["skip_generation"] = True
          kwargs["validation_metric"] = "loss"
          kwargs["model_file"] = os.path.join(tmpdir, "model")
          kwargs[
              "dict_file"
          ] = "zoo:unittest/transformer_generator2/model.dict"
          kwargs["log_every_n_steps"] = 1
          kwargs["validation_every_n_steps"] = 10
          kwargs["max_train_steps"] = 100
          kwargs["save_after_valid"] = True
          kwargs["learningrate"] = 1
          opt = TrainModel.setup_args().parse_kwargs(**kwargs)
  
          logs_first = []
          for i, train_step_log in enumerate(
              TrainLoop(opt).train_steps(), 1
          ):
              logs_first.append(train_step_log)
              if i >= self.PREEMPT - 2:
                  # simulate preemption
                  break
  
          # resume training
          logs_second = []
          for train_step_log in TrainLoop(opt).train_steps():
              logs_second.append(train_step_log)
  
          # check correctness
          assert (
              logs_first[20]["total_train_updates"]
              == logs_second[0]["total_train_updates"]
          )
          assert logs_first[20]["lr"] == logs_second[0]["lr"]
  
          if "warmup_updates" in kwargs:
              full_logs = logs_first[:20] + logs_second
              assert (
                  abs(
                      1.0
                      - full_logs[kwargs["warmup_updates"] - 1]["lr"]
                  )
                  < 0.04
              )
  
          return logs_first, logs_second
  
imports: import os
  import unittest
  import torch
  import parlai.nn.lr_scheduler
  import parlai.utils.testing
label_desc: Use image node_id 1 for calling the TestLRIntegration obj's underlying member method code with example usage: obj._test_scheduler() and returns: logs_first,  logs_second