<commit_before>import os
import random
from collections import defaultdict
from pprint import pprint

import numpy as np
import torch
import torch.optim as optim

from metal.logging import Checkpointer, LogWriter, TensorBoardWriter
from metal.mmtl.mmtl_logger import Logger  # NOTE: we load special MTL logger
from metal.utils import recursive_merge_dicts

# Import tqdm_notebook if in Jupyter notebook
try:
    from IPython import get_ipython

    if "IPKernelApp" not in get_ipython().config:
        raise ImportError("console")
except (AttributeError, ImportError):
    from tqdm import tqdm
else:
    # Only use tqdm notebook if not in travis testing
    if "CI" not in os.environ:
        from tqdm import tqdm_notebook as tqdm
    else:
        from tqdm import tqdm

trainer_config = {
    "verbose": True,
    # Display
    "progress_bar": False,
    # Dataloader
    # TODO: Restore the option for them to pass in raw simple data which we wrap up
    # "data_loader_config": {"batch_size": 32, "num_workers": 1, "shuffle": True},
    # Loss weights
    # TODO: Restore ability to weight losses by class and/or task
    # "loss_weights": None,
    # Train Loop
    "n_epochs": 10,
    # 'grad_clip': 0.0,
    "l2": 0.0,
    # Evaluate dev for during training every this many epochs
    # Optimizer
    "optimizer_config": {
        "optimizer": "adam",
        "optimizer_common": {"lr": 0.01},
        # Optimizer - SGD
        "sgd_config": {"momentum": 0.9},
        # Optimizer - Adam
        "adam_config": {"betas": (0.9, 0.999)},
        # Optimizer - RMSProp
        "rmsprop_config": {},  # Use defaults
    },
    # LR Scheduler (for learning rate)
    "lr_scheduler": "reduce_on_plateau",
    # [None, 'exponential', 'reduce_on_plateau']
    # 'reduce_on_plateau' uses checkpoint_metric to assess plateaus
    "lr_scheduler_config": {
        # Freeze learning rate initially this many epochs
        "lr_freeze": 0,
        # Scheduler - exponential
        "exponential_config": {"gamma": 0.9},  # decay rate
        # Scheduler - reduce_on_plateau
        "plateau_config": {
            "factor": 0.5,
            "patience": 10,
            "threshold": 0.0001,
            "min_lr": 1e-4,
        },
    },
    # Logger (see metal/logging/logger.py for descriptions)
    "logger": True,
    "logger_config": {
        "log_unit": "epochs",  # ['seconds', 'examples', 'batches', 'epochs']
        # Report loss every this many log_units
        "log_every": 1,
        # Calculate and report metrics every this many log_units:
        #   None: default to log_every
        #   0: do not calculate or log metrics
        #   otherwise: must be a multiple of log_every
        "score_every": None,
        # If non-None, only calculate and report these metrics every `score_every`
        # units (this can include the names of built-in and user-defined metrics);
        # otherwise, include all metrics returned by task Scorers.
        # TODO: "metrics_filter": None,
        # TODO: "score_limit": None,  # Evaluate scorer on only this many examples
    },
    # LogWriter/Tensorboard (see metal/logging/writer.py for descriptions)
    "writer": None,  # [None, "json", "tensorboard"]
    "writer_config": {  # Log (or event) file stored at log_dir/run_dir/run_name
        "log_dir": None,
        "run_dir": None,
        "run_name": None,
        "writer_metrics": [],  # May specify a subset of metrics in metrics_dict to be written
        "include_config": True,  # If True, include model config in log
    },
    # Checkpointer (see metal/logging/checkpointer.py for descriptions)
    "checkpoint": True,  # If True, checkpoint models when certain conditions are met
    "checkpoint_config": {
        # TODO: unify checkpoint=['every', 'best', 'final']; specify one strategy
        "checkpoint_every": 0,  # Save a model checkpoint every this many log_units
        # If checkpoint_best, also save the "best" model according to some metric
        # The "best" model will have the ['max', 'min'] value of checkpoint_metric
        # This metric must be produced by one of the task Scorer objects so it will be
        # available for lookup; assumes valid split unless appended with "train/"
        "checkpoint_best": False,
        # "checkpoint_final": False,  # Save a model checkpoint at the end of training
        "checkpoint_metric": "train/loss",
        "checkpoint_metric_mode": "min",
        "checkpoint_dir": f"{os.environ['METALHOME']}/checkpoints",
        "checkpoint_runway": 0,
    },
}

class MultitaskTrainer(object):
    """Driver for the MTL training process"""

    def __init__(self, config={}):
        self.config = recursive_merge_dicts(trainer_config, config)
        self._normalize_metric_names()

    def train_model(self, model, tasks, **kwargs):
        self.config = recursive_merge_dicts(self.config, kwargs)
        self.task_names = [task.name for task in tasks]

        # Calculate epoch statistics
        examples_per_epoch = sum([len(t.data_loaders["train"].dataset) for t in tasks])
        batches_per_epoch = sum([len(t.data_loaders["train"]) for t in tasks])
        if self.config["verbose"]:
            print(f"Beginning train loop.")
            print(
                f"Expecting a total of {examples_per_epoch} examples and "
                f"{batches_per_epoch} batches per epoch from {len(tasks)} tasks."
            )

        # Set training components
        self._set_writer()
        self._set_logger(batches_per_epoch)
        self._set_checkpointer()
        self._set_optimizer(model)
        # TODO: Accept training matrix to give more fine-tuned training commands
        self._set_scheduler()

        # TODO: Restore the ability to resume training from a given epoch
        # That code goes here

        # Train the model
        # TODO: Allow other ways to train besides 1 epoch of all datasets
        model.train()
        # Dict metrics_hist contains the most recently recorded value of all metrics
        self.metrics_hist = {}
        self._reset_losses()
        for epoch in range(self.config["n_epochs"]):
            progress_bar = self.config["progress_bar"] and self.config["verbose"]
            t = tqdm(
                enumerate(self._get_train_batches(tasks)),
                total=batches_per_epoch,
                disable=(not progress_bar),
            )
            for batch_num, (task_name, batch) in t:
                # NOTE: actual batch_size may not equal config's target batch_size
                batch_size = len(batch[0])

                # Zero the parameter gradients
                self.optimizer.zero_grad()

                # Forward pass to calculate the average loss per example
                loss = model.calculate_loss(*batch, [task_name])[task_name]
                if torch.isnan(loss):
                    msg = "Loss is NaN. Consider reducing learning rate."
                    raise Exception(msg)

                # Backward pass to calculate gradients
                # Loss is an average loss per example
                loss.backward()

                # Perform optimizer step
                self.optimizer.step()

                # Update loss
                self.running_losses[task_name] += loss.item() * batch_size
                self.running_examples[task_name] += batch_size

                # Calculate metrics, log, and checkpoint as necessary
                metrics_dict = self._execute_logging(model, tasks, batch_size)

                # Update most recently seen value of each metric
                self.metrics_hist.update(metrics_dict)

                # tqdm output
                if len(tasks) == 1:
                    t.set_postfix(loss=metrics_dict["train/loss"])
                else:
                    losses = {}
                    for key in metrics_dict:
                        if "loss" in key:
                            losses[key] = metrics_dict[key]
                    t.set_postfix(losses)

            # Apply learning rate scheduler
            # self._update_scheduler(epoch, metrics_hist)

        model.eval()

        # Restore best model if applicable
        if self.checkpointer and self.checkpointer.checkpoint_best:
            self.checkpointer.load_best_model(model=model)

        # Write log if applicable
        if self.writer:
            if self.writer.include_config:
                self.writer.add_config(self.config)
            self.writer.close()

        # Print final performance values
        if self.config["verbose"]:
            print("Finished Training")
            metrics_dict = self.calculate_metrics(model, tasks)
            pprint(metrics_dict)

    def _execute_logging(self, model, tasks, batch_size):
        model.eval()
        metrics_dict = {}
        metrics_dict.update(self.calculate_losses(tasks))
        # HACK: This exposes more of the logger than it should; abstract away
        self.logger.increment(batch_size)
        if self.logger.loss_time():
            self._reset_losses()
            self.logger.loss_ticks += 1
        if self.logger.metrics_time():
            metrics_dict.update(self.calculate_metrics(model, tasks))
            self.logger.loss_ticks = 0
        if self.logger.loss_time() or self.logger.metrics_time():
            # Log to screen/file/TensorBoard
            self.logger.log(metrics_dict)
            # Save best model if applicable
            self._checkpoint(model, metrics_dict)
        model.train()
        return metrics_dict

    def calculate_losses(self, tasks):
        """Calculate the average loss for each task since the last calculation

        If no examples of a certain task have been seen since the losses were reset,
        use the most recently reported value again (stored in metrics_hist).
        If the loss for a certain task has never been reported, report it as None.
        """
        metrics_dict = {}
        for task in tasks:
            if self.running_examples[task.name]:
                loss = self.running_losses[task.name] / self.running_examples[task.name]
            elif self.metrics_hist.get(f"{task.name}/loss"):
                loss = self.metrics_hist[f"{task.name}/loss"]
            else:
                loss = None
            metrics_dict[f"{task.name}/loss"] = loss
        # Report micro average of losses
        total_loss = sum(self.running_losses.values())
        total_examples = sum(self.running_examples.values())
        # TODO: Don't report task loss and "overall" loss if there is only one task?
        # But they may be planning on their named task loss being in the metrics_dict...
        metrics_dict["train/loss"] = total_loss / total_examples
        return metrics_dict

    def calculate_metrics(self, model, tasks):
        metrics_dict = {}
        for task in tasks:
            task_metrics = task.scorer.score(model, task)
            metrics_dict.update(task_metrics)
        return metrics_dict

    def _get_train_batches(self, tasks):
        """Yields batches one at a time sampled from tasks with some strategy"""
        # TODO: Allow more involved strategies for sampling from tasks
        # For now, just use proportional sampling
        # Length of a dataloader is the number of batches it contains
        approx_batch_counts = [len(t.data_loaders["train"]) for t in tasks]
        batch_assignments = []
        for task_idx, task in enumerate(tasks):
            batch_assignments.extend([task_idx] * approx_batch_counts[task_idx])
        random.shuffle(batch_assignments)
        train_loaders = [iter(t.data_loaders["train"]) for t in tasks]

        for task_idx in batch_assignments:
            yield (tasks[task_idx].name, next(train_loaders[task_idx]))

    def _checkpoint(self, model, metrics_dict):
        if self.checkpointer is None:
            return
        iteration = self.logger.unit_total
        self.checkpointer.checkpoint(
            metrics_dict, iteration, model, self.optimizer, self.lr_scheduler
        )

    def _reset_losses(self):
        self.running_losses = defaultdict(float)
        self.running_examples = defaultdict(int)

    def _set_writer(self):
        if self.config["writer"] is None:
            self.writer = None
        elif self.config["writer"] == "json":
            self.writer = LogWriter(**(self.config["writer_config"]))
        elif self.config["writer"] == "tensorboard":
            self.writer = TensorBoardWriter(**(self.config["writer_config"]))
        else:
            raise Exception(f"Unrecognized writer: {self.config['writer']}")

    def _set_logger(self, batches_per_epoch):
        # If not provided, set score_every to log_every
        logger_config = self.config["logger_config"]
        if logger_config["score_every"] is None:
            logger_config["score_every"] = logger_config["log_every"]
        self.logger = Logger(
            logger_config,
            batches_per_epoch,
            self.writer,
            verbose=self.config["verbose"],
        )

    def _set_checkpointer(self):
        if self.config["checkpoint"]:
            checkpoint_metric = self.config["checkpoint_config"]["checkpoint_metric"]
            if (
                len(self.task_names) > 1
                and checkpoint_metric != "train/loss"
                and not any(
                    task_name in checkpoint_metric for task_name in self.task_names
                )
            ):
                raise Exception(
                    "When len(tasks) > 1, checkpoint_metric must be train/loss or else must include task name; e.g., task/split/metric or task/metric (with assumed split='valid')"
                )
            self.checkpointer = Checkpointer(
                self.config["checkpoint_config"], verbose=self.config["verbose"]
            )
        else:
            self.checkpointer = None

    def _set_optimizer(self, model):
        optimizer_config = self.config["optimizer_config"]
        opt = optimizer_config["optimizer"]

        parameters = filter(lambda p: p.requires_grad, model.parameters())
        if opt == "sgd":
            optimizer = optim.SGD(
                parameters,
                **optimizer_config["optimizer_common"],
                **optimizer_config["sgd_config"],
                weight_decay=self.config["l2"],
            )
        elif opt == "rmsprop":
            optimizer = optim.RMSprop(
                parameters,
                **optimizer_config["optimizer_common"],
                **optimizer_config["rmsprop_config"],
                weight_decay=self.config["l2"],
            )
        elif opt == "adam":
            optimizer = optim.Adam(
                parameters,
                **optimizer_config["optimizer_common"],
                **optimizer_config["adam_config"],
                weight_decay=self.config["l2"],
            )
        elif opt == "adamax":
            optimizer = optim.Adamax(
                parameters,
                **optimizer_config["optimizer_common"],
                **optimizer_config["adam_config"],
                weight_decay=self.config["l2"],
            )
        elif opt == "sparseadam":
            optimizer = optim.SparseAdam(
                parameters,
                **optimizer_config["optimizer_common"],
                **optimizer_config["adam_config"],
            )
            if self.config["l2"]:
                raise Exception(
                    "SparseAdam optimizer does not support weight_decay (l2 penalty)."
                )
        else:
            raise ValueError(f"Did not recognize optimizer option '{opt}'")
        self.optimizer = optimizer

    def _set_scheduler(self):
        lr_scheduler = self.config["lr_scheduler"]
        if lr_scheduler is None:
            lr_scheduler = None
        else:
            lr_scheduler_config = self.config["lr_scheduler_config"]
            if lr_scheduler == "exponential":
                lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
                    self.optimizer, **lr_scheduler_config["exponential_config"]
                )
            elif lr_scheduler == "reduce_on_plateau":
                lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                    self.optimizer, **lr_scheduler_config["plateau_config"]
                )
            else:
                raise ValueError(
                    f"Did not recognize lr_scheduler option '{lr_scheduler}'"
                )
        self.lr_scheduler = lr_scheduler

    def _normalize_metric_names(self):
        # TODO: expand and check metric names in log_[valid/train]_metrics, checkpoint_metric, and writer_metrics
        # add train/valid based on log_train_metrics or log_valid_metrics
        # check if task names are required and not included (e.g., if multitask)
        pass

    def _expand_metric_name(self, metric, split, task):
        return f"task/split/{metric.split('/')[-1]}"
