<commit_before>import os
import romanesco
from . import pyspark_executor, spark

SC_KEY = '_romanesco_spark_context'


def setup_pyspark_task(event):
    """
    This is executed before a task execution. If it is a pyspark task, we
    create the spark context here so it can be used for any input conversion.
    """
    info = event.info
    if info['mode'] == 'spark.python' and SC_KEY not in info['kwargs']:
        spark_conf = info['task'].get('spark_conf', {})
        info['kwargs'][SC_KEY] = spark.create_spark_context(spark_conf)


def pyspark_run_cleanup(event):
    if SC_KEY in event.info['kwargs']:
        event.info['kwargs'][SC_KEY].stop()


def load(params):
    # If we have a spark config section then try to setup spark environment
    if romanesco.config.has_section('spark') or 'SPARK_HOME' in os.environ:
        spark.setup_spark_env()

    romanesco.register_executor('spark.python', pyspark_executor.run)

    romanesco.events.bind('run.before', 'spark', setup_pyspark_task)
    romanesco.events.bind('run.finally', 'spark', pyspark_run_cleanup)
<commit_msg>Fix spark mode (faulty shutdown conditional logic)<commit_after>import os
import romanesco
from . import pyspark_executor, spark

SC_KEY = '_romanesco_spark_context'


def setup_pyspark_task(event):
    """
    This is executed before a task execution. If it is a pyspark task, we
    create the spark context here so it can be used for any input conversion.
    """
    info = event.info
    if info['mode'] == 'spark.python' and SC_KEY not in info['kwargs']:
        spark_conf = info['task'].get('spark_conf', {})
        info['kwargs'][SC_KEY] = spark.create_spark_context(spark_conf)
        info['cleanup_spark'] = True


def pyspark_run_cleanup(event):
    if event.info.get('cleanup_spark'):
        event.info['kwargs'][SC_KEY].stop()


def load(params):
    # If we have a spark config section then try to setup spark environment
    if romanesco.config.has_section('spark') or 'SPARK_HOME' in os.environ:
        spark.setup_spark_env()

    romanesco.register_executor('spark.python', pyspark_executor.run)

    romanesco.events.bind('run.before', 'spark', setup_pyspark_task)
    romanesco.events.bind('run.finally', 'spark', pyspark_run_cleanup)

{'fraction_non_alphanumeric': 0.11492204899777284, 'fraction_numerical': 0.0, 'mean_word_length': 4.0699774266365685, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2851862', 'n_tokens_mistral': 746, 'n_tokens_neox': 717, 'n_words': 188}