<commit_before>"""

Pre-process DICOM files in a study folder

"""

import logging
import os

from datetime import datetime, timedelta
from functools import partial
from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from airflow.operators.python_operator import PythonOperator
from airflow_spm.operators import SpmPipelineOperator
from airflow_freespace.operators import FreeSpaceSensor
from airflow import configuration

from util import dicom_import
from util import nifti_import

# constants

DAG_NAME = 'mri_pre_process_dicom'

pipelines_path = str(configuration.get('mri', 'PIPELINES_PATH'))
protocols_file = str(configuration.get('mri', 'PROTOCOLS_FILE'))
min_free_space_local_folder = float(
    configuration.get('mri', 'MIN_FREE_SPACE_LOCAL_FOLDER'))
dicom_local_folder = str(
    configuration.get('mri', 'DICOM_LOCAL_FOLDER'))
dicom_to_nifti_local_folder = str(
    configuration.get('mri', 'NIFTI_LOCAL_FOLDER'))
dicom_to_nifti_server_folder = str(
    configuration.get('mri', 'NIFTI_SERVER_FOLDER'))
dicom_to_nifti_pipeline_path = pipelines_path + '/Nifti_Conversion_Pipeline'
neuro_morphometric_atlas_local_folder = str(
    configuration.get('mri', 'NEURO_MORPHOMETRIC_ATLAS_LOCAL_FOLDER'))
neuro_morphometric_atlas_server_folder = str(
    configuration.get('mri', 'NEURO_MORPHOMETRIC_ATLAS_SERVER_FOLDER'))
neuro_morphometric_atlas_pipeline_path = pipelines_path + \
    '/NeuroMorphometric_Pipeline/NeuroMorphometric_tbx/label'
mpm_maps_local_folder = str(
    configuration.get('mri', 'MPM_MAPS_LOCAL_FOLDER'))
mpm_maps_server_folder = str(
    configuration.get('mri', 'MPM_MAPS_SERVER_FOLDER'))
mpm_maps_pipeline_path = pipelines_path + '/MPMs_Pipeline'
misc_library_path = pipelines_path + '/../Miscellaneous&Others'

# functions

# Extract the information from DICOM files located inside a folder.
# The folder information should be given in the configuration parameter
# 'folder' of the DAG run
def extract_dicom_info_fn(**kwargs):
    ti = kwargs['task_instance']
    dr = kwargs['dag_run']
    session_id = dr.conf['session_id']
    input_data_folder = dicom_local_folder + '/' + session_id

    logging.info('folder %s, session_id %s' % (input_data_folder, session_id))

    (participant_id, scan_date) = dicom_import.visit_info(input_data_folder)
    dicom_import.dicom2db(input_data_folder)

    ti.xcom_push(key='folder', value=input_data_folder)
    ti.xcom_push(key='session_id', value=session_id)
    ti.xcom_push(key='participant_id', value=participant_id)
    ti.xcom_push(key='scan_date', value=scan_date)

    return "ok"

# Prepare the arguments for conversion pipeline from DICOM to Nifti format.
# It converts all files located in the folder 'folder'

def dicom_to_nifti_arguments_fn(input_data_folder, session_id, participant_id, scan_date, **kwargs):
    parent_data_folder = os.path.abspath(input_data_folder + '/..')

    return [parent_data_folder,
            session_id,
            dicom_to_nifti_local_folder,
            dicom_to_nifti_server_folder,
            protocols_file]

# Prepare the arguments for the pipeline that builds a Neuro morphometric
# atlas from the Nitfi files located in the folder 'folder'

def neuro_morphometric_atlas_arguments_fn(input_data_folder, session_id, participant_id, scan_date, **kwargs):
    parent_data_folder = os.path.abspath(input_data_folder + '/..')
    table_format = 'csv'

    return [session_id,
            parent_data_folder,
            neuro_morphometric_atlas_local_folder,
            neuro_morphometric_atlas_server_folder,
            protocols_file,
            table_format]

# Pipeline that builds the MPM maps from the Nitfi files located in the
# folder 'folder'

def mpm_maps_arguments_fn(input_data_folder, session_id, participant_id, scan_date, **kwargs):
    parent_data_folder = os.path.abspath(input_data_folder + '/..')
    pipeline_params_config_file = 'Preproc_mpm_maps_pipeline_config.txt'

    return [parent_data_folder,
            session_id,
            mpm_maps_local_folder,
            protocols_file,
            pipeline_params_config_file,
            mpm_maps_server_folder]

# Extract information from the Nifti files located in the folder 'folder'
# parent_task should contain XCOM keys 'folder' and 'session_id'

def extract_nifti_info_fn(parent_task, **kwargs):
    ti = kwargs['task_instance']
    input_data_folder = ti.xcom_pull(key='folder', task_ids=parent_task)
    session_id = ti.xcom_pull(key='session_id', task_ids=parent_task)
    participant_id = ti.xcom_pull(key='participant_id', task_ids=parent_task)
    scan_date = ti.xcom_pull(key='scan_date', task_ids=parent_task)

    logging.info("NIFTI extract: session_id=%s, input_folder=%s" %
                 (session_id, input_data_folder))

    nifti_import.nifti2db(input_data_folder, participant_id, scan_date)

    ti.xcom_push(key='folder', value=input_data_folder)
    ti.xcom_push(key='session_id', value=session_id)
    ti.xcom_push(key='participant_id', value=participant_id)
    ti.xcom_push(key='scan_date', value=scan_date)

    return "ok"

# Define the DAG

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime.now(),
    'retries': 1,
    'retry_delay': timedelta(seconds=120),
    'email': 'ludovic.claude@chuv.ch',
    'email_on_failure': True,
    'email_on_retry': True
}

dag = DAG(
    dag_id=DAG_NAME,
    default_args=default_args,
    schedule_interval=None)

check_free_space = FreeSpaceSensor(
    task_id='check_free_space',
    path=dicom_local_folder,
    free_disk_threshold=min_free_space_local_folder,
    retry_delay=timedelta(hours=1),
    retries=24 * 7,
    pool='remote_file_copy',
    dag=dag
)

check_free_space.doc_md = """\
# Check free space

Check that there is enough free space on the local disk for processing, wait otherwise.
"""

copy_dicom_to_local_cmd = """
    used="$(df -h /home | grep '/' | grep -Po '[^ ]*(?=%)')"
    if (( 101 - used < {{ min_free_space_local_folder|float * 100 }} )); then
      echo "Not enough space left, cannot continue"
      ecit 1
    fi
    rsync -av {{ dag_run.conf["folder"] }}/ {{ params["local_output_folder"] }}/{{ dag_run.conf["session_id"] }}
"""

copy_dicom_to_local = BashOperator(
    task_id='copy_dicom_to_local',
    bash_command=copy_dicom_to_local_cmd,
    params={'local_output_folder': dicom_local_folder},
    pool='remote_file_copy',
    priority_weight=10,
    execution_timeout=timedelta(hours=3),
    dag=dag
)
copy_dicom_to_local.set_upstream(check_free_space)

copy_dicom_to_local.doc_md = """\
# Copy DICOM files to a local drive

Speed-up the processing of DICOM files by first copying them from a shared folder to the local hard-drive.
"""

extract_dicom_info = PythonOperator(
    task_id='extract_dicom_info',
    python_callable=extract_dicom_info_fn,
    provide_context=True,
    execution_timeout=timedelta(hours=3),
    dag=dag
)
extract_dicom_info.set_upstream(copy_dicom_to_local)

extract_dicom_info.doc_md = """\
# Extract DICOM information

Read DICOM information from the files stored in the session folder and store that information in the database.
"""

dicom_to_nifti_pipeline = SpmPipelineOperator(
    task_id='dicom_to_nifti_pipeline',
    spm_function='DCM2NII_LREN',
    spm_arguments_callable=dicom_to_nifti_arguments_fn,
    provide_context=True,
    matlab_paths=[misc_library_path, dicom_to_nifti_pipeline_path],
    output_folder_callable=lambda session_id, **kwargs: dicom_to_nifti_local_folder + '/' + session_id,
    pool='image_preprocessing',
    parent_task='extract_dicom_info',
    execution_timeout=timedelta(hours=24),
    dag=dag
)

dicom_to_nifti_pipeline.set_upstream(extract_dicom_info)

dicom_to_nifti_pipeline.doc_md = """\
# DICOM to Nitfi Pipeline

This function convert the dicom files to Nifti format using the SPM tools and dcm2nii tool developed by Chris Rorden.

Webpage: http://www.mccauslandcenter.sc.edu/mricro/mricron/dcm2nii.html

"""

cleanup_local_dicom_cmd = """
    rm -rf {{ params["local_folder"] }}/{{ dag_run.conf["session_id"] }}
"""

cleanup_local_dicom = BashOperator(
    task_id='cleanup_local_dicom',
    bash_command=copy_dicom_to_local_cmd,
    params={'local_folder': dicom_local_folder},
    priority_weight=20,
    execution_timeout=timedelta(hours=1),
    dag=dag
)
cleanup_local_dicom.set_upstream(dicom_to_nifti_pipeline)

cleanup_local_dicom.doc_md = """\
# Cleanup local DICOM files

Remove locally stored DICOM files as they have been processed already.
"""

extract_nifti_info = PythonOperator(
    task_id='extract_nifti_info',
    python_callable=partial(extract_nifti_info_fn, 'dicom_to_nifti_pipeline'),
    provide_context=True,
    execution_timeout=timedelta(hours=3),
    dag=dag
)

extract_nifti_info.set_upstream(dicom_to_nifti_pipeline)

extract_nifti_info.doc_md = """\
# Extract information from NIFTI files converted from DICOM

Read NIFTI information from a directory tree of nifti files freshly converted from DICOM and store that information in the database.
"""

mpm_maps_pipeline = SpmPipelineOperator(
    task_id='mpm_maps_pipeline',
    spm_function='Preproc_mpm_maps',
    spm_arguments_callable=mpm_maps_arguments_fn,
    provide_context=True,
    matlab_paths=[misc_library_path, mpm_maps_pipeline_path],
    output_folder_callable=lambda session_id, **kwargs: mpm_maps_local_folder + '/' + session_id,
    execution_timeout=timedelta(hours=3),
    pool='image_preprocessing',
    parent_task='extract_nifti_info',
    dag=dag
)

mpm_maps_pipeline.set_upstream(extract_nifti_info)

mpm_maps_pipeline.doc_md = """\
# MPM Maps Pipeline

This function computes the Multiparametric Maps (MPMs) (R2*, R1, MT, PD) and brain segmentation in different tissue maps.
All computation was programmed based on the LREN database structure. The MPMs are calculated locally in 'OutputFolder' and finally copied to 'ServerFolder'.

"""