<commit_msg>fix checkpoint needing task name even in single task case bug<commit_after>import os
import random
from collections import defaultdict
from pprint import pprint

import numpy as np
import torch
import torch.optim as optim

from metal.logging import Checkpointer, LogWriter, TensorBoardWriter
from metal.mmtl.mmtl_logger import Logger  # NOTE: we load special MTL logger
from metal.utils import recursive_merge_dicts

# Import tqdm_notebook if in Jupyter notebook
try:
    from IPython import get_ipython

    if "IPKernelApp" not in get_ipython().config:
        raise ImportError("console")
except (AttributeError, ImportError):
    from tqdm import tqdm
else:
    # Only use tqdm notebook if not in travis testing
    if "CI" not in os.environ:
        from tqdm import tqdm_notebook as tqdm
    else:
        from tqdm import tqdm

trainer_config = {
    "verbose": True,
    # Display
    "progress_bar": False,
    # Dataloader
    # TODO: Restore the option for them to pass in raw simple data which we wrap up
    # "data_loader_config": {"batch_size": 32, "num_workers": 1, "shuffle": True},
    # Loss weights
    # TODO: Restore ability to weight losses by class and/or task
    # "loss_weights": None,
    # Train Loop
    "n_epochs": 10,
    # 'grad_clip': 0.0,
    "l2": 0.0,
    # Evaluate dev for during training every this many epochs
    # Optimizer
    "optimizer_config": {
        "optimizer": "adam",
        "optimizer_common": {"lr": 0.01},
        # Optimizer - SGD
        "sgd_config": {"momentum": 0.9},
        # Optimizer - Adam
        "adam_config": {"betas": (0.9, 0.999)},
        # Optimizer - RMSProp
        "rmsprop_config": {},  # Use defaults
    },
    # LR Scheduler (for learning rate)
    "lr_scheduler": "reduce_on_plateau",
    # [None, 'exponential', 'reduce_on_plateau']
    # 'reduce_on_plateau' uses checkpoint_metric to assess plateaus
    "lr_scheduler_config": {
        # Freeze learning rate initially this many epochs
        "lr_freeze": 0,
        # Scheduler - exponential
        "exponential_config": {"gamma": 0.9},  # decay rate
        # Scheduler - reduce_on_plateau
        "plateau_config": {
            "factor": 0.5,
            "patience": 10,
            "threshold": 0.0001,
            "min_lr": 1e-4,
        },
    },
    # Logger (see metal/logging/logger.py for descriptions)
    "logger": True,
    "logger_config": {
        "log_unit": "epochs",  # ['seconds', 'examples', 'batches', 'epochs']
        # Report loss every this many log_units
        "log_every": 1,
        # Calculate and report metrics every this many log_units:
        #   None: default to log_every
        #   0: do not calculate or log metrics
        #   otherwise: must be a multiple of log_every
        "score_every": None,
        # If non-None, only calculate and report these metrics every `score_every`
        # units (this can include the names of built-in and user-defined metrics);
        # otherwise, include all metrics returned by task Scorers.
        # TODO: "metrics_filter": None,
        # TODO: "score_limit": None,  # Evaluate scorer on only this many examples
    },
    # LogWriter/Tensorboard (see metal/logging/writer.py for descriptions)
    "writer": None,  # [None, "json", "tensorboard"]
    "writer_config": {  # Log (or event) file stored at log_dir/run_dir/run_name
        "log_dir": None,
        "run_dir": None,
        "run_name": None,
        "writer_metrics": [],  # May specify a subset of metrics in metrics_dict to be written
        "include_config": True,  # If True, include model config in log
    },
    # Checkpointer (see metal/logging/checkpointer.py for descriptions)
    "checkpoint": True,  # If True, checkpoint models when certain conditions are met
    "checkpoint_config": {
        # TODO: unify checkpoint=['every', 'best', 'final']; specify one strategy
        "checkpoint_every": 0,  # Save a model checkpoint every this many log_units
        # If checkpoint_best, also save the "best" model according to some metric
        # The "best" model will have the ['max', 'min'] value of checkpoint_metric
        # This metric must be produced by one of the task Scorer objects so it will be
        # available for lookup; assumes valid split unless appended with "train/"
        "checkpoint_best": False,
        # "checkpoint_final": False,  # Save a model checkpoint at the end of training
        "checkpoint_metric": "train/loss",
        "checkpoint_metric_mode": "min",
        "checkpoint_dir": f"{os.environ['METALHOME']}/checkpoints",
        "checkpoint_runway": 0,
    },
}

