<commit_before>import baseline as bl
import baseline.tf.embeddings
import baseline.tf.lm
import numpy as np
import os
import codecs
import argparse
import logging
import tensorflow as tf
import glob

SOS = "<GO>"
EOS = "<EOS>"
LOGGER = logging.getLogger('baseline.timing')
NUM_PRE = 8
NUM_PARALLEL_CALLS = 1
SHUF_BUF_SZ = 1000
EMBED_TYPE = 'default'

def get_logging_level(ll):
    ll = ll.lower()
    if ll == 'debug':
        return logging.DEBUG
    if ll == 'info':
        return logging.INFO
    return logging.WARNING

def get_tf_logging_level(ll):
    ll = ll.lower()
    if ll == 'debug':
        return tf.logging.DEBUG
    if ll == 'info':
        return logging.INFO
    return tf.logging.WARN

parser = argparse.ArgumentParser(description='Train a Baseline Language Model with TensorFlow Estimator API')
parser.add_argument('--export_dir', help='Directory for TF export (for serving)', default='./models', type=str)
parser.add_argument('--checkpoint_dir', help='Directory for model checkpoints', default='./tmp/lm', type=str)
# Unused currently, just using default
# parser.add_argument('--model_type', help='What type of model to build', type=str, default='default')
parser.add_argument('--hsz', help='How many hidden units', type=int, default=100)
parser.add_argument('--dsz', help='How many embedding dimensions', type=int, default=150)
parser.add_argument('--steps', help='Number of steps to train', type=int, default=1500)
parser.add_argument('--eval_steps', help='Number of steps to eval', type=int, default=500)
parser.add_argument('--checkpoint_steps', help='Number of steps at which to checkpoint', type=int, default=500)
parser.add_argument('--batchsz', help='Batch size', type=int, default=20)
parser.add_argument('--nctx', help='Size of context or number of steps per line', type=int, default=20)
# Unused currently
# parser.add_argument('--filts', help='Parallel convolution filter widths (if default model)',
#                    type=int,
#                    default=[3, 4, 5],
#                    nargs='+')
parser.add_argument('--train', help='Training file', default="./training-monolingual.tokenized.shuffled/news.en-*")
parser.add_argument('--valid',
                    help='Validation file',
                    default='./heldout-monolingual.tokenized.shuffled/news.en.heldout*')
parser.add_argument('--ll', help='Log level', type=str, default='info')
parser.add_argument('--tf_ll', help='TensorFlow Log level', type=str, default='warn')
parser.add_argument('--lr', help='Learning rate', type=float, default=0.001)
parser.add_argument('--optim', help='Optimizer (sgd, adam) (default is adam)', type=str, default='adam')
parser.add_argument('--gpus', help='Number of GPUs to use', default=1, type=int)
parser.add_argument('--vocab_file', help='Location of a vocab file', default='./vocab-2016-09-10.txt')
args = parser.parse_known_args()[0]
train_batchsz = args.batchsz // args.gpus

logging.basicConfig(level=get_logging_level(args.ll))
tf.compat.v1.logging.set_verbosity(get_tf_logging_level(args.tf_ll))

@bl.str_file
def read_vocab(f):
    return {x.strip(): 1 for x in f.readlines()}

class MultiFileLMDataset(object):
    """Minimally modified from https://github.com/rafaljozefowicz/lm
    """
    def __init__(self, vocab, file_pattern):
        """Constructor

        :param vocab: A dictionary mapping words to indices, <PAD>=0, and assumes a <GO> and <EOS> token defined
        :param file_pattern: This is any pattern that can be used to glob for files
        """
        self._vocab = vocab
        self._file_pattern = file_pattern

    def _parse_sentence(self, line):
        return [self._vocab.get(SOS)] + \
               [self._vocab.get(word, 0) for word in line.strip().split()] + \
               [self._vocab.get(EOS)]

    def _parse_file(self, file_name):
        print("Processing file: %s" % file_name)
        with codecs.open(file_name, "r", "utf-8") as f:
            lines = [line.strip() for line in f]
            for line in lines:
                yield self._parse_sentence(line)

    def _sentence_stream(self, file_stream):
        for file_name in file_stream:
            for sentence in self._parse_file(file_name):
                yield sentence

    def _iterate(self, sentences, batch_size, num_steps):
        streams = [None] * batch_size
        x = np.zeros([batch_size, num_steps], np.int32)
        y = np.zeros([batch_size, num_steps], np.int32)
        w = np.zeros([batch_size, num_steps], np.uint8)
        while True:
            x[:] = 0
            y[:] = 0
            w[:] = 0
            for i in range(batch_size):
                tokens_filled = 0
                try:
                    while tokens_filled < num_steps:
                        if streams[i] is None or len(streams[i]) <= 1:
                            streams[i] = next(sentences)
                        num_tokens = min(len(streams[i]) - 1, num_steps - tokens_filled)
                        x[i, tokens_filled:tokens_filled+num_tokens] = streams[i][:num_tokens]
                        y[i, tokens_filled:tokens_filled + num_tokens] = streams[i][1:num_tokens+1]
                        w[i, tokens_filled:tokens_filled + num_tokens] = 1
                        streams[i] = streams[i][num_tokens:]
                        tokens_filled += num_tokens
                except StopIteration:
                    pass
            if not np.any(w):
                return

            yield x, y, w

    def iterate_once(self, batch_size, num_steps):
        """A generator that does a single pass over the data

        :param batch_size: The number of lines in the batch
        :param num_steps: The number of timesteps in each line to process
        :return: The next batch
        """
        def file_stream():
            for file_name in glob.glob(self._file_pattern):
                yield file_name
        for value in self._iterate(self._sentence_stream(file_stream()), batch_size, num_steps):
            yield value

    def iterate_forever(self, batch_size, num_steps):
        """A generator that does runs over the data infinitely

        :param batch_size: The number of lines in the batch
        :param num_steps: The number of timesteps in each line to process
        :return: The next batch
        """
        def file_stream():
            while True:
                file_patterns = glob.glob(self._file_pattern)
                for file_name in file_patterns:
                    yield file_name
        for value in self._iterate(self._sentence_stream(file_stream()), batch_size, num_steps):
            yield value

# Read the vocabulary file
vocab = read_vocab(args.vocab_file)

embedding_bundle = bl.embeddings.load_embeddings('x',
                                                 dsz=args.dsz,
                                                 known_vocab=vocab,
                                                 embed_type=EMBED_TYPE)
embeddings_x = embedding_bundle['embeddings']
vocab = embedding_bundle['vocab']

idx2word = {v: k for k, v in vocab.items()}

def word_vecs(x):
    print('[')
    for i in range(len(x)):
        print('\t{}'.format('\t'.join([idx2word[xv] for xv in x[i]])))
    print(']')

def train_input_fn():

    d = MultiFileLMDataset(vocab, args.train)

    def gen():
        return d.iterate_forever(train_batchsz, args.nctx)

    dataset = tf.data.Dataset.from_generator(gen,
                                             (tf.int32, tf.int32, tf.int8),
                                             (tf.TensorShape([None, None]),
                                              tf.TensorShape([None, None]),
                                              tf.TensorShape([None, None])))

    dataset = dataset.prefetch(NUM_PRE)
    it = dataset.make_one_shot_iterator()
    x, y, lengths = it.get_next()
    return {'x': x, 'lengths': lengths}, y

def eval_input_fn():

    d = MultiFileLMDataset(vocab, args.valid)

    def gen():
        return d.iterate_once(args.batchsz, args.nctx)

    dataset = tf.data.Dataset.from_generator(gen,
                                             (tf.int32, tf.int32, tf.int8),
                                             (tf.TensorShape([None, None]),
                                              tf.TensorShape([None, None]),
                                              tf.TensorShape([None, None])))

    dataset = dataset.prefetch(NUM_PRE)
    it = dataset.make_one_shot_iterator()
    x, y, lengths = it.get_next()
    return {'x': x, 'lengths': lengths}, y

def model_fn(features, labels, mode, params):

    if mode == tf.estimator.ModeKeys.EVAL:
        bl.tf.SET_TRAIN_FLAG(False)
        model = bl.create_lang_model({'x': embeddings_x},
                                     x=features['x'],
                                     y=labels,
                                     hsz=args.hsz,
                                     sess=None,
                                     tgt_key='x')
        loss = model.create_test_loss()
        (log_perplexity_tensor,
         log_perplexity_update) = tf.metrics.mean(loss)

        perplexity_tensor = tf.exp(log_perplexity_tensor)
        eval_metric_ops = {
            'perplexity': (perplexity_tensor, log_perplexity_update),
        }
        return tf.estimator.EstimatorSpec(mode=mode, predictions=model.logits, loss=loss, eval_metric_ops=eval_metric_ops)
    bl.tf.SET_TRAIN_FLAG(True)
    model = bl.create_lang_model({'x': embeddings_x},
                                 x=features['x'],
                                 y=labels,
                                 hsz=args.hsz,
                                 sess=None,
                                 src_keys=['x'],
                                 tgt_key='x')
    loss = model.create_loss()
    colocate = True if args.gpus > 1 else False
    global_step, train_op = bl.tf.optz.optimizer(loss, optim=args.optim, eta=args.lr, colocate_gradients_with_ops=colocate)

    return tf.estimator.EstimatorSpec(mode=mode, predictions=model.logits,
                                      loss=loss,
                                      train_op=train_op)

checkpoint_dir = '{}-{}'.format(args.checkpoint_dir, os.getpid())
if args.gpus > 1:
    config = tf.estimator.RunConfig(model_dir=checkpoint_dir,
                                    save_checkpoints_steps=args.checkpoint_steps,
                                    train_distribute=tf.contrib.distribute.MirroredStrategy(num_gpus=args.gpus))
else:
    config = tf.estimator.RunConfig(model_dir=checkpoint_dir,
                                    save_checkpoints_steps=args.checkpoint_steps)

estimator = tf.estimator.Estimator(model_fn=model_fn, config=config, params={})
train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, max_steps=args.steps)
eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn, steps=args.eval_steps)
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

<commit_msg>remove this example<commit_after><|endoftext|>


<commit_before># Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Executes RetinaNet benchmarks and accuracy tests."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# pylint: disable=g-bad-import-order
import copy
import json
import os
import time

from absl import flags
from absl import logging
from absl.testing import flagsaver
import tensorflow as tf
# pylint: enable=g-bad-import-order

