KaneAI - World’s First E2E Software Testing Agent.Book a Demo Join TestMu Conference

How to use sess method in Pytest

Best Python code snippet using pytest

basic_session_run_hooks_test.py

Source:basic_session_run_hooks_test.py

1# pylint: disable=g-bad-file-header2# Copyright 2016 The TensorFlow Authors. All Rights Reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15# ==============================================================================16"""Tests for basic_session_run_hooks."""17from __future__ import absolute_import18from __future__ import division19from __future__ import print_function20import shutil21import tempfile22import threading23import time24from tensorflow.contrib.framework.python.framework import checkpoint_utils25from tensorflow.contrib.framework.python.ops import variables26from tensorflow.contrib.testing.python.framework import fake_summary_writer27from tensorflow.python.client import session as session_lib28from tensorflow.python.framework import constant_op29from tensorflow.python.framework import dtypes30from tensorflow.python.framework import meta_graph31from tensorflow.python.framework import ops32from tensorflow.python.ops import array_ops33from tensorflow.python.ops import control_flow_ops34from tensorflow.python.ops import state_ops35from tensorflow.python.ops import variable_scope36from tensorflow.python.ops import variables as variables_lib37import tensorflow.python.ops.nn_grad  # pylint: disable=unused-import38from tensorflow.python.platform import test39from tensorflow.python.platform import tf_logging40from tensorflow.python.summary import summary as summary_lib41from tensorflow.python.summary.writer import writer_cache42from tensorflow.python.training import basic_session_run_hooks43from tensorflow.python.training import monitored_session44from tensorflow.python.training import session_run_hook45class MockCheckpointSaverListener(46    basic_session_run_hooks.CheckpointSaverListener):47  def __init__(self):48    self.begin_count = 049    self.before_save_count = 050    self.after_save_count = 051    self.end_count = 052  def begin(self):53    self.begin_count += 154  def before_save(self, session, global_step):55    self.before_save_count += 156  def after_save(self, session, global_step):57    self.after_save_count += 158  def end(self, session, global_step):59    self.end_count += 160  def get_counts(self):61    return {62        'begin': self.begin_count,63        'before_save': self.before_save_count,64        'after_save': self.after_save_count,65        'end': self.end_count66    }67class SecondOrStepTimerTest(test.TestCase):68  def test_raise_in_both_secs_and_steps(self):69    with self.assertRaises(ValueError):70      basic_session_run_hooks.SecondOrStepTimer(every_secs=2.0, every_steps=10)71  def test_raise_in_none_secs_and_steps(self):72    with self.assertRaises(ValueError):73      basic_session_run_hooks.SecondOrStepTimer()74  def test_every_secs(self):75    timer = basic_session_run_hooks.SecondOrStepTimer(every_secs=1.0)76    self.assertTrue(timer.should_trigger_for_step(1))77    timer.update_last_triggered_step(1)78    self.assertFalse(timer.should_trigger_for_step(1))79    self.assertFalse(timer.should_trigger_for_step(2))80    time.sleep(1.0)81    self.assertFalse(timer.should_trigger_for_step(1))82    self.assertTrue(timer.should_trigger_for_step(2))83  def test_every_steps(self):84    timer = basic_session_run_hooks.SecondOrStepTimer(every_steps=3)85    self.assertTrue(timer.should_trigger_for_step(1))86    timer.update_last_triggered_step(1)87    self.assertFalse(timer.should_trigger_for_step(1))88    self.assertFalse(timer.should_trigger_for_step(2))89    self.assertFalse(timer.should_trigger_for_step(3))90    self.assertTrue(timer.should_trigger_for_step(4))91  def test_update_last_triggered_step(self):92    timer = basic_session_run_hooks.SecondOrStepTimer(every_steps=1)93    elapsed_secs, elapsed_steps = timer.update_last_triggered_step(1)94    self.assertEqual(None, elapsed_secs)95    self.assertEqual(None, elapsed_steps)96    elapsed_secs, elapsed_steps = timer.update_last_triggered_step(5)97    self.assertLess(0, elapsed_secs)98    self.assertEqual(4, elapsed_steps)99    elapsed_secs, elapsed_steps = timer.update_last_triggered_step(7)100    self.assertLess(0, elapsed_secs)101    self.assertEqual(2, elapsed_steps)102class StopAtStepTest(test.TestCase):103  def test_raise_in_both_last_step_and_num_steps(self):104    with self.assertRaises(ValueError):105      basic_session_run_hooks.StopAtStepHook(num_steps=10, last_step=20)106  def test_stop_based_on_last_step(self):107    h = basic_session_run_hooks.StopAtStepHook(last_step=10)108    with ops.Graph().as_default():109      global_step = variables.get_or_create_global_step()110      no_op = control_flow_ops.no_op()111      h.begin()112      with session_lib.Session() as sess:113        mon_sess = monitored_session._HookedSession(sess, [h])114        sess.run(state_ops.assign(global_step, 5))115        h.after_create_session(sess, None)116        mon_sess.run(no_op)117        self.assertFalse(mon_sess.should_stop())118        sess.run(state_ops.assign(global_step, 9))119        mon_sess.run(no_op)120        self.assertFalse(mon_sess.should_stop())121        sess.run(state_ops.assign(global_step, 10))122        mon_sess.run(no_op)123        self.assertTrue(mon_sess.should_stop())124        sess.run(state_ops.assign(global_step, 11))125        mon_sess._should_stop = False126        mon_sess.run(no_op)127        self.assertTrue(mon_sess.should_stop())128  def test_stop_based_on_num_step(self):129    h = basic_session_run_hooks.StopAtStepHook(num_steps=10)130    with ops.Graph().as_default():131      global_step = variables.get_or_create_global_step()132      no_op = control_flow_ops.no_op()133      h.begin()134      with session_lib.Session() as sess:135        mon_sess = monitored_session._HookedSession(sess, [h])136        sess.run(state_ops.assign(global_step, 5))137        h.after_create_session(sess, None)138        mon_sess.run(no_op)139        self.assertFalse(mon_sess.should_stop())140        sess.run(state_ops.assign(global_step, 13))141        mon_sess.run(no_op)142        self.assertFalse(mon_sess.should_stop())143        sess.run(state_ops.assign(global_step, 14))144        mon_sess.run(no_op)145        self.assertFalse(mon_sess.should_stop())146        sess.run(state_ops.assign(global_step, 15))147        mon_sess.run(no_op)148        self.assertTrue(mon_sess.should_stop())149        sess.run(state_ops.assign(global_step, 16))150        mon_sess._should_stop = False151        mon_sess.run(no_op)152        self.assertTrue(mon_sess.should_stop())153  def test_stop_based_with_multiple_steps(self):154    h = basic_session_run_hooks.StopAtStepHook(num_steps=10)155    with ops.Graph().as_default():156      global_step = variables.get_or_create_global_step()157      no_op = control_flow_ops.no_op()158      h.begin()159      with session_lib.Session() as sess:160        mon_sess = monitored_session._HookedSession(sess, [h])161        sess.run(state_ops.assign(global_step, 5))162        h.after_create_session(sess, None)163        mon_sess.run(no_op)164        self.assertFalse(mon_sess.should_stop())165        sess.run(state_ops.assign(global_step, 15))166        mon_sess.run(no_op)167        self.assertTrue(mon_sess.should_stop())168class LoggingTensorHookTest(test.TestCase):169  def setUp(self):170    # Mock out logging calls so we can verify whether correct tensors are being171    # monitored.172    self._actual_log = tf_logging.info173    self.logged_message = None174    def mock_log(*args, **kwargs):175      self.logged_message = args176      self._actual_log(*args, **kwargs)177    tf_logging.info = mock_log178  def tearDown(self):179    tf_logging.info = self._actual_log180  def test_illegal_args(self):181    with self.assertRaisesRegexp(ValueError, 'nvalid every_n_iter'):182      basic_session_run_hooks.LoggingTensorHook(tensors=['t'], every_n_iter=0)183    with self.assertRaisesRegexp(ValueError, 'nvalid every_n_iter'):184      basic_session_run_hooks.LoggingTensorHook(tensors=['t'], every_n_iter=-10)185    with self.assertRaisesRegexp(ValueError, 'xactly one of'):186      basic_session_run_hooks.LoggingTensorHook(187          tensors=['t'], every_n_iter=5, every_n_secs=5)188    with self.assertRaisesRegexp(ValueError, 'xactly one of'):189      basic_session_run_hooks.LoggingTensorHook(tensors=['t'])190  def test_print_at_end_only(self):191    with ops.Graph().as_default(), session_lib.Session() as sess:192      t = constant_op.constant(42.0, name='foo')193      train_op = constant_op.constant(3)194      hook = basic_session_run_hooks.LoggingTensorHook(195          tensors=[t.name], at_end=True)196      hook.begin()197      mon_sess = monitored_session._HookedSession(sess, [hook])198      sess.run(variables_lib.global_variables_initializer())199      self.logged_message = ''200      for _ in range(3):201        mon_sess.run(train_op)202        # assertNotRegexpMatches is not supported by python 3.1 and later203        self.assertEqual(str(self.logged_message).find(t.name), -1)204      hook.end(sess)205      self.assertRegexpMatches(str(self.logged_message), t.name)206  def _validate_print_every_n_steps(self, sess, at_end):207    t = constant_op.constant(42.0, name='foo')208    train_op = constant_op.constant(3)209    hook = basic_session_run_hooks.LoggingTensorHook(210        tensors=[t.name], every_n_iter=10, at_end=at_end)211    hook.begin()212    mon_sess = monitored_session._HookedSession(sess, [hook])213    sess.run(variables_lib.global_variables_initializer())214    mon_sess.run(train_op)215    self.assertRegexpMatches(str(self.logged_message), t.name)216    for _ in range(3):217      self.logged_message = ''218      for _ in range(9):219        mon_sess.run(train_op)220        # assertNotRegexpMatches is not supported by python 3.1 and later221        self.assertEqual(str(self.logged_message).find(t.name), -1)222      mon_sess.run(train_op)223      self.assertRegexpMatches(str(self.logged_message), t.name)224    # Add additional run to verify proper reset when called multiple times.225    self.logged_message = ''226    mon_sess.run(train_op)227    # assertNotRegexpMatches is not supported by python 3.1 and later228    self.assertEqual(str(self.logged_message).find(t.name), -1)229    self.logged_message = ''230    hook.end(sess)231    if at_end:232      self.assertRegexpMatches(str(self.logged_message), t.name)233    else:234      # assertNotRegexpMatches is not supported by python 3.1 and later235      self.assertEqual(str(self.logged_message).find(t.name), -1)236  def test_print_every_n_steps(self):237    with ops.Graph().as_default(), session_lib.Session() as sess:238      self._validate_print_every_n_steps(sess, at_end=False)239      # Verify proper reset.240      self._validate_print_every_n_steps(sess, at_end=False)241  def test_print_every_n_steps_and_end(self):242    with ops.Graph().as_default(), session_lib.Session() as sess:243      self._validate_print_every_n_steps(sess, at_end=True)244      # Verify proper reset.245      self._validate_print_every_n_steps(sess, at_end=True)246  def test_print_first_step(self):247    # if it runs every iteration, first iteration has None duration.248    with ops.Graph().as_default(), session_lib.Session() as sess:249      t = constant_op.constant(42.0, name='foo')250      train_op = constant_op.constant(3)251      hook = basic_session_run_hooks.LoggingTensorHook(252          tensors={'foo': t}, every_n_iter=1)253      hook.begin()254      mon_sess = monitored_session._HookedSession(sess, [hook])255      sess.run(variables_lib.global_variables_initializer())256      mon_sess.run(train_op)257      self.assertRegexpMatches(str(self.logged_message), 'foo')258      # in first run, elapsed time is None.259      self.assertEqual(str(self.logged_message).find('sec'), -1)260  def _validate_print_every_n_secs(self, sess, at_end):261    t = constant_op.constant(42.0, name='foo')262    train_op = constant_op.constant(3)263    hook = basic_session_run_hooks.LoggingTensorHook(264        tensors=[t.name], every_n_secs=1.0, at_end=at_end)265    hook.begin()266    mon_sess = monitored_session._HookedSession(sess, [hook])267    sess.run(variables_lib.global_variables_initializer())268    mon_sess.run(train_op)269    self.assertRegexpMatches(str(self.logged_message), t.name)270    # assertNotRegexpMatches is not supported by python 3.1 and later271    self.logged_message = ''272    mon_sess.run(train_op)273    self.assertEqual(str(self.logged_message).find(t.name), -1)274    time.sleep(1.0)275    self.logged_message = ''276    mon_sess.run(train_op)277    self.assertRegexpMatches(str(self.logged_message), t.name)278    self.logged_message = ''279    hook.end(sess)280    if at_end:281      self.assertRegexpMatches(str(self.logged_message), t.name)282    else:283      # assertNotRegexpMatches is not supported by python 3.1 and later284      self.assertEqual(str(self.logged_message).find(t.name), -1)285  def test_print_every_n_secs(self):286    with ops.Graph().as_default(), session_lib.Session() as sess:287      self._validate_print_every_n_secs(sess, at_end=False)288      # Verify proper reset.289      self._validate_print_every_n_secs(sess, at_end=False)290  def test_print_every_n_secs_and_end(self):291    with ops.Graph().as_default(), session_lib.Session() as sess:292      self._validate_print_every_n_secs(sess, at_end=True)293      # Verify proper reset.294      self._validate_print_every_n_secs(sess, at_end=True)295  def test_print_formatter(self):296    with ops.Graph().as_default(), session_lib.Session() as sess:297      t = constant_op.constant(42.0, name='foo')298      train_op = constant_op.constant(3)299      hook = basic_session_run_hooks.LoggingTensorHook(300          tensors=[t.name], every_n_iter=10,301          formatter=lambda items: 'qqq=%s' % items[t.name])302      hook.begin()303      mon_sess = monitored_session._HookedSession(sess, [hook])304      sess.run(variables_lib.global_variables_initializer())305      mon_sess.run(train_op)306      self.assertEqual(self.logged_message[0], 'qqq=42.0')307class CheckpointSaverHookTest(test.TestCase):308  def setUp(self):309    self.model_dir = tempfile.mkdtemp()310    self.graph = ops.Graph()311    with self.graph.as_default():312      self.scaffold = monitored_session.Scaffold()313      self.global_step = variables.get_or_create_global_step()314      self.train_op = state_ops.assign_add(self.global_step, 1)315  def tearDown(self):316    shutil.rmtree(self.model_dir, ignore_errors=True)317  def test_raise_when_saver_and_scaffold_both_missing(self):318    with self.assertRaises(ValueError):319      basic_session_run_hooks.CheckpointSaverHook(self.model_dir)320  def test_raise_when_saver_and_scaffold_both_present(self):321    with self.assertRaises(ValueError):322      basic_session_run_hooks.CheckpointSaverHook(323          self.model_dir, saver=self.scaffold.saver, scaffold=self.scaffold)324  def test_raise_in_both_secs_and_steps(self):325    with self.assertRaises(ValueError):326      basic_session_run_hooks.CheckpointSaverHook(327          self.model_dir, save_secs=10, save_steps=20)328  def test_raise_in_none_secs_and_steps(self):329    with self.assertRaises(ValueError):330      basic_session_run_hooks.CheckpointSaverHook(self.model_dir)331  def test_save_secs_saves_in_first_step(self):332    with self.graph.as_default():333      hook = basic_session_run_hooks.CheckpointSaverHook(334          self.model_dir, save_secs=2, scaffold=self.scaffold)335      hook.begin()336      self.scaffold.finalize()337      with session_lib.Session() as sess:338        sess.run(self.scaffold.init_op)339        mon_sess = monitored_session._HookedSession(sess, [hook])340        mon_sess.run(self.train_op)341        self.assertEqual(1,342                         checkpoint_utils.load_variable(self.model_dir,343                                                        self.global_step.name))344  def test_save_secs_calls_listeners_at_begin_and_end(self):345    with self.graph.as_default():346      listener = MockCheckpointSaverListener()347      hook = basic_session_run_hooks.CheckpointSaverHook(348          self.model_dir,349          save_secs=2,350          scaffold=self.scaffold,351          listeners=[listener])352      hook.begin()353      self.scaffold.finalize()354      with session_lib.Session() as sess:355        sess.run(self.scaffold.init_op)356        mon_sess = monitored_session._HookedSession(sess, [hook])357        mon_sess.run(self.train_op)  # hook runs here358        mon_sess.run(self.train_op)  # hook won't run here, so it does at end359        hook.end(sess)  # hook runs here360      self.assertEqual({361          'begin': 1,362          'before_save': 2,363          'after_save': 2,364          'end': 1365      }, listener.get_counts())366  def test_listener_with_monitored_session(self):367    with ops.Graph().as_default():368      scaffold = monitored_session.Scaffold()369      global_step = variables.get_or_create_global_step()370      train_op = state_ops.assign_add(global_step, 1)371      listener = MockCheckpointSaverListener()372      hook = basic_session_run_hooks.CheckpointSaverHook(373          self.model_dir,374          save_steps=1,375          scaffold=scaffold,376          listeners=[listener])377      with monitored_session.SingularMonitoredSession(378          hooks=[hook],379          scaffold=scaffold,380          checkpoint_dir=self.model_dir) as sess:381        sess.run(train_op)382        sess.run(train_op)383        global_step_val = sess.run(global_step)384      listener_counts = listener.get_counts()385    self.assertEqual(2, global_step_val)386    self.assertEqual({387        'begin': 1,388        'before_save': 2,389        'after_save': 2,390        'end': 1391    }, listener_counts)392  def test_listener_with_default_saver(self):393    with ops.Graph().as_default():394      global_step = variables.get_or_create_global_step()395      train_op = state_ops.assign_add(global_step, 1)396      listener = MockCheckpointSaverListener()397      hook = basic_session_run_hooks.CheckpointSaverHook(398          self.model_dir,399          save_steps=1,400          listeners=[listener])401      with monitored_session.SingularMonitoredSession(402          hooks=[hook],403          checkpoint_dir=self.model_dir) as sess:404        sess.run(train_op)405        sess.run(train_op)406        global_step_val = sess.run(global_step)407      listener_counts = listener.get_counts()408    self.assertEqual(2, global_step_val)409    self.assertEqual({410        'begin': 1,411        'before_save': 2,412        'after_save': 2,413        'end': 1414    }, listener_counts)415    with ops.Graph().as_default():416      global_step = variables.get_or_create_global_step()417      with monitored_session.SingularMonitoredSession(418          checkpoint_dir=self.model_dir) as sess2:419        global_step_saved_val = sess2.run(global_step)420    self.assertEqual(2, global_step_saved_val)421  def test_two_listeners_with_default_saver(self):422    with ops.Graph().as_default():423      global_step = variables.get_or_create_global_step()424      train_op = state_ops.assign_add(global_step, 1)425      listener1 = MockCheckpointSaverListener()426      listener2 = MockCheckpointSaverListener()427      hook = basic_session_run_hooks.CheckpointSaverHook(428          self.model_dir,429          save_steps=1,430          listeners=[listener1, listener2])431      with monitored_session.SingularMonitoredSession(432          hooks=[hook],433          checkpoint_dir=self.model_dir) as sess:434        sess.run(train_op)435        sess.run(train_op)436        global_step_val = sess.run(global_step)437      listener1_counts = listener1.get_counts()438      listener2_counts = listener2.get_counts()439    self.assertEqual(2, global_step_val)440    self.assertEqual({441        'begin': 1,442        'before_save': 2,443        'after_save': 2,444        'end': 1445    }, listener1_counts)446    self.assertEqual(listener1_counts, listener2_counts)447    with ops.Graph().as_default():448      global_step = variables.get_or_create_global_step()449      with monitored_session.SingularMonitoredSession(450          checkpoint_dir=self.model_dir) as sess2:451        global_step_saved_val = sess2.run(global_step)452    self.assertEqual(2, global_step_saved_val)453  @test.mock.patch('time.time')454  def test_save_secs_saves_periodically(self, mock_time):455    # Let's have a realistic start time456    current_time = 1484695987.209386457    with self.graph.as_default():458      mock_time.return_value = current_time459      hook = basic_session_run_hooks.CheckpointSaverHook(460          self.model_dir, save_secs=2, scaffold=self.scaffold)461      hook.begin()462      self.scaffold.finalize()463      with session_lib.Session() as sess:464        sess.run(self.scaffold.init_op)465        mon_sess = monitored_session._HookedSession(sess, [hook])466        mock_time.return_value = current_time467        mon_sess.run(self.train_op)  # Saved.468        mock_time.return_value = current_time + 0.5469        mon_sess.run(self.train_op)  # Not saved.470        self.assertEqual(1,471                         checkpoint_utils.load_variable(self.model_dir,472                                                        self.global_step.name))473        # Simulate 2.5 seconds of sleep.474        mock_time.return_value = current_time + 2.5475        mon_sess.run(self.train_op)  # Saved.476        mock_time.return_value = current_time + 2.6477        mon_sess.run(self.train_op)  # Not saved.478        mock_time.return_value = current_time + 2.7479        mon_sess.run(self.train_op)  # Not saved.480        self.assertEqual(3,481                         checkpoint_utils.load_variable(self.model_dir,482                                                        self.global_step.name))483        # Simulate 7.5 more seconds of sleep (10 seconds from start.484        mock_time.return_value = current_time + 10485        mon_sess.run(self.train_op)  # Saved.486        self.assertEqual(6,487                         checkpoint_utils.load_variable(self.model_dir,488                                                        self.global_step.name))489  # Flaky because of time.sleep()490  def DISABLED_test_save_secs_calls_listeners_periodically(self):491    with self.graph.as_default():492      listener = MockCheckpointSaverListener()493      hook = basic_session_run_hooks.CheckpointSaverHook(494          self.model_dir,495          save_secs=2,496          scaffold=self.scaffold,497          listeners=[listener])498      hook.begin()499      self.scaffold.finalize()500      with session_lib.Session() as sess:501        sess.run(self.scaffold.init_op)502        mon_sess = monitored_session._HookedSession(sess, [hook])503        mon_sess.run(self.train_op)  # hook runs here504        mon_sess.run(self.train_op)505        time.sleep(2.5)506        mon_sess.run(self.train_op)  # hook runs here507        mon_sess.run(self.train_op)508        mon_sess.run(self.train_op)509        time.sleep(2.5)510        mon_sess.run(self.train_op)  # hook runs here511        mon_sess.run(self.train_op)  # hook won't run here, so it does at end512        hook.end(sess)  # hook runs here513      self.assertEqual({514          'begin': 1,515          'before_save': 4,516          'after_save': 4,517          'end': 1518      }, listener.get_counts())519  def test_save_steps_saves_in_first_step(self):520    with self.graph.as_default():521      hook = basic_session_run_hooks.CheckpointSaverHook(522          self.model_dir, save_steps=2, scaffold=self.scaffold)523      hook.begin()524      self.scaffold.finalize()525      with session_lib.Session() as sess:526        sess.run(self.scaffold.init_op)527        mon_sess = monitored_session._HookedSession(sess, [hook])528        mon_sess.run(self.train_op)529        self.assertEqual(1,530                         checkpoint_utils.load_variable(self.model_dir,531                                                        self.global_step.name))532  def test_save_steps_saves_periodically(self):533    with self.graph.as_default():534      hook = basic_session_run_hooks.CheckpointSaverHook(535          self.model_dir, save_steps=2, scaffold=self.scaffold)536      hook.begin()537      self.scaffold.finalize()538      with session_lib.Session() as sess:539        sess.run(self.scaffold.init_op)540        mon_sess = monitored_session._HookedSession(sess, [hook])541        mon_sess.run(self.train_op)542        mon_sess.run(self.train_op)543        # Not saved544        self.assertEqual(1,545                         checkpoint_utils.load_variable(self.model_dir,546                                                        self.global_step.name))547        mon_sess.run(self.train_op)548        # saved549        self.assertEqual(3,550                         checkpoint_utils.load_variable(self.model_dir,551                                                        self.global_step.name))552        mon_sess.run(self.train_op)553        # Not saved554        self.assertEqual(3,555                         checkpoint_utils.load_variable(self.model_dir,556                                                        self.global_step.name))557        mon_sess.run(self.train_op)558        # saved559        self.assertEqual(5,560                         checkpoint_utils.load_variable(self.model_dir,561                                                        self.global_step.name))562  def test_save_saves_at_end(self):563    with self.graph.as_default():564      hook = basic_session_run_hooks.CheckpointSaverHook(565          self.model_dir, save_secs=2, scaffold=self.scaffold)566      hook.begin()567      self.scaffold.finalize()568      with session_lib.Session() as sess:569        sess.run(self.scaffold.init_op)570        mon_sess = monitored_session._HookedSession(sess, [hook])571        mon_sess.run(self.train_op)572        mon_sess.run(self.train_op)573        hook.end(sess)574        self.assertEqual(2,575                         checkpoint_utils.load_variable(self.model_dir,576                                                        self.global_step.name))577  def test_summary_writer_defs(self):578    fake_summary_writer.FakeSummaryWriter.install()579    writer_cache.FileWriterCache.clear()580    summary_writer = writer_cache.FileWriterCache.get(self.model_dir)581    with self.graph.as_default():582      hook = basic_session_run_hooks.CheckpointSaverHook(583          self.model_dir, save_steps=2, scaffold=self.scaffold)584      hook.begin()585      self.scaffold.finalize()586      with session_lib.Session() as sess:587        sess.run(self.scaffold.init_op)588        mon_sess = monitored_session._HookedSession(sess, [hook])589        mon_sess.run(self.train_op)590      summary_writer.assert_summaries(591          test_case=self,592          expected_logdir=self.model_dir,593          expected_added_meta_graphs=[594              meta_graph.create_meta_graph_def(595                  graph_def=self.graph.as_graph_def(add_shapes=True),596                  saver_def=self.scaffold.saver.saver_def)597          ])598    fake_summary_writer.FakeSummaryWriter.uninstall()599class ResourceCheckpointSaverHookTest(test.TestCase):600  def setUp(self):601    self.model_dir = tempfile.mkdtemp()602    self.graph = ops.Graph()603    with self.graph.as_default():604      self.scaffold = monitored_session.Scaffold()605      with variable_scope.variable_scope('foo', use_resource=True):606        self.global_step = variables.get_or_create_global_step()607      self.train_op = state_ops.assign_add(self.global_step, 1)608  def test_save_steps_saves_periodically(self):609    with self.graph.as_default():610      hook = basic_session_run_hooks.CheckpointSaverHook(611          self.model_dir, save_steps=2, scaffold=self.scaffold)612      hook.begin()613      self.scaffold.finalize()614      with session_lib.Session() as sess:615        sess.run(self.scaffold.init_op)616        mon_sess = monitored_session._HookedSession(sess, [hook])617        mon_sess.run(self.train_op)618        mon_sess.run(self.train_op)619        # Not saved620        self.assertEqual(1,621                         checkpoint_utils.load_variable(self.model_dir,622                                                        self.global_step.name))623        mon_sess.run(self.train_op)624        # saved625        self.assertEqual(3,626                         checkpoint_utils.load_variable(self.model_dir,627                                                        self.global_step.name))628        mon_sess.run(self.train_op)629        # Not saved630        self.assertEqual(3,631                         checkpoint_utils.load_variable(self.model_dir,632                                                        self.global_step.name))633        mon_sess.run(self.train_op)634        # saved635        self.assertEqual(5,636                         checkpoint_utils.load_variable(self.model_dir,637                                                        self.global_step.name))638class StepCounterHookTest(test.TestCase):639  def setUp(self):640    self.log_dir = tempfile.mkdtemp()641  def tearDown(self):642    shutil.rmtree(self.log_dir, ignore_errors=True)643  def test_step_counter_every_n_steps(self):644    with ops.Graph().as_default() as g, session_lib.Session() as sess:645      global_step = variables.get_or_create_global_step()646      train_op = state_ops.assign_add(global_step, 1)647      summary_writer = fake_summary_writer.FakeSummaryWriter(self.log_dir, g)648      hook = basic_session_run_hooks.StepCounterHook(649          summary_writer=summary_writer, every_n_steps=10)650      hook.begin()651      sess.run(variables_lib.global_variables_initializer())652      mon_sess = monitored_session._HookedSession(sess, [hook])653      for _ in range(30):654        time.sleep(0.01)655        mon_sess.run(train_op)656      hook.end(sess)657      summary_writer.assert_summaries(658          test_case=self,659          expected_logdir=self.log_dir,660          expected_graph=g,661          expected_summaries={})662      self.assertItemsEqual([11, 21], summary_writer.summaries.keys())663      for step in [11, 21]:664        summary_value = summary_writer.summaries[step][0].value[0]665        self.assertEqual('global_step/sec', summary_value.tag)666        self.assertGreater(summary_value.simple_value, 0)667  def test_step_counter_every_n_secs(self):668    with ops.Graph().as_default() as g, session_lib.Session() as sess:669      global_step = variables.get_or_create_global_step()670      train_op = state_ops.assign_add(global_step, 1)671      summary_writer = fake_summary_writer.FakeSummaryWriter(self.log_dir, g)672      hook = basic_session_run_hooks.StepCounterHook(673          summary_writer=summary_writer, every_n_steps=None, every_n_secs=0.1)674      hook.begin()675      sess.run(variables_lib.global_variables_initializer())676      mon_sess = monitored_session._HookedSession(sess, [hook])677      mon_sess.run(train_op)678      time.sleep(0.2)679      mon_sess.run(train_op)680      time.sleep(0.2)681      mon_sess.run(train_op)682      hook.end(sess)683      summary_writer.assert_summaries(684          test_case=self,685          expected_logdir=self.log_dir,686          expected_graph=g,687          expected_summaries={})688      self.assertTrue(summary_writer.summaries, 'No summaries were created.')689      self.assertItemsEqual([2, 3], summary_writer.summaries.keys())690      for summary in summary_writer.summaries.values():691        summary_value = summary[0].value[0]692        self.assertEqual('global_step/sec', summary_value.tag)693        self.assertGreater(summary_value.simple_value, 0)694  def test_global_step_name(self):695    with ops.Graph().as_default() as g, session_lib.Session() as sess:696      with variable_scope.variable_scope('bar'):697        foo_step = variable_scope.get_variable(698            'foo',699            initializer=0,700            trainable=False,701            collections=[702                ops.GraphKeys.GLOBAL_STEP, ops.GraphKeys.GLOBAL_VARIABLES703            ])704      train_op = state_ops.assign_add(foo_step, 1)705      summary_writer = fake_summary_writer.FakeSummaryWriter(self.log_dir, g)706      hook = basic_session_run_hooks.StepCounterHook(707          summary_writer=summary_writer, every_n_steps=1, every_n_secs=None)708      hook.begin()709      sess.run(variables_lib.global_variables_initializer())710      mon_sess = monitored_session._HookedSession(sess, [hook])711      mon_sess.run(train_op)712      mon_sess.run(train_op)713      hook.end(sess)714      summary_writer.assert_summaries(715          test_case=self,716          expected_logdir=self.log_dir,717          expected_graph=g,718          expected_summaries={})719      self.assertTrue(summary_writer.summaries, 'No summaries were created.')720      self.assertItemsEqual([2], summary_writer.summaries.keys())721      summary_value = summary_writer.summaries[2][0].value[0]722      self.assertEqual('bar/foo/sec', summary_value.tag)723class SummarySaverHookTest(test.TestCase):724  def setUp(self):725    test.TestCase.setUp(self)726    self.log_dir = 'log/dir'727    self.summary_writer = fake_summary_writer.FakeSummaryWriter(self.log_dir)728    var = variables_lib.Variable(0.0)729    tensor = state_ops.assign_add(var, 1.0)730    tensor2 = tensor * 2731    self.summary_op = summary_lib.scalar('my_summary', tensor)732    self.summary_op2 = summary_lib.scalar('my_summary2', tensor2)733    global_step = variables.get_or_create_global_step()734    self.train_op = state_ops.assign_add(global_step, 1)735  def test_raise_when_scaffold_and_summary_op_both_missing(self):736    with self.assertRaises(ValueError):737      basic_session_run_hooks.SummarySaverHook()738  def test_raise_when_scaffold_and_summary_op_both_present(self):739    with self.assertRaises(ValueError):740      basic_session_run_hooks.SummarySaverHook(741          scaffold=monitored_session.Scaffold(), summary_op=self.summary_op)742  def test_raise_in_both_secs_and_steps(self):743    with self.assertRaises(ValueError):744      basic_session_run_hooks.SummarySaverHook(745          save_secs=10, save_steps=20, summary_writer=self.summary_writer)746  def test_raise_in_none_secs_and_steps(self):747    with self.assertRaises(ValueError):748      basic_session_run_hooks.SummarySaverHook(749          save_secs=None, save_steps=None, summary_writer=self.summary_writer)750  def test_save_steps(self):751    hook = basic_session_run_hooks.SummarySaverHook(752        save_steps=8,753        summary_writer=self.summary_writer,754        summary_op=self.summary_op)755    with self.test_session() as sess:756      hook.begin()757      sess.run(variables_lib.global_variables_initializer())758      mon_sess = monitored_session._HookedSession(sess, [hook])759      for _ in range(30):760        mon_sess.run(self.train_op)761      hook.end(sess)762    self.summary_writer.assert_summaries(763        test_case=self,764        expected_logdir=self.log_dir,765        expected_summaries={766            1: {767                'my_summary': 1.0768            },769            9: {770                'my_summary': 2.0771            },772            17: {773                'my_summary': 3.0774            },775            25: {776                'my_summary': 4.0777            },778        })779  def test_multiple_summaries(self):780    hook = basic_session_run_hooks.SummarySaverHook(781        save_steps=8,782        summary_writer=self.summary_writer,783        summary_op=[self.summary_op, self.summary_op2])784    with self.test_session() as sess:785      hook.begin()786      sess.run(variables_lib.global_variables_initializer())787      mon_sess = monitored_session._HookedSession(sess, [hook])788      for _ in range(10):789        mon_sess.run(self.train_op)790      hook.end(sess)791    self.summary_writer.assert_summaries(792        test_case=self,793        expected_logdir=self.log_dir,794        expected_summaries={795            1: {796                'my_summary': 1.0,797                'my_summary2': 2.0798            },799            9: {800                'my_summary': 2.0,801                'my_summary2': 4.0802            },803        })804  def test_save_secs_saving_once_every_step(self):805    hook = basic_session_run_hooks.SummarySaverHook(806        save_secs=0.5,807        summary_writer=self.summary_writer,808        summary_op=self.summary_op)809    with self.test_session() as sess:810      hook.begin()811      sess.run(variables_lib.global_variables_initializer())812      mon_sess = monitored_session._HookedSession(sess, [hook])813      for _ in range(4):814        mon_sess.run(self.train_op)815        time.sleep(0.5)816      hook.end(sess)817    self.summary_writer.assert_summaries(818        test_case=self,819        expected_logdir=self.log_dir,820        expected_summaries={821            1: {822                'my_summary': 1.0823            },824            2: {825                'my_summary': 2.0826            },827            3: {828                'my_summary': 3.0829            },830            4: {831                'my_summary': 4.0832            },833        })834  def test_save_secs_saving_once_every_three_steps(self):835    hook = basic_session_run_hooks.SummarySaverHook(836        save_secs=0.9,837        summary_writer=self.summary_writer,838        summary_op=self.summary_op)839    with self.test_session() as sess:840      hook.begin()841      sess.run(variables_lib.global_variables_initializer())842      mon_sess = monitored_session._HookedSession(sess, [hook])843      for _ in range(8):844        mon_sess.run(self.train_op)845        time.sleep(0.3)846      hook.end(sess)847    self.summary_writer.assert_summaries(848        test_case=self,849        expected_logdir=self.log_dir,850        expected_summaries={851            1: {852                'my_summary': 1.0853            },854            4: {855                'my_summary': 2.0856            },857            7: {858                'my_summary': 3.0859            },860        })861class GlobalStepWaiterHookTest(test.TestCase):862  def test_not_wait_for_step_zero(self):863    with ops.Graph().as_default():864      variables.get_or_create_global_step()865      hook = basic_session_run_hooks.GlobalStepWaiterHook(wait_until_step=0)866      hook.begin()867      with session_lib.Session() as sess:868        # Before run should return without waiting gstep increment.869        hook.before_run(870            session_run_hook.SessionRunContext(871                original_args=None, session=sess))872  def test_wait_for_step(self):873    with ops.Graph().as_default():874      gstep = variables.get_or_create_global_step()875      hook = basic_session_run_hooks.GlobalStepWaiterHook(wait_until_step=1000)876      hook.begin()877      with session_lib.Session() as sess:878        sess.run(variables_lib.global_variables_initializer())879        waiter = threading.Thread(880            target=hook.before_run,881            args=(session_run_hook.SessionRunContext(882                original_args=None, session=sess),))883        waiter.daemon = True884        waiter.start()885        time.sleep(1.0)886        self.assertTrue(waiter.is_alive())887        sess.run(state_ops.assign(gstep, 500))888        time.sleep(1.0)889        self.assertTrue(waiter.is_alive())890        sess.run(state_ops.assign(gstep, 1100))891        time.sleep(1.2)892        self.assertFalse(waiter.is_alive())893class FinalOpsHookTest(test.TestCase):894  def test_final_ops_is_scalar_tensor(self):895    with ops.Graph().as_default():896      expected_value = 4897      final_ops = constant_op.constant(expected_value)898      hook = basic_session_run_hooks.FinalOpsHook(final_ops)899      hook.begin()900      with session_lib.Session() as session:901        hook.end(session)902        self.assertEqual(expected_value,903                         hook.final_ops_values)904  def test_final_ops_is_tensor(self):905    with ops.Graph().as_default():906      expected_values = [1, 6, 3, 5, 2, 4]907      final_ops = constant_op.constant(expected_values)908      hook = basic_session_run_hooks.FinalOpsHook(final_ops)909      hook.begin()910      with session_lib.Session() as session:911        hook.end(session)912        self.assertListEqual(expected_values,913                             hook.final_ops_values.tolist())914  def test_final_ops_with_dictionary(self):915    with ops.Graph().as_default():916      expected_values = [4, -3]917      final_ops = array_ops.placeholder(dtype=dtypes.float32)918      final_ops_feed_dict = {final_ops: expected_values}919      hook = basic_session_run_hooks.FinalOpsHook(920          final_ops, final_ops_feed_dict)921      hook.begin()922      with session_lib.Session() as session:923        hook.end(session)924        self.assertListEqual(expected_values,925                             hook.final_ops_values.tolist())926class ResourceSummarySaverHookTest(test.TestCase):927  def setUp(self):928    test.TestCase.setUp(self)929    self.log_dir = 'log/dir'930    self.summary_writer = fake_summary_writer.FakeSummaryWriter(self.log_dir)931    var = variable_scope.get_variable('var', initializer=0.0, use_resource=True)932    tensor = state_ops.assign_add(var, 1.0)933    self.summary_op = summary_lib.scalar('my_summary', tensor)934    with variable_scope.variable_scope('foo', use_resource=True):935      global_step = variables.get_or_create_global_step()936    self.train_op = state_ops.assign_add(global_step, 1)937  def test_save_steps(self):938    hook = basic_session_run_hooks.SummarySaverHook(939        save_steps=8,940        summary_writer=self.summary_writer,941        summary_op=self.summary_op)942    with self.test_session() as sess:943      hook.begin()944      sess.run(variables_lib.global_variables_initializer())945      mon_sess = monitored_session._HookedSession(sess, [hook])946      for _ in range(30):947        mon_sess.run(self.train_op)948      hook.end(sess)949    self.summary_writer.assert_summaries(950        test_case=self,951        expected_logdir=self.log_dir,952        expected_summaries={953            1: {954                'my_summary': 1.0955            },956            9: {957                'my_summary': 2.0958            },959            17: {960                'my_summary': 3.0961            },962            25: {963                'my_summary': 4.0964            },965        })966class FeedFnHookTest(test.TestCase):967  def test_feeding_placeholder(self):968    with ops.Graph().as_default(), session_lib.Session() as sess:969      x = array_ops.placeholder(dtype=dtypes.float32)970      y = x + 1971      hook = basic_session_run_hooks.FeedFnHook(972          feed_fn=lambda: {x: 1.0})973      hook.begin()974      mon_sess = monitored_session._HookedSession(sess, [hook])975      self.assertEqual(mon_sess.run(y), 2)976if __name__ == '__main__':...

