<commit_before>import os
import time
from io import StringIO
from random import randint
from shlex import quote as shell_quote
from urllib.parse import quote as url_quote
from string import Template
from collections import namedtuple

import paramiko.client as sshclient
import tornado.ioloop
from tornado import gen
from tornado.httpclient import AsyncHTTPClient

from loadsbroker import logger
from loadsbroker.exceptions import LoadsException
from loadsbroker.dockerctrl import DockerDaemon
from loadsbroker.ssh import makedirs
from loadsbroker.util import join_host_port

# Default ping request options.
_PING_DEFAULTS = {
    "method": "HEAD",
    "headers": {"Connection": "close"},
    "follow_redirects": False
}

# The Heka configuration file template. Heka containers on each instance
# forward messages to a central Heka server via TcpOutput.
HEKA_CONFIG_PATH = os.path.join(os.path.dirname(__file__), "heka",
                                "config.src.toml")

with open(HEKA_CONFIG_PATH, "r") as f:
    HEKA_CONFIG_TEMPLATE = Template(f.read())


class Ping:
    def __init__(self, io_loop=None):
        self._loop = io_loop or tornado.ioloop.IOLoop.instance()
        self._ping_client = AsyncHTTPClient(io_loop=self._loop,
                                            defaults=_PING_DEFAULTS)

    @gen.coroutine
    def ping(self, instance, url, attempts=5, delay=0.5, max_jitter=0.2,
             max_delay=15, **options):
        """Attempts to load a URL to verify its reachable."""
        attempt = 1
        while True:
            try:
                yield self._ping_client.fetch(url, **options)
                return True
            except ConnectionError:
                jitter = randint(0, max_jitter * 100) / 100
                yield gen.Task(self._loop.add_timeout,
                               time.time() + delay + jitter)
                attempt += 1
                delay = min(delay * 2, max_delay)
                if attempt >= attempts:
                    raise


class SSH:
    """SSH client to communicate with instances."""
    def __init__(self, ssh_keyfile):
        self._ssh_keyfile = ssh_keyfile

    def connect(self, instance):
        """Opens an SSH connection to this instance."""
        client = sshclient.SSHClient()
        client.set_missing_host_key_policy(sshclient.AutoAddPolicy())
        client.connect(instance.ip_address, username="core",
                       key_filename=self._ssh_keyfile)
        return client

    def _send_file(self, sftp, local_obj, remote_file):
        # Ensure the base directory for the remote file exists
        base_dir = os.path.dirname(remote_file)
        makedirs(sftp, base_dir)

        # Copy the local file to the remote location.
        sftp.putfo(local_obj, remote_file)

    def upload_file(self, instance, local_obj, remote_file):
        """Upload a file to an instance. Blocks."""
        client = self.connect(instance)
        try:
            sftp = client.open_sftp()
            try:
                self._send_file(sftp, local_obj, remote_file)
            finally:
                sftp.close()
        finally:
            client.close()

    @gen.coroutine
    def reload_sysctl(self, collection):
        def _reload(inst):
            client = self.connect(inst.instance)
            try:
                stdin, stdout, stderr = client.exec_command(
                    "sudo sysctl -p /etc/sysctl.conf")
                output = stdout.channel.recv(4096)
                stdin.close()
                stdout.close()
                stderr.close()
                return output
            finally:
                client.close()
        yield collection.map(_reload)