extract_nifti_mpm_info = PythonOperator(
    task_id='extract_nifti_mpm_info',
    python_callable=partial(extract_nifti_info_fn, 'mpm_maps_pipeline'),
    provide_context=True,
    execution_timeout=timedelta(hours=3),
    dag=dag
)

extract_nifti_mpm_info.set_upstream(mpm_maps_pipeline)

extract_nifti_mpm_info.doc_md = """\
# Extract information from NIFTI files generated by MPM pipeline

Read NIFTI information from a directory tree containing the Nifti files created by MPM pipeline and store that information in the database.
"""

neuro_morphometric_atlas_pipeline = SpmPipelineOperator(
    task_id='neuro_morphometric_atlas_pipeline',
    spm_function='NeuroMorphometric_pipeline',
    spm_arguments_callable=neuro_morphometric_atlas_arguments_fn,
    provide_context=True,
    matlab_paths=[misc_library_path, neuro_morphometric_atlas_pipeline_path],
    output_folder_callable=lambda session_id, **kwargs: neuro_morphometric_atlas_local_folder + '/' + session_id,
    pool='image_preprocessing',
    parent_task='mpm_maps_pipeline',
    execution_timeout=timedelta(hours=24),
    dag=dag
)

neuro_morphometric_atlas_pipeline.set_upstream(mpm_maps_pipeline)

neuro_morphometric_atlas_pipeline.doc_md = """\
# NeuroMorphometric Pipeline

This function computes an individual Atlas based on the NeuroMorphometrics Atlas. This is based on the NeuroMorphometrics Toolbox.
This delivers three files:

1. Atlas File (*.nii);
2. Volumes of the Morphometric Atlas structures (*.txt);
3. CSV File (.csv) containing the volume, globals, and Multiparametric Maps (R2*, R1, MT, PD) for each structure defined in the Subject Atlas.

"""

extract_nifti_atlas_info = PythonOperator(
    task_id='extract_nifti_atlas_info',
    python_callable=partial(extract_nifti_info_fn,
                            'neuro_morphometric_atlas_pipeline'),
    provide_context=True,
    execution_timeout=timedelta(hours=3),
    dag=dag
)

extract_nifti_atlas_info.set_upstream(neuro_morphometric_atlas_pipeline)

extract_nifti_atlas_info.doc_md = """\
# Extract information from NIFTI files generated by Neuro Morphometrics Atlas pipeline

Read NIFTI information from a directory tree containing the Nifti files created by Neuro Morphometrics Atlas pipeline and store that information in the database.
"""
<commit_msg>Fix check in copy task<commit_after>"""

Pre-process DICOM files in a study folder

"""

import logging
import os

from datetime import datetime, timedelta
from functools import partial
from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from airflow.operators.python_operator import PythonOperator
from airflow_spm.operators import SpmPipelineOperator
from airflow_freespace.operators import FreeSpaceSensor
from airflow import configuration

from util import dicom_import
from util import nifti_import

# constants

DAG_NAME = 'mri_pre_process_dicom'

pipelines_path = str(configuration.get('mri', 'PIPELINES_PATH'))
protocols_file = str(configuration.get('mri', 'PROTOCOLS_FILE'))
min_free_space_local_folder = float(
    configuration.get('mri', 'MIN_FREE_SPACE_LOCAL_FOLDER'))
dicom_local_folder = str(
    configuration.get('mri', 'DICOM_LOCAL_FOLDER'))
dicom_to_nifti_local_folder = str(
    configuration.get('mri', 'NIFTI_LOCAL_FOLDER'))
dicom_to_nifti_server_folder = str(
    configuration.get('mri', 'NIFTI_SERVER_FOLDER'))
dicom_to_nifti_pipeline_path = pipelines_path + '/Nifti_Conversion_Pipeline'
neuro_morphometric_atlas_local_folder = str(
    configuration.get('mri', 'NEURO_MORPHOMETRIC_ATLAS_LOCAL_FOLDER'))
neuro_morphometric_atlas_server_folder = str(
    configuration.get('mri', 'NEURO_MORPHOMETRIC_ATLAS_SERVER_FOLDER'))
neuro_morphometric_atlas_pipeline_path = pipelines_path + \
    '/NeuroMorphometric_Pipeline/NeuroMorphometric_tbx/label'
mpm_maps_local_folder = str(
    configuration.get('mri', 'MPM_MAPS_LOCAL_FOLDER'))
mpm_maps_server_folder = str(
    configuration.get('mri', 'MPM_MAPS_SERVER_FOLDER'))
mpm_maps_pipeline_path = pipelines_path + '/MPMs_Pipeline'
misc_library_path = pipelines_path + '/../Miscellaneous&Others'

# functions

# Extract the information from DICOM files located inside a folder.
# The folder information should be given in the configuration parameter
# 'folder' of the DAG run
def extract_dicom_info_fn(**kwargs):
    ti = kwargs['task_instance']
    dr = kwargs['dag_run']
    session_id = dr.conf['session_id']
    input_data_folder = dicom_local_folder + '/' + session_id

    logging.info('folder %s, session_id %s' % (input_data_folder, session_id))

    (participant_id, scan_date) = dicom_import.visit_info(input_data_folder)
    dicom_import.dicom2db(input_data_folder)

    ti.xcom_push(key='folder', value=input_data_folder)
    ti.xcom_push(key='session_id', value=session_id)
    ti.xcom_push(key='participant_id', value=participant_id)
    ti.xcom_push(key='scan_date', value=scan_date)

    return "ok"

# Prepare the arguments for conversion pipeline from DICOM to Nifti format.
# It converts all files located in the folder 'folder'

def dicom_to_nifti_arguments_fn(input_data_folder, session_id, participant_id, scan_date, **kwargs):
    parent_data_folder = os.path.abspath(input_data_folder + '/..')

    return [parent_data_folder,
            session_id,
            dicom_to_nifti_local_folder,
            dicom_to_nifti_server_folder,
            protocols_file]

# Prepare the arguments for the pipeline that builds a Neuro morphometric
# atlas from the Nitfi files located in the folder 'folder'

def neuro_morphometric_atlas_arguments_fn(input_data_folder, session_id, participant_id, scan_date, **kwargs):
    parent_data_folder = os.path.abspath(input_data_folder + '/..')
    table_format = 'csv'

    return [session_id,
            parent_data_folder,
            neuro_morphometric_atlas_local_folder,
            neuro_morphometric_atlas_server_folder,
            protocols_file,
            table_format]

# Pipeline that builds the MPM maps from the Nitfi files located in the
# folder 'folder'

def mpm_maps_arguments_fn(input_data_folder, session_id, participant_id, scan_date, **kwargs):
    parent_data_folder = os.path.abspath(input_data_folder + '/..')
    pipeline_params_config_file = 'Preproc_mpm_maps_pipeline_config.txt'

    return [parent_data_folder,
            session_id,
            mpm_maps_local_folder,
            protocols_file,
            pipeline_params_config_file,
            mpm_maps_server_folder]

# Extract information from the Nifti files located in the folder 'folder'
# parent_task should contain XCOM keys 'folder' and 'session_id'

def extract_nifti_info_fn(parent_task, **kwargs):
    ti = kwargs['task_instance']
    input_data_folder = ti.xcom_pull(key='folder', task_ids=parent_task)
    session_id = ti.xcom_pull(key='session_id', task_ids=parent_task)
    participant_id = ti.xcom_pull(key='participant_id', task_ids=parent_task)
    scan_date = ti.xcom_pull(key='scan_date', task_ids=parent_task)

    logging.info("NIFTI extract: session_id=%s, input_folder=%s" %
                 (session_id, input_data_folder))

    nifti_import.nifti2db(input_data_folder, participant_id, scan_date)

    ti.xcom_push(key='folder', value=input_data_folder)
    ti.xcom_push(key='session_id', value=session_id)
    ti.xcom_push(key='participant_id', value=participant_id)
    ti.xcom_push(key='scan_date', value=scan_date)

    return "ok"

# Define the DAG

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime.now(),
    'retries': 1,
    'retry_delay': timedelta(seconds=120),
    'email': 'ludovic.claude@chuv.ch',
    'email_on_failure': True,
    'email_on_retry': True
}

dag = DAG(
    dag_id=DAG_NAME,
    default_args=default_args,
    schedule_interval=None)

check_free_space = FreeSpaceSensor(
    task_id='check_free_space',
    path=dicom_local_folder,
    free_disk_threshold=min_free_space_local_folder,
    retry_delay=timedelta(hours=1),
    retries=24 * 7,
    pool='remote_file_copy',
    dag=dag
)

check_free_space.doc_md = """\
# Check free space

Check that there is enough free space on the local disk for processing, wait otherwise.
"""

copy_dicom_to_local_cmd = """
    used="$(df -h /home | grep '/' | grep -Po '[^ ]*(?=%)')"
    if (( 101 - used < {{ params['min_free_space_local_folder']|float * 100 }} )); then
      echo "Not enough space left, cannot continue"
      ecit 1
    fi
    rsync -av {{ dag_run.conf["folder"] }}/ {{ params["local_output_folder"] }}/{{ dag_run.conf["session_id"] }}
"""

copy_dicom_to_local = BashOperator(
    task_id='copy_dicom_to_local',
    bash_command=copy_dicom_to_local_cmd,
    params={'local_output_folder': dicom_local_folder, 'min_free_space_local_folder': min_free_space_local_folder},
    pool='remote_file_copy',
    priority_weight=10,
    execution_timeout=timedelta(hours=3),
    dag=dag
)
copy_dicom_to_local.set_upstream(check_free_space)

copy_dicom_to_local.doc_md = """\
# Copy DICOM files to a local drive

Speed-up the processing of DICOM files by first copying them from a shared folder to the local hard-drive.
"""

extract_dicom_info = PythonOperator(
    task_id='extract_dicom_info',
    python_callable=extract_dicom_info_fn,
    provide_context=True,
    execution_timeout=timedelta(hours=3),
    dag=dag
)
extract_dicom_info.set_upstream(copy_dicom_to_local)

extract_dicom_info.doc_md = """\
# Extract DICOM information

Read DICOM information from the files stored in the session folder and store that information in the database.
"""

dicom_to_nifti_pipeline = SpmPipelineOperator(
    task_id='dicom_to_nifti_pipeline',
    spm_function='DCM2NII_LREN',
    spm_arguments_callable=dicom_to_nifti_arguments_fn,
    provide_context=True,
    matlab_paths=[misc_library_path, dicom_to_nifti_pipeline_path],
    output_folder_callable=lambda session_id, **kwargs: dicom_to_nifti_local_folder + '/' + session_id,
    pool='image_preprocessing',
    parent_task='extract_dicom_info',
    execution_timeout=timedelta(hours=24),
    dag=dag
)

dicom_to_nifti_pipeline.set_upstream(extract_dicom_info)