class MultitaskTrainer(object):
    """Driver for the MTL training process"""

    def __init__(self, config={}):
        self.config = recursive_merge_dicts(trainer_config, config)
        self._normalize_metric_names()

    def train_model(self, model, tasks, **kwargs):
        self.config = recursive_merge_dicts(self.config, kwargs)
        self.task_names = [task.name for task in tasks]

        # Calculate epoch statistics
        examples_per_epoch = sum([len(t.data_loaders["train"].dataset) for t in tasks])
        batches_per_epoch = sum([len(t.data_loaders["train"]) for t in tasks])
        if self.config["verbose"]:
            print(f"Beginning train loop.")
            print(
                f"Expecting a total of {examples_per_epoch} examples and "
                f"{batches_per_epoch} batches per epoch from {len(tasks)} tasks."
            )

        # Set training components
        self._set_writer()
        self._set_logger(batches_per_epoch)
        self._set_checkpointer()
        self._set_optimizer(model)
        # TODO: Accept training matrix to give more fine-tuned training commands
        self._set_scheduler()

        # TODO: Restore the ability to resume training from a given epoch
        # That code goes here

        # Train the model
        # TODO: Allow other ways to train besides 1 epoch of all datasets
        model.train()
        # Dict metrics_hist contains the most recently recorded value of all metrics
        self.metrics_hist = {}
        self._reset_losses()
        for epoch in range(self.config["n_epochs"]):
            progress_bar = self.config["progress_bar"] and self.config["verbose"]
            t = tqdm(
                enumerate(self._get_train_batches(tasks)),
                total=batches_per_epoch,
                disable=(not progress_bar),
            )
            for batch_num, (task_name, batch) in t:
                # NOTE: actual batch_size may not equal config's target batch_size
                batch_size = len(batch[0])

                # Zero the parameter gradients
                self.optimizer.zero_grad()

                # Forward pass to calculate the average loss per example
                loss = model.calculate_loss(*batch, [task_name])[task_name]
                if torch.isnan(loss):
                    msg = "Loss is NaN. Consider reducing learning rate."
                    raise Exception(msg)

                # Backward pass to calculate gradients
                # Loss is an average loss per example
                loss.backward()

                # Perform optimizer step
                self.optimizer.step()

                # Update loss
                self.running_losses[task_name] += loss.item() * batch_size
                self.running_examples[task_name] += batch_size

                # Calculate metrics, log, and checkpoint as necessary
                metrics_dict = self._execute_logging(model, tasks, batch_size)

                # Update most recently seen value of each metric
                self.metrics_hist.update(metrics_dict)

                # tqdm output
                if len(tasks) == 1:
                    t.set_postfix(loss=metrics_dict["train/loss"])
                else:
                    losses = {}
                    for key in metrics_dict:
                        if "loss" in key:
                            losses[key] = metrics_dict[key]
                    t.set_postfix(losses)

            # Apply learning rate scheduler
            # self._update_scheduler(epoch, metrics_hist)

        model.eval()

        # Restore best model if applicable
        if self.checkpointer and self.checkpointer.checkpoint_best:
            self.checkpointer.load_best_model(model=model)

        # Write log if applicable
        if self.writer:
            if self.writer.include_config:
                self.writer.add_config(self.config)
            self.writer.close()

        # Print final performance values
        if self.config["verbose"]:
            print("Finished Training")
            metrics_dict = self.calculate_metrics(model, tasks)
            pprint(metrics_dict)

    def _execute_logging(self, model, tasks, batch_size):
        model.eval()
        metrics_dict = {}
        metrics_dict.update(self.calculate_losses(tasks))
        # HACK: This exposes more of the logger than it should; abstract away
        self.logger.increment(batch_size)
        if self.logger.loss_time():
            self._reset_losses()
            self.logger.loss_ticks += 1
        if self.logger.metrics_time():
            metrics_dict.update(self.calculate_metrics(model, tasks))
            self.logger.loss_ticks = 0
        if self.logger.loss_time() or self.logger.metrics_time():
            # Log to screen/file/TensorBoard
            self.logger.log(metrics_dict)
            # Save best model if applicable
            self._checkpoint(model, metrics_dict)
        model.train()
        return metrics_dict

    def calculate_losses(self, tasks):
        """Calculate the average loss for each task since the last calculation

        If no examples of a certain task have been seen since the losses were reset,
        use the most recently reported value again (stored in metrics_hist).
        If the loss for a certain task has never been reported, report it as None.
        """
        metrics_dict = {}
        for task in tasks:
            if self.running_examples[task.name]:
                loss = self.running_losses[task.name] / self.running_examples[task.name]
            elif self.metrics_hist.get(f"{task.name}/loss"):
                loss = self.metrics_hist[f"{task.name}/loss"]
            else:
                loss = None
            metrics_dict[f"{task.name}/loss"] = loss
        # Report micro average of losses
        total_loss = sum(self.running_losses.values())
        total_examples = sum(self.running_examples.values())
        # TODO: Don't report task loss and "overall" loss if there is only one task?
        # But they may be planning on their named task loss being in the metrics_dict...
        metrics_dict["train/loss"] = total_loss / total_examples
        return metrics_dict

    def calculate_metrics(self, model, tasks):
        metrics_dict = {}
        for task in tasks:
            task_metrics = task.scorer.score(model, task)
            metrics_dict.update(task_metrics)
        return metrics_dict

    def _get_train_batches(self, tasks):
        """Yields batches one at a time sampled from tasks with some strategy"""
        # TODO: Allow more involved strategies for sampling from tasks
        # For now, just use proportional sampling
        # Length of a dataloader is the number of batches it contains
        approx_batch_counts = [len(t.data_loaders["train"]) for t in tasks]
        batch_assignments = []
        for task_idx, task in enumerate(tasks):
            batch_assignments.extend([task_idx] * approx_batch_counts[task_idx])
        random.shuffle(batch_assignments)
        train_loaders = [iter(t.data_loaders["train"]) for t in tasks]

        for task_idx in batch_assignments:
            yield (tasks[task_idx].name, next(train_loaders[task_idx]))

    def _checkpoint(self, model, metrics_dict):
        if self.checkpointer is None:
            return
        iteration = self.logger.unit_total
        self.checkpointer.checkpoint(
            metrics_dict, iteration, model, self.optimizer, self.lr_scheduler
        )

    def _reset_losses(self):
        self.running_losses = defaultdict(float)
        self.running_examples = defaultdict(int)

    def _set_writer(self):
        if self.config["writer"] is None:
            self.writer = None
        elif self.config["writer"] == "json":
            self.writer = LogWriter(**(self.config["writer_config"]))
        elif self.config["writer"] == "tensorboard":
            self.writer = TensorBoardWriter(**(self.config["writer_config"]))
        else:
            raise Exception(f"Unrecognized writer: {self.config['writer']}")

    def _set_logger(self, batches_per_epoch):
        # If not provided, set score_every to log_every
        logger_config = self.config["logger_config"]
        if logger_config["score_every"] is None:
            logger_config["score_every"] = logger_config["log_every"]
        self.logger = Logger(
            logger_config,
            batches_per_epoch,
            self.writer,
            verbose=self.config["verbose"],
        )

    def _set_checkpointer(self):
        if self.config["checkpoint"]:
            checkpoint_metric = self.config["checkpoint_config"]["checkpoint_metric"]
            if checkpoint_metric != "train/loss" and not any(
                task_name in checkpoint_metric for task_name in self.task_names
            ):
                raise Exception(
                    "checkpoint_metric must be train/loss or else must include task name; e.g., task/split/metric or task/metric (with assumed split='valid')"
                )
            self.checkpointer = Checkpointer(
                self.config["checkpoint_config"], verbose=self.config["verbose"]
            )
        else:
            self.checkpointer = None

    def _set_optimizer(self, model):
        optimizer_config = self.config["optimizer_config"]
        opt = optimizer_config["optimizer"]

        parameters = filter(lambda p: p.requires_grad, model.parameters())
        if opt == "sgd":
            optimizer = optim.SGD(
                parameters,
                **optimizer_config["optimizer_common"],
                **optimizer_config["sgd_config"],
                weight_decay=self.config["l2"],
            )
        elif opt == "rmsprop":
            optimizer = optim.RMSprop(
                parameters,
                **optimizer_config["optimizer_common"],
                **optimizer_config["rmsprop_config"],
                weight_decay=self.config["l2"],
            )
        elif opt == "adam":
            optimizer = optim.Adam(
                parameters,
                **optimizer_config["optimizer_common"],
                **optimizer_config["adam_config"],
                weight_decay=self.config["l2"],
            )
        elif opt == "adamax":
            optimizer = optim.Adamax(
                parameters,
                **optimizer_config["optimizer_common"],
                **optimizer_config["adam_config"],
                weight_decay=self.config["l2"],
            )
        elif opt == "sparseadam":
            optimizer = optim.SparseAdam(
                parameters,
                **optimizer_config["optimizer_common"],
                **optimizer_config["adam_config"],
            )
            if self.config["l2"]:
                raise Exception(
                    "SparseAdam optimizer does not support weight_decay (l2 penalty)."
                )
        else:
            raise ValueError(f"Did not recognize optimizer option '{opt}'")
        self.optimizer = optimizer

    def _set_scheduler(self):
        lr_scheduler = self.config["lr_scheduler"]
        if lr_scheduler is None:
            lr_scheduler = None
        else:
            lr_scheduler_config = self.config["lr_scheduler_config"]
            if lr_scheduler == "exponential":
                lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
                    self.optimizer, **lr_scheduler_config["exponential_config"]
                )
            elif lr_scheduler == "reduce_on_plateau":
                lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                    self.optimizer, **lr_scheduler_config["plateau_config"]
                )
            else:
                raise ValueError(
                    f"Did not recognize lr_scheduler option '{lr_scheduler}'"
                )
        self.lr_scheduler = lr_scheduler

    def _normalize_metric_names(self):
        # TODO: expand and check metric names in log_[valid/train]_metrics, checkpoint_metric, and writer_metrics
        # add train/valid based on log_train_metrics or log_valid_metrics
        # check if task names are required and not included (e.g., if multitask)
        pass

    def _expand_metric_name(self, metric, split, task):
        return f"task/split/{metric.split('/')[-1]}"
<|endoftext|>


<commit_before># copyright: (c) 2012 by Hansel Dunlop.
# license: ISC, see LICENSE for more details.
#
# kubrick.api
#
# This module implements the Kubrick api.
#

import os

from boto.ec2.connection import EC2Connection
from boto.ec2 import get_region
from boto.s3.connection import S3Connection
from boto.s3.key import Key
from fabric.api import env, put, run
from fabric.contrib.files import append
from fabric.state import connections
import sys
import time

import config

class Server(object):
    install_command_map = {
        'apt': 'sudo apt-get install',
        'pip': 'pip install --upgrade',
        'easy_install': 'easy_install --upgrade',
    }

class AWSServer(Server):

    def __init__(self, config, identifier):
        self.key_id = config.AWS_ACCESS_KEY_ID
        self.secret_key = config.AWS_SECRET_ACCESS_KEY
        self.ami_id = config.BASE_AMI
        self.config = config
        self.key_filename = config.KEY_FILENAME
        self.identifier = identifier
        self.conn = EC2Connection(self.key_id, self.secret_key)
        self.instance = None
        self.root_account = config.USERNAME
        self.operating_system = config.OPERATING_SYSTEM
        if self.operating_system == 'Ubuntu':
            self.sudo_required = True
        else:
            self.sudo_required = False

    @property
    def host_string(self):
        return self.root_account + '@' + self.instance.public_dns_name

    def run(self, command, warn_only=False):
        if not self.instance:
            raise Exception('No machine attached to this server instance')
        env.host_string = self.host_string
        env.disable_known_hosts = True
        env.user = self.config.USERNAME
        env.connection_attempts = 6
        env.warn_only = warn_only
        env.key_filename = self.key_filename
        if self.sudo_required:
            command = 'sudo ' + command
        run(command)

    def append(self, filename, text):
        if not self.instance:
            raise Exception('No machine attached to this server instance')
        env.host_string = self.host_string
        env.disable_known_hosts = True
        env.user = self.config.USERNAME
        env.connection_attempts = 6
        env.key_filename = self.key_filename
        if self.sudo_required:
            append(filename, text, use_sudo=True)
        else:
            append(filename, text)

    def put(self, src, dest):
        env.host_string = self.host_string
        env.disable_known_hosts = True
        env.user = self.config.USERNAME
        env.connection_attempts = 6
        env.key_filename = self.key_filename
        if self.sudo_required:
            put(src, dest, use_sudo=True, mirror_local_mode=True)
            return
        put(src, dest, mirror_local_mode=True)

    def start_server(self):
        self.instance = start_server_from_ami(self)

    def destroy(self):
        if self.instance:
            print 'Destroying AWS server', self.instance
            self.instance.terminate()
        else:
            print 'No instance on this server object'

    def reboot(self):
        self.run('reboot')
        time.sleep(30)
        print "Reconnecting",
        sys.stdout.flush()
        for retry in range(60):
            try:
                print ".",
                sys.stdout.flush()
                env.host_string = self.host_string
                connections.connect(env.host_string)
                break
            except:
                print "-",
                sys.stdout.flush()
                time.sleep(3)
        print