class Docker:
    def __init__(self, ssh):
        self.sshclient = ssh

    @gen.coroutine
    def setup_collection(self, collection):
        def setup_docker(ec2_instance):
            instance, state = ec2_instance
            if instance.ip_address is None:
                docker_host = 'tcp://0.0.0.0:7890'
            else:
                docker_host = "tcp://%s:2375" % instance.ip_address

            if not hasattr(state, "docker"):
                state.docker = DockerDaemon(host=docker_host)
        yield collection.map(setup_docker)

    @staticmethod
    def not_responding_instances(collection):
        return [x for x in collection.instances
                if not x.state.docker.responded]

    @gen.coroutine
    def wait(self, collection, interval=5, timeout=600):
        """Waits till docker is available on every instance in the
        collection."""
        end = time.time() + timeout

        not_responded = self.not_responding_instances(collection)

        def get_container(inst):
            try:
                inst.state.docker.get_containers()
                inst.state.docker.responded = True
            except Exception:
                pass

        # Attempt to fetch until they've all responded
        while not_responded and time.time() < end:
            yield [collection.execute(get_container, x) for x in
                   not_responded]

            # Update the not_responded
            not_responded = self.not_responding_instances(collection)

            if not_responded:
                yield collection.wait(interval)

        # Prune the non-responding
        logger.debug("Pruning %d non-responding instances.",
                     len(not_responded))
        yield collection.remove_instances(not_responded)

    @gen.coroutine
    def is_running(self, collection, container_name):
        """Checks running instances in a collection to see if the provided
        container_name is running on the instance."""
        def has_container(instance):
            try:
                all_containers = instance.state.docker.get_containers()
                for _, container in all_containers.items():
                    if container_name in container["Image"]:
                        return True
            except:
                logger.debug("Lost contact with a container, marking dead.")
                instance.state.nonresponsive = True
            return False

        results = yield [collection.execute(has_container, x) for x in
                         collection.running_instances()]
        return any(results)

    @gen.coroutine
    def load_containers(self, collection, container_name, container_url):
        """Loads's a container of the provided name to the instance."""
        def load(instance, tries=0):
            docker = instance.state.docker

            has_container = docker.has_image(container_name)
            if has_container:
                return

            if container_url:
                client = self.sshclient.connect(instance.instance)
                try:
                    output = docker.import_container(client, container_url)
                finally:
                    client.close()
            else:
                output = docker.pull_container(container_name)

            if not docker.has_image(container_name):
                if tries > 3:
                    logger.debug("Can't load container, retries exceeded.")
                    return False

                logger.debug("Unable to load container: %s. Retrying.",
                             output)
                return load(instance, tries+1)
            return output

        yield collection.map(load)

    @gen.coroutine
    def run_containers(self, collection, container_name, env, command_args,
                       volumes={}, ports={}, local_dns=None):
        """Run a container of the provided name with the env/command
        args supplied."""
        env = env or ""

        if local_dns is not None:
            local_dns = collection.local_dns

        if isinstance(ports, str):
            port_list = [x.split(":") for x in ports.split(",")]
            ports = {x[0]: x[1] for x in port_list if x and len(x) == 2}

        if isinstance(volumes, str):
            volume_list = [x.split(":") for x in volumes.split(",")]
            volumes = {x[1]: {"bind": x[0], "ro": len(x) < 3 or x[2] == "ro"}
                for x in volume_list if x and len(x) >= 2}

        def run(instance, tries=0):
            dns = getattr(instance.state, "dns_server", [])
            docker = instance.state.docker
            added_env = "\n".join([
                "HOST_IP=%s" % instance.instance.ip_address,
                "PRIVATE_IP=%s" % instance.instance.private_ip_address,
                "STATSD_HOST=%s" % instance.instance.private_ip_address,
                "STATSD_PORT=8125"])
            if env:
                _env = env + "\n" + added_env
            else:
                _env = added_env
            _env = self.substitute_names(_env, _env)
            container_env = _env.split("\n")
            container_args = self.substitute_names(command_args, _env)

            container_volumes = {}
            for host, volume in volumes.items():
                binding = volume.copy()
                binding["bind"] = self.substitute_names(
                    binding.get("bind", host), _env)
                container_volumes[self.substitute_names(host, _env)] = binding

            try:
                return docker.run_container(
                    container_name, container_env, container_args,
                    container_volumes, ports, dns=dns)
            except Exception as exc:
                logger.debug("Exception with run_container: %s", exc)
                if tries > 3:
                    logger.debug("Giving up on running container.")
                    return False
                docker.stop_container(container_name)
                return run(instance, tries=tries+1)
        results = yield collection.map(run)
        return results

    @gen.coroutine
    def kill_containers(self, collection, container_name):
        """Kill the container with the provided name."""
        def kill(instance):<commit_msg>Ensure no empty env vars.<commit_after>import os
import time
from io import StringIO
from random import randint
from shlex import quote as shell_quote
from urllib.parse import quote as url_quote
from string import Template
from collections import namedtuple

