climate_ref.solver #

def extract_covered_datasets(
    data_catalog: pd.DataFrame, requirement: DataRequirement
) -> dict[Selector, pd.DataFrame]:
    """
    Determine the different diagnostic executions that should be performed with the current data catalog
    """
    if len(data_catalog) == 0:
        logger.error(f"No datasets found in the data catalog: {requirement.source_type.value}")
        return {}

    subset = requirement.apply_filters(data_catalog)

    if len(subset) == 0:
        logger.debug(f"No datasets found for requirement {requirement}")
        return {}

    if requirement.group_by is None:
        # Use a single group
        groups = [((), subset)]
    else:
        groups = list(subset.groupby(list(requirement.group_by)))

    results = {}

    for name, group in groups:
        if requirement.group_by is None:
            assert len(groups) == 1
            group_keys: Selector = ()
        else:
            group_keys = tuple(zip(requirement.group_by, name))
        constrained_group = _process_group_constraints(data_catalog, group, requirement)

        if constrained_group is not None:
            results[group_keys] = constrained_group

    return results

def matches_filter(diagnostic: Diagnostic, filters: SolveFilterOptions | None) -> bool:
    """
    Check if a diagnostic matches the provided filters

    Each filter is optional and a diagnostic will match if it satisfies all the provided filters.
    i.e. the filters are ANDed together.

    Parameters
    ----------
    diagnostic
        Diagnostic to check against the filters
    filters
        Collection of filters to apply to the diagnostic

        If no filters are provided, the diagnostic is considered to match

    Returns
    -------
        True if the diagnostic matches the filters, False otherwise
    """
    if filters is None:
        return True

    diagnostic_slug = diagnostic.slug
    provider_slug = diagnostic.provider.slug

    if filters.provider and not any([f.lower() in provider_slug for f in filters.provider]):
        return False

    if filters.diagnostic and not any([f.lower() in diagnostic_slug for f in filters.diagnostic]):
        return False

    return True

def solve_executions(
    data_catalog: dict[SourceDatasetType, pd.DataFrame], diagnostic: Diagnostic, provider: DiagnosticProvider
) -> typing.Generator["DiagnosticExecution", None, None]:
    """
    Calculate the diagnostic executions that need to be performed for a given diagnostic

    Parameters
    ----------
    data_catalog
        Data catalogs for each source dataset type
    diagnostic
        Diagnostic of interest
    provider
        Provider of the diagnostic

    Returns
    -------
    :
        A generator that yields the diagnostic executions that need to be performed

    """
    if not diagnostic.data_requirements:
        raise ValueError(f"Diagnostic {diagnostic.slug!r} has no data requirements")

    first_item = next(iter(diagnostic.data_requirements))

    if isinstance(first_item, DataRequirement):
        # We have a single collection of data requirements
        yield from _solve_from_data_requirements(
            data_catalog,
            diagnostic,
            typing.cast(Sequence[DataRequirement], diagnostic.data_requirements),
            provider,
        )
    elif isinstance(first_item, Sequence):
        # We have a sequence of collections of data requirements
        for requirement_collection in diagnostic.data_requirements:
            if not isinstance(requirement_collection, Sequence):
                raise TypeError(f"Expected a sequence of DataRequirement, got {type(requirement_collection)}")
            yield from _solve_from_data_requirements(
                data_catalog, diagnostic, requirement_collection, provider
            )
    else:
        raise TypeError(f"Expected a DataRequirement, got {type(first_item)}")

def solve_required_executions(  # noqa: PLR0912, PLR0913
    db: Database,
    dry_run: bool = False,
    execute: bool = True,
    solver: ExecutionSolver | None = None,
    config: Config | None = None,
    timeout: int = 60,
    one_per_provider: bool = False,
    one_per_diagnostic: bool = False,
    filters: SolveFilterOptions | None = None,
) -> None:
    """
    Solve for executions that require recalculation

    This may trigger a number of additional calculations depending on what data has been ingested
    since the last solve.

    Raises
    ------
    TimeoutError
        If the execution isn't completed within the specified timeout
    """
    if config is None:
        config = Config.default()
    if solver is None:
        solver = ExecutionSolver.build_from_db(config, db)

    logger.info("Solving for diagnostics that require recalculation...")

    executor = config.executor.build(config, db)

    diagnostic_count = {}
    provider_count = {}

    for potential_execution in solver.solve(filters):
        # The diagnostic output is first written to the scratch directory
        definition = potential_execution.build_execution_definition(output_root=config.paths.scratch)

        logger.debug(
            f"Identified candidate execution {definition.key} "
            f"for {potential_execution.diagnostic.full_slug()}"
        )

        if potential_execution.provider.slug not in provider_count:
            provider_count[potential_execution.provider.slug] = 0
        if potential_execution.diagnostic.full_slug() not in diagnostic_count:
            diagnostic_count[potential_execution.diagnostic.full_slug()] = 0

        if dry_run:
            provider_count[potential_execution.provider.slug] += 1
            diagnostic_count[potential_execution.diagnostic.full_slug()] += 1
            continue

        # Use a transaction to make sure that the models
        # are created correctly before potentially executing out of process
        with db.session.begin():
            diagnostic = (
                db.session.query(DiagnosticModel)
                .join(DiagnosticModel.provider)
                .filter(
                    ProviderModel.slug == potential_execution.provider.slug,
                    DiagnosticModel.slug == potential_execution.diagnostic.slug,
                )
                .one()
            )
            execution_group, created = db.get_or_create(
                ExecutionGroup,
                key=definition.key,
                diagnostic_id=diagnostic.id,
                defaults={
                    "selectors": potential_execution.selectors,
                    "dirty": True,
                },
            )

            if created:
                logger.info(f"Created new execution group: {potential_execution.execution_slug()!r}")
                db.session.flush()

            # TODO: Move this logic to the solver
            # Check if we should run given the one_per_provider or one_per_diagnostic flags
            one_of_check_failed = (
                one_per_provider and provider_count.get(diagnostic.provider.slug, 0) > 0
            ) or (one_per_diagnostic and diagnostic_count.get(diagnostic.full_slug(), 0) > 0)

            logger.debug(
                f"one_per_provider={one_per_provider}, one_per_diagnostic={one_per_diagnostic}, "
                f"one_of_check_failed={one_of_check_failed}, diagnostic_count={diagnostic_count}, "
                f"provider_count={provider_count}"
            )

            if execution_group.should_run(definition.datasets.hash):
                if (one_per_provider or one_per_diagnostic) and one_of_check_failed:
                    logger.info(
                        f"Skipping execution due to one-of check: {potential_execution.execution_slug()!r}"
                    )
                    continue

                logger.info(
                    f"Running new execution for execution group: {potential_execution.execution_slug()!r}"
                )
                execution = Execution(
                    execution_group=execution_group,
                    dataset_hash=definition.datasets.hash,
                    output_fragment=str(definition.output_fragment()),
                )
                db.session.add(execution)
                db.session.flush()

                # Add links to the datasets used in the execution
                execution.register_datasets(db, definition.datasets)

                if execute:
                    executor.run(
                        definition=definition,
                        execution=execution,
                    )

                provider_count[diagnostic.provider.slug] += 1
                diagnostic_count[diagnostic.full_slug()] += 1

    logger.info("Solve complete")
    logger.info(f"Found {sum(diagnostic_count.values())} new executions")
    for diag, count in diagnostic_count.items():
        logger.info(f"  {diag}: {count} new executions")
    for prov, count in provider_count.items():
        logger.info(f"  {prov}: {count} new executions")

    if timeout > 0:
        executor.join(timeout=timeout)
        logger.info("All executions complete")