<commit_before>import os
import sys
import math
import select
import socket
import getpass
import logging
import textwrap
from six import iteritems
from six.moves import urllib
from pymesos import Scheduler, MesosSchedulerDriver
from tfmesos.utils import send, recv, setup_logger


FOREVER = 0xFFFFFFFF
logger = logging.getLogger(__name__)


class Job(object):

    def __init__(self, name, num, cpus=1.0, mem=1024.0,
                 gpus=0, cmd=None, start=0):

        self.name = name
        self.num = num
        self.cpus = cpus
        self.gpus = gpus
        self.mem = mem
        self.cmd = cmd
        self.start = start


class Task(object):

    def __init__(self, mesos_task_id, job_name, task_index,
                 cpus=1.0, mem=1024.0, gpus=0, cmd=None, volumes={}):
        self.mesos_task_id = mesos_task_id
        self.job_name = job_name
        self.task_index = task_index

        self.cpus = cpus
        self.gpus = gpus
        self.mem = mem
        self.cmd = cmd
        self.volumes = volumes
        self.offered = False

        self.addr = None
        self.connection = None
        self.initalized = False

    def __str__(self):
        return textwrap.dedent('''
        <Task
          mesos_task_id=%s
          addr=%s
        >''' % (self.mesos_task_id, self.addr))

    def to_task_info(self, offer, master_addr, gpu_uuids=[],
                     gpu_resource_type=None):
        resources = [
            dict(
                name='cpus',
                type='SCALAR',
                scalar=dict(value=self.cpus),
            ),
            dict(
                name='mem',
                type='SCALAR',
                scalar=dict(value=self.mem),
            )
        ]

        cmd = ' '.join([
            sys.executable, '-m', '%s.server' % __package__,
            str(self.mesos_task_id), master_addr
        ])

        ti = dict(
            task_id=dict(
                value=str(self.mesos_task_id)
            ),
            agent_id=offer['agent_id'],
            name='/job:%s/task:%s' % (self.job_name, self.task_index),
            resources=resources,
            command=dict(
                shell=True,
                value=cmd,
                environment=dict(
                    variables=[
                        dict(
                            name='PYTHONPATH',
                            value=':'.join(sys.path)
                        )
                    ]
                )
            ),
        )

        image = os.environ.get('DOCKER_IMAGE')

        if image is not None:
            volumes = []
            ti.update(dict(
                container=dict(
                    type='DOCKER',
                    docker=dict(image=image),
                    volumes=volumes
                ),
            ))

            for path in ['/etc/passwd', '/etc/group']:
                volumes.append(dict(
                    host_path=path,
                    container_path=path,
                    mode='RO'
                ))

            for src, dst in iteritems(self.volumes):
                volumes.append(dict(
                    host_path=src,
                    container_path=dst,
                    mode='RW'
                ))

            if self.gpus and gpu_uuids and gpu_resource_type is not None:
                if gpu_resource_type == 'SET':
                    hostname = offer['hostname']
                    url = 'http://%s:3476/docker/cli?dev=%s' % (
                        hostname, urllib.parse.quote(
                            ' '.join(gpu_uuids)
                        )
                    )

                    try:
                        docker_args = urllib.request.urlopen(url).read()
                        for arg in docker_args.split():
                            k, v = arg.split('=')
                            assert k.startswith('--')
                            k = k[2:]
                            p = ti.container.docker.parameters.add()
                            p.key = k
                            p.value = v

                        resources.append(dict(
                            name='gpus',
                            type='SET',
                            set=dict(item=[gpu_uuids])
                        ))

                    except Exception:
                        logger.exception(
                            'fail to determine remote device parameter,'
                            ' disable gpu resources'
                        )
                else:
                    resources.append(dict(
                        name='gpus',
                        type='SCALAR',
                        scalar=dict(value=len(gpu_uuids))
                    ))

        else:
            if self.gpus and gpu_uuids and gpu_resource_type is not None:
                if gpu_resource_type == 'SET':
                    resources.append(dict(
                        name='gpus',
                        type='SET',
                        set=dict(item=[gpu_uuids])
                    ))
                else:
                    resources.append(dict(
                        name='gpus',
                        type='SCALAR',
                        scalar=len(gpu_uuids)<commit_msg>fix nvidia-docker support<commit_after>import os
import sys
import math
import select
import socket
import getpass
import logging
import textwrap
from six import iteritems
from six.moves import urllib
from pymesos import Scheduler, MesosSchedulerDriver
from tfmesos.utils import send, recv, setup_logger


FOREVER = 0xFFFFFFFF
logger = logging.getLogger(__name__)


class Job(object):

    def __init__(self, name, num, cpus=1.0, mem=1024.0,
                 gpus=0, cmd=None, start=0):

        self.name = name
        self.num = num
        self.cpus = cpus
        self.gpus = gpus
        self.mem = mem
        self.cmd = cmd
        self.start = start


class Task(object):

    def __init__(self, mesos_task_id, job_name, task_index,
                 cpus=1.0, mem=1024.0, gpus=0, cmd=None, volumes={}):
        self.mesos_task_id = mesos_task_id
        self.job_name = job_name
        self.task_index = task_index

        self.cpus = cpus
        self.gpus = gpus
        self.mem = mem
        self.cmd = cmd
        self.volumes = volumes
        self.offered = False

        self.addr = None
        self.connection = None
        self.initalized = False

    def __str__(self):
        return textwrap.dedent('''
        <Task
          mesos_task_id=%s
          addr=%s
        >''' % (self.mesos_task_id, self.addr))

    def to_task_info(self, offer, master_addr, gpu_uuids=[],
                     gpu_resource_type=None):
        resources = [
            dict(
                name='cpus',
                type='SCALAR',
                scalar=dict(value=self.cpus),
            ),
            dict(
                name='mem',
                type='SCALAR',
                scalar=dict(value=self.mem),
            )
        ]

        cmd = ' '.join([
            sys.executable, '-m', '%s.server' % __package__,
            str(self.mesos_task_id), master_addr
        ])

        ti = dict(
            task_id=dict(
                value=str(self.mesos_task_id)
            ),
            agent_id=offer['agent_id'],
            name='/job:%s/task:%s' % (self.job_name, self.task_index),
            resources=resources,
            command=dict(
                shell=True,
                value=cmd,
                environment=dict(
                    variables=[
                        dict(
                            name='PYTHONPATH',
                            value=':'.join(sys.path)
                        )
                    ]
                )
            ),
        )

        image = os.environ.get('DOCKER_IMAGE')

        if image is not None:
            volumes = []
            ti.update(dict(
                container=dict(
                    type='DOCKER',
                    docker=dict(image=image),
                    volumes=volumes
                ),
            ))

            for path in ['/etc/passwd', '/etc/group']:
                volumes.append(dict(
                    host_path=path,
                    container_path=path,
                    mode='RO'
                ))

            for src, dst in iteritems(self.volumes):
                volumes.append(dict(
                    host_path=src,
                    container_path=dst,
                    mode='RW'
                ))

            if self.gpus and gpu_uuids and gpu_resource_type is not None:
                if gpu_resource_type == 'SET':
                    hostname = offer['hostname']
                    url = 'http://%s:3476/docker/cli?dev=%s' % (
                        hostname, urllib.parse.quote(
                            ' '.join(gpu_uuids)
                        )
                    )

                    try:
                        parameters = []
                        ti['container']['docker']['parameters'] = parameters
                        docker_args = urllib.request.urlopen(url).read()
                        for arg in docker_args.split():
                            k, v = arg.split('=')
                            assert k.startswith('--')
                            k = k[2:]
                            parameters.append(dict(
                                key=k,
                                value=v,
                            ))

                        resources.append(dict(
                            name='gpus',
                            type='SET',
                            set=dict(item=gpu_uuids)
                        ))

                    except Exception:
                        logger.exception(
                            'fail to determine remote device parameter,'
                            ' disable gpu resources'
                        )
                else:
                    resources.append(dict(
                        name='gpus',
                        type='SCALAR',
                        scalar=dict(value=len(gpu_uuids))
                    ))

        else:
            if self.gpus and gpu_uuids and gpu_resource_type is not None:
                if gpu_resource_type == 'SET':
                    resources.append(dict(
                        name='gpus',
                        type='SET',
                        set=dict(item=[gpu_uuids])
                    ))
                else:
                    resources.append(dict(
                        name='gpus',
                        type='SCALAR',
                        scalar=len(gpu_uuids)
{'fraction_non_alphanumeric': 0.09382322616927546, 'fraction_numerical': 0.004311556846939732, 'mean_word_length': 1.007148231753198, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4079904', 'n_tokens_mistral': 2889, 'n_tokens_neox': 2690, 'n_words': 525}