<commit_before>import boto3
import csv
import os
import shutil
import subprocess

from typing import Dict

from data_refinery_common.logging import get_and_configure_logger
from data_refinery_common.models import (
    ComputationalResult,
    ComputedFile,
    Pipeline,
    Processor,
    SampleAnnotation,
    SampleComputedFileAssociation,
    SampleResultAssociation,
)
from data_refinery_common.utils import get_env_variable, get_internal_microarray_accession
from data_refinery_workers.processors import utils

logger = get_and_configure_logger(__name__)
LOCAL_ROOT_DIR = get_env_variable("LOCAL_ROOT_DIR", "/home/user/data_store")
S3_BUCKET_NAME = get_env_variable("S3_BUCKET_NAME", "data-refinery")

def _prepare_files(job_context: Dict) -> Dict:
    """A processor which takes externally-processed sample data and makes it smashable.
    """
    try:
        original_file = job_context["original_files"][0]
        sample0 = job_context['samples'][0]
        if sample0.manufacturer == 'ILLUMINA':
            job_context["is_illumina"] = True
        else:
            job_context["is_illumina"] = False

        # All files for the job are in the same directory.
        job_context["work_dir"] = LOCAL_ROOT_DIR + "/" + "processor_job_" + str(job_context["job_id"]) + "/"
        try:
            os.makedirs(job_context["work_dir"])
        except Exception as e:
            logger.exception("Could not create work directory for processor job.",
                             job_context=job_context)
            job_context["job"].failure_reason = str(e)
            job_context["success"] = False
            return job_context

        # Create the output directory and path
        job_context["input_file_path"] = original_file.absolute_file_path
        new_filename = "gene_converted_" + original_file.filename
        job_context["output_file_path"] = job_context["work_dir"] + new_filename

        try:
            # Make sure header row is correct
            with open(job_context["input_file_path"], 'r', encoding='utf-8') as tsv_in:
                tsv_in = csv.reader(tsv_in, delimiter='\t')
                for line in tsv_in:
                    row = line
                    joined = ''.join(row)
                    break
        except Exception as e:
            logger.exception("Expected text file, was this in fact a binary file?",
                             job_context=job_context,
                             file=job_context["input_file_path"])

            job_context["job"].failure_reason = str(e)
            job_context["job"].no_retry = True
            job_context["success"] = True
            return job_context

        # We want to make sure that all of our columns for conversion and smashing
        # use the same column name for gene identifiers for later lookup. ID_REF
        # is the most common, so we use that.
        if 'ID_REF' not in joined and not job_context["is_illumina"]:
            try:
                float(row[1])
                # Okay, there's no header so can just prepend to the file.
                with open(job_context["input_file_path"], 'r', encoding='utf-8') as f:
                    all_content = f.read()

                job_context["input_file_path"] = job_context["input_file_path"] + ".fixed"
                with open(job_context["input_file_path"], 'w+', encoding='utf-8') as f:
                    f.seek(0, 0)
                    f.write('ID_REF\tVALUE' + '\n' + all_content)
            except ValueError:
                # There is already a header row. Let's replace it with ID_Ref
                with open(job_context["input_file_path"], 'r', encoding='utf-8') as f:
                    all_content = f.read()

                job_context["input_file_path"] = job_context["input_file_path"] + ".fixed"
                with open(job_context["input_file_path"], 'w+', encoding='utf-8') as f:
                    f.seek(0, 0)
                    f.write('ID_REF\tVALUE' + '\n' + ("\n".join(all_content[1:])))
            except Exception as e:
                logger.exception("Unable to read input file or header row.",
                    input_file_path=job_context["input_file_path"])
                job_context['job'].failure_reason = str(e)
                job_context['success'] = False
                job_context["job"].no_retry = True
                return job_context
        else:
            if job_context["is_illumina"]:
                try:
                    float(row[1])
                    # If we're here, there's no header. We're gonna have to fudge it a little bit.

                    # These exist, but they're too hard for us to handle right now.
                    if len(row) > 3:
                        e_msg = "We found an Illumina file to NO_OP that we're not set up to process yet. Tell Rich!"
                        logger.error(e_msg,
                            job_id = job_context["job"].pk,
                            file=job_context["input_file_path"]
                        )
                        job_context['success'] = False
                        job_context["job"].failure_reason = str(e_msg)
                        job_context["job"].no_retry = True
                        return job_context

                    # Okay, there's no header so can just prepend to the file.
                    with open(job_context["input_file_path"], 'r', encoding='utf-8') as f:
                        all_content = f.read()
                    job_context["input_file_path"] = job_context["input_file_path"] + ".fixed"
                    with open(job_context["input_file_path"], 'w+', encoding='utf-8') as f:
                        f.seek(0, 0)
                        if len(row) == 2:
                            f.write('Reporter Identifier\tVALUE' + '\n' + all_content)
                        elif len(row) == 3:
                            f.write('Reporter Identifier\tVALUE\tDetection Pval' + '\n' + all_content)

                    job_context['column_name'] = 'Reporter Identifier'
                except ValueError:
                    # Okay, there's a header column, we're good.
                    job_context['column_name'] = row[0]

        # Platform
        job_context["platform"] = job_context["samples"][0].platform_accession_code
        job_context["internal_accession"] = get_internal_microarray_accession(job_context["platform"])
        if not job_context["internal_accession"]:
            logger.error("Failed to find internal accession for code", code=job_context["platform"])
            job_context["job"].failure_reason = "Failed to find internal accession for code" + job_context["platform"]
            job_context['success'] = False
            job_context["job"].no_retry = True
    except Exception as e:
        logger.exception("Failed to prepare for NO_OP job.", job_id=job_context['job'].id)
        job_context["job"].failure_reason = str(e)
        job_context['success'] = False
        job_context["job"].no_retry = True

    return job_context

def _convert_genes(job_context: Dict) -> Dict:
    """ Dispatches to the appropriate gene converter"""

    sample0 = job_context['samples'][0]
    if sample0.manufacturer == 'ILLUMINA':
        return _convert_illumina_genes(job_context)
    else:
        return _convert_affy_genes(job_context)

def _convert_affy_genes(job_context: Dict) -> Dict:
    """ Convert to Ensembl genes if we can"""

    if 'internal_accession' not in job_context:
        error_msg = "Told to convert AFFY genes without an internal_accession - how did this happen?"
        logger.error(error_msg, job_context=job_context)
        job_context["job"].failure_reason = str(error_msg)
        job_context['success'] = False
        job_context["job"].no_retry = True
        return job_context

    gene_index_path = "/home/user/gene_indexes/" + job_context["internal_accession"] + ".tsv.gz"
    if not os.path.exists(gene_index_path):
        logger.error("Missing gene index file for platform!",
            platform=job_context["internal_accession"],
            job_id=job_context["job_id"])
        job_context["job"].failure_reason = "Missing gene index for " + job_context['internal_accession']
        job_context["job"].no_retry = True
        job_context["success"] = False
        return job_context

    job_context['script_name'] = "gene_convert.R"
    try:
        result = subprocess.check_output([
            "/usr/bin/Rscript",
            "--vanilla", # Shut up Rscript
            "/home/user/data_refinery_workers/processors/" + job_context['script_name'],
            "--geneIndexPath", gene_index_path,
            "--inputFile", job_context['input_file_path'],
            "--outputFile", job_context['output_file_path']
        ], stderr=subprocess.PIPE)
    except subprocess.CalledProcessError as e:
        error_template = "Status code {0} from {1}: {2}"
        error_message = error_template.format(e.returncode, job_context['script_name'], e.stderr)
        logger.error(error_message, job_context=job_context)
        job_context["job"].failure_reason = error_message
        job_context["success"] = False
        job_context["job"].no_retry = True
        return job_context
    except Exception as e:
        error_template = ("Encountered error in R code while running {0}"
                          " pipeline during processing of {1}: {2}")
        error_message = error_template.format(job_context['script_name'],
            job_context['input_file_path'], str(e))
        logger.error(error_message, job_context=job_context)
        job_context["job"].failure_reason = error_message
        job_context["success"] = False
        job_context["job"].no_retry = True
        return job_context

    # Quality control!
    # Related: https://github.com/AlexsLemonade/refinebio/issues/614
    # Related: GSM102671
    is_good_quality = check_output_quality(job_context['output_file_path'])
    if not is_good_quality:
        job_context["success"] = False
        job_context["job"].failure_reason = "NO_OP output failed quality control check."
        job_context["job"].no_retry = True
        return job_context
    else:
        job_context["success"] = True
        return job_context