local_cli_wrapper_test.py

Source:local_cli_wrapper_test.py

1# Copyright 2016 The TensorFlow Authors. All Rights Reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14# ==============================================================================15"""Unit tests for local command-line-interface debug wrapper session."""16from __future__ import absolute_import17from __future__ import division18from __future__ import print_function19import os20import shutil21import tempfile22from tensorflow.core.protobuf import config_pb223from tensorflow.python.client import session24from tensorflow.python.debug.cli import cli_shared25from tensorflow.python.debug.cli import debugger_cli_common26from tensorflow.python.debug.wrappers import local_cli_wrapper27from tensorflow.python.framework import constant_op28from tensorflow.python.framework import dtypes29from tensorflow.python.framework import errors30from tensorflow.python.framework import ops31from tensorflow.python.framework import test_util32from tensorflow.python.ops import array_ops33from tensorflow.python.ops import control_flow_ops34from tensorflow.python.ops import math_ops35# Import resource_variable_ops for the variables-to-tensor implicit conversion.36from tensorflow.python.ops import resource_variable_ops  # pylint: disable=unused-import37from tensorflow.python.ops import state_ops38from tensorflow.python.ops import variables39from tensorflow.python.platform import googletest40class LocalCLIDebuggerWrapperSessionForTest(41    local_cli_wrapper.LocalCLIDebugWrapperSession):42  """Subclasses the wrapper class for testing.43  Overrides its CLI-related methods for headless testing environments.44  Inserts observer variables for assertions.45  """46  def __init__(self,47               command_args_sequence,48               sess,49               dump_root=None):50    """Constructor of the for-test subclass.51    Args:52      command_args_sequence: (list of list of str) A list of arguments for the53        "run" command.54      sess: See the doc string of LocalCLIDebugWrapperSession.__init__.55      dump_root: See the doc string of LocalCLIDebugWrapperSession.__init__.56    """57    local_cli_wrapper.LocalCLIDebugWrapperSession.__init__(58        self, sess, dump_root=dump_root, log_usage=False)59    self._command_args_sequence = command_args_sequence60    self._response_pointer = 061    # Observer variables.62    self.observers = {63        "debug_dumps": [],64        "tf_errors": [],65        "run_start_cli_run_numbers": [],66        "run_end_cli_run_numbers": [],67        "profiler_py_graphs": [],68        "profiler_run_metadata": [],69    }70  def _prep_cli_for_run_start(self):71    pass72  def _prep_debug_cli_for_run_end(self, debug_dump, tf_error, passed_filter):73    self.observers["debug_dumps"].append(debug_dump)74    self.observers["tf_errors"].append(tf_error)75  def _prep_profile_cli_for_run_end(self, py_graph, run_metadata):76    self.observers["profiler_py_graphs"].append(py_graph)77    self.observers["profiler_run_metadata"].append(run_metadata)78  def _launch_cli(self):79    if self._is_run_start:80      self.observers["run_start_cli_run_numbers"].append(self._run_call_count)81    else:82      self.observers["run_end_cli_run_numbers"].append(self._run_call_count)83    command_args = self._command_args_sequence[self._response_pointer]84    self._response_pointer += 185    try:86      self._run_handler(command_args)87    except debugger_cli_common.CommandLineExit as e:88      response = e.exit_token89    return response90class LocalCLIDebugWrapperSessionTest(test_util.TensorFlowTestCase):91  def setUp(self):92    self._tmp_dir = tempfile.mktemp()93    self.v = variables.Variable(10.0, name="v")94    self.w = variables.Variable(21.0, name="w")95    self.delta = constant_op.constant(1.0, name="delta")96    self.inc_v = state_ops.assign_add(self.v, self.delta, name="inc_v")97    self.w_int = control_flow_ops.with_dependencies(98        [self.inc_v],99        math_ops.cast(self.w, dtypes.int32, name="w_int_inner"),100        name="w_int_outer")101    self.ph = array_ops.placeholder(dtypes.float32, name="ph")102    self.xph = array_ops.transpose(self.ph, name="xph")103    self.m = constant_op.constant(104        [[0.0, 1.0, 2.0], [-4.0, -1.0, 0.0]], dtype=dtypes.float32, name="m")105    self.y = math_ops.matmul(self.m, self.xph, name="y")106    self.sess = session.Session()107    # Initialize variable.108    self.sess.run(variables.global_variables_initializer())109  def tearDown(self):110    ops.reset_default_graph()111    if os.path.isdir(self._tmp_dir):112      shutil.rmtree(self._tmp_dir)113  def testConstructWrapper(self):114    local_cli_wrapper.LocalCLIDebugWrapperSession(115        session.Session(), log_usage=False)116  def testConstructWrapperWithExistingEmptyDumpRoot(self):117    os.mkdir(self._tmp_dir)118    self.assertTrue(os.path.isdir(self._tmp_dir))119    local_cli_wrapper.LocalCLIDebugWrapperSession(120        session.Session(), dump_root=self._tmp_dir, log_usage=False)121  def testConstructWrapperWithExistingNonEmptyDumpRoot(self):122    os.mkdir(self._tmp_dir)123    dir_path = os.path.join(self._tmp_dir, "foo")124    os.mkdir(dir_path)125    self.assertTrue(os.path.isdir(dir_path))126    with self.assertRaisesRegexp(127        ValueError, "dump_root path points to a non-empty directory"):128      local_cli_wrapper.LocalCLIDebugWrapperSession(129          session.Session(), dump_root=self._tmp_dir, log_usage=False)130  def testConstructWrapperWithExistingFileDumpRoot(self):131    os.mkdir(self._tmp_dir)132    file_path = os.path.join(self._tmp_dir, "foo")133    open(file_path, "a").close()  # Create the file134    self.assertTrue(os.path.isfile(file_path))135    with self.assertRaisesRegexp(ValueError, "dump_root path points to a file"):136      local_cli_wrapper.LocalCLIDebugWrapperSession(137          session.Session(), dump_root=file_path, log_usage=False)138  def testRunsUnderDebugMode(self):139    # Test command sequence: run; run; run;140    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(141        [[], [], []], self.sess, dump_root=self._tmp_dir)142    # run under debug mode twice.143    wrapped_sess.run(self.inc_v)144    wrapped_sess.run(self.inc_v)145    # Verify that the assign_add op did take effect.146    self.assertAllClose(12.0, self.sess.run(self.v))147    # Assert correct run call numbers for which the CLI has been launched at148    # run-start and run-end.149    self.assertEqual([1], wrapped_sess.observers["run_start_cli_run_numbers"])150    self.assertEqual([1, 2], wrapped_sess.observers["run_end_cli_run_numbers"])151    # Verify that the dumps have been generated and picked up during run-end.152    self.assertEqual(2, len(wrapped_sess.observers["debug_dumps"]))153    # Verify that the TensorFlow runtime errors are picked up and in this case,154    # they should be both None.155    self.assertEqual([None, None], wrapped_sess.observers["tf_errors"])156  def testRunsWithEmptyStringDumpRootWorks(self):157    # Test command sequence: run, run158    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(159        [[], []], self.sess, dump_root="")160    # run under debug mode.161    wrapped_sess.run(self.inc_v)162    self.assertAllClose(11.0, self.sess.run(self.v))163  def testRunInfoOutputAtRunEndIsCorrect(self):164    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(165        [[], [], []], self.sess, dump_root=self._tmp_dir)166    wrapped_sess.run(self.inc_v)167    run_info_output = wrapped_sess._run_info_handler([])168    tfdbg_logo = cli_shared.get_tfdbg_logo()169    # The run_info output in the first run() call should contain the tfdbg logo.170    self.assertEqual(tfdbg_logo.lines,171                     run_info_output.lines[:len(tfdbg_logo.lines)])172    menu = run_info_output.annotations[debugger_cli_common.MAIN_MENU_KEY]173    self.assertIn("list_tensors", menu.captions())174    wrapped_sess.run(self.inc_v)175    run_info_output = wrapped_sess._run_info_handler([])176    # The run_info output in the second run() call should NOT contain the logo.177    self.assertNotEqual(tfdbg_logo.lines,178                        run_info_output.lines[:len(tfdbg_logo.lines)])179    menu = run_info_output.annotations[debugger_cli_common.MAIN_MENU_KEY]180    self.assertIn("list_tensors", menu.captions())181  def testRunsUnderNonDebugMode(self):182    # Test command sequence: run -n; run -n; run -n;183    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(184        [["-n"], ["-n"], ["-n"]], self.sess, dump_root=self._tmp_dir)185    # run three times.186    wrapped_sess.run(self.inc_v)187    wrapped_sess.run(self.inc_v)188    wrapped_sess.run(self.inc_v)189    self.assertAllClose(13.0, self.sess.run(self.v))190    self.assertEqual([1, 2, 3],191                     wrapped_sess.observers["run_start_cli_run_numbers"])192    self.assertEqual([], wrapped_sess.observers["run_end_cli_run_numbers"])193  def testRunsUnderNonDebugThenDebugMode(self):194    # Test command sequence: run -n; run -n; run; run;195    # Do two NON_DEBUG_RUNs, followed by DEBUG_RUNs.196    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(197        [["-n"], ["-n"], [], []], self.sess, dump_root=self._tmp_dir)198    # run three times.199    wrapped_sess.run(self.inc_v)200    wrapped_sess.run(self.inc_v)201    wrapped_sess.run(self.inc_v)202    self.assertAllClose(13.0, self.sess.run(self.v))203    self.assertEqual([1, 2, 3],204                     wrapped_sess.observers["run_start_cli_run_numbers"])205    # Here, the CLI should have been launched only under the third run,206    # because the first and second runs are NON_DEBUG.207    self.assertEqual([3], wrapped_sess.observers["run_end_cli_run_numbers"])208    self.assertEqual(1, len(wrapped_sess.observers["debug_dumps"]))209    self.assertEqual([None], wrapped_sess.observers["tf_errors"])210  def testRunMultipleTimesWithinLimit(self):211    # Test command sequence: run -t 3; run;212    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(213        [["-t", "3"], []], self.sess, dump_root=self._tmp_dir)214    # run three times.215    wrapped_sess.run(self.inc_v)216    wrapped_sess.run(self.inc_v)217    wrapped_sess.run(self.inc_v)218    self.assertAllClose(13.0, self.sess.run(self.v))219    self.assertEqual([1], wrapped_sess.observers["run_start_cli_run_numbers"])220    self.assertEqual([3], wrapped_sess.observers["run_end_cli_run_numbers"])221    self.assertEqual(1, len(wrapped_sess.observers["debug_dumps"]))222    self.assertEqual([None], wrapped_sess.observers["tf_errors"])223  def testRunMultipleTimesOverLimit(self):224    # Test command sequence: run -t 3;225    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(226        [["-t", "3"]], self.sess, dump_root=self._tmp_dir)227    # run twice, which is less than the number of times specified by the228    # command.229    wrapped_sess.run(self.inc_v)230    wrapped_sess.run(self.inc_v)231    self.assertAllClose(12.0, self.sess.run(self.v))232    self.assertEqual([1], wrapped_sess.observers["run_start_cli_run_numbers"])233    self.assertEqual([], wrapped_sess.observers["run_end_cli_run_numbers"])234    self.assertEqual(0, len(wrapped_sess.observers["debug_dumps"]))235    self.assertEqual([], wrapped_sess.observers["tf_errors"])236  def testRunMixingDebugModeAndMultpleTimes(self):237    # Test command sequence: run -n; run -t 2; run; run;238    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(239        [["-n"], ["-t", "2"], [], []], self.sess, dump_root=self._tmp_dir)240    # run four times.241    wrapped_sess.run(self.inc_v)242    wrapped_sess.run(self.inc_v)243    wrapped_sess.run(self.inc_v)244    wrapped_sess.run(self.inc_v)245    self.assertAllClose(14.0, self.sess.run(self.v))246    self.assertEqual([1, 2],247                     wrapped_sess.observers["run_start_cli_run_numbers"])248    self.assertEqual([3, 4], wrapped_sess.observers["run_end_cli_run_numbers"])249    self.assertEqual(2, len(wrapped_sess.observers["debug_dumps"]))250    self.assertEqual([None, None], wrapped_sess.observers["tf_errors"])251  def testDebuggingMakeCallableTensorRunnerWorks(self):252    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(253        [[], []], self.sess, dump_root=self._tmp_dir)254    v = variables.Variable(42)255    tensor_runner = wrapped_sess.make_callable(v)256    self.sess.run(v.initializer)257    self.assertAllClose(42, tensor_runner())258    self.assertEqual(1, len(wrapped_sess.observers["debug_dumps"]))259  def testDebuggingMakeCallableTensorRunnerWithCustomRunOptionsWorks(self):260    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(261        [[], []], self.sess, dump_root=self._tmp_dir)262    a = constant_op.constant(42)263    tensor_runner = wrapped_sess.make_callable(a)264    run_options = config_pb2.RunOptions(265        trace_level=config_pb2.RunOptions.FULL_TRACE)266    run_metadata = config_pb2.RunMetadata()267    self.assertAllClose(268        42, tensor_runner(options=run_options, run_metadata=run_metadata))269    self.assertEqual(1, len(wrapped_sess.observers["debug_dumps"]))270    self.assertGreater(len(run_metadata.step_stats.dev_stats), 0)271  def testDebuggingMakeCallableOperationRunnerWorks(self):272    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(273        [[], []], self.sess, dump_root=self._tmp_dir)274    v = variables.Variable(10.0)275    inc_v = state_ops.assign_add(v, 1.0)276    op_runner = wrapped_sess.make_callable(inc_v.op)277    self.sess.run(v.initializer)278    op_runner()279    self.assertEqual(1, len(wrapped_sess.observers["debug_dumps"]))280    self.assertEqual(11.0, self.sess.run(v))281  def testDebuggingMakeCallableRunnerWithFeedListWorks(self):282    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(283        [[], []], self.sess, dump_root=self._tmp_dir)284    ph1 = array_ops.placeholder(dtypes.float32)285    ph2 = array_ops.placeholder(dtypes.float32)286    a = math_ops.add(ph1, ph2)287    tensor_runner = wrapped_sess.make_callable(a, feed_list=[ph1, ph2])288    self.assertAllClose(42.0, tensor_runner(41.0, 1.0))289    self.assertEqual(1, len(wrapped_sess.observers["debug_dumps"]))290  def testRuntimeErrorShouldBeCaught(self):291    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(292        [[], []], self.sess, dump_root=self._tmp_dir)293    # Do a run that should lead to an TensorFlow runtime error.294    wrapped_sess.run(self.y, feed_dict={self.ph: [[0.0], [1.0], [2.0]]})295    self.assertEqual([1], wrapped_sess.observers["run_start_cli_run_numbers"])296    self.assertEqual([1], wrapped_sess.observers["run_end_cli_run_numbers"])297    self.assertEqual(1, len(wrapped_sess.observers["debug_dumps"]))298    # Verify that the runtime error is caught by the wrapped session properly.299    self.assertEqual(1, len(wrapped_sess.observers["tf_errors"]))300    tf_error = wrapped_sess.observers["tf_errors"][0]301    self.assertEqual("y", tf_error.op.name)302  def testRuntimeErrorBeforeGraphExecutionIsRaised(self):303    # Use an impossible device name to cause an error before graph execution.304    with ops.device("/gpu:1337"):305      w = variables.Variable([1.0] * 10, name="w")306    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(307        [[]], self.sess, dump_root=self._tmp_dir)308    with self.assertRaisesRegexp(errors.OpError, r".*[Dd]evice.*1337.*"):309      wrapped_sess.run(w)310  def testRunTillFilterPassesShouldLaunchCLIAtCorrectRun(self):311    # Test command sequence:312    #   run -f greater_than_twelve; run -f greater_than_twelve; run;313    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(314        [["-f", "v_greater_than_twelve"], ["-f", "v_greater_than_twelve"], []],315        self.sess,316        dump_root=self._tmp_dir)317    def v_greater_than_twelve(datum, tensor):318      return datum.node_name == "v" and tensor > 12.0319    wrapped_sess.add_tensor_filter("v_greater_than_twelve",320                                   v_greater_than_twelve)321    # run five times.322    wrapped_sess.run(self.inc_v)323    wrapped_sess.run(self.inc_v)324    wrapped_sess.run(self.inc_v)325    wrapped_sess.run(self.inc_v)326    wrapped_sess.run(self.inc_v)327    self.assertAllClose(15.0, self.sess.run(self.v))328    self.assertEqual([1], wrapped_sess.observers["run_start_cli_run_numbers"])329    # run-end CLI should NOT have been launched for run #2 and #3, because only330    # starting from run #4 v becomes greater than 12.0.331    self.assertEqual([4, 5], wrapped_sess.observers["run_end_cli_run_numbers"])332    self.assertEqual(2, len(wrapped_sess.observers["debug_dumps"]))333    self.assertEqual([None, None], wrapped_sess.observers["tf_errors"])334  def testRunsUnderDebugModeWithWatchFnFilteringNodeNames(self):335    # Test command sequence:336    #   run --node_name_filter inc.*337    #   run --node_name_filter delta338    #   run339    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(340        [["--node_name_filter", "inc.*"], ["--node_name_filter", "delta"], []],341        self.sess, dump_root=self._tmp_dir)342    # run under debug mode twice.343    wrapped_sess.run(self.inc_v)344    wrapped_sess.run(self.inc_v)345    # Verify that the assign_add op did take effect.346    self.assertAllClose(12.0, self.sess.run(self.v))347    # Verify that the dumps have been generated and picked up during run-end.348    self.assertEqual(2, len(wrapped_sess.observers["debug_dumps"]))349    dumps = wrapped_sess.observers["debug_dumps"][0]350    self.assertEqual(1, dumps.size)351    self.assertEqual("inc_v", dumps.dumped_tensor_data[0].node_name)352    dumps = wrapped_sess.observers["debug_dumps"][1]353    self.assertEqual(1, dumps.size)354    self.assertEqual("delta", dumps.dumped_tensor_data[0].node_name)355  def testRunsUnderDebugModeWithWatchFnFilteringOpTypes(self):356    # Test command sequence:357    #   run --node_name_filter delta358    #   run --op_type_filter AssignAdd359    #   run360    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(361        [["--node_name_filter", "delta"],362         ["--op_type_filter", "AssignAdd"],363         []],364        self.sess, dump_root=self._tmp_dir)365    # run under debug mode twice.366    wrapped_sess.run(self.inc_v)367    wrapped_sess.run(self.inc_v)368    # Verify that the assign_add op did take effect.369    self.assertAllClose(12.0, self.sess.run(self.v))370    # Verify that the dumps have been generated and picked up during run-end.371    self.assertEqual(2, len(wrapped_sess.observers["debug_dumps"]))372    dumps = wrapped_sess.observers["debug_dumps"][0]373    self.assertEqual(1, dumps.size)374    self.assertEqual("delta", dumps.dumped_tensor_data[0].node_name)375    dumps = wrapped_sess.observers["debug_dumps"][1]376    self.assertEqual(1, dumps.size)377    self.assertEqual("inc_v", dumps.dumped_tensor_data[0].node_name)378  def testRunsUnderDebugModeWithWatchFnFilteringTensorDTypes(self):379    # Test command sequence:380    #   run --op_type_filter Variable.*381    #   run --dtype_filter int32382    #   run383    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(384        [["--op_type_filter", "Variable.*"],385         ["--tensor_dtype_filter", "int32"], []],386        self.sess, dump_root=self._tmp_dir)387    # run under debug mode twice.388    wrapped_sess.run(self.w_int)389    wrapped_sess.run(self.w_int)390    # Verify that the dumps have been generated and picked up during run-end.391    self.assertEqual(2, len(wrapped_sess.observers["debug_dumps"]))392    dumps = wrapped_sess.observers["debug_dumps"][0]393    self.assertEqual(2, dumps.size)394    self.assertItemsEqual(395        ["v", "w"], [dumps.dumped_tensor_data[i].node_name for i in [0, 1]])396    dumps = wrapped_sess.observers["debug_dumps"][1]397    self.assertEqual(2, dumps.size)398    self.assertEqual(399        ["w_int_inner", "w_int_outer"],400        [dumps.dumped_tensor_data[i].node_name for i in [0, 1]])401  def testRunsUnderDebugModeWithWatchFnFilteringOpTypesAndTensorDTypes(self):402    # Test command sequence:403    #   run --op_type_filter Cast --dtype_filter int32404    #   run405    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(406        [["--op_type_filter", "Cast", "--tensor_dtype_filter", "int32"], []],407        self.sess, dump_root=self._tmp_dir)408    # run under debug mode twice.409    wrapped_sess.run(self.w_int)410    # Verify that the dumps have been generated and picked up during run-end.411    self.assertEqual(1, len(wrapped_sess.observers["debug_dumps"]))412    dumps = wrapped_sess.observers["debug_dumps"][0]413    self.assertEqual(1, dumps.size)414    self.assertEqual("w_int_inner", dumps.dumped_tensor_data[0].node_name)415  def testRunUnderProfilerModeWorks(self):416    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(417        [["-p"], []], self.sess)418    wrapped_sess.run(self.w_int)419    self.assertEqual(1, len(wrapped_sess.observers["profiler_run_metadata"]))420    self.assertTrue(421        wrapped_sess.observers["profiler_run_metadata"][0].step_stats)422    self.assertEqual(1, len(wrapped_sess.observers["profiler_py_graphs"]))423    self.assertIsInstance(424        wrapped_sess.observers["profiler_py_graphs"][0], ops.Graph)425if __name__ == "__main__":...