class UbuntuMixin(object):
    '''A mixin for Ubuntu based servers'''

    apt_updated = False
    apt_command = 'apt-get -qq -y'

    def make_swap(self, size):
        swap_location = '/var/swap'
        self.run('dd if=/dev/zero of=%s bs=1M count=%d' % (swap_location, size,))
        self.run('mkswap ' + swap_location)
        self.run('swapon ' + swap_location)
        self.append(
            '/etc/fstab', '%s swap swap defaults 0 0' % (swap_location,)
        )

    def purge_apt_packages(self):
        self.update_apt_if_necessary()
        for package in self.APT_PURGES:
            self.run(self.apt_command + ' --purge remove ' + package)

    def update_apt_if_necessary(self):
        if not self.apt_updated:
            self.run(self.apt_command + ' update')
            self.apt_updated = True

    def upgrade_installed_packages(self):
        self.update_apt_if_necessary()
        self.run(self.apt_command + ' upgrade')
        self.apt_updated = False

    def install_apt_packages(self):
        self.update_apt_if_necessary()
        self.run(self.apt_command + ' install ' + ' '.join(self.APT_INSTALLS))

    def install_python_modules(self):
        self.run('easy_install -U distribute')
        for package in self.PIP_INSTALLS:
            self.run('pip install %s' % (package,))

    def configure_locales(self):
        # prevents certain warning messages
        self.run('locale-gen en_US en_US.UTF-8 en_GB.UTF-8')
        self.run('dpkg-reconfigure locales')
        self.run('mkdir -p /etc/ssl/intermediate_and_root_certs/')

    def add_non_root_user(self, user, uid, gid, groups, shell='/bin/bash', base_dir='/home'):
        groups_arg = ','.join(groups)

        self.run('groupadd --gid %d %s' % (gid, user))
        self.run(
            'useradd -m -b %s -s %s --uid %d --gid %d -G %s %s' % (
                base_dir, shell, uid, gid, groups_arg, user
            )
        )

    def put_static_config_files(self):
        # recursively put all the files in the server_config dir
        config_files = []
        for dirname, dirs, files in os.walk(self.CONFIG_LOCATION):
            for path in files:
                full_path = os.path.join(dirname, path)
                config_files.append(
                    (full_path, full_path.replace(self.CONFIG_LOCATION, ''))
                )

        for src, dest in config_files:
            self.run('mkdir -p ' + os.path.dirname(dest))
            self.put(src, dest)

    def setup_cron(self):
        self.run('chmod 600 /etc/crontab')
        self.run('chgrp root /etc/crontab')
        self.run('chown root /etc/crontab')

    def set_hostname(self, hostname):
        self.run('hostname %s' % (hostname,))

def get_region_from_name(zone_name=config.DEFAULT_ZONE):
    return get_region(
        zone_name,
        aws_access_key_id=config.AWS_ACCESS_KEY_ID,
        aws_secret_access_key=config.AWS_SECRET_ACCESS_KEY,
    )

def create_aws_connection():
    return EC2Connection(
        config.AWS_ACCESS_KEY_ID,
        config.AWS_SECRET_ACCESS_KEY,
        region=get_region_from_name()
    )

def create_storage_connection():
    return S3Connection(
        config.AWS_ACCESS_KEY_ID,
        config.AWS_SECRET_ACCESS_KEY,
    )

def put_file_into_bucket_with_key(path, bucketname, key):
    conn = create_storage_connection()
    buckets = conn.get_all_buckets()
    bucket = None
    for b in buckets:
        if b.name == bucketname:
            bucket = b
    if not bucket:
        raise Exception('Bucket with name ' + bucketname + ' not found')
    k = Key(bucket)
    k.key = key
    k.set_contents_from_filename(path)
    print path, 'has been saved in', bucket.name, 'with key:', key

def get_instance(instance_id):
    conn = create_aws_connection()
    reservations = conn.get_all_instances([instance_id])[0]
    return reservations.instances[0]

def start_server_from_ami(server):
    """Starts a VM on AWS using a given AMI and name"""
    if not server.conn:
        server.conn = create_aws_connection()
    reservation = server.conn.run_instances(
        server.ami_id, security_groups=server.config.SECURITY_GROUPS,
        key_name=server.config.KEY_NAME,
        instance_type=server.config.DEFAULT_INSTANCE_TYPE,
        placement=server.config.DEFAULT_ZONE,
    )

    instance = reservation.instances[0]

    print 'Waiting for instance to come up...'
    # Give the instance a chance to exist
    time.sleep(1.5)
    while instance.state != 'running':
        print '.',
        sys.stdout.flush()
        time.sleep(2)
        instance.update()
    print
    server.conn.create_tags([instance.id], {'Name': server.identifier})
    return instance

<commit_msg>Makes the apt-get update/upgrade part fault tolerant<commit_after># copyright: (c) 2012 by Hansel Dunlop.
# license: ISC, see LICENSE for more details.
#
# kubrick.api
#
# This module implements the Kubrick api.
#

import os

from boto.ec2.connection import EC2Connection
from boto.ec2 import get_region
from boto.s3.connection import S3Connection
from boto.s3.key import Key
from fabric.api import env, put, run
from fabric.contrib.files import append
from fabric.state import connections
import sys
import time

import config

class Server(object):
    install_command_map = {
        'apt': 'sudo apt-get install',
        'pip': 'pip install --upgrade',
        'easy_install': 'easy_install --upgrade',
    }

class AWSServer(Server):

    def __init__(self, config, identifier):
        self.key_id = config.AWS_ACCESS_KEY_ID
        self.secret_key = config.AWS_SECRET_ACCESS_KEY
        self.ami_id = config.BASE_AMI
        self.config = config
        self.key_filename = config.KEY_FILENAME
        self.identifier = identifier
        self.conn = EC2Connection(self.key_id, self.secret_key)
        self.instance = None
        self.root_account = config.USERNAME
        self.operating_system = config.OPERATING_SYSTEM
        if self.operating_system == 'Ubuntu':
            self.sudo_required = True
        else:
            self.sudo_required = False

    @property
    def host_string(self):
        return self.root_account + '@' + self.instance.public_dns_name

    def run(self, command, warn_only=False):
        if not self.instance:
            raise Exception('No machine attached to this server instance')
        env.host_string = self.host_string
        env.disable_known_hosts = True
        env.user = self.config.USERNAME
        env.connection_attempts = 6
        env.warn_only = warn_only
        env.key_filename = self.key_filename
        if self.sudo_required:
            command = 'sudo ' + command
        run(command)

    def append(self, filename, text):
        if not self.instance:
            raise Exception('No machine attached to this server instance')
        env.host_string = self.host_string
        env.disable_known_hosts = True
        env.user = self.config.USERNAME
        env.connection_attempts = 6
        env.key_filename = self.key_filename
        if self.sudo_required:
            append(filename, text, use_sudo=True)
        else:
            append(filename, text)

    def put(self, src, dest):
        env.host_string = self.host_string
        env.disable_known_hosts = True
        env.user = self.config.USERNAME
        env.connection_attempts = 6
        env.key_filename = self.key_filename
        if self.sudo_required:
            put(src, dest, use_sudo=True, mirror_local_mode=True)
            return
        put(src, dest, mirror_local_mode=True)

    def start_server(self):
        self.instance = start_server_from_ami(self)

    def destroy(self):
        if self.instance:
            print 'Destroying AWS server', self.instance
            self.instance.terminate()
        else:
            print 'No instance on this server object'

    def reboot(self):
        self.run('reboot')
        time.sleep(30)
        print "Reconnecting",
        sys.stdout.flush()
        for retry in range(60):
            try:
                print ".",
                sys.stdout.flush()
                env.host_string = self.host_string
                connections.connect(env.host_string)
                break
            except:
                print "-",
                sys.stdout.flush()
                time.sleep(3)
        print

class UbuntuMixin(object):
    '''A mixin for Ubuntu based servers'''

    apt_updated = False
    apt_command = 'apt-get -qq -y'

    def make_swap(self, size):
        swap_location = '/var/swap'
        self.run('dd if=/dev/zero of=%s bs=1M count=%d' % (swap_location, size,))
        self.run('mkswap ' + swap_location)
        self.run('swapon ' + swap_location)
        self.append(
            '/etc/fstab', '%s swap swap defaults 0 0' % (swap_location,)
        )

    def purge_apt_packages(self):
        self.update_apt_if_necessary()
        for package in self.APT_PURGES:
            self.run(
                self.apt_command + ' --purge remove ' + package, warn_only=True
            )

    def update_apt_if_necessary(self):
        if not self.apt_updated:
            self.run(self.apt_command + ' update', warn_only=True)
            self.apt_updated = True

    def upgrade_installed_packages(self):
        self.update_apt_if_necessary()
        self.run(self.apt_command + ' upgrade', warn_only=True)
        self.apt_updated = False

    def install_apt_packages(self):
        self.update_apt_if_necessary()
        self.run(self.apt_command + ' install ' + ' '.join(self.APT_INSTALLS))

    def install_python_modules(self):
        self.run('easy_install -U distribute')
        for package in self.PIP_INSTALLS:
            self.run('pip install %s' % (package,))

    def configure_locales(self):
        # prevents certain warning messages
        self.run('locale-gen en_US en_US.UTF-8 en_GB.UTF-8')
        self.run('dpkg-reconfigure locales')
        self.run('mkdir -p /etc/ssl/intermediate_and_root_certs/')

    def add_non_root_user(self, user, uid, gid, groups, shell='/bin/bash', base_dir='/home'):
        groups_arg = ','.join(groups)

        self.run('groupadd --gid %d %s' % (gid, user))
        self.run(
            'useradd -m -b %s -s %s --uid %d --gid %d -G %s %s' % (
                base_dir, shell, uid, gid, groups_arg, user
            )
        )

    def put_static_config_files(self):
        # recursively put all the files in the server_config dir
        config_files = []
        for dirname, dirs, files in os.walk(self.CONFIG_LOCATION):
            for path in files:
                full_path = os.path.join(dirname, path)
                config_files.append(
                    (full_path, full_path.replace(self.CONFIG_LOCATION, ''))
                )

        for src, dest in config_files:
            self.run('mkdir -p ' + os.path.dirname(dest))
            self.put(src, dest)

    def setup_cron(self):
        self.run('chmod 600 /etc/crontab')
        self.run('chgrp root /etc/crontab')
        self.run('chown root /etc/crontab')

    def set_hostname(self, hostname):
        self.run('hostname %s' % (hostname,))

def get_region_from_name(zone_name=config.DEFAULT_ZONE):
    return get_region(
        zone_name,
        aws_access_key_id=config.AWS_ACCESS_KEY_ID,
        aws_secret_access_key=config.AWS_SECRET_ACCESS_KEY,
    )

def create_aws_connection():
    return EC2Connection(
        config.AWS_ACCESS_KEY_ID,
        config.AWS_SECRET_ACCESS_KEY,
        region=get_region_from_name()
    )

def create_storage_connection():
    return S3Connection(
        config.AWS_ACCESS_KEY_ID,
        config.AWS_SECRET_ACCESS_KEY,
    )