def _convert_illumina_genes(job_context: Dict) -> Dict:
    """ Convert to Ensembl genes if we can"""

    all_databases = {
        'HOMO_SAPIENS': [
            'illuminaHumanv1',
            'illuminaHumanv2',
            'illuminaHumanv3',
            'illuminaHumanv4',
        ],
        'MUS_MUSCULUS': [
            'illuminaMousev1',
            'illuminaMousev1p1',
            'illuminaMousev2',
        ],
        'RATTUS_NORVEGICUS': [
            'illuminaRatv1'
        ]
    }

    sample0 = job_context['samples'][0]
    databases = all_databases[sample0.organism.name]

    # Loop over all of the possible platforms and find the one with the best match.
    highest = 0.0
    high_mapped_percent = 0.0
    high_db = None
    for platform in databases:
        try:
            result = subprocess.check_output([
                    "/usr/bin/Rscript",
                    "--vanilla",
                    "/home/user/data_refinery_workers/processors/detect_database.R",
                    "--platform", platform,
                    "--inputFile", job_context['input_file_path'],
                    "--column", job_context.get('column_name', "Reporter Identifier")
                ])

            results = result.decode().split('\n')
            cleaned_result = float(results[0].strip())

            if cleaned_result > highest:
                highest = cleaned_result
                high_db = platform
                high_mapped_percent = float(results[1].strip())

        except Exception as e:
            logger.exception("Could not detect database for file!",
                             platform=platform,
                             job_context=job_context)
            continue

    # Record our sample detection outputs for every sample.
    for sample in job_context['samples']:
        sa = SampleAnnotation()
        sa.sample = sample
        sa.data = {
            "detected_platform": high_db,
            "detection_percentage": highest,
            "mapped_percentage": high_mapped_percent
        }
        sa.save()

    job_context['script_name'] = "gene_convert_illumina.R"
    try:
        result = subprocess.check_output([
                "/usr/bin/Rscript",
                "--vanilla",
                "/home/user/data_refinery_workers/processors/" + job_context['script_name'],
                "--platform", high_db,
                "--inputFile", job_context['input_file_path'],
                "--outputFile", job_context['output_file_path']
            ], stderr=subprocess.PIPE)
    except subprocess.CalledProcessError as e:
        error_template = "Status code {0} from {1}: {2}"
        error_message = error_template.format(e.returncode, job_context['script_name'], e.stderr)
        logger.error(error_message, job_context=job_context)
        job_context["job"].failure_reason = error_message
        job_context["success"] = False
        job_context["job"].no_retry = True
        return job_context
    except Exception as e:
        error_template = ("Encountered error in R code while running {0}"
                          " pipeline during processing of {1}: {2}")
        error_message = error_template.format(job_context['script_name'],
            job_context['input_file_path'], str(e))
        logger.error(error_message, job_context=job_context)
        job_context["job"].failure_reason = error_message
        job_context["success"] = False
        job_context["job"].no_retry = True
        return job_context

    job_context["success"] = True
    return job_context

def _create_result(job_context: Dict) -> Dict:
    """ Create the actual Result object"""

    # This is a NO-OP, but we make a ComputationalResult regardless.
    result = ComputationalResult()
    result.commands.append(job_context['script_name'])
    result.is_ccdl = True
    try:
        processor_key = "SUBMITTER_PROCESSED"
        result.processor = utils.find_processor(processor_key)
    except Exception as e:
        return utils.handle_processor_exception(job_context, processor_key, e)

    result.save()
    job_context['pipeline'].steps.append(result.id)

    # Create a ComputedFile for the computed file,
    # sync it S3 and save it.
    computed_file = ComputedFile()
    computed_file.absolute_file_path = job_context["output_file_path"]
    computed_file.filename = job_context['output_file_path'].split('/')[-1]
    computed_file.calculate_sha1()
    computed_file.calculate_size()
    computed_file.result = result
    computed_file.is_smashable = True
    computed_file.is_qc = False
    computed_file.save()

    # utils.end_job will sync this to S3 for us.
    job_context["computed_files"] = [computed_file]

    for sample in job_context['samples']:
        assoc = SampleResultAssociation()
        assoc.sample = sample
        assoc.result = result
        assoc.save()

        SampleComputedFileAssociation.objects.get_or_create(
            sample=sample,
            computed_file=computed_file)

    logger.debug("Created %s", result)
    job_context["success"] = True
    return job_context

def check_output_quality(output_file_path: str):
    """ Verify our output file meets spec """
    try:
        with open(output_file_path, 'r') as tsv_in:
            tsv_in = csv.reader(tsv_in, delimiter='\t')
            for i, row in enumerate(tsv_in):
                # We could change this to a header column checker
                if i == 0:
                    continue

                # If there are more than 2 columns,
                # we consider this bad data. (Likely old machine/processing!)
                if len(row) > 2:
                    return False
    except Exception as e:
        return False

    return True

def no_op_processor(job_id: int) -> None:
    pipeline = Pipeline(name=utils.PipelineEnum.NO_OP.value)
    return utils.run_pipeline({"job_id": job_id, "pipeline": pipeline},
                              [utils.start_job,
                               _prepare_files,
                               _convert_genes,
                               _create_result,
                               utils.end_job])
<commit_msg>address comment<commit_after>import boto3
import csv
import os
import shutil
import subprocess

from typing import Dict

from data_refinery_common.logging import get_and_configure_logger
from data_refinery_common.models import (
    ComputationalResult,
    ComputedFile,
    Pipeline,
    Processor,
    SampleAnnotation,
    SampleComputedFileAssociation,
    SampleResultAssociation,
)
from data_refinery_common.utils import get_env_variable, get_internal_microarray_accession
from data_refinery_workers.processors import utils

logger = get_and_configure_logger(__name__)
LOCAL_ROOT_DIR = get_env_variable("LOCAL_ROOT_DIR", "/home/user/data_store")
S3_BUCKET_NAME = get_env_variable("S3_BUCKET_NAME", "data-refinery")

def _prepare_files(job_context: Dict) -> Dict:
    """A processor which takes externally-processed sample data and makes it smashable.
    """
    try:
        original_file = job_context["original_files"][0]
        sample0 = job_context['samples'][0]
        if sample0.manufacturer == 'ILLUMINA':
            job_context["is_illumina"] = True
        else:
            job_context["is_illumina"] = False

        # All files for the job are in the same directory.
        job_context["work_dir"] = LOCAL_ROOT_DIR + "/" + "processor_job_" + str(job_context["job_id"]) + "/"
        try:
            os.makedirs(job_context["work_dir"])
        except Exception as e:
            logger.exception("Could not create work directory for processor job.",
                             job_context=job_context)
            job_context["job"].failure_reason = str(e)
            job_context["success"] = False
            return job_context

        # Create the output directory and path
        job_context["input_file_path"] = original_file.absolute_file_path
        new_filename = "gene_converted_" + original_file.filename
        job_context["output_file_path"] = job_context["work_dir"] + new_filename

        try:
            # Make sure header row is correct
            with open(job_context["input_file_path"], 'r', encoding='utf-8') as tsv_in:
                tsv_in = csv.reader(tsv_in, delimiter='\t')
                for line in tsv_in:
                    row = line
                    joined = ''.join(row)
                    break
        except Exception as e:
            logger.exception("Expected text file, was this in fact a binary file?",
                             job_context=job_context,
                             file=job_context["input_file_path"])

            job_context["job"].failure_reason = str(e)
            job_context["job"].no_retry = True
            job_context["success"] = True
            return job_context

        # We want to make sure that all of our columns for conversion and smashing
        # use the same column name for gene identifiers for later lookup. ID_REF
        # is the most common, so we use that.
        if 'ID_REF' not in joined and not job_context["is_illumina"]:
            try:
                float(row[1])
                # Okay, there's no header so can just prepend to the file.
                with open(job_context["input_file_path"], 'r', encoding='utf-8') as f:
                    all_content = f.read()

                job_context["input_file_path"] = job_context["input_file_path"] + ".fixed"
                with open(job_context["input_file_path"], 'w+', encoding='utf-8') as f:
                    f.seek(0, 0)
                    f.write('ID_REF\tVALUE' + '\n' + all_content)
            except ValueError:
                # There is already a header row. Let's replace it with ID_Ref
                with open(job_context["input_file_path"], 'r', encoding='utf-8') as f:
                    all_content = f.read()

                job_context["input_file_path"] = job_context["input_file_path"] + ".fixed"
                with open(job_context["input_file_path"], 'w+', encoding='utf-8') as f:
                    f.seek(0, 0)
                    f.write('ID_REF\tVALUE' + '\n' + ("\n".join(all_content[1:])))
            except Exception as e:
                logger.exception("Unable to read input file or header row.",
                    input_file_path=job_context["input_file_path"])
                job_context['job'].failure_reason = str(e)
                job_context['success'] = False
                job_context["job"].no_retry = True
                return job_context
        else:
            if job_context["is_illumina"]:
                try:
                    float(row[1])
                    # If we're here, there's no header. We're gonna have to fudge it a little bit.

                    # These exist, but they're too hard for us to handle right now.
                    if len(row) > 3:
                        e_msg = "We found an Illumina file to NO_OP that we're not set up to process yet. Tell Rich!"
                        logger.error(e_msg,
                            job_id = job_context["job"].pk,
                            file=job_context["input_file_path"]
                        )
                        job_context['success'] = False
                        job_context["job"].failure_reason = str(e_msg)
                        job_context["job"].no_retry = True
                        return job_context

                    # Okay, there's no header so can just prepend to the file.
                    with open(job_context["input_file_path"], 'r', encoding='utf-8') as f:
                        all_content = f.read()
                    job_context["input_file_path"] = job_context["input_file_path"] + ".fixed"
                    with open(job_context["input_file_path"], 'w+', encoding='utf-8') as f:
                        f.seek(0, 0)
                        if len(row) == 2:
                            f.write('Reporter Identifier\tVALUE' + '\n' + all_content)
                        elif len(row) == 3:
                            f.write('Reporter Identifier\tVALUE\tDetection Pval' + '\n' + all_content)

                    job_context['column_name'] = 'Reporter Identifier'
                except ValueError:
                    # Okay, there's a header column, we're good.
                    job_context['column_name'] = row[0]

        # Platform
        job_context["platform"] = job_context["samples"][0].platform_accession_code
        job_context["internal_accession"] = get_internal_microarray_accession(job_context["platform"])
        if not job_context["internal_accession"]:
            logger.error("Failed to find internal accession for code", code=job_context["platform"])
            job_context["job"].failure_reason = "Failed to find internal accession for code" + job_context["platform"]
            job_context['success'] = False
            job_context["job"].no_retry = True
    except Exception as e:
        logger.exception("Failed to prepare for NO_OP job.", job_id=job_context['job'].id)
        job_context["job"].failure_reason = str(e)
        job_context['success'] = False
        job_context["job"].no_retry = True

    return job_context

