<commit_before>import collections.abc as collections_abc
import collections
import logging

from bndl.execute.job import Job, Stage, Task
from functools import partial
import traceback


logger = logging.getLogger(__name__)


def schedule_job(dset, workers=None):
    '''
    Schedule a job for a data set
    :param dset:
        The data set to schedule
    '''

    ctx = dset.ctx
    assert ctx.running, 'context of dataset is not running'

    if not workers:
        ctx._await_workers()
        workers = ctx.workers[:]

    job = Job(ctx, *_job_calling_info())

    stage = Stage(None, job)
    schedule_stage(stage, workers, dset)
    job.stages.insert(0, stage)

    def _cleaner(dset, job):
        if job.stopped:
            dset.cleanup(job)

    while dset:
        if dset.cleanup:
            job.add_listener(partial(_cleaner, dset))
        if isinstance(dset.src, collections_abc.Iterable):
            for src in dset.src:
                branch = schedule_job(src)

                for task in branch.stages[-1].tasks:
                    task.args[1] = False

                for l in branch.listeners:
                    job.add_listener(l)

                if dset.sync_required:
                    branch_stages = branch.stages
                elif len(branch.stages) > 1:
                    branch_stages = branch.stages[:-1]
                else:
                    continue

                # branch_stages[-1].next_stage = job.stages[-1]
                for stage in reversed(branch_stages):
                    stage.job = job
                    stage.is_last = False
                    job.stages.insert(0, stage)

            break

        elif dset.sync_required:
            stage = stage.prev_stage = Stage(None, job)  # , next_stage=stage)
            schedule_stage(stage, workers, dset)
            job.stages.insert(0, stage)

        dset = dset.src

    # Since stages are added in reverse, setting the ids in execution order
    # later in execution order gives a clearer picture to users
    for idx, stage in enumerate(job.stages):
        stage.id = idx

    for task in job.stages[-1].tasks:
        task.args[1] = True

    return job


def _job_calling_info():
    name = None
    desc = None
    for file, lineno, func, text in reversed(traceback.extract_stack()):
        if 'bndl/' in file and func[0] != '_':
            name = func
        desc = file, lineno, func, text
        if not 'bndl/' in file:
            break
    return name, desc



def get_assignments(dset, workers):
    assignments = collections.defaultdict(lambda: [])
    assignment_count = lambda worker: (len(assignments[worker]), worker.name)

    for part in dset.parts():
        allowed_workers = list(part.allowed_workers(workers) or [])
        preferred_workers = list(part.preferred_workers(allowed_workers or workers) or [])

        # sort the allowed and preferred workers, least assigned first
        allowed_workers.sort(key=assignment_count)
        preferred_workers.sort(key=assignment_count)

        # select one worker to execute the task on
        selected_worker = (preferred_workers or allowed_workers or [None])[0]
        if not selected_worker:
            selected_worker = sorted(workers, key=assignment_count)[0]

        # use it as preferred worker if no provided by partition
        if not preferred_workers:
            preferred_workers = [selected_worker]

        # and update assignments
        assignments[selected_worker].append((part, (allowed_workers, preferred_workers)))

    return assignments


def schedule_stage(stage, workers, dset):
    '''
    Schedule a stage for a data set.
    
    It is assumed that all source data sets (and their parts) are materialized when
    this data set is materialized. (i.e. parts call materialize on their sources,
    if any).
    
    Also it is assumed that stages are scheduled backwards. Specifically if
    stage.is_last when this function is called it will remain that way ...
    
    :param stage: Stage
        stage to add tasks to
    :param workers: iterable
        Workers to schedule the data set on
    :param dset:
        The data set to schedule
    '''
    stage.name = dset.__class__.__name__

    assignments = get_assignments(dset, workers)

    # loop over assignments per worker
    for parts in assignments.values():
        # and task per assigned part
        for part, (allowed_workers, preferred_workers) in parts:
            stage.tasks.append(Task(
                (part.dset.id, part.idx),
                stage,
                materialize_partition, [part, False], None,
                preferred_workers, allowed_workers
            ))

    # sort the tasks by their id
    stage.tasks.sort(key=lambda t: t.id)



