<commit_before>"""
"""

import logging
import os
import mesos.interface
import progressbar
import sys
import StringIO
import tarfile
import tempfile
import threading
import traceback
import uuid

from fnmatch import fnmatch
from functools import partial
from fs.opener import opener
from pesos.vendor.mesos import mesos_pb2
from urlparse import urlparse

from portainer.proto import portainer_pb2
from portainer.util.parser import parse_dockerfile, parse_dockerignore

logger = logging.getLogger("portainer.scheduler")


class TaskContextException(Exception):
    pass


class StagingSystemRequiredException(Exception):
    pass


class Scheduler(mesos.interface.Scheduler):
    """Mesos scheduler that is responsible for launching the builder tasks."""

    def __init__(self, tasks, executor_uri, cpu_limit, mem_limit, push_registry,
                 staging_uri, stream=False, verbose=False, repository=None,
                 pull_registry=None, docker_host=None, container_image=None,
                 insecure_registries=False):

        self.executor_uri = executor_uri
        self.cpu = float(cpu_limit)
        self.mem = int(mem_limit)
        self.push_registry = push_registry
        self.pull_registry = pull_registry
        self.staging_uri = staging_uri
        self.stream = stream
        self.verbose = verbose
        self.repository = repository
        self.docker_host = docker_host
        self.container_image = container_image
        self.insecure_registries = insecure_registries

        self.queued_tasks = []
        for path, tags in tasks:
            dockerfile = parse_dockerfile(path, registry=pull_registry)
            self.queued_tasks.append((path, dockerfile, tags))

        self.pending = len(self.queued_tasks)
        self.running = 0
        self.finished = 0
        self.failed = 0
        self.task_ids = {}

        self.processing_offers = threading.Lock()

        # Ensure the staging directory exists
        self.filesystem = None
        if self.staging_uri:
            staging_uri = urlparse(self.staging_uri)
            staging_fs = opener.opendir(
                "%s://%s/" % (staging_uri.scheme, staging_uri.netloc)
            )

            staging_fs.makedir(
                staging_uri.path.lstrip("/"),
                recursive=True,
                allow_recreate=True
            )

            self.filesystem = opener.opendir(self.staging_uri)

    def registered(self, driver, frameworkId, masterInfo):
        host = masterInfo.hostname or masterInfo.ip
        master = "http://%s:%s" % (host, masterInfo.port)
        logger.info("Framework %s registered to %s", frameworkId.value, master)

    def disconnected(self, driver):
        logger.warning("Framework disconnected from the mesos master")

    def reregistered(self, driver, masterInfo):
        host = masterInfo.hostname or masterInfo.ip
        master = "http://%s:%s" % (host, masterInfo.port)
        logger.info("Framework re-registered to %s", master)

    def error(self, driver, message):
        logger.error("Framework error: %s", message)

    def resource_offers(self, driver, offers):

        # Spawn another thread to handle offer processing to free up the driver
        t = threading.Thread(target=partial(
            self._handle_offers,
            driver,
            offers
        ))

        t.setDaemon(True)
        t.start()

    def _handle_offers(self, driver, offers):

        # We only want to process offers one set at a time
        with self.processing_offers:
            tasks_to_launch = []

            if not self.pending:
                for offer in offers:
                    driver.declineOffer(offer.id)
            else:
                for offer in offers:
                    offer_cpu = 0.0
                    offer_mem = 0

                    # Extract the important resources from the offer
                    for resource in offer.resources:
                        if resource.name == "cpus":
                            offer_cpu = float(resource.scalar.value)
                        if resource.name == "mem":
                            offer_mem = int(resource.scalar.value)

                    logger.debug("Received offer for cpus:%f mem:%d", offer_cpu, offer_mem)

                    # Look for a task in the queue that fits the bill
                    for idx, (path, dockerfile, tags) in enumerate(self.queued_tasks):
                        cpu = float(dockerfile.get("BUILD_CPU", [self.cpu]).next()[0])
                        mem = int(dockerfile.get("BUILD_MEM", [self.mem]).next()[0])

                        if cpu <= offer_cpu and mem <= offer_mem:
                            self.queued_tasks[idx] = None  # Remove the task from the queue
                            self.pending -= 1
                            self.running += 1
                            tasks_to_launch.append((offer, path, dockerfile,
                                                    tags, cpu, mem))
                            # TODO: No support for multiple tasks per offer yet
                            break
                    else:
                        logger.debug("Ignoring offer %r", offer)
                        driver.declineOffer(offer.id)

                    # Remove all of the tasks that are about to be launched
                    self.queued_tasks = filter(None, self.queued_tasks)

            # Launch the build tasks on the mesos cluster
            for offer, path, dockerfile, tags, cpu, mem in tasks_to_launch:
                # Generate a task ID
                task_id = str(uuid.uuid1())

                try:
                    tasks = [self._prepare_task(
                        driver=driver,
                        task_id=task_id,
                        path=path,
                        dockerfile=dockerfile,
                        tags=tags,
                        offer=offer,
                        cpu=cpu,
                        mem=mem
                    )]
                except TaskContextException as e:
                    logger.error("Caught exception: %s", e.message)
                    self.failed += 1
                    self.running -= 1
                    tasks = []
                except StagingSystemRequiredException as e:
                    logger.error("Caught exception: %s", e.message)
                    self.failed += 1
                    self.running -= 1
                    tasks = []

                if not tasks:
                    logger.error("Task %s failed to launch", task_id)

                    # If there's no pending tasks or any tasks running, stop
                    # the driver.
                    if (self.pending + self.running) == 0:
                        driver.stop()
                else:
                    logger.info("Launching %d tasks", len(tasks))
                    driver.launchTasks(offer.id, tasks)

    def status_update(self, driver, update):
        """Called when a status update is received from the mesos cluster."""

        finished = False
        failed = False
        task_id = None

        if update.task_id.value in self.task_ids:
            build_task = self.task_ids[update.task_id.value]
            task_id = build_task.image.repository
        else:
            task_id = update.task_id.value
            logger.error("Task update for unknown task! %s", task_id)

        if update.state == mesos_pb2.TASK_STARTING:
            logger.info("Task update %s : STARTING", task_id)
        if update.state == mesos_pb2.TASK_RUNNING:
            logger.info("Task update %s : RUNNING", task_id)
        if update.state == mesos_pb2.TASK_FAILED:
            logger.info("Task update %s : FAILED", task_id)
            if update.message and update.data:
                logger.info("Exception caught while building image: \n\n%s", update.data)
            failed = True
        elif update.state == mesos_pb2.TASK_FINISHED:<commit_msg>Send list of offer IDs:<commit_after>"""
"""