def _convert_genes(job_context: Dict) -> Dict:
    """ Dispatches to the appropriate gene converter"""

    sample0 = job_context['samples'][0]
    if sample0.manufacturer == 'ILLUMINA':
        return _convert_illumina_genes(job_context)
    else:
        return _convert_affy_genes(job_context)

def _convert_affy_genes(job_context: Dict) -> Dict:
    """ Convert to Ensembl genes if we can"""

    if 'internal_accession' not in job_context or not job_context['internal_accession']:
        error_msg = "Told to convert AFFY genes without an internal_accession - how did this happen?"
        logger.error(error_msg, job_context=job_context)
        job_context["job"].failure_reason = str(error_msg)
        job_context['success'] = False
        job_context["job"].no_retry = True
        return job_context

    gene_index_path = "/home/user/gene_indexes/" + job_context["internal_accession"] + ".tsv.gz"
    if not os.path.exists(gene_index_path):
        logger.error("Missing gene index file for platform!",
            platform=job_context["internal_accession"],
            job_id=job_context["job_id"])
        job_context["job"].failure_reason = "Missing gene index for " + job_context['internal_accession']
        job_context["job"].no_retry = True
        job_context["success"] = False
        return job_context

    job_context['script_name'] = "gene_convert.R"
    try:
        result = subprocess.check_output([
            "/usr/bin/Rscript",
            "--vanilla", # Shut up Rscript
            "/home/user/data_refinery_workers/processors/" + job_context['script_name'],
            "--geneIndexPath", gene_index_path,
            "--inputFile", job_context['input_file_path'],
            "--outputFile", job_context['output_file_path']
        ], stderr=subprocess.PIPE)
    except subprocess.CalledProcessError as e:
        error_template = "Status code {0} from {1}: {2}"
        error_message = error_template.format(e.returncode, job_context['script_name'], e.stderr)
        logger.error(error_message, job_context=job_context)
        job_context["job"].failure_reason = error_message
        job_context["success"] = False
        job_context["job"].no_retry = True
        return job_context
    except Exception as e:
        error_template = ("Encountered error in R code while running {0}"
                          " pipeline during processing of {1}: {2}")
        error_message = error_template.format(job_context['script_name'],
            job_context['input_file_path'], str(e))
        logger.error(error_message, job_context=job_context)
        job_context["job"].failure_reason = error_message
        job_context["success"] = False
        job_context["job"].no_retry = True
        return job_context

    # Quality control!
    # Related: https://github.com/AlexsLemonade/refinebio/issues/614
    # Related: GSM102671
    is_good_quality = check_output_quality(job_context['output_file_path'])
    if not is_good_quality:
        job_context["success"] = False
        job_context["job"].failure_reason = "NO_OP output failed quality control check."
        job_context["job"].no_retry = True
        return job_context
    else:
        job_context["success"] = True
        return job_context

def _convert_illumina_genes(job_context: Dict) -> Dict:
    """ Convert to Ensembl genes if we can"""

    all_databases = {
        'HOMO_SAPIENS': [
            'illuminaHumanv1',
            'illuminaHumanv2',
            'illuminaHumanv3',
            'illuminaHumanv4',
        ],
        'MUS_MUSCULUS': [
            'illuminaMousev1',
            'illuminaMousev1p1',
            'illuminaMousev2',
        ],
        'RATTUS_NORVEGICUS': [
            'illuminaRatv1'
        ]
    }

    sample0 = job_context['samples'][0]
    databases = all_databases[sample0.organism.name]

    # Loop over all of the possible platforms and find the one with the best match.
    highest = 0.0
    high_mapped_percent = 0.0
    high_db = None
    for platform in databases:
        try:
            result = subprocess.check_output([
                    "/usr/bin/Rscript",
                    "--vanilla",
                    "/home/user/data_refinery_workers/processors/detect_database.R",
                    "--platform", platform,
                    "--inputFile", job_context['input_file_path'],
                    "--column", job_context.get('column_name', "Reporter Identifier")
                ])

            results = result.decode().split('\n')
            cleaned_result = float(results[0].strip())

            if cleaned_result > highest:
                highest = cleaned_result
                high_db = platform
                high_mapped_percent = float(results[1].strip())

        except Exception as e:
            logger.exception("Could not detect database for file!",
                             platform=platform,
                             job_context=job_context)
            continue

    # Record our sample detection outputs for every sample.
    for sample in job_context['samples']:
        sa = SampleAnnotation()
        sa.sample = sample
        sa.data = {
            "detected_platform": high_db,
            "detection_percentage": highest,
            "mapped_percentage": high_mapped_percent
        }
        sa.save()

    job_context['script_name'] = "gene_convert_illumina.R"
    try:
        result = subprocess.check_output([
                "/usr/bin/Rscript",
                "--vanilla",
                "/home/user/data_refinery_workers/processors/" + job_context['script_name'],
                "--platform", high_db,
                "--inputFile", job_context['input_file_path'],
                "--outputFile", job_context['output_file_path']
            ], stderr=subprocess.PIPE)
    except subprocess.CalledProcessError as e:
        error_template = "Status code {0} from {1}: {2}"
        error_message = error_template.format(e.returncode, job_context['script_name'], e.stderr)
        logger.error(error_message, job_context=job_context)
        job_context["job"].failure_reason = error_message
        job_context["success"] = False
        job_context["job"].no_retry = True
        return job_context
    except Exception as e:
        error_template = ("Encountered error in R code while running {0}"
                          " pipeline during processing of {1}: {2}")
        error_message = error_template.format(job_context['script_name'],
            job_context['input_file_path'], str(e))
        logger.error(error_message, job_context=job_context)
        job_context["job"].failure_reason = error_message
        job_context["success"] = False
        job_context["job"].no_retry = True
        return job_context

    job_context["success"] = True
    return job_context

def _create_result(job_context: Dict) -> Dict:
    """ Create the actual Result object"""

    # This is a NO-OP, but we make a ComputationalResult regardless.
    result = ComputationalResult()
    result.commands.append(job_context['script_name'])
    result.is_ccdl = True
    try:
        processor_key = "SUBMITTER_PROCESSED"
        result.processor = utils.find_processor(processor_key)
    except Exception as e:
        return utils.handle_processor_exception(job_context, processor_key, e)

    result.save()
    job_context['pipeline'].steps.append(result.id)

    # Create a ComputedFile for the computed file,
    # sync it S3 and save it.
    computed_file = ComputedFile()
    computed_file.absolute_file_path = job_context["output_file_path"]
    computed_file.filename = job_context['output_file_path'].split('/')[-1]
    computed_file.calculate_sha1()
    computed_file.calculate_size()
    computed_file.result = result
    computed_file.is_smashable = True
    computed_file.is_qc = False
    computed_file.save()

    # utils.end_job will sync this to S3 for us.
    job_context["computed_files"] = [computed_file]

    for sample in job_context['samples']:
        assoc = SampleResultAssociation()
        assoc.sample = sample
        assoc.result = result
        assoc.save()

        SampleComputedFileAssociation.objects.get_or_create(
            sample=sample,
            computed_file=computed_file)

    logger.debug("Created %s", result)
    job_context["success"] = True
    return job_context

def check_output_quality(output_file_path: str):
    """ Verify our output file meets spec """
    try:
        with open(output_file_path, 'r') as tsv_in:
            tsv_in = csv.reader(tsv_in, delimiter='\t')
            for i, row in enumerate(tsv_in):
                # We could change this to a header column checker
                if i == 0:
                    continue

                # If there are more than 2 columns,
                # we consider this bad data. (Likely old machine/processing!)
                if len(row) > 2:
                    return False
    except Exception as e:
        return False

    return True