def materialize_partition(worker, part, return_data):
    try:
        ctx = part.dset.ctx
        _set_worker(part.dset, worker)

        data = part.materialize(ctx)
        if return_data and data is not None:
            # TODO This rule is supposed to catch generators, islices, map
            # objects and the lot. They aren't serializable unless materialized<commit_msg>don't overcomplicate scheduling (handle this at the execute layer)<commit_after>import collections.abc as collections_abc
import collections
import logging

from bndl.execute.job import Job, Stage, Task
from functools import partial
import traceback


logger = logging.getLogger(__name__)


def schedule_job(dset, workers=None):
    '''
    Schedule a job for a data set
    :param dset:
        The data set to schedule
    '''

    ctx = dset.ctx
    assert ctx.running, 'context of dataset is not running'

    if not workers:
        ctx._await_workers()
        workers = ctx.workers[:]

    job = Job(ctx, *_job_calling_info())

    stage = Stage(None, job)
    schedule_stage(stage, workers, dset)
    job.stages.insert(0, stage)

    def _cleaner(dset, job):
        if job.stopped:
            dset.cleanup(job)

    while dset:
        if dset.cleanup:
            job.add_listener(partial(_cleaner, dset))
        if isinstance(dset.src, collections_abc.Iterable):
            for src in dset.src:
                branch = schedule_job(src)

                for task in branch.stages[-1].tasks:
                    task.args[1] = False

                for l in branch.listeners:
                    job.add_listener(l)

                if dset.sync_required:
                    branch_stages = branch.stages
                elif len(branch.stages) > 1:
                    branch_stages = branch.stages[:-1]
                else:
                    continue

                # branch_stages[-1].next_stage = job.stages[-1]
                for stage in reversed(branch_stages):
                    stage.job = job
                    stage.is_last = False
                    job.stages.insert(0, stage)

            break

        elif dset.sync_required:
            stage = stage.prev_stage = Stage(None, job)  # , next_stage=stage)
            schedule_stage(stage, workers, dset)
            job.stages.insert(0, stage)

        dset = dset.src

    # Since stages are added in reverse, setting the ids in execution order
    # later in execution order gives a clearer picture to users
    for idx, stage in enumerate(job.stages):
        stage.id = idx

    for task in job.stages[-1].tasks:
        task.args[1] = True

    return job


def _job_calling_info():
    name = None
    desc = None
    for file, lineno, func, text in reversed(traceback.extract_stack()):
        if 'bndl/' in file and func[0] != '_':
            name = func
        desc = file, lineno, func, text
        if not 'bndl/' in file:
            break
    return name, desc



def schedule_stage(stage, workers, dset):
    '''
    Schedule a stage for a data set.
    
    It is assumed that all source data sets (and their parts) are materialized when
    this data set is materialized. (i.e. parts call materialize on their sources,
    if any).
    
    Also it is assumed that stages are scheduled backwards. Specifically if
    stage.is_last when this function is called it will remain that way ...
    
    :param stage: Stage
        stage to add tasks to
    :param workers: list or set
        Workers to schedule the data set on
    :param dset:
        The data set to schedule
    '''
    stage.name = dset.__class__.__name__

    for part in dset.parts():
        allowed_workers = list(part.allowed_workers(workers) or [])
        preferred_workers = list(part.preferred_workers(allowed_workers or workers) or [])

        stage.tasks.append(Task(
            (part.dset.id, part.idx),
            stage,
            materialize_partition, [part, False], None,
            preferred_workers, allowed_workers
        ))

    # sort the tasks by their id
    stage.tasks.sort(key=lambda t: t.id)



def materialize_partition(worker, part, return_data):
    try:
        ctx = part.dset.ctx
        _set_worker(part.dset, worker)

        data = part.materialize(ctx)
        if return_data and data is not None:
            # TODO This rule is supposed to catch generators, islices, map
            # objects and the lot. They aren't serializable unless materialized
{'fraction_non_alphanumeric': 0.08001734793451155, 'fraction_numerical': 0.0028190393581264232, 'mean_word_length': 2.0859819337571093, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '6578211', 'n_tokens_mistral': 2637, 'n_tokens_neox': 2526, 'n_words': 824}