Mitigate the trade-off between scalability and ease of use

from typing import List

from flytekit import task, workflow


@task
def mean(x: List[float]) -> float:
    return sum(x) / len(x)


@task
def variance(x: List[int], mu: float) -> float:
    return sum([(val - mu) ** 2 for val in x]) / len(x)


@workflow
def variance_wf(x: List[float] = [1.3, 2.4, 3.9]) -> float:
    return variance(x=x, mu=mean(x=x))

@task
def mappable_task(string_item: str) -> dict[str, int]:
    return dict(Counter(string_item.split()))


@workflow
def wf(string_list: list[str] = ["Deer Bear", "Car Deer"]) -> dict[str, int]:
    mapped_out = map_task(mappable_task)(string_item=string_list).with_overrides(
        retries=1
    )

    return coalesce(string_list=mapped_out)

@dynamic
def evaluate_model(
    models: List[JoblibSerializedFile],
    X: List[List[float]],
    y: List[float],
) -> List[List[str]]:
    final_result = []
    for model_ser in models:
        model = joblib.load(model_ser)
        cv = RepeatedStratifiedKFold(n_splits=7, n_repeats=2, random_state=1)
        scores = cross_val_score(model, X, y, scoring="accuracy", cv=cv, n_jobs=-1)
        final_result.append(
            evaluate_flyte_task(name=str(model_ser), scores=scores.tolist())
        )
    return final_result

@workflow
def conditional_wf():
    conditional("stop_if_fails").if_(validated.is_false()).then(
        summary_drift_report(
            new_data=target_df,
            reference_data=get_reference_data(),
        )
    ).else_().then(make_predictions(input_data=target_df, output_path="./data"))

@task
def normalize_columns(
    csv_url: FlyteFile,
    column_names: List[str],
) -> FlyteFile:
    with open(csv_url, newline="\n") as input_file:
        ...
    return FlyteFile(path=out_path, remote_path=output_location)


@task
def download_files(csv_urls: List[str]) -> FlyteDirectory:
    ...
    return FlyteDirectory(path=local_dir)

@task
def analyze_df(
   df: vaex.dataframe.DataFrameLocal,
) -> Annotated[StructuredDataset, kwtypes(age=float)]:
    result_df = df.groupby(by="animal").agg({"age": "mean"})
    return StructuredDataset(result_df)


@task
def read_df(df: StructuredDataset) -> float:
    pandas_df = df.open(pd.DataFrame).all()
    return pandas_df["age"].min()


@workflow
def sd_wf(
    vaex_df: vaex.dataframe.DataFrameLocal = vaex.from_arrays(
        animal=["dog", "cat", "guinea pig"],
        age=[2, 5, 1],
    )
) -> float:
    return read_df(df=analyze_df(df=vaex_df))

@workflow
def gate_nodes_wf():
    # sleep 
    sleeping = sleep(timedelta(seconds=10))
    # wait for an external input
    title_input = wait_for_input("title", timeout=timedelta(hours=1), expected_type=str)
    # wait for explicit approval
    approve(final_report, "approve-final-report", timeout=timedelta(hours=2))

pandas_image_spec = ImageSpec(
    base_image="ghcr.io/flyteorg/flytekit:py3.8-1.6.2",
    packages=["pandas", "numpy"],
    python_version="3.8",
    apt_packages=["git"],
    env={"Debug": "True"},
    registry="ghcr.io/flyteorg",
)

@task(container_image=pandas_image_spec)
def get_pandas_dataframe() -> typing.Tuple[pd.DataFrame, pd.Series]:
    df = pd.read_csv("https://storage.googleapis.com/download.tensorflow.org/data/heart.csv")
    return df[["age", "thalach", "trestbps", "chol", "oldpeak"]], df.pop("target")

@task(cache=True, cache_version="1.0")
def train_mnist_pipeline(hp: Hyperparameters) -> train_outputs:
    kwargs = {"num_workers": 1, "pin_memory": True} if torch.cuda.is_available() else {}
    training_data_loader = mnist_dataloader(hp.batch_size, train=True, **kwargs)
    ...

@task
def train_model():
    try:
        checkpoint = flytekit.current_context().checkpoint
        prev_checkpoint = checkpoint.read()
    except (NotImplementedError, ValueError):
        checkpoint, prev_checkpoint = None, False

    if prev_checkpoint:
        start_epoch, model, opt = torch.load(BytesIO(prev_checkpoint))
    else:
        start_epoch = 0
        model = ...
        opt = ...

    for epoch in range(start_epoch, n_epochs):
        ...

        if checkpoint:
            model_io = BytesIO()
            torch.save((epoch, model, opt), model_io)
            model_io.seek(0)
            checkpoint.write(model_io.read())

@reference_task(
    project="flytesnacks",
    domain="development",
    name="core.flyte_basics.files.normalize_columns",
    version="2NRd4zAkKIf9hpqBAMLjzw==",
)
def normalize_columns(
    csv_url: FlyteFile,
    column_names: List[str],
    columns_to_normalize: List[str],

@task(timeout=datetime.timedelta(hours=1))
def train_mnist_pipeline(hp: Hyperparameters) -> train_outputs:
    kwargs = {"num_workers": 1, "pin_memory": True} if torch.cuda.is_available() else {}
    training_data_loader = mnist_dataloader(hp.batch_size, train=True, **kwargs)
    ...

@task(interruptible=True)
def data_processing(google_trend_csv: pyspark.sql.DataFrame):
    google_trend_all = google_trend_csv.withColumn(
        "Date", F.regexp_extract(google_trend_csv.week, "(.*?) -", 1)
    ).withColumn(
        "State", F.regexp_extract(google_trend_csv.file, "Rossmann_DE_(.*)", 1)
    )
    ...

cron_lp = LaunchPlan.get_or_create(
    name="fetch_data_lp",
    workflow=fetch_data_wf,
    schedule=CronSchedule(
        schedule="0 1 * * *",
        kickoff_time_input_arg="kickoff_time",
    ),
)

@task(
    secret_requests=[
        Secret(
            group=SECRET_GROUP,
            key=SECRET_NAME,
            mount_requirement=Secret.MountType.ENV_VAR,
        )
    ]
)
def secret_file_task() -> str:
    ...

@task(requests=Resources(gpu="2"), limits=Resources(gpu="4"))
def train_mnist_pipeline(hp: Hyperparameters) -> train_outputs:
    kwargs = {"num_workers": 1, "pin_memory": True} if torch.cuda.is_available() else {}
    training_data_loader = mnist_dataloader(hp.batch_size, train=True, **kwargs)
    ...

@task(container_image="ghcr.io/flyteorg/flytecookbook:core-latest")
def get_data() -> Tuple[np.ndarray, np.ndarray]:
    from sklearn import datasets

    iris = datasets.load_iris()
    X, y = (iris.data[:, :2], iris.target)
    return (X, y)


@task(container_image="ghcr.io/flyteorg/flytekit:py3.9-1.2.7")
def normalize(X: np.ndarray) -> np.ndarray:
    return (X - X.mean(axis=0)) / X.std(axis=0)


@workflow
def wf() -> np.ndarray:
    X, _ = get_data()
    X = normalize(X=X)
    return X

@task
def mean(x: List[float]) -> float:
    return sum(x) / len(x)


@dynamic
def dynamic_wf(x: List[float], mem: str) -> float:
    return mean(x=x).with_overrides(requests=Resources(mem=mem))


@workflow
def mean_wf(x: List[float] = [1.3, 2.4, 3.9], mem: str = "600Mi") -> float:
    return dynamic_wf(x=x, mem=mem)