<commit_before>from __future__ import print_function
import pickle
import shutil
import os
import ntpath
from docker import Client
from . import cfg
import io
import tarfile
import uuid
import stat

def docker_copy(docker_client, container_id, path, target="."):
    """Copy file from docker container to host machine.

    Parameters
    ----------
    docker_client : docker.Client object
        The connected Docker client.
    container_id : str
        ID of the container to copy from.
    path : str
        Path to the file in the container.
    target : str
        Folder where to put the file.

    """
    response = docker_client.copy(container_id, path)
    buffer = io.BytesIO()
    buffer.write(response.data)
    buffer.seek(0)
    tar = tarfile.open(fileobj=buffer, mode='r|')
    tar.extractall(path=target)


def copy_data_files_featurize_prep(args_dict):
    """Copy needed files into expected location for featurization.

    Parameters
    ----------
    args_dict : dict
        Dictionary of local variable names and associated values (produced
        by calling the built-in method `locals()`).

    Returns
    -------
    list of str
        List of paths (strings) to all copied/created temporary files to
        be removed upon completion of featurization.

    """
    tmp_files = []
    if os.path.isfile(str(args_dict['headerfile_path'])):
        copied_headerfile_path = os.path.join(
            args_dict['copied_data_dir'],
            ntpath.basename(args_dict['headerfile_path']))
        tmp_files.append(copied_headerfile_path)
        shutil.copy(args_dict['headerfile_path'], copied_headerfile_path)
        args_dict["headerfile_path"] = os.path.join(
            "/data",
            ntpath.basename(args_dict['headerfile_path']))
    if os.path.isfile(str(args_dict['zipfile_path'])):
        copied_zipfile_path = os.path.join(
            args_dict['copied_data_dir'],
            ntpath.basename(args_dict['zipfile_path']))
        tmp_files.append(copied_zipfile_path)
        shutil.copy(args_dict['zipfile_path'], copied_zipfile_path)
        args_dict["zipfile_path"] = os.path.join(
            "/data",
            ntpath.basename(args_dict['zipfile_path']))
    if os.path.isfile(str(args_dict['custom_script_path'])):
        copied_custom_script_path = os.path.join(
            args_dict['copied_data_dir'], "custom_feature_defs.py")
        tmp_files.append(copied_custom_script_path)
        shutil.copy(args_dict['custom_script_path'], copied_custom_script_path)
        args_dict["custom_script_path"] = os.path.join(
            "/data", "custom_feature_defs.py")
        # Create __init__.py file so that custom feats script can be imported
        open(os.path.join(args_dict['copied_data_dir'], "__init__.py"),
             "w").close()
    args_dict["path_map"] = {args_dict['copied_data_dir']:
                             "/data"}
    function_args_path = os.path.join(args_dict['copied_data_dir'],
                                      "function_args.pkl")
    tmp_files.append(function_args_path)
    with open(function_args_path, "wb") as f:
        pickle.dump(args_dict, f, protocol=2)
    # Was getting permissions errors from Disco workers
    os.chmod(args_dict['copied_data_dir'], 0777)
    return tmp_files


def spin_up_and_run_container(task_name, tmp_data_dir):
    """Spin up and run a Docker container.

    Parameters
    ----------
    task_name : str
        Name/tag of the Docker image to be used, e.g. "mltsp/featurize".
    tmp_dir : str
        Path to directory containing data files to be mounted into
        container.

    Returns
    -------
    tuple
       Two-element tuple containing docker.Client object, and the ID of
       the newly-created container (str).

    """
    # Instantiate Docker client
    client = Client(base_url='unix://var/run/docker.sock',
                    version='1.14')
    # Create container
    cont_id = client.create_container(
        image="mltsp/base_disco",
        command="python /home/mltsp/mltsp/run_script_in_container.py --%s" %\
        task_name,
        tty=True,
        volumes={"/home/mltsp/mltsp": "",
                 "/data": ""})["Id"]
    # Start container
    client.start(cont_id,
                 binds={cfg.PROJECT_PATH: {"bind": "/home/mltsp/mltsp", "ro": True},
                        tmp_data_dir: {"bind": "/data",
                                       "ro": True}},
                 privileged=True)
    # Wait for process to complete
    client.wait(cont_id)
    stdout = client.logs(container=cont_id, stdout=True)
    stderr = client.logs(container=cont_id, stderr=True)
    print("\n\ndocker container stdout:\n\n", str(stdout),
          "\n\ndocker container stderr:\n\n", str(stderr), "\n\n")
    # Clean up unwanted copied files
    for fname in ["custom_feature_defs.py", "custom_feature_defs.pyc",
                  "__init__.pyc"]:
        fpath = os.path.join(cfg.TMP_CUSTOM_FEATS_FOLDER, fname)
        if os.path.exists(fpath):
            os.remove(fpath)
    return (client, cont_id)


