flyte

def GPU(
    device: typing.Literal['T4', 'A100', 'A100 80G', 'H100', 'L4', 'L40s'],
    quantity: typing.Literal[1, 2, 3, 4, 5, 6, 7, 8],
    partition: typing.Union[typing.Literal['1g.5gb', '2g.10gb', '3g.20gb', '4g.20gb', '7g.40gb'], typing.Literal['1g.10gb', '2g.20gb', '3g.40gb', '4g.40gb', '7g.80gb'], NoneType],
) -> flyte._resources.Device

def TPU(
    device: typing.Literal['V5P', 'V6E'],
    partition: typing.Union[typing.Literal['2x2x1', '2x2x2', '2x4x4', '4x4x4', '4x4x8', '4x8x8', '8x8x8', '8x8x16', '8x16x16', '16x16x16', '16x16x24'], typing.Literal['1x1', '2x2', '2x4', '4x4', '4x8', '8x8', '8x16', '16x16'], NoneType],
)

def build(
    image: Image,
) -> str

def build_images(
    envs: Environment,
) -> ImageCache

def ctx()

def deploy(
    envs: Environment,
    dryrun: bool,
    version: str | None,
    interactive_mode: bool | None,
    copy_style: CopyFiles,
) -> Deployment

def group(
    name: str,
)

@task
async def my_task():
    ...
    with group("my_group"):
        t1(x,y)  # tasks in this block will be grouped under "my_group"
    ...

def init(
    org: str | None,
    project: str | None,
    domain: str | None,
    root_dir: Path | None,
    log_level: int | None,
    endpoint: str | None,
    headless: bool,
    insecure: bool,
    insecure_skip_verify: bool,
    ca_cert_file_path: str | None,
    auth_type: AuthType,
    command: List[str] | None,
    proxy_command: List[str] | None,
    api_key: str | None,
    client_id: str | None,
    client_credentials_secret: str | None,
    auth_client_config: ClientConfig | None,
    rpc_retries: int,
    http_proxy_url: str | None,
    storage: Storage | None,
    batch_size: int,
    image_builder: ImageBuildEngine.ImageBuilderType,
)

def init_from_config(
    path_or_config: str | Config | None,
    root_dir: Path | None,
    log_level: int | None,
)

def run(
    task: TaskTemplate[P, R],
    args: *args,
    kwargs: **kwargs,
) -> Union[R, Run]

def trace(
    func: typing.Callable[..., ~T],
) -> typing.Callable[..., ~T]

def with_runcontext(
    mode: Mode | None,
    name: Optional[str],
    service_account: Optional[str],
    version: Optional[str],
    copy_style: CopyFiles,
    dry_run: bool,
    copy_bundle_to: pathlib.Path | None,
    interactive_mode: bool | None,
    raw_data_path: str | None,
    run_base_dir: str | None,
    overwrite_cache: bool,
    project: str | None,
    domain: str | None,
    env: Dict[str, str] | None,
    labels: Dict[str, str] | None,
    annotations: Dict[str, str] | None,
    interruptible: bool,
    log_level: int | None,
) -> _Runner

import flyte
env = flyte.TaskEnvironment("example")

@env.task
async def example_task(x: int, y: str) -> str:
    return f"{x} {y}"

if __name__ == "__main__":
    flyte.with_runcontext(name="example_run_id").run(example_task, 1, y="hello")

class Cache(
    behavior: typing.Literal['auto', 'override', 'disable', 'enabled'],
    version_override: typing.Optional[str],
    serialize: bool,
    ignored_inputs: typing.Union[typing.Tuple[str, ...], str],
    salt: str,
    policies: typing.Union[typing.List[flyte._cache.cache.CachePolicy], flyte._cache.cache.CachePolicy, NoneType],
)

def get_ignored_inputs()

def get_version(
    params: typing.Optional[flyte._cache.cache.VersionParameters],
) -> str

def is_enabled()

protocol CachePolicy()