dicom_to_nifti_pipeline.doc_md = """\
# DICOM to Nitfi Pipeline

This function convert the dicom files to Nifti format using the SPM tools and dcm2nii tool developed by Chris Rorden.

Webpage: http://www.mccauslandcenter.sc.edu/mricro/mricron/dcm2nii.html

"""

cleanup_local_dicom_cmd = """
    rm -rf {{ params["local_folder"] }}/{{ dag_run.conf["session_id"] }}
"""

cleanup_local_dicom = BashOperator(
    task_id='cleanup_local_dicom',
    bash_command=copy_dicom_to_local_cmd,
    params={'local_folder': dicom_local_folder},
    priority_weight=20,
    execution_timeout=timedelta(hours=1),
    dag=dag
)
cleanup_local_dicom.set_upstream(dicom_to_nifti_pipeline)

cleanup_local_dicom.doc_md = """\
# Cleanup local DICOM files

Remove locally stored DICOM files as they have been processed already.
"""

extract_nifti_info = PythonOperator(
    task_id='extract_nifti_info',
    python_callable=partial(extract_nifti_info_fn, 'dicom_to_nifti_pipeline'),
    provide_context=True,
    execution_timeout=timedelta(hours=3),
    dag=dag
)

extract_nifti_info.set_upstream(dicom_to_nifti_pipeline)

extract_nifti_info.doc_md = """\
# Extract information from NIFTI files converted from DICOM

Read NIFTI information from a directory tree of nifti files freshly converted from DICOM and store that information in the database.
"""

mpm_maps_pipeline = SpmPipelineOperator(
    task_id='mpm_maps_pipeline',
    spm_function='Preproc_mpm_maps',
    spm_arguments_callable=mpm_maps_arguments_fn,
    provide_context=True,
    matlab_paths=[misc_library_path, mpm_maps_pipeline_path],
    output_folder_callable=lambda session_id, **kwargs: mpm_maps_local_folder + '/' + session_id,
    execution_timeout=timedelta(hours=3),
    pool='image_preprocessing',
    parent_task='extract_nifti_info',
    dag=dag
)

mpm_maps_pipeline.set_upstream(extract_nifti_info)

mpm_maps_pipeline.doc_md = """\
# MPM Maps Pipeline

This function computes the Multiparametric Maps (MPMs) (R2*, R1, MT, PD) and brain segmentation in different tissue maps.
All computation was programmed based on the LREN database structure. The MPMs are calculated locally in 'OutputFolder' and finally copied to 'ServerFolder'.

"""

extract_nifti_mpm_info = PythonOperator(
    task_id='extract_nifti_mpm_info',
    python_callable=partial(extract_nifti_info_fn, 'mpm_maps_pipeline'),
    provide_context=True,
    execution_timeout=timedelta(hours=3),
    dag=dag
)

extract_nifti_mpm_info.set_upstream(mpm_maps_pipeline)

extract_nifti_mpm_info.doc_md = """\
# Extract information from NIFTI files generated by MPM pipeline

Read NIFTI information from a directory tree containing the Nifti files created by MPM pipeline and store that information in the database.
"""

neuro_morphometric_atlas_pipeline = SpmPipelineOperator(
    task_id='neuro_morphometric_atlas_pipeline',
    spm_function='NeuroMorphometric_pipeline',
    spm_arguments_callable=neuro_morphometric_atlas_arguments_fn,
    provide_context=True,
    matlab_paths=[misc_library_path, neuro_morphometric_atlas_pipeline_path],
    output_folder_callable=lambda session_id, **kwargs: neuro_morphometric_atlas_local_folder + '/' + session_id,
    pool='image_preprocessing',
    parent_task='mpm_maps_pipeline',
    execution_timeout=timedelta(hours=24),
    dag=dag
)

neuro_morphometric_atlas_pipeline.set_upstream(mpm_maps_pipeline)

neuro_morphometric_atlas_pipeline.doc_md = """\
# NeuroMorphometric Pipeline

This function computes an individual Atlas based on the NeuroMorphometrics Atlas. This is based on the NeuroMorphometrics Toolbox.
This delivers three files:

1. Atlas File (*.nii);
2. Volumes of the Morphometric Atlas structures (*.txt);
3. CSV File (.csv) containing the volume, globals, and Multiparametric Maps (R2*, R1, MT, PD) for each structure defined in the Subject Atlas.

"""

extract_nifti_atlas_info = PythonOperator(
    task_id='extract_nifti_atlas_info',
    python_callable=partial(extract_nifti_info_fn,
                            'neuro_morphometric_atlas_pipeline'),
    provide_context=True,
    execution_timeout=timedelta(hours=3),
    dag=dag
)

extract_nifti_atlas_info.set_upstream(neuro_morphometric_atlas_pipeline)

extract_nifti_atlas_info.doc_md = """\
# Extract information from NIFTI files generated by Neuro Morphometrics Atlas pipeline

Read NIFTI information from a directory tree containing the Nifti files created by Neuro Morphometrics Atlas pipeline and store that information in the database.
"""
<|endoftext|>


<commit_before># Script to create app cache manifest files for the Dragonfly project.
# Takes very much in regard how the dragonfly.opera.com server works.
# There are two main paths on https://dragonfly.opera.com,
# /app/ and /app/cutting-edge/.
# If the request has no language cookie then a small document to
# set the language cookie is served.
# If the language cookie is set the version with the correct core version
# (depending on the user agent string) is returned. That means
#
#     /app/
#     /app/<core version>/
#     /app/<core version>/client-<language code>.xml
#
# or
#
#     /app/cutting-edge/
#     /app/<core version>/cutting-edge/
#     /app/<core version>/cutting-edge/client-<language code>.xml
#
# are all pointing to the same resource.
# All these possible entry paths are added to the mainifest.
# Additionally all linked resources are added. All paths are absolut do to
# all the rewrites.

from __future__ import with_statement
import re
import os
import sys
from urllib import quote, unquote
from time import gmtime, strftime, mktime, strptime, time

MANIFEST_DIR = "manifests"
_re_client = re.compile("client-(?P<lang>[^.]+)\.xml")
_re_resource = re.compile("(?:<script[^>]*src=\"([^\"]*))|(?:<link[^>]*href=\"([^\"]*))")

def get_timestamp(path = None):
    return strftime("%a, %d %b %Y %H:%M:%S GMT", gmtime(path and stat(path).st_mtime or None))

def get_resources(os_path, web_path, file_name):
    resources = []
    with open(os.path.join(os_path, file_name), 'r') as f:
        content = f.read()
        for match in _re_resource.finditer(content):
            resource = filter(bool, match.groups())[0]
            if resource.startswith("/"):
                pass
            elif resource.startswith("./"):
                resource = web_path + resource[1:]
            # everything else resolved by the browser
            else:
                resourc = web_path + "/" + resource
            resources.append(resource)
    return resources

def write_maifest(os_path, file_name, resources, tag=""):
    with open(os.path.join(os_path, MANIFEST_DIR, file_name), "wb") as f:
        content = ["CACHE MANIFEST", "# created %s %s" % (get_timestamp(), tag)]
        content.extend(resources)
        content.append("")
        f.write("\n".join(content))

def add_manifest(os_path, web_path, client_file, manifest):
    content = None
    manifest_path = "/".join([web_path, MANIFEST_DIR, manifest])
    with open(os.path.join(os_path, client_file), "rb") as f:
        content = f.read()
    if content and not ' manifest="' in content:
        with open(os.path.join(os_path, client_file), "wb") as f:
            f.write(content.replace(">", ' manifest="%s">' % manifest_path, 1))

def main(argv=sys.argv):
    import optparse
    usage = """%prog [options] [path]

    Create a manifest for all files in path with the pattern client-<lang code>.xml.
    """
    parser = optparse.OptionParser(usage)
    parser.add_option("-d", "--domain", type="string", dest="domain",
                      default=None,
                      help="The domain token of the url.")
    parser.add_option("-t", "--tag", dest="tag", default="",
                      help="A tag to be used in the manifest comment.")
    options, args = parser.parse_args()
    os_path = os.path.abspath(len(args) > 0 and args[0] or ".")
    web_path = os_path
    if options.domain:
        pos = web_path.find(options.domain)
        if pos > -1:
            web_path = web_path[pos + len(options.domain):]
    web_path = "/".join(map(lambda p: quote(p), web_path.split(os.path.sep)))
    if not os.path.exists(os.path.join(os_path, MANIFEST_DIR)):
        os.makedirs(os.path.join(os_path, MANIFEST_DIR))
    for name in os.listdir(os_path):
         match = _re_client.search(name)
         if match:
            resources = [web_path + "/", web_path + "/" + name]
            manifest = "dragonfly-%s.manifest" % match.group("lang")
            if web_path.startswith("/app"):
                if "cutting-edge" in web_path:
                    resources.append("/app/cutting-edge/")
                else:
                    resources.append("/app/")
            resources.extend(get_resources(os_path, web_path, name))
            write_maifest(os_path, manifest, resources, options.tag)
            add_manifest(os_path, web_path, name, manifest)

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Added some comments to the manifest script.<commit_after># Script to create app cache manifest files for the Dragonfly project.
# Takes very much in regard how the dragonfly.opera.com server works.
# There are two main paths on https://dragonfly.opera.com,
# /app/ and /app/cutting-edge/.
# If the request has no language cookie then a small document to
# set the language cookie is served.
# If the language cookie is set the version with the correct core version
# (depending on the user agent string) is returned. That means
#
#   /app/
#   /app/<core version>/
#   /app/<core version>/client-<language code>.xml
#
# or
#
#   /app/cutting-edge/
#   /app/<core version>/cutting-edge/
#   /app/<core version>/cutting-edge/client-<language code>.xml
#
# are all pointing to the same resource.
# All these possible entry paths are added to the mainifest.
# Additionally all linked resources are added. All paths are absolut do to
# all the rewrites.

from __future__ import with_statement
import re
import os
import sys
from urllib import quote, unquote
from time import gmtime, strftime, mktime, strptime, time

MANIFEST_DIR = "manifests"
_re_client = re.compile("client-(?P<lang>[^.]+)\.xml")
_re_resource = re.compile("(?:<script[^>]*src=\"([^\"]*))|(?:<link[^>]*href=\"([^\"]*))")

def get_timestamp(path = None):
    return strftime("%a, %d %b %Y %H:%M:%S GMT", gmtime(path and stat(path).st_mtime or None))

def get_resources(os_path, web_path, file_name):
    resources = []
    with open(os.path.join(os_path, file_name), 'r') as f:
        content = f.read()
        for match in _re_resource.finditer(content):
            resource = filter(bool, match.groups())[0]
            if resource.startswith("/"):
                pass
            elif resource.startswith("./"):
                resource = web_path + resource[1:]
            # everything else resolved by the browser
            else:
                resourc = web_path + "/" + resource
            resources.append(resource)
    return resources