def put_file_into_bucket_with_key(path, bucketname, key):
    conn = create_storage_connection()
    buckets = conn.get_all_buckets()
    bucket = None
    for b in buckets:
        if b.name == bucketname:
            bucket = b
    if not bucket:
        raise Exception('Bucket with name ' + bucketname + ' not found')
    k = Key(bucket)
    k.key = key
    k.set_contents_from_filename(path)
    print path, 'has been saved in', bucket.name, 'with key:', key

def get_instance(instance_id):
    conn = create_aws_connection()
    reservations = conn.get_all_instances([instance_id])[0]
    return reservations.instances[0]

def start_server_from_ami(server):
    """Starts a VM on AWS using a given AMI and name"""
    if not server.conn:
        server.conn = create_aws_connection()
    reservation = server.conn.run_instances(
        server.ami_id, security_groups=server.config.SECURITY_GROUPS,
        key_name=server.config.KEY_NAME,
        instance_type=server.config.DEFAULT_INSTANCE_TYPE,
        placement=server.config.DEFAULT_ZONE,
    )

    instance = reservation.instances[0]

    print 'Waiting for instance to come up...'
    # Give the instance a chance to exist
    time.sleep(1.5)
    while instance.state != 'running':
        print '.',
        sys.stdout.flush()
        time.sleep(2)
        instance.update()
    print
    server.conn.create_tags([instance.id], {'Name': server.identifier})
    return instance

<|endoftext|>


<commit_before># killableprocess - subprocesses which can be reliably killed
#
# Parts of this module are copied from the subprocess.py file contained
# in the Python distribution.
#
# Copyright (c) 2003-2004 by Peter Astrand <astrand@lysator.liu.se>
#
# Additions and modifications written by Benjamin Smedberg
# <benjamin@smedbergs.us> are Copyright (c) 2006 by the Mozilla Foundation
# <http://www.mozilla.org/>
#
# Additions and modifications written by Luke McCarthy <luke@iogopro.co.uk>
# are Copyright (c) 2012 by Luke McCarthy
#
# By obtaining, using, and/or copying this software and/or its
# associated documentation, you agree that you have read, understood,
# and will comply with the following terms and conditions:
#
# Permission to use, copy, modify, and distribute this software and
# its associated documentation for any purpose and without fee is
# hereby granted, provided that the above copyright notice appears in
# all copies, and that both that copyright notice and this permission
# notice appear in supporting documentation, and that the name of the
# author not be used in advertising or publicity pertaining to
# distribution of the software without specific, written prior
# permission.
#
# THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,
# INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.
# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR
# CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
# OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

r"""killableprocess - Subprocesses which can be reliably killed

This module is a subclass of the builtin "subprocess" module. It allows
processes that launch subprocesses to be reliably killed on Windows (via the Popen.kill() method.

It also adds a timeout argument to Wait() for a limited period of time before
forcefully killing the process.

Note: On Windows, this module requires Windows 2000 or higher (no support for
Windows 95, 98, or NT 4.0). It also requires ctypes, which is bundled with
Python 2.5+ or available from http://python.net/crew/theller/ctypes/
"""

import subprocess
import sys
import os
import signal
import time
import types

try:
    from subprocess import CalledProcessError
except ImportError:
    # Python 2.4 doesn't implement CalledProcessError
    class CalledProcessError(Exception):
        """This exception is raised when a process run by check_call() returns
        a non-zero exit status. The exit status will be stored in the
        returncode attribute."""
        def __init__(self, returncode, cmd):
            self.returncode = returncode
            self.cmd = cmd
        def __str__(self):
            return "Command '%s' returned non-zero exit status %d" % (self.cmd, self.returncode)

mswindows = (sys.platform == "win32")

if mswindows:
    import winprocess

def call(*args, **kwargs):
    waitargs = {}
    if "timeout" in kwargs:
        waitargs["timeout"] = kwargs.pop("timeout")

    return Popen(*args, **kwargs).wait(**waitargs)

def check_call(*args, **kwargs):
    """Call a program with an optional timeout. If the program has a non-zero
    exit status, raises a CalledProcessError."""

    retcode = call(*args, **kwargs)
    if retcode:
        cmd = kwargs.get("args")
        if cmd is None:
            cmd = args[0]
        raise CalledProcessError(retcode, cmd)

if not mswindows:
    def DoNothing(*args):
        pass

class Popen(subprocess.Popen):
    if not mswindows:
        # Override __init__ to set a preexec_fn
        def __init__(self, *args, **kwargs):
            if len(args) >= 7:
                raise Exception("Arguments preexec_fn and after must be passed by keyword.")

            real_preexec_fn = kwargs.pop("preexec_fn", None)
            def setpgid_preexec_fn():
                os.setpgid(0, 0)
                if real_preexec_fn:
                    apply(real_preexec_fn)

            kwargs['preexec_fn'] = setpgid_preexec_fn

            subprocess.Popen.__init__(self, *args, **kwargs)

    if mswindows:
        def _execute_child(self, *args_tuple):
            # workaround for bug 958609
            if sys.hexversion < 0x02070600: # prior to 2.7.6
                (args, executable, preexec_fn, close_fds,
                    cwd, env, universal_newlines, startupinfo,
                    creationflags, shell,
                    p2cread, p2cwrite,
                    c2pread, c2pwrite,
                    errread, errwrite) = args_tuple
                to_close = set()
            else: # 2.7.6 and later
                (args, executable, preexec_fn, close_fds,
                    cwd, env, universal_newlines, startupinfo,
                    creationflags, shell, to_close,
                    p2cread, p2cwrite,
                    c2pread, c2pwrite,
                    errread, errwrite) = args_tuple

            if not isinstance(args, types.StringTypes):
                args = subprocess.list2cmdline(args)

            if startupinfo is None:
                startupinfo = winprocess.STARTUPINFO()

            if None not in (p2cread, c2pwrite, errwrite):
                startupinfo.dwFlags |= winprocess.STARTF_USESTDHANDLES

                startupinfo.hStdInput = int(p2cread)
                startupinfo.hStdOutput = int(c2pwrite)
                startupinfo.hStdError = int(errwrite)
            if shell:
                startupinfo.dwFlags |= winprocess.STARTF_USESHOWWINDOW
                startupinfo.wShowWindow = winprocess.SW_HIDE
                comspec = os.environ.get("COMSPEC", "cmd.exe")
                args = comspec + " /c " + args

            # We create a new job for this process, so that we can kill
            # the process and any sub-processes
            self._job = winprocess.CreateJobObject()

            creationflags |= winprocess.CREATE_SUSPENDED
            creationflags |= winprocess.CREATE_UNICODE_ENVIRONMENT

            hp, ht, pid, tid = winprocess.CreateProcess(
                executable, args,
                None, None, # No special security
                1, # Must inherit handles!
                creationflags,
                winprocess.EnvironmentBlock(env),
                cwd, startupinfo)

            self._child_created = True
            self._handle = hp
            self._thread = ht
            self.pid = pid

            winprocess.AssignProcessToJobObject(self._job, hp)
            winprocess.ResumeThread(ht)

            if p2cread is not None:
                p2cread.Close()
            if c2pwrite is not None:
                c2pwrite.Close()
            if errwrite is not None:
                errwrite.Close()

    def send_signal(self, sig, group=True):
        if mswindows:
            if group:
                winprocess.TerminateJobObject(self._job, sig)
            else:
                winprocess.TerminateProcess(self._handle, sig)
        else:
            if group:
                os.killpg(self.pid, sig)
            else:
                os.kill(self.pid, sig)

    def kill(self, group=True):
        """Kill the process. If group=True, all sub-processes will also be killed."""
        sig = signal.CTRL_BREAK_EVENT if mswindows else signal.SIGKILL
        self.send_signal(sig, group)

    def terminate(self, group=True):
        """Terminate the process. If group=True, all sub-processes will also be killed."""
        sig = signal.CTRL_C_EVENT if mswindows else signal.SIGTERM
        self.send_signal(sig, group)

    def wait(self, timeout=-1, group=True):
        """Wait for the process to terminate. Returns returncode attribute.
        If timeout seconds are reached and the process has not terminated,
        it will be forcefully killed. If timeout is -1, wait will not
        time out."""

        if self.returncode is not None:
            return self.returncode

        if mswindows:
            if timeout != -1:
                timeout = timeout * 1000
            rc = winprocess.WaitForSingleObject(self._handle, timeout)
            if rc == winprocess.WAIT_TIMEOUT:
                self.kill(group)
            else:
                self.returncode = winprocess.GetExitCodeProcess(self._handle)
        else:
            if timeout == -1:
                subprocess.Popen.wait(self)
                return self.returncode

            starttime = time.time()

            # Make sure there is a signal handler for SIGCHLD installed
            oldsignal = signal.signal(signal.SIGCHLD, DoNothing)

            while time.time() < starttime + timeout - 0.01:
                pid, sts = os.waitpid(self.pid, os.WNOHANG)
                if pid != 0:
                    self._handle_exitstatus(sts)
                    signal.signal(signal.SIGCHLD, oldsignal)
                    return self.returncode

                # time.sleep is interrupted by signals (good!)
                newtimeout = timeout - time.time() + starttime
                time.sleep(newtimeout)

            self.kill(group)
            signal.signal(signal.SIGCHLD, oldsignal)
            subprocess.Popen.wait(self)

        return self.returncode
<commit_msg>Apparently these flags are needed now or else we get an Access Denied error.<commit_after># killableprocess - subprocesses which can be reliably killed
#
# Parts of this module are copied from the subprocess.py file contained
# in the Python distribution.
#
# Copyright (c) 2003-2004 by Peter Astrand <astrand@lysator.liu.se>
#
# Additions and modifications written by Benjamin Smedberg
# <benjamin@smedbergs.us> are Copyright (c) 2006 by the Mozilla Foundation
# <http://www.mozilla.org/>
#
# Additions and modifications written by Luke McCarthy <luke@iogopro.co.uk>
# are Copyright (c) 2012 by Luke McCarthy
#
# By obtaining, using, and/or copying this software and/or its
# associated documentation, you agree that you have read, understood,
# and will comply with the following terms and conditions:
#
# Permission to use, copy, modify, and distribute this software and
# its associated documentation for any purpose and without fee is
# hereby granted, provided that the above copyright notice appears in
# all copies, and that both that copyright notice and this permission
# notice appear in supporting documentation, and that the name of the
# author not be used in advertising or publicity pertaining to
# distribution of the software without specific, written prior
# permission.
#
# THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,
# INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.
# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR
# CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
# OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

r"""killableprocess - Subprocesses which can be reliably killed

This module is a subclass of the builtin "subprocess" module. It allows
processes that launch subprocesses to be reliably killed on Windows (via the Popen.kill() method.

It also adds a timeout argument to Wait() for a limited period of time before
forcefully killing the process.

Note: On Windows, this module requires Windows 2000 or higher (no support for
Windows 95, 98, or NT 4.0). It also requires ctypes, which is bundled with
Python 2.5+ or available from http://python.net/crew/theller/ctypes/
"""