backend.py

Source:backend.py

1import asyncio, time2from collections import defaultdict3from enum import IntFlag4from util.misc import gen_uuid, EMPTY_SET, run_loop5from .user import UserService6from .auth import AuthService7from .stats import Stats8from .models import User, Group, Lst, Contact, UserStatus9from . import error, event10class Ack(IntFlag):11	Zero = 012	NAK = 113	ACK = 214	Full = 315class Backend:16	def __init__(self, loop, *, user_service = None, auth_service = None):17		self._loop = loop18		self._user_service = user_service or UserService()19		self._auth_service = auth_service or AuthService()20		self._stats = Stats()21		22		self._sc = _SessionCollection()23		# Dict[User.uuid, User]24		self._user_by_uuid = {}25		# Dict[User, UserDetail]26		self._unsynced_db = {}27		28		# Dict[chatid, Chat]29		self._chats = {}30		31		self._runners = []32		33		loop.create_task(self._sync_db())34		loop.create_task(self._clean_sessions())35		loop.create_task(self._sync_stats())36	37	def add_runner(self, runner):38		self._runners.append(runner)39	40	def run_forever(self):41		run_loop(self._loop, self._runners)42	43	def on_leave(self, sess):44		user = sess.user45		if user is None: return46		self._stats.on_logout()47		self._sc.remove_session(sess)48		if self._sc.get_sessions_by_user(user):49			# There are still other people logged in as this user,50			# so don't send offline notifications.51			return52		# User is offline, send notifications53		user.detail = None54		self._sync_contact_statuses()55		self._generic_notify(sess)56	57	def login_md5_get_salt(self, email):58		return self._user_service.get_md5_salt(email)59	60	def login_md5_verify(self, sess, email, md5_hash):61		uuid = self._user_service.login_md5(email, md5_hash)62		return self._login_common(sess, uuid, email)63	64	def login_twn_start(self, email, password):65		uuid = self._user_service.login(email, password)66		if uuid is None: return None67		return self._auth_service.create_token('nb/login', uuid)68	69	def login_twn_verify(self, sess, email, token):70		uuid = self._auth_service.pop_token('nb/login', token)71		return self._login_common(sess, uuid, email)72	73	def login_IKWIAD(self, sess, email):74		uuid = self.util_get_uuid_from_email(email)75		return self._login_common(sess, uuid, email)76	77	def _login_common(self, sess, uuid, email):78		if uuid is None: return None79		self._user_service.update_date_login(uuid)80		user = self._load_user_record(uuid)81		sess.user = user82		self._stats.on_login()83		self._stats.on_user_active(user, sess.client)84		self._sc.add_session(sess)85		user.detail = self._load_detail(user)86		return user87	88	def _load_user_record(self, uuid):89		if uuid not in self._user_by_uuid:90			user = self._user_service.get(uuid)91			if user is None: return None92			self._user_by_uuid[uuid] = user93		return self._user_by_uuid[uuid]94	95	def _load_detail(self, user):96		if user.detail: return user.detail97		return self._user_service.get_detail(user.uuid)98	99	def _generic_notify(self, sess):100		# Notify relevant `Session`s of status, name, message, media101		user = sess.user102		if user is None: return103		# TODO: This does a lot of work, iterating through _every_ session.104		# If RL is set up properly, could iterate through `user.detail.contacts`.105		for sess_other in self._sc.iter_sessions():106			if sess_other == sess: continue107			user_other = sess_other.user108			if user_other is None: continue109			if user_other.detail is None: continue110			ctc = user_other.detail.contacts.get(user.uuid)111			if ctc is None: continue112			sess_other.send_event(event.PresenceNotificationEvent(ctc))113	114	def _sync_contact_statuses(self):115		# Recompute all `Contact.status`'s116		for user in self._user_by_uuid.values():117			detail = user.detail118			if detail is None: continue119			for ctc in detail.contacts.values():120				ctc.compute_visible_status(user)121	122	def _mark_modified(self, user, *, detail = None):123		ud = user.detail or detail124		if detail: assert ud is detail125		assert ud is not None126		self._unsynced_db[user] = ud127	128	def sb_token_create(self, sess, *, extra_data = None):129		if extra_data is None:130			extra_data = {}131		extra_data['client'] = sess.client132		return self._auth_service.create_token('sb/xfr', { 'uuid': sess.user.uuid, 'extra_data': extra_data })133	134	def me_update(self, sess, fields):135		user = sess.user136		137		if 'message' in fields:138			user.status.message = fields['message']139		if 'media' in fields:140			user.status.media = fields['media']141		if 'name' in fields:142			user.status.name = fields['name']143		if 'gtc' in fields:144			user.detail.settings['gtc'] = fields['gtc']145		if 'blp' in fields:146			user.detail.settings['blp'] = fields['blp']147		if 'substatus' in fields:148			user.status.substatus = fields['substatus']149		150		self._mark_modified(user)151		self._sync_contact_statuses()152		self._generic_notify(sess)153	154	def me_group_add(self, sess, name, *, is_favorite = None):155		if len(name) > MAX_GROUP_NAME_LENGTH:156			raise error.GroupNameTooLong()157		user = sess.user158		group = Group(_gen_group_id(user.detail), name, is_favorite = is_favorite)159		user.detail.groups[group.id] = group160		self._mark_modified(user)161		return group162	163	def me_group_remove(self, sess, group_id):164		if group_id == '0':165			raise error.CannotRemoveSpecialGroup()166		user = sess.user167		try:168			del user.detail.groups[group_id]169		except KeyError:170			raise error.GroupDoesNotExist()171		for ctc in user.detail.contacts.values():172			ctc.groups.discard(group_id)173		self._mark_modified(user)174	175	def me_group_edit(self, sess, group_id, new_name, *, is_favorite = None):176		user = sess.user177		g = user.detail.groups.get(group_id)178		if g is None:179			raise error.GroupDoesNotExist()180		if new_name is not None:181			if len(new_name) > MAX_GROUP_NAME_LENGTH:182				raise error.GroupNameTooLong()183			g.new_name = new_name184		if is_favorite is not None:185			g.is_favorite = is_favorite186		self._mark_modified(user)187	188	def me_group_contact_add(self, sess, group_id, contact_uuid):189		if group_id == '0': return190		user = sess.user191		detail = user.detail192		if group_id not in detail.groups:193			raise error.GroupDoesNotExist()194		ctc = detail.contacts.get(contact_uuid)195		if ctc is None:196			raise error.ContactDoesNotExist()197		if group_id in ctc.groups:198			raise error.ContactAlreadyOnList()199		ctc.groups.add(group_id)200		self._mark_modified(user)201	202	def me_group_contact_remove(self, sess, group_id, contact_uuid):203		user = sess.user204		detail = user.detail205		ctc = detail.contacts.get(contact_uuid)206		if ctc is None:207			raise error.ContactDoesNotExist()208		if group_id not in detail.groups and group_id != '0':209			raise error.GroupDoesNotExist()210		try:211			ctc.groups.remove(group_id)212		except KeyError:213			if group_id == '0':214				raise error.ContactNotOnList()215		self._mark_modified(user)216	217	def me_contact_add(self, sess, contact_uuid, lst, name):218		ctc_head = self._load_user_record(contact_uuid)219		if ctc_head is None:220			raise error.UserDoesNotExist()221		user = sess.user222		ctc = self._add_to_list(user, ctc_head, lst, name)223		if lst is Lst.FL:224			# FL needs a matching RL on the contact225			self._add_to_list(ctc_head, user, Lst.RL, user.status.name)226			self._notify_reverse_add(sess, ctc_head)227		self._sync_contact_statuses()228		self._generic_notify(sess)229		return ctc, ctc_head230	231	def _notify_reverse_add(self, sess, user_added):232		user_adder = sess.user233		# `user_added` was added to `user_adder`'s RL234		for sess_added in self._sc.get_sessions_by_user(user_added):235			if sess_added == sess: continue236			sess_added.send_event(event.AddedToListEvent(Lst.RL, user_adder))237	238	def me_contact_edit(self, sess, contact_uuid, *, is_messenger_user = None):239		user = sess.user240		ctc = user.detail.contacts.get(contact_uuid)241		if ctc is None:242			raise error.ContactDoesNotExist()243		if is_messenger_user is not None:244			ctc.is_messenger_user = is_messenger_user245		self._mark_modified(user)246	247	def me_contact_remove(self, sess, contact_uuid, lst):248		user = sess.user249		ctc = user.detail.contacts.get(contact_uuid)250		if ctc is None:251			raise error.ContactDoesNotExist()252		if lst is Lst.FL:253			# Remove from FL254			self._remove_from_list(user, ctc.head, Lst.FL)255			# Remove matching RL256			self._remove_from_list(ctc.head, user, Lst.RL)257		else:258			assert lst is not Lst.RL259			ctc.lists &= ~lst260		self._mark_modified(user)261		self._sync_contact_statuses()262	263	def _add_to_list(self, user, ctc_head, lst, name):264		# Add `ctc_head` to `user`'s `lst`265		detail = self._load_detail(user)266		contacts = detail.contacts267		if ctc_head.uuid not in contacts:268			contacts[ctc_head.uuid] = Contact(ctc_head, set(), 0, UserStatus(name))269		ctc = contacts.get(ctc_head.uuid)270		if ctc.status.name is None:271			ctc.status.name = name272		ctc.lists |= lst273		self._mark_modified(user, detail = detail)274		return ctc275	276	def _remove_from_list(self, user, ctc_head, lst):277		# Remove `ctc_head` from `user`'s `lst`278		detail = self._load_detail(user)279		contacts = detail.contacts280		ctc = contacts.get(ctc_head.uuid)281		if ctc is None: return282		ctc.lists &= ~lst283		if not ctc.lists:284			del contacts[ctc_head.uuid]285		self._mark_modified(user, detail = detail)286	287	def me_pop_boot_others(self, sess):288		for sess_other in self._sc.get_sessions_by_user(sess.user):289			if sess is sess_other: continue290			sess_other.send_event(event.POPBootEvent())291	292	def me_pop_notify_others(self, sess):293		for sess_other in self._sc.get_sessions_by_user(sess.user):294			if sess is sess_other: continue295			sess_other.send_event(event.POPNotifyEvent())296	297	def login_xfr(self, sess, email, token):298		(user, extra_data) = self._load_user('sb/xfr', token)299		if user is None: return None300		if user.email != email: return None301		sess.user = user302		sess.client = extra_data['client']303		chat = Chat(self._stats)304		self._chats[chat.id] = chat305		chat.add_session(sess)306		return chat, extra_data307	308	def login_cal(self, sess, email, token, chatid):309		(user, extra_data) = self._load_user('sb/cal', token)310		if user is None: return None311		if user.email != email: return None312		sess.user = user313		sess.client = extra_data['client']314		chat = self._chats.get(chatid)315		if chat is None: return None316		chat.add_session(sess)317		return chat, extra_data318	319	def _load_user(self, purpose, token):320		data = self._auth_service.pop_token(purpose, token)321		if data is None: return (None, None)322		return (self._user_service.get(data['uuid']), data['extra_data'])323	324	def util_get_uuid_from_email(self, email):325		return self._user_service.get_uuid(email)326	327	def util_set_sess_token(self, sess, token):328		self._sc.set_nc_by_token(sess, token)329	330	def util_get_sess_by_token(self, token):331		return self._sc.get_nc_by_token(token)332	333	def util_get_sessions_by_user(self, user):334		return self._sc.get_sessions_by_user(user)335	336	def notify_call(self, caller_uuid, callee_email, chatid):337		caller = self._user_by_uuid.get(caller_uuid)338		if caller is None: raise error.ServerError()339		if caller.detail is None: raise error.ServerError()340		callee_uuid = self._user_service.get_uuid(callee_email)341		if callee_uuid is None: raise error.UserDoesNotExist()342		ctc = caller.detail.contacts.get(callee_uuid)343		if ctc is None:344			if callee_uuid != caller_uuid: raise error.ContactDoesNotExist()345			ctc_user = caller346		else:347			if ctc.status.is_offlineish(): raise error.ContactNotOnline()348			ctc_user = ctc.head349		ctc_sessions = self._sc.get_sessions_by_user(ctc_user)350		if not ctc_sessions: raise error.ContactNotOnline()351		352		for ctc_sess in ctc_sessions:353			extra_data = ctc_sess.state.get_sb_extra_data() or {}354			extra_data['client'] = ctc_sess.client355			token = self._auth_service.create_token('sb/cal', { 'uuid': ctc_user.uuid, 'extra_data': extra_data })356			ctc_sess.send_event(event.InvitedToChatEvent(chatid, token, caller))357	358	async def _sync_db(self):359		while True:360			await asyncio.sleep(1)361			self._sync_db_impl()362	363	def _sync_db_impl(self):364		if not self._unsynced_db: return365		try:366			users = list(self._unsynced_db.keys())[:100]367			batch = []368			for user in users:369				detail = self._unsynced_db.pop(user, None)370				if not detail: continue371				batch.append((user, detail))372			self._user_service.save_batch(batch)373		except Exception:374			import traceback375			traceback.print_exc()376	377	async def _clean_sessions(self):378		from .session import PollingSession379		while True:380			await asyncio.sleep(10)381			now = time.time()382			closed = []383			384			try:385				for sess in self._sc.iter_sessions():386					if sess.closed:387						closed.append(sess)388						continue389					if isinstance(sess, PollingSession):390						if now >= sess.time_last_connect + sess.timeout:391							sess.close()392							closed.append(sess)393			except Exception:394				import traceback395				traceback.print_exc()396			397			for sess in closed:398				self._sc.remove_session(sess)399	400	async def _sync_stats(self):401		while True:402			await asyncio.sleep(60)403			try:404				self._stats.flush()405			except Exception:406				import traceback407				traceback.print_exc()408class _SessionCollection:409	def __init__(self):410		# Set[Session]411		self._sessions = set()412		# Dict[User, Set[Session]]413		self._sessions_by_user = defaultdict(set)414		# Dict[str, Session]415		self._sess_by_token = {}416		# Dict[Session, Set[str]]417		self._tokens_by_sess = defaultdict(set)418	419	def get_sessions_by_user(self, user):420		if user not in self._sessions_by_user:421			return EMPTY_SET422		return self._sessions_by_user[user]423	424	def iter_sessions(self):425		yield from self._sessions426	427	def set_nc_by_token(self, sess, token: str):428		self._sess_by_token[token] = sess429		self._tokens_by_sess[sess].add(sess)430		self._sessions.add(sess)431	432	def get_nc_by_token(self, token: str):433		return self._sess_by_token.get(token)434	435	def add_session(self, sess):436		if sess.user:437			self._sessions_by_user[sess.user].add(sess)438		self._sessions.add(sess)439	440	def remove_session(self, sess):441		if sess in self._tokens_by_sess:442			tokens = self._tokens_by_sess.pop(sess)443			for token in tokens:444				self._sess_by_token.pop(token, None)445		self._sessions.discard(sess)446		if sess.user in self._sessions_by_user:447			self._sessions_by_user[sess.user].discard(sess)448class Chat:449	def __init__(self, stats):450		self.id = gen_uuid()451		# Dict[Session, User]452		self._users_by_sess = {}453		self._stats = stats454	455	def add_session(self, sess):456		self._users_by_sess[sess] = sess.user457	458	def send_message_to_everyone(self, sess_sender, data):459		self._stats.on_message_sent(sess_sender.user, sess_sender.client)460		self._stats.on_user_active(sess_sender.user, sess_sender.client)461		su_sender = self._users_by_sess[sess_sender]462		for sess in self._users_by_sess.keys():463			if sess == sess_sender: continue464			sess.send_event(event.ChatMessage(su_sender, data))465			self._stats.on_message_received(sess.user, sess.client)466	467	def get_roster(self, sess):468		roster = []469		for sess1, su1 in self._users_by_sess.items():470			if sess1 == sess: continue471			roster.append((sess1, su1))472		return roster473	474	def send_participant_joined(self, sess):475		for sc, _ in self.get_roster(self):476			sc.send_event(event.ChatParticipantJoined(sess))477	478	def on_leave(self, sess):479		su = self._users_by_sess.pop(sess, None)480		if su is None: return481		# Notify others that `sess` has left482		for sess1, su1 in self._users_by_sess.items():483			if sess1 == sess: continue484			sess1.send_event(event.ChatParticipantLeft(su))485def _gen_group_id(detail):486	id = 1487	s = str(id)488	while s in detail.groups:489		id += 1490		s = str(id)491	return s...