def write_maifest(os_path, file_name, resources, tag=""):
    with open(os.path.join(os_path, MANIFEST_DIR, file_name), "wb") as f:
        content = ["CACHE MANIFEST", "# created %s %s" % (get_timestamp(), tag)]
        content.extend(resources)
        content.append("")
        f.write("\n".join(content))

def add_manifest(os_path, web_path, client_file, manifest):
    content = None
    manifest_path = "/".join([web_path, MANIFEST_DIR, manifest])
    with open(os.path.join(os_path, client_file), "rb") as f:
        content = f.read()
    if content and not ' manifest="' in content:
        with open(os.path.join(os_path, client_file), "wb") as f:
            f.write(content.replace(">", ' manifest="%s">' % manifest_path, 1))

def main(argv=sys.argv):
    import optparse
    usage = """%prog [options] [path]

    Create a manifest for all files in path with the pattern client-<lang code>.xml.
    """
    parser = optparse.OptionParser(usage)
    parser.add_option("-d", "--domain", type="string", dest="domain",
                      default=None,
                      help="The domain token of the url.")
    parser.add_option("-t", "--tag", dest="tag", default="",
                      help="A tag to be used in the manifest comment.")
    options, args = parser.parse_args()
    os_path = os.path.abspath(len(args) > 0 and args[0] or ".")
    web_path = os_path
    if options.domain:
        pos = web_path.find(options.domain)
        if pos > -1:
            web_path = web_path[pos + len(options.domain):]
    web_path = "/".join(map(lambda p: quote(p), web_path.split(os.path.sep)))
    if not os.path.exists(os.path.join(os_path, MANIFEST_DIR)):
        os.makedirs(os.path.join(os_path, MANIFEST_DIR))
    for name in os.listdir(os_path):
         match = _re_client.search(name)
         if match:
            resources = [web_path + "/", web_path + "/" + name]
            manifest = "dragonfly-%s.manifest" % match.group("lang")
            if web_path.startswith("/app"):
                if "cutting-edge" in web_path:
                    resources.append("/app/cutting-edge/")
                else:
                    resources.append("/app/")
            resources.extend(get_resources(os_path, web_path, name))
            write_maifest(os_path, manifest, resources, options.tag)
            add_manifest(os_path, web_path, name, manifest)

if __name__ == "__main__":
    sys.exit(main())
<|endoftext|>


<commit_before># Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities for API compatibility between TensorFlow release versions.