def get_version(
    salt: str,
    params: flyte._cache.cache.VersionParameters,
) -> str

class Device(
    quantity: int,
    device: str | None,
    partition: str | None,
)

class Environment(
    name: str,
    depends_on: List[Environment],
    pod_template: Optional[Union[str, 'V1PodTemplate']],
    description: Optional[str],
    secrets: Optional[SecretRequest],
    env: Optional[Dict[str, str]],
    resources: Optional[Resources],
    image: Union[str, Image, Literal['auto']],
)

def add_dependency(
    env: Environment,
)

def clone_with(
    name: str,
    image: Optional[Union[str, Image, Literal['auto']]],
    resources: Optional[Resources],
    env: Optional[Dict[str, str]],
    secrets: Optional[SecretRequest],
    depends_on: Optional[List[Environment]],
    kwargs: **kwargs,
) -> Environment

class Image(
    base_image: Optional[str],
    dockerfile: Optional[Path],
    registry: Optional[str],
    name: Optional[str],
    platform: Tuple[Architecture, ...],
    python_version: Tuple[int, int],
    _layers: Tuple[Layer, ...],
)

def clone(
    registry: Optional[str],
    name: Optional[str],
    python_version: Optional[Tuple[int, int]],
    addl_layer: Optional[Layer],
) -> Image

def from_base(
    image_uri: str,
) -> Image

def from_debian_base(
    python_version: Optional[Tuple[int, int]],
    flyte_version: Optional[str],
    install_flyte: bool,
    registry: Optional[str],
    name: Optional[str],
    platform: Optional[Tuple[Architecture, ...]],
) -> Image

def from_dockerfile(
    file: Path,
    registry: str,
    name: str,
    platform: Union[Architecture, Tuple[Architecture, ...], None],
) -> Image

def from_uv_script(
    script: Path | str,
    name: str,
    registry: str | None,
    python_version: Optional[Tuple[int, int]],
    index_url: Optional[str],
    extra_index_urls: Union[str, List[str], Tuple[str, ...], None],
    pre: bool,
    extra_args: Optional[str],
    platform: Optional[Tuple[Architecture, ...]],
) -> Image

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = ["httpx"]
# ///

def validate()

def with_apt_packages(
    packages: str,
) -> Image

def with_commands(
    commands: List[str],
) -> Image

def with_dockerignore(
    path: Path,
) -> Image

def with_env_vars(
    env_vars: Dict[str, str],
) -> Image

def with_local_v2()

def with_pip_packages(
    packages: str,
    index_url: Optional[str],
    extra_index_urls: Union[str, List[str], Tuple[str, ...], None],
    pre: bool,
    extra_args: Optional[str],
) -> Image

@flyte.task(image=(flyte.Image
                .ubuntu_python()
                .with_pip_packages("requests", "numpy")))
def my_task(x: int) -> int:
    import numpy as np
    return np.sum([x, 1])

def with_requirements(
    file: str | Path,
) -> Image

def with_source_file(
    src: Path,
    dst: str,
) -> Image

def with_source_folder(
    src: Path,
    dst: str,
) -> Image

def with_uv_project(
    pyproject_file: Path,
    index_url: Optional[str],
    extra_index_urls: Union[str, List[str], Tuple[str, ...], None],
    pre: bool,
    extra_args: Optional[str],
) -> Image

def with_workdir(
    workdir: str,
) -> Image

class PodTemplate(
    pod_spec: typing.Optional[ForwardRef('V1PodSpec')],
    primary_container_name: str,
    labels: typing.Optional[typing.Dict[str, str]],
    annotations: typing.Optional[typing.Dict[str, str]],
)

def to_k8s_pod()

@task(resources=Resources(cpu=1, memory="1GiB", gpu="T4:1"))
def my_task() -> int:
    return 42

@task(resources=Resources(cpu=(1, 2), memory="2GiB", gpu="A100:8", disk="10GiB"))
def my_task() -> int:
    return 42

