pyflyte register → FlyteAdmin (stores Protobuf specs)
                         ↓
                   FlytePropeller (K8s operator, watches FlyteWorkflow CRDs)
                         ↓
                   Creates K8s Pod per task (or delegates to operator plugin)
                         ↓
                   Pod runs: pyflyte-execute --task-module my.module --task-name my_task
                         ↓
                   Task function executes, outputs written to blob store
                         ↓
                   FlytePropeller reads outputs, schedules downstream tasks

from flytekit import ImageSpec, task
from flytekit.image_spec import CopyFileDetection

# GPU training image
training_image = ImageSpec(
    name="training",
    base_image="nvcr.io/nvidia/pytorch:24.01-py3",
    packages=[
        "transformers==4.46.0",
        "datasets==3.0.0",
        "accelerate==1.0.0",
        "wandb",
    ],
    pip_extra_index_url=["https://download.pytorch.org/whl/cu121"],
    apt_packages=["git"],
    env={"WANDB_PROJECT": "my-project"},
    registry="ghcr.io/my-org",
    source_copy_mode=CopyFileDetection.ALL,  # copy local source into image
)

# Lightweight CPU image for data preprocessing
preprocess_image = ImageSpec(
    name="preprocess",
    packages=["pandas", "pyarrow", "datasets"],
    registry="ghcr.io/my-org",
)

@task(container_image=training_image)
def train_model(config: dict) -> str:
    ...

@task(container_image=preprocess_image)
def preprocess_data(path: str) -> str:
    ...

training_image = ImageSpec(packages=["torch", "transformers"], ...)

if training_image.is_container():
    import torch
    from transformers import AutoModelForCausalLM

@task(container_image=training_image)
def train(model_name: str) -> float:
    model = AutoModelForCausalLM.from_pretrained(model_name)
    ...

from flytekit import task, Resources

@task(
    requests=Resources(cpu="4", mem="16Gi", gpu="1", ephemeral_storage="50Gi"),
    limits=Resources(cpu="8", mem="32Gi", gpu="1"),
    accelerator=GPUAccelerator("nvidia-tesla-a100"),
    timeout=timedelta(hours=12),
    retries=2,
    interruptible=True,  # allow scheduling on spot/preemptible nodes
)
def train_model(data_path: str) -> float:
    ...

@task(
    retries=3,                    # total attempts = retries + 1 on user errors
    interruptible=True,           # retries on preemption don't count against retries budget
)
def training_task(config: dict) -> float:
    ...

from flytekit import task, PodTemplate
from kubernetes.client import V1PodSpec, V1Container, V1VolumeMount, V1Volume, V1EmptyDirVolumeSource

custom_pod = PodTemplate(
    pod_spec=V1PodSpec(
        containers=[
            V1Container(
                name="primary",
                volume_mounts=[
                    V1VolumeMount(name="shm", mount_path="/dev/shm"),
                    V1VolumeMount(name="data", mount_path="/data"),
                ],
            ),
        ],
        volumes=[
            V1Volume(name="shm", empty_dir=V1EmptyDirVolumeSource(medium="Memory", size_limit="16Gi")),
            V1Volume(name="data", persistent_volume_claim={"claimName": "training-data"}),
        ],
        node_selector={"nvidia.com/gpu.product": "NVIDIA-A100-SXM4-80GB"},
    ),
)

@task(pod_template=custom_pod, requests=Resources(gpu="1"))
def gpu_training(config: dict) -> str:
    ...

from flytekit.types.structured import StructuredDataset
import pandas as pd

@task
def generate_data() -> StructuredDataset:
    df = pd.DataFrame({"col": [1, 2, 3]})
    return StructuredDataset(dataframe=df)

@task
def consume_data(ds: StructuredDataset) -> float:
    df = ds.open(pd.DataFrame).all()  # reads from blob store
    return df["col"].mean()

from flytekit.types.file import FlyteFile
from flytekit.types.directory import FlyteDirectory

@task
def train(data_dir: FlyteDirectory, config: FlyteFile) -> FlyteFile:
    # data_dir and config are automatically downloaded to local paths
    local_data = data_dir.download()
    model_path = "/tmp/model.pt"
    torch.save(model.state_dict(), model_path)
    return FlyteFile(model_path)  # automatically uploaded to blob store

from flytekit import workflow

@workflow
def training_pipeline(model_name: str, data_path: str) -> float:
    processed = preprocess_data(path=data_path)
    model = train_model(data_path=processed)
    metrics = evaluate_model(model=model)
    return metrics

from flytekit import conditional

@workflow
def pipeline(data_path: str, use_large_model: bool) -> float:
    return (
        conditional("model_size")
        .if_(use_large_model.is_true())
        .then(train_large(data_path=data_path))
        .else_()
        .then(train_small(data_path=data_path))
    )

from flytekit import dynamic

@dynamic
def hyperparameter_search(configs: list[dict]) -> list[float]:
    results = []
    for config in configs:  # Python loops work here
        result = train_model(config=config)
        results.append(result)
    return results

from flytekit import eager

@eager
async def adaptive_training(model_name: str) -> float:
    metrics = await initial_train(model_name=model_name)
    if metrics < 0.9:  # real Python if/else
        metrics = await extended_train(model_name=model_name)
    return metrics

from flytekit import map_task

@workflow
def parallel_eval(models: list[str]) -> list[float]:
    return map_task(evaluate_model)(model_path=models)

from flytekit import LaunchPlan, CronSchedule

# Scheduled execution
nightly_train = LaunchPlan.get_or_create(
    training_pipeline,
    name="nightly-training",
    schedule=CronSchedule(schedule="0 2 * * *"),  # 2 AM daily
    fixed_inputs={"model_name": "llama-8b"},
    default_inputs={"data_path": "/data/latest"},
    max_parallelism=10,
    labels={"team": "ml"},
)

# Fixed config for production
prod_launch = LaunchPlan.get_or_create(
    training_pipeline,
    name="prod-training",
    fixed_inputs={"model_name": "llama-70b", "data_path": "/data/prod"},
)

from flytekit import task, Secret

@task(
    secret_requests=[
        Secret(group="wandb", key="api_key", mount_requirement=Secret.MountType.ENV_VAR),
        Secret(group="hf", key="token"),
    ],
)
def train_with_secrets(config: dict) -> str:
    import os
    wandb_key = os.environ["WANDB_API_KEY"]  # ENV_VAR mount
    # or: flytekit.current_context().secrets.get("wandb", "api_key")
    ...

kubectl create secret generic wandb --from-literal=api_key=<key> -n <flyte-ns>

from flytekit import task, current_context

@task(retries=3)
def long_training(epochs: int) -> float:
    ctx = current_context()
    checkpoint = ctx.checkpoint

    # Try to restore
    prev = checkpoint.read()
    start_epoch = 0
    if prev:
        state = pickle.loads(prev)
        model.load_state_dict(state["model"])
        start_epoch = state["epoch"] + 1

    for epoch in range(start_epoch, epochs):
        train_one_epoch(model)
        checkpoint.write(pickle.dumps({"model": model.state_dict(), "epoch": epoch}))

    return evaluate(model)

# Tasks are regular Python functions locally
result = train_model(config={"lr": 1e-4})

# Workflows too
metrics = training_pipeline(model_name="test", data_path="/local/data")

# Register all workflows/tasks in current directory to Flyte cluster
# Run as a CI/CD step or Job — not interactively
# pyflyte register --project my-project --domain development .

# Package to a tarball (for offline registration)
# pyflyte package --image ghcr.io/my-org/training:latest -o workflows.tgz