See [Version
Compatibility](https://tensorflow.org/guide/version_compat#backward_forward)
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import datetime
import os

from tensorflow.python.util import tf_contextlib
from tensorflow.python.util.tf_export import tf_export

# This value changes every day with an automatic CL. It can be modified in code
# via `forward_compatibility_horizon()` or with the environment variable
# TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 4, 13)
_FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
_FORWARD_COMPATIBILITY_DATE_NUMBER = None

def _date_to_date_number(year, month, day):
  return (year << 9) | (month << 5) | day

def _update_forward_compatibility_date_number(date_to_override=None):
  """Update the base date to compare in forward_compatible function."""

  global _FORWARD_COMPATIBILITY_DATE_NUMBER

  if date_to_override:
    date = date_to_override
  else:
    date = _FORWARD_COMPATIBILITY_HORIZON
    delta_days = os.getenv(_FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME)
    if delta_days:
      date += datetime.timedelta(days=int(delta_days))

  _FORWARD_COMPATIBILITY_DATE_NUMBER = _date_to_date_number(
      date.year, date.month, date.day)

_update_forward_compatibility_date_number()

@tf_export("compat.forward_compatible")
def forward_compatible(year, month, day):
  """Return true if the forward compatibility window has expired.

  See [Version
  compatibility](https://tensorflow.org/guide/version_compat#backward_forward).

  Forward-compatibility refers to scenarios where the producer of a TensorFlow
  model (a GraphDef or SavedModel) is compiled against a version of the
  TensorFlow library newer than what the consumer was compiled against. The
  "producer" is typically a Python program that constructs and trains a model
  while the "consumer" is typically another program that loads and serves the
  model.

  TensorFlow has been supporting a 3 week forward-compatibility window for
  programs compiled from source at HEAD.

  For example, consider the case where a new operation `MyNewAwesomeAdd` is
  created with the intent of replacing the implementation of an existing Python
  wrapper - `tf.add`.  The Python wrapper implementation should change from
  something like:

  ```python
  def add(inputs, name=None):
    return gen_math_ops.add(inputs, name)
  ```

  to:

  ```python
  from tensorflow.python.compat import compat

  def add(inputs, name=None):
    if compat.forward_compatible(year, month, day):
      # Can use the awesome new implementation.
      return gen_math_ops.my_new_awesome_add(inputs, name)
    # To maintain forward compatibility, use the old implementation.
    return gen_math_ops.add(inputs, name)
  ```

  Where `year`, `month`, and `day` specify the date beyond which binaries
  that consume a model are expected to have been updated to include the
  new operations. This date is typically at least 3 weeks beyond the date
  the code that adds the new operation is committed.

  Args:
    year:  A year (e.g., 2018). Must be an `int`.
    month: A month (1 <= month <= 12) in year. Must be an `int`.
    day:   A day (1 <= day <= 31, or 30, or 29, or 28) in month. Must be an
      `int`.

  Returns:
    True if the caller can expect that serialized TensorFlow graphs produced
    can be consumed by programs that are compiled with the TensorFlow library
    source code after (year, month, day).
  """
  return _FORWARD_COMPATIBILITY_DATE_NUMBER > _date_to_date_number(
      year, month, day)

@tf_export("compat.forward_compatibility_horizon")
@tf_contextlib.contextmanager
def forward_compatibility_horizon(year, month, day):
  """Context manager for testing forward compatibility of generated graphs.

  See [Version
  compatibility](https://tensorflow.org/guide/version_compat#backward_forward).

  To ensure forward compatibility of generated graphs (see `forward_compatible`)
  with older binaries, new features can be gated with:

  ```python
  if compat.forward_compatible(year=2018, month=08, date=01):
    generate_graph_with_new_features()
  else:
    generate_graph_so_older_binaries_can_consume_it()
  ```

  However, when adding new features, one may want to unittest it before
  the forward compatibility window expires. This context manager enables
  such tests. For example:

  ```python
  from tensorflow.python.compat import compat

  def testMyNewFeature(self):
    with compat.forward_compatibility_horizon(2018, 08, 02):
       # Test that generate_graph_with_new_features() has an effect
  ```

  Args:
    year:  A year (e.g., 2018). Must be an `int`.
    month: A month (1 <= month <= 12) in year. Must be an `int`.
    day:   A day (1 <= day <= 31, or 30, or 29, or 28) in month. Must be an
      `int`.

  Yields:
    Nothing.
  """
  try:
    _update_forward_compatibility_date_number(datetime.date(year, month, day))
    yield
  finally:
    _update_forward_compatibility_date_number()
<commit_msg>compat: Update forward compatibility horizon to 2020-04-14<commit_after># Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities for API compatibility between TensorFlow release versions.

See [Version
Compatibility](https://tensorflow.org/guide/version_compat#backward_forward)
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import datetime
import os

from tensorflow.python.util import tf_contextlib
from tensorflow.python.util.tf_export import tf_export

# This value changes every day with an automatic CL. It can be modified in code
# via `forward_compatibility_horizon()` or with the environment variable
# TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 4, 14)
_FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
_FORWARD_COMPATIBILITY_DATE_NUMBER = None

def _date_to_date_number(year, month, day):
  return (year << 9) | (month << 5) | day

def _update_forward_compatibility_date_number(date_to_override=None):
  """Update the base date to compare in forward_compatible function."""

  global _FORWARD_COMPATIBILITY_DATE_NUMBER

  if date_to_override:
    date = date_to_override
  else:
    date = _FORWARD_COMPATIBILITY_HORIZON
    delta_days = os.getenv(_FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME)
    if delta_days:
      date += datetime.timedelta(days=int(delta_days))

  _FORWARD_COMPATIBILITY_DATE_NUMBER = _date_to_date_number(
      date.year, date.month, date.day)

_update_forward_compatibility_date_number()

@tf_export("compat.forward_compatible")
def forward_compatible(year, month, day):
  """Return true if the forward compatibility window has expired.

  See [Version
  compatibility](https://tensorflow.org/guide/version_compat#backward_forward).

  Forward-compatibility refers to scenarios where the producer of a TensorFlow
  model (a GraphDef or SavedModel) is compiled against a version of the
  TensorFlow library newer than what the consumer was compiled against. The
  "producer" is typically a Python program that constructs and trains a model
  while the "consumer" is typically another program that loads and serves the
  model.

  TensorFlow has been supporting a 3 week forward-compatibility window for
  programs compiled from source at HEAD.

  For example, consider the case where a new operation `MyNewAwesomeAdd` is
  created with the intent of replacing the implementation of an existing Python
  wrapper - `tf.add`.  The Python wrapper implementation should change from
  something like:

  ```python
  def add(inputs, name=None):
    return gen_math_ops.add(inputs, name)
  ```

  to:

  ```python
  from tensorflow.python.compat import compat

  def add(inputs, name=None):
    if compat.forward_compatible(year, month, day):
      # Can use the awesome new implementation.
      return gen_math_ops.my_new_awesome_add(inputs, name)
    # To maintain forward compatibility, use the old implementation.
    return gen_math_ops.add(inputs, name)
  ```

  Where `year`, `month`, and `day` specify the date beyond which binaries
  that consume a model are expected to have been updated to include the
  new operations. This date is typically at least 3 weeks beyond the date
  the code that adds the new operation is committed.

  Args:
    year:  A year (e.g., 2018). Must be an `int`.
    month: A month (1 <= month <= 12) in year. Must be an `int`.
    day:   A day (1 <= day <= 31, or 30, or 29, or 28) in month. Must be an
      `int`.

  Returns:
    True if the caller can expect that serialized TensorFlow graphs produced
    can be consumed by programs that are compiled with the TensorFlow library
    source code after (year, month, day).
  """
  return _FORWARD_COMPATIBILITY_DATE_NUMBER > _date_to_date_number(
      year, month, day)

@tf_export("compat.forward_compatibility_horizon")
@tf_contextlib.contextmanager
def forward_compatibility_horizon(year, month, day):
  """Context manager for testing forward compatibility of generated graphs.

  See [Version
  compatibility](https://tensorflow.org/guide/version_compat#backward_forward).

  To ensure forward compatibility of generated graphs (see `forward_compatible`)
  with older binaries, new features can be gated with:

  ```python
  if compat.forward_compatible(year=2018, month=08, date=01):
    generate_graph_with_new_features()
  else:
    generate_graph_so_older_binaries_can_consume_it()
  ```

  However, when adding new features, one may want to unittest it before
  the forward compatibility window expires. This context manager enables
  such tests. For example:

  ```python
  from tensorflow.python.compat import compat

  def testMyNewFeature(self):
    with compat.forward_compatibility_horizon(2018, 08, 02):
       # Test that generate_graph_with_new_features() has an effect
  ```

  Args:
    year:  A year (e.g., 2018). Must be an `int`.
    month: A month (1 <= month <= 12) in year. Must be an `int`.
    day:   A day (1 <= day <= 31, or 30, or 29, or 28) in month. Must be an
      `int`.

  Yields:
    Nothing.
  """
  try:
    _update_forward_compatibility_date_number(datetime.date(year, month, day))
    yield
  finally:
    _update_forward_compatibility_date_number()
<|endoftext|>


<commit_before>import re
import os

from django.template.loader import render_to_string
from django.template import Template, Context
from django.contrib.auth.models import SiteProfileNotAvailable
from django.core.exceptions import ObjectDoesNotExist
from django.conf import settings

from doc_builder.base import BaseBuilder
from projects.utils import safe_write, run

RTD_CONF_ADDITIONS = """
#Add RTD Template Path.
if 'templates_path' in locals():
    templates_path.insert(0, '{{ template_path }}')
else:
    templates_path = ['{{ template_path }}', 'templates', '_templates', '.templates']

#Add RTD Static Path. Add to the end because it overwrites previous files.
if 'html_static_path' in locals():
    html_static_path.append('{{ static_path }}')
else:
    html_static_path = ['_static', '{{ static_path }}']

#Add RTD CSS File only if they aren't overriding it already
using_rtd_theme = False
if 'html_theme' in locals():
    if html_theme in ['default']:
        if not 'html_style' in locals():
            html_style = 'rtd.css'
            html_theme = 'default'
            html_theme_options = {}
            using_rtd_theme = True

#Add sponsorship and project information to the template context.
context = {
    'using_theme': using_rtd_theme,
    'versions': [{% for version in verisons %}
    ("{{ version.slug }}", "{{ version.get_absolute_url }}"),{% endfor %}
    ],
    'slug': '{{ project.slug }}',
    'badge_revsys': {{ project.sponsored }}
}
if 'html_context' in locals():
    html_context.update(context)
else:
    html_context = context
"""

TEMPLATE_DIR = '%s/readthedocs/templates/sphinx' % settings.SITE_ROOT
STATIC_DIR = '%s/_static' % TEMPLATE_DIR

class Builder(BaseBuilder):

    def _whitelisted(self, version):
        """Modify the given ``conf.py`` file from a whitelisted user's project.
        For now, this just adds the RTD template directory to ``templates_path``.
        """
        project = version.project
        #Open file for appending.
        try:
            outfile = open(project.conf_file(version.slug), 'a')
        except IndexError:
            print "Nothing to white list, no conf.py"
            return
        outfile.write("\n")
        rtd_ctx = Context({
            'verisons': project.active_versions(),
            'project': project,
            'static_path': STATIC_DIR,
            'template_path': TEMPLATE_DIR,
        })
        rtd_string = Template(RTD_CONF_ADDITIONS).render(rtd_ctx)
        outfile.write(rtd_string)
        outfile.close()

    def _sanitize(self, version):
        project = version.project
        conf_template =  render_to_string('sphinx/conf.py.conf',
                                          {'project': project,
                                           'template_dir': TEMPLATE_DIR,
                                            'badge': project.sponsored
                                            })
        rtd_ctx = Context({
            'verisons': project.active_versions(),
            'project': project,
            'static_path': STATIC_DIR,
            'template_path': TEMPLATE_DIR,
        })
        rtd_string = Template(RTD_CONF_ADDITIONS).render(rtd_ctx)
        conf_template = conf_template + "\n" + rtd_string
        safe_write(project.conf_file(version.slug), conf_template)

    def clean(self, version):
        try:
            if version.project.whitelisted:
                print "Project whitelisted"
                self._whitelisted(version)
            else:
                print "Writing conf to disk"
                self._sanitize(version)
        except (OSError, SiteProfileNotAvailable, ObjectDoesNotExist):
            try:
                print "Writing conf to disk on error."
                self._sanitize(version)
            except (OSError, IOError):
                print "Conf file not found. Error writing to disk."
                return ('','Conf file not found. Error writing to disk.',-1)

    def build(self, version):
        project = version.project
        os.chdir(version.project.conf_dir(version.slug))
        if project.use_virtualenv and project.whitelisted:
            build_command = '%s -b html . _build/html' % project.venv_bin(version=version.slug, bin='sphinx-build')
        else:
            build_command = "sphinx-build -b html . _build/html"
        build_results = run(build_command)
        return build_results
<commit_msg>Theme docs that don't have a theme set.<commit_after>import re
import os

from django.template.loader import render_to_string
from django.template import Template, Context
from django.contrib.auth.models import SiteProfileNotAvailable
from django.core.exceptions import ObjectDoesNotExist
from django.conf import settings

from doc_builder.base import BaseBuilder
from projects.utils import safe_write, run

RTD_CONF_ADDITIONS = """
#Add RTD Template Path.
if 'templates_path' in locals():
    templates_path.insert(0, '{{ template_path }}')
else:
    templates_path = ['{{ template_path }}', 'templates', '_templates', '.templates']

#Add RTD Static Path. Add to the end because it overwrites previous files.
if 'html_static_path' in locals():
    html_static_path.append('{{ static_path }}')
else:
    html_static_path = ['_static', '{{ static_path }}']

#Add RTD CSS File only if they aren't overriding it already
using_rtd_theme = False
if 'html_theme' in locals():
    if html_theme in ['default']:
        if not 'html_style' in locals():
            html_style = 'rtd.css'
            html_theme = 'default'
            html_theme_options = {}
            using_rtd_theme = True
else:
    html_style = 'rtd.css'
    html_theme = 'default'
    html_theme_options = {}
    using_rtd_theme = True

#Add sponsorship and project information to the template context.
context = {
    'using_theme': using_rtd_theme,
    'versions': [{% for version in verisons %}
    ("{{ version.slug }}", "{{ version.get_absolute_url }}"),{% endfor %}
    ],
    'slug': '{{ project.slug }}',
    'badge_revsys': {{ project.sponsored }}
}
if 'html_context' in locals():
    html_context.update(context)
else:
    html_context = context
"""

TEMPLATE_DIR = '%s/readthedocs/templates/sphinx' % settings.SITE_ROOT
STATIC_DIR = '%s/_static' % TEMPLATE_DIR

class Builder(BaseBuilder):

    def _whitelisted(self, version):
        """Modify the given ``conf.py`` file from a whitelisted user's project.
        For now, this just adds the RTD template directory to ``templates_path``.
        """
        project = version.project
        #Open file for appending.
        try:
            outfile = open(project.conf_file(version.slug), 'a')
        except IndexError:
            print "Nothing to white list, no conf.py"
            return
        outfile.write("\n")
        rtd_ctx = Context({
            'verisons': project.active_versions(),
            'project': project,
            'static_path': STATIC_DIR,
            'template_path': TEMPLATE_DIR,
        })
        rtd_string = Template(RTD_CONF_ADDITIONS).render(rtd_ctx)
        outfile.write(rtd_string)
        outfile.close()

    def _sanitize(self, version):
        project = version.project
        conf_template =  render_to_string('sphinx/conf.py.conf',
                                          {'project': project,
                                           'template_dir': TEMPLATE_DIR,
                                            'badge': project.sponsored
                                            })
        rtd_ctx = Context({
            'verisons': project.active_versions(),
            'project': project,
            'static_path': STATIC_DIR,
            'template_path': TEMPLATE_DIR,
        })
        rtd_string = Template(RTD_CONF_ADDITIONS).render(rtd_ctx)
        conf_template = conf_template + "\n" + rtd_string
        safe_write(project.conf_file(version.slug), conf_template)

    def clean(self, version):
        try:
            if version.project.whitelisted:
                print "Project whitelisted"
                self._whitelisted(version)
            else:
                print "Writing conf to disk"
                self._sanitize(version)
        except (OSError, SiteProfileNotAvailable, ObjectDoesNotExist):
            try:
                print "Writing conf to disk on error."
                self._sanitize(version)
            except (OSError, IOError):
                print "Conf file not found. Error writing to disk."
                return ('','Conf file not found. Error writing to disk.',-1)

    def build(self, version):
        project = version.project
        os.chdir(version.project.conf_dir(version.slug))
        if project.use_virtualenv and project.whitelisted:
            build_command = '%s -b html . _build/html' % project.venv_bin(version=version.slug, bin='sphinx-build')
        else:
            build_command = "sphinx-build -b html . _build/html"
        build_results = run(build_command)
        return build_results
<|endoftext|>


<commit_before># Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities for API compatibility between TensorFlow release versions.

See [Version
Compatibility](https://tensorflow.org/guide/version_compat#backward_forward)
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import datetime

from tensorflow.python import tf2
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import variable_scope

from tensorflow.python.util import tf_contextlib
from tensorflow.python.util.tf_export import tf_export

_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2018, 12, 12)

@tf_export("compat.forward_compatible")
def forward_compatible(year, month, day):
  """Return true if the forward compatibility window has expired.

  See [Version
  compatibility](https://tensorflow.org/guide/version_compat#backward_forward).

  Forward-compatibility refers to scenarios where the producer of a TensorFlow
  model (a GraphDef or SavedModel) is compiled against a version of the
  TensorFlow library newer than what the consumer was compiled against. The
  "producer" is typically a Python program that constructs and trains a model
  while the "consumer" is typically another program that loads and serves the
  model.

  TensorFlow has been supporting a 3 week forward-compatibility window for
  programs compiled from source at HEAD.

  For example, consider the case where a new operation `MyNewAwesomeAdd` is
  created with the intent of replacing the implementation of an existing Python
  wrapper - `tf.add`.  The Python wrapper implementation should change from
  something like:

  ```python
  def add(inputs, name=None):
    return gen_math_ops.add(inputs, name)
  ```

  to:

  ```python
  from tensorflow.python.compat import compat

  def add(inputs, name=None):
    if compat.forward_compatible(year, month, day):
      # Can use the awesome new implementation.
      return gen_math_ops.my_new_awesome_add(inputs, name)
    # To maintain forward compatibiltiy, use the old implementation.
    return gen_math_ops.add(inputs, name)
  ```

  Where `year`, `month`, and `day` specify the date beyond which binaries
  that consume a model are expected to have been updated to include the
  new operations. This date is typically at least 3 weeks beyond the date
  the code that adds the new operation is committed.

  Args:
    year:  A year (e.g., 2018).
    month: A month (1 <= month <= 12) in year.
    day:   A day (1 <= day <= 31, or 30, or 29, or 28) in month.

  Returns:
    True if the caller can expect that serialized TensorFlow graphs produced
    can be consumed by programs that are compiled with the TensorFlow library
    source code after (year, month, day).
  """
  return _FORWARD_COMPATIBILITY_HORIZON > datetime.date(year, month, day)

@tf_export("compat.forward_compatibility_horizon")
@tf_contextlib.contextmanager
def forward_compatibility_horizon(year, month, day):
  """Context manager for testing forward compatibility of generated graphs.

  See [Version
  compatibility](https://tensorflow.org/guide/version_compat#backward_forward).

  To ensure forward compatibility of generated graphs (see `forward_compatible`)
  with older binaries, new features can be gated with:

  ```python
  if compat.forward_compatible(year=2018, month=08, date=01):
    generate_graph_with_new_features()
  else:
    generate_graph_so_older_binaries_can_consume_it()
  ```

  However, when adding new features, one may want to unittest it before
  the forward compatibility window expires. This context manager enables
  such tests. For example:

  ```python
  from tensorflow.python.compat import compat

  def testMyNewFeature(self):
    with compat.forward_compatibility_horizon(2018, 08, 02):
       # Test that generate_graph_with_new_features() has an effect
  ```

  Args :
    year:  A year (e.g. 2018).
    month: A month (1 <= month <= 12) in year.
    day:   A day (1 <= day <= 31, or 30, or 29, or 28) in month.

  Yields:
    Nothing.
  """
  global _FORWARD_COMPATIBILITY_HORIZON
  try:
    old_compat_date = _FORWARD_COMPATIBILITY_HORIZON
    _FORWARD_COMPATIBILITY_HORIZON = datetime.date(year, month, day)
    yield
  finally:
    _FORWARD_COMPATIBILITY_HORIZON = old_compat_date

@tf_export(v1=["enable_v2_behavior"])
def enable_v2_behavior():
  """Enables TensorFlow 2.x behaviors.

  This function can be called at the beginning of the program (before `Tensors`,
  `Graphs` or other structures have been created, and before devices have been
  initialized. It switches all global behaviors that are different between
  TensorFlow 1.x and 2.x to behave as intended for 2.x.

  This function is called in the main TensorFlow `__init__.py` file, user should
  not need to call it, except during complex migrations.
  """
  tf2.enable()  # Switches TensorArrayV2 and control flow V2
  ops.enable_eager_execution()
  tensor_shape.enable_v2_tensorshape()  # Also switched by tf2
  variable_scope.enable_resource_variables()

@tf_export(v1=["disable_v2_behavior"])
def disable_v2_behavior():
  """Enables TensorFlow 2.x behaviors.

  This function can be called at the beginning of the program (before `Tensors`,
  `Graphs` or other structures have been created, and before devices have been
  initialized. It switches all global behaviors that are different between
  TensorFlow 1.x and 2.x to behave as intended for 1.x.

  User can call this function to disable 2.x behavior during complex migrations.
  """
  tf2.disable()  # Switches TensorArrayV2 and control flow V2
  ops.disable_eager_execution()
  tensor_shape.disable_v2_tensorshape()  # Also switched by tf2
  variable_scope.disable_resource_variables()

<commit_msg>compat: Update forward compatibility horizon to 2018-12-13<commit_after># Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities for API compatibility between TensorFlow release versions.

See [Version
Compatibility](https://tensorflow.org/guide/version_compat#backward_forward)
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import datetime

from tensorflow.python import tf2
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import variable_scope

from tensorflow.python.util import tf_contextlib
from tensorflow.python.util.tf_export import tf_export

_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2018, 12, 13)

@tf_export("compat.forward_compatible")
def forward_compatible(year, month, day):
  """Return true if the forward compatibility window has expired.

  See [Version
  compatibility](https://tensorflow.org/guide/version_compat#backward_forward).

  Forward-compatibility refers to scenarios where the producer of a TensorFlow
  model (a GraphDef or SavedModel) is compiled against a version of the
  TensorFlow library newer than what the consumer was compiled against. The
  "producer" is typically a Python program that constructs and trains a model
  while the "consumer" is typically another program that loads and serves the
  model.

  TensorFlow has been supporting a 3 week forward-compatibility window for
  programs compiled from source at HEAD.

  For example, consider the case where a new operation `MyNewAwesomeAdd` is
  created with the intent of replacing the implementation of an existing Python
  wrapper - `tf.add`.  The Python wrapper implementation should change from
  something like:

  ```python
  def add(inputs, name=None):
    return gen_math_ops.add(inputs, name)
  ```

  to:

  ```python
  from tensorflow.python.compat import compat

  def add(inputs, name=None):
    if compat.forward_compatible(year, month, day):
      # Can use the awesome new implementation.
      return gen_math_ops.my_new_awesome_add(inputs, name)
    # To maintain forward compatibiltiy, use the old implementation.
    return gen_math_ops.add(inputs, name)
  ```

  Where `year`, `month`, and `day` specify the date beyond which binaries
  that consume a model are expected to have been updated to include the
  new operations. This date is typically at least 3 weeks beyond the date
  the code that adds the new operation is committed.

  Args:
    year:  A year (e.g., 2018).
    month: A month (1 <= month <= 12) in year.
    day:   A day (1 <= day <= 31, or 30, or 29, or 28) in month.

  Returns:
    True if the caller can expect that serialized TensorFlow graphs produced
    can be consumed by programs that are compiled with the TensorFlow library
    source code after (year, month, day).
  """
  return _FORWARD_COMPATIBILITY_HORIZON > datetime.date(year, month, day)

@tf_export("compat.forward_compatibility_horizon")
@tf_contextlib.contextmanager
def forward_compatibility_horizon(year, month, day):
  """Context manager for testing forward compatibility of generated graphs.

  See [Version
  compatibility](https://tensorflow.org/guide/version_compat#backward_forward).

  To ensure forward compatibility of generated graphs (see `forward_compatible`)
  with older binaries, new features can be gated with:

  ```python
  if compat.forward_compatible(year=2018, month=08, date=01):
    generate_graph_with_new_features()
  else:
    generate_graph_so_older_binaries_can_consume_it()
  ```

  However, when adding new features, one may want to unittest it before
  the forward compatibility window expires. This context manager enables
  such tests. For example:

  ```python
  from tensorflow.python.compat import compat

  def testMyNewFeature(self):
    with compat.forward_compatibility_horizon(2018, 08, 02):
       # Test that generate_graph_with_new_features() has an effect
  ```

  Args :
    year:  A year (e.g. 2018).
    month: A month (1 <= month <= 12) in year.
    day:   A day (1 <= day <= 31, or 30, or 29, or 28) in month.

  Yields:
    Nothing.
  """
  global _FORWARD_COMPATIBILITY_HORIZON
  try:
    old_compat_date = _FORWARD_COMPATIBILITY_HORIZON
    _FORWARD_COMPATIBILITY_HORIZON = datetime.date(year, month, day)
    yield
  finally:
    _FORWARD_COMPATIBILITY_HORIZON = old_compat_date

@tf_export(v1=["enable_v2_behavior"])
def enable_v2_behavior():
  """Enables TensorFlow 2.x behaviors.

  This function can be called at the beginning of the program (before `Tensors`,
  `Graphs` or other structures have been created, and before devices have been
  initialized. It switches all global behaviors that are different between
  TensorFlow 1.x and 2.x to behave as intended for 2.x.

  This function is called in the main TensorFlow `__init__.py` file, user should
  not need to call it, except during complex migrations.
  """
  tf2.enable()  # Switches TensorArrayV2 and control flow V2
  ops.enable_eager_execution()
  tensor_shape.enable_v2_tensorshape()  # Also switched by tf2
  variable_scope.enable_resource_variables()

@tf_export(v1=["disable_v2_behavior"])
def disable_v2_behavior():
  """Enables TensorFlow 2.x behaviors.

  This function can be called at the beginning of the program (before `Tensors`,
  `Graphs` or other structures have been created, and before devices have been
  initialized. It switches all global behaviors that are different between
  TensorFlow 1.x and 2.x to behave as intended for 1.x.

  User can call this function to disable 2.x behavior during complex migrations.
  """
  tf2.disable()  # Switches TensorArrayV2 and control flow V2
  ops.disable_eager_execution()
  tensor_shape.disable_v2_tensorshape()  # Also switched by tf2
  variable_scope.disable_resource_variables()

<|endoftext|>


<commit_before>"""Django admin interface for `~builds.models.Build` and related models."""

from django.contrib import admin, messages
from guardian.admin import GuardedModelAdmin
from polymorphic.admin import (
    PolymorphicChildModelAdmin,
    PolymorphicParentModelAdmin,
)

from readthedocs.builds.models import (
    Build,
    BuildCommandResult,
    RegexAutomationRule,
    Version,
    VersionAutomationRule,
)
from readthedocs.core.utils import trigger_build
from readthedocs.core.utils.general import wipe_version_via_slugs
from readthedocs.projects.models import HTMLFile
from readthedocs.search.utils import _indexing_helper

class BuildCommandResultInline(admin.TabularInline):
    model = BuildCommandResult
    fields = ('command', 'exit_code', 'output')

class BuildAdmin(admin.ModelAdmin):
    fields = (
        'project',
        'version',
        'type',
        'state',
        'error',
        'success',
        'length',
        'cold_storage',
    )
    list_display = (
        'id',
        'project',
        'version_name',
        'success',
        'type',
        'state',
        'date',
    )
    list_filter = ('type', 'state', 'success')
    list_select_related = ('project', 'version')
    raw_id_fields = ('project', 'version')
    inlines = (BuildCommandResultInline,)
    search_fields = ('project__name', 'version__slug')

    def version_name(self, obj):
        return obj.version.verbose_name

class VersionAdmin(GuardedModelAdmin):
    list_display = (
        'slug',
        'type',
        'project',
        'privacy_level',
        'active',
        'built',
    )
    list_filter = ('type', 'privacy_level', 'active', 'built')
    search_fields = ('slug', 'project__slug', 'project__name')
    raw_id_fields = ('project',)
    actions = ['build_version', 'reindex_version', 'wipe_version', 'wipe_selected_versions']

    def wipe_selected_versions(self, request, queryset):
        """Wipes the selected versions."""
        for version in queryset:
            wipe_version_via_slugs(
                version_slug=version.slug,
                project_slug=version.project.slug
            )
            self.message_user(
                request,
                'Wiped {}.'.format(version.slug),
                level=messages.SUCCESS
            )

    wipe_selected_versions.short_description = 'Wipe selected versions'

    def build_version(self, request, queryset):
        """Trigger a build for the project version."""
        total = 0
        for version in queryset:
            trigger_build(
                project=version.project,
                version=version,
            )
            total += 1
        messages.add_message(
            request,
            messages.INFO,
            'Triggered builds for {} version(s).'.format(total),
        )

    build_version.short_description = 'Build version'

    def reindex_version(self, request, queryset):
        """Reindexes all selected versions to ES."""
        html_objs_qs = []
        for version in queryset.iterator():
            html_objs = HTMLFile.objects.filter(project=version.project, version=version)

            if html_objs.exists():
                html_objs_qs.append(html_objs)

        if html_objs_qs:
            _indexing_helper(html_objs_qs, wipe=False)

        self.message_user(
            request,
            'Task initiated successfully.',
            messages.SUCCESS
        )

    reindex_version.short_description = 'Reindex version to ES'

    def wipe_version(self, request, queryset):
        """Wipe selected versions from ES."""
        html_objs_qs = []
        for version in queryset.iterator():
            html_objs = HTMLFile.objects.filter(project=version.project, version=version)

            if html_objs.exists():
                html_objs_qs.append(html_objs)

        if html_objs_qs:
            _indexing_helper(html_objs_qs, wipe=True)

        self.message_user(
            request,
            'Task initiated successfully',
            messages.SUCCESS,
        )

    wipe_version.short_description = 'Wipe version from ES'

@admin.register(RegexAutomationRule)
class RegexAutomationRuleAdmin(PolymorphicChildModelAdmin, admin.ModelAdmin):
    base_model = RegexAutomationRule

@admin.register(VersionAutomationRule)
class VersionAutomationRuleAdmin(PolymorphicParentModelAdmin, admin.ModelAdmin):
    base_model = VersionAutomationRule
    list_display = (
        'project',
        'priority',
        'match_arg',
        'action',
        'version_type',
    )
    child_models = (RegexAutomationRule,)

admin.site.register(Build, BuildAdmin)
admin.site.register(Version, VersionAdmin)
<commit_msg>Don't search for project.name<commit_after>"""Django admin interface for `~builds.models.Build` and related models."""

from django.contrib import admin, messages
from guardian.admin import GuardedModelAdmin
from polymorphic.admin import (
    PolymorphicChildModelAdmin,
    PolymorphicParentModelAdmin,
)

from readthedocs.builds.models import (
    Build,
    BuildCommandResult,
    RegexAutomationRule,
    Version,
    VersionAutomationRule,
)
from readthedocs.core.utils import trigger_build
from readthedocs.core.utils.general import wipe_version_via_slugs
from readthedocs.projects.models import HTMLFile
from readthedocs.search.utils import _indexing_helper

class BuildCommandResultInline(admin.TabularInline):
    model = BuildCommandResult
    fields = ('command', 'exit_code', 'output')

class BuildAdmin(admin.ModelAdmin):
    fields = (
        'project',
        'version',
        'type',
        'state',
        'error',
        'success',
        'length',
        'cold_storage',
    )
    list_display = (
        'id',
        'project',
        'version_name',
        'success',
        'type',
        'state',
        'date',
    )
    list_filter = ('type', 'state', 'success')
    list_select_related = ('project', 'version')
    raw_id_fields = ('project', 'version')
    inlines = (BuildCommandResultInline,)
    search_fields = ('project__name', 'version__slug')

    def version_name(self, obj):
        return obj.version.verbose_name

class VersionAdmin(GuardedModelAdmin):
    list_display = (
        'slug',
        'type',
        'project',
        'privacy_level',
        'active',
        'built',
    )
    list_filter = ('type', 'privacy_level', 'active', 'built')
    search_fields = ('slug', 'project__slug')
    raw_id_fields = ('project',)
    actions = ['build_version', 'reindex_version', 'wipe_version', 'wipe_selected_versions']

    def wipe_selected_versions(self, request, queryset):
        """Wipes the selected versions."""
        for version in queryset:
            wipe_version_via_slugs(
                version_slug=version.slug,
                project_slug=version.project.slug
            )
            self.message_user(
                request,
                'Wiped {}.'.format(version.slug),
                level=messages.SUCCESS
            )

    wipe_selected_versions.short_description = 'Wipe selected versions'

    def build_version(self, request, queryset):
        """Trigger a build for the project version."""
        total = 0
        for version in queryset:
            trigger_build(
                project=version.project,
                version=version,
            )
            total += 1
        messages.add_message(
            request,
            messages.INFO,
            'Triggered builds for {} version(s).'.format(total),
        )

    build_version.short_description = 'Build version'

    def reindex_version(self, request, queryset):
        """Reindexes all selected versions to ES."""
        html_objs_qs = []
        for version in queryset.iterator():
            html_objs = HTMLFile.objects.filter(project=version.project, version=version)

            if html_objs.exists():
                html_objs_qs.append(html_objs)

        if html_objs_qs:
            _indexing_helper(html_objs_qs, wipe=False)

        self.message_user(
            request,
            'Task initiated successfully.',
            messages.SUCCESS
        )

    reindex_version.short_description = 'Reindex version to ES'

    def wipe_version(self, request, queryset):
        """Wipe selected versions from ES."""
        html_objs_qs = []
        for version in queryset.iterator():
            html_objs = HTMLFile.objects.filter(project=version.project, version=version)

            if html_objs.exists():
                html_objs_qs.append(html_objs)

        if html_objs_qs:
            _indexing_helper(html_objs_qs, wipe=True)

        self.message_user(
            request,
            'Task initiated successfully',
            messages.SUCCESS,
        )

    wipe_version.short_description = 'Wipe version from ES'

@admin.register(RegexAutomationRule)
class RegexAutomationRuleAdmin(PolymorphicChildModelAdmin, admin.ModelAdmin):
    base_model = RegexAutomationRule

@admin.register(VersionAutomationRule)
class VersionAutomationRuleAdmin(PolymorphicParentModelAdmin, admin.ModelAdmin):
    base_model = VersionAutomationRule
    list_display = (
        'project',
        'priority',
        'match_arg',
        'action',
        'version_type',
    )
    child_models = (RegexAutomationRule,)

admin.site.register(Build, BuildAdmin)
admin.site.register(Version, VersionAdmin)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# This file is part of REANA.
# Copyright (C) 2017 CERN.
#
# REANA 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.
#
# REANA 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
# REANA; if not, write to the Free Software Foundation, Inc., 59 Temple Place,
# Suite 330, Boston, MA 02111-1307, USA.
#
# In applying this license, CERN does not waive the privileges and immunities
# granted to it by virtue of its status as an Intergovernmental Organization or
# submit itself to any jurisdiction.

"""Kubernetes wrapper."""

import json
import logging
import os
import time
import traceback

from flask import current_app as app
from kubernetes import client
from kubernetes import config as k8s_config
from kubernetes import watch
from kubernetes.client.models.v1_delete_options import V1DeleteOptions
from reana_commons.database import Session
from reana_commons.models import Job

from reana_job_controller import config, volume_templates

def create_api_client(api='BatchV1'):
    """Create pykube HTTPClient using config.

    :param api: String which represents which Kubernetes API to spawn. By
        default BatchV1.
    :returns: Kubernetes python client object for a specific API i.e. BatchV1.
    """
    k8s_config.load_incluster_config()
    api_configuration = client.Configuration()
    api_configuration.verify_ssl = False
    if api == 'CoreV1':
        api_client = client.CoreV1Api()
    else:
        api_client = client.BatchV1Api()
    return api_client

def add_shared_volume(job):
    """Add shared CephFS volume to a given job spec.

    :param job: Kubernetes job spec.
    """
    storage_backend = os.getenv('REANA_STORAGE_BACKEND', 'LOCAL')
    if storage_backend == 'CEPHFS':
        volume = volume_templates.get_k8s_cephfs_volume(
            job['metadata']['namespace'])
    else:
        volume = volume_templates.get_k8s_hostpath_volume(
            job['metadata']['namespace'])
    mount_path = volume_templates.REANA_STORAGE_MOUNT_PATH

    job['spec']['template']['spec']['containers'][0]['volumeMounts'].append(
        {'name': volume['name'], 'mountPath': mount_path}
    )
    job['spec']['template']['spec']['volumes'].append(volume)

def instantiate_job(job_id, docker_img, cmd, cvmfs_repos, env_vars, namespace,
                    shared_file_system, job_type):
    """Create Kubernetes job.

    :param job_id: Job uuid.
    :param docker_img: Docker image to run the job.
    :param cmd: Command provided to the docker container.
    :param cvmfs_repos: List of CVMFS repository names.
    :param env_vars: Dictionary representing environment variables
        as {'var_name': 'var_value'}.
    :param namespace: Job's namespace.
    :shared_file_system: Boolean which represents whether the job
        should have a shared file system mounted.
    :returns: Kubernetes job object if the job was successfuly created,
        None if not.
    """
    job = {
        'kind': 'Job',
        'apiVersion': 'batch/v1',
        'metadata': {
            'name': job_id,
            'namespace': namespace
        },
        'spec': {
            'backoffLimit': app.config['MAX_JOB_RESTARTS'],
            'autoSelector': True,
            'template': {
                'metadata': {
                    'name': job_id
                },
                'spec': {
                    'containers': [
                        {
                            'name': job_id,
                            'image': docker_img,
                            'env': [],
                            'volumeMounts': []
                        },
                    ],
                    'volumes': [],
                    'restartPolicy': 'Never'
                }
            }
        }
    }

    if cmd:
        import shlex
        (job['spec']['template']['spec']['containers']
         [0]['command']) = shlex.split(cmd)

    if env_vars:
        for var, value in env_vars.items():
            job['spec']['template']['spec']['containers'][0]['env'].append(
                {'name': var, 'value': value}
            )

    if shared_file_system:
        add_shared_volume(job)

    if cvmfs_repos:
        for num, repo in enumerate(cvmfs_repos):
            volume = volume_templates.get_k8s_cvmfs_volume(namespace, repo)
            mount_path = volume_templates.get_cvmfs_mount_point(repo)

            volume['name'] += '-{}'.format(num)
            (job['spec']['template']['spec']['containers'][0]
                ['volumeMounts'].append(
                    {'name': volume['name'], 'mountPath': mount_path}
            ))
            job['spec']['template']['spec']['volumes'].append(volume)

    # add better handling
    try:
        api_response = \
            app.config['KUBERNETES_CLIENT'].create_namespaced_job(
                namespace=namespace, body=job)
        return api_response.to_str()
    except client.rest.ApiException as e:
        logging.debug("Error while connecting to Kubernetes API: {}".format(e))
    except Exception as e:
        logging.error(traceback.format_exc())
        logging.debug("Unexpected error: {}".format(e))

def watch_jobs(job_db):
    """Open stream connection to k8s apiserver to watch all jobs status.

    :param job_db: Dictionary which contains all current jobs.
    :param config: configuration to connect to k8s apiserver.
    """
    batchv1_api_client = create_api_client()
    corev1_api_client = create_api_client('CoreV1')
    while True:
        logging.debug('Starting a new stream request to watch Jobs')
        try:
            w = watch.Watch()
            for event in w.stream(
                    batchv1_api_client.list_job_for_all_namespaces,
                    _request_timeout=60):
                logging.info(
                    'New Job event received: {0}'.format(event['type']))
                job = event['object']

                # Taking note of the remaining jobs since deletion might not
                # happend straight away.
                remaining_jobs = [j for j in job_db.keys()
                                  if not job_db[j]['deleted']]
                if (not job_db.get(job.metadata.name) or
                        job.metadata.name not in remaining_jobs):
                    # Ignore jobs not created by this specific instance
                    # or already deleted jobs.
                    continue
                elif job.status.succeeded:
                    logging.info(
                        'Job {} succeeded.'.format(
                            job.metadata.name)
                    )
                    job_db[job.metadata.name]['status'] = 'succeeded'
                elif (job.status.failed and
                      job.status.failed >= config.MAX_JOB_RESTARTS):
                    logging.info('Job {} failed.'.format(
                        job.metadata.name))
                    job_db[job.metadata.name]['status'] = 'failed'
                else:
                    continue
                # Grab logs when job either succeeds or fails.
                logging.info('Getting last spawned pod for job {}'.format(
                    job.metadata.name))
                last_spawned_pod = corev1_api_client.list_namespaced_pod(
                    job.metadata.namespace,
                    label_selector='job-name={job_name}'.format(
                        job_name=job.metadata.name)).items[-1]
                logging.info('Grabbing pod {} logs...'.format(
                    last_spawned_pod.metadata.name))
                job_db[job.metadata.name]['log'] = \
                    corev1_api_client.read_namespaced_pod_log(
                        namespace=last_spawned_pod.metadata.namespace,
                        name=last_spawned_pod.metadata.name)
                # Store job logs
                try:
                    logging.info('Storing job logs: {}'.
                                 format(job_db[job.metadata.name]['log']))
                    Session.query(Job).filter_by(id_=job.metadata.name).\
                        update(dict(logs=job_db[job.metadata.name]['log']))
                    Session.commit()

                except Exception as e:
                    logging.debug('Could not retrieve'
                                  ' logs for object: {}'.
                                  format(last_spawned_pod))
                    logging.debug('Exception: {}'.format(str(e)))

                logging.info('Cleaning job {} ...'.format(
                    job.metadata.name))
                # Delete all depending pods.
                delete_options = V1DeleteOptions(
                    propagation_policy='Foreground')
                batchv1_api_client.delete_namespaced_job(
                    job.metadata.name, job.metadata.namespace, delete_options)
                job_db[job.metadata.name]['deleted'] = True
        except client.rest.ApiException as e:
            logging.debug(
                "Error while connecting to Kubernetes API: {}".format(e))
        except Exception as e:
            logging.error(traceback.format_exc())
            logging.debug("Unexpected error: {}".format(e))
<commit_msg>k8s: fix timeout problem<commit_after># -*- coding: utf-8 -*-
#
# This file is part of REANA.
# Copyright (C) 2017 CERN.
#
# REANA 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.
#
# REANA 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
# REANA; if not, write to the Free Software Foundation, Inc., 59 Temple Place,
# Suite 330, Boston, MA 02111-1307, USA.
#
# In applying this license, CERN does not waive the privileges and immunities
# granted to it by virtue of its status as an Intergovernmental Organization or
# submit itself to any jurisdiction.

"""Kubernetes wrapper."""

import json
import logging
import os
import time
import traceback

from flask import current_app as app
from kubernetes import client
from kubernetes import config as k8s_config
from kubernetes import watch
from kubernetes.client.models.v1_delete_options import V1DeleteOptions
from reana_commons.database import Session
from reana_commons.models import Job

from reana_job_controller import config, volume_templates

def create_api_client(api='BatchV1'):
    """Create pykube HTTPClient using config.

    :param api: String which represents which Kubernetes API to spawn. By
        default BatchV1.
    :returns: Kubernetes python client object for a specific API i.e. BatchV1.
    """
    k8s_config.load_incluster_config()
    api_configuration = client.Configuration()
    api_configuration.verify_ssl = False
    if api == 'CoreV1':
        api_client = client.CoreV1Api()
    else:
        api_client = client.BatchV1Api()
    return api_client

def add_shared_volume(job):
    """Add shared CephFS volume to a given job spec.

    :param job: Kubernetes job spec.
    """
    storage_backend = os.getenv('REANA_STORAGE_BACKEND', 'LOCAL')
    if storage_backend == 'CEPHFS':
        volume = volume_templates.get_k8s_cephfs_volume(
            job['metadata']['namespace'])
    else:
        volume = volume_templates.get_k8s_hostpath_volume(
            job['metadata']['namespace'])
    mount_path = volume_templates.REANA_STORAGE_MOUNT_PATH

    job['spec']['template']['spec']['containers'][0]['volumeMounts'].append(
        {'name': volume['name'], 'mountPath': mount_path}
    )
    job['spec']['template']['spec']['volumes'].append(volume)

def instantiate_job(job_id, docker_img, cmd, cvmfs_repos, env_vars, namespace,
                    shared_file_system, job_type):
    """Create Kubernetes job.

    :param job_id: Job uuid.
    :param docker_img: Docker image to run the job.
    :param cmd: Command provided to the docker container.
    :param cvmfs_repos: List of CVMFS repository names.
    :param env_vars: Dictionary representing environment variables
        as {'var_name': 'var_value'}.
    :param namespace: Job's namespace.
    :shared_file_system: Boolean which represents whether the job
        should have a shared file system mounted.
    :returns: Kubernetes job object if the job was successfuly created,
        None if not.
    """
    job = {
        'kind': 'Job',
        'apiVersion': 'batch/v1',
        'metadata': {
            'name': job_id,
            'namespace': namespace
        },
        'spec': {
            'backoffLimit': app.config['MAX_JOB_RESTARTS'],
            'autoSelector': True,
            'template': {
                'metadata': {
                    'name': job_id
                },
                'spec': {
                    'containers': [
                        {
                            'name': job_id,
                            'image': docker_img,
                            'env': [],
                            'volumeMounts': []
                        },
                    ],
                    'volumes': [],
                    'restartPolicy': 'Never'
                }
            }
        }
    }

    if cmd:
        import shlex
        (job['spec']['template']['spec']['containers']
         [0]['command']) = shlex.split(cmd)

    if env_vars:
        for var, value in env_vars.items():
            job['spec']['template']['spec']['containers'][0]['env'].append(
                {'name': var, 'value': value}
            )

    if shared_file_system:
        add_shared_volume(job)

    if cvmfs_repos:
        for num, repo in enumerate(cvmfs_repos):
            volume = volume_templates.get_k8s_cvmfs_volume(namespace, repo)
            mount_path = volume_templates.get_cvmfs_mount_point(repo)

            volume['name'] += '-{}'.format(num)
            (job['spec']['template']['spec']['containers'][0]
                ['volumeMounts'].append(
                    {'name': volume['name'], 'mountPath': mount_path}
            ))
            job['spec']['template']['spec']['volumes'].append(volume)

    # add better handling
    try:
        api_response = \
            app.config['KUBERNETES_CLIENT'].create_namespaced_job(
                namespace=namespace, body=job)
        return api_response.to_str()
    except client.rest.ApiException as e:
        logging.debug("Error while connecting to Kubernetes API: {}".format(e))
    except Exception as e:
        logging.error(traceback.format_exc())
        logging.debug("Unexpected error: {}".format(e))

def watch_jobs(job_db):
    """Open stream connection to k8s apiserver to watch all jobs status.

    :param job_db: Dictionary which contains all current jobs.
    :param config: configuration to connect to k8s apiserver.
    """
    batchv1_api_client = create_api_client()
    corev1_api_client = create_api_client('CoreV1')
    while True:
        logging.debug('Starting a new stream request to watch Jobs')
        try:
            w = watch.Watch()
            for event in w.stream(
                    batchv1_api_client.list_job_for_all_namespaces):
                logging.info(
                    'New Job event received: {0}'.format(event['type']))
                job = event['object']

                # Taking note of the remaining jobs since deletion might not
                # happend straight away.
                remaining_jobs = [j for j in job_db.keys()
                                  if not job_db[j]['deleted']]
                if (not job_db.get(job.metadata.name) or
                        job.metadata.name not in remaining_jobs):
                    # Ignore jobs not created by this specific instance
                    # or already deleted jobs.
                    continue
                elif job.status.succeeded:
                    logging.info(
                        'Job {} succeeded.'.format(
                            job.metadata.name)
                    )
                    job_db[job.metadata.name]['status'] = 'succeeded'
                elif (job.status.failed and
                      job.status.failed >= config.MAX_JOB_RESTARTS):
                    logging.info('Job {} failed.'.format(
                        job.metadata.name))
                    job_db[job.metadata.name]['status'] = 'failed'
                else:
                    continue
                # Grab logs when job either succeeds or fails.
                logging.info('Getting last spawned pod for job {}'.format(
                    job.metadata.name))
                last_spawned_pod = corev1_api_client.list_namespaced_pod(
                    job.metadata.namespace,
                    label_selector='job-name={job_name}'.format(
                        job_name=job.metadata.name)).items[-1]
                logging.info('Grabbing pod {} logs...'.format(
                    last_spawned_pod.metadata.name))
                job_db[job.metadata.name]['log'] = \
                    corev1_api_client.read_namespaced_pod_log(
                        namespace=last_spawned_pod.metadata.namespace,
                        name=last_spawned_pod.metadata.name)
                # Store job logs
                try:
                    logging.info('Storing job logs: {}'.
                                 format(job_db[job.metadata.name]['log']))
                    Session.query(Job).filter_by(id_=job.metadata.name).\
                        update(dict(logs=job_db[job.metadata.name]['log']))
                    Session.commit()

                except Exception as e:
                    logging.debug('Could not retrieve'
                                  ' logs for object: {}'.
                                  format(last_spawned_pod))
                    logging.debug('Exception: {}'.format(str(e)))

                logging.info('Cleaning job {} ...'.format(
                    job.metadata.name))
                # Delete all depending pods.
                delete_options = V1DeleteOptions(
                    propagation_policy='Background')
                batchv1_api_client.delete_namespaced_job(
                    job.metadata.name, job.metadata.namespace, delete_options)
                job_db[job.metadata.name]['deleted'] = True
        except client.rest.ApiException as e:
            logging.debug(
                "Error while connecting to Kubernetes API: {}".format(e))
        except Exception as e:
            logging.error(traceback.format_exc())
            logging.debug("Unexpected error: {}".format(e))
<|endoftext|>