import paramiko.client as sshclient
import tornado.ioloop
from tornado import gen
from tornado.httpclient import AsyncHTTPClient

from loadsbroker import logger
from loadsbroker.exceptions import LoadsException
from loadsbroker.dockerctrl import DockerDaemon
from loadsbroker.ssh import makedirs
from loadsbroker.util import join_host_port

# Default ping request options.
_PING_DEFAULTS = {
    "method": "HEAD",
    "headers": {"Connection": "close"},
    "follow_redirects": False
}

# The Heka configuration file template. Heka containers on each instance
# forward messages to a central Heka server via TcpOutput.
HEKA_CONFIG_PATH = os.path.join(os.path.dirname(__file__), "heka",
                                "config.src.toml")

with open(HEKA_CONFIG_PATH, "r") as f:
    HEKA_CONFIG_TEMPLATE = Template(f.read())


class Ping:
    def __init__(self, io_loop=None):
        self._loop = io_loop or tornado.ioloop.IOLoop.instance()
        self._ping_client = AsyncHTTPClient(io_loop=self._loop,
                                            defaults=_PING_DEFAULTS)

    @gen.coroutine
    def ping(self, instance, url, attempts=5, delay=0.5, max_jitter=0.2,
             max_delay=15, **options):
        """Attempts to load a URL to verify its reachable."""
        attempt = 1
        while True:
            try:
                yield self._ping_client.fetch(url, **options)
                return True
            except ConnectionError:
                jitter = randint(0, max_jitter * 100) / 100
                yield gen.Task(self._loop.add_timeout,
                               time.time() + delay + jitter)
                attempt += 1
                delay = min(delay * 2, max_delay)
                if attempt >= attempts:
                    raise


class SSH:
    """SSH client to communicate with instances."""
    def __init__(self, ssh_keyfile):
        self._ssh_keyfile = ssh_keyfile

    def connect(self, instance):
        """Opens an SSH connection to this instance."""
        client = sshclient.SSHClient()
        client.set_missing_host_key_policy(sshclient.AutoAddPolicy())
        client.connect(instance.ip_address, username="core",
                       key_filename=self._ssh_keyfile)
        return client

    def _send_file(self, sftp, local_obj, remote_file):
        # Ensure the base directory for the remote file exists
        base_dir = os.path.dirname(remote_file)
        makedirs(sftp, base_dir)

        # Copy the local file to the remote location.
        sftp.putfo(local_obj, remote_file)

    def upload_file(self, instance, local_obj, remote_file):
        """Upload a file to an instance. Blocks."""
        client = self.connect(instance)
        try:
            sftp = client.open_sftp()
            try:
                self._send_file(sftp, local_obj, remote_file)
            finally:
                sftp.close()
        finally:
            client.close()

    @gen.coroutine
    def reload_sysctl(self, collection):
        def _reload(inst):
            client = self.connect(inst.instance)
            try:
                stdin, stdout, stderr = client.exec_command(
                    "sudo sysctl -p /etc/sysctl.conf")
                output = stdout.channel.recv(4096)
                stdin.close()
                stdout.close()
                stderr.close()
                return output
            finally:
                client.close()
        yield collection.map(_reload)


