diff --git a/.github/workflows/skill-evals.yml b/.github/workflows/skill-evals.yml index d40004c..f7966d9 100644 --- a/.github/workflows/skill-evals.yml +++ b/.github/workflows/skill-evals.yml @@ -1,8 +1,8 @@ name: skill-evals -# Runs the flyte-skills eval harness. On PRs it runs only the scenarios affected -# by the changed files (see evals/select.py); nightly it runs the full matrix -# including the `real` tier. +# Runs the flyte-agent-plugin skills eval harness. On PRs it runs only the +# scenarios affected by the changed files (see evals/select.py); nightly it runs +# the full matrix including the `real` tier. on: pull_request: paths: diff --git a/README.md b/README.md index 5d915b8..5eae412 100644 --- a/README.md +++ b/README.md @@ -134,6 +134,20 @@ pi install git:github.com/flyteorg/flyte-agent-plugins@ # pinned to | [`flyte-sdk-data`](plugins/flyte/skills/flyte-sdk-data) | Handles data engineering patterns: ETL pipelines, data processing, data quality checks, fanout/map tasks, conditions, dynamic workflows, and batch data transformations. For: ETL, Parquet, CSV, JsonlFile/Dir, schema validation. | | [`flyte-sdk-ml`](plugins/flyte/skills/flyte-sdk-ml) | Handles ML workload patterns: model training, hyperparameter optimization, experiment tracking, model evaluation and selection, batch inference, real-time serving, and model monitoring. For: PyTorch, scikit-learn, HuggingFace, GPU, drift detection. | +### Migration (Flyte 1 → 2) + +Convert existing Flyte 1 (`flytekit`) code to Flyte 2. Distilled from the official +[Flyte 1 → 2 migration guide](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/). + +| Skill | Description | +|-------|-------------| +| [`flyte-migrate`](plugins/flyte/skills/flyte-migrate) | Start-here migration orchestrator: the `flytekit`→`flyte` shift, the terminology/concept mapping, the two mechanical changes, an incremental migration strategy, hybrid v1/v2 pipelines during transition, and the gotchas — routes to the specific skills below. | +| [`flyte-migrate-tasks-workflows`](plugins/flyte/skills/flyte-migrate-tasks-workflows) | Migrate `@task`/`@workflow`/`@dynamic` into a single `@env.task` on a `TaskEnvironment`; sequential ordering without `>>`, nested "subworkflows" as tasks, and the parameter-mapping table. | +| [`flyte-migrate-config`](plugins/flyte/skills/flyte-migrate-config) | Migrate task configuration (images `ImageSpec`→`flyte.Image`, resources/GPUs, `cache_version`→`cache`, secrets, `LaunchPlan`/`CronSchedule`→`Trigger`/`Cron`) and the `pyflyte`→`flyte` CLI / config files. | +| [`flyte-migrate-control-flow`](plugins/flyte/skills/flyte-migrate-control-flow) | Replace `conditional()` with native `if`/`else`, `@dynamic` with plain Python loops, `on_failure` with `try`/`except`, and `map_task` with `flyte.map` / `asyncio.gather`. | +| [`flyte-migrate-data-io`](plugins/flyte/skills/flyte-migrate-data-io) | Migrate data types & I/O: `FlyteFile`/`FlyteDirectory`→`flyte.io.File`/`Dir`, `StructuredDataset`→`flyte.io.DataFrame`, dataclasses/Pydantic as task I/O. | +| [`flyte-migrate-ml`](plugins/flyte/skills/flyte-migrate-ml) | Migrate ML workloads (training, HPO, GPU/deep learning, batch inference, end-to-end pipelines) and the new-in-v2 patterns (real-time serving, apps, sandboxed execution) they unlock. | + Example: ``` @@ -233,8 +247,7 @@ scripts/smoke_test_mcp.py # end-to-end check of the local MCP Each harness consumes a different part of this. Claude Code and Codex read the plugin manifests, so the **plugin name** matters to them. Hermes, opencode, and pi install skills -by **directory path**, so `plugins/flyte/skills/…` is their interface — which is why the -rename from `flyte-skills` touched both. +by **directory path**, so `plugins/flyte/skills/…` is their interface. The `.mcp.json` server is Claude Code-specific; the skills themselves stay portable across harnesses. diff --git a/evals/README.md b/evals/README.md index 80e89da..197b123 100644 --- a/evals/README.md +++ b/evals/README.md @@ -1,6 +1,6 @@ -# flyte-skills eval harness +# flyte agent plugin eval harness -Automated testing & evaluation for the `flyte-skills` agent skills. It runs a real +Automated testing & evaluation for the `flyte` agent skills. It runs a real agent harness (**opencode / pi / hermes**) against the union-hosted **GLM** endpoint, hands it a skill + a task, and scores what it produces — orchestrated as **Flyte workflows on `demo.hosted.unionai.cloud`**, with path-based selection so only the diff --git a/evals/harness/run.py b/evals/harness/run.py index 441841b..0fd0a6f 100644 --- a/evals/harness/run.py +++ b/evals/harness/run.py @@ -68,7 +68,7 @@ def _print_table(results: list[ScenarioResult]) -> int: def main(argv: list[str] | None = None) -> int: - ap = argparse.ArgumentParser(prog="flyte-evals", description="Run flyte-skills evals locally") + ap = argparse.ArgumentParser(prog="flyte-evals", description="Run flyte-agent-plugin evals locally") ap.add_argument("--scenario", help="scenario id") ap.add_argument("--skill", help="filter by skill name") ap.add_argument("--tier", choices=["static", "trajectory", "real"], help="filter by tier") diff --git a/evals/manifest.yaml b/evals/manifest.yaml index 5b5a899..d54708b 100644 --- a/evals/manifest.yaml +++ b/evals/manifest.yaml @@ -1,4 +1,4 @@ -# Top-level configuration for the flyte-skills eval harness. +# Top-level configuration for the flyte-agent-plugin eval harness. # Scenario -> skill mapping is derived by scanning `scenarios_dir`; each scenario # file declares its own `skill`. This manifest holds the cross-cutting config. diff --git a/evals/pyproject.toml b/evals/pyproject.toml index 191e0f7..01a09b5 100644 --- a/evals/pyproject.toml +++ b/evals/pyproject.toml @@ -1,7 +1,7 @@ [project] -name = "flyte-skills-evals" +name = "flyte-agent-plugin-evals" version = "0.0.1" -description = "Testing & eval harness for the Flyte plugin skills." +description = "Testing & eval harness for the Flyte agent plugin." requires-python = ">=3.10" dependencies = [ "pyyaml>=6.0", diff --git a/evals/report.py b/evals/report.py index 78ec266..026c590 100644 --- a/evals/report.py +++ b/evals/report.py @@ -22,7 +22,7 @@ def to_markdown(results: list[dict]) -> str: total = len(results) failed = [r for r in results if not r["passed"]] lines = [ - f"### flyte-skills evals — {total - len(failed)}/{total} passing", + f"### flyte-agent-plugin evals — {total - len(failed)}/{total} passing", "", "| scenario | skill | harness | tier | pass | treat | ctrl | lift |", "|---|---|---|---|:--:|--:|--:|--:|", @@ -71,14 +71,14 @@ def to_html(results: list[dict]) -> str: ) passed = sum(1 for r in results if r["passed"]) return f""" -flyte-skills evals +flyte-agent-plugin evals -

flyte-skills evals — {passed}/{len(results)} passing

+

flyte-agent-plugin evals — {passed}/{len(results)} passing