import subprocess
import sys
import os
import signal
import time
import types

try:
    from subprocess import CalledProcessError
except ImportError:
    # Python 2.4 doesn't implement CalledProcessError
    class CalledProcessError(Exception):
        """This exception is raised when a process run by check_call() returns
        a non-zero exit status. The exit status will be stored in the
        returncode attribute."""
        def __init__(self, returncode, cmd):
            self.returncode = returncode
            self.cmd = cmd
        def __str__(self):
            return "Command '%s' returned non-zero exit status %d" % (self.cmd, self.returncode)

mswindows = (sys.platform == "win32")

if mswindows:
    import winprocess

def call(*args, **kwargs):
    waitargs = {}
    if "timeout" in kwargs:
        waitargs["timeout"] = kwargs.pop("timeout")

    return Popen(*args, **kwargs).wait(**waitargs)

def check_call(*args, **kwargs):
    """Call a program with an optional timeout. If the program has a non-zero
    exit status, raises a CalledProcessError."""

    retcode = call(*args, **kwargs)
    if retcode:
        cmd = kwargs.get("args")
        if cmd is None:
            cmd = args[0]
        raise CalledProcessError(retcode, cmd)

if not mswindows:
    def DoNothing(*args):
        pass

class Popen(subprocess.Popen):
    if not mswindows:
        # Override __init__ to set a preexec_fn
        def __init__(self, *args, **kwargs):
            if len(args) >= 7:
                raise Exception("Arguments preexec_fn and after must be passed by keyword.")

            real_preexec_fn = kwargs.pop("preexec_fn", None)
            def setpgid_preexec_fn():
                os.setpgid(0, 0)
                if real_preexec_fn:
                    apply(real_preexec_fn)

            kwargs['preexec_fn'] = setpgid_preexec_fn

            subprocess.Popen.__init__(self, *args, **kwargs)

    if mswindows:
        def _execute_child(self, *args_tuple):
            # workaround for bug 958609
            if sys.hexversion < 0x02070600: # prior to 2.7.6
                (args, executable, preexec_fn, close_fds,
                    cwd, env, universal_newlines, startupinfo,
                    creationflags, shell,
                    p2cread, p2cwrite,
                    c2pread, c2pwrite,
                    errread, errwrite) = args_tuple
                to_close = set()
            else: # 2.7.6 and later
                (args, executable, preexec_fn, close_fds,
                    cwd, env, universal_newlines, startupinfo,
                    creationflags, shell, to_close,
                    p2cread, p2cwrite,
                    c2pread, c2pwrite,
                    errread, errwrite) = args_tuple

            if not isinstance(args, types.StringTypes):
                args = subprocess.list2cmdline(args)

            creationflags = winprocess.CREATE_NEW_PROCESS_GROUP

            if startupinfo is None:
                startupinfo = winprocess.STARTUPINFO()

            if None not in (p2cread, c2pwrite, errwrite):
                startupinfo.dwFlags |= winprocess.STARTF_USESTDHANDLES

                startupinfo.hStdInput = int(p2cread)
                startupinfo.hStdOutput = int(c2pwrite)
                startupinfo.hStdError = int(errwrite)
            if shell:
                startupinfo.dwFlags |= winprocess.STARTF_USESHOWWINDOW
                startupinfo.wShowWindow = winprocess.SW_HIDE
                comspec = os.environ.get("COMSPEC", "cmd.exe")
                args = comspec + " /c " + args

            # We create a new job for this process, so that we can kill
            # the process and any sub-processes
            self._job = winprocess.CreateJobObject()

            creationflags |= winprocess.CREATE_SUSPENDED
            creationflags |= winprocess.CREATE_UNICODE_ENVIRONMENT
            creationflags |= winprocess.CREATE_BREAKAWAY_FROM_JOB

            hp, ht, pid, tid = winprocess.CreateProcess(
                executable, args,
                None, None, # No special security
                1, # Must inherit handles!
                creationflags,
                winprocess.EnvironmentBlock(env),
                cwd, startupinfo)

            self._child_created = True
            self._handle = hp
            self._thread = ht
            self.pid = pid

            winprocess.AssignProcessToJobObject(self._job, hp)
            winprocess.ResumeThread(ht)

            if p2cread is not None:
                p2cread.Close()
            if c2pwrite is not None:
                c2pwrite.Close()
            if errwrite is not None:
                errwrite.Close()

    def send_signal(self, sig, group=True):
        if mswindows:
            if group:
                winprocess.TerminateJobObject(self._job, sig)
            else:
                winprocess.TerminateProcess(self._handle, sig)
        else:
            if group:
                os.killpg(self.pid, sig)
            else:
                os.kill(self.pid, sig)

    def kill(self, group=True):
        """Kill the process. If group=True, all sub-processes will also be killed."""
        sig = signal.CTRL_BREAK_EVENT if mswindows else signal.SIGKILL
        self.send_signal(sig, group)

    def terminate(self, group=True):
        """Terminate the process. If group=True, all sub-processes will also be killed."""
        sig = signal.CTRL_C_EVENT if mswindows else signal.SIGTERM
        self.send_signal(sig, group)

    def wait(self, timeout=-1, group=True):
        """Wait for the process to terminate. Returns returncode attribute.
        If timeout seconds are reached and the process has not terminated,
        it will be forcefully killed. If timeout is -1, wait will not
        time out."""

        if self.returncode is not None:
            return self.returncode

        if mswindows:
            if timeout != -1:
                timeout = timeout * 1000
            rc = winprocess.WaitForSingleObject(self._handle, timeout)
            if rc == winprocess.WAIT_TIMEOUT:
                self.kill(group)
            else:
                self.returncode = winprocess.GetExitCodeProcess(self._handle)
        else:
            if timeout == -1:
                subprocess.Popen.wait(self)
                return self.returncode

            starttime = time.time()

            # Make sure there is a signal handler for SIGCHLD installed
            oldsignal = signal.signal(signal.SIGCHLD, DoNothing)

            while time.time() < starttime + timeout - 0.01:
                pid, sts = os.waitpid(self.pid, os.WNOHANG)
                if pid != 0:
                    self._handle_exitstatus(sts)
                    signal.signal(signal.SIGCHLD, oldsignal)
                    return self.returncode

                # time.sleep is interrupted by signals (good!)
                newtimeout = timeout - time.time() + starttime
                time.sleep(newtimeout)

            self.kill(group)
            signal.signal(signal.SIGCHLD, oldsignal)
            subprocess.Popen.wait(self)

        return self.returncode
<|endoftext|>


<commit_before>
__version__ = "0.10.6"
<commit_msg>Version bump.<commit_after>
__version__ = "0.11"
<|endoftext|>


<commit_before>import logging, os, redis, string, subprocess
from apscheduler.scheduler import Scheduler
from datetime import datetime, timedelta
from multiprocessing import Pool
from plumbum import SshMachine
from rpyc.utils.factory import ssh_connect
from utilities import distance, threaded_map

process_pool_size = 60
sched = Scheduler(standalone=True)
logging.basicConfig()

all_dns = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=0))
open_resolvers = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=1))
geoip = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=2))
pl_hosts = [line.split(' ')[0:4] for line in map(string.strip,open('pl-host-list-geo').readlines())]

from sqlalchemy import create_engine
from sqlalchemy import Column, Integer, PickleType
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base

engine = create_engine('sqlite:///latency.db', echo=True)
Base = declarative_base(bind=engine)

class DataPoint(Base):
    __tablename__ = 'data'

    id = Column(Integer, primary_key=True)
    target1 = Column(PickleType)
    target2 = Column(PickleType)
    start = Column(PickleType)
    end = Column(PickleType)
    pings = Column(PickleType)
    address = Column(PickleType)

    def __init__(self, target1, target2, start, end, pings, address):
        self.target1 = target1
        self.target2 = target2
        self.start = start
        self.end = end
        self.pings = pings
        self.address = address

Base.metadata.create_all()
Session = sessionmaker(bind=engine)
s = Session()

class PlanetLabNode(object):
    def __init__(self, host, ip, lat, lon):
        self.host = host
        self.ip = ip
        self.lat = float(lat)
        self.lon = float(lon)
        self.connected = False

    def connect(self):
        try:
            self.connectPL()
        except Exception, e:
            try:
                self.restartPL()
                self.connectPL()
            except Exception, e:
                self.connected = False

    def connectPL(self):
        rem = SshMachine(self.host, user='ucb_268_measure', keyfile='~/.ssh/id_rsa',
                         ssh_opts=["-o StrictHostKeyChecking no",
                                   "-o UserKnownHostsFile=/dev/null"])
        conn = ssh_connect(rem, 18861, config={'allow_pickle' : True})
        assert conn.root.test() == 1, "%s: RPC Failure" % self.host
        self.conn = conn
        self.connected = True

    def restartPL(self):
        FNULL = open(os.devnull, 'w')
        subprocess.call(["ssh", "-t", "-i", "~/.ssh/id_rsa", "-o StrictHostKeyChecking no",
                         "-o UserKnownHostsFile=/dev/null", "ucb_268_measure@%s" % self.host,
                         "sudo tking-server stop; sudo tking-server start;"],
                         stdout=FNULL, stderr=FNULL)

    def handleConnExceptions(fn):
        def wrapped(self, *args, **kwargs):
            if not self.connected:
                self.connect()

            if self.connected:
                try:
                    return fn(self, *args, **kwargs)
                except Exception, e:
                    self.connected = False
                    try:
                        self.connect()
                        return fn(self, *args, **kwargs)
                    except Exception, e:
                        self.connected = False
                        return None
            else:
                return None
        return wrapped

    @handleConnExceptions
    def get_latency(self, target1, ip1, target2, ip2):
        return self.conn.root.get_latency(target1, ip1, target2, ip2)

    @handleConnExceptions
    def get_k(self, target, ip):
        return self.conn.get_k(target, ip)

    @handleConnExceptions
    def get_distance(self, lat, lon):
        return distance((self.lat, self.lon), (lat, lon))

def select_random_points():
    target1 = open_resolvers.randomkey()
    target2 = geoip.randomkey()
    while not geoip.exists(target1):
        target1 = open_resolvers.randomkey()

    ip1, coord1 = list(all_dns.smembers(target1))[0], eval(list(geoip.smembers(target1))[0])[1:]
    ip2, coord2 = list(all_dns.smembers(target2))[0], eval(list(geoip.smembers(target2))[0])[1:]

    return (target1, ip1, coord1), (target2, ip2, coord2)

def query_latency(target1, target2):
    name1, ip1, coord1 = target1
    name2, ip2, coord2 = target2

    # Get closest 4 PL Nodes
    distances = map(lambda node: (distance(coord1, (node.lat, node.lon)), node), pl_nodes)
    distances.sort()
    distances = distances[:4]

    results = threaded_map(lambda (dist, node): (node.host, node.get_latency(name1, ip1, name2, ip2)), distances, timeout=10.0)
    return results