def copy_results_files_featurize(featureset_key, client, cont_id):
    """Copy features results files from Docker container to host machine.

    Parameters
    ----------
    featureset_key : str
        Key/ID identifying the relevant feature set.
    client : docker.Client object<commit_msg>Do not run in privileged mode<commit_after>from __future__ import print_function
import pickle
import shutil
import os
import ntpath
from docker import Client
from . import cfg
import io
import tarfile
import uuid
import stat

def docker_copy(docker_client, container_id, path, target="."):
    """Copy file from docker container to host machine.

    Parameters
    ----------
    docker_client : docker.Client object
        The connected Docker client.
    container_id : str
        ID of the container to copy from.
    path : str
        Path to the file in the container.
    target : str
        Folder where to put the file.

    """
    response = docker_client.copy(container_id, path)
    buffer = io.BytesIO()
    buffer.write(response.data)
    buffer.seek(0)
    tar = tarfile.open(fileobj=buffer, mode='r|')
    tar.extractall(path=target)


def copy_data_files_featurize_prep(args_dict):
    """Copy needed files into expected location for featurization.

    Parameters
    ----------
    args_dict : dict
        Dictionary of local variable names and associated values (produced
        by calling the built-in method `locals()`).

    Returns
    -------
    list of str
        List of paths (strings) to all copied/created temporary files to
        be removed upon completion of featurization.

    """
    tmp_files = []
    if os.path.isfile(str(args_dict['headerfile_path'])):
        copied_headerfile_path = os.path.join(
            args_dict['copied_data_dir'],
            ntpath.basename(args_dict['headerfile_path']))
        tmp_files.append(copied_headerfile_path)
        shutil.copy(args_dict['headerfile_path'], copied_headerfile_path)
        args_dict["headerfile_path"] = os.path.join(
            "/data",
            ntpath.basename(args_dict['headerfile_path']))
    if os.path.isfile(str(args_dict['zipfile_path'])):
        copied_zipfile_path = os.path.join(
            args_dict['copied_data_dir'],
            ntpath.basename(args_dict['zipfile_path']))
        tmp_files.append(copied_zipfile_path)
        shutil.copy(args_dict['zipfile_path'], copied_zipfile_path)
        args_dict["zipfile_path"] = os.path.join(
            "/data",
            ntpath.basename(args_dict['zipfile_path']))
    if os.path.isfile(str(args_dict['custom_script_path'])):
        copied_custom_script_path = os.path.join(
            args_dict['copied_data_dir'], "custom_feature_defs.py")
        tmp_files.append(copied_custom_script_path)
        shutil.copy(args_dict['custom_script_path'], copied_custom_script_path)
        args_dict["custom_script_path"] = os.path.join(
            "/data", "custom_feature_defs.py")
        # Create __init__.py file so that custom feats script can be imported
        open(os.path.join(args_dict['copied_data_dir'], "__init__.py"),
             "w").close()
    args_dict["path_map"] = {args_dict['copied_data_dir']:
                             "/data"}
    function_args_path = os.path.join(args_dict['copied_data_dir'],
                                      "function_args.pkl")
    tmp_files.append(function_args_path)
    with open(function_args_path, "wb") as f:
        pickle.dump(args_dict, f, protocol=2)
    # Was getting permissions errors from Disco workers
    os.chmod(args_dict['copied_data_dir'], 0777)
    return tmp_files


def spin_up_and_run_container(task_name, tmp_data_dir):
    """Spin up and run a Docker container.

    Parameters
    ----------
    task_name : str
        Name/tag of the Docker image to be used, e.g. "mltsp/featurize".
    tmp_dir : str
        Path to directory containing data files to be mounted into
        container.

    Returns
    -------
    tuple
       Two-element tuple containing docker.Client object, and the ID of
       the newly-created container (str).

    """
    # Instantiate Docker client
    client = Client(base_url='unix://var/run/docker.sock',
                    version='1.14')
    # Create container
    cont_id = client.create_container(
        image="mltsp/base_disco",
        command="python /home/mltsp/mltsp/run_script_in_container.py --%s" %\
        task_name,
        tty=True,
        volumes={"/home/mltsp/mltsp": "",
                 "/data": ""})["Id"]
    # Start container
    client.start(cont_id,
                 binds={cfg.PROJECT_PATH: {"bind": "/home/mltsp/mltsp", "ro": True},
                        tmp_data_dir: {"bind": "/data",
                                       "ro": True}},
                 privileged=False)
    # Wait for process to complete
    client.wait(cont_id)
    stdout = client.logs(container=cont_id, stdout=True)
    stderr = client.logs(container=cont_id, stderr=True)
    print("\n\ndocker container stdout:\n\n", str(stdout),
          "\n\ndocker container stderr:\n\n", str(stderr), "\n\n")
    # Clean up unwanted copied files
    for fname in ["custom_feature_defs.py", "custom_feature_defs.pyc",
                  "__init__.pyc"]:
        fpath = os.path.join(cfg.TMP_CUSTOM_FEATS_FOLDER, fname)
        if os.path.exists(fpath):
            os.remove(fpath)
    return (client, cont_id)


def copy_results_files_featurize(featureset_key, client, cont_id):
    """Copy features results files from Docker container to host machine.

    Parameters
    ----------
    featureset_key : str
        Key/ID identifying the relevant feature set.
    client : docker.Client object
{'fraction_non_alphanumeric': 0.11364270436157352, 'fraction_numerical': 0.0016739514554077932, 'mean_word_length': 2.5515191545574636, 'pattern_counts': {'":': 12, '<': 3, '<?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': '646752', 'n_tokens_mistral': 3321, 'n_tokens_neox': 3235, 'n_words': 732}