from official.benchmark import bert_benchmark_utils as benchmark_utils
from official.utils.flags import core as flags_core
from official.utils.testing import benchmark_wrappers
from official.vision.detection import main as detection

TMP_DIR = os.getenv('TMPDIR')
FLAGS = flags.FLAGS

# pylint: disable=line-too-long
COCO_TRAIN_DATA = 'gs://tf-perfzero-data/coco/train*'
COCO_EVAL_DATA = 'gs://tf-perfzero-data/coco/val*'
COCO_EVAL_JSON = 'gs://tf-perfzero-data/coco/instances_val2017.json'
RESNET_CHECKPOINT_PATH = 'gs://cloud-tpu-checkpoints/retinanet/resnet50-checkpoint-2018-02-07'
# pylint: enable=line-too-long

class DetectionBenchmarkBase(tf.test.Benchmark):
  """Base class to hold methods common to test classes."""
  local_flags = None

  def __init__(self, output_dir=None):
    self.num_gpus = 8

    if not output_dir:
      output_dir = '/tmp'
    self.output_dir = output_dir
    self.timer_callback = None

  def _get_model_dir(self, folder_name):
    """Returns directory to store info, e.g. saved model and event log."""
    return os.path.join(self.output_dir, folder_name)

  def _setup(self):
    """Sets up and resets flags before each test."""
    self.timer_callback = benchmark_utils.BenchmarkTimerCallback()

    if DetectionBenchmarkBase.local_flags is None:
      # Loads flags to get defaults to then override. List cannot be empty.
      flags.FLAGS(['foo'])
      saved_flag_values = flagsaver.save_flag_values()
      DetectionBenchmarkBase.local_flags = saved_flag_values
    else:
      flagsaver.restore_flag_values(DetectionBenchmarkBase.local_flags)

  def _report_benchmark(self,
                        stats,
                        wall_time_sec,
                        min_ap,
                        max_ap,
                        train_batch_size=None):
    """Report benchmark results by writing to local protobuf file.

    Args:
      stats: dict returned from Detection models with known entries.
      wall_time_sec: the during of the benchmark execution in seconds
      min_ap: Minimum detection AP constraint to verify correctness of the
        model.
      max_ap: Maximum detection AP accuracy constraint to verify correctness of
        the model.
      train_batch_size: Train batch size. It is needed for computing
        exp_per_second.
    """
    metrics = [{
        'name': 'total_loss',
        'value': stats['total_loss'],
    }]
    if self.timer_callback:
      metrics.append({
          'name': 'exp_per_second',
          'value': self.timer_callback.get_examples_per_sec(train_batch_size)
      })
    else:
      metrics.append({
          'name': 'exp_per_second',
          'value': 0.0,
      })

    if 'eval_metrics' in stats:
      metrics.append({
          'name': 'AP',
          'value': stats['AP'],
          'min_value': min_ap,
          'max_value': max_ap,
      })
    flags_str = flags_core.get_nondefault_flags_as_str()
    self.report_benchmark(
        iters=stats['total_steps'],
        wall_time=wall_time_sec,
        metrics=metrics,
        extras={'flags': flags_str})

class RetinanetBenchmarkBase(DetectionBenchmarkBase):
  """Base class to hold methods common to test classes in the module."""

  def __init__(self, output_dir=None, **kwargs):
    self.train_data_path = COCO_TRAIN_DATA
    self.eval_data_path = COCO_EVAL_DATA
    self.eval_json_path = COCO_EVAL_JSON
    self.resnet_checkpoint_path = RESNET_CHECKPOINT_PATH

    super(RetinanetBenchmarkBase, self).__init__(output_dir=output_dir)

  def _run_detection_main(self):
    """Starts detection job."""
    if self.timer_callback:
      return detection.run(callbacks=[self.timer_callback])
    else:
      return detection.run()

class RetinanetAccuracy(RetinanetBenchmarkBase):
  """Accuracy test for RetinaNet model.

  Tests RetinaNet detection task model accuracy. The naming
  convention of below test cases follow
  `benchmark_(number of gpus)_gpu_(dataset type)` format.
  """

  def __init__(self, output_dir=TMP_DIR, **kwargs):
    super(RetinanetAccuracy, self).__init__(output_dir=output_dir)

  @benchmark_wrappers.enable_runtime_flags
  def _run_and_report_benchmark(self, min_ap=0.325, max_ap=0.35):
    """Starts RetinaNet accuracy benchmark test."""

    start_time_sec = time.time()
    FLAGS.mode = 'train'
    summary, _ = self._run_detection_main()
    wall_time_sec = time.time() - start_time_sec

    FLAGS.mode = 'eval'
    eval_metrics = self._run_detection_main()
    summary.update(eval_metrics)

    summary['train_batch_size'] = self.params_override['train']['batch_size']
    summary['total_steps'] = self.params_override['train']['total_steps']
    super(RetinanetAccuracy, self)._report_benchmark(
        stats=summary,
        wall_time_sec=wall_time_sec,
        min_ap=min_ap,
        max_ap=max_ap,
        train_batch_size=self.params_override['train']['batch_size'])

  def _setup(self):
    super(RetinanetAccuracy, self)._setup()
    FLAGS.strategy_type = 'mirrored'
    FLAGS.model = 'retinanet'

    self.params_override = {
        'train': {
            'batch_size': 64,
            'iterations_per_loop': 100,
            'total_steps': 22500,
            'train_file_pattern': self.train_data_path,
            'checkpoint': {
                'path': self.resnet_checkpoint_path,
                'prefix': 'resnet50/'
            },
        },
        'eval': {
            'batch_size': 8,
            'eval_samples': 5000,
            'val_json_file': self.eval_json_path,
            'eval_file_pattern': self.eval_data_path,
        },
    }

  @flagsaver.flagsaver
  def benchmark_8_gpu_coco(self):
    """Run RetinaNet model accuracy test with 8 GPUs."""
    self._setup()
    params = copy.deepcopy(self.params_override)
    FLAGS.params_override = json.dumps(params)
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_coco')
    # Sets timer_callback to None as we do not use it now.
    self.timer_callback = None

    self._run_and_report_benchmark()

class RetinanetBenchmarkReal(RetinanetAccuracy):
  """Short benchmark performance tests for RetinaNet model.

  Tests RetinaNet performance in different GPU configurations.
  The naming convention of below test cases follow
  `benchmark_(number of gpus)_gpu` format.
  """

  def __init__(self, output_dir=TMP_DIR, **kwargs):
    super(RetinanetBenchmarkReal, self).__init__(output_dir=output_dir)

  @flagsaver.flagsaver
  def benchmark_8_gpu_coco(self):
    """Run RetinaNet model accuracy test with 8 GPUs."""
    self._setup()
    params = copy.deepcopy(self.params_override)
    params['train']['total_steps'] = 1875  # One epoch.
    # The iterations_per_loop must be one, otherwise the number of examples per
    # second would be wrong. Currently only support calling callback per batch
    # when each loop only runs on one batch, i.e. host loop for one step. The
    # performance of this situation might be lower than the case of
    # iterations_per_loop > 1.
    # Related bug: b/135933080
    params['train']['iterations_per_loop'] = 1
    params['eval']['eval_samples'] = 8
    FLAGS.params_override = json.dumps(params)
    FLAGS.model_dir = self._get_model_dir('real_benchmark_8_gpu_coco')
    # Use negative value to avoid saving checkpoints.
    FLAGS.save_checkpoint_freq = -1
    if self.timer_callback is None:
      logging.error('Cannot measure performance without timer callback')
    else:
      self._run_and_report_benchmark()

  @flagsaver.flagsaver
  def benchmark_1_gpu_coco(self):
    """Run RetinaNet model accuracy test with 1 GPU."""
    self.num_gpus = 1
    self._setup()
    params = copy.deepcopy(self.params_override)
    params['train']['batch_size'] = 8
    params['train']['total_steps'] = 200
    params['train']['iterations_per_loop'] = 1
    params['eval']['eval_samples'] = 8
    FLAGS.params_override = json.dumps(params)
    FLAGS.model_dir = self._get_model_dir('real_benchmark_1_gpu_coco')
    FLAGS.strategy_type = 'one_device_gpu'
    # Use negative value to avoid saving checkpoints.
    FLAGS.save_checkpoint_freq = -1
    if self.timer_callback is None:
      logging.error('Cannot measure performance without timer callback')
    else:
      self._run_and_report_benchmark()

  @flagsaver.flagsaver
  def benchmark_xla_1_gpu_coco(self):
    """Run RetinaNet model accuracy test with 1 GPU and XLA enabled."""
    self.num_gpus = 1
    self._setup()
    params = copy.deepcopy(self.params_override)
    params['train']['batch_size'] = 8
    params['train']['total_steps'] = 200
    params['train']['iterations_per_loop'] = 1
    params['eval']['eval_samples'] = 8
    FLAGS.params_override = json.dumps(params)
    FLAGS.model_dir = self._get_model_dir('real_benchmark_1_gpu_coco')
    FLAGS.strategy_type = 'one_device_gpu'
    FLAGS.enable_xla = True
    # Use negative value to avoid saving checkpoints.
    FLAGS.save_checkpoint_freq = -1
    if self.timer_callback is None:
      logging.error('Cannot measure performance without timer callback')
    else:
      self._run_and_report_benchmark()

if __name__ == '__main__':
  tf.test.main()
<commit_msg>Internal cleanup.<commit_after># Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Executes RetinaNet benchmarks and accuracy tests."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# pylint: disable=g-bad-import-order
import copy
import json
import os
import time

from absl import flags
from absl import logging
from absl.testing import flagsaver
import tensorflow as tf
# pylint: enable=g-bad-import-order

from official.benchmark import bert_benchmark_utils as benchmark_utils
from official.utils.flags import core as flags_core
from official.utils.testing import benchmark_wrappers
from official.vision.detection import main as detection

TMP_DIR = os.getenv('TMPDIR')
FLAGS = flags.FLAGS

# pylint: disable=line-too-long
COCO_TRAIN_DATA = 'gs://tf-perfzero-data/coco/train*'
COCO_EVAL_DATA = 'gs://tf-perfzero-data/coco/val*'
COCO_EVAL_JSON = 'gs://tf-perfzero-data/coco/instances_val2017.json'
RESNET_CHECKPOINT_PATH = 'gs://cloud-tpu-checkpoints/retinanet/resnet50-checkpoint-2018-02-07'
# pylint: enable=line-too-long

