entrypoints.py

import errno
import importlib
import os
import traceback
from datetime import datetime
from typing import List, Optional, Tuple, Type

import flyte.errors
from flyte._code_bundle import download_bundle
from flyte._context import contextual_run
from flyte._internal import Controller
from flyte._internal.imagebuild.image_builder import ImageCache
from flyte._logging import log, logger
from flyte._metrics import Stopwatch
from flyte._task import TaskTemplate
from flyte.models import ActionID, CheckpointPaths, CodeBundle, RawDataPath

from ..._utils import adjust_sys_path
from .convert import Error, Inputs, Outputs, convert_from_native_to_error
from .taskrunner import (
    convert_and_run,
    extract_download_run_upload,
)


async def direct_dispatch(
    task: TaskTemplate,
    *,
    action: ActionID,
    raw_data_path: RawDataPath,
    controller: Controller,
    version: str,
    output_path: str,
    run_base_dir: str,
    checkpoint_paths: CheckpointPaths | None = None,
    code_bundle: CodeBundle | None = None,
    inputs: Inputs | None = None,
    run_start_time: Optional[datetime] = None,
) -> Tuple[Optional[Outputs], Optional[Error]]:
    """
    This method is used today by the local_controller and is positioned to be used by a rust core in the future.
    The caller, loads the task and invokes this method. This method is used to convert the inputs to native types,
    The reason for this is that the rust entrypoint will not have access to the python context, and
    will not be able to run the tasks in the context tree.
    """
    return await contextual_run(
        convert_and_run,
        task=task,
        inputs=inputs or Inputs.empty(),
        action=action,
        raw_data_path=raw_data_path,
        checkpoint_paths=checkpoint_paths,
        code_bundle=code_bundle,
        controller=controller,
        version=version,
        output_path=output_path,
        run_base_dir=run_base_dir,
        run_start_time=run_start_time,
    )


def load_class(qualified_name) -> Type:
    """
    Load a class from a qualified name. The qualified name should be in the format 'module.ClassName'.

    Args:
        qualified_name: The qualified name of the class to load.

    Returns:
        The class object.
    """
    module_name, class_name = qualified_name.rsplit(".", 1)  # Split module and class
    module = importlib.import_module(module_name)  # Import the module
    return getattr(module, class_name)  # Retrieve the class


_SKIP_DIRS = frozenset(
    {
        ".git",
        ".hg",
        ".svn",
        ".venv",
        "venv",
        "env",
        ".local",
        ".cache",
        ".uv",
        "__pycache__",
        "node_modules",
        ".mypy_cache",
        ".pytest_cache",
        ".ruff_cache",
        ".idea",
        ".tox",
        ".nox",
        ".eggs",
        "site-packages",
        "dist-packages",
        "dist",
        "build",
    }
)


def _list_user_files(cwd: str) -> list[str]:
    """List user-relevant files under cwd, filtering out virtual envs, caches, and other non-user directories."""
    files: list[str] = []
    try:
        for root, dirs, filenames in os.walk(cwd):
            # Prune irrelevant directories in-place so os.walk won't descend into them
            dirs[:] = [d for d in dirs if d not in _SKIP_DIRS]
            for name in filenames:
                rel_path = os.path.relpath(os.path.join(root, name), cwd)
                files.append(rel_path)
    except Exception as list_err:
        files = [f"(Failed to list directory: {list_err})"]
    return files


def load_task(resolver: str, *resolver_args: str) -> TaskTemplate:
    """
    Load a task from a resolver. This is a placeholder function.

    Args:
        resolver: The resolver to use to load the task.
        resolver_args: Arguments to pass to the resolver.

    Returns:
        The loaded task.
    """
    resolver_class = load_class(resolver)
    resolver_instance = resolver_class()
    try:
        return resolver_instance.load_task(resolver_args)
    except ModuleNotFoundError as e:
        cwd = os.getcwd()
        files = _list_user_files(cwd)

        msg = (
            "\n\nFull traceback:\n" + "".join(traceback.format_exc()) + f"\n[ImportError Diagnostics]\n"
            f"Module '{e.name}' not found in either the Python virtual environment or the current working directory.\n"
            f"Current working directory: {cwd}\n"
            f"Files found under current directory:\n" + "\n".join(f"  - {f}" for f in files)
        )
        raise ModuleNotFoundError(msg) from e


def _classify_load_error(err: Exception) -> Exception:
    if isinstance(err, OSError) and err.errno == errno.EIO:
        return flyte.errors.RuntimeSystemError(
            "FilesystemIOError",
            f"Filesystem I/O error while loading task: {err}",
        )
    return err


def load_pkl_task(code_bundle: CodeBundle) -> TaskTemplate:
    """
    Loads a task from a pickled code bundle.

    Args:
        code_bundle: The code bundle to load the task from.

    Returns:
        The loaded task template.
    """
    logger.debug(f"Loading task from pkl: {code_bundle.downloaded_path}")
    try:
        import gzip

        import cloudpickle

        with gzip.open(str(code_bundle.downloaded_path), "rb") as f:
            return cloudpickle.load(f)
    except Exception as e:
        logger.exception(f"Failed to load pickled task from {code_bundle.downloaded_path}. Reason: {e!s}")
        raise