diff --git a/evals/scenarios/flyte-migrate-config/resources.yaml b/evals/scenarios/flyte-migrate-config/resources.yaml new file mode 100644 index 0000000..18ebd18 --- /dev/null +++ b/evals/scenarios/flyte-migrate-config/resources.yaml @@ -0,0 +1,55 @@ +id: migrate-config-resources +skill: flyte-migrate-config +tier: trajectory +control: true +prompt: | + Migrate this Flyte 1 (flytekit) module to Flyte 2. Move the per-task image, + resources, and caching onto a shared TaskEnvironment. Write the result to + `migrated.py`. Do not run anything. + + ```python + import flytekit + from flytekit import task, workflow, Resources, ImageSpec + + image = ImageSpec(packages=["scikit-learn", "pandas"], python_version="3.11") + + @task(container_image=image, requests=Resources(cpu="2", mem="4Gi"), + cache=True, cache_version="1.0") + def train(n: int) -> float: + return float(n) + + @workflow + def main(n: int) -> float: + return train(n=n) + ``` +setup: + workspace: empty +checks: + - kind: file_glob + glob: "migrated.py" + - kind: python_parses + file_glob: "*.py" + - kind: python_imports + module: flyte + file_glob: "*.py" + - kind: contains_regex + file_glob: "*.py" + pattern: "TaskEnvironment" + - kind: contains_regex + file_glob: "*.py" + pattern: "Image|Resources" + # v1 image/caching constructs must be gone. + - kind: not_contains_regex + file_glob: "*.py" + pattern: "ImageSpec|cache_version" + - kind: not_contains_regex + file_glob: "*.py" + pattern: "import flytekit" +judge: + rubric: | + correctness: image/resources/cache moved onto a flyte.TaskEnvironment (image via + flyte.Image, cache="auto" or a CachePolicy replacing cache_version), train kept as + @env.task. skill_adherence: config declared once on the environment, not per-task. + completeness: no ImageSpec/cache_version/flytekit leftovers. + weights: {correctness: 0.5, skill_adherence: 0.3, completeness: 0.2} + pass_threshold: 0.7 diff --git a/evals/scenarios/flyte-migrate-config/static.yaml b/evals/scenarios/flyte-migrate-config/static.yaml new file mode 100644 index 0000000..85d5cce --- /dev/null +++ b/evals/scenarios/flyte-migrate-config/static.yaml @@ -0,0 +1,5 @@ +id: flyte-migrate-config-static +skill: flyte-migrate-config +tier: static +# Static tier lints the SKILL.md itself (frontmatter, code fences, no secrets). +# No agent, no LLM. Runs on every change. diff --git a/evals/scenarios/flyte-migrate-control-flow/conditional.yaml b/evals/scenarios/flyte-migrate-control-flow/conditional.yaml new file mode 100644 index 0000000..960cd67 --- /dev/null +++ b/evals/scenarios/flyte-migrate-control-flow/conditional.yaml @@ -0,0 +1,63 @@ +id: migrate-control-flow-conditional +skill: flyte-migrate-control-flow +tier: trajectory +control: true +prompt: | + Migrate this Flyte 1 (flytekit) module that uses `conditional()` branching to + Flyte 2, replacing the DSL with native Python control flow. Write the result to + `migrated.py`. Do not run anything. + + ```python + import flytekit + from flytekit import conditional + + @flytekit.task + def is_positive(x: int) -> bool: + return x > 0 + + @flytekit.task + def double(x: int) -> int: + return x * 2 + + @flytekit.task + def negate(x: int) -> int: + return -x + + @flytekit.workflow + def main(x: int) -> int: + return ( + conditional("check") + .if_(is_positive(x=x).is_true()) + .then(double(x=x)) + .else_() + .then(negate(x=x)) + ) + ``` +setup: + workspace: empty +checks: + - kind: file_glob + glob: "migrated.py" + - kind: python_parses + file_glob: "*.py" + - kind: python_imports + module: flyte + file_glob: "*.py" + - kind: contains_regex + file_glob: "*.py" + pattern: "\\bif\\b" + # v1 branching DSL must be gone. + - kind: not_contains_regex + file_glob: "*.py" + pattern: "conditional\\(" + - kind: not_contains_regex + file_glob: "*.py" + pattern: "import flytekit" +judge: + rubric: | + correctness: does migrated.py express the branch as native Python if/else + (double when positive, negate otherwise) inside an orchestrating @env.task? + skill_adherence: no conditional()/if_()/then() DSL, uses TaskEnvironment + + @env.task. completeness: no leftover flytekit constructs. + weights: {correctness: 0.5, skill_adherence: 0.3, completeness: 0.2} + pass_threshold: 0.7 diff --git a/evals/scenarios/flyte-migrate-control-flow/static.yaml b/evals/scenarios/flyte-migrate-control-flow/static.yaml new file mode 100644 index 0000000..b86788d --- /dev/null +++ b/evals/scenarios/flyte-migrate-control-flow/static.yaml @@ -0,0 +1,5 @@ +id: flyte-migrate-control-flow-static +skill: flyte-migrate-control-flow +tier: static +# Static tier lints the SKILL.md itself (frontmatter, code fences, no secrets). +# No agent, no LLM. Runs on every change. diff --git a/evals/scenarios/flyte-migrate-data-io/flytefile.yaml b/evals/scenarios/flyte-migrate-data-io/flytefile.yaml new file mode 100644 index 0000000..13d139a --- /dev/null +++ b/evals/scenarios/flyte-migrate-data-io/flytefile.yaml @@ -0,0 +1,57 @@ +id: migrate-data-io-flytefile +skill: flyte-migrate-data-io +tier: trajectory +control: true +prompt: | + Migrate this Flyte 1 (flytekit) module to Flyte 2, converting the FlyteFile and + StructuredDataset I/O to the Flyte 2 equivalents. Write the result to + `migrated.py`. Do not run anything. + + ```python + import flytekit + from flytekit import task, workflow + from flytekit.types.file import FlyteFile + from flytekit.types.structured import StructuredDataset + + @task + def load(path: str) -> FlyteFile: + return FlyteFile(path) + + @task + def to_table(f: FlyteFile) -> StructuredDataset: + import pandas as pd + return StructuredDataset(dataframe=pd.read_csv(f.path)) + + @workflow + def main(path: str) -> StructuredDataset: + return to_table(f=load(path=path)) + ``` +setup: + workspace: empty +checks: + - kind: file_glob + glob: "migrated.py" + - kind: python_parses + file_glob: "*.py" + - kind: python_imports + module: flyte + file_glob: "*.py" + # v2 offloaded types. + - kind: contains_regex + file_glob: "*.py" + pattern: "File|DataFrame" + # v1 types must be gone. + - kind: not_contains_regex + file_glob: "*.py" + pattern: "FlyteFile|StructuredDataset" + - kind: not_contains_regex + file_glob: "*.py" + pattern: "import flytekit" +judge: + rubric: | + correctness: FlyteFile -> flyte.io.File and StructuredDataset -> flyte.io.DataFrame, + with the load -> to_table -> main flow preserved as @env.task functions. + skill_adherence: uses TaskEnvironment + @env.task + flyte.io types. completeness: + no leftover flytekit / FlyteFile / StructuredDataset references. + weights: {correctness: 0.5, skill_adherence: 0.3, completeness: 0.2} + pass_threshold: 0.7 diff --git a/evals/scenarios/flyte-migrate-data-io/static.yaml b/evals/scenarios/flyte-migrate-data-io/static.yaml new file mode 100644 index 0000000..d0f1883 --- /dev/null +++ b/evals/scenarios/flyte-migrate-data-io/static.yaml @@ -0,0 +1,5 @@ +id: flyte-migrate-data-io-static +skill: flyte-migrate-data-io +tier: static +# Static tier lints the SKILL.md itself (frontmatter, code fences, no secrets). +# No agent, no LLM. Runs on every change. diff --git a/evals/scenarios/flyte-migrate-ml/static.yaml b/evals/scenarios/flyte-migrate-ml/static.yaml new file mode 100644 index 0000000..f81b7a6 --- /dev/null +++ b/evals/scenarios/flyte-migrate-ml/static.yaml @@ -0,0 +1,5 @@ +id: flyte-migrate-ml-static +skill: flyte-migrate-ml +tier: static +# Static tier lints the SKILL.md itself (frontmatter, code fences, no secrets). +# No agent, no LLM. Runs on every change. diff --git a/evals/scenarios/flyte-migrate-ml/train-gpu.yaml b/evals/scenarios/flyte-migrate-ml/train-gpu.yaml new file mode 100644 index 0000000..1186cca --- /dev/null +++ b/evals/scenarios/flyte-migrate-ml/train-gpu.yaml @@ -0,0 +1,50 @@ +id: migrate-ml-train-gpu +skill: flyte-migrate-ml +tier: trajectory +control: true +prompt: | + Migrate this Flyte 1 (flytekit) GPU training module to Flyte 2. Write the result + to `migrated.py`. Do not run anything. + + ```python + import flytekit + from flytekit import task, workflow, Resources, ImageSpec + + image = ImageSpec(packages=["torch", "torchvision"]) + + @task(container_image=image, requests=Resources(cpu="4", mem="16Gi", gpu="1")) + def train(epochs: int) -> str: + import torch + return f"trained {epochs} epochs on {torch.cuda.device_count()} gpus" + + @workflow + def main(epochs: int) -> str: + return train(epochs=epochs) + ``` +setup: + workspace: empty +checks: + - kind: file_glob + glob: "migrated.py" + - kind: python_parses + file_glob: "*.py" + - kind: python_imports + module: flyte + file_glob: "*.py" + - kind: contains_regex + file_glob: "*.py" + pattern: "TaskEnvironment" + - kind: contains_regex + file_glob: "*.py" + pattern: "gpu|GPU" + - kind: not_contains_regex + file_glob: "*.py" + pattern: "ImageSpec|import flytekit" +judge: + rubric: | + correctness: training task ported to @env.task with GPU resources expressed the + Flyte 2 way (e.g. a "T4:1"-style gpu string / flyte.Resources on the + TaskEnvironment), image via flyte.Image. skill_adherence: image/resources set once + on the environment. completeness: no ImageSpec/flytekit leftovers. + weights: {correctness: 0.5, skill_adherence: 0.3, completeness: 0.2} + pass_threshold: 0.7 diff --git a/evals/scenarios/flyte-migrate-tasks-workflows/hello-world.yaml b/evals/scenarios/flyte-migrate-tasks-workflows/hello-world.yaml new file mode 100644 index 0000000..5dc2236 --- /dev/null +++ b/evals/scenarios/flyte-migrate-tasks-workflows/hello-world.yaml @@ -0,0 +1,55 @@ +id: migrate-tasks-hello-world +skill: flyte-migrate-tasks-workflows +tier: trajectory +control: true +prompt: | + Migrate this Flyte 1 (flytekit) module to Flyte 2. Write the migrated code to + `migrated.py`. Do not run anything. + + ```python + import flytekit + + @flytekit.task + def say_hello(name: str) -> str: + return f"Hello, {name}!" + + @flytekit.task + def to_upper(greeting: str) -> str: + return greeting.upper() + + @flytekit.workflow + def main(name: str) -> str: + greeting = say_hello(name=name) + return to_upper(greeting=greeting) + ``` +setup: + workspace: empty +checks: + - kind: file_glob + glob: "migrated.py" + - kind: python_parses + file_glob: "*.py" + - kind: python_imports + module: flyte + file_glob: "*.py" + - kind: contains_regex + file_glob: "*.py" + pattern: "TaskEnvironment" + - kind: contains_regex + file_glob: "*.py" + pattern: "@\\w+\\.task" + # The whole point of the migration: v1 constructs must be gone. + - kind: not_contains_regex + file_glob: "*.py" + pattern: "import flytekit" + - kind: not_contains_regex + file_glob: "*.py" + pattern: "@\\w*\\.?workflow" +judge: + rubric: | + correctness: is migrated.py a faithful Flyte 2 port — say_hello + to_upper as + @env.task, and main as an orchestrating task that calls them (no @workflow)? + skill_adherence: TaskEnvironment created once, @env.task used, sequential calls + without the `>>` operator. completeness: no leftover flytekit imports/decorators. + weights: {correctness: 0.5, skill_adherence: 0.3, completeness: 0.2} + pass_threshold: 0.7 diff --git a/evals/scenarios/flyte-migrate-tasks-workflows/static.yaml b/evals/scenarios/flyte-migrate-tasks-workflows/static.yaml new file mode 100644 index 0000000..c3702e8 --- /dev/null +++ b/evals/scenarios/flyte-migrate-tasks-workflows/static.yaml @@ -0,0 +1,5 @@ +id: flyte-migrate-tasks-workflows-static +skill: flyte-migrate-tasks-workflows +tier: static +# Static tier lints the SKILL.md itself (frontmatter, code fences, no secrets). +# No agent, no LLM. Runs on every change. diff --git a/evals/scenarios/flyte-migrate/end-to-end.yaml b/evals/scenarios/flyte-migrate/end-to-end.yaml new file mode 100644 index 0000000..80d6a90 --- /dev/null +++ b/evals/scenarios/flyte-migrate/end-to-end.yaml @@ -0,0 +1,60 @@ +id: migrate-overview-end-to-end +skill: flyte-migrate +tier: trajectory +control: true +prompt: | + Migrate this small Flyte 1 (flytekit) module to Flyte 2. Apply the standard + changes: rename imports, move config to a TaskEnvironment, and turn the workflow + into an orchestrating task. Write the result to `migrated.py`. Do not run anything. + + ```python + import flytekit + from flytekit import task, workflow, Resources, map_task + + @flytekit.task(requests=Resources(cpu="1", mem="2Gi")) + def square(x: int) -> int: + return x * x + + @flytekit.task + def total(xs: list[int]) -> int: + return sum(xs) + + @flytekit.workflow + def main(xs: list[int]) -> int: + squared = map_task(square)(x=xs) + return total(xs=squared) + ``` +setup: + workspace: empty +checks: + - kind: file_glob + glob: "migrated.py" + - kind: python_parses + file_glob: "*.py" + - kind: python_imports + module: flyte + file_glob: "*.py" + - kind: contains_regex + file_glob: "*.py" + pattern: "TaskEnvironment" + - kind: contains_regex + file_glob: "*.py" + pattern: "@\\w+\\.task" + # v1 constructs across several categories must all be gone. + - kind: not_contains_regex + file_glob: "*.py" + pattern: "import flytekit" + - kind: not_contains_regex + file_glob: "*.py" + pattern: "map_task" + - kind: not_contains_regex + file_glob: "*.py" + pattern: "@\\w*\\.?workflow" +judge: + rubric: | + correctness: full port — flytekit->flyte imports, Resources on a TaskEnvironment, + square/total as @env.task, main as an orchestrating task, and map_task replaced by + flyte.map or asyncio.gather. skill_adherence: applies the two mechanical changes + from the migration overview. completeness: no flytekit/map_task/@workflow leftovers. + weights: {correctness: 0.5, skill_adherence: 0.3, completeness: 0.2} + pass_threshold: 0.7 diff --git a/evals/scenarios/flyte-migrate/static.yaml b/evals/scenarios/flyte-migrate/static.yaml new file mode 100644 index 0000000..4acebd2 --- /dev/null +++ b/evals/scenarios/flyte-migrate/static.yaml @@ -0,0 +1,5 @@ +id: flyte-migrate-static +skill: flyte-migrate +tier: static +# Static tier lints the SKILL.md itself (frontmatter, code fences, no secrets). +# No agent, no LLM. Runs on every change. diff --git a/evals/workflows/eval_wf.py b/evals/workflows/eval_wf.py index 698b50a..223097c 100644 --- a/evals/workflows/eval_wf.py +++ b/evals/workflows/eval_wf.py @@ -1,4 +1,4 @@ -"""Flyte orchestration of the flyte-skills eval harness. +"""Flyte orchestration of the flyte-agent-plugin eval harness. Runs on demo.hosted.unionai.cloud (org demo / project flytesnacks / domain development, remote image builder — see evals/config/flyte.yaml). The workflow @@ -21,7 +21,7 @@ from evals.workflows.images import eval_image env = flyte.TaskEnvironment( - name="flyte-skills-evals", + name="flyte-agent-plugin-evals", image=eval_image, resources=flyte.Resources(cpu="2", memory="4Gi"), secrets=[flyte.Secret(key="glm-api-key", as_env_var="GLM_API_KEY")], diff --git a/plugins/flyte/README.md b/plugins/flyte/README.md index ddab8a4..850a945 100644 --- a/plugins/flyte/README.md +++ b/plugins/flyte/README.md @@ -40,6 +40,24 @@ two bundled MCP servers. - **`flyte-sdk-ml`** — ML workload patterns (training, HPO, experiment tracking, evaluation, batch/real-time inference, monitoring). +### Migration (Flyte 1 → 2) + +Convert existing Flyte 1 (`flytekit`) code to Flyte 2, distilled from the official +[migration guide](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/). + +- **`flyte-migrate`** — start-here orchestrator: the `flytekit`→`flyte` shift, concept + mapping, the two mechanical changes, incremental strategy, hybrid v1/v2 pipelines, gotchas. +- **`flyte-migrate-tasks-workflows`** — `@task`/`@workflow`/`@dynamic` → a single `@env.task` + on a `TaskEnvironment`; ordering without `>>`; subworkflows as tasks. +- **`flyte-migrate-config`** — images (`ImageSpec`→`flyte.Image`), resources/GPUs, caching, + secrets, scheduling (`LaunchPlan`/`CronSchedule`→`Trigger`/`Cron`), and `pyflyte`→`flyte`. +- **`flyte-migrate-control-flow`** — `conditional()`→`if`/`else`, `@dynamic`→plain loops, + `on_failure`→`try`/`except`, `map_task`→`flyte.map`/`asyncio.gather`. +- **`flyte-migrate-data-io`** — `FlyteFile`/`FlyteDirectory`→`flyte.io.File`/`Dir`, + `StructuredDataset`→`flyte.io.DataFrame`, dataclasses/Pydantic I/O. +- **`flyte-migrate-ml`** — training, HPO, GPU/deep learning, batch inference, and the + new-in-v2 patterns (serving, apps, sandboxed execution). + ## Bundled MCP servers **Claude Code only.** The servers live in `.mcp.json`, which Claude Code reads by diff --git a/plugins/flyte/skills/flyte-migrate-config/SKILL.md b/plugins/flyte/skills/flyte-migrate-config/SKILL.md new file mode 100644 index 0000000..61cf157 --- /dev/null +++ b/plugins/flyte/skills/flyte-migrate-config/SKILL.md @@ -0,0 +1,452 @@ +--- +name: flyte-migrate-config +description: Migrates Flyte 1 task configuration, container images, resources, caching, secrets, scheduling, and CLI/config-file usage to Flyte 2 equivalents. Use when migrating Flyte 1 task configuration, images, resources, secrets, scheduling, or CLI/config-file usage to Flyte 2. Trigger words include resources, ImageSpec, cache_version, secrets, LaunchPlan, CronSchedule, pyflyte, config, register, and deploy. +--- + +# Flyte 1 to Flyte 2 Migration: Task Configuration and CLI/Config + +In Flyte 1, image, resources, caching, secrets, and scheduling were configured per-task on the `@task` decorator or per-workflow on a `LaunchPlan`. In Flyte 2 most of this moves to the `flyte.TaskEnvironment`, so it is declared once and shared. The CLI is renamed from `pyflyte` to `flyte` and the config file is trimmed down. This skill covers migrating those settings and commands. + +## Grounding References + +| Resource | URL | +|---|---| +| Migration guide (Task configuration) | https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/configuration/ | +| Migration guide (CLI) | https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/cli-and-configuration/ | +| Official docs | https://www.union.ai/docs/v2/flyte | +| Docs index (LLMs) | https://www.union.ai/docs/v2/flyte/llms.txt | +| SDK API reference | https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/ | +| CLI API reference | https://www.union.ai/docs/v2/union/api-reference/flyte-cli/ | +| Example code | https://github.com/unionai/unionai-examples | + +## Image, resources, and caching move to the TaskEnvironment + +Image, resources, and caching move from the `@task` decorator to the `TaskEnvironment`. Per-task settings like `retries` and `timeout` stay on `@env.task`. Note that `mem` is renamed to `memory`, and there are no separate `requests`/`limits` — a single `Resources` value serves as both. + +### Flyte 1 + +```python +from datetime import timedelta + +import flytekit +from flytekit import Resources + +image = flytekit.ImageSpec( + name="training-image", + packages=["scikit-learn", "pandas"], +) + +@flytekit.task( + container_image=image, + requests=Resources(cpu="2", mem="4Gi"), + limits=Resources(cpu="4", mem="8Gi"), + cache=True, + cache_version="1.0", + retries=3, + timeout=timedelta(minutes=30), +) +def train_epoch(step: int) -> float: + # A stand-in for a training step that returns the current loss. + return 1.0 / (step + 1) + +@flytekit.workflow +def main(step: int) -> float: + return train_epoch(step=step) +``` + +### Flyte 2 + +```python +from datetime import timedelta + +import flyte + +# Image, resources, and caching move to the TaskEnvironment, so they are declared +# once and shared by every task in the environment. +env = flyte.TaskEnvironment( + name="training", + image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn", "pandas"), + resources=flyte.Resources(cpu="2", memory="4Gi"), # "memory", not "mem" + cache="auto", +) + +# retries and timeout stay on the task decorator. +@env.task(retries=3, timeout=timedelta(minutes=30)) +def train_epoch(step: int) -> float: + # A stand-in for a training step that returns the current loss. + return 1.0 / (step + 1) + +@env.task +def main(step: int) -> float: + return train_epoch(step) +``` + +## Container images: ImageSpec to flyte.Image + +Flyte 1's `ImageSpec` is replaced by Flyte 2's `flyte.Image` with a fluent builder API. Instead of one constructor with many arguments, you start from a base and chain builder methods. + +```python +from flyte import Image + +image = ( + Image.from_debian_base(name="my-image", registry="ghcr.io/myorg", python_version=(3, 11)) + .with_pip_packages("pandas", "numpy") + .with_apt_packages("curl", "git") + .with_env_vars({"MY_VAR": "value"}) +) +``` + +| Constructor | Use case | +|---|---| +| `Image.from_debian_base()` | Most common; includes the Flyte SDK | +| `Image.from_base(image_uri)` | Start from any existing image | +| `Image.from_dockerfile(path)` | Complex custom builds | +| `Image.from_uv_script(path)` | UV-based projects | + +Common chainable builder methods: `.with_pip_packages(...)`, `.with_requirements(path)`, `.with_uv_project(path)`, `.with_apt_packages(...)`, `.with_commands([...])`, `.with_source_file(path, dst=...)`, `.with_source_folder(path, dst=...)`, `.with_env_vars({...})`, and `.with_workdir(...)`. + +| Flyte 1 `ImageSpec` | Flyte 2 `Image` | Notes | +|---|---|---| +| `name` | `name` (constructor) | Same | +| `registry` | `registry` (constructor) | Same | +| `python_version` | `python_version` (tuple) | `"3.11"` becomes `(3, 11)` | +| `packages` | `.with_pip_packages()` | Method instead of param | +| `apt_packages` | `.with_apt_packages()` | Method instead of param | +| `requirements` | `.with_requirements()` | Supports txt, poetry.lock, uv.lock | +| `env` | `.with_env_vars()` | Method instead of param | +| `commands` | `.with_commands()` | Method instead of param | +| `copy` / `source_root` | `.with_source_file()` / `.with_source_folder()` | More explicit methods | +| `base_image` | `Image.from_base()` | Different constructor | +| `builder` | Config file or `flyte.init()` | Global setting | +| `platform` | `platform` (constructor) | Tuple: `("linux/amd64", "linux/arm64")` | + +For a private registry, create an image-pull secret and reference it: + +```bash +flyte create secret --type image_pull my-registry-secret --from-file ~/.docker/config.json +``` + +```python +image = Image.from_debian_base( + registry="private.registry.com", + name="my-image", + registry_secret="my-registry-secret", +) +``` + +## Resources and GPUs + +A single `flyte.Resources` value serves as both request and limit — there are no separate `requests`/`limits`. Several parameters were renamed. + +| Flyte 1 | Flyte 2 | Notes | +|---|---|---| +| `cpu="1"` | `cpu="1"` | Same | +| `mem="2Gi"` | `memory="2Gi"` | Renamed | +| `gpu="1"` | `gpu="A100:1"` | `Type:count` format | +| `ephemeral_storage="10Gi"` | `disk="10Gi"` | Renamed | +| N/A | `shm="auto"` | New: shared memory | + +GPU type and count are combined into one string, replacing the separate Flyte 1 `accelerator=` argument: + +```python +env = flyte.TaskEnvironment( + name="gpu_env", + resources=flyte.Resources( + cpu="4", + memory="32Gi", + gpu="A100:2", # Type:count + # gpu="A100 80G:1" # 80GB variant + # gpu=flyte.GPU("A100", count=1, partition="1g.5gb") # MIG partition + ), +) +``` + +Supported GPU types include A10, A10G, A100, A100 80G, B200, H100, H200, L4, L40s, T4, V100, RTX PRO 6000, and GB10. + +## Caching: cache_version to cache="auto" / CachePolicy + +Caching is enabled at the env level with `cache="auto"` (or per-task on `@env.task`). The explicit `cache_version` string moves into a `flyte.Cache` object. + +| Behavior | Description | +|---|---| +| `"auto"` | Cache results and reuse if available | +| `"override"` | Always execute and overwrite the cache | +| `"disable"` | No caching (default for a `TaskEnvironment`) | + +```python +# Flyte 1: @task(cache=True, cache_version="1.0") +# Flyte 2: +@env.task(cache="auto") +def cached_task(x: int) -> int: + return x * 2 + +# Advanced control (replaces cache_version, serialize, ignored_inputs, ...) +@env.task(cache=flyte.Cache( + behavior="auto", + version_override="v1.0", + serialize=True, + ignored_inputs=("debug",), +)) +def advanced(x: int, debug: bool = False) -> int: + return x * 2 +``` + +## Secrets: current_context().secrets to env vars + +Secrets move from `secret_requests` on the task to `secrets` on the `TaskEnvironment`, and you read them from environment variables instead of `current_context().secrets` — for example, an API key for a model registry or hosted LLM. + +### Flyte 1 + +```python +from flytekit import task, workflow, Secret, current_context + +@task(secret_requests=[Secret(group="openai", key="api_key")]) +def call_api() -> str: + token = current_context().secrets.get(group="openai", key="api_key") + return f"token has {len(token)} chars" + +@workflow +def main() -> str: + return call_api() +``` + +### Flyte 2 + +```python +import os + +import flyte + +# Secrets are declared on the TaskEnvironment and injected as environment +# variables (instead of read through current_context().secrets). +env = flyte.TaskEnvironment( + name="secrets", + secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")], +) + +@env.task +def call_api() -> str: + token = os.getenv("OPENAI_API_KEY", "") + return f"token has {len(token)} chars" + +@env.task +def main() -> str: + return call_api() +``` + +A `flyte.Secret` can be mounted as an environment variable or as a file, and the access convention changes: + +```python +flyte.Secret(key="openai-key", as_env_var="OPENAI_API_KEY") # mount as env var +flyte.Secret(key="access-key", group="aws") # env var: AWS_ACCESS_KEY +flyte.Secret(key="ssl-cert", mount="/etc/flyte/secrets") # mount as a file +``` + +| Flyte 1 pattern | Flyte 2 pattern | +|---|---| +| `ctx.secrets.get(key="mykey", group="mygroup")` | `os.environ["MYGROUP_MYKEY"]` (auto-named) | +| `ctx.secrets.get(key="mykey")` | `os.environ["MY_SECRET"]` (with `as_env_var="MY_SECRET"`) | + +Create and manage secrets from the CLI: + +```bash +flyte create secret MY_SECRET_KEY --value my_secret_value +flyte create secret MY_SECRET_KEY --from-file /path/to/secret +flyte get secret +flyte delete secret MY_SECRET_KEY +``` + +## Scheduling: LaunchPlan + CronSchedule to flyte.Trigger + flyte.Cron + +A `LaunchPlan` with a `CronSchedule` (say, a nightly retraining job) becomes a `flyte.Trigger` attached directly to the task. Use `flyte.TriggerTime` to bind the scheduled fire time to an input, and deploy the trigger with `flyte deploy`. + +### Flyte 1 + +```python +from flytekit import task, workflow, LaunchPlan, CronSchedule + +@task +def retrain(kickoff_time: str) -> str: + return f"retrained model at {kickoff_time}" + +@workflow +def main(kickoff_time: str) -> str: + return retrain(kickoff_time=kickoff_time) + +# A LaunchPlan attaches a schedule (and default inputs) to a workflow. +nightly_retrain = LaunchPlan.get_or_create( + workflow=main, + name="nightly_retrain", + schedule=CronSchedule( + schedule="0 2 * * *", # 2 AM daily + kickoff_time_input_arg="kickoff_time", + ), +) +``` + +### Flyte 2 + +```python +from datetime import datetime + +import flyte + +env = flyte.TaskEnvironment(name="scheduling") + +# A Trigger replaces LaunchPlan + CronSchedule. It is attached directly to the +# task and deployed with it (flyte deploy). flyte.TriggerTime binds the +# scheduled fire time to a task input. +nightly_retrain = flyte.Trigger( + name="nightly_retrain", + automation=flyte.Cron("0 2 * * *"), # 2 AM daily + inputs={"trigger_time": flyte.TriggerTime}, + auto_activate=True, +) + +@env.task(triggers=nightly_retrain) +def main(trigger_time: datetime = datetime(2024, 1, 1, 2, 0)) -> str: + return f"retrained model at {trigger_time.isoformat()}" +``` + +Triggers support `flyte.Cron("0 9 * * *", timezone="America/New_York")` and `flyte.FixedRate(timedelta(hours=1))` as automations, plus convenience constructors like `flyte.Trigger.hourly()` and `flyte.Trigger.daily()`. + +## CLI command mapping: pyflyte to flyte + +The command-line tool is renamed from `pyflyte` to `flyte`, and remote is now the default. + +| Flyte 1 | Flyte 2 | Notes | +|---|---|---| +| `pyflyte run` | `flyte run` | Similar, different flags | +| `pyflyte run --remote` | `flyte run` | Remote is the default in Flyte 2 | +| `pyflyte run` (local) | `flyte run --local` | Local execution is now explicit | +| `pyflyte register` | `flyte deploy` | Different concept | +| `pyflyte package` | N/A | Not needed in Flyte 2 | +| `pyflyte serialize` | N/A | Not needed in Flyte 2 | + +### Running tasks — Flyte 1 + +```bash +# Local +pyflyte run my_module.py my_workflow --arg1 value1 + +# Remote +pyflyte --config config.yaml run --remote my_module.py my_workflow --arg1 value1 +``` + +### Running tasks — Flyte 2 + +```bash +# Remote (default) +flyte run my_module.py my_task --arg1 value1 + +# Local +flyte run --local my_module.py my_task --arg1 value1 + +# With an explicit config file +flyte --config config.yaml run my_module.py my_task --arg1 value1 +``` + +### Deploying (register to deploy) + +In Flyte 1 you registered a module; in Flyte 2 you deploy task environments. + +#### Flyte 1 + +```bash +pyflyte register my_module.py -p my-project -d development +``` + +#### Flyte 2 + +```bash +# Deploy a task environment +flyte deploy my_module.py my_env --project my-project --domain development + +# Deploy all environments in a file +flyte deploy --all my_module.py + +# Deploy with an explicit version, or recursively +flyte deploy --version v1.0.0 my_module.py my_env +flyte deploy --recursive --all ./src +``` + +### Key flag differences + +| Flyte 1 flag | Flyte 2 flag | Notes | +|---|---|---| +| `--remote` | (default) | Remote is the default | +| `--copy-all` | `--copy-style all` | File copying | +| N/A | `--copy-style loaded_modules` | Default: only imported modules | +| `-p, --project` | `--project` | Same | +| `-d, --domain` | `--domain` | Same | +| `-i, --image` | `--image` | Same format | +| N/A | `--follow, -f` | Follow execution logs | + +## Configuration files + +The config file lives in the same place (`~/.flyte/config.yaml`), but the environment variable changes from `FLYTECTL_CONFIG` to `FLYTE_CONFIG`, and the format is simpler. + +### Flyte 1 + +```yaml +admin: + endpoint: dns:///your-cluster.hosted.unionai.cloud + insecure: false + authType: Pkce +``` + +### Flyte 2 + +```yaml +admin: + endpoint: dns:///your-cluster.hosted.unionai.cloud + +image: + builder: remote # or "local" + +task: + domain: development + org: your-org + project: your-project +``` + +| Setting | Flyte 1 | Flyte 2 | +|---|---|---| +| Endpoint | `admin.endpoint` | `admin.endpoint` | +| Auth type | `admin.authType` | Auto-detected (PKCE default) | +| Project | CLI flag `-p` | `task.project` (default) | +| Domain | CLI flag `-d` | `task.domain` (default) | +| Organization | CLI flag `--org` | `task.org` (default) | +| Image builder | N/A | `image.builder` (`local` or `remote`) | + +### Configuring in code + +```python +import flyte + +# From a config file (auto-discovers, or pass a path) +flyte.init_from_config() +flyte.init_from_config("path/to/config.yaml") + +# Programmatically +flyte.init( + endpoint="flyte.example.com", + project="my-project", + domain="development", +) +``` + +For API-key authentication in non-interactive environments, use `flyte.init_from_api_key()`. + +## Anti-Patterns + +1. **Don't keep `image`, `resources`, and `cache` on `@env.task`** — move them onto the shared `flyte.TaskEnvironment`; only per-task settings like `retries` and `timeout` stay on `@env.task`. +2. **Don't use `mem`, `ephemeral_storage`, or separate `requests`/`limits`** — use `memory`, `disk`, and a single `flyte.Resources` value that serves as both. +3. **Don't pass GPUs with `gpu="1"` plus `accelerator=`** — combine type and count into one `Type:count` string like `gpu="A100:2"`. +4. **Don't rebuild `ImageSpec`'s many constructor args** — start from a base (`Image.from_debian_base()`) and chain `.with_*` builder methods. +5. **Don't keep `cache_version="1.0"`** — use `cache="auto"` for the common case, or `flyte.Cache(version_override=...)` for advanced control. +6. **Don't read secrets via `current_context().secrets.get(...)`** — declare them on the `TaskEnvironment` and read the injected environment variable with `os.environ` / `os.getenv`. +7. **Don't recreate `LaunchPlan` + `CronSchedule`** — use `flyte.Trigger` with `flyte.Cron` attached to the task, and deploy it with `flyte deploy`. +8. **Don't run `pyflyte ... --remote`** — `flyte run` is remote by default; add `--local` explicitly for in-process runs. +9. **Don't use `pyflyte register`** — use `flyte deploy` to deploy task environments. +10. **Don't set `FLYTECTL_CONFIG` or rely on `admin.authType`** — use `FLYTE_CONFIG` and the simpler config format with auto-detected auth. diff --git a/plugins/flyte/skills/flyte-migrate-control-flow/SKILL.md b/plugins/flyte/skills/flyte-migrate-control-flow/SKILL.md new file mode 100644 index 0000000..13dcab8 --- /dev/null +++ b/plugins/flyte/skills/flyte-migrate-control-flow/SKILL.md @@ -0,0 +1,356 @@ +--- +name: flyte-migrate-control-flow +description: "Migrates Flyte 1 branching, dynamic workflows, failure handling, and fan-out to native Flyte 2 Python. Use when migrating Flyte 1 branching, dynamic workflows, failure handling, or map_task/fan-out to Flyte 2. Trigger words: conditional, @dynamic, map_task, on_failure, branching, parallelism, fan-out, flyte.map, asyncio.gather." +--- + +# Flyte 1 to 2 Migration: Control Flow and Parallelism + +Flyte 1 expressed branching, dynamic fan-out, and failure handling through DSL constructs (`conditional()`, `@dynamic`, `@workflow(on_failure=...)`) and `map_task`. In Flyte 2 these are all ordinary Python, because orchestration runs as real Python at runtime. Native `if`/`elif`/`else` replaces the conditional DSL, plain task loops replace `@dynamic`, `try`/`except` replaces `on_failure`, and `flyte.map` / `asyncio.gather` replace `map_task`. + +## Grounding References + +| Resource | URL | +|---|---| +| Migration guide (Control flow) | https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/control-flow/ | +| Migration guide (Parallelism) | https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/parallelism/ | +| Official docs | https://www.union.ai/docs/v2/flyte | +| Docs index (LLMs) | https://www.union.ai/docs/v2/flyte/llms.txt | +| SDK API reference | https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/ | +| Example code | https://github.com/unionai/unionai-examples | +| Flyte MCP tools | Available via `flyte-mcp` server | + +## Conditional Execution + +The `conditional()` DSL becomes ordinary Python `if` / `elif` / `else` — for example, choosing a model based on dataset size. + +### Flyte 1 + +```python +from flytekit import task, workflow, conditional + +@task +def train_gradient_boosting(n_rows: int) -> str: + return f"trained gradient boosting on {n_rows} rows" + +@task +def train_logistic_regression(n_rows: int) -> str: + return f"trained logistic regression on {n_rows} rows" + +@workflow +def main(n_rows: int) -> str: + # Pick the model based on dataset size. + return ( + conditional("model_choice") + .if_(n_rows > 10_000) + .then(train_gradient_boosting(n_rows=n_rows)) + .else_() + .then(train_logistic_regression(n_rows=n_rows)) + ) +``` + +### Flyte 2 + +```python +import flyte + +env = flyte.TaskEnvironment(name="conditional") + +@env.task +def train_gradient_boosting(n_rows: int) -> str: + return f"trained gradient boosting on {n_rows} rows" + +@env.task +def train_logistic_regression(n_rows: int) -> str: + return f"trained logistic regression on {n_rows} rows" + +# Branching is now ordinary Python control flow -- no conditional() DSL. +@env.task +def main(n_rows: int) -> str: + if n_rows > 10_000: + return train_gradient_boosting(n_rows) + return train_logistic_regression(n_rows) +``` + +## Dynamic Workflows + +`@dynamic` existed so a task could generate a variable number of subtask calls at runtime (e.g. one per data partition discovered at runtime). In Flyte 2 every task can do this natively, so `@dynamic` simply disappears — loop over runtime data in an ordinary `@env.task`. + +### Flyte 1 + +```python +from flytekit import task, workflow, dynamic + +@task +def list_partitions(n: int) -> list[int]: + return list(range(n)) + +@task +def process_partition(partition_id: int) -> int: + # Aggregate one data partition. + return partition_id * 2 + +@dynamic +def process_all(partitions: list[int]) -> list[int]: + results = [] + for partition_id in partitions: + results.append(process_partition(partition_id=partition_id)) + return results + +@workflow +def main(n: int) -> list[int]: + partitions = list_partitions(n=n) + return process_all(partitions=partitions) +``` + +### Flyte 2 + +```python +import flyte + +env = flyte.TaskEnvironment(name="dynamic") + +@env.task +def process_partition(partition_id: int) -> int: + # Aggregate one data partition. + return partition_id * 2 + +# No @dynamic decorator needed: a plain task can loop over runtime data (e.g. a +# variable number of partitions discovered at runtime) and call other tasks. +@env.task +def main(n: int) -> list[int]: + return [process_partition(partition_id) for partition_id in range(n)] +``` + +## Error Handling + +Flyte 1's `@workflow(on_failure=...)` handler becomes ordinary Python `try` / `except` — catch a failed training run, run cleanup, and recover or re-raise. + +### Flyte 1 + +```python +from flytekit import task, workflow + +@task +def train_fold(max_depth: int) -> float: + if max_depth <= 0: + raise ValueError("max_depth must be positive") + # Return validation accuracy for this hyperparameter. + return 0.90 + 0.001 * max_depth + +@task +def notify_failure() -> None: + print("training run failed -- sending alert") + +# The on_failure handler runs if any node in the workflow fails. There is no +# try/except inside a Flyte 1 workflow. +@workflow(on_failure=notify_failure) +def main(max_depth: int) -> float: + return train_fold(max_depth=max_depth) +``` + +### Flyte 2 + +```python +import flyte + +env = flyte.TaskEnvironment(name="error_handling") + +@env.task +async def train_fold(max_depth: int) -> float: + if max_depth <= 0: + raise ValueError("max_depth must be positive") + return 0.90 + 0.001 * max_depth + +# Failure handling is ordinary Python try/except -- no on_failure handler. +@env.task +async def main(max_depth: int) -> float: + try: + return await train_fold(max_depth) + except ValueError as e: + print(f"invalid hyperparameter ({e}); falling back to a safe default") + # Recover with a safe default instead of failing the whole run. + return await train_fold(max_depth=6) +``` + +Flyte 2 also exposes typed errors, so you can catch a specific failure and retry with more resources — a common need for memory-hungry training jobs: + +```python +try: + return await train_fold(sample_size) +except flyte.errors.OOMError: + # Retry the same task with a larger memory request. + return await train_fold.override( + resources=flyte.Resources(memory="16Gi") + )(sample_size) +``` + +## Fan-out: map_task + +`map_task()` becomes `flyte.map()`, a near drop-in replacement. The one catch: `flyte.map` returns a generator, so wrap it in `list()`. For new code, the idiomatic approach is Python `async`/`await` with `asyncio.gather()`, which gives finer control over concurrency and error handling. + +### Flyte 1 + +```python +from functools import partial + +from flytekit import task, workflow, map_task + +@task +def get_shards(n: int) -> list[int]: + return list(range(n)) + +@task +def score_shard(shard_id: int, model_version: int) -> int: + # Score one shard of records with the given model version. + return shard_id * model_version + +@workflow +def main(n: int, model_version: int) -> list[int]: + shards = get_shards(n=n) + return map_task( + partial(score_shard, model_version=model_version), + concurrency=10, + )(shard_id=shards) +``` + +### Flyte 2 (flyte.map) + +```python +import flyte +from functools import partial + +env = flyte.TaskEnvironment(name="map_task") + +@env.task +def score_shard(shard_id: int, model_version: int) -> int: + # Score one shard of records with the given model version. + return shard_id * model_version + +@env.task +def main(n: int, model_version: int) -> list[int]: + bound = partial(score_shard, model_version=model_version) + # flyte.map is a drop-in for map_task, but it returns a generator, so wrap + # it in list() to materialize the results. + return list(flyte.map(bound, range(n), concurrency=10)) +``` + +### Flyte 2 (asyncio.gather) + +```python +import asyncio + +import flyte + +env = flyte.TaskEnvironment(name="map_task") + +@env.task +async def score_shard_async(shard_id: int, model_version: int) -> int: + return shard_id * model_version + +@env.task +async def main_async(n: int, model_version: int) -> list[int]: + # asyncio.gather is the idiomatic Flyte 2 way to fan out. + coros = [score_shard_async(i, model_version) for i in range(n)] + return list(await asyncio.gather(*coros)) +``` + +### Choosing flyte.map vs asyncio.gather + +| Feature | `flyte.map` (sync) | `asyncio.gather` (async) | +|---|---|---| +| Syntax | `list(flyte.map(fn, items))` | `await asyncio.gather(*tasks)` | +| Concurrency limit | Built-in `concurrency=N` | Use `asyncio.Semaphore` | +| Streaming / as-completed | No | Yes, via `asyncio.as_completed()` | +| Error handling | `return_exceptions=True` | Check return type | + +Use `flyte.map` for the smallest change from Flyte 1 `map_task`, or when stuck in synchronous code. Use `asyncio.gather` for new code where you want streaming results or fine-grained concurrency control. + +### Concurrency Control and Error Handling + +`map_task`'s `concurrency` and `min_success_ratio` become an `asyncio.Semaphore` and `return_exceptions=True`: + +```python +import asyncio + +@env.task +async def main(items: list[int], max_concurrent: int = 5) -> list[str]: + sem = asyncio.Semaphore(max_concurrent) + + async def process_with_limit(item: int) -> str: + async with sem: + return await process_item(item) + + tasks = [process_with_limit(i) for i in items] + results = await asyncio.gather(*tasks, return_exceptions=True) + + return [r for r in results if not isinstance(r, Exception)] +``` + +## Data Backfills + +Reprocessing a range of dates is a textbook `@dynamic` use case in Flyte 1, because the number of days is only known at runtime. In Flyte 2 it's a plain task that builds the date range and fans the days out with `asyncio.gather`. + +### Flyte 1 + +```python +from datetime import date, timedelta + +from flytekit import task, workflow, dynamic + +@task +def process_day(day: str) -> int: + # Reprocess a single day's partition; return the row count. + return len(day) + +# @dynamic is needed because the number of days is only known at runtime. +@dynamic +def backfill(start: str, days: int) -> list[int]: + base = date.fromisoformat(start) + results = [] + for i in range(days): + day = (base + timedelta(days=i)).isoformat() + results.append(process_day(day=day)) + return results + +@workflow +def main(start: str, days: int) -> list[int]: + return backfill(start=start, days=days) +``` + +### Flyte 2 + +```python +import asyncio +from datetime import date, timedelta + +import flyte + +env = flyte.TaskEnvironment(name="data_backfill") + +@env.task +async def process_day(day: str) -> int: + # Reprocess a single day's partition; return the row count. + return len(day) + +# A plain task builds the date range at runtime and fans the days out in +# parallel with asyncio.gather -- no @dynamic and no map_task needed. +@env.task +async def main(start: str, days: int) -> list[int]: + base = date.fromisoformat(start) + coros = [ + process_day((base + timedelta(days=i)).isoformat()) + for i in range(days) + ] + return list(await asyncio.gather(*coros)) +``` + +## Anti-Patterns + +1. **Don't import `conditional`, `dynamic`, or `map_task` from `flytekit`** — none exist in Flyte 2. Branching is native `if`/`elif`/`else`, dynamic fan-out is a plain task loop, and `map_task` becomes `flyte.map`. +2. **Don't keep the `conditional().if_().then().else_()` DSL** — rewrite it as ordinary Python control flow inside an `@env.task`. +3. **Don't reach for `@dynamic`** — every Flyte 2 task can loop over runtime data and call other tasks, so drop the decorator entirely. +4. **Don't pass `on_failure=...` to `@workflow`** — there is no workflow decorator in Flyte 2; handle failures with ordinary `try`/`except` inside a task. +5. **Don't forget to `list()` a `flyte.map` result** — it returns a generator, not a materialized list. +6. **Don't forget to `await` async fan-out** — `asyncio.gather(*coros)` returns a coroutine; without `await` you get a coroutine object instead of results. +7. **Don't drop concurrency limits** — port `concurrency=N` to `flyte.map(..., concurrency=N)` or an `asyncio.Semaphore`, and `min_success_ratio` to `return_exceptions=True` with filtering. +8. **Don't use Union-only features** — avoid `ReusePolicy` and other Union-specific APIs. diff --git a/plugins/flyte/skills/flyte-migrate-data-io/SKILL.md b/plugins/flyte/skills/flyte-migrate-data-io/SKILL.md new file mode 100644 index 0000000..fa83f34 --- /dev/null +++ b/plugins/flyte/skills/flyte-migrate-data-io/SKILL.md @@ -0,0 +1,309 @@ +--- +name: flyte-migrate-data-io +description: Migrates Flyte 1 data types and offloaded I/O to Flyte 2. Use when migrating Flyte 1 data types and I/O (files, directories, dataframes, dataclasses) to Flyte 2, converting FlyteFile, FlyteDirectory, or StructuredDataset to flyte.io.File, flyte.io.Dir, and flyte.io.DataFrame. Trigger words are FlyteFile, FlyteDirectory, StructuredDataset, DataFrame, dataclass, Pydantic, type, I/O, and serialization. +--- + +# Flyte 1 to 2 Migration: Data Types and I/O + +Flyte 2 renames the offloaded-data types and makes their I/O `async`, but the mental model is the same: pass lightweight references to large data between tasks, not the materialized bytes. `FlyteFile`, `FlyteDirectory`, and `StructuredDataset` become `flyte.io.File`, `flyte.io.Dir`, and `flyte.io.DataFrame`. Plain dataclasses and Pydantic `BaseModel`s work directly as task I/O with no JSON mixin. + +## Grounding References + +| Resource | URL | +|---|---| +| Migration guide | https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/data-io/ | +| Official docs | https://www.union.ai/docs/v2/flyte | +| Docs index (LLMs) | https://www.union.ai/docs/v2/flyte/llms.txt | +| SDK API reference | https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/ | +| Example code | https://github.com/unionai/unionai-examples | +| Flyte MCP tools | Available via `flyte-mcp` server | + +## Type Mapping + +| Flyte 1 | Flyte 2 | Notes | +|---|---|---| +| `flytekit.types.file.FlyteFile` | `flyte.io.File` | I/O is `async` | +| `flytekit.types.directory.FlyteDirectory` | `flyte.io.Dir` | I/O is `async` | +| `flytekit.types.structured.StructuredDataset` | `flyte.io.DataFrame` | build with `from_df`, read with `open(...).all()` | +| `@dataclass_json` + `@dataclass` | plain `@dataclass` | no mixin needed | +| Pydantic `BaseModel` (+ config) | plain Pydantic `BaseModel` | works directly as task I/O | + +## Offloaded Data: The Mental Model + +`File`, `Dir`, and `DataFrame` are lightweight references (pointers) to data offloaded in blob storage — not the materialized bytes. In Flyte 2 the read/write operations are `async`: upload with `await File.from_local(local_path)`, read with `async with f.open("rb") as fh: await fh.read()`, build a frame with `flyte.io.DataFrame.from_df(df)` (sync constructor), and read it with `await fdf.open(pandas.DataFrame).all()`. + +## Files and Directories + +`FlyteFile` and `FlyteDirectory` become `flyte.io.File` and `flyte.io.Dir` — the way you pass model artifacts and datasets between tasks. Use `await File.from_local(...)` to upload and `async with file.open(...)` to read. + +### Flyte 1 + +```python +import os + +from flytekit import task, workflow, current_context +from flytekit.types.file import FlyteFile + +@task +def write_file(content: str) -> FlyteFile: + path = os.path.join(current_context().working_directory, "out.txt") + with open(path, "w") as f: + f.write(content) + return FlyteFile(path=path) + +@task +def read_file(f: FlyteFile) -> str: + with open(f.download()) as fh: + return fh.read() + +@workflow +def main(content: str) -> str: + f = write_file(content=content) + return read_file(f=f) +``` + +### Flyte 2 + +```python +import flyte +from flyte.io import File + +env = flyte.TaskEnvironment(name="files") + +@env.task +async def write_file(content: str) -> File: + with open("out.txt", "w") as f: + f.write(content) + # File.from_local uploads the file to blob storage and returns a reference + # (a lightweight pointer, not the materialized bytes). + return await File.from_local("out.txt") + +@env.task +async def read_file(f: File) -> str: + async with f.open("rb") as fh: + return (await fh.read()).decode("utf-8") + +@env.task +async def main(content: str) -> str: + f = await write_file(content) + return await read_file(f) +``` + +Directories follow the same pattern: import `Dir` from `flyte.io` and use its `async` upload/read methods in place of `FlyteDirectory`. See [Files and directories](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/files-and-directories) for more. + +## DataFrames + +`StructuredDataset` becomes `flyte.io.DataFrame`. Construct one with `flyte.io.DataFrame.from_df(df)` and read it back with `await df.open(pandas.DataFrame).all()`. + +### Flyte 1 + +```python +import pandas as pd +from flytekit import task, workflow +from flytekit.types.structured import StructuredDataset + +@task +def make_df() -> StructuredDataset: + df = pd.DataFrame({"employee_id": [1, 2, 3], "salary": [50000, 60000, 70000]}) + return StructuredDataset(dataframe=df) + +@task +def total_payroll(sd: StructuredDataset) -> float: + df = sd.open(pd.DataFrame).all() + return float(df["salary"].sum()) + +@workflow +def main() -> float: + return total_payroll(sd=make_df()) +``` + +### Flyte 2 + +```python +import pandas as pd +import flyte +import flyte.io + +env = flyte.TaskEnvironment( + name="dataframe", + image=flyte.Image.from_debian_base().with_pip_packages("pandas", "pyarrow"), +) + +@env.task +async def make_df() -> flyte.io.DataFrame: + df = pd.DataFrame({"employee_id": [1, 2, 3], "salary": [50000, 60000, 70000]}) + # StructuredDataset becomes flyte.io.DataFrame. + return flyte.io.DataFrame.from_df(df) + +@env.task +async def total_payroll(fdf: flyte.io.DataFrame) -> float: + df = await fdf.open(pd.DataFrame).all() + return float(df["salary"].sum()) + +@env.task +async def main() -> float: + return await total_payroll(await make_df()) +``` + +Add the dataframe dependencies (for example `pandas` and `pyarrow`) to the `TaskEnvironment` image. See [DataFrames](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/dataframes) for more. + +## Dataclasses and Structured Types + +Flyte 1 required a `@dataclass_json` mixin for dataclass I/O. In Flyte 2, plain dataclasses (and Pydantic `BaseModel`s) work directly as task inputs and outputs — handy for passing around a training config. + +### Flyte 1 + +```python +from dataclasses import dataclass + +from dataclasses_json import dataclass_json +from flytekit import task, workflow + +@dataclass_json +@dataclass +class TrainingConfig: + learning_rate: float + n_estimators: int + max_depth: int = 6 + +@task +def make_config(learning_rate: float, n_estimators: int) -> TrainingConfig: + return TrainingConfig(learning_rate=learning_rate, n_estimators=n_estimators) + +@task +def train(config: TrainingConfig) -> str: + return ( + f"trained with lr={config.learning_rate}, " + f"n_estimators={config.n_estimators}, max_depth={config.max_depth}" + ) + +@workflow +def main(learning_rate: float, n_estimators: int) -> str: + config = make_config(learning_rate=learning_rate, n_estimators=n_estimators) + return train(config=config) +``` + +### Flyte 2 + +```python +from dataclasses import dataclass + +import flyte + +env = flyte.TaskEnvironment(name="dataclasses") + +# Plain dataclasses work directly as task I/O -- no @dataclass_json mixin needed. +# Pydantic BaseModels work the same way. +@dataclass +class TrainingConfig: + learning_rate: float + n_estimators: int + max_depth: int = 6 + +@env.task +def make_config(learning_rate: float, n_estimators: int) -> TrainingConfig: + return TrainingConfig(learning_rate=learning_rate, n_estimators=n_estimators) + +@env.task +def train(config: TrainingConfig) -> str: + return ( + f"trained with lr={config.learning_rate}, " + f"n_estimators={config.n_estimators}, max_depth={config.max_depth}" + ) + +@env.task +def main(learning_rate: float, n_estimators: int) -> str: + config = make_config(learning_rate, n_estimators) + return train(config) +``` + +## Data ETL: Putting It Together + +Extract, clean, aggregate, and write out a feature table. `StructuredDataset` becomes `flyte.io.DataFrame`, and the tasks become `async`. + +### Flyte 1 + +```python +import pandas as pd +from flytekit import task, workflow +from flytekit.types.structured import StructuredDataset + +@task +def extract() -> pd.DataFrame: + # Read raw transaction records (stand-in for a real source). + return pd.DataFrame( + { + "user_id": [1, 1, 2, 3, 3, 3], + "amount": [10.0, 5.0, 20.0, 7.5, 2.5, 1.0], + } + ) + +@task +def transform(df: pd.DataFrame) -> StructuredDataset: + # Clean and aggregate into a per-user feature table. + df = df[df["amount"] > 0] + agg = df.groupby("user_id", as_index=False)["amount"].sum() + return StructuredDataset(dataframe=agg) + +@task +def load(sd: StructuredDataset) -> int: + df = sd.open(pd.DataFrame).all() + return len(df) + +@workflow +def main() -> int: + raw = extract() + features = transform(df=raw) + return load(sd=features) +``` + +### Flyte 2 + +```python +import pandas as pd +import flyte +import flyte.io + +env = flyte.TaskEnvironment( + name="data_etl", + image=flyte.Image.from_debian_base().with_pip_packages("pandas", "pyarrow"), +) + +@env.task +async def extract() -> pd.DataFrame: + # Read raw transaction records (stand-in for a real source). + return pd.DataFrame( + { + "user_id": [1, 1, 2, 3, 3, 3], + "amount": [10.0, 5.0, 20.0, 7.5, 2.5, 1.0], + } + ) + +@env.task +async def transform(df: pd.DataFrame) -> flyte.io.DataFrame: + # Clean and aggregate into a per-user feature table. + df = df[df["amount"] > 0] + agg = df.groupby("user_id", as_index=False)["amount"].sum() + # StructuredDataset becomes flyte.io.DataFrame. + return flyte.io.DataFrame.from_df(agg) + +@env.task +async def load(sd: flyte.io.DataFrame) -> int: + df = await sd.open(pd.DataFrame).all() + return len(df) + +@env.task +async def main() -> int: + raw = await extract() + features = await transform(raw) + return await load(features) +``` + +## Anti-Patterns + +1. **Don't call the offloaded-data I/O synchronously** — `File.from_local`, `file.open(...).read()`, and `DataFrame.open(...).all()` are `async` in Flyte 2; `await` them inside `async` tasks. +2. **Don't keep the `@dataclass_json` mixin** — plain `@dataclass` and Pydantic `BaseModel`s serialize as task I/O directly; drop `dataclasses_json`. +3. **Don't return `StructuredDataset(dataframe=df)`** — use `flyte.io.DataFrame.from_df(df)` instead. +4. **Don't materialize large data into task outputs** — return `File`, `Dir`, or `DataFrame` references, not the raw bytes or full frames. +5. **Don't forget the dataframe dependencies** — add `pandas` and `pyarrow` (or your engine) to the `TaskEnvironment` image so DataFrame I/O works remotely. +6. **Don't import from `flytekit.types.*`** — import `File` and `Dir` from `flyte.io`, and use `flyte.io.DataFrame`. diff --git a/plugins/flyte/skills/flyte-migrate-ml/SKILL.md b/plugins/flyte/skills/flyte-migrate-ml/SKILL.md new file mode 100644 index 0000000..01b64dc --- /dev/null +++ b/plugins/flyte/skills/flyte-migrate-ml/SKILL.md @@ -0,0 +1,611 @@ +--- +name: flyte-migrate-ml +description: Migrates Flyte 1 machine learning code to Flyte 2 and unlocks net-new v2 patterns. Use when migrating Flyte 1 ML workloads (training, HPO, GPU/deep learning, batch inference) to Flyte 2, specifying GPU resources, or building the end-to-end pipeline pattern. Trigger words - migrate training, HPO, GPU, deep learning, batch inference, model serving, pytorch. +--- + +# Flyte 1 to Flyte 2 ML Migration Skill + +Migrate existing Flyte 1 ML workloads — small-model training, hyperparameter optimization, deep learning on GPUs, and batch inference — to Flyte 2, then take advantage of patterns that were not possible in Flyte 1 (real-time serving, apps, sandboxed execution). + +This skill is specifically about **migrating existing v1 ML code**. For greenfield authoring in Flyte 2, use the companion skills: + +- `flyte-sdk-ml` — writing new ML training / inference tasks in Flyte 2. +- `flyte-sdk-app` — writing new apps and serving endpoints. +- `flyte-sdk-agent` — writing new agents and sandboxed / code-mode workloads. + +## Grounding References + +| Resource | URL | +|---|---| +| Migration guide (ML workloads) | https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/ml-workloads/ | +| Migration guide (New in Flyte 2) | https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/new-in-flyte-2/ | +| Official docs | https://www.union.ai/docs/v2/flyte | +| Docs index (LLMs) | https://www.union.ai/docs/v2/flyte/llms.txt | +| SDK API reference | https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/ | +| Example code | https://github.com/unionai/unionai-examples | +| Flyte MCP tools | Available via `flyte-mcp` server | + +## Migration Cheat Sheet + +| Flyte 1 | Flyte 2 | +|---|---| +| `ImageSpec(name=..., packages=[...])` | `flyte.Image.from_debian_base().with_pip_packages(...)` | +| `@task(container_image=..., requests=..., cache=...)` | Set `image`, `resources`, `cache` once on `flyte.TaskEnvironment`, then `@env.task` | +| `Resources(cpu=..., mem=...)` | `flyte.Resources(cpu=..., memory=...)` (note `mem` becomes `memory`) | +| `Resources(gpu="1")` + `accelerator=T4` | `flyte.Resources(gpu="T4:1")` | +| `FlyteFile` / `FlyteFile(path=...)` | `flyte.io.File` / `await File.from_local(...)` | +| `model_file.download()` | `await model_file.download()` | +| `current_context().working_directory` | `os.getcwd()` | +| `@workflow` | An orchestrating `@env.task` (plain `async` Python) | +| `map_task(fn)(x=xs)` | `await asyncio.gather(*[fn(x) for x in xs])` | +| A "pick the best" task | Plain Python after `gather` | + +## Small model training (scikit-learn / XGBoost) + +Train a model, persist it as a `File`, and evaluate it. Image, resources, and caching move to the `TaskEnvironment`; `FlyteFile` becomes `flyte.io.File`. + +### Flyte 1 + +```python +import os + +import joblib +from flytekit import task, workflow, ImageSpec, Resources, current_context +from flytekit.types.file import FlyteFile +from sklearn.datasets import load_breast_cancer +from sklearn.model_selection import train_test_split +from xgboost import XGBClassifier + +image = ImageSpec( + name="xgb-image", + packages=["xgboost", "scikit-learn", "joblib"], +) + +@task(container_image=image, requests=Resources(cpu="2", mem="4Gi")) +def train_model(n_estimators: int, max_depth: int) -> FlyteFile: + data = load_breast_cancer() + X_train, _, y_train, _ = train_test_split(data.data, data.target, random_state=42) + model = XGBClassifier(n_estimators=n_estimators, max_depth=max_depth) + model.fit(X_train, y_train) + + model_path = os.path.join(current_context().working_directory, "model.json") + joblib.dump(model, model_path) + return FlyteFile(path=model_path) + +@task(container_image=image) +def evaluate(model_file: FlyteFile) -> float: + model = joblib.load(model_file.download()) + data = load_breast_cancer() + _, X_test, _, y_test = train_test_split(data.data, data.target, random_state=42) + return float(model.score(X_test, y_test)) + +@workflow +def main(n_estimators: int, max_depth: int) -> float: + model = train_model(n_estimators=n_estimators, max_depth=max_depth) + return evaluate(model_file=model) +``` + +### Flyte 2 + +```python +import os + +import joblib +import flyte +from flyte.io import File +from sklearn.datasets import load_breast_cancer +from sklearn.model_selection import train_test_split +from xgboost import XGBClassifier + +env = flyte.TaskEnvironment( + name="train_xgboost", + image=flyte.Image.from_debian_base().with_pip_packages( + "xgboost", "scikit-learn", "joblib" + ), + resources=flyte.Resources(cpu="2", memory="4Gi"), +) + +@env.task +async def train_model(n_estimators: int, max_depth: int) -> File: + data = load_breast_cancer() + X_train, _, y_train, _ = train_test_split(data.data, data.target, random_state=42) + model = XGBClassifier(n_estimators=n_estimators, max_depth=max_depth) + model.fit(X_train, y_train) + + model_path = os.path.join(os.getcwd(), "model.json") + joblib.dump(model, model_path) + return await File.from_local(model_path) + +@env.task +async def evaluate(model_file: File) -> float: + local_path = await model_file.download() + model = joblib.load(local_path) + data = load_breast_cancer() + _, X_test, _, y_test = train_test_split(data.data, data.target, random_state=42) + return float(model.score(X_test, y_test)) + +@env.task +async def main(n_estimators: int, max_depth: int) -> float: + model = await train_model(n_estimators, max_depth) + return await evaluate(model) +``` + +## Hyperparameter optimization + +Fan out one training run per hyperparameter, then pick the best. In Flyte 1 the grid search runs through `map_task` and the "pick the best" step must itself be a task. In Flyte 2 you `gather` the runs and select the winner in plain Python. + +### Flyte 1 + +```python +from flytekit import task, workflow, map_task +from sklearn.datasets import load_iris +from sklearn.ensemble import RandomForestClassifier +from sklearn.model_selection import cross_val_score + +@task +def get_grid() -> list[int]: + return [2, 4, 8, 16] + +@task +def train_eval(max_depth: int) -> float: + data = load_iris() + model = RandomForestClassifier(max_depth=max_depth, random_state=42) + scores = cross_val_score(model, data.data, data.target, cv=3) + return float(scores.mean()) + +@task +def best_score(scores: list[float]) -> float: + return max(scores) + +@workflow +def main() -> float: + grid = get_grid() + # Fan out one training run per hyperparameter value. + scores = map_task(train_eval)(max_depth=grid) + return best_score(scores=scores) +``` + +### Flyte 2 + +```python +import asyncio + +import flyte +from sklearn.datasets import load_iris +from sklearn.ensemble import RandomForestClassifier +from sklearn.model_selection import cross_val_score + +env = flyte.TaskEnvironment( + name="hpo", + image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn"), +) + +@env.task +async def train_eval(max_depth: int) -> float: + data = load_iris() + model = RandomForestClassifier(max_depth=max_depth, random_state=42) + scores = cross_val_score(model, data.data, data.target, cv=3) + return float(scores.mean()) + +@env.task +async def main() -> dict: + grid = [2, 4, 8, 16] + # Fan out one training run per hyperparameter value... + scores = await asyncio.gather(*[train_eval(d) for d in grid]) + # ...then pick the best in plain Python (impossible in a Flyte 1 workflow). + best_idx = max(range(len(scores)), key=lambda i: scores[i]) + return {"best_max_depth": grid[best_idx], "best_score": scores[best_idx]} +``` + +## Large model training (deep learning) + +GPU configuration moves to the `TaskEnvironment`: the Flyte 1 `Resources(gpu="1")` plus a separate `accelerator=T4` become a single `gpu="T4:1"` string on `flyte.Resources`. + +### Flyte 1 + +```python +from flytekit import task, workflow, ImageSpec, Resources +from flytekit.extras.accelerators import T4 +import torch +import torch.nn as nn + +image = ImageSpec( + name="dl-image", + packages=["torch"], +) + +@task( + container_image=image, + requests=Resources(cpu="4", mem="16Gi", gpu="1"), + accelerator=T4, +) +def train(epochs: int) -> float: + device = "cuda" if torch.cuda.is_available() else "cpu" + model = nn.Linear(10, 1).to(device) + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + loss_fn = nn.MSELoss() + + X = torch.randn(128, 10, device=device) + y = torch.randn(128, 1, device=device) + + loss = torch.tensor(0.0) + for _ in range(epochs): + optimizer.zero_grad() + loss = loss_fn(model(X), y) + loss.backward() + optimizer.step() + return float(loss.item()) + +@workflow +def main(epochs: int) -> float: + return train(epochs=epochs) +``` + +### Flyte 2 + +```python +import flyte +import torch +import torch.nn as nn + +# GPU type and count go in a single "T4:1"-style string. For multi-node +# distributed training, wrap the training task with the torch elastic plugin. +env = flyte.TaskEnvironment( + name="train_deep_learning", + image=flyte.Image.from_debian_base().with_pip_packages("torch"), + resources=flyte.Resources(cpu="4", memory="16Gi", gpu="T4:1"), +) + +@env.task +async def train(epochs: int) -> float: + device = "cuda" if torch.cuda.is_available() else "cpu" + model = nn.Linear(10, 1).to(device) + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + loss_fn = nn.MSELoss() + + X = torch.randn(128, 10, device=device) + y = torch.randn(128, 1, device=device) + + loss = torch.tensor(0.0) + for _ in range(epochs): + optimizer.zero_grad() + loss = loss_fn(model(X), y) + loss.backward() + optimizer.step() + return float(loss.item()) + +@env.task +async def main(epochs: int) -> float: + return await train(epochs) +``` + +For multi-node distributed training (PyTorch elastic, etc.), wrap the training task with the torch elastic plugin. See the Resources docs and plugin integrations at https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/resources. + +## Batch inference + +Load a trained model once and score many batches in parallel. `map_task` with a `partial`-bound model becomes `asyncio.gather` over the batches, reusing the same model reference. + +### Flyte 1 + +```python +import os +from functools import partial + +import joblib +from flytekit import task, workflow, map_task, ImageSpec, current_context +from flytekit.types.file import FlyteFile +from sklearn.datasets import load_iris +from sklearn.ensemble import RandomForestClassifier + +image = ImageSpec(name="inference-image", packages=["scikit-learn", "joblib"]) + +@task(container_image=image) +def train_model() -> FlyteFile: + data = load_iris() + model = RandomForestClassifier().fit(data.data, data.target) + model_path = os.path.join(current_context().working_directory, "model.joblib") + joblib.dump(model, model_path) + return FlyteFile(path=model_path) + +@task(container_image=image) +def get_batches() -> list[list[list[float]]]: + data = load_iris() + rows = data.data.tolist() + # Split the rows into batches of 30. + return [rows[i : i + 30] for i in range(0, len(rows), 30)] + +@task(container_image=image) +def score_batch(model_file: FlyteFile, batch: list[list[float]]) -> list[int]: + model = joblib.load(model_file.download()) + return [int(p) for p in model.predict(batch)] + +@workflow +def main() -> list[list[int]]: + model = train_model() + batches = get_batches() + return map_task(partial(score_batch, model_file=model))(batch=batches) +``` + +### Flyte 2 + +```python +import asyncio +import os + +import joblib +import flyte +from flyte.io import File +from sklearn.datasets import load_iris +from sklearn.ensemble import RandomForestClassifier + +env = flyte.TaskEnvironment( + name="batch_inference", + image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn", "joblib"), +) + +@env.task +async def train_model() -> File: + data = load_iris() + model = RandomForestClassifier().fit(data.data, data.target) + model_path = os.path.join(os.getcwd(), "model.joblib") + joblib.dump(model, model_path) + return await File.from_local(model_path) + +@env.task +async def score_batch(model_file: File, batch: list[list[float]]) -> list[int]: + local_path = await model_file.download() + model = joblib.load(local_path) + return [int(p) for p in model.predict(batch)] + +@env.task +async def main() -> list[list[int]]: + model = await train_model() + rows = load_iris().data.tolist() + batches = [rows[i : i + 30] for i in range(0, len(rows), 30)] + # Score every batch in parallel, reusing the same model reference. + coros = [score_batch(model, batch) for batch in batches] + return list(await asyncio.gather(*coros)) +``` + +## A complete example: end-to-end ML pipeline + +Putting it together — a load / train / evaluate pipeline shows the image, resources, caching, file I/O, and orchestration changes in one place. Image, resources, and cache are set **once** on the `TaskEnvironment`, and the "workflow" is just an orchestrating task. + +### Flyte 1 + +```python +import os + +import joblib +import pandas as pd +from flytekit import task, workflow, ImageSpec, Resources, current_context +from flytekit.types.file import FlyteFile +from sklearn.datasets import load_iris +from sklearn.ensemble import RandomForestClassifier + +image = ImageSpec( + name="ml-image", + packages=["pandas", "scikit-learn", "joblib"], +) + +@task( + container_image=image, + requests=Resources(cpu="2", mem="4Gi"), + cache=True, + cache_version="1.0", +) +def load_data() -> pd.DataFrame: + data = load_iris(as_frame=True) + df = data.frame + df["species"] = data.target + return df + +@task(container_image=image) +def train_model(data: pd.DataFrame) -> FlyteFile: + model = RandomForestClassifier() + X = data.drop("species", axis=1) + y = data["species"] + model.fit(X, y) + + model_path = os.path.join(current_context().working_directory, "model.joblib") + joblib.dump(model, model_path) + return FlyteFile(path=model_path) + +@task(container_image=image) +def evaluate(model_file: FlyteFile, data: pd.DataFrame) -> float: + model = joblib.load(model_file.download()) + X = data.drop("species", axis=1) + y = data["species"] + return float(model.score(X, y)) + +@workflow +def main() -> float: + data = load_data() + model = train_model(data=data) + return evaluate(model_file=model, data=data) +``` + +### Flyte 2 + +```python +import os + +import joblib +import pandas as pd +import flyte +from flyte.io import File +from sklearn.datasets import load_iris +from sklearn.ensemble import RandomForestClassifier + +# Image, resources, and cache are set once on the TaskEnvironment. +env = flyte.TaskEnvironment( + name="ml_pipeline", + image=flyte.Image.from_debian_base().with_pip_packages( + "pandas", "scikit-learn", "joblib" + ), + resources=flyte.Resources(cpu="2", memory="4Gi"), + cache="auto", +) + +@env.task +async def load_data() -> pd.DataFrame: + data = load_iris(as_frame=True) + df = data.frame + df["species"] = data.target + return df + +@env.task +async def train_model(data: pd.DataFrame) -> File: + model = RandomForestClassifier() + X = data.drop("species", axis=1) + y = data["species"] + model.fit(X, y) + + model_path = os.path.join(os.getcwd(), "model.joblib") + joblib.dump(model, model_path) + return await File.from_local(model_path) + +@env.task +async def evaluate(model_file: File, data: pd.DataFrame) -> float: + local_path = await model_file.download() + model = joblib.load(local_path) + X = data.drop("species", axis=1) + y = data["species"] + return float(model.score(X, y)) + +# The "workflow" is just an orchestrating task. +@env.task +async def main() -> float: + data = await load_data() + model = await train_model(data) + return await evaluate(model, data) +``` + +## New in Flyte 2 + +Flyte 1 was a batch orchestration system: everything ran as a finite DAG that started, did work, and finished. Flyte 2 keeps all of that and adds long-running services, high-throughput batch inference, and sandboxed code execution — so the same project that trains your model can also serve it, host a dashboard, saturate a GPU, or safely run LLM-generated code. There is no v1 counterpart to migrate here; these are net-new capabilities that your migrated training code unlocks. For greenfield authoring of these, see the `flyte-sdk-app` and `flyte-sdk-agent` skills. + +### Real-time inference and model serving + +Instead of scoring a batch and exiting, you can stand up an always-on REST endpoint from a `FastAPIAppEnvironment` and deploy it with `flyte.deploy`. The app can load a model artifact produced by one of your migrated training tasks. + +```python +app = FastAPI(title="ML Model API") + +# Define request/response models +class PredictionRequest(BaseModel): + feature1: float + feature2: float + feature3: float + +class PredictionResponse(BaseModel): + prediction: float + probability: float + +# Load model (you would typically load this from storage) +model = None + +@asynccontextmanager +async def lifespan(app: FastAPI): + global model + model_path = os.getenv("MODEL_PATH", "/app/models/model.joblib") + # In production, load from your storage + if os.path.exists(model_path): + with open(model_path, "rb") as f: + model = joblib.load(f) + yield + +@app.post("/predict", response_model=PredictionResponse) +async def predict(request: PredictionRequest): + # Make prediction + # prediction = model.predict([[request.feature1, request.feature2, request.feature3]]) + + # Dummy prediction for demo + prediction = 0.85 + probability = 0.92 + + return PredictionResponse( + prediction=prediction, + probability=probability, + ) + +env = FastAPIAppEnvironment( + name="ml-model-api", + app=app, + image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages( + "fastapi", + "uvicorn", + "scikit-learn", + "pydantic", + "joblib", + ), + parameters=[ + flyte.app.Parameter( + name="model_file", + value=flyte.io.File.from_existing_remote("s3://bucket/models/model.joblib"), + mount="/app/models", + env_var="MODEL_PATH", + ), + ], + resources=flyte.Resources(cpu=2, memory="2Gi"), + requires_auth=False, +) +``` + +For serving large language models, the `flyteplugins-vllm` integration gives you a production-grade vLLM server (with autoscaling to zero) via `VLLMAppEnvironment`. Any web app — a Streamlit dashboard, a Gradio demo, a Flask backend — runs as a `flyte.app.AppEnvironment` that you configure with image, resources, port, autoscaling, and a custom subdomain, then `flyte.serve`. + +### Dynamic batching for GPU inference + +For in-process batch inference, `DynamicBatcher` from `flyte.extras` keeps an expensive GPU saturated: async producers load and preprocess data concurrently while a single consumer feeds the model in optimally-sized batches, with built-in backpressure. This replaces the Flyte 1 pattern of standing up a separate inference server just to get request batching. + +```python +import asyncio +from flyte.extras import DynamicBatcher + +async with DynamicBatcher( + process_fn=run_inference, # takes a batch, returns results in the same order + target_batch_cost=1000, # cost budget per batch + max_batch_size=64, # hard cap on records per batch + batch_timeout_s=0.05, # max wait before dispatching a partial batch +) as batcher: + futures = [await batcher.submit(record) for record in records] + results = await asyncio.gather(*futures) +``` + +`submit()` is non-blocking and returns a `Future`; when the queue is full it applies backpressure automatically. See the batch inference docs for `TokenBatcher` (token-aware LLM batching). + +### Sandboxed code execution + +`flyte.sandbox.create()` runs arbitrary Python code or shell commands inside an ephemeral, single-use Docker container — built on demand from declared dependencies, executed once, then discarded. Only declared inputs go in and only declared outputs come back, which makes it the safe way to run untrusted code, most importantly code generated by an LLM. + +```python +# sandbox_environment provides the base runtime for code sandboxes. +# Include it in depends_on so the sandbox runtime is available when tasks execute. +env = flyte.TaskEnvironment( + name="sandbox-demo", + image=flyte.Image.from_debian_base(name="sandbox-demo"), + depends_on=[sandbox_environment], +) + +# Auto-IO mode: pure computation. The code string runs in an isolated sandbox; +# only the declared inputs go in and only the declared outputs come back. +sum_sandbox = flyte.sandbox.create( + name="sum-to-n", + code="total = sum(range(n + 1)) if conditional else 0", + inputs={"n": int, "conditional": bool}, + outputs={"total": int}, +) +``` + +Call it from a task with `await sum_sandbox.run.aio(n=10, conditional=True)`. This also powers **code mode** (programmatic tool calling), where an agent writes a whole program instead of emitting one tool call at a time. + +## Anti-Patterns + +1. **Don't keep `@task` / `@workflow` per-task config** — move `image`, `resources`, and `cache` onto a single `flyte.TaskEnvironment` and decorate with `@env.task`. +2. **Don't leave a separate "pick the best" task** — after `asyncio.gather`, select the winner in plain Python inside the orchestrating task. +3. **Don't carry `map_task` + `partial` into v2** — fan out with `asyncio.gather` over coroutines, reusing the same model reference. +4. **Don't split GPU type and count** — replace `Resources(gpu="1")` + `accelerator=T4` with a single `gpu="T4:1"` string on `flyte.Resources`. +5. **Don't use `mem=` or `current_context().working_directory`** — use `memory=` on `flyte.Resources` and `os.getcwd()` for local paths. +6. **Don't forget `await`** — `File.from_local`, `download`, and task calls are all async in v2. +7. **Don't hand-roll a serving container or a request-batching server** — use a `FastAPIAppEnvironment` / `AppEnvironment` for serving and `DynamicBatcher` for GPU batching. +8. **Don't run untrusted or LLM-generated code inline** — use `flyte.sandbox.create()` with `sandbox_environment` in `depends_on`. diff --git a/plugins/flyte/skills/flyte-migrate-tasks-workflows/SKILL.md b/plugins/flyte/skills/flyte-migrate-tasks-workflows/SKILL.md new file mode 100644 index 0000000..6e4f767 --- /dev/null +++ b/plugins/flyte/skills/flyte-migrate-tasks-workflows/SKILL.md @@ -0,0 +1,258 @@ +--- +name: flyte-migrate-tasks-workflows +description: >- + Migrates Flyte 1 tasks and workflows to Flyte 2, where the @task, @workflow, + and @dynamic decorators collapse into a single @env.task on a + flyte.TaskEnvironment and a workflow becomes a task that calls other tasks. + Use when the user is migrating Flyte 1 tasks/workflows to Flyte 2. Trigger + words: migrate task, migrate workflow, @task, @workflow, @dynamic, + TaskEnvironment, env.task. +--- + +# Flyte 1 to 2 Migration: Tasks and Workflows + +The biggest structural change in Flyte 2 is that everything is a task. The Flyte 1 `@task`, `@workflow`, and `@dynamic` decorators all collapse into a single `@env.task` on a `flyte.TaskEnvironment`, and a "workflow" is now just a task that calls other tasks. This skill covers the structural shift, sequential ordering without the `>>` operator, nested subworkflows, TaskEnvironment configuration basics, and the full `@task` parameter mapping. + +## Grounding References + +| Resource | URL | +|---|---| +| Migration guide | https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/tasks-and-workflows/ | +| Official docs | https://www.union.ai/docs/v2/flyte | +| Docs index (LLMs) | https://www.union.ai/docs/v2/flyte/llms.txt | +| SDK API reference | https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/ | +| Example code | https://github.com/unionai/unionai-examples | +| Flyte MCP tools | Available via `flyte-mcp` server | + +## The Structural Shift + +In Flyte 1 you decorated units of work with `@task`, composed them with `@workflow`, and used `@dynamic` for runtime-generated graphs. In Flyte 2 you create one `flyte.TaskEnvironment` that carries the configuration, then decorate every function — leaf tasks and orchestrating "workflows" alike — with `@env.task`. There is no separate `@workflow` decorator: the entrypoint is just a task that calls other tasks. + +## Hello World: Tasks and Workflows + +A `@task` plus `@workflow` becomes two `@env.task`s, where the entrypoint task calls the others. Sequential calls are naturally ordered — no `>>` operator required. + +### Flyte 1 + +```python +import flytekit + +@flytekit.task +def say_hello(name: str) -> str: + return f"Hello, {name}!" + +@flytekit.task +def to_upper(greeting: str) -> str: + return greeting.upper() + +@flytekit.workflow +def main(name: str) -> str: + greeting = say_hello(name=name) + return to_upper(greeting=greeting) +``` + +### Flyte 2 + +```python +import flyte + +env = flyte.TaskEnvironment(name="hello_world") + +@env.task +def say_hello(name: str) -> str: + return f"Hello, {name}!" + +@env.task +def to_upper(greeting: str) -> str: + return greeting.upper() + +# The "workflow" is now just a task that calls other tasks. +@env.task +def main(name: str) -> str: + greeting = say_hello(name) + return to_upper(greeting) +``` + +Note that Flyte 2 task calls pass arguments positionally (`say_hello(name)`) rather than requiring keyword arguments as in Flyte 1 (`say_hello(name=name)`). + +## Chaining and Ordering + +In Flyte 1 you sometimes used `>>` to force ordering between tasks with no data dependency. In Flyte 2, sequential (synchronous) calls run in the order they are written, and `await`ing async tasks in sequence does the same. The `>>` operator is gone. + +### Flyte 1 + +```python +from flytekit import task, workflow + +@task +def clear_staging_table() -> None: + # Side effect only: truncate the staging table. + print("cleared staging table") + +@task +def load_into_staging() -> None: + # Side effect only: load fresh rows into staging. + print("loaded staging table") + +@task +def publish_to_prod() -> None: + # Side effect only: swap staging into the production table. + print("published to prod") + +@workflow +def main() -> None: + clear = clear_staging_table() + load = load_into_staging() + publish = publish_to_prod() + + # These tasks pass no data between them, so use the >> operator to force + # ordering: clear must finish before load, which must finish before publish. + clear >> load >> publish +``` + +### Flyte 2 + +```python +import flyte + +env = flyte.TaskEnvironment(name="staging_publish") + +@env.task +def clear_staging_table() -> None: + print("cleared staging table") + +@env.task +def load_into_staging() -> None: + print("loaded staging table") + +@env.task +def publish_to_prod() -> None: + print("published to prod") + +# Sequential (synchronous) calls run in the order they're written, even when no +# data flows between them. The Flyte 1 `>>` ordering operator is gone. +@env.task +def main() -> None: + clear_staging_table() + load_into_staging() + publish_to_prod() +``` + +## Subworkflows + +A `@workflow` invoked by another `@workflow` (for example, a reusable preprocessing pipeline) becomes a task that calls other tasks — nest them as deeply as you like. + +### Flyte 1 + +```python +from flytekit import task, workflow + +@task +def impute(value: float) -> float: + # Replace missing/negative sentinel values with 0. + return value if value >= 0 else 0.0 + +@task +def scale(value: float) -> float: + return value / 100.0 + +@workflow +def preprocess(value: float) -> float: + imputed = impute(value=value) + return scale(value=imputed) + +@workflow +def main(raw_value: float) -> float: + return preprocess(value=raw_value) +``` + +### Flyte 2 + +```python +import flyte + +env = flyte.TaskEnvironment(name="subworkflow") + +@env.task +def impute(value: float) -> float: + # Replace missing/negative sentinel values with 0. + return value if value >= 0 else 0.0 + +@env.task +def scale(value: float) -> float: + return value / 100.0 + +# A preprocessing "subworkflow" is just a task that calls other tasks. +@env.task +def preprocess(value: float) -> float: + imputed = impute(value) + return scale(imputed) + +@env.task +def main(raw_value: float) -> float: + return preprocess(raw_value) +``` + +## TaskEnvironment Configuration + +The `TaskEnvironment` holds the configuration that Flyte 1 spread across the `@task` decorator. The task decorator can still override a few settings per-task. + +```python +import flyte + +env = flyte.TaskEnvironment( + name="my_env", # Required: unique name + image=flyte.Image.from_debian_base(), # Or a string, or "auto" + resources=flyte.Resources( + cpu="2", + memory="4Gi", + gpu="A100:1", + disk="10Gi", + ), + env_vars={"LOG_LEVEL": "INFO"}, + secrets=[flyte.Secret(key="api-key", as_env_var="API_KEY")], + cache="auto", # "auto", "override", "disable", or a Cache object + reusable=flyte.ReusePolicy(replicas=5, idle_ttl=60), + interruptible=True, +) + +# The task decorator can override some settings: +@env.task( + short_name="my_task", # Display name + cache="disable", # Override cache + retries=3, # Retry count + timeout=3600, # Seconds or a timedelta + report=True, # Generate an HTML report +) +def my_task(x: int) -> int: + return x +``` + +## Parameter Mapping: `@task` to `TaskEnvironment` + `@env.task` + +| Flyte 1 `@task` parameter | Flyte 2 location | Notes | +|---|---|---| +| `container_image` | `TaskEnvironment(image=...)` | Env-level only | +| `requests` | `TaskEnvironment(resources=...)` | Env-level only | +| `limits` | `TaskEnvironment(resources=...)` | Combined with requests (single value) | +| `environment` | `TaskEnvironment(env_vars=...)` | Env-level only | +| `secret_requests` | `TaskEnvironment(secrets=...)` | Env-level only | +| `cache` | Both | Can override at task level | +| `cache_version` | `flyte.Cache(version_override=...)` | In a `Cache` object | +| `retries` | `@env.task(retries=...)` | Task-level only | +| `timeout` | `@env.task(timeout=...)` | Task-level only | +| `interruptible` | Both | Can override at task level | +| `pod_template` | Both | Can override at task level | +| `deprecated` | N/A | Not in Flyte 2 | +| `docs` | `@env.task(docs=...)` | Task-level only | + +For image, resource, secret, and caching detail, see the Task configuration migration page. + +## Anti-Patterns + +1. **Don't reach for `>>`** — the ordering operator is gone. Sequential synchronous calls already run in written order; `await` async tasks in sequence for the same effect. +2. **Don't look for a `@workflow` decorator** — there isn't one. The orchestrating entrypoint is just another `@env.task` that calls other tasks. +3. **Don't look for a `@dynamic` decorator** — dynamic graphs also collapse into ordinary `@env.task` functions that call other tasks at runtime. +4. **Don't put heavy compute in the orchestrating task** — keep the entrypoint task focused on calling other tasks; push CPU/GPU/memory-intensive work into leaf tasks whose resources you can tune per environment. +5. **Don't set image, resources, or secrets on the `@env.task` decorator** — those are env-level and belong on `TaskEnvironment`. Only per-task overrides like `retries`, `timeout`, `cache`, and `short_name` go on `@env.task`. +6. **Don't keep passing every argument by keyword** — Flyte 2 task calls accept positional arguments (`say_hello(name)`). diff --git a/plugins/flyte/skills/flyte-migrate/SKILL.md b/plugins/flyte/skills/flyte-migrate/SKILL.md new file mode 100644 index 0000000..51116f7 --- /dev/null +++ b/plugins/flyte/skills/flyte-migrate/SKILL.md @@ -0,0 +1,353 @@ +--- +name: flyte-migrate +description: Entry-point orchestrator for porting Flyte 1 (flytekit) code to Flyte 2 (flyte). Explains the v1 to v2 shift, the terminology mapping, a recommended migration strategy, hybrid v1/v2 pipelines, and routes to sibling migration skills. Use when the user wants to migrate, port, or upgrade Flyte 1 (flytekit) code to Flyte 2. Trigger words are migrate, flytekit, v1 to v2, port, upgrade, convert workflow. +--- + +# Flyte 1 to 2 Migration Skill + +This is the entry point for migrating a Flyte 1 (`flytekit`) codebase to Flyte 2 (`flyte`). Flyte 2 is a fundamental shift: there is no `@workflow` decorator, everything is a `@env.task`, orchestration runs as real Python at runtime, and parallelism is expressed with `asyncio`. This skill explains the overall shift, gives a recommended migration strategy, covers hybrid v1/v2 pipelines during the transition, and routes to the sibling skills that handle each theme in depth. + +## Grounding References + +| Resource | URL | +|---|---| +| Migration guide | https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/ | +| Official docs | https://www.union.ai/docs/v2/flyte | +| Docs index (LLMs) | https://www.union.ai/docs/v2/flyte/llms.txt | +| SDK API reference | https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/ | +| Example code | https://github.com/unionai/unionai-examples | +| Flyte MCP tools | Available via `flyte-mcp` server | + +## The overall v1 to v2 shift + +Two conceptual shifts motivate almost every change — **pure Python execution** and the **asynchronous model** — after which most migrations come down to a couple of mechanical moves. + +- **`flytekit` (package) becomes `flyte`.** Imports change from `import flytekit` to `import flyte`. +- **`pyflyte` (CLI) becomes `flyte`.** The command-line tool was renamed. +- **Everything is a task.** In Flyte 1, `@workflow` functions were constrained to a DSL subset of Python that compiled to a static DAG. In Flyte 2 there is **no `@workflow` decorator**: everything is a `@env.task`, and a "workflow" is simply a task that calls other tasks. Loops, conditionals, and `try`/`except` work anywhere. +- **Async is the parallelism model.** Flyte 2 is built on `asyncio`, with the Flyte orchestrator acting as the event loop, scheduling awaited tasks across distributed infrastructure. `await` signals where a task can be scheduled in parallel, and `asyncio.gather` tells the orchestrator that a set of tasks are independent. + +### Simplified API mapping + +| Use case | Flyte 1 | Flyte 2 | +| --- | --- | --- | +| Environment management | `N/A` | `TaskEnvironment` | +| Perform basic computation | `@task` | `@env.task` | +| Combine tasks into a workflow | `@workflow` | `@env.task` | +| Create dynamic workflows | `@dynamic` | `@env.task` | +| Fanout parallelism | `flytekit.map` | Python `for` loop with `asyncio.gather` | +| Conditional execution | `flytekit.conditional` | Python `if-elif-else` | +| Catching workflow failures | `@workflow(on_failure=...)` | Python `try-except` | + +## Terminology and concept mapping + +Several Flyte 1 concepts were renamed or reshaped in Flyte 2. The table below maps the ones you'll meet most often. + +| Flyte 1 | Flyte 2 | Notes | +|---|---|---| +| `flytekit` (package) | `flyte` (package) | The Python SDK was renamed; imports change from `import flytekit` to `import flyte`. | +| `pyflyte` (CLI) | `flyte` (CLI) | The command-line tool was renamed. | +| `@task` / `@workflow` / `@dynamic` | `@env.task` | A single task decorator off a `flyte.TaskEnvironment`. Workflows and dynamic tasks are no longer distinct constructs: everything is a task, and orchestration is plain Python. | +| `map_task()` | `flyte.map()` | Plus `asyncio.gather()` for async fan-out. | +| `conditional()` | native `if` / `elif` / `else` | Branching is now ordinary Python control flow, not a DSL. | +| `ImageSpec` | `flyte.Image` | Container image definition. | +| `current_context()` | `flyte.ctx()` | Runtime context access. | +| `FlyteFile` / `FlyteDirectory` | `flyte.io.File` / `flyte.io.Dir` | Offloaded file and directory references. | +| `StructuredDataset` | `flyte.io.DataFrame` | Offloaded tabular data. | +| `LaunchPlan` | `flyte.Trigger` | Scheduling and parameterized entry points. | +| `CronSchedule` | `flyte.Cron` | Cron-based scheduling, used with a `flyte.Trigger`. | +| Decks (`enable_deck=True`) | Reports (`report=True`) | Custom HTML rendered in the UI during/after a run. | + +## The two mechanical changes behind (almost) every migration + +Most of a migration comes down to two moves. + +### 1. Move task configuration into a `TaskEnvironment` + +Instead of configuring the image, resources, and caching on each task decorator, configure them once on a `flyte.TaskEnvironment` and share it across tasks: + +```python +env = flyte.TaskEnvironment( + name="training", + image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn", "pandas"), + resources=flyte.Resources(cpu="2", memory="4Gi"), + cache="auto", +) +``` + +### 2. Replace `@task` / `@workflow` / `@dynamic` with `@env.task` + +Every decorated function becomes an `@env.task`. There is no separate workflow or dynamic construct: a "workflow" is simply a task that calls other tasks, and orchestration is plain Python. The `env` in `@env.task` is just the variable you assigned your `TaskEnvironment` to — name it whatever you like. + +## Package imports + +The package is renamed from `flytekit` to `flyte`, and the workflow/dynamic/map_task imports disappear. + +### Flyte 1 + +```python +import flytekit +from flytekit import task, workflow, dynamic, map_task +from flytekit import ImageSpec, Resources, Secret +from flytekit import current_context, LaunchPlan, CronSchedule +``` + +### Flyte 2 + +```python +import flyte +from flyte import TaskEnvironment, Resources, Secret +from flyte import Image, Trigger, Cron +``` + +## Before and after: pure Python execution + +### Flyte 1 + +```python +import flytekit + +image = flytekit.ImageSpec( + name="hello-world-image", + packages=["requests"], +) + +@flytekit.task(container_image=image) +def mean(data: list[float]) -> float: + return sum(list) / len(list) + +@flytekit.workflow +def main(data: list[float]) -> float: + output = mean(data) + + # ❌ performing trivial operations in a workflow is not allowed + # output = output / 100 + + # ❌ if/else is not allowed + # if output < 0: + # raise ValueError("Output cannot be negative") + + return output +``` + +### Flyte 2 + +```python +import flyte + +env = flyte.TaskEnvironment( + "hello_world", + image=flyte.Image.from_debian_base().with_pip_packages("requests"), +) + +@env.task +def mean(data: list[float]) -> float: + return sum(data) / len(data) + +@env.task +def main(data: list[float]) -> float: + output = mean(data) + + # ✅ performing trivial operations in a workflow is allowed + output = output / 100 + + # ✅ if/else is allowed + if output < 0: + raise ValueError("Output cannot be negative") + + return output +``` + +## Quick reference: minimal Flyte 2 module + +```python +import asyncio +import flyte + +# 1. Define an image +image = ( + flyte.Image.from_debian_base(python_version=(3, 11)) + .with_pip_packages("pandas", "numpy") +) + +# 2. Create a TaskEnvironment +env = flyte.TaskEnvironment( + name="my_env", + image=image, + resources=flyte.Resources(cpu="1", memory="2Gi"), +) + +# 3. Define tasks +@env.task +async def process(x: int) -> int: + return x * 2 + +# 4. Define the entrypoint task +@env.task +async def main(items: list[int]) -> list[int]: + results = await asyncio.gather(*[process(x) for x in items]) + return list(results) + +# 5. Run it +if __name__ == "__main__": + flyte.init_from_config() + run = flyte.run(main, items=[1, 2, 3, 4, 5]) + print(run.url) + run.wait() +``` + +```bash +# CLI +flyte run my_module.py main --items '[1,2,3,4,5]' # remote (default) +flyte run --local my_module.py main --items '[1,2,3,4,5]' +flyte deploy my_module.py my_env +``` + +## Recommended migration strategy + +Migrations rarely happen all at once. Work incrementally and lean on hybrid pipelines while the transition is in progress. + +1. **Assess the codebase.** Inventory every `@task`, `@workflow`, `@dynamic`, and `map_task`; the images (`ImageSpec`), resources, and secrets; the control-flow constructs (`conditional`, `on_failure`, `>>`); the data types (`FlyteFile`, `FlyteDirectory`, `StructuredDataset`); and any schedules (`LaunchPlan`, `CronSchedule`). +2. **Establish the `TaskEnvironment`(s).** Group tasks by their image/resource/cache needs and define a `flyte.TaskEnvironment` for each group. This is mechanical change #1 and unblocks everything else. +3. **Port leaf tasks first, then orchestration.** Convert atomic compute tasks (`@task` → `@env.task`), then rebuild the `@workflow`/`@dynamic` orchestration as plain-Python driver tasks that call them. +4. **Migrate control flow and I/O.** Replace `conditional()` with `if`/`elif`/`else`, `on_failure` with `try`/`except`, `map_task` with `flyte.map` / `asyncio.gather`, and the `FlyteFile`/`FlyteDirectory`/`StructuredDataset` types with their `flyte.io` equivalents. +5. **Update config, CLI, and schedules.** Swap `pyflyte` for `flyte`, migrate config files, and convert `LaunchPlan`/`CronSchedule` to `flyte.Trigger`/`flyte.Cron`. +6. **Run hybrid during the transition.** Keep unported v1 workflows callable via bridge tasks (see below) until every piece is on v2. + +### Sibling skills to route to + +Migrate by theme. Start with tasks and workflows, then jump to whatever the workload needs: + +- **`flyte-migrate-tasks-workflows`** — the structural shift: `@task`/`@workflow` → `@env.task`, sequential ordering, nested "subworkflows", and the `@task` → `TaskEnvironment` parameter mapping. +- **`flyte-migrate-config`** — moving image/resources/cache to the `TaskEnvironment`, GPUs, secrets, caching, scheduling with triggers, and the `pyflyte` → `flyte` command/config-file changes. +- **`flyte-migrate-control-flow`** — `conditional()` and `@dynamic` become plain Python `if`/loops, `on_failure` becomes `try`/`except`, and `map_task` → `flyte.map` / `asyncio.gather`. +- **`flyte-migrate-data-io`** — `FlyteFile`/`FlyteDirectory` → `flyte.io.File`/`Dir`, `StructuredDataset` → `flyte.io.DataFrame`, dataclasses, and ETL patterns. +- **`flyte-migrate-ml`** — small-model training, hyperparameter optimization, deep learning, batch inference, and end-to-end pipelines. + +## Hybrid v1 and v2 pipelines + +For a while you'll have Flyte 1 and Flyte 2 workloads running side by side, and you'll want them to call each other: a Flyte 1 workflow that kicks off a newly ported Flyte 2 task, or a Flyte 2 task that triggers a workflow that hasn't been migrated yet. + +You can bridge the two in both directions. The idea is the same each way: one task installs **both** SDKs, authenticates to the **other** control plane, fetches the entity it wants to run, and launches it. Keep the bridging task lightweight and focused on orchestration. + +### Running a Flyte 2 task from a Flyte 1 workflow + +The bridge is a single Flyte 1 task that runs the Flyte 2 client. Give it an image with **both** `flytekit` and `flyte` installed, provide a Flyte 2 API key as a secret, authenticate inside the task with `flyte.init_from_api_key()`, fetch the deployed task with `flyte.remote.Task.get(...)`, and run it. + +```python +import flytekit +from flytekit import task, workflow, ImageSpec, Secret, current_context + +# The bridge image needs BOTH the v1 (flytekit) and v2 (flyte) SDKs. +bridge_image = ImageSpec( + name="v1-to-v2-bridge", + packages=["flytekit", "flyte"], +) + +@task( + container_image=bridge_image, + secret_requests=[Secret(group="flyte", key="flyte_api_key")], +) +def launch_v2_from_v1(x: int) -> str: + import flyte + import flyte.remote + + # Authenticate to the Flyte 2 control plane with the API key. + # Option A: read the mounted secret and pass it explicitly. + api_key = current_context().secrets.get(group="flyte", key="flyte_api_key") + flyte.init_from_api_key(api_key=api_key) + + # Option B: if FLYTE_API_KEY is set as an env var, no argument is needed: + # flyte.init_from_api_key() + + # Fetch the deployed Flyte 2 task and run it. + remote_v2_task = flyte.remote.Task.get( + "my_v2_env.process", + auto_version="latest", + ) + run = flyte.run(remote_v2_task, x=x) + run.wait() # optional: block until the v2 run finishes + return run.url + +@workflow +def main(x: int) -> str: + return launch_v2_from_v1(x=x) +``` + +The referenced Flyte 2 task (`my_v2_env.process` above) must be **deployed** before the bridge runs. Use `flyte.init_from_api_key()` here — do **not** use `flyte.init_from_config()`, which reads a `config.yaml` that has no API-key field. + +### Running a Flyte 1 workflow from a Flyte 2 task + +The reverse works the same way: a Flyte 2 task installs the Flyte 1 client and uses `FlyteRemote` to launch a Flyte 1 workflow. + +```python +import flyte + +env = flyte.TaskEnvironment( + name="v2_to_v1_bridge", + # The image needs the Flyte 1 client installed. + image=flyte.Image.from_debian_base().with_pip_packages("flytekit"), + # Supply credentials for the Flyte 1 control plane (config or API key). + secrets=[flyte.Secret(key="v1_client_secret", as_env_var="V1_CLIENT_SECRET")], +) + +@env.task +async def launch_v1_from_v2(x: int) -> str: + from flytekit.remote import FlyteRemote + from flytekit.configuration import Config + + # Point the client at your Flyte 1 cluster. + remote = FlyteRemote( + config=Config.for_endpoint(endpoint="my-v1-cluster.example.com"), + default_project="flytesnacks", + default_domain="development", + ) + + # Fetch the deployed Flyte 1 workflow and execute it. + wf = remote.fetch_workflow(name="my_v1_module.main", version="v1.2.3") + execution = remote.execute(wf, inputs={"x": x}, wait=True) + return execution.id.name +``` + +### Hybrid considerations + +- **Both SDKs in one image.** The bridging task installs `flytekit` and `flyte` together. Pin versions and watch for dependency conflicts; keep the bridge image minimal. +- **Deploy the callee first.** For v1→v2, the Flyte 2 task must be deployed (`flyte deploy`) before `flyte.remote.Task.get()` can resolve it. For v2→v1, the Flyte 1 workflow must be registered on its cluster. +- **Wait vs. fire-and-forget.** Both `run.wait()` (v2) and `execute(..., wait=True)` (v1) block until the launched run finishes. Omit them to launch and return immediately. +- **Credentials cross a boundary.** The bridge authenticates to a *different* control plane than the one it runs on. Store the API key or client credentials as a secret — never hard-code them. +- **Keep the bridge lightweight.** Like any orchestrating task, it should mostly launch and assemble results rather than do heavy compute. + +## Gotchas + +Flyte 2 lets each Python task act as its own engine, launching sub-tasks and assembling their outputs. That flexibility warrants some caveats. + +### Common gotchas + +- **`flyte.map` returns a generator.** Wrap it in `list()` to materialize results, unlike `map_task` which returned a list directly. +- **`memory`, not `mem`.** The `Resources` parameter was renamed, and there are no separate `requests`/`limits` — a single value serves as both. +- **GPUs use a `"T4:1"` string.** Type and count are combined; the separate `accelerator=` argument is gone. +- **Image, resources, and cache live on the `TaskEnvironment`.** Set them once at the env level instead of repeating them on every task decorator. +- **`current_context()` is gone.** Read secrets from environment variables and use `flyte.ctx()` for runtime context. +- **The `>>` ordering operator is gone.** Sequential (sync) calls and sequential `await`s are naturally ordered. +- **Retries no longer have a platform cap.** In Flyte 1 the control plane capped attempts at 3; in Flyte 2 total attempts equal `retries + 1`. Audit any large `retries` values before deploying. +- **You can only `await` async tasks.** Call a sync task from an async context with `.aio()`. +- **Pick an entrypoint task name.** There's no `@workflow`, so the top-level task is just a task (commonly `main`); run it with `flyte run module.py main`. +- **Type annotations are more lenient.** Flyte 2 will pickle untyped I/O rather than rejecting it at registration. +- **Keep orchestration lightweight.** A task that calls other tasks acts as a driver pod. Avoid heavy CPU work in it. + +## Anti-Patterns + +1. **Don't introduce non-determinism into orchestration.** When a task launches another task, a new Action ID is determined as a hash of the inputs and task definition — consistent hashing is what makes recovery and replay work. Branching on `datetime.now()` or other non-deterministic values breaks that guarantee: on retry, a *different* downstream task may get kicked off. If non-determinism is unavoidable, decorate sub-task functions with `@trace` for fine-grained checkpointing and observability. +2. **Don't do heavy compute in a driver task.** When a task runs other tasks and assembles their outputs, it becomes a driver pod (work that Flyte Propeller did in v1). A CPU-bound function between two `await`s makes the driver pod hang and slows downstream kickoff. Keep parent tasks focused on orchestration: + +```python +@env.task +async def t_main(): + await t1() + local_cpu_intensive_function() # ❌ blocks the driver pod between t1 and t2 + await t2() +``` + +3. **Don't rely on global state across tasks.** Each task runs in its own isolated container; globals are not carried across task containers. Any state that must persist has to be reconstructable through repeated deterministic execution. +4. **Don't materialize huge in-memory I/O between tasks.** Outputs are materialized in the parent pod's memory, so passing a 1 GB `list[float]` requires the pod to hold all of it, risking OOM. Use `flyte.io.File`, `flyte.io.Dir`, and `flyte.io.DataFrame` — they're materialized only as pointers to offloaded data, so their memory footprint stays low. +5. **Don't skip type hints at the "workflow" level.** The top-level task now runs at runtime, so the system can't guarantee type safety across the DAG the way the v1 DSL did. Use Python type hints and a type checker like `mypy` at all levels, including the top-most task.
scenarioskillharnesstier passtreatmentcontrollift