class Docker:
    def __init__(self, ssh):
        self.sshclient = ssh

    @gen.coroutine
    def setup_collection(self, collection):
        def setup_docker(ec2_instance):
            instance, state = ec2_instance
            if instance.ip_address is None:
                docker_host = 'tcp://0.0.0.0:7890'
            else:
                docker_host = "tcp://%s:2375" % instance.ip_address

            if not hasattr(state, "docker"):
                state.docker = DockerDaemon(host=docker_host)
        yield collection.map(setup_docker)

    @staticmethod
    def not_responding_instances(collection):
        return [x for x in collection.instances
                if not x.state.docker.responded]

    @gen.coroutine
    def wait(self, collection, interval=5, timeout=600):
        """Waits till docker is available on every instance in the
        collection."""
        end = time.time() + timeout

        not_responded = self.not_responding_instances(collection)

        def get_container(inst):
            try:
                inst.state.docker.get_containers()
                inst.state.docker.responded = True
            except Exception:
                pass

        # Attempt to fetch until they've all responded
        while not_responded and time.time() < end:
            yield [collection.execute(get_container, x) for x in
                   not_responded]

            # Update the not_responded
            not_responded = self.not_responding_instances(collection)

            if not_responded:
                yield collection.wait(interval)

        # Prune the non-responding
        logger.debug("Pruning %d non-responding instances.",
                     len(not_responded))
        yield collection.remove_instances(not_responded)

    @gen.coroutine
    def is_running(self, collection, container_name):
        """Checks running instances in a collection to see if the provided
        container_name is running on the instance."""
        def has_container(instance):
            try:
                all_containers = instance.state.docker.get_containers()
                for _, container in all_containers.items():
                    if container_name in container["Image"]:
                        return True
            except:
                logger.debug("Lost contact with a container, marking dead.")
                instance.state.nonresponsive = True
            return False

        results = yield [collection.execute(has_container, x) for x in
                         collection.running_instances()]
        return any(results)

    @gen.coroutine
    def load_containers(self, collection, container_name, container_url):
        """Loads's a container of the provided name to the instance."""
        def load(instance, tries=0):
            docker = instance.state.docker

            has_container = docker.has_image(container_name)
            if has_container:
                return

            if container_url:
                client = self.sshclient.connect(instance.instance)
                try:
                    output = docker.import_container(client, container_url)
                finally:
                    client.close()
            else:
                output = docker.pull_container(container_name)

            if not docker.has_image(container_name):
                if tries > 3:
                    logger.debug("Can't load container, retries exceeded.")
                    return False

                logger.debug("Unable to load container: %s. Retrying.",
                             output)
                return load(instance, tries+1)
            return output

        yield collection.map(load)

    @gen.coroutine
    def run_containers(self, collection, container_name, env, command_args,
                       volumes={}, ports={}, local_dns=None):
        """Run a container of the provided name with the env/command
        args supplied."""
        env = env or ""

        if local_dns is not None:
            local_dns = collection.local_dns

        if isinstance(ports, str):
            port_list = [x.split(":") for x in ports.split(",")]
            ports = {x[0]: x[1] for x in port_list if x and len(x) == 2}

        if isinstance(volumes, str):
            volume_list = [x.split(":") for x in volumes.split(",")]
            volumes = {x[1]: {"bind": x[0], "ro": len(x) < 3 or x[2] == "ro"}
                for x in volume_list if x and len(x) >= 2}

        def run(instance, tries=0):
            dns = getattr(instance.state, "dns_server", [])
            docker = instance.state.docker
            added_env = "\n".join([
                "HOST_IP=%s" % instance.instance.ip_address,
                "PRIVATE_IP=%s" % instance.instance.private_ip_address,
                "STATSD_HOST=%s" % instance.instance.private_ip_address,
                "STATSD_PORT=8125"])
            if env:
                _env = env + "\n" + added_env
            else:
                _env = added_env
            _env = self.substitute_names(_env, _env)
            container_env = [x for x in _env.split("\n") if x]
            container_args = self.substitute_names(command_args, _env)

            container_volumes = {}
            for host, volume in volumes.items():
                binding = volume.copy()
                binding["bind"] = self.substitute_names(
                    binding.get("bind", host), _env)
                container_volumes[self.substitute_names(host, _env)] = binding

            try:
                return docker.run_container(
                    container_name, container_env, container_args,
                    container_volumes, ports, dns=dns)
            except Exception as exc:
                logger.debug("Exception with run_container: %s", exc)
                if tries > 3:
                    logger.debug("Giving up on running container.")
                    return False
                docker.stop_container(container_name)
                return run(instance, tries=tries+1)
        results = yield collection.map(run)
        return results

    @gen.coroutine
    def kill_containers(self, collection, container_name):
        """Kill the container with the provided name."""
        def kill(instance):
{'fraction_non_alphanumeric': 0.0875250501002004, 'fraction_numerical': 0.005711422845691383, 'mean_word_length': 1.8837041317538283, 'pattern_counts': {'":': 16, '<': 7, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6405700', 'n_tokens_mistral': 5576, 'n_tokens_neox': 5184, 'n_words': 1419}