async def download_code_bundle(code_bundle: CodeBundle) -> CodeBundle:
    """
    Downloads the code bundle if it is not already downloaded.

    Args:
        code_bundle: The code bundle to download.

    Returns:
        The code bundle with the downloaded path.
    """
    adjust_sys_path([str(code_bundle.destination)])
    logger.debug(f"Downloading {code_bundle}")
    sw = Stopwatch("download_code_bundle")
    sw.start()
    downloaded_path = await download_bundle(code_bundle)
    sw.stop()
    return code_bundle.with_downloaded_path(downloaded_path)


async def _download_and_load_task(
    code_bundle: CodeBundle | None, resolver: str | None = None, resolver_args: List[str] | None = None
) -> TaskTemplate:
    if code_bundle and (code_bundle.tgz or code_bundle.pkl):
        logger.debug(f"Downloading {code_bundle}")
        code_bundle = await download_code_bundle(code_bundle)
        if code_bundle.pkl:
            sw = Stopwatch("load_pkl_task")
            sw.start()
            result = load_pkl_task(code_bundle)
            sw.stop()
            return result

        if not resolver or not resolver_args:
            raise flyte.errors.RuntimeSystemError(
                "MalformedCommand", "Resolver and resolver args are required. for task"
            )
        logger.debug(
            f"Loading task from tgz: {code_bundle.downloaded_path}, resolver: {resolver}, args: {resolver_args}"
        )
        sw = Stopwatch("load_task_from_tgz")
        sw.start()
        result = load_task(resolver, *resolver_args)
        sw.stop()
        return result
    if not resolver or not resolver_args:
        raise flyte.errors.RuntimeSystemError("MalformedCommand", "Resolver and resolver args are required. for task")
    logger.debug(f"No code bundle provided, loading task from resolver: {resolver}, args: {resolver_args}")
    sw = Stopwatch("load_task_from_resolver")
    sw.start()
    result = load_task(resolver, *resolver_args)
    sw.stop()
    return result


@log
async def load_and_run_task(
    action: ActionID,
    raw_data_path: RawDataPath,
    output_path: str,
    run_base_dir: str,
    version: str,
    controller: Optional[Controller],
    resolver: str,
    resolver_args: List[str],
    checkpoint_paths: CheckpointPaths | None = None,
    code_bundle: CodeBundle | None = None,
    input_path: str | None = None,
    image_cache: ImageCache | None = None,
    interactive_mode: bool = False,
    run_start_time: Optional[datetime] = None,
):
    """
    This method is invoked from the runtime/CLI and is used to run a task. This creates the context tree,
    for the tasks to run in. It also handles the loading of the task.

    Args:
        controller: Controller to use for the task.
        resolver: The resolver to use to load the task.
        resolver_args: The arguments to pass to the resolver.
        action: The ActionID to use for the task.
        raw_data_path: The raw data path to use for the task.
        output_path: The output path to use for the task.
        run_base_dir: Base output directory to pass down to child tasks.
        version: The version of the task to run.
        checkpoint_paths: The checkpoint paths to use for the task.
        code_bundle: The code bundle to use for the task.
        input_path: The input path to use for the task.
        image_cache: Mappings of Image identifiers to image URIs.
        interactive_mode: Whether to run the task in interactive mode.
    """
    sw = Stopwatch("load_and_run_task_total")
    sw.start()
    if output_path:
        # Clustered rank-0 only; a no-op for regular tasks. Must precede every upload_error site (the load
        # failure below, taskrunner.extract_download_run_upload, runtime._run_and_stop) and the task body.
        # rusty.run_task bypasses this function but never hosts clustered workers (clustered.py execs a0).
        from .io import clear_stale_clustered_error

        await clear_stale_clustered_error(output_path)
    try:
        task = await _download_and_load_task(code_bundle, resolver, resolver_args)
    except Exception as e:
        classified_error = _classify_load_error(e)
        # Import/load failures happen before the contextual_run wrapper below, so they must
        # also be uploaded to the error file -- otherwise the UI shows an empty message.
        logger.exception(f"Failed to load task before execution: {e}")
        if output_path:
            from .io import upload_error

            error = convert_from_native_to_error(classified_error)
            await upload_error(error.err, output_path, recoverable=error.recoverable)
        if classified_error is not e:
            raise classified_error from e
        raise

    await contextual_run(
        extract_download_run_upload,
        task,
        action=action,
        version=version,
        controller=controller,
        raw_data_path=raw_data_path,
        output_path=output_path,
        run_base_dir=run_base_dir,
        checkpoint_paths=checkpoint_paths,
        code_bundle=code_bundle,
        input_path=input_path,
        image_cache=image_cache,
        interactive_mode=interactive_mode,
        run_start_time=run_start_time,
    )
    sw.stop()