def one_round(x):
    target1, target2 = select_random_points()
    results = query_latency(target1, target2)
    return target1, target2, results

pl_nodes = map(lambda args: PlanetLabNode(*args), pl_hosts)

def task():
    print 'Task'
    p = Pool(process_pool_size)
    results = p.map(one_round, range(process_pool_size))
    for target1, target2, result_set in filter(None,results):
        for host, info in filter(None,result_set):
            if info:
                (end_time, start_time, ping_times, address) = info
                point = DataPoint(target1, target2, start_time, end_time, ping_times, address)
                s.add(point)
    p.close()
    s.commit()

sched.add_interval_job(task, minutes=1)
sched.start()
<commit_msg>hopefully got last error<commit_after>import logging, os, redis, string, subprocess
from apscheduler.scheduler import Scheduler
from datetime import datetime, timedelta
from multiprocessing import Pool
from plumbum import SshMachine
from rpyc.utils.factory import ssh_connect
from utilities import distance, threaded_map

process_pool_size = 60
sched = Scheduler(standalone=True)
logging.basicConfig()

all_dns = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=0))
open_resolvers = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=1))
geoip = redis.Redis(connection_pool=redis.ConnectionPool(host='localhost', port=6379, db=2))
pl_hosts = [line.split(' ')[0:4] for line in map(string.strip,open('pl-host-list-geo').readlines())]

from sqlalchemy import create_engine
from sqlalchemy import Column, Integer, PickleType
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base

engine = create_engine('sqlite:///latency.db', echo=True)
Base = declarative_base(bind=engine)

class DataPoint(Base):
    __tablename__ = 'data'

    id = Column(Integer, primary_key=True)
    target1 = Column(PickleType)
    target2 = Column(PickleType)
    start = Column(PickleType)
    end = Column(PickleType)
    pings = Column(PickleType)
    address = Column(PickleType)

    def __init__(self, target1, target2, start, end, pings, address):
        self.target1 = target1
        self.target2 = target2
        self.start = start
        self.end = end
        self.pings = pings
        self.address = address

Base.metadata.create_all()
Session = sessionmaker(bind=engine)
s = Session()

class PlanetLabNode(object):
    def __init__(self, host, ip, lat, lon):
        self.host = host
        self.ip = ip
        self.lat = float(lat)
        self.lon = float(lon)
        self.connected = False

    def connect(self):
        try:
            self.connectPL()
        except Exception, e:
            try:
                self.restartPL()
                self.connectPL()
            except Exception, e:
                self.connected = False

    def connectPL(self):
        rem = SshMachine(self.host, user='ucb_268_measure', keyfile='~/.ssh/id_rsa',
                         ssh_opts=["-o StrictHostKeyChecking no",
                                   "-o UserKnownHostsFile=/dev/null"])
        conn = ssh_connect(rem, 18861, config={'allow_pickle' : True})
        assert conn.root.test() == 1, "%s: RPC Failure" % self.host
        self.conn = conn
        self.connected = True

    def restartPL(self):
        FNULL = open(os.devnull, 'w')
        subprocess.call(["ssh", "-t", "-i", "~/.ssh/id_rsa", "-o StrictHostKeyChecking no",
                         "-o UserKnownHostsFile=/dev/null", "ucb_268_measure@%s" % self.host,
                         "sudo tking-server stop; sudo tking-server start;"],
                         stdout=FNULL, stderr=FNULL)

    def handleConnExceptions(fn):
        def wrapped(self, *args, **kwargs):
            if not self.connected:
                self.connect()

            if self.connected:
                try:
                    return fn(self, *args, **kwargs)
                except Exception, e:
                    self.connected = False
                    try:
                        self.connect()
                        return fn(self, *args, **kwargs)
                    except Exception, e:
                        self.connected = False
                        return None
            else:
                return None
        return wrapped

    @handleConnExceptions
    def get_latency(self, target1, ip1, target2, ip2):
        return self.conn.root.get_latency(target1, ip1, target2, ip2)

    @handleConnExceptions
    def get_k(self, target, ip):
        return self.conn.get_k(target, ip)

    @handleConnExceptions
    def get_distance(self, lat, lon):
        return distance((self.lat, self.lon), (lat, lon))

def select_random_points():
    target1 = open_resolvers.randomkey()
    target2 = geoip.randomkey()
    while not geoip.exists(target1):
        target1 = open_resolvers.randomkey()

    ip1, coord1 = list(all_dns.smembers(target1))[0], eval(list(geoip.smembers(target1))[0])[1:]
    ip2, coord2 = list(all_dns.smembers(target2))[0], eval(list(geoip.smembers(target2))[0])[1:]

    return (target1, ip1, coord1), (target2, ip2, coord2)

def query_latency(target1, target2):
    name1, ip1, coord1 = target1
    name2, ip2, coord2 = target2

    # Get closest 4 PL Nodes
    distances = map(lambda node: (distance(coord1, (node.lat, node.lon)), node), pl_nodes)
    distances.sort()
    distances = distances[:4]

    results = threaded_map(lambda (dist, node): (node.host, node.get_latency(name1, ip1, name2, ip2)), distances, timeout=10.0)
    return results

def one_round(x):
    target1, target2 = select_random_points()
    results = query_latency(target1, target2)
    return target1, target2, results

pl_nodes = map(lambda args: PlanetLabNode(*args), pl_hosts)

def task():
    print 'Task'
    p = Pool(process_pool_size)
    results = p.map(one_round, range(process_pool_size))
    for target1, target2, result_set in filter(None,results):
        for host, info in filter(None,result_set):
            if info:
                try:
                    (end_time, start_time, ping_times, address) = info
                    point = DataPoint(target1, target2, start_time, end_time, ping_times, address)
                    s.add(point)
                except Exception, e:
                    pass
    p.close()
    s.commit()

sched.add_interval_job(task, minutes=1)
sched.start()
<|endoftext|>


<commit_before>
__version__ = "1.0.5"
<commit_msg>Version bump to 1.0.6<commit_after>
__version__ = "1.0.6"
<|endoftext|>


<commit_before>#-*- coding=utf-8 -*-

from collections import defaultdict

from django.db import models
#from django.contrib.auth.models import User

from exceptions import OverlappingReservations, CapacityExceeded

class BaseSubject(models.Model):
    class Meta:
        abstract = True

    def reserve(self, **kwargs):
        Reservation = self._meta.get_field_by_name('reservations')[0].model
        return Reservation.objects.create(subject=self, **kwargs)

class GroupableSubject(BaseSubject):
    class Meta:
        abstract = True

    cluster = models.ForeignKey('Cluster', null=True)

class ExclusiveSubject(BaseSubject):
    class Meta:
        abstract = True

    def reserve(self, start, end, **kwargs):
        overlapping_reservations = self.reservations.filter(start__lt=end,
                                                            end__gt=start)
        if overlapping_reservations:
            raise OverlappingReservations(overlapping_reservations)
        else:
            return super(ExclusiveSubject, self).reserve(start=start, end=end,
                                                         **kwargs)

class FiniteCapacitySubject(BaseSubject):
    '''
    This mixin requires capacity property. Available e.g. in
    VariableCapacitySubject and FixedCapacitySubject
    '''
    class Meta:
        abstract = True

    def reserve(self, start, end, size, **kwargs):
        overlapping_reservations = self.reservations.filter(start__lt=end,
                                                            end__gt=start)
        if size > self.capacity:
            raise CapacityExceeded

        dates = defaultdict(lambda: 0)
        for r in overlapping_reservations:
            # mark when usage of subject changes
            dates[r.start] += getattr(r, 'size', 1)
            dates[r.end] -= getattr(r, 'size', 1)
        balance = 0
        for date, delta in sorted(dates.iteritems()):
            balance += delta
            if balance + size > self.capacity:
                raise CapacityExceeded

        super(FiniteCapacitySubject, self).reserve(start=start, end=end,
                size=size, **kwargs)

class VariableCapacitySubject(FiniteCapacitySubject):
    class Meta:
        abstract = True

    capacity = models.PositiveIntegerField()

class FixedCapacitySubject(FiniteCapacitySubject):
    '''
    You can inherit from this class like this:
    class YourClass(FixedCapacitySubject.with_capacity(5)
        pass
    '''
    class Meta:
        abstract = True

    @classmethod
    def with_capacity(cls, capacity):
        cls._capacity = capacity
        return cls

    @property
    def capacity(self):
        return self._capacity
<commit_msg>no log<commit_after>#-*- coding=utf-8 -*-

from collections import defaultdict

from django.db import models

from kitabu.exceptions import OverlappingReservations, CapacityExceeded

class BaseSubject(models.Model):
    class Meta:
        abstract = True

    def reserve(self, **kwargs):
        Reservation = self._meta.get_field_by_name('reservations')[0].model
        return Reservation.objects.create(subject=self, **kwargs)

class ExclusiveSubject(BaseSubject):
    class Meta:
        abstract = True

    def reserve(self, start, end, **kwargs):
        overlapping_reservations = self.reservations.filter(start__lt=end,
                                                            end__gt=start)
        if overlapping_reservations:
            raise OverlappingReservations(overlapping_reservations)
        else:
            return super(ExclusiveSubject, self).reserve(start=start, end=end,
                                                         **kwargs)

class FiniteCapacitySubject(BaseSubject):
    '''
    This mixin requires capacity property. Available e.g. in
    VariableCapacitySubject and FixedCapacitySubject
    '''
    class Meta:
        abstract = True

    def reserve(self, start, end, size, **kwargs):
        overlapping_reservations = self.reservations.filter(start__lt=end,
                                                            end__gt=start)
        if size > self.capacity:
            raise CapacityExceeded

        dates = defaultdict(lambda: 0)
        for r in overlapping_reservations:
            # mark when usage of subject changes
            dates[r.start] += getattr(r, 'size', 1)
            dates[r.end] -= getattr(r, 'size', 1)
        balance = 0
        for date, delta in sorted(dates.iteritems()):
            balance += delta
            if balance + size > self.capacity:
                raise CapacityExceeded

        super(FiniteCapacitySubject, self).reserve(start=start, end=end,
                size=size, **kwargs)

class VariableCapacitySubject(FiniteCapacitySubject):
    class Meta:
        abstract = True

    capacity = models.PositiveIntegerField()

class FixedCapacitySubject(FiniteCapacitySubject):
    '''
    You can inherit from this class like this:
    class YourClass(FixedCapacitySubject.with_capacity(5)
        pass
    '''
    class Meta:
        abstract = True

    @classmethod
    def with_capacity(cls, capacity):
        cls._capacity = capacity
        return cls

    @property
    def capacity(self):
        return self._capacity
<|endoftext|>


<commit_before>"""
  Grab the torrent from bttt99 website.
"""
import re
import urllib.parse
import concurrent.futures
import itertools
import reprlib
from collections import namedtuple

import lxml.html
import requests
import pyperclip