class Resources(
    cpu: typing.Union[int, float, str, typing.Tuple[int | float | str, int | float | str], NoneType],
    memory: typing.Union[str, typing.Tuple[str, str], NoneType],
    gpu: typing.Union[typing.Literal['T4:1', 'T4:2', 'T4:3', 'T4:4', 'T4:5', 'T4:6', 'T4:7', 'T4:8', 'L4:1', 'L4:2', 'L4:3', 'L4:4', 'L4:5', 'L4:6', 'L4:7', 'L4:8', 'L40s:1', 'L40s:2', 'L40s:3', 'L40s:4', 'L40s:5', 'L40s:6', 'L40s:7', 'L40s:8', 'A100:1', 'A100:2', 'A100:3', 'A100:4', 'A100:5', 'A100:6', 'A100:7', 'A100:8', 'A100 80G:1', 'A100 80G:2', 'A100 80G:3', 'A100 80G:4', 'A100 80G:5', 'A100 80G:6', 'A100 80G:7', 'A100 80G:8', 'H100:1', 'H100:2', 'H100:3', 'H100:4', 'H100:5', 'H100:6', 'H100:7', 'H100:8'], int, flyte._resources.Device, NoneType],
    disk: typing.Optional[str],
    shm: typing.Union[str, typing.Literal['auto'], NoneType],
)

def get_device()

def get_shared_memory()

class RetryStrategy(
    count: int,
    backoff: typing.Union[float, datetime.timedelta, NoneType],
    backoff_factor: typing.Union[int, float, NoneType],
)

class ReusePolicy(
    replicas: typing.Union[int, typing.Tuple[int, int]],
    idle_ttl: typing.Union[int, datetime.timedelta, NoneType],
    reuse_salt: str | None,
    concurrency: int,
)

@task(secrets="MY_SECRET")
async def my_task():
    os.environ["MY_SECRET"]  # This will be set to the value of the secret

@task(secrets=Secret("MY_SECRET", mount="/path/to/secret"))
async def my_task2():
    async with open("/path/to/secret") as f:
        secret_value = f.read()

class Secret(
    key: str,
    group: typing.Optional[str],
    mount: pathlib._local.Path | None,
    as_env_var: typing.Optional[str],
)

def stable_hash()

env = flyte.TaskEnvironment(name="my_env", image="my_image", resources=Resources(cpu="1", memory="1Gi"))

@env.task
async def my_task():
    pass

class TaskEnvironment(
    name: str,
    depends_on: List[Environment],
    pod_template: Optional[Union[str, 'V1PodTemplate']],
    description: Optional[str],
    secrets: Optional[SecretRequest],
    env: Optional[Dict[str, str]],
    resources: Optional[Resources],
    image: Union[str, Image, Literal['auto']],
    cache: Union[CacheRequest],
    reusable: ReusePolicy | None,
    plugin_config: Optional[Any],
)

def add_dependency(
    env: Environment,
)

def add_task(
    task: TaskTemplate,
) -> TaskTemplate

def clone_with(
    name: str,
    image: Optional[Union[str, Image, Literal['auto']]],
    resources: Optional[Resources],
    env: Optional[Dict[str, str]],
    secrets: Optional[SecretRequest],
    depends_on: Optional[List[Environment]],
    kwargs: **kwargs,
) -> TaskEnvironment

def task(
    _func,
    name: Optional[str],
    cache: Union[CacheRequest] | None,
    retries: Union[int, RetryStrategy],
    timeout: Union[timedelta, int],
    docs: Optional[Documentation],
    secrets: Optional[SecretRequest],
    pod_template: Optional[Union[str, 'V1PodTemplate']],
    report: bool,
) -> Union[AsyncFunctionTaskTemplate, Callable[P, R]]

timeout = Timeout(max_runtime=timedelta(minutes=5), max_queued_time=timedelta(minutes=10))
@env.task(timeout=timeout)
async def my_task():
    pass

class Timeout(
    max_runtime: datetime.timedelta | int,
    max_queued_time: datetime.timedelta | int | None,
)