class DetectionBenchmarkBase(tf.test.Benchmark):
  """Base class to hold methods common to test classes."""
  local_flags = None

  def __init__(self, output_dir=None):
    self.num_gpus = 8

    if not output_dir:
      output_dir = '/tmp'
    self.output_dir = output_dir
    self.timer_callback = None

  def _get_model_dir(self, folder_name):
    """Returns directory to store info, e.g. saved model and event log."""
    return os.path.join(self.output_dir, folder_name)

  def _setup(self):
    """Sets up and resets flags before each test."""
    self.timer_callback = benchmark_utils.BenchmarkTimerCallback()

    if DetectionBenchmarkBase.local_flags is None:
      # Loads flags to get defaults to then override. List cannot be empty.
      flags.FLAGS(['foo'])
      saved_flag_values = flagsaver.save_flag_values()
      DetectionBenchmarkBase.local_flags = saved_flag_values
    else:
      flagsaver.restore_flag_values(DetectionBenchmarkBase.local_flags)

  def _report_benchmark(self,
                        stats,
                        wall_time_sec,
                        min_ap,
                        max_ap,
                        train_batch_size=None):
    """Report benchmark results by writing to local protobuf file.

    Args:
      stats: dict returned from Detection models with known entries.
      wall_time_sec: the during of the benchmark execution in seconds
      min_ap: Minimum detection AP constraint to verify correctness of the
        model.
      max_ap: Maximum detection AP accuracy constraint to verify correctness of
        the model.
      train_batch_size: Train batch size. It is needed for computing
        exp_per_second.
    """
    metrics = [{
        'name': 'total_loss',
        'value': stats['total_loss'],
    }]
    if self.timer_callback:
      metrics.append({
          'name': 'exp_per_second',
          'value': self.timer_callback.get_examples_per_sec(train_batch_size)
      })
    else:
      metrics.append({
          'name': 'exp_per_second',
          'value': 0.0,
      })

    if 'eval_metrics' in stats:
      metrics.append({
          'name': 'AP',
          'value': stats['AP'],
          'min_value': min_ap,
          'max_value': max_ap,
      })
    flags_str = flags_core.get_nondefault_flags_as_str()
    self.report_benchmark(
        iters=stats['total_steps'],
        wall_time=wall_time_sec,
        metrics=metrics,
        extras={'flags': flags_str})

class RetinanetBenchmarkBase(DetectionBenchmarkBase):
  """Base class to hold methods common to test classes in the module."""

  def __init__(self, output_dir=None, **kwargs):
    self.train_data_path = COCO_TRAIN_DATA
    self.eval_data_path = COCO_EVAL_DATA
    self.eval_json_path = COCO_EVAL_JSON
    self.resnet_checkpoint_path = RESNET_CHECKPOINT_PATH

    super(RetinanetBenchmarkBase, self).__init__(output_dir=output_dir)

  def _run_detection_main(self):
    """Starts detection job."""
    if self.timer_callback:
      return detection.run(callbacks=[self.timer_callback])
    else:
      return detection.run()

class RetinanetAccuracy(RetinanetBenchmarkBase):
  """Accuracy test for RetinaNet model.

  Tests RetinaNet detection task model accuracy. The naming
  convention of below test cases follow
  `benchmark_(number of gpus)_gpu_(dataset type)` format.
  """

  def __init__(self, output_dir=TMP_DIR, **kwargs):
    super(RetinanetAccuracy, self).__init__(output_dir=output_dir)

  @benchmark_wrappers.enable_runtime_flags
  def _run_and_report_benchmark(self, min_ap=0.325, max_ap=0.35):
    """Starts RetinaNet accuracy benchmark test."""

    start_time_sec = time.time()
    FLAGS.mode = 'train'
    summary, _ = self._run_detection_main()
    wall_time_sec = time.time() - start_time_sec

    FLAGS.mode = 'eval'
    eval_metrics = self._run_detection_main()
    summary.update(eval_metrics)

    summary['train_batch_size'] = self.params_override['train']['batch_size']
    summary['total_steps'] = self.params_override['train']['total_steps']
    super(RetinanetAccuracy, self)._report_benchmark(
        stats=summary,
        wall_time_sec=wall_time_sec,
        min_ap=min_ap,
        max_ap=max_ap,
        train_batch_size=self.params_override['train']['batch_size'])

  def _setup(self):
    super(RetinanetAccuracy, self)._setup()
    FLAGS.strategy_type = 'mirrored'
    FLAGS.model = 'retinanet'

    self.params_override = {
        'train': {
            'batch_size': 64,
            'iterations_per_loop': 100,
            'total_steps': 22500,
            'train_file_pattern': self.train_data_path,
            'checkpoint': {
                'path': self.resnet_checkpoint_path,
                'prefix': 'resnet50/'
            },
        },
        'eval': {
            'batch_size': 8,
            'eval_samples': 5000,
            'val_json_file': self.eval_json_path,
            'eval_file_pattern': self.eval_data_path,
        },
    }

  @flagsaver.flagsaver
  def benchmark_8_gpu_coco(self):
    """Run RetinaNet model accuracy test with 8 GPUs."""
    self._setup()
    params = copy.deepcopy(self.params_override)
    FLAGS.params_override = json.dumps(params)
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_coco')
    # Sets timer_callback to None as we do not use it now.
    self.timer_callback = None

    self._run_and_report_benchmark()

class RetinanetBenchmarkReal(RetinanetAccuracy):
  """Short benchmark performance tests for RetinaNet model.

  Tests RetinaNet performance in different GPU configurations.
  The naming convention of below test cases follow
  `benchmark_(number of gpus)_gpu` format.
  """

  def __init__(self, output_dir=TMP_DIR, **kwargs):
    super(RetinanetBenchmarkReal, self).__init__(output_dir=output_dir)

  @flagsaver.flagsaver
  def benchmark_8_gpu_coco(self):
    """Run RetinaNet model accuracy test with 8 GPUs."""
    self.num_gpus = 8
    self._setup()
    params = copy.deepcopy(self.params_override)
    params['train']['total_steps'] = 1875  # One epoch.
    # The iterations_per_loop must be one, otherwise the number of examples per
    # second would be wrong. Currently only support calling callback per batch
    # when each loop only runs on one batch, i.e. host loop for one step. The
    # performance of this situation might be lower than the case of
    # iterations_per_loop > 1.
    # Related bug: b/135933080
    params['train']['iterations_per_loop'] = 1
    params['eval']['eval_samples'] = 8
    FLAGS.num_gpus = self.num_gpus
    FLAGS.params_override = json.dumps(params)
    FLAGS.model_dir = self._get_model_dir('real_benchmark_8_gpu_coco')
    # Use negative value to avoid saving checkpoints.
    FLAGS.save_checkpoint_freq = -1
    if self.timer_callback is None:
      logging.error('Cannot measure performance without timer callback')
    else:
      self._run_and_report_benchmark()

  @flagsaver.flagsaver
  def benchmark_1_gpu_coco(self):
    """Run RetinaNet model accuracy test with 1 GPU."""
    self.num_gpus = 1
    self._setup()
    params = copy.deepcopy(self.params_override)
    params['train']['batch_size'] = 8
    params['train']['total_steps'] = 200
    params['train']['iterations_per_loop'] = 1
    params['eval']['eval_samples'] = 8
    FLAGS.num_gpus = self.num_gpus
    FLAGS.params_override = json.dumps(params)
    FLAGS.model_dir = self._get_model_dir('real_benchmark_1_gpu_coco')
    FLAGS.strategy_type = 'one_device_gpu'
    # Use negative value to avoid saving checkpoints.
    FLAGS.save_checkpoint_freq = -1
    if self.timer_callback is None:
      logging.error('Cannot measure performance without timer callback')
    else:
      self._run_and_report_benchmark()

  @flagsaver.flagsaver
  def benchmark_xla_1_gpu_coco(self):
    """Run RetinaNet model accuracy test with 1 GPU and XLA enabled."""
    self.num_gpus = 1
    self._setup()
    params = copy.deepcopy(self.params_override)
    params['train']['batch_size'] = 8
    params['train']['total_steps'] = 200
    params['train']['iterations_per_loop'] = 1
    params['eval']['eval_samples'] = 8
    FLAGS.num_gpus = self.num_gpus
    FLAGS.params_override = json.dumps(params)
    FLAGS.model_dir = self._get_model_dir('real_benchmark_1_gpu_coco')
    FLAGS.strategy_type = 'one_device_gpu'
    FLAGS.enable_xla = True
    # Use negative value to avoid saving checkpoints.
    FLAGS.save_checkpoint_freq = -1
    if self.timer_callback is None:
      logging.error('Cannot measure performance without timer callback')
    else:
      self._run_and_report_benchmark()

if __name__ == '__main__':
  tf.test.main()
<|endoftext|>


<commit_before>import sys
from pyinstrument import processors
from pyinstrument.frame import Frame, SelfTimeFrame
from pytest import approx

all_processors = [
    processors.aggregate_repeated_calls,
    processors.group_library_frames_processor,
    processors.merge_consecutive_self_time,
    processors.remove_importlib,
    processors.remove_unnecessary_self_time_nodes,
    processors.remove_irrelevant_nodes,
]

def test_remove_importlib():
    frame = Frame(
        identifier='<module>\x00sympy/__init__.py\x0012',
        children=[
            Frame(
                identifier='_handle_fromlist\x00../<frozen importlib._bootstrap>\x00997',
                self_time=0.1,
                children=[
                    Frame(
                        identifier='_find_and_load\x00../<frozen importlib._bootstrap>\x00997',
                        self_time=0.1,
                        children=[
                            Frame(
                                identifier='<module>\x00sympy/polys/polyfuncs.py\x001',
                                self_time=0.05,
                            ),
                            Frame(
                                identifier='<module>\x00sympy/polys/partfrac.py\x001',
                                self_time=0.2,
                            )
                        ]
                    ),
                    Frame(
                        identifier='<module>\x00sympy/polys/numberfields.py\x001',
                        self_time=0.05,
                    )
                ]
            )
        ]
    )

    assert frame.self_time == 0.0
    assert frame.time() == approx(0.5)

    frame = processors.remove_importlib(frame, options={})

    assert frame.self_time == approx(0.2)  # the root gets the self_time from the importlib
    assert frame.time() == approx(0.5)
    assert len(frame.children) == 3
    assert frame.children[0].file_path == 'sympy/polys/polyfuncs.py'
    assert frame.children[1].file_path == 'sympy/polys/partfrac.py'
    assert frame.children[2].file_path == 'sympy/polys/numberfields.py'