BASE_URL = 'http://www.bttt99.com/'
TAGS_URL = 'http://www.bttt99.com/tag/'
SEARCH_URL = 'http://www.bttt99.com/s'
UA = 'Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/51.0.2704.106 Safari/537.36'

torrents_re = re.compile('(.*?)')

session = requests.Session()
session.headers.update({'user-agent': UA})

Torrent = namedtuple('Torrent', 'link title description rate')

def get_torrents(url):
    _id = url.rstrip('/').rsplit('/', 1)[-1]
    page = session.get(url).content.decode('utf8')
    doc = lxml.html.fromstring(page)
    for a in doc.xpath('//div[@class="download"]/ul/li/a'):
        link = a.attrib['href']
        title = a.text_content()
        description = doc.xpath('//div[@class="summary"]')[0].text_content().strip()
        rate = doc.xpath('//span[@class="rate"]')[0].text_content()
        yield Torrent(link=link, title=title, description=description, rate=rate)

def copy_to_clipboard(text):
    pyperclip.copy(text)

def grab():
    r = session.get(BASE_URL)
    futures = []
    if r.ok:
        html = r.text
        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
            for link in set(re.findall(r'/tv/\d+/', html)):
                url = urllib.parse.urljoin(BASE_URL, link)
                print(url)
                future = executor.submit(get_torrent, url)
                futures.append(future)
            for fu in concurrent.futures.as_completed(futures):
                result = fu.result()
                print(result)

def search(name):
    r = session.get(SEARCH_URL, params={'q': name})
    doc = lxml.html.fromstring(r.text)
    xp = doc.xpath('//div[@class="taginfo"]//a')
    if xp:
        founds = [(a.attrib['href'], a.text_content()) for a in xp]
        while 1:
            for i, (url, title) in itertools.chain(enumerate(founds, 1),  [('q', ('Quit', 'Quit'))]):
                print('{}: {}'.format(i, title))

            choice = input('Please have a choice: ')
            if choice == 'q':
                exit(0)
            try:
                uri, _ = founds[int(choice) - 1]
                url = urllib.parse.urljoin(BASE_URL, uri)
                torrents = list(get_torrents(url))
                if torrents:
                    for i, tor in enumerate(torrents, 1):
                        print('\n{}: {}\n{}'.format(i , tor.link, tor.title))
                    print('简介:', torrents[0].description)
                    print('Rate: ', torrents[0].rate)
                    which = int(input('Which one?'))
                    torrent = torrents[which - 1]
                    copy_to_clipboard(torrent.link)
                    print('Copied to clipboard')
                    return torrent
                else:
                    print("未找到查询的种子")
            except (IndexError, ValueError):
                print('Input invalid')
    else:
        print('No result')

if __name__ == '__main__':
    search('rounders')
<commit_msg>remove test code<commit_after>"""
  Grab the torrent from bttt99 website.
"""
import re
import urllib.parse
import concurrent.futures
import itertools
import reprlib
from collections import namedtuple

import lxml.html
import requests
import pyperclip

BASE_URL = 'http://www.bttt99.com/'
TAGS_URL = 'http://www.bttt99.com/tag/'
SEARCH_URL = 'http://www.bttt99.com/s'
UA = 'Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/51.0.2704.106 Safari/537.36'

torrents_re = re.compile('(.*?)')

session = requests.Session()
session.headers.update({'user-agent': UA})

Torrent = namedtuple('Torrent', 'link title description rate')

def get_torrents(url):
    _id = url.rstrip('/').rsplit('/', 1)[-1]
    page = session.get(url).content.decode('utf8')
    doc = lxml.html.fromstring(page)
    for a in doc.xpath('//div[@class="download"]/ul/li/a'):
        link = a.attrib['href']
        title = a.text_content()
        description = doc.xpath('//div[@class="summary"]')[0].text_content().strip()
        rate = doc.xpath('//span[@class="rate"]')[0].text_content()
        yield Torrent(link=link, title=title, description=description, rate=rate)

def copy_to_clipboard(text):
    pyperclip.copy(text)

def grab():
    r = session.get(BASE_URL)
    futures = []
    if r.ok:
        html = r.text
        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
            for link in set(re.findall(r'/tv/\d+/', html)):
                url = urllib.parse.urljoin(BASE_URL, link)
                print(url)
                future = executor.submit(get_torrent, url)
                futures.append(future)
            for fu in concurrent.futures.as_completed(futures):
                result = fu.result()
                print(result)

def search(name):
    r = session.get(SEARCH_URL, params={'q': name})
    doc = lxml.html.fromstring(r.text)
    xp = doc.xpath('//div[@class="taginfo"]//a')
    if xp:
        founds = [(a.attrib['href'], a.text_content()) for a in xp]
        while 1:
            for i, (url, title) in itertools.chain(enumerate(founds, 1),  [('q', ('Quit', 'Quit'))]):
                print('{}: {}'.format(i, title))

            choice = input('Please have a choice: ')
            if choice == 'q':
                exit(0)
            try:
                uri, _ = founds[int(choice) - 1]
                url = urllib.parse.urljoin(BASE_URL, uri)
                torrents = list(get_torrents(url))
                if torrents:
                    for i, tor in enumerate(torrents, 1):
                        print('\n{}: {}\n{}'.format(i , tor.link, tor.title))
                    print('简介:', torrents[0].description)
                    print('Rate: ', torrents[0].rate)
                    which = int(input('Which one?'))
                    torrent = torrents[which - 1]
                    copy_to_clipboard(torrent.link)
                    print('Copied to clipboard')
                    return torrent
                else:
                    print("未找到查询的种子")
            except (IndexError, ValueError):
                print('Input invalid')
    else:
        print('No result')
<|endoftext|>


<commit_before>from six import binary_type
import logging
from kafka_influxdb.encoder.escape_functions import influxdb_tag_escaper

class Encoder(object):
    """
    An encoder for the Collectd Graphite ASCII format
    See https://collectd.org/wiki/index.php/Graphite

    Sample measurements this encoder can handle:
    [prefix.]host.plugin.measurement[.postfix] value timestamp

    26f2fc918f50.load.load.shortterm 0.05 1436357630
    26f2fc918f50.load.load.midterm 0.05 1436357630
    26f2fc918f50.load.load.longterm 0.05 1436357630

    26f2fc918f50.cpu-0.cpu-user 30364 1436357630

    26f2fc918f50.memory.memory-buffered 743657472 1436357630

    The optional prefix and postifx can be set in the collectd plugin:
    <Plugin write_kafka>
        Property "metadata.broker.list" "kafka:9092"
        <Topic "metrics">
            Format Graphite
            GraphitePrefix "myprefix"
            GraphitePostfix "mypostfix"
        </Topic>
    </Plugin>
    """

    def __init__(self):
        self.escape_tag = influxdb_tag_escaper()

    def encode(self,
               msg,
               delimiter='.',
               prefix='',
               prefix_tag=None,
               postfix='',
               postfix_tag=None,
               ):
        """
        :param msg: Payload from reader
        :param delimiter: Delimiter between Graphite series parts
        :param prefix: Graphite prefix string
        :param prefix_tag: Tag to use for Graphite prefix
        :param postfix: Graphite postfix string
        :param postfix_tag: Tag to use for Graphite postfix
        :return: A list of encoded messages
        """
        # One message could consist of several measurements
        measurements = []

        for line in msg.decode().split("\n"):
            try:
                series, value, timestamp = line.split()
            except ValueError as e:
                logging.debug("Error in encoder: %s", e)
                continue
            # Strip prefix and postfix:
            series = series[len(prefix):len(series) - len(postfix)]
            # Split into tags
            hostname, measurement = series.split(delimiter, 1)
            measurement = measurement.replace(delimiter, '_')

            tags = {
                "host": hostname
            }
            if prefix_tag:
                if prefix.endswith(delimiter):
                    prefix = prefix[:-len(delimiter)]
                tags[prefix_tag] = prefix
            if postfix_tag:
                if postfix.endswith(delimiter):
                    postfix = postfix[:-len(delimiter)]
                tags[postfix_tag] = postfix

            encoded = ''.join([
                str(measurement),
                ',',
                ','.join('{}={}'.format(self.escape_tag(k), self.escape_tag(tags[k])) for k in tags),
                ' value=',
                str(value),
                ' ',
                timestamp
            ])
            measurements.append(encoded)
        return measurements
<commit_msg>Remove unused dependency to six<commit_after>import logging
from kafka_influxdb.encoder.escape_functions import influxdb_tag_escaper

class Encoder(object):
    """
    An encoder for the Collectd Graphite ASCII format
    See https://collectd.org/wiki/index.php/Graphite

    Sample measurements this encoder can handle:
    [prefix.]host.plugin.measurement[.postfix] value timestamp

    26f2fc918f50.load.load.shortterm 0.05 1436357630
    26f2fc918f50.load.load.midterm 0.05 1436357630
    26f2fc918f50.load.load.longterm 0.05 1436357630

    26f2fc918f50.cpu-0.cpu-user 30364 1436357630

    26f2fc918f50.memory.memory-buffered 743657472 1436357630

    The optional prefix and postifx can be set in the collectd plugin:
    <Plugin write_kafka>
        Property "metadata.broker.list" "kafka:9092"
        <Topic "metrics">
            Format Graphite
            GraphitePrefix "myprefix"
            GraphitePostfix "mypostfix"
        </Topic>
    </Plugin>
    """

    def __init__(self):
        self.escape_tag = influxdb_tag_escaper()

    def encode(self,
               msg,
               delimiter='.',
               prefix='',
               prefix_tag=None,
               postfix='',
               postfix_tag=None,
               ):
        """
        :param msg: Payload from reader
        :param delimiter: Delimiter between Graphite series parts
        :param prefix: Graphite prefix string
        :param prefix_tag: Tag to use for Graphite prefix
        :param postfix: Graphite postfix string
        :param postfix_tag: Tag to use for Graphite postfix
        :return: A list of encoded messages
        """
        # One message could consist of several measurements
        measurements = []

        for line in msg.decode().split("\n"):
            try:
                series, value, timestamp = line.split()
            except ValueError as e:
                logging.debug("Error in encoder: %s", e)
                continue
            # Strip prefix and postfix:
            series = series[len(prefix):len(series) - len(postfix)]
            # Split into tags
            hostname, measurement = series.split(delimiter, 1)
            measurement = measurement.replace(delimiter, '_')

            tags = {
                "host": hostname
            }
            if prefix_tag:
                if prefix.endswith(delimiter):
                    prefix = prefix[:-len(delimiter)]
                tags[prefix_tag] = prefix
            if postfix_tag:
                if postfix.endswith(delimiter):
                    postfix = postfix[:-len(delimiter)]
                tags[postfix_tag] = postfix

            encoded = ''.join([
                str(measurement),
                ',',
                ','.join('{}={}'.format(self.escape_tag(k), self.escape_tag(tags[k])) for k in tags),
                ' value=',
                str(value),
                ' ',
                timestamp
            ])
            measurements.append(encoded)
        return measurements
<|endoftext|>