import logging
import os
import mesos.interface
import progressbar
import sys
import StringIO
import tarfile
import tempfile
import threading
import traceback
import uuid

from fnmatch import fnmatch
from functools import partial
from fs.opener import opener
from pesos.vendor.mesos import mesos_pb2
from urlparse import urlparse

from portainer.proto import portainer_pb2
from portainer.util.parser import parse_dockerfile, parse_dockerignore

logger = logging.getLogger("portainer.scheduler")


class TaskContextException(Exception):
    pass


class StagingSystemRequiredException(Exception):
    pass


class Scheduler(mesos.interface.Scheduler):
    """Mesos scheduler that is responsible for launching the builder tasks."""

    def __init__(self, tasks, executor_uri, cpu_limit, mem_limit, push_registry,
                 staging_uri, stream=False, verbose=False, repository=None,
                 pull_registry=None, docker_host=None, container_image=None,
                 insecure_registries=False):

        self.executor_uri = executor_uri
        self.cpu = float(cpu_limit)
        self.mem = int(mem_limit)
        self.push_registry = push_registry
        self.pull_registry = pull_registry
        self.staging_uri = staging_uri
        self.stream = stream
        self.verbose = verbose
        self.repository = repository
        self.docker_host = docker_host
        self.container_image = container_image
        self.insecure_registries = insecure_registries

        self.queued_tasks = []
        for path, tags in tasks:
            dockerfile = parse_dockerfile(path, registry=pull_registry)
            self.queued_tasks.append((path, dockerfile, tags))

        self.pending = len(self.queued_tasks)
        self.running = 0
        self.finished = 0
        self.failed = 0
        self.task_ids = {}

        self.processing_offers = threading.Lock()

        # Ensure the staging directory exists
        self.filesystem = None
        if self.staging_uri:
            staging_uri = urlparse(self.staging_uri)
            staging_fs = opener.opendir(
                "%s://%s/" % (staging_uri.scheme, staging_uri.netloc)
            )

            staging_fs.makedir(
                staging_uri.path.lstrip("/"),
                recursive=True,
                allow_recreate=True
            )

            self.filesystem = opener.opendir(self.staging_uri)

    def registered(self, driver, frameworkId, masterInfo):
        host = masterInfo.hostname or masterInfo.ip
        master = "http://%s:%s" % (host, masterInfo.port)
        logger.info("Framework %s registered to %s", frameworkId.value, master)

    def disconnected(self, driver):
        logger.warning("Framework disconnected from the mesos master")

    def reregistered(self, driver, masterInfo):
        host = masterInfo.hostname or masterInfo.ip
        master = "http://%s:%s" % (host, masterInfo.port)
        logger.info("Framework re-registered to %s", master)

    def error(self, driver, message):
        logger.error("Framework error: %s", message)

    def resource_offers(self, driver, offers):

        # Spawn another thread to handle offer processing to free up the driver
        t = threading.Thread(target=partial(
            self._handle_offers,
            driver,
            offers
        ))

        t.setDaemon(True)
        t.start()

    def _handle_offers(self, driver, offers):

        # We only want to process offers one set at a time
        with self.processing_offers:
            tasks_to_launch = []

            if not self.pending:
                for offer in offers:
                    driver.declineOffer(offer.id)
            else:
                for offer in offers:
                    offer_cpu = 0.0
                    offer_mem = 0

                    # Extract the important resources from the offer
                    for resource in offer.resources:
                        if resource.name == "cpus":
                            offer_cpu = float(resource.scalar.value)
                        if resource.name == "mem":
                            offer_mem = int(resource.scalar.value)

                    logger.debug("Received offer for cpus:%f mem:%d", offer_cpu, offer_mem)

                    # Look for a task in the queue that fits the bill
                    for idx, (path, dockerfile, tags) in enumerate(self.queued_tasks):
                        cpu = float(dockerfile.get("BUILD_CPU", [self.cpu]).next()[0])
                        mem = int(dockerfile.get("BUILD_MEM", [self.mem]).next()[0])

                        if cpu <= offer_cpu and mem <= offer_mem:
                            self.queued_tasks[idx] = None  # Remove the task from the queue
                            self.pending -= 1
                            self.running += 1
                            tasks_to_launch.append((offer, path, dockerfile,
                                                    tags, cpu, mem))
                            # TODO: No support for multiple tasks per offer yet
                            break
                    else:
                        logger.debug("Ignoring offer %r", offer)
                        driver.declineOffer(offer.id)

                    # Remove all of the tasks that are about to be launched
                    self.queued_tasks = filter(None, self.queued_tasks)

            # Launch the build tasks on the mesos cluster
            for offer, path, dockerfile, tags, cpu, mem in tasks_to_launch:
                # Generate a task ID
                task_id = str(uuid.uuid1())

                try:
                    tasks = [self._prepare_task(
                        driver=driver,
                        task_id=task_id,
                        path=path,
                        dockerfile=dockerfile,
                        tags=tags,
                        offer=offer,
                        cpu=cpu,
                        mem=mem
                    )]
                except TaskContextException as e:
                    logger.error("Caught exception: %s", e.message)
                    self.failed += 1
                    self.running -= 1
                    tasks = []
                except StagingSystemRequiredException as e:
                    logger.error("Caught exception: %s", e.message)
                    self.failed += 1
                    self.running -= 1
                    tasks = []

                if not tasks:
                    logger.error("Task %s failed to launch", task_id)

                    # If there's no pending tasks or any tasks running, stop
                    # the driver.
                    if (self.pending + self.running) == 0:
                        driver.stop()
                else:
                    logger.info("Launching %d tasks", len(tasks))
                    driver.launchTasks([offer.id], tasks)

    def status_update(self, driver, update):
        """Called when a status update is received from the mesos cluster."""

        finished = False
        failed = False
        task_id = None

        if update.task_id.value in self.task_ids:
            build_task = self.task_ids[update.task_id.value]
            task_id = build_task.image.repository
        else:
            task_id = update.task_id.value
            logger.error("Task update for unknown task! %s", task_id)

        if update.state == mesos_pb2.TASK_STARTING:
            logger.info("Task update %s : STARTING", task_id)
        if update.state == mesos_pb2.TASK_RUNNING:
            logger.info("Task update %s : RUNNING", task_id)
        if update.state == mesos_pb2.TASK_FAILED:
            logger.info("Task update %s : FAILED", task_id)
            if update.message and update.data:
                logger.info("Exception caught while building image: \n\n%s", update.data)
            failed = True
        elif update.state == mesos_pb2.TASK_FINISHED:
{'fraction_non_alphanumeric': 0.08138077935095853, 'fraction_numerical': 0.002756546798646786, 'mean_word_length': 1.6285196772600032, 'pattern_counts': {'":': 4, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, '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': '7558358', 'n_tokens_mistral': 4279, 'n_tokens_neox': 3997, 'n_words': 1095}