def test_frame_passthrough_none():
    for processor in all_processors:
        assert processor(None, options={}) is None

def test_merge_consecutive_self_time():
    frame = Frame(
        identifier='<module>\x00cibuildwheel/__init__.py\x0012',
        children=[
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.1),
            SelfTimeFrame(
                self_time=0.2,
            ),
            SelfTimeFrame(
                self_time=0.1,
            ),
            Frame(
                identifier='calculate_metrics\x00cibuildwheel/utils.py\x007',
                self_time=0.1),
            SelfTimeFrame(
                self_time=0.05,
            ),
        ]
    )

    assert frame.time() == approx(0.55)

    frame = processors.merge_consecutive_self_time(frame, options={})

    assert frame.time() == approx(0.55)
    assert len(frame.children) == 4
    assert frame.children[0].self_time == approx(0.1)
    assert frame.children[1].self_time == approx(0.3)
    assert isinstance(frame.children[1], SelfTimeFrame)
    assert frame.children[2].self_time == approx(0.1)
    assert frame.children[3].self_time == approx(0.05)
    assert isinstance(frame.children[3], SelfTimeFrame)

def test_aggregate_repeated_calls():
    frame = Frame(
        identifier='<module>\x00cibuildwheel/__init__.py\x0012',
        children=[
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.1,
                children=[
                    Frame(
                        identifier='scan_string\x00cibuildwheel/utils.py\x0054',
                        self_time=0.2,
                    ),
                ]
            ),
            SelfTimeFrame(
                self_time=0.1,
            ),
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.1,
            ),
            SelfTimeFrame(
                self_time=0.2,
            ),
            Frame(
                identifier='calculate_metrics\x00cibuildwheel/utils.py\x007',
                self_time=0.1,
            ),
            SelfTimeFrame(
                self_time=0.05,
            ),
        ]
    )

    assert frame.time() == approx(0.85)

    frame = processors.aggregate_repeated_calls(frame, options={})

    assert frame.time() == approx(0.85)
    # children should be sorted by time
    assert len(frame.children) == 3
    assert frame.children[0].function == 'strip_newlines'
    assert frame.children[0].time() == 0.4
    assert frame.children[0].children[0].function == 'scan_string'
    assert isinstance(frame.children[1], SelfTimeFrame)
    assert frame.children[1].time() == approx(0.35)

def test_remove_irrelevant_nodes():
    frame = Frame(
        identifier='<module>\x00cibuildwheel/__init__.py\x0012',
        children=[
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                children=[
                    Frame(
                        identifier='scan_string\x00cibuildwheel/utils.py\x0054',
                        self_time=10,
                    ),
                ]
            ),
            SelfTimeFrame(
                self_time=0.5,
            ),
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.5,
            ),
            Frame(
                identifier='calculate_metrics\x00cibuildwheel/utils.py\x007',
                self_time=0.01,
            ),
        ]
    )

    assert frame.time() == approx(11.01)

    frame = processors.remove_irrelevant_nodes(frame, options={})

    assert frame.time() == approx(11.01)
    # check the calculate metrics function was deleted
    assert len(frame.children) == 3
    assert 'calculate_metrics' not in [f.function for f in frame.children]

def test_remove_unnecessary_self_time_nodes():
    frame = Frame(
        identifier='<module>\x00cibuildwheel/__init__.py\x0012',
        children=[
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                children=[
                    SelfTimeFrame(
                        self_time=0.2,
                    ),
                ]
            ),
            SelfTimeFrame(
                self_time=0.5,
            ),
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.5,
            ),
            Frame(
                identifier='calculate_metrics\x00cibuildwheel/utils.py\x007',
                self_time=0.1,
            ),
        ]
    )

    assert frame.time() == approx(1.3)

    frame = processors.remove_unnecessary_self_time_nodes(frame, options={})

    assert frame.time() == approx(1.3)
    assert len(frame.children) == 4
    # check the self time node was deleted
    strip_newlines_frame = frame.children[0]
    assert strip_newlines_frame.function == 'strip_newlines'
    assert len(strip_newlines_frame.children) == 0
    assert strip_newlines_frame.self_time == 0.2

def test_group_library_frames_processor():
    frame = Frame(
        identifier='<module>\x00cibuildwheel/__init__.py\x0012',
        children=[
            Frame(
                identifier='library_function\x00env/lib/python3.6/django/__init__.py\x00997',
                children=[
                    Frame(
                        identifier='library_inner\x00env/lib/python3.6/django/http.py\x0054',
                        children=[
                            Frame(
                                identifier='library_callback\x00env/lib/python3.6/django/views.py\x0054',
                                children=[
                                    Frame(
                                        identifier='<module>\x00cibuildwheel/views.py\x0012',
                                        self_time=0.3,
                                    ),
                                ]
                            ),
                        ]
                    ),
                ]
            ),
            SelfTimeFrame(
                self_time=0.5,
            ),
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.5,
            ),
            Frame(
                identifier='calculate_metrics\x00cibuildwheel/utils.py\x007',
                self_time=0.1,
            ),
        ]
    )

    assert frame.time() == approx(1.4)

    frame = processors.group_library_frames_processor(frame, options={})

    assert frame.time() == approx(1.4)
    group_root = frame.children[0]

    group = group_root.group
    assert group.root == group_root

    for frame in group.frames:
        assert frame.group == group

    assert group_root in group.frames
    assert group_root.children[0] in group.frames
    assert group_root.children[0].children[0] in group.frames
    assert group_root.children[0].children[0] in group.exit_frames
    assert group_root.children[0].children[0].children[0] not in group.frames

    old_sys_path = sys.path[:]
    sys.path.append('env/lib/python3.6')
    assert group.libraries == ['django']
    sys.path[:] = old_sys_path

<commit_msg>Tidy up<commit_after>import sys
from pyinstrument import processors
from pyinstrument.frame import Frame, SelfTimeFrame
from pytest import approx

all_processors = [
    processors.aggregate_repeated_calls,
    processors.group_library_frames_processor,
    processors.merge_consecutive_self_time,
    processors.remove_importlib,
    processors.remove_unnecessary_self_time_nodes,
    processors.remove_irrelevant_nodes,
]

def test_frame_passthrough_none():
    for processor in all_processors:
        assert processor(None, options={}) is None

def test_remove_importlib():
    frame = Frame(
        identifier='<module>\x00sympy/__init__.py\x0012',
        children=[
            Frame(
                identifier='_handle_fromlist\x00../<frozen importlib._bootstrap>\x00997',
                self_time=0.1,
                children=[
                    Frame(
                        identifier='_find_and_load\x00../<frozen importlib._bootstrap>\x00997',
                        self_time=0.1,
                        children=[
                            Frame(
                                identifier='<module>\x00sympy/polys/polyfuncs.py\x001',
                                self_time=0.05,
                            ),
                            Frame(
                                identifier='<module>\x00sympy/polys/partfrac.py\x001',
                                self_time=0.2,
                            )
                        ]
                    ),
                    Frame(
                        identifier='<module>\x00sympy/polys/numberfields.py\x001',
                        self_time=0.05,
                    )
                ]
            )
        ]
    )

    assert frame.self_time == 0.0
    assert frame.time() == approx(0.5)

    frame = processors.remove_importlib(frame, options={})

    assert frame.self_time == approx(0.2)  # the root gets the self_time from the importlib
    assert frame.time() == approx(0.5)
    assert len(frame.children) == 3
    assert frame.children[0].file_path == 'sympy/polys/polyfuncs.py'
    assert frame.children[1].file_path == 'sympy/polys/partfrac.py'
    assert frame.children[2].file_path == 'sympy/polys/numberfields.py'

def test_merge_consecutive_self_time():
    frame = Frame(
        identifier='<module>\x00cibuildwheel/__init__.py\x0012',
        children=[
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.1),
            SelfTimeFrame(
                self_time=0.2,
            ),
            SelfTimeFrame(
                self_time=0.1,
            ),
            Frame(
                identifier='calculate_metrics\x00cibuildwheel/utils.py\x007',
                self_time=0.1),
            SelfTimeFrame(
                self_time=0.05,
            ),
        ]
    )

    assert frame.time() == approx(0.55)

    frame = processors.merge_consecutive_self_time(frame, options={})

    assert frame.time() == approx(0.55)
    assert len(frame.children) == 4
    assert frame.children[0].self_time == approx(0.1)
    assert frame.children[1].self_time == approx(0.3)
    assert isinstance(frame.children[1], SelfTimeFrame)
    assert frame.children[2].self_time == approx(0.1)
    assert frame.children[3].self_time == approx(0.05)
    assert isinstance(frame.children[3], SelfTimeFrame)

def test_aggregate_repeated_calls():
    frame = Frame(
        identifier='<module>\x00cibuildwheel/__init__.py\x0012',
        children=[
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.1,
                children=[
                    Frame(
                        identifier='scan_string\x00cibuildwheel/utils.py\x0054',
                        self_time=0.2,
                    ),
                ]
            ),
            SelfTimeFrame(
                self_time=0.1,
            ),
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.1,
            ),
            SelfTimeFrame(
                self_time=0.2,
            ),
            Frame(
                identifier='calculate_metrics\x00cibuildwheel/utils.py\x007',
                self_time=0.1,
            ),
            SelfTimeFrame(
                self_time=0.05,
            ),
        ]
    )

    assert frame.time() == approx(0.85)

    frame = processors.aggregate_repeated_calls(frame, options={})

    assert frame.time() == approx(0.85)
    # children should be sorted by time
    assert len(frame.children) == 3
    assert frame.children[0].function == 'strip_newlines'
    assert frame.children[0].time() == 0.4
    assert frame.children[0].children[0].function == 'scan_string'
    assert isinstance(frame.children[1], SelfTimeFrame)
    assert frame.children[1].time() == approx(0.35)

def test_remove_irrelevant_nodes():
    frame = Frame(
        identifier='<module>\x00cibuildwheel/__init__.py\x0012',
        children=[
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                children=[
                    Frame(
                        identifier='scan_string\x00cibuildwheel/utils.py\x0054',
                        self_time=10,
                    ),
                ]
            ),
            SelfTimeFrame(
                self_time=0.5,
            ),
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.5,
            ),
            Frame(
                identifier='calculate_metrics\x00cibuildwheel/utils.py\x007',
                self_time=0.01,
            ),
        ]
    )

    assert frame.time() == approx(11.01)

    frame = processors.remove_irrelevant_nodes(frame, options={})

    assert frame.time() == approx(11.01)
    # check the calculate metrics function was deleted
    assert len(frame.children) == 3
    assert 'calculate_metrics' not in [f.function for f in frame.children]