gen_dsin_input.py

Source:gen_dsin_input.py

1import os2import numpy as np3import pandas as pd4from deepctr.utils import SingleFeat5from sklearn.preprocessing import LabelEncoder, StandardScaler6from tensorflow.python.keras.preprocessing.sequence import pad_sequences7from tqdm import tqdm8from config import DSIN_SESS_COUNT, DSIN_SESS_MAX_LEN, FRAC9FRAC = FRAC10SESS_COUNT = DSIN_SESS_COUNT11def gen_sess_feature_dsin(row):12    sess_count = DSIN_SESS_COUNT13    sess_max_len = DSIN_SESS_MAX_LEN14    sess_input_dict = {}15    sess_input_length_dict = {}16    for i in range(sess_count):17        sess_input_dict['sess_' + str(i)] = {'cate_id': [], 'brand': []}18        sess_input_length_dict['sess_' + str(i)] = 019    sess_length = 020    user, time_stamp = row[1]['user'], row[1]['time_stamp']21    # sample_time = pd.to_datetime(timestamp_datetime(time_stamp ))22    if user not in user_hist_session:23        for i in range(sess_count):24            sess_input_dict['sess_' + str(i)]['cate_id'] = [0]25            sess_input_dict['sess_' + str(i)]['brand'] = [0]26            sess_input_length_dict['sess_' + str(i)] = 027        sess_length = 028    else:29        valid_sess_count = 030        last_sess_idx = len(user_hist_session[user]) - 131        for i in reversed(range(len(user_hist_session[user]))):32            cur_sess = user_hist_session[user][i]33            if cur_sess[0][2] < time_stamp:34                in_sess_count = 135                for j in range(1, len(cur_sess)):36                    if cur_sess[j][2] < time_stamp:37                        in_sess_count += 138                if in_sess_count > 2:39                    sess_input_dict['sess_0']['cate_id'] = [e[0] for e in cur_sess[max(0,40                                                                                       in_sess_count - sess_max_len):in_sess_count]]41                    sess_input_dict['sess_0']['brand'] = [e[1] for e in42                                                          cur_sess[max(0, in_sess_count - sess_max_len):in_sess_count]]43                    sess_input_length_dict['sess_0'] = min(44                        sess_max_len, in_sess_count)45                    last_sess_idx = i46                    valid_sess_count += 147                    break48        for i in range(1, sess_count):49            if last_sess_idx - i >= 0:50                cur_sess = user_hist_session[user][last_sess_idx - i]51                sess_input_dict['sess_' + str(i)]['cate_id'] = [e[0]52                                                                for e in cur_sess[-sess_max_len:]]53                sess_input_dict['sess_' + str(i)]['brand'] = [e[1]54                                                              for e in cur_sess[-sess_max_len:]]55                sess_input_length_dict['sess_' +56                                       str(i)] = min(sess_max_len, len(cur_sess))57                valid_sess_count += 158            else:59                sess_input_dict['sess_' + str(i)]['cate_id'] = [0]60                sess_input_dict['sess_' + str(i)]['brand'] = [0]61                sess_input_length_dict['sess_' + str(i)] = 062        sess_length = valid_sess_count63    return sess_input_dict, sess_input_length_dict, sess_length64if __name__ == "__main__":65    user_hist_session = {}66    FILE_NUM = len(67        list(filter(lambda x: x.startswith('user_hist_session_' + str(FRAC) + '_dsin_'),68                    os.listdir('../sampled_data/'))))69    print('total', FILE_NUM, 'files')70    for i in range(FILE_NUM):71        user_hist_session_ = pd.read_pickle(72            '../sampled_data/user_hist_session_' + str(FRAC) + '_dsin_' + str(i) + '.pkl')  # 19,3473        user_hist_session.update(user_hist_session_)74        del user_hist_session_75    sample_sub = pd.read_pickle(76        '../sampled_data/raw_sample_' + str(FRAC) + '.pkl')77    index_list = []78    sess_input_dict = {}79    sess_input_length_dict = {}80    for i in range(SESS_COUNT):81        sess_input_dict['sess_' + str(i)] = {'cate_id': [], 'brand': []}82        sess_input_length_dict['sess_' + str(i)] = []83    sess_length_list = []84    for row in tqdm(sample_sub[['user', 'time_stamp']].iterrows()):85        sess_input_dict_, sess_input_length_dict_, sess_length = gen_sess_feature_dsin(86            row)87        # index_list.append(index)88        for i in range(SESS_COUNT):89            sess_name = 'sess_' + str(i)90            sess_input_dict[sess_name]['cate_id'].append(91                sess_input_dict_[sess_name]['cate_id'])92            sess_input_dict[sess_name]['brand'].append(93                sess_input_dict_[sess_name]['brand'])94            sess_input_length_dict[sess_name].append(95                sess_input_length_dict_[sess_name])96        sess_length_list.append(sess_length)97    print('done')98    user = pd.read_pickle('../sampled_data/user_profile_' + str(FRAC) + '.pkl')99    ad = pd.read_pickle('../sampled_data/ad_feature_enc_' + str(FRAC) + '.pkl')100    user = user.fillna(-1)101    user.rename(102        columns={'new_user_class_level ': 'new_user_class_level'}, inplace=True)103    sample_sub = pd.read_pickle(104        '../sampled_data/raw_sample_' + str(FRAC) + '.pkl')105    sample_sub.rename(columns={'user': 'userid'}, inplace=True)106    data = pd.merge(sample_sub, user, how='left', on='userid', )107    data = pd.merge(data, ad, how='left', on='adgroup_id')108    sparse_features = ['userid', 'adgroup_id', 'pid', 'cms_segid', 'cms_group_id', 'final_gender_code', 'age_level',109                       'pvalue_level', 'shopping_level', 'occupation', 'new_user_class_level', 'campaign_id',110                       'customer'] # sparse feature for user and ads111    dense_features = ['price'] # dense feature for user and ads112    for feat in tqdm(sparse_features):113        lbe = LabelEncoder()  # or Hash114        data[feat] = lbe.fit_transform(data[feat]) # å°ä¸åçåå¼è½¬æ¢ä¸ºå¯¹åºçç¼å·115    mms = StandardScaler()116    data[dense_features] = mms.fit_transform(data[dense_features])117    # class SingleFeat(namedtuple('SingleFeat', ['name', 'dimension', 'hash_flag', 'dtype'])):118    sparse_feature_list = [SingleFeat(feat, data[feat].nunique(119    ) + 1) for feat in sparse_features + ['cate_id', 'brand']]120    dense_feature_list = [SingleFeat(feat, 1) for feat in dense_features]121    sess_feature = ['cate_id', 'brand'] # sess feature for ad122    sess_input = []123    sess_input_length = []124    for i in tqdm(range(SESS_COUNT)):125        sess_name = 'sess_' + str(i)126        for feat in sess_feature:127            sess_input.append(pad_sequences(128                sess_input_dict[sess_name][feat], maxlen=SESS_COUNT, padding='post'))129        sess_input_length.append(sess_input_length_dict[sess_name])130    model_input = [data[feat.name].values for feat in sparse_feature_list] + \131                  [data[feat.name].values for feat in dense_feature_list]132    sess_lists = sess_input + [np.array(sess_length_list)]133    model_input += sess_lists134    if not os.path.exists('../model_input/'):135        os.mkdir('../model_input/')136    pd.to_pickle(model_input, '../model_input/dsin_input_' +137                 str(FRAC) + '_' + str(SESS_COUNT) + '.pkl')138    pd.to_pickle(data['clk'].values, '../model_input/dsin_label_' +139                 str(FRAC) + '_' + str(SESS_COUNT) + '.pkl')140    pd.to_pickle({'sparse': sparse_feature_list, 'dense': dense_feature_list},141                 '../model_input/dsin_fd_' + str(FRAC) + '_' + str(SESS_COUNT) + '.pkl')...

Pytest Tutorial

Looking for an in-depth tutorial around pytest? LambdaTest covers the detailed pytest tutorial that has everything related to the pytest, from setting up the pytest framework to automation testing. Delve deeper into pytest testing by exploring advanced use cases like parallel testing, pytest fixtures, parameterization, executing multiple test cases from a single file, and more.

Chapters

What is pytest
Pytest installation: Want to start pytest from scratch? See how to install and configure pytest for Python automation testing.
Run first test with pytest framework: Follow this step-by-step tutorial to write and run your first pytest script.
Parallel testing with pytest: A hands-on guide to parallel testing with pytest to improve the scalability of your test automation.
Generate pytest reports: Reports make it easier to understand the results of pytest-based test runs. Learn how to generate pytest reports.
Pytest Parameterized tests: Create and run your pytest scripts while avoiding code duplication and increasing test coverage with parameterization.
Pytest Fixtures: Check out how to implement pytest fixtures for your end-to-end testing needs.
Execute Multiple Test Cases: Explore different scenarios for running multiple test cases in pytest from a single file.
Stop Test Suite after N Test Failures: See how to stop your test suite after n test failures in pytest using the @pytest.mark.incremental decorator and maxfail command-line option.