def no_op_processor(job_id: int) -> None:
    pipeline = Pipeline(name=utils.PipelineEnum.NO_OP.value)
    return utils.run_pipeline({"job_id": job_id, "pipeline": pipeline},
                              [utils.start_job,
                               _prepare_files,
                               _convert_genes,
                               _create_result,
                               utils.end_job])
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
新版本打包自动化
'''

import argparse, codecs, importlib, os, re, time
import plib.utils as utils
import plib.git as git
from plib.environment import get_current_packet_id, get_packet_path, set_packet_as_cwd
from plib.language.converter import Converter
import plib.environment as env

def __compress(addon):
	'''
	Compress and concat addon source into one file.

	Args:
		addon: Addon name
	'''
	print('--------------------------------')
	print('Compressing: %s' % addon)
	file_count = 0
	converter = Converter('zh-TW')
	# Remove debug codes in source
	for line in open('%s/info.ini' % addon):
		parts = line.strip().split('=')
		if parts[0].find('lua_') == 0:
			source_file = os.path.join(addon, parts[1])
			source_code = codecs.open(source_file,'r',encoding='gbk').read()
			source_code = re.sub(r'(?is)[^\n]*--\[\[#DEBUG LINE\]\][^\n]*\n?', '', source_code)
			source_code = re.sub(r'(?is)\n?--\[\[#DEBUG BEGIN\]\].*?--\[\[#DEBUG END\]\]\n?', '', source_code)
			codecs.open(source_file,'w',encoding='gbk').write(source_code)
	# Generate squishy file and execute squish
	with open('squishy', 'w') as squishy:
		squishy.write('Output "./%s/src.lua"\n' % addon)
		for line in open('%s/info.ini' % addon):
			parts = line.strip().split('=')
			if parts[0].find('lua_') == 0:
				if parts[1] == 'src.lua':
					print('Already compressed...')
					return
				file_path = os.path.join('.', addon, parts[1]).replace('\\', '/')
				file_count = file_count + 1
				squishy.write('Module "%d" "%s"\n' % (file_count, file_path))
	os.popen('lua "./!src-dist/tools/react/squish" --minify-level=full').read()
	os.remove('squishy')
	# Modify dist file for loading modules
	with open('./%s/src.lua' % addon, 'r+') as src:
		content = src.read()
		src.seek(0, 0)
		src.write('local package={preload={}}\n' + content)
	with open('./%s/src.lua' % addon, 'a') as src:
		src.write('\nfor _, k in ipairs({')
		for i in range(1, file_count + 1):
			src.write('\'%d\',' % i)
		src.write('}) do package.preload[k]() end')
	print('Compress done...')
	# Modify info.ini file
	info_content = ''
	for _, line in enumerate(codecs.open('%s/info.ini' % addon,'r',encoding='gbk')):
		parts = line.split('=')
		if parts[0].find('lua_') == 0:
			if parts[0] == 'lua_0':
				info_content = info_content + 'lua_0=src.lua\n'
		else:
			info_content = info_content + line
	with codecs.open('%s/info.ini' % addon,'w',encoding='gbk') as f:
		f.write(info_content)
	with codecs.open('%s/info.ini.zh_TW' % addon,'w',encoding='utf8') as f:
		f.write(converter.convert(info_content))
	print('Update info done...')

def __get_version_info():
	'''Get version information'''
	# Read version from Base.lua
	current_version = 0
	for line in open('%s_!Base/src/lib/Base.lua' % get_current_packet_id()):
		if line[6:23] == '_NATURAL_VERSION_':
			current_version = int(line.split(' ')[-1])
	# Read max and previous release commit
	commit_list = os.popen('git log --grep Release --pretty=format:"%s|%p"').read().split('\n')
	max_version, prev_version, prev_version_message, prev_version_hash = 0, 0, '', ''
	for commit in commit_list:
		try:
			info = commit.split('|')
			version = int(info[0][9:])
			if version < current_version and version > prev_version:
				prev_version = version
				prev_version_message = info[0]
				prev_version_hash = info[1]
			if version > max_version:
				max_version = version
		except:
			pass
	return { 'current': current_version, 'max': max_version, 'previous': prev_version, 'previous_message': prev_version_message, 'previous_hash': prev_version_hash }

def __7zip(file_name, base_message, base_hash):
	cmd_suffix = ''
	if base_hash != '':
		# Generate file change list since previous release commit
		def pathToModule(path):
			return re.sub('(?:^\\!src-dist/data/|["/].*$)', '', path)
		paths = {
			'package.ini': True,
			'package.ini.*': True,
		}
		print('File change list:')
		filelist = os.popen('git diff ' + base_hash + ' HEAD --name-status').read().strip().split('\n')
		for file in filelist:
			lst = file.split('\t')
			if lst[0] == 'A' or lst[0] == 'M' or lst[0] == 'D':
				paths[pathToModule(lst[1])] = True
			elif lst[0][0] == 'R':
				paths[pathToModule(lst[1])] = True
				paths[pathToModule(lst[2])] = True
			print(file)
		print('')
		# Print addon change list
		print('Subpath change list:')
		for path in paths:
			print('/' + path)
			cmd_suffix = cmd_suffix + ' "' + path + '"'
		print('')

	# Prepare for 7z compressing
	print('zippping...')
	os.system('start /wait /b ./!src-dist/bin/7z.exe a -t7z ' + file_name + ' -xr!manifest.dat -xr!manifest.key -xr!publisher.key -x@.7zipignore' + cmd_suffix)
	print('File(s) compressing acomplished!')
	print('Url: ' + file_name)
	print('Based on git commit "%s(%s)".' % (base_message, base_hash) if base_hash != '' else 'Full package.')

def run(is_full, is_source):
	print('> RELEASE MODE: %s, %s.' % ('full' if is_full else 'diff', 'source' if is_source else 'dist'))
	version_info = __get_version_info()

	if not is_source:
		# Merge master into prelease
		if git.get_current_branch() == 'master':
			utils.assert_exit(git.is_clean(), 'Error: master branch has uncommited file change(s)!')
			os.system('git checkout prelease || git checkout -b prelease')
			os.system('git rebase master')
			os.system('code "./%s_!Base/src/lib/Base.lua"' % get_current_packet_id())
			os.system('code "./CHANGELOG.md"')
			utils.exit_with_message('Switched to prelease branch. Please commit release info and then run this script again!')

		# Merge prelease into stable
		if git.get_current_branch() == 'prelease':
			os.system('git reset master')
			version_info = __get_version_info()
			utils.assert_exit(version_info.get('current') > version_info.get('max'),
				'Error: current version(%s) must be larger than max history version(%d)!' % (version_info.get('current'), version_info.get('max')))
			os.system('git add * && git commit -m "Release V%s"' % version_info.get('current'))
			os.system('git checkout stable || git checkout -b stable')
			os.system('git reset origin/stable --hard')
			os.system('git rebase prelease')

		# Check if branch
		utils.assert_exit(git.is_clean(), 'Error: resolve conflict and remove uncommited changes first!')
		utils.assert_exit(git.get_current_branch() == 'stable', 'Error: current branch is not on stable!')

	# Compress and concat source file
	for addon in os.listdir('./'):
		if os.path.exists(os.path.join('./', addon, 'info.ini')):
			__compress(addon)

	# Package files
	file_name = '!src-dist/dist/%s_%s_v%s.7z' % (get_current_packet_id(), time.strftime('%Y%m%d%H%M%S', time.localtime()), version_info.get('current'))
	base_message = ''
	base_hash = ''
	if not is_full and version_info.get('current') != '' and version_info.get('previous_hash') != '':
		base_message = version_info.get('previous_message')
		base_hash = version_info.get('previous_hash')
	__7zip(file_name, base_message, base_hash)

	# Revert source code modify by compressing
	if not is_source:
		os.system('git reset HEAD --hard')
		os.system('git checkout master')
	time.sleep(5)
	print('Exiting...')

if __name__ == '__main__':
	parser = argparse.ArgumentParser(description='One-key release packet product helper.')
	parser.add_argument('--full', action='store_true', help='Package full plugin packet.')
	parser.add_argument('--no-build', action='store_true', help='Package source code.')
	args = parser.parse_args()

	env.set_packet_as_cwd()

	run(args.full, args.no_build)
<commit_msg>fix: 修复打包工具相对路径问题<commit_after># -*- coding: utf-8 -*-
'''
新版本打包自动化
'''

import argparse, codecs, importlib, os, re, time
import plib.utils as utils
import plib.git as git
from plib.environment import get_current_packet_id, get_packet_path, set_packet_as_cwd
from plib.language.converter import Converter
import plib.environment as env

def __compress(addon):
	'''
	Compress and concat addon source into one file.

	Args:
		addon: Addon name
	'''
	print('--------------------------------')
	print('Compressing: %s' % addon)
	file_count = 0
	converter = Converter('zh-TW')
	# Remove debug codes in source
	for line in open('%s/info.ini' % addon):
		parts = line.strip().split('=')
		if parts[0].find('lua_') == 0:
			source_file = os.path.join(addon, parts[1])
			source_code = codecs.open(source_file,'r',encoding='gbk').read()
			source_code = re.sub(r'(?is)[^\n]*--\[\[#DEBUG LINE\]\][^\n]*\n?', '', source_code)
			source_code = re.sub(r'(?is)\n?--\[\[#DEBUG BEGIN\]\].*?--\[\[#DEBUG END\]\]\n?', '', source_code)
			codecs.open(source_file,'w',encoding='gbk').write(source_code)
	# Generate squishy file and execute squish
	with open('squishy', 'w') as squishy:
		squishy.write('Output "./%s/src.lua"\n' % addon)
		for line in open('%s/info.ini' % addon):
			parts = line.strip().split('=')
			if parts[0].find('lua_') == 0:
				if parts[1] == 'src.lua':
					print('Already compressed...')
					return
				file_path = os.path.join('.', addon, parts[1]).replace('\\', '/')
				file_count = file_count + 1
				squishy.write('Module "%d" "%s"\n' % (file_count, file_path))
	os.popen('lua "./!src-dist/tools/react/squish" --minify-level=full').read()
	os.remove('squishy')
	# Modify dist file for loading modules
	with open('./%s/src.lua' % addon, 'r+') as src:
		content = src.read()
		src.seek(0, 0)
		src.write('local package={preload={}}\n' + content)
	with open('./%s/src.lua' % addon, 'a') as src:
		src.write('\nfor _, k in ipairs({')
		for i in range(1, file_count + 1):
			src.write('\'%d\',' % i)
		src.write('}) do package.preload[k]() end')
	print('Compress done...')
	# Modify info.ini file
	info_content = ''
	for _, line in enumerate(codecs.open('%s/info.ini' % addon,'r',encoding='gbk')):
		parts = line.split('=')
		if parts[0].find('lua_') == 0:
			if parts[0] == 'lua_0':
				info_content = info_content + 'lua_0=src.lua\n'
		else:
			info_content = info_content + line
	with codecs.open('%s/info.ini' % addon,'w',encoding='gbk') as f:
		f.write(info_content)
	with codecs.open('%s/info.ini.zh_TW' % addon,'w',encoding='utf8') as f:
		f.write(converter.convert(info_content))
	print('Update info done...')

def __get_version_info():
	'''Get version information'''
	# Read version from Base.lua
	current_version = 0
	for line in open('%s_!Base/src/lib/Base.lua' % get_current_packet_id()):
		if line[6:23] == '_NATURAL_VERSION_':
			current_version = int(line.split(' ')[-1])
	# Read max and previous release commit
	commit_list = os.popen('git log --grep Release --pretty=format:"%s|%p"').read().split('\n')
	max_version, prev_version, prev_version_message, prev_version_hash = 0, 0, '', ''
	for commit in commit_list:
		try:
			info = commit.split('|')
			version = int(info[0][9:])
			if version < current_version and version > prev_version:
				prev_version = version
				prev_version_message = info[0]
				prev_version_hash = info[1]
			if version > max_version:
				max_version = version
		except:
			pass
	return { 'current': current_version, 'max': max_version, 'previous': prev_version, 'previous_message': prev_version_message, 'previous_hash': prev_version_hash }

def __7zip(file_name, base_message, base_hash):
	cmd_suffix = ''
	if base_hash != '':
		# Generate file change list since previous release commit
		def pathToModule(path):
			return re.sub('(?:^\\!src-dist/data/|["/].*$)', '', path)
		paths = {
			'package.ini': True,
			'package.ini.*': True,
		}
		print('File change list:')
		filelist = os.popen('git diff ' + base_hash + ' HEAD --name-status').read().strip().split('\n')
		for file in filelist:
			lst = file.split('\t')
			if lst[0] == 'A' or lst[0] == 'M' or lst[0] == 'D':
				paths[pathToModule(lst[1])] = True
			elif lst[0][0] == 'R':
				paths[pathToModule(lst[1])] = True
				paths[pathToModule(lst[2])] = True
			print(file)
		print('')
		# Print addon change list
		print('Subpath change list:')
		for path in paths:
			print('/' + path)
			cmd_suffix = cmd_suffix + ' "' + path + '"'
		print('')

	# Prepare for 7z compressing
	print('zippping...')
	os.system('start /wait /b ./!src-dist/bin/7z.exe a -t7z ' + file_name + ' -xr!manifest.dat -xr!manifest.key -xr!publisher.key -x@.7zipignore' + cmd_suffix)
	print('File(s) compressing acomplished!')
	print('Url: ' + file_name)
	print('Based on git commit "%s(%s)".' % (base_message, base_hash) if base_hash != '' else 'Full package.')

def run(is_full, is_source):
	print('> RELEASE MODE: %s, %s.' % ('full' if is_full else 'diff', 'source' if is_source else 'dist'))
	version_info = __get_version_info()

	if not is_source:
		# Merge master into prelease
		if git.get_current_branch() == 'master':
			utils.assert_exit(git.is_clean(), 'Error: master branch has uncommited file change(s)!')
			os.system('git checkout prelease || git checkout -b prelease')
			os.system('git rebase master')
			os.system('code "%s"' % os.path.join(get_packet_path(), './%s_!Base/src/lib/Base.lua' % get_current_packet_id()))
			os.system('code "%s"' % os.path.join(get_packet_path(), './CHANGELOG.md'))
			utils.exit_with_message('Switched to prelease branch. Please commit release info and then run this script again!')

		# Merge prelease into stable
		if git.get_current_branch() == 'prelease':
			os.system('git reset master')
			version_info = __get_version_info()
			utils.assert_exit(version_info.get('current') > version_info.get('max'),
				'Error: current version(%s) must be larger than max history version(%d)!' % (version_info.get('current'), version_info.get('max')))
			os.system('git add * && git commit -m "Release V%s"' % version_info.get('current'))
			os.system('git checkout stable || git checkout -b stable')
			os.system('git reset origin/stable --hard')
			os.system('git rebase prelease')

		# Check if branch
		utils.assert_exit(git.is_clean(), 'Error: resolve conflict and remove uncommited changes first!')
		utils.assert_exit(git.get_current_branch() == 'stable', 'Error: current branch is not on stable!')

	# Compress and concat source file
	for addon in os.listdir('./'):
		if os.path.exists(os.path.join('./', addon, 'info.ini')):
			__compress(addon)

	# Package files
	file_name = '!src-dist/dist/%s_%s_v%s.7z' % (get_current_packet_id(), time.strftime('%Y%m%d%H%M%S', time.localtime()), version_info.get('current'))
	base_message = ''
	base_hash = ''
	if not is_full and version_info.get('current') != '' and version_info.get('previous_hash') != '':
		base_message = version_info.get('previous_message')
		base_hash = version_info.get('previous_hash')
	__7zip(file_name, base_message, base_hash)

	# Revert source code modify by compressing
	if not is_source:
		os.system('git reset HEAD --hard')
		os.system('git checkout master')
	time.sleep(5)
	print('Exiting...')

if __name__ == '__main__':
	parser = argparse.ArgumentParser(description='One-key release packet product helper.')
	parser.add_argument('--full', action='store_true', help='Package full plugin packet.')
	parser.add_argument('--no-build', action='store_true', help='Package source code.')
	args = parser.parse_args()

	env.set_packet_as_cwd()

	run(args.full, args.no_build)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
skeppa
------------
:copyright: (c) 2016 Martin Sandström
:license: MIT, see LICENSE for more details.
"""

__title__ = "skeppa"
__version__ = "1.3.0"
__build__ = 130
__license__ = "MIT"
__copyright__ = "Copyright 2016 Martin Sandström"
<commit_msg>Bumped version<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
skeppa
------------
:copyright: (c) 2016 Martin Sandström
:license: MIT, see LICENSE for more details.
"""

__title__ = "skeppa"
__version__ = "1.4.0"
__build__ = 140
__license__ = "MIT"
__copyright__ = "Copyright 2016 Martin Sandström"
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import nested_scopes
from __future__ import print_function

import logging
import pickle
import select
import socket
import struct
import threading
import time

from hops import util

MAX_RETRIES = 3
BUFSIZE = 1024*2

class Reservations:
  """Thread-safe store for node reservations.

  Args:
    :required: expected number of nodes in the cluster.
  """

  def __init__(self, required):
    self.required = required
    self.lock = threading.RLock()
    self.reservations = []
    self.cluster_spec = {}
    self.check_done = False

  def add(self, meta):
    """Add a reservation.

    Args:
      :meta: a dictonary of metadata about a node
    """
    with self.lock:
      self.reservations.append(meta)

      if self.remaining() == 0:
        gpus_present = False

        for entry in self.reservations:
          if entry["gpus_present"] == True:
            gpus_present = True
            break

        cluster_spec = {"chief": [], "ps": [], "worker": []}
        if not gpus_present:
          added_chief=False
          for entry in self.reservations:
            if entry["task_type"] == "ps":
              cluster_spec["ps"].append(entry["host_port"])
            elif added_chief == False and entry["task_type"] == "worker":
              cluster_spec["chief"].append(entry["host_port"])
              added_chief = True
            else:
              cluster_spec["worker"].append(entry["host_port"])
        else:
          added_chief=False
          for entry in self.reservations:
            if entry["task_type"] == "worker" and entry["gpus_present"] == True and added_chief == False:
              added_chief=True
              cluster_spec["chief"].append(entry["host_port"])
            elif entry["task_type"] == "ps" and entry["gpus_present"] == True:
              for possible_entry_to_switch in self.reservations:
                if possible_entry_to_switch["task_type"] == "worker" and possible_entry_to_switch["gpus_present"] == False:
                  entry["task_type"] = "worker"
                  possible_entry_to_switch["task_type"] = "ps"
                  break
              cluster_spec[entry["task_type"]].append(entry["host_port"])
            else:
              cluster_spec[entry["task_type"]].append(entry["host_port"])
        self.cluster_spec = cluster_spec

        self.check_done = True

  def done(self):
    """Returns True if the ``required`` number of reservations have been fulfilled."""
    with self.lock:
      return self.check_done

  def get(self):
    """Get the list of current reservations."""
    with self.lock:
      return self.cluster_spec

  def remaining(self):
    """Get a count of remaining/unfulfilled reservations."""
    with self.lock:
      num_registered = len(self.reservations)
      return self.required - num_registered

class MessageSocket(object):
  """Abstract class w/ length-prefixed socket send/receive functions."""

  def receive(self, sock):
    """Receive a message on ``sock``."""
    msg = None
    data = b''
    recv_done = False
    recv_len = -1
    while not recv_done:
      buf = sock.recv(BUFSIZE)
      if buf is None or len(buf) == 0:
        raise Exception("socket closed")
      if recv_len == -1:
        recv_len = struct.unpack('>I', buf[:4])[0]
        data += buf[4:]
        recv_len -= len(data)
      else:
        data += buf
        recv_len -= len(buf)
      recv_done = (recv_len == 0)

    msg = pickle.loads(data)
    return msg

  def send(self, sock, msg):
    """Send ``msg`` to destination ``sock``."""
    data = pickle.dumps(msg)
    buf = struct.pack('>I', len(data)) + data
    sock.sendall(buf)

class Server(MessageSocket):
  """Simple socket server with length prefixed pickle messages"""
  reservations = None
  done = False

  def __init__(self, count):
    assert count > 0
    self.reservations = Reservations(count)

  def await_reservations(self, sc, status={}, timeout=600):
    """Block until all reservations are received."""
    timespent = 0
    while not self.reservations.done():
      logging.info("waiting for {0} reservations".format(self.reservations.remaining()))
      # check status flags for any errors
      if 'error' in status:
        sc.cancelAllJobs()
        #sc.stop()
        #sys.exit(1)
      time.sleep(1)
      timespent += 1
      if (timespent > timeout):
        raise Exception("timed out waiting for reservations to complete")
    logging.info("all reservations completed")
    return self.reservations.get()

  def _handle_message(self, sock, msg):
    logging.debug("received: {0}".format(msg))
    msg_type = msg['type']
    if msg_type == 'REG':
      self.reservations.add(msg['data'])
      MessageSocket.send(self, sock, 'OK')
    elif msg_type == 'QUERY':
      MessageSocket.send(self, sock, self.reservations.done())
    elif msg_type == 'QINFO':
      rinfo = self.reservations.get()
      MessageSocket.send(self, sock, rinfo)
    elif msg_type == 'STOP':
      logging.info("setting server.done")
      MessageSocket.send(self, sock, 'OK')
      self.done = True
    else:
      MessageSocket.send(self, sock, 'ERR')

  def start(self):
    """Start listener in a background thread

    Returns:
      address of the Server as a tuple of (host, port)
    """
    server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_sock.bind(('', 0))
    server_sock.listen(10)

    # hostname may not be resolvable but IP address probably will be
    host = util.get_ip_address()
    port = server_sock.getsockname()[1]
    addr = (host,port)

    def _listen(self, sock):
      CONNECTIONS = []
      CONNECTIONS.append(sock)

      while not self.done:
        read_socks, write_socks, err_socks = select.select(CONNECTIONS, [], [], 60)
        for sock in read_socks:
          if sock == server_sock:
            client_sock, client_addr = sock.accept()
            CONNECTIONS.append(client_sock)
            logging.debug("client connected from {0}".format(client_addr))
          else:
            try:
              msg = self.receive(sock)
              self._handle_message(sock, msg)
            except Exception as e:
              logging.debug(e)
              sock.close()
              CONNECTIONS.remove(sock)

      server_sock.close()

    t = threading.Thread(target=_listen, args=(self, server_sock))
    t.daemon = True
    t.start()

    return addr

  def stop(self):
    """Stop the Server's socket listener."""
    self.done = True

class Client(MessageSocket):
  """Client to register and await node reservations.

  Args:
    :server_addr: a tuple of (host, port) pointing to the Server.
  """
  sock = None                   #: socket to server TCP connection
  server_addr = None            #: address of server

  def __init__(self, server_addr):
    self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    self.sock.connect(server_addr)
    self.server_addr = server_addr
    logging.info("connected to server at {0}".format(server_addr))

  def _request(self, msg_type, msg_data=None):
    """Helper function to wrap msg w/ msg_type."""
    msg = {}
    msg['type'] = msg_type
    if msg_data or ((msg_data == True) or (msg_data == False)):
      msg['data'] = msg_data

    done = False
    tries = 0
    while not done and tries < MAX_RETRIES:
      try:
        MessageSocket.send(self, self.sock, msg)
        done = True
      except socket.error as e:
        tries += 1
        if tries >= MAX_RETRIES:
          raise
        print("Socket error: {}".format(e))
        self.sock.close()
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect(self.server_addr)

    logging.debug("sent: {0}".format(msg))
    resp = MessageSocket.receive(self, self.sock)
    logging.debug("received: {0}".format(resp))
    return resp

  def close(self):
    """Close the client socket."""
    self.sock.close()

  def register(self, reservation):
    """Register ``reservation`` with server."""
    resp = self._request('REG', reservation)
    return resp

  def get_reservations(self):
    """Get current list of reservations."""
    cluster_info = self._request('QINFO')
    return cluster_info

  def await_reservations(self):
    """Poll until all reservations completed, then return cluster_info."""
    done = False
    while not done:
      done = self._request('QUERY')
      time.sleep(1)
    reservations = self.get_reservations()
    return reservations

  def request_stop(self):
    """Request server stop."""
    resp = self._request('STOP')
    return resp
<commit_msg>ParameterServerStrategy fixes<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import nested_scopes
from __future__ import print_function

import logging
import pickle
import select
import socket
import struct
import threading
import time

from hops import util

MAX_RETRIES = 3
BUFSIZE = 1024*2

class Reservations:
  """Thread-safe store for node reservations.

  Args:
    :required: expected number of nodes in the cluster.
  """

  def __init__(self, required):
    self.required = required
    self.lock = threading.RLock()
    self.reservations = []
    self.cluster_spec = {}
    self.check_done = False

  def add(self, meta):
    """Add a reservation.

    Args:
      :meta: a dictonary of metadata about a node
    """
    with self.lock:
      self.reservations.append(meta)

      if self.remaining() == 0:
        gpus_present = False

        for entry in self.reservations:
          if entry["gpus_present"] == True:
            gpus_present = True
            break

        cluster_spec = {"chief": [], "ps": [], "worker": []}
        if not gpus_present:
          added_chief=False
          for entry in self.reservations:
            if entry["task_type"] == "ps":
              cluster_spec["ps"].append(entry["host_port"])
            elif added_chief == False and entry["task_type"] == "worker":
              cluster_spec["chief"].append(entry["host_port"])
              added_chief = True
            else:
              cluster_spec["worker"].append(entry["host_port"])
        else:
          added_chief=False

          # switch Worker without GPU with PS with GPU
          for possible_switch in self.reservations:
            if possible_switch["task_type"] == "worker" and possible_switch["gpus_present"] == False:
              for candidate in self.reservations:
                if candidate["task_type"] == "ps" and candidate["gpus_present"] == True:
                  candidate["task_type"] = "worker"
                  possible_switch["task_type"] = "ps"
                  break

          for entry in self.reservations:
            if entry["task_type"] == "worker" and entry["gpus_present"] == True and added_chief == False:
              added_chief=True
              cluster_spec["chief"].append(entry["host_port"])
            elif entry["task_type"] == "worker" and entry["gpus_present"] == True:
              cluster_spec["worker"].append(entry["host_port"])
            elif entry["task_type"] == "ps" and entry["gpus_present"] == False:
              cluster_spec["ps"].append(entry["host_port"])

        self.cluster_spec = cluster_spec

        self.check_done = True

  def done(self):
    """Returns True if the ``required`` number of reservations have been fulfilled."""
    with self.lock:
      return self.check_done

  def get(self):
    """Get the list of current reservations."""
    with self.lock:
      return self.cluster_spec

  def remaining(self):
    """Get a count of remaining/unfulfilled reservations."""
    with self.lock:
      num_registered = len(self.reservations)
      return self.required - num_registered

class MessageSocket(object):
  """Abstract class w/ length-prefixed socket send/receive functions."""

  def receive(self, sock):
    """Receive a message on ``sock``."""
    msg = None
    data = b''
    recv_done = False
    recv_len = -1
    while not recv_done:
      buf = sock.recv(BUFSIZE)
      if buf is None or len(buf) == 0:
        raise Exception("socket closed")
      if recv_len == -1:
        recv_len = struct.unpack('>I', buf[:4])[0]
        data += buf[4:]
        recv_len -= len(data)
      else:
        data += buf
        recv_len -= len(buf)
      recv_done = (recv_len == 0)

    msg = pickle.loads(data)
    return msg

  def send(self, sock, msg):
    """Send ``msg`` to destination ``sock``."""
    data = pickle.dumps(msg)
    buf = struct.pack('>I', len(data)) + data
    sock.sendall(buf)

class Server(MessageSocket):
  """Simple socket server with length prefixed pickle messages"""
  reservations = None
  done = False

  def __init__(self, count):
    assert count > 0
    self.reservations = Reservations(count)

  def await_reservations(self, sc, status={}, timeout=600):
    """Block until all reservations are received."""
    timespent = 0
    while not self.reservations.done():
      logging.info("waiting for {0} reservations".format(self.reservations.remaining()))
      # check status flags for any errors
      if 'error' in status:
        sc.cancelAllJobs()
        #sc.stop()
        #sys.exit(1)
      time.sleep(1)
      timespent += 1
      if (timespent > timeout):
        raise Exception("timed out waiting for reservations to complete")
    logging.info("all reservations completed")
    return self.reservations.get()

  def _handle_message(self, sock, msg):
    logging.debug("received: {0}".format(msg))
    msg_type = msg['type']
    if msg_type == 'REG':
      self.reservations.add(msg['data'])
      MessageSocket.send(self, sock, 'OK')
    elif msg_type == 'QUERY':
      MessageSocket.send(self, sock, self.reservations.done())
    elif msg_type == 'QINFO':
      rinfo = self.reservations.get()
      MessageSocket.send(self, sock, rinfo)
    elif msg_type == 'STOP':
      logging.info("setting server.done")
      MessageSocket.send(self, sock, 'OK')
      self.done = True
    else:
      MessageSocket.send(self, sock, 'ERR')

  def start(self):
    """Start listener in a background thread

    Returns:
      address of the Server as a tuple of (host, port)
    """
    server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_sock.bind(('', 0))
    server_sock.listen(10)

    # hostname may not be resolvable but IP address probably will be
    host = util.get_ip_address()
    port = server_sock.getsockname()[1]
    addr = (host,port)

    def _listen(self, sock):
      CONNECTIONS = []
      CONNECTIONS.append(sock)

      while not self.done:
        read_socks, write_socks, err_socks = select.select(CONNECTIONS, [], [], 60)
        for sock in read_socks:
          if sock == server_sock:
            client_sock, client_addr = sock.accept()
            CONNECTIONS.append(client_sock)
            logging.debug("client connected from {0}".format(client_addr))
          else:
            try:
              msg = self.receive(sock)
              self._handle_message(sock, msg)
            except Exception as e:
              logging.debug(e)
              sock.close()
              CONNECTIONS.remove(sock)

      server_sock.close()

    t = threading.Thread(target=_listen, args=(self, server_sock))
    t.daemon = True
    t.start()

    return addr

  def stop(self):
    """Stop the Server's socket listener."""
    self.done = True

class Client(MessageSocket):
  """Client to register and await node reservations.

  Args:
    :server_addr: a tuple of (host, port) pointing to the Server.
  """
  sock = None                   #: socket to server TCP connection
  server_addr = None            #: address of server

  def __init__(self, server_addr):
    self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    self.sock.connect(server_addr)
    self.server_addr = server_addr
    logging.info("connected to server at {0}".format(server_addr))

  def _request(self, msg_type, msg_data=None):
    """Helper function to wrap msg w/ msg_type."""
    msg = {}
    msg['type'] = msg_type
    if msg_data or ((msg_data == True) or (msg_data == False)):
      msg['data'] = msg_data

    done = False
    tries = 0
    while not done and tries < MAX_RETRIES:
      try:
        MessageSocket.send(self, self.sock, msg)
        done = True
      except socket.error as e:
        tries += 1
        if tries >= MAX_RETRIES:
          raise
        print("Socket error: {}".format(e))
        self.sock.close()
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect(self.server_addr)

    logging.debug("sent: {0}".format(msg))
    resp = MessageSocket.receive(self, self.sock)
    logging.debug("received: {0}".format(resp))
    return resp

  def close(self):
    """Close the client socket."""
    self.sock.close()

  def register(self, reservation):
    """Register ``reservation`` with server."""
    resp = self._request('REG', reservation)
    return resp

  def get_reservations(self):
    """Get current list of reservations."""
    cluster_info = self._request('QINFO')
    return cluster_info

  def await_reservations(self):
    """Poll until all reservations completed, then return cluster_info."""
    done = False
    while not done:
      done = self._request('QUERY')
      time.sleep(1)
    reservations = self.get_reservations()
    return reservations

  def request_stop(self):
    """Request server stop."""
    resp = self._request('STOP')
    return resp
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
##
## This file is part of CDS Invenio.
## Copyright (C) 2002, 2003, 2004, 2005, 2006, 2007, 2008 CERN.
##
## CDS Invenio is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 2 of the
## License, or (at your option) any later version.
##
## CDS Invenio is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with CDS Invenio; if not, write to the Free Software Foundation, Inc.,
## 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.
"""
WebJournal widget - Display Indico seminars
"""
from invenio.config import CFG_CACHEDIR, CFG_SITE_LANG
from urllib2 import urlopen
from xml.dom import minidom
import time
import socket

from invenio.webjournal_utils import \
     parse_url_string
from invenio.messages import gettext_set_language

update_frequency = 3600 # in seconds

def format(bfo, indico_seminar_xml="http://indico.cern.ch/tools/export.py?fid=1l7&date=today&days=1&of=xml"):
    """
    Display the list of seminar from the given Indico XML URL

    @param indico_seminar_xml: the URL to the XML generated by an Indico instance
    """
    args = parse_url_string(bfo.user_info['uri'])
    journal_name = args["journal_name"]
    cached_filename = "webjournal_widget_seminars_%s.xml" % journal_name
    out = get_widget_html(bfo, indico_seminar_xml, cached_filename, args['ln'])
    return out

def escape_values(bfo):
    """
    Called by BibFormat in order to check if output of this element
    should be escaped.
    """
    return 0

def get_widget_html(bfo, indico_seminar_xml, cached_filename,
                    ln=CFG_SITE_LANG):
    """
    Indico seminars of the day service
    Gets seminars of the day from CERN Indico every 60 minutes and displays
    them in a widget.
    """
    _ = gettext_set_language(ln)
    try:
        seminar_xml = minidom.parse('%s/%s' % (CFG_CACHEDIR, cached_filename))
    except:
        try:
            _update_seminars(indico_seminar_xml, cached_filename)
            seminar_xml = minidom.parse('%s/%s' % (CFG_CACHEDIR, cached_filename))
        except:
            return "

" + _("No information available") + "

"
    try:
        timestamp = seminar_xml.firstChild.getAttribute("time")
    except:
        timestamp = time.struct_time()

    last_update = time.mktime(time.strptime(timestamp,
                                            "%a, %d %b %Y %H:%M:%S %Z"))
    now = time.mktime(time.gmtime())
    if last_update + update_frequency < now:
        try:
            _update_seminars(indico_seminar_xml, cached_filename)
            seminar_xml = minidom.parse('%s/%s' % (CFG_CACHEDIR, cached_filename))
        except:
            return "

" + _("No information available") + "

"
    html = ""
    seminars = seminar_xml.getElementsByTagName("seminar")
    if len(seminars) == 0:
        return "

" + _("No seminars today") + "

"
    for seminar in seminars:
        html += "
"
        try:
            seminar_time = seminar.getElementsByTagName("start_time")[0].firstChild.toxml()
        except:
            seminar_time = ""
        try:
            category = seminar.getElementsByTagName("category")[0].firstChild.toxml()
        except:
            category = "Seminar"
        html += '%s %s
' % (seminar_time, category)
        try:
            title = seminar.getElementsByTagName("title")[0].firstChild.toxml()
        except:
            title = ""
        try:
            url = seminar.getElementsByTagName("url")[0].firstChild.toxml()
        except:
            url = "#"
        try:
            speaker = seminar.getElementsByTagName("speaker")[0].firstChild.toxml()
        except:
            speaker = ""
        if (title != ""):
            html += '%s, %s
' % (url, title, speaker)
        try:
            room = seminar.getElementsByTagName("room")[0].firstChild.toxml()
        except:
            room = ""
        html += room

        html += "
"

    html = '
' + html + '
'
    return html.encode('utf-8')

def _update_seminars(indico_seminar_xml, cached_filename):
    """
    helper function that gets the xml data source from CERN Indico and creates
    a dedicated xml file in the cache for easy use in the widget.
    """
    default_timeout = socket.getdefaulttimeout()
    socket.setdefaulttimeout(2.0)
    try:
        try:
            indico_xml = urlopen(indico_seminar_xml)
        except:
            return
    finally:
        socket.setdefaulttimeout(default_timeout)
    xml_file_handler = minidom.parseString(indico_xml.read())
    seminar_xml = ['<Indico_Seminars time="%s">' % time.strftime("%a, %d %b %Y %H:%M:%S GMT", time.gmtime()), ]
    agenda_items = xml_file_handler.getElementsByTagName("agenda_item")
    for item in agenda_items:
        seminar_xml.extend(["<seminar>", ])
        try:
            start_time = item.getElementsByTagName("start_time")[0].firstChild.toxml()
        except:
            start_time = ""
        seminar_xml.extend(["<start_time>%s</start_time>" % start_time, ])
        try:
            category = item.getElementsByTagName("category")[0].firstChild.toxml()
            category = category.split("/")[-1]
            category = category.replace("&", "")
            category = category.replace("nbsp;", "")
            category = category.replace(" ", "")
        except:
            category = ""
        seminar_xml.extend(["<category>%s</category>" % category, ])
        try:
            title = item.getElementsByTagName("title")[0].firstChild.toxml()
        except:
            title = ""
        seminar_xml.extend(["<title>%s</title>" % title, ])
        try:
            url = item.getElementsByTagName("agenda_url")[0].firstChild.toxml()
        except:
            url = "#"
        seminar_xml.extend(["<url>%s</url>" % url, ])
        try:
            speaker = item.getElementsByTagName("speaker")[0].firstChild.toxml()
        except:
            speaker = ""
        seminar_xml.extend(["<speaker>%s</speaker>" % speaker, ])
        try:
            room = item.getElementsByTagName("room")[0].firstChild.toxml()
        except:
            room = ""
        seminar_xml.extend(["<room>%s</room>" % room, ])
        seminar_xml.extend(["</seminar>", ])
    seminar_xml.extend(["</Indico_Seminars>", ])
    # write the created file to cache
    fptr = open("%s/%s" % (CFG_CACHEDIR, cached_filename), "w")
    fptr.write(("\n".join(seminar_xml)).encode('utf-8'))
    fptr.close()

_ = gettext_set_language('en')
dummy = _("What's on today")
dummy = _("Seminars of the week")
<commit_msg>WebJournal: fixed sample data source URL<commit_after># -*- coding: utf-8 -*-
##
## This file is part of CDS Invenio.
## Copyright (C) 2002, 2003, 2004, 2005, 2006, 2007, 2008 CERN.
##
## CDS Invenio is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 2 of the
## License, or (at your option) any later version.
##
## CDS Invenio is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with CDS Invenio; if not, write to the Free Software Foundation, Inc.,
## 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.
"""
WebJournal widget - Display Indico seminars
"""
from invenio.config import CFG_CACHEDIR, CFG_SITE_LANG
from urllib2 import urlopen
from xml.dom import minidom
import time
import socket

from invenio.webjournal_utils import \
     parse_url_string
from invenio.messages import gettext_set_language

update_frequency = 3600 # in seconds

def format(bfo, indico_seminar_xml="http://indico.cern.ch/tools/export.py?fid=1l7&date=today&days=1&of=xml"):
    """
    Display the list of seminar from the given Indico XML URL

    @param indico_seminar_xml: the URL to the XML generated by an Indico instance
    """
    args = parse_url_string(bfo.user_info['uri'])
    journal_name = args["journal_name"]
    cached_filename = "webjournal_widget_seminars_%s.xml" % journal_name
    out = get_widget_html(bfo, indico_seminar_xml, cached_filename, args['ln'])
    return out

def escape_values(bfo):
    """
    Called by BibFormat in order to check if output of this element
    should be escaped.
    """
    return 0

def get_widget_html(bfo, indico_seminar_xml, cached_filename,
                    ln=CFG_SITE_LANG):
    """
    Indico seminars of the day service
    Gets seminars of the day from CERN Indico every 60 minutes and displays
    them in a widget.
    """
    _ = gettext_set_language(ln)
    try:
        seminar_xml = minidom.parse('%s/%s' % (CFG_CACHEDIR, cached_filename))
    except:
        try:
            _update_seminars(indico_seminar_xml, cached_filename)
            seminar_xml = minidom.parse('%s/%s' % (CFG_CACHEDIR, cached_filename))
        except:
            return "

" + _("No information available") + "

"
    try:
        timestamp = seminar_xml.firstChild.getAttribute("time")
    except:
        timestamp = time.struct_time()

    last_update = time.mktime(time.strptime(timestamp,
                                            "%a, %d %b %Y %H:%M:%S %Z"))
    now = time.mktime(time.gmtime())
    if last_update + update_frequency < now:
        try:
            _update_seminars(indico_seminar_xml, cached_filename)
            seminar_xml = minidom.parse('%s/%s' % (CFG_CACHEDIR, cached_filename))
        except:
            return "

" + _("No information available") + "

"
    html = ""
    seminars = seminar_xml.getElementsByTagName("seminar")
    if len(seminars) == 0:
        return "

" + _("No seminars today") + "

"
    for seminar in seminars:
        html += "
"
        try:
            seminar_time = seminar.getElementsByTagName("start_time")[0].firstChild.toxml()
        except:
            seminar_time = ""
        try:
            category = seminar.getElementsByTagName("category")[0].firstChild.toxml()
        except:
            category = "Seminar"
        html += '%s %s
' % (seminar_time, category)
        try:
            title = seminar.getElementsByTagName("title")[0].firstChild.toxml()
        except:
            title = ""
        try:
            url = seminar.getElementsByTagName("url")[0].firstChild.toxml()
        except:
            url = "#"
        try:
            speaker = seminar.getElementsByTagName("speaker")[0].firstChild.toxml()
        except:
            speaker = ""
        if (title != ""):
            html += '%s, %s
' % (url, title, speaker)
        try:
            room = seminar.getElementsByTagName("room")[0].firstChild.toxml()
        except:
            room = ""
        html += room

        html += "
"

    html = '
' + html + '
'
    return html.encode('utf-8')

def _update_seminars(indico_seminar_xml, cached_filename):
    """
    helper function that gets the xml data source from CERN Indico and creates
    a dedicated xml file in the cache for easy use in the widget.
    """
    default_timeout = socket.getdefaulttimeout()
    socket.setdefaulttimeout(2.0)
    try:
        try:
            indico_xml = urlopen(indico_seminar_xml)
        except:
            return
    finally:
        socket.setdefaulttimeout(default_timeout)
    xml_file_handler = minidom.parseString(indico_xml.read())
    seminar_xml = ['<Indico_Seminars time="%s">' % time.strftime("%a, %d %b %Y %H:%M:%S GMT", time.gmtime()), ]
    agenda_items = xml_file_handler.getElementsByTagName("agenda_item")
    for item in agenda_items:
        seminar_xml.extend(["<seminar>", ])
        try:
            start_time = item.getElementsByTagName("start_time")[0].firstChild.toxml()
        except:
            start_time = ""
        seminar_xml.extend(["<start_time>%s</start_time>" % start_time, ])
        try:
            category = item.getElementsByTagName("category")[0].firstChild.toxml()
            category = category.split("/")[-1]
            category = category.replace("&", "")
            category = category.replace("nbsp;", "")
            category = category.replace(" ", "")
        except:
            category = ""
        seminar_xml.extend(["<category>%s</category>" % category, ])
        try:
            title = item.getElementsByTagName("title")[0].firstChild.toxml()
        except:
            title = ""
        seminar_xml.extend(["<title>%s</title>" % title, ])
        try:
            url = item.getElementsByTagName("agenda_url")[0].firstChild.toxml()
        except:
            url = "#"
        seminar_xml.extend(["<url>%s</url>" % url, ])
        try:
            speaker = item.getElementsByTagName("speaker")[0].firstChild.toxml()
        except:
            speaker = ""
        seminar_xml.extend(["<speaker>%s</speaker>" % speaker, ])
        try:
            room = item.getElementsByTagName("room")[0].firstChild.toxml()
        except:
            room = ""
        seminar_xml.extend(["<room>%s</room>" % room, ])
        seminar_xml.extend(["</seminar>", ])
    seminar_xml.extend(["</Indico_Seminars>", ])
    # write the created file to cache
    fptr = open("%s/%s" % (CFG_CACHEDIR, cached_filename), "w")
    fptr.write(("\n".join(seminar_xml)).encode('utf-8'))
    fptr.close()

_ = gettext_set_language('en')
dummy = _("What's on today")
dummy = _("Seminars of the week")
<|endoftext|>