def test_remove_unnecessary_self_time_nodes():
    frame = Frame(
        identifier='<module>\x00cibuildwheel/__init__.py\x0012',
        children=[
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                children=[
                    SelfTimeFrame(
                        self_time=0.2,
                    ),
                ]
            ),
            SelfTimeFrame(
                self_time=0.5,
            ),
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.5,
            ),
            Frame(
                identifier='calculate_metrics\x00cibuildwheel/utils.py\x007',
                self_time=0.1,
            ),
        ]
    )

    assert frame.time() == approx(1.3)

    frame = processors.remove_unnecessary_self_time_nodes(frame, options={})

    assert frame.time() == approx(1.3)
    assert len(frame.children) == 4
    # check the self time node was deleted
    strip_newlines_frame = frame.children[0]
    assert strip_newlines_frame.function == 'strip_newlines'
    assert len(strip_newlines_frame.children) == 0
    assert strip_newlines_frame.self_time == 0.2

def test_group_library_frames_processor():
    frame = Frame(
        identifier='<module>\x00cibuildwheel/__init__.py\x0012',
        children=[
            Frame(
                identifier='library_function\x00env/lib/python3.6/django/__init__.py\x00997',
                children=[
                    Frame(
                        identifier='library_inner\x00env/lib/python3.6/django/http.py\x0054',
                        children=[
                            Frame(
                                identifier='library_callback\x00env/lib/python3.6/django/views.py\x0054',
                                children=[
                                    Frame(
                                        identifier='<module>\x00cibuildwheel/views.py\x0012',
                                        self_time=0.3,
                                    ),
                                ]
                            ),
                        ]
                    ),
                ]
            ),
            SelfTimeFrame(
                self_time=0.5,
            ),
            Frame(
                identifier='strip_newlines\x00cibuildwheel/utils.py\x00997',
                self_time=0.5,
            ),
            Frame(
                identifier='calculate_metrics\x00cibuildwheel/utils.py\x007',
                self_time=0.1,
            ),
        ]
    )

    assert frame.time() == approx(1.4)

    frame = processors.group_library_frames_processor(frame, options={})

    assert frame.time() == approx(1.4)
    group_root = frame.children[0]

    group = group_root.group
    assert group.root == group_root

    for frame in group.frames:
        assert frame.group == group

    assert group_root in group.frames
    assert group_root.children[0] in group.frames
    assert group_root.children[0].children[0] in group.frames
    assert group_root.children[0].children[0] in group.exit_frames
    assert group_root.children[0].children[0].children[0] not in group.frames

    old_sys_path = sys.path[:]
    sys.path.append('env/lib/python3.6')
    assert group.libraries == ['django']
    sys.path[:] = old_sys_path

<|endoftext|>


<commit_before># Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Helper functions for running models in a distributed setting."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

def get_distribution_strategy(num_gpus, all_reduce_alg=None):
  """Return a DistributionStrategy for running the model.

  Args:
    num_gpus: Number of GPUs to run this model.
    all_reduce_alg: Specify which algorithm to use when performing all-reduce.
      See tf.contrib.distribute.AllReduceCrossDeviceOps for available algorithms.
      If None, DistributionStrategy will choose based on device topology.

  Returns:
    tf.contrib.distribute.DistibutionStrategy object.
  """
  if num_gpus == 0:
    return tf.contrib.distribute.OneDeviceStrategy("device:CPU:0")
  elif num_gpus == 1:
    return tf.contrib.distribute.OneDeviceStrategy("device:GPU:0")
  else:
    if all_reduce_alg:
      return tf.contrib.distribute.MirroredStrategy(
          num_gpus=num_gpus,
          cross_tower_ops=tf.contrib.distribute.AllReduceCrossDeviceOps(
              all_reduce_alg, num_packs=2))
    else:
      return tf.contrib.distribute.MirroredStrategy(num_gpus=num_gpus)

def per_device_batch_size(batch_size, num_gpus):
  """For multi-gpu, batch-size must be a multiple of the number of GPUs.

  Note that this should eventually be handled by DistributionStrategies
  directly. Multi-GPU support is currently experimental, however,
  so doing the work here until that feature is in place.

  Args:
    batch_size: Global batch size to be divided among devices. This should be
      equal to num_gpus times the single-GPU batch_size for multi-gpu training.
    num_gpus: How many GPUs are used with DistributionStrategies.

  Returns:
    Batch size per device.

  Raises:
    ValueError: if batch_size is not divisible by number of devices
  """
  if num_gpus <= 1:
    return batch_size

  remainder = batch_size % num_gpus
  if remainder:
    err = ("When running with multiple GPUs, batch size "
           "must be a multiple of the number of available GPUs. Found {} "
           "GPUs with a batch size of {}; try --batch_size={} instead."
          ).format(num_gpus, batch_size, batch_size - remainder)
    raise ValueError(err)
  return int(batch_size / num_gpus)
<commit_msg>Update distribution_utils.py<commit_after># Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Helper functions for running models in a distributed setting."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

def get_distribution_strategy(num_gpus, all_reduce_alg=None):
  """Return a DistributionStrategy for running the model.

  Args:
    num_gpus: Number of GPUs to run this model.
    all_reduce_alg: Specify which algorithm to use when performing all-reduce.
      See tf.contrib.distribute.AllReduceCrossDeviceOps for available
      algorithms. If None, DistributionStrategy will choose based on device
      topology.

  Returns:
    tf.contrib.distribute.DistibutionStrategy object.
  """
  if num_gpus == 0:
    return tf.contrib.distribute.OneDeviceStrategy("device:CPU:0")
  elif num_gpus == 1:
    return tf.contrib.distribute.OneDeviceStrategy("device:GPU:0")
  else:
    if all_reduce_alg:
      return tf.contrib.distribute.MirroredStrategy(
          num_gpus=num_gpus,
          cross_tower_ops=tf.contrib.distribute.AllReduceCrossDeviceOps(
              all_reduce_alg, num_packs=2))
    else:
      return tf.contrib.distribute.MirroredStrategy(num_gpus=num_gpus)

def per_device_batch_size(batch_size, num_gpus):
  """For multi-gpu, batch-size must be a multiple of the number of GPUs.

  Note that this should eventually be handled by DistributionStrategies
  directly. Multi-GPU support is currently experimental, however,
  so doing the work here until that feature is in place.

  Args:
    batch_size: Global batch size to be divided among devices. This should be
      equal to num_gpus times the single-GPU batch_size for multi-gpu training.
    num_gpus: How many GPUs are used with DistributionStrategies.

  Returns:
    Batch size per device.

  Raises:
    ValueError: if batch_size is not divisible by number of devices
  """
  if num_gpus <= 1:
    return batch_size

  remainder = batch_size % num_gpus
  if remainder:
    err = ("When running with multiple GPUs, batch size "
           "must be a multiple of the number of available GPUs. Found {} "
           "GPUs with a batch size of {}; try --batch_size={} instead."
          ).format(num_gpus, batch_size, batch_size - remainder)
    raise ValueError(err)
  return int(batch_size / num_gpus)
<|endoftext|>


<commit_before>#!./venv/bin/python

from functools import partial
import hashlib
import json
import operator
import subprocess
import sys

def _create_root_node(node_class, *args, **kwargs):
    root = node_class(*args, **kwargs)
    root._update_hash()
    return root

def _create_child_node(node_class, parent, *args, **kwargs):
    child = node_class(parent, *args, **kwargs)
    child._update_hash()
    return child

class _Node(object):
    def __init__(self, parents):
        parent_hashes = [parent.hash for parent in parents]
        assert len(parent_hashes) == len(set(parent_hashes)), 'Same node cannot be included as parent multiple times'
        self.parents = parents

    @classmethod
    def _add_operator(cls, node_class):
        if not getattr(node_class, 'STATIC', False):
            def func(self, *args, **kwargs):
                return _create_child_node(node_class, self, *args, **kwargs)
            setattr(cls, node_class.NAME, func)

    @classmethod
    def _add_operators(cls, node_classes):
        [cls._add_operator(node_class) for node_class in node_classes]

    @property
    def _props(self):
        return {k: v for k, v in self.__dict__.items() if k not in ['parents', 'hash']}

    def __repr__(self):
        # TODO: exclude default values.
        props = self._props
        formatted_props = ['{}={!r}'.format(key, props[key]) for key in sorted(self._props)]
        return '{}({})'.format(self.NAME, ','.join(formatted_props))

    def __eq__(self, other):
        return self.hash == other.hash

    def _update_hash(self):
        my_hash = hashlib.md5(json.dumps(self._props)).hexdigest()
        parent_hashes = [parent.hash for parent in self.parents]
        hashes = parent_hashes + [my_hash]
        self.hash = hashlib.md5(','.join(hashes)).hexdigest()

class _InputNode(_Node):
    pass

class _FileInputNode(_InputNode):
    NAME = 'file_input'
    STATIC = True

    def __init__(self, filename):
        super(_FileInputNode, self).__init__(parents=[])
        self.filename = filename

class _FilterNode(_Node):
    def _get_filter(self):
        raise NotImplementedError()

class _TrimNode(_FilterNode):
    NAME = 'trim'

    def __init__(self, parent, start_frame, end_frame, setpts='PTS-STARTPTS'):
        super(_TrimNode, self).__init__([parent])
        self.start_frame = start_frame
        self.end_frame = end_frame
        self.setpts = setpts

    def _get_filter(self):
        return 'trim=start_frame={}:end_frame={},setpts={}'.format(self.start_frame, self.end_frame, self.setpts)

class _OverlayNode(_FilterNode):
    NAME = 'overlay'

    def __init__(self, main_parent, overlay_parent, eof_action='repeat'):
        super(_OverlayNode, self).__init__([main_parent, overlay_parent])
        self.eof_action = eof_action

    def _get_filter(self):
        return 'overlay=eof_action={}'.format(self.eof_action)

class _HFlipNode(_FilterNode):
    NAME = 'hflip'

    def __init__(self, parent):
        super(_HFlipNode, self).__init__([parent])

    def _get_filter(self):
        return 'hflip'

class _DrawBoxNode(_FilterNode):
    NAME = 'drawbox'

    def __init__(self, parent, x, y, width, height, color, thickness=1):
        super(_DrawBoxNode, self).__init__([parent])
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.color = color
        self.thickness = thickness

    def _get_filter(self):
        return 'drawbox={}:{}:{}:{}:{}:t={}'.format(self.x, self.y, self.width, self.height, self.color, self.thickness)

class _ConcatNode(_Node):
    NAME = 'concat'
    STATIC = True

    def __init__(self, *parents):
        super(_ConcatNode, self).__init__(parents)

    def _get_filter(self):
        return 'concat=n={}'.format(len(self.parents))

class _OutputNode(_Node):
    @classmethod
    def _get_stream_name(cls, name):
        return '[{}]'.format(name)

    @classmethod
    def _get_input_args(cls, input_node):
        if isinstance(input_node, _FileInputNode):
            args = ['-i', input_node.filename]
        else:
            assert False, 'Unsupported input node: {}'.format(input_node)
        return args

    @classmethod
    def _topo_sort(cls, start_node):
        marked_nodes = []
        sorted_nodes = []
        child_map = {}
        def visit(node, child):
            assert node not in marked_nodes, 'Graph is not a DAG'
            if child is not None:
                if node not in child_map:
                    child_map[node] = []
                child_map[node].append(child)
            if node not in sorted_nodes:
                marked_nodes.append(node)
                [visit(parent, node) for parent in node.parents]
                marked_nodes.remove(node)
                sorted_nodes.append(node)
        unmarked_nodes = [start_node]
        while unmarked_nodes:
            visit(unmarked_nodes.pop(), None)
        return sorted_nodes, child_map

    @classmethod
    def _get_filter_spec(cls, i, node, stream_name_map):
        stream_name = cls._get_stream_name('v{}'.format(i))
        stream_name_map[node] = stream_name
        inputs = [stream_name_map[parent] for parent in node.parents]
        filter_spec = '{}{}{}'.format(''.join(inputs), node._get_filter(), stream_name)
        return filter_spec

    @classmethod
    def _get_filter_arg(cls, filter_nodes, stream_name_map):
        filter_specs = [cls._get_filter_spec(i, node, stream_name_map) for i, node in enumerate(filter_nodes)]
        return ';'.join(filter_specs)

    @classmethod
    def _get_global_args(cls, node):
        if isinstance(node, _OverwriteOutputNode):
            return ['-y']
        else:
            assert False, 'Unsupported global node: {}'.format(node)

    @classmethod
    def _get_output_args(cls, node, stream_name_map):
        args = []
        if not isinstance(node, _MergeOutputsNode):
            stream_name = stream_name_map[node.parents[0]]
            if stream_name != '[0]':
                args += ['-map', stream_name]
            if isinstance(node, _FileOutputNode):
                args += [node.filename]
            else:
                assert False, 'Unsupported output node: {}'.format(node)
        return args

    def get_args(self):
        args = []
        # TODO: group nodes together, e.g. `-i somefile -r somerate`.
        sorted_nodes, child_map = self._topo_sort(self)
        input_nodes = [node for node in sorted_nodes if isinstance(node, _InputNode)]
        output_nodes = [node for node in sorted_nodes if isinstance(node, _OutputNode) and not
            isinstance(node, _GlobalNode)]
        global_nodes = [node for node in sorted_nodes if isinstance(node, _GlobalNode)]
        filter_nodes = [node for node in sorted_nodes if node not in (input_nodes + output_nodes + global_nodes)]
        stream_name_map = {node: self._get_stream_name(i) for i, node in enumerate(input_nodes)}
        filter_arg = self._get_filter_arg(filter_nodes, stream_name_map)
        args += reduce(operator.add, [self._get_input_args(node) for node in input_nodes])
        if filter_arg:
            args += ['-filter_complex', filter_arg]
        args += reduce(operator.add, [self._get_output_args(node, stream_name_map) for node in output_nodes])
        args += reduce(operator.add, [self._get_global_args(node) for node in global_nodes], [])
        return args

    def run(self):
        args = ['ffmpeg'] + self.get_args()
        subprocess.check_call(args)

class _GlobalNode(_OutputNode):
    def __init__(self, parent):
        assert isinstance(parent, _OutputNode), 'Global nodes can only be attached after output nodes'
        super(_GlobalNode, self).__init__([parent])

class _OverwriteOutputNode(_GlobalNode):
    NAME = 'overwrite_output'

class _MergeOutputsNode(_OutputNode):
    NAME = 'merge_outputs'
    STATIC = True

    def __init__(self, *parents):
        assert not any([not isinstance(parent, _OutputNode) for parent in parents]), 'Can only merge output streams'
        super(_MergeOutputsNode, self).__init__(*parents)

class _FileOutputNode(_OutputNode):
    NAME = 'file_output'

    def __init__(self, parent, filename):
        super(_FileOutputNode, self).__init__([parent])
        self.filename = filename

NODE_CLASSES = [
    _HFlipNode,
    _DrawBoxNode,
    _ConcatNode,
    _FileInputNode,
    _FileOutputNode,
    _OverlayNode,
    _OverwriteOutputNode,
    _TrimNode,
]

_Node._add_operators(NODE_CLASSES)

_module = sys.modules[__name__]
for _node_class in NODE_CLASSES:
    if getattr(_node_class, 'STATIC', False):
        func = _create_root_node
    else:
        func = _create_child_node
    func = partial(func, _node_class)
    setattr(_module, _node_class.NAME, func)

def get_args(node):
    assert isinstance(node, _OutputNode), 'Cannot generate ffmpeg args for non-output node'
    return node.get_args()

def run(node):
    assert isinstance(node, _OutputNode), 'Cannot run ffmpeg on non-output node'
    return node.run()
<commit_msg>add more filters, options<commit_after>#!./venv/bin/python

from functools import partial
import hashlib
import json
import operator
import subprocess
import sys

def _create_root_node(node_class, *args, **kwargs):
    root = node_class(*args, **kwargs)
    root._update_hash()
    return root

def _create_child_node(node_class, parent, *args, **kwargs):
    child = node_class(parent, *args, **kwargs)
    child._update_hash()
    return child

class _Node(object):
    def __init__(self, parents):
        parent_hashes = [parent.hash for parent in parents]
        assert len(parent_hashes) == len(set(parent_hashes)), 'Same node cannot be included as parent multiple times'
        self.parents = parents

    @classmethod
    def _add_operator(cls, node_class):
        if not getattr(node_class, 'STATIC', False):
            def func(self, *args, **kwargs):
                return _create_child_node(node_class, self, *args, **kwargs)
            setattr(cls, node_class.NAME, func)

    @classmethod
    def _add_operators(cls, node_classes):
        [cls._add_operator(node_class) for node_class in node_classes]

    @property
    def _props(self):
        return {k: v for k, v in self.__dict__.items() if k not in ['parents', 'hash']}

    def __repr__(self):
        # TODO: exclude default values.
        props = self._props
        formatted_props = ['{}={!r}'.format(key, props[key]) for key in sorted(self._props)]
        return '{}({})'.format(self.NAME, ','.join(formatted_props))

    def __eq__(self, other):
        return self.hash == other.hash

    def _update_hash(self):
        my_hash = hashlib.md5(json.dumps(self._props)).hexdigest()
        parent_hashes = [parent.hash for parent in self.parents]
        hashes = parent_hashes + [my_hash]
        self.hash = hashlib.md5(','.join(hashes)).hexdigest()

class _InputNode(_Node):
    pass

class _FileInputNode(_InputNode):
    NAME = 'file_input'
    STATIC = True

    def __init__(self, filename):
        super(_FileInputNode, self).__init__(parents=[])
        self.filename = filename

class _FilterNode(_Node):
    def _get_filter(self):
        raise NotImplementedError()

    def _get_params_from_dict(self, d):
        params = ""
        for k in self.kwargs:
            params += k + "={}:".format(self.kwargs[k])
        if len(params) > 0:
            params = params[:-1]
        return params

    def _get_params_from_list(self, l):
        return ":".join(["{}".format(i) for i in l])

    def _get_filter_from_dict(self, d):
        p = self._get_params_from_dict(d)
        if len(p) > 0:
            return self.NAME + "=" + p
        return self.NAME

class _TrimNode(_FilterNode):
    NAME = 'trim'

    def __init__(self, parent, **kwargs):
        super(_TrimNode, self).__init__([parent])
        # if "setpts" not in kwargs:
        #     kwargs["setpts"] = 'PTS-STARTPTS'
        self.kwargs = kwargs

    def _get_filter(self):
        params = ""
        for k in self.kwargs:
            if k == "setpts":
                continue
            params += k
            params += "={}:".format(self.kwargs[k])
        if len(params) > 0:
            params = params[:-1]

        if "setpts" in self.kwargs:
            params += "setpts={}".format(self.kwargs["setpts"])

        return self.NAME + '=' + params

class _OverlayNode(_FilterNode):
    NAME = 'overlay'

    def __init__(self, main_parent, overlay_parent, **kwargs):
        super(_OverlayNode, self).__init__([main_parent, overlay_parent])
        self.eof_action = eof_action
        self.kwargs = kwargs

    def _get_filter(self):
        return self._get_filter_from_dict(self.kwargs)

class _HFlipNode(_FilterNode):
    NAME = 'hflip'

    def __init__(self, parent):
        super(_HFlipNode, self).__init__([parent])

    def _get_filter(self):
        return self.NAME

class _VFlipNode(_FilterNode):
    NAME = 'vflip'

    def __init__(self, parent):
        super(_VFlipNode, self).__init__([parent])

    def _get_filter(self):
        return self.NAME

class _DrawBoxNode(_FilterNode):
    NAME = 'drawbox'

    def __init__(self, parent, x, y, width, height, color, **kwargs):
        super(_DrawBoxNode, self).__init__([parent])
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.color = color
        self.kwargs = kwargs

    def _get_filter(self):
        f = 'drawbox={}:{}:{}:{}:{}'.format(self.x, self.y, self.width, self.height, self.color)
        p = self._get_params_from_dict(self.kwargs)
        if len(p) > 0:
            return f + ":" + p
        return f

class _ConcatNode(_Node):
    NAME = 'concat'
    STATIC = True

    def __init__(self, *parents):
        super(_ConcatNode, self).__init__(parents)

    def _get_filter(self):
        return 'concat=n={}'.format(len(self.parents))

class _ZoomPanNode(_FilterNode):
    NAME = 'zoompan'

    def __init__(self, parent, **kwargs):
        super(_ZoomPanNode, self).__init__([parent])
        self.kwargs = kwargs

    def _get_filter(self):
        return self._get_filter_from_dict(self.kwargs)

class _HueNode(_FilterNode):
    NAME = 'hue'

    def __init__(self, parent, **kwargs):
        super(_HueNode, self).__init__([parent])
        self.kwargs = kwargs

    def _get_filter(self):
        return self._get_filter_from_dict(self.kwargs)

class _ColorChannelMixerNode(_FilterNode):
    NAME = 'colorchannelmixer'

    def __init__(self, parent, *args, **kwargs):
        super(_ColorChannelMixerNode, self).__init__([parent])
        self.args = args
        self.kwargs = kwargs

    def _get_filter(self):
        f = self.NAME + "="
        if self.args:
            f += self._get_params_from_list(self.args)
            if self.kwargs:
                f += ":"
        if self.kwargs:
            f += self._get_params_from_dict(self.kwargs)
        return f

class _OutputNode(_Node):
    @classmethod
    def _get_stream_name(cls, name):
        return '[{}]'.format(name)

    @classmethod
    def _get_input_args(cls, input_node):
        if isinstance(input_node, _FileInputNode):
            args = ['-i', input_node.filename]
        else:
            assert False, 'Unsupported input node: {}'.format(input_node)
        return args

    @classmethod
    def _topo_sort(cls, start_node):
        marked_nodes = []
        sorted_nodes = []
        child_map = {}
        def visit(node, child):
            assert node not in marked_nodes, 'Graph is not a DAG'
            if child is not None:
                if node not in child_map:
                    child_map[node] = []
                child_map[node].append(child)
            if node not in sorted_nodes:
                marked_nodes.append(node)
                [visit(parent, node) for parent in node.parents]
                marked_nodes.remove(node)
                sorted_nodes.append(node)
        unmarked_nodes = [start_node]
        while unmarked_nodes:
            visit(unmarked_nodes.pop(), None)
        return sorted_nodes, child_map

    @classmethod
    def _get_filter_spec(cls, i, node, stream_name_map):
        stream_name = cls._get_stream_name('v{}'.format(i))
        stream_name_map[node] = stream_name
        inputs = [stream_name_map[parent] for parent in node.parents]
        filter_spec = '{}{}{}'.format(''.join(inputs), node._get_filter(), stream_name)
        return filter_spec

    @classmethod
    def _get_filter_arg(cls, filter_nodes, stream_name_map):
        filter_specs = [cls._get_filter_spec(i, node, stream_name_map) for i, node in enumerate(filter_nodes)]
        return ';'.join(filter_specs)

    @classmethod
    def _get_global_args(cls, node):
        if isinstance(node, _OverwriteOutputNode):
            return ['-y']
        else:
            assert False, 'Unsupported global node: {}'.format(node)

    @classmethod
    def _get_output_args(cls, node, stream_name_map):
        args = []
        if not isinstance(node, _MergeOutputsNode):
            stream_name = stream_name_map[node.parents[0]]
            if stream_name != '[0]':
                args += ['-map', stream_name]
            if isinstance(node, _FileOutputNode):
                args += [node.filename]
            else:
                assert False, 'Unsupported output node: {}'.format(node)
        return args

    def get_args(self):
        args = []
        # TODO: group nodes together, e.g. `-i somefile -r somerate`.
        sorted_nodes, child_map = self._topo_sort(self)
        input_nodes = [node for node in sorted_nodes if isinstance(node, _InputNode)]
        output_nodes = [node for node in sorted_nodes if isinstance(node, _OutputNode) and not
            isinstance(node, _GlobalNode)]
        global_nodes = [node for node in sorted_nodes if isinstance(node, _GlobalNode)]
        filter_nodes = [node for node in sorted_nodes if node not in (input_nodes + output_nodes + global_nodes)]
        stream_name_map = {node: self._get_stream_name(i) for i, node in enumerate(input_nodes)}
        filter_arg = self._get_filter_arg(filter_nodes, stream_name_map)
        args += reduce(operator.add, [self._get_input_args(node) for node in input_nodes])
        if filter_arg:
            args += ['-filter_complex', filter_arg]
        args += reduce(operator.add, [self._get_output_args(node, stream_name_map) for node in output_nodes])
        args += reduce(operator.add, [self._get_global_args(node) for node in global_nodes], [])
        return args

    def run(self):
        args = ['ffmpeg'] + self.get_args()
        subprocess.check_call(args)

class _GlobalNode(_OutputNode):
    def __init__(self, parent):
        assert isinstance(parent, _OutputNode), 'Global nodes can only be attached after output nodes'
        super(_GlobalNode, self).__init__([parent])

class _OverwriteOutputNode(_GlobalNode):
    NAME = 'overwrite_output'

class _MergeOutputsNode(_OutputNode):
    NAME = 'merge_outputs'
    STATIC = True

    def __init__(self, *parents):
        assert not any([not isinstance(parent, _OutputNode) for parent in parents]), 'Can only merge output streams'
        super(_MergeOutputsNode, self).__init__(*parents)

class _FileOutputNode(_OutputNode):
    NAME = 'file_output'

    def __init__(self, parent, filename):
        super(_FileOutputNode, self).__init__([parent])
        self.filename = filename

NODE_CLASSES = [
    _HFlipNode,
    _DrawBoxNode,
    _ConcatNode,
    _FileInputNode,
    _FileOutputNode,
    _OverlayNode,
    _OverwriteOutputNode,
    _TrimNode,
]

_Node._add_operators(NODE_CLASSES)

_module = sys.modules[__name__]
for _node_class in NODE_CLASSES:
    if getattr(_node_class, 'STATIC', False):
        func = _create_root_node
    else:
        func = _create_child_node
    func = partial(func, _node_class)
    setattr(_module, _node_class.NAME, func)

def get_args(node):
    assert isinstance(node, _OutputNode), 'Cannot generate ffmpeg args for non-output node'
    return node.get_args()

def run(node):
    assert isinstance(node, _OutputNode), 'Cannot run ffmpeg on non-output node'
    return node.run()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Copyright 2014 Olivier Cortès <oc@1flow.io>.

This file is part of the 1flow project.

1flow is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of
the License, or (at your option) any later version.

1flow is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public
License along with 1flow.  If not, see http://www.gnu.org/licenses/

"""
import uuid
import logging

from constance import config

from json_field import JSONField

# from django.conf import settings
from django.db import models
from django.utils.translation import ugettext_lazy as _
# from django.db.models.signals import post_save, pre_save  # , pre_delete

from polymorphic import (
    PolymorphicQuerySet,
    PolymorphicManager,
    PolymorphicModel,
)

from sparks.django.models import DiffMixin

from oneflow.base.utils.dateutils import now, timedelta

from oneflow.base.utils import (
    # register_task_method,
    RedisExpiringLock,
)

from ..common import User

LOGGER = logging.getLogger(__name__)

__all__ = [
    'BaseAccount',
    'BaseAccountManager',
    'BaseAccountQuerySet',
    # 'baseaccount_pre_save',
]

# ————————————————————————————————————————— Utils / redis descriptors functions

def get_account_image_upload_path(instance, filename):

    if not filename.strip():
        filename = uuid.uuid4()

    # The filename will be used in a shell command later. In case the
    # user/admin forgets the '"' in the configuration, avoid problems.
    filename = filename.replace(u' ', u'_')

    if instance:
        return 'user/{0}/account/{1}/images/{1}'.format(
            instance.user.username, instance.id, filename)

    return u'images/%Y/%m/%d/{0}'.format(filename)

# ———————————————————————————————————————————————————————————————————— Managers

class BaseAccountQuerySet(PolymorphicQuerySet):

    """ Account based queryset.

    .. note:: this query set will be patched by subclasses.
    """

    def active(self):
        return self.filter(is_active=True)

    def usable(self):
        return self.active().filter(is_usable=True)

    def unusable(self):
        return self.active().filter(is_usable=False)

class BaseAccountManager(PolymorphicManager):

    """ A manager that adds some things. """

    use_for_related_fields = True
    queryset_class = BaseAccountQuerySet

# ——————————————————————————————————————————————————————————————————————— Model

class BaseAccount(PolymorphicModel, DiffMixin):

    """ Base 1flow account. """

    class Meta:
        app_label = 'core'
        verbose_name = _(u'Base account')
        verbose_name_plural = _(u'Base accounts')

    # Can be overriden by inherited classes.
    usable_start_task_name   = None
    reset_unusable_task_name = None

    # ———————————————————————————————————————————————————————————————— Managers

    objects = BaseAccountManager()

    # —————————————————————————————————————————————————————————————— Attributes

    # NOTE: keep ID a simple integer/auto
    # field, this surely helps JOIN operations.
    # id             = models.UUIDField(primary_key=True,
    #                                   default=uuid.uuid4, editable=False)

    user = models.ForeignKey(User, verbose_name=_(u'Owner'),
                             related_name='accounts')

    name = models.CharField(verbose_name=_(u'Account name'),
                            max_length=128, blank=True)

    slug = models.CharField(verbose_name=_(u'slug'),
                            max_length=255,
                            null=True, blank=True)

    # items = models.ManyToManyField(BaseItem, blank=True, null=True,
    #                               verbose_name=_(u'Feed items'),
    #                               related_name='feeds')

    date_created = models.DateTimeField(auto_now_add=True, blank=True,
                                        verbose_name=_(u'Date created'))

    date_last_conn = models.DateTimeField(
        verbose_name=_(u'Last connection date'),
        null=True, blank=True)

    conn_error = models.CharField(
        verbose_name=_(u'Last connection error'),
        max_length=255, null=True, blank=True)

    is_usable = models.BooleanField(
        verbose_name=_(u'usable'),
        default=True, blank=True,
        help_text=_(u'Is account connectivity good or not?'))

    is_active = models.BooleanField(
        verbose_name=_(u'active'),
        default=True, blank=True,
        help_text=_(u'Is account enabled for fetching?'))

    options = JSONField(default=dict, blank=True)

    notes = models.TextField(
        verbose_name=_(u'Notes'), null=True, blank=True,
        help_text=_(u'User notes for this account.'))

    image = models.ImageField(
        verbose_name=_(u'Thumbnail'), null=True, blank=True,
        upload_to=get_account_image_upload_path, max_length=256,
        help_text=_(u'Use either image when 1flow instance hosts the '
                    u'image, or image_url when hosted elsewhere. If '
                    u'both are filled, image takes precedence.'))

    image_url = models.URLField(
        verbose_name=_(u'Thumbnail URL'), null=True, blank=True, max_length=384,
        help_text=_(u'Full URL of the image displayed in the feed '
                    u'selector. Can be hosted outside of 1flow.'))

    # —————————————————————————————————————————————————————— Django & Grappelli

    def __unicode__(self):
        """ Hello, pep257. I love you so. """

        return _(u'BaseAccount {0} ({1}) of {2}').format(
            self.name, self.id, self.user.username)

    @staticmethod
    def autocomplete_search_fields():
        """ grappelli auto-complete method. """

        return ('name__icontains', )

    # —————————————————————————————————————————————————————————————— Properties

    @property
    def config_account_refresh_period(self):
        """ Return the default account refresh period.

        .. note:: it's recommended to override this
            property in polymorphic inheritant models.
        """

        return config.ACCOUNT_REFRESH_PERIOD_DEFAULT

    @property
    def refresh_lock(self):
        try:
            return self.__refresh_lock

        except AttributeError:
            self.__refresh_lock = RedisExpiringLock(
                self, lock_name='account_fetch',
            )
            return self.__refresh_lock

    @property
    def recently_usable(self):
        """ Return True if the account has been tested/connected recently. """

        return self.is_usable and (
            now() - self.date_last_conn
            < timedelta(seconds=self.config_account_refresh_period))

    # ————————————————————————————————————————————————————————————————— Methods

    def reset_unusable(self, commit=True):
        """ Mark the current instance needing to test usability.

        This is called typically when a connection parameter has changed,
        and the current account connectivity needs to be tested again to
        validate them all.
        """

        self.date_last_conn = None
        self.is_usable = False

        if commit:
            self.save()

        try:
            reset_unusable_task = self._meta.model.reset_unusable_task

        except:
            pass

        else:
            reset_unusable_task.apply_async(
                args=(self.id, ), countdown=3)

    def mark_unusable(self, message, args=(), exc=None, commit=True):
        """ Mark account unsable with date & message, log exception if any. """

        if exc is not None:
            if args:
                message = message % args

            message = u'{0} ({1})'.format(message, unicode(exc))
            LOGGER.exception(u'%s unusable: %s', self, message)

        self.date_last_conn = now()
        self.conn_error = message
        self.is_usable = False

        if commit:
            self.save()

    def mark_usable(self, commit=True, verbose=True):
        """ Mark the account usable and clear error. """

        if verbose:
            LOGGER.info(u'%s is now considered usable.', self)

        if self.is_usable:
            start_task = False

        else:
            start_task = True

        self.date_last_conn = now()
        self.conn_error = None
        self.is_usable = True

        if commit:
            self.save()

        try:
            usable_start_task = self._meta.model.usable_start_task

        except:
            pass

        else:
            if start_task:
                usable_start_task.delay(self.id)

    def has_option(self, option):
        """ True if option in self.options. """

        return option in self.options
<commit_msg>add inactive() query selector.<commit_after># -*- coding: utf-8 -*-
"""
Copyright 2014 Olivier Cortès <oc@1flow.io>.

This file is part of the 1flow project.

1flow is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of
the License, or (at your option) any later version.

1flow is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public
License along with 1flow.  If not, see http://www.gnu.org/licenses/

"""
import uuid
import logging

from constance import config

from json_field import JSONField

# from django.conf import settings
from django.db import models
from django.utils.translation import ugettext_lazy as _
# from django.db.models.signals import post_save, pre_save  # , pre_delete

from polymorphic import (
    PolymorphicQuerySet,
    PolymorphicManager,
    PolymorphicModel,
)

from sparks.django.models import DiffMixin

from oneflow.base.utils.dateutils import now, timedelta

from oneflow.base.utils import (
    # register_task_method,
    RedisExpiringLock,
)

from ..common import User

LOGGER = logging.getLogger(__name__)

__all__ = [
    'BaseAccount',
    'BaseAccountManager',
    'BaseAccountQuerySet',
    # 'baseaccount_pre_save',
]

# ————————————————————————————————————————— Utils / redis descriptors functions

def get_account_image_upload_path(instance, filename):

    if not filename.strip():
        filename = uuid.uuid4()

    # The filename will be used in a shell command later. In case the
    # user/admin forgets the '"' in the configuration, avoid problems.
    filename = filename.replace(u' ', u'_')

    if instance:
        return 'user/{0}/account/{1}/images/{1}'.format(
            instance.user.username, instance.id, filename)

    return u'images/%Y/%m/%d/{0}'.format(filename)

# ———————————————————————————————————————————————————————————————————— Managers

class BaseAccountQuerySet(PolymorphicQuerySet):

    """ Account based queryset.

    .. note:: this query set will be patched by subclasses.
    """

    def inactive(self):
        return self.filter(is_active=False)

    def active(self):
        return self.filter(is_active=True)

    def usable(self):
        return self.active().filter(is_usable=True)

    def unusable(self):
        return self.active().filter(is_usable=False)

class BaseAccountManager(PolymorphicManager):

    """ A manager that adds some things. """

    use_for_related_fields = True
    queryset_class = BaseAccountQuerySet

# ——————————————————————————————————————————————————————————————————————— Model

class BaseAccount(PolymorphicModel, DiffMixin):

    """ Base 1flow account. """

    class Meta:
        app_label = 'core'
        verbose_name = _(u'Base account')
        verbose_name_plural = _(u'Base accounts')

    # Can be overriden by inherited classes.
    usable_start_task_name   = None
    reset_unusable_task_name = None

    # ———————————————————————————————————————————————————————————————— Managers

    objects = BaseAccountManager()

    # —————————————————————————————————————————————————————————————— Attributes

    # NOTE: keep ID a simple integer/auto
    # field, this surely helps JOIN operations.
    # id             = models.UUIDField(primary_key=True,
    #                                   default=uuid.uuid4, editable=False)

    user = models.ForeignKey(User, verbose_name=_(u'Owner'),
                             related_name='accounts')

    name = models.CharField(verbose_name=_(u'Account name'),
                            max_length=128, blank=True)

    slug = models.CharField(verbose_name=_(u'slug'),
                            max_length=255,
                            null=True, blank=True)

    # items = models.ManyToManyField(BaseItem, blank=True, null=True,
    #                               verbose_name=_(u'Feed items'),
    #                               related_name='feeds')

    date_created = models.DateTimeField(auto_now_add=True, blank=True,
                                        verbose_name=_(u'Date created'))

    date_last_conn = models.DateTimeField(
        verbose_name=_(u'Last connection date'),
        null=True, blank=True)

    conn_error = models.CharField(
        verbose_name=_(u'Last connection error'),
        max_length=255, null=True, blank=True)

    is_usable = models.BooleanField(
        verbose_name=_(u'usable'),
        default=True, blank=True,
        help_text=_(u'Is account connectivity good or not?'))

    is_active = models.BooleanField(
        verbose_name=_(u'active'),
        default=True, blank=True,
        help_text=_(u'Is account enabled for fetching?'))

    options = JSONField(default=dict, blank=True)

    notes = models.TextField(
        verbose_name=_(u'Notes'), null=True, blank=True,
        help_text=_(u'User notes for this account.'))

    image = models.ImageField(
        verbose_name=_(u'Thumbnail'), null=True, blank=True,
        upload_to=get_account_image_upload_path, max_length=256,
        help_text=_(u'Use either image when 1flow instance hosts the '
                    u'image, or image_url when hosted elsewhere. If '
                    u'both are filled, image takes precedence.'))

    image_url = models.URLField(
        verbose_name=_(u'Thumbnail URL'), null=True, blank=True, max_length=384,
        help_text=_(u'Full URL of the image displayed in the feed '
                    u'selector. Can be hosted outside of 1flow.'))

    # —————————————————————————————————————————————————————— Django & Grappelli

    def __unicode__(self):
        """ Hello, pep257. I love you so. """

        return _(u'BaseAccount {0} ({1}) of {2}').format(
            self.name, self.id, self.user.username)

    @staticmethod
    def autocomplete_search_fields():
        """ grappelli auto-complete method. """

        return ('name__icontains', )

    # —————————————————————————————————————————————————————————————— Properties

    @property
    def config_account_refresh_period(self):
        """ Return the default account refresh period.

        .. note:: it's recommended to override this
            property in polymorphic inheritant models.
        """

        return config.ACCOUNT_REFRESH_PERIOD_DEFAULT

    @property
    def refresh_lock(self):
        try:
            return self.__refresh_lock

        except AttributeError:
            self.__refresh_lock = RedisExpiringLock(
                self, lock_name='account_fetch',
            )
            return self.__refresh_lock

    @property
    def recently_usable(self):
        """ Return True if the account has been tested/connected recently. """

        return self.is_usable and (
            now() - self.date_last_conn
            < timedelta(seconds=self.config_account_refresh_period))

    # ————————————————————————————————————————————————————————————————— Methods

    def reset_unusable(self, commit=True):
        """ Mark the current instance needing to test usability.

        This is called typically when a connection parameter has changed,
        and the current account connectivity needs to be tested again to
        validate them all.
        """

        self.date_last_conn = None
        self.is_usable = False

        if commit:
            self.save()

        try:
            reset_unusable_task = self._meta.model.reset_unusable_task

        except:
            pass

        else:
            reset_unusable_task.apply_async(
                args=(self.id, ), countdown=3)

    def mark_unusable(self, message, args=(), exc=None, commit=True):
        """ Mark account unsable with date & message, log exception if any. """

        if exc is not None:
            if args:
                message = message % args

            message = u'{0} ({1})'.format(message, unicode(exc))
            LOGGER.exception(u'%s unusable: %s', self, message)

        self.date_last_conn = now()
        self.conn_error = message
        self.is_usable = False

        if commit:
            self.save()

    def mark_usable(self, commit=True, verbose=True):
        """ Mark the account usable and clear error. """

        if verbose:
            LOGGER.info(u'%s is now considered usable.', self)

        if self.is_usable:
            start_task = False

        else:
            start_task = True

        self.date_last_conn = now()
        self.conn_error = None
        self.is_usable = True

        if commit:
            self.save()

        try:
            usable_start_task = self._meta.model.usable_start_task

        except:
            pass

        else:
            if start_task:
                usable_start_task.delay(self.id)

    def has_option(self, option):
        """ True if option in self.options. """

        return option in self.options
<|endoftext|>