# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
"""Auto-tuning Task Scheduler"""
from typing import Callable, List, Optional, Union

# isort: off
from typing_extensions import Literal

# isort: on

from tvm._ffi import register_object
from tvm.runtime import Object

from .. import _ffi_api
from ..builder import Builder, BuilderResult
from ..cost_model import CostModel
from ..database import Database
from ..logging import get_logger, get_logging_func
from ..measure_callback import MeasureCallback
from ..runner import Runner, RunnerResult
from ..search_strategy import MeasureCandidate
from ..tune_context import TuneContext

logger = get_logger(__name__)  # pylint: disable=invalid-name


@register_object("meta_schedule.TaskRecord")
class TaskRecord(Object):
    """The running record of a task."""

    ctx: TuneContext
    task_weight: float
    flop: float
    is_terminated: bool
    build_error_count: int
    run_error_count: int
    measure_candidates: List[MeasureCandidate]
    builder_results: List[BuilderResult]
    runner_results: List[RunnerResult]



[文档]
@register_object("meta_schedule.TaskScheduler")
class TaskScheduler(Object):
    """The abstract task scheduler interface."""

    tasks_: List[TaskRecord]
    measure_callbacks_: List[MeasureCallback]
    database_: Optional[Database]
    cost_model_: Optional[CostModel]
    remaining_tasks_: int

    TaskSchedulerType = Union["TaskScheduler", Literal["gradient", "round-robin"]]


[文档]
    def next_task_id(self) -> int:
        """Fetch the next task id.

        Returns
        -------
        next_task_id : int
            The next task id.
        """
        return _ffi_api.TaskSchedulerNextTaskId(self)  # type: ignore # pylint: disable=no-member



[文档]
    def join_running_task(self, task_id: int) -> List[RunnerResult]:
        """Wait until the task is finished.

        Parameters
        ----------
        task_id : int
            The task id to be joined.

        Returns
        -------
        results : List[RunnerResult]
            The list of results.
        """
        return _ffi_api.TaskSchedulerJoinRunningTask(self, task_id)  # type: ignore # pylint: disable=no-member



[文档]
    def tune(
        self,
        tasks: List[TuneContext],
        task_weights: List[float],
        max_trials_global: int,
        max_trials_per_task: int,
        num_trials_per_iter: int,
        builder: Builder,
        runner: Runner,
        measure_callbacks: List[MeasureCallback],
        database: Optional[Database],
        cost_model: Optional[CostModel],
    ) -> None:
        """Auto-tuning.

        Parameters
        ----------
        tasks : List[TuneContext]
            The list of tuning contexts as tasks.
        task_weights : List[float]
            The list of task weights.
        max_trials_global : int
            The maximum number of trials globally.
        max_trials_per_task : int
            The maximum number of trials per task.
        num_trials_per_iter : int
            The number of trials per iteration.
        builder : Builder
            The builder.
        runner : Runner
            The runner.
        measure_callbacks : List[MeasureCallback]
            The list of measure callbacks.
        database : Optional[Database]
            The database.
        cost_model : Optional[CostModel]
            The cost model.
        """
        task_weights = [float(w) for w in task_weights]
        _ffi_api.TaskSchedulerTune(  # type: ignore # pylint: disable=no-member
            self,
            tasks,
            task_weights,
            max_trials_global,
            max_trials_per_task,
            num_trials_per_iter,
            builder,
            runner,
            measure_callbacks,
            database,
            cost_model,
        )



[文档]
    def terminate_task(self, task_id: int) -> None:
        """Terminate the task

        Parameters
        ----------
        task_id : int
            The task id to be terminated.
        """
        _ffi_api.TaskSchedulerTerminateTask(self, task_id)  # type: ignore # pylint: disable=no-member



[文档]
    def touch_task(self, task_id: int) -> None:
        """Touch the task and update its status

        Parameters
        ----------
        task_id : int
            The task id to be checked.
        """
        _ffi_api.TaskSchedulerTouchTask(self, task_id)  # type: ignore # pylint: disable=no-member



[文档]
    def print_tuning_statistics(self) -> None:
        """Print out a human-readable format of the tuning statistics."""
        return _ffi_api.TaskSchedulerPrintTuningStatistics(self)  # type: ignore # pylint: disable=no-member



[文档]
    @staticmethod
    def create(  # pylint: disable=keyword-arg-before-vararg
        kind: Literal["round-robin", "gradient"] = "gradient",
        *args,
        **kwargs,
    ) -> "TaskScheduler":
        """Create a task scheduler."""
        from . import (  # pylint: disable=import-outside-toplevel
            GradientBased,
            RoundRobin,
        )

        if kind == "round-robin":
            return RoundRobin(*args, **kwargs)  # type: ignore
        if kind == "gradient":
            return GradientBased(*args, **kwargs)
        raise ValueError(f"Unknown TaskScheduler name: {kind}")




create = TaskScheduler.create  # pylint: disable=invalid-name


@register_object("meta_schedule.PyTaskScheduler")
class _PyTaskScheduler(TaskScheduler):
    """
    A TVM object task scheduler to support customization on the python side.
    This is NOT the user facing class for function overloading inheritance.

    See also: PyTaskScheduler
    """

    def __init__(
        self,
        f_next_task_id: Callable,
        f_join_running_task: Callable,
        f_tune: Callable,
    ):
        """Constructor."""

        self.__init_handle_by_constructor__(
            _ffi_api.TaskSchedulerPyTaskScheduler,  # type: ignore # pylint: disable=no-member
            get_logging_func(logger),
            f_next_task_id,
            f_join_running_task,
            f_tune,
        )


class PyTaskScheduler:
    """
    An abstract task scheduler with customized methods on the python-side.
    This is the user facing class for function overloading inheritance.

    Note: @derived_object is required for proper usage of any inherited class.
    """

    _tvm_metadata = {
        "cls": _PyTaskScheduler,
        "fields": [],
        "methods": ["next_task_id", "join_running_task", "tune"],
    }

    def __init__(self):
        ...

    def tune(
        self,
        tasks: List[TuneContext],
        task_weights: List[float],
        max_trials_global: int,
        max_trials_per_task: int,
        builder: Builder,
        runner: Runner,
        measure_callbacks: List[MeasureCallback],
        database: Optional[Database],
        cost_model: Optional[CostModel],
    ) -> None:
        """Auto-tuning."""
        # Using self._outer to replace the self pointer
        _ffi_api.TaskSchedulerTune(  # type: ignore # pylint: disable=no-member
            self._outer(),  # type: ignore # pylint: disable=no-member
            tasks,
            task_weights,
            max_trials_global,
            max_trials_per_task,
            builder,
            runner,
            measure_callbacks,
            database,
            cost_model,
        )

    def next_task_id(self) -> int:
        """Fetch the next task id.

        Returns
        -------
        next_task_id : int
            The next task id.
        """
        raise NotImplementedError

    def join_running_task(self, task_id: int) -> List[RunnerResult]:
        """Wait until the task is finished.

        Parameters
        ----------
        task_id : int
            The task id to be joined.
        """
        # Using self._outer to replace the self pointer
        return _ffi_api.TaskSchedulerJoinRunningTask(self._outer(), task_id)  # type: ignore # pylint: disable=no-member