# Table of Contents - [User guide | Union.ai Docs](#user-guide-union-ai-docs) - [Unknown](#unknown) - [Unknown](#unknown) - [Overview | Union.ai Docs](#overview-union-ai-docs) - [Build tasks | Union.ai Docs](#build-tasks-union-ai-docs) - [Build apps | Union.ai Docs](#build-apps-union-ai-docs) - [Configure apps | Union.ai Docs](#configure-apps-union-ai-docs) - [Tutorials | Union.ai Docs](#tutorials-union-ai-docs) - [Home | Union.ai Docs](#home-union-ai-docs) - [Run and deploy tasks | Union.ai Docs](#run-and-deploy-tasks-union-ai-docs) - [Quickstart | Union.ai Docs](#quickstart-union-ai-docs) - [Reference | Union.ai Docs](#reference-union-ai-docs) - [Build an MCP | Union.ai Docs](#build-an-mcp-union-ai-docs) - [Scale your runs | Union.ai Docs](#scale-your-runs-union-ai-docs) - [Run modes | Union.ai Docs](#run-modes-union-ai-docs) - [Agent framework integrations | Union.ai Docs](#agent-framework-integrations-union-ai-docs) - [Biotech & Healthcare | Union.ai Docs](#biotech-healthcare-union-ai-docs) - [Computer Vision | Union.ai Docs](#computer-vision-union-ai-docs) - [Context Engineering | Union.ai Docs](#context-engineering-union-ai-docs) - [Sandboxing | Union.ai Docs](#sandboxing-union-ai-docs) - [Model Training | Union.ai Docs](#model-training-union-ai-docs) - [Frontier AI | Union.ai Docs](#frontier-ai-union-ai-docs) - [Contributing code | Union.ai Docs](#contributing-code-union-ai-docs) - [Geospatial | Union.ai Docs](#geospatial-union-ai-docs) - [Agents | Union.ai Docs](#agents-union-ai-docs) - [Serve and deploy apps | Union.ai Docs](#serve-and-deploy-apps-union-ai-docs) - [Core concepts | Union.ai Docs](#core-concepts-union-ai-docs) - [Native app integrations | Union.ai Docs](#native-app-integrations-union-ai-docs) - [Financial Services & Fintech | Union.ai Docs](#financial-services-fintech-union-ai-docs) - [Project patterns | Union.ai Docs](#project-patterns-union-ai-docs) - [Anthropic | Union.ai Docs](#anthropic-union-ai-docs) - [Gemini | Union.ai Docs](#gemini-union-ai-docs) - [Weights & Biases | Union.ai Docs](#weights-biases-union-ai-docs) - [BigQuery | Union.ai Docs](#bigquery-union-ai-docs) - [Dask | Union.ai Docs](#dask-union-ai-docs) - [OpenAI | Union.ai Docs](#openai-union-ai-docs) - [Joining the community | Union.ai Docs](#joining-the-community-union-ai-docs) - [Migration | Union.ai Docs](#migration-union-ai-docs) - [Advanced project | Union.ai Docs](#advanced-project-union-ai-docs) - [Contributing docs and examples | Union.ai Docs](#contributing-docs-and-examples-union-ai-docs) - [Configure tasks | Union.ai Docs](#configure-tasks-union-ai-docs) - [Databricks | Union.ai Docs](#databricks-union-ai-docs) - [Spark | Union.ai Docs](#spark-union-ai-docs) - [PyTorch | Union.ai Docs](#pytorch-union-ai-docs) - [MLflow | Union.ai Docs](#mlflow-union-ai-docs) - [Ray | Union.ai Docs](#ray-union-ai-docs) - [Community | Union.ai Docs](#community-union-ai-docs) - [Flyte SDK | Union.ai Docs](#flyte-sdk-union-ai-docs) - [Integrations | Union.ai Docs](#integrations-union-ai-docs) - [Build an agent | Union.ai Docs](#build-an-agent-union-ai-docs) - [Migration from Flyte 1 | Union.ai Docs](#migration-from-flyte-1-union-ai-docs) - [Snowflake | Union.ai Docs](#snowflake-union-ai-docs) - [Pandera | Union.ai Docs](#pandera-union-ai-docs) - [Papermill | Union.ai Docs](#papermill-union-ai-docs) - [Hydra | Union.ai Docs](#hydra-union-ai-docs) - [Code generation | Union.ai Docs](#code-generation-union-ai-docs) - [OmegaConf | Union.ai Docs](#omegaconf-union-ai-docs) - [LLM-optimized documentation | Union.ai Docs](#llm-optimized-documentation-union-ai-docs) - [Integrations | Union.ai Docs](#integrations-union-ai-docs) - [Flyte CLI | Union.ai Docs](#flyte-cli-union-ai-docs) - [Home | Union.ai Docs](#home-union-ai-docs) - [Home | Union.ai Docs](#home-union-ai-docs) - [User guide | Union.ai Docs](#user-guide-union-ai-docs) - [Unknown](#unknown) - [Unknown](#unknown) - [Unknown](#unknown) - [Unknown](#unknown) - [Unknown](#unknown) - [Unknown](#unknown) - [Unknown](#unknown) - [Unknown](#unknown) - [Unknown](#unknown) - [Unknown](#unknown) - [Unknown](#unknown) - [Unknown](#unknown) - [Tutorials | Union.ai Docs](#tutorials-union-ai-docs) - [Architecture | Union.ai Docs](#architecture-union-ai-docs) - [Reference | Union.ai Docs](#reference-union-ai-docs) - [Platform deployment | Union.ai Docs](#platform-deployment-union-ai-docs) - [Bioinformatics | Union.ai Docs](#bioinformatics-union-ai-docs) - [Programming | Union.ai Docs](#programming-union-ai-docs) - [Flytelab | Union.ai Docs](#flytelab-union-ai-docs) - [Feature engineering | Union.ai Docs](#feature-engineering-union-ai-docs) - [Integrations | Union.ai Docs](#integrations-union-ai-docs) - [Model training | Union.ai Docs](#model-training-union-ai-docs) - [Data input/output | Union.ai Docs](#data-input-output-union-ai-docs) - [Flyte deployment | Union.ai Docs](#flyte-deployment-union-ai-docs) - [Connector setup | Union.ai Docs](#connector-setup-union-ai-docs) - [Getting support | Union.ai Docs](#getting-support-union-ai-docs) - [Community | Union.ai Docs](#community-union-ai-docs) - [Platform configuration | Union.ai Docs](#platform-configuration-union-ai-docs) - [Configuration reference | Union.ai Docs](#configuration-reference-union-ai-docs) - [Plugins setup | Union.ai Docs](#plugins-setup-union-ai-docs) - [Native backend plugins | Union.ai Docs](#native-backend-plugins-union-ai-docs) - [Executions | Union.ai Docs](#executions-union-ai-docs) - [Registration | Union.ai Docs](#registration-union-ai-docs) - [Flytekit plugins | Union.ai Docs](#flytekit-plugins-union-ai-docs) - [Deprecated integrations | Union.ai Docs](#deprecated-integrations-union-ai-docs) - [Flyte operators | Union.ai Docs](#flyte-operators-union-ai-docs) - [External service backend plugins | Union.ai Docs](#external-service-backend-plugins-union-ai-docs) - [Connectors | Union.ai Docs](#connectors-union-ai-docs) - [Workflow timeline | Union.ai Docs](#workflow-timeline-union-ai-docs) - [Control Plane | Union.ai Docs](#control-plane-union-ai-docs) - [Workflow state transitions | Union.ai Docs](#workflow-state-transitions-union-ai-docs) - [Development cycle | Union.ai Docs](#development-cycle-union-ai-docs) - [Versions | Union.ai Docs](#versions-union-ai-docs) - [Data catalog | Union.ai Docs](#data-catalog-union-ai-docs) - [Data handling | Union.ai Docs](#data-handling-union-ai-docs) - [Joining the community | Union.ai Docs](#joining-the-community-union-ai-docs) - [Component Architecture | Union.ai Docs](#component-architecture-union-ai-docs) - [Contributing docs and examples | Union.ai Docs](#contributing-docs-and-examples-union-ai-docs) - [Flytekit SDK | Union.ai Docs](#flytekit-sdk-union-ai-docs) - [Workflow lifecycle | Union.ai Docs](#workflow-lifecycle-union-ai-docs) - [Extending Flyte | Union.ai Docs](#extending-flyte-union-ai-docs) - [Roadmap | Union.ai Docs](#roadmap-union-ai-docs) - [FlyteKit Plugins | Union.ai Docs](#flytekit-plugins-union-ai-docs) - [Troubeshooting | Union.ai Docs](#troubeshooting-union-ai-docs) - [Core concepts | Union.ai Docs](#core-concepts-union-ai-docs) - [Pyflyte CLI | Union.ai Docs](#pyflyte-cli-union-ai-docs) - [Contributing code | Union.ai Docs](#contributing-code-union-ai-docs) - [Flytectl CLI | Union.ai Docs](#flytectl-cli-union-ai-docs) - [LLM-optimized documentation | Union.ai Docs](#llm-optimized-documentation-union-ai-docs) - [Getting started | Union.ai Docs](#getting-started-union-ai-docs) - [Platform deployment | Union.ai Docs](#platform-deployment-union-ai-docs) - [Reference | Union.ai Docs](#reference-union-ai-docs) - [Introduction | Union.ai Docs](#introduction-union-ai-docs) - [User guide | Union.ai Docs](#user-guide-union-ai-docs) - [Tutorials | Union.ai Docs](#tutorials-union-ai-docs) - [Run modes | Union.ai Docs](#run-modes-union-ai-docs) - [Core concepts | Union.ai Docs](#core-concepts-union-ai-docs) - [Build tasks | Union.ai Docs](#build-tasks-union-ai-docs) - [Build an MCP | Union.ai Docs](#build-an-mcp-union-ai-docs) - [Biotech & Healthcare | Union.ai Docs](#biotech-healthcare-union-ai-docs) - [Native app integrations | Union.ai Docs](#native-app-integrations-union-ai-docs) - [Serve and deploy apps | Union.ai Docs](#serve-and-deploy-apps-union-ai-docs) - [Project patterns | Union.ai Docs](#project-patterns-union-ai-docs) - [Scale your runs | Union.ai Docs](#scale-your-runs-union-ai-docs) - [Geospatial | Union.ai Docs](#geospatial-union-ai-docs) - [Computer Vision | Union.ai Docs](#computer-vision-union-ai-docs) - [Frontier AI | Union.ai Docs](#frontier-ai-union-ai-docs) - [Financial Services & Fintech | Union.ai Docs](#financial-services-fintech-union-ai-docs) - [Agents | Union.ai Docs](#agents-union-ai-docs) - [Build apps | Union.ai Docs](#build-apps-union-ai-docs) - [Configure tasks | Union.ai Docs](#configure-tasks-union-ai-docs) - [Agent framework integrations | Union.ai Docs](#agent-framework-integrations-union-ai-docs) - [Run and deploy tasks | Union.ai Docs](#run-and-deploy-tasks-union-ai-docs) - [Configure apps | Union.ai Docs](#configure-apps-union-ai-docs) - [Context Engineering | Union.ai Docs](#context-engineering-union-ai-docs) - [Sandboxing | Union.ai Docs](#sandboxing-union-ai-docs) - [Model Training | Union.ai Docs](#model-training-union-ai-docs) - [User management | Union.ai Docs](#user-management-union-ai-docs) - [Advanced project | Union.ai Docs](#advanced-project-union-ai-docs) - [Migration | Union.ai Docs](#migration-union-ai-docs) - [Data Processing | Union.ai Docs](#data-processing-union-ai-docs) - [Anthropic | Union.ai Docs](#anthropic-union-ai-docs) - [BigQuery | Union.ai Docs](#bigquery-union-ai-docs) - [Authenticating | Union.ai Docs](#authenticating-union-ai-docs) - [Dask | Union.ai Docs](#dask-union-ai-docs) - [Gemini | Union.ai Docs](#gemini-union-ai-docs) - [Build an agent | Union.ai Docs](#build-an-agent-union-ai-docs) - [Databricks | Union.ai Docs](#databricks-union-ai-docs) - [Quickstart | Union.ai Docs](#quickstart-union-ai-docs) - [OpenAI | Union.ai Docs](#openai-union-ai-docs) - [Spark | Union.ai Docs](#spark-union-ai-docs) - [PyTorch | Union.ai Docs](#pytorch-union-ai-docs) --- # User guide | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Flyte OSS ========= Flyte is a free and open source platform that provides a full suite of powerful features for orchestrating AI workflows. Flyte empowers AI development teams to rapidly ship high-quality code to production by offering optimized performance, unparalleled resource efficiency, and a delightful workflow authoring experience. You deploy and manage Flyte yourself, on your own cloud infrastructure. These are the Flyte **2.0** docs. To switch to [version 1.0](https://www.union.ai/docs/v1/flyte) or to the commercial product, [**Union.ai**](https://www.union.ai/docs/v2/union) , use the selectors above. [Basics](https://www.union.ai/docs/v2/flyte/user-guide/#basics) ----------------------------------------------------------------- Learn the basics of Flyte, covering all the core concepts around tasks and apps. Flyte 2 Build AI workflows in pure Python with built-in durability, reproducibility, and recovery. Quickstart Install the SDK and run your first workflow locally in a few minutes. Core concepts The building blocks of every Flyte program: TaskEnvironments, tasks, runs, actions, and apps. Run modes Run the same task code locally, on a devbox, or on a remote cluster. [Tasks](https://www.union.ai/docs/v2/flyte/user-guide/#tasks) --------------------------------------------------------------- Build durable, scalable, and reproducible batch workloads. Configure tasks Define [`TaskEnvironment`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/taskenvironment/) s for container images, resources, secrets, caching, retries, and more; use triggers for schedules. Build tasks Compose tasks with fanout, parallelism, error handling, traces, files, and DataFrames. Run and deploy tasks Use `flyte run` for iteration or `flyte deploy` to register a stable task version. [Apps](https://www.union.ai/docs/v2/flyte/user-guide/#apps) ------------------------------------------------------------- Create long-running services to host dashboards, APIs, and model endpoints. Configure apps Define [`AppEnvironment`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/appenvironment/) s with ports, autoscaling, custom domains, and authentication. Build apps Build dashboards, REST APIs, and model endpoints with FastAPI, Streamlit, vLLM, and more. Native app integrations Use pre-built environments for popular frameworks like Streamlit, FastAPI, vLLM, and SGLang. Serve and deploy apps Use `flyte serve` for fast iteration or `flyte deploy` for production deployments. [Agents](https://www.union.ai/docs/v2/flyte/user-guide/#agents) ----------------------------------------------------------------- Build durable, self-healing agents using tasks and apps as building blocks. Build agents Implement ReAct, Plan-and-Execute, and other agent patterns with full observability. Agent framework integrations Integrate with third-party agent frameworks like LangGraph, PydanticAI, and OpenAI Agents SDK. Sandboxing Safely execute LLM-generated code with workflow sandboxes or ephemeral containers. Build an MCP Serve Model Context Protocol servers for AI assistants to interact with, hosted on Flyte. [Advanced Guides](https://www.union.ai/docs/v2/flyte/user-guide/#advanced-guides) ----------------------------------------------------------------------------------- Organize your codebase, optimize performance for production, and migrate from other workflow orchestrators. Project patterns Patterns for BYO images, monorepos with uv, CI/CD, and multi-team resource management. Run scaling Tune task overhead, batching, reusable containers, and fanout to scale your workflows. Advanced project An advanced guide for building an LLM reporting agent on Flyte. Migration Port a Flyte 1 codebase to Flyte 2, or map Airflow concepts to their Flyte 2 equivalents. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 1. [Flyte OSS](https://www.union.ai/docs/v2/flyte/user-guide/#flyte-oss) 1. [Basics](https://www.union.ai/docs/v2/flyte/user-guide/#basics) 2. [Tasks](https://www.union.ai/docs/v2/flyte/user-guide/#tasks) 3. [Apps](https://www.union.ai/docs/v2/flyte/user-guide/#apps) 4. [Agents](https://www.union.ai/docs/v2/flyte/user-guide/#agents) 5. [Advanced Guides](https://www.union.ai/docs/v2/flyte/user-guide/#advanced-guides) 404 Page not found Showing closest match [![](https://www.union.ai/docs/v2/flyte/images/icon-logo-flyte.svg)](https://www.union.ai/docs/v2/flyte/) [Flyte OSS Docs](https://www.union.ai/docs/v2/flyte/) --- # Unknown \# Flyte Open Source Documentation > Full documentation (single file): https://www.union.ai/docs/v2/flyte/llms-full.txt > Site: https://www.union.ai/docs/v2/flyte Each entry below is \`- \[Page title\](URL)\` followed by the H2/H3 headings found on that page. Pages link to individual \`page.md\` files. Sections marked with a "Section bundle" link have a \`section.md\` that concatenates all pages in the section into a single file — use it to load an entire section into context at once. ## User guide - \[Overview\](https://www.union.ai/docs/v2/flyte/user-guide/overview/page.md) - Pure Python, no DSL - Durability - Reproducibility - Recoverability - Built for scale - What this means in practice - \[Quickstart\](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/page.md) - What you'll need - Install the SDK - Configure - Write your first workflow - Run it - See the results - Next steps - \[Core concepts\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/page.md) - How Flyte works > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/section.md - \[Core concepts > TaskEnvironment\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/task-environment/page.md) - A minimal example - What TaskEnvironment controls - Configuring resources - Configuring container images - Multiple tasks, one environment - Multiple environments - Next steps - \[Core concepts > Tasks\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/tasks/page.md) - Defining a task - Type hints are required - Tasks calling tasks - The top-level task - Running tasks locally - Running tasks remotely - Next steps - \[Core concepts > Runs and actions\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/runs-and-actions/page.md) - What is a run? - What is an action? - Runs vs actions in practice - Viewing runs in the UI - Understanding the execution graph - Checking run status - Next steps - \[Core concepts > Where your data lives\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/where-data-lives/page.md) - The two stores - What goes in the database - What goes in the bucket - What "metadata" means - 1. "Metadata" as in the control-plane database (Flyte's usage) - 2. "Metadata bucket" (a deployment/ops term you may see) - Per-run customization: \`raw\_data\_path\` - What happens if the bucket is purged - The short version - \[Core concepts > Apps\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/introducing-apps/page.md) - Tasks vs apps - AppEnvironment - A hello world app - Understanding the code - Serving the app - When to use apps vs tasks - Common patterns - Next steps - \[Core concepts > Projects and domains\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/projects-and-domains/page.md) - How projects and domains are used - Managing projects via CLI - Create a project - List projects - Update a project - Archive a project - Unarchive a project - Listing projects programmatically - Managing projects via the UI - Domains - Targeting a domain - \[Core concepts > Key capabilities\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/key-capabilities/page.md) - Environment and resources - Deployment - Data handling - Parallelism and composition - Security and automation - Durability and reliability - Apps and serving - Notebooks - Next steps - \[Core concepts > Basic project: RAG\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/basic-project/page.md) - Concepts covered - Part 1: The embedding pipeline - Setting up the environment - Fetching data - Creating embeddings - Orchestrating the pipeline - Running the pipeline - Part 2: The serving application - App environment configuration - The Streamlit application - Deploying the app - Key takeaways - \[Run modes\](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/page.md) - \[Local\](running-locally/page.md) - \[Devbox\](running-devbox/page.md) - \[Run modes > Run locally\](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-locally/page.md) - Getting started - Running tasks locally - Terminal UI - Exploring past runs - What works locally - Local to devbox/remote - Next steps - \[Run modes > Run on the devbox\](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox/page.md) - What you'll need - Install the SDK - Start the devbox - CPU - GPU - Configure - Run a workflow on the devbox - View results in the UI - Stop the devbox - Inline configuration - Programmatic - CLI - Delete the devbox - Using a CUDA-enabled GPU host - Next steps - \[Configure tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/page.md) - Task configuration levels - Example - Task configuration parameters > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/section.md - \[Configure tasks > Container images\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/container-images/page.md) - Specifying your own image directly - Specifying your own image with the \`flyte.Image\` object - Example: Defining a custom image with \`Image.from\_debian\_base\` - Example: Defining an image based on uv script metadata - Image building - Configuring the \`builder\` - Local image building - Remote \`ImageBuilder\` - Install private PyPI packages - \[Configure tasks > Resources\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/resources/page.md) - Resources data class - Examples - Usage in TaskEnvironment - Usage in a task-specific override - \[Configure tasks > Secrets\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets/page.md) - Creating a literal string secret - Creating a file secret - Scoping secrets - Listing secrets - Deleting secrets - Using a literal string secret - Using a file secret - \[Configure tasks > Caching\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/caching/page.md) - Overview - Basic caching usage - \`"auto"\` - Automatic versioning - \`"override"\` - \`"disable"\` - No caching - Advanced caching configuration - Ignoring specific inputs - Cache serialization - Salt for cache key variation - Cache policies - Function body policy (default) - Custom cache policies - Caching configuration at different levels - \`TaskEnvironment\` Level - \`@env.task\` decorator level - \`task.override\` level - Runtime cache control - Project and domain cache isolation - Local development caching - \[Configure tasks > Reusable containers\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/reusable-containers/page.md) - \[Configure tasks > Pod templates\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/pod-templates/page.md) - How it works - Requirements - Basic usage - PodTemplate parameters - Volume mounts - GCS/S3 volume mounts - Sidecar containers - Image pull secrets - Cluster-specific configuration - Important notes - Best practices - Learn more - \[Configure tasks > Multiple environments\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/multiple-environments/page.md) - Constraints on multiple environments - Task \`depends\_on\` constraints - Dependency inclusion constraints - Example - \[Configure tasks > Retries and timeouts\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/retries-and-timeouts/page.md) - The action lifecycle - Retries - Retry count - Retries with exponential backoff - Skip retries for failures that can't be fixed - System retries - Timeouts - \`max\_runtime\` — bound a single attempt's execution - \`max\_queued\_time\` — fail fast when capacity isn't available - \`deadline\` — bound the total wall-clock - Combining the bounds - Combining retries and timeouts - \[Configure tasks > Triggers\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/triggers/page.md) - Triggers are set in the task decorator - \`flyte.Trigger\` - The \`automation\` parameter with \`flyte.FixedRate\` - Examples - The \`automation\` parameter with \`flyte.Cron\` - Examples - The \`inputs\` parameter - Basic Usage - Using \`flyte.TriggerTime\` - Required vs optional parameters - Complex input types - Predefined schedule triggers - Available Predefined Triggers - Trigger time in predefined triggers - Multiple triggers per task - Notifications - Execution phases - Template variables - Slack notifications - Email notifications - Microsoft Teams notifications - Custom webhook notifications - Deploying a task with triggers - Activating and deactivating triggers - Trigger run timing - Cron-based triggers - Fixed-rate triggers without \`start\_time\` - Fixed-rate triggers with \`start\_time\` - Deleting triggers - Schedule time zones - Setting time zone for a Cron schedule - \`flyte.TriggerTime\` is always in UTC - Daylight Savings Time behavior - \[Configure tasks > Interruptible tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/interruptible-tasks-and-queues/page.md) - Setting at different levels - Behavior on preemption - Spot to on-demand fallback - \[Configure tasks > Task plugins\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/task-plugins/page.md) - Default Execution: Containers - Compute Plugins - Available Compute Plugins - How Compute Plugins Work - Using Compute Plugins - Using Plugins on Union - Backend Integrations - Next Steps - \[Configure tasks > Additional task settings\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/additional-task-settings/page.md) - Naming and metadata - \`name\` - \`short\_name\` - \`description\` - \`docs\` - \`report\` - Source-code link (automatic) - \`links\` - Default inputs - Environment variables - Inline I/O threshold - \[Build tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/page.md) - What you'll learn - When to use these patterns > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/task-programming/section.md - \[Build tasks > Files and directories\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/files-and-directories/page.md) - Example usage - JSONL files - Setup - JsonlFile - Compression - JsonlDir - Error handling - Batch iteration - \[Build tasks > Data classes and structures\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/dataclasses-and-structures/page.md) - Example: Combining Dataclasses and Pydantic Models - \[Build tasks > DataFrames\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/dataframes/page.md) - Setting up the environment and sample data - Create a raw DataFrame - Create a flyte.io.DataFrame - Automatically convert between types - Downloading DataFrames - Run the example - Polars DataFrames - Setup - Eager DataFrames - Lazy DataFrames - Run the example - \[Build tasks > Custom types\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/handling-custom-types/page.md) - Types of extensions - Creating a type transformer - Step 1: Define your custom type - Step 2: Create the type transformer - Step 3: Register the transformer - Distributing type plugins - Configure pyproject.toml - Automatic loading - Controlling plugin loading - Using custom types in tasks - DataFrame extensions - Best practices - \[Build tasks > Custom context\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/custom-context/page.md) - Overview - When to use it and when not to - Setting custom context - Run-level context - Overriding inside a task (local override that affects nested tasks) - Adding new keys for nested tasks - Accessing custom context - \[Build tasks > Abort and cancel actions\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/abort-tasks/page.md) - Action lifetime - Canceling actions programmatically - External abort - Aborting via the CLI - Handling external aborts - \[Build tasks > Run a bioinformatics tool\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/container-tasks/page.md) - What are Container Tasks? - How Data Flows In and Out - Basic Usage - Template Syntax for Inputs - Using Container Tasks in Workflows - Advanced: Passing Files and Directories - Use Case: Agentic Sandbox Execution - Use Case: Legacy and Specialized Containers - Use Case: Multi-Language Workflows - Configuration Options - ContainerTask Parameters - Supported Input/Output Types - Best Practices - Local Execution - When to Use Container Tasks - \[Build tasks > Links\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/links/page.md) - Creating a link - Using execution metadata - Dynamic links with override - \[Build tasks > Reports\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/reports/page.md) - A simple example - A more complex example - Streaming example - \[Build tasks > Notebooks\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/notebooks/page.md) - Iterating on and running a workflow - Accessing runs and downloading logs - \[Build tasks > Remote tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/remote-tasks/page.md) - Prerequisites - Basic usage - Understanding lazy loading - When tasks are fetched - Benefits of lazy loading - Error handling - Eager fetching with \`fetch()\` - Module-level vs dynamic loading - Complete example - Team A: Spark environment - Team B: ML environment - Team C: Orchestration - Invoke remote tasks in a script. - Why use remote tasks? - When to use remote tasks - How remote tasks work - Security model - Type system - Versioning options - Customizing remote tasks - Available overrides - Override examples - Chain overrides - Best practices - 1. Use meaningful task names - 2. Document task interfaces - 3. Prefer module-level loading - 4. Handle versioning thoughtfully - 5. Deploy remote tasks first - Limitations - Next steps - \[Build tasks > Error handling\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/error-handling/page.md) - \[Build tasks > Traces\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/traces/page.md) - What are traced functions for? - What Gets Traced - Errors are not recorded - Supported Function Types - Task Orchestration Pattern - Relationship to Caching and Checkpointing - How They Work Together - Execution Flow - Error Handling and Observability - Examples in Practice - LLM Pipeline with Traces - \[Build tasks > Grouping actions\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/grouping-actions/page.md) - What are groups? - The problem groups solve - How groups work - Common grouping patterns - Sequential operations - Parallel processing with groups - Multi-phase workflows - Conditional grouping - Key insights - \[Build tasks > Fanout\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/fanout/page.md) - Understanding fanout - Example - Parallel execution - Running the example - How Flyte handles concurrency and parallelism - \[Build tasks > Controlling parallel execution\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/controlling-parallelism/page.md) - The problem: unbounded parallelism - Using asyncio.Semaphore - Using flyte.map with concurrency - Running the example - When to use each approach - \[Build tasks > Human-in-the-loop\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/human-in-the-loop/page.md) - Setup - Automated task - Requesting human input - Wiring it together - Event options - Submitting input programmatically - \[Build tasks > Test business logic directly\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/unit-testing/page.md) - Understanding Task Invocation - Direct Function Invocation - Using \`flyte.run()\` - Testing Business Logic - Testing Async Tasks - Testing Nested Tasks - Testing Type Transformations and Serialization - Testing Type Restrictions - Testing Nested Tasks with Serialization - Testing Traced Functions - Best Practices - Quick Reference - Example Test Suite - Future Improvements - \[Build tasks > Regular async function (not a task)\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/other-features/page.md) - Task Forwarding - Passing Tasks and Functions as Arguments - Custom Action Names - Set at Task Definition - Override at Call Time - Invoking Async Functions from Sync Tasks - Async and Sync Task Interoperability - Calling Sync Tasks from Async Tasks - Using with \`flyte.map.aio()\` - Using AnyIO in Async Tasks - \[Run and deploy tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/page.md) - Ephemeral deployment and immediate execution - Programmatic - CLI - Persistent deployment - Programmatic - CLI - Running already deployed tasks - Programmatic - CLI - Configuring runs with \`flyte.with\_runcontext()\` > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/section.md - \[Run and deploy tasks > How task run works\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/how-task-run-works/page.md) - Ephemeral deployment + run: The development shortcut - Programmatic - CLI - Running deployed tasks - Programmatic - CLI - Local execution - Programmatic - CLI - Running tasks through the Union UI - Accessing task execution in the Union UI - Execution flow and architecture - Fast registration architecture - Ephemeral preparation logic - Execution modes comparison - \[Run and deploy tasks > Interact with runs and actions\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/interacting-with-runs/page.md) - Understanding runs and actions - Key concepts - Working with runs - Retrieving a run - Programmatic - CLI - Watching run progress - Getting detailed run information - Working with actions - Retrieving an action - Programmatic - CLI - Nested actions - Getting detailed action information - Retrieving inputs and outputs - Programmatic - CLI - Handling failures - Understanding data storage - Accessing large data from cloud storage - S3 storage access - GCS storage access - Azure Blob Storage access - Complete example - API reference - Key classes - CLI commands - Storage configuration - \[Run and deploy tasks > Work with local data\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/work-with-local-data/page.md) - Local execution - Uploading local data to remote runs - Uploading DataFrames - Uploading files - Uploading directories - Passing outputs between runs - Performance considerations - Summary - \[Run and deploy tasks > Run command options\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/run-command-options/page.md) - \`--project\`, \`--domain\` - \`--run-project\`, \`--run-domain\` - \`--local\` - When to use local execution - \`--copy-style\` - Copy style options - \`--root-dir\` - \`--raw-data-path\` - Use cases - \`--service-account\` - Use cases - \`--name\` - Benefits of custom names - \`--follow\` - Behavior - \`--image\` - Image mapping formats - \`--no-sync-local-sys-paths\` - Task argument passing - SDK options - \[Run and deploy tasks > How task deployment works\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/how-task-deployment-works/page.md) - 1. Module loading and task environment discovery - Single file (default) - \`--all\` option - \`--recursive\` option - 2. Task analysis and serialization - 3. Task environment dependency resolution - 4. Code bundle creation and upload - \`--copy\_style loaded\_modules\` (default) - \`--copy\_style all\` - \`--copy\_style none\` - \`--root-dir\` option - 5. Image building - Local image building - Remote image building - 6. Source-code link discovery - How the link is built - Conditions - Understanding option relationships - \[Run and deploy tasks > Deploy command options\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/deploy-command-options/page.md) - \`--project\`, \`--domain\` - \`--version\` - When versions are used - \`--dry-run\` - \`--all\` and \`--recursive\` - \`--copy-style\` - \`--copy-style loaded\_modules\` (default) - \`--copy-style all\` - \`--copy-style none\` - \`--root-dir\` - Default behavior (without \`--root-dir\`) - Common use cases - How it works - Example with complex project structure - \`--image\` - Named image mappings - Default image mapping - How it works - \`--ignore-load-errors\` - \`--no-sync-local-sys-paths\` - Default behavior (path synchronization enabled) - When to disable path synchronization - Use cases for disabling - How it works - SDK deployment options - \[Run and deploy tasks > Code packaging for remote execution\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/packaging/page.md) - Quick comparison - Code bundling - How it works - Automatic code bundling - Manual code bundling - Including additional files with \`include\` - Controlling the root directory - Code bundling examples - When to use code bundling - Container-based deployment - How it works - Configuration - Image source copying methods - Complete container-based example - Using externally built images - Container-based best practices - When to use container-based deployment - Choosing the right approach - Decision tree - Hybrid approach - Troubleshooting - Import errors - Code changes not reflected - Files missing in container - Container build failures - Version conflicts - \[Run and deploy tasks > Deployment patterns\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/deployment-patterns/page.md) - Overview of deployment patterns - Simple file deployment - Example structure - Deployment commands - When to use - Custom Dockerfile deployment - Example structure - Alternative: Dockerfile in different directory - Key considerations - When to use - PyProject package deployment - Example structure - Business logic modules - Flyte orchestration layer - Entrypoint configuration - Dependencies and configuration - Key features - Key learning points - Usage patterns - What this example demonstrates - When to use - Package structure deployment - Example structure - Key concepts - Running with root directory - How \`--root-dir\` works - Alternative: Using a Python project - When to use - Full build deployment - Overview - Key configuration - Local dependency example - Critical configuration components - Configuration options - Version management best practices - Performance considerations - When to use - Troubleshooting - Python path deployment - Example structure - Implementation - Task environment dependencies - Key considerations - CLI vs Direct Python execution - Best practices - Common pitfalls - When to use - Dynamic environment deployment - Domain-based environment selection - Why this pattern works - How it works - Important constraints - Alternative: Environment variable approach - Usage patterns - When to use dynamic environments - Best practices - Project organization - Image management - Configuration management - Development workflow - Choosing the right pattern - \[Run and deploy tasks > Run context\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/run-context/page.md) - Configuring a run with \`flyte.with\_runcontext()\` - Execution target - Storage - Caching - Identity and resources - Logging - Code bundling - Context propagation - Reading context inside a task with \`flyte.ctx()\` - \`TaskContext\` fields - \`ActionID\` fields - Naming external resources - \[Run and deploy tasks > Entrypoint tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/entrypoints/page.md) - When to use entrypoints - Mark a task as an entrypoint - Discover entrypoint tasks - CLI - Programmatic - UI - \[Run and deploy tasks > Run a Python script\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/run-python-script/page.md) - When to use it - Quickstart - What you can specify - Handling dependencies - 1. Standard library only - 2. Pip packages on top of the default image - 3. Bring your own image - Handling inputs and outputs - Passing inputs - Capturing outputs - Bundling: what gets uploaded - Python API - \[Run and deploy tasks > Run with notifications\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/run-with-notifications/page.md) - Execution phases - Template variables - Slack notifications - Email notifications - \[Configure apps\](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/page.md) - Hello World example - Using fserve args - Using @app\_env.server - Differences from TaskEnvironment - Configuration topics > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/section.md - \[Configure apps > App environment settings\](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/app-environment-settings/page.md) - Shared environment settings - App-specific environment settings - Environment variable substitution in \`args\` - App startup - Server decorator via \`@app\_env.server\` - Container command via \`command\` vs \`args\` - Shared settings - \[Configure apps > Including additional files\](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/including-additional-files/page.md) - How include works - When to use include - Examples - Multi-file Streamlit app - Multi-file FastAPI app - App with configuration files - File discovery - Path resolution - Best practices - Limitations - \[Configure apps > Passing parameters into app environments\](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/passing-parameters/page.md) - Parameter types overview - Basic parameter types - Accessing parameters in your app - \`mount\` - \`env\_var\` - \`get\_parameter\` - Full example - Delayed values - RunOutput - AppEndpoint - Overriding parameters at serve time - Example: FastAPI app with configurable model - Example: Using RunOutput for model serving - Best practices - Limitations - \[Configure apps > /// script\](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/auto-scaling-apps/page.md) - Autoscaling apps - Scaling configuration - Basic scaling example - Scaling patterns - Idle TTL (Time To Live) - Autoscaling best practices - Autoscaling limitations - Autoscaling troubleshooting - \[Configure apps > Apps depending on other environments\](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/apps-depending-on-environments/page.md) - Basic usage - Example: App calling another app - Dependency chain - Multiple dependencies - Using AppEndpoint for dependency URLs - Deployment behavior - Task environment dependencies - Best practices - Example: A/B testing with dependencies - Limitations - \[Build apps\](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/page.md) - App types - Usage patterns - Next steps > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/build-apps/section.md - \[Build apps > Single-script apps\](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/single-script-apps/page.md) - Plain Python HTTP server - Streamlit app - FastAPI app - Running single-script apps - When to use single-script apps - \[Build apps > Multi-script apps\](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/multi-script-apps/page.md) - FastAPI multi-script app - Project structure - Example: Multi-file FastAPI app - Automatic file discovery - Streamlit multi-script app - Project structure - Example: Multi-file Streamlit app - Deploying multi-file Streamlit app - Complex multi-file example - Project structure - Example code - Deploying complex app - Best practices - Troubleshooting - \[Build apps > Serving graphs\](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/serving-graphs/page.md) - Core concepts: a minimal two-app chain - Deploying multiple apps together with \`depends\_on\` - Getting an upstream app's endpoint - Sizing each node independently - Example: CPU / GPU inference split - Disjoint images per node - GPU app: model.forward only - CPU app: pre/postprocess + call GPU - Deploy - Example: A/B testing with Statsig - Statsig client singleton - Variant apps - Root app with Statsig in its lifespan - App environments - Routing endpoint - Deploy - When to split into a serving graph - Best practices - \[Build apps > Hybrid app-task graphs\](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/hybrid-graphs/page.md) - Call app from task - Example: FastAPI app called from a task - Example: Call a model inference service from a task - Call task from app (webhooks / APIs) - Example: Basic webhook app - Advanced webhook patterns - Webhook security and best practices - Example: GitHub webhook - Gradio agent UI - Best practices - \[Build apps > WebSocket apps\](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/websocket-apps/page.md) - Example: Basic WebSocket app - WebSocket patterns - Using WebSockets with Flyte tasks - WebSocket client example - Best practices - \[Build apps > Browser apps\](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/browser-apps/page.md) - Accessing browser-based apps - Common browser-based app types - Streamlit apps - Gradio apps - Custom HTML/JS apps - Best practices - \[Build apps > Secret-based authentication\](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/secret-based-authentication/page.md) - Create the secret - Define the FastAPI app - Deploy the FastAPI app - Invoke the endpoint - Authentication for vLLM and SGLang apps - Create the authentication secret - Deploy vLLM app with authentication - Deploy SGLang app with authentication - Invoke authenticated LLM endpoints - Accessing Swagger documentation - Security best practices - Troubleshooting - Next steps - \[Native app integrations\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/page.md) - When to use a native integration - Available integrations - Next steps > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/section.md - \[Native app integrations > Streamlit app\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/streamlit-app/page.md) - Basic Streamlit app - Single-file Streamlit app - Multi-file Streamlit app - Example: Data visualization dashboard - Best practices - Troubleshooting - \[Native app integrations > FastAPI app\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/fastapi-app/page.md) - Basic FastAPI app - Serving a machine learning model - Accessing Swagger documentation - Example: REST API with multiple endpoints - Multi-file FastAPI app - Local-to-remote model serving - Best practices - Advanced features - Troubleshooting - \[Native app integrations > vLLM app\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/vllm-app/page.md) - Installation - Basic vLLM app - Using prefetched models - Model streaming - Custom vLLM arguments - Using the OpenAI-compatible API - Multi-GPU inference (Tensor Parallelism) - Model sharding with prefetch - Autoscaling - Best practices - Troubleshooting - \[Native app integrations > SGLang app\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/sglang-app/page.md) - Installation - Basic SGLang app - Using prefetched models - Model streaming - Custom SGLang arguments - Using the OpenAI-compatible API - Multi-GPU inference (Tensor Parallelism) - Model sharding with prefetch - Autoscaling - Structured generation - Best practices - Troubleshooting - \[Native app integrations > Flyte webhook\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/flyte-webhook/page.md) - Available endpoints - Basic usage - Filtering endpoints - Endpoint groups - Individual endpoints - Allow-listing - Task allow-list - App allow-list - Trigger allow-list - Calling the webhook - Authentication - Self-reference protection - \[Serve and deploy apps\](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/page.md) - Serve vs Deploy - \`flyte serve\` - \`flyte deploy\` - Using Python SDK - Serve - Deploy - Using the CLI - Serve - Deploy - Next steps > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/section.md - \[Serve and deploy apps > How app serving works\](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/how-app-serving-works/page.md) - Overview - Using the Python SDK - Overriding parameters - Advanced serving options - Using CLI - Return value - Best practices - Troubleshooting - \[Serve and deploy apps > How app custom domains work\](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/how-app-custom-domain-works/page.md) - \[Serve and deploy apps > How app deployment works\](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/how-app-deployment-works/page.md) - Overview - Using the Python SDK - Deployment plan - Overriding App configuration at deployment time - Activation/deactivation - Using the CLI - Example: Full deployment configuration - Best practices - Deployment status and return value - Troubleshooting - \[Serve and deploy apps > Activating and deactivating apps\](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/activating-and-deactivating-apps/page.md) - Activation - Activate after deployment - Activate an app - Check activation status - Deactivation - Lifecycle management - Typical deployment workflow - Blue-green deployment - Using CLI - Activate - Deactivate - Check status - Best practices - Automatic activation with serve - Example: Complete deployment and activation - Troubleshooting - \[Serve and deploy apps > Prefetching models\](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/prefetching-models/page.md) - Why prefetch? - Basic prefetch - Using Python SDK - Using CLI - Using prefetched models - Prefetch options - Custom artifact name - With HuggingFace token - With resources - Sharding models for multi-GPU - vLLM sharding - Using shard config via CLI - Using prefetched sharded models - CLI options - Complete example - Best practices - Troubleshooting - \[Build an agent\](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/page.md) - How Flyte maps to the agentic world - Ways to build an agent - Deploying an agent - Related > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/build-agent/section.md - \[Build an agent > Pure Python agents\](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/python-agents/page.md) - ReAct pattern: Reason, Act, Observe (no framework needed) - Plan-and-Execute with parallel fan-out - More agentic patterns - How Flyte's primitives map to the agent stack - Next steps - \[Build an agent > Flyte-native agents\](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/flyte-agents/page.md) - How it works - Sync vs async - A minimal agent - Tools - Customizing a tool with \`tool(...)\` - MCP integration - Skills - Observability - Extending the Agent class - \`run\` is sync-by-default - Strategy 1: wrap the built-in loop - Strategy 2: implement \`run\` from scratch - Choosing between subclassing and composition - Next steps - \[Build an agent > Agent memory\](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/agent-memory/page.md) - What a \`MemoryStore\` holds - Sync vs async - Keyed stores: the easy path - Path-addressed artifacts - Optimistic concurrency - Optional capabilities - Lower-level usage (without a key) - Next steps - \[Build an agent > Agent chat UI\](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/agent-chat-ui/page.md) - Option 1: the built-in chat UI - Option 2: a custom FastAPI chat app - Next steps - \[Build an agent > Deploy an agent as a service\](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/deploy-agent-as-service/page.md) - As a task - As a scheduled task (via \`Trigger\`) - Behind a webhook (\`AppEnvironment\`) - Chat and other app patterns - \[Agent framework integrations\](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/page.md) - How much control does the framework give you? - Supported frameworks - Next steps - \[Agent framework integrations > LangGraph agents\](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/langgraph/page.md) - A single LangGraph agent in a task - Plan-and-Execute: fan out LangGraph agents in parallel - Next steps - \[Agent framework integrations > PydanticAI agents\](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/pydantic-ai/page.md) - A PydanticAI agent in a task - Parallel agents - Next steps - \[Agent framework integrations > OpenAI agents\](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/openai-agents-sdk/page.md) - Tools that are also durable tasks - Next steps - \[Agent framework integrations > Bring your own framework\](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/bring-your-own-framework/page.md) - The core pattern - Make tools durable - Trace the framework's internals - Fan out across containers - Checklist - Next steps - \[Build an MCP\](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/page.md) - HTTP layout - Quickstart - \[Build an MCP > User-defined MCP server\](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/mcp\_server/page.md) - When to use it - How it works - Basic example - HTTP layout - Choosing a transport - Connecting a client - Claude Code - OpenCode - Configuration tips - Best practices - \[Build an MCP > Flyte MCP server\](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/flyte\_mcp\_server/page.md) - When to use it - How to run it - Running locally with \`uvx\` - Deploying remotely - Basic example - Scoping the server - 1. Tool groups - 2. Individual tools - 3. Allowlists - Enabling the search tools - Putting it together: a filtered server - Connecting a client - Claude Code — local (stdio) - Claude Code — remote (HTTP) - OpenCode — local - OpenCode — remote - Best practices - MCP tools reference - \[Sandboxing\](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/page.md) - Why sandboxing matters for AI - Types of sandboxes - What Flyte offers - Workflow sandbox (Monty) - Code sandbox (container) - When to use which - Learn more > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/section.md - \[Sandboxing > Workflow sandboxing in Flyte\](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/workflow-sandboxing-flyte/page.md) - Why workflow sandboxing? - How it works - Example: sandboxed orchestrator - Example: dynamic code execution - Reusable task from a code string - One-shot local execution - Parameterized code generation - Building agents with programmatic tool calling - Syntax restrictions - Allowed - Not allowed - Type restrictions - Security model - \[Sandboxing > Programmatic tool calling for agents\](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/code-mode/page.md) - Programmatic tool calling vs sequential tool calling - Why programmatic tool calling is powerful - Token efficiency - Performance - Natural programming patterns - Progressive tool discovery - Data privacy - Example: sequential vs programmatic tool calling - Sequential tool calling approach - Programmatic tool calling approach - Example: defining tools - Example: programmatic tool-calling agent - Example: chat app - Example: durable agent - References - \[Sandboxing > Code sandboxing\](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/code-sandboxing/page.md) - Execution modes - Auto-IO mode - Verbatim mode - Command mode - Executing a sandbox - Error handling - Supported types - How types are handled - Configuring the container image - Python packages - System packages - Additional Dockerfile commands - Pre-built images - Image configuration - Runtime configuration - Resources - Retries - Timeout - Environment variables - Secrets - Caching - Deploying a sandbox as a task - End-to-end example - API reference - \`flyte.sandbox.create()\` - Sandbox methods - \[Project patterns\](https://www.union.ai/docs/v2/flyte/user-guide/project-patterns/page.md) - \[Bring your own image (BYOI)\](bring-your-own-image/page.md) - \[Monorepo with uv\](monorepo-with-uv/page.md) - \[CI/CD deployments\](cicd) - \[Resource management and multi-team scaling\](resource-management) - \[Project patterns > Bring your own image\](https://www.union.ai/docs/v2/flyte/user-guide/project-patterns/bring-your-own-image/page.md) - The multi-team problem - Pattern 1: Pure BYOI - Dockerfiles - Build and push - Python code - Run and deploy - Pattern 2: Remote Builder - The base images - Adapting with \`flyte.Image\` - Task definitions - Entry point - Run and deploy - Decision matrix - \[Project patterns > Structuring Flyte projects with uv\](https://www.union.ai/docs/v2/flyte/user-guide/project-patterns/monorepo-with-uv/page.md) - The two layers - How the image gets built - \`with\_code\_bundle()\` — one image for dev and prod - \`root\_dir\` - Monorepo patterns - Pattern A: Shared lockfile (recommended) - Pattern B: Independent packages - The full build path (production) - \[Scale your runs\](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/page.md) - Understanding Flyte execution - Performance optimization - Key concepts for scaling > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/section.md - \[Scale your runs > Data flow\](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/data-flow/page.md) - Overview - Data types and transport - Passed by reference - Passed by value (inline I/O) - Task execution and data flow - Input download - Output upload - Task-to-task data flow - Caching and data hashing - Cache key computation - Inline data caching - Reference data hashing - Cache control - Traces and data flow - Object stores and latency considerations - Configuring data storage - Organization and project level - Per-run configuration - \[Scale your runs > Life of a run\](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/life-of-a-run/page.md) - Overview - Phase 1: Code analysis and preparation - Phase 2: Image building - Phase 3: Code bundling - Default: \`copy\_style="loaded\_modules"\` - Alternative: \`copy\_style="none"\` - Phase 4: Upload code bundle - Phase 5: Run creation and queuing - Phase 6: Task execution in data plane - Container startup - Invoking downstream tasks - Execution flow diagram - Action identifiers and crash recovery - Downstream task execution - Reusable containers - Reusable container execution flow - State replication and visualization - Queue Service to Run Service - UI limitations - Optimization opportunities - \[Scale your runs > Scale your workflows\](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/scale-your-workflows/page.md) - Understanding performance dimensions - Latency - Throughput - Task execution overhead - The overhead principle - System architecture and data flow - Optimization strategies - 1. Use reusable containers for concurrency - 2. Batch workloads to reduce overhead - 3. Use traces for lightweight operations - 4. Limit fanout for system stability - 5. Optimize data transfer - 6. Leverage caching - 7. Parallelize with \`flyte.map\` - Performance tuning workflow - Real-world example: PyIceberg batch processing - Example: Optimizing a data pipeline - Before optimization - After optimization - When to contact the Union team - \[Scale your runs > Maximize GPU utilization for batch inference\](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/batch-inference/page.md) - Why GPU utilization drops - Serving vs in-process batch inference - Solution: \`DynamicBatcher\` - Basic usage - Cost estimation - \`TokenBatcher\` for LLM inference - Combining with app environments - Example: batch LLM inference with vLLM behind a FastAPI app - Monitoring utilization - \[Advanced project: LLM reporting agent\](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/page.md) - What you'll build - Concepts covered - Architecture - Prerequisites - Parts - Key takeaways > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/section.md - \[Advanced project: LLM reporting agent > Resilient generation\](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/resilient-generation/page.md) - Two environments - Reusable environment for LLM work - ReusePolicy parameters - Standard environment for orchestration - Traced LLM calls - Benefits of tracing - When to use @flyte.trace - Traced helper functions - Retry strategies - Configuring retries - Combining tracing with retries - Structured prompts - Pydantic models for structured output - Next steps - \[Advanced project: LLM reporting agent > Agentic refinement\](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/agentic-refinement/page.md) - The agentic pattern - Critique function - Revise function - The refinement loop - How it works - Early exit - Grouping iterations with flyte.group - Why use flyte.group? - Group context - Configuring the loop - Choosing thresholds - Best practices for agentic loops - Next steps - \[Advanced project: LLM reporting agent > Parallel outputs\](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/parallel-outputs/page.md) - The formatting functions - When to trace and when not to - Parallel execution with asyncio.gather - How asyncio.gather works - When to use asyncio.gather - Grouping parallel operations - Collecting outputs in a directory - The batch pipeline - Pipeline flow - Running the pipeline - Cost optimization tips - 1. Choose the right model - 2. Tune iteration parameters - 3. Use caching effectively - 4. Scale the batch - Next steps - \[Advanced project: LLM reporting agent > Serving app\](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/serving-app/page.md) - App environment configuration - Key configuration - Connecting to pipeline output with RunOutput - The Streamlit application - Displaying multiple reports - Generation instructions - Deploying the app - Workflow: Generate then view - Automatic updates with RunOutput - Complete example structure - Running the complete example - Summary - \[Migration\](https://www.union.ai/docs/v2/flyte/user-guide/migration/page.md) - \[From Flyte 1 to 2\](flyte-2/page.md) - \[From Airflow to Flyte\](from-airflow/page.md) - \[Migration > From Flyte 1 to 2\](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/page.md) - Pure Python execution - Sync Python - Async Python - Simplified API - Fine-grained reproducibility and recoverability - Improved remote functionality - Native Notebook support > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/section.md - \[Migration > From Flyte 1 to 2 > Pure Python\](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/pure-python/page.md) - From \`@workflow\` DSL to pure Python - Flyte 1 - Flyte 2 - \[Migration > From Flyte 1 to 2 > Asynchronous model\](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/async/page.md) - Why we need an async model - Understanding concurrency vs. parallelism - Python's async evolution - Parallelism in Flyte 1 vs Flyte 2 - Core async concepts - True parallelism for all workloads - Calling sync tasks from async tasks - Synchronous task support - The \`flyte.map\` function: Familiar patterns - Sync Map - Async Map - \[Migration > From Flyte 1 to 2 > Migration from Flyte 1 to Flyte 2\](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/migration/page.md) - 1. Move task configuration to a \`TaskEnvironment\` object - 2. Replace workflow decorators - Flyte 1 - Flyte 2 Sync - Flyte 2 Async - \[Migration > From Flyte 1 to 2 > Considerations\](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/considerations/page.md) - Non-deterministic behavior - Dealing with non-determinism - Type safety - No global state - Driver pod requirements - OOM risk from materialized I/O - \[Migration > From Airflow to Flyte\](https://www.union.ai/docs/v2/flyte/user-guide/migration/from-airflow/page.md) - \[Part 1 — Vanilla Operators\](part-1-vanilla-operators/page.md) - \[Migration > From Airflow to Flyte > Part 1 — Vanilla Operators\](https://www.union.ai/docs/v2/flyte/user-guide/migration/from-airflow/part-1-vanilla-operators/page.md) - 1. Where dependencies are specified - 2. The driver task (in place of DAGs) - 3. Triggers (in place of schedules) - 4. PythonOperator to \`@env.task\` - File and Dir — for data that doesn't fit in a return value - 5. TaskFlow to \`@env.task\` - TaskFlow decorator variants - 6. BashOperator to ContainerTask - When to use \`ContainerTask\` - How the arguments map - 7. KubernetesPodOperator to TaskEnvironment + PodTemplate - Where every KPO knob lands - What a fully-specified task looks like - 8. Orchestration: parallelism, conditionals, error handling - Parallelism - Dynamic mapping - Conditionals - Error handling - What's next - Caching - Reusable containers - Reports - Apps --- ## Tutorials - \[Biotech & Healthcare\](https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/page.md) - \[Genomic alignment\](genomic-alignment/page.md) - \[Brain tumor MRI classification\](tumor-detection/page.md) - \[Biotech & Healthcare > Genomic alignment\](https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/genomic-alignment/page.md) - Define the container image - Define the task environments - Define the data classes - Fetch assets - Quality filtering with fastp - Build the Bowtie 2 index - Align reads - Orchestrate the workflow - Run the workflow - \[Biotech & Healthcare > Brain tumor MRI classification\](https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/tumor-detection/page.md) - Define the container image - Define the task environments - Configure training - Load the dataset - Train the model - Resumable checkpointing - Generate the report - Orchestrate the pipeline - Run the pipeline - \[Geospatial\](https://www.union.ai/docs/v2/flyte/tutorials/geospatial/page.md) - \[GPU-accelerated climate modeling\](climate-modeling/page.md) - \[Satellite image classification\](satellite\_image\_classification) - \[Geospatial > GPU-accelerated climate modeling\](https://www.union.ai/docs/v2/flyte/tutorials/geospatial/climate-modeling/page.md) - Overview - Implementation - Dependencies and container image - Simulation parameters and data structures - Task environments - Data ingestion: multiple sources in parallel - Preprocessing with Dask - GPU-accelerated atmospheric simulation - Distributing across multiple GPUs - The main workflow - Running the pipeline - Key concepts - Ensemble forecasting - Adaptive mesh refinement - Real-time event detection - Where to go next - \[Financial Services & Fintech\](https://www.union.ai/docs/v2/flyte/tutorials/financial-services/page.md) - \[Financial research agent\](financial-research-agent/page.md) - \[Multi-agent trading simulation\](trading-agents/page.md) - \[Financial Services & Fintech > Multi-agent trading simulation\](https://www.union.ai/docs/v2/flyte/tutorials/financial-services/trading-agents/page.md) - TL;DR - What is an agent, anyway? - What's different here? - How it works: step-by-step walkthrough - Entry point - Analyst agents - Research agents - Trading agent - Risk agents - Retaining agent memory with S3 vectors - Running the simulation - Why Flyte? \_(A quick note before you go)\_ - \[Financial Services & Fintech > Financial research agent\](https://www.union.ai/docs/v2/flyte/tutorials/financial-services/financial-research-agent/page.md) - Setting up the environment - Data types - You.com Research and Search APIs - Synthesize briefings with Claude - Research one company - Orchestration - Run the agent - Create secrets - Run locally or remotely - \[Frontier AI\](https://www.union.ai/docs/v2/flyte/tutorials/frontier-ai/page.md) - \[Distributed LLM pretraining\](distributed-pretraining/page.md) - \[Frontier AI > Distributed LLM pretraining\](https://www.union.ai/docs/v2/flyte/tutorials/frontier-ai/distributed-pretraining/page.md) - Overview - Implementation - Setting up the environment - Declaring resource requirements - Model configurations - Building the GPT model - The Lightning training module - Checkpointing for fault tolerance - Real-time metrics with Flyte Reports - Streaming data at scale - Distributed training with FSDP - Tying it together - Running the pipeline - Going further - \[Computer Vision\](https://www.union.ai/docs/v2/flyte/tutorials/computer-vision/page.md) - \[Fine-tuning a VLM\](qwen-vl-finetuning/page.md) - \[Multimodal retrieval evaluation\](multimodal-retrieval-evaluation/page.md) - \[Computer Vision > Fine-tuning a VLM\](https://www.union.ai/docs/v2/flyte/tutorials/computer-vision/qwen-vl-finetuning/page.md) - Overview - Implementation - Setting up the environment - Preparing the dataset - The adapter - Multi-node training with DeepSpeed - Fault tolerance and recovery - Live observability - Evaluation - Putting it all together - Running the tutorial - Going further - \[Computer Vision > Multimodal retrieval evaluation\](https://www.union.ai/docs/v2/flyte/tutorials/computer-vision/multimodal-retrieval-evaluation/page.md) - Define the container image - Define the task environments - Configuration and data types - Loading, indexing, and search - Run one experiment - Compare experiments - Run the evaluation - \[Agents\](https://www.union.ai/docs/v2/flyte/tutorials/agents/page.md) - \[Autoresearch agent\](autoresearch/page.md) - \[Coding agent\](code-agent/page.md) - \[Competitive intelligence agent\](competitive-intelligence-agent/page.md) - \[Compliance monitoring agent\](compliance-monitoring-agent/page.md) - \[Deep research\](deep-research/page.md) - \[Field data enrichment agent\](field-data-enrichment-agent/page.md) - \[MLE Bot: autonomous ML engineer\](mle-bot/page.md) - \[Support resolution agent\](support-resolution-agent/page.md) - \[Agents > Autoresearch agent\](https://www.union.ai/docs/v2/flyte/tutorials/agents/autoresearch/page.md) - Define the container image - Define the task environment - What changed - All changed files - Model the result - What changed - All changed files - The autoresearch task - What changed - All changed files - What changed - All changed files - Run the agent - Create secrets - Prepare the research repository - Run remotely - \[Agents > Coding agent\](https://www.union.ai/docs/v2/flyte/tutorials/agents/code-agent/page.md) - What this example demonstrates - Setting up the agent environment - Retrieve docs - Code generation - Running the code agent - \[Agents > MLE Bot: an autonomous ML engineer\](https://www.union.ai/docs/v2/flyte/tutorials/agents/mle-bot/page.md) - TL;DR - The problem with LLMs and ML pipelines - How it works - What to expect - Declaring task environments - Building durable tool functions - Guiding the LLM with domain knowledge - Dataset context - General ML best practices - The agent loop: profile, design, execute, iterate - Dataset context - General ML best practices — apply these based on the dataset context above - Available tools - Monty sandbox restrictions - Critical patterns for using tool results - When fixing a previous error - Pipeline design — you decide the structure - Response format - Reasoning - Code - Running LLM-generated code in Flyte's sandbox - Dataset context - General ML best practices — apply these based on the dataset context above - Available tools - Monty sandbox restrictions - Critical patterns for using tool results - When fixing a previous error - Pipeline design — you decide the structure - Response format - Reasoning - Code - Streaming results to a live report - Dataset context - General ML best practices — apply these based on the dataset context above - Available tools - Monty sandbox restrictions - Critical patterns for using tool results - When fixing a previous error - Pipeline design — you decide the structure - Response format - Reasoning - Code - Running it - Why Flyte? - \[Agents > Deep research\](https://www.union.ai/docs/v2/flyte/tutorials/agents/deep-research/page.md) - Setting up the environment - Generate research queries - Search and summarize - Evaluate research completeness - Filter results - Generate the final answer - Orchestration - Run the deep research agent - Evaluate with Weights & Biases Weave - \[Agents > Competitive intelligence agent\](https://www.union.ai/docs/v2/flyte/tutorials/agents/competitive-intelligence-agent/page.md) - Setting up the environment - Data types - Search with the You.com Search API - Extract deltas with Claude - Watch one competitor - Orchestration - Run the agent - Create secrets - Run locally or remotely - \[Agents > Compliance monitoring agent\](https://www.union.ai/docs/v2/flyte/tutorials/agents/compliance-monitoring-agent/page.md) - Setting up the environment - Data types - Research with the You.com Research API - Triage findings with Claude - Monitor one watch item - Orchestration - Run the agent - Create secrets - Run locally or remotely - \[Agents > Field data enrichment agent\](https://www.union.ai/docs/v2/flyte/tutorials/agents/field-data-enrichment-agent/page.md) - Setting up the environment - Data types - Search with the You.com Search API - Enrich one event - Orchestration - Run the agent - Create secrets - Run locally or remotely - \[Agents > Support resolution agent\](https://www.union.ai/docs/v2/flyte/tutorials/agents/support-resolution-agent/page.md) - Setting up the environment - Data types - Ground answers with the You.com Research API - Ground one ticket - Draft a customer-ready reply - Resolve one ticket - Orchestration - Run the agent - Create secrets - Run locally or remotely - \[Context Engineering\](https://www.union.ai/docs/v2/flyte/tutorials/context-engineering/page.md) - \[Automatic prompt engineering\](auto\_prompt\_engineering/page.md) - \[Text-to-SQL\](text\_to\_sql/page.md) - \[Context Engineering > Text-to-SQL prompt optimization\](https://www.union.ai/docs/v2/flyte/tutorials/context-engineering/text\_to\_sql/page.md) - Ingesting data - From question to SQL - Vector indexing - Table retrieval and context building - SQL generation and response synthesis - Building the QA dataset - Schema extraction and chunking - Question and SQL generation - Validation and quality control - Optimizing prompts - Evaluation pipeline - Iterative optimization - Run it - What we observed - The bigger lesson - \[Context Engineering > Automatic prompt engineering\](https://www.union.ai/docs/v2/flyte/tutorials/context-engineering/auto\_prompt\_engineering/page.md) - Set up the environment - Prepare the evaluation dataset - Define models - Evaluate prompts - Optimize prompts - Build the full pipeline - Run it - Why this matters - Next steps - \[Model Training\](https://www.union.ai/docs/v2/flyte/tutorials/model-training/page.md) - \[Hyperparameter optimization\](hpo/page.md) - \[Model Training > Hyperparameter optimization\](https://www.union.ai/docs/v2/flyte/tutorials/model-training/hpo/page.md) - A better way to run HPO - Declare dependencies - Define the task environment - Define the optimizer - Define the objective function - Define the main optimization loop - Run the experiment --- ## Integrations - \[Anthropic\](https://www.union.ai/docs/v2/flyte/integrations/anthropic/page.md) - Installation - Quick start - API - \`function\_tool\` - \`Agent\` - \`run\_agent\` - Secrets - API reference - \[BigQuery\](https://www.union.ai/docs/v2/flyte/integrations/bigquery/page.md) - Installation - Quick start - Configuration - \`BigQueryConfig\` parameters - \`BigQueryTask\` parameters - Authentication - Query templating - Supported input types - Parameterized query example - Retrieving query results - API reference - \[Code generation\](https://www.union.ai/docs/v2/flyte/integrations/codegen/page.md) - Installation - Quick start - Two execution backends - LiteLLM (default) - Agent (Claude) - Providing data - Sample data - Schema and constraints - Inputs and outputs - Running generated code - One-shot execution with \`result.run()\` - Reusable task with \`result.as\_task()\` - Error diagnosis - Durable execution - Replay logs - Caching - Non-determinism in Agent mode - Observability - LiteLLM backend - Agent backend - Examples - Processing CSVs with different schemas - DataFrame analysis with constraints - Agent mode - Configuration - LiteLLM parameters - Image configuration - Skipping tests - Base packages - Best practices - API reference - \`AutoCoderAgent\` constructor - \`generate()\` parameters - \`CodeGenEvalResult\` fields - \`CodeGenEvalResult\` methods - \[Dask\](https://www.union.ai/docs/v2/flyte/integrations/dask/page.md) - When to use this plugin - Installation - Configuration - \`Dask\` parameters - \`Scheduler\` parameters - \`WorkerGroup\` parameters - Accessing the Dask client - Example - API reference - \[Databricks\](https://www.union.ai/docs/v2/flyte/integrations/databricks/page.md) - Installation - Quick start - Configuration - Spark fields (inherited) - Databricks-specific fields - \`databricks\_conf\` structure - Authentication - Accessing the Spark session - API reference - \[Gemini\](https://www.union.ai/docs/v2/flyte/integrations/gemini/page.md) - Installation - Quick start - API - \`function\_tool\` - \`Agent\` - \`run\_agent\` - Secrets - API reference - \[Hydra\](https://www.union.ai/docs/v2/flyte/integrations/hydra/page.md) - Installation - Requirements on tasks - A walkthrough config - Execution mode - Hydra launcher (\`@hydra.main\` scripts) - Python SDK - Single run - Grid sweep - Custom sweepers - Forwarding \`flyte.with\_runcontext\` options - Flyte CLI (\`flyte hydra run\`) - Single run - Grid sweep - App-level vs Hydra-namespace overrides - \`--follow\` and \`--no-wait\` - Shell completion - Override grammar - Sweeps - Grid sweeps (BasicSweeper) - Bayesian / TPE sweeps (Optuna) - Sweep output directories - Task environment overrides - Prebuilt images - Applying overrides to child tasks - Renaming the task-env key - What \`task\_env\` should not model - Structured configs (without YAML) - \[MLflow\](https://www.union.ai/docs/v2/flyte/integrations/mlflow/page.md) - Installation - Quick start - Autologging - Generic autologging - Framework-specific autologging - Run modes - Sharing a run across tasks - Creating independent runs - Nested runs - Workflow-level configuration - Per-task overrides - Configuration priority - Distributed training - MLflow UI links - Setup - Custom URL templates - Explicit links - Link behavior by run mode - Automatic Flyte tags - API reference - \`mlflow\_run\` and \`mlflow\_config\` - \`get\_mlflow\_run\` - \`get\_mlflow\_context\` - \`Mlflow\` - \[OmegaConf\](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/page.md) - Installation - Quick start - When to use this plugin - Building a DictConfig - From a plain dict - From a YAML file - From a dataclass (structured config) - From a base config plus overrides - Variable interpolation - Nested and deeply structured configs - DictConfigs that contain lists - ListConfig as input and output - Lists of primitives - Building a schedule from another task - Nested lists (list of lists) - Lists of DictConfigs - Lists of dataclass instances - Structured configs - Basic structured config - Schema reconstruction in the receiving task - Required (\`MISSING\`) fields - Advanced field types - Merging overrides on top of a structured base - Embedding rich Python values inside a plain DictConfig - Reserved-looking keys - YAML reports - Wire format - End-to-end example - \[OpenAI\](https://www.union.ai/docs/v2/flyte/integrations/openai/page.md) - When to use this plugin - Installation - Usage - \`function\_tool\` - Basic pattern - Secrets - Example - API reference - \[OpenAI > Agent tools\](https://www.union.ai/docs/v2/flyte/integrations/openai/agent\_tools/page.md) - Define the tools - Define the agent - Run the agent - Conclusion - \[Pandera\](https://www.union.ai/docs/v2/flyte/integrations/pandera/page.md) - When to use this plugin - Installation - pandas - Polars - PySpark SQL - Defining schemas - Using schemas in tasks - Error handling with \`ValidationConfig\` - Image configuration - Pandas - Polars - PySpark SQL - Polars lazy frames - Examples - pandas - Polars - PySpark SQL - \[Papermill\](https://www.union.ai/docs/v2/flyte/integrations/papermill/page.md) - When to use this plugin - Installation - Quick start - Notebook setup - \`parameters\` cell - \`outputs\` cell - Inputs and outputs - Supported input types - Complex types: File, Dir, DataFrame - Outputs: single, multiple, none - Calling Flyte tasks from notebooks - Workflow patterns - Chaining notebooks - Mixing notebooks with regular tasks - Inline definition - Calling from sync vs. async tasks - Running a NotebookTask directly as the entrypoint - Reports and notebook artifacts - HTML report (default) - Notebook artifacts - Clean reports - Failure reports - Spark notebooks - Local testing - Execution options - \`NotebookTask\` reference - Helper functions - \[PyTorch\](https://www.union.ai/docs/v2/flyte/integrations/pytorch/page.md) - When to use this plugin - Installation - Configuration - \`Elastic\` parameters - \`RunPolicy\` parameters - NCCL tuning parameters - Writing a distributed training task - Example - API reference - \[Ray\](https://www.union.ai/docs/v2/flyte/integrations/ray/page.md) - When to use this plugin - Installation - Configuration - \`RayJobConfig\` parameters - \`WorkerNodeConfig\` parameters - \`HeadNodeConfig\` parameters - Connecting to an existing cluster - Examples - API reference - \[Snowflake\](https://www.union.ai/docs/v2/flyte/integrations/snowflake/page.md) - Installation - Quick start - Configuration - Required fields - Additional connection parameters - Authentication - Key-pair authentication - Password authentication - OAuth authentication - Query templating - Supported input types - Batched \`INSERT\` with list inputs - Parameterized \`SELECT\` - Multiple inputs - Retrieving query results - End-to-end example - \[Spark\](https://www.union.ai/docs/v2/flyte/integrations/spark/page.md) - When to use this plugin - Installation - Configuration - \`Spark\` parameters - Accessing the Spark session - Overriding configuration at runtime - Example - API reference - \[Weights & Biases\](https://www.union.ai/docs/v2/flyte/integrations/wandb/page.md) - Installation - Quick start - What's next - \[Weights & Biases > Experiments\](https://www.union.ai/docs/v2/flyte/integrations/wandb/experiments/page.md) - Basic usage - Accessing the run object - Parent-child task relationships - Run modes - Using \`run\_mode="new"\` for independent runs - Using \`run\_mode="shared"\` for explicit sharing - Configuration with \`wandb\_config\` - Workflow-level configuration - Overriding configuration for child tasks - Using traces with W&B runs - \[Weights & Biases > Distributed training\](https://www.union.ai/docs/v2/flyte/integrations/wandb/distributed\_training/page.md) - Quick start - Run modes in distributed training - Single-node behavior - Multi-node behavior - Choosing run mode and rank scope - Single-node multi-GPU - Basic example with \`auto\` mode - Using \`shared\` mode for per-rank metrics - Using \`new\` mode for per-rank runs - Multi-node training with \`Elastic\` - Global scope (default): Single run across all nodes - Worker scope: One run per node - Shared mode: All ranks log to the same run - New mode: Separate run per rank - How it works - Run ID patterns - \[Weights & Biases > Sweeps\](https://www.union.ai/docs/v2/flyte/integrations/wandb/sweeps/page.md) - Creating a sweep - Running parallel agents - Writing objective functions - \[Weights & Biases > Downloading logs\](https://www.union.ai/docs/v2/flyte/integrations/wandb/downloading\_logs/page.md) - Automatic download - Accessing run directories during execution - \[Weights & Biases > Constraints and best practices\](https://www.union.ai/docs/v2/flyte/integrations/wandb/constraints\_and\_best\_practices/page.md) - Decorator ordering - Traces cannot use decorators - Maximum sweep agents - Configuration priority - Run ID generation - Sync delay for local files - Shared run mode requirements - Objective functions for sweeps - Error handling - \[Weights & Biases > Manual integration\](https://www.union.ai/docs/v2/flyte/integrations/wandb/manual/page.md) - Using the Wandb link class - With a custom run ID - Adding links at runtime with override - Using the \`WandbSweep\` link class --- ## Reference - \[LLM-optimized documentation\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/page.md) - Per-page Markdown (\`page.md\`) - Section bundles (\`section.md\`) - Page index (\`llms.txt\`) - Full documentation (\`llms-full.txt\`) - \[Migration from Flyte 1 to Flyte 2\](https://www.union.ai/docs/v2/flyte/api-reference/migration/page.md) - Key API changes at a glance - Topics - \[Philosophy and imports\](overview/page.md) - \[Container images\](images/page.md) - \[Configuration and CLI\](configuration-and-cli/page.md) - \[Tasks and workflows\](tasks-and-workflows/page.md) - \[Secrets, resources, and caching\](secrets-resources-caching/page.md) - \[Parallelism and async\](parallelism-and-async/page.md) - \[Triggers and dynamic workflows\](triggers-and-dynamic/page.md) - \[Examples and common gotchas\](examples-and-gotchas/page.md) > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/api-reference/migration/section.md - \[Migration from Flyte 1 to Flyte 2 > Philosophy and imports\](https://www.union.ai/docs/v2/flyte/api-reference/migration/overview/page.md) - Key paradigm shifts - What Flyte 2 eliminates - What Flyte 2 introduces - Package imports - Basic import changes - Flyte 1 - Flyte 2 - Import mapping table - \[Migration from Flyte 1 to Flyte 2 > Container images\](https://www.union.ai/docs/v2/flyte/api-reference/migration/images/page.md) - Basic migration - Flyte 1 - Flyte 2 - Image constructor methods - Image builder methods (chainable) - Builder configuration (local vs remote) - Private registry with secrets - Flyte 1 - Flyte 2 - Parameter mapping - \[Migration from Flyte 1 to Flyte 2 > Configuration and CLI\](https://www.union.ai/docs/v2/flyte/api-reference/migration/configuration-and-cli/page.md) - Configuration files - Config file location - Config format - Flyte 1 - Flyte 2 - Key config differences - Specifying config via CLI - Flyte 1 - Flyte 2 - Specifying config in code - CLI commands - Command mapping - Running tasks - Flyte 1 - Flyte 2 - Key CLI flag differences - Deploying - Flyte 1 - Flyte 2 - Running deployed tasks - Complete Flyte 2 CLI options - \[Migration from Flyte 1 to Flyte 2 > Tasks and workflows\](https://www.union.ai/docs/v2/flyte/api-reference/migration/tasks-and-workflows/page.md) - Basic task migration - Flyte 1 - Flyte 2 - Workflow to task migration - Flyte 1 - Flyte 2 Sync - Flyte 2 Async - TaskEnvironment configuration - Parameter mapping: @task to TaskEnvironment + @env.task - \[Migration from Flyte 1 to Flyte 2 > Secrets, resources, and caching\](https://www.union.ai/docs/v2/flyte/api-reference/migration/secrets-resources-caching/page.md) - Secrets - Declaring and accessing secrets - Flyte 1 - Flyte 2 - Secret configuration options - Secret name convention changes - Creating secrets via CLI - Resources - Basic resource configuration - Flyte 1 - Flyte 2 - GPU configuration - Flyte 1 - Flyte 2 - Supported GPU types (Flyte 2) - Resource parameter mapping - Caching - Basic caching - Flyte 1 - Flyte 2 - Cache behavior options (Flyte 2) - \[Migration from Flyte 1 to Flyte 2 > Parallelism and async\](https://www.union.ai/docs/v2/flyte/api-reference/migration/parallelism-and-async/page.md) - Basic map\_task migration - Flyte 1 - Flyte 2 - map\_task with concurrency - Flyte 1 - Flyte 2 - Async parallel execution with asyncio.gather - Concurrency control with semaphore - Error handling with asyncio.gather - flyte.map vs asyncio.gather comparison - Recommended pattern selection - Sync and async task patterns - Sync tasks calling sync tasks - Async tasks calling async tasks - Sequential execution with await - Flyte 1 - Flyte 2 - \[Migration from Flyte 1 to Flyte 2 > Triggers and dynamic workflows\](https://www.union.ai/docs/v2/flyte/api-reference/migration/triggers-and-dynamic/page.md) - LaunchPlan to Trigger migration - Flyte 1 - Flyte 2 - Trigger options - Deploying triggers - Dynamic workflows - @dynamic to regular tasks - Flyte 1 - Flyte 2 Sync - Flyte 2 Async - Conditional execution - Flyte 1 - Flyte 2 - Subworkflows to nested tasks - Flyte 1 - Flyte 2 - \[Migration from Flyte 1 to Flyte 2 > Examples and common gotchas\](https://www.union.ai/docs/v2/flyte/api-reference/migration/examples-and-gotchas/page.md) - Complete migration examples - Example 1: Simple ML pipeline - Flyte 1 - Flyte 2 - Example 2: Parallel processing with map\_task - Flyte 1 - Flyte 2 Sync - Flyte 2 Async - Common gotchas - 1. current\_context() is replaced - 2. Workflow >> ordering notation is gone - 3. flyte.map returns a generator - 4. Image configuration location - 5. Resource configuration - 6. Cache version - 7. Entrypoint task naming - 8. Memory parameter name - 9. Retries have no platform cap - 10. Type annotations - Quick reference cheat sheet - \[Flyte CLI\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/page.md) - flyte - flyte abort - flyte build - flyte create - flyte delete - flyte deploy - flyte edit - flyte gen - flyte get - flyte prefetch - flyte run - flyte serve - flyte signal - flyte start - flyte stop - flyte update - flyte whoami - \[Flyte SDK\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/page.md) - \[Flyte SDK > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/classes/page.md) - \[Flyte SDK > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/page.md) - \[Flyte SDK > Packages > flyte\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/page.md) - Directory - Classes - Protocols - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte > AppHandle\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/apphandle/page.md) - Properties - Methods - activate() - deactivate() - ephemeral\_ctx() - ephemeral\_ctx\_sync() - is\_active() - is\_deactivated() - \[Flyte SDK > Packages > flyte > Backoff\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/backoff/page.md) - Parameters - Methods - compute\_delay() - \[Flyte SDK > Packages > flyte > BaseCheckpoint\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/basecheckpoint/page.md) - Properties - Methods - load() - load\_sync() - prev\_exists() - save() - save\_sync() - \[Flyte SDK > Packages > flyte > Cache\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/cache/page.md) - Parameters - Methods - get\_ignored\_inputs() - get\_version() - is\_enabled() - \[Flyte SDK > Packages > flyte > CachePolicy\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/cachepolicy/page.md) - Methods - get\_version() - \[Flyte SDK > Packages > flyte > Checkpoint\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/checkpoint/page.md) - Parameters - Properties - Methods - load() - load\_sync() - prev\_exists() - save() - save\_sync() - \[Flyte SDK > Packages > flyte > Cron\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/cron/page.md) - Parameters - Properties - \[Flyte SDK > Packages > flyte > Device\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/device/page.md) - Parameters - \[Flyte SDK > Packages > flyte > Environment\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/environment/page.md) - Parameters - Methods - add\_dependency() - clone\_with() - \[Flyte SDK > Packages > flyte > EventWebhook\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/eventwebhook/page.md) - Parameters - \[Flyte SDK > Packages > flyte > FixedRate\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/fixedrate/page.md) - Parameters - \[Flyte SDK > Packages > flyte > Image\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/image/page.md) - Parameters - Properties - Methods - clone() - from\_base() - from\_debian\_base() - from\_dockerfile() - from\_ref\_name() - from\_uv\_script() - validate() - with\_apt\_packages() - with\_code\_bundle() - with\_commands() - with\_dockerignore() - with\_env\_vars() - with\_local\_rs\_controller() - with\_local\_v2() - with\_local\_v2\_plugins() - with\_pip\_packages() - with\_poetry\_project() - with\_requirements() - with\_source\_file() - with\_source\_folder() - with\_uv\_project() - with\_workdir() - \[Flyte SDK > Packages > flyte > ImageBuild\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/imagebuild/page.md) - Parameters - \[Flyte SDK > Packages > flyte > Link\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/link/page.md) - Methods - get\_link() - \[Flyte SDK > Packages > flyte > PodTemplate\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/podtemplate/page.md) - Parameters - Methods - to\_k8s\_pod() - \[Flyte SDK > Packages > flyte > Resources\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/resources/page.md) - Parameters - Methods - get\_device() - get\_shared\_memory() - \[Flyte SDK > Packages > flyte > RetryStrategy\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/retrystrategy/page.md) - Parameters - \[Flyte SDK > Packages > flyte > ReusePolicy\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/reusepolicy/page.md) - Parameters - Properties - Methods - get\_scaledown\_ttl() - \[Flyte SDK > Packages > flyte > Secret\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/secret/page.md) - Parameters - Methods - stable\_hash() - \[Flyte SDK > Packages > flyte > TaskEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/taskenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - from\_task() - task() - \[Flyte SDK > Packages > flyte > Timeout\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/timeout/page.md) - Parameters - \[Flyte SDK > Packages > flyte > Trigger\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/trigger/page.md) - Parameters - Methods - daily() - hourly() - minutely() - monthly() - weekly() - \[Flyte SDK > Packages > flyte.ai.agents\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/page.md) - Directory - Classes - Protocols - Errors - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.ai.agents > AccessDenied\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/accessdenied/page.md) - \[Flyte SDK > Packages > flyte.ai.agents > Agent\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/agent/page.md) - Parameters - Properties - Methods - add\_tool() - approval\_callback() - call\_llm() - run() - tool\_descriptions() - \[Flyte SDK > Packages > flyte.ai.agents > AgentEvent\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/agentevent/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > AgentProtocol\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/agentprotocol/page.md) - Methods - run() - tool\_descriptions() - \[Flyte SDK > Packages > flyte.ai.agents > AgentResult\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/agentresult/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > AgentTool\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/agenttool/page.md) - Parameters - Methods - to\_openai\_format() - \[Flyte SDK > Packages > flyte.ai.agents > CodeModeAgent\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/codemodeagent/page.md) - Parameters - Methods - run() - tool\_descriptions() - uses\_flyte\_tools() - \[Flyte SDK > Packages > flyte.ai.agents > ConcurrencyError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/concurrencyerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > LLMMessage\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/llmmessage/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > MCPServerSpec\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/mcpserverspec/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > MemoryMeta\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/memorymeta/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > MemoryStore\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/memorystore/page.md) - Parameters - Methods - append() - audit\_tail() - audit\_tail\_sync() - create() - current\_sha() - exists() - extend() - flush\_messages() - flush\_messages\_sync() - get\_meta() - get\_or\_create() - list\_paths() - read\_json() - read\_text() - remote\_path\_for\_key() - save() - write\_json() - write\_text() - \[Flyte SDK > Packages > flyte.ai.agents > MemoryStoreError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents/memorystoreerror/page.md) - \[Flyte SDK > Packages > flyte.ai.agents.agent\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.agent/page.md) - Directory - Classes - Variables - \[Flyte SDK > Packages > flyte.ai.agents.agent > Agent\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.agent/agent/page.md) - Parameters - Properties - Methods - add\_tool() - approval\_callback() - call\_llm() - run() - tool\_descriptions() - \[Flyte SDK > Packages > flyte.ai.agents.agent > AgentEvent\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.agent/agentevent/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents.codemode\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.codemode/page.md) - Directory - Classes - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.ai.agents.codemode > CodeModeAgent\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.codemode/codemodeagent/page.md) - Parameters - Methods - run() - tool\_descriptions() - uses\_flyte\_tools() - \[Flyte SDK > Packages > flyte.ai.agents.memory\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/page.md) - Directory - Classes - Errors - Variables - \[Flyte SDK > Packages > flyte.ai.agents.memory > AccessDenied\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/accessdenied/page.md) - \[Flyte SDK > Packages > flyte.ai.agents.memory > ConcurrencyError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/concurrencyerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents.memory > MemoryMeta\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/memorymeta/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents.memory > MemoryStore\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/memorystore/page.md) - Parameters - Methods - append() - audit\_tail() - audit\_tail\_sync() - create() - current\_sha() - exists() - extend() - flush\_messages() - flush\_messages\_sync() - get\_meta() - get\_or\_create() - list\_paths() - read\_json() - read\_text() - remote\_path\_for\_key() - save() - write\_json() - write\_text() - \[Flyte SDK > Packages > flyte.ai.agents.memory > MemoryStoreError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/memorystoreerror/page.md) - \[Flyte SDK > Packages > flyte.ai.agents.protocol\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.protocol/page.md) - Directory - Classes - Protocols - \[Flyte SDK > Packages > flyte.ai.agents.protocol > AgentProtocol\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.protocol/agentprotocol/page.md) - Methods - run() - tool\_descriptions() - \[Flyte SDK > Packages > flyte.ai.agents.protocol > AgentResult\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.agents.protocol/agentresult/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.chat\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.chat/page.md) - Directory - Classes - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.ai.chat > AgentChatAppEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.chat/agentchatappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - build\_fastapi\_app() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.ai.chat > CustomTheme\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.chat/customtheme/page.md) - Parameters - Methods - to\_css() - \[Flyte SDK > Packages > flyte.ai.chat.app\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.chat.app/page.md) - Directory - Classes - Variables - \[Flyte SDK > Packages > flyte.ai.chat.app > AgentChatAppEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.chat.app/agentchatappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - build\_fastapi\_app() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.ai.chat.app > CustomTheme\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.chat.app/customtheme/page.md) - Parameters - Methods - to\_css() - \[Flyte SDK > Packages > flyte.ai.mcp\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.mcp/page.md) - Directory - Classes - \[Flyte SDK > Packages > flyte.ai.mcp > FlyteMCPAppEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.mcp/flytemcpappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.ai.mcp > MCPAppEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.ai.mcp/mcpappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.app\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.app > AppEndpoint\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/appendpoint/page.md) - Parameters - Methods - check\_type() - get() - materialize() - \[Flyte SDK > Packages > flyte.app > AppEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/appenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.app > ConnectorEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/connectorenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.app > Domain\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/domain/page.md) - Parameters - \[Flyte SDK > Packages > flyte.app > Link\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/link/page.md) - Parameters - \[Flyte SDK > Packages > flyte.app > Parameter\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/parameter/page.md) - Parameters - \[Flyte SDK > Packages > flyte.app > Port\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/port/page.md) - Parameters - \[Flyte SDK > Packages > flyte.app > RunOutput\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/runoutput/page.md) - Parameters - Methods - check\_type() - get() - materialize() - \[Flyte SDK > Packages > flyte.app > Scaling\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/scaling/page.md) - Parameters - Methods - get\_replicas() - \[Flyte SDK > Packages > flyte.app > Timeouts\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app/timeouts/page.md) - Parameters - \[Flyte SDK > Packages > flyte.app.extras\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app.extras/page.md) - Directory - Classes - \[Flyte SDK > Packages > flyte.app.extras > FastAPIAppEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app.extras/fastapiappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.app.extras > FastAPIPassthroughAuthMiddleware\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app.extras/fastapipassthroughauthmiddleware/page.md) - Parameters - Methods - dispatch() - extract\_authorization\_header() - extract\_cookie\_header() - extract\_custom\_header() - \[Flyte SDK > Packages > flyte.app.extras > FlyteWebhookAppEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.app.extras/flytewebhookappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.config\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.config/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.config > Config\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.config/config/page.md) - Parameters - Methods - auto() - with\_params() - \[Flyte SDK > Packages > flyte.connectors\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.connectors/page.md) - Directory - Classes - \[Flyte SDK > Packages > flyte.connectors > AsyncConnector\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.connectors/asyncconnector/page.md) - Methods - create() - delete() - get() - get\_logs() - get\_metrics() - \[Flyte SDK > Packages > flyte.connectors > AsyncConnectorExecutorMixin\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.connectors/asyncconnectorexecutormixin/page.md) - Methods - execute() - \[Flyte SDK > Packages > flyte.connectors > ConnectorRegistry\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.connectors/connectorregistry/page.md) - Methods - get\_connector() - register() - \[Flyte SDK > Packages > flyte.connectors > ConnectorService\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.connectors/connectorservice/page.md) - Methods - run() - \[Flyte SDK > Packages > flyte.connectors > Resource\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.connectors/resource/page.md) - Parameters - \[Flyte SDK > Packages > flyte.connectors > ResourceMeta\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.connectors/resourcemeta/page.md) - Parameters - Methods - decode() - encode() - \[Flyte SDK > Packages > flyte.connectors.utils\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.connectors.utils/page.md) - Directory - Methods - Methods - \[Flyte SDK > Packages > flyte.durable\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.durable/page.md) - Directory - Methods - Methods - \[Flyte SDK > Packages > flyte.errors\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/page.md) - Directory - Errors - Methods - Methods - \[Flyte SDK > Packages > flyte.errors > ActionAbortedError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/actionabortederror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > ActionNotFoundError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/actionnotfounderror/page.md) - \[Flyte SDK > Packages > flyte.errors > BaseRuntimeError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/baseruntimeerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > CodeBundleError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/codebundleerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > CustomError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/customerror/page.md) - Parameters - Methods - from\_exception() - \[Flyte SDK > Packages > flyte.errors > DeploymentError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/deploymenterror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > EventAlreadyExistsError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/eventalreadyexistserror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > EventFailedError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/eventfailederror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > EventNotFoundError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/eventnotfounderror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > EventTimedoutError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/eventtimedouterror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > ImageBuildError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/imagebuilderror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > ImagePullBackOffError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/imagepullbackofferror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > InitializationError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/initializationerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > InlineIOMaxBytesBreached\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/inlineiomaxbytesbreached/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > InvalidImageNameError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/invalidimagenameerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > InvalidPackageError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/invalidpackageerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > LogsNotYetAvailableError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/logsnotyetavailableerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > ModuleLoadError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/moduleloaderror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > NonRecoverableError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/nonrecoverableerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > NotInTaskContextError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/notintaskcontexterror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > OnlyAsyncIOSupportedError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/onlyasynciosupportederror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > OOMError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/oomerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > ParameterMaterializationError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/parametermaterializationerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > PrimaryContainerNotFoundError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/primarycontainernotfounderror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RemoteTaskNotFoundError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/remotetasknotfounderror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RemoteTaskUsageError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/remotetaskusageerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RestrictedTypeError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/restrictedtypeerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RetriesExhaustedError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/retriesexhaustederror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RuntimeDataValidationError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/runtimedatavalidationerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RuntimeSystemError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/runtimesystemerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RuntimeUnknownError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/runtimeunknownerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RuntimeUserError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/runtimeusererror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > SlowDownError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/slowdownerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > TaskInterruptedError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/taskinterruptederror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > TaskTimeoutError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/tasktimeouterror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > TraceDoesNotAllowNestedTasksError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/tracedoesnotallownestedtaskserror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > UnionRpcError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.errors/unionrpcerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.extend\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extend/page.md) - Directory - Classes - Protocols - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.extend > AsyncFunctionTaskTemplate\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extend/asyncfunctiontasktemplate/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.extend > ImageBuildEngine\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extend/imagebuildengine/page.md) - \[Flyte SDK > Packages > flyte.extend > ImageBuilder\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extend/imagebuilder/page.md) - Methods - build\_image() - get\_checkers() - \[Flyte SDK > Packages > flyte.extend > ImageChecker\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extend/imagechecker/page.md) - Methods - image\_exists() - \[Flyte SDK > Packages > flyte.extend > TaskTemplate\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extend/tasktemplate/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.extras\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras/page.md) - Directory - Classes - Protocols - \[Flyte SDK > Packages > flyte.extras > BatchStats\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras/batchstats/page.md) - Parameters - Properties - \[Flyte SDK > Packages > flyte.extras > ContainerTask\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras/containertask/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.extras > CostEstimator\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras/costestimator/page.md) - Methods - estimate\_cost() - \[Flyte SDK > Packages > flyte.extras > DynamicBatcher\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras/dynamicbatcher/page.md) - Parameters - Properties - Methods - start() - stop() - submit() - submit\_batch() - \[Flyte SDK > Packages > flyte.extras > Prompt\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras/prompt/page.md) - Parameters - Methods - estimate\_tokens() - \[Flyte SDK > Packages > flyte.extras > Sleep\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras/sleep/page.md) - Parameters - \[Flyte SDK > Packages > flyte.extras > SleepTask\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras/sleeptask/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.extras > TokenBatcher\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras/tokenbatcher/page.md) - Parameters - Properties - Methods - start() - stop() - submit() - submit\_batch() - \[Flyte SDK > Packages > flyte.extras > TokenEstimator\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras/tokenestimator/page.md) - Methods - estimate\_tokens() - \[Flyte SDK > Packages > flyte.extras.shell\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras.shell/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.extras.shell > FlagSpec\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras.shell/flagspec/page.md) - Parameters - Methods - coerce() - \[Flyte SDK > Packages > flyte.extras.shell > Glob\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras.shell/glob/page.md) - Parameters - \[Flyte SDK > Packages > flyte.extras.shell > Stderr\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras.shell/stderr/page.md) - Parameters - \[Flyte SDK > Packages > flyte.extras.shell > Stdout\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.extras.shell/stdout/page.md) - Parameters - \[Flyte SDK > Packages > flyte.git\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.git/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.git > GitStatus\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.git/gitstatus/page.md) - Parameters - Methods - build\_url() - from\_current\_repo() - \[Flyte SDK > Packages > flyte.io\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.io/page.md) - IO data types - Directory - Classes - Variables - \[Flyte SDK > Packages > flyte.io > DataFrame\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.io/dataframe/page.md) - Parameters - Properties - Methods - all() - all\_sync() - column\_names() - columns() - deserialize\_dataframe() - from\_df() - from\_existing\_remote() - from\_local() - from\_local\_sync() - iter() - model\_post\_init() - open() - schema\_match() - serialize\_dataframe() - set\_literal() - wrap\_df() - \[Flyte SDK > Packages > flyte.io > Dir\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.io/dir/page.md) - Parameters - Properties - Methods - download() - download\_sync() - empty() - exists() - exists\_sync() - from\_existing\_remote() - from\_local() - from\_local\_sync() - get\_file() - get\_file\_sync() - list\_files() - list\_files\_sync() - model\_post\_init() - new\_remote() - pre\_init() - schema\_match() - walk() - walk\_sync() - \[Flyte SDK > Packages > flyte.io > EmptyDir\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.io/emptydir/page.md) - Parameters - Properties - Methods - download() - download\_sync() - empty() - exists() - exists\_sync() - from\_existing\_remote() - from\_local() - from\_local\_sync() - get\_file() - get\_file\_sync() - list\_files() - list\_files\_sync() - model\_post\_init() - new\_remote() - pre\_init() - schema\_match() - walk() - walk\_sync() - \[Flyte SDK > Packages > flyte.io > File\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.io/file/page.md) - Parameters - Properties - Methods - download() - download\_sync() - exists() - exists\_sync() - from\_existing\_remote() - from\_local() - from\_local\_sync() - model\_post\_init() - named\_remote() - new\_remote() - open() - open\_sync() - pre\_init() - schema\_match() - \[Flyte SDK > Packages > flyte.io > HashFunction\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.io/hashfunction/page.md) - Parameters - Methods - from\_fn() - reset() - result() - update() - \[Flyte SDK > Packages > flyte.io.extend\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.io.extend/page.md) - Directory - Classes - \[Flyte SDK > Packages > flyte.io.extend > DataFrameDecoder\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.io.extend/dataframedecoder/page.md) - Parameters - Properties - Methods - decode() - \[Flyte SDK > Packages > flyte.io.extend > DataFrameEncoder\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.io.extend/dataframeencoder/page.md) - Parameters - Properties - Methods - encode() - \[Flyte SDK > Packages > flyte.io.extend > DataFrameTransformerEngine\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.io.extend/dataframetransformerengine/page.md) - Parameters - Properties - Methods - assert\_type() - encode() - from\_binary\_idl() - get\_decoder() - get\_encoder() - get\_literal\_type() - get\_structured\_dataset\_type() - guess\_python\_type() - isinstance\_generic() - iter\_as() - open\_as() - register() - register\_for\_protocol() - register\_renderer() - schema\_match() - to\_html() - to\_literal() - to\_python\_value() - \[Flyte SDK > Packages > flyte.models\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/page.md) - Directory - Classes - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.models > ActionID\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/actionid/page.md) - Parameters - Methods - create\_random() - new\_sub\_action() - new\_sub\_action\_from() - unique\_id\_str() - \[Flyte SDK > Packages > flyte.models > ActionPhase\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/actionphase/page.md) - Parameters - \[Flyte SDK > Packages > flyte.models > CheckpointPaths\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/checkpointpaths/page.md) - Parameters - \[Flyte SDK > Packages > flyte.models > CodeBundle\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/codebundle/page.md) - Parameters - Methods - with\_downloaded\_path() - \[Flyte SDK > Packages > flyte.models > GroupData\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/groupdata/page.md) - Parameters - \[Flyte SDK > Packages > flyte.models > NativeInterface\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/nativeinterface/page.md) - Parameters - Properties - Methods - convert\_to\_kwargs() - from\_callable() - from\_types() - get\_input\_types() - has\_outputs() - num\_required\_inputs() - required\_inputs() - \[Flyte SDK > Packages > flyte.models > PathRewrite\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/pathrewrite/page.md) - Parameters - Methods - from\_str() - \[Flyte SDK > Packages > flyte.models > RawDataPath\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/rawdatapath/page.md) - Parameters - Methods - from\_local\_folder() - get\_random\_remote\_path() - \[Flyte SDK > Packages > flyte.models > SerializationContext\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/serializationcontext/page.md) - Parameters - Methods - get\_entrypoint\_path() - \[Flyte SDK > Packages > flyte.models > TaskContext\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.models/taskcontext/page.md) - Parameters - Properties - Methods - is\_in\_cluster() - replace() - \[Flyte SDK > Packages > flyte.notify\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.notify/page.md) - Directory - Classes - \[Flyte SDK > Packages > flyte.notify > Email\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.notify/email/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > NamedDelivery\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.notify/nameddelivery/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > NamedRule\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.notify/namedrule/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > Notification\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.notify/notification/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > Slack\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.notify/slack/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > Teams\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.notify/teams/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > Webhook\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.notify/webhook/page.md) - Parameters - \[Flyte SDK > Packages > flyte.prefetch\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.prefetch/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.prefetch > HuggingFaceModelInfo\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.prefetch/huggingfacemodelinfo/page.md) - Parameters - \[Flyte SDK > Packages > flyte.prefetch > ShardConfig\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.prefetch/shardconfig/page.md) - Parameters - \[Flyte SDK > Packages > flyte.prefetch > StoredModelInfo\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.prefetch/storedmodelinfo/page.md) - Parameters - \[Flyte SDK > Packages > flyte.prefetch > VLLMShardArgs\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.prefetch/vllmshardargs/page.md) - Parameters - Methods - get\_vllm\_args() - \[Flyte SDK > Packages > flyte.remote\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.remote > Action\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/action/page.md) - Parameters - Properties - Methods - abort() - details() - done() - get() - get\_logs() - listall() - show\_logs() - sync() - to\_dict() - to\_json() - wait() - watch() - \[Flyte SDK > Packages > flyte.remote > ActionDetails\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/actiondetails/page.md) - Parameters - Properties - Methods - done() - get() - get\_details() - get\_phase\_transitions() - inputs() - logs\_available() - outputs() - to\_dict() - to\_json() - watch() - watch\_updates() - \[Flyte SDK > Packages > flyte.remote > ActionInputs\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/actioninputs/page.md) - Parameters - Methods - clear() - copy() - fromkeys() - get() - items() - keys() - pop() - popitem() - setdefault() - to\_dict() - to\_json() - update() - values() - \[Flyte SDK > Packages > flyte.remote > ActionOutputs\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/actionoutputs/page.md) - Parameters - Properties - Methods - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > App\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/app/page.md) - Parameters - Properties - Methods - activate() - create() - deactivate() - delete() - ephemeral\_ctx() - ephemeral\_ctx\_sync() - get() - is\_active() - is\_deactivated() - listall() - replace() - to\_dict() - to\_json() - update() - watch() - \[Flyte SDK > Packages > flyte.remote > Event\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/event/page.md) - Parameters - Properties - Methods - get() - listall() - signal() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > Project\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/project/page.md) - Parameters - Methods - archive() - create() - get() - listall() - to\_dict() - to\_json() - unarchive() - update() - \[Flyte SDK > Packages > flyte.remote > Run\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/run/page.md) - Parameters - Properties - Methods - abort() - details() - done() - get() - get\_debug\_url() - get\_logs() - inputs() - listall() - outputs() - show\_logs() - sync() - to\_dict() - to\_json() - wait() - watch() - \[Flyte SDK > Packages > flyte.remote > RunDetails\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/rundetails/page.md) - Parameters - Properties - Methods - done() - get() - get\_details() - inputs() - outputs() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > Secret\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/secret/page.md) - Parameters - Properties - Methods - create() - delete() - get() - listall() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > Settings\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/settings/page.md) - Parameters - Methods - available\_keys() - effective\_values() - get\_settings\_for\_edit() - local\_overrides() - parse\_yaml() - scope\_description() - to\_dict() - to\_json() - to\_yaml() - to\_yaml\_sections() - update\_settings() - \[Flyte SDK > Packages > flyte.remote > Task\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/task/page.md) - Parameters - Properties - Methods - get() - listall() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > TaskDetails\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/taskdetails/page.md) - Parameters - Properties - Methods - fetch() - get() - override() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > TimeFilter\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/timefilter/page.md) - Parameters - \[Flyte SDK > Packages > flyte.remote > Trigger\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/trigger/page.md) - Parameters - Properties - Methods - create() - delete() - get() - get\_details() - listall() - to\_dict() - to\_json() - update() - \[Flyte SDK > Packages > flyte.remote > User\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.remote/user/page.md) - Parameters - Methods - get() - name() - subject() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.report\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.report/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.report > Report\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.report/report/page.md) - Parameters - Methods - get\_final\_report() - get\_tab() - \[Flyte SDK > Packages > flyte.sandbox\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.sandbox/page.md) - Directory - Classes - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.sandbox > CodeTaskTemplate\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.sandbox/codetasktemplate/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.sandbox > ImageConfig\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.sandbox/imageconfig/page.md) - Parameters - \[Flyte SDK > Packages > flyte.sandbox > SandboxedConfig\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.sandbox/sandboxedconfig/page.md) - Parameters - \[Flyte SDK > Packages > flyte.sandbox > SandboxedTaskTemplate\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.sandbox/sandboxedtasktemplate/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.storage\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.storage/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.storage > ABFS\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.storage/abfs/page.md) - Parameters - Methods - auto() - get\_fsspec\_kwargs() - \[Flyte SDK > Packages > flyte.storage > GCS\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.storage/gcs/page.md) - Parameters - Methods - auto() - get\_fsspec\_kwargs() - \[Flyte SDK > Packages > flyte.storage > S3\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.storage/s3/page.md) - Parameters - Methods - auto() - for\_sandbox() - get\_fsspec\_kwargs() - \[Flyte SDK > Packages > flyte.storage > Storage\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.storage/storage/page.md) - Parameters - Methods - auto() - get\_fsspec\_kwargs() - \[Flyte SDK > Packages > flyte.syncify\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.syncify/page.md) - Creating a Syncify Instance - How does it work? - Directory - Classes - \[Flyte SDK > Packages > flyte.syncify > Syncify\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.syncify/syncify/page.md) - Parameters - \[Flyte SDK > Packages > flyte.types\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.types/page.md) - Directory - Classes - Protocols - Errors - Methods - Methods - \[Flyte SDK > Packages > flyte.types > FlytePickle\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.types/flytepickle/page.md) - Methods - from\_pickle() - python\_type() - to\_pickle() - \[Flyte SDK > Packages > flyte.types > Renderable\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.types/renderable/page.md) - Methods - to\_html() - \[Flyte SDK > Packages > flyte.types > TypeEngine\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.types/typeengine/page.md) - Methods - dict\_to\_literal\_map() - get\_available\_transformers() - get\_transformer() - guess\_python\_type() - guess\_python\_types() - lazy\_import\_transformers() - literal\_map\_to\_kwargs() - named\_tuple\_to\_variable\_map() - register() - register\_additional\_type() - register\_restricted\_type() - to\_html() - to\_literal() - to\_literal\_checks() - to\_literal\_type() - to\_python\_value() - unwrap\_offloaded\_literal() - \[Flyte SDK > Packages > flyte.types > TypeTransformer\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.types/typetransformer/page.md) - Parameters - Properties - Methods - assert\_type() - from\_binary\_idl() - get\_literal\_type() - guess\_python\_type() - isinstance\_generic() - schema\_match() - to\_html() - to\_literal() - to\_python\_value() - \[Flyte SDK > Packages > flyte.types > TypeTransformerFailedError\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte.types/typetransformerfailederror/page.md) - \[Integrations\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/page.md) - \[Integrations > Anthropic\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/anthropic/page.md) - \[Integrations > Anthropic > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/anthropic/classes/page.md) - \[Integrations > Anthropic > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/anthropic/packages/page.md) - \[Integrations > Anthropic > Packages > flyteplugins.anthropic\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/anthropic/packages/flyteplugins.anthropic/page.md) - Directory - Classes - Methods - Methods - \[Integrations > Anthropic > Packages > flyteplugins.anthropic > Agent\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/anthropic/packages/flyteplugins.anthropic/agent/page.md) - Parameters - Methods - get\_anthropic\_tools() - \[Integrations > BigQuery\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/bigquery/page.md) - \[Integrations > BigQuery > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/bigquery/classes/page.md) - \[Integrations > BigQuery > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/bigquery/packages/page.md) - \[Integrations > BigQuery > Packages > flyteplugins.bigquery\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/bigquery/packages/flyteplugins.bigquery/page.md) - Directory - Classes - \[Integrations > BigQuery > Packages > flyteplugins.bigquery > BigQueryConfig\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/bigquery/packages/flyteplugins.bigquery/bigqueryconfig/page.md) - Parameters - \[Integrations > BigQuery > Packages > flyteplugins.bigquery > BigQueryConnector\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/bigquery/packages/flyteplugins.bigquery/bigqueryconnector/page.md) - Methods - create() - delete() - get() - get\_logs() - get\_metrics() - \[Integrations > BigQuery > Packages > flyteplugins.bigquery > BigQueryTask\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/bigquery/packages/flyteplugins.bigquery/bigquerytask/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Integrations > Code generation\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/codegen/page.md) - \[Integrations > Code generation > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/codegen/classes/page.md) - \[Integrations > Code generation > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/codegen/packages/page.md) - \[Integrations > Code generation > Packages > flyteplugins.codegen\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/codegen/packages/flyteplugins.codegen/page.md) - Directory - Classes - \[Integrations > Code generation > Packages > flyteplugins.codegen > AutoCoderAgent\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/codegen/packages/flyteplugins.codegen/autocoderagent/page.md) - Parameters - Methods - generate() - \[Integrations > Code generation > Packages > flyteplugins.codegen > CodeGenEvalResult\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/codegen/packages/flyteplugins.codegen/codegenevalresult/page.md) - Parameters - Methods - as\_task() - run() - \[Integrations > Code generation > Packages > flyteplugins.codegen > CodePlan\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/codegen/packages/flyteplugins.codegen/codeplan/page.md) - Parameters - \[Integrations > Code generation > Packages > flyteplugins.codegen > CodeSolution\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/codegen/packages/flyteplugins.codegen/codesolution/page.md) - Parameters - Methods - normalize\_language() - \[Integrations > Code generation > Packages > flyteplugins.codegen > ErrorDiagnosis\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/codegen/packages/flyteplugins.codegen/errordiagnosis/page.md) - Parameters - \[Integrations > Code generation > Packages > flyteplugins.codegen > ImageConfig\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/codegen/packages/flyteplugins.codegen/imageconfig/page.md) - Parameters - \[Integrations > Dask\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/dask/page.md) - \[Integrations > Dask > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/dask/classes/page.md) - \[Integrations > Dask > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/dask/packages/page.md) - \[Integrations > Dask > Packages > flyteplugins.dask\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/dask/packages/flyteplugins.dask/page.md) - Directory - Classes - \[Integrations > Dask > Packages > flyteplugins.dask > Dask\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/dask/packages/flyteplugins.dask/dask/page.md) - Parameters - \[Integrations > Dask > Packages > flyteplugins.dask > Scheduler\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/dask/packages/flyteplugins.dask/scheduler/page.md) - Parameters - \[Integrations > Dask > Packages > flyteplugins.dask > WorkerGroup\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/dask/packages/flyteplugins.dask/workergroup/page.md) - Parameters - \[Integrations > Databricks\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/databricks/page.md) - \[Integrations > Databricks > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/databricks/classes/page.md) - \[Integrations > Databricks > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/databricks/packages/page.md) - \[Integrations > Databricks > Packages > flyteplugins.databricks\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/databricks/packages/flyteplugins.databricks/page.md) - Directory - Classes - \[Integrations > Databricks > Packages > flyteplugins.databricks > Databricks\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/databricks/packages/flyteplugins.databricks/databricks/page.md) - Parameters - \[Integrations > Databricks > Packages > flyteplugins.databricks > DatabricksConnector\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/databricks/packages/flyteplugins.databricks/databricksconnector/page.md) - Methods - create() - delete() - get() - get\_logs() - get\_metrics() - \[Integrations > Gemini\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/gemini/page.md) - \[Integrations > Gemini > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/gemini/classes/page.md) - \[Integrations > Gemini > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/gemini/packages/page.md) - \[Integrations > Gemini > Packages > flyteplugins.gemini\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/gemini/packages/flyteplugins.gemini/page.md) - Directory - Classes - Methods - Methods - \[Integrations > Gemini > Packages > flyteplugins.gemini > Agent\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/gemini/packages/flyteplugins.gemini/agent/page.md) - Parameters - Methods - get\_gemini\_tools() - \[Integrations > Human-in-the-Loop\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/hitl/page.md) - \[Integrations > Human-in-the-Loop > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/hitl/classes/page.md) - \[Integrations > Human-in-the-Loop > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/hitl/packages/page.md) - \[Integrations > Human-in-the-Loop > Packages > flyteplugins.hitl\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/hitl/packages/flyteplugins.hitl/page.md) - Basic usage: - Features: - Directory - Classes - Methods - Variables - Methods - \[Integrations > Human-in-the-Loop > Packages > flyteplugins.hitl > Event\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/hitl/packages/flyteplugins.hitl/event/page.md) - Parameters - Properties - Methods - create() - wait() - \[Integrations > Hydra\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/hydra/page.md) - \[Integrations > JSONL\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/jsonl/page.md) - \[Integrations > JSONL > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/jsonl/classes/page.md) - \[Integrations > JSONL > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/jsonl/packages/page.md) - \[Integrations > JSONL > Packages > flyteplugins.jsonl\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/jsonl/packages/flyteplugins.jsonl/page.md) - Directory - Classes - \[Integrations > JSONL > Packages > flyteplugins.jsonl > JsonlDir\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/jsonl/packages/flyteplugins.jsonl/jsonldir/page.md) - Parameters - Properties - Methods - download() - download\_sync() - empty() - exists() - exists\_sync() - from\_existing\_remote() - from\_local() - from\_local\_sync() - get\_file() - get\_file\_sync() - iter\_arrow\_batches() - iter\_arrow\_batches\_sync() - iter\_batches() - iter\_batches\_sync() - iter\_records() - iter\_records\_sync() - list\_files() - list\_files\_sync() - model\_post\_init() - new\_remote() - pre\_init() - schema\_match() - walk() - walk\_sync() - writer() - writer\_sync() - \[Integrations > JSONL > Packages > flyteplugins.jsonl > JsonlFile\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/jsonl/packages/flyteplugins.jsonl/jsonlfile/page.md) - Parameters - Properties - Methods - download() - download\_sync() - exists() - exists\_sync() - from\_existing\_remote() - from\_local() - from\_local\_sync() - iter\_arrow\_batches() - iter\_arrow\_batches\_sync() - iter\_records() - iter\_records\_sync() - model\_post\_init() - named\_remote() - new\_remote() - open() - open\_sync() - pre\_init() - schema\_match() - writer() - writer\_sync() - \[Integrations > MLflow\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/mlflow/page.md) - \[Integrations > MLflow > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/mlflow/classes/page.md) - \[Integrations > MLflow > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/mlflow/packages/page.md) - \[Integrations > MLflow > Packages > flyteplugins.mlflow\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/mlflow/packages/flyteplugins.mlflow/page.md) - Key features: - Basic usage: - Directory - Classes - Methods - Methods - \[Integrations > MLflow > Packages > flyteplugins.mlflow > Mlflow\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/mlflow/packages/flyteplugins.mlflow/mlflow/page.md) - Parameters - Methods - get\_link() - \[Integrations > OmegaConf\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/omegaconf/page.md) - \[Integrations > OpenAI\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/openai/page.md) - \[Integrations > Papermill\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/papermill/page.md) - \[Integrations > Papermill > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/papermill/classes/page.md) - \[Integrations > Papermill > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/papermill/packages/page.md) - \[Integrations > Papermill > Packages > flyteplugins.papermill\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/papermill/packages/flyteplugins.papermill/page.md) - Directory - Classes - Methods - Methods - \[Integrations > Papermill > Packages > flyteplugins.papermill > NotebookTask\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/papermill/packages/flyteplugins.papermill/notebooktask/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Integrations > Polars\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/polars/page.md) - \[Integrations > Polars > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/polars/classes/page.md) - \[Integrations > Polars > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/polars/packages/page.md) - \[Integrations > Polars > Packages > flyteplugins.polars.df\_transformer\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/polars/packages/flyteplugins.polars.df\_transformer/page.md) - Directory - Classes - Methods - Variables - Methods - \[Integrations > Polars > Packages > flyteplugins.polars.df\_transformer > ParquetToPolarsDecodingHandler\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/polars/packages/flyteplugins.polars.df\_transformer/parquettopolarsdecodinghandler/page.md) - Parameters - Properties - Methods - decode() - \[Integrations > Polars > Packages > flyteplugins.polars.df\_transformer > ParquetToPolarsLazyFrameDecodingHandler\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/polars/packages/flyteplugins.polars.df\_transformer/parquettopolarslazyframedecodinghandler/page.md) - Parameters - Properties - Methods - decode() - \[Integrations > Polars > Packages > flyteplugins.polars.df\_transformer > PolarsLazyFrameToParquetEncodingHandler\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/polars/packages/flyteplugins.polars.df\_transformer/polarslazyframetoparquetencodinghandler/page.md) - Parameters - Properties - Methods - encode() - \[Integrations > Polars > Packages > flyteplugins.polars.df\_transformer > PolarsToParquetEncodingHandler\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/polars/packages/flyteplugins.polars.df\_transformer/polarstoparquetencodinghandler/page.md) - Parameters - Properties - Methods - encode() - \[Integrations > PyTorch\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/pytorch/page.md) - \[Integrations > PyTorch > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/pytorch/classes/page.md) - \[Integrations > PyTorch > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/pytorch/packages/page.md) - \[Integrations > PyTorch > Packages > flyteplugins.pytorch\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/pytorch/packages/flyteplugins.pytorch/page.md) - Directory - Classes - \[Integrations > PyTorch > Packages > flyteplugins.pytorch > Elastic\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/pytorch/packages/flyteplugins.pytorch/elastic/page.md) - Parameters - \[Integrations > Ray\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/ray/page.md) - \[Integrations > Ray > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/ray/classes/page.md) - \[Integrations > Ray > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/ray/packages/page.md) - \[Integrations > Ray > Packages > flyteplugins.ray\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/ray/packages/flyteplugins.ray/page.md) - Directory - Classes - \[Integrations > Ray > Packages > flyteplugins.ray > HeadNodeConfig\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/ray/packages/flyteplugins.ray/headnodeconfig/page.md) - Parameters - \[Integrations > Ray > Packages > flyteplugins.ray > RayJobConfig\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/ray/packages/flyteplugins.ray/rayjobconfig/page.md) - Parameters - \[Integrations > Ray > Packages > flyteplugins.ray > WorkerNodeConfig\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/ray/packages/flyteplugins.ray/workernodeconfig/page.md) - Parameters - \[Integrations > SGLang\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/sglang/page.md) - \[Integrations > SGLang > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/sglang/classes/page.md) - \[Integrations > SGLang > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/sglang/packages/page.md) - \[Integrations > SGLang > Packages > flyteplugins.sglang\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/sglang/packages/flyteplugins.sglang/page.md) - Directory - Classes - Variables - \[Integrations > SGLang > Packages > flyteplugins.sglang > SGLangAppEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/sglang/packages/flyteplugins.sglang/sglangappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Integrations > Snowflake\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/snowflake/page.md) - \[Integrations > Snowflake > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/snowflake/classes/page.md) - \[Integrations > Snowflake > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/snowflake/packages/page.md) - \[Integrations > Snowflake > Packages > flyteplugins.snowflake\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/snowflake/packages/flyteplugins.snowflake/page.md) - Directory - Classes - \[Integrations > Snowflake > Packages > flyteplugins.snowflake > Snowflake\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/snowflake/packages/flyteplugins.snowflake/snowflake/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Integrations > Snowflake > Packages > flyteplugins.snowflake > SnowflakeConfig\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/snowflake/packages/flyteplugins.snowflake/snowflakeconfig/page.md) - Parameters - \[Integrations > Snowflake > Packages > flyteplugins.snowflake > SnowflakeConnector\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/snowflake/packages/flyteplugins.snowflake/snowflakeconnector/page.md) - Methods - create() - delete() - get() - get\_logs() - get\_metrics() - \[Integrations > Spark\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/spark/page.md) - \[Integrations > Spark > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/spark/classes/page.md) - \[Integrations > Spark > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/spark/packages/page.md) - \[Integrations > Spark > Packages > flyteplugins.spark\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/spark/packages/flyteplugins.spark/page.md) - Directory - Classes - \[Integrations > Spark > Packages > flyteplugins.spark > ParquetToSparkDecoder\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/spark/packages/flyteplugins.spark/parquettosparkdecoder/page.md) - Parameters - Properties - Methods - decode() - \[Integrations > Spark > Packages > flyteplugins.spark > Spark\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/spark/packages/flyteplugins.spark/spark/page.md) - Parameters - \[Integrations > Spark > Packages > flyteplugins.spark > SparkToParquetEncoder\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/spark/packages/flyteplugins.spark/sparktoparquetencoder/page.md) - Parameters - Properties - Methods - encode() - \[Integrations > Union\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/page.md) - \[Integrations > Union > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/classes/page.md) - \[Integrations > Union > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/page.md) - \[Integrations > Union > Packages > flyteplugins.union.cli\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.cli/page.md) - Directory - Methods - Methods - \[Integrations > Union > Packages > flyteplugins.union.cli.queue\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.cli.queue/page.md) - Directory - Variables - \[Integrations > Union > Packages > flyteplugins.union.internal.validate.validate.validate\_pb2\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.internal.validate.validate.validate\_pb2/page.md) - Directory - Variables - \[Integrations > Union > Packages > flyteplugins.union.remote\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.remote/page.md) - Directory - Classes - \[Integrations > Union > Packages > flyteplugins.union.remote > ApiKey\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.remote/apikey/page.md) - Parameters - Properties - Methods - create() - delete() - get() - listall() - to\_dict() - to\_json() - update() - \[Integrations > Union > Packages > flyteplugins.union.remote > Assignment\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.remote/assignment/page.md) - Parameters - Properties - Methods - create() - get() - listall() - to\_dict() - to\_json() - unassign() - \[Integrations > Union > Packages > flyteplugins.union.remote > Cluster\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.remote/cluster/page.md) - Parameters - Properties - Methods - get() - listall() - to\_dict() - to\_json() - \[Integrations > Union > Packages > flyteplugins.union.remote > Member\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.remote/member/page.md) - Parameters - Properties - Methods - listall() - to\_dict() - to\_json() - \[Integrations > Union > Packages > flyteplugins.union.remote > Policy\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.remote/policy/page.md) - Parameters - Properties - Methods - create() - delete() - get() - listall() - to\_dict() - to\_json() - update() - \[Integrations > Union > Packages > flyteplugins.union.remote > Queue\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.remote/queue/page.md) - Parameters - Properties - Methods - activate() - create() - details() - drain() - get() - listall() - to\_dict() - to\_json() - update() - watch() - \[Integrations > Union > Packages > flyteplugins.union.remote > Role\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.remote/role/page.md) - Parameters - Properties - Methods - create() - delete() - get() - listall() - to\_dict() - to\_json() - update() - \[Integrations > Union > Packages > flyteplugins.union.remote > User\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.remote/user/page.md) - Parameters - Properties - Methods - create() - delete() - get() - listall() - to\_dict() - to\_json() - \[Integrations > Union > Packages > flyteplugins.union.utils.auth\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.utils.auth/page.md) - Directory - Classes - Methods - Methods - \[Integrations > Union > Packages > flyteplugins.union.utils.auth > AppClientCredentials\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/union/packages/flyteplugins.union.utils.auth/appclientcredentials/page.md) - Parameters - \[Integrations > vLLM\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/vllm/page.md) - \[Integrations > vLLM > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/vllm/classes/page.md) - \[Integrations > vLLM > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/vllm/packages/page.md) - \[Integrations > vLLM > Packages > flyteplugins.vllm\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/vllm/packages/flyteplugins.vllm/page.md) - Directory - Classes - Variables - \[Integrations > vLLM > Packages > flyteplugins.vllm > VLLMAppEnvironment\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/vllm/packages/flyteplugins.vllm/vllmappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Integrations > Weights & Biases\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/wandb/page.md) - \[Integrations > Weights & Biases > Classes\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/wandb/classes/page.md) - \[Integrations > Weights & Biases > Packages\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/wandb/packages/page.md) - \[Integrations > Weights & Biases > Packages > flyteplugins.wandb\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/wandb/packages/flyteplugins.wandb/page.md) - Key features: - Basic usage: - Directory - Classes - Methods - Methods - \[Integrations > Weights & Biases > Packages > flyteplugins.wandb > Wandb\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/wandb/packages/flyteplugins.wandb/wandb/page.md) - Parameters - Methods - get\_link() - \[Integrations > Weights & Biases > Packages > flyteplugins.wandb > WandbSweep\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/wandb/packages/flyteplugins.wandb/wandbsweep/page.md) - Parameters - Methods - get\_link() --- ## Community - \[Joining the community\](https://www.union.ai/docs/v2/flyte/community/joining-the-community/page.md) - Community sync - Contributor's sync - Newsletter - Slack guidelines - Abide by the \[LF's Code of Conduct\](https://lfprojects.org/policies/code-of-conduct/) - Avoid using DMs and @mentions - Make use of threads - Do not post the same question across multiple channels - Do not solicit members of our Slack - \[Contributing code\](https://www.union.ai/docs/v2/flyte/community/contributing-code/page.md) - Flyte 2 - Becoming a contributor - Before submitting your PR - 🐞 File an issue - Component Reference - \`flyte\` - \`flyteidl\` - \`flytepropeller\` - \`flyteadmin\` - \`flytekit\` - \`flyteconsole\` - \`datacatalog\` - \`flyteplugins\` - \`flytestdlib\` - \`flytectl\` - Development Environment Setup Guide - Requirements - Content - How to setup dev environment for flyteidl, flyteadmin, flyteplugins, flytepropeller, datacatalog and flytestdlib? - How to setup dev environment for flytekit? - How to setup dev environment for flyteconsole? - How to access Flyte UI, minio, postgres, k3s, and endpoints? - \[Contributing docs and examples\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/page.md) - The combined Flyte and Union docs site - Versions - Common build infrastructure - Variants - Both Flyte and Union docs are open source > Section bundle (all pages): https://www.union.ai/docs/v2/flyte/community/contributing-docs/section.md - \[Contributing docs and examples > Quick start\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/quick-start/page.md) - Prerequisites - Clone the repository - Live preview - Distribution build - \[Contributing docs and examples > Variants\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/variants/page.md) - Variants at the whole-page level - Conditional rendering within a page - {{}} - {{}} - Full example - Adding a new variant - Location - Creating a new variant - Testing the new variant - Building (just) the variant - \[Contributing docs and examples > Versions\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/versions/page.md) - Versions are branches - How to create an archive version - How to create an archive version - Publishing an archive version - \[Contributing docs and examples > Authoring\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/authoring/page.md) - Getting started - Target the right branch - Live preview - Pull Requests + Site Preview - Page Visibility - Page order - Page settings - Conditional Content - Linking to the API reference - Sigils for special cases - Warnings and Notices - Special Content Generation - Python Generated Content - Run on Union Instructions - Jupyter Notebooks - Mapped Keys (\`{{}}\`) - Mermaid Graphs - \[Contributing docs and examples > Shortcodes\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/shortcodes/page.md) - How to specify a "shortcode" - Variants - Component Library - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` and \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\`, \`{{}}\`, and \`{{}}\` - \`{{}}\` - \`{{}}\` - \[Contributing docs and examples > Redirects\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/redirects/page.md) - \`docs.union.ai\` redirects - \`docs.flyte.org\` redirects - \[Contributing docs and examples > API docs\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/api-docs/page.md) - API naming convention - Package Resource Resolution - Tips and Tricks - Auto-linking - Short vs. fully-qualified names - How auto-linking works - Magic-marker syntax for inline code - \[Contributing docs and examples > LLM-optimized documentation\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/llm-docs/page.md) - Output files - Discovery hierarchy - How \`page.md\` files are generated - Enabling section bundles - The \`llms-full.txt\` link conversion - Regenerating - \[Contributing docs and examples > Publishing\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/publishing/page.md) - Requirements - Managing the Tutorial Pages - Building and running locally - Developer Experience - Controlling Development Environment - Changing 'variants' - Troubleshootting - Identifying Problems: Missing Content - Identifying Problems: Page Visibility - Building Production - Testing Production Build --- # Unknown \# Flyte OSS Flyte is a free and open source platform that provides a full suite of powerful features for orchestrating AI workflows. Flyte empowers AI development teams to rapidly ship high-quality code to production by offering optimized performance, unparalleled resource efficiency, and a delightful workflow authoring experience. You deploy and manage Flyte yourself, on your own cloud infrastructure. > \[!NOTE\] > These are the Flyte \*\*2.0\*\* docs. > To switch to \[version 1.0\](https://www.union.ai/docs/v1/flyte/) or to the commercial product, \[\*\*Union.ai\*\*\](https://www.union.ai/docs/v2/union/), use the selectors above. ## Basics Learn the basics of Flyte, covering all the core concepts around tasks and apps. ### \[Flyte 2\](https://www.union.ai/docs/v2/flyte/user-guide/overview/page.md) Build AI workflows in pure Python with built-in durability, reproducibility, and recovery. ### \[Quickstart\](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/page.md) Install the SDK and run your first workflow locally in a few minutes. ### \[Core concepts\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/page.md) The building blocks of every Flyte program: TaskEnvironments, tasks, runs, actions, and apps. ### \[Run modes\](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/page.md) Run the same task code locally, on a devbox, or on a remote cluster. ## Tasks Build durable, scalable, and reproducible batch workloads. ### \[Configure tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/page.md) Define \`TaskEnvironment\`s for container images, resources, secrets, caching, retries, and more; use triggers for schedules. ### \[Build tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/page.md) Compose tasks with fanout, parallelism, error handling, traces, files, and DataFrames. ### \[Run and deploy tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/page.md) Use \`flyte run\` for iteration or \`flyte deploy\` to register a stable task version. ## Apps Create long-running services to host dashboards, APIs, and model endpoints. ### \[Configure apps\](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/page.md) Define \`AppEnvironment\`s with ports, autoscaling, custom domains, and authentication. ### \[Build apps\](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/page.md) Build dashboards, REST APIs, and model endpoints with FastAPI, Streamlit, vLLM, and more. ### \[Native app integrations\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/page.md) Use pre-built environments for popular frameworks like Streamlit, FastAPI, vLLM, and SGLang. ### \[Serve and deploy apps\](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/page.md) Use \`flyte serve\` for fast iteration or \`flyte deploy\` for production deployments. ## Agents Build durable, self-healing agents using tasks and apps as building blocks. ### \[Build agents\](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/page.md) Implement ReAct, Plan-and-Execute, and other agent patterns with full observability. ### \[Agent framework integrations\](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/page.md) Integrate with third-party agent frameworks like LangGraph, PydanticAI, and OpenAI Agents SDK. ### \[Sandboxing\](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/page.md) Safely execute LLM-generated code with workflow sandboxes or ephemeral containers. ### \[Build an MCP\](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/page.md) Serve Model Context Protocol servers for AI assistants to interact with, hosted on Flyte. ## Advanced Guides Organize your codebase, optimize performance for production, and migrate from other workflow orchestrators. ### \[Project patterns\](https://www.union.ai/docs/v2/flyte/user-guide/project-patterns/page.md) Patterns for BYO images, monorepos with uv, CI/CD, and multi-team resource management. ### \[Run scaling\](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/page.md) Tune task overhead, batching, reusable containers, and fanout to scale your workflows. ### \[Advanced project\](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/page.md) An advanced guide for building an LLM reporting agent on Flyte. ### \[Migration\](https://www.union.ai/docs/v2/flyte/user-guide/migration/page.md) Port a Flyte 1 codebase to Flyte 2, or map Airflow concepts to their Flyte 2 equivalents. ## Subpages - \[Overview\](https://www.union.ai/docs/v2/flyte/user-guide/overview/page.md) - Pure Python, no DSL - Durability - Reproducibility - Recoverability - Built for scale - What this means in practice - \[Quickstart\](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/page.md) - What you'll need - Install the SDK - Configure - Write your first workflow - Run it - See the results - Next steps - \[Core concepts\](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/page.md) - How Flyte works - \[Run modes\](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/page.md) - \[Local\](https://www.union.ai/docs/v2/flyte/user-guide/running-locally/page.md) - \[Devbox\](https://www.union.ai/docs/v2/flyte/user-guide/running-devbox/page.md) - \[Configure tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/page.md) - Task configuration levels - Example - Task configuration parameters - \[Build tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/page.md) - What you'll learn - When to use these patterns - \[Run and deploy tasks\](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/page.md) - Ephemeral deployment and immediate execution - Programmatic - CLI - Persistent deployment - Programmatic - CLI - Running already deployed tasks - Programmatic - CLI - Configuring runs with \`flyte.with\_runcontext()\` - \[Configure apps\](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/page.md) - Hello World example - Using fserve args - Using @app\_env.server - Differences from TaskEnvironment - Configuration topics - \[Build apps\](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/page.md) - App types - Usage patterns - Next steps - \[Native app integrations\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/page.md) - When to use a native integration - Available integrations - Next steps - \[Serve and deploy apps\](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/page.md) - Serve vs Deploy - \`flyte serve\` - \`flyte deploy\` - Using Python SDK - Serve - Deploy - Using the CLI - Serve - Deploy - Next steps - \[Build an agent\](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/page.md) - How Flyte maps to the agentic world - Ways to build an agent - Deploying an agent - Related - \[Agent framework integrations\](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/page.md) - How much control does the framework give you? - Supported frameworks - Next steps - \[Build an MCP\](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/page.md) - HTTP layout - Quickstart - \[Sandboxing\](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/page.md) - Why sandboxing matters for AI - Types of sandboxes - What Flyte offers - Workflow sandbox (Monty) - Code sandbox (container) - When to use which - Learn more - \[Project patterns\](https://www.union.ai/docs/v2/flyte/user-guide/project-patterns/page.md) - \[Bring your own image (BYOI)\](https://www.union.ai/docs/v2/flyte/user-guide/bring-your-own-image/page.md) - \[Monorepo with uv\](https://www.union.ai/docs/v2/flyte/user-guide/monorepo-with-uv/page.md) - \[CI/CD deployments\](https://www.union.ai/docs/v2/flyte/user-guide/cicd) - \[Resource management and multi-team scaling\](https://www.union.ai/docs/v2/flyte/user-guide/resource-management) - \[Scale your runs\](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/page.md) - Understanding Flyte execution - Performance optimization - Key concepts for scaling - \[Advanced project\](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/page.md) - What you'll build - Concepts covered - Architecture - Prerequisites - Parts - Key takeaways - \[Migration\](https://www.union.ai/docs/v2/flyte/user-guide/migration/page.md) - \[From Flyte 1 to 2\](https://www.union.ai/docs/v2/flyte/user-guide/flyte-2/page.md) - \[From Airflow to Flyte\](https://www.union.ai/docs/v2/flyte/user-guide/from-airflow/page.md) --- \*\*Source\*\*: https://github.com/unionai/unionai-docs/blob/main/content/user-guide/\_index.md \*\*HTML\*\*: https://www.union.ai/docs/v2/flyte/user-guide/ --- # Overview | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Overview ======== In this guide we cover how to build AI applications, data pipelines, and ML workflows using the Flyte 2 SDK. Programs written using the Flyte 2 SDK can run on either a Union.ai or Flyte OSS back-end. This guide applies to both. [Pure Python, no DSL](https://www.union.ai/docs/v2/flyte/user-guide/overview/#pure-python-no-dsl) --------------------------------------------------------------------------------------------------- Flyte lets you write workflows in standard Python—no domain-specific language, no special syntax, no restrictions. Your “workflow” is simply a task that calls other tasks: @env.task() async def my_workflow(data: list[str]) -> list[str]: results = [] for item in data: if should_process(item): result = await process_item(item) results.append(result) return results You can use everything Python offers: * **Loops and conditionals** — standard `for`, `while`, `if-elif-else` * **Error handling** — `try/except` blocks work as expected * **Async/await** — native Python concurrency model * **Any library** — import and use whatever you need This means no learning curve beyond Python itself, and no fighting a DSL when your requirements don’t fit its constraints. [Durability](https://www.union.ai/docs/v2/flyte/user-guide/overview/#durability) ---------------------------------------------------------------------------------- Every task execution in Flyte is automatically persisted. Inputs, outputs, and intermediate results are stored in an object store, giving you: * **Full observability** — see exactly what data flowed through each step * **Audit trail** — track what ran, when, and with what parameters * **Data lineage** — trace outputs back to their inputs This persistence happens automatically. You don’t need to add logging or manually save state—Flyte handles it. [Reproducibility](https://www.union.ai/docs/v2/flyte/user-guide/overview/#reproducibility) -------------------------------------------------------------------------------------------- Flyte ensures that runs can be reproduced exactly: * **Deterministic execution** — same inputs produce same outputs * **Caching** — task results are cached and reused when inputs match * **Versioned containers** — code runs in the same environment every time Caching is configurable per task: @env.task(cache="auto") async def expensive_computation(data: str) -> str: # This result will be cached and reused for identical inputs ... When you rerun a workflow, Flyte serves cached results for unchanged tasks rather than recomputing them. [Recoverability](https://www.union.ai/docs/v2/flyte/user-guide/overview/#recoverability) ------------------------------------------------------------------------------------------ When something fails, Flyte doesn’t make you start over. Failed workflows can resume from where they left off: * **Completed tasks are preserved** — successful outputs remain cached * **Retry from failure point** — no need to re-execute what already succeeded * **Fine-grained checkpoints** — the `@flyte.trace` decorator creates checkpoints within tasks This reduces wasted compute and speeds up debugging. When a task fails after hours of prior computation, you fix the issue and continue—not restart. [Built for scale](https://www.union.ai/docs/v2/flyte/user-guide/overview/#built-for-scale) -------------------------------------------------------------------------------------------- Flyte handles the hard parts of distributed execution: * **Parallel execution** — express parallelism with `asyncio.gather()`, Flyte handles the rest * **Dynamic workflows** — construct workflows based on runtime data, not just static definitions * **Fast scheduling** — reusable containers achieve millisecond-level task startup * **Resource management** — specify CPU, memory, and GPU requirements per task [What this means in practice](https://www.union.ai/docs/v2/flyte/user-guide/overview/#what-this-means-in-practice) -------------------------------------------------------------------------------------------------------------------- Consider a data pipeline that processes thousands of files, trains a model, and deploys it: * If file processing fails on item 847, you fix the issue and resume from item 847 * If training succeeds, but deployment fails, you redeploy without retraining * If you rerun next week with the same data, cached results skip redundant computation * If you need to audit what happened, every step is recorded Flyte gives you the flexibility of Python scripts with the reliability of a production system. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/overview/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Build tasks | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Build tasks =========== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section covers the essential programming patterns and techniques for developing robust Flyte workflows. Once you understand the basics of task configuration, these guides will help you build sophisticated, production-ready data pipelines and machine learning workflows. [What you’ll learn](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/#what-youll-learn) ------------------------------------------------------------------------------------------------------- The task programming section covers key patterns for building effective Flyte workflows: **Data handling and types** * [**Files and directories**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/files-and-directories) : Work with large datasets using Flyte’s efficient file and directory types that automatically handle data upload, storage, and transfer between tasks. * [**DataFrames**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/dataframes) : Pass DataFrames between tasks without downloading data into memory, with support for Pandas, Polars, PyArrow, Dask, and other DataFrame backends. * [**Data classes and structures**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/dataclasses-and-structures) : Use Python data classes and Pydantic models as task inputs and outputs to create well-structured, type-safe workflows. * [**Custom context**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/custom-context) : Use custom context to pass metadata through your task execution hierarchy without adding parameters to every task. **Execution patterns** * [**Fanout**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/fanout) : Scale your workflows by running many tasks in parallel, perfect for processing large datasets or running hyperparameter sweeps. * [**Controlling parallel execution**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/controlling-parallelism) : Limit concurrent task executions using semaphores or `flyte.map` concurrency for rate-limited APIs, GPU quotas, and resource-constrained workflows. * [**Human-in-the-loop**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/human-in-the-loop) : Pause workflow execution at a checkpoint and wait for a human to provide input or approval before continuing. * [**Grouping actions**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/grouping-actions) : Organize related task executions into logical groups for better visualization and management in the UI. * [**Container tasks**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/container-tasks) : Run arbitrary containers in any language without the Flyte SDK installed, using Flyte’s copilot sidecar for seamless data flow. * [**Remote tasks**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/remote-tasks) : Use previously deployed tasks without importing their code or dependencies, enabling team collaboration and task reuse. * [**Pod templates**](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/pod-templates) : Extend tasks with Kubernetes pod templates to add sidecars, volume mounts, and advanced Kubernetes configurations. * [**Abort and cancel actions**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/abort-tasks) : Stop in-progress actions automatically, programmatically, or manually via the CLI and UI. * [**Other features**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/other-features) : Advanced patterns like task forwarding and other specialized task execution techniques. **Development and debugging** * [**Notebooks**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/notebooks) : Write and iterate on workflows directly in Jupyter notebooks for interactive development and experimentation. * [**Unit testing**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/unit-testing) : Test your Flyte tasks using direct invocation for business logic or `flyte.run()` for Flyte-specific features. * [**Links**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/links) : Add clickable URLs to tasks in the Flyte UI, connecting them to external tools like experiment trackers and monitoring dashboards. * [**Reports**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/reports) : Generate custom HTML reports during task execution to display progress, results, and visualizations in the UI. * [**Traces**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/traces) : Add fine-grained observability to helper functions within your tasks for better debugging and resumption capabilities. * [**Error handling**](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/error-handling) : Implement robust error recovery strategies, including automatic resource scaling and graceful failure handling. [When to use these patterns](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/#when-to-use-these-patterns) -------------------------------------------------------------------------------------------------------------------------- These programming patterns become essential as your workflows grow in complexity: * Use **fanout** when you need to process multiple items concurrently or run parameter sweeps. Use **controlling parallel execution** when you need to limit how many run at the same time. * Implement **error handling** for production workflows that need to recover from infrastructure failures. * Apply **grouping** to organize complex workflows with many task executions. * Leverage **files and directories** when working with large datasets that don’t fit in memory. * Use **DataFrames** to efficiently pass tabular data between tasks across different processing engines. * Choose **container tasks** when you need to run code in non-Python languages, use legacy containers, or execute AI-generated code in sandboxes. * Use **remote tasks** to reuse tasks deployed by other teams without managing their dependencies. * Apply **pod templates** when you need advanced Kubernetes features like sidecars or specialized storage configurations. * Use **traces** to debug non-deterministic operations like API calls or ML inference. * Use **links** to connect tasks to external tools like Weights & Biases, Grafana, or custom dashboards directly from the Flyte UI. * Create **reports** to monitor long-running workflows and share results with stakeholders. * Use **custom context** when you need lightweight, cross-cutting metadata to flow through your task hierarchy without becoming part of the task’s logical inputs. * Write **unit tests** to validate your task logic and ensure type transformations work correctly before deployment. * Use **abort and cancel** to stop unnecessary actions when conditions change, such as early convergence in HPO or manual intervention. * Use **human-in-the-loop** to insert approval gates or data collection checkpoints into automated workflows. Each guide includes practical examples and best practices to help you implement these patterns effectively in your own workflows. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Build apps | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Build apps ========== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section covers how to build different types of apps with Flyte, from single-script apps to multi-file projects, common usage patterns, and authentication. Go to [Introducing apps](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/introducing-apps) for an overview of apps and a quick example. For pre-built environments for popular frameworks like Streamlit, FastAPI, vLLM, and SGLang, see [Native app integrations](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations) . [App types](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/#app-types) ---------------------------------------------------------------------------------- Flyte supports various types of apps: * **UI dashboard apps**: Interactive web dashboards and data visualization tools like Streamlit and Gradio * **Web API apps**: REST APIs, webhooks, and backend services like FastAPI and Flask * **Model serving apps**: High-performance LLM serving with vLLM and SGLang For ready-to-use environments for these frameworks, see [Native app integrations](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations) . [Usage patterns](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/#usage-patterns) -------------------------------------------------------------------------------------------- Apps and tasks can interact in various ways: calling each other via HTTP, webhooks, WebSockets, or direct browser usage. | Pattern | Use Case | Implementation | | --- | --- | --- | | App | Stand-alone serving app | HTTP requests from arbitrary clients | | App → App | Microservices, proxies, agent routers, LLM routers | HTTP requests between apps | | App → Task | Webhooks, APIs triggering workflows | Flyte SDK in app | | Task → App | Batch processing using inference services | HTTP requests from task | | Browser app | User-facing dashboards (e.g. Streamlit, Gradio) | Direct browser access | [Next steps](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/#next-steps) ------------------------------------------------------------------------------------ * [**Single-script apps**](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/single-script-apps) : The simplest way to build and deploy apps in a single Python script * [**Multi-script apps**](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/multi-script-apps) : Build FastAPI and Streamlit apps with multiple files * [**Serving graphs**](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/serving-graphs) : Apps calling other apps for microservice architectures * [**Hybrid graphs**](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/hybrid-graphs) : Tasks calling apps and apps calling tasks (webhooks, APIs) * [**WebSocket apps**](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/websocket-apps) : Real-time, bidirectional communication with WebSockets * [**Browser apps**](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/browser-apps) : User-facing dashboards and UIs * [**Secret-based authentication**](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/secret-based-authentication) : Authenticate FastAPI apps using Flyte secrets LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Configure apps | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Configure apps ============== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. `[[AppEnvironment]]`s allows you to configure the environment in which your app runs, including the container image, compute resources, secrets, domains, scaling behavior, and more. Similar to `[[TaskEnvironment]]`, configuration can be set when creating the `[[AppEnvironment]]` object. Unlike tasks, apps are long-running services, so they have additional configuration options specific to web services: * `port`: What port the app listens on * `command` and `args`: How to start the app * `scaling`: Autoscaling configuration for handling variable load * `domain`: Custom domains and subdomains for your app * `requires_auth`: Whether the app requires authentication to access * `depends_on`: Other app or task environments that the app depends on [Hello World example](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/#hello-world-example) ---------------------------------------------------------------------------------------------------------- Here’s a complete example of deploying a simple Streamlit “hello world” app with a custom subdomain. There are two ways to build apps in Flyte: 1. Defining `AppEnvironment(.., args=[...])` to run the app with the underlying `fserve` command. 2. Defining `@app_env.server` to run the app with a custom server function. Using fserve argsUsing @app\_env.server hello-world-app.py [](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/hello-world-app.py "View source on GitHub") # /// script # requires-python = "==3.13" # dependencies = [\ # "flyte>=2.0.0b52",\ # ] # /// import flyte import flyte.app image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("streamlit==1.41.1") app_env = flyte.app.AppEnvironment( name="hello-world-app", image=image, args=["streamlit", "hello", "--server.port", "8080"], port=8080, resources=flyte.Resources(cpu="1", memory="1Gi"), requires_auth=False, domain=flyte.app.Domain(subdomain="hello"), ) if __name__ == "__main__": flyte.init_from_config() # Deploy the app app = flyte.serve(app_env) print(f"App served at: {app.url}") This example demonstrates: * Creating a custom Docker image with Streamlit * Setting the `args` to run the Streamlit hello app, which uses the underlying `fserve` command to run the app. * Configuring the port * Setting resource limits * Disabling authentication (for public access) * Using a custom subdomain hello-world-app-server.py [](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/configure-apps/hello-world-app-server.py "View source on GitHub") # /// script # requires-python = "==3.13" # dependencies = [\ # "flyte>=2.0.0b52",\ # ] # /// import flyte import flyte.app image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("streamlit==1.41.1") app_env = flyte.app.AppEnvironment( name="hello-world-app-server", image=image, port=8080, resources=flyte.Resources(cpu="1", memory="1Gi"), requires_auth=False, domain=flyte.app.Domain(subdomain="hello-server"), ) @app_env.server def server(): import subprocess subprocess.run(["streamlit", "hello", "--server.port", "8080"], check=False) if __name__ == "__main__": flyte.init_from_config() # Deploy the app app = flyte.serve(app_env) print(f"App served at: {app.url}") This example demonstrates: * Creating a custom Docker image with Streamlit * Using the `@app_env.server` decorator to define a server function that runs the Streamlit hello app. * Configuring the port * Setting resource limits * Disabling authentication (for public access) * Using a custom subdomain Once deployed, your app will be accessible at the generated URL or your custom subdomain. [Differences from TaskEnvironment](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/#differences-from-taskenvironment) ------------------------------------------------------------------------------------------------------------------------------------ While `AppEnvironment` inherits from `Environment` (the same base class as `TaskEnvironment`), it has several app-specific parameters: | Parameter | AppEnvironment | TaskEnvironment | Description | | --- | --- | --- | --- | | `type` | ✅ | ❌ | Type of app (e.g., “FastAPI”, “Streamlit”) | | `port` | ✅ | ❌ | Port the app listens on | | `args` | ✅ | ❌ | Arguments to pass to the app | | `command` | ✅ | ❌ | Command to run the app | | `requires_auth` | ✅ | ❌ | Whether app requires authentication | | `scaling` | ✅ | ❌ | Autoscaling configuration | | `domain` | ✅ | ❌ | Custom domain/subdomain | | `links` | ✅ | ❌ | Links to include in the App UI page | | `include` | ✅ | ❌ | Files to include in app | | `parameters` | ✅ | ❌ | Parameters to pass to app | | `cluster_pool` | ✅ | ❌ | Cluster pool for deployment | Parameters like `image`, `resources`, `secrets`, `env_vars`, and `depends_on` are shared between both environment types. See the [task configuration](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration) docs for details on these shared parameters. [Configuration topics](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/#configuration-topics) ------------------------------------------------------------------------------------------------------------ Learn more about configuring apps: * [**Environment settings**](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/app-environment-settings) : Images, resources, secrets, and app-specific settings like `type`, `port`, `args`, `requires_auth` * [**App startup**](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/app-environment-settings#app-startup) : Understanding the difference between `args` and `command` * [**Including additional files**](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/including-additional-files) : How to include additional files needed by your app * [**App parameters**](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/passing-parameters) : Pass parameters to your app at deployment time * [**Autoscaling apps**](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/auto-scaling-apps) : Configure scaling up and down based on traffic with idle TTL * [**App depending on other environments**](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/apps-depending-on-environments) : Use `depends_on` to deploy dependent apps together LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Tutorials | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Tutorials ========= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/tutorials/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section contains tutorials that showcase relevant use cases and provide step-by-step instructions on how to implement various features using Flyte and Union. Tutorials are organized by **industry vertical** and by **technical topic**. [Industry verticals](https://www.union.ai/docs/v2/flyte/tutorials/#industry-verticals) ---------------------------------------------------------------------------------------- Biotech & Healthcare Bioinformatics, medical imaging, and other life-sciences workloads. Geospatial Satellite imagery, remote sensing, and earth and atmospheric modeling workloads. Financial Services & Fintech Financial research, trading, and other fintech workloads. Frontier AI Frontier-model pretraining, automated experimentation, and large-scale AI workloads. [Technical topics](https://www.union.ai/docs/v2/flyte/tutorials/#technical-topics) ------------------------------------------------------------------------------------ Computer Vision Image and vision-language model workloads. Agents Agentic workflows and autonomous LLM-powered systems. Context Engineering Prompt engineering, prompt optimization, and context construction. Model Training Training, fine-tuning, and hyperparameter optimization of models at scale. Data Processing Large-scale data processing and batching strategies. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/tutorials/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/tutorials/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Home | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Run and deploy tasks | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Run and deploy tasks ==================== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. You have seen how to configure and build the tasks that compose your project. Now you need to decide how to execute them on your Flyte backend. Flyte offers two distinct approaches for getting your tasks onto the backend: **Use `flyte run` when you’re iterating and experimenting:** * Quickly test changes during development * Try different parameters or code modifications * Debug issues without creating permanent artifacts * Prototype new ideas rapidly **Use `flyte deploy` when your project is ready to be formalized:** * Freeze a stable version of your tasks for repeated use * Share tasks with team members or across environments * Move from experimentation to a more structured workflow * Create a permanent reference point (not necessarily production-ready) This section explains both approaches and when to use each one. [Ephemeral deployment and immediate execution](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/#ephemeral-deployment-and-immediate-execution) ------------------------------------------------------------------------------------------------------------------------------------------------------------- The `flyte run` CLI command and the `flyte.run()` SDK function are used to **ephemerally deploy** and **immediately execute** a task on the backend in a single step. The task can be re-run and its execution and outputs can be observed in the **Runs list** UI, but it is not permanently added to the **Tasks list** on the backend. Let’s say you have the following file called `greeting.py`: # greeting.py import flyte env = flyte.TaskEnvironment(name="greeting_env") @env.task async def greet(message: str) -> str: return f"{message}!" ProgrammaticCLI You can run the task programmatically using the `flyte.run()` function: # greeting.py import flyte env = flyte.TaskEnvironment(name="greeting_env") @env.task async def greet(message: str) -> str: return f"{message}!" if __name__ == "__main__": flyte.init_from_config() result = flyte.run(greet, message="Good morning!") print(f"Result: {result}") Here we add a `__main__` block to the `greeting.py` file that initializes the Flyte SDK from the configuration file and then calls `flyte.run()` with the `greet` task and its argument. Now you can run the `greet` task on the backend just by executing the `greeting.py` file locally as a script: python greeting.py The general form of the command for running a task from a local file is: flyte run So, to run the `greet` task defined in the `greeting.py` file, you would run: flyte run greeting.py greet --message "Good morning!" This command: 1. **Temporarily deploys** the task environment named `greeting_env` (held by the variable `env`) that contains the `greet` task. 2. **Executes** the `greet` function with argument `message` set to `"Good morning!"`. Note that `message` is the actual parameter name defined in the function signature. 3. **Returns** the execution results and displays them in the terminal. For more details on how `flyte run` and `flyte.run()` work under the hood, see [How Run Works](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/how-task-run-works) . [Persistent deployment](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/#persistent-deployment) --------------------------------------------------------------------------------------------------------------- The `flyte deploy` CLI command and the `flyte.deploy()` SDK function are used to **persistently deploy** a task environment (and all its contained tasks) to the backend. The tasks within the deployed environment will appear in the **Tasks list** UI on the backend and can then be executed multiple times without needing to redeploy them. ProgrammaticCLI You can deploy programmatically using the `flyte.deploy()` function: # greeting.py import flyte env = flyte.TaskEnvironment(name="greeting_env") @env.task async def greet(message: str) -> str: return f"{message}!" if __name__ == "__main__": flyte.init_from_config() deployments = flyte.deploy(env) print(deployments[0].summary_repr()) Now you can deploy the `greeting_env` task environment (and therefore the `greet()` task) just by executing the `greeting.py` file locally as a script. python greeting.py The general form of the command for deploying a task environment from a local file is: flyte deploy So, using the same `greeting.py` file as before, you can deploy the `greeting_env` task environment like this: flyte deploy greeting.py env This command deploys the task environment _assigned to the variable `env`_ in the `greeting.py` file, which is the `TaskEnvironment` named `greeting_env`. Notice that you must specify the _variable_ to which the `TaskEnvironment` is assigned (`env` in this case), not the name of the environment itself (`greeting_env`). Deploying a task environment deploys all tasks defined within it. Here, that means all functions decorated with `@env.task`. In this case there is just one: `greet()`. For more details on how `flyte deploy` and `flyte.deploy()` work under the hood, see [How Deployment Works](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/how-task-deployment-works) . [Running already deployed tasks](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/#running-already-deployed-tasks) --------------------------------------------------------------------------------------------------------------------------------- If you have already deployed your task environment, you can run its tasks without redeploying by using the `flyte run` CLI command or the `flyte.run()` SDK function in a slightly different way. Alternatively, you can always initiate execution of a deployed task from the UI. ProgrammaticCLI You can run already-deployed tasks programmatically using the `flyte.run()` function. For example, to run the previously deployed `greet` task from the `greeting_env` environment: # greeting.py import flyte env = flyte.TaskEnvironment(name="greeting_env") @env.task async def greet(message: str) -> str: return f"{message}!" if __name__ == "__main__": flyte.init_from_config() flyte.deploy(env) task = flyte.remote.Task.get("greeting_env.greet", auto_version="latest") result = flyte.run(task, message="Good morning!") print(f"Result: {result}") When you execute this script locally, it will: * Deploy the `greeting_env` task environment as before. * Retrieve the already-deployed `greet` task using `flyte.remote.Task.get()`, specifying its full task reference as a string: `"greeting_env.greet"`. * Call `flyte.run()` with the retrieved task and its argument. For more details on how running already-deployed tasks works, see [How task Run works > SDK: Running deployed tasks](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/how-task-run-works#running-deployed-tasks) . To run a permanently deployed task using the `flyte run` CLI command, use the special `deployed-task` keyword followed by the task reference in the format `{environment_name}.{task_name}`. For example, to run the previously deployed `greet` task from the `greeting_env` environment: flyte run deployed-task greeting_env.greet --message "World" Notice that now that the task environment is deployed, you use its name (`greeting_env`), not by the variable name to which it was assigned in source code (`env`). The task environment name plus the task name (`greet`) are combined with a dot (`.`) to form the full task reference: `greeting_env.greet`. The special `deployed-task` keyword tells the CLI that you are referring to a task that has already been deployed. In effect, it replaces the file path argument used for ephemeral runs. When executed, this command will run the already-deployed `greet` task with argument `message` set to `"World"`. You will see the result printed in the terminal. You can also, of course, observe the execution in the **Runs list** UI. To execute a deployed task in a different project or domain than your configured defaults, use `--run-project` and `--run-domain`: flyte run --run-project prod-project --run-domain production deployed-task greeting_env.greet --message "World" For all `flyte run` options, see [Run command options](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/run-command-options) . [Configuring runs with `flyte.with_runcontext()`](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/#configuring-runs-with-flytewith_runcontext) -------------------------------------------------------------------------------------------------------------------------------------------------------------- Both `flyte run` and `flyte.run()` accept a range of invocation-time parameters that control where the run executes, where outputs are stored, caching behavior, and more. Programmatically, these are set with `flyte.with_runcontext()` before calling `.run()`. Inside a running task, `flyte.ctx()` provides read access to the same context. For the full parameter reference, see [Run context](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/run-context) . LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Quickstart | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Quickstart ========== Let’s get you up and running with your first workflow on your local machine. Want to try Flyte without installing anything? [Try Flyte 2 in your browser](https://flyte2intro.apps.demo.hosted.unionai.cloud/) . [What you’ll need](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/#what-youll-need) ----------------------------------------------------------------------------------------------- * Python 3.10+ in a virtual environment [Install the SDK](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/#install-the-sdk) ---------------------------------------------------------------------------------------------- Install the `flyte` package: pip install 'flyte[tui]' We also install the `tui` extra to enable the terminal user interface. Verify it worked: flyte --version Output: Flyte SDK version: 2.*.* [Configure](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/#configure) ---------------------------------------------------------------------------------- Create a config file for local execution. Runs will be persisted locally in a SQLite database. flyte create config --local-persistence This creates `.flyte/config.yaml` in your current directory. See [Setting up a configuration file](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox#configure) for more options when connecting to a cluster. Run `flyte get config` to check which configuration is currently active. [Write your first workflow](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/#write-your-first-workflow) ------------------------------------------------------------------------------------------------------------------ Create `hello.py`: hello.py [](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/getting-started/hello.py "View source on GitHub") # hello.py import flyte # The `hello_env` TaskEnvironment is assigned to the variable `env`. # It is then used in the `@env.task` decorator to define tasks. # The environment groups configuration for all tasks defined within it. env = flyte.TaskEnvironment(name="hello_env") # We use the `@env.task` decorator to define a task called `fn`. @env.task def fn(x: int) -> int: # Type annotations are required slope, intercept = 2, 5 return slope * x + intercept # We also use the `@env.task` decorator to define another task called `main`. # This is the entrypoint task of the workflow. # It calls the `fn` task defined above multiple times using `flyte.map`. @env.task def main(x_list: list[int] = list(range(10))) -> float: y_list = list(flyte.map(fn, x_list)) # flyte.map is like Python map, but runs in parallel. y_mean = sum(y_list) / len(y_list) return y_mean Here’s what’s happening: * **`TaskEnvironment`** specifies configuration for your tasks (container image, resources, etc.) * **`@env.task`** turns Python functions into tasks that run remotely * Both tasks share the same `env`, so they’ll have identical configurations [Run it](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/#run-it) ---------------------------------------------------------------------------- Create a project directory and place your files there: . ├── hello.py └── .flyte └── config.yaml Do not run `flyte run` from your home directory. Flyte packages the current directory when running remotely, so running from `$HOME` would attempt to bundle your entire home folder. Always work from a dedicated project directory. Run the workflow: flyte run --local hello.py main This executes the workflow locally on your machine. [See the results](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/#see-the-results) ---------------------------------------------------------------------------------------------- You can see the run in the TUI by running: flyte start tui The TUI will open into the explorer view ![Explorer View](https://www.union.ai/docs/v2/flyte/_static/images/user-guide/quickstart/explorer-tui.png) To navigate to the run details, double-click it or press `Enter` to view the run details. ![Run Details View](https://www.union.ai/docs/v2/flyte/_static/images/user-guide/quickstart/run-tui.png) [Next steps](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/#next-steps) ------------------------------------------------------------------------------------ Now that you’ve run your first workflow: * [**Core concepts**](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts) : Understand the core concepts of Flyte programming * [**Run locally**](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-locally) : Learn about the TUI, caching, and other features that work locally * [**Run on the devbox**](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) : Learn about the devbox cluster and how to run workflows on it LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/quickstart/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Reference | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Reference ========= This section provides the reference material for the Flyte SDK and CLI. To get started, add `flyte` to your project $ uv pip install --no-cache --upgrade flyte This will install both the Flyte SDK and CLI. Flyte SDK The Flyte SDK provides the core Python API for building workflows and apps on your Union instance. Flyte CLI The Flyte CLI is the command-line interface for interacting with your Union instance. Migration from Flyte 1 Comprehensive reference for migrating Flyte 1 workflows to Flyte 2. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/api-reference/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Build an MCP | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Build an MCP ============ Flyte supports serving [Model Context Protocol (MCP)](https://modelcontextprotocol.io/) servers over HTTP. There are two main MCP environment types: | Environment | Purpose | | --- | --- | | **`MCPAppEnvironment`** | Serve any FastMCP instance with custom tools | | **`FlyteMCPAppEnvironment`** | Flyte-specific server that exposes Flyte operations as tools | See the sub-pages for detailed guides: * [User-defined MCP servers](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/mcp_server) : Build and deploy your own FastMCP instances * [Flyte MCP servers](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/flyte_mcp_server) : Use Flyte-specific tools to interact with your cluster [HTTP layout](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/#http-layout) ------------------------------------------------------------------------------------- All MCP app environments expose the same HTTP endpoints: * `GET /health` — Liveness/readiness check (`{"status": "healthy"}`) * `POST {mcp_mount_path}/mcp` or `{mcp_mount_path}/sse` — MCP protocol endpoint (default: `/mcp` for generic, `/flyte-mcp` for Flyte) [Quickstart](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/#quickstart) ----------------------------------------------------------------------------------- The fastest way to try Flyte MCP is locally — no deployment needed: uvx --from "flyte[mcp]" flyte-mcp For client setup, tool selection, allowlists, and remote deployment, see [Flyte MCP server](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/flyte_mcp_server) . LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Scale your runs | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Scale your runs =============== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This guide helps you understand and optimize the performance of your Flyte workflows. Whether you’re building latency-sensitive applications or high-throughput data pipelines, these docs will help you make the right architectural choices. [Understanding Flyte execution](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/#understanding-flyte-execution) --------------------------------------------------------------------------------------------------------------------------- Before optimizing performance, it’s important to understand how Flyte executes your workflows: * **[Data flow](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/data-flow) **: Learn how data moves between tasks, including inline vs. reference data types, caching mechanisms, and storage configuration. * **[Life of a run](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/life-of-a-run) **: Understand what happens when you invoke `flyte.run()`, from code analysis and image building to task execution and state management. [Performance optimization](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/#performance-optimization) ----------------------------------------------------------------------------------------------------------------- Once you understand the fundamentals, dive into performance tuning: * **[Scale your workflows](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/scale-your-workflows) **: A comprehensive guide to optimizing workflow performance, covering latency vs. throughput, task overhead analysis, batching strategies, reusable containers, and more. [Key concepts for scaling](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/#key-concepts-for-scaling) ----------------------------------------------------------------------------------------------------------------- When scaling your workflows, keep these principles in mind: 1. **Task overhead matters**: The overhead of creating a task (uploading data, enqueuing, creating containers) should be much smaller than the task runtime. 2. **Batch for throughput**: For large-scale data processing, batch multiple items into single tasks to reduce overhead. 3. **Reusable containers**: Eliminate container startup overhead and enable concurrent execution with reusable containers. 4. **Traces for lightweight ops**: Use traces instead of tasks for lightweight operations that need checkpointing. 5. **Limit fanout**: Keep the total number of actions per run below 50k (target 10k-20k for best performance). 6. **Choose the right data types**: Use reference types (files, directories, DataFrames) for large data and inline types for small data. For detailed guidance on each of these topics, see [Scale your workflows](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/scale-your-workflows) . LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Run modes | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Run modes ========= Flyte OSS currently supports two execution modes — Local and Devbox (Remote execution is coming soon): Local Run tasks and apps directly in your local Python process with no K8s cluster or Docker required. Ideal for rapid iteration and debugging. Devbox Run tasks and apps in a lightweight Flyte cluster using Docker. Get the full Flyte UI and backend experience on your machine. | Aspect | Local (`--local`) | Devbox | | --- | --- | --- | | **⚡️ Execution** | In-process Python | Containerized, local Docker | | **🐳 Docker required** | No | Yes | | **💻 Flyte UI** | No (TUI only) | Yes (`localhost:30080`) | | **📦 Container images** | Ignored | Built locally | | **🔀 Parallelism** | Sequential | Cluster-level | | **⭐️ Best for** | Fast iteration, debugging | Testing container builds, full Flyte features | The same task code runs unchanged in both modes. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Agent framework integrations | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Agent framework integrations ============================ **Any Python-based agent framework works with Flyte.** Flyte doesn’t replace your framework — it provides the production layer around it. You write your agent with whatever framework you prefer, then invoke it from inside an `@env.task`, where it runs in a container with durable checkpointing and full observability. Each LLM call, tool call, and routing decision can be captured as a span in the Flyte dashboard. Because the framework drives the loop and Flyte wraps it, you don’t need a dedicated plugin for a framework to use it — if it runs in Python, it runs on Flyte. The `flyte` SDK provides a [native agent harness](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/flyte-agents) that you can use to build your own agent loop. [How much control does the framework give you?](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/#how-much-control-does-the-framework-give-you) --------------------------------------------------------------------------------------------------------------------------------------------------------------------------- Frameworks differ in how much of the agent loop they own, which determines how you integrate them with Flyte: | Level of control | What it means | Example | Integration pattern | | --- | --- | --- | --- | | **You own the loop** | The framework gives you primitives (graph nodes, tools) and you wire the control flow | [LangGraph](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/langgraph) | Decorate nodes with `@flyte.trace`; run the compiled graph inside a task | | **The framework owns the loop, you own the tools** | The framework runs the tool-calling loop; you provide tools as plain functions | [PydanticAI](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/pydantic-ai) | Have tools delegate to durable `@env.task`s | | **First-party tool adapter** | Flyte ships a decorator that turns a task into a framework tool | [OpenAI Agents SDK](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/openai-agents-sdk) | Stack `function_tool` on `@env.task` | Whichever model your framework uses, the integration is the same in spirit: the framework decides _what_ the agent does next, and Flyte decides _where and how durably_ each step runs. [Supported frameworks](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/#supported-frameworks) -------------------------------------------------------------------------------------------------------------------------- * [**LangGraph**](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/langgraph) — run compiled graphs inside tasks and fan them out in parallel. * [**PydanticAI**](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/pydantic-ai) — type-safe agents whose tools delegate to durable tasks. * [**OpenAI Agents SDK**](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/openai-agents-sdk) — expose durable tasks as Agents SDK tools with `flyteplugins-openai`. Don’t see your framework? The same pattern — invoke the framework from inside an `@env.task` and trace its calls — applies to any Python agent library. See [Bring your own framework](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/bring-your-own-framework) for a framework-agnostic template. [Next steps](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/#next-steps) ------------------------------------------------------------------------------------------------------ * [Deploy an agent as a service](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/deploy-agent-as-service) : run your agent on a schedule or behind a webhook. * [Build an agent with pure Python](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/python-agents) : roll the loop yourself with no framework at all. * [The Flyte Agent harness](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/flyte-agents) : the built-in, batteries-included agent loop. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Biotech & Healthcare | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Biotech & Healthcare ==================== Tutorials for bioinformatics, medical imaging, and other life-sciences workloads. Genomic alignment Align sequencing reads to a reference genome with a cached, parallel Bowtie 2 pipeline. Brain tumor MRI classification Classify brain MRI scans with a two-phase EfficientNet-B4 pipeline featuring resumable GPU checkpointing and in-UI reports. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Computer Vision | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Computer Vision =============== Tutorials for image and vision-language model workloads. Fine-tuning a VLM Adapt Qwen2.5-VL to occluded image classification by training a 10K-parameter adapter with multi-node DeepSpeed, automatic recovery, and live training dashboards. Multimodal retrieval evaluation Benchmark ColPali, SigLIP, and OCR+BM25 visual document retrieval on ViDoRe with warm GPU containers, dynamic batching, and an interactive report. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/tutorials/computer-vision/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Context Engineering | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Context Engineering =================== Tutorials for prompt engineering, prompt optimization, and context construction. Automatic prompt engineering Easily run prompt optimization with real-time observability, traceability, and automatic recovery. Text-to-SQL Learn how to turn natural language questions into SQL queries with Flyte and LlamaIndex, and explore prompt optimization in practice. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/tutorials/context-engineering/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Sandboxing | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Sandboxing ========== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. A **sandbox** is an isolated, secure environment where code can run without affecting the host system. Sandboxes restrict what the executing code can do — limiting filesystem access, blocking network calls, and preventing arbitrary system operations — so that even malicious or buggy code cannot cause harm. The exact restrictions depend on the sandboxing approach: some sandboxes eliminate dangerous operations entirely, while others provide full capabilities within an isolated, disposable container. [Why sandboxing matters for AI](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/#why-sandboxing-matters-for-ai) -------------------------------------------------------------------------------------------------------------------------- LLM-generated code is inherently untrusted. The model may produce code that is correct and useful, but it can also produce code that is dangerous — and it does so without intent or awareness. | Risk | Example | | --- | --- | | Data destruction | `DELETE FROM orders WHERE 1=1` — wipes an entire table | | Credential exfiltration | Reads environment variables and sends API keys to an external endpoint | | Infinite loops | `while True: pass` — consumes CPU indefinitely | | Resource abuse | Spawns thousands of threads or allocates unbounded memory | | Filesystem damage | `rm -rf /` or overwrites critical configuration files | | Network abuse | Makes unauthorized API calls, sends spam, or joins a botnet | Running LLM-generated code without a sandbox means trusting the model to never make these mistakes. Sandboxing eliminates this trust requirement by making dangerous operations structurally impossible. [Types of sandboxes](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/#types-of-sandboxes) ---------------------------------------------------------------------------------------------------- There are three broad approaches to sandboxing LLM-generated code, each with different tradeoffs: | Type | How it works | Tradeoffs | Examples | | --- | --- | --- | --- | | **One-shot execution** | Code runs to completion in a disposable container, then the container is discarded. Stdout, stderr, and outputs are captured. | Simple, no state reuse. Good for single-turn tasks. | Container tasks, serverless functions | | **Interactive sessions** | A persistent VM or container where you send commands incrementally and observe results between steps. Sessions last for the lifetime of the VM. | Flexible and multi-turn, but heavier to provision and manage. | E2B, Daytona, fly.io | | **Programmatic tool calling** | The LLM generates orchestration code that calls a predefined set of tools. The orchestration code runs in a sandbox while the tools run in full containers. | Durable, observable, and secure. Tools are known ahead of time. | Flyte workflow sandboxing | [What Flyte offers](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/#what-flyte-offers) -------------------------------------------------------------------------------------------------- Flyte provides two complementary sandboxing approaches: ### [Workflow sandbox (Monty)](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/#workflow-sandbox-monty) A **sandboxed orchestrator** built on [Monty](https://github.com/pydantic/pydantic-monty) , a Rust-based sandboxed Python interpreter. The sandbox starts in microseconds, runs pure Python control flow, and dispatches heavy work to full container tasks through the Flyte controller. This enables the **programmatic tool calling** pattern (also known as code mode): LLMs generate Python orchestration code that invokes registered tools, and Flyte executes it safely with full durability, observability, and type checking. ### [Code sandbox (container)](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/#code-sandbox-container) A **stateless code sandbox** that runs arbitrary Python scripts or shell commands inside an ephemeral Docker container. The container is built on demand from declared dependencies, executed once, and discarded. This is the right choice when you need full Python capabilities — third-party packages, file I/O, shell commands, or any computation that goes beyond pure control flow. ### [When to use which](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/#when-to-use-which) | | Workflow sandbox | Code sandbox | | --- | --- | --- | | **Runtime** | Monty (Rust-based Python interpreter) | Ephemeral Docker container | | **Startup** | Microseconds | Seconds (image build + container spin-up) | | **Capabilities** | Pure Python control flow only — no imports, no I/O, no network | Full Python environment — any package, any library, full I/O | | **Use case** | LLM-generated orchestration logic that calls registered tools | Arbitrary computation — data processing, test execution, ETL, shell pipelines | | **State** | Runs within a worker container process | Stateless — fresh container per invocation | | **Security model** | Dangerous operations are structurally impossible | Isolated container | * Use the **workflow sandbox** when you need to run untrusted control flow (loops, conditionals, routing) that dispatches work to known tasks. It starts in microseconds and provides the strongest isolation guarantees. * Use the **code sandbox** when you need full Python capabilities — third-party packages, file I/O, shell commands, or any computation that goes beyond pure control flow. ### [Learn more](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/#learn-more) * [**Workflow sandboxing**](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/workflow-sandboxing-flyte) — How the Monty-based sandboxed orchestrator works, with examples * [**Programmatic tool calling for agents**](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/code-mode) — The concept behind programmatic tool calling and how to build agents that use it * [**Code sandboxing**](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/code-sandboxing) — Running arbitrary code and commands in ephemeral containers with `flyte.sandbox.create()` LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Model Training | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Model Training ============== Tutorials for training, fine-tuning, and hyperparameter optimization of models at scale. Hyperparameter optimization Run large-scale HPO experiments with zero manual tracking, deterministic results, and automatic recovery. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/tutorials/model-training/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Frontier AI | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Frontier AI =========== Tutorials for frontier-model pretraining, automated experimentation, and large-scale AI workloads. Distributed LLM pretraining Pretrain large language models at scale with PyTorch Lightning, FSDP, and H200 GPUs, featuring streaming data and real-time metrics. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/tutorials/frontier-ai/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Contributing code | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Contributing code ================= Thank you for your interest in Flyte! This page is part of the Flyte 2 documentation. If you are interested in contributing code to for Flyte 1, switch the selector at the top of the page to \*v1\*\*. [Flyte 2](https://www.union.ai/docs/v2/flyte/community/contributing-code/#flyte-2) ------------------------------------------------------------------------------------ The Flyte 2 SDK source code is available on [GitHub](https://github.com/flyteorg/flyte-sdk) under the same Apache license as the original Flyte 1. You are welcome to take a look, [download the package](https://pypi.org/project/flyte/#history) and try running code locally. The Flyte 2 backend is not yet available as open source, (but it will be soon!) To run Flyte 2 code now you can apply for a [beta preview of the Union 2 backend](https://www.union.ai/beta) . When the Flyte 2 backend is released we will roll out a full contributor program just as we have for Flyte 1. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/community/contributing-code/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Geospatial | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Geospatial ========== Tutorials for satellite imagery, remote sensing, and earth and atmospheric modeling workloads. GPU-accelerated climate modeling Run ensemble atmospheric simulations on H200 GPUs with multi-source data ingestion and real-time extreme event detection. Satellite image classification Build a production-grade EfficientNet pipeline for land-use classification with caching, experiment tracking, and reporting. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/tutorials/geospatial/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Agents | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Agents ====== Tutorials for building agentic workflows and autonomous LLM-powered systems. Autoresearch agent Run an autonomous research loop that drives Claude Code in a GPU container to run experiments, then commits results and opens a pull request. Coding agent Securely execute and iterate on LLM-generated code using a code agent with error reflection and retry logic. Competitive intelligence agent Fan out across competitors, extract source-cited market deltas with the You.com Search API, and build a knowledge-graph-ready intelligence table. Compliance monitoring agent Monitor trusted regulatory sources with the You.com Research API and route citation-precise findings to the right team. Deep research Build an agentic workflow for deep research with multi-step reasoning and evaluation. Field data enrichment agent Enrich geo-tagged operational events with real-world public context using the You.com Search API with country and freshness targeting. MLE Bot: autonomous ML engineer An autonomous ML agent that designs, runs, and iterates on experiments using Flyte’s durable sandbox for safe LLM-generated code execution. Support resolution agent Ground support tickets in fresh public sources via the You.com Research API and draft cited, customer-ready replies for human review. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/tutorials/agents/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Serve and deploy apps | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Serve and deploy apps ===================== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. Flyte provides two main ways to deploy apps: **serve** (for development) and **deploy** (for production). This section covers both methods and their differences. [Serve vs Deploy](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/#serve-vs-deploy) --------------------------------------------------------------------------------------------------------- ### [`flyte serve`](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/#flyte-serve) Serving is designed for development and iteration: * **Dynamic parameter modification**: You can override app parameters when serving * **Quick iteration**: Faster feedback loop for development * **Interactive**: Better suited for testing and experimentation ### [`flyte deploy`](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/#flyte-deploy) Deployment is designed for production use: * **Immutable**: Apps are deployed with fixed configurations * **Production-ready**: Optimized for stability and reproducibility [Using Python SDK](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/#using-python-sdk) ----------------------------------------------------------------------------------------------------------- ServeDeployserve\_and\_deploy\_examples.py [](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/serve_and_deploy_examples.py "View source on GitHub") app_env = flyte.app.AppEnvironment( name="my-app", image=flyte.app.Image.from_debian_base().with_pip_packages("streamlit==1.41.1"), args=["streamlit", "hello", "--server.port", "8080"], port=8080, resources=flyte.Resources(cpu="1", memory="1Gi"), ) if __name__ == "__main__": flyte.init_from_config() app = flyte.serve(app_env) print(f"Served at: {app.url}") serve\_and\_deploy\_examples.py [](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/serve-and-deploy-apps/serve_and_deploy_examples.py "View source on GitHub") app_env = flyte.app.AppEnvironment( name="my-app", image=flyte.app.Image.from_debian_base().with_pip_packages("streamlit==1.41.1"), args=["streamlit", "hello", "--server.port", "8080"], port=8080, resources=flyte.Resources(cpu="1", memory="1Gi"), ) if __name__ == "__main__": flyte.init_from_config() deployments = flyte.deploy(app_env) # Access deployed app URL from the deployment for deployed_env in deployments[0].envs.values(): print(f"Deployed: {deployed_env.deployed_app.url}") [Using the CLI](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/#using-the-cli) ----------------------------------------------------------------------------------------------------- ServeDeploy flyte serve path/to/app.py app_env flyte deploy path/to/app.py app_env [Next steps](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/#next-steps) ----------------------------------------------------------------------------------------------- * [**How app serving works**](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/how-app-serving-works) : Understanding the serve process and configuration options * [**How app deployment works**](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/how-app-deployment-works) : Understanding the deploy process and configuration options * [**Activating and deactivating apps**](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/activating-and-deactivating-apps) : Managing app lifecycle * [**Basic project**](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/basic-project/) : Build a RAG embedding pipeline and semantic search app with Streamlit * [**Prefetching models**](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/prefetching-models) : Download and shard HuggingFace models for vLLM and SGLang LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Core concepts | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Core concepts ============= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. Now that you’ve completed the [Quickstart](https://www.union.ai/docs/v2/flyte/user-guide/quickstart) , let’s explore Flyte’s core concepts through working examples. By the end of this section, you’ll understand: * **TaskEnvironment**: The container configuration that defines where and how your code runs * **Tasks**: Python functions that execute remotely in containers * **Runs and Actions**: How Flyte tracks and manages your executions * **Apps**: Long-running services for APIs, dashboards, and inference endpoints Each concept is introduced with a practical example you can run yourself. [How Flyte works](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/#how-flyte-works) ------------------------------------------------------------------------------------------------- When you run code with Flyte, here’s what happens: 1. You define a **TaskEnvironment** that specifies the container image and resources 2. You decorate Python functions with `@env.task` to create **tasks** 3. When you execute a task, Flyte creates a **run** that tracks the execution 4. Each task execution within a run is an **action** Let’s explore each of these in detail. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Native app integrations | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Native app integrations ======================= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. Flyte ships with a set of pre-built [`AppEnvironment`](https://www.union.ai/docs/v2/flyte/user-guide/build-apps) integrations that wrap popular frameworks and serving runtimes, so you can deploy common app types without writing the integration glue yourself. Each integration provides a ready-to-use environment class — just configure your app, image, resources, and scaling, and Flyte handles the rest. If you’re new to apps in Flyte, start with [Introducing apps](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/introducing-apps) for an overview, then see [Build apps](https://www.union.ai/docs/v2/flyte/user-guide/build-apps) to learn how to build custom app environments from scratch. [When to use a native integration](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/#when-to-use-a-native-integration) --------------------------------------------------------------------------------------------------------------------------------------------- Use a native integration when your app fits one of the supported frameworks and you want: * **A minimal, opinionated setup** — sensible defaults for the framework, no boilerplate * **First-class support** — features like model streaming, OpenAI-compatible APIs, and passthrough auth wired in for you * **Faster time-to-deploy** — focus on your app logic, not on packaging and serving plumbing For app types not covered here, build a custom [`AppEnvironment`](https://www.union.ai/docs/v2/flyte/user-guide/build-apps) using the patterns in the [Build apps](https://www.union.ai/docs/v2/flyte/user-guide/build-apps) section. [Available integrations](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/#available-integrations) ------------------------------------------------------------------------------------------------------------------------- | Integration | Framework | Typical use case | | --- | --- | --- | | [Streamlit app](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/streamlit-app) | [Streamlit](https://streamlit.io/) | Interactive dashboards and data apps | | [FastAPI app](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/fastapi-app) | [FastAPI](https://fastapi.tiangolo.com/) | REST APIs, webhooks, and backend services | | [vLLM app](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/vllm-app) | [vLLM](https://docs.vllm.ai/) | High-throughput LLM inference with an OpenAI-compatible API | | [SGLang app](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/sglang-app) | [SGLang](https://docs.sglang.io/) | Structured generation and LLM serving with an OpenAI-compatible API | | [Flyte webhook](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/flyte-webhook) | [FastAPI](https://fastapi.tiangolo.com/) | Pre-built HTTP endpoints for common Flyte control-plane operations | [Next steps](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/#next-steps) ------------------------------------------------------------------------------------------------- * [**Streamlit app**](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/streamlit-app) : Build interactive Streamlit dashboards * [**FastAPI app**](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/fastapi-app) : Create REST APIs and backend services * [**vLLM app**](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/vllm-app) : Serve large language models with vLLM * [**SGLang app**](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/sglang-app) : Serve LLMs with SGLang for structured generation * [**Flyte webhook**](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/flyte-webhook) : Pre-built webhook for common Flyte operations LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Financial Services & Fintech | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Financial Services & Fintech ============================ Tutorials for financial research, trading, and other fintech workloads. Financial research agent Prep equity briefings for the earnings cycle with grounded You.com Research synthesis and fresh news from the Search API. Multi-agent trading simulation A multi-agent trading simulation, modeling how agents within a firm might interact, strategize, and make trades collaboratively. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/tutorials/financial-services/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Project patterns | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Project patterns ================ The rest of the user guide explains what Flyte can do. This section explains how we recommend you structure real projects. These are opinionated guides. They represent patterns we’ve seen work well across many teams and production deployments. If you’re starting a new project or scaling an existing one, start here. Bring your own image (BYOI) Two patterns for teams that own their Docker images and want Flyte for orchestration without handing over their build pipeline. Monorepo with uv How to structure Flyte projects with uv, from single-package setups to multi-team monorepos with shared and independent lockfiles. CI/CD deployments How to deploy a Flyte project from CI. Uses GitHub Actions as the reference, but the building blocks — API key, `flyte deploy`, commit-pinned versions — translate to any runner. Resource management and multi-team scaling Projects, domains, quotas, RBAC, and secrets — the primitives to set up before you have ten teams and a noisy-neighbor problem. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/project-patterns/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Anthropic | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Anthropic ========= The Anthropic plugin lets you build agentic workflows with [Claude](https://www.anthropic.com/) on Flyte. It provides a `function_tool` decorator that wraps Flyte tasks as tools that Claude can call, and a `run_agent` function that drives the agent conversation loop. When Claude calls a tool, the call executes as a Flyte task with full observability, retries, and caching. [Installation](https://www.union.ai/docs/v2/flyte/integrations/anthropic/#installation) ----------------------------------------------------------------------------------------- pip install flyteplugins-anthropic Requires `anthropic >= 0.40.0`. [Quick start](https://www.union.ai/docs/v2/flyte/integrations/anthropic/#quick-start) --------------------------------------------------------------------------------------- import flyte from flyteplugins.anthropic import function_tool, run_agent env = flyte.TaskEnvironment( name="claude-agent", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from_uv_script(__file__, name="anthropic_agent"), secrets=flyte.Secret("anthropic_api_key", as_env_var="ANTHROPIC_API_KEY"), ) @function_tool @env.task async def get_weather(city: str) -> str: """Get the current weather for a city.""" return f"The weather in {city} is sunny, 72F" @env.task async def main(prompt: str) -> str: tools = [get_weather] return await run_agent(prompt=prompt, tools=tools) [API](https://www.union.ai/docs/v2/flyte/integrations/anthropic/#api) ----------------------------------------------------------------------- ### [`function_tool`](https://www.union.ai/docs/v2/flyte/integrations/anthropic/#function_tool) Converts a Flyte task, `@flyte.trace`\-decorated function, or plain callable into a tool that Claude can invoke. @function_tool @env.task async def my_tool(param: str) -> str: """Tool description sent to Claude.""" ... Can also be called with optional overrides: @function_tool(name="custom_name", description="Custom description") @env.task async def my_tool(param: str) -> str: ... Parameters: | Parameter | Type | Description | | --- | --- | --- | | `func` | callable | The function to wrap | | `name` | `str` | Override the tool name (defaults to the function name) | | `description` | `str` | Override the tool description (defaults to the docstring) | The docstring on each `@function_tool` task is sent to Claude as the tool description. Write clear, concise docstrings. ### [`Agent`](https://www.union.ai/docs/v2/flyte/integrations/anthropic/#agent) A dataclass for bundling agent configuration: from flyteplugins.anthropic import Agent agent = Agent( name="my-agent", instructions="You are a helpful assistant.", model="claude-sonnet-4-20250514", tools=[get_weather], max_tokens=4096, max_iterations=10, ) | Field | Type | Default | Description | | --- | --- | --- | --- | | `name` | `str` | `"assistant"` | Agent name | | `instructions` | `str` | `"You are a helpful assistant."` | System prompt | | `model` | `str` | `"claude-sonnet-4-20250514"` | Claude model ID | | `tools` | `list[FunctionTool]` | `[]` | Tools available to the agent | | `max_tokens` | `int` | `4096` | Maximum tokens per response | | `max_iterations` | `int` | `10` | Maximum tool-call loop iterations | ### [`run_agent`](https://www.union.ai/docs/v2/flyte/integrations/anthropic/#run_agent) Runs a Claude conversation loop, dispatching tool calls to Flyte tasks until Claude returns a final response. result = await run_agent( prompt="What's the weather in Tokyo?", tools=[get_weather], model="claude-sonnet-4-20250514", ) You can also pass an `Agent` object: result = await run_agent(prompt="What's the weather?", agent=agent) | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `prompt` | `str` | required | User message | | `tools` | `list[FunctionTool]` | `None` | Tools available to the agent | | `agent` | `Agent` | `None` | Agent config (overrides individual params) | | `model` | `str` | `"claude-sonnet-4-20250514"` | Claude model ID | | `system` | `str` | `None` | System prompt | | `max_tokens` | `int` | `4096` | Maximum tokens per response | | `max_iterations` | `int` | `10` | Maximum iterations (prevents infinite loops) | | `api_key` | `str` | `None` | API key (falls back to `ANTHROPIC_API_KEY` env var) | [Secrets](https://www.union.ai/docs/v2/flyte/integrations/anthropic/#secrets) ------------------------------------------------------------------------------- Store your Anthropic API key as a Flyte secret and expose it as an environment variable: secrets=flyte.Secret("anthropic_api_key", as_env_var="ANTHROPIC_API_KEY") [API reference](https://www.union.ai/docs/v2/flyte/integrations/anthropic/#api-reference) ------------------------------------------------------------------------------------------- See the [Anthropic API reference](https://www.union.ai/docs/v2/flyte/api-reference/integrations/anthropic) for full details. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/anthropic/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Gemini | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Gemini ====== The Gemini plugin lets you build agentic workflows with [Gemini](https://ai.google.dev/) on Flyte. It provides a `function_tool` decorator that wraps Flyte tasks as tools that Gemini can call, and a `run_agent` function that drives the agent conversation loop. When Gemini calls a tool, the call executes as a Flyte task with full observability, retries, and caching. Gemini’s native parallel function calling is supported: multiple tool calls in a single turn are all dispatched and their results bundled into one response. [Installation](https://www.union.ai/docs/v2/flyte/integrations/gemini/#installation) -------------------------------------------------------------------------------------- pip install flyteplugins-gemini Requires `google-genai >= 1.0.0`. [Quick start](https://www.union.ai/docs/v2/flyte/integrations/gemini/#quick-start) ------------------------------------------------------------------------------------ import flyte from flyteplugins.gemini import function_tool, run_agent env = flyte.TaskEnvironment( name="gemini-agent", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from_uv_script(__file__, name="gemini_agent"), secrets=flyte.Secret("google_api_key", as_env_var="GOOGLE_API_KEY"), ) @function_tool @env.task async def get_weather(city: str) -> str: """Get the current weather for a city.""" return f"The weather in {city} is sunny, 72F" @env.task async def main(prompt: str) -> str: tools = [get_weather] return await run_agent(prompt=prompt, tools=tools) [API](https://www.union.ai/docs/v2/flyte/integrations/gemini/#api) -------------------------------------------------------------------- ### [`function_tool`](https://www.union.ai/docs/v2/flyte/integrations/gemini/#function_tool) Converts a Flyte task, `@flyte.trace`\-decorated function, or plain callable into a tool that Gemini can invoke. @function_tool @env.task async def my_tool(param: str) -> str: """Tool description sent to Gemini.""" ... Can also be called with optional overrides: @function_tool(name="custom_name", description="Custom description") @env.task async def my_tool(param: str) -> str: ... Parameters: | Parameter | Type | Description | | --- | --- | --- | | `func` | callable | The function to wrap | | `name` | `str` | Override the tool name (defaults to the function name) | | `description` | `str` | Override the tool description (defaults to the docstring) | The docstring on each `@function_tool` task is sent to Gemini as the tool description. Write clear, concise docstrings. ### [`Agent`](https://www.union.ai/docs/v2/flyte/integrations/gemini/#agent) A dataclass for bundling agent configuration: from flyteplugins.gemini import Agent agent = Agent( name="my-agent", instructions="You are a helpful assistant.", model="gemini-2.5-flash", tools=[get_weather], max_output_tokens=8192, max_iterations=10, ) | Field | Type | Default | Description | | --- | --- | --- | --- | | `name` | `str` | `"assistant"` | Agent name | | `instructions` | `str` | `"You are a helpful assistant."` | System prompt | | `model` | `str` | `"gemini-2.5-flash"` | Gemini model ID | | `tools` | `list[FunctionTool]` | `[]` | Tools available to the agent | | `max_output_tokens` | `int` | `8192` | Maximum tokens per response | | `max_iterations` | `int` | `10` | Maximum tool-call loop iterations | ### [`run_agent`](https://www.union.ai/docs/v2/flyte/integrations/gemini/#run_agent) Runs a Gemini conversation loop, dispatching tool calls to Flyte tasks until Gemini returns a final response. result = await run_agent( prompt="What's the weather in Tokyo?", tools=[get_weather], model="gemini-2.5-flash", ) You can also pass an `Agent` object: result = await run_agent(prompt="What's the weather?", agent=agent) | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `prompt` | `str` | required | User message | | `tools` | `list[FunctionTool]` | `None` | Tools available to the agent | | `agent` | `Agent` | `None` | Agent config (overrides individual params) | | `model` | `str` | `"gemini-2.5-flash"` | Gemini model ID | | `system` | `str` | `None` | System prompt | | `max_output_tokens` | `int` | `8192` | Maximum tokens per response | | `max_iterations` | `int` | `10` | Maximum iterations (prevents infinite loops) | | `api_key` | `str` | `None` | API key (falls back to `GOOGLE_API_KEY` env var) | [Secrets](https://www.union.ai/docs/v2/flyte/integrations/gemini/#secrets) ---------------------------------------------------------------------------- Store your Google API key as a Flyte secret and expose it as an environment variable: secrets=flyte.Secret("google_api_key", as_env_var="GOOGLE_API_KEY") [API reference](https://www.union.ai/docs/v2/flyte/integrations/gemini/#api-reference) ---------------------------------------------------------------------------------------- See the [Gemini API reference](https://www.union.ai/docs/v2/flyte/api-reference/integrations/gemini) for full details. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/gemini/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Weights & Biases | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Weights & Biases ================ [Weights & Biases](https://wandb.ai/) (W&B) is a platform for tracking machine learning experiments, visualizing metrics and optimizing hyperparameters. This plugin integrates W&B with Flyte, enabling you to: * Automatically initialize W&B runs in your tasks without boilerplate * Link directly from the Flyte UI to your W&B runs and sweeps * Share W&B runs across parent and child tasks * Track distributed training jobs across multiple GPUs and nodes * Run hyperparameter sweeps with parallel agents [Installation](https://www.union.ai/docs/v2/flyte/integrations/wandb/#installation) ------------------------------------------------------------------------------------- pip install flyteplugins-wandb You also need a W&B API key. Store it as a Flyte secret so your tasks can authenticate with W&B. [Quick start](https://www.union.ai/docs/v2/flyte/integrations/wandb/#quick-start) ----------------------------------------------------------------------------------- Here’s a minimal example that logs metrics to W&B from a Flyte task: quick\_start.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/quick_start.py "View source on GitHub") import flyte from flyteplugins.wandb import get_wandb_run, wandb_config, wandb_init env = flyte.TaskEnvironment( name="wandb-example", image=flyte.Image.from_debian_base(name="wandb-example").with_pip_packages( "flyteplugins-wandb" ), secrets=[flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY")], ) @wandb_init @env.task async def train_model() -> str: wandb_run = get_wandb_run() # Your training code here for epoch in range(10): loss = 1.0 / (epoch + 1) wandb_run.log({"epoch": epoch, "loss": loss}) return "Training complete" if __name__ == "__main__": flyte.init_from_config() r = flyte.with_runcontext( custom_context=wandb_config( project="my-project", entity="my-team", ), ).run(train_model) print(f"run url: {r.url}") This example demonstrates the core pattern: 1. **Define a task environment** with the plugin installed and your W&B API key as a secret 2. **Decorate your task** with `@wandb_init` (must be the outermost decorator, above `@env.task`) 3. **Access the run** with `get_wandb_run()` to log metrics 4. **Provide configuration** via `wandb_config()` when running the task The plugin handles calling `wandb.init()` and `wandb.finish()` for you, and automatically adds a link to the W&B run in the Flyte UI. ![UI](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/wandb/ui.png) [What’s next](https://www.union.ai/docs/v2/flyte/integrations/wandb/#whats-next) ---------------------------------------------------------------------------------- This integration guide is split into focused sections, depending on how you want to use Weights & Biases with Flyte: * **[Experiments](https://www.union.ai/docs/v2/flyte/integrations/wandb/experiments) **: Create and manage W&B runs from Flyte tasks. * **[Distributed training](https://www.union.ai/docs/v2/flyte/integrations/wandb/distributed_training) **: Track experiments across multi-GPU and multi-node training jobs. * **[Sweeps](https://www.union.ai/docs/v2/flyte/integrations/wandb/sweeps) **: Run hyperparameter searches and manage sweep execution from Flyte tasks. * **[Downloading logs](https://www.union.ai/docs/v2/flyte/integrations/wandb/downloading_logs) **: Download logs and execution metadata from Weights & Biases. * **[Constraints and best practices](https://www.union.ai/docs/v2/flyte/integrations/wandb/constraints_and_best_practices) **: Learn about limitations, edge cases and recommended patterns. * **[Manual integration](https://www.union.ai/docs/v2/flyte/integrations/wandb/manual) **: Use Weights & Biases directly in Flyte tasks without decorators or helpers. We’ve included additional examples developed while testing edge cases of the plugin [here](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/wandb/examples) . LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/wandb/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # BigQuery | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) BigQuery ======== The BigQuery connector lets you run SQL queries against [Google BigQuery](https://cloud.google.com/bigquery) directly from Flyte tasks. Queries are submitted asynchronously via the BigQuery Jobs API and polled for completion, so they don’t block a worker while waiting for results. The connector supports: * Parameterized SQL queries with typed inputs * Google Cloud service account authentication * Returns query results as DataFrames * Query cancellation on task abort [Installation](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#installation) ---------------------------------------------------------------------------------------- pip install flyteplugins-bigquery This installs the Google Cloud BigQuery client libraries. [Quick start](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#quick-start) -------------------------------------------------------------------------------------- Here’s a minimal example that runs a SQL query on BigQuery: from flyte.io import DataFrame from flyteplugins.bigquery import BigQueryConfig, BigQueryTask config = BigQueryConfig( ProjectID="my-gcp-project", Location="US", ) count_users = BigQueryTask( name="count_users", query_template="SELECT COUNT(*) FROM dataset.users", plugin_config=config, output_dataframe_type=DataFrame, ) This defines a task called `count_users` that runs the query on the configured BigQuery instance. When executed, the connector: 1. Connects to BigQuery using the provided configuration 2. Submits the query asynchronously via the Jobs API 3. Polls until the query completes or fails To run the task, create a `TaskEnvironment` from it and execute it locally or remotely: import flyte bigquery_env = flyte.TaskEnvironment.from_task("bigquery_env", count_users) if __name__ == "__main__": flyte.init_from_config() # Run locally (connector runs in-process, requires credentials locally) run = flyte.with_runcontext(mode="local").run(count_users) # Run remotely (connector runs as a service in your data plane) run = flyte.with_runcontext(mode="remote").run(count_users) print(run.url) The `TaskEnvironment` created by `from_task` does not need an image or pip packages. BigQuery tasks are connector tasks, which means the query executes on the connector service, not in your task container. In `local` mode, the connector runs in-process and requires `flyteplugins-bigquery` and credentials to be available on your machine. [Configuration](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#configuration) ------------------------------------------------------------------------------------------ ### [`BigQueryConfig` parameters](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#bigqueryconfig-parameters) | Field | Type | Required | Description | | --- | --- | --- | --- | | `ProjectID` | `str` | Yes | GCP project ID | | `Location` | `str` | No | BigQuery region (e.g., `"US"`, `"EU"`) | | `QueryJobConfig` | `bigquery.QueryJobConfig` | No | Native BigQuery [QueryJobConfig](https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJobConfig)
object for advanced settings | ### [`BigQueryTask` parameters](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#bigquerytask-parameters) | Parameter | Type | Description | | --- | --- | --- | | `name` | `str` | Unique task name | | `query_template` | `str` | SQL query (whitespace is normalized before execution) | | `plugin_config` | `BigQueryConfig` | Connection configuration | | `inputs` | `Dict[str, Type]` | Named typed inputs bound as query parameters | | `output_dataframe_type` | `Type[DataFrame]` | If set, query results are returned as a `DataFrame` | | `google_application_credentials` | `str` | Name of the Flyte secret containing the GCP service account JSON key | [Authentication](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#authentication) -------------------------------------------------------------------------------------------- Pass the name of a Flyte secret containing your GCP service account JSON key: query = BigQueryTask( name="secure_query", query_template="SELECT * FROM dataset.sensitive_data", plugin_config=config, google_application_credentials="my-gcp-sa-key", ) [Query templating](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#query-templating) ------------------------------------------------------------------------------------------------ Use the `inputs` parameter to define typed inputs for your query. Input values are bound as BigQuery `ScalarQueryParameter` values. ### [Supported input types](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#supported-input-types) | Python type | BigQuery type | | --- | --- | | `int` | `INT64` | | `float` | `FLOAT64` | | `str` | `STRING` | | `bool` | `BOOL` | | `bytes` | `BYTES` | | `datetime` | `DATETIME` | | `list` | `ARRAY` | ### [Parameterized query example](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#parameterized-query-example) from flyte.io import DataFrame events_by_region = BigQueryTask( name="events_by_region", query_template="SELECT * FROM dataset.events WHERE region = @region AND score > @min_score", plugin_config=config, inputs={"region": str, "min_score": float}, output_dataframe_type=DataFrame, ) The query template is normalized before execution: newlines and tabs are replaced with spaces and consecutive whitespace is collapsed. You can format your queries across multiple lines for readability without affecting execution. [Retrieving query results](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#retrieving-query-results) ---------------------------------------------------------------------------------------------------------------- Set `output_dataframe_type` to capture results as a DataFrame: from flyte.io import DataFrame top_customers = BigQueryTask( name="top_customers", query_template=""" SELECT customer_id, SUM(amount) AS total_spend FROM dataset.orders GROUP BY customer_id ORDER BY total_spend DESC LIMIT 100 """, plugin_config=config, output_dataframe_type=DataFrame, ) If you don’t need query results (for example, DDL statements or INSERT queries), omit `output_dataframe_type`. [API reference](https://www.union.ai/docs/v2/flyte/integrations/bigquery/#api-reference) ------------------------------------------------------------------------------------------ See the [BigQuery API reference](https://www.union.ai/docs/v2/flyte/api-reference/integrations/bigquery) for full details. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/bigquery/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Dask | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Dask ==== The Dask plugin lets you run [Dask](https://www.dask.org/) jobs natively on Kubernetes. Flyte provisions a transient Dask cluster for each task execution using the [Dask Kubernetes Operator](https://kubernetes.dask.org/en/latest/operator.html) and tears it down on completion. [When to use this plugin](https://www.union.ai/docs/v2/flyte/integrations/dask/#when-to-use-this-plugin) ---------------------------------------------------------------------------------------------------------- * Parallel Python workloads that outgrow a single machine * Distributed DataFrame operations on large datasets * Workloads that use Dask’s task scheduler for arbitrary computation graphs * Jobs that need to scale NumPy, pandas, or scikit-learn workflows across multiple nodes [Installation](https://www.union.ai/docs/v2/flyte/integrations/dask/#installation) ------------------------------------------------------------------------------------ pip install flyteplugins-dask Your task image must also include the Dask distributed scheduler: image = flyte.Image.from_debian_base(name="dask").with_pip_packages("flyteplugins-dask") [Configuration](https://www.union.ai/docs/v2/flyte/integrations/dask/#configuration) -------------------------------------------------------------------------------------- Create a `Dask` configuration and pass it as `plugin_config` to a `TaskEnvironment`: from flyteplugins.dask import Dask, Scheduler, WorkerGroup dask_config = Dask( scheduler=Scheduler(), workers=WorkerGroup(number_of_workers=4), ) dask_env = flyte.TaskEnvironment( name="dask_env", plugin_config=dask_config, image=image, ) ### [`Dask` parameters](https://www.union.ai/docs/v2/flyte/integrations/dask/#dask-parameters) | Parameter | Type | Description | | --- | --- | --- | | `scheduler` | `Scheduler` | Scheduler pod configuration (defaults to `Scheduler()`) | | `workers` | `WorkerGroup` | Worker group configuration (defaults to `WorkerGroup()`) | ### [`Scheduler` parameters](https://www.union.ai/docs/v2/flyte/integrations/dask/#scheduler-parameters) | Parameter | Type | Description | | --- | --- | --- | | `image` | `str` | Custom scheduler image (must include `dask[distributed]`) | | `resources` | `Resources` | Resource requests for the scheduler pod | ### [`WorkerGroup` parameters](https://www.union.ai/docs/v2/flyte/integrations/dask/#workergroup-parameters) | Parameter | Type | Description | | --- | --- | --- | | `number_of_workers` | `int` | Number of worker pods (default: `1`) | | `image` | `str` | Custom worker image (must include `dask[distributed]`) | | `resources` | `Resources` | Resource requests per worker pod | The scheduler and all workers should use the same Python environment to avoid serialization issues. ### [Accessing the Dask client](https://www.union.ai/docs/v2/flyte/integrations/dask/#accessing-the-dask-client) Inside a Dask task, create a `distributed.Client()` with no arguments. It automatically connects to the provisioned cluster: from distributed import Client @dask_env.task async def my_dask_task(n: int) -> list: client = Client() futures = client.map(lambda x: x + 1, range(n)) return client.gather(futures) [Example](https://www.union.ai/docs/v2/flyte/integrations/dask/#example) -------------------------------------------------------------------------- dask\_example.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/dask/dask_example.py "View source on GitHub") # /// script # requires-python = "==3.13" # dependencies = [\ # "flyte>=2.0.0b52",\ # "flyteplugins-dask",\ # "distributed"\ # ] # main = "hello_dask_nested" # params = "" # /// import asyncio import typing from distributed import Client from flyteplugins.dask import Dask, Scheduler, WorkerGroup import flyte.remote import flyte.storage from flyte import Resources image = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("flyteplugins-dask") dask_config = Dask( scheduler=Scheduler(), workers=WorkerGroup(number_of_workers=4), ) task_env = flyte.TaskEnvironment( name="hello_dask", resources=Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image ) dask_env = flyte.TaskEnvironment( name="dask_env", plugin_config=dask_config, image=image, resources=Resources(cpu="1", memory="1Gi"), depends_on=[task_env], ) @task_env.task() async def hello_dask(): await asyncio.sleep(5) print("Hello from the Dask task!") @dask_env.task async def hello_dask_nested(n: int = 3) -> typing.List[int]: print("running dask task") t = asyncio.create_task(hello_dask()) client = Client() futures = client.map(lambda x: x + 1, range(n)) res = client.gather(futures) await t return res if __name__ == "__main__": flyte.init_from_config() r = flyte.run(hello_dask_nested) print(r.name) print(r.url) r.wait() [API reference](https://www.union.ai/docs/v2/flyte/integrations/dask/#api-reference) -------------------------------------------------------------------------------------- See the [Dask API reference](https://www.union.ai/docs/v2/flyte/api-reference/integrations/dask) for full details. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/dask/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # OpenAI | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) OpenAI ====== The OpenAI plugin provides a drop-in replacement for the [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/) `function_tool` decorator. It lets you use Flyte tasks as tools in agentic workflows so that tool calls run as tracked, reproducible Flyte task executions. [When to use this plugin](https://www.union.ai/docs/v2/flyte/integrations/openai/#when-to-use-this-plugin) ------------------------------------------------------------------------------------------------------------ * Building agentic workflows with the OpenAI Agents SDK on Flyte * You want tool calls to run as Flyte tasks with full observability, retries, and caching * You want to combine LLM agents with existing Flyte pipelines [Installation](https://www.union.ai/docs/v2/flyte/integrations/openai/#installation) -------------------------------------------------------------------------------------- pip install flyteplugins-openai Requires `openai-agents >= 0.2.4`. [Usage](https://www.union.ai/docs/v2/flyte/integrations/openai/#usage) ------------------------------------------------------------------------ The plugin provides a single decorator, `function_tool`, that wraps Flyte tasks as OpenAI agent tools. ### [`function_tool`](https://www.union.ai/docs/v2/flyte/integrations/openai/#function_tool) When applied to a Flyte task (a function decorated with `@env.task`), `function_tool` makes that task available as an OpenAI `FunctionTool`. The agent can call it like any other tool, and the call executes as a Flyte task. When applied to a regular function or a `@flyte.trace`\-decorated function, it delegates directly to the OpenAI Agents SDK’s built-in `function_tool`. ### [Basic pattern](https://www.union.ai/docs/v2/flyte/integrations/openai/#basic-pattern) 1. Define a `TaskEnvironment` with your image and secrets 2. Decorate your task functions with `@function_tool` and `@env.task` 3. Pass the tools to an `Agent` 4. Run the agent from another Flyte task from agents import Agent, Runner from flyteplugins.openai.agents import function_tool env = flyte.TaskEnvironment( name="openai_agents", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from_uv_script(__file__, name="openai_agents_image"), secrets=flyte.Secret("openai_api_key", as_env_var="OPENAI_API_KEY"), ) @function_tool @env.task async def get_weather(city: str) -> Weather: """Get the weather for a given city.""" return Weather(city=city, temperature_range="14-20C", conditions="Sunny") agent = Agent( name="Weather Agent", instructions="You are a helpful agent.", tools=[get_weather], ) @env.task async def main() -> str: result = await Runner.run(agent, input="What's the weather in Tokyo?") return result.final_output The docstring on each `@function_tool` task is sent to the LLM as the tool description. Write clear, concise docstrings that describe what the tool does and what its parameters mean. ### [Secrets](https://www.union.ai/docs/v2/flyte/integrations/openai/#secrets) Store your OpenAI API key as a Flyte secret and expose it as an environment variable: secrets=flyte.Secret("openai_api_key", as_env_var="OPENAI_API_KEY") [Example](https://www.union.ai/docs/v2/flyte/integrations/openai/#example) ---------------------------------------------------------------------------- agents\_tools.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/openai/openai/agents_tools.py "View source on GitHub") """OpenAI Agents with Flyte, basic tool example. Usage: Create secret: ``` flyte create secret openai_api_key uv run agents_tools.py ``` """ # /// script # requires-python = "==3.13" # dependencies = [\ # "flyte>=2.0.0b52",\ # "flyteplugins-openai>=2.0.0b7",\ # "openai-agents>=0.2.4",\ # "pydantic>=2.10.6",\ # ] # main = "main" # params = "" # /// from agents import Agent, Runner from pydantic import BaseModel import flyte from flyteplugins.openai.agents import function_tool env = flyte.TaskEnvironment( name="openai_agents_tools", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from_uv_script(__file__, name="openai_agents_image"), secrets=flyte.Secret("openai_api_key", as_env_var="OPENAI_API_KEY"), ) class Weather(BaseModel): city: str temperature_range: str conditions: str @function_tool @env.task async def get_weather(city: str) -> Weather: """Get the weather for a given city.""" return Weather(city=city, temperature_range="14-20C", conditions="Sunny with wind.") agent = Agent( name="Hello world", instructions="You are a helpful agent.", tools=[get_weather], ) @env.task async def main() -> str: result = await Runner.run(agent, input="What's the weather in Tokyo?") print(result.final_output) return result.final_output if __name__ == "__main__": flyte.init_from_config() run = flyte.run(main) print(run.url) run.wait() [API reference](https://www.union.ai/docs/v2/flyte/integrations/openai/#api-reference) ---------------------------------------------------------------------------------------- See the [OpenAI API reference](https://www.union.ai/docs/v2/flyte/api-reference/integrations/openai) for full details. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/openai/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Joining the community | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Joining the community ===================== Keeping the lines of communication open is important in growing and maintain the Flyte community. Please join us on: [![Flyte Slack](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/community/joining-the-community/slack-chat-pink.svg)](https://slack.flyte.org/) [![GitHub Discussion](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/community/joining-the-community/github-discussion-badge.svg)](https://github.com/flyteorg/flyte/discussions) [![Twitter](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/community/joining-the-community/twitter-social-blue.svg)](https://twitter.com/flyteorg) [![LinkedIn](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/community/joining-the-community/linkedin-social-lightblue.svg)](https://www.linkedin.com/groups/13962256) [Community sync](https://www.union.ai/docs/v2/flyte/community/joining-the-community/#community-sync) ------------------------------------------------------------------------------------------------------ 1. **When**: First Tuesday of every month, 9:00 AM Pacific Time. 2. **Where**: Live streamed on [YouTube](https://www.youtube.com/@flyteorg/streams) and [LinkedIn](https://www.linkedin.com/company/union-ai/events/) . 3. **Watch the recordings**: [here](https://www.youtube.com/live/d81Jd4rfmzw?feature=shared) . 4. **Import the public calendar**: [here](https://lists.lfaidata.foundation/g/flyte-announce/ics/12031983/2145304139/feed.ics) to not miss any event. 5. **Want to present?** Fill out [this form](https://tally.so/r/wgN8LM) . We’re eager to learn from you! You’re welcome to join and learn from other community members sharing their experiences with Flyte or any other technology from the AI ecosystem. [Contributor’s sync](https://www.union.ai/docs/v2/flyte/community/joining-the-community/#contributors-sync) ------------------------------------------------------------------------------------------------------------- 1. **When**: Every 2 weeks on Thursdays. Alternating schedule between 11:00 AM PT and 7:00 AM PT. 2. **Where**: Live on [Zoom](https://zoom-lfx.platform.linuxfoundation.org/meeting/92309721545?password=c93d76a7-801a-47c6-9916-08e38e5a5c1f) . 3. **Purpose**: Address questions from new contributors, discuss active initiatives, and RFCs. 4. **Import the public calendar**: [here](https://lists.lfaidata.foundation/g/flyte-announce/ics/12031983/2145304139/feed.ics) to not miss any event. [Newsletter](https://www.union.ai/docs/v2/flyte/community/joining-the-community/#newsletter) ---------------------------------------------------------------------------------------------- [Join the Flyte mailing list](https://lists.lfaidata.foundation/g/flyte-announce/join) to receive the monthly newsletter. [Slack guidelines](https://www.union.ai/docs/v2/flyte/community/joining-the-community/#slack-guidelines) ---------------------------------------------------------------------------------------------------------- Flyte strives to build and maintain an open, inclusive, productive, and self-governing open source community. In consequence, we expect all community members to respect the following guidelines: ### [Abide by the](https://www.union.ai/docs/v2/flyte/community/joining-the-community/#abide-by-the-lf) [LF’s Code of Conduct](https://lfprojects.org/policies/code-of-conduct/) As a Linux Foundation project, we must enforce the rules that govern professional and positive open source communities. ### [Avoid using DMs and @mentions](https://www.union.ai/docs/v2/flyte/community/joining-the-community/#avoid-using-dms-and-mentions) Whenever possible, post your questions and responses in public channels so other community members can benefit from the conversation and outcomes. Exceptions to this are when you need to share private or sensitive information. In such a case, the outcome should still be shared publicly. Limit the use of `@mentions` of other community members to be considerate of notification noise. ### [Make use of threads](https://www.union.ai/docs/v2/flyte/community/joining-the-community/#make-use-of-threads) Threads help us keep conversations contained and organized, reducing the time it takes to give you the support you need. **Thread best practices:** * Don’t break your question into multiple messages. Put everything in one. * For long questions, write a few sentences in the first message, and put the rest in a thread. * If there’s a code snippet (more than 5 lines of code), put it inside the thread. * Avoid using the “Also send to channel” feature unless it’s really necessary. * If your question contains multiple questions, make sure to break them into multiple messages, so each could be answered in a separate thread. ### [Do not post the same question across multiple channels](https://www.union.ai/docs/v2/flyte/community/joining-the-community/#do-not-post-the-same-question-across-multiple-channels) If you consider that a question needs to be shared on other channels, ask it once and then indicate explicitly that you’re cross-posting. If you’re having a tough time getting the support you need (or aren’t sure where to go!), please DM `@David Espejo` or `@Samhita Alla` for support. ### [Do not solicit members of our Slack](https://www.union.ai/docs/v2/flyte/community/joining-the-community/#do-not-solicit-members-of-our-slack) The Flyte community exists to collaborate with, learn from, and support one another. It is not a space to pitch your products or services directly to our members via public channels, private channels, or direct messages. We are excited to have a growing presence from vendors to help answer questions from community members as they may arise, but we have a strict 3-strike policy against solicitation: * **First occurrence**: We’ll give you a friendly but public reminder that the behavior is inappropriate according to our guidelines. * **Second occurrence**: We’ll send you a DM warning that any additional violations will result in removal from the community. * **Third occurrence**: We’ll delete or ban your account. We reserve the right to ban users without notice if they are clearly spamming our community members. If you want to promote a product or service, go to the `#shameless-promotion` channel and make sure to follow these rules: * Don’t post more than two promotional posts per week. * Non-relevant topics aren’t allowed. Messages that don’t follow these rules will be deleted. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/community/joining-the-community/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Migration | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Migration ========= Guides for migrating to Flyte 2 from other systems. From Flyte 1 to 2 What’s new in Flyte 2 — pure Python execution, simplified API, fine-grained reproducibility — and how to port a Flyte 1 codebase. From Airflow to Flyte Mapping from Airflow concepts (DAGs, operators, schedules, XCom, trigger rules) to their Flyte 2 equivalents. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/migration/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Advanced project | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Advanced project: LLM reporting agent ===================================== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This example demonstrates a resilient agentic report generator that showcases Flyte 2.0’s advanced features for building production-grade AI workflows. [What you’ll build](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/#what-youll-build) ------------------------------------------------------------------------------------------------------- A batch report generator that: 1. Processes multiple topics in parallel 2. Iteratively critiques and refines each report until it meets a quality threshold 3. Produces multiple output formats (Markdown, HTML, summary) for each report 4. Serves results through an interactive UI [Concepts covered](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/#concepts-covered) ------------------------------------------------------------------------------------------------------ | Feature | Description | | --- | --- | | `ReusePolicy` | Keep containers warm for high-throughput batch processing | | `@flyte.trace` | Checkpoint LLM calls for recovery and observability | | `RetryStrategy` | Handle transient API failures gracefully | | `flyte.group` | Organize parallel batches and iterations in the UI | | `asyncio.gather` | Fan out to process multiple topics concurrently | | Pydantic models | Structured LLM outputs | | `AppEnvironment` | Deploy interactive Streamlit apps | | `RunOutput` | Connect apps to pipeline outputs | [Architecture](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/#architecture) ---------------------------------------------------------------------------------------------- flowchart TD A\[Topics List\] --> B B\["report\_batch\_pipeline
driver\_env"\] subgraph B1 \["refine\_all (parallel)"\] direction LR R1\["refine\_report
topic 1"\] R2\["refine\_report
topic 2"\] R3\["refine\_report
topic N"\] end B --> B1 subgraph B2 \["format\_all (parallel)"\] direction LR F1\["format\_outputs
report 1"\] F2\["format\_outputs
report 2"\] F3\["format\_outputs
report N"\] end B1 --> B2 B2 --> C\["Output: List of Dirs"\] Each `refine_report` task runs in a reusable container (`llm_env`) and performs multiple LLM calls through traced functions: flowchart TD A\[Topic\] --> B\["generate\_initial\_draft
@flyte.trace"\] B --> C subgraph C \["refinement\_loop"\] direction TB D\["critique\_content
@flyte.trace"\] -->|score >= threshold| E\[exit loop\] D -->|score < threshold| F\["revise\_content
@flyte.trace"\] F --> D end C --> G\[Refined Report\] [Prerequisites](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/#prerequisites) ------------------------------------------------------------------------------------------------ * A Flyte account with an active project * An OpenAI API key stored as a secret named `openai-api-key` To create the secret: flyte secret create openai-api-key [Parts](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/#parts) -------------------------------------------------------------------------------- 1. **[Resilient generation](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/resilient-generation) **: Set up reusable environments, traced LLM calls, and retry strategies 2. **[Agentic refinement](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/agentic-refinement) **: Build the iterative critique-and-revise loop 3. **[Parallel outputs](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/parallel-outputs) **: Generate multiple formats concurrently 4. **[Serving app](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/serving-app) **: Deploy an interactive UI for report generation [Key takeaways](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/#key-takeaways) ------------------------------------------------------------------------------------------------ 1. **Reusable environments for batch processing**: `ReusePolicy` keeps containers warm, enabling efficient processing of multiple topics without cold start overhead. With 5 topics × ~7 LLM calls each, the reusable pool handles ~35 calls efficiently. 2. **Checkpointed LLM calls**: `@flyte.trace` provides automatic checkpointing at the function level, enabling recovery without re-running expensive API calls. 3. **Agentic patterns**: The generate-critique-revise loop demonstrates how to build self-improving AI workflows with clear observability through `flyte.group`. 4. **Parallel fan-out**: `asyncio.gather` processes multiple topics concurrently, maximizing throughput by running refinement tasks in parallel across the batch. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Contributing docs and examples | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Contributing docs and examples ============================== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/community/contributing-docs/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. We welcome contributions to the docs and examples for both Flyte and Union. This section will explain how the docs site works, how to author and build it locally, and how to publish your changes. [The combined Flyte and Union docs site](https://www.union.ai/docs/v2/flyte/community/contributing-docs/#the-combined-flyte-and-union-docs-site) -------------------------------------------------------------------------------------------------------------------------------------------------- As the primary maintainer and contributor of the open-source Flyte project, Union AI is responsible for hosting the Flyte documentation. Additionally, Union AI is also the company behind the commercial Union.ai product, which is based on Flyte. Since Flyte and Union.ai share a lot of common functionality, much of the documentation content is common between the two. However, there are some significant differences between not only Flyte and Union.ai but also among the different Union.ai product offering (Serverless, BYOC, and Self-managed). To effectively and efficiently maintain the documentation for all of these variants, we employ a single-source-of-truth approach where: * All content is stored in a single GitHub repository, [`unionai/unionai-docs`](https://github.com/unionai/unionai-docs) * All content is published on a single website, [`www.union.ai/docs`](https://www.union.ai/docs/v2) . * The website has a variant selector at the top of the page that lets you choose which variant you want to view: * Flyte OSS * Union Serverless * Union BYOC * Union Self-managed * There is also version selector. Currently two versions are available: * v1 (the original docs for Flyte/Union 1.x) * v2 (the new docs for Flyte/Union 2.0, which is the one you are currently viewing) [Versions](https://www.union.ai/docs/v2/flyte/community/contributing-docs/#versions) -------------------------------------------------------------------------------------- The two versions of the docs are stored in separate branches of the GitHub repository: * [`v1` branch](https://github.com/unionai/unionai-docs/tree/v1) for the v1 docs. * [`main` branch](https://github.com/unionai/unionai-docs) for the v2 docs. See [Versions](https://www.union.ai/docs/v2/flyte/community/contributing-docs/versions) for more details. [Common build infrastructure](https://www.union.ai/docs/v2/flyte/community/contributing-docs/#common-build-infrastructure) ---------------------------------------------------------------------------------------------------------------------------- The build infrastructure for the docs site (Hugo configuration, layouts, themes, build scripts, and Python tools) is maintained in a separate repository, [`unionai/unionai-docs-infra`](https://github.com/unionai/unionai-docs-infra) , which is imported as a [git submodule](https://git-scm.com/book/en/v2/Git-Tools-Submodules) at `unionai-docs-infra/` in the `unionai-docs` repository. This means both the `main` (v2) and `v1` content branches share the same build infrastructure. Changes to the build system are made once in `unionai-docs-infra` and are picked up by both branches, keeping them in sync without duplicating build logic. [Variants](https://www.union.ai/docs/v2/flyte/community/contributing-docs/#variants) -------------------------------------------------------------------------------------- Within each branch the multiple variants are supported by using conditional rendering: * Each page of content has a `variants` front matter field that specifies which variants the page is applicable to. * Within each page, rendering logic can be used to include or exclude content based on the selected variant. The result is that: * Content that is common to all variants is authored and stored once. There is no need to keep multiple copies of the same content in-sync. * Content specific to a variant is conditionally rendered based on the selected variant. See [Variants](https://www.union.ai/docs/v2/flyte/community/contributing-docs/variants) for more details. [Both Flyte and Union docs are open source](https://www.union.ai/docs/v2/flyte/community/contributing-docs/#both-flyte-and-union-docs-are-open-source) -------------------------------------------------------------------------------------------------------------------------------------------------------- Since the docs are now combined in one repository, and the Flyte docs are open source, the Union docs are also open source. All the docs are available for anyone to contribute to: Flyte contributors, Union customers, and Union employees. If you are a Flyte contributor, you will be contributing docs related to Flyte features and functionality, but in many cases these features and functionality will also be available in Union. Because the docs site is a single source for all the documentation, when you make changes related to Flyte that are also valid for Union you do so in the same place. This is by design and is a key feature of the docs site. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/community/contributing-docs/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/community/contributing-docs/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Configure tasks | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Configure tasks =============== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. As we saw in [**Quickstart**](https://www.union.ai/docs/v2/flyte/user-guide/quickstart) , you can run any Python function as a task in Flyte just by decorating it with `@env.task`. This allows you to run your Python code in a distributed manner, with each function running in its own container. Flyte manages the spinning up of the containers, the execution of the code, and the passing of data between the tasks. The simplest possible case is a `TaskEnvironment` with only a `name` parameter, and an `env.task` decorator, with no parameters: task\_config.py [](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/task_config.py "View source on GitHub") env = flyte.TaskEnvironment(name="my_env") @env.task async def my_task(name:str) -> str: return f"Hello {name}!" Notice how the `TaskEnvironment` is assigned to the variable `env` and then that variable is used in the `@env.task`. This is what connects the `TaskEnvironment` to the task definition. In the following we will often use `@env.task` generically to refer to the decorator, but it is important to remember that it is actually a decorator attached to a specific `TaskEnvironment` object, and the `env` part can be any variable name you like. This will run your task in the default container environment with default settings. But, of course, one of the key advantages of Flyte is the ability to control the software environment, hardware environment, and other execution parameters for each task, right in your Python code. In this section we will explore the various configuration options available for tasks in Flyte. [Task configuration levels](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/#task-configuration-levels) -------------------------------------------------------------------------------------------------------------------------- Task configuration is done at three levels. From most general to most specific, they are: * The `TaskEnvironment` level: setting parameters when defining the `TaskEnvironment` object. * The `@env.task` decorator level: Setting parameters in the `@env.task` decorator when defining a task function. * The task invocation level: Using the `task.override()` method when invoking task execution. Each level has its own set of parameters, and some parameters are shared across levels. For shared parameters, the more specific level will override the more general one. ### [Example](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/#example) Here is an example of how these levels work together, showing each level with all available parameters: task\_config.py [](https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/task_config.py "View source on GitHub") # Level 1: TaskEnvironment - Base configuration env_2 = flyte.TaskEnvironment( name="data_processing_env", image=flyte.Image.from_debian_base(), resources=flyte.Resources(cpu=1, memory="512Mi"), env_vars={"MY_VAR": "value"}, # secrets=flyte.Secret(key="openapi_key", as_env_var="MY_API_KEY"), cache="disable", # pod_template=my_pod_template, # reusable=flyte.ReusePolicy(replicas=2, idle_ttl=300), depends_on=[another_env], description="Data processing task environment", # plugin_config=my_plugin_config ) # Level 2: Decorator - Override some environment settings @env_2.task( short_name="process", # secrets=flyte.Secret(key="openapi_key", as_env_var="MY_API_KEY_2"), cache="auto", # pod_template=my_pod_template, report=True, max_inline_io_bytes=100 * 1024, retries=3, timeout=60, docs="This task processes data and generates a report." ) async def process_data(data_path: str) -> str: return f"Processed {data_path}" @env_2.task async def invoke_process_data() -> str: result = await process_data.override( resources=flyte.Resources(cpu=4, memory="2Gi"), env_vars={"MY_VAR": "new_value"}, # secrets=flyte.Secret(key="openapi_key", as_env_var="MY_API_KEY_3"), cache="auto", max_inline_io_bytes=100 * 1024, retries=3, timeout=60 )("input.csv") return result [Task configuration parameters](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/#task-configuration-parameters) ---------------------------------------------------------------------------------------------------------------------------------- Each parameter is documented in detail on its dedicated page in this section. For the complete parameter interaction matrix showing which parameters can be set at which level, and for full type signatures and constraints, see the [`TaskEnvironment` API reference](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/taskenvironment) . | Parameter | Set at | Details | | --- | --- | --- | | **name** | `TaskEnvironment` only | [Additional task settings](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/additional-task-settings)
• [`TaskEnvironment` API ref](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/taskenvironment) | | **image** | `TaskEnvironment` only | [Container images](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/container-images)
• [`Image` API ref](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/image) | | **depends\_on** | `TaskEnvironment` only | [Multiple environments](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/multiple-environments) | | **description** | `TaskEnvironment` only | [Additional task settings](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/additional-task-settings) | | **plugin\_config** | `TaskEnvironment` only | [Task plugins](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/task-plugins) | | **resources** | `TaskEnvironment`, `override`\* | [Resources](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/resources)
• [`Resources` API ref](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/resources) | | **env\_vars** | `TaskEnvironment`, `override`\* | [Additional task settings](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/additional-task-settings#environment-variables) | | **secrets** | `TaskEnvironment`, `override`\* | [Secrets](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets)
• [`Secret` API ref](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/secret) | | **cache** | All three levels | [Caching](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/caching)
• [`Cache` API ref](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/cache) | | **pod\_template** | All three levels | [Pod templates](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/pod-templates)
• [`PodTemplate` API ref](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/podtemplate) | | **reusable** | `TaskEnvironment`, `override` | [Reusable containers](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/reusable-containers)
• [`ReusePolicy` API ref](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/reusepolicy) | | **interruptible** | All three levels | [Interruptible tasks](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/interruptible-tasks-and-queues) | | **queue** | All three levels | [Queues](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/queues) | | **short\_name** | `@env.task`, `override` | [Additional task settings](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/additional-task-settings) | | **retries** | `@env.task`, `override` | [Retries and timeouts](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/retries-and-timeouts)
• [`RetryStrategy` API ref](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/retrystrategy) | | **timeout** | `@env.task`, `override` | [Retries and timeouts](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/retries-and-timeouts)
• [`Timeout` API ref](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/timeout) | | **max\_inline\_io\_bytes** | `@env.task`, `override` | [Additional task settings](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/additional-task-settings#inline-io-threshold) | | **links** | `@env.task`, `override` | [Additional task settings](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/additional-task-settings#links) | | **report** | `@env.task` only | [Additional task settings](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/additional-task-settings#report) | | **triggers** | `@env.task` only | [Triggers](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/triggers)
• [`Trigger` API ref](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/packages/flyte/trigger) | | **docs** | `@env.task` only | [Additional task settings](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/additional-task-settings#docs) | \*When `reusable` is set, `resources`, `env_vars`, and `secrets` can only be overridden via `task.override()` with `reusable="off"` in the same call. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Databricks | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Databricks ========== The Databricks plugin lets you run PySpark jobs on [Databricks](https://www.databricks.com/) clusters directly from Flyte tasks. You write normal PySpark code in a Flyte task, and the plugin submits it to Databricks via the [Jobs API 2.1](https://docs.databricks.com/api/workspace/jobs/submit) . The connector handles job submission, polling, and cancellation. The plugin supports: * Running PySpark tasks on new or existing Databricks clusters * Full Spark configuration (driver/executor memory, cores, instances) * Databricks cluster auto-scaling * API token-based authentication [Installation](https://www.union.ai/docs/v2/flyte/integrations/databricks/#installation) ------------------------------------------------------------------------------------------ pip install flyteplugins-databricks This also installs `flyteplugins-spark` as a dependency, since the Databricks plugin extends the Spark plugin. [Quick start](https://www.union.ai/docs/v2/flyte/integrations/databricks/#quick-start) ---------------------------------------------------------------------------------------- Create a `Databricks` configuration and pass it as `plugin_config` to a `TaskEnvironment`: from flyteplugins.databricks import Databricks import flyte image = ( flyte.Image.from_base("databricksruntime/standard:16.4-LTS") .clone(name="spark", registry="ghcr.io/flyteorg", extendable=True) .with_env_vars({"UV_PYTHON": "/databricks/python3/bin/python"}) .with_pip_packages("flyteplugins-databricks", pre=True) ) databricks_conf = Databricks( spark_conf={ "spark.driver.memory": "2000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", }, executor_path="/databricks/python3/bin/python", databricks_conf={ "run_name": "flyte databricks plugin", "new_cluster": { "spark_version": "13.3.x-scala2.12", "node_type_id": "m6i.large", "autoscale": {"min_workers": 1, "max_workers": 2}, }, "timeout_seconds": 3600, "max_retries": 1, }, databricks_instance="myaccount.cloud.databricks.com", databricks_token="DATABRICKS_TOKEN", ) databricks_env = flyte.TaskEnvironment( name="databricks_env", resources=flyte.Resources(cpu=(1, 2), memory=("3000Mi", "5000Mi")), plugin_config=databricks_conf, image=image, ) Then use the environment to decorate your task: @databricks_env.task async def hello_databricks() -> float: spark = flyte.ctx().data["spark_session"] # Use spark as a normal SparkSession count = spark.sparkContext.parallelize(range(100)).count() return float(count) [Configuration](https://www.union.ai/docs/v2/flyte/integrations/databricks/#configuration) -------------------------------------------------------------------------------------------- The `Databricks` config extends the [Spark](https://www.union.ai/docs/v2/flyte/integrations/spark) config with Databricks-specific fields. ### [Spark fields (inherited)](https://www.union.ai/docs/v2/flyte/integrations/databricks/#spark-fields-inherited) | Parameter | Type | Description | | --- | --- | --- | | `spark_conf` | `Dict[str, str]` | Spark configuration key-value pairs | | `hadoop_conf` | `Dict[str, str]` | Hadoop configuration key-value pairs | | `executor_path` | `str` | Path to the Python binary on the Databricks cluster (e.g., `/databricks/python3/bin/python`) | | `applications_path` | `str` | Path to the main application file | ### [Databricks-specific fields](https://www.union.ai/docs/v2/flyte/integrations/databricks/#databricks-specific-fields) | Parameter | Type | Description | | --- | --- | --- | | `databricks_conf` | `Dict[str, Union[str, dict]]` | Databricks [run-submit](https://docs.databricks.com/api/workspace/jobs/submit)
job configuration. Must contain either `existing_cluster_id` or `new_cluster` | | `databricks_instance` | `str` | Your workspace domain (e.g., `myaccount.cloud.databricks.com`). Can also be set via the `FLYTE_DATABRICKS_INSTANCE` env var on the connector | | `databricks_token` | `str` | Name of the Flyte secret containing the Databricks API token | ### [`databricks_conf` structure](https://www.union.ai/docs/v2/flyte/integrations/databricks/#databricks_conf-structure) The `databricks_conf` dict maps to the Databricks run-submit API payload. Key fields: | Field | Description | | --- | --- | | `new_cluster` | Cluster spec with `spark_version`, `node_type_id`, `autoscale`, etc. | | `existing_cluster_id` | ID of an existing cluster to use instead of creating a new one | | `run_name` | Display name in the Databricks UI | | `timeout_seconds` | Maximum job duration | | `max_retries` | Number of retries before marking the job as failed | The connector automatically injects the Docker image, Spark configuration, and environment variables from the task container into the cluster spec. [Authentication](https://www.union.ai/docs/v2/flyte/integrations/databricks/#authentication) ---------------------------------------------------------------------------------------------- Store your Databricks API token as a Flyte secret. The `databricks_token` parameter specifies the secret name: databricks_conf = Databricks( # ... databricks_token="DATABRICKS_TOKEN", ) [Accessing the Spark session](https://www.union.ai/docs/v2/flyte/integrations/databricks/#accessing-the-spark-session) ------------------------------------------------------------------------------------------------------------------------ Inside a Databricks task, the `SparkSession` is available through the task context, just like the [Spark plugin](https://www.union.ai/docs/v2/flyte/integrations/spark) : @databricks_env.task async def my_databricks_task() -> float: spark = flyte.ctx().data["spark_session"] df = spark.read.parquet("s3://my-bucket/data.parquet") return float(df.count()) [API reference](https://www.union.ai/docs/v2/flyte/integrations/databricks/#api-reference) -------------------------------------------------------------------------------------------- See the [Databricks API reference](https://www.union.ai/docs/v2/flyte/api-reference/integrations/databricks) for full details. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/databricks/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Spark | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Spark ===== The Spark plugin lets you run [Apache Spark](https://spark.apache.org/) jobs natively on Kubernetes. Flyte manages the full cluster lifecycle: provisioning a transient Spark cluster for each task execution, running the job, and tearing the cluster down on completion. Under the hood, the plugin uses the [Spark on Kubernetes Operator](https://github.com/GoogleCloudPlatform/spark-on-k8s-operator) to create and manage Spark applications. No external Spark service or long-running cluster is required. [When to use this plugin](https://www.union.ai/docs/v2/flyte/integrations/spark/#when-to-use-this-plugin) ----------------------------------------------------------------------------------------------------------- * Large-scale data processing and ETL pipelines * Jobs that benefit from Spark’s distributed execution engine (Spark SQL, PySpark, Spark MLlib) * Workloads that need Hadoop-compatible storage access (S3, GCS, HDFS) [Installation](https://www.union.ai/docs/v2/flyte/integrations/spark/#installation) ------------------------------------------------------------------------------------- pip install flyteplugins-spark [Configuration](https://www.union.ai/docs/v2/flyte/integrations/spark/#configuration) --------------------------------------------------------------------------------------- Create a `Spark` configuration and pass it as `plugin_config` to a `TaskEnvironment`: from flyteplugins.spark import Spark spark_config = Spark( spark_conf={ "spark.driver.memory": "3000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", }, ) spark_env = flyte.TaskEnvironment( name="spark_env", plugin_config=spark_config, image=image, ) ### [`Spark` parameters](https://www.union.ai/docs/v2/flyte/integrations/spark/#spark-parameters) | Parameter | Type | Description | | --- | --- | --- | | `spark_conf` | `Dict[str, str]` | Spark configuration key-value pairs (e.g., executor memory, cores, instances) | | `hadoop_conf` | `Dict[str, str]` | Hadoop configuration key-value pairs (e.g., S3/GCS access settings) | | `executor_path` | `str` | Path to the Python binary for PySpark executors | | `applications_path` | `str` | Path to the main Spark application file | | `driver_pod` | `PodTemplate` | Pod template for the Spark driver pod | | `executor_pod` | `PodTemplate` | Pod template for the Spark executor pods | ### [Accessing the Spark session](https://www.union.ai/docs/v2/flyte/integrations/spark/#accessing-the-spark-session) Inside a Spark task, the `SparkSession` is available through the task context: from flyte._context import internal_ctx @spark_env.task async def my_spark_task() -> float: ctx = internal_ctx() spark = ctx.data.task_context.data["spark_session"] # Use spark as a normal SparkSession df = spark.read.parquet("s3://my-bucket/data.parquet") return df.count() ### [Overriding configuration at runtime](https://www.union.ai/docs/v2/flyte/integrations/spark/#overriding-configuration-at-runtime) You can override Spark configuration for individual task calls using `.override()`: from copy import deepcopy updated_config = deepcopy(spark_config) updated_config.spark_conf["spark.executor.instances"] = "4" result = await my_spark_task.override(plugin_config=updated_config)() [Example](https://www.union.ai/docs/v2/flyte/integrations/spark/#example) --------------------------------------------------------------------------- spark\_example.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/spark/spark_example.py "View source on GitHub") # /// script # requires-python = "==3.13" # dependencies = [\ # "flyte>=2.0.0b52",\ # "flyteplugins-spark"\ # ] # main = "hello_spark_nested" # params = "3" # /// import random from copy import deepcopy from operator import add from flyteplugins.spark.task import Spark import flyte.remote from flyte._context import internal_ctx image = ( flyte.Image.from_base("apache/spark-py:v3.4.0") .clone(name="spark", python_version=(3, 10), registry="ghcr.io/flyteorg") .with_pip_packages("flyteplugins-spark", pre=True) ) task_env = flyte.TaskEnvironment( name="get_pi", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image ) spark_conf = Spark( spark_conf={ "spark.driver.memory": "3000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", "spark.kubernetes.file.upload.path": "/opt/spark/work-dir", "spark.jars": "https://storage.googleapis.com/hadoop-lib/gcs/gcs-connector-hadoop3-latest.jar,https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/3.2.2/hadoop-aws-3.2.2.jar,https://repo1.maven.org/maven2/com/amazonaws/aws-java-sdk-bundle/1.12.262/aws-java-sdk-bundle-1.12.262.jar", }, ) spark_env = flyte.TaskEnvironment( name="spark_env", resources=flyte.Resources(cpu=(1, 2), memory=("3000Mi", "5000Mi")), plugin_config=spark_conf, image=image, depends_on=[task_env], ) def f(_): x = random.random() * 2 - 1 y = random.random() * 2 - 1 return 1 if x**2 + y**2 <= 1 else 0 @task_env.task async def get_pi(count: int, partitions: int) -> float: return 4.0 * count / partitions @spark_env.task async def hello_spark_nested(partitions: int = 3) -> float: n = 1 * partitions ctx = internal_ctx() spark = ctx.data.task_context.data["spark_session"] count = spark.sparkContext.parallelize(range(1, n + 1), partitions).map(f).reduce(add) return await get_pi(count, partitions) @task_env.task async def spark_overrider(executor_instances: int = 3, partitions: int = 4) -> float: updated_spark_conf = deepcopy(spark_conf) updated_spark_conf.spark_conf["spark.executor.instances"] = str(executor_instances) return await hello_spark_nested.override(plugin_config=updated_spark_conf)(partitions=partitions) if __name__ == "__main__": flyte.init_from_config() r = flyte.run(hello_spark_nested) print(r.name) print(r.url) r.wait() [API reference](https://www.union.ai/docs/v2/flyte/integrations/spark/#api-reference) --------------------------------------------------------------------------------------- See the [Spark API reference](https://www.union.ai/docs/v2/flyte/api-reference/integrations/spark) for full details. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/spark/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # PyTorch | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) PyTorch ======= The PyTorch plugin lets you run distributed [PyTorch](https://pytorch.org/) training jobs natively on Kubernetes. It uses the [Kubeflow Training Operator](https://github.com/kubeflow/training-operator) to manage multi-node training with PyTorch’s elastic launch (`torchrun`). [When to use this plugin](https://www.union.ai/docs/v2/flyte/integrations/pytorch/#when-to-use-this-plugin) ------------------------------------------------------------------------------------------------------------- * Single-node or multi-node distributed training with `DistributedDataParallel` (DDP) * Elastic training that can scale up and down during execution * Any workload that uses `torch.distributed` for data-parallel or model-parallel training [Installation](https://www.union.ai/docs/v2/flyte/integrations/pytorch/#installation) --------------------------------------------------------------------------------------- pip install flyteplugins-pytorch [Configuration](https://www.union.ai/docs/v2/flyte/integrations/pytorch/#configuration) ----------------------------------------------------------------------------------------- Create an `Elastic` configuration and pass it as `plugin_config` to a `TaskEnvironment`: from flyteplugins.pytorch import Elastic torch_env = flyte.TaskEnvironment( name="torch_env", resources=flyte.Resources(cpu=(1, 2), memory=("1Gi", "2Gi")), plugin_config=Elastic( nnodes=2, nproc_per_node=1, ), image=image, ) ### [`Elastic` parameters](https://www.union.ai/docs/v2/flyte/integrations/pytorch/#elastic-parameters) | Parameter | Type | Description | | --- | --- | --- | | `nnodes` | `int` or `str` | **Required.** Number of nodes. Use an int for a fixed count or a range string (e.g., `"2:4"`) for elastic training | | `nproc_per_node` | `int` | **Required.** Number of processes (workers) per node | | `rdzv_backend` | `str` | Rendezvous backend: `"c10d"` (default), `"etcd"`, or `"etcd-v2"` | | `max_restarts` | `int` | Maximum worker group restarts (default: `3`) | | `monitor_interval` | `int` | Agent health check interval in seconds (default: `3`) | | `run_policy` | `RunPolicy` | Job run policy (cleanup, TTL, deadlines, retries) | ### [`RunPolicy` parameters](https://www.union.ai/docs/v2/flyte/integrations/pytorch/#runpolicy-parameters) | Parameter | Type | Description | | --- | --- | --- | | `clean_pod_policy` | `str` | Pod cleanup policy: `"None"`, `"all"`, or `"Running"` | | `ttl_seconds_after_finished` | `int` | Seconds to keep pods after job completion | | `active_deadline_seconds` | `int` | Maximum time the job can run (seconds) | | `backoff_limit` | `int` | Number of retries before marking the job as failed | ### [NCCL tuning parameters](https://www.union.ai/docs/v2/flyte/integrations/pytorch/#nccl-tuning-parameters) The plugin includes built-in NCCL timeout tuning to reduce failure-detection latency (PyTorch defaults to 1800 seconds): | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `nccl_heartbeat_timeout_sec` | `int` | `300` | NCCL heartbeat timeout (seconds) | | `nccl_async_error_handling` | `bool` | `False` | Enable async NCCL error handling | | `nccl_collective_timeout_sec` | `int` | `None` | Timeout for NCCL collective operations | | `nccl_enable_monitoring` | `bool` | `True` | Enable NCCL monitoring | ### [Writing a distributed training task](https://www.union.ai/docs/v2/flyte/integrations/pytorch/#writing-a-distributed-training-task) Tasks using this plugin do not need to be `async`. Initialize the process group and use `DistributedDataParallel` as you normally would with `torchrun`: import torch import torch.distributed from torch.nn.parallel import DistributedDataParallel as DDP @torch_env.task def train(epochs: int) -> float: torch.distributed.init_process_group("gloo") model = DDP(MyModel()) # ... training loop ... return final_loss When `nnodes=1`, the task runs as a regular Python task (no Kubernetes training job is created). Set `nnodes >= 2` for multi-node distributed training. [Example](https://www.union.ai/docs/v2/flyte/integrations/pytorch/#example) ----------------------------------------------------------------------------- pytorch\_example.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pytorch/pytorch_example.py "View source on GitHub") # /// script # requires-python = "==3.13" # dependencies = [\ # "flyte>=2.0.0b52",\ # "flyteplugins-pytorch",\ # "torch"\ # ] # main = "torch_distributed_train" # params = "3" # /// import typing import torch import torch.distributed import torch.nn as nn import torch.optim as optim from flyteplugins.pytorch.task import Elastic from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler, TensorDataset import flyte image = flyte.Image.from_debian_base(name="torch").with_pip_packages("flyteplugins-pytorch", pre=True) torch_env = flyte.TaskEnvironment( name="torch_env", resources=flyte.Resources(cpu=(1, 2), memory=("1Gi", "2Gi")), plugin_config=Elastic( nproc_per_node=1, # if you want to do local testing set nnodes=1 nnodes=2, ), image=image, ) class LinearRegressionModel(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(1, 1) def forward(self, x): return self.linear(x) def prepare_dataloader(rank: int, world_size: int, batch_size: int = 2) -> DataLoader: """ Prepare a DataLoader with a DistributedSampler so each rank gets a shard of the dataset. """ # Dummy dataset x_train = torch.tensor([[1.0], [2.0], [3.0], [4.0]]) y_train = torch.tensor([[3.0], [5.0], [7.0], [9.0]]) dataset = TensorDataset(x_train, y_train) # Distributed-aware sampler sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True) return DataLoader(dataset, batch_size=batch_size, sampler=sampler) def train_loop(epochs: int = 3) -> float: """ A simple training loop for linear regression. """ torch.distributed.init_process_group("gloo") model = DDP(LinearRegressionModel()) rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() dataloader = prepare_dataloader( rank=rank, world_size=world_size, batch_size=64, ) criterion = nn.MSELoss() optimizer = optim.SGD(model.parameters(), lr=0.01) final_loss = 0.0 for _ in range(epochs): for x, y in dataloader: outputs = model(x) loss = criterion(outputs, y) optimizer.zero_grad() loss.backward() optimizer.step() final_loss = loss.item() if torch.distributed.get_rank() == 0: print(f"Loss: {final_loss}") return final_loss @torch_env.task def torch_distributed_train(epochs: int) -> typing.Optional[float]: """ A nested task that sets up a simple distributed training job using PyTorch's """ print("starting launcher") loss = train_loop(epochs=epochs) print("Training complete") return loss if __name__ == "__main__": flyte.init_from_config() r = flyte.run(torch_distributed_train, epochs=3) print(r.name) print(r.url) r.wait() [API reference](https://www.union.ai/docs/v2/flyte/integrations/pytorch/#api-reference) ----------------------------------------------------------------------------------------- See the [PyTorch API reference](https://www.union.ai/docs/v2/flyte/api-reference/integrations/pytorch) for full details. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/pytorch/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # MLflow | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) MLflow ====== The MLflow plugin integrates [MLflow](https://mlflow.org/) experiment tracking with Flyte. It provides a `@mlflow_run` decorator that automatically manages MLflow runs within Flyte tasks, with support for autologging, parent-child run sharing, distributed training, and auto-generated UI links. The decorator works with both sync and async tasks. [Installation](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#installation) -------------------------------------------------------------------------------------- pip install flyteplugins-mlflow Requires `mlflow` and `flyte`. [Quick start](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#quick-start) ------------------------------------------------------------------------------------ import flyte import mlflow from flyteplugins.mlflow import mlflow_run, get_mlflow_run env = flyte.TaskEnvironment( name="mlflow-tracking", resources=flyte.Resources(cpu=1, memory="500Mi"), image=flyte.Image.from_debian_base(name="mlflow_example").with_pip_packages( "flyteplugins-mlflow" ), ) @mlflow_run( tracking_uri="http://localhost:5000", experiment_name="my-experiment", ) @env.task async def train_model(learning_rate: float) -> str: mlflow.log_param("lr", learning_rate) mlflow.log_metric("loss", 0.42) run = get_mlflow_run() return run.info.run_id ![Link](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/link.png) ![Mlflow UI](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/mlflow_dashboard.png) `@mlflow_run` must be the outermost decorator, before `@env.task`: @mlflow_run # outermost @env.task # innermost async def my_task(): ... [Autologging](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#autologging) ------------------------------------------------------------------------------------ Enable MLflow’s autologging to automatically capture parameters, metrics, and models without manual `mlflow.log_*` calls. ### [Generic autologging](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#generic-autologging) @mlflow_run(autolog=True) @env.task async def train(): from sklearn.linear_model import LogisticRegression model = LogisticRegression() model.fit(X, y) # Parameters, metrics, and model are logged automatically ### [Framework-specific autologging](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#framework-specific-autologging) Pass `framework` to use a framework-specific autolog implementation: @mlflow_run( autolog=True, framework="sklearn", log_models=True, log_datasets=False, ) @env.task async def train_sklearn(): from sklearn.ensemble import RandomForestClassifier model = RandomForestClassifier(n_estimators=100) model.fit(X_train, y_train) Supported frameworks include any framework with an `mlflow.{framework}.autolog()` function. You can find the full list of supported frameworks [here](https://mlflow.org/docs/latest/ml/tracking/autolog/#supported-libraries) . You can pass additional autolog parameters via `autolog_kwargs`: @mlflow_run( autolog=True, framework="pytorch", autolog_kwargs={"log_every_n_epoch": 5}, ) @env.task async def train_pytorch(): ... ![Autolog](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/autolog.png) [Run modes](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#run-modes) -------------------------------------------------------------------------------- The `run_mode` parameter controls how MLflow runs are created and shared across tasks: | Mode | Behavior | | --- | --- | | `"auto"` (default) | Reuse the parent’s run if one exists, otherwise create a new run | | `"new"` | Always create a new independent run | | `"nested"` | Create a new run nested under the parent via `mlflow.parentRunId` tag | ### [Sharing a run across tasks](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#sharing-a-run-across-tasks) With `run_mode="auto"` (the default), child tasks reuse the parent’s MLflow run: @mlflow_run @env.task async def parent_task(): mlflow.log_param("stage", "parent") await child_task() # Shares the same MLflow run @mlflow_run @env.task async def child_task(): mlflow.log_metric("child_metric", 1.0) # Logged to the parent's run ### [Creating independent runs](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#creating-independent-runs) Use `run_mode="new"` when a task should always create its own top-level MLflow run, completely independent of any parent: @mlflow_run(run_mode="new") @env.task async def standalone_experiment(): mlflow.log_param("experiment_type", "baseline") mlflow.log_metric("accuracy", 0.95) ### [Nested runs](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#nested-runs) Use `run_mode="nested"` to create a child run that appears under the parent in the MLflow UI. This works across processes and containers via the `mlflow.parentRunId` tag. ![Nested runs](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/mlflow_hpo.png) This is the recommended pattern for hyperparameter optimization, where each trial should be tracked as a child of the parent study run: from flyteplugins.mlflow import Mlflow @mlflow_run(run_mode="nested") @env.task(links=[Mlflow()]) async def run_trial(trial_number: int, n_estimators: int, max_depth: int) -> float: """Each trial creates a nested MLflow run under the parent.""" mlflow.log_params({"n_estimators": n_estimators, "max_depth": max_depth}) mlflow.log_param("trial_number", trial_number) model = RandomForestRegressor(n_estimators=n_estimators, max_depth=max_depth) model.fit(X_train, y_train) rmse = float(np.sqrt(mean_squared_error(y_val, model.predict(X_val)))) mlflow.log_metric("rmse", rmse) return rmse @mlflow_run @env.task async def hpo_search(n_trials: int = 30) -> str: """Parent run tracks the overall study.""" run = get_mlflow_run() mlflow.log_param("n_trials", n_trials) # Run trials in parallel — each gets a nested MLflow run rmses = await asyncio.gather( *(run_trial(trial_number=i, **params) for i, params in enumerate(trial_params)) ) mlflow.log_metric("best_rmse", min(rmses)) return run.info.run_id ![HPO](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/hpo.png) [Workflow-level configuration](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#workflow-level-configuration) ---------------------------------------------------------------------------------------------------------------------- Use `mlflow_config()` with `flyte.with_runcontext()` to set MLflow configuration for an entire workflow. All `@mlflow_run`\-decorated tasks in the workflow inherit these settings: from flyteplugins.mlflow import mlflow_config r = flyte.with_runcontext( custom_context=mlflow_config( tracking_uri="http://localhost:5000", experiment_id="846992856162999", tags={"team": "ml"}, ) ).run(train_model, learning_rate=0.001) This eliminates the need to repeat `tracking_uri` and experiment settings on every `@mlflow_run` decorator. ### [Per-task overrides](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#per-task-overrides) Use `mlflow_config()` as a context manager inside a task to override configuration for specific child tasks: @mlflow_run @env.task async def parent_task(): await shared_child() # Inherits parent config with mlflow_config(run_mode="new", tags={"role": "independent"}): await independent_child() # Gets its own run ### [Configuration priority](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#configuration-priority) Settings are resolved in priority order: 1. Explicit `@mlflow_run` decorator arguments 2. `mlflow_config()` context configuration 3. Environment variables (for `tracking_uri`) 4. MLflow defaults [Distributed training](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#distributed-training) ------------------------------------------------------------------------------------------------------ In distributed training, only rank 0 logs to MLflow by default. The plugin detects rank automatically from the `RANK` environment variable: @mlflow_run @env.task async def distributed_train(): # Only rank 0 creates an MLflow run and logs metrics. # Other ranks execute the task function directly without # creating an MLflow run or incurring any MLflow overhead. ... On non-rank-0 workers, no MLflow run is created and `get_mlflow_run()` returns `None`. The task function still executes normally — only the MLflow instrumentation is skipped. ![Distributed training](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/distributed_training.png) You can also set rank explicitly: @mlflow_run(rank=0) @env.task async def train(): ... [MLflow UI links](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#mlflow-ui-links) -------------------------------------------------------------------------------------------- The `Mlflow` link class displays links to the MLflow UI in the Flyte UI. Since the MLflow run is created inside the task at execution time, the run URL cannot be determined before the task starts. Links are only shown when a run URL is already available from context, either because a parent task created the run, or because an explicit URL is provided. The recommended pattern is for the parent task to create the MLflow run, and child tasks that inherit the run (via `run_mode="auto"`) display the link to that run. For nested runs (`run_mode="nested"`), children display a link to the parent run. ### [Setup](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#setup) Set `link_host` via `mlflow_config()` and attach `Mlflow()` links to child tasks: from flyteplugins.mlflow import Mlflow, mlflow_config @mlflow_run @env.task(links=[Mlflow()]) async def child_task(): ... # Link points to the parent's MLflow run @mlflow_run @env.task async def parent_task(): await child_task() if __name__ == "__main__": r = flyte.with_runcontext( custom_context=mlflow_config( tracking_uri="http://localhost:5000", link_host="http://localhost:5000", ) ).run(parent_task) `Mlflow()` is instantiated without a `link` argument because the URL is auto-generated at runtime. When the parent task creates an MLflow run, the plugin builds the URL from `link_host` and the run’s experiment/run IDs, then propagates it to child tasks via the Flyte context. Passing an explicit `link` would bypass this auto-generation. ### [Custom URL templates](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#custom-url-templates) The default link format is: {host}/#/experiments/{experiment_id}/runs/{run_id} For platforms like Databricks that use a different URL structure, provide a custom template: mlflow_config( link_host="https://dbc-xxx.cloud.databricks.com", link_template="{host}/ml/experiments/{experiment_id}/runs/{run_id}", ) ### [Explicit links](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#explicit-links) If you know the run URL ahead of time, you can set it directly: @env.task(links=[Mlflow(link="https://mlflow.example.com/#/experiments/1/runs/abc123")]) async def my_task(): ... ### [Link behavior by run mode](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#link-behavior-by-run-mode) | Run mode | Link behavior | | --- | --- | | `"auto"` | Parent link propagates to child tasks sharing the run | | `"new"` | Parent link is cleared; no link is shown until the task’s own run is available to its children | | `"nested"` | Parent link is kept and renamed to “MLflow (parent)” | [Automatic Flyte tags](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#automatic-flyte-tags) ------------------------------------------------------------------------------------------------------ When running inside Flyte, the plugin automatically tags MLflow runs with execution metadata: | Tag | Description | | --- | --- | | `flyte.action_name` | Task action name | | `flyte.run_name` | Flyte run name | | `flyte.project` | Flyte project | | `flyte.domain` | Flyte domain | These tags are merged with any user-provided tags. [API reference](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#api-reference) ---------------------------------------------------------------------------------------- ### [`mlflow_run` and `mlflow_config`](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#mlflow_run-and-mlflow_config) `mlflow_run` is a decorator that manages MLflow runs for Flyte tasks. `mlflow_config` creates workflow-level configuration or per-task overrides. Both accept the same core parameters: | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `run_mode` | `str` | `"auto"` | `"auto"`, `"new"`, or `"nested"` | | `tracking_uri` | `str` | `None` | MLflow tracking server URL | | `experiment_name` | `str` | `None` | MLflow experiment name (raises `ValueError` if combined with `experiment_id`) | | `experiment_id` | `str` | `None` | MLflow experiment ID (raises `ValueError` if combined with `experiment_name`) | | `run_name` | `str` | `None` | Human-readable run name (raises `ValueError` if combined with `run_id`) | | `run_id` | `str` | `None` | Explicit MLflow run ID (raises `ValueError` if combined with `run_name`) | | `tags` | `dict[str, str]` | `None` | Tags for the run | | `autolog` | `bool` | `False` | Enable MLflow autologging | | `framework` | `str` | `None` | Framework for autolog (e.g. `"sklearn"`, `"pytorch"`) | | `log_models` | `bool` | `None` | Log models automatically (requires `autolog`) | | `log_datasets` | `bool` | `None` | Log datasets automatically (requires `autolog`) | | `autolog_kwargs` | `dict` | `None` | Extra parameters for `mlflow.autolog()` | Additional keyword arguments are passed to `mlflow.start_run()`. `mlflow_run` also accepts: | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `rank` | `int` | `None` | Process rank for distributed training (only rank 0 logs) | `mlflow_config` also accepts: | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `link_host` | `str` | `None` | MLflow UI host for auto-generating links | | `link_template` | `str` | `None` | Custom URL template (placeholders: `{host}`, `{experiment_id}`, `{run_id}`) | ### [`get_mlflow_run`](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#get_mlflow_run) Returns the current `mlflow.ActiveRun` if within a `@mlflow_run`\-decorated task. Returns `None` otherwise. from flyteplugins.mlflow import get_mlflow_run run = get_mlflow_run() if run: print(run.info.run_id) ### [`get_mlflow_context`](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#get_mlflow_context) Returns the current `mlflow_config` settings from the Flyte context, or `None` if no MLflow configuration is set. Useful for inspecting the inherited configuration inside a task: from flyteplugins.mlflow import get_mlflow_context @mlflow_run @env.task async def my_task(): config = get_mlflow_context() if config: print(config.tracking_uri, config.experiment_id) ### [`Mlflow`](https://www.union.ai/docs/v2/flyte/integrations/mlflow/#mlflow-1) Link class for displaying MLflow UI links in the Flyte console. | Field | Type | Default | Description | | --- | --- | --- | --- | | `name` | `str` | `"MLflow"` | Display name for the link | | `link` | `str` | `""` | Explicit URL (bypasses auto-generation) | LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/mlflow/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Ray | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Ray === The Ray plugin lets you run [Ray](https://www.ray.io/) jobs natively on Kubernetes. Flyte provisions a transient Ray cluster for each task execution using [KubeRay](https://github.com/ray-project/kuberay) and tears it down on completion. [When to use this plugin](https://www.union.ai/docs/v2/flyte/integrations/ray/#when-to-use-this-plugin) --------------------------------------------------------------------------------------------------------- * Distributed Python workloads (parallel computation, data processing) * ML training with Ray Train or hyperparameter tuning with Ray Tune * Ray Serve inference workloads * Any workload that benefits from Ray’s actor model or task parallelism [Installation](https://www.union.ai/docs/v2/flyte/integrations/ray/#installation) ----------------------------------------------------------------------------------- pip install flyteplugins-ray Your task image must also include a compatible version of Ray: image = ( flyte.Image.from_debian_base(name="ray") .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray") ) [Configuration](https://www.union.ai/docs/v2/flyte/integrations/ray/#configuration) ------------------------------------------------------------------------------------- Create a `RayJobConfig` and pass it as `plugin_config` to a `TaskEnvironment`: from flyteplugins.ray import HeadNodeConfig, RayJobConfig, WorkerNodeConfig ray_config = RayJobConfig( head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}), worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)], runtime_env={"pip": ["numpy", "pandas"]}, enable_autoscaling=False, shutdown_after_job_finishes=True, ttl_seconds_after_finished=300, ) ray_env = flyte.TaskEnvironment( name="ray_env", plugin_config=ray_config, image=image, ) ### [`RayJobConfig` parameters](https://www.union.ai/docs/v2/flyte/integrations/ray/#rayjobconfig-parameters) | Parameter | Type | Description | | --- | --- | --- | | `worker_node_config` | `List[WorkerNodeConfig]` | **Required.** List of worker group configurations | | `head_node_config` | `HeadNodeConfig` | Head node configuration (optional) | | `enable_autoscaling` | `bool` | Enable Ray autoscaler (default: `False`) | | `runtime_env` | `dict` | Ray runtime environment (pip packages, env vars, etc.) | | `address` | `str` | Connect to an existing Ray cluster instead of provisioning one | | `shutdown_after_job_finishes` | `bool` | Shut down the cluster after the job completes (default: `False`) | | `ttl_seconds_after_finished` | `int` | Seconds to keep the cluster after completion before cleanup | ### [`WorkerNodeConfig` parameters](https://www.union.ai/docs/v2/flyte/integrations/ray/#workernodeconfig-parameters) | Parameter | Type | Description | | --- | --- | --- | | `group_name` | `str` | **Required.** Name of this worker group | | `replicas` | `int` | **Required.** Number of worker replicas | | `min_replicas` | `int` | Minimum replicas (for autoscaling) | | `max_replicas` | `int` | Maximum replicas (for autoscaling) | | `ray_start_params` | `Dict[str, str]` | Ray start parameters for workers | | `requests` | `Resources` | Resource requests per worker | | `limits` | `Resources` | Resource limits per worker | | `pod_template` | `PodTemplate` | Full pod template (mutually exclusive with `requests`/`limits`) | ### [`HeadNodeConfig` parameters](https://www.union.ai/docs/v2/flyte/integrations/ray/#headnodeconfig-parameters) | Parameter | Type | Description | | --- | --- | --- | | `ray_start_params` | `Dict[str, str]` | Ray start parameters for the head node | | `requests` | `Resources` | Resource requests for the head node | | `limits` | `Resources` | Resource limits for the head node | | `pod_template` | `PodTemplate` | Full pod template (mutually exclusive with `requests`/`limits`) | ### [Connecting to an existing cluster](https://www.union.ai/docs/v2/flyte/integrations/ray/#connecting-to-an-existing-cluster) To connect to an existing Ray cluster instead of provisioning a new one, set the `address` parameter: ray_config = RayJobConfig( worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)], address="ray://existing-cluster:10001", ) [Examples](https://www.union.ai/docs/v2/flyte/integrations/ray/#examples) --------------------------------------------------------------------------- The following example shows how to configure Ray in a `TaskEnvironment`. Flyte automatically provisions a Ray cluster for each task using this configuration: ray\_example.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/ray/ray_example.py "View source on GitHub") # /// script # requires-python = "==3.13" # dependencies = [\ # "flyte>=2.0.0b52",\ # "flyteplugins-ray",\ # "ray[default]==2.46.0"\ # ] # main = "hello_ray_nested" # params = "3" # /// import asyncio import typing import ray from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig import flyte.remote import flyte.storage @ray.remote def f(x): return x * x ray_config = RayJobConfig( head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}), worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)], runtime_env={"pip": ["numpy", "pandas"]}, enable_autoscaling=False, shutdown_after_job_finishes=True, ttl_seconds_after_finished=300, ) image = ( flyte.Image.from_debian_base(name="ray") .with_apt_packages("wget") .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray", "pip", "mypy") ) task_env = flyte.TaskEnvironment( name="hello_ray", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image ) ray_env = flyte.TaskEnvironment( name="ray_env", plugin_config=ray_config, image=image, resources=flyte.Resources(cpu=(3, 4), memory=("3000Mi", "5000Mi")), depends_on=[task_env], ) @task_env.task() async def hello_ray(): await asyncio.sleep(20) print("Hello from the Ray task!") @ray_env.task async def hello_ray_nested(n: int = 3) -> typing.List[int]: print("running ray task") t = asyncio.create_task(hello_ray()) futures = [f.remote(i) for i in range(n)] res = ray.get(futures) await t return res if __name__ == "__main__": flyte.init_from_config() r = flyte.run(hello_ray_nested) print(r.name) print(r.url) r.wait() The next example demonstrates how Flyte can create ephemeral Ray clusters and run a subtask that connects to an existing Ray cluster: ray\_existing\_example.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/ray/ray_existing_example.py "View source on GitHub") # /// script # requires-python = "==3.13" # dependencies = [\ # "flyte>=2.0.0b52",\ # "flyteplugins-ray",\ # "ray[default]==2.46.0"\ # ] # main = "create_ray_cluster" # params = "" # /// import os import typing import ray from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig import flyte.storage @ray.remote def f(x): return x * x ray_config = RayJobConfig( head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}), worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)], enable_autoscaling=False, shutdown_after_job_finishes=True, ttl_seconds_after_finished=3600, ) image = ( flyte.Image.from_debian_base(name="ray") .with_apt_packages("wget") .with_pip_packages("ray[default]==2.46.0", "flyteplugins-ray") ) task_env = flyte.TaskEnvironment( name="ray_client", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image ) ray_env = flyte.TaskEnvironment( name="ray_cluster", plugin_config=ray_config, image=image, resources=flyte.Resources(cpu=(2, 4), memory=("2000Mi", "4000Mi")), depends_on=[task_env], ) @task_env.task() async def hello_ray(cluster_ip: str) -> typing.List[int]: """ Run a simple Ray task that connects to an existing Ray cluster. """ ray.init(address=f"ray://{cluster_ip}:10001") futures = [f.remote(i) for i in range(5)] res = ray.get(futures) return res @ray_env.task async def create_ray_cluster() -> str: """ Create a Ray cluster and return the head node IP address. """ print("creating ray cluster") cluster_ip = os.getenv("MY_POD_IP") if cluster_ip is None: raise ValueError("MY_POD_IP environment variable is not set") return f"{cluster_ip}" if __name__ == "__main__": flyte.init_from_config() run = flyte.run(create_ray_cluster) run.wait() print("run url:", run.url) print("cluster created, running ray task") print("ray address:", run.outputs()[0]) run = flyte.run(hello_ray, cluster_ip=run.outputs()[0]) print("run url:", run.url) [API reference](https://www.union.ai/docs/v2/flyte/integrations/ray/#api-reference) ------------------------------------------------------------------------------------- See the [Ray API reference](https://www.union.ai/docs/v2/flyte/api-reference/integrations/ray) for full details. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/ray/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Community | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Community ========= Flyte is an open source project that is built and maintained by a community of contributors. Union AI is the primary maintainer of Flyte and developer of Union.ai, a closed source commercial product that is built on top of Flyte. Since the success of Flyte is essential to the success of Union.ai, the company is dedicated to building and expanding the Flyte open source project and community. For information on how to get involved and how to keep in touch, see [Joining the community](https://www.union.ai/docs/v2/flyte/community/joining-the-community) . [Contributing to the codebase](https://www.union.ai/docs/v2/flyte/community/#contributing-to-the-codebase) ------------------------------------------------------------------------------------------------------------ The full Flyte codebase is open source and available on GitHub. If you are interested, see [Contributing code](https://www.union.ai/docs/v2/flyte/community/contributing-code) . [Contributing to documentation](https://www.union.ai/docs/v2/flyte/community/#contributing-to-documentation) -------------------------------------------------------------------------------------------------------------- Union AI maintains and hosts both Flyte and Union documentation at [www.union.ai/docs](https://www.union.ai/docs/v2) . The two sets of documentation are deeply integrated, as the Union product is built on top of Flyte. To better maintain both sets of docs, they are hosted in the same repository (but rendered so that you can choose to view one or the other). Both the Flyte and Union documentation are open source. Flyte community members and Union customers are both welcome to contribute to the documentation. If you are interested, see [Contributing documentation and examples](https://www.union.ai/docs/v2/flyte/community/contributing-docs) . LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/community/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Flyte SDK | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) 2.4.0 Flyte SDK ========= These are the docs for Flyte SDK version 2.0 Flyte is the core Python SDK for the Union and Flyte platforms. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Integrations | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Integrations ============ API reference for Flyte integration plugins. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/api-reference/integrations/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Build an agent | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Build an agent ============== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section covers how to build, deploy, and run agentic AI applications on Flyte. Building an agent on Flyte breaks down into two **orthogonal** choices: 1. **How you build the agent loop** — plain Python, the built-in `flyte.ai.agents.Agent` harness, or a third-party framework (LangGraph, PydanticAI, OpenAI Agents SDK). 2. **How you deploy and run it** — as a task you invoke on demand, as a scheduled task driven by a `flyte.Trigger`, or as a long-running app (e.g. a webhook or chat UI). Any agent from axis (1) can be deployed via any pattern in axis (2). The two are independent, so you can start with a pure-Python loop run on demand and later move it behind a schedule or a webhook without rewriting the agent. [How Flyte maps to the agentic world](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/#how-hahahugoshortcode63s3hbhb-maps-to-the-agentic-world) ----------------------------------------------------------------------------------------------------------------------------------------------------------- * **`TaskEnvironment`**: The sandboxed execution environment for your agent steps. It configures the container image, hardware resources (CPU, GPU), and secrets (API keys). Think of it as defining “where this code runs.” * **`@env.task`**: Turns any Python function into a remotely-executed step. Each task runs in its own container with the resources you specified. This is the equivalent of a node in LangGraph or n8n. * **Tasks calling tasks**: A task can `await` other tasks, and each called task gets its own container automatically. No separate workflow decorator needed. The calling task IS your workflow, this is how you build multi-step agentic pipelines. * **`@flyte.trace`**: Marks helper functions inside a task for fine-grained observability and caching. Each traced call appears as a span in the Flyte dashboard, with its inputs and outputs captured and checkpointed. Use this on your LLM calls, tool executions, and routing decisions to get full visibility into every turn of the agent loop. See the [Flyte Quickstart](https://www.union.ai/docs/v2/flyte/user-guide/quickstart) for a hands-on walkthrough. [Ways to build an agent](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/#ways-to-build-an-agent) ------------------------------------------------------------------------------------------------------------- | Approach | When to use it | Guide | | --- | --- | --- | | **Pure Python** | You want full control over the loop and the lightest possible dependency footprint | [Build an agent with pure Python](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/python-agents) | | **The `Agent` harness** | You want a batteries-included tool-use loop with tools, MCP servers, memory, and HITL built in | [The Flyte Agent harness](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/flyte-agents) | | **Third-party frameworks** | You already have agents written with LangGraph, PydanticAI, or the OpenAI Agents SDK | [Agent framework integrations](https://www.union.ai/docs/v2/flyte/user-guide/agent-framework-integrations) | The `Agent` harness has a few dedicated guides of its own: * [**Extend the Agent class**](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/flyte-agents#extending-the-agent-class) : customize the loop by overriding `run`. * [**Agent memory**](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/agent-memory) : persist conversation history and artifacts across runs with `MemoryStore`. * [**Add a chat UI**](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/agent-chat-ui) : give any agent a hosted chat interface. [Deploying an agent](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/#deploying-an-agent) ----------------------------------------------------------------------------------------------------- Once you’ve built an agent, [**Deploy an agent as a service**](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/deploy-agent-as-service) covers running it as a task, on a schedule via `flyte.Trigger`, and behind an `AppEnvironment` webhook. [Related](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/#related) ------------------------------------------------------------------------------- * [**Sandboxing**](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing) : safely execute LLM-generated code. * [**Build an MCP server**](https://www.union.ai/docs/v2/flyte/user-guide/build-mcp) : serve Model Context Protocol servers for AI assistants to interact with Flyte. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Migration from Flyte 1 | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Migration from Flyte 1 to Flyte 2 ================================= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/api-reference/migration/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section provides a comprehensive reference for migrating Flyte 1 (flytekit) workflows to Flyte 2 (flyte SDK). For a quick-start overview of the migration process, see [Migration](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/migration) in the User Guide. [Key API changes at a glance](https://www.union.ai/docs/v2/flyte/api-reference/migration/#key-api-changes-at-a-glance) ------------------------------------------------------------------------------------------------------------------------ | 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` | [Topics](https://www.union.ai/docs/v2/flyte/api-reference/migration/#topics) ------------------------------------------------------------------------------ Philosophy and imports Key paradigm shifts and package import mapping from flytekit to flyte. Container images Migrate from ImageSpec to flyte.Image with the fluent builder API. Configuration and CLI Config file format changes and CLI command mapping. Tasks and workflows Migrate @task, @workflow, and @dynamic to TaskEnvironment and @env.task. Secrets, resources, and caching Updated patterns for secrets access, resource configuration, and caching. Parallelism and async Migrate map\_task to flyte.map and asyncio.gather patterns. Triggers and dynamic workflows Migrate LaunchPlan schedules to Triggers and @dynamic to regular tasks. Examples and common gotchas Complete migration examples and common pitfalls to avoid. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/api-reference/migration/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/api-reference/migration/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Snowflake | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Snowflake ========= The Snowflake connector lets you run SQL queries against [Snowflake](https://www.snowflake.com/) directly from Flyte tasks. Queries are submitted asynchronously and polled for completion, so they don’t block a worker while waiting for results. The connector supports: * Parameterized SQL queries with typed inputs * Key-pair and password-based authentication * Returns query results as DataFrames * Automatic links to the Snowflake query dashboard in the Flyte UI * Query cancellation on task abort [Installation](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#installation) ----------------------------------------------------------------------------------------- pip install flyteplugins-snowflake This installs the Snowflake Python connector and the `cryptography` library for key-pair authentication. [Quick start](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#quick-start) --------------------------------------------------------------------------------------- Here’s a minimal example that runs a SQL query on Snowflake: from flyte.io import DataFrame from flyteplugins.connectors.snowflake import Snowflake, SnowflakeConfig config = SnowflakeConfig( account="myorg-myaccount", user="flyte_user", database="ANALYTICS", schema="PUBLIC", warehouse="COMPUTE_WH", ) count_users = Snowflake( name="count_users", query_template="SELECT COUNT(*) FROM users", plugin_config=config, output_dataframe_type=DataFrame, ) This defines a task called `count_users` that runs `SELECT COUNT(*) FROM users` on the configured Snowflake instance. When executed, the connector: 1. Connects to Snowflake using the provided configuration 2. Submits the query asynchronously 3. Polls until the query completes or fails 4. Provides a link to the query in the Snowflake dashboard ![Snowflake Link](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/snowflake/ui.png) To run the task, create a `TaskEnvironment` from it and execute it locally or remotely: import flyte snowflake_env = flyte.TaskEnvironment.from_task("snowflake_env", count_users) if __name__ == "__main__": flyte.init_from_config() # Run locally (connector runs in-process, requires credentials and packages locally) run = flyte.with_runcontext(mode="local").run(count_users) # Run remotely (connector runs as a service in your data plane) run = flyte.with_runcontext(mode="remote").run(count_users) print(run.url) The `TaskEnvironment` created by `from_task` does not need an image or pip packages. Snowflake tasks are connector tasks, which means the query executes on the connector service, not in your task container. In `local` mode, the connector runs in-process and requires `flyteplugins-snowflake` and credentials to be available on your machine. In `remote` mode, the connector runs as a service in your data plane. [Configuration](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#configuration) ------------------------------------------------------------------------------------------- The `SnowflakeConfig` dataclass defines the connection settings for your Snowflake instance. ### [Required fields](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#required-fields) | Field | Type | Description | | --- | --- | --- | | `account` | `str` | Snowflake account identifier (e.g. `"myorg-myaccount"`) | | `database` | `str` | Target database name | | `schema` | `str` | Target schema name (e.g. `"PUBLIC"`) | | `warehouse` | `str` | Compute warehouse to use for query execution | | `user` | `str` | Snowflake username | ### [Additional connection parameters](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#additional-connection-parameters) Use `connection_kwargs` to pass any additional parameters supported by the [Snowflake Python connector](https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api) . This is a dictionary that gets forwarded directly to `snowflake.connector.connect()`. Common options include: | Parameter | Type | Description | | --- | --- | --- | | `role` | `str` | Snowflake role to use for the session | | `authenticator` | `str` | Authentication method (e.g. `"snowflake"`, `"externalbrowser"`, `"oauth"`) | | `token` | `str` | OAuth token when using `authenticator="oauth"` | | `login_timeout` | `int` | Timeout in seconds for the login request | Example with a role: config = SnowflakeConfig( account="myorg-myaccount", user="flyte_user", database="ANALYTICS", schema="PUBLIC", warehouse="COMPUTE_WH", connection_kwargs={ "role": "DATA_ANALYST", }, ) [Authentication](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#authentication) --------------------------------------------------------------------------------------------- The connector supports two authentication approaches: key-pair authentication, and password-based or other authentication methods provided through `connection_kwargs`. ### [Key-pair authentication](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#key-pair-authentication) Key-pair authentication is the recommended approach for automated workloads. Pass the names of the Flyte secrets containing the private key and optional passphrase: query = Snowflake( name="secure_query", query_template="SELECT * FROM sensitive_data", plugin_config=config, snowflake_private_key="my-snowflake-private-key", snowflake_private_key_passphrase="my-snowflake-pk-passphrase", ) The `snowflake_private_key` parameter is the name of the secret (or secret key) that contains your PEM-encoded private key. The `snowflake_private_key_passphrase` parameter is the name of the secret (or secret key) that contains the passphrase, if your key is encrypted. If your key is not encrypted, omit the passphrase parameter. The connector decodes the PEM key and converts it to DER format for Snowflake authentication. If your credentials are stored in a secret group, you can pass `secret_group` to the `Snowflake` task. The plugin expects `snowflake_private_key` and `snowflake_private_key_passphrase` to be keys within the same secret group. ### [Password authentication](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#password-authentication) Send the password via `connection_kwargs`: config = SnowflakeConfig( account="myorg-myaccount", user="flyte_user", database="ANALYTICS", schema="PUBLIC", warehouse="COMPUTE_WH", connection_kwargs={ "password": "my-password", }, ) ### [OAuth authentication](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#oauth-authentication) For OAuth-based authentication, specify the authenticator and token: config = SnowflakeConfig( account="myorg-myaccount", user="flyte_user", database="ANALYTICS", schema="PUBLIC", warehouse="COMPUTE_WH", connection_kwargs={ "authenticator": "oauth", "token": "", }, ) [Query templating](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#query-templating) ------------------------------------------------------------------------------------------------- Use the `inputs` parameter to define typed inputs for your query. Input values are bound using the `%(param)s` syntax supported by the [Snowflake Python connector](https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api) , which handles type conversion and escaping automatically. ### [Supported input types](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#supported-input-types) The `inputs` dictionary maps parameter names to Python values. Supported scalar types include `str`, `int`, `float`, and `bool`. To insert multiple rows in a single query, you can also provide lists as input values. When using list inputs, be sure to set `batch=True` on the `Snowflake` task. This enables automatic batching, where the inputs are expanded and sent as a single multi-row query instead of you having to write multiple individual statements. ### [Batched `INSERT` with list inputs](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#batched-insert-with-list-inputs) When `batch=True` is enabled, a parameterized `INSERT` query with list inputs is automatically expanded into a multi-row `VALUES` statement. Example: query = "INSERT INTO t (a, b) VALUES (%(a)s, %(b)s)" inputs = {"a": [1, 2], "b": ["x", "y"]} This is expanded into: INSERT INTO t (a, b) VALUES (%(a_0)s, %(b_0)s), (%(a_1)s, %(b_1)s) with the following flattened parameters: flat_params = { "a_0": 1, "b_0": "x", "a_1": 2, "b_1": "y", } #### [Constraints](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#constraints) * The query must contain exactly one `VALUES (...)` clause. * All list inputs must have the same non-zero length. ### [Parameterized `SELECT`](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#parameterized-select) from flyte.io import DataFrame events_by_date = Snowflake( name="events_by_date", query_template="SELECT * FROM events WHERE event_date = %(event_date)s", plugin_config=config, inputs={"event_date": str}, output_dataframe_type=DataFrame, ) You can call the task with the required inputs: @env.task async def fetch_events() -> DataFrame: return await events_by_date(event_date="2024-01-15") ### [Multiple inputs](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#multiple-inputs) You can define multiple input parameters of different types: filtered_events = Snowflake( name="filtered_events", query_template=""" SELECT * FROM events WHERE event_date >= %(start_date)s AND event_date <= %(end_date)s AND region = %(region)s AND score > %(min_score)s """, plugin_config=config, inputs={ "start_date": str, "end_date": str, "region": str, "min_score": float, }, output_dataframe_type=DataFrame, ) The query template is normalized before execution: newlines and tabs are replaced with spaces, and consecutive whitespace is collapsed. You can format your queries across multiple lines for readability without affecting execution. [Retrieving query results](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#retrieving-query-results) ----------------------------------------------------------------------------------------------------------------- If your query produces output, set `output_dataframe_type` to capture the results. `output_dataframe_type` accepts `DataFrame` from `flyte.io`. This is a meta-wrapper type that represents tabular results and can be materialized into a concrete DataFrame implementation using `open()` where you specify the target type and `all()`. from flyte.io import DataFrame top_customers = Snowflake( name="top_customers", query_template=""" SELECT customer_id, SUM(amount) AS total_spend FROM orders GROUP BY customer_id ORDER BY total_spend DESC LIMIT 100 """, plugin_config=config, output_dataframe_type=DataFrame, ) At present, only `pandas.DataFrame` is supported. The results are returned directly when you call the task: import pandas as pd @env.task async def analyze_top_customers() -> dict: df = await top_customers() pandas_df = await df.open(pd.DataFrame).all() total_spend = pandas_df["total_spend"].sum() return {"total_spend": float(total_spend)} If you specify `pandas.DataFrame` as the `output_dataframe_type`, you do not need to call `open()` and `all()` to materialize the results. import pandas as pd top_customers = Snowflake( name="top_customers", query_template=""" SELECT customer_id, SUM(amount) AS total_spend FROM orders GROUP BY customer_id ORDER BY total_spend DESC LIMIT 100 """, plugin_config=config, output_dataframe_type=pd.DataFrame, ) @env.task async def analyze_top_customers() -> dict: df = await top_customers() total_spend = df["total_spend"].sum() return {"total_spend": float(total_spend)} Be sure to inject the `SNOWFLAKE_PRIVATE_KEY` and `SNOWFLAKE_PRIVATE_KEY_PASSPHRASE` environment variables as secrets into your downstream tasks, as they must have access to Snowflake credentials in order to retrieve the DataFrame results. More on how you can create secrets [here](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets) . If you don’t need query results (for example, `DDL` statements or `INSERT` queries), omit `output_dataframe_type`. [End-to-end example](https://www.union.ai/docs/v2/flyte/integrations/snowflake/#end-to-end-example) ----------------------------------------------------------------------------------------------------- Here’s a complete workflow that uses the Snowflake connector as part of a data pipeline. The workflow creates a staging table, inserts records, queries aggregated results and processes them in a downstream task. example.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/snowflake/example.py "View source on GitHub") import flyte from flyte.io import DataFrame from flyteplugins.connectors.snowflake import Snowflake, SnowflakeConfig config = SnowflakeConfig( account="myorg-myaccount", user="flyte_user", database="ANALYTICS", schema="PUBLIC", warehouse="COMPUTE_WH", connection_kwargs={ "role": "ETL_ROLE", }, ) # Step 1: Create the staging table if it doesn't exist create_staging = Snowflake( name="create_staging", query_template=""" CREATE TABLE IF NOT EXISTS staging.daily_events ( event_id STRING, event_date DATE, user_id STRING, event_type STRING, payload VARIANT ) """, plugin_config=config, snowflake_private_key="snowflake", snowflake_private_key_passphrase="snowflake_passphrase", ) # Step 2: Insert a record into the staging table insert_events = Snowflake( name="insert_event", query_template=""" INSERT INTO staging.daily_events (event_id, event_date, user_id, event_type) VALUES (%(event_id)s, %(event_date)s, %(user_id)s, %(event_type)s) """, plugin_config=config, inputs={ "event_id": list[str], "event_date": list[str], "user_id": list[str], "event_type": list[str], }, snowflake_private_key="snowflake", snowflake_private_key_passphrase="snowflake_passphrase", batch=True, ) # Step 3: Query aggregated results for a given date daily_summary = Snowflake( name="daily_summary", query_template=""" SELECT event_type, COUNT(*) AS event_count, COUNT(DISTINCT user_id) AS unique_users FROM staging.daily_events WHERE event_date = %(report_date)s GROUP BY event_type ORDER BY event_count DESC """, plugin_config=config, inputs={"report_date": str}, output_dataframe_type=DataFrame, snowflake_private_key="snowflake", snowflake_private_key_passphrase="snowflake_passphrase", ) # Create environments for all Snowflake tasks snowflake_env = flyte.TaskEnvironment.from_task( "snowflake_env", create_staging, insert_events, daily_summary ) # Main pipeline environment that depends on the Snowflake task environments env = flyte.TaskEnvironment( name="analytics_env", resources=flyte.Resources(memory="512Mi"), image=flyte.Image.from_debian_base(name="analytics").with_pip_packages( "flyteplugins-snowflake", pre=True ), secrets=[\ flyte.Secret(key="snowflake", as_env_var="SNOWFLAKE_PRIVATE_KEY"),\ flyte.Secret(\ key="snowflake_passphrase", as_env_var="SNOWFLAKE_PRIVATE_KEY_PASSPHRASE"\ ),\ ], depends_on=[snowflake_env], ) # Step 4: Process the results in Python @env.task async def generate_report(summary: DataFrame) -> dict: import pandas as pd df = await summary.open(pd.DataFrame).all() total_events = df["event_count"].sum() top_event = df.iloc[0]["event_type"] return { "total_events": int(total_events), "top_event_type": top_event, "event_types_count": len(df), } # Compose the pipeline @env.task async def run_daily_pipeline( event_ids: list[str], event_dates: list[str], user_ids: list[str], event_types: list[str], ) -> dict: await create_staging() await insert_events( event_id=event_ids, event_date=event_dates, user_id=user_ids, event_type=event_types, ) summary = await daily_summary(report_date=event_dates[0]) return await generate_report(summary=summary) if __name__ == "__main__": flyte.init_from_config() # Run locally run = flyte.with_runcontext(mode="local").run( run_daily_pipeline, event_ids=["event-1", "event-2"], event_dates=["2023-01-01", "2023-01-02"], user_ids=["user-1", "user-2"], event_types=["click", "view"], ) # Or run remotely run = flyte.with_runcontext(mode="remote").run( run_daily_pipeline, event_ids=["event-1", "event-2"], event_dates=["2023-01-01", "2023-01-02"], user_ids=["user-1", "user-2"], event_types=["click", "view"], ) print(run.url) LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/snowflake/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Pandera | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Pandera ======= The [Pandera](https://pandera.readthedocs.io/en/latest/) plugin validates dataframes at task boundaries using [`DataFrameModel`](https://pandera.readthedocs.io/en/latest/dataframe_models.html) schemas. When a task receives or returns a pandera-typed dataframe, the plugin automatically validates the data, raises or warns on schema violations, and writes an HTML validation report to the Flyte deck. Pandera supports multiple dataframe backends. The `flyteplugins-pandera` plugin handles: | Pandera typing module | DataFrame library | Additional plugin | | --- | --- | --- | | `pandera.typing.pandas` | pandas | — | | `pandera.typing.polars` | Polars (eager and lazy) | `flyteplugins-polars` | | `pandera.typing.pyspark_sql` | PySpark SQL | `flyteplugins-spark` | [When to use this plugin](https://www.union.ai/docs/v2/flyte/integrations/pandera/#when-to-use-this-plugin) ------------------------------------------------------------------------------------------------------------- * You want compile-time-style guarantees that data flowing between tasks conforms to a declared schema * You need column-level type, constraint, and statistical checks on task inputs and outputs * You want automatic validation reports visible in the Flyte UI [Installation](https://www.union.ai/docs/v2/flyte/integrations/pandera/#installation) --------------------------------------------------------------------------------------- Install the plugin with the pandera extras for your dataframe backend: pandasPolarsPySpark SQL pip install flyteplugins-pandera 'pandera[pandas]' pip install flyteplugins-pandera flyteplugins-polars 'pandera[polars]' pip install flyteplugins-pandera flyteplugins-spark 'pandera[pyspark]' [Defining schemas](https://www.union.ai/docs/v2/flyte/integrations/pandera/#defining-schemas) ----------------------------------------------------------------------------------------------- Schemas are defined as Python classes that inherit from pandera’s `DataFrameModel`. Each field declares a column name, type, and optional constraints: import pandera.pandas as pa class EmployeeSchema(pa.DataFrameModel): employee_id: int = pa.Field(ge=0) name: str class EmployeeSchemaWithStatus(EmployeeSchema): status: str = pa.Field(isin=["active", "inactive"]) Schemas compose through inheritance: `EmployeeSchemaWithStatus` includes all columns from `EmployeeSchema` plus the `status` column. For full details on schema definition—including custom checks, regex column matching, and `Config` options—see the [pandera DataFrameModel documentation](https://pandera.readthedocs.io/en/latest/dataframe_models.html) . [Using schemas in tasks](https://www.union.ai/docs/v2/flyte/integrations/pandera/#using-schemas-in-tasks) ----------------------------------------------------------------------------------------------------------- Annotate task inputs and outputs with pandera’s generic `DataFrame` type. The plugin validates data on every encode (output) and decode (input): import pandera.typing.pandas as pt @env.task(report=True) async def build_employees() -> pt.DataFrame[EmployeeSchema]: return pd.DataFrame({ "employee_id": [1, 2, 3], "name": ["Ada", "Grace", "Barbara"], }) @env.task(report=True) async def add_status( df: pt.DataFrame[EmployeeSchema], ) -> pt.DataFrame[EmployeeSchemaWithStatus]: return df.assign(status="active") Setting `report=True` on the task makes validation reports visible as deck tabs in the Flyte UI. [Error handling with `ValidationConfig`](https://www.union.ai/docs/v2/flyte/integrations/pandera/#error-handling-with-validationconfig) ----------------------------------------------------------------------------------------------------------------------------------------- By default, a validation failure raises an exception and fails the task. To downgrade failures to warnings instead, annotate the parameter with `ValidationConfig(on_error="warn")`: from typing import Annotated from flyteplugins.pandera import ValidationConfig @env.task(report=True) async def lenient_pass_through( df: Annotated[pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")], ) -> Annotated[pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="warn")]: ... | `on_error` value | Behavior | | --- | --- | | `"raise"` (default) | Validation failure raises `pandera.errors.SchemaError` and the task fails | | `"warn"` | Validation failure logs a warning and writes the report, but the task continues | You can mix `"raise"` and `"warn"` across inputs and outputs of the same task. For example, use `"warn"` on inputs to accept best-effort data while still enforcing strict output contracts. [Image configuration](https://www.union.ai/docs/v2/flyte/integrations/pandera/#image-configuration) ----------------------------------------------------------------------------------------------------- Include the plugin in your task image. The exact setup depends on your dataframe backend: PandasPolarsPySpark SQL import flyte img = flyte.Image.from_debian_base( python_version=(3, 12), ).with_pip_packages("flyteplugins-pandera") env = flyte.TaskEnvironment( "pandera_pandas", image=img, resources=flyte.Resources(cpu="1", memory="2Gi"), ) import flyte img = ( flyte.Image.from_debian_base(python_version=(3, 12)) .with_pip_packages("flyteplugins-polars", "pandera[polars]") ) env = flyte.TaskEnvironment( "pandera_polars", image=img, resources=flyte.Resources(cpu="1", memory="2Gi"), ) import flyte from flyteplugins.spark.task import Spark image = ( flyte.Image.from_base("apache/spark-py:v3.4.0") .clone(name="pandera-pyspark-sql", python_version=(3, 10), extendable=True) .with_pip_packages("flyteplugins-spark", "pandera[pyspark]") ) spark_conf = Spark( spark_conf={ "spark.driver.memory": "1000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", }, ) env = flyte.TaskEnvironment( name="pandera_pyspark", plugin_config=spark_conf, image=image, resources=flyte.Resources(cpu="1", memory="2Gi"), ) [Polars lazy frames](https://www.union.ai/docs/v2/flyte/integrations/pandera/#polars-lazy-frames) --------------------------------------------------------------------------------------------------- The Polars backend supports both `pt.DataFrame` (eager) and `pt.LazyFrame` (lazy). With lazy frames, pandera validates the data when the frame is materialized at task I/O boundaries: import pandera.typing.polars as pt import polars as pl @env.task(report=True) async def create_lazy() -> pt.LazyFrame[MetricsSchema]: return pl.LazyFrame({"item": ["x", "y"], "value": [3.0, 4.0]}) @env.task(report=True) async def consume_lazy( lf: pt.LazyFrame[MetricsSchema], ) -> pt.DataFrame[MetricsSchema]: return lf.filter(pl.col("value") > 0.0).collect() [Examples](https://www.union.ai/docs/v2/flyte/integrations/pandera/#examples) ------------------------------------------------------------------------------- pandasPolarsPySpark SQLpandas\_schema.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pandera/pandas_schema.py "View source on GitHub") # /// script # requires-python = ">=3.12" # dependencies = [\ # "flyte",\ # "flyteplugins-pandera",\ # "pandera[pandas]",\ # ] # main = "main" # /// from __future__ import annotations from typing import Annotated import pandas as pd import pandera.pandas as pa import pandera.typing.pandas as pt from flyteplugins.pandera import ValidationConfig import flyte img = flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages( "flyteplugins-pandera", "pandera[pandas]" ) env = flyte.TaskEnvironment( "pandera_pandas_schema", image=img, resources=flyte.Resources(cpu="1", memory="2Gi"), ) class EmployeeSchema(pa.DataFrameModel): employee_id: int = pa.Field(ge=0) name: str class EmployeeSchemaWithStatus(EmployeeSchema): status: str = pa.Field(isin=["active", "inactive"]) @env.task(report=True) async def build_valid_employees() -> pt.DataFrame[EmployeeSchema]: return pd.DataFrame( { "employee_id": [1, 2, 3], "name": ["Ada", "Grace", "Barbara"], } ) @env.task(report=True) async def pass_through( df: pt.DataFrame[EmployeeSchema], ) -> pt.DataFrame[EmployeeSchemaWithStatus]: return df.assign(status="active") @env.task(report=True) async def pass_through_with_error_warn( df: Annotated[\ pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")\ ], ) -> Annotated[\ pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="warn")\ ]: del df["name"] return df @env.task(report=True) async def pass_through_with_error_raise( df: Annotated[\ pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")\ ], ) -> Annotated[\ pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="raise")\ ]: del df["name"] return df @env.task(report=True) async def main() -> pt.DataFrame[EmployeeSchemaWithStatus]: df = await build_valid_employees() df2 = await pass_through(df) await pass_through_with_error_warn(df.drop(["employee_id"], axis="columns")) await pass_through_with_error_warn(df.assign(employee_id=-1)) try: await pass_through_with_error_raise(df) except Exception as exc: print(exc) return df2 if __name__ == "__main__": flyte.init_from_config() run = flyte.run(main) print(run.url) run.wait() print("pandas pandera example OK:", run.outputs()[0]) polars\_schema.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pandera/polars_schema.py "View source on GitHub") # /// script # requires-python = ">=3.12" # dependencies = [\ # "flyte>=2.0.0b52",\ # "flyteplugins-pandera",\ # "flyteplugins-polars",\ # "pandera[polars]",\ # ] # main = "main" # /// from __future__ import annotations from typing import Annotated import pandera.polars as pa import pandera.typing.polars as pt import polars as pl from flyteplugins.pandera import ValidationConfig import flyte img = ( flyte.Image.from_debian_base(python_version=(3, 12)) .with_pip_packages("flyteplugins-pandera", "flyteplugins-polars", "pandera[polars]") ) env = flyte.TaskEnvironment( "pandera_polars_schema", image=img, resources=flyte.Resources(cpu="1", memory="2Gi"), ) class EmployeeSchema(pa.DataFrameModel): employee_id: int = pa.Field(ge=0) name: str class EmployeeSchemaWithStatus(EmployeeSchema): status: str = pa.Field(isin=["active", "inactive"]) class MetricsSchema(pa.DataFrameModel): item: str value: float @env.task(report=True) async def build_valid_employees() -> pt.DataFrame[EmployeeSchema]: return pl.DataFrame( { "employee_id": [1, 2, 3], "name": ["Ada", "Grace", "Barbara"], } ) @env.task(report=True) async def pass_through( df: pt.DataFrame[EmployeeSchema], ) -> pt.DataFrame[EmployeeSchemaWithStatus]: return df.with_columns(pl.lit("active").alias("status")) @env.task(report=True) async def pass_through_with_error_warn( df: Annotated[\ pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")\ ], ) -> Annotated[\ pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="warn")\ ]: return df.drop("name") @env.task(report=True) async def pass_through_with_error_raise( df: Annotated[\ pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")\ ], ) -> Annotated[\ pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="raise")\ ]: return df.drop("name") @env.task(report=True) async def metrics_eager() -> pt.DataFrame[MetricsSchema]: return pl.DataFrame({"item": ["a", "b"], "value": [1.0, 2.0]}) @env.task(report=True) async def metrics_lazy() -> pt.LazyFrame[MetricsSchema]: return pl.LazyFrame({"item": ["x", "y"], "value": [3.0, 4.0]}) @env.task(report=True) async def filter_metrics( lf: pt.LazyFrame[MetricsSchema], ) -> pt.DataFrame[MetricsSchema]: return lf.filter(pl.col("value") > 0.0).collect() @env.task(report=True) async def main() -> pt.DataFrame[EmployeeSchemaWithStatus]: df = await build_valid_employees() df2 = await pass_through(df) await pass_through_with_error_warn(df.drop("employee_id")) await pass_through_with_error_warn( df.with_columns(pl.lit(-1).alias("employee_id")) ) try: await pass_through_with_error_raise(df) except Exception as exc: print(exc) _ = await metrics_eager() lazy = await metrics_lazy() _ = await filter_metrics(lazy) return df2 if __name__ == "__main__": flyte.init_from_config() run = flyte.run(main) print(run.url) run.wait() print("polars pandera example OK:", run.outputs()[0]) pyspark\_sql\_schema.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pandera/pyspark_sql_schema.py "View source on GitHub") # /// script # requires-python = ">=3.10" # dependencies = [\ # "flyte>=2.0.0b52",\ # "flyteplugins-pandera",\ # "flyteplugins-spark",\ # "pandera[pyspark]",\ # ] # main = "main" # /// from __future__ import annotations from typing import Annotated, cast import pandera.typing.pyspark_sql as pt import pyspark.sql.types as T from flyteplugins.pandera import ValidationConfig from flyteplugins.spark.task import Spark from pandera.pyspark import DataFrameModel, Field from pyspark.sql import SparkSession from pyspark.sql import functions as F import flyte image = ( flyte.Image.from_base("apache/spark-py:v3.4.0") .clone(name="pandera-pyspark-sql", python_version=(3, 10), extendable=True) .with_pip_packages( "flyteplugins-pandera", "flyteplugins-spark", "pandera[pyspark]", ) ) spark_conf = Spark( spark_conf={ "spark.driver.memory": "1000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", "spark.kubernetes.file.upload.path": "/opt/spark/work-dir", "spark.jars": ( "https://storage.googleapis.com/hadoop-lib/gcs/" "gcs-connector-hadoop3-latest.jar," "https://repo1.maven.org/maven2/org/apache/hadoop/" "hadoop-aws/3.2.2/hadoop-aws-3.2.2.jar," "https://repo1.maven.org/maven2/com/amazonaws/" "aws-java-sdk-bundle/1.12.262/aws-java-sdk-bundle-1.12.262.jar" ), }, ) env = flyte.TaskEnvironment( name="pandera_pyspark_sql_schema", plugin_config=spark_conf, image=image, resources=flyte.Resources(cpu="1", memory="2Gi"), ) class EmployeeSchema(DataFrameModel): employee_id: int = Field(ge=0) name: str = Field() job_title: str = Field() class EmployeeSchemaWithStatus(EmployeeSchema): status: str = Field(isin=["active", "inactive"]) @env.task(report=True) async def build_valid_employees() -> pt.DataFrame[EmployeeSchema]: spark = cast(SparkSession, flyte.ctx().data["spark_session"]) data = [\ (1, "Ada", "Engineer"),\ (2, "Grace", "Mathematician"),\ (3, "Barbara", "Computer scientist"),\ ] schema = T.StructType( [\ T.StructField("employee_id", T.IntegerType(), False),\ T.StructField("name", T.StringType(), False),\ T.StructField("job_title", T.StringType(), False),\ ] ) return spark.createDataFrame(data, schema=schema) @env.task(report=True) async def pass_through( df: pt.DataFrame[EmployeeSchema], ) -> pt.DataFrame[EmployeeSchemaWithStatus]: return df.withColumn("status", F.lit("active")) @env.task(report=True) async def pass_through_with_error_warn( df: Annotated[\ pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")\ ], ) -> Annotated[\ pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="warn")\ ]: return df.drop("name") @env.task(report=True) async def pass_through_with_error_raise( df: Annotated[\ pt.DataFrame[EmployeeSchema], ValidationConfig(on_error="warn")\ ], ) -> Annotated[\ pt.DataFrame[EmployeeSchemaWithStatus], ValidationConfig(on_error="raise")\ ]: return df.drop("name") @env.task(report=True) async def main() -> pt.DataFrame[EmployeeSchemaWithStatus]: df = await build_valid_employees() df2 = await pass_through(df) await pass_through_with_error_warn(df.drop("employee_id")) await pass_through_with_error_warn(df.withColumn("employee_id", F.lit(-1))) try: await pass_through_with_error_raise(df) except Exception as exc: print(exc) return df2 if __name__ == "__main__": flyte.init_from_config() run = flyte.run(main) print(run.url) run.wait() print("pyspark_sql pandera example OK:", run.outputs()[0]) LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/pandera/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Papermill | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Papermill ========= The Papermill plugin lets you run Jupyter notebooks as Flyte tasks. It uses [papermill](https://papermill.readthedocs.io/) to parameterize and execute `.ipynb` files, capture their outputs as typed Flyte values, and render the executed notebook as an HTML report visible in the Flyte UI. A `NotebookTask` behaves like any other Flyte task: it has typed inputs and outputs, participates in workflows, runs remotely, integrates with the Flyte type system (including `File`, `Dir`, and `DataFrame`), and can call other Flyte tasks from within the notebook. [When to use this plugin](https://www.union.ai/docs/v2/flyte/integrations/papermill/#when-to-use-this-plugin) --------------------------------------------------------------------------------------------------------------- * Productionizing exploratory notebooks without rewriting them as Python modules * Generating cell-by-cell HTML reports as task artifacts (charts, tables, narrative analysis) * Letting data scientists iterate in notebooks while platform teams orchestrate them * Running notebooks on Spark or with GPU/CPU resources configured on the task environment [Installation](https://www.union.ai/docs/v2/flyte/integrations/papermill/#installation) ----------------------------------------------------------------------------------------- pip install flyteplugins-papermill The plugin must also be installed in the task image. For example: import flyte image = flyte.Image.from_debian_base(name="papermill-env").with_pip_packages( "flyteplugins-papermill" ) env = flyte.TaskEnvironment(name="papermill_env", image=image) [Quick start](https://www.union.ai/docs/v2/flyte/integrations/papermill/#quick-start) --------------------------------------------------------------------------------------- from flyteplugins.papermill import NotebookTask import flyte env = flyte.TaskEnvironment( name="my_env", image=flyte.Image.from_debian_base(name="my-env").with_pip_packages("flyteplugins-papermill"), ) add_numbers = NotebookTask( name="add_numbers", notebook_path="notebooks/basic_math.ipynb", task_environment=env, inputs={"x": int, "y": float}, outputs={"result": float}, ) @env.task def workflow(x: int = 5, y: float = 3.14) -> float: return add_numbers(x=x, y=y) `notebook_path` may be relative (resolved against the calling file’s directory) or absolute. [Notebook setup](https://www.union.ai/docs/v2/flyte/integrations/papermill/#notebook-setup) --------------------------------------------------------------------------------------------- Each notebook driven by a `NotebookTask` needs two specially tagged cells. ### [`parameters` cell](https://www.union.ai/docs/v2/flyte/integrations/papermill/#parameters-cell) Tag a cell with `parameters` and assign default values matching the names declared in `inputs={...}`. Papermill injects the actual values into a cell appended right after this one at execution time. # tagged: parameters x = 0 y = 0.0 ### [`outputs` cell](https://www.union.ai/docs/v2/flyte/integrations/papermill/#outputs-cell) Tag a cell with `outputs` and call `record_outputs(...)` as the last expression of the cell. The function returns a serialized representation of the values, which Jupyter captures as the cell’s displayed output. `NotebookTask` then reads that captured output from the executed notebook to recover the typed values. # tagged: outputs from flyteplugins.papermill import record_outputs record_outputs(result=x + y) `record_outputs` accepts any value that the Flyte type system supports such as primitives, `File`, `Dir`, `DataFrame`, dataclasses, etc. The output names and types must match the `outputs={...}` declaration on the `NotebookTask`. Inputs and outputs have different type rules. Inputs are restricted to JSON-serializable primitives plus `File`/`Dir`/`DataFrame` because papermill’s parameter mechanism is JSON-only. Outputs go through the full Flyte type engine inside the notebook via `record_outputs`, so dataclasses and any other Flyte-supported type work there. If a notebook has no outputs, omit the `outputs` cell and don’t pass `outputs` to `NotebookTask`. The notebook still runs and its HTML report is rendered, but no values are returned. [Inputs and outputs](https://www.union.ai/docs/v2/flyte/integrations/papermill/#inputs-and-outputs) ----------------------------------------------------------------------------------------------------- ### [Supported input types](https://www.union.ai/docs/v2/flyte/integrations/papermill/#supported-input-types) Notebook parameters are passed through papermill, which only accepts JSON-serializable values. The plugin allows: * Primitives: `int`, `float`, `str`, `bool`, `list`, `dict`, `None` * Flyte I/O types: `flyte.io.File`, `flyte.io.Dir`, `flyte.io.DataFrame` (serialized to their path/URI strings) Passing any other type raises `TypeError` at call time. Wrap unsupported values in a dataclass and serialize them to a primitive container, or write them to a `File`/`Dir` first. ### [Complex types: File, Dir, DataFrame](https://www.union.ai/docs/v2/flyte/integrations/papermill/#complex-types-file-dir-dataframe) `File`, `Dir` and `DataFrame` are passed to the notebook as plain path/URI strings. Reconstruct them inside the notebook with the provided helpers: from flyteplugins.papermill import load_file, load_dir, load_dataframe # input_file, input_dir, input_df were injected as strings by papermill f = load_file(input_file) # -> flyte.io.File d = load_dir(input_dir) # -> flyte.io.Dir df = load_dataframe(input_df) # -> flyte.io.DataFrame (parquet by default) `load_dataframe` accepts a `fmt` argument (default `"parquet"`) for non-parquet storage formats. Jupyter supports top-level `await`, so use it directly for async I/O: import pandas as pd from flyte.io import DataFrame pdf = await df.open(pd.DataFrame).all() output_df = await DataFrame.from_local(pdf) To return a `DataFrame` from a notebook, materialize it as a `flyte.io.DataFrame` and pass it to `record_outputs`: # tagged: outputs import pandas as pd from flyte.io import DataFrame from flyteplugins.papermill import record_outputs result_df = pd.DataFrame({"name": ["alice", "bob"], "score": [90, 75]}) output = await DataFrame.from_local(result_df) record_outputs(filtered_df=output, row_count=len(result_df)) The same pattern applies to `File` (`await File.from_local(...)`) and `Dir` (`await Dir.from_local(...)`). ### [Outputs: single, multiple, none](https://www.union.ai/docs/v2/flyte/integrations/papermill/#outputs-single-multiple-none) A `NotebookTask` returns: * A single value when `outputs` has one entry * A tuple in the order declared in `outputs` when there are multiple entries * `None` when `outputs` is omitted # Multiple outputs text_analysis = NotebookTask( name="text_analysis", notebook_path="notebooks/text.ipynb", task_environment=env, inputs={"text": str, "n": int}, outputs={"repeated": str, "word_count": int, "char_count": int}, ) @env.task def workflow(text: str, n: int) -> tuple[str, int, int]: repeated, word_count, char_count = text_analysis(text=text, n=n) return repeated, word_count, char_count # No outputs — useful for side-effect-only notebooks (reports, exports) printer = NotebookTask( name="printer", notebook_path="notebooks/print_report.ipynb", task_environment=env, inputs={"message": str}, ) @env.task def report_workflow(message: str = "hello"): printer(message=message) If a declared output is missing from `record_outputs(...)`, `NotebookTask` raises `TypeError` listing the missing names. [Calling Flyte tasks from notebooks](https://www.union.ai/docs/v2/flyte/integrations/papermill/#calling-flyte-tasks-from-notebooks) ------------------------------------------------------------------------------------------------------------------------------------- You can call other Flyte tasks directly from inside a notebook. The plugin injects the parent task’s runtime context into the notebook kernel at the start of execution, so task calls are routed through the Flyte controller automatically, so no manual setup required. When running remotely, each task call is submitted to Flyte and appears as a separate node in the run graph. When running locally, the calls execute in-process as regular Python functions. # Inside a notebook cell from my_tasks import expensive_task result = await expensive_task(data=42) Sync tasks can be called the same way: from my_tasks import compute_total total = compute_total(values=[1, 2, 3]) The setup cell that initializes the runtime context is injected automatically and stripped from the rendered HTML report and the uploaded `.ipynb` files, so it never shows up to users. [Workflow patterns](https://www.union.ai/docs/v2/flyte/integrations/papermill/#workflow-patterns) --------------------------------------------------------------------------------------------------- ### [Chaining notebooks](https://www.union.ai/docs/v2/flyte/integrations/papermill/#chaining-notebooks) Outputs from one `NotebookTask` can feed directly into another: @env.task def chained_workflow(a: int, b: float, c: float) -> float: intermediate = step1_add(x=a, y=b) final = step2_add(x=int(intermediate), y=c) return final ### [Mixing notebooks with regular tasks](https://www.union.ai/docs/v2/flyte/integrations/papermill/#mixing-notebooks-with-regular-tasks) `NotebookTask` composes with `@env.task` functions in either direction: @env.task def mixed_workflow(n: int) -> float: doubled = double(n=n) # regular task nb_result = notebook_add(x=doubled, y=100.0) # notebook task return add(a=nb_result, b=0.5) # regular task ### [Inline definition](https://www.union.ai/docs/v2/flyte/integrations/papermill/#inline-definition) `NotebookTask` can be created inside a task function rather than at module scope. The resolver bakes the notebook path and type schemas into the task spec at registration time, so no module-level reference is required at execution. @env.task def workflow(x: int = 3, y: float = 1.5) -> int: from flyteplugins.papermill import NotebookTask nb = NotebookTask( name="add_numbers", notebook_path="notebooks/basic_math.ipynb", task_environment=env, inputs={"x": int, "y": float}, outputs={"result": float}, ) return nb(x=x, y=y) ### [Calling from sync vs. async tasks](https://www.union.ai/docs/v2/flyte/integrations/papermill/#calling-from-sync-vs-async-tasks) `NotebookTask` is internally synchronous. Papermill blocks while the notebook runs. Call it directly from a sync task or use `.aio()` from an async task: @env.task def sync_parent(x: int) -> float: return notebook(x=x) @env.task async def async_parent(x: int) -> float: return await notebook.aio(x=x) ### [Running a NotebookTask directly as the entrypoint](https://www.union.ai/docs/v2/flyte/integrations/papermill/#running-a-notebooktask-directly-as-the-entrypoint) A `NotebookTask` can be the workflow entrypoint without wrapping it in another task: nb = NotebookTask( name="add_numbers", notebook_path="notebooks/basic_math.ipynb", task_environment=env, inputs={"x": int, "y": float}, outputs={"result": float}, ) if __name__ == "__main__": flyte.init_from_config() run = flyte.with_runcontext(mode="remote", copy_style="all").run(nb, x=3, y=1.5) print(run.url) [Reports and notebook artifacts](https://www.union.ai/docs/v2/flyte/integrations/papermill/#reports-and-notebook-artifacts) ----------------------------------------------------------------------------------------------------------------------------- ### [HTML report (default)](https://www.union.ai/docs/v2/flyte/integrations/papermill/#html-report-default) Every `NotebookTask` execution renders the executed notebook to HTML and logs it to the Flyte Report tab for that task. This happens whether the notebook succeeds or fails — see [Failure reports](https://www.union.ai/docs/v2/flyte/integrations/papermill/#failure-reports) below. The report is on by default and requires no configuration. ![HTML Report](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/papermill/default_report.png) ### [Notebook artifacts](https://www.union.ai/docs/v2/flyte/integrations/papermill/#notebook-artifacts) By default the executed notebook lives only inside the rendered HTML report. To get the source and executed `.ipynb` files as typed Flyte outputs (so downstream tasks can read them or so they show up as artifacts in the run UI), set `output_notebooks=True`: notebook = NotebookTask( name="analysis", notebook_path="notebooks/analysis.ipynb", task_environment=env, inputs={"x": int}, outputs={"result": float}, output_notebooks=True, ) @env.task def workflow(x: int = 5) -> tuple[float, File, File]: result, source_nb, executed_nb = notebook(x=x) return result, source_nb, executed_nb When enabled, two outputs are appended to the task’s interface automatically: * `output_notebook`: The source `.ipynb` (no executed cell outputs) * `output_notebook_executed`: The executed `.ipynb` (with cell outputs) The names `output_notebook` and `output_notebook_executed` are reserved when `output_notebooks=True`. Don’t use them as your own user output names. ### [Clean reports](https://www.union.ai/docs/v2/flyte/integrations/papermill/#clean-reports) `report_mode=True` tells papermill to mark input cells with a `source_hidden` flag during execution. The plugin then strips those input cells from both the rendered HTML report and the uploaded `.ipynb` files, so only cell outputs (charts, tables, text) remain. This produces a clean stakeholder-facing report without exposing the underlying code. notebook = NotebookTask( ... report_mode=True, output_notebooks=True, ) ![Clean Report](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/papermill/clean_report.png) ### [Failure reports](https://www.union.ai/docs/v2/flyte/integrations/papermill/#failure-reports) The HTML report is rendered even when the notebook fails. Papermill writes the output notebook cell-by-cell as it executes, so the partial notebook is on disk when an exception propagates out. The plugin renders this partial notebook to HTML and flushes it to the Flyte Report before re-raising the error, giving full visibility into which cell failed and what output the earlier cells produced. This is especially useful for long-running notebooks: you can inspect partial results without re-running the whole pipeline. ![Failed Report](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/papermill/failed_report.png) [Spark notebooks](https://www.union.ai/docs/v2/flyte/integrations/papermill/#spark-notebooks) ----------------------------------------------------------------------------------------------- Pass `plugin_config=Spark(...)` to run a notebook inside a Spark driver pod managed by the Spark on Kubernetes Operator: from flyteplugins.papermill import NotebookTask from flyteplugins.spark import Spark spark_nb = NotebookTask( name="spark_analysis", notebook_path="notebooks/spark_analysis.ipynb", task_environment=env, plugin_config=Spark( spark_conf={ "spark.executor.instances": "2", "spark.executor.memory": "2g", "spark.executor.cores": "1", "spark.driver.memory": "1g", "spark.driver.cores": "1", }, ), inputs={"data": list}, outputs={"total": int, "count": int}, ) Inside the notebook, build the `SparkSession` directly: from pyspark.sql import SparkSession spark = SparkSession.builder.appName("FlyteSpark").getOrCreate() `SparkContext.addPyFile()` is not called for notebook tasks. The notebook kernel runs in a subprocess that cannot share state with the parent task process, so dynamic code distribution via `addPyFile` is not supported. Executor pods use the same Docker image as the driver, so any package needed in UDFs must be installed in the image. See the [Spark plugin](https://www.union.ai/docs/v2/flyte/integrations/spark) page for the full `Spark` configuration reference. [Local testing](https://www.union.ai/docs/v2/flyte/integrations/papermill/#local-testing) ------------------------------------------------------------------------------------------- Calling a `NotebookTask` as a regular Python function outside any Flyte runner executes the notebook synchronously through papermill and returns Python values: result = add_numbers(x=1, y=2.5) In this mode: * The notebook runs in-process (no remote submission) * No HTML report is rendered (no task context) * `File` and `Dir` outputs created inside the notebook resolve to local paths * No plugin lifecycle hooks fire (so no Spark cluster is provisioned, etc.) This makes iteration on notebook logic fast. You can run the task from a script, REPL or test without going through Flyte at all. [Execution options](https://www.union.ai/docs/v2/flyte/integrations/papermill/#execution-options) --------------------------------------------------------------------------------------------------- `NotebookTask` exposes the full set of papermill execution knobs. The snippet below shows example values. See [the reference table](https://www.union.ai/docs/v2/flyte/integrations/papermill/#notebooktask-reference) for defaults. NotebookTask( name="all_options", notebook_path="notebooks/basic_math.ipynb", task_environment=env, inputs={"x": int, "y": float}, outputs={"result": float}, kernel_name="python3", # default None - use kernel from notebook metadata language=None, # rarely needed; overrides notebook language execution_timeout=300, # default None - no per-cell timeout start_timeout=120, # default 60 seconds to wait for kernel startup log_output=True, # default False; stream cell output to task log progress_bar=True, # default True; tqdm-style progress in logs report_mode=False, # default False; True hides input cells in report request_save_on_cell_execute=True, # default True; save after every cell (nbclient) engine_name=None, # default None - nbclient engine_kwargs={"autosave_cell_every": 30}, # extra kwargs forwarded to engine ) `request_save_on_cell_execute` is largely redundant in remote execution: the plugin always renders and uploads the partial notebook on failure, so crash diagnostics don’t depend on it. Leave it on its default unless using a custom engine that requires it. [`NotebookTask` reference](https://www.union.ai/docs/v2/flyte/integrations/papermill/#notebooktask-reference) --------------------------------------------------------------------------------------------------------------- | Parameter | Default | Description | | --- | --- | --- | | `name` | — | Task name | | `notebook_path` | — | Path to the `.ipynb`, relative to the calling file or absolute | | `task_environment` | — | `TaskEnvironment` for registration and remote execution | | `inputs` | `None` | `{name: type}` dict of notebook inputs | | `outputs` | `None` | `{name: type}` dict of notebook outputs | | `plugin_config` | `None` | Plugin config — currently only `Spark(...)` is supported. Sets the task type accordingly. | | `kernel_name` | `None` | Jupyter kernel name; `None` uses the kernel from notebook metadata | | `engine_name` | `None` | Papermill engine; `None` uses the default `nbclient` engine | | `log_output` | `False` | Stream cell output to the task log | | `start_timeout` | `60` | Seconds to wait for kernel startup | | `execution_timeout` | `None` | Per-cell timeout in seconds; `None` means no timeout | | `report_mode` | `False` | Strip input cells from the report and uploaded `.ipynb` | | `request_save_on_cell_execute` | `True` | Save notebook after every cell (nbclient engine only) | | `progress_bar` | `True` | Show a tqdm-style progress bar during execution | | `language` | `None` | Override notebook language (rarely needed) | | `engine_kwargs` | `{}` | Extra kwargs forwarded to the papermill engine | | `output_notebooks` | `False` | Upload source and executed `.ipynb` as `File` task outputs | [Helper functions](https://www.union.ai/docs/v2/flyte/integrations/papermill/#helper-functions) ------------------------------------------------------------------------------------------------- These are imported from `flyteplugins.papermill` and called from inside the notebook. | Function | Purpose | | --- | --- | | `record_outputs(**kwargs)` | Records outputs from the `outputs`\-tagged cell. Must be the cell’s last expression. Accepts any Flyte-typed values. | | `load_file(path)` | Reconstructs a `flyte.io.File` from the path string injected by papermill. | | `load_dir(path)` | Reconstructs a `flyte.io.Dir` from the path string injected by papermill. | | `load_dataframe(uri, fmt="parquet")` | Reconstructs a `flyte.io.DataFrame` from the URI string injected by papermill. | LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/papermill/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Hydra | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Hydra ===== [Hydra](https://hydra.cc/) is a framework for composing and overriding configuration trees from YAML files, dataclasses and the command line. The `flyteplugins-hydra` plugin makes Hydra a first-class submission layer for Flyte, so you can compose a config exactly as you would in any other Hydra app and have each composed run executed as a Flyte task, locally or as a remote execution on a Flyte cluster. The plugin offers three complementary entry points that share a single launcher implementation: | Entry point | Use it when | | --- | --- | | `hydra/launcher=flyte` (Hydra Launcher plugin) | You already have a `@hydra.main` script and want standard Hydra CLI ergonomics, including `--multirun` and custom sweepers. | | `flyte hydra run` (Flyte CLI extension) | You want a Flyte-style CLI that imports a task from a Python file and composes a Hydra config without requiring a `@hydra.main` wrapper. | | `hydra_run` / `hydra_sweep` (Python SDK) | You want to submit runs directly from Python – notebooks, tests, examples or another orchestration script. | All three paths converge on the same `FlyteLauncher`. [Installation](https://www.union.ai/docs/v2/flyte/integrations/hydra/#installation) ------------------------------------------------------------------------------------- pip install flyteplugins-hydra The plugin depends on `flyteplugins-omegaconf`, which is installed automatically and provides the `DictConfig`/`ListConfig` type transformers that allow Hydra-composed configs to flow into Flyte tasks. Both packages must be available in the same environment as `flyte`. If you call `apply_task_env` for child tasks (see [Task environment overrides](https://www.union.ai/docs/v2/flyte/integrations/hydra/#task-environment-overrides) ), include `flyteplugins-hydra` in the task image as well. [Requirements on tasks](https://www.union.ai/docs/v2/flyte/integrations/hydra/#requirements-on-tasks) ------------------------------------------------------------------------------------------------------- Every task launched through this plugin must accept an OmegaConf `DictConfig` input. Any other parameters are passed through as ordinary task arguments. from omegaconf import DictConfig @env.task async def pipeline(cfg: DictConfig, dataset: str) -> float: ... The plugin auto-detects the `DictConfig` parameter name. If your parameter is `cfg`, app-level overrides are passed through `--cfg` on the CLI; if it is `config`, they are passed through `--config`; and so on. [A walkthrough config](https://www.union.ai/docs/v2/flyte/integrations/hydra/#a-walkthrough-config) ----------------------------------------------------------------------------------------------------- The examples in this page assume a small project layout: project/ ├── train.py └── conf/ ├── training.yaml ├── model/ │ ├── resnet.yaml │ └── vit.yaml ├── optimizer/ │ ├── adam.yaml │ └── sgd.yaml └── task_env/ ├── a100.yaml └── prebuilt_image.yaml `conf/training.yaml`: defaults: - optimizer: adam - model: resnet - _self_ data: path: s3://my-bucket/imagenet dataset: imagenet training: epochs: 30 batch_size: 64 `train.py` (abbreviated): import flyte from omegaconf import DictConfig from flyteplugins.hydra import apply_task_env env = flyte.TaskEnvironment(name="training", image=...) @env.task async def preprocess(cfg: DictConfig) -> flyte.io.Dir: ... @env.task async def train_model(cfg: DictConfig, data: flyte.io.Dir) -> tuple[flyte.io.Dir, float]: ... @env.task async def pipeline(cfg: DictConfig, dataset: str) -> float: data = await preprocess(cfg) train_task = apply_task_env(train_model, cfg) _, val_loss = await train_task(cfg, data) return val_loss The same `pipeline` task is the target of every example below. `config_path` is resolved relative to the current working directory. If you submit runs from a directory other than `project/`, pass an absolute path (or an absolute path on the CLI via `--config-path /abs/path/to/conf`). For structured-config-only setups (no YAML files), omit `config_path` / `--config-path` entirely. [Execution mode](https://www.union.ai/docs/v2/flyte/integrations/hydra/#execution-mode) ----------------------------------------------------------------------------------------- Remote execution is the default. Every entry point exposes an explicit knob: | Surface | Local | Remote | | --- | --- | --- | | `@hydra.main` launcher | `hydra.launcher.mode=local` | `hydra.launcher.mode=remote` (default) | | `flyte hydra run` | `--local` | `--mode remote` (default) | | Python SDK | `mode="local"` | `mode="remote"` (default) | For the `@hydra.main` launcher, the default applies as soon as `hydra/launcher=flyte` is selected. Remote runs print the Flyte run URL immediately after submission, before any waiting. By default the plugin then waits for every submitted run to reach a terminal phase, capped at 32 worker threads. To tune or disable waiting: | Surface | Tune wait threads | Fire and forget | | --- | --- | --- | | `@hydra.main` launcher | `hydra.launcher.wait_max_workers=64` | `hydra.launcher.wait=false` | | `flyte hydra run` | `--wait-max-workers 64` | `--no-wait` | | Python SDK | `wait_max_workers=64` | `wait=False` | For a sweep, every job is submitted first, and then the plugin waits for all runs concurrently. Submission is not blocked by earlier runs reaching a terminal phase. [Hydra launcher (`@hydra.main` scripts)](https://www.union.ai/docs/v2/flyte/integrations/hydra/#hydra-launcher-hydramain-scripts) ----------------------------------------------------------------------------------------------------------------------------------- Use this path when your script already has a `@hydra.main` entry point. Selecting `hydra/launcher=flyte` swaps Hydra’s built-in `BasicLauncher` for `FlyteLauncher`. Single remote run: python train.py hydra/launcher=flyte hydra.launcher.mode=remote Single local run: python train.py hydra/launcher=flyte hydra.launcher.mode=local Remote grid sweep submission: Each comma-separated value expands into a separate Flyte execution; six executions in this example: python train.py --multirun \ hydra/launcher=flyte hydra.launcher.mode=remote \ hydra.launcher.wait_max_workers=64 \ optimizer.lr=0.001,0.01,0.1 training.epochs=10,20 Fire-and-forget sweep submission: python train.py --multirun \ hydra/launcher=flyte hydra.launcher.wait=false \ optimizer.lr=0.001,0.01,0.1 Custom sweepers (Optuna) work exactly as they do with the BasicLauncher. Selecting `hydra/sweeper=...` activates the sweeper and `FlyteLauncher` runs each trial as a Flyte execution: python train.py --multirun \ hydra/launcher=flyte hydra.launcher.mode=remote \ hydra/sweeper=optuna hydra.sweeper.n_trials=20 \ hydra.sweeper.n_jobs=4 \ "optimizer.lr=interval(1e-4,1e-1)" Inside `@hydra.main`, the standard pattern is: import flyte import hydra from omegaconf import DictConfig from flyteplugins.hydra import apply_task_env @hydra.main(version_base=None, config_path="conf", config_name="training") def main(cfg: DictConfig): flyte.init_from_config() entry_task = apply_task_env(pipeline, cfg) return flyte.run(entry_task, cfg=cfg, dataset=cfg.data.dataset) if __name__ == "__main__": main() [Python SDK](https://www.union.ai/docs/v2/flyte/integrations/hydra/#python-sdk) --------------------------------------------------------------------------------- `hydra_run` composes one config and runs the task once. `hydra_sweep` expands sweep overrides and runs the task once per combination. ### [Single run](https://www.union.ai/docs/v2/flyte/integrations/hydra/#single-run) from flyteplugins.hydra import hydra_run run = hydra_run( pipeline, config_path="conf", config_name="training", overrides=["optimizer.lr=0.01"], dataset="s3://my-bucket/imagenet", mode="remote", wait=True, wait_max_workers=64, ) For a remote run with `wait=True`, the return value is a wrapper exposing both `run.url` and `run.value` (the resolved task output). The wrapper is `float()`\-castable so Hydra sweepers such as Optuna can consume scalar objectives directly. With `wait=False`, the return value is the underlying `flyte.remote.Run`. ### [Grid sweep](https://www.union.ai/docs/v2/flyte/integrations/hydra/#grid-sweep) from flyteplugins.hydra import hydra_sweep runs = hydra_sweep( pipeline, config_path="conf", config_name="training", overrides=["optimizer.lr=0.001,0.01,0.1", "training.epochs=10,20"], dataset="s3://my-bucket/imagenet", mode="remote", ) Six executions are submitted (3 × 2). `runs` is a list aligned with the Cartesian-product order Hydra’s `BasicSweeper` produces. ### [Custom sweepers](https://www.union.ai/docs/v2/flyte/integrations/hydra/#custom-sweepers) Custom sweeper plugins are activated by passing their selection in `overrides`: runs = hydra_sweep( pipeline, config_path="conf", config_name="training", overrides=[\ "hydra/sweeper=optuna",\ "hydra.sweeper.n_trials=20",\ "hydra.sweeper.n_jobs=4",\ "optimizer.lr=interval(1e-4,1e-1)",\ ], dataset="s3://my-bucket/imagenet", mode="remote", ) Whenever an override starts with `hydra/`, the plugin invokes the full Hydra runtime so plugin discovery (sweepers, launchers, callbacks) can run. Pure value overrides on the `hydra.*` namespace (for example `hydra.run.dir=...`) do not need the full runtime and are applied per-job by the launcher directly. ### [Forwarding `flyte.with_runcontext` options](https://www.union.ai/docs/v2/flyte/integrations/hydra/#forwarding-flytewith_runcontext-options) Use `run_options` to pass Flyte runtime options through to every job: runs = hydra_sweep( pipeline, config_path="conf", config_name="training", overrides=["optimizer.lr=0.001,0.01,0.1"], dataset="s3://my-bucket/imagenet", mode="remote", run_options={ "name": "my-training-sweep", "service_account": "default", "copy_style": "all", "raw_data_path": "s3://my-bucket/raw-data", "debug": True, }, ) [Flyte CLI (`flyte hydra run`)](https://www.union.ai/docs/v2/flyte/integrations/hydra/#flyte-cli-flyte-hydra-run) ------------------------------------------------------------------------------------------------------------------- `flyte hydra run` is registered through the `flyte.plugins.cli.commands` entry point. It loads a task from a Python file, composes a Hydra config, and runs the task without requiring the script to have its own `@hydra.main` function. It also inherits the relevant flags from `flyte run` (`--project`, `--domain`, `--image`, `--name`, `--service-account`, `--raw-data-path`, `--copy-style`, `--debug`, `--local`, `--follow`). ### [Single run](https://www.union.ai/docs/v2/flyte/integrations/hydra/#single-run-1) Remote (default): flyte hydra run --config-path conf --config-name training \ train.py pipeline --dataset s3://my-bucket/imagenet Forced local: flyte hydra run --local --config-path conf --config-name training \ train.py pipeline --dataset s3://my-bucket/imagenet ### [Grid sweep](https://www.union.ai/docs/v2/flyte/integrations/hydra/#grid-sweep-1) flyte hydra run --multirun --config-path conf --config-name training \ --wait-max-workers 64 \ train.py pipeline --dataset s3://my-bucket/imagenet \ --cfg "optimizer.lr=0.001,0.01,0.1" --cfg "training.epochs=10,20" ### [App-level vs Hydra-namespace overrides](https://www.union.ai/docs/v2/flyte/integrations/hydra/#app-level-vs-hydra-namespace-overrides) The CLI keeps app-level overrides separate from Hydra runtime overrides so they do not collide with ordinary Flyte task arguments. App-level overrides target the composed config and are passed through the **task’s `DictConfig` parameter name**. For `pipeline(cfg: DictConfig, ...)`, use `--cfg`. For `pipeline_with_config(config: DictConfig, ...)`, use `--config`: flyte hydra run --config-path conf --config-name training \ train.py pipeline \ --cfg optimizer.lr=0.01 \ --cfg training.epochs=20 flyte hydra run --config-path conf --config-name training \ train.py pipeline_with_config \ --config optimizer.lr=0.01 Hydra runtime overrides: Anything in the `hydra.*` or `hydra/*` namespace go through `--hydra-override`: flyte hydra run --config-path conf --config-name training \ train.py pipeline \ --hydra-override hydra.run.dir=./outputs/exp1 \ --hydra-override hydra/launcher=flyte Custom sweepers combine the two: flyte hydra run --multirun --config-path conf --config-name training \ train.py pipeline --dataset s3://my-bucket/imagenet \ --hydra-override hydra/sweeper=optuna \ --hydra-override hydra.sweeper.n_trials=20 \ --hydra-override hydra.sweeper.n_jobs=4 \ --cfg "optimizer.lr=interval(1e-4,1e-1)" \ --cfg "training.epochs=choice(10,20,50)" ### [`--follow` and `--no-wait`](https://www.union.ai/docs/v2/flyte/integrations/hydra/#--follow-and---no-wait) `--follow` streams logs from the launched run after submission; it implies waiting and cannot be combined with `--no-wait`. `--no-wait` returns immediately after submission and skips log streaming. ### [Shell completion](https://www.union.ai/docs/v2/flyte/integrations/hydra/#shell-completion) Install Click’s completion hook for the `flyte` executable. For zsh: eval "$(_FLYTE_COMPLETE=zsh_source flyte)" For bash: eval "$(_FLYTE_COMPLETE=bash_source flyte)" Once installed, `flyte hydra run` adds Hydra-aware completion after `SCRIPT TASK_NAME`. The command imports the script, inspects the task signature, and suggests: * The app override flag matching the task’s `DictConfig` parameter (`--cfg`, `--config`, …). * Override values for that flag and `--hydra-override` via Hydra’s own completion engine, including config keys, config-group selections and sweep functions. flyte hydra run --config-path conf --config-name training \ train.py pipeline --cfg optimizer. # suggests optimizer.lr=, optimizer.weight_decay=, ... flyte hydra run --config-path conf --config-name training \ train.py pipeline --hydra-override hydra/launcher= # suggests hydra launcher choices Because completion has to import the target script, keep task definitions and `ConfigStore` registration import-safe, and avoid expensive top-level work in scripts you reach via `flyte hydra run`. ![Auto Completion](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/hydra/auto_complete.gif) [Override grammar](https://www.union.ai/docs/v2/flyte/integrations/hydra/#override-grammar) --------------------------------------------------------------------------------------------- The override grammar is identical to standard Hydra; what differs is only how you pass the strings (positional in `python train.py ...`, list entries in `overrides=[...]`, repeated `--cfg`/`--hydra-override` on the Flyte CLI). | Form | Meaning | | --- | --- | | `optimizer.lr=0.01` | Set an existing key. | | `optimizer=sgd` | Select a config group (replaces the `optimizer` subtree with `conf/optimizer/sgd.yaml`). | | `+task_env=a100` | Append a config group whose key is not currently in the config. | | `+training.grad_clip=1.0` | Append a key that does not exist. | | `++optimizer.lr=0.05` | Force-set a key, creating it if missing and overriding strict-schema errors. | | `~training.warmup_steps` | Delete a key from the composed config. | | `optimizer.lr=0.001,0.01,0.1` | Sweep value (with `--multirun`); expanded into one job per element. | | `optimizer.lr=interval(1e-4,1e-1)` | Continuous sweep range; consumed by samplers like Optuna. | | `optimizer=choice(adam,sgd)` | Categorical sweep; consumed by samplers. | | `hydra.run.dir=./outputs/exp1` | Hydra-namespace value override (single run output dir). | | `hydra.sweep.dir=./outputs/sweep1` | Hydra-namespace sweep output dir. | | `hydra/sweeper=optuna` | Hydra-namespace config group selection (activates the Optuna sweeper plugin). | [Sweeps](https://www.union.ai/docs/v2/flyte/integrations/hydra/#sweeps) ------------------------------------------------------------------------- ### [Grid sweeps (BasicSweeper)](https://www.union.ai/docs/v2/flyte/integrations/hydra/#grid-sweeps-basicsweeper) Comma-separated overrides expand into a Cartesian product. The plugin uses Hydra’s `BasicSweeper` to expand them, then submits one Flyte execution per combination. from flyteplugins.hydra import hydra_sweep runs = hydra_sweep( pipeline, config_path="conf", config_name="training", overrides=["model=resnet,vit", "optimizer.lr=0.001,0.01,0.1"], dataset="s3://my-bucket/imagenet", mode="remote", ) # 6 executions flyte hydra run --multirun --config-path conf --config-name training \ train.py pipeline --dataset s3://my-bucket/imagenet \ --cfg "model=resnet,vit" --cfg "optimizer.lr=0.001,0.01,0.1" Hardware presets can sweep alongside hyperparameters: flyte hydra run --multirun --config-path conf --config-name training \ train.py pipeline --dataset s3://my-bucket/imagenet \ --cfg "+task_env=a10g,a100" --cfg "optimizer.lr=0.001,0.01,0.1" ### [Bayesian / TPE sweeps (Optuna)](https://www.union.ai/docs/v2/flyte/integrations/hydra/#bayesian--tpe-sweeps-optuna) Install the sweeper, then activate it via `hydra/sweeper=optuna`. Continuous parameters use `interval(...)`; categorical parameters use `choice(...)`. pip install hydra-optuna-sweeper flyte hydra run --multirun --config-path conf --config-name training \ train.py pipeline --dataset s3://my-bucket/imagenet \ --hydra-override "hydra/sweeper=optuna" \ --hydra-override "hydra.sweeper.n_trials=30" \ --hydra-override "hydra.sweeper.n_jobs=5" \ --cfg "optimizer.lr=interval(1e-4,1e-1)" \ --cfg "optimizer.weight_decay=interval(1e-6,1e-2)" \ --cfg "model=choice(resnet,vit)" When `wait=True`, each remote run’s wrapped result exposes the task output as a float (via `__float__`), so Optuna can use it directly as the trial objective. With `wait=False`, the sweeper sees the run URL but cannot read objective values; use this only for fire-and-forget submission. Other sweepers that respect Hydra’s plugin protocol are activated the same way: install the package, select `hydra/sweeper=`, and set the sweeper’s parameters under `hydra.sweeper.*`. ### [Sweep output directories](https://www.union.ai/docs/v2/flyte/integrations/hydra/#sweep-output-directories) Hydra-namespace overrides redirect where Hydra writes per-job logs and config snapshots: flyte hydra run --multirun --config-path conf --config-name training \ train.py pipeline --dataset s3://my-bucket/imagenet \ --hydra-override "hydra.sweep.dir=./outputs/sweep1" \ --hydra-override "hydra.sweep.subdir=\${hydra.job.num}" \ --cfg "optimizer.lr=0.001,0.01,0.1" [Task environment overrides](https://www.union.ai/docs/v2/flyte/integrations/hydra/#task-environment-overrides) ----------------------------------------------------------------------------------------------------------------- Hydra is good at composing flat YAML; Flyte tasks need richer settings such as resources and container images. The plugin reserves a config key named `task_env` by default that maps task names to `task.override` kwargs. task_env: pipeline: resources: cpu: "2" memory: 8Gi train_model: resources: cpu: "16" memory: 64Gi gpu: "A100:1" When the plugin launches a task, it looks up `task_env[]` (`pipeline` in this example) and applies the values via `task.override(...)`. Resource mappings are converted into `flyte.Resources(**values)` automatically. ### [Prebuilt images](https://www.union.ai/docs/v2/flyte/integrations/hydra/#prebuilt-images) To run a task in a prebuilt container image, set `image` (and optionally `primary_container_name`): task_env: pipeline: image: ghcr.io/acme/flyte-training:latest primary_container_name: main resources: cpu: "4" memory: 16Gi `task.override` does not accept `image` directly. The task image is part of the task definition. Instead, the plugin lowers the override to a `flyte.PodTemplate` whose primary container uses the requested image: * If the task has no inline pod template, a new one is created. * If the task already has an inline `flyte.PodTemplate`, the plugin deep-copies it and sets only the image on the primary container. * If the task references a pod template by name (a string), the plugin raises an error. You must patch a string-named template by editing it in cluster config rather than at submission time. ### [Applying overrides to child tasks](https://www.union.ai/docs/v2/flyte/integrations/hydra/#applying-overrides-to-child-tasks) The launcher only controls the entry task it submits. Child tasks called from within the entry task are not patched automatically. Use `apply_task_env` to apply the same `resources`/`image` handling to a child task before invoking it: from flyteplugins.hydra import apply_task_env @env.task async def pipeline(cfg: DictConfig, dataset: str) -> float: data = await preprocess(cfg) train_task = apply_task_env(train_model, cfg) _, val_loss = await train_task(cfg, data) return val_loss This keeps the override knobs in YAML/CLI surfaces while leaving each task in control of which children it patches. ### [Renaming the task-env key](https://www.union.ai/docs/v2/flyte/integrations/hydra/#renaming-the-task-env-key) If your config uses a different name for the task-env subtree, pass it explicitly: hydra_run(..., task_env_key="task_environment") flyte hydra run --task-env-key task_environment ... ### [What `task_env` should not model](https://www.union.ai/docs/v2/flyte/integrations/hydra/#what-task_env-should-not-model) The YAML schema intentionally omits the full Kubernetes `V1PodSpec`. Keep advanced pod configuration (volumes, init containers, node selectors, etc.) in Python task/environment code where you have a real type. Use Hydra `task_env` presets for the common knobs only: image, primary container name and resources. [Structured configs (without YAML)](https://www.union.ai/docs/v2/flyte/integrations/hydra/#structured-configs-without-yaml) ----------------------------------------------------------------------------------------------------------------------------- Structured configs work with this plugin as long as they are registered before the launcher composes the config. `flyte hydra run` imports the script first, so top-level `ConfigStore.instance().store(...)` calls run before composition. from dataclasses import dataclass, field from hydra.core.config_store import ConfigStore from omegaconf import DictConfig @dataclass class TrainingConf: epochs: int = 30 batch_size: int = 64 @dataclass class RootConf: training: TrainingConf = field(default_factory=TrainingConf) ConfigStore.instance().store(name="structured_training", node=RootConf) Run a fully-structured config without YAML: flyte hydra run --config-name structured_training \ train.py pipeline --dataset s3://my-bucket/imagenet The same config also works through `@hydra.main`: python train.py --config-name structured_training If the structured config still references YAML config groups, keep `--config-path conf`. If everything is registered in `ConfigStore`, omit `--config-path`. Do not register structured configs only inside `if __name__ == "__main__":` or inside the `@hydra.main` function body. `flyte hydra run` and shell completion inspect the script at import time, before either of those blocks runs, and registrations placed there will not be visible. Structured configs sweep just like YAML configs: runs = hydra_sweep( pipeline, config_path=None, config_name="structured_training", overrides=["training.epochs=10,20", "training.batch_size=32,64"], dataset="s3://my-bucket/imagenet", mode="remote", ) LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/hydra/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Code generation | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Code generation =============== The code generation plugin turns natural-language prompts into tested, production-ready Python code. You describe what the code should do, along with sample data, schema definitions, constraints, and typed inputs/outputs, and the plugin handles the rest: generating code, writing tests, building an isolated [code sandbox](https://www.union.ai/docs/v2/union/user-guide/sandboxing/code-sandboxing) with the right dependencies, running the tests, diagnosing failures, and iterating until everything passes. The result is a validated script you can execute against real data or deploy as a reusable Flyte task. [Installation](https://www.union.ai/docs/v2/flyte/integrations/codegen/#installation) --------------------------------------------------------------------------------------- pip install flyteplugins-codegen # For Agent mode (Claude-only) pip install flyteplugins-codegen[agent] [Quick start](https://www.union.ai/docs/v2/flyte/integrations/codegen/#quick-start) ------------------------------------------------------------------------------------- import flyte from flyte.io import File from flyte.sandbox import sandbox_environment from flyteplugins.codegen import AutoCoderAgent agent = AutoCoderAgent(model="gpt-4.1", name="summarize-sales") env = flyte.TaskEnvironment( name="my-env", secrets=[flyte.Secret(key="openai_key", as_env_var="OPENAI_API_KEY")], image=flyte.Image.from_debian_base().with_pip_packages( "flyteplugins-codegen", ), depends_on=[sandbox_environment], ) @env.task async def process_data(csv_file: File) -> tuple[float, int, int]: result = await agent.generate.aio( prompt="Read the CSV and compute total_revenue, total_units and row_count.", samples={"sales": csv_file}, outputs={"total_revenue": float, "total_units": int, "row_count": int}, ) return await result.run.aio() The `depends_on=[sandbox_environment]` declaration is required. It ensures the sandbox runtime is available when dynamically-created sandboxes execute. ![Sandbox](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/codegen/sandbox.png) [Two execution backends](https://www.union.ai/docs/v2/flyte/integrations/codegen/#two-execution-backends) ----------------------------------------------------------------------------------------------------------- The plugin supports two backends for generating and validating code. Both share the same `AutoCoderAgent` interface and produce the same `CodeGenEvalResult`. ### [LiteLLM (default)](https://www.union.ai/docs/v2/flyte/integrations/codegen/#litellm-default) Uses structured-output LLM calls to generate code, detect packages, build sandbox images, run tests, diagnose failures, and iterate. Works with any model that supports structured outputs (GPT-4, Claude, Gemini, etc. via LiteLLM). agent = AutoCoderAgent( name="my-task", model="gpt-4.1", max_iterations=10, ) The LiteLLM backend follows a fixed pipeline: flowchart TD A\["prompt + samples"\] --> B\["generate\_plan"\] B --> C\["generate\_code"\] C --> D\["detect\_packages"\] D --> E\["build\_image"\] E --> F{skip\_tests?} F -- yes --> G\["return result"\] F -- no --> H\["generate\_tests"\] H --> I\["execute\_tests"\] I --> J{pass?} J -- yes --> G J -- no --> K\["diagnose\_error"\] K --> L{error type?} L -- "logic error" --> M\["regenerate code"\] L -- "environment error" --> N\["add packages, rebuild image"\] L -- "test error" --> O\["fix test expectations"\] M --> I N --> I O --> I The loop continues until tests pass or `max_iterations` is reached. ![LiteLLM](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/codegen/litellm.png) ### [Agent (Claude)](https://www.union.ai/docs/v2/flyte/integrations/codegen/#agent-claude) Uses the Claude Agent SDK to autonomously generate, test, and fix code. The agent has access to `Bash`, `Read`, `Write`, and `Edit` tools and decides what to do at each step. Test execution commands (`pytest`) are intercepted and run inside isolated sandboxes. agent = AutoCoderAgent( name="my-task", model="claude-sonnet-4-5-20250929", backend="claude", ) Agent mode requires `ANTHROPIC_API_KEY` as a Flyte secret and is Claude-only. **Key differences from LiteLLM:** | | LiteLLM | Agent | | --- | --- | --- | | **Execution** | Fixed generate-test-fix pipeline | Autonomous agent decides actions | | **Model support** | Any model with structured outputs | Claude only | | **Iteration control** | `max_iterations` | `agent_max_turns` | | **Test execution** | Direct sandbox execution | `pytest` commands intercepted via hooks | | **Tool safety** | N/A | Commands classified as safe/denied/intercepted | | **Observability** | Logs + token counts | Full tool call tracing in Flyte UI | In Agent mode, Bash commands are classified before execution: * **Safe** (`ls`, `cat`, `grep`, `head`, etc.) — allowed to run directly * **Intercepted** (`pytest`) — routed to sandbox execution * **Denied** (`apt`, `pip install`, `curl`, etc.) — blocked for safety [Providing data](https://www.union.ai/docs/v2/flyte/integrations/codegen/#providing-data) ------------------------------------------------------------------------------------------- ### [Sample data](https://www.union.ai/docs/v2/flyte/integrations/codegen/#sample-data) Pass sample data via `samples` as `File` objects or pandas `DataFrame`s. The plugin automatically: 1. Converts DataFrames to CSV files 2. Infers [Pandera](https://pandera.readthedocs.io/) schemas from the data — column types, nullability 3. Parses natural-language `constraints` into Pandera checks (e.g., `"quantity must be positive"` becomes `pa.Check.gt(0)`) 4. Extracts data context — column statistics, distributions, patterns, sample rows 5. Injects all of this into the LLM prompt so the generated code is aware of the exact data structure Pandera is used purely for prompt enrichment, not runtime validation. The generated code does not import Pandera — it benefits from the LLM knowing the precise data structure. The generated schemas are stored on `result.generated_schemas` for inspection. result = await agent.generate.aio( prompt="Clean and validate the data, remove duplicates", samples={"orders": orders_df, "products": products_file}, constraints=["quantity must be positive", "price between 0 and 10000"], outputs={"cleaned_orders": File}, ) ### [Schema and constraints](https://www.union.ai/docs/v2/flyte/integrations/codegen/#schema-and-constraints) Use `schema` to provide free-form context about data formats or target structures (e.g., a database schema). Use `constraints` to declare business rules that the generated code must respect: result = await agent.generate.aio( prompt=prompt, samples={"readings": sensor_df}, schema="""Output JSON schema for report_json: { "sensor_id": str, "avg_temp": float, "min_temp": float, "max_temp": float, "avg_humidity": float, } """, constraints=[\ "Temperature values must be between -40 and 60 Celsius",\ "Humidity values must be between 0 and 100 percent",\ "Output report must have one row per unique sensor_id",\ ], outputs={ "report_json": str, "total_anomalies": int, }, ) ![Pandera Constraints](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/codegen/pandera_constraints.png) ### [Inputs and outputs](https://www.union.ai/docs/v2/flyte/integrations/codegen/#inputs-and-outputs) Declare `inputs` for non-sample arguments (e.g., thresholds, flags) and `outputs` for the expected result types. Supported output types: `str`, `int`, `float`, `bool`, `datetime.datetime`, `datetime.timedelta`, `File`. Sample entries are automatically added as `File` inputs — you do not need to redeclare them. result = await agent.generate.aio( prompt="Filter transactions above the threshold", samples={"transactions": tx_file}, inputs={"threshold": float, "include_pending": bool}, outputs={"filtered": File, "count": int}, ) [Running generated code](https://www.union.ai/docs/v2/flyte/integrations/codegen/#running-generated-code) ----------------------------------------------------------------------------------------------------------- `agent.generate()` returns a `CodeGenEvalResult`. If `result.success` is `True`, the generated code passed all tests and you can execute it against real data. If `max_iterations` (LiteLLM) or `agent_max_turns` (Agent) is reached without tests passing, `result.success` is `False` and `result.error` contains the failure details. Both `run()` and `as_task()` return output values as a tuple in the order declared in `outputs`. If there is a single output, the value is returned directly (not wrapped in a tuple). ### [One-shot execution with `result.run()`](https://www.union.ai/docs/v2/flyte/integrations/codegen/#one-shot-execution-with-resultrun) Runs the generated code in a sandbox. If samples were provided during `generate()`, they are used as default inputs. # Use sample data as defaults total_revenue, total_units, count = await result.run.aio() # Override specific inputs total_revenue, total_units, count = await result.run.aio(threshold=0.5) # Sync version total_revenue, total_units, count = result.run() `result.run()` accepts optional configuration: total_revenue, total_units, count = await result.run.aio( name="execute-on-data", resources=flyte.Resources(cpu=2, memory="4Gi"), retries=2, timeout=600, cache="auto", ) ### [Reusable task with `result.as_task()`](https://www.union.ai/docs/v2/flyte/integrations/codegen/#reusable-task-with-resultas_task) Creates a callable sandbox task from the generated code. Useful when you want to run the same generated code against different data. task = result.as_task( name="run-sensor-analysis", resources=flyte.Resources(cpu=1, memory="512Mi"), ) # Call with sample defaults report, total_anomalies = await task.aio() # Call with different data report, total_anomalies = await task.aio(readings=new_data_file) [Error diagnosis](https://www.union.ai/docs/v2/flyte/integrations/codegen/#error-diagnosis) --------------------------------------------------------------------------------------------- The LiteLLM backend classifies test failures into three categories and applies targeted fixes: | Error type | Meaning | Action | | --- | --- | --- | | `logic` | Bug in the generated code | Regenerate code with specific patch instructions | | `environment` | Missing package or dependency | Add the package and rebuild the sandbox image | | `test_error` | Bug in the generated test | Fix the test expectations | If the same error persists after a fix, the plugin reclassifies it (e.g., `logic` to `test_error`) to try the other approach. In Agent mode, the agent diagnoses and fixes issues autonomously based on error output. [Durable execution](https://www.union.ai/docs/v2/flyte/integrations/codegen/#durable-execution) ------------------------------------------------------------------------------------------------- Code generation is expensive — it involves multiple LLM calls, image builds, and sandbox executions. Without durability, a transient failure in the pipeline (network blip, OOM, downstream service error) would force the entire process to restart from scratch: regenerating code, rebuilding images, re-running sandboxes, making additional LLM calls. Flyte solves this through two complementary mechanisms: **replay logs** and **caching**. ### [Replay logs](https://www.union.ai/docs/v2/flyte/integrations/codegen/#replay-logs) Flyte maintains a replay log that records every trace and task execution within a run. When a task crashes and retries, the system replays the log from the previous attempt rather than recomputing everything: * No additional model calls * No code regeneration * No sandbox re-execution * No container rebuilds The workflow breezes through the earlier steps and resumes from the failure point. This applies as long as the traces and tasks execute in the same order and use the same inputs as the first attempt. ### [Caching](https://www.union.ai/docs/v2/flyte/integrations/codegen/#caching) Separately, Flyte can cache task results across runs. With `cache="auto"`, sandbox executions (image builds, test runs, code execution) are cached. This is useful when you re-run the same pipeline — not just when recovering from a crash, but across entirely separate invocations with the same inputs. Together, replay logs handle crash recovery within a run, and caching avoids redundant work across runs. ### [Non-determinism in Agent mode](https://www.union.ai/docs/v2/flyte/integrations/codegen/#non-determinism-in-agent-mode) One challenge with agents is that they are inherently non-deterministic — the sequence of actions can vary between runs, which could break replay. In practice, the codegen agent follows a predictable pattern (write code, generate tests, run tests, inspect results), which works in replay’s favor. The plugin also embeds logic that instructs the agent not to regenerate or re-execute steps that already completed successfully in the first run. This acts as an additional safety check alongside the replay log to account for non-determinism. ![Agent](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/codegen/agent.png) On the first attempt, the full pipeline runs. If a transient failure occurs, the system instantly replays the traces (which track model calls) and sandbox executions, allowing the pipeline to resume from the point of failure. ![Durability](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/codegen/durability.png) [Observability](https://www.union.ai/docs/v2/flyte/integrations/codegen/#observability) ----------------------------------------------------------------------------------------- ### [LiteLLM backend](https://www.union.ai/docs/v2/flyte/integrations/codegen/#litellm-backend) * Logs every iteration with attempt count, error type, and package changes * Tracks total input/output tokens across all LLM calls (available on `result.total_input_tokens` and `result.total_output_tokens`) * Results include full conversation history for debugging (`result.conversation_history`) ### [Agent backend](https://www.union.ai/docs/v2/flyte/integrations/codegen/#agent-backend) * Traces each tool call (name + input) via `PostToolUse` hooks * Traces tool failures via `PostToolUseFailure` hooks * Traces a summary when the agent finishes (total tool calls, tool distribution, final image/packages) * Classifies Bash commands as safe, denied, or intercepted (for sandbox execution) * All traces appear in the Flyte UI [Examples](https://www.union.ai/docs/v2/flyte/integrations/codegen/#examples) ------------------------------------------------------------------------------- ### [Processing CSVs with different schemas](https://www.union.ai/docs/v2/flyte/integrations/codegen/#processing-csvs-with-different-schemas) Generate code that handles varying CSV formats, then run on real data: from flyteplugins.codegen import AutoCoderAgent agent = AutoCoderAgent( name="sales-processor", model="gpt-4.1", max_iterations=5, resources=flyte.Resources(cpu=1, memory="512Mi"), litellm_params={"temperature": 0.2, "max_tokens": 4096}, ) @env.task async def process_sales(csv_file: File) -> dict[str, float | int]: result = await agent.generate.aio( prompt="Read the CSV and compute total_revenue, total_units, and transaction_count.", samples={"csv_data": csv_file}, outputs={ "total_revenue": float, "total_units": int, "transaction_count": int, }, ) if not result.success: raise RuntimeError(f"Code generation failed: {result.error}") total_revenue, total_units, transaction_count = await result.run.aio() return { "total_revenue": total_revenue, "total_units": total_units, "transaction_count": transaction_count, } ### [DataFrame analysis with constraints](https://www.union.ai/docs/v2/flyte/integrations/codegen/#dataframe-analysis-with-constraints) Pass DataFrames directly and enforce business rules with constraints: agent = AutoCoderAgent( model="gpt-4.1", name="sensor-analysis", base_packages=["numpy"], max_sample_rows=30, ) @env.task async def analyze_sensors(sensor_df: pd.DataFrame) -> tuple[File, int]: result = await agent.generate.aio( prompt="""Analyze IoT sensor data. For each sensor, calculate mean/min/max temperature, mean humidity, and count warnings. Output a summary CSV.""", samples={"readings": sensor_df}, constraints=[\ "Temperature values must be between -40 and 60 Celsius",\ "Humidity values must be between 0 and 100 percent",\ "Output report must have one row per unique sensor_id",\ ], outputs={ "report": File, "total_anomalies": int, }, ) if not result.success: raise RuntimeError(f"Code generation failed: {result.error}") task = result.as_task( name="run-sensor-analysis", resources=flyte.Resources(cpu=1, memory="512Mi"), ) return await task.aio(readings=result.original_samples["readings"]) ### [Agent mode](https://www.union.ai/docs/v2/flyte/integrations/codegen/#agent-mode) The same task using Claude as an autonomous agent: agent = AutoCoderAgent( name="sales-agent", backend="claude", model="claude-sonnet-4-5-20250929", resources=flyte.Resources(cpu=1, memory="512Mi"), ) @env.task async def process_sales_with_agent(csv_file: File) -> dict[str, float | int]: result = await agent.generate.aio( prompt="Read the CSV and compute total_revenue, total_units, and transaction_count.", samples={"csv_data": csv_file}, outputs={ "total_revenue": float, "total_units": int, "transaction_count": int, }, ) if not result.success: raise RuntimeError(f"Agent code generation failed: {result.error}") total_revenue, total_units, transaction_count = await result.run.aio() return { "total_revenue": total_revenue, "total_units": total_units, "transaction_count": transaction_count, } [Configuration](https://www.union.ai/docs/v2/flyte/integrations/codegen/#configuration) ----------------------------------------------------------------------------------------- ### [LiteLLM parameters](https://www.union.ai/docs/v2/flyte/integrations/codegen/#litellm-parameters) Tune model behavior with `litellm_params`: agent = AutoCoderAgent( name="my-task", model="anthropic/claude-sonnet-4-20250514", api_key="ANTHROPIC_API_KEY", litellm_params={ "temperature": 0.3, "max_tokens": 4000, }, ) ### [Image configuration](https://www.union.ai/docs/v2/flyte/integrations/codegen/#image-configuration) Control the registry and Python version for sandbox images: from flyte.sandbox import ImageConfig agent = AutoCoderAgent( name="my-task", model="gpt-4.1", image_config=ImageConfig( registry="my-registry.io", registry_secret="registry-creds", python_version=(3, 12), ), ) ### [Skipping tests](https://www.union.ai/docs/v2/flyte/integrations/codegen/#skipping-tests) Set `skip_tests=True` to skip test generation and execution. The agent still generates code, detects packages, and builds the sandbox image, but does not generate or run tests. agent = AutoCoderAgent( name="my-task", model="gpt-4.1", skip_tests=True, ) `skip_tests` only applies to LiteLLM mode. In Agent mode, the agent autonomously decides when to test. ### [Base packages](https://www.union.ai/docs/v2/flyte/integrations/codegen/#base-packages) Ensure specific packages are always installed in every sandbox: agent = AutoCoderAgent( name="my-task", model="gpt-4.1", base_packages=["numpy", "pandas"], ) [Best practices](https://www.union.ai/docs/v2/flyte/integrations/codegen/#best-practices) ------------------------------------------------------------------------------------------- * **One agent per task.** Each `generate()` call builds its own sandbox image and manages its own package state. Running multiple agents in the same task can cause resource contention and makes failures harder to diagnose. * **Keep `cache="auto"` (the default).** Caching flows to all internal sandboxes, making retries near-instant. Use `"disable"` during development if you want fresh executions, or `"override"` to force re-execution and update the cached result. * **Set `max_iterations` conservatively.** Start with 5-10 iterations. If the model cannot produce correct code in that budget, the prompt or constraints likely need refinement. * **Provide constraints for data-heavy tasks.** Explicit constraints (e.g., `"quantity must be positive"`) produce better schemas and better generated code. * **Inspect `result.generated_schemas`.** Review the inferred Pandera schemas to verify the model understood your data structure correctly. [API reference](https://www.union.ai/docs/v2/flyte/integrations/codegen/#api-reference) ----------------------------------------------------------------------------------------- ### [`AutoCoderAgent` constructor](https://www.union.ai/docs/v2/flyte/integrations/codegen/#autocoderagent-constructor) | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `name` | `str` | `"auto-coder"` | Unique name for tracking and image naming | | `model` | `str` | `"gpt-4.1"` | LiteLLM model identifier | | `backend` | `str` | `"litellm"` | Execution backend: `"litellm"` or `"claude"` | | `system_prompt` | `str` | `None` | Custom system prompt override | | `api_key` | `str` | `None` | Name of the environment variable containing the LLM API key (e.g., `"OPENAI_API_KEY"`) | | `api_base` | `str` | `None` | Custom API base URL | | `litellm_params` | `dict` | `None` | Extra LiteLLM params (temperature, max\_tokens, etc.) | | `base_packages` | `list[str]` | `None` | Always-install pip packages | | `resources` | `flyte.Resources` | `None` | Resources for sandbox execution (default: 1 CPU, 1Gi) | | `image_config` | `ImageConfig` | `None` | Registry, secret, and Python version | | `max_iterations` | `int` | `10` | Max generate-test-fix iterations (LiteLLM mode) | | `max_sample_rows` | `int` | `100` | Rows to sample from data for LLM context | | `skip_tests` | `bool` | `False` | Skip test generation and execution (LiteLLM mode) | | `sandbox_retries` | `int` | `0` | Flyte task-level retries for each sandbox execution | | `timeout` | `int` | `None` | Timeout in seconds for sandboxes | | `env_vars` | `dict[str, str]` | `None` | Environment variables for sandboxes | | `secrets` | `list[Secret]` | `None` | Flyte secrets for sandboxes | | `cache` | `str` | `"auto"` | Cache behavior: `"auto"`, `"override"`, or `"disable"` | | `agent_max_turns` | `int` | `50` | Max turns when `backend="claude"` | ### [`generate()` parameters](https://www.union.ai/docs/v2/flyte/integrations/codegen/#generate-parameters) | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `prompt` | `str` | required | Natural-language task description | | `schema` | `str` | `None` | Free-form context about data formats or target structures | | `constraints` | `list[str]` | `None` | Natural-language constraints (e.g., `"quantity must be positive"`) | | `samples` | `dict[str, File \| DataFrame]` | `None` | Sample data. DataFrames are auto-converted to CSV files. | | `inputs` | `dict[str, type]` | `None` | Non-sample input types (e.g., `{"threshold": float}`) | | `outputs` | `dict[str, type]` | `None` | Output types. Supported: `str`, `int`, `float`, `bool`, `datetime`, `timedelta`, `File` | ### [`CodeGenEvalResult` fields](https://www.union.ai/docs/v2/flyte/integrations/codegen/#codegenevalresult-fields) | Field | Type | Description | | --- | --- | --- | | `success` | `bool` | Whether tests passed | | `solution` | `CodeSolution` | Generated code (`.code`, `.language`, `.system_packages`) | | `tests` | `str` | Generated test code | | `output` | `str` | Test output | | `exit_code` | `int` | Test exit code | | `error` | `str \| None` | Error message if failed | | `attempts` | `int` | Number of iterations used | | `image` | `str` | Built sandbox image with all dependencies | | `detected_packages` | `list[str]` | Pip packages detected | | `detected_system_packages` | `list[str]` | Apt packages detected | | `generated_schemas` | `dict[str, str] \| None` | Pandera schemas as Python code strings | | `data_context` | `str \| None` | Extracted data context | | `original_samples` | `dict[str, File] \| None` | Sample data as Files (defaults for `run()`/`as_task()`) | | `total_input_tokens` | `int` | Total input tokens across all LLM calls | | `total_output_tokens` | `int` | Total output tokens across all LLM calls | | `conversation_history` | `list[dict]` | Full LLM conversation history for debugging | ### [`CodeGenEvalResult` methods](https://www.union.ai/docs/v2/flyte/integrations/codegen/#codegenevalresult-methods) | Method | Description | | --- | --- | | `result.run(**overrides)` | Execute generated code in a sandbox. Sample data used as defaults. | | `await result.run.aio(**overrides)` | Async version of `run()`. | | `result.as_task(name, ...)` | Create a reusable callable sandbox task from the generated code. | Both `run()` and `as_task()` accept optional `name`, `resources`, `retries`, `timeout`, `env_vars`, `secrets`, and `cache` parameters. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/codegen/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # OmegaConf | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) OmegaConf ========= [OmegaConf](https://omegaconf.readthedocs.io/) is a hierarchical configuration system used by many ML frameworks (and the foundation of [Hydra](https://www.union.ai/docs/v2/flyte/integrations/hydra) ). The `flyteplugins-omegaconf` plugin makes OmegaConf’s `DictConfig` and `ListConfig` first-class types in Flyte tasks, so you can pass entire configs like plain dicts, YAML files or dataclass-backed structured configs between tasks without flattening them into individual scalar arguments. The plugin enables: * `DictConfig` and `ListConfig` as native task input and output types * Round-tripping of structured configs (dataclass schemas) across task boundaries * Preservation of OmegaConf-specific values: `MISSING` sentinels, `Enum`s, `pathlib.Path`s, `tuple`s, and `bytes` * Resolved variable interpolations on the wire * A YAML-rendered Flyte report tab for human-readable config inspection [Installation](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#installation) ----------------------------------------------------------------------------------------- pip install flyteplugins-omegaconf Installing the package automatically registers `DictConfig` and `ListConfig` with Flyte’s `TypeEngine`. No manual setup is required. If you are using the [Hydra plugin](https://www.union.ai/docs/v2/flyte/integrations/hydra) , `flyteplugins-omegaconf` is installed as a transitive dependency. [Quick start](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#quick-start) --------------------------------------------------------------------------------------- import flyte from omegaconf import DictConfig, OmegaConf env = flyte.TaskEnvironment(name="training", image=...) @env.task async def train(cfg: DictConfig) -> float: return run_experiment(cfg.optimizer.lr, cfg.training.epochs) @env.task async def pipeline() -> float: cfg = OmegaConf.create( {"optimizer": {"lr": 0.001}, "training": {"epochs": 10}} ) return await train(cfg) The config is serialized when `train` is invoked and reconstructed as a `DictConfig` inside the task. No type registration, manual encoding or schema declaration is required. [When to use this plugin](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#when-to-use-this-plugin) --------------------------------------------------------------------------------------------------------------- Use `flyteplugins-omegaconf` when: * You already use OmegaConf. For example, you have YAML configs, dataclass-based config trees or a Hydra app, and want to keep that representation intact across task boundaries. * You want to pass a single composed config object instead of widening task signatures with dozens of scalar arguments. * You want to enforce schema validation at the task entry point via dataclass-backed structured configs. * You want resolved interpolations (`${other.value}`) to be materialized at submission time rather than at task runtime. If you do not use OmegaConf elsewhere, prefer plain dataclasses, `pydantic.BaseModel` or `dict` for task inputs as they are supported by Flyte natively without an extra dependency. [Building a DictConfig](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#building-a-dictconfig) ----------------------------------------------------------------------------------------------------------- Any of the standard OmegaConf construction methods produce a value the plugin can serialize. ### [From a plain dict](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#from-a-plain-dict) cfg = OmegaConf.create( {"optimizer": {"lr": 0.001}, "training": {"epochs": 10}} ) flyte.run(train, cfg=cfg) ### [From a YAML file](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#from-a-yaml-file) cfg = OmegaConf.load("configs/training.yaml") flyte.run(train, cfg=cfg) The file is read locally on the submitter, not on the worker. If the YAML lives in your project tree and needs to be packaged into the task image, use `flyte.with_runcontext(copy_style="all").run(...)`. ### [From a dataclass (structured config)](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#from-a-dataclass-structured-config) from dataclasses import dataclass @dataclass class TrainConf: lr: float = 0.001 epochs: int = 10 cfg = OmegaConf.structured(TrainConf()) flyte.run(train, cfg=cfg) Structured configs are covered in detail in [Structured configs](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#structured-configs) below. ### [From a base config plus overrides](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#from-a-base-config-plus-overrides) base = OmegaConf.load("configs/training.yaml") override = OmegaConf.create({"optimizer": {"lr": 0.01}}) cfg = OmegaConf.merge(base, override) flyte.run(train, cfg=cfg) This is the same pattern Hydra uses internally. See the [Hydra integration](https://www.union.ai/docs/v2/flyte/integrations/hydra) for a full composition layer on top of this plugin. [Variable interpolation](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#variable-interpolation) ------------------------------------------------------------------------------------------------------------- OmegaConf supports `${...}` interpolations that resolve relative to the config tree: cfg = OmegaConf.create( { "base_lr": 0.01, "optimizer": {"lr": "${base_lr}", "momentum": 0.9}, } ) flyte.run(train, cfg=cfg) Interpolations are resolved at serialization time. By the time the task runs, `cfg.optimizer.lr` is the concrete float `0.01`, not the string `"${base_lr}"`. This means: * The receiving task does not need any context that only existed in the submitter’s environment. * Resolved values appear in the Flyte I/O panel. * A reference that fails to resolve at submission time fails fast, before any task runs. If you need lazy resolution on the worker, resolve the reference yourself inside the task or pass the unresolved string through a normal `str` input. [Nested and deeply structured configs](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#nested-and-deeply-structured-configs) ----------------------------------------------------------------------------------------------------------------------------------------- Nested configs are supported, including deeply structured OmegaConf objects. cfg = OmegaConf.create( { "experiment": { "model": { "encoder": { "attention": {"num_heads": 8, "head_dim": 64}, "ffn": {"hidden_dim": 2048, "activation": "gelu"}, }, "decoder": {"num_layers": 6}, } } } ) @env.task async def extract_leaf(cfg: DictConfig) -> int: return int(cfg.experiment.model.encoder.attention.num_heads) [DictConfigs that contain lists](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#dictconfigs-that-contain-lists) ----------------------------------------------------------------------------------------------------------------------------- A `DictConfig` may hold list values; they are reconstructed as nested `ListConfig`s on the receiving side. cfg = OmegaConf.create( { "model": { "layer_sizes": [64, 128, 256, 512], "activations": ["relu", "relu", "relu", "sigmoid"], }, "data": { "augmentations": ["random_flip", "random_crop", "color_jitter"], "input_size": [224, 224], }, } ) @env.task async def double_layer_sizes(cfg: DictConfig) -> DictConfig: doubled = [size * 2 for size in cfg.model.layer_sizes] return OmegaConf.merge(cfg, {"model": {"layer_sizes": doubled}}) [ListConfig as input and output](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#listconfig-as-input-and-output) ----------------------------------------------------------------------------------------------------------------------------- `ListConfig` is symmetric with `DictConfig` and supports the same construction patterns. ### [Lists of primitives](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#lists-of-primitives) @env.task async def scale_values(values: ListConfig, factor: float) -> ListConfig: return OmegaConf.create([v * factor for v in values]) ### [Building a schedule from another task](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#building-a-schedule-from-another-task) @env.task async def build_lr_schedule(base_lr: float, num_stages: int) -> ListConfig: return OmegaConf.create([base_lr * (0.5 ** i) for i in range(num_stages)]) @env.task async def train_with_schedule(cfg: DictConfig, lr_schedule: ListConfig) -> float: final_lr = float(lr_schedule[-1]) ... ### [Nested lists (list of lists)](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#nested-lists-list-of-lists) grid = OmegaConf.create([[0.001, 0.01, 0.1], [10, 20, 50]]) @env.task async def flatten_grid(grid: ListConfig) -> ListConfig: flat = [item for sublist in OmegaConf.to_container(grid) for item in sublist] return OmegaConf.create(flat) ### [Lists of DictConfigs](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#lists-of-dictconfigs) configs = OmegaConf.create( [\ {"optimizer": {"lr": 0.001}, "training": {"epochs": 10}},\ {"optimizer": {"lr": 0.01}, "training": {"epochs": 20}},\ {"optimizer": {"lr": 0.1}, "training": {"epochs": 5}},\ ] ) @env.task async def select_best_config(configs: ListConfig) -> DictConfig: best = max(OmegaConf.to_container(configs), key=lambda c: c["optimizer"]["lr"]) return OmegaConf.create(best) ### [Lists of dataclass instances](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#lists-of-dataclass-instances) @dataclass class LayerConf: name: str width: int activation: str layers = OmegaConf.create( [\ LayerConf(name="encoder", width=768, activation="gelu"),\ LayerConf(name="bottleneck", width=128, activation="relu"),\ LayerConf(name="decoder", width=768, activation="linear"),\ ] ) Each element round-trips as a typed `DictConfig` backed by `LayerConf`, so the receiving task can call `OmegaConf.get_type(layers[0])` and access fields with attribute notation. ListConfig is always plain. Even when its elements are dataclass-backed, the outer `ListConfig` does not carry a list-level schema as there is no structured (typed-element) `ListConfig` in OmegaConf. This affects only the outer container; nested elements retain their schemas. [Structured configs](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#structured-configs) ----------------------------------------------------------------------------------------------------- A structured config is a `DictConfig` that is bound to a Python dataclass. The dataclass acts as a schema: assigning a value of the wrong type raises `omegaconf.ValidationError`, and merging unknown keys raises an error instead of silently extending the config. ### [Basic structured config](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#basic-structured-config) from dataclasses import dataclass, field from omegaconf import OmegaConf, DictConfig @dataclass class OptimizerConf: lr: float = 0.001 weight_decay: float = 1e-4 @dataclass class TrainConf: optimizer: OptimizerConf = field(default_factory=OptimizerConf) epochs: int = 10 cfg = OmegaConf.structured(TrainConf()) flyte.run(train, cfg=cfg) # cfg.optimizer.lr = "oops" # raises omegaconf.ValidationError ### [Schema reconstruction in the receiving task](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#schema-reconstruction-in-the-receiving-task) When a structured `DictConfig` is deserialized in a downstream task, the plugin operates in **Auto mode**: it reads the originating dataclass name from the wire payload and tries to import it. Two outcomes are possible: * Dataclass importable in the receiving task: `cfg` is reconstructed as a `TrainConf`\-backed `DictConfig`. `OmegaConf.get_type(cfg)` returns `TrainConf`, and type validation is enforced. * Dataclass not importable: `cfg` falls back to a plain `DictConfig` carrying the raw values. `OmegaConf.get_type(cfg)` returns `dict`. The values are intact but the schema is lost. To keep schemas across task hops, define dataclasses in modules that are importable from every task in the pipeline (for example, in a shared `configs.py` module bundled into the task image). ### [Required (`MISSING`) fields](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#required-missing-fields) OmegaConf’s `MISSING` sentinel marks a required field that has no default: from omegaconf import MISSING @dataclass class TrainConf: data_path: str = MISSING epochs: int = 10 # Pass with MISSING still unset — serialization succeeds. cfg = OmegaConf.structured(TrainConf()) flyte.run(train, cfg=cfg) # Or fill it before passing. cfg = OmegaConf.structured(TrainConf(data_path="/data/imagenet")) flyte.run(train, cfg=cfg) A config with an unset `MISSING` field serializes and deserializes successfully as the sentinel is preserved on the wire. Accessing the field on the receiving side raises `MissingMandatoryValue`. Type annotations are preserved only in Auto mode. When the dataclass is importable on the receiving side, an unfilled `MISSING` field still carries its declared type (e.g. `StringNode` for `str`). When the plugin falls back to a plain `DictConfig` because the dataclass is not importable, the field becomes an `AnyNode` where the value is preserved, but the type annotation is not. ### [Advanced field types](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#advanced-field-types) Beyond primitives and nested dataclasses, structured configs may declare fields of these types and they will round-trip with their schemas intact: * `Enum` subclasses * `pathlib.Path` * `Optional[T]` * `bytes` * `dict[str, T]` where `T` is a dataclass * `list[T]` where `T` is a dataclass from enum import Enum from pathlib import Path from typing import Optional class RunMode(Enum): TRAIN = "train" EVAL = "eval" @dataclass class CallbackConf: name: str = "early_stop" patience: int = 3 monitor: str = MISSING @dataclass class AdvancedTrainConf: mode: RunMode = RunMode.TRAIN checkpoint_dir: Path = Path("/tmp/checkpoints") maybe_seed: Optional[int] = None payload: bytes = b"default-token" callbacks_by_name: dict[str, CallbackConf] = field( default_factory=lambda: { "early_stop": CallbackConf(name="early_stop", patience=3), "checkpoint": CallbackConf(name="checkpoint", monitor="val_loss"), } ) callbacks: list[CallbackConf] = field( default_factory=lambda: [\ CallbackConf(name="lr_monitor", patience=2, monitor="lr"),\ CallbackConf(name="nan_guard", patience=1, monitor="loss"),\ ] ) Inside a downstream task: @env.task async def inspect(cfg: DictConfig) -> str: assert OmegaConf.get_type(cfg) == AdvancedTrainConf assert OmegaConf.get_type(cfg.callbacks[0]) == CallbackConf assert isinstance(cfg.mode, RunMode) assert isinstance(cfg.checkpoint_dir, Path) assert isinstance(cfg.payload, bytes) return cfg.mode.value ### [Merging overrides on top of a structured base](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#merging-overrides-on-top-of-a-structured-base) @env.task async def structured_merge_pipeline() -> str: base = OmegaConf.structured(TrainConf()) overrides = OmegaConf.create( { "optimizer": {"lr": 0.05}, "training": {"epochs": 100}, "experiment_name": "sweep-run-1", } ) cfg = OmegaConf.merge(base, overrides) return await validate_config(cfg) Merging an unknown key against a structured config raises an error, so define every key the override layer might supply on the dataclass. [Embedding rich Python values inside a plain DictConfig](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#embedding-rich-python-values-inside-a-plain-dictconfig) ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- A plain `DictConfig` (one not bound to a dataclass) can still hold Python values that OmegaConf does not natively model. The plugin preserves the following types end-to-end whether they appear in plain or structured configs: * `pathlib.Path` and any subclass of `pathlib.PurePath` * `enum.Enum` members * `tuple` (round-trips as `tuple`, not `list`) * `bytes` cfg = OmegaConf.create({"model_path": Path("/opt/models/model.bin")}) @env.task async def use_path(cfg: DictConfig) -> str: assert isinstance(cfg.model_path, Path) return f"model_path={cfg.model_path}" If an `Enum`’s class cannot be imported in the receiving environment, the value is returned as the underlying primitive (`int`, `str`, …) instead of the enum member. [Reserved-looking keys](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#reserved-looking-keys) ----------------------------------------------------------------------------------------------------------- The plugin’s wire format uses an internal payload marker (`__flyte_omegaconf__`), which means user-facing keys named `kind`, `values`, `name`, `value`, `type`, or `schema` round-trip unchanged: cfg = OmegaConf.create({"kind": "training-job", "values": {"lr": 0.001}}) @env.task async def use_payload_shaped_config(cfg: DictConfig) -> str: # cfg.values resolves to DictConfig.values() — use bracket notation # to reach the user key named "values". return f"kind={cfg.kind} lr={cfg['values'].lr}" The only practical consideration is Python’s normal attribute-vs-method conflict: `cfg.values` is the `.values()` method, so reach for `cfg["values"]` when your config has a key with that name. [YAML reports](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#yaml-reports) ----------------------------------------------------------------------------------------- The Flyte I/O panel displays the literal wire representation of a `DictConfig`. ![Wire Representation](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/omegaconf/input.png) For a YAML view, enable a Flyte report on the task and log the config with `log_yaml`: from flyteplugins.omegaconf import log_yaml @env.task(report=True) async def train(cfg: DictConfig) -> DictConfig: await log_yaml.aio(cfg, title="Input config") ... ![YAML Report](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/omegaconf/yaml_repr.png) The plugin also exposes: * `to_yaml(cfg)`: render an OmegaConf container as a YAML string. * `to_html(cfg, title=...)`: wrap the YAML in escaped HTML for embedding in a custom report. * `replace_yaml(cfg, ...)`: replace the contents of a report tab instead of appending. from flyteplugins.omegaconf.report import to_yaml, replace_yaml text = to_yaml(cfg) await replace_yaml.aio(cfg, tab="Final config") `MISSING` fields appear as `???` in the YAML output, matching OmegaConf’s own convention. [Wire format](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#wire-format) --------------------------------------------------------------------------------------- Both `DictConfig` and `ListConfig` are serialized as MessagePack blobs with the literal representation: Literal(scalar=Scalar(binary=Binary(value=, tag="msgpack"))) The msgpack payload uses an internal tagged structure to distinguish OmegaConf-specific concepts from raw values: * A `DictConfig` payload includes the originating dataclass name (`builtins.dict` for plain configs) plus its values. * `MISSING`, `Enum`, `Path`, and `tuple` values carry tagged shapes so they can be reconstructed faithfully. You normally do not need to inspect this format. It is documented here because: * The plugin serializes with `resolve=True`, so the wire representation always contains concrete values for `${...}` interpolations. * Cache-key metadata is set via Flyte’s `MESSAGEPACK` serialization format, so two tasks given equivalent configs hit the same cache entry. [End-to-end example](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/#end-to-end-example) ----------------------------------------------------------------------------------------------------- The example below ties the pieces together: a structured `DictConfig` is created in a parent task, flows through several child tasks that read and modify it, and a `ListConfig` produced midway is consumed by a later stage. Each hop serializes and deserializes the config; the dataclass schema is recovered on the receiving side because `TrainConf` (and friends) are importable in every task in the pipeline. example.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/omegaconf/example.py "View source on GitHub") from dataclasses import dataclass, field import flyte from omegaconf import DictConfig, ListConfig, OmegaConf env = flyte.TaskEnvironment( name="omegaconf-pipeline-example", image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-omegaconf"), ) @dataclass class OptimizerConf: lr: float = 0.001 weight_decay: float = 1e-4 @dataclass class DataConf: path: str = "" preprocessed: bool = False @dataclass class ResultsConf: val_loss: float = 0.0 final_lr: float = 0.0 num_lr_steps: int = 0 @dataclass class TrainConf: optimizer: OptimizerConf = field(default_factory=OptimizerConf) data: DataConf = field(default_factory=DataConf) results: ResultsConf = field(default_factory=ResultsConf) epochs: int = 10 batch_size: int = 32 experiment: str = "baseline" @env.task async def preprocess(cfg: DictConfig, dataset: str) -> DictConfig: """First stage: fills in the data section of cfg.""" return OmegaConf.merge(cfg, {"data": {"path": dataset, "preprocessed": True}}) @env.task async def build_schedule(cfg: DictConfig) -> ListConfig: """Produces an LR schedule from cfg as a ListConfig.""" lrs = [cfg.optimizer.lr * (0.5**i) for i in range(cfg.epochs)] return OmegaConf.create(lrs) @env.task async def train(cfg: DictConfig, lr_schedule: ListConfig) -> tuple[DictConfig, float]: """Simulates training. Returns the final cfg (with results filled in) and val loss.""" final_lr = float(lr_schedule[-1]) val_loss = final_lr * 10 # placeholder result_cfg = OmegaConf.merge( cfg, { "results": { "val_loss": val_loss, "final_lr": final_lr, "num_lr_steps": len(lr_schedule), } }, ) return result_cfg, val_loss @env.task async def evaluate(result_cfg: DictConfig, val_loss: float) -> str: """Final stage: formats a report from the result config.""" return ( f"experiment={result_cfg.experiment} " f"data={result_cfg.data.path} " f"val_loss={val_loss:.6f} " f"final_lr={result_cfg.results.final_lr:.6f} " f"lr_steps={result_cfg.results.num_lr_steps}" ) @env.task async def training_pipeline(dataset: str) -> str: """Full pipeline: cfg flows preprocess, build_schedule, train and evaluate.""" cfg = OmegaConf.structured( TrainConf( optimizer=OptimizerConf(lr=0.01, weight_decay=1e-5), epochs=5, batch_size=64, experiment="structured-cfg-pipeline", ) ) preprocessed_cfg = await preprocess(cfg, dataset=dataset) lr_schedule = await build_schedule(preprocessed_cfg) result_cfg, val_loss = await train(preprocessed_cfg, lr_schedule=lr_schedule) return await evaluate(result_cfg, val_loss=val_loss) if __name__ == "__main__": flyte.init_from_config() run = flyte.run(training_pipeline, dataset="s3://my-bucket/imagenet") print(f"Run URL: {run.url}") print(f"Outputs: {run.outputs()}") For more focused examples such as plain `DictConfig` patterns, advanced `ListConfig` shapes, all `MISSING`/`Enum`/`Path`/`bytes` cases, see the [plugin repository](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/omegaconf/examples) . LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # LLM-optimized documentation | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) LLM-optimized documentation =========================== This site provides LLM-optimized documentation at four levels of granularity, designed for use by AI coding agents such as [Claude Code](https://docs.anthropic.com/en/docs/claude-code) , [Cursor](https://www.cursor.com/) , [Windsurf](https://windsurf.com/) , and similar tools. These files also follow the [`llms.txt` convention](https://llmstxt.org/) , making them discoverable by AI search engines. Every page on the site also has an **LLM-optimized** section in the right-hand sidebar that points to: * This “LLM-optimized documentation” page (for explanation). * An LLM-optimized version of that page. * An LLM-optimized single file containing the whole section (only on top pages of key sections). * The full site index for LLMs. All links within LLM-optimized files use absolute URLs (`https://www.union.ai/docs/...`), so files work correctly when copied locally and used outside the docs site. [Per-page Markdown (`page.md`)](https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/#per-page-markdown-pagemd) --------------------------------------------------------------------------------------------------------------------------- Every page on this site has a parallel LLM-optimized version in clean Markdown, accessible at the same URL path with `/page.md` appended and via the “**This page**” link in the “**LLM-optimized**” section of the right sidebar. For example, this page is at: * [`https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/) and its LLM-optimized version is at: * [`https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/page.md`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/page) Section landing pages include a `## Subpages` table listing child pages with their H2/H3 headings, making it easy to identify the right page to fetch. [Section bundles (`section.md`)](https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/#section-bundles-sectionmd) ----------------------------------------------------------------------------------------------------------------------------- For key documentation sections, a curated bundle file concatenates all pages in the section into a single `section.md` file. These are accessible at the same URL path as the top page of the section, with `/section.md` appended and via the “**This section in one file**” link in the “**LLM-optimized**” section of the right sidebar. These `section.md` files are sized to fit within modern LLM context windows and are ideal for pasting into a prompt or adding to project context. Available bundle files: **User guide:** * [`Core concepts`](https://www.union.ai/docs/v2/flyte/user-guide/core-concepts/section.md) (~12K tokens) * [`Configure tasks`](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/section.md) (~25K tokens) * [`Build tasks`](https://www.union.ai/docs/v2/flyte/user-guide/task-programming/section.md) (~28K tokens) * [`Run and deploy tasks`](https://www.union.ai/docs/v2/flyte/user-guide/task-deployment/section.md) (~39K tokens) * [`Configure apps`](https://www.union.ai/docs/v2/flyte/user-guide/configure-apps/section.md) (~7K tokens) * [`Build apps`](https://www.union.ai/docs/v2/flyte/user-guide/build-apps/section.md) (~12K tokens) * [`Native app integrations`](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/section.md) (~6K tokens) * [`Serve and deploy apps`](https://www.union.ai/docs/v2/flyte/user-guide/serve-and-deploy-apps/section.md) (~4K tokens) * [`Build an agent`](https://www.union.ai/docs/v2/flyte/user-guide/build-agent/section.md) (~10K tokens) * [`Sandboxing`](https://www.union.ai/docs/v2/flyte/user-guide/sandboxing/section.md) (~11K tokens) * [`Scale your runs`](https://www.union.ai/docs/v2/flyte/user-guide/run-scaling/section.md) (~14K tokens) * [`Advanced project`](https://www.union.ai/docs/v2/flyte/user-guide/advanced-project/section.md) (~6K tokens) * [`From Flyte 1 to 2`](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/section.md) (~5K tokens) **Tutorials:** * [`Tutorials`](https://www.union.ai/docs/v2/flyte/tutorials/section.md) (~56K tokens) **Reference:** * [`Migration from Flyte 1`](https://www.union.ai/docs/v2/flyte/api-reference/migration/section.md) (~12K tokens) **Integrations:** * [`Integrations`](https://www.union.ai/docs/v2/flyte/integrations/section.md) (~53K tokens) **Community:** * [`Contributing docs and examples`](https://www.union.ai/docs/v2/flyte/community/contributing-docs/section.md) (~12K tokens) [Page index (`llms.txt`)](https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/#page-index-llmstxt) --------------------------------------------------------------------------------------------------------------- The `llms.txt` file is a compact index of all LLM-optimized pages, organized by section. Each page entry includes the H2/H3 headings found on that page, so an agent can identify the right page to fetch without downloading it first. Sections that have a `section.md` bundle are marked in the index. Download it and append its contents to the `AGENTS.md`, `CLAUDE.md` or similar file in your project root. Make sure you append the index into a file that is **loaded into context by default** by your coding tool. Adding it as a skill or tool is less effective because the agent must decide to load it rather than having the information always available. * [`llms.txt`](https://www.union.ai/docs/v2/flyte/llms.txt) (~32K tokens) You are viewing the **Flyte OSS** docs. To get the `llms.txt` for a different product variant, use the variant selector at the top of the page. [Full documentation (`llms-full.txt`)](https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/#full-documentation-llms-fulltxt) ----------------------------------------------------------------------------------------------------------------------------------------- The `llms-full.txt` file contains the entire Flyte version 2.0 documentation as a single Markdown file. This file is very large and is not suitable for direct inclusion in an LLM context window, but it may be useful for RAG-based tools. * [`llms-full.txt`](https://www.union.ai/docs/v2/flyte/llms-full.txt) (~1.4M tokens) You are viewing the **Flyte OSS** docs. To get the `llms-full.txt` for a different product variant, use the variant selector at the top of the page. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Integrations | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) Integrations ============ An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v2/flyte/integrations/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. Flyte 2 is designed to be extensible by default. While the core platform covers the most common orchestration needs, many production workloads require specialized infrastructure, external services or execution semantics that go beyond the core runtime. Flyte 2 exposes these capabilities through integrations. Under the hood, integrations are implemented using Flyte 2’s plugin system, which provides a consistent way to extend the platform without modifying core execution logic. An integration allows you to declaratively enable new capabilities such as distributed compute frameworks or third-party services without manually managing infrastructure. You specify what you need, and Flyte takes care of how it is provisioned, used and cleaned up. This page covers: * The types of integrations Flyte 2 supports today * How integrations fit into Flyte 2’s execution model * How to use integrations in your tasks * The integrations available out of the box If you need functionality that doesn’t exist yet, Flyte 2’s plugin system is intentionally open-ended. You can build and register your own integrations using the same architecture described here. [Integration categories](https://www.union.ai/docs/v2/flyte/integrations/#integration-categories) --------------------------------------------------------------------------------------------------- Flyte 2 integrations fall into the following categories: 1. **Distributed compute**: Provision transient compute clusters to run tasks across multiple nodes, with automatic lifecycle management. 2. **Agentic AI**: Support for various common aspects of agentic AI applications. 3. **Configuration**: Compose and pass hierarchical configuration objects between tasks, with type-safe schemas and CLI/YAML composition. 4. **Experiment tracking**: Integrate with experiment tracking platforms for logging metrics, parameters, and artifacts. 5. **Data validation**: Enforce schema contracts on dataframes flowing between tasks, with automatic validation reports. 6. **Connectors**: Stateless, long-running services that receive execution requests via gRPC and then submit work to external (or internal) systems. 7. **LLM Serving**: Deploy and serve large language models with an OpenAI-compatible API. 8. **Notebook execution**: Run parameterized Jupyter notebooks as typed Flyte tasks with cell-level reports. [Distributed compute](https://www.union.ai/docs/v2/flyte/integrations/#distributed-compute) --------------------------------------------------------------------------------------------- Distributed compute integrations allow tasks to run on dynamically provisioned clusters. These clusters are created just-in-time, scoped to the task execution and torn down automatically when the task completes. This enables large-scale parallelism without requiring users to operate or maintain long-running infrastructure. ### [Supported distributed compute integrations](https://www.union.ai/docs/v2/flyte/integrations/#supported-distributed-compute-integrations) | Plugin | Description | Common use cases | | --- | --- | --- | | [Ray](https://www.union.ai/docs/v2/flyte/integrations/ray) | Provisions Ray clusters via KubeRay | Distributed Python, ML training, hyperparameter tuning | | [Spark](https://www.union.ai/docs/v2/flyte/integrations/spark) | Provisions Spark clusters via Spark Operator | Large-scale data processing, ETL pipelines | | [Dask](https://www.union.ai/docs/v2/flyte/integrations/dask) | Provisions Dask clusters via Dask Operator | Parallel Python workloads, dataframe operations | | [PyTorch](https://www.union.ai/docs/v2/flyte/integrations/pytorch) | Distributed PyTorch training with elastic launch | Single-node and multi-node training | Each plugin encapsulates: * Cluster provisioning * Resource configuration * Networking and service discovery * Lifecycle management and teardown From the task author’s perspective, these details are abstracted away. ### [How the plugin system works](https://www.union.ai/docs/v2/flyte/integrations/#how-the-plugin-system-works) At a high level, Flyte 2’s distributed compute plugin architecture follows a simple and consistent pattern. #### [1\. Registration](https://www.union.ai/docs/v2/flyte/integrations/#1-registration) Each plugin registers itself with Flyte 2’s core plugin registry: * **`TaskPluginRegistry`**: The central registry for all distributed compute plugins * Each plugin declares: * Its configuration schema * How that configuration maps to execution behavior This registration step makes the plugin discoverable by the runtime. #### [2\. Task environments and plugin configuration](https://www.union.ai/docs/v2/flyte/integrations/#2-task-environments-and-plugin-configuration) Integrations are activated through a `TaskEnvironment`. A `TaskEnvironment` bundles: * A container image * Execution settings * A plugin configuration object enabled with `plugin_config` The plugin configuration describes _what_ infrastructure or integration the task requires. #### [3\. Automatic provisioning and execution](https://www.union.ai/docs/v2/flyte/integrations/#3-automatic-provisioning-and-execution) When a task associated with a `TaskEnvironment` runs: 1. Flyte inspects the environment’s plugin configuration 2. The plugin provisions the required infrastructure or integration 3. The task executes with access to that capability 4. Flyte cleans up all transient resources after completion ### [Example: Using the Dask plugin](https://www.union.ai/docs/v2/flyte/integrations/#example-using-the-dask-plugin) Below is a complete example showing how a task gains access to a Dask cluster simply by running inside an environment configured with the Dask plugin. from flyteplugins.dask import Dask, WorkerGroup import flyte # Define the Dask cluster configuration dask_config = Dask( workers=WorkerGroup(number_of_workers=4) ) # Create a task environment that enables Dask env = flyte.TaskEnvironment( name="dask_env", plugin_config=dask_config, image=image, ) # Any task in this environment has access to the Dask cluster @env.task async def process_data(data: list) -> list: from distributed import Client client = Client() # Automatically connects to the provisioned cluster futures = client.map(transform, data) return client.gather(futures) When `process_data` executes, Flyte performs the following steps: 1. Provisions a Dask cluster with 4 workers 2. Executes the task with network access to the cluster 3. Tears down the cluster once the task completes No cluster management logic appears in the task code. The task only expresses intent. ### [Key design principle](https://www.union.ai/docs/v2/flyte/integrations/#key-design-principle) All distributed compute integrations follow the same mental model: * You declare the required capability via configuration * You attach that configuration to a task environment * Tasks decorated with that environment automatically gain access to the capability This makes it easy to swap execution backends or introduce distributed compute incrementally without rewriting workflows. [Agentic AI](https://www.union.ai/docs/v2/flyte/integrations/#agentic-ai) --------------------------------------------------------------------------- Agentic AI integrations provide drop-in replacements for LLM provider SDKs. They let you use Flyte tasks as agent tools so that tool calls run with full Flyte observability, retries, and caching. ### [Supported agentic AI integrations](https://www.union.ai/docs/v2/flyte/integrations/#supported-agentic-ai-integrations) | Plugin | Description | Common use cases | | --- | --- | --- | | [OpenAI](https://www.union.ai/docs/v2/flyte/integrations/openai) | Drop-in replacement for OpenAI Agents SDK `function_tool` | Agentic workflows with OpenAI models | | [Anthropic](https://www.union.ai/docs/v2/flyte/integrations/anthropic) | Agent loop and `function_tool` for the Anthropic Claude SDK | Agentic workflows with Claude | | [Gemini](https://www.union.ai/docs/v2/flyte/integrations/gemini) | Agent loop and `function_tool` for the Google Gemini SDK | Agentic workflows with Gemini | | [Code generation](https://www.union.ai/docs/v2/flyte/integrations/codegen) | LLM-driven code generation with automatic testing in sandboxes | Data processing, ETL, analysis pipelines | [Experiment tracking](https://www.union.ai/docs/v2/flyte/integrations/#experiment-tracking) --------------------------------------------------------------------------------------------- Experiment tracking integrations let you log metrics, parameters, and artifacts to external tracking platforms during Flyte task execution. ### [Supported experiment tracking integrations](https://www.union.ai/docs/v2/flyte/integrations/#supported-experiment-tracking-integrations) | Plugin | Description | Common use cases | | --- | --- | --- | | [MLflow](https://www.union.ai/docs/v2/flyte/integrations/mlflow) | MLflow experiment tracking | Experiment tracking, autologging, model registry | | [Weights and Biases](https://www.union.ai/docs/v2/flyte/integrations/wandb) | Weights & Biases integration | Experiment tracking and hyperparameter tuning | [Configuration](https://www.union.ai/docs/v2/flyte/integrations/#configuration) --------------------------------------------------------------------------------- Configuration integrations let you compose and pass hierarchical configuration objects between Flyte tasks, with type-safe schemas and CLI/YAML composition. ### [Supported configuration integrations](https://www.union.ai/docs/v2/flyte/integrations/#supported-configuration-integrations) | Plugin | Description | Common use cases | | --- | --- | --- | | [OmegaConf](https://www.union.ai/docs/v2/flyte/integrations/omegaconf) | `DictConfig` / `ListConfig` as native task input and output types | Passing composed configs between tasks, structured configs, YAML-driven pipelines | | [Hydra](https://www.union.ai/docs/v2/flyte/integrations/hydra) | Hydra config composition and sweep submission for Flyte tasks | YAML-driven experiment composition, grid and Bayesian sweeps, hardware presets | [Data validation](https://www.union.ai/docs/v2/flyte/integrations/#data-validation) ------------------------------------------------------------------------------------- Data validation integrations enforce schema contracts on the dataframes flowing between tasks. They validate data at task boundaries, catch type and constraint violations early, and produce HTML reports visible in the Flyte UI. ### [Supported data validation integrations](https://www.union.ai/docs/v2/flyte/integrations/#supported-data-validation-integrations) | Plugin | Description | Common use cases | | --- | --- | --- | | [Pandera](https://www.union.ai/docs/v2/flyte/integrations/pandera) | Validates dataframes with pandera `DataFrameModel` schemas | Schema enforcement, data quality checks, validation reports | [Connectors](https://www.union.ai/docs/v2/flyte/integrations/#connectors) --------------------------------------------------------------------------- Connectors are stateless, long‑running services that receive execution requests via gRPC and then submit work to external (or internal) systems. Each connector runs as its own Kubernetes deployment, and is triggered when a Flyte task of the matching type is executed. Although they normally run inside the data plane, you can also run connectors locally as long as the required secrets/credentials are present locally. This is useful because connectors are just Python services that can be spawned in‑process. Connectors are designed to scale horizontally and reduce load on the core Flyte backend because they execute _outside_ the core system. This decoupling makes connectors efficient, resilient, and easy to iterate on. You can even test them locally without modifying backend configuration, which reduces friction during development. ### [Supported connectors](https://www.union.ai/docs/v2/flyte/integrations/#supported-connectors) | Connector | Description | Common use cases | | --- | --- | --- | | [Snowflake](https://www.union.ai/docs/v2/flyte/integrations/snowflake) | Run SQL queries on Snowflake asynchronously | Data warehousing, ETL, analytics queries | | [BigQuery](https://www.union.ai/docs/v2/flyte/integrations/bigquery) | Run SQL queries on Google BigQuery | Data warehousing, ETL, analytics queries | | [Databricks](https://www.union.ai/docs/v2/flyte/integrations/databricks) | Run PySpark jobs on Databricks clusters | Large-scale data processing, Spark ETL | ### [Creating a new connector](https://www.union.ai/docs/v2/flyte/integrations/#creating-a-new-connector) If none of the existing connectors meet your needs, you can build your own. Connectors communicate via Protobuf, so in theory they can be implemented in any language. Today, only **Python** connectors are supported. ### [Async connector interface](https://www.union.ai/docs/v2/flyte/integrations/#async-connector-interface) To implement a new async connector, extend `AsyncConnector` and implement the following methods, all of which must be idempotent: | Method | Purpose | | --- | --- | | `create` | Launch the external job (via REST, gRPC, SDK, or other API) | | `get` | Fetch current job state (return job status or output) | | `delete` | Delete / cancel the external job | | `get_logs` | Stream paginated log lines to the Flyte UI | To test the connector locally, the connector task should inherit from [AsyncConnectorExecutorMixin](https://github.com/flyteorg/flyte-sdk/blob/1d49299294cd5e15385fe8c48089b3454b7a4cd1/src/flyte/connectors/_connector.py#L206) . This mixin simulates how the Flyte 2 system executes asynchronous connector tasks, making it easier to validate your connector implementation before deploying it. ### [Example: Batch job connector](https://www.union.ai/docs/v2/flyte/integrations/#example-batch-job-connector) The following example implements a connector that simulates submitting and polling an external batch job. Replace the mock logic with real API calls for your use case. **Connector** (`my_connector/connector.py`): connector.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/batch_job/connector.py "View source on GitHub") import time import uuid from dataclasses import dataclass from typing import Any, Dict, Optional from flyteidl2.connector.connector_pb2 import ( GetTaskLogsResponse, GetTaskLogsResponseBody, GetTaskLogsResponseHeader, ) from flyteidl2.core.execution_pb2 import TaskExecution from flyteidl2.logs.dataplane.payload_pb2 import LogLine, LogLineOriginator from google.protobuf.timestamp_pb2 import Timestamp from flyte import logger from flyte.connectors import AsyncConnector, ConnectorRegistry, Resource, ResourceMeta @dataclass class BatchJobMetadata(ResourceMeta): job_id: str created_at: float class BatchJobConnector(AsyncConnector): name = "Batch Job Connector" task_type_name = "batch_job" metadata_type = BatchJobMetadata async def create(self, task_template, inputs: Optional[Dict[str, Any]] = None, **kwargs) -> BatchJobMetadata: job_id = str(uuid.uuid4())[:8] logger.info(f"Submitted batch job {job_id}") return BatchJobMetadata(job_id=job_id, created_at=time.time()) async def get(self, resource_meta: BatchJobMetadata, **kwargs) -> Resource: elapsed = time.time() - resource_meta.created_at if elapsed < 5: return Resource(phase=TaskExecution.RUNNING, message="Job in progress") return Resource( phase=TaskExecution.SUCCEEDED, message="Job completed", outputs={"result": f"output-from-{resource_meta.job_id}"}, ) async def delete(self, resource_meta: BatchJobMetadata, **kwargs): logger.info(f"Cancelled job {resource_meta.job_id}") async def get_logs(self, resource_meta: BatchJobMetadata, token: str = "", **kwargs): def line(message: str, ts: float) -> LogLine: t = Timestamp() t.FromSeconds(int(ts)) return LogLine(timestamp=t, message=message, originator=LogLineOriginator.USER) start = resource_meta.created_at job_id = resource_meta.job_id pages = { "": GetTaskLogsResponseBody(lines=[\ line(f"[INFO] Job {job_id} submitted", start),\ line(f"[INFO] Job {job_id} started", start + 1),\ ]), "page-2": GetTaskLogsResponseBody(lines=[\ line(f"[INFO] Job {job_id} finished", start + 5),\ ]), } next_tokens = {"": "page-2", "page-2": ""} yield GetTaskLogsResponse(body=pages.get(token, GetTaskLogsResponseBody(lines=[]))) next_token = next_tokens.get(token, "") if next_token: yield GetTaskLogsResponse(header=GetTaskLogsResponseHeader(token=next_token)) ConnectorRegistry.register(BatchJobConnector()) **Task plugin** (`my_connector/task.py`): task.py [](https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/batch_job/task.py "View source on GitHub") from dataclasses import dataclass from typing import Any, Dict, Optional, Type from flyte.connectors import AsyncConnectorExecutorMixin from flyte.extend import TaskTemplate from flyte.models import NativeInterface, SerializationContext @dataclass class BatchJobConfig: timeout_seconds: int = 300 class BatchJobTask(AsyncConnectorExecutorMixin, TaskTemplate): _TASK_TYPE = "batch_job" def __init__(self, name: str, plugin_config: BatchJobConfig, inputs: Optional[Dict[str, Type]] = None, outputs: Optional[Dict[str, Type]] = None, **kwargs): super().__init__( name=name, interface=NativeInterface( {k: (v, None) for k, v in inputs.items()} if inputs else {}, outputs or {}, ), task_type=self._TASK_TYPE, image=None, **kwargs, ) self.plugin_config = plugin_config def custom_config(self, sctx: SerializationContext) -> Optional[Dict[str, Any]]: return {"timeout_seconds": self.plugin_config.timeout_seconds} **Usage**: import flyte from my_connector.task import BatchJobConfig, BatchJobTask batch_job = BatchJobTask( name="my_batch_job", plugin_config=BatchJobConfig(timeout_seconds=60), inputs={"name": str}, outputs={"result": str}, ) flyte.TaskEnvironment.from_task("batch-job-env", batch_job) ### [Connector-level secrets](https://www.union.ai/docs/v2/flyte/integrations/#connector-level-secrets) If your connector needs credentials (API keys, tokens) shared across all tasks, pass them as environment variables into the connector process. Set environment variables on the connector Kubernetes deployment: kubectl set env deployment/ MY_API_KEY= -n Inside the connector, read the secret from the environment: import os api_key = os.environ["MY_API_KEY"] See [Secrets](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets) for how to store and manage secrets. ### [Deploy a custom connector](https://www.union.ai/docs/v2/flyte/integrations/#deploy-a-custom-connector) Deploying a connector requires two steps: building a Docker image that contains your connector code and then patching the connector Kubernetes deployment to use it. **Step 1: Build the connector image** import asyncio from flyte import Image from flyte.extend import ImageBuildEngine async def build_connector_image(registry: str, name: str, builder: str = "local"): image = Image.from_debian_base( registry=registry, name=name ).with_pip_packages("flyte[connector]", "my-connector-package") await ImageBuildEngine.build(image, builder=builder) if __name__ == "__main__": asyncio.run( build_connector_image( registry="", name="my-connector", builder="local" ) ) **Step 2: Override the connector deployment image** Once the image is pushed, patch the connector Kubernetes deployment to use it: kubectl set image deployment/ \ connector=/my-connector: \ -n Replace `` with the name of your connector deployment (e.g. `flyte-connector`), and `` with the namespace where Flyte is installed (typically `flyte`). [LLM Serving](https://www.union.ai/docs/v2/flyte/integrations/#llm-serving) ----------------------------------------------------------------------------- LLM serving integrations let you deploy and serve large language models as Flyte apps with an OpenAI-compatible API. They handle model loading, GPU management, and autoscaling. ### [Supported LLM serving integrations](https://www.union.ai/docs/v2/flyte/integrations/#supported-llm-serving-integrations) | Plugin | Description | Common use cases | | --- | --- | --- | | [SGLang](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/sglang-app) | Deploy models with SGLang’s high-throughput runtime | LLM inference, model serving | | [vLLM](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/vllm-app) | Deploy models with vLLM’s PagedAttention engine | LLM inference, model serving | For full setup instructions including multi-GPU deployment, model prefetching, and autoscaling, see the [SGLang app](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/sglang-app) and [vLLM app](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/vllm-app) pages. [Notebook execution](https://www.union.ai/docs/v2/flyte/integrations/#notebook-execution) ------------------------------------------------------------------------------------------- Notebook execution integrations let you run Jupyter notebooks as first-class Flyte tasks with typed inputs and outputs, HTML reports surfaced in the Flyte UI, and the ability to call other Flyte tasks from within the notebook. ### [Supported notebook execution integrations](https://www.union.ai/docs/v2/flyte/integrations/#supported-notebook-execution-integrations) | Plugin | Description | Common use cases | | --- | --- | --- | | [Papermill](https://www.union.ai/docs/v2/flyte/integrations/papermill) | Parameterize and execute `.ipynb` files via [papermill](https://papermill.readthedocs.io/) | Productionizing exploratory notebooks, cell-by-cell HTML reports, notebook-driven analysis pipelines | LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/integrations/page.md) [This section in one file](https://www.union.ai/docs/v2/flyte/integrations/section.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Flyte CLI | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) [Flyte 2 Devbox](https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox) is available today to run a full Flyte backend and UI locally. [Preview Flyte 2 for production](https://www.union.ai/try-flyte-2) , hosted on [Union.ai](https://www.union.ai/) 2.4.0 Flyte CLI ========= This is the command line interface for Flyte. | Object | Action | | --- | --- | | `action` | [`abort`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-abort-action)
, [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-action) | | `run` | [`abort`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-abort-run)
, [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-run) | | `config` | [`create`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-config)
, [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-config) | | `project` | [`create`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-project)
, [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-project)
, [`update`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-update-project) | | `secret` | [`create`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-secret)
, [`delete`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-secret)
, [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-secret) | | `trigger` | [`create`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-trigger)
, [`delete`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-trigger)
, [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-trigger)
, [`update`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-update-trigger) | | `app` | [`delete`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-app)
, [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-app)
, [`update`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-update-app) | | `devbox` | [`delete`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-devbox)
, [`start`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-start-devbox)
, [`stop`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-stop-devbox) | | `settings` | [`edit`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-edit-settings)
, [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-settings) | | `docs` | [`gen`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-gen-docs) | | `event` | [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-event)
, [`signal`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-signal-event) | | `io` | [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-io) | | `logs` | [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-logs) | | `task` | [`get`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-task) | | `hf-model` | [`prefetch`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-prefetch-hf-model) | | `deployed-task` | [`run`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-run-deployed-task) | | `tui` | [`start`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-start-tui) | | Action | On | | --- | --- | | `abort` | [`action`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-abort-action)
, [`run`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-abort-run) | | [`build`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-build) | \- | | `create` | [`config`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-config)
, [`project`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-project)
, [`secret`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-secret)
, [`trigger`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-trigger) | | `delete` | [`app`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-app)
, [`devbox`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-devbox)
, [`secret`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-secret)
, [`trigger`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-trigger) | | [`deploy`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-deploy) | \- | | `edit` | [`settings`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-edit-settings) | | `gen` | [`docs`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-gen-docs) | | `get` | [`action`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-action)
, [`app`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-app)
, [`config`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-config)
, [`event`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-event)
, [`io`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-io)
, [`logs`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-logs)
, [`project`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-project)
, [`run`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-run)
, [`secret`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-secret)
, [`settings`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-settings)
, [`task`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-task)
, [`trigger`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-trigger) | | `prefetch` | [`hf-model`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-prefetch-hf-model) | | `run` | [`deployed-task`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-run-deployed-task) | | [`serve`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-serve) | \- | | `signal` | [`event`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-signal-event) | | `start` | [`devbox`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-start-devbox)
, [`tui`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-start-tui) | | `stop` | [`devbox`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-stop-devbox) | | `update` | [`app`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-update-app)
, [`project`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-update-project)
, [`trigger`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-update-trigger) | | [`whoami`](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-whoami) | \- | [flyte](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte) ---------------------------------------------------------------------------- **`flyte [OPTIONS] COMMAND [ARGS]...`** The Flyte CLI is the command line interface for working with the Flyte SDK and backend. It follows a simple verb/noun structure, where the top-level commands are verbs that describe the action to be taken, and the subcommands are nouns that describe the object of the action. The root command can be used to configure the CLI for persistent settings, such as the endpoint, organization, and verbosity level. Set endpoint and organization: $ flyte --endpoint --org get project Increase verbosity level (This is useful for debugging, this will show more logs and exception traces): $ flyte -vvv get logs Override the default config file: $ flyte --config /path/to/config.yaml run ... * [Documentation](https://www.union.ai/docs/flyte/user-guide/) * [GitHub](https://github.com/flyteorg/flyte) : Please leave a star if you like Flyte! * [Slack](https://slack.flyte.org/) : Join the community and ask questions. * [Issues](https://github.com/flyteorg/flyte/issues) | Option | Type | Default | Description | | --- | --- | --- | --- | | `--version` | `boolean` | `False` | Show the version and exit. | | `--endpoint` | `text` | `Sentinel.UNSET` | The endpoint to connect to. This will override any configuration file and simply use `pkce` to connect. | | `--insecure` | `boolean` | | Use an insecure connection to the endpoint. If not specified, the CLI will use TLS. | | `--image-builder`
`--builder` | `choice` | | Image builder to use for building images. Overrides the config file setting. If not specified, the builder from the config file (image.builder) is used, falling back to ’local'. | | `--auth-type` | `choice` | | Authentication type to use for the Flyte backend. Defaults to ‘pkce’. | | `-v`
`--verbose` | `integer` | `0` | Show verbose messages and exception traces. Repeating multiple times increases the verbosity (e.g., -vvv). | | `--org` | `text` | `Sentinel.UNSET` | The organization to which the command applies. | | `-c`
`--config` | `file` | `Sentinel.UNSET` | Path to the configuration file to use. If not specified, the default configuration file is used. | | `--output-format`
`-of` | `choice` | `table` | Output format for commands that support it. Defaults to ’table'. | | `--log-format` | `choice` | `console` | Formatting for logs, defaults to ‘console’ which is meant to be human readable. ‘json’ is meant for machine parsing. | | `--user-log-level` | `choice` | `info` | Log level for user task logs. Independent of the internal Flyte log level (-v). | | `--reset-root-logger` | `boolean` | `False` | If set, the root logger will be reset to use Flyte logging style | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte abort](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-abort) **`flyte abort COMMAND [ARGS]...`** Abort an ongoing process. #### [flyte abort action](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-abort-action) **`flyte abort action [OPTIONS] RUN_NAME ACTION_NAME`** Abort an action associated with a run. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--reason` | `text` | `Manually aborted from the CLI` | The reason to abort the run. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte abort run](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-abort-run) **`flyte abort run [OPTIONS] RUN_NAME`** Abort a run. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--reason` | `text` | `Manually aborted from the CLI` | The reason to abort the run. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte build](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-build) **`flyte build [OPTIONS] COMMAND [ARGS]...`** Build the environments defined in a python file or directory. This will build the images associated with the environments. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--copy-style` | `choice` | `loaded_modules` | Copy style of the eventual deploy. Must match the deploy’s –copy-style so the image content hash — and therefore the registry tag — lines up. | | `--root-dir` | `text` | `Sentinel.UNSET` | Override the root source directory, helpful when working with monorepos. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte create](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create) **`flyte create COMMAND [ARGS]...`** Create resources in a Flyte deployment. #### [flyte create config](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-config) **`flyte create config [OPTIONS]`** Creates a configuration file for Flyte CLI. If the `--output` option is not specified, it will create a file named `config.yaml` in the current directory. If the file already exists, it will raise an error unless the `--force` option is used. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--endpoint` | `text` | `Sentinel.UNSET` | Endpoint of the Flyte backend. | | `--insecure` | `boolean` | `False` | Use an insecure connection to the Flyte backend. | | `--org` | `text` | `Sentinel.UNSET` | Organization to use. This will override the organization in the configuration file. | | `-o`
`--output` | `path` | `.flyte/config.yaml` | Path to the output directory where the configuration will be saved. Defaults to current directory. | | `--force` | `boolean` | `False` | Force overwrite of the configuration file if it already exists. | | `--image-builder`
`--builder` | `choice` | `local` | Image builder to use for building images. Defaults to ’local’. | | `--auth-type` | `choice` | | Authentication type to use for the Flyte backend. Defaults to ‘pkce’. | | `--local-persistence` | `boolean` | `False` | Enable SQLite persistence for local run metadata, allowing past runs to be browsed via ‘flyte start tui’. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte create project](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-project) **`flyte create project [OPTIONS]`** Create a new project.  Example usage: flyte create project --id my_project_id --name "My Project" flyte create project --id my_project_id --name "My Project" --description "My project" -l team=ml -l env=prod | Option | Type | Default | Description | | --- | --- | --- | --- | | `--id` | `text` | `Sentinel.UNSET` | Unique identifier for the project (immutable). | | `--name` | `text` | `Sentinel.UNSET` | Display name for the project. | | `--description` | `text` | \`\` | Description for the project. | | `--label`
`-l` | `text` | `Sentinel.UNSET` | Labels as key=value pairs. Can be specified multiple times. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte create secret](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-secret) **`flyte create secret [OPTIONS] NAME`** Create a new secret. The name of the secret is required. For example: $ flyte create secret my_secret --value my_value If you don’t provide a `--value` flag, you will be prompted to enter the secret value in the terminal. $ flyte create secret my_secret Enter secret value: If `--from-file` is specified, the value will be read from the file instead of being provided directly: $ flyte create secret my_secret --from-file /path/to/secret_file The `--type` option can be used to create specific types of secrets. Either `regular` or `image_pull` can be specified. Secrets intended to access container images should be specified as `image_pull`. Other secrets should be specified as `regular`. If no type is specified, `regular` is assumed. For image pull secrets, you have several options: 1. Interactive mode (prompts for registry, username, password): $ flyte create secret my_secret --type image_pull 2. With explicit credentials: $ flyte create secret my_secret --type image_pull --registry ghcr.io --username myuser 3. Lastly, you can create a secret from your existing Docker installation (i.e., you’ve run `docker login` in the past) and you just want to pull from those credentials. Since you may have logged in to multiple registries, you can specify which registries to include. If no registries are specified, all registries are added. $ flyte create secret my_secret --type image_pull --from-docker-config --registries ghcr.io,docker.io | Option | Type | Default | Description | | --- | --- | --- | --- | | `--value` | `text` | `Sentinel.UNSET` | Secret value Mutually exclusive with from\_file, from\_docker\_config, registry. | | `--from-file` | `path` | `Sentinel.UNSET` | Path to the file with the binary secret. Mutually exclusive with value, from\_docker\_config, registry. | | `--type` | `choice` | `regular` | Type of the secret. | | `--from-docker-config` | `boolean` | `False` | Create image pull secret from Docker config file (only for –type image\_pull). Mutually exclusive with value, from\_file, registry, username, password. | | `--docker-config-path` | `path` | `Sentinel.UNSET` | Path to Docker config file (defaults to ~/.docker/config.json or $DOCKER\_CONFIG). Requires from\_docker\_config. | | `--registries` | `text` | `Sentinel.UNSET` | Comma-separated list of registries to include (only with –from-docker-config). | | `--registry` | `text` | `Sentinel.UNSET` | Registry hostname (e.g., ghcr.io, docker.io) for explicit credentials (only for –type image\_pull). Mutually exclusive with value, from\_file, from\_docker\_config. | | `--username` | `text` | `Sentinel.UNSET` | Username for the registry (only with –registry). | | `--password` | `text` | `Sentinel.UNSET` | Password for the registry (only with –registry). If not provided, will prompt. | | `--cluster-pool` | `text` | | Scope the secret to a cluster pool. Mutually exclusive with –project and –domain. Mutually exclusive with project, domain. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte create trigger](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-create-trigger) **`flyte create trigger [OPTIONS] TASK_NAME NAME`** Create a new trigger for a task. The task name and trigger name are required. Example: $ flyte create trigger my_task my_trigger --schedule "0 0 * * *" This will create a trigger that runs every day at midnight. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--schedule` | `text` | `Sentinel.UNSET` | Cron schedule for the trigger. Defaults to every minute. | | `--description` | `text` | \`\` | Description of the trigger. | | `--auto-activate` | `boolean` | `True` | Whether the trigger should not be automatically activated. Defaults to True. | | `--trigger-time-var` | `text` | `trigger_time` | Variable name for the trigger time in the task inputs. Defaults to ’trigger\_time’. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte delete](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete) **`flyte delete COMMAND [ARGS]...`** Remove resources from a Flyte deployment. #### [flyte delete app](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-app) **`flyte delete app [OPTIONS] NAME`** Delete apps from a Flyte deployment. | Option | Type | Default | Description | | --- | --- | --- | --- | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte delete devbox](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-devbox) **`flyte delete devbox [OPTIONS]`** Stop and remove the local Flyte devbox cluster container. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--volume` | `boolean` | `False` | Also delete the Docker volume used for persistent storage. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte delete secret](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-secret) **`flyte delete secret [OPTIONS] NAME`** Delete a secret. The name of the secret is required. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--cluster-pool` | `text` | | Scope the secret to a cluster pool. Mutually exclusive with –project and –domain. Mutually exclusive with project, domain. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte delete trigger](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-delete-trigger) **`flyte delete trigger [OPTIONS] NAME TASK_NAME`** Delete a trigger. The name of the trigger is required. | Option | Type | Default | Description | | --- | --- | --- | --- | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte deploy](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-deploy) **`flyte deploy [OPTIONS] COMMAND [ARGS]...`** Deploy one or more environments from a python file. This command will create or update environments in the Flyte system, registering all tasks and their dependencies. Example usage: flyte deploy hello.py my_env Arguments to the deploy command are provided right after the `deploy` command and before the file name. To deploy all environments in a file, use the `--all` flag: flyte deploy --all hello.py To recursively deploy all environments in a directory and its subdirectories, use the `--recursive` flag: flyte deploy --recursive ./src You can combine `--all` and `--recursive` to deploy everything: flyte deploy --all --recursive ./src You can provide image mappings with `--image` flag. This allows you to specify the image URI for the task environment during CLI execution without changing the code. Any images defined with `Image.from_ref_name("name")` will resolve to the corresponding URIs you specify here. flyte deploy --image my_image=ghcr.io/myorg/my-image:v1.0 hello.py my_env If the image name is not provided, it is regarded as a default image and will be used when no image is specified in TaskEnvironment: flyte deploy --image ghcr.io/myorg/default-image:latest hello.py my_env You can specify multiple image arguments: flyte deploy --image ghcr.io/org/default:latest --image gpu=ghcr.io/org/gpu:v2.0 hello.py my_env To deploy a specific version, use the `--version` flag: flyte deploy --version v1.0.0 hello.py my_env To preview what would be deployed without actually deploying, use the `--dry-run` flag: flyte deploy --dry-run hello.py my_env You can specify the `--config` flag to point to a specific Flyte cluster: flyte --config my-config.yaml deploy hello.py my_env You can override the default configured project and domain: flyte deploy --project my-project --domain development hello.py my_env If loading some files fails during recursive deployment, you can use the `--ignore-load-errors` flag to continue deploying the environments that loaded successfully: flyte deploy --recursive --ignore-load-errors ./src Other arguments to the deploy command are listed below. To see the environments available in a file, use `--help` after the file name: flyte deploy hello.py --help | Option | Type | Default | Description | | --- | --- | --- | --- | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--version` | `text` | `Sentinel.UNSET` | Version of the environment to deploy | | `--dry-run`
`--dryrun` | `boolean` | `False` | Dry run. Do not actually call the backend service. | | `--copy-style` | `choice` | `loaded_modules` | Copy style to use when running the task | | `--root-dir` | `text` | `Sentinel.UNSET` | Override the root source directory, helpful when working with monorepos. | | `--recursive`
`-r` | `boolean` | `False` | Recursively deploy all environments in the current directory | | `--all` | `boolean` | `False` | Deploy all environments in the current directory, ignoring the file name | | `--ignore-load-errors`
`-i` | `boolean` | `False` | Ignore errors when loading environments especially when using –recursive or –all. | | `--no-sync-local-sys-paths` | `boolean` | `False` | Disable synchronization of local sys.path entries under the root directory to the remote container. | | `--image` | `text` | `Sentinel.UNSET` | Image to be used in the run. Format: imagename=imageuri. Can be specified multiple times. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte edit](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-edit) **`flyte edit COMMAND [ARGS]...`** #### [flyte edit settings](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-edit-settings) **`flyte edit settings [OPTIONS]`** Edit hierarchical settings interactively — or apply a YAML file directly. **Interactive mode** (default). Opens settings in your ``$EDITOR``. Three comment tiers appear: - ``###`` section headers and the scope line - ``##`` per-field descriptions and inline metadata - ``#`` inactive settings (uncomment the single ``#`` to activate) If the edited YAML fails to parse, the editor reopens with an error header so you can fix the syntax without losing your edits. If you decline to reopen — or if the server rejects the update — your buffer is saved under ``~/.flyte/settings-edit-.yaml``. **Non-interactive mode**: pass ``--from-file `` to skip the editor entirely. The file's contents are parsed, the diff is printed, and the overrides are applied without a confirmation prompt. Ideal for CI/automation. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--from-file`
`-f` | `file` | | Apply overrides from a YAML file and skip the editor. The file can be produced by `flyte get settings` (comment markers are honoured) or be a plain YAML mapping of flat dot-notation keys to values. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte gen](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-gen) **`flyte gen COMMAND [ARGS]...`** Generate documentation. #### [flyte gen docs](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-gen-docs) **`flyte gen docs [OPTIONS]`** Generate documentation. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--type` | `text` | `Sentinel.UNSET` | Type of documentation (valid: markdown) | | `--plugin-variants` | `text` | | Hugo variant names for plugin commands (e.g., ‘union’). When set, plugin command sections and index entries are wrapped in {{< variant >}} shortcodes. Core commands appear unconditionally. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte get](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get) **`flyte get COMMAND [ARGS]...`** Retrieve resources from a Flyte deployment. You can get information about projects, runs, tasks, actions, secrets, logs and input/output values. Each command supports optional parameters to filter or specify the resource you want to retrieve. Using a `get` subcommand without any arguments will retrieve a list of available resources to get. For example: * `get project` (without specifying a project), will list all projects. * `get project my_project` will return the details of the project named `my_project`. In some cases, a partially specified command will act as a filter and return available further parameters. For example: * `get action my_run` will return all actions for the run named `my_run`. * `get action my_run my_action` will return the details of the action named `my_action` for the run `my_run`. #### [flyte get action](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-action) **`flyte get action [OPTIONS] RUN_NAME [ACTION_NAME]`** Get all actions for a run or details for a specific action. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--in-phase` | `choice` | `Sentinel.UNSET` | Filter actions by their phase. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte get app](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-app) **`flyte get app [OPTIONS] [NAME]`** Get a list of all apps, or details of a specific app by name. Apps are long-running services deployed on the Flyte platform. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--limit` | `integer` | `100` | Limit the number of apps to fetch when listing. | | `--only-mine` | `boolean` | `False` | Show only apps created by the current user (you). | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte get config](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-config) **`flyte get config`** Shows the automatically detected configuration to connect with the remote backend. The configuration will include the endpoint, organization, and other settings that are used by the CLI. #### [flyte get event](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-event) **`flyte get event [OPTIONS] RUN_NAME [ACTION_NAME]`** List events (paused condition actions) for a run, optionally filtered to a specific parent action. Each event corresponds to a condition action registered via `flyte.new_event(...)` from a workflow. Use `flyte signal event` to resolve one. | Option | Type | Default | Description | | --- | --- | --- | --- | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte get io](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-io) **`flyte get io [OPTIONS] RUN_NAME [ACTION_NAME]`** Get the inputs and outputs of a run or action. If only the run name is provided, it will show the inputs and outputs of the root action of that run. If an action name is provided, it will show the inputs and outputs for that action. If `--inputs-only` or `--outputs-only` is specified, it will only show the inputs or outputs respectively. Examples: $ flyte get io my_run $ flyte get io my_run my_action | Option | Type | Default | Description | | --- | --- | --- | --- | | `--inputs-only`
`-i` | `boolean` | `False` | Show only inputs | | `--outputs-only`
`-o` | `boolean` | `False` | Show only outputs | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte get logs](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-logs) **`flyte get logs [OPTIONS] RUN_NAME [ACTION_NAME]`** Stream logs for the provided run or action. If only the run is provided, only the logs for the parent action will be streamed: $ flyte get logs my_run If you want to see the logs for a specific action, you can provide the action name as well: $ flyte get logs my_run my_action By default, logs will be shown in the raw format and will scroll the terminal. If automatic scrolling and only tailing `--lines` number of lines is desired, use the `--pretty` flag: $ flyte get logs my_run my_action --pretty --lines 50 | Option | Type | Default | Description | | --- | --- | --- | --- | | `--lines`
`-l` | `integer` | `30` | Number of lines to show, only useful for –pretty | | `--show-ts` | `boolean` | `False` | Show timestamps | | `--pretty` | `boolean` | `False` | Show logs in an auto-scrolling box, where number of lines is limited to `--lines` | | `--attempt`
`-a` | `integer` | | Attempt number to show logs for, defaults to the latest attempt. | | `--filter-system` | `boolean` | `False` | Filter all system logs from the output. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte get project](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-project) **`flyte get project [OPTIONS] [NAME]`** Get a list of all projects, or details of a specific project by name. By default, only active (unarchived) projects are shown. Use `--archived` to show archived projects instead. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--archived` | `boolean` | `False` | Show archived projects instead of active ones. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte get run](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-run) **`flyte get run [OPTIONS] [NAME]`** Get a list of all runs, or details of a specific run by name. The run details will include information about the run, its status, but only the root action will be shown. If you want to see the actions for a run, use `get action `. You can filter runs by task name and optionally task version: $ flyte get run --task-name my_task $ flyte get run --task-name my_task --task-version v1.0 | Option | Type | Default | Description | | --- | --- | --- | --- | | `--limit` | `integer` | `100` | Limit the number of runs to fetch when listing. | | `--in-phase` | `choice` | `Sentinel.UNSET` | Filter runs by their status. | | `--only-mine` | `boolean` | `False` | Show only runs created by the current user (you). | | `--task-name` | `text` | | Filter runs by task name. | | `--task-version` | `text` | | Filter runs by task version. | | `--created-after` | `datetime` | | Show runs created at or after this datetime (UTC). Accepts ISO dates, ’now’, ’today’, or ’now - 1 day’. | | `--created-before` | `datetime` | | Show runs created before this datetime (UTC). | | `--updated-after` | `datetime` | | Show runs updated at or after this datetime (UTC). Accepts ISO dates, ’now’, ’today’, or ’now - 1 day'. | | `--updated-before` | `datetime` | | Show runs updated before this datetime (UTC). | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte get secret](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-secret) **`flyte get secret [OPTIONS] [NAME]`** Get a list of all secrets, or details of a specific secret by name. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--cluster-pool` | `text` | | Scope the secret to a cluster pool. Mutually exclusive with –project and –domain. Mutually exclusive with project, domain. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte get settings](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-settings) **`flyte get settings [OPTIONS]`** Get settings for a scope as editable YAML. Renders three sections:  * Local overrides — uncommented, applied at this scope. * Inherited settings — commented, with the scope they come from. * Available settings — commented placeholders for every key that isn’t set anywhere yet, so you can see what can be configured.  Examples: # Get ORG-level settings flyte get settings # Get settings for a domain flyte get settings --domain production # Get settings for a project (inherits from domain, which inherits from org) flyte get settings --domain production --project ml-pipeline # Dump to a file, edit it, then apply non-interactively flyte get settings --domain production -o prod.yaml # ...edit prod.yaml... flyte edit settings --domain production --from-file prod.yaml Use `flyte edit settings` to interactively modify these values. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--to-file`
`-o` | `file` | | Write the scope’s YAML to this file instead of printing it. The file round-trips through `flyte edit settings --from-file`. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte get task](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-task) **`flyte get task [OPTIONS] [NAME] [VERSION]`** Retrieve a list of all tasks, or details of a specific task by name and version. Currently, both `name` and `version` are required to get a specific task. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--limit` | `integer` | `100` | Limit the number of tasks to fetch. | | `--entrypoint` | `boolean` | `False` | Show only entrypoint tasks. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte get trigger](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-get-trigger) **`flyte get trigger [OPTIONS] [TASK_NAME] [NAME]`** Get a list of all triggers, or details of a specific trigger by name. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--limit` | `integer` | `100` | Limit the number of triggers to fetch. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte prefetch](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-prefetch) **`flyte prefetch COMMAND [ARGS]...`** Prefetch artifacts from remote registries. These commands help you download and prefetch artifacts like HuggingFace models to your Flyte storage for faster access during task execution. #### [flyte prefetch hf-model](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-prefetch-hf-model) **`flyte prefetch hf-model [OPTIONS] REPO`** Prefetch a HuggingFace model to Flyte storage. Downloads a model from the HuggingFace Hub and prefetches it to your configured Flyte storage backend. This is useful for: * Pre-fetching large models before running inference tasks * Sharding models for tensor-parallel inference * Avoiding repeated downloads during development **Basic Usage:** $ flyte prefetch hf-model meta-llama/Llama-2-7b-hf --hf-token-key HF_TOKEN **With Sharding:** Create a shard config file (shard\_config.yaml): engine: vllm args: tensor_parallel_size: 8 dtype: auto trust_remote_code: true Then run: $ flyte prefetch hf-model meta-llama/Llama-2-70b-hf \ --shard-config shard_config.yaml \ --accelerator A100:8 \ --hf-token-key HF_TOKEN **Wait for Completion:** $ flyte prefetch hf-model meta-llama/Llama-2-7b-hf --wait | Option | Type | Default | Description | | --- | --- | --- | --- | | `--raw-data-path` | `text` | | Object store path to store the model. If not provided, the model will be stored using the default path generated by Flyte storage layer. | | `--artifact-name` | `text` | | Artifact name to use for the stored model. Must only contain alphanumeric characters, underscores, and hyphens. If not provided, the repo name will be used (replacing ‘.’ with ‘-’). | | `--architecture` | `text` | `Sentinel.UNSET` | Model architecture, as given in HuggingFace config.json. | | `--task` | `text` | `auto` | Model task, e.g., ‘generate’, ‘classify’, ’embed’, ‘score’, etc. Refer to vLLM docs. ‘auto’ will try to discover this automatically. | | `--modality` | `text` | `('text',)` | Modalities supported by the model, e.g., ’text’, ‘image’, ‘audio’, ‘video’. Can be specified multiple times. | | `--format` | `text` | `Sentinel.UNSET` | Model serialization format, e.g., safetensors, onnx, torchscript, joblib, etc. | | `--model-type` | `text` | `Sentinel.UNSET` | Model type, e.g., ’transformer’, ‘xgboost’, ‘custom’, etc. For HuggingFace models, this is auto-determined from config.json\[‘model\_type’\]. | | `--short-description` | `text` | `Sentinel.UNSET` | Short description of the model. | | `--force` | `integer` | `0` | Force store of the model. Increment value (–force=1, –force=2, …) to force a new store. | | `--wait` | `boolean` | `False` | Wait for the model to be stored before returning. | | `--hf-token-key` | `text` | `HF_TOKEN` | Name of the Flyte secret containing your HuggingFace token. Note: This is not the HuggingFace token itself, but the name of the secret in the Flyte secret store. | | `--cpu` | `text` | `2` | CPU request for the prefetch task (e.g., ‘2’, ‘4’, ‘2,4’ for 2-4 CPUs). | | `--mem` | `text` | `8Gi` | Memory request for the prefetch task (e.g., ‘16Gi’, ‘64Gi’, ‘16Gi,64Gi’ for 16-64GB). | | `--gpu` | `choice` | | The gpu to use for downloading and (optionally) sharding the model. Format: ‘{type}:{quantity}’ (e.g., ‘A100:8’, ‘L4:1’). | | `--disk` | `text` | `50Gi` | Disk storage request for the prefetch task (e.g., ‘100Gi’, ‘500Gi’). | | `--shm` | `text` | | Shared memory request for the prefetch task (e.g., ‘100Gi’, ‘auto’). | | `--shard-config` | `path` | `Sentinel.UNSET` | Path to a YAML file containing sharding configuration. The file should have ’engine’ (currently only ‘vllm’) and ‘args’ keys. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte run](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-run) **`flyte run [OPTIONS] COMMAND [ARGS]...`** Run a task from a python file or deployed task. Example usage: flyte run hello.py my_task --arg1 value1 --arg2 value2 Arguments to the run command are provided right after the `run` command and before the file name. Arguments for the task itself are provided after the task name. To run a task locally, use the `--local` flag. This will run the task in the local environment instead of the remote Flyte environment: flyte run --local hello.py my_task --arg1 value1 --arg2 value2 You can provide image mappings with `--image` flag. This allows you to specify the image URI for the task environment during CLI execution without changing the code. Any images defined with `Image.from_ref_name("name")` will resolve to the corresponding URIs you specify here. flyte run --image my_image=ghcr.io/myorg/my-image:v1.0 hello.py my_task If the image name is not provided, it is regarded as a default image and will be used when no image is specified in TaskEnvironment: flyte run --image ghcr.io/myorg/default-image:latest hello.py my_task You can specify multiple image arguments: flyte run --image ghcr.io/org/default:latest --image gpu=ghcr.io/org/gpu:v2.0 hello.py my_task To run tasks that you’ve already deployed to Flyte, use the deployed-task command: flyte run deployed-task my_env.my_task --arg1 value1 --arg2 value2 To run a specific version of a deployed task, use the `env.task:version` syntax: flyte run deployed-task my_env.my_task:xyz123 --arg1 value1 --arg2 value2 You can specify the `--config` flag to point to a specific Flyte cluster: flyte run --config my-config.yaml deployed-task ... You can override the default configured project and domain: flyte run --project my-project --domain development hello.py my_task You can discover what deployed tasks are available by running: flyte run deployed-task To run an arbitrary Python script on a remote cluster (without defining a task), use `python-script`: flyte run python-script script.py --gpu 1 --gpu-type A100 --memory 64Gi You can also install extra packages and wait for completion: flyte run --follow python-script train.py --packages torch,transformers Other arguments to the run command are listed below. Arguments for the task itself are provided after the task name and can be retrieved using `--help`. For example: flyte run hello.py my_task --help | Option | Type | Default | Description | | --- | --- | --- | --- | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--local` | `boolean` | `False` | Run the task locally | | `--copy-style` | `choice` | `loaded_modules` | Copy style to use when running the task | | `--root-dir` | `text` | `Sentinel.UNSET` | Override the root source directory, helpful when working with monorepos. | | `--raw-data-path` | `text` | `Sentinel.UNSET` | Override the output prefix used to store offloaded data types. e.g. s3://bucket/ | | `--service-account` | `text` | `Sentinel.UNSET` | Kubernetes service account. If not provided, the configured default will be used | | `--name` | `text` | `Sentinel.UNSET` | Name of the run. If not provided, a random name will be generated. | | `--follow`
`-f` | `boolean` | `False` | Wait and watch logs for the parent action. If not provided, the CLI will exit after successfully launching a remote execution with a link to the UI. | | `--tui` | `boolean` | `False` | Show interactive TUI for local execution (requires flyte\[tui\]). | | `--image` | `text` | `Sentinel.UNSET` | Image to be used in the run. Format: imagename=imageuri. Can be specified multiple times. | | `--no-sync-local-sys-paths` | `boolean` | `False` | Disable synchronization of local sys.path entries under the root directory to the remote container. | | `--run-project` | `text` | | Run the remote task in this project, only applicable when using `deployed-task` subcommand. | | `--run-domain` | `text` | | Run the remote task in this domain, only applicable when using `deployed-task` subcommand. | | `--debug` | `boolean` | `False` | Run the task as a VSCode debug task. Starts a code-server in the container so you can connect via the UI to interactively debug/run the task. | | `--env`
`-e` | `text` | `Sentinel.UNSET` | Environment variable to set on the run context. Format: KEY=VALUE. Can be specified multiple times, e.g. `-e LOG_LEVEL=debug -e FOO=bar`. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte run deployed-task](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-run-deployed-task) **`flyte run deployed-task [OPTIONS] COMMAND [ARGS]...`** Run remote task from the Flyte backend | Option | Type | Default | Description | | --- | --- | --- | --- | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte serve](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-serve) **`flyte serve [OPTIONS] COMMAND [ARGS]...`** Serve an app from a Python file using flyte.serve(). This command allows you to serve apps defined with `flyte.app.AppEnvironment` in your Python files. The serve command will deploy the app to the Flyte backend and start it, making it accessible via a URL. Example usage: flyte serve examples/apps/basic_app.py app_env **Local serving:** Use the `--local` flag to serve the app on localhost without deploying to the Flyte backend. This is useful for local development and testing: flyte serve --local examples/apps/single_script_fastapi.py env Arguments to the serve command are provided right after the `serve` command and before the file name. To follow the logs of the served app, use the `--follow` flag: flyte serve --follow examples/apps/basic_app.py app_env Note: Log streaming is not yet fully implemented and will be added in a future release. You can provide image mappings with `--image` flag. This allows you to specify the image URI for the app environment during CLI execution without changing the code. Any images defined with `Image.from_ref_name("name")` will resolve to the corresponding URIs you specify here. flyte serve --image my_image=ghcr.io/myorg/my-image:v1.0 examples/apps/basic_app.py app_env If the image name is not provided, it is regarded as a default image and will be used when no image is specified in AppEnvironment: flyte serve --image ghcr.io/myorg/default-image:latest examples/apps/basic_app.py app_env You can specify multiple image arguments: flyte serve --image ghcr.io/org/default:latest --image gpu=ghcr.io/org/gpu:v2.0 examples/apps/basic_app.py app_env You can specify the `--config` flag to point to a specific Flyte cluster: flyte serve --config my-config.yaml examples/apps/basic_app.py app_env You can override the default configured project and domain: flyte serve --project my-project --domain development examples/apps/basic_app.py app_env Other arguments to the serve command are listed below. Note: This pattern is primarily useful for serving apps defined in tasks. Serving deployed apps is not currently supported through this CLI command. | Option | Type | Default | Description | | --- | --- | --- | --- | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--copy-style` | `choice` | `loaded_modules` | Copy style to use when serving the app | | `--root-dir` | `text` | `Sentinel.UNSET` | Override the root source directory, helpful when working with monorepos. | | `--service-account` | `text` | `Sentinel.UNSET` | Kubernetes service account. If not provided, the configured default will be used | | `--name` | `text` | `Sentinel.UNSET` | Name of the app deployment. If not provided, the app environment name will be used. | | `--follow`
`-f` | `boolean` | `False` | Wait and watch logs for the app. If not provided, the CLI will exit after successfully deploying the app with a link to the UI. | | `--image` | `text` | `Sentinel.UNSET` | Image to be used in the serve. Format: imagename=imageuri. Can be specified multiple times. | | `--no-sync-local-sys-paths` | `boolean` | `False` | Disable synchronization of local sys.path entries under the root directory to the remote container. | | `--env-var`
`-e` | `text` | `Sentinel.UNSET` | Environment variable to set in the app. Format: KEY=VALUE. Can be specified multiple times. Example: –env-var LOG\_LEVEL=DEBUG –env-var DATABASE\_URL=postgresql://… | | `--local` | `boolean` | `False` | Serve the app locally on localhost instead of deploying to the Flyte backend. The app will be served on the port defined in the AppEnvironment. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte signal](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-signal) **`flyte signal COMMAND [ARGS]...`** Signal an event waiting on a paused condition action. #### [flyte signal event](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-signal-event) **`flyte signal event [OPTIONS] RUN_NAME ACTION_NAME [VALUE]`** Signal a paused condition action. The condition’s declared payload type and prompt are read from the backend. If VALUE is omitted the condition’s prompt is displayed and a typed interactive prompt is shown to collect the payload. When VALUE is provided it’s coerced to the expected type (`true`/`false` for bool, integer literals for int, decimal literals for float, any string for str). | Option | Type | Default | Description | | --- | --- | --- | --- | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte start](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-start) **`flyte start COMMAND [ARGS]...`** Start various Flyte services. #### [flyte start devbox](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-start-devbox) **`flyte start devbox [OPTIONS]`** Start a local Flyte devbox cluster. | Option | Type | Default | Description | | --- | --- | --- | --- | | `--image` | `text` | | Docker image to use for the devbox cluster. | | `--dev` | `boolean` | `False` | Enable dev mode inside the devbox cluster (sets FLYTE\_DEV=True). | | `--gpu` | `boolean` | `False` | Pass host GPUs into the devbox container (adds –gpus all to docker run). Requires an NVIDIA-enabled host. Defaults –image to a GPU-capable image if –image is not explicitly set. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte start tui](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-start-tui) **`flyte start tui`** Launch TUI explore mode to browse past local runs. To use the TUI install `pip install flyte[tui]` TUI, allows you to explore all your local runs if you have persistence enabled. Persistence can be enabled in 2 ways, 1. By setting it in the config to record every local run flyte create config --endpoint ... --local-persistence 2. By passing it in flyte.init(local\_persistence=True) This will record all `flyte.run` runs, that are local and are within the flyte.init being active. ### [flyte stop](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-stop) **`flyte stop COMMAND [ARGS]...`** Stop various Flyte services. #### [flyte stop devbox](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-stop-devbox) **`flyte stop devbox`** Pause the local Flyte devbox cluster without removing it. ### [flyte update](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-update) **`flyte update COMMAND [ARGS]...`** Update various flyte entities. #### [flyte update app](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-update-app) **`flyte update app [OPTIONS] NAME`** Update an app by starting or stopping it.  Example usage: flyte update app --activate | --deactivate [--wait] [--project ] [--domain ] | Option | Type | Default | Description | | --- | --- | --- | --- | | `--activate`
`--deactivate` | `boolean` | | Activate or deactivate app. | | `--wait` | `boolean` | `False` | Wait for the app to reach the desired state. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte update project](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-update-project) **`flyte update project [OPTIONS] ID`** Update a project’s name, description, labels, or archive state.  Example usage: flyte update project my_project --archive flyte update project my_project --unarchive flyte update project my_project --description "New description" flyte update project my_project --name "New Display Name" flyte update project my_project --label team=ml --label env=prod | Option | Type | Default | Description | | --- | --- | --- | --- | | `--name` | `text` | | Update the project display name. | | `--description` | `text` | | Update the project description. | | `--label`
`-l` | `text` | `Sentinel.UNSET` | Set labels as key=value pairs. Can be specified multiple times. Replaces all existing labels. | | `--archive`
`--unarchive` | `boolean` | | Archive or unarchive the project. | | `--help` | `boolean` | `False` | Show this message and exit. | #### [flyte update trigger](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-update-trigger) **`flyte update trigger [OPTIONS] NAME TASK_NAME`** Update a trigger.  Example usage: flyte update trigger --activate | --deactivate [--project --domain ] | Option | Type | Default | Description | | --- | --- | --- | --- | | `--activate`
`--deactivate` | `boolean` | `Sentinel.UNSET` | Activate or deactivate the trigger. | | `-p`
`--project` | `text` | | Project to which this command applies. | | `-d`
`--domain` | `text` | | Domain to which this command applies. | | `--help` | `boolean` | `False` | Show this message and exit. | ### [flyte whoami](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/#flyte-whoami) **`flyte whoami`** Display the current user information. LLM-optimized [This page](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/page.md) [Full docs index](https://www.union.ai/docs/v2/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Home | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Home | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) LLM-optimized [This page](https://www.union.ai/docs/v2/union/page.md) [Full docs index](https://www.union.ai/docs/v2/union/llms.txt) On this page 404 Page not found Showing closest match --- # User guide | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Flyte OSS ========= Flyte is a free and open source platform that provides a full suite of powerful features for orchestrating AI workflows. Flyte empowers AI development teams to rapidly ship high-quality code to production by offering optimized performance, unparalleled resource efficiency, and a delightful workflow authoring experience. You deploy and manage Flyte yourself, on your own cloud infrastructure. This documentation for open-source Flyte is maintained by Union.ai. You can switch to the documentation for the commercial versions with the selector above. Introduction Flyte is the leading open-source Kubernetes-native workflow orchestrator. Getting started Build your first Flyte workflow, exploring the major features of the platform along the way. Core concepts Understand the core concepts of the Flyte platform. Development cycle Explore the Flyte development cycle from experimentation to production. Data input/output Manage the input and output of data in your Flyte workflow. Programming Learn about Flyte-specific programming constructs. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/user-guide/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 1. [Flyte OSS](https://www.union.ai/docs/v1/flyte/user-guide/#flyte-oss) 404 Page not found Showing closest match [![](https://www.union.ai/docs/v1/flyte/images/icon-logo-flyte.svg)](https://www.union.ai/docs/v1/flyte/) [Flyte OSS Docs](https://www.union.ai/docs/v1/flyte/) --- # Unknown \# Tutorials > \*\*📝 Note\*\* > > An LLM-optimized bundle of this entire section is available at \[\`section.md\`\](https://www.union.ai/docs/v2/flyte/tutorials/section.md). > This single file contains all pages in this section, optimized for AI coding agent context. This section contains tutorials that showcase relevant use cases and provide step-by-step instructions on how to implement various features using Flyte and Union. Tutorials are organized by \*\*industry vertical\*\* and by \*\*technical topic\*\*. ## Industry verticals ### \[Biotech & Healthcare\](https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/page.md) Bioinformatics, medical imaging, and other life-sciences workloads. ### \[Geospatial\](https://www.union.ai/docs/v2/flyte/tutorials/geospatial/page.md) Satellite imagery, remote sensing, and earth and atmospheric modeling workloads. ### \[Financial Services & Fintech\](https://www.union.ai/docs/v2/flyte/tutorials/financial-services/page.md) Financial research, trading, and other fintech workloads. ### \[Frontier AI\](https://www.union.ai/docs/v2/flyte/tutorials/frontier-ai/page.md) Frontier-model pretraining, automated experimentation, and large-scale AI workloads. ## Technical topics ### \[Computer Vision\](https://www.union.ai/docs/v2/flyte/tutorials/computer-vision/page.md) Image and vision-language model workloads. ### \[Agents\](https://www.union.ai/docs/v2/flyte/tutorials/agents/page.md) Agentic workflows and autonomous LLM-powered systems. ### \[Context Engineering\](https://www.union.ai/docs/v2/flyte/tutorials/context-engineering/page.md) Prompt engineering, prompt optimization, and context construction. ### \[Model Training\](https://www.union.ai/docs/v2/flyte/tutorials/model-training/page.md) Training, fine-tuning, and hyperparameter optimization of models at scale. ### \[Data Processing\](https://www.union.ai/docs/v2/flyte/tutorials/data-processing) Large-scale data processing and batching strategies. ## Subpages - \[Biotech & Healthcare\](https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/page.md) - \[Genomic alignment\](https://www.union.ai/docs/v2/flyte/tutorials/genomic-alignment/page.md) - \[Brain tumor MRI classification\](https://www.union.ai/docs/v2/flyte/tutorials/tumor-detection/page.md) - \[Geospatial\](https://www.union.ai/docs/v2/flyte/tutorials/geospatial/page.md) - \[GPU-accelerated climate modeling\](https://www.union.ai/docs/v2/flyte/tutorials/climate-modeling/page.md) - \[Satellite image classification\](https://www.union.ai/docs/v2/flyte/tutorials/satellite\_image\_classification) - \[Financial Services & Fintech\](https://www.union.ai/docs/v2/flyte/tutorials/financial-services/page.md) - \[Financial research agent\](https://www.union.ai/docs/v2/flyte/tutorials/financial-research-agent/page.md) - \[Multi-agent trading simulation\](https://www.union.ai/docs/v2/flyte/tutorials/trading-agents/page.md) - \[Frontier AI\](https://www.union.ai/docs/v2/flyte/tutorials/frontier-ai/page.md) - \[Distributed LLM pretraining\](https://www.union.ai/docs/v2/flyte/tutorials/distributed-pretraining/page.md) - \[Computer Vision\](https://www.union.ai/docs/v2/flyte/tutorials/computer-vision/page.md) - \[Fine-tuning a VLM\](https://www.union.ai/docs/v2/flyte/tutorials/qwen-vl-finetuning/page.md) - \[Multimodal retrieval evaluation\](https://www.union.ai/docs/v2/flyte/tutorials/multimodal-retrieval-evaluation/page.md) - \[Agents\](https://www.union.ai/docs/v2/flyte/tutorials/agents/page.md) - \[Autoresearch agent\](https://www.union.ai/docs/v2/flyte/tutorials/autoresearch/page.md) - \[Coding agent\](https://www.union.ai/docs/v2/flyte/tutorials/code-agent/page.md) - \[Competitive intelligence agent\](https://www.union.ai/docs/v2/flyte/tutorials/competitive-intelligence-agent/page.md) - \[Compliance monitoring agent\](https://www.union.ai/docs/v2/flyte/tutorials/compliance-monitoring-agent/page.md) - \[Deep research\](https://www.union.ai/docs/v2/flyte/tutorials/deep-research/page.md) - \[Field data enrichment agent\](https://www.union.ai/docs/v2/flyte/tutorials/field-data-enrichment-agent/page.md) - \[MLE Bot: autonomous ML engineer\](https://www.union.ai/docs/v2/flyte/tutorials/mle-bot/page.md) - \[Support resolution agent\](https://www.union.ai/docs/v2/flyte/tutorials/support-resolution-agent/page.md) - \[Context Engineering\](https://www.union.ai/docs/v2/flyte/tutorials/context-engineering/page.md) - \[Automatic prompt engineering\](https://www.union.ai/docs/v2/flyte/tutorials/auto\_prompt\_engineering/page.md) - \[Text-to-SQL\](https://www.union.ai/docs/v2/flyte/tutorials/text\_to\_sql/page.md) - \[Model Training\](https://www.union.ai/docs/v2/flyte/tutorials/model-training/page.md) - \[Hyperparameter optimization\](https://www.union.ai/docs/v2/flyte/tutorials/hpo/page.md) --- \*\*Source\*\*: https://github.com/unionai/unionai-docs/blob/main/content/tutorials/\_index.md \*\*HTML\*\*: https://www.union.ai/docs/v2/flyte/tutorials/ --- # Unknown \# Integrations > \*\*📝 Note\*\* > > An LLM-optimized bundle of this entire section is available at \[\`section.md\`\](https://www.union.ai/docs/v2/flyte/integrations/section.md). > This single file contains all pages in this section, optimized for AI coding agent context. Flyte 2 is designed to be extensible by default. While the core platform covers the most common orchestration needs, many production workloads require specialized infrastructure, external services or execution semantics that go beyond the core runtime. Flyte 2 exposes these capabilities through integrations. Under the hood, integrations are implemented using Flyte 2's plugin system, which provides a consistent way to extend the platform without modifying core execution logic. An integration allows you to declaratively enable new capabilities such as distributed compute frameworks or third-party services without manually managing infrastructure. You specify what you need, and Flyte takes care of how it is provisioned, used and cleaned up. This page covers: - The types of integrations Flyte 2 supports today - How integrations fit into Flyte 2's execution model - How to use integrations in your tasks - The integrations available out of the box If you need functionality that doesn't exist yet, Flyte 2's plugin system is intentionally open-ended. You can build and register your own integrations using the same architecture described here. ## Integration categories Flyte 2 integrations fall into the following categories: 1. \*\*Distributed compute\*\*: Provision transient compute clusters to run tasks across multiple nodes, with automatic lifecycle management. 2. \*\*Agentic AI\*\*: Support for various common aspects of agentic AI applications. 3. \*\*Configuration\*\*: Compose and pass hierarchical configuration objects between tasks, with type-safe schemas and CLI/YAML composition. 4. \*\*Experiment tracking\*\*: Integrate with experiment tracking platforms for logging metrics, parameters, and artifacts. 5. \*\*Data validation\*\*: Enforce schema contracts on dataframes flowing between tasks, with automatic validation reports. 6. \*\*Connectors\*\*: Stateless, long-running services that receive execution requests via gRPC and then submit work to external (or internal) systems. 7. \*\*LLM Serving\*\*: Deploy and serve large language models with an OpenAI-compatible API. 8. \*\*Notebook execution\*\*: Run parameterized Jupyter notebooks as typed Flyte tasks with cell-level reports. ## Distributed compute Distributed compute integrations allow tasks to run on dynamically provisioned clusters. These clusters are created just-in-time, scoped to the task execution and torn down automatically when the task completes. This enables large-scale parallelism without requiring users to operate or maintain long-running infrastructure. ### Supported distributed compute integrations | Plugin | Description | Common use cases | | --------------------------- | ------------------------------------------------ | ------------------------------------------------------ | | \[Ray\](https://www.union.ai/docs/v2/flyte/integrations/ray/\_index) | Provisions Ray clusters via KubeRay | Distributed Python, ML training, hyperparameter tuning | | \[Spark\](https://www.union.ai/docs/v2/flyte/integrations/spark/\_index) | Provisions Spark clusters via Spark Operator | Large-scale data processing, ETL pipelines | | \[Dask\](https://www.union.ai/docs/v2/flyte/integrations/dask/\_index) | Provisions Dask clusters via Dask Operator | Parallel Python workloads, dataframe operations | | \[PyTorch\](https://www.union.ai/docs/v2/flyte/integrations/pytorch/\_index) | Distributed PyTorch training with elastic launch | Single-node and multi-node training | Each plugin encapsulates: - Cluster provisioning - Resource configuration - Networking and service discovery - Lifecycle management and teardown From the task author's perspective, these details are abstracted away. ### How the plugin system works At a high level, Flyte 2's distributed compute plugin architecture follows a simple and consistent pattern. #### 1. Registration Each plugin registers itself with Flyte 2's core plugin registry: - \*\*\`TaskPluginRegistry\`\*\*: The central registry for all distributed compute plugins - Each plugin declares: - Its configuration schema - How that configuration maps to execution behavior This registration step makes the plugin discoverable by the runtime. #### 2. Task environments and plugin configuration Integrations are activated through a \`TaskEnvironment\`. A \`TaskEnvironment\` bundles: - A container image - Execution settings - A plugin configuration object enabled with \`plugin\_config\` The plugin configuration describes \_what\_ infrastructure or integration the task requires. #### 3. Automatic provisioning and execution When a task associated with a \`TaskEnvironment\` runs: 1. Flyte inspects the environment's plugin configuration 2. The plugin provisions the required infrastructure or integration 3. The task executes with access to that capability 4. Flyte cleans up all transient resources after completion ### Example: Using the Dask plugin Below is a complete example showing how a task gains access to a Dask cluster simply by running inside an environment configured with the Dask plugin. \`\`\`python from flyteplugins.dask import Dask, WorkerGroup import flyte # Define the Dask cluster configuration dask\_config = Dask( workers=WorkerGroup(number\_of\_workers=4) ) # Create a task environment that enables Dask env = flyte.TaskEnvironment( name="dask\_env", plugin\_config=dask\_config, image=image, ) # Any task in this environment has access to the Dask cluster @env.task async def process\_data(data: list) -> list: from distributed import Client client = Client() # Automatically connects to the provisioned cluster futures = client.map(transform, data) return client.gather(futures) \`\`\` When \`process\_data\` executes, Flyte performs the following steps: 1. Provisions a Dask cluster with 4 workers 2. Executes the task with network access to the cluster 3. Tears down the cluster once the task completes No cluster management logic appears in the task code. The task only expresses intent. ### Key design principle All distributed compute integrations follow the same mental model: - You declare the required capability via configuration - You attach that configuration to a task environment - Tasks decorated with that environment automatically gain access to the capability This makes it easy to swap execution backends or introduce distributed compute incrementally without rewriting workflows. ## Agentic AI Agentic AI integrations provide drop-in replacements for LLM provider SDKs. They let you use Flyte tasks as agent tools so that tool calls run with full Flyte observability, retries, and caching. ### Supported agentic AI integrations | Plugin | Description | Common use cases | | ----------------------------------- | -------------------------------------------------------------- | ---------------------------------------- | | \[OpenAI\](https://www.union.ai/docs/v2/flyte/integrations/openai/\_index) | Drop-in replacement for OpenAI Agents SDK \`function\_tool\` | Agentic workflows with OpenAI models | | \[Anthropic\](https://www.union.ai/docs/v2/flyte/integrations/anthropic/\_index) | Agent loop and \`function\_tool\` for the Anthropic Claude SDK | Agentic workflows with Claude | | \[Gemini\](https://www.union.ai/docs/v2/flyte/integrations/gemini/\_index) | Agent loop and \`function\_tool\` for the Google Gemini SDK | Agentic workflows with Gemini | | \[Code generation\](https://www.union.ai/docs/v2/flyte/integrations/codegen/\_index) | LLM-driven code generation with automatic testing in sandboxes | Data processing, ETL, analysis pipelines | ## Experiment tracking Experiment tracking integrations let you log metrics, parameters, and artifacts to external tracking platforms during Flyte task execution. ### Supported experiment tracking integrations | Plugin | Description | Common use cases | | ------------------------------------ | ---------------------------- | ------------------------------------------------ | | \[MLflow\](https://www.union.ai/docs/v2/flyte/integrations/mlflow/\_index) | MLflow experiment tracking | Experiment tracking, autologging, model registry | | \[Weights and Biases\](https://www.union.ai/docs/v2/flyte/integrations/wandb/\_index) | Weights & Biases integration | Experiment tracking and hyperparameter tuning | ## Configuration Configuration integrations let you compose and pass hierarchical configuration objects between Flyte tasks, with type-safe schemas and CLI/YAML composition. ### Supported configuration integrations | Plugin | Description | Common use cases | | ------------------------------- | ----------------------------------------------------------------- | --------------------------------------------------------------------------------- | | \[OmegaConf\](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/\_index) | \`DictConfig\` / \`ListConfig\` as native task input and output types | Passing composed configs between tasks, structured configs, YAML-driven pipelines | | \[Hydra\](https://www.union.ai/docs/v2/flyte/integrations/hydra/\_index) | Hydra config composition and sweep submission for Flyte tasks | YAML-driven experiment composition, grid and Bayesian sweeps, hardware presets | ## Data validation Data validation integrations enforce schema contracts on the dataframes flowing between tasks. They validate data at task boundaries, catch type and constraint violations early, and produce HTML reports visible in the Flyte UI. ### Supported data validation integrations | Plugin | Description | Common use cases | | --------------------------- | ---------------------------------------------------------- | ----------------------------------------------------------- | | \[Pandera\](https://www.union.ai/docs/v2/flyte/integrations/pandera/\_index) | Validates dataframes with pandera \`DataFrameModel\` schemas | Schema enforcement, data quality checks, validation reports | ## Connectors Connectors are stateless, long‑running services that receive execution requests via gRPC and then submit work to external (or internal) systems. Each connector runs as its own Kubernetes deployment, and is triggered when a Flyte task of the matching type is executed. Although they normally run inside the data plane, you can also run connectors locally as long as the required secrets/credentials are present locally. This is useful because connectors are just Python services that can be spawned in‑process. Connectors are designed to scale horizontally and reduce load on the core Flyte backend because they execute \_outside\_ the core system. This decoupling makes connectors efficient, resilient, and easy to iterate on. You can even test them locally without modifying backend configuration, which reduces friction during development. ### Supported connectors | Connector | Description | Common use cases | | --------------------------------- | ------------------------------------------- | ---------------------------------------- | | \[Snowflake\](https://www.union.ai/docs/v2/flyte/integrations/snowflake/\_index) | Run SQL queries on Snowflake asynchronously | Data warehousing, ETL, analytics queries | | \[BigQuery\](https://www.union.ai/docs/v2/flyte/integrations/bigquery/\_index) | Run SQL queries on Google BigQuery | Data warehousing, ETL, analytics queries | | \[Databricks\](https://www.union.ai/docs/v2/flyte/integrations/databricks/\_index) | Run PySpark jobs on Databricks clusters | Large-scale data processing, Spark ETL | ### Creating a new connector If none of the existing connectors meet your needs, you can build your own. > \[!NOTE\] > Connectors communicate via Protobuf, so in theory they can be implemented in any language. > Today, only \*\*Python\*\* connectors are supported. ### Async connector interface To implement a new async connector, extend \`AsyncConnector\` and implement the following methods, all of which must be idempotent: | Method | Purpose | | ---------- | ----------------------------------------------------------- | | \`create\` | Launch the external job (via REST, gRPC, SDK, or other API) | | \`get\` | Fetch current job state (return job status or output) | | \`delete\` | Delete / cancel the external job | | \`get\_logs\` | Stream paginated log lines to the Flyte UI | To test the connector locally, the connector task should inherit from \[AsyncConnectorExecutorMixin\](https://github.com/flyteorg/flyte-sdk/blob/1d49299294cd5e15385fe8c48089b3454b7a4cd1/src/flyte/connectors/\_connector.py#L206). This mixin simulates how the Flyte 2 system executes asynchronous connector tasks, making it easier to validate your connector implementation before deploying it. ### Example: Batch job connector The following example implements a connector that simulates submitting and polling an external batch job. Replace the mock logic with real API calls for your use case. \*\*Connector\*\* (\`my\_connector/connector.py\`): \`\`\` import time import uuid from dataclasses import dataclass from typing import Any, Dict, Optional from flyteidl2.connector.connector\_pb2 import ( GetTaskLogsResponse, GetTaskLogsResponseBody, GetTaskLogsResponseHeader, ) from flyteidl2.core.execution\_pb2 import TaskExecution from flyteidl2.logs.dataplane.payload\_pb2 import LogLine, LogLineOriginator from google.protobuf.timestamp\_pb2 import Timestamp from flyte import logger from flyte.connectors import AsyncConnector, ConnectorRegistry, Resource, ResourceMeta @dataclass class BatchJobMetadata(ResourceMeta): job\_id: str created\_at: float class BatchJobConnector(AsyncConnector): name = "Batch Job Connector" task\_type\_name = "batch\_job" metadata\_type = BatchJobMetadata async def create(self, task\_template, inputs: Optional\[Dict\[str, Any\]\] = None, \*\*kwargs) -> BatchJobMetadata: job\_id = str(uuid.uuid4())\[:8\] logger.info(f"Submitted batch job {job\_id}") return BatchJobMetadata(job\_id=job\_id, created\_at=time.time()) async def get(self, resource\_meta: BatchJobMetadata, \*\*kwargs) -> Resource: elapsed = time.time() - resource\_meta.created\_at if elapsed < 5: return Resource(phase=TaskExecution.RUNNING, message="Job in progress") return Resource( phase=TaskExecution.SUCCEEDED, message="Job completed", outputs={"result": f"output-from-{resource\_meta.job\_id}"}, ) async def delete(self, resource\_meta: BatchJobMetadata, \*\*kwargs): logger.info(f"Cancelled job {resource\_meta.job\_id}") async def get\_logs(self, resource\_meta: BatchJobMetadata, token: str = "", \*\*kwargs): def line(message: str, ts: float) -> LogLine: t = Timestamp() t.FromSeconds(int(ts)) return LogLine(timestamp=t, message=message, originator=LogLineOriginator.USER) start = resource\_meta.created\_at job\_id = resource\_meta.job\_id pages = { "": GetTaskLogsResponseBody(lines=\[\ line(f"\[INFO\] Job {job\_id} submitted", start),\ line(f"\[INFO\] Job {job\_id} started", start + 1),\ \]), "page-2": GetTaskLogsResponseBody(lines=\[\ line(f"\[INFO\] Job {job\_id} finished", start + 5),\ \]), } next\_tokens = {"": "page-2", "page-2": ""} yield GetTaskLogsResponse(body=pages.get(token, GetTaskLogsResponseBody(lines=\[\]))) next\_token = next\_tokens.get(token, "") if next\_token: yield GetTaskLogsResponse(header=GetTaskLogsResponseHeader(token=next\_token)) ConnectorRegistry.register(BatchJobConnector()) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/batch\_job/connector.py\* \*\*Task plugin\*\* (\`my\_connector/task.py\`): \`\`\` from dataclasses import dataclass from typing import Any, Dict, Optional, Type from flyte.connectors import AsyncConnectorExecutorMixin from flyte.extend import TaskTemplate from flyte.models import NativeInterface, SerializationContext @dataclass class BatchJobConfig: timeout\_seconds: int = 300 class BatchJobTask(AsyncConnectorExecutorMixin, TaskTemplate): \_TASK\_TYPE = "batch\_job" def \_\_init\_\_(self, name: str, plugin\_config: BatchJobConfig, inputs: Optional\[Dict\[str, Type\]\] = None, outputs: Optional\[Dict\[str, Type\]\] = None, \*\*kwargs): super().\_\_init\_\_( name=name, interface=NativeInterface( {k: (v, None) for k, v in inputs.items()} if inputs else {}, outputs or {}, ), task\_type=self.\_TASK\_TYPE, image=None, \*\*kwargs, ) self.plugin\_config = plugin\_config def custom\_config(self, sctx: SerializationContext) -> Optional\[Dict\[str, Any\]\]: return {"timeout\_seconds": self.plugin\_config.timeout\_seconds} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/batch\_job/task.py\* \*\*Usage\*\*: \`\`\`python import flyte from my\_connector.task import BatchJobConfig, BatchJobTask batch\_job = BatchJobTask( name="my\_batch\_job", plugin\_config=BatchJobConfig(timeout\_seconds=60), inputs={"name": str}, outputs={"result": str}, ) flyte.TaskEnvironment.from\_task("batch-job-env", batch\_job) \`\`\` ### Connector-level secrets If your connector needs credentials (API keys, tokens) shared across all tasks, pass them as environment variables into the connector process. Set environment variables on the connector Kubernetes deployment: \`\`\`bash kubectl set env deployment/ MY\_API\_KEY= -n \`\`\` Inside the connector, read the secret from the environment: \`\`\`python import os api\_key = os.environ\["MY\_API\_KEY"\] \`\`\` See \[Secrets\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets/page.md) for how to store and manage secrets. ### Deploy a custom connector Deploying a connector requires two steps: building a Docker image that contains your connector code and then patching the connector Kubernetes deployment to use it. \*\*Step 1: Build the connector image\*\* \`\`\`python import asyncio from flyte import Image from flyte.extend import ImageBuildEngine async def build\_connector\_image(registry: str, name: str, builder: str = "local"): image = Image.from\_debian\_base( registry=registry, name=name ).with\_pip\_packages("flyte\[connector\]", "my-connector-package") await ImageBuildEngine.build(image, builder=builder) if \_\_name\_\_ == "\_\_main\_\_": asyncio.run( build\_connector\_image( registry="", name="my-connector", builder="local" ) ) \`\`\` \*\*Step 2: Override the connector deployment image\*\* Once the image is pushed, patch the connector Kubernetes deployment to use it: \`\`\`bash kubectl set image deployment/ \\ connector=/my-connector: \\ -n \`\`\` Replace \`\` with the name of your connector deployment (e.g. \`flyte-connector\`), and \`\` with the namespace where Flyte is installed (typically \`flyte\`). ## LLM Serving LLM serving integrations let you deploy and serve large language models as Flyte apps with an OpenAI-compatible API. They handle model loading, GPU management, and autoscaling. ### Supported LLM serving integrations | Plugin | Description | Common use cases | | --------------------------------------------- | --------------------------------------------------- | ---------------------------- | | \[SGLang\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/sglang-app/page.md) | Deploy models with SGLang's high-throughput runtime | LLM inference, model serving | | \[vLLM\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/vllm-app/page.md) | Deploy models with vLLM's PagedAttention engine | LLM inference, model serving | For full setup instructions including multi-GPU deployment, model prefetching, and autoscaling, see the \[SGLang app\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/sglang-app/page.md) and \[vLLM app\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/vllm-app/page.md) pages. ## Notebook execution Notebook execution integrations let you run Jupyter notebooks as first-class Flyte tasks with typed inputs and outputs, HTML reports surfaced in the Flyte UI, and the ability to call other Flyte tasks from within the notebook. ### Supported notebook execution integrations | Plugin | Description | Common use cases | | ------------------------------- | ------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------- | | \[Papermill\](https://www.union.ai/docs/v2/flyte/integrations/papermill/\_index) | Parameterize and execute \`.ipynb\` files via \[papermill\](https://papermill.readthedocs.io/) | Productionizing exploratory notebooks, cell-by-cell HTML reports, notebook-driven analysis pipelines | ## Subpages - \[Anthropic\](https://www.union.ai/docs/v2/flyte/integrations/anthropic/page.md) - Installation - Quick start - API - \`function\_tool\` - \`Agent\` - \`run\_agent\` - Secrets - API reference - \[BigQuery\](https://www.union.ai/docs/v2/flyte/integrations/bigquery/page.md) - Installation - Quick start - Configuration - \`BigQueryConfig\` parameters - \`BigQueryTask\` parameters - Authentication - Query templating - Supported input types - Parameterized query example - Retrieving query results - API reference - \[Code generation\](https://www.union.ai/docs/v2/flyte/integrations/codegen/page.md) - Installation - Quick start - Two execution backends - LiteLLM (default) - Agent (Claude) - Providing data - Sample data - Schema and constraints - Inputs and outputs - Running generated code - One-shot execution with \`result.run()\` - Reusable task with \`result.as\_task()\` - Error diagnosis - Durable execution - Replay logs - Caching - Non-determinism in Agent mode - Observability - LiteLLM backend - Agent backend - Examples - Processing CSVs with different schemas - DataFrame analysis with constraints - Agent mode - Configuration - LiteLLM parameters - Image configuration - Skipping tests - Base packages - Best practices - API reference - \`AutoCoderAgent\` constructor - \`generate()\` parameters - \`CodeGenEvalResult\` fields - \`CodeGenEvalResult\` methods - \[Dask\](https://www.union.ai/docs/v2/flyte/integrations/dask/page.md) - When to use this plugin - Installation - Configuration - \`Dask\` parameters - \`Scheduler\` parameters - \`WorkerGroup\` parameters - Accessing the Dask client - Example - API reference - \[Databricks\](https://www.union.ai/docs/v2/flyte/integrations/databricks/page.md) - Installation - Quick start - Configuration - Spark fields (inherited) - Databricks-specific fields - \`databricks\_conf\` structure - Authentication - Accessing the Spark session - API reference - \[Gemini\](https://www.union.ai/docs/v2/flyte/integrations/gemini/page.md) - Installation - Quick start - API - \`function\_tool\` - \`Agent\` - \`run\_agent\` - Secrets - API reference - \[Hydra\](https://www.union.ai/docs/v2/flyte/integrations/hydra/page.md) - Installation - Requirements on tasks - A walkthrough config - Execution mode - Hydra launcher (\`@hydra.main\` scripts) - Python SDK - Single run - Grid sweep - Custom sweepers - Forwarding \`flyte.with\_runcontext\` options - Flyte CLI (\`flyte hydra run\`) - Single run - Grid sweep - App-level vs Hydra-namespace overrides - \`--follow\` and \`--no-wait\` - Shell completion - Override grammar - Sweeps - Grid sweeps (BasicSweeper) - Bayesian / TPE sweeps (Optuna) - Sweep output directories - Task environment overrides - Prebuilt images - Applying overrides to child tasks - Renaming the task-env key - What \`task\_env\` should not model - Structured configs (without YAML) - \[MLflow\](https://www.union.ai/docs/v2/flyte/integrations/mlflow/page.md) - Installation - Quick start - Autologging - Generic autologging - Framework-specific autologging - Run modes - Sharing a run across tasks - Creating independent runs - Nested runs - Workflow-level configuration - Per-task overrides - Configuration priority - Distributed training - MLflow UI links - Setup - Custom URL templates - Explicit links - Link behavior by run mode - Automatic Flyte tags - API reference - \`mlflow\_run\` and \`mlflow\_config\` - \`get\_mlflow\_run\` - \`get\_mlflow\_context\` - \`Mlflow\` - \[OmegaConf\](https://www.union.ai/docs/v2/flyte/integrations/omegaconf/page.md) - Installation - Quick start - When to use this plugin - Building a DictConfig - From a plain dict - From a YAML file - From a dataclass (structured config) - From a base config plus overrides - Variable interpolation - Nested and deeply structured configs - DictConfigs that contain lists - ListConfig as input and output - Lists of primitives - Building a schedule from another task - Nested lists (list of lists) - Lists of DictConfigs - Lists of dataclass instances - Structured configs - Basic structured config - Schema reconstruction in the receiving task - Required (\`MISSING\`) fields - Advanced field types - Merging overrides on top of a structured base - Embedding rich Python values inside a plain DictConfig - Reserved-looking keys - YAML reports - Wire format - End-to-end example - \[OpenAI\](https://www.union.ai/docs/v2/flyte/integrations/openai/page.md) - When to use this plugin - Installation - Usage - \`function\_tool\` - Basic pattern - Secrets - Example - API reference - \[Pandera\](https://www.union.ai/docs/v2/flyte/integrations/pandera/page.md) - When to use this plugin - Installation - pandas - Polars - PySpark SQL - Defining schemas - Using schemas in tasks - Error handling with \`ValidationConfig\` - Image configuration - Pandas - Polars - PySpark SQL - Polars lazy frames - Examples - pandas - Polars - PySpark SQL - \[Papermill\](https://www.union.ai/docs/v2/flyte/integrations/papermill/page.md) - When to use this plugin - Installation - Quick start - Notebook setup - \`parameters\` cell - \`outputs\` cell - Inputs and outputs - Supported input types - Complex types: File, Dir, DataFrame - Outputs: single, multiple, none - Calling Flyte tasks from notebooks - Workflow patterns - Chaining notebooks - Mixing notebooks with regular tasks - Inline definition - Calling from sync vs. async tasks - Running a NotebookTask directly as the entrypoint - Reports and notebook artifacts - HTML report (default) - Notebook artifacts - Clean reports - Failure reports - Spark notebooks - Local testing - Execution options - \`NotebookTask\` reference - Helper functions - \[PyTorch\](https://www.union.ai/docs/v2/flyte/integrations/pytorch/page.md) - When to use this plugin - Installation - Configuration - \`Elastic\` parameters - \`RunPolicy\` parameters - NCCL tuning parameters - Writing a distributed training task - Example - API reference - \[Ray\](https://www.union.ai/docs/v2/flyte/integrations/ray/page.md) - When to use this plugin - Installation - Configuration - \`RayJobConfig\` parameters - \`WorkerNodeConfig\` parameters - \`HeadNodeConfig\` parameters - Connecting to an existing cluster - Examples - API reference - \[Snowflake\](https://www.union.ai/docs/v2/flyte/integrations/snowflake/page.md) - Installation - Quick start - Configuration - Required fields - Additional connection parameters - Authentication - Key-pair authentication - Password authentication - OAuth authentication - Query templating - Supported input types - Batched \`INSERT\` with list inputs - Parameterized \`SELECT\` - Multiple inputs - Retrieving query results - End-to-end example - \[Spark\](https://www.union.ai/docs/v2/flyte/integrations/spark/page.md) - When to use this plugin - Installation - Configuration - \`Spark\` parameters - Accessing the Spark session - Overriding configuration at runtime - Example - API reference - \[Weights & Biases\](https://www.union.ai/docs/v2/flyte/integrations/wandb/page.md) - Installation - Quick start - What's next --- \*\*Source\*\*: https://github.com/unionai/unionai-docs/blob/main/content/integrations/\_index.md \*\*HTML\*\*: https://www.union.ai/docs/v2/flyte/integrations/ --- # Unknown \# Tutorials > This bundle contains all pages in the Tutorials section. > Source: https://www.union.ai/docs/v2/flyte/tutorials/ === PAGE: https://www.union.ai/docs/v2/flyte/tutorials === # Tutorials > \*\*📝 Note\*\* > > An LLM-optimized bundle of this entire section is available at \[\`section.md\`\](section.md). > This single file contains all pages in this section, optimized for AI coding agent context. This section contains tutorials that showcase relevant use cases and provide step-by-step instructions on how to implement various features using Flyte and Union. Tutorials are organized by \*\*industry vertical\*\* and by \*\*technical topic\*\*. ## Industry verticals ### \*\*Biotech & Healthcare\*\* Bioinformatics, medical imaging, and other life-sciences workloads. ### \*\*Geospatial\*\* Satellite imagery, remote sensing, and earth and atmospheric modeling workloads. ### \*\*Financial Services & Fintech\*\* Financial research, trading, and other fintech workloads. ### \*\*Frontier AI\*\* Frontier-model pretraining, automated experimentation, and large-scale AI workloads. ## Technical topics ### \*\*Computer Vision\*\* Image and vision-language model workloads. ### \*\*Agents\*\* Agentic workflows and autonomous LLM-powered systems. ### \*\*Context Engineering\*\* Prompt engineering, prompt optimization, and context construction. ### \*\*Model Training\*\* Training, fine-tuning, and hyperparameter optimization of models at scale. ### \[Data Processing\](data-processing) Large-scale data processing and batching strategies. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare === # Biotech & Healthcare Tutorials for bioinformatics, medical imaging, and other life-sciences workloads. ### \*\*Biotech & Healthcare > Genomic alignment\*\* Align sequencing reads to a reference genome with a cached, parallel Bowtie 2 pipeline. ### \*\*Biotech & Healthcare > Brain tumor MRI classification\*\* Classify brain MRI scans with a two-phase EfficientNet-B4 pipeline featuring resumable GPU checkpointing and in-UI reports. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/genomic-alignment === # Genomic alignment > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/genomic\_alignment). This tutorial builds a bioinformatics pipeline that aligns raw sequencing reads to a reference genome. The workflow downloads a reference genome and paired-end sequencing data, performs quality filtering, builds a reference index, and aligns the filtered reads with the \[Bowtie 2\](https://bowtie-bio.sourceforge.net/bowtie2/index.shtml) aligner — running each sample in parallel. It's a good showcase of how Flyte handles real bioinformatics workloads: - \*\*Per-task resources\*\* so quality filtering, indexing, and alignment each request exactly the CPU and memory they need. - \*\*\`cache="auto"\`\*\* on the download and indexing steps, so re-runs skip work that hasn't changed. - \*\*Fan-out parallelism\*\* across samples with \`asyncio.gather\`. - \*\*System dependencies\*\* (\`fastp\`, \`bowtie2\`) installed into the container image with \`apt\`. ## Define the container image Because the pipeline shells out to bioinformatics tools, we build a custom image with \`flyte.Image.from\_uv\_script\` and install \`fastp\` (quality filtering) and \`bowtie2\` (alignment) via \`apt\`. \`\`\` # # Genomic Alignment # # This tutorial demonstrates how to use Flyte to build a workflow that # performs genomic alignment on sequencing data. The workflow takes as input # a reference genome and raw sequencing data, performs quality filtering and # preprocessing on the raw data, generates an index for the reference genome, # and aligns the filtered data to the reference genome using the Bowtie 2 aligner. # {{run-on-union}} # The tutorial is divided into the following sections: # 1. Define the container image # 2. Define the data classes # 3. Define the tasks # 4. Define the workflow # /// script # requires-python = "3.12" # dependencies = \[\ # "flyte",\ # "requests",\ # \] # main = "alignment\_wf" # params = "" # /// import asyncio import subprocess import tempfile from dataclasses import dataclass from pathlib import Path from typing import List import requests import flyte from flyte.io import Dir, File # ## Defining a Container Image # # We define a custom container image using \`flyte.Image\`. Since we need bioinformatics # tools — \`fastp\` for quality filtering and \`bowtie2\` for alignment — we install them # via apt. This approach replaces the v1 \`ImageSpec\` with conda channels. # {{docs-fragment image}} main\_img = ( flyte.Image.from\_uv\_script( \_\_file\_\_, name="alignment-tutorial", ) .with\_apt\_packages("fastp", "bowtie2") ) # {{/docs-fragment image}} # We define per-task environments with different resource requirements, then a # top-level \`base\_env\` that declares all of them as dependencies (required because # \`alignment\_wf\` and \`bowtie2\_align\_samples\` call tasks that live in those environments). # {{docs-fragment envs}} fetch\_env = flyte.TaskEnvironment( name="alignment-tutorial-fetch", image=main\_img, cache="auto", ) fastp\_env = flyte.TaskEnvironment( name="alignment-tutorial-fastp", image=main\_img, resources=flyte.Resources(memory="2Gi"), ) index\_env = flyte.TaskEnvironment( name="alignment-tutorial-index", image=main\_img, resources=flyte.Resources(memory="10Gi"), cache="auto", ) align\_env = flyte.TaskEnvironment( name="alignment-tutorial-align", image=main\_img, resources=flyte.Resources(cpu=2, memory="10Gi"), ) base\_env = flyte.TaskEnvironment( name="alignment-tutorial", image=main\_img, depends\_on=\[fetch\_env, fastp\_env, index\_env, align\_env\], ) # {{/docs-fragment envs}} # ## Defining Data Classes # # We define three data classes to represent the reference genome, sequencing reads, # and alignment results. We'll first define a convenience function to download files, # which we'll use within the fetch task to materialize assets from their remote locations. def fetch\_file(url: str, local\_dir: str) -> Path: """ Downloads a file from the specified URL. Args: url (str): The URL of the file to download. local\_dir (str): The directory where you would like this file saved. Returns: Path: The local path to the file. Raises: requests.HTTPError: If an HTTP error occurs while downloading the file. """ url\_parts = url.split("/") fname = url\_parts\[-1\] local\_path = Path(local\_dir) / fname response = requests.get(url) with open(local\_path, "wb") as file: file.write(response.content) return local\_path # Reference genomes are used extensively throughout bioinformatics workflows. We define a # \`Reference\` data class to represent a reference genome and its associated index files. # {{docs-fragment dataclasses}} @dataclass class Reference: """ Represents a reference FASTA and associated index files. Attributes: ref\_name (str): Name or identifier of the reference file. ref\_dir (Dir): Directory containing the reference and any index files. index\_name (str): Index string to pass to tools requiring it. indexed\_with (str): Name of tool used to create the index. """ ref\_name: str ref\_dir: Dir index\_name: str | None = None indexed\_with: str | None = None # Sequencing reads are the raw data generated from a sequencing experiment. @dataclass class Reads: """ Represents a sequencing reads sample via its associated FastQ files. Attributes: sample (str): The name or identifier of the raw sequencing sample. read1 (File): A File object representing the path to the raw R1 read file. read2 (File): A File object representing the path to the raw R2 read file. """ sample: str read1: File | None = None read2: File | None = None def get\_read\_fnames(self): return ( f"{self.sample}\_1.fastq.gz", f"{self.sample}\_2.fastq.gz", ) # Finally, we define an \`Alignment\` data class to represent an alignment file. @dataclass class Alignment: """ Represents an alignment file and its associated sample. Attributes: sample (str): The name or identifier of the sample. aligner (str): The name of the aligner used to generate the alignment file. format (str): The format of the alignment file (e.g., SAM, BAM). alignment (File): A File object representing the path to the alignment file. """ sample: str aligner: str format: str | None = None alignment: File | None = None def get\_alignment\_fname(self): return f"{self.sample}\_{self.aligner}\_aligned.{self.format}" # {{/docs-fragment dataclasses}} # ## Tasks # # We define a series of tasks to perform the following operations: # 1. Fetch assets from remote URLs # 2. Perform quality filtering and preprocessing using FastP # 3. Generate Bowtie2 index files from a reference genome # 4. Perform alignment using Bowtie2 on a filtered sample # # The first task fetches the reference genome and sequencing reads. It is cached # so that re-runs skip the download step. # {{docs-fragment fetch\_assets}} @fetch\_env.task async def fetch\_assets( ref\_url: str, read\_urls: List\[str\] ) -> tuple\[Reference, List\[Reads\]\]: """ Fetch assets from remote URLs. """ # Download reference genome ref\_dir = Path("/tmp/reference\_genome") ref\_dir.mkdir(exist\_ok=True, parents=True) ref = fetch\_file(ref\_url, str(ref\_dir)) ref\_obj = Reference( ref\_name=ref.name, ref\_dir=await Dir.from\_local(str(ref\_dir)), ) # Download sequencing reads dl\_loc = Path("/tmp/reads") dl\_loc.mkdir(exist\_ok=True, parents=True) samples: dict\[str, Reads\] = {} for url in read\_urls: fp = fetch\_file(url, str(dl\_loc)) sample = fp.stem.split("\_")\[0\] if sample not in samples: samples\[sample\] = Reads(sample=sample) if ".fastq.gz" in fp.name or "fasta" in fp.name: mate = fp.name.strip(".fastq.gz").strip(".filt").split("\_")\[-1\] if "1" in mate: samples\[sample\].read1 = await File.from\_local(str(fp)) elif "2" in mate: samples\[sample\].read2 = await File.from\_local(str(fp)) return ref\_obj, list(samples.values()) # {{/docs-fragment fetch\_assets}} # The second task performs quality filtering and preprocessing using FastP on a Reads object. # FastP is a performant tool for removing duplicate or low-quality reads. We increase # the memory request for this task so FastP can efficiently process reads from larger files. # {{docs-fragment pyfastp}} @fastp\_env.task async def pyfastp(rs: Reads) -> Reads: """ Perform quality filtering and preprocessing using Fastp on a Reads object. Args: rs (Reads): A Reads object containing raw sequencing data to be processed. Returns: Reads: A Reads object representing the filtered and preprocessed data. """ ldir = Path(tempfile.mkdtemp()) samp = Reads(rs.sample) o1, o2 = samp.get\_read\_fnames() o1p = ldir / o1 o2p = ldir / o2 assert rs.read1 is not None and rs.read2 is not None r1 = await rs.read1.download() r2 = await rs.read2.download() cmd = \[\ "fastp",\ "-i", str(r1),\ "-I", str(r2),\ "-o", str(o1p),\ "-O", str(o2p),\ \] subprocess.run(cmd, check=True) samp.read1 = await File.from\_local(str(o1p)) samp.read2 = await File.from\_local(str(o2p)) return samp # {{/docs-fragment pyfastp}} # Next, we define a task to generate Bowtie2 index files from a reference genome. As the index # for a given tool and reference seldom changes, we cache this task. # {{docs-fragment bowtie2\_index}} @index\_env.task async def bowtie2\_index(ref: Reference) -> Reference: """ Generate Bowtie2 index files from a reference genome. Args: ref (Reference): A Reference object representing the reference genome. Returns: Reference: The same reference object with the index\_name and indexed\_with attributes set. """ ref\_dir = await ref.ref\_dir.download() idx\_name = "bt2\_idx" cmd = \[\ "bowtie2-build",\ str(Path(str(ref\_dir)) / ref.ref\_name),\ str(Path(str(ref\_dir)) / idx\_name),\ \] subprocess.run(cmd, check=True) return Reference( ref.ref\_name, await Dir.from\_local(str(ref\_dir)), idx\_name, "bowtie2", ) # {{/docs-fragment bowtie2\_index}} # The next task performs paired-end alignment using Bowtie 2 on a single sample. # {{docs-fragment bowtie2\_align}} @align\_env.task async def bowtie2\_align\_paired\_reads(idx: Reference, fs: Reads) -> Alignment: """ Perform paired-end alignment using Bowtie 2 on a filtered sample. Args: idx (Reference): A Reference object containing the Bowtie 2 index. fs (Reads): A filtered Reads object containing sample data to be aligned. Returns: Alignment: An Alignment object representing the alignment result. """ assert idx.indexed\_with == "bowtie2", "Reference index must be generated with bowtie2" assert idx.index\_name is not None assert fs.read1 is not None and fs.read2 is not None ref\_dir = await idx.ref\_dir.download() r1 = await fs.read1.download() r2 = await fs.read2.download() ldir = Path(tempfile.mkdtemp()) alignment = Alignment(fs.sample, "bowtie2", "sam") al = ldir / alignment.get\_alignment\_fname() cmd = \[\ "bowtie2",\ "-x", str(Path(str(ref\_dir)) / idx.index\_name),\ "-1", str(r1),\ "-2", str(r2),\ "-S", str(al),\ \] subprocess.run(cmd, check=True) alignment.alignment = await File.from\_local(str(al)) return alignment # {{/docs-fragment bowtie2\_align}} # In place of the v1 \`@dynamic\` workflow, we use a plain async task with \`asyncio.gather\` # to run alignments for all samples in parallel. @base\_env.task async def bowtie2\_align\_samples( idx: Reference, samples: List\[Reads\] ) -> List\[Alignment\]: """ Process samples through bowtie2 in parallel. Args: idx (Reference): A Reference object containing the Bowtie 2 index. samples (List\[Reads\]): A list of Reads objects to be aligned. Returns: List\[Alignment\]: A list of Alignment objects representing the alignment results. """ tasks = \[bowtie2\_align\_paired\_reads(idx=idx, fs=sample) for sample in samples\] return list(await asyncio.gather(\*tasks)) # ## End-to-End Workflow # # We tie everything together in a final task that fetches assets, filters them, generates # an index, and aligns the samples. In place of the v1 \`@workflow\`, we use a top-level # \`@base\_env.task\`. Parallelism across samples is achieved with \`asyncio.gather\`. # {{docs-fragment workflow}} @base\_env.task async def alignment\_wf() -> List\[Alignment\]: # Prepare raw samples from remote URLs ref, samples = await fetch\_assets( ref\_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38\_short.fasta", read\_urls=\[\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_1.fastq.gz",\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_2.fastq.gz",\ \], ) # Filter all samples in parallel filtered\_samples = list( await asyncio.gather(\*\[pyfastp(rs=s) for s in samples\]) ) # Generate a bowtie2 index or load it from cache bowtie2\_idx = await bowtie2\_index(ref=ref) # Generate alignments using bowtie2 sams = await bowtie2\_align\_samples(idx=bowtie2\_idx, samples=filtered\_samples) return sams # {{/docs-fragment workflow}} # You can now run the workflow using the command in the dropdown at the top of the page! if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(alignment\_wf) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic\_alignment/genomic\_alignment.py\* The Python dependencies are declared at the top of the file using the \`uv\` script style: \`\`\` # /// script # requires-python = "3.12" # dependencies = \[\ # "flyte",\ # "requests",\ # \] # main = "alignment\_wf" # /// \`\`\` ## Define the task environments Each stage runs in its own \`TaskEnvironment\` with tailored resources. The top-level \`base\_env\` declares the others as \`depends\_on\` so the tasks it calls are available at run time. \`\`\` # # Genomic Alignment # # This tutorial demonstrates how to use Flyte to build a workflow that # performs genomic alignment on sequencing data. The workflow takes as input # a reference genome and raw sequencing data, performs quality filtering and # preprocessing on the raw data, generates an index for the reference genome, # and aligns the filtered data to the reference genome using the Bowtie 2 aligner. # {{run-on-union}} # The tutorial is divided into the following sections: # 1. Define the container image # 2. Define the data classes # 3. Define the tasks # 4. Define the workflow # /// script # requires-python = "3.12" # dependencies = \[\ # "flyte",\ # "requests",\ # \] # main = "alignment\_wf" # params = "" # /// import asyncio import subprocess import tempfile from dataclasses import dataclass from pathlib import Path from typing import List import requests import flyte from flyte.io import Dir, File # ## Defining a Container Image # # We define a custom container image using \`flyte.Image\`. Since we need bioinformatics # tools — \`fastp\` for quality filtering and \`bowtie2\` for alignment — we install them # via apt. This approach replaces the v1 \`ImageSpec\` with conda channels. # {{docs-fragment image}} main\_img = ( flyte.Image.from\_uv\_script( \_\_file\_\_, name="alignment-tutorial", ) .with\_apt\_packages("fastp", "bowtie2") ) # {{/docs-fragment image}} # We define per-task environments with different resource requirements, then a # top-level \`base\_env\` that declares all of them as dependencies (required because # \`alignment\_wf\` and \`bowtie2\_align\_samples\` call tasks that live in those environments). # {{docs-fragment envs}} fetch\_env = flyte.TaskEnvironment( name="alignment-tutorial-fetch", image=main\_img, cache="auto", ) fastp\_env = flyte.TaskEnvironment( name="alignment-tutorial-fastp", image=main\_img, resources=flyte.Resources(memory="2Gi"), ) index\_env = flyte.TaskEnvironment( name="alignment-tutorial-index", image=main\_img, resources=flyte.Resources(memory="10Gi"), cache="auto", ) align\_env = flyte.TaskEnvironment( name="alignment-tutorial-align", image=main\_img, resources=flyte.Resources(cpu=2, memory="10Gi"), ) base\_env = flyte.TaskEnvironment( name="alignment-tutorial", image=main\_img, depends\_on=\[fetch\_env, fastp\_env, index\_env, align\_env\], ) # {{/docs-fragment envs}} # ## Defining Data Classes # # We define three data classes to represent the reference genome, sequencing reads, # and alignment results. We'll first define a convenience function to download files, # which we'll use within the fetch task to materialize assets from their remote locations. def fetch\_file(url: str, local\_dir: str) -> Path: """ Downloads a file from the specified URL. Args: url (str): The URL of the file to download. local\_dir (str): The directory where you would like this file saved. Returns: Path: The local path to the file. Raises: requests.HTTPError: If an HTTP error occurs while downloading the file. """ url\_parts = url.split("/") fname = url\_parts\[-1\] local\_path = Path(local\_dir) / fname response = requests.get(url) with open(local\_path, "wb") as file: file.write(response.content) return local\_path # Reference genomes are used extensively throughout bioinformatics workflows. We define a # \`Reference\` data class to represent a reference genome and its associated index files. # {{docs-fragment dataclasses}} @dataclass class Reference: """ Represents a reference FASTA and associated index files. Attributes: ref\_name (str): Name or identifier of the reference file. ref\_dir (Dir): Directory containing the reference and any index files. index\_name (str): Index string to pass to tools requiring it. indexed\_with (str): Name of tool used to create the index. """ ref\_name: str ref\_dir: Dir index\_name: str | None = None indexed\_with: str | None = None # Sequencing reads are the raw data generated from a sequencing experiment. @dataclass class Reads: """ Represents a sequencing reads sample via its associated FastQ files. Attributes: sample (str): The name or identifier of the raw sequencing sample. read1 (File): A File object representing the path to the raw R1 read file. read2 (File): A File object representing the path to the raw R2 read file. """ sample: str read1: File | None = None read2: File | None = None def get\_read\_fnames(self): return ( f"{self.sample}\_1.fastq.gz", f"{self.sample}\_2.fastq.gz", ) # Finally, we define an \`Alignment\` data class to represent an alignment file. @dataclass class Alignment: """ Represents an alignment file and its associated sample. Attributes: sample (str): The name or identifier of the sample. aligner (str): The name of the aligner used to generate the alignment file. format (str): The format of the alignment file (e.g., SAM, BAM). alignment (File): A File object representing the path to the alignment file. """ sample: str aligner: str format: str | None = None alignment: File | None = None def get\_alignment\_fname(self): return f"{self.sample}\_{self.aligner}\_aligned.{self.format}" # {{/docs-fragment dataclasses}} # ## Tasks # # We define a series of tasks to perform the following operations: # 1. Fetch assets from remote URLs # 2. Perform quality filtering and preprocessing using FastP # 3. Generate Bowtie2 index files from a reference genome # 4. Perform alignment using Bowtie2 on a filtered sample # # The first task fetches the reference genome and sequencing reads. It is cached # so that re-runs skip the download step. # {{docs-fragment fetch\_assets}} @fetch\_env.task async def fetch\_assets( ref\_url: str, read\_urls: List\[str\] ) -> tuple\[Reference, List\[Reads\]\]: """ Fetch assets from remote URLs. """ # Download reference genome ref\_dir = Path("/tmp/reference\_genome") ref\_dir.mkdir(exist\_ok=True, parents=True) ref = fetch\_file(ref\_url, str(ref\_dir)) ref\_obj = Reference( ref\_name=ref.name, ref\_dir=await Dir.from\_local(str(ref\_dir)), ) # Download sequencing reads dl\_loc = Path("/tmp/reads") dl\_loc.mkdir(exist\_ok=True, parents=True) samples: dict\[str, Reads\] = {} for url in read\_urls: fp = fetch\_file(url, str(dl\_loc)) sample = fp.stem.split("\_")\[0\] if sample not in samples: samples\[sample\] = Reads(sample=sample) if ".fastq.gz" in fp.name or "fasta" in fp.name: mate = fp.name.strip(".fastq.gz").strip(".filt").split("\_")\[-1\] if "1" in mate: samples\[sample\].read1 = await File.from\_local(str(fp)) elif "2" in mate: samples\[sample\].read2 = await File.from\_local(str(fp)) return ref\_obj, list(samples.values()) # {{/docs-fragment fetch\_assets}} # The second task performs quality filtering and preprocessing using FastP on a Reads object. # FastP is a performant tool for removing duplicate or low-quality reads. We increase # the memory request for this task so FastP can efficiently process reads from larger files. # {{docs-fragment pyfastp}} @fastp\_env.task async def pyfastp(rs: Reads) -> Reads: """ Perform quality filtering and preprocessing using Fastp on a Reads object. Args: rs (Reads): A Reads object containing raw sequencing data to be processed. Returns: Reads: A Reads object representing the filtered and preprocessed data. """ ldir = Path(tempfile.mkdtemp()) samp = Reads(rs.sample) o1, o2 = samp.get\_read\_fnames() o1p = ldir / o1 o2p = ldir / o2 assert rs.read1 is not None and rs.read2 is not None r1 = await rs.read1.download() r2 = await rs.read2.download() cmd = \[\ "fastp",\ "-i", str(r1),\ "-I", str(r2),\ "-o", str(o1p),\ "-O", str(o2p),\ \] subprocess.run(cmd, check=True) samp.read1 = await File.from\_local(str(o1p)) samp.read2 = await File.from\_local(str(o2p)) return samp # {{/docs-fragment pyfastp}} # Next, we define a task to generate Bowtie2 index files from a reference genome. As the index # for a given tool and reference seldom changes, we cache this task. # {{docs-fragment bowtie2\_index}} @index\_env.task async def bowtie2\_index(ref: Reference) -> Reference: """ Generate Bowtie2 index files from a reference genome. Args: ref (Reference): A Reference object representing the reference genome. Returns: Reference: The same reference object with the index\_name and indexed\_with attributes set. """ ref\_dir = await ref.ref\_dir.download() idx\_name = "bt2\_idx" cmd = \[\ "bowtie2-build",\ str(Path(str(ref\_dir)) / ref.ref\_name),\ str(Path(str(ref\_dir)) / idx\_name),\ \] subprocess.run(cmd, check=True) return Reference( ref.ref\_name, await Dir.from\_local(str(ref\_dir)), idx\_name, "bowtie2", ) # {{/docs-fragment bowtie2\_index}} # The next task performs paired-end alignment using Bowtie 2 on a single sample. # {{docs-fragment bowtie2\_align}} @align\_env.task async def bowtie2\_align\_paired\_reads(idx: Reference, fs: Reads) -> Alignment: """ Perform paired-end alignment using Bowtie 2 on a filtered sample. Args: idx (Reference): A Reference object containing the Bowtie 2 index. fs (Reads): A filtered Reads object containing sample data to be aligned. Returns: Alignment: An Alignment object representing the alignment result. """ assert idx.indexed\_with == "bowtie2", "Reference index must be generated with bowtie2" assert idx.index\_name is not None assert fs.read1 is not None and fs.read2 is not None ref\_dir = await idx.ref\_dir.download() r1 = await fs.read1.download() r2 = await fs.read2.download() ldir = Path(tempfile.mkdtemp()) alignment = Alignment(fs.sample, "bowtie2", "sam") al = ldir / alignment.get\_alignment\_fname() cmd = \[\ "bowtie2",\ "-x", str(Path(str(ref\_dir)) / idx.index\_name),\ "-1", str(r1),\ "-2", str(r2),\ "-S", str(al),\ \] subprocess.run(cmd, check=True) alignment.alignment = await File.from\_local(str(al)) return alignment # {{/docs-fragment bowtie2\_align}} # In place of the v1 \`@dynamic\` workflow, we use a plain async task with \`asyncio.gather\` # to run alignments for all samples in parallel. @base\_env.task async def bowtie2\_align\_samples( idx: Reference, samples: List\[Reads\] ) -> List\[Alignment\]: """ Process samples through bowtie2 in parallel. Args: idx (Reference): A Reference object containing the Bowtie 2 index. samples (List\[Reads\]): A list of Reads objects to be aligned. Returns: List\[Alignment\]: A list of Alignment objects representing the alignment results. """ tasks = \[bowtie2\_align\_paired\_reads(idx=idx, fs=sample) for sample in samples\] return list(await asyncio.gather(\*tasks)) # ## End-to-End Workflow # # We tie everything together in a final task that fetches assets, filters them, generates # an index, and aligns the samples. In place of the v1 \`@workflow\`, we use a top-level # \`@base\_env.task\`. Parallelism across samples is achieved with \`asyncio.gather\`. # {{docs-fragment workflow}} @base\_env.task async def alignment\_wf() -> List\[Alignment\]: # Prepare raw samples from remote URLs ref, samples = await fetch\_assets( ref\_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38\_short.fasta", read\_urls=\[\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_1.fastq.gz",\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_2.fastq.gz",\ \], ) # Filter all samples in parallel filtered\_samples = list( await asyncio.gather(\*\[pyfastp(rs=s) for s in samples\]) ) # Generate a bowtie2 index or load it from cache bowtie2\_idx = await bowtie2\_index(ref=ref) # Generate alignments using bowtie2 sams = await bowtie2\_align\_samples(idx=bowtie2\_idx, samples=filtered\_samples) return sams # {{/docs-fragment workflow}} # You can now run the workflow using the command in the dropdown at the top of the page! if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(alignment\_wf) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic\_alignment/genomic\_alignment.py\* ## Define the data classes We model the reference genome, sequencing reads, and alignment results as dataclasses. \`flyte.io.File\` and \`flyte.io.Dir\` reference offloaded data in blob storage, so large genomic files are passed between tasks by reference rather than copied through the orchestrator. \`\`\` # # Genomic Alignment # # This tutorial demonstrates how to use Flyte to build a workflow that # performs genomic alignment on sequencing data. The workflow takes as input # a reference genome and raw sequencing data, performs quality filtering and # preprocessing on the raw data, generates an index for the reference genome, # and aligns the filtered data to the reference genome using the Bowtie 2 aligner. # {{run-on-union}} # The tutorial is divided into the following sections: # 1. Define the container image # 2. Define the data classes # 3. Define the tasks # 4. Define the workflow # /// script # requires-python = "3.12" # dependencies = \[\ # "flyte",\ # "requests",\ # \] # main = "alignment\_wf" # params = "" # /// import asyncio import subprocess import tempfile from dataclasses import dataclass from pathlib import Path from typing import List import requests import flyte from flyte.io import Dir, File # ## Defining a Container Image # # We define a custom container image using \`flyte.Image\`. Since we need bioinformatics # tools — \`fastp\` for quality filtering and \`bowtie2\` for alignment — we install them # via apt. This approach replaces the v1 \`ImageSpec\` with conda channels. # {{docs-fragment image}} main\_img = ( flyte.Image.from\_uv\_script( \_\_file\_\_, name="alignment-tutorial", ) .with\_apt\_packages("fastp", "bowtie2") ) # {{/docs-fragment image}} # We define per-task environments with different resource requirements, then a # top-level \`base\_env\` that declares all of them as dependencies (required because # \`alignment\_wf\` and \`bowtie2\_align\_samples\` call tasks that live in those environments). # {{docs-fragment envs}} fetch\_env = flyte.TaskEnvironment( name="alignment-tutorial-fetch", image=main\_img, cache="auto", ) fastp\_env = flyte.TaskEnvironment( name="alignment-tutorial-fastp", image=main\_img, resources=flyte.Resources(memory="2Gi"), ) index\_env = flyte.TaskEnvironment( name="alignment-tutorial-index", image=main\_img, resources=flyte.Resources(memory="10Gi"), cache="auto", ) align\_env = flyte.TaskEnvironment( name="alignment-tutorial-align", image=main\_img, resources=flyte.Resources(cpu=2, memory="10Gi"), ) base\_env = flyte.TaskEnvironment( name="alignment-tutorial", image=main\_img, depends\_on=\[fetch\_env, fastp\_env, index\_env, align\_env\], ) # {{/docs-fragment envs}} # ## Defining Data Classes # # We define three data classes to represent the reference genome, sequencing reads, # and alignment results. We'll first define a convenience function to download files, # which we'll use within the fetch task to materialize assets from their remote locations. def fetch\_file(url: str, local\_dir: str) -> Path: """ Downloads a file from the specified URL. Args: url (str): The URL of the file to download. local\_dir (str): The directory where you would like this file saved. Returns: Path: The local path to the file. Raises: requests.HTTPError: If an HTTP error occurs while downloading the file. """ url\_parts = url.split("/") fname = url\_parts\[-1\] local\_path = Path(local\_dir) / fname response = requests.get(url) with open(local\_path, "wb") as file: file.write(response.content) return local\_path # Reference genomes are used extensively throughout bioinformatics workflows. We define a # \`Reference\` data class to represent a reference genome and its associated index files. # {{docs-fragment dataclasses}} @dataclass class Reference: """ Represents a reference FASTA and associated index files. Attributes: ref\_name (str): Name or identifier of the reference file. ref\_dir (Dir): Directory containing the reference and any index files. index\_name (str): Index string to pass to tools requiring it. indexed\_with (str): Name of tool used to create the index. """ ref\_name: str ref\_dir: Dir index\_name: str | None = None indexed\_with: str | None = None # Sequencing reads are the raw data generated from a sequencing experiment. @dataclass class Reads: """ Represents a sequencing reads sample via its associated FastQ files. Attributes: sample (str): The name or identifier of the raw sequencing sample. read1 (File): A File object representing the path to the raw R1 read file. read2 (File): A File object representing the path to the raw R2 read file. """ sample: str read1: File | None = None read2: File | None = None def get\_read\_fnames(self): return ( f"{self.sample}\_1.fastq.gz", f"{self.sample}\_2.fastq.gz", ) # Finally, we define an \`Alignment\` data class to represent an alignment file. @dataclass class Alignment: """ Represents an alignment file and its associated sample. Attributes: sample (str): The name or identifier of the sample. aligner (str): The name of the aligner used to generate the alignment file. format (str): The format of the alignment file (e.g., SAM, BAM). alignment (File): A File object representing the path to the alignment file. """ sample: str aligner: str format: str | None = None alignment: File | None = None def get\_alignment\_fname(self): return f"{self.sample}\_{self.aligner}\_aligned.{self.format}" # {{/docs-fragment dataclasses}} # ## Tasks # # We define a series of tasks to perform the following operations: # 1. Fetch assets from remote URLs # 2. Perform quality filtering and preprocessing using FastP # 3. Generate Bowtie2 index files from a reference genome # 4. Perform alignment using Bowtie2 on a filtered sample # # The first task fetches the reference genome and sequencing reads. It is cached # so that re-runs skip the download step. # {{docs-fragment fetch\_assets}} @fetch\_env.task async def fetch\_assets( ref\_url: str, read\_urls: List\[str\] ) -> tuple\[Reference, List\[Reads\]\]: """ Fetch assets from remote URLs. """ # Download reference genome ref\_dir = Path("/tmp/reference\_genome") ref\_dir.mkdir(exist\_ok=True, parents=True) ref = fetch\_file(ref\_url, str(ref\_dir)) ref\_obj = Reference( ref\_name=ref.name, ref\_dir=await Dir.from\_local(str(ref\_dir)), ) # Download sequencing reads dl\_loc = Path("/tmp/reads") dl\_loc.mkdir(exist\_ok=True, parents=True) samples: dict\[str, Reads\] = {} for url in read\_urls: fp = fetch\_file(url, str(dl\_loc)) sample = fp.stem.split("\_")\[0\] if sample not in samples: samples\[sample\] = Reads(sample=sample) if ".fastq.gz" in fp.name or "fasta" in fp.name: mate = fp.name.strip(".fastq.gz").strip(".filt").split("\_")\[-1\] if "1" in mate: samples\[sample\].read1 = await File.from\_local(str(fp)) elif "2" in mate: samples\[sample\].read2 = await File.from\_local(str(fp)) return ref\_obj, list(samples.values()) # {{/docs-fragment fetch\_assets}} # The second task performs quality filtering and preprocessing using FastP on a Reads object. # FastP is a performant tool for removing duplicate or low-quality reads. We increase # the memory request for this task so FastP can efficiently process reads from larger files. # {{docs-fragment pyfastp}} @fastp\_env.task async def pyfastp(rs: Reads) -> Reads: """ Perform quality filtering and preprocessing using Fastp on a Reads object. Args: rs (Reads): A Reads object containing raw sequencing data to be processed. Returns: Reads: A Reads object representing the filtered and preprocessed data. """ ldir = Path(tempfile.mkdtemp()) samp = Reads(rs.sample) o1, o2 = samp.get\_read\_fnames() o1p = ldir / o1 o2p = ldir / o2 assert rs.read1 is not None and rs.read2 is not None r1 = await rs.read1.download() r2 = await rs.read2.download() cmd = \[\ "fastp",\ "-i", str(r1),\ "-I", str(r2),\ "-o", str(o1p),\ "-O", str(o2p),\ \] subprocess.run(cmd, check=True) samp.read1 = await File.from\_local(str(o1p)) samp.read2 = await File.from\_local(str(o2p)) return samp # {{/docs-fragment pyfastp}} # Next, we define a task to generate Bowtie2 index files from a reference genome. As the index # for a given tool and reference seldom changes, we cache this task. # {{docs-fragment bowtie2\_index}} @index\_env.task async def bowtie2\_index(ref: Reference) -> Reference: """ Generate Bowtie2 index files from a reference genome. Args: ref (Reference): A Reference object representing the reference genome. Returns: Reference: The same reference object with the index\_name and indexed\_with attributes set. """ ref\_dir = await ref.ref\_dir.download() idx\_name = "bt2\_idx" cmd = \[\ "bowtie2-build",\ str(Path(str(ref\_dir)) / ref.ref\_name),\ str(Path(str(ref\_dir)) / idx\_name),\ \] subprocess.run(cmd, check=True) return Reference( ref.ref\_name, await Dir.from\_local(str(ref\_dir)), idx\_name, "bowtie2", ) # {{/docs-fragment bowtie2\_index}} # The next task performs paired-end alignment using Bowtie 2 on a single sample. # {{docs-fragment bowtie2\_align}} @align\_env.task async def bowtie2\_align\_paired\_reads(idx: Reference, fs: Reads) -> Alignment: """ Perform paired-end alignment using Bowtie 2 on a filtered sample. Args: idx (Reference): A Reference object containing the Bowtie 2 index. fs (Reads): A filtered Reads object containing sample data to be aligned. Returns: Alignment: An Alignment object representing the alignment result. """ assert idx.indexed\_with == "bowtie2", "Reference index must be generated with bowtie2" assert idx.index\_name is not None assert fs.read1 is not None and fs.read2 is not None ref\_dir = await idx.ref\_dir.download() r1 = await fs.read1.download() r2 = await fs.read2.download() ldir = Path(tempfile.mkdtemp()) alignment = Alignment(fs.sample, "bowtie2", "sam") al = ldir / alignment.get\_alignment\_fname() cmd = \[\ "bowtie2",\ "-x", str(Path(str(ref\_dir)) / idx.index\_name),\ "-1", str(r1),\ "-2", str(r2),\ "-S", str(al),\ \] subprocess.run(cmd, check=True) alignment.alignment = await File.from\_local(str(al)) return alignment # {{/docs-fragment bowtie2\_align}} # In place of the v1 \`@dynamic\` workflow, we use a plain async task with \`asyncio.gather\` # to run alignments for all samples in parallel. @base\_env.task async def bowtie2\_align\_samples( idx: Reference, samples: List\[Reads\] ) -> List\[Alignment\]: """ Process samples through bowtie2 in parallel. Args: idx (Reference): A Reference object containing the Bowtie 2 index. samples (List\[Reads\]): A list of Reads objects to be aligned. Returns: List\[Alignment\]: A list of Alignment objects representing the alignment results. """ tasks = \[bowtie2\_align\_paired\_reads(idx=idx, fs=sample) for sample in samples\] return list(await asyncio.gather(\*tasks)) # ## End-to-End Workflow # # We tie everything together in a final task that fetches assets, filters them, generates # an index, and aligns the samples. In place of the v1 \`@workflow\`, we use a top-level # \`@base\_env.task\`. Parallelism across samples is achieved with \`asyncio.gather\`. # {{docs-fragment workflow}} @base\_env.task async def alignment\_wf() -> List\[Alignment\]: # Prepare raw samples from remote URLs ref, samples = await fetch\_assets( ref\_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38\_short.fasta", read\_urls=\[\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_1.fastq.gz",\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_2.fastq.gz",\ \], ) # Filter all samples in parallel filtered\_samples = list( await asyncio.gather(\*\[pyfastp(rs=s) for s in samples\]) ) # Generate a bowtie2 index or load it from cache bowtie2\_idx = await bowtie2\_index(ref=ref) # Generate alignments using bowtie2 sams = await bowtie2\_align\_samples(idx=bowtie2\_idx, samples=filtered\_samples) return sams # {{/docs-fragment workflow}} # You can now run the workflow using the command in the dropdown at the top of the page! if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(alignment\_wf) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic\_alignment/genomic\_alignment.py\* ## Fetch assets The first task downloads the reference genome and paired-end reads from remote URLs and materializes them as \`File\`/\`Dir\` objects. It's cached, so repeat runs skip the download. \`\`\` # # Genomic Alignment # # This tutorial demonstrates how to use Flyte to build a workflow that # performs genomic alignment on sequencing data. The workflow takes as input # a reference genome and raw sequencing data, performs quality filtering and # preprocessing on the raw data, generates an index for the reference genome, # and aligns the filtered data to the reference genome using the Bowtie 2 aligner. # {{run-on-union}} # The tutorial is divided into the following sections: # 1. Define the container image # 2. Define the data classes # 3. Define the tasks # 4. Define the workflow # /// script # requires-python = "3.12" # dependencies = \[\ # "flyte",\ # "requests",\ # \] # main = "alignment\_wf" # params = "" # /// import asyncio import subprocess import tempfile from dataclasses import dataclass from pathlib import Path from typing import List import requests import flyte from flyte.io import Dir, File # ## Defining a Container Image # # We define a custom container image using \`flyte.Image\`. Since we need bioinformatics # tools — \`fastp\` for quality filtering and \`bowtie2\` for alignment — we install them # via apt. This approach replaces the v1 \`ImageSpec\` with conda channels. # {{docs-fragment image}} main\_img = ( flyte.Image.from\_uv\_script( \_\_file\_\_, name="alignment-tutorial", ) .with\_apt\_packages("fastp", "bowtie2") ) # {{/docs-fragment image}} # We define per-task environments with different resource requirements, then a # top-level \`base\_env\` that declares all of them as dependencies (required because # \`alignment\_wf\` and \`bowtie2\_align\_samples\` call tasks that live in those environments). # {{docs-fragment envs}} fetch\_env = flyte.TaskEnvironment( name="alignment-tutorial-fetch", image=main\_img, cache="auto", ) fastp\_env = flyte.TaskEnvironment( name="alignment-tutorial-fastp", image=main\_img, resources=flyte.Resources(memory="2Gi"), ) index\_env = flyte.TaskEnvironment( name="alignment-tutorial-index", image=main\_img, resources=flyte.Resources(memory="10Gi"), cache="auto", ) align\_env = flyte.TaskEnvironment( name="alignment-tutorial-align", image=main\_img, resources=flyte.Resources(cpu=2, memory="10Gi"), ) base\_env = flyte.TaskEnvironment( name="alignment-tutorial", image=main\_img, depends\_on=\[fetch\_env, fastp\_env, index\_env, align\_env\], ) # {{/docs-fragment envs}} # ## Defining Data Classes # # We define three data classes to represent the reference genome, sequencing reads, # and alignment results. We'll first define a convenience function to download files, # which we'll use within the fetch task to materialize assets from their remote locations. def fetch\_file(url: str, local\_dir: str) -> Path: """ Downloads a file from the specified URL. Args: url (str): The URL of the file to download. local\_dir (str): The directory where you would like this file saved. Returns: Path: The local path to the file. Raises: requests.HTTPError: If an HTTP error occurs while downloading the file. """ url\_parts = url.split("/") fname = url\_parts\[-1\] local\_path = Path(local\_dir) / fname response = requests.get(url) with open(local\_path, "wb") as file: file.write(response.content) return local\_path # Reference genomes are used extensively throughout bioinformatics workflows. We define a # \`Reference\` data class to represent a reference genome and its associated index files. # {{docs-fragment dataclasses}} @dataclass class Reference: """ Represents a reference FASTA and associated index files. Attributes: ref\_name (str): Name or identifier of the reference file. ref\_dir (Dir): Directory containing the reference and any index files. index\_name (str): Index string to pass to tools requiring it. indexed\_with (str): Name of tool used to create the index. """ ref\_name: str ref\_dir: Dir index\_name: str | None = None indexed\_with: str | None = None # Sequencing reads are the raw data generated from a sequencing experiment. @dataclass class Reads: """ Represents a sequencing reads sample via its associated FastQ files. Attributes: sample (str): The name or identifier of the raw sequencing sample. read1 (File): A File object representing the path to the raw R1 read file. read2 (File): A File object representing the path to the raw R2 read file. """ sample: str read1: File | None = None read2: File | None = None def get\_read\_fnames(self): return ( f"{self.sample}\_1.fastq.gz", f"{self.sample}\_2.fastq.gz", ) # Finally, we define an \`Alignment\` data class to represent an alignment file. @dataclass class Alignment: """ Represents an alignment file and its associated sample. Attributes: sample (str): The name or identifier of the sample. aligner (str): The name of the aligner used to generate the alignment file. format (str): The format of the alignment file (e.g., SAM, BAM). alignment (File): A File object representing the path to the alignment file. """ sample: str aligner: str format: str | None = None alignment: File | None = None def get\_alignment\_fname(self): return f"{self.sample}\_{self.aligner}\_aligned.{self.format}" # {{/docs-fragment dataclasses}} # ## Tasks # # We define a series of tasks to perform the following operations: # 1. Fetch assets from remote URLs # 2. Perform quality filtering and preprocessing using FastP # 3. Generate Bowtie2 index files from a reference genome # 4. Perform alignment using Bowtie2 on a filtered sample # # The first task fetches the reference genome and sequencing reads. It is cached # so that re-runs skip the download step. # {{docs-fragment fetch\_assets}} @fetch\_env.task async def fetch\_assets( ref\_url: str, read\_urls: List\[str\] ) -> tuple\[Reference, List\[Reads\]\]: """ Fetch assets from remote URLs. """ # Download reference genome ref\_dir = Path("/tmp/reference\_genome") ref\_dir.mkdir(exist\_ok=True, parents=True) ref = fetch\_file(ref\_url, str(ref\_dir)) ref\_obj = Reference( ref\_name=ref.name, ref\_dir=await Dir.from\_local(str(ref\_dir)), ) # Download sequencing reads dl\_loc = Path("/tmp/reads") dl\_loc.mkdir(exist\_ok=True, parents=True) samples: dict\[str, Reads\] = {} for url in read\_urls: fp = fetch\_file(url, str(dl\_loc)) sample = fp.stem.split("\_")\[0\] if sample not in samples: samples\[sample\] = Reads(sample=sample) if ".fastq.gz" in fp.name or "fasta" in fp.name: mate = fp.name.strip(".fastq.gz").strip(".filt").split("\_")\[-1\] if "1" in mate: samples\[sample\].read1 = await File.from\_local(str(fp)) elif "2" in mate: samples\[sample\].read2 = await File.from\_local(str(fp)) return ref\_obj, list(samples.values()) # {{/docs-fragment fetch\_assets}} # The second task performs quality filtering and preprocessing using FastP on a Reads object. # FastP is a performant tool for removing duplicate or low-quality reads. We increase # the memory request for this task so FastP can efficiently process reads from larger files. # {{docs-fragment pyfastp}} @fastp\_env.task async def pyfastp(rs: Reads) -> Reads: """ Perform quality filtering and preprocessing using Fastp on a Reads object. Args: rs (Reads): A Reads object containing raw sequencing data to be processed. Returns: Reads: A Reads object representing the filtered and preprocessed data. """ ldir = Path(tempfile.mkdtemp()) samp = Reads(rs.sample) o1, o2 = samp.get\_read\_fnames() o1p = ldir / o1 o2p = ldir / o2 assert rs.read1 is not None and rs.read2 is not None r1 = await rs.read1.download() r2 = await rs.read2.download() cmd = \[\ "fastp",\ "-i", str(r1),\ "-I", str(r2),\ "-o", str(o1p),\ "-O", str(o2p),\ \] subprocess.run(cmd, check=True) samp.read1 = await File.from\_local(str(o1p)) samp.read2 = await File.from\_local(str(o2p)) return samp # {{/docs-fragment pyfastp}} # Next, we define a task to generate Bowtie2 index files from a reference genome. As the index # for a given tool and reference seldom changes, we cache this task. # {{docs-fragment bowtie2\_index}} @index\_env.task async def bowtie2\_index(ref: Reference) -> Reference: """ Generate Bowtie2 index files from a reference genome. Args: ref (Reference): A Reference object representing the reference genome. Returns: Reference: The same reference object with the index\_name and indexed\_with attributes set. """ ref\_dir = await ref.ref\_dir.download() idx\_name = "bt2\_idx" cmd = \[\ "bowtie2-build",\ str(Path(str(ref\_dir)) / ref.ref\_name),\ str(Path(str(ref\_dir)) / idx\_name),\ \] subprocess.run(cmd, check=True) return Reference( ref.ref\_name, await Dir.from\_local(str(ref\_dir)), idx\_name, "bowtie2", ) # {{/docs-fragment bowtie2\_index}} # The next task performs paired-end alignment using Bowtie 2 on a single sample. # {{docs-fragment bowtie2\_align}} @align\_env.task async def bowtie2\_align\_paired\_reads(idx: Reference, fs: Reads) -> Alignment: """ Perform paired-end alignment using Bowtie 2 on a filtered sample. Args: idx (Reference): A Reference object containing the Bowtie 2 index. fs (Reads): A filtered Reads object containing sample data to be aligned. Returns: Alignment: An Alignment object representing the alignment result. """ assert idx.indexed\_with == "bowtie2", "Reference index must be generated with bowtie2" assert idx.index\_name is not None assert fs.read1 is not None and fs.read2 is not None ref\_dir = await idx.ref\_dir.download() r1 = await fs.read1.download() r2 = await fs.read2.download() ldir = Path(tempfile.mkdtemp()) alignment = Alignment(fs.sample, "bowtie2", "sam") al = ldir / alignment.get\_alignment\_fname() cmd = \[\ "bowtie2",\ "-x", str(Path(str(ref\_dir)) / idx.index\_name),\ "-1", str(r1),\ "-2", str(r2),\ "-S", str(al),\ \] subprocess.run(cmd, check=True) alignment.alignment = await File.from\_local(str(al)) return alignment # {{/docs-fragment bowtie2\_align}} # In place of the v1 \`@dynamic\` workflow, we use a plain async task with \`asyncio.gather\` # to run alignments for all samples in parallel. @base\_env.task async def bowtie2\_align\_samples( idx: Reference, samples: List\[Reads\] ) -> List\[Alignment\]: """ Process samples through bowtie2 in parallel. Args: idx (Reference): A Reference object containing the Bowtie 2 index. samples (List\[Reads\]): A list of Reads objects to be aligned. Returns: List\[Alignment\]: A list of Alignment objects representing the alignment results. """ tasks = \[bowtie2\_align\_paired\_reads(idx=idx, fs=sample) for sample in samples\] return list(await asyncio.gather(\*tasks)) # ## End-to-End Workflow # # We tie everything together in a final task that fetches assets, filters them, generates # an index, and aligns the samples. In place of the v1 \`@workflow\`, we use a top-level # \`@base\_env.task\`. Parallelism across samples is achieved with \`asyncio.gather\`. # {{docs-fragment workflow}} @base\_env.task async def alignment\_wf() -> List\[Alignment\]: # Prepare raw samples from remote URLs ref, samples = await fetch\_assets( ref\_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38\_short.fasta", read\_urls=\[\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_1.fastq.gz",\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_2.fastq.gz",\ \], ) # Filter all samples in parallel filtered\_samples = list( await asyncio.gather(\*\[pyfastp(rs=s) for s in samples\]) ) # Generate a bowtie2 index or load it from cache bowtie2\_idx = await bowtie2\_index(ref=ref) # Generate alignments using bowtie2 sams = await bowtie2\_align\_samples(idx=bowtie2\_idx, samples=filtered\_samples) return sams # {{/docs-fragment workflow}} # You can now run the workflow using the command in the dropdown at the top of the page! if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(alignment\_wf) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic\_alignment/genomic\_alignment.py\* ## Quality filtering with fastp \`pyfastp\` removes duplicate and low-quality reads. It requests extra memory so it can process larger read files efficiently. \`\`\` # # Genomic Alignment # # This tutorial demonstrates how to use Flyte to build a workflow that # performs genomic alignment on sequencing data. The workflow takes as input # a reference genome and raw sequencing data, performs quality filtering and # preprocessing on the raw data, generates an index for the reference genome, # and aligns the filtered data to the reference genome using the Bowtie 2 aligner. # {{run-on-union}} # The tutorial is divided into the following sections: # 1. Define the container image # 2. Define the data classes # 3. Define the tasks # 4. Define the workflow # /// script # requires-python = "3.12" # dependencies = \[\ # "flyte",\ # "requests",\ # \] # main = "alignment\_wf" # params = "" # /// import asyncio import subprocess import tempfile from dataclasses import dataclass from pathlib import Path from typing import List import requests import flyte from flyte.io import Dir, File # ## Defining a Container Image # # We define a custom container image using \`flyte.Image\`. Since we need bioinformatics # tools — \`fastp\` for quality filtering and \`bowtie2\` for alignment — we install them # via apt. This approach replaces the v1 \`ImageSpec\` with conda channels. # {{docs-fragment image}} main\_img = ( flyte.Image.from\_uv\_script( \_\_file\_\_, name="alignment-tutorial", ) .with\_apt\_packages("fastp", "bowtie2") ) # {{/docs-fragment image}} # We define per-task environments with different resource requirements, then a # top-level \`base\_env\` that declares all of them as dependencies (required because # \`alignment\_wf\` and \`bowtie2\_align\_samples\` call tasks that live in those environments). # {{docs-fragment envs}} fetch\_env = flyte.TaskEnvironment( name="alignment-tutorial-fetch", image=main\_img, cache="auto", ) fastp\_env = flyte.TaskEnvironment( name="alignment-tutorial-fastp", image=main\_img, resources=flyte.Resources(memory="2Gi"), ) index\_env = flyte.TaskEnvironment( name="alignment-tutorial-index", image=main\_img, resources=flyte.Resources(memory="10Gi"), cache="auto", ) align\_env = flyte.TaskEnvironment( name="alignment-tutorial-align", image=main\_img, resources=flyte.Resources(cpu=2, memory="10Gi"), ) base\_env = flyte.TaskEnvironment( name="alignment-tutorial", image=main\_img, depends\_on=\[fetch\_env, fastp\_env, index\_env, align\_env\], ) # {{/docs-fragment envs}} # ## Defining Data Classes # # We define three data classes to represent the reference genome, sequencing reads, # and alignment results. We'll first define a convenience function to download files, # which we'll use within the fetch task to materialize assets from their remote locations. def fetch\_file(url: str, local\_dir: str) -> Path: """ Downloads a file from the specified URL. Args: url (str): The URL of the file to download. local\_dir (str): The directory where you would like this file saved. Returns: Path: The local path to the file. Raises: requests.HTTPError: If an HTTP error occurs while downloading the file. """ url\_parts = url.split("/") fname = url\_parts\[-1\] local\_path = Path(local\_dir) / fname response = requests.get(url) with open(local\_path, "wb") as file: file.write(response.content) return local\_path # Reference genomes are used extensively throughout bioinformatics workflows. We define a # \`Reference\` data class to represent a reference genome and its associated index files. # {{docs-fragment dataclasses}} @dataclass class Reference: """ Represents a reference FASTA and associated index files. Attributes: ref\_name (str): Name or identifier of the reference file. ref\_dir (Dir): Directory containing the reference and any index files. index\_name (str): Index string to pass to tools requiring it. indexed\_with (str): Name of tool used to create the index. """ ref\_name: str ref\_dir: Dir index\_name: str | None = None indexed\_with: str | None = None # Sequencing reads are the raw data generated from a sequencing experiment. @dataclass class Reads: """ Represents a sequencing reads sample via its associated FastQ files. Attributes: sample (str): The name or identifier of the raw sequencing sample. read1 (File): A File object representing the path to the raw R1 read file. read2 (File): A File object representing the path to the raw R2 read file. """ sample: str read1: File | None = None read2: File | None = None def get\_read\_fnames(self): return ( f"{self.sample}\_1.fastq.gz", f"{self.sample}\_2.fastq.gz", ) # Finally, we define an \`Alignment\` data class to represent an alignment file. @dataclass class Alignment: """ Represents an alignment file and its associated sample. Attributes: sample (str): The name or identifier of the sample. aligner (str): The name of the aligner used to generate the alignment file. format (str): The format of the alignment file (e.g., SAM, BAM). alignment (File): A File object representing the path to the alignment file. """ sample: str aligner: str format: str | None = None alignment: File | None = None def get\_alignment\_fname(self): return f"{self.sample}\_{self.aligner}\_aligned.{self.format}" # {{/docs-fragment dataclasses}} # ## Tasks # # We define a series of tasks to perform the following operations: # 1. Fetch assets from remote URLs # 2. Perform quality filtering and preprocessing using FastP # 3. Generate Bowtie2 index files from a reference genome # 4. Perform alignment using Bowtie2 on a filtered sample # # The first task fetches the reference genome and sequencing reads. It is cached # so that re-runs skip the download step. # {{docs-fragment fetch\_assets}} @fetch\_env.task async def fetch\_assets( ref\_url: str, read\_urls: List\[str\] ) -> tuple\[Reference, List\[Reads\]\]: """ Fetch assets from remote URLs. """ # Download reference genome ref\_dir = Path("/tmp/reference\_genome") ref\_dir.mkdir(exist\_ok=True, parents=True) ref = fetch\_file(ref\_url, str(ref\_dir)) ref\_obj = Reference( ref\_name=ref.name, ref\_dir=await Dir.from\_local(str(ref\_dir)), ) # Download sequencing reads dl\_loc = Path("/tmp/reads") dl\_loc.mkdir(exist\_ok=True, parents=True) samples: dict\[str, Reads\] = {} for url in read\_urls: fp = fetch\_file(url, str(dl\_loc)) sample = fp.stem.split("\_")\[0\] if sample not in samples: samples\[sample\] = Reads(sample=sample) if ".fastq.gz" in fp.name or "fasta" in fp.name: mate = fp.name.strip(".fastq.gz").strip(".filt").split("\_")\[-1\] if "1" in mate: samples\[sample\].read1 = await File.from\_local(str(fp)) elif "2" in mate: samples\[sample\].read2 = await File.from\_local(str(fp)) return ref\_obj, list(samples.values()) # {{/docs-fragment fetch\_assets}} # The second task performs quality filtering and preprocessing using FastP on a Reads object. # FastP is a performant tool for removing duplicate or low-quality reads. We increase # the memory request for this task so FastP can efficiently process reads from larger files. # {{docs-fragment pyfastp}} @fastp\_env.task async def pyfastp(rs: Reads) -> Reads: """ Perform quality filtering and preprocessing using Fastp on a Reads object. Args: rs (Reads): A Reads object containing raw sequencing data to be processed. Returns: Reads: A Reads object representing the filtered and preprocessed data. """ ldir = Path(tempfile.mkdtemp()) samp = Reads(rs.sample) o1, o2 = samp.get\_read\_fnames() o1p = ldir / o1 o2p = ldir / o2 assert rs.read1 is not None and rs.read2 is not None r1 = await rs.read1.download() r2 = await rs.read2.download() cmd = \[\ "fastp",\ "-i", str(r1),\ "-I", str(r2),\ "-o", str(o1p),\ "-O", str(o2p),\ \] subprocess.run(cmd, check=True) samp.read1 = await File.from\_local(str(o1p)) samp.read2 = await File.from\_local(str(o2p)) return samp # {{/docs-fragment pyfastp}} # Next, we define a task to generate Bowtie2 index files from a reference genome. As the index # for a given tool and reference seldom changes, we cache this task. # {{docs-fragment bowtie2\_index}} @index\_env.task async def bowtie2\_index(ref: Reference) -> Reference: """ Generate Bowtie2 index files from a reference genome. Args: ref (Reference): A Reference object representing the reference genome. Returns: Reference: The same reference object with the index\_name and indexed\_with attributes set. """ ref\_dir = await ref.ref\_dir.download() idx\_name = "bt2\_idx" cmd = \[\ "bowtie2-build",\ str(Path(str(ref\_dir)) / ref.ref\_name),\ str(Path(str(ref\_dir)) / idx\_name),\ \] subprocess.run(cmd, check=True) return Reference( ref.ref\_name, await Dir.from\_local(str(ref\_dir)), idx\_name, "bowtie2", ) # {{/docs-fragment bowtie2\_index}} # The next task performs paired-end alignment using Bowtie 2 on a single sample. # {{docs-fragment bowtie2\_align}} @align\_env.task async def bowtie2\_align\_paired\_reads(idx: Reference, fs: Reads) -> Alignment: """ Perform paired-end alignment using Bowtie 2 on a filtered sample. Args: idx (Reference): A Reference object containing the Bowtie 2 index. fs (Reads): A filtered Reads object containing sample data to be aligned. Returns: Alignment: An Alignment object representing the alignment result. """ assert idx.indexed\_with == "bowtie2", "Reference index must be generated with bowtie2" assert idx.index\_name is not None assert fs.read1 is not None and fs.read2 is not None ref\_dir = await idx.ref\_dir.download() r1 = await fs.read1.download() r2 = await fs.read2.download() ldir = Path(tempfile.mkdtemp()) alignment = Alignment(fs.sample, "bowtie2", "sam") al = ldir / alignment.get\_alignment\_fname() cmd = \[\ "bowtie2",\ "-x", str(Path(str(ref\_dir)) / idx.index\_name),\ "-1", str(r1),\ "-2", str(r2),\ "-S", str(al),\ \] subprocess.run(cmd, check=True) alignment.alignment = await File.from\_local(str(al)) return alignment # {{/docs-fragment bowtie2\_align}} # In place of the v1 \`@dynamic\` workflow, we use a plain async task with \`asyncio.gather\` # to run alignments for all samples in parallel. @base\_env.task async def bowtie2\_align\_samples( idx: Reference, samples: List\[Reads\] ) -> List\[Alignment\]: """ Process samples through bowtie2 in parallel. Args: idx (Reference): A Reference object containing the Bowtie 2 index. samples (List\[Reads\]): A list of Reads objects to be aligned. Returns: List\[Alignment\]: A list of Alignment objects representing the alignment results. """ tasks = \[bowtie2\_align\_paired\_reads(idx=idx, fs=sample) for sample in samples\] return list(await asyncio.gather(\*tasks)) # ## End-to-End Workflow # # We tie everything together in a final task that fetches assets, filters them, generates # an index, and aligns the samples. In place of the v1 \`@workflow\`, we use a top-level # \`@base\_env.task\`. Parallelism across samples is achieved with \`asyncio.gather\`. # {{docs-fragment workflow}} @base\_env.task async def alignment\_wf() -> List\[Alignment\]: # Prepare raw samples from remote URLs ref, samples = await fetch\_assets( ref\_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38\_short.fasta", read\_urls=\[\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_1.fastq.gz",\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_2.fastq.gz",\ \], ) # Filter all samples in parallel filtered\_samples = list( await asyncio.gather(\*\[pyfastp(rs=s) for s in samples\]) ) # Generate a bowtie2 index or load it from cache bowtie2\_idx = await bowtie2\_index(ref=ref) # Generate alignments using bowtie2 sams = await bowtie2\_align\_samples(idx=bowtie2\_idx, samples=filtered\_samples) return sams # {{/docs-fragment workflow}} # You can now run the workflow using the command in the dropdown at the top of the page! if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(alignment\_wf) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic\_alignment/genomic\_alignment.py\* ## Build the Bowtie 2 index A reference index rarely changes, so this task is cached. \`\`\` # # Genomic Alignment # # This tutorial demonstrates how to use Flyte to build a workflow that # performs genomic alignment on sequencing data. The workflow takes as input # a reference genome and raw sequencing data, performs quality filtering and # preprocessing on the raw data, generates an index for the reference genome, # and aligns the filtered data to the reference genome using the Bowtie 2 aligner. # {{run-on-union}} # The tutorial is divided into the following sections: # 1. Define the container image # 2. Define the data classes # 3. Define the tasks # 4. Define the workflow # /// script # requires-python = "3.12" # dependencies = \[\ # "flyte",\ # "requests",\ # \] # main = "alignment\_wf" # params = "" # /// import asyncio import subprocess import tempfile from dataclasses import dataclass from pathlib import Path from typing import List import requests import flyte from flyte.io import Dir, File # ## Defining a Container Image # # We define a custom container image using \`flyte.Image\`. Since we need bioinformatics # tools — \`fastp\` for quality filtering and \`bowtie2\` for alignment — we install them # via apt. This approach replaces the v1 \`ImageSpec\` with conda channels. # {{docs-fragment image}} main\_img = ( flyte.Image.from\_uv\_script( \_\_file\_\_, name="alignment-tutorial", ) .with\_apt\_packages("fastp", "bowtie2") ) # {{/docs-fragment image}} # We define per-task environments with different resource requirements, then a # top-level \`base\_env\` that declares all of them as dependencies (required because # \`alignment\_wf\` and \`bowtie2\_align\_samples\` call tasks that live in those environments). # {{docs-fragment envs}} fetch\_env = flyte.TaskEnvironment( name="alignment-tutorial-fetch", image=main\_img, cache="auto", ) fastp\_env = flyte.TaskEnvironment( name="alignment-tutorial-fastp", image=main\_img, resources=flyte.Resources(memory="2Gi"), ) index\_env = flyte.TaskEnvironment( name="alignment-tutorial-index", image=main\_img, resources=flyte.Resources(memory="10Gi"), cache="auto", ) align\_env = flyte.TaskEnvironment( name="alignment-tutorial-align", image=main\_img, resources=flyte.Resources(cpu=2, memory="10Gi"), ) base\_env = flyte.TaskEnvironment( name="alignment-tutorial", image=main\_img, depends\_on=\[fetch\_env, fastp\_env, index\_env, align\_env\], ) # {{/docs-fragment envs}} # ## Defining Data Classes # # We define three data classes to represent the reference genome, sequencing reads, # and alignment results. We'll first define a convenience function to download files, # which we'll use within the fetch task to materialize assets from their remote locations. def fetch\_file(url: str, local\_dir: str) -> Path: """ Downloads a file from the specified URL. Args: url (str): The URL of the file to download. local\_dir (str): The directory where you would like this file saved. Returns: Path: The local path to the file. Raises: requests.HTTPError: If an HTTP error occurs while downloading the file. """ url\_parts = url.split("/") fname = url\_parts\[-1\] local\_path = Path(local\_dir) / fname response = requests.get(url) with open(local\_path, "wb") as file: file.write(response.content) return local\_path # Reference genomes are used extensively throughout bioinformatics workflows. We define a # \`Reference\` data class to represent a reference genome and its associated index files. # {{docs-fragment dataclasses}} @dataclass class Reference: """ Represents a reference FASTA and associated index files. Attributes: ref\_name (str): Name or identifier of the reference file. ref\_dir (Dir): Directory containing the reference and any index files. index\_name (str): Index string to pass to tools requiring it. indexed\_with (str): Name of tool used to create the index. """ ref\_name: str ref\_dir: Dir index\_name: str | None = None indexed\_with: str | None = None # Sequencing reads are the raw data generated from a sequencing experiment. @dataclass class Reads: """ Represents a sequencing reads sample via its associated FastQ files. Attributes: sample (str): The name or identifier of the raw sequencing sample. read1 (File): A File object representing the path to the raw R1 read file. read2 (File): A File object representing the path to the raw R2 read file. """ sample: str read1: File | None = None read2: File | None = None def get\_read\_fnames(self): return ( f"{self.sample}\_1.fastq.gz", f"{self.sample}\_2.fastq.gz", ) # Finally, we define an \`Alignment\` data class to represent an alignment file. @dataclass class Alignment: """ Represents an alignment file and its associated sample. Attributes: sample (str): The name or identifier of the sample. aligner (str): The name of the aligner used to generate the alignment file. format (str): The format of the alignment file (e.g., SAM, BAM). alignment (File): A File object representing the path to the alignment file. """ sample: str aligner: str format: str | None = None alignment: File | None = None def get\_alignment\_fname(self): return f"{self.sample}\_{self.aligner}\_aligned.{self.format}" # {{/docs-fragment dataclasses}} # ## Tasks # # We define a series of tasks to perform the following operations: # 1. Fetch assets from remote URLs # 2. Perform quality filtering and preprocessing using FastP # 3. Generate Bowtie2 index files from a reference genome # 4. Perform alignment using Bowtie2 on a filtered sample # # The first task fetches the reference genome and sequencing reads. It is cached # so that re-runs skip the download step. # {{docs-fragment fetch\_assets}} @fetch\_env.task async def fetch\_assets( ref\_url: str, read\_urls: List\[str\] ) -> tuple\[Reference, List\[Reads\]\]: """ Fetch assets from remote URLs. """ # Download reference genome ref\_dir = Path("/tmp/reference\_genome") ref\_dir.mkdir(exist\_ok=True, parents=True) ref = fetch\_file(ref\_url, str(ref\_dir)) ref\_obj = Reference( ref\_name=ref.name, ref\_dir=await Dir.from\_local(str(ref\_dir)), ) # Download sequencing reads dl\_loc = Path("/tmp/reads") dl\_loc.mkdir(exist\_ok=True, parents=True) samples: dict\[str, Reads\] = {} for url in read\_urls: fp = fetch\_file(url, str(dl\_loc)) sample = fp.stem.split("\_")\[0\] if sample not in samples: samples\[sample\] = Reads(sample=sample) if ".fastq.gz" in fp.name or "fasta" in fp.name: mate = fp.name.strip(".fastq.gz").strip(".filt").split("\_")\[-1\] if "1" in mate: samples\[sample\].read1 = await File.from\_local(str(fp)) elif "2" in mate: samples\[sample\].read2 = await File.from\_local(str(fp)) return ref\_obj, list(samples.values()) # {{/docs-fragment fetch\_assets}} # The second task performs quality filtering and preprocessing using FastP on a Reads object. # FastP is a performant tool for removing duplicate or low-quality reads. We increase # the memory request for this task so FastP can efficiently process reads from larger files. # {{docs-fragment pyfastp}} @fastp\_env.task async def pyfastp(rs: Reads) -> Reads: """ Perform quality filtering and preprocessing using Fastp on a Reads object. Args: rs (Reads): A Reads object containing raw sequencing data to be processed. Returns: Reads: A Reads object representing the filtered and preprocessed data. """ ldir = Path(tempfile.mkdtemp()) samp = Reads(rs.sample) o1, o2 = samp.get\_read\_fnames() o1p = ldir / o1 o2p = ldir / o2 assert rs.read1 is not None and rs.read2 is not None r1 = await rs.read1.download() r2 = await rs.read2.download() cmd = \[\ "fastp",\ "-i", str(r1),\ "-I", str(r2),\ "-o", str(o1p),\ "-O", str(o2p),\ \] subprocess.run(cmd, check=True) samp.read1 = await File.from\_local(str(o1p)) samp.read2 = await File.from\_local(str(o2p)) return samp # {{/docs-fragment pyfastp}} # Next, we define a task to generate Bowtie2 index files from a reference genome. As the index # for a given tool and reference seldom changes, we cache this task. # {{docs-fragment bowtie2\_index}} @index\_env.task async def bowtie2\_index(ref: Reference) -> Reference: """ Generate Bowtie2 index files from a reference genome. Args: ref (Reference): A Reference object representing the reference genome. Returns: Reference: The same reference object with the index\_name and indexed\_with attributes set. """ ref\_dir = await ref.ref\_dir.download() idx\_name = "bt2\_idx" cmd = \[\ "bowtie2-build",\ str(Path(str(ref\_dir)) / ref.ref\_name),\ str(Path(str(ref\_dir)) / idx\_name),\ \] subprocess.run(cmd, check=True) return Reference( ref.ref\_name, await Dir.from\_local(str(ref\_dir)), idx\_name, "bowtie2", ) # {{/docs-fragment bowtie2\_index}} # The next task performs paired-end alignment using Bowtie 2 on a single sample. # {{docs-fragment bowtie2\_align}} @align\_env.task async def bowtie2\_align\_paired\_reads(idx: Reference, fs: Reads) -> Alignment: """ Perform paired-end alignment using Bowtie 2 on a filtered sample. Args: idx (Reference): A Reference object containing the Bowtie 2 index. fs (Reads): A filtered Reads object containing sample data to be aligned. Returns: Alignment: An Alignment object representing the alignment result. """ assert idx.indexed\_with == "bowtie2", "Reference index must be generated with bowtie2" assert idx.index\_name is not None assert fs.read1 is not None and fs.read2 is not None ref\_dir = await idx.ref\_dir.download() r1 = await fs.read1.download() r2 = await fs.read2.download() ldir = Path(tempfile.mkdtemp()) alignment = Alignment(fs.sample, "bowtie2", "sam") al = ldir / alignment.get\_alignment\_fname() cmd = \[\ "bowtie2",\ "-x", str(Path(str(ref\_dir)) / idx.index\_name),\ "-1", str(r1),\ "-2", str(r2),\ "-S", str(al),\ \] subprocess.run(cmd, check=True) alignment.alignment = await File.from\_local(str(al)) return alignment # {{/docs-fragment bowtie2\_align}} # In place of the v1 \`@dynamic\` workflow, we use a plain async task with \`asyncio.gather\` # to run alignments for all samples in parallel. @base\_env.task async def bowtie2\_align\_samples( idx: Reference, samples: List\[Reads\] ) -> List\[Alignment\]: """ Process samples through bowtie2 in parallel. Args: idx (Reference): A Reference object containing the Bowtie 2 index. samples (List\[Reads\]): A list of Reads objects to be aligned. Returns: List\[Alignment\]: A list of Alignment objects representing the alignment results. """ tasks = \[bowtie2\_align\_paired\_reads(idx=idx, fs=sample) for sample in samples\] return list(await asyncio.gather(\*tasks)) # ## End-to-End Workflow # # We tie everything together in a final task that fetches assets, filters them, generates # an index, and aligns the samples. In place of the v1 \`@workflow\`, we use a top-level # \`@base\_env.task\`. Parallelism across samples is achieved with \`asyncio.gather\`. # {{docs-fragment workflow}} @base\_env.task async def alignment\_wf() -> List\[Alignment\]: # Prepare raw samples from remote URLs ref, samples = await fetch\_assets( ref\_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38\_short.fasta", read\_urls=\[\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_1.fastq.gz",\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_2.fastq.gz",\ \], ) # Filter all samples in parallel filtered\_samples = list( await asyncio.gather(\*\[pyfastp(rs=s) for s in samples\]) ) # Generate a bowtie2 index or load it from cache bowtie2\_idx = await bowtie2\_index(ref=ref) # Generate alignments using bowtie2 sams = await bowtie2\_align\_samples(idx=bowtie2\_idx, samples=filtered\_samples) return sams # {{/docs-fragment workflow}} # You can now run the workflow using the command in the dropdown at the top of the page! if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(alignment\_wf) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic\_alignment/genomic\_alignment.py\* ## Align reads Each sample is aligned to the indexed reference with Bowtie 2, producing a SAM file. \`\`\` # # Genomic Alignment # # This tutorial demonstrates how to use Flyte to build a workflow that # performs genomic alignment on sequencing data. The workflow takes as input # a reference genome and raw sequencing data, performs quality filtering and # preprocessing on the raw data, generates an index for the reference genome, # and aligns the filtered data to the reference genome using the Bowtie 2 aligner. # {{run-on-union}} # The tutorial is divided into the following sections: # 1. Define the container image # 2. Define the data classes # 3. Define the tasks # 4. Define the workflow # /// script # requires-python = "3.12" # dependencies = \[\ # "flyte",\ # "requests",\ # \] # main = "alignment\_wf" # params = "" # /// import asyncio import subprocess import tempfile from dataclasses import dataclass from pathlib import Path from typing import List import requests import flyte from flyte.io import Dir, File # ## Defining a Container Image # # We define a custom container image using \`flyte.Image\`. Since we need bioinformatics # tools — \`fastp\` for quality filtering and \`bowtie2\` for alignment — we install them # via apt. This approach replaces the v1 \`ImageSpec\` with conda channels. # {{docs-fragment image}} main\_img = ( flyte.Image.from\_uv\_script( \_\_file\_\_, name="alignment-tutorial", ) .with\_apt\_packages("fastp", "bowtie2") ) # {{/docs-fragment image}} # We define per-task environments with different resource requirements, then a # top-level \`base\_env\` that declares all of them as dependencies (required because # \`alignment\_wf\` and \`bowtie2\_align\_samples\` call tasks that live in those environments). # {{docs-fragment envs}} fetch\_env = flyte.TaskEnvironment( name="alignment-tutorial-fetch", image=main\_img, cache="auto", ) fastp\_env = flyte.TaskEnvironment( name="alignment-tutorial-fastp", image=main\_img, resources=flyte.Resources(memory="2Gi"), ) index\_env = flyte.TaskEnvironment( name="alignment-tutorial-index", image=main\_img, resources=flyte.Resources(memory="10Gi"), cache="auto", ) align\_env = flyte.TaskEnvironment( name="alignment-tutorial-align", image=main\_img, resources=flyte.Resources(cpu=2, memory="10Gi"), ) base\_env = flyte.TaskEnvironment( name="alignment-tutorial", image=main\_img, depends\_on=\[fetch\_env, fastp\_env, index\_env, align\_env\], ) # {{/docs-fragment envs}} # ## Defining Data Classes # # We define three data classes to represent the reference genome, sequencing reads, # and alignment results. We'll first define a convenience function to download files, # which we'll use within the fetch task to materialize assets from their remote locations. def fetch\_file(url: str, local\_dir: str) -> Path: """ Downloads a file from the specified URL. Args: url (str): The URL of the file to download. local\_dir (str): The directory where you would like this file saved. Returns: Path: The local path to the file. Raises: requests.HTTPError: If an HTTP error occurs while downloading the file. """ url\_parts = url.split("/") fname = url\_parts\[-1\] local\_path = Path(local\_dir) / fname response = requests.get(url) with open(local\_path, "wb") as file: file.write(response.content) return local\_path # Reference genomes are used extensively throughout bioinformatics workflows. We define a # \`Reference\` data class to represent a reference genome and its associated index files. # {{docs-fragment dataclasses}} @dataclass class Reference: """ Represents a reference FASTA and associated index files. Attributes: ref\_name (str): Name or identifier of the reference file. ref\_dir (Dir): Directory containing the reference and any index files. index\_name (str): Index string to pass to tools requiring it. indexed\_with (str): Name of tool used to create the index. """ ref\_name: str ref\_dir: Dir index\_name: str | None = None indexed\_with: str | None = None # Sequencing reads are the raw data generated from a sequencing experiment. @dataclass class Reads: """ Represents a sequencing reads sample via its associated FastQ files. Attributes: sample (str): The name or identifier of the raw sequencing sample. read1 (File): A File object representing the path to the raw R1 read file. read2 (File): A File object representing the path to the raw R2 read file. """ sample: str read1: File | None = None read2: File | None = None def get\_read\_fnames(self): return ( f"{self.sample}\_1.fastq.gz", f"{self.sample}\_2.fastq.gz", ) # Finally, we define an \`Alignment\` data class to represent an alignment file. @dataclass class Alignment: """ Represents an alignment file and its associated sample. Attributes: sample (str): The name or identifier of the sample. aligner (str): The name of the aligner used to generate the alignment file. format (str): The format of the alignment file (e.g., SAM, BAM). alignment (File): A File object representing the path to the alignment file. """ sample: str aligner: str format: str | None = None alignment: File | None = None def get\_alignment\_fname(self): return f"{self.sample}\_{self.aligner}\_aligned.{self.format}" # {{/docs-fragment dataclasses}} # ## Tasks # # We define a series of tasks to perform the following operations: # 1. Fetch assets from remote URLs # 2. Perform quality filtering and preprocessing using FastP # 3. Generate Bowtie2 index files from a reference genome # 4. Perform alignment using Bowtie2 on a filtered sample # # The first task fetches the reference genome and sequencing reads. It is cached # so that re-runs skip the download step. # {{docs-fragment fetch\_assets}} @fetch\_env.task async def fetch\_assets( ref\_url: str, read\_urls: List\[str\] ) -> tuple\[Reference, List\[Reads\]\]: """ Fetch assets from remote URLs. """ # Download reference genome ref\_dir = Path("/tmp/reference\_genome") ref\_dir.mkdir(exist\_ok=True, parents=True) ref = fetch\_file(ref\_url, str(ref\_dir)) ref\_obj = Reference( ref\_name=ref.name, ref\_dir=await Dir.from\_local(str(ref\_dir)), ) # Download sequencing reads dl\_loc = Path("/tmp/reads") dl\_loc.mkdir(exist\_ok=True, parents=True) samples: dict\[str, Reads\] = {} for url in read\_urls: fp = fetch\_file(url, str(dl\_loc)) sample = fp.stem.split("\_")\[0\] if sample not in samples: samples\[sample\] = Reads(sample=sample) if ".fastq.gz" in fp.name or "fasta" in fp.name: mate = fp.name.strip(".fastq.gz").strip(".filt").split("\_")\[-1\] if "1" in mate: samples\[sample\].read1 = await File.from\_local(str(fp)) elif "2" in mate: samples\[sample\].read2 = await File.from\_local(str(fp)) return ref\_obj, list(samples.values()) # {{/docs-fragment fetch\_assets}} # The second task performs quality filtering and preprocessing using FastP on a Reads object. # FastP is a performant tool for removing duplicate or low-quality reads. We increase # the memory request for this task so FastP can efficiently process reads from larger files. # {{docs-fragment pyfastp}} @fastp\_env.task async def pyfastp(rs: Reads) -> Reads: """ Perform quality filtering and preprocessing using Fastp on a Reads object. Args: rs (Reads): A Reads object containing raw sequencing data to be processed. Returns: Reads: A Reads object representing the filtered and preprocessed data. """ ldir = Path(tempfile.mkdtemp()) samp = Reads(rs.sample) o1, o2 = samp.get\_read\_fnames() o1p = ldir / o1 o2p = ldir / o2 assert rs.read1 is not None and rs.read2 is not None r1 = await rs.read1.download() r2 = await rs.read2.download() cmd = \[\ "fastp",\ "-i", str(r1),\ "-I", str(r2),\ "-o", str(o1p),\ "-O", str(o2p),\ \] subprocess.run(cmd, check=True) samp.read1 = await File.from\_local(str(o1p)) samp.read2 = await File.from\_local(str(o2p)) return samp # {{/docs-fragment pyfastp}} # Next, we define a task to generate Bowtie2 index files from a reference genome. As the index # for a given tool and reference seldom changes, we cache this task. # {{docs-fragment bowtie2\_index}} @index\_env.task async def bowtie2\_index(ref: Reference) -> Reference: """ Generate Bowtie2 index files from a reference genome. Args: ref (Reference): A Reference object representing the reference genome. Returns: Reference: The same reference object with the index\_name and indexed\_with attributes set. """ ref\_dir = await ref.ref\_dir.download() idx\_name = "bt2\_idx" cmd = \[\ "bowtie2-build",\ str(Path(str(ref\_dir)) / ref.ref\_name),\ str(Path(str(ref\_dir)) / idx\_name),\ \] subprocess.run(cmd, check=True) return Reference( ref.ref\_name, await Dir.from\_local(str(ref\_dir)), idx\_name, "bowtie2", ) # {{/docs-fragment bowtie2\_index}} # The next task performs paired-end alignment using Bowtie 2 on a single sample. # {{docs-fragment bowtie2\_align}} @align\_env.task async def bowtie2\_align\_paired\_reads(idx: Reference, fs: Reads) -> Alignment: """ Perform paired-end alignment using Bowtie 2 on a filtered sample. Args: idx (Reference): A Reference object containing the Bowtie 2 index. fs (Reads): A filtered Reads object containing sample data to be aligned. Returns: Alignment: An Alignment object representing the alignment result. """ assert idx.indexed\_with == "bowtie2", "Reference index must be generated with bowtie2" assert idx.index\_name is not None assert fs.read1 is not None and fs.read2 is not None ref\_dir = await idx.ref\_dir.download() r1 = await fs.read1.download() r2 = await fs.read2.download() ldir = Path(tempfile.mkdtemp()) alignment = Alignment(fs.sample, "bowtie2", "sam") al = ldir / alignment.get\_alignment\_fname() cmd = \[\ "bowtie2",\ "-x", str(Path(str(ref\_dir)) / idx.index\_name),\ "-1", str(r1),\ "-2", str(r2),\ "-S", str(al),\ \] subprocess.run(cmd, check=True) alignment.alignment = await File.from\_local(str(al)) return alignment # {{/docs-fragment bowtie2\_align}} # In place of the v1 \`@dynamic\` workflow, we use a plain async task with \`asyncio.gather\` # to run alignments for all samples in parallel. @base\_env.task async def bowtie2\_align\_samples( idx: Reference, samples: List\[Reads\] ) -> List\[Alignment\]: """ Process samples through bowtie2 in parallel. Args: idx (Reference): A Reference object containing the Bowtie 2 index. samples (List\[Reads\]): A list of Reads objects to be aligned. Returns: List\[Alignment\]: A list of Alignment objects representing the alignment results. """ tasks = \[bowtie2\_align\_paired\_reads(idx=idx, fs=sample) for sample in samples\] return list(await asyncio.gather(\*tasks)) # ## End-to-End Workflow # # We tie everything together in a final task that fetches assets, filters them, generates # an index, and aligns the samples. In place of the v1 \`@workflow\`, we use a top-level # \`@base\_env.task\`. Parallelism across samples is achieved with \`asyncio.gather\`. # {{docs-fragment workflow}} @base\_env.task async def alignment\_wf() -> List\[Alignment\]: # Prepare raw samples from remote URLs ref, samples = await fetch\_assets( ref\_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38\_short.fasta", read\_urls=\[\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_1.fastq.gz",\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_2.fastq.gz",\ \], ) # Filter all samples in parallel filtered\_samples = list( await asyncio.gather(\*\[pyfastp(rs=s) for s in samples\]) ) # Generate a bowtie2 index or load it from cache bowtie2\_idx = await bowtie2\_index(ref=ref) # Generate alignments using bowtie2 sams = await bowtie2\_align\_samples(idx=bowtie2\_idx, samples=filtered\_samples) return sams # {{/docs-fragment workflow}} # You can now run the workflow using the command in the dropdown at the top of the page! if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(alignment\_wf) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic\_alignment/genomic\_alignment.py\* ## Orchestrate the workflow The top-level task fetches the assets, filters every sample in parallel, builds the index, and aligns all samples. Parallelism across samples is achieved with \`asyncio.gather\` rather than a separate \`@dynamic\` decorator. \`\`\` # # Genomic Alignment # # This tutorial demonstrates how to use Flyte to build a workflow that # performs genomic alignment on sequencing data. The workflow takes as input # a reference genome and raw sequencing data, performs quality filtering and # preprocessing on the raw data, generates an index for the reference genome, # and aligns the filtered data to the reference genome using the Bowtie 2 aligner. # {{run-on-union}} # The tutorial is divided into the following sections: # 1. Define the container image # 2. Define the data classes # 3. Define the tasks # 4. Define the workflow # /// script # requires-python = "3.12" # dependencies = \[\ # "flyte",\ # "requests",\ # \] # main = "alignment\_wf" # params = "" # /// import asyncio import subprocess import tempfile from dataclasses import dataclass from pathlib import Path from typing import List import requests import flyte from flyte.io import Dir, File # ## Defining a Container Image # # We define a custom container image using \`flyte.Image\`. Since we need bioinformatics # tools — \`fastp\` for quality filtering and \`bowtie2\` for alignment — we install them # via apt. This approach replaces the v1 \`ImageSpec\` with conda channels. # {{docs-fragment image}} main\_img = ( flyte.Image.from\_uv\_script( \_\_file\_\_, name="alignment-tutorial", ) .with\_apt\_packages("fastp", "bowtie2") ) # {{/docs-fragment image}} # We define per-task environments with different resource requirements, then a # top-level \`base\_env\` that declares all of them as dependencies (required because # \`alignment\_wf\` and \`bowtie2\_align\_samples\` call tasks that live in those environments). # {{docs-fragment envs}} fetch\_env = flyte.TaskEnvironment( name="alignment-tutorial-fetch", image=main\_img, cache="auto", ) fastp\_env = flyte.TaskEnvironment( name="alignment-tutorial-fastp", image=main\_img, resources=flyte.Resources(memory="2Gi"), ) index\_env = flyte.TaskEnvironment( name="alignment-tutorial-index", image=main\_img, resources=flyte.Resources(memory="10Gi"), cache="auto", ) align\_env = flyte.TaskEnvironment( name="alignment-tutorial-align", image=main\_img, resources=flyte.Resources(cpu=2, memory="10Gi"), ) base\_env = flyte.TaskEnvironment( name="alignment-tutorial", image=main\_img, depends\_on=\[fetch\_env, fastp\_env, index\_env, align\_env\], ) # {{/docs-fragment envs}} # ## Defining Data Classes # # We define three data classes to represent the reference genome, sequencing reads, # and alignment results. We'll first define a convenience function to download files, # which we'll use within the fetch task to materialize assets from their remote locations. def fetch\_file(url: str, local\_dir: str) -> Path: """ Downloads a file from the specified URL. Args: url (str): The URL of the file to download. local\_dir (str): The directory where you would like this file saved. Returns: Path: The local path to the file. Raises: requests.HTTPError: If an HTTP error occurs while downloading the file. """ url\_parts = url.split("/") fname = url\_parts\[-1\] local\_path = Path(local\_dir) / fname response = requests.get(url) with open(local\_path, "wb") as file: file.write(response.content) return local\_path # Reference genomes are used extensively throughout bioinformatics workflows. We define a # \`Reference\` data class to represent a reference genome and its associated index files. # {{docs-fragment dataclasses}} @dataclass class Reference: """ Represents a reference FASTA and associated index files. Attributes: ref\_name (str): Name or identifier of the reference file. ref\_dir (Dir): Directory containing the reference and any index files. index\_name (str): Index string to pass to tools requiring it. indexed\_with (str): Name of tool used to create the index. """ ref\_name: str ref\_dir: Dir index\_name: str | None = None indexed\_with: str | None = None # Sequencing reads are the raw data generated from a sequencing experiment. @dataclass class Reads: """ Represents a sequencing reads sample via its associated FastQ files. Attributes: sample (str): The name or identifier of the raw sequencing sample. read1 (File): A File object representing the path to the raw R1 read file. read2 (File): A File object representing the path to the raw R2 read file. """ sample: str read1: File | None = None read2: File | None = None def get\_read\_fnames(self): return ( f"{self.sample}\_1.fastq.gz", f"{self.sample}\_2.fastq.gz", ) # Finally, we define an \`Alignment\` data class to represent an alignment file. @dataclass class Alignment: """ Represents an alignment file and its associated sample. Attributes: sample (str): The name or identifier of the sample. aligner (str): The name of the aligner used to generate the alignment file. format (str): The format of the alignment file (e.g., SAM, BAM). alignment (File): A File object representing the path to the alignment file. """ sample: str aligner: str format: str | None = None alignment: File | None = None def get\_alignment\_fname(self): return f"{self.sample}\_{self.aligner}\_aligned.{self.format}" # {{/docs-fragment dataclasses}} # ## Tasks # # We define a series of tasks to perform the following operations: # 1. Fetch assets from remote URLs # 2. Perform quality filtering and preprocessing using FastP # 3. Generate Bowtie2 index files from a reference genome # 4. Perform alignment using Bowtie2 on a filtered sample # # The first task fetches the reference genome and sequencing reads. It is cached # so that re-runs skip the download step. # {{docs-fragment fetch\_assets}} @fetch\_env.task async def fetch\_assets( ref\_url: str, read\_urls: List\[str\] ) -> tuple\[Reference, List\[Reads\]\]: """ Fetch assets from remote URLs. """ # Download reference genome ref\_dir = Path("/tmp/reference\_genome") ref\_dir.mkdir(exist\_ok=True, parents=True) ref = fetch\_file(ref\_url, str(ref\_dir)) ref\_obj = Reference( ref\_name=ref.name, ref\_dir=await Dir.from\_local(str(ref\_dir)), ) # Download sequencing reads dl\_loc = Path("/tmp/reads") dl\_loc.mkdir(exist\_ok=True, parents=True) samples: dict\[str, Reads\] = {} for url in read\_urls: fp = fetch\_file(url, str(dl\_loc)) sample = fp.stem.split("\_")\[0\] if sample not in samples: samples\[sample\] = Reads(sample=sample) if ".fastq.gz" in fp.name or "fasta" in fp.name: mate = fp.name.strip(".fastq.gz").strip(".filt").split("\_")\[-1\] if "1" in mate: samples\[sample\].read1 = await File.from\_local(str(fp)) elif "2" in mate: samples\[sample\].read2 = await File.from\_local(str(fp)) return ref\_obj, list(samples.values()) # {{/docs-fragment fetch\_assets}} # The second task performs quality filtering and preprocessing using FastP on a Reads object. # FastP is a performant tool for removing duplicate or low-quality reads. We increase # the memory request for this task so FastP can efficiently process reads from larger files. # {{docs-fragment pyfastp}} @fastp\_env.task async def pyfastp(rs: Reads) -> Reads: """ Perform quality filtering and preprocessing using Fastp on a Reads object. Args: rs (Reads): A Reads object containing raw sequencing data to be processed. Returns: Reads: A Reads object representing the filtered and preprocessed data. """ ldir = Path(tempfile.mkdtemp()) samp = Reads(rs.sample) o1, o2 = samp.get\_read\_fnames() o1p = ldir / o1 o2p = ldir / o2 assert rs.read1 is not None and rs.read2 is not None r1 = await rs.read1.download() r2 = await rs.read2.download() cmd = \[\ "fastp",\ "-i", str(r1),\ "-I", str(r2),\ "-o", str(o1p),\ "-O", str(o2p),\ \] subprocess.run(cmd, check=True) samp.read1 = await File.from\_local(str(o1p)) samp.read2 = await File.from\_local(str(o2p)) return samp # {{/docs-fragment pyfastp}} # Next, we define a task to generate Bowtie2 index files from a reference genome. As the index # for a given tool and reference seldom changes, we cache this task. # {{docs-fragment bowtie2\_index}} @index\_env.task async def bowtie2\_index(ref: Reference) -> Reference: """ Generate Bowtie2 index files from a reference genome. Args: ref (Reference): A Reference object representing the reference genome. Returns: Reference: The same reference object with the index\_name and indexed\_with attributes set. """ ref\_dir = await ref.ref\_dir.download() idx\_name = "bt2\_idx" cmd = \[\ "bowtie2-build",\ str(Path(str(ref\_dir)) / ref.ref\_name),\ str(Path(str(ref\_dir)) / idx\_name),\ \] subprocess.run(cmd, check=True) return Reference( ref.ref\_name, await Dir.from\_local(str(ref\_dir)), idx\_name, "bowtie2", ) # {{/docs-fragment bowtie2\_index}} # The next task performs paired-end alignment using Bowtie 2 on a single sample. # {{docs-fragment bowtie2\_align}} @align\_env.task async def bowtie2\_align\_paired\_reads(idx: Reference, fs: Reads) -> Alignment: """ Perform paired-end alignment using Bowtie 2 on a filtered sample. Args: idx (Reference): A Reference object containing the Bowtie 2 index. fs (Reads): A filtered Reads object containing sample data to be aligned. Returns: Alignment: An Alignment object representing the alignment result. """ assert idx.indexed\_with == "bowtie2", "Reference index must be generated with bowtie2" assert idx.index\_name is not None assert fs.read1 is not None and fs.read2 is not None ref\_dir = await idx.ref\_dir.download() r1 = await fs.read1.download() r2 = await fs.read2.download() ldir = Path(tempfile.mkdtemp()) alignment = Alignment(fs.sample, "bowtie2", "sam") al = ldir / alignment.get\_alignment\_fname() cmd = \[\ "bowtie2",\ "-x", str(Path(str(ref\_dir)) / idx.index\_name),\ "-1", str(r1),\ "-2", str(r2),\ "-S", str(al),\ \] subprocess.run(cmd, check=True) alignment.alignment = await File.from\_local(str(al)) return alignment # {{/docs-fragment bowtie2\_align}} # In place of the v1 \`@dynamic\` workflow, we use a plain async task with \`asyncio.gather\` # to run alignments for all samples in parallel. @base\_env.task async def bowtie2\_align\_samples( idx: Reference, samples: List\[Reads\] ) -> List\[Alignment\]: """ Process samples through bowtie2 in parallel. Args: idx (Reference): A Reference object containing the Bowtie 2 index. samples (List\[Reads\]): A list of Reads objects to be aligned. Returns: List\[Alignment\]: A list of Alignment objects representing the alignment results. """ tasks = \[bowtie2\_align\_paired\_reads(idx=idx, fs=sample) for sample in samples\] return list(await asyncio.gather(\*tasks)) # ## End-to-End Workflow # # We tie everything together in a final task that fetches assets, filters them, generates # an index, and aligns the samples. In place of the v1 \`@workflow\`, we use a top-level # \`@base\_env.task\`. Parallelism across samples is achieved with \`asyncio.gather\`. # {{docs-fragment workflow}} @base\_env.task async def alignment\_wf() -> List\[Alignment\]: # Prepare raw samples from remote URLs ref, samples = await fetch\_assets( ref\_url="https://github.com/unionai-oss/unionbio/raw/main/tests/assets/references/GRCh38\_short.fasta", read\_urls=\[\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_1.fastq.gz",\ "https://github.com/unionai-oss/unionbio/raw/main/tests/assets/sequences/raw/ERR250683-tiny\_2.fastq.gz",\ \], ) # Filter all samples in parallel filtered\_samples = list( await asyncio.gather(\*\[pyfastp(rs=s) for s in samples\]) ) # Generate a bowtie2 index or load it from cache bowtie2\_idx = await bowtie2\_index(ref=ref) # Generate alignments using bowtie2 sams = await bowtie2\_align\_samples(idx=bowtie2\_idx, samples=filtered\_samples) return sams # {{/docs-fragment workflow}} # You can now run the workflow using the command in the dropdown at the top of the page! if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(alignment\_wf) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/genomic\_alignment/genomic\_alignment.py\* ## Run the workflow This example has no secrets or external API keys — it pulls public test data from GitHub. From the \[example directory\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/genomic\_alignment), run it as a \`uv\` script: \`\`\` cd v2/tutorials/genomic\_alignment uv run --script genomic\_alignment.py \`\`\` Or submit it with the Flyte CLI: \`\`\` flyte run genomic\_alignment.py alignment\_wf \`\`\` When the run completes, each returned \`Alignment\` points to a SAM file in blob storage that you can download from the run's outputs in the UI. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/biotech-healthcare/tumor-detection === # Brain tumor MRI classification > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/tumor\_detection). This tutorial builds a medical-imaging pipeline that classifies brain MRI scans into four categories — Glioma, Meningioma, No Tumor, and Pituitary — using a two-phase EfficientNet-B4 transfer-learning strategy. The pipeline downloads the dataset, trains on a GPU with fault-tolerant checkpointing, and renders training curves and a confusion matrix directly in the Flyte UI. The example is split into focused modules: - \`config.py\` — container image, task environments, and the \`TrainingConfig\` hyperparameters. - \`dataset.py\` — downloads the Hugging Face dataset, builds class-balanced data loaders. - \`model.py\` / \`training.py\` — the Lightning module and the two-phase training loop. - \`utils.py\` — plotting helpers for the report. - \`run.py\` — the three Flyte tasks and the pipeline driver. Flyte handles the production concerns: - \*\*Per-task resources\*\*: CPU for download/reporting, a GPU for training. - \*\*\`cache="auto"\`\*\* on dataset download and training, so reruns with the same data and config are free. - \*\*\`retries=3\`\*\* plus \*\*Flyte checkpointing\*\* on the training task so a preempted GPU job resumes from the last epoch. - \*\*Built-in reports\*\* to visualize metrics without separate dashboard infrastructure. ## Define the container image A single GPU-ready image is shared by all tasks. \`with\_source\_folder\` copies the local modules (\`dataset.py\`, \`model.py\`, etc.) into the image. \`\`\` """ Configuration for brain tumor MRI classification pipeline. Defines task environments, resource requirements, and training hyperparameters. """ import pathlib import flyte # {{docs-fragment image}} image = flyte.Image.from\_debian\_base( name="tumor\_detection\_gpu" ).with\_pip\_packages( "torch", "lightning", "torchvision", "timm", "pillow", "scikit-learn", "plotly", "numpy", "pandas", "torchmetrics", "datasets", "typing\_extensions", ).with\_source\_folder( pathlib.Path(\_\_file\_\_).parent, copy\_contents\_only=True, ) # {{/docs-fragment image}} # {{docs-fragment envs}} # Downloads raw MRI JPEG files — CPU only, no auth needed, result is cached dataset\_env = flyte.TaskEnvironment( name="tumor\_dataset", image=image, resources=flyte.Resources(cpu=2, memory="4Gi", disk="8Gi"), cache="auto", ) # GPU training — result is cached so re-running with the same data + config is free training\_env = flyte.TaskEnvironment( name="tumor\_gpu\_training", image=image, resources=flyte.Resources( cpu=8, memory="32Gi", gpu="T4:1", disk="100Gi", ), env\_vars={ "CUDA\_VISIBLE\_DEVICES": "0", "CUDA\_LAUNCH\_BLOCKING": "1", "TORCH\_CUDA\_MEMORY\_FRACTION": "1.0", "PYTORCH\_CUDA\_ALLOC\_CONF": "expandable\_segments:True", }, cache="auto", ) # Report generation — CPU only, reads training results and renders Union UI panels report\_env = flyte.TaskEnvironment( name="tumor\_report", image=image, resources=flyte.Resources(cpu=2, memory="4Gi"), ) # Pipeline driver — lightweight orchestrator that calls the three tasks above pipeline\_env = flyte.TaskEnvironment( name="tumor\_pipeline", image=image, resources=flyte.Resources(cpu=2, memory="4Gi"), depends\_on=\[dataset\_env, training\_env, report\_env\], ) # {{/docs-fragment envs}} class TrainingConfig: """Unified training configuration for brain tumor MRI classification.""" def \_\_init\_\_( self, image\_size: int = 380, num\_classes: int = 4, model\_name: str = "efficientnet\_b4", pretrained: bool = True, phase1\_epochs: int = 8, phase1\_lr: float = 1e-3, phase1\_freeze\_backbone: bool = True, phase2\_epochs: int = 25, phase2\_lr: float = 5e-5, batch\_size: int = 16, num\_workers: int = 0, val\_split: float = 0.2, weight\_decay: float = 1e-4, warmup\_steps: int = 200, focal\_gamma: float = 2.0, mixup\_alpha: float = 0.0, log\_interval: int = 50, ): self.image\_size = image\_size self.num\_classes = num\_classes self.model\_name = model\_name self.pretrained = pretrained self.phase1\_epochs = phase1\_epochs self.phase1\_lr = phase1\_lr self.phase1\_freeze\_backbone = phase1\_freeze\_backbone self.phase2\_epochs = phase2\_epochs self.phase2\_lr = phase2\_lr self.batch\_size = batch\_size self.num\_workers = num\_workers self.val\_split = val\_split self.weight\_decay = weight\_decay self.warmup\_steps = warmup\_steps self.focal\_gamma = focal\_gamma self.mixup\_alpha = mixup\_alpha self.log\_interval = log\_interval def to\_dict(self) -> dict: return self.\_\_dict\_\_ \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor\_detection/config.py\* ## Define the task environments Each stage declares the resources it needs. The lightweight \`pipeline\_env\` orchestrates the others via \`depends\_on\`. \`\`\` """ Configuration for brain tumor MRI classification pipeline. Defines task environments, resource requirements, and training hyperparameters. """ import pathlib import flyte # {{docs-fragment image}} image = flyte.Image.from\_debian\_base( name="tumor\_detection\_gpu" ).with\_pip\_packages( "torch", "lightning", "torchvision", "timm", "pillow", "scikit-learn", "plotly", "numpy", "pandas", "torchmetrics", "datasets", "typing\_extensions", ).with\_source\_folder( pathlib.Path(\_\_file\_\_).parent, copy\_contents\_only=True, ) # {{/docs-fragment image}} # {{docs-fragment envs}} # Downloads raw MRI JPEG files — CPU only, no auth needed, result is cached dataset\_env = flyte.TaskEnvironment( name="tumor\_dataset", image=image, resources=flyte.Resources(cpu=2, memory="4Gi", disk="8Gi"), cache="auto", ) # GPU training — result is cached so re-running with the same data + config is free training\_env = flyte.TaskEnvironment( name="tumor\_gpu\_training", image=image, resources=flyte.Resources( cpu=8, memory="32Gi", gpu="T4:1", disk="100Gi", ), env\_vars={ "CUDA\_VISIBLE\_DEVICES": "0", "CUDA\_LAUNCH\_BLOCKING": "1", "TORCH\_CUDA\_MEMORY\_FRACTION": "1.0", "PYTORCH\_CUDA\_ALLOC\_CONF": "expandable\_segments:True", }, cache="auto", ) # Report generation — CPU only, reads training results and renders Union UI panels report\_env = flyte.TaskEnvironment( name="tumor\_report", image=image, resources=flyte.Resources(cpu=2, memory="4Gi"), ) # Pipeline driver — lightweight orchestrator that calls the three tasks above pipeline\_env = flyte.TaskEnvironment( name="tumor\_pipeline", image=image, resources=flyte.Resources(cpu=2, memory="4Gi"), depends\_on=\[dataset\_env, training\_env, report\_env\], ) # {{/docs-fragment envs}} class TrainingConfig: """Unified training configuration for brain tumor MRI classification.""" def \_\_init\_\_( self, image\_size: int = 380, num\_classes: int = 4, model\_name: str = "efficientnet\_b4", pretrained: bool = True, phase1\_epochs: int = 8, phase1\_lr: float = 1e-3, phase1\_freeze\_backbone: bool = True, phase2\_epochs: int = 25, phase2\_lr: float = 5e-5, batch\_size: int = 16, num\_workers: int = 0, val\_split: float = 0.2, weight\_decay: float = 1e-4, warmup\_steps: int = 200, focal\_gamma: float = 2.0, mixup\_alpha: float = 0.0, log\_interval: int = 50, ): self.image\_size = image\_size self.num\_classes = num\_classes self.model\_name = model\_name self.pretrained = pretrained self.phase1\_epochs = phase1\_epochs self.phase1\_lr = phase1\_lr self.phase1\_freeze\_backbone = phase1\_freeze\_backbone self.phase2\_epochs = phase2\_epochs self.phase2\_lr = phase2\_lr self.batch\_size = batch\_size self.num\_workers = num\_workers self.val\_split = val\_split self.weight\_decay = weight\_decay self.warmup\_steps = warmup\_steps self.focal\_gamma = focal\_gamma self.mixup\_alpha = mixup\_alpha self.log\_interval = log\_interval def to\_dict(self) -> dict: return self.\_\_dict\_\_ \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor\_detection/config.py\* ## Configure training Hyperparameters are gathered in a single \`TrainingConfig\`, serialized to JSON, and passed into the training task so the exact configuration is captured alongside the run. \`\`\` """ Flyte/Union pipeline for brain tumor MRI classification. Three-task pipeline: 1. load\_dataset — download Brain Tumor MRI from Hugging Face, cache as Dir (CPU) 2. train\_model — two-phase EfficientNet-B4 training with focal loss (GPU) 3. create\_report — render training curves and confusion matrix in the Union UI (CPU) """ import json import flyte from flyte.io import Dir from config import TrainingConfig, dataset\_env, pipeline\_env, report\_env, training\_env from dataset import download\_tumor\_dataset # {{docs-fragment config}} TRAINING\_CONFIG = TrainingConfig( phase1\_epochs=8, phase2\_epochs=25, phase1\_lr=1e-3, phase2\_lr=5e-5, batch\_size=16, num\_workers=0, log\_interval=50, mixup\_alpha=0.0, image\_size=380, focal\_gamma=3.0, ) # {{/docs-fragment config}} # {{docs-fragment load\_dataset}} @dataset\_env.task async def load\_dataset() -> Dir: """ Download raw Brain Tumor MRI JPEG files from Hugging Face and cache as flyte.io.Dir. Runs once — result is reused on subsequent pipeline runs (cache="auto"). """ return await download\_tumor\_dataset() # {{/docs-fragment load\_dataset}} # {{docs-fragment train\_model}} @training\_env.task(retries=3) async def train\_model(dataset\_dir: Dir, config\_json: str) -> Dir: """ Download the raw dataset Dir, run two-phase training, and return training metrics and final predictions as a Dir for the report task. """ from pathlib import Path local\_dir = Path("/tmp/tumor\_local") local\_dir.mkdir(parents=True, exist\_ok=True) await dataset\_dir.download(local\_path=str(local\_dir)) from training import train\_tumor\_classifier config = TrainingConfig(\*\*json.loads(config\_json)) result = train\_tumor\_classifier(config=config, dataset\_path=str(local\_dir)) output\_dir = Path("/tmp/training\_results") output\_dir.mkdir(parents=True, exist\_ok=True) (output\_dir / "metrics.json").write\_text(json.dumps(result\["metrics"\])) (output\_dir / "predictions.json").write\_text(json.dumps({ "preds": result\["final\_preds"\], "targets": result\["final\_targets"\], })) return await Dir.from\_local(str(output\_dir)) # {{/docs-fragment train\_model}} # {{docs-fragment create\_report}} @report\_env.task(report=True) async def create\_report(results\_dir: Dir) -> None: """ Download training metrics and render loss/accuracy curves, confusion matrix, and per-class F1 chart in the Union UI report panel. """ import numpy as np from pathlib import Path from utils import create\_confusion\_matrix\_plot, create\_metrics\_plots, create\_per\_class\_f1\_plot local\_dir = Path("/tmp/tumor\_report") local\_dir.mkdir(parents=True, exist\_ok=True) await results\_dir.download(local\_path=str(local\_dir)) matches = list(local\_dir.glob("\*\*/metrics.json")) if not matches: raise RuntimeError(f"metrics.json not found under {local\_dir}") local\_path = matches\[0\].parent history = json.loads((local\_path / "metrics.json").read\_text()) predictions = json.loads((local\_path / "predictions.json").read\_text()) preds = np.array(predictions\["preds"\]) targets = np.array(predictions\["targets"\]) loss\_fig, acc\_fig = create\_metrics\_plots(history) cm\_fig = create\_confusion\_matrix\_plot(preds, targets) f1\_fig = create\_per\_class\_f1\_plot(preds, targets) combined\_html = ( acc\_fig.to\_html(include\_plotlyjs=True, full\_html=False) + loss\_fig.to\_html(include\_plotlyjs=False, full\_html=False) + cm\_fig.to\_html(include\_plotlyjs=False, full\_html=False) + f1\_fig.to\_html(include\_plotlyjs=False, full\_html=False) ) flyte.report.log(combined\_html, do\_flush=True) # {{/docs-fragment create\_report}} # {{docs-fragment pipeline}} @pipeline\_env.task async def tumor\_detection\_pipeline() -> None: """Orchestrate dataset loading, GPU training, and report generation.""" dataset\_dir = await load\_dataset() results\_dir = await train\_model( dataset\_dir=dataset\_dir, config\_json=json.dumps(TRAINING\_CONFIG.to\_dict()), ) await create\_report(results\_dir=results\_dir) # {{/docs-fragment pipeline}} if \_\_name\_\_ == "\_\_main\_\_": import pathlib flyte.init\_from\_config(root\_dir=pathlib.Path(\_\_file\_\_).parent) run = flyte.with\_runcontext().run(tumor\_detection\_pipeline) print(f"\\n✓ Pipeline submitted!") print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor\_detection/run.py\* ## Load the dataset The first task downloads the public \[Brain Tumor MRI dataset\](https://huggingface.co/datasets/AIOmarRehan/Brain\_Tumor\_MRI\_Dataset) from Hugging Face (no auth required) and stores it as a \`flyte.io.Dir\`. It's cached, so subsequent runs reuse it. \`\`\` """ Flyte/Union pipeline for brain tumor MRI classification. Three-task pipeline: 1. load\_dataset — download Brain Tumor MRI from Hugging Face, cache as Dir (CPU) 2. train\_model — two-phase EfficientNet-B4 training with focal loss (GPU) 3. create\_report — render training curves and confusion matrix in the Union UI (CPU) """ import json import flyte from flyte.io import Dir from config import TrainingConfig, dataset\_env, pipeline\_env, report\_env, training\_env from dataset import download\_tumor\_dataset # {{docs-fragment config}} TRAINING\_CONFIG = TrainingConfig( phase1\_epochs=8, phase2\_epochs=25, phase1\_lr=1e-3, phase2\_lr=5e-5, batch\_size=16, num\_workers=0, log\_interval=50, mixup\_alpha=0.0, image\_size=380, focal\_gamma=3.0, ) # {{/docs-fragment config}} # {{docs-fragment load\_dataset}} @dataset\_env.task async def load\_dataset() -> Dir: """ Download raw Brain Tumor MRI JPEG files from Hugging Face and cache as flyte.io.Dir. Runs once — result is reused on subsequent pipeline runs (cache="auto"). """ return await download\_tumor\_dataset() # {{/docs-fragment load\_dataset}} # {{docs-fragment train\_model}} @training\_env.task(retries=3) async def train\_model(dataset\_dir: Dir, config\_json: str) -> Dir: """ Download the raw dataset Dir, run two-phase training, and return training metrics and final predictions as a Dir for the report task. """ from pathlib import Path local\_dir = Path("/tmp/tumor\_local") local\_dir.mkdir(parents=True, exist\_ok=True) await dataset\_dir.download(local\_path=str(local\_dir)) from training import train\_tumor\_classifier config = TrainingConfig(\*\*json.loads(config\_json)) result = train\_tumor\_classifier(config=config, dataset\_path=str(local\_dir)) output\_dir = Path("/tmp/training\_results") output\_dir.mkdir(parents=True, exist\_ok=True) (output\_dir / "metrics.json").write\_text(json.dumps(result\["metrics"\])) (output\_dir / "predictions.json").write\_text(json.dumps({ "preds": result\["final\_preds"\], "targets": result\["final\_targets"\], })) return await Dir.from\_local(str(output\_dir)) # {{/docs-fragment train\_model}} # {{docs-fragment create\_report}} @report\_env.task(report=True) async def create\_report(results\_dir: Dir) -> None: """ Download training metrics and render loss/accuracy curves, confusion matrix, and per-class F1 chart in the Union UI report panel. """ import numpy as np from pathlib import Path from utils import create\_confusion\_matrix\_plot, create\_metrics\_plots, create\_per\_class\_f1\_plot local\_dir = Path("/tmp/tumor\_report") local\_dir.mkdir(parents=True, exist\_ok=True) await results\_dir.download(local\_path=str(local\_dir)) matches = list(local\_dir.glob("\*\*/metrics.json")) if not matches: raise RuntimeError(f"metrics.json not found under {local\_dir}") local\_path = matches\[0\].parent history = json.loads((local\_path / "metrics.json").read\_text()) predictions = json.loads((local\_path / "predictions.json").read\_text()) preds = np.array(predictions\["preds"\]) targets = np.array(predictions\["targets"\]) loss\_fig, acc\_fig = create\_metrics\_plots(history) cm\_fig = create\_confusion\_matrix\_plot(preds, targets) f1\_fig = create\_per\_class\_f1\_plot(preds, targets) combined\_html = ( acc\_fig.to\_html(include\_plotlyjs=True, full\_html=False) + loss\_fig.to\_html(include\_plotlyjs=False, full\_html=False) + cm\_fig.to\_html(include\_plotlyjs=False, full\_html=False) + f1\_fig.to\_html(include\_plotlyjs=False, full\_html=False) ) flyte.report.log(combined\_html, do\_flush=True) # {{/docs-fragment create\_report}} # {{docs-fragment pipeline}} @pipeline\_env.task async def tumor\_detection\_pipeline() -> None: """Orchestrate dataset loading, GPU training, and report generation.""" dataset\_dir = await load\_dataset() results\_dir = await train\_model( dataset\_dir=dataset\_dir, config\_json=json.dumps(TRAINING\_CONFIG.to\_dict()), ) await create\_report(results\_dir=results\_dir) # {{/docs-fragment pipeline}} if \_\_name\_\_ == "\_\_main\_\_": import pathlib flyte.init\_from\_config(root\_dir=pathlib.Path(\_\_file\_\_).parent) run = flyte.with\_runcontext().run(tumor\_detection\_pipeline) print(f"\\n✓ Pipeline submitted!") print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor\_detection/run.py\* ## Train the model The training task downloads the dataset \`Dir\`, runs two-phase training (frozen backbone, then full fine-tuning), and writes metrics and predictions to an output \`Dir\`. It sets \`retries=3\` so a preempted GPU node restarts the task. \`\`\` """ Flyte/Union pipeline for brain tumor MRI classification. Three-task pipeline: 1. load\_dataset — download Brain Tumor MRI from Hugging Face, cache as Dir (CPU) 2. train\_model — two-phase EfficientNet-B4 training with focal loss (GPU) 3. create\_report — render training curves and confusion matrix in the Union UI (CPU) """ import json import flyte from flyte.io import Dir from config import TrainingConfig, dataset\_env, pipeline\_env, report\_env, training\_env from dataset import download\_tumor\_dataset # {{docs-fragment config}} TRAINING\_CONFIG = TrainingConfig( phase1\_epochs=8, phase2\_epochs=25, phase1\_lr=1e-3, phase2\_lr=5e-5, batch\_size=16, num\_workers=0, log\_interval=50, mixup\_alpha=0.0, image\_size=380, focal\_gamma=3.0, ) # {{/docs-fragment config}} # {{docs-fragment load\_dataset}} @dataset\_env.task async def load\_dataset() -> Dir: """ Download raw Brain Tumor MRI JPEG files from Hugging Face and cache as flyte.io.Dir. Runs once — result is reused on subsequent pipeline runs (cache="auto"). """ return await download\_tumor\_dataset() # {{/docs-fragment load\_dataset}} # {{docs-fragment train\_model}} @training\_env.task(retries=3) async def train\_model(dataset\_dir: Dir, config\_json: str) -> Dir: """ Download the raw dataset Dir, run two-phase training, and return training metrics and final predictions as a Dir for the report task. """ from pathlib import Path local\_dir = Path("/tmp/tumor\_local") local\_dir.mkdir(parents=True, exist\_ok=True) await dataset\_dir.download(local\_path=str(local\_dir)) from training import train\_tumor\_classifier config = TrainingConfig(\*\*json.loads(config\_json)) result = train\_tumor\_classifier(config=config, dataset\_path=str(local\_dir)) output\_dir = Path("/tmp/training\_results") output\_dir.mkdir(parents=True, exist\_ok=True) (output\_dir / "metrics.json").write\_text(json.dumps(result\["metrics"\])) (output\_dir / "predictions.json").write\_text(json.dumps({ "preds": result\["final\_preds"\], "targets": result\["final\_targets"\], })) return await Dir.from\_local(str(output\_dir)) # {{/docs-fragment train\_model}} # {{docs-fragment create\_report}} @report\_env.task(report=True) async def create\_report(results\_dir: Dir) -> None: """ Download training metrics and render loss/accuracy curves, confusion matrix, and per-class F1 chart in the Union UI report panel. """ import numpy as np from pathlib import Path from utils import create\_confusion\_matrix\_plot, create\_metrics\_plots, create\_per\_class\_f1\_plot local\_dir = Path("/tmp/tumor\_report") local\_dir.mkdir(parents=True, exist\_ok=True) await results\_dir.download(local\_path=str(local\_dir)) matches = list(local\_dir.glob("\*\*/metrics.json")) if not matches: raise RuntimeError(f"metrics.json not found under {local\_dir}") local\_path = matches\[0\].parent history = json.loads((local\_path / "metrics.json").read\_text()) predictions = json.loads((local\_path / "predictions.json").read\_text()) preds = np.array(predictions\["preds"\]) targets = np.array(predictions\["targets"\]) loss\_fig, acc\_fig = create\_metrics\_plots(history) cm\_fig = create\_confusion\_matrix\_plot(preds, targets) f1\_fig = create\_per\_class\_f1\_plot(preds, targets) combined\_html = ( acc\_fig.to\_html(include\_plotlyjs=True, full\_html=False) + loss\_fig.to\_html(include\_plotlyjs=False, full\_html=False) + cm\_fig.to\_html(include\_plotlyjs=False, full\_html=False) + f1\_fig.to\_html(include\_plotlyjs=False, full\_html=False) ) flyte.report.log(combined\_html, do\_flush=True) # {{/docs-fragment create\_report}} # {{docs-fragment pipeline}} @pipeline\_env.task async def tumor\_detection\_pipeline() -> None: """Orchestrate dataset loading, GPU training, and report generation.""" dataset\_dir = await load\_dataset() results\_dir = await train\_model( dataset\_dir=dataset\_dir, config\_json=json.dumps(TRAINING\_CONFIG.to\_dict()), ) await create\_report(results\_dir=results\_dir) # {{/docs-fragment pipeline}} if \_\_name\_\_ == "\_\_main\_\_": import pathlib flyte.init\_from\_config(root\_dir=pathlib.Path(\_\_file\_\_).parent) run = flyte.with\_runcontext().run(tumor\_detection\_pipeline) print(f"\\n✓ Pipeline submitted!") print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor\_detection/run.py\* ### Resumable checkpointing To make retries cheap, training mirrors its Lightning checkpoint directory to a \`flyte.Checkpoint\` after every epoch, and resumes from the latest checkpoint when the task restarts. \`\`\` """ Training pipeline for brain tumor MRI classification. Implements two-phase training: - Phase 1: Frozen backbone (feature extractor), train classification head - Phase 2: Fine-tune full model with differential LRs + cosine annealing """ from config import TrainingConfig from dataset import compute\_class\_weights, create\_data\_loaders from model import TumorClassifierLightningModule from utils import get\_model\_size, get\_trainable\_params def train\_tumor\_classifier( config: TrainingConfig, dataset\_path: str, ) -> dict: """ Run two-phase training on the preprocessed dataset and return metrics + final predictions. dataset\_path: local directory where the flyte.io.Dir was downloaded by the training task. """ import pathlib import flyte import lightning as L import torch from lightning.pytorch.callbacks import ModelCheckpoint from typing\_extensions import override # {{docs-fragment flyte\_checkpoint}} class FlyteLightningCheckpointCallback(ModelCheckpoint): """Mirrors the checkpoint directory to Flyte after every epoch so retries can resume.""" def \_\_init\_\_(self, flyte\_checkpoint: flyte.Checkpoint, \*, dirpath: str, \*\*kwargs): super().\_\_init\_\_(dirpath=dirpath, \*\*kwargs) self.\_flyte\_checkpoint = flyte\_checkpoint @override def on\_train\_epoch\_end(self, trainer: L.Trainer, pl\_module: L.LightningModule) -> None: super().on\_train\_epoch\_end(trainer, pl\_module) if self.dirpath: self.\_flyte\_checkpoint.save\_sync(pathlib.Path(self.dirpath)) # {{/docs-fragment flyte\_checkpoint}} class MetricsLoggerCallback(L.Callback): def \_\_init\_\_(self, phase1\_epochs: int): super().\_\_init\_\_() self.phase1\_epochs = phase1\_epochs self.history = \[\] def on\_validation\_epoch\_end(self, trainer, \_pl\_module): epoch = trainer.current\_epoch metrics = trainer.callback\_metrics self.history.append({ "epoch": epoch, "phase": 1 if epoch < self.phase1\_epochs else 2, "train\_loss": float(metrics.get("train/loss\_epoch", 0)), "val\_loss": float(metrics.get("val/loss", 0)), "val\_acc": float(metrics.get("val/acc", 0)), "macro\_f1": float(metrics.get("val/macro\_f1", 0)), }) class PhaseChangeCallback(L.Callback): def \_\_init\_\_(self, phase1\_epochs: int, phase2\_lr: float): super().\_\_init\_\_() self.phase1\_epochs = phase1\_epochs self.phase2\_lr = phase2\_lr self.phase\_changed = False def on\_train\_epoch\_end(self, trainer, pl\_module): if not self.phase\_changed and (trainer.current\_epoch + 1) == self.phase1\_epochs: print("\\n" + "=" \* 80) print("TRANSITIONING TO PHASE 2: UNFREEZING BACKBONE AND ADJUSTING LR") print("=" \* 80 + "\\n") pl\_module.model.unfreeze\_backbone() for param\_group in trainer.optimizers\[0\].param\_groups: param\_group\["lr"\] = self.phase2\_lr # Add backbone params to optimizer with 10x lower LR. # Backbone was excluded at init because it was frozen. backbone\_lr = self.phase2\_lr \* 0.1 backbone\_decay, backbone\_no\_decay = \[\], \[\] for param in pl\_module.model.backbone.parameters(): if param.ndim >= 2: backbone\_decay.append(param) else: backbone\_no\_decay.append(param) optimizer = trainer.optimizers\[0\] optimizer.add\_param\_group({"params": backbone\_decay, "lr": backbone\_lr, "weight\_decay": pl\_module.weight\_decay}) optimizer.add\_param\_group({"params": backbone\_no\_decay, "lr": backbone\_lr, "weight\_decay": 0.0}) # Fresh cosine schedule over remaining Phase 2 steps to avoid # the Phase 1 schedule arriving near-zero before Phase 2 begins. steps\_remaining = trainer.estimated\_stepping\_batches - trainer.global\_step new\_scheduler = torch.optim.lr\_scheduler.CosineAnnealingLR( trainer.optimizers\[0\], T\_max=max(1, steps\_remaining), eta\_min=1e-6, ) for lr\_scheduler\_config in trainer.lr\_scheduler\_configs: lr\_scheduler\_config.scheduler = new\_scheduler print(f"Phase 2 started: lr={self.phase2\_lr}") print(f"Total parameters: {get\_model\_size(pl\_module.model):,}") print(f"Trainable parameters: {get\_trainable\_params(pl\_module.model):,}") self.phase\_changed = True print("\\n" + "=" \* 80) print("BRAIN TUMOR MRI CLASSIFICATION WITH EFFICIENTNET-B4") print("=" \* 80) print(f"Config: {config.to\_dict()}\\n") device = torch.device("cuda" if torch.cuda.is\_available() else "cpu") print(f"Using device: {device}") if torch.cuda.is\_available(): print(f"GPU: {torch.cuda.get\_device\_name(0)}") print(f"GPU Memory: {torch.cuda.get\_device\_properties(0).total\_memory / 1e9:.2f} GB") print("\\nLoading MRI images...") train\_loader, val\_loader = create\_data\_loaders( dataset\_path=dataset\_path, image\_size=config.image\_size, batch\_size=config.batch\_size, num\_workers=config.num\_workers, val\_split=config.val\_split, ) print(f"Data loaders created: {len(train\_loader)} train batches, {len(val\_loader)} val batches") print("\\nComputing class weights for focal loss...") class\_weights = compute\_class\_weights(dataset\_path) print(f"Class weights: {class\_weights.tolist()}") # Per-class gamma: Meningioma gets 7.0, all others 3.0. # CLASS\_NAMES alphabetical order: Glioma=0, Meningioma=1, No Tumor=2, Pituitary=3 gamma\_per\_class = torch.tensor(\[3.0, 7.0, 3.0, 3.0\]) print("\\nInitializing model...") model = TumorClassifierLightningModule( num\_classes=config.num\_classes, model\_name=config.model\_name, pretrained=config.pretrained, learning\_rate=config.phase1\_lr, freeze\_backbone=config.phase1\_freeze\_backbone, weight\_decay=config.weight\_decay, warmup\_steps=config.warmup\_steps, max\_epochs=config.phase1\_epochs + config.phase2\_epochs, focal\_gamma=config.focal\_gamma, mixup\_alpha=config.mixup\_alpha, class\_weights=class\_weights, gamma\_per\_class=gamma\_per\_class, ) print(f"Model: {config.model\_name}") print(f"Total parameters: {get\_model\_size(model.model):,}") print(f"Trainable parameters: {get\_trainable\_params(model.model):,}") from pathlib import Path checkpoint\_dir = Path("/tmp/tumor\_checkpoints") checkpoint\_dir.mkdir(parents=True, exist\_ok=True) # {{docs-fragment resume}} # --- Flyte checkpoint: resume from previous attempt if one exists --- resume\_ckpt: str | None = None ctx = flyte.ctx() flyte\_checkpoint = getattr(ctx, "checkpoint", None) if ctx else None if flyte\_checkpoint: prev\_path = flyte\_checkpoint.load\_sync() if prev\_path: last = flyte.latest\_checkpoint(prev\_path) if last: ck = torch.load(str(last), map\_location="cpu", weights\_only=False) epoch\_start = int(ck.get("epoch", 0)) resume\_ckpt = str(last) print(f"Resuming from epoch {epoch\_start}, checkpoint: {last}") # -------------------------------------------------------------------- # {{/docs-fragment resume}} metrics\_cb = MetricsLoggerCallback(phase1\_epochs=config.phase1\_epochs) resume\_callback = ( FlyteLightningCheckpointCallback( flyte\_checkpoint, dirpath=str(checkpoint\_dir), filename="last", save\_last=True, save\_top\_k=1, ) if flyte\_checkpoint else ModelCheckpoint( dirpath=str(checkpoint\_dir), filename="best-{epoch:03d}-{val\_acc:.3f}", monitor="val/acc", mode="max", save\_top\_k=3, verbose=True, auto\_insert\_metric\_name=False, ) ) callbacks = \[\ resume\_callback,\ metrics\_cb,\ PhaseChangeCallback(\ phase1\_epochs=config.phase1\_epochs,\ phase2\_lr=config.phase2\_lr,\ ),\ \] trainer = L.Trainer( max\_epochs=config.phase1\_epochs + config.phase2\_epochs, accelerator="gpu" if torch.cuda.is\_available() else "cpu", devices=1, precision="16-mixed", callbacks=callbacks, enable\_progress\_bar=True, enable\_model\_summary=True, log\_every\_n\_steps=config.log\_interval, gradient\_clip\_val=1.0, ) trainer.fit(model, train\_loader, val\_loader, ckpt\_path=resume\_ckpt) best\_checkpoint = trainer.checkpoint\_callback.best\_model\_path print(f"\\n✓ Training complete!") print(f"Best checkpoint: {best\_checkpoint}") # Final inference with TTA (test-time augmentation): average logits over # original + h-flip + v-flip + 90° rotations for a free accuracy boost. print("\\nRunning final inference with TTA for confusion matrix...") import numpy as np import torchvision.transforms.functional as TF model.eval() model.to(device) all\_preds, all\_targets = \[\], \[\] with torch.no\_grad(): for images, labels in val\_loader: images = images.to(device) aug\_logits = \[\ model.model(images),\ model.model(TF.hflip(images)),\ model.model(TF.vflip(images)),\ model.model(torch.rot90(images, k=1, dims=\[2, 3\])),\ model.model(torch.rot90(images, k=3, dims=\[2, 3\])),\ \] avg\_logits = torch.stack(aug\_logits).mean(dim=0) all\_preds.append(avg\_logits.argmax(dim=1).cpu()) all\_targets.append(labels.cpu()) final\_preds = torch.cat(all\_preds).numpy() final\_targets = torch.cat(all\_targets).numpy() return { "best\_checkpoint": best\_checkpoint, "metrics": metrics\_cb.history, "final\_preds": final\_preds.tolist(), "final\_targets": final\_targets.tolist(), } \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor\_detection/training.py\* On startup, the training loop looks for a checkpoint from a previous attempt and resumes from it if present: \`\`\` """ Training pipeline for brain tumor MRI classification. Implements two-phase training: - Phase 1: Frozen backbone (feature extractor), train classification head - Phase 2: Fine-tune full model with differential LRs + cosine annealing """ from config import TrainingConfig from dataset import compute\_class\_weights, create\_data\_loaders from model import TumorClassifierLightningModule from utils import get\_model\_size, get\_trainable\_params def train\_tumor\_classifier( config: TrainingConfig, dataset\_path: str, ) -> dict: """ Run two-phase training on the preprocessed dataset and return metrics + final predictions. dataset\_path: local directory where the flyte.io.Dir was downloaded by the training task. """ import pathlib import flyte import lightning as L import torch from lightning.pytorch.callbacks import ModelCheckpoint from typing\_extensions import override # {{docs-fragment flyte\_checkpoint}} class FlyteLightningCheckpointCallback(ModelCheckpoint): """Mirrors the checkpoint directory to Flyte after every epoch so retries can resume.""" def \_\_init\_\_(self, flyte\_checkpoint: flyte.Checkpoint, \*, dirpath: str, \*\*kwargs): super().\_\_init\_\_(dirpath=dirpath, \*\*kwargs) self.\_flyte\_checkpoint = flyte\_checkpoint @override def on\_train\_epoch\_end(self, trainer: L.Trainer, pl\_module: L.LightningModule) -> None: super().on\_train\_epoch\_end(trainer, pl\_module) if self.dirpath: self.\_flyte\_checkpoint.save\_sync(pathlib.Path(self.dirpath)) # {{/docs-fragment flyte\_checkpoint}} class MetricsLoggerCallback(L.Callback): def \_\_init\_\_(self, phase1\_epochs: int): super().\_\_init\_\_() self.phase1\_epochs = phase1\_epochs self.history = \[\] def on\_validation\_epoch\_end(self, trainer, \_pl\_module): epoch = trainer.current\_epoch metrics = trainer.callback\_metrics self.history.append({ "epoch": epoch, "phase": 1 if epoch < self.phase1\_epochs else 2, "train\_loss": float(metrics.get("train/loss\_epoch", 0)), "val\_loss": float(metrics.get("val/loss", 0)), "val\_acc": float(metrics.get("val/acc", 0)), "macro\_f1": float(metrics.get("val/macro\_f1", 0)), }) class PhaseChangeCallback(L.Callback): def \_\_init\_\_(self, phase1\_epochs: int, phase2\_lr: float): super().\_\_init\_\_() self.phase1\_epochs = phase1\_epochs self.phase2\_lr = phase2\_lr self.phase\_changed = False def on\_train\_epoch\_end(self, trainer, pl\_module): if not self.phase\_changed and (trainer.current\_epoch + 1) == self.phase1\_epochs: print("\\n" + "=" \* 80) print("TRANSITIONING TO PHASE 2: UNFREEZING BACKBONE AND ADJUSTING LR") print("=" \* 80 + "\\n") pl\_module.model.unfreeze\_backbone() for param\_group in trainer.optimizers\[0\].param\_groups: param\_group\["lr"\] = self.phase2\_lr # Add backbone params to optimizer with 10x lower LR. # Backbone was excluded at init because it was frozen. backbone\_lr = self.phase2\_lr \* 0.1 backbone\_decay, backbone\_no\_decay = \[\], \[\] for param in pl\_module.model.backbone.parameters(): if param.ndim >= 2: backbone\_decay.append(param) else: backbone\_no\_decay.append(param) optimizer = trainer.optimizers\[0\] optimizer.add\_param\_group({"params": backbone\_decay, "lr": backbone\_lr, "weight\_decay": pl\_module.weight\_decay}) optimizer.add\_param\_group({"params": backbone\_no\_decay, "lr": backbone\_lr, "weight\_decay": 0.0}) # Fresh cosine schedule over remaining Phase 2 steps to avoid # the Phase 1 schedule arriving near-zero before Phase 2 begins. steps\_remaining = trainer.estimated\_stepping\_batches - trainer.global\_step new\_scheduler = torch.optim.lr\_scheduler.CosineAnnealingLR( trainer.optimizers\[0\], T\_max=max(1, steps\_remaining), eta\_min=1e-6, ) for lr\_scheduler\_config in trainer.lr\_scheduler\_configs: lr\_scheduler\_config.scheduler = new\_scheduler print(f"Phase 2 started: lr={self.phase2\_lr}") print(f"Total parameters: {get\_model\_size(pl\_module.model):,}") print(f"Trainable parameters: {get\_trainable\_params(pl\_module.model):,}") self.phase\_changed = True print("\\n" + "=" \* 80) print("BRAIN TUMOR MRI CLASSIFICATION WITH EFFICIENTNET-B4") print("=" \* 80) print(f"Config: {config.to\_dict()}\\n") device = torch.device("cuda" if torch.cuda.is\_available() else "cpu") print(f"Using device: {device}") if torch.cuda.is\_available(): print(f"GPU: {torch.cuda.get\_device\_name(0)}") print(f"GPU Memory: {torch.cuda.get\_device\_properties(0).total\_memory / 1e9:.2f} GB") print("\\nLoading MRI images...") train\_loader, val\_loader = create\_data\_loaders( dataset\_path=dataset\_path, image\_size=config.image\_size, batch\_size=config.batch\_size, num\_workers=config.num\_workers, val\_split=config.val\_split, ) print(f"Data loaders created: {len(train\_loader)} train batches, {len(val\_loader)} val batches") print("\\nComputing class weights for focal loss...") class\_weights = compute\_class\_weights(dataset\_path) print(f"Class weights: {class\_weights.tolist()}") # Per-class gamma: Meningioma gets 7.0, all others 3.0. # CLASS\_NAMES alphabetical order: Glioma=0, Meningioma=1, No Tumor=2, Pituitary=3 gamma\_per\_class = torch.tensor(\[3.0, 7.0, 3.0, 3.0\]) print("\\nInitializing model...") model = TumorClassifierLightningModule( num\_classes=config.num\_classes, model\_name=config.model\_name, pretrained=config.pretrained, learning\_rate=config.phase1\_lr, freeze\_backbone=config.phase1\_freeze\_backbone, weight\_decay=config.weight\_decay, warmup\_steps=config.warmup\_steps, max\_epochs=config.phase1\_epochs + config.phase2\_epochs, focal\_gamma=config.focal\_gamma, mixup\_alpha=config.mixup\_alpha, class\_weights=class\_weights, gamma\_per\_class=gamma\_per\_class, ) print(f"Model: {config.model\_name}") print(f"Total parameters: {get\_model\_size(model.model):,}") print(f"Trainable parameters: {get\_trainable\_params(model.model):,}") from pathlib import Path checkpoint\_dir = Path("/tmp/tumor\_checkpoints") checkpoint\_dir.mkdir(parents=True, exist\_ok=True) # {{docs-fragment resume}} # --- Flyte checkpoint: resume from previous attempt if one exists --- resume\_ckpt: str | None = None ctx = flyte.ctx() flyte\_checkpoint = getattr(ctx, "checkpoint", None) if ctx else None if flyte\_checkpoint: prev\_path = flyte\_checkpoint.load\_sync() if prev\_path: last = flyte.latest\_checkpoint(prev\_path) if last: ck = torch.load(str(last), map\_location="cpu", weights\_only=False) epoch\_start = int(ck.get("epoch", 0)) resume\_ckpt = str(last) print(f"Resuming from epoch {epoch\_start}, checkpoint: {last}") # -------------------------------------------------------------------- # {{/docs-fragment resume}} metrics\_cb = MetricsLoggerCallback(phase1\_epochs=config.phase1\_epochs) resume\_callback = ( FlyteLightningCheckpointCallback( flyte\_checkpoint, dirpath=str(checkpoint\_dir), filename="last", save\_last=True, save\_top\_k=1, ) if flyte\_checkpoint else ModelCheckpoint( dirpath=str(checkpoint\_dir), filename="best-{epoch:03d}-{val\_acc:.3f}", monitor="val/acc", mode="max", save\_top\_k=3, verbose=True, auto\_insert\_metric\_name=False, ) ) callbacks = \[\ resume\_callback,\ metrics\_cb,\ PhaseChangeCallback(\ phase1\_epochs=config.phase1\_epochs,\ phase2\_lr=config.phase2\_lr,\ ),\ \] trainer = L.Trainer( max\_epochs=config.phase1\_epochs + config.phase2\_epochs, accelerator="gpu" if torch.cuda.is\_available() else "cpu", devices=1, precision="16-mixed", callbacks=callbacks, enable\_progress\_bar=True, enable\_model\_summary=True, log\_every\_n\_steps=config.log\_interval, gradient\_clip\_val=1.0, ) trainer.fit(model, train\_loader, val\_loader, ckpt\_path=resume\_ckpt) best\_checkpoint = trainer.checkpoint\_callback.best\_model\_path print(f"\\n✓ Training complete!") print(f"Best checkpoint: {best\_checkpoint}") # Final inference with TTA (test-time augmentation): average logits over # original + h-flip + v-flip + 90° rotations for a free accuracy boost. print("\\nRunning final inference with TTA for confusion matrix...") import numpy as np import torchvision.transforms.functional as TF model.eval() model.to(device) all\_preds, all\_targets = \[\], \[\] with torch.no\_grad(): for images, labels in val\_loader: images = images.to(device) aug\_logits = \[\ model.model(images),\ model.model(TF.hflip(images)),\ model.model(TF.vflip(images)),\ model.model(torch.rot90(images, k=1, dims=\[2, 3\])),\ model.model(torch.rot90(images, k=3, dims=\[2, 3\])),\ \] avg\_logits = torch.stack(aug\_logits).mean(dim=0) all\_preds.append(avg\_logits.argmax(dim=1).cpu()) all\_targets.append(labels.cpu()) final\_preds = torch.cat(all\_preds).numpy() final\_targets = torch.cat(all\_targets).numpy() return { "best\_checkpoint": best\_checkpoint, "metrics": metrics\_cb.history, "final\_preds": final\_preds.tolist(), "final\_targets": final\_targets.tolist(), } \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor\_detection/training.py\* ## Generate the report The reporting task reads the metrics and predictions, then renders accuracy/loss curves, a confusion matrix, and a per-class F1 chart with Plotly. \`report=True\` surfaces the HTML directly in the run's report panel. \`\`\` """ Flyte/Union pipeline for brain tumor MRI classification. Three-task pipeline: 1. load\_dataset — download Brain Tumor MRI from Hugging Face, cache as Dir (CPU) 2. train\_model — two-phase EfficientNet-B4 training with focal loss (GPU) 3. create\_report — render training curves and confusion matrix in the Union UI (CPU) """ import json import flyte from flyte.io import Dir from config import TrainingConfig, dataset\_env, pipeline\_env, report\_env, training\_env from dataset import download\_tumor\_dataset # {{docs-fragment config}} TRAINING\_CONFIG = TrainingConfig( phase1\_epochs=8, phase2\_epochs=25, phase1\_lr=1e-3, phase2\_lr=5e-5, batch\_size=16, num\_workers=0, log\_interval=50, mixup\_alpha=0.0, image\_size=380, focal\_gamma=3.0, ) # {{/docs-fragment config}} # {{docs-fragment load\_dataset}} @dataset\_env.task async def load\_dataset() -> Dir: """ Download raw Brain Tumor MRI JPEG files from Hugging Face and cache as flyte.io.Dir. Runs once — result is reused on subsequent pipeline runs (cache="auto"). """ return await download\_tumor\_dataset() # {{/docs-fragment load\_dataset}} # {{docs-fragment train\_model}} @training\_env.task(retries=3) async def train\_model(dataset\_dir: Dir, config\_json: str) -> Dir: """ Download the raw dataset Dir, run two-phase training, and return training metrics and final predictions as a Dir for the report task. """ from pathlib import Path local\_dir = Path("/tmp/tumor\_local") local\_dir.mkdir(parents=True, exist\_ok=True) await dataset\_dir.download(local\_path=str(local\_dir)) from training import train\_tumor\_classifier config = TrainingConfig(\*\*json.loads(config\_json)) result = train\_tumor\_classifier(config=config, dataset\_path=str(local\_dir)) output\_dir = Path("/tmp/training\_results") output\_dir.mkdir(parents=True, exist\_ok=True) (output\_dir / "metrics.json").write\_text(json.dumps(result\["metrics"\])) (output\_dir / "predictions.json").write\_text(json.dumps({ "preds": result\["final\_preds"\], "targets": result\["final\_targets"\], })) return await Dir.from\_local(str(output\_dir)) # {{/docs-fragment train\_model}} # {{docs-fragment create\_report}} @report\_env.task(report=True) async def create\_report(results\_dir: Dir) -> None: """ Download training metrics and render loss/accuracy curves, confusion matrix, and per-class F1 chart in the Union UI report panel. """ import numpy as np from pathlib import Path from utils import create\_confusion\_matrix\_plot, create\_metrics\_plots, create\_per\_class\_f1\_plot local\_dir = Path("/tmp/tumor\_report") local\_dir.mkdir(parents=True, exist\_ok=True) await results\_dir.download(local\_path=str(local\_dir)) matches = list(local\_dir.glob("\*\*/metrics.json")) if not matches: raise RuntimeError(f"metrics.json not found under {local\_dir}") local\_path = matches\[0\].parent history = json.loads((local\_path / "metrics.json").read\_text()) predictions = json.loads((local\_path / "predictions.json").read\_text()) preds = np.array(predictions\["preds"\]) targets = np.array(predictions\["targets"\]) loss\_fig, acc\_fig = create\_metrics\_plots(history) cm\_fig = create\_confusion\_matrix\_plot(preds, targets) f1\_fig = create\_per\_class\_f1\_plot(preds, targets) combined\_html = ( acc\_fig.to\_html(include\_plotlyjs=True, full\_html=False) + loss\_fig.to\_html(include\_plotlyjs=False, full\_html=False) + cm\_fig.to\_html(include\_plotlyjs=False, full\_html=False) + f1\_fig.to\_html(include\_plotlyjs=False, full\_html=False) ) flyte.report.log(combined\_html, do\_flush=True) # {{/docs-fragment create\_report}} # {{docs-fragment pipeline}} @pipeline\_env.task async def tumor\_detection\_pipeline() -> None: """Orchestrate dataset loading, GPU training, and report generation.""" dataset\_dir = await load\_dataset() results\_dir = await train\_model( dataset\_dir=dataset\_dir, config\_json=json.dumps(TRAINING\_CONFIG.to\_dict()), ) await create\_report(results\_dir=results\_dir) # {{/docs-fragment pipeline}} if \_\_name\_\_ == "\_\_main\_\_": import pathlib flyte.init\_from\_config(root\_dir=pathlib.Path(\_\_file\_\_).parent) run = flyte.with\_runcontext().run(tumor\_detection\_pipeline) print(f"\\n✓ Pipeline submitted!") print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor\_detection/run.py\* ## Orchestrate the pipeline The driver task wires the three steps together. \`\`\` """ Flyte/Union pipeline for brain tumor MRI classification. Three-task pipeline: 1. load\_dataset — download Brain Tumor MRI from Hugging Face, cache as Dir (CPU) 2. train\_model — two-phase EfficientNet-B4 training with focal loss (GPU) 3. create\_report — render training curves and confusion matrix in the Union UI (CPU) """ import json import flyte from flyte.io import Dir from config import TrainingConfig, dataset\_env, pipeline\_env, report\_env, training\_env from dataset import download\_tumor\_dataset # {{docs-fragment config}} TRAINING\_CONFIG = TrainingConfig( phase1\_epochs=8, phase2\_epochs=25, phase1\_lr=1e-3, phase2\_lr=5e-5, batch\_size=16, num\_workers=0, log\_interval=50, mixup\_alpha=0.0, image\_size=380, focal\_gamma=3.0, ) # {{/docs-fragment config}} # {{docs-fragment load\_dataset}} @dataset\_env.task async def load\_dataset() -> Dir: """ Download raw Brain Tumor MRI JPEG files from Hugging Face and cache as flyte.io.Dir. Runs once — result is reused on subsequent pipeline runs (cache="auto"). """ return await download\_tumor\_dataset() # {{/docs-fragment load\_dataset}} # {{docs-fragment train\_model}} @training\_env.task(retries=3) async def train\_model(dataset\_dir: Dir, config\_json: str) -> Dir: """ Download the raw dataset Dir, run two-phase training, and return training metrics and final predictions as a Dir for the report task. """ from pathlib import Path local\_dir = Path("/tmp/tumor\_local") local\_dir.mkdir(parents=True, exist\_ok=True) await dataset\_dir.download(local\_path=str(local\_dir)) from training import train\_tumor\_classifier config = TrainingConfig(\*\*json.loads(config\_json)) result = train\_tumor\_classifier(config=config, dataset\_path=str(local\_dir)) output\_dir = Path("/tmp/training\_results") output\_dir.mkdir(parents=True, exist\_ok=True) (output\_dir / "metrics.json").write\_text(json.dumps(result\["metrics"\])) (output\_dir / "predictions.json").write\_text(json.dumps({ "preds": result\["final\_preds"\], "targets": result\["final\_targets"\], })) return await Dir.from\_local(str(output\_dir)) # {{/docs-fragment train\_model}} # {{docs-fragment create\_report}} @report\_env.task(report=True) async def create\_report(results\_dir: Dir) -> None: """ Download training metrics and render loss/accuracy curves, confusion matrix, and per-class F1 chart in the Union UI report panel. """ import numpy as np from pathlib import Path from utils import create\_confusion\_matrix\_plot, create\_metrics\_plots, create\_per\_class\_f1\_plot local\_dir = Path("/tmp/tumor\_report") local\_dir.mkdir(parents=True, exist\_ok=True) await results\_dir.download(local\_path=str(local\_dir)) matches = list(local\_dir.glob("\*\*/metrics.json")) if not matches: raise RuntimeError(f"metrics.json not found under {local\_dir}") local\_path = matches\[0\].parent history = json.loads((local\_path / "metrics.json").read\_text()) predictions = json.loads((local\_path / "predictions.json").read\_text()) preds = np.array(predictions\["preds"\]) targets = np.array(predictions\["targets"\]) loss\_fig, acc\_fig = create\_metrics\_plots(history) cm\_fig = create\_confusion\_matrix\_plot(preds, targets) f1\_fig = create\_per\_class\_f1\_plot(preds, targets) combined\_html = ( acc\_fig.to\_html(include\_plotlyjs=True, full\_html=False) + loss\_fig.to\_html(include\_plotlyjs=False, full\_html=False) + cm\_fig.to\_html(include\_plotlyjs=False, full\_html=False) + f1\_fig.to\_html(include\_plotlyjs=False, full\_html=False) ) flyte.report.log(combined\_html, do\_flush=True) # {{/docs-fragment create\_report}} # {{docs-fragment pipeline}} @pipeline\_env.task async def tumor\_detection\_pipeline() -> None: """Orchestrate dataset loading, GPU training, and report generation.""" dataset\_dir = await load\_dataset() results\_dir = await train\_model( dataset\_dir=dataset\_dir, config\_json=json.dumps(TRAINING\_CONFIG.to\_dict()), ) await create\_report(results\_dir=results\_dir) # {{/docs-fragment pipeline}} if \_\_name\_\_ == "\_\_main\_\_": import pathlib flyte.init\_from\_config(root\_dir=pathlib.Path(\_\_file\_\_).parent) run = flyte.with\_runcontext().run(tumor\_detection\_pipeline) print(f"\\n✓ Pipeline submitted!") print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/tumor\_detection/run.py\* ## Run the pipeline This example has no secrets — the dataset is public. Because the pipeline imports sibling modules and uses \`with\_source\_folder\`, run it from inside the example directory so the local files are picked up. From the \[example directory\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/tumor\_detection): \`\`\` cd v2/tutorials/tumor\_detection python run.py \`\`\` Or submit it with the Flyte CLI from the same directory: \`\`\` flyte run run.py tumor\_detection\_pipeline \`\`\` When the run completes, open the \`create\_report\` task in the UI to view the training curves, confusion matrix, and per-class F1 scores. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/geospatial === # Geospatial Tutorials for satellite imagery, remote sensing, and earth and atmospheric modeling workloads. ### \*\*Geospatial > GPU-accelerated climate modeling\*\* Run ensemble atmospheric simulations on H200 GPUs with multi-source data ingestion and real-time extreme event detection. ### \[Satellite image classification\](satellite\_image\_classification) Build a production-grade EfficientNet pipeline for land-use classification with caching, experiment tracking, and reporting. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/geospatial/climate-modeling === # GPU-accelerated climate modeling Climate modeling is hard for two reasons: data and compute. Satellite imagery arrives continuously from multiple sources. Reanalysis datasets have to be pulled from remote APIs. Weather station data shows up in different formats and schemas. And once all of that is finally in one place, running atmospheric physics simulations demands serious GPU compute. In practice, many climate workflows are held together with scripts, cron jobs, and a lot of manual babysitting. Data ingestion breaks without warning. GPU jobs run overnight with little visibility into what's happening. When something interesting shows up in a simulation, like a developing hurricane, no one notices until the job finishes hours later. In this tutorial, we build a production-grade climate modeling pipeline using Flyte. We ingest data from three different sources in parallel, combine it with Dask, run ensemble atmospheric simulations on H200 GPUs, detect extreme weather events as they emerge, and visualize everything in a live dashboard. The entire pipeline is orchestrated, cached, and fault-tolerant, so it can run reliably at scale. !\[Report\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/tutorials/climate-modeling/report.png) > \[!NOTE\] > Full code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/climate\_modeling/simulation.py). ## Overview We're building an ensemble weather forecasting system. Ensemble forecasting runs the same simulation multiple times with slightly different initial conditions. This quantifies forecast uncertainty. Instead of saying "the temperature will be 25°C", we can say "the temperature will be 24-26°C with 90% confidence". The pipeline has five stages: 1. \*\*Data ingestion\*\*: Pull satellite imagery from NOAA GOES, reanalysis data from ERA5, and surface observations from weather stations in parallel. 2. \*\*Preprocessing\*\*: Fuse the datasets, interpolate to a common grid, and run quality control using Dask for distributed computation. 3. \*\*GPU simulation\*\*: Run ensemble atmospheric physics on H200 GPUs. Each ensemble member evolves independently. PyTorch handles the tensor operations; \`torch.compile\` optimizes the kernels. 4. \*\*Event detection\*\*: Monitor for hurricanes (high wind + low pressure) and heatwaves during simulation. When extreme events are detected, the pipeline can adaptively refine the grid resolution. 5. \*\*Real-time reporting\*\*: Stream metrics to a live Flyte Reports dashboard showing convergence and detected events. This workflow is a good example of where Flyte shines! - \*\*Parallel data ingestion\*\*: Three different data sources, three different APIs, all running concurrently. Flyte's async task execution handles this naturally. - \*\*Resource heterogeneity\*\*: Data ingestion needs CPU and network. Preprocessing needs a Dask cluster. Simulation needs GPUs. Flyte provisions exactly what each stage needs. - \*\*Caching\*\*: ERA5 data fetches can take minutes. Run the pipeline twice with the same date range, and Flyte skips the fetch entirely. - \*\*Adaptive workflows\*\*: When a hurricane is detected, we can dynamically refine the simulation. Flyte makes this kind of conditional logic straightforward. ## Implementation ### Dependencies and container image \`\`\` import asyncio import gc import io import json import os import tempfile from dataclasses import dataclass from datetime import datetime, timedelta from typing import Literal import flyte import numpy as np import pandas as pd import xarray as xr from flyte.io import File from flyteplugins.dask import Dask, Scheduler, WorkerGroup \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* The key imports include \`xarray\` for multi-dimensional climate data, \`flyteplugins.dask\` for distributed preprocessing, and \`flyte\` for orchestration. \`\`\` climate\_image = ( flyte.Image.from\_debian\_base(name="climate\_modeling\_h200") .with\_apt\_packages( "libnetcdf-dev", # NetCDF for climate data "libhdf5-dev", # HDF5 for large datasets "libeccodes-dev", # GRIB format support (ECMWF's native format) "libudunits2-dev", # Unit conversions ) .with\_pip\_packages( "numpy==2.3.5", "pandas==2.3.3", "xarray==2025.11.0", "torch==2.9.1", "netCDF4==1.7.3", "s3fs==2025.10.0", "aiohttp==3.13.2", "ecmwf-datastores-client==0.4.1", "h5netcdf==1.7.3", "cfgrib==0.9.15.1", "pyarrow==22.0.0", "scipy==1.15.1", "flyteplugins-dask>=2.0.0b33", "nvidia-ml-py3==7.352.0", ) .with\_env\_vars({"PYTORCH\_CUDA\_ALLOC\_CONF": "max\_split\_size\_mb:512"}) ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* Climate data comes in specialized formats such as NetCDF, HDF5, and GRIB. The container image includes libraries to work with all of them, along with PyTorch for GPU computation and the ECMWF client for accessing ERA5 data. ### Simulation parameters and data structures \`\`\` @dataclass class SimulationParams: grid\_resolution\_km: float = 10.0 time\_step\_minutes: int = 10 simulation\_hours: int = 240 physics\_model: Literal\["WRF", "MPAS", "CAM"\] = "WRF" boundary\_layer\_scheme: str = "YSU" microphysics\_scheme: str = "Thompson" radiation\_scheme: str = "RRTMG" # Ensemble forecasting parameters ensemble\_size: int = 800 perturbation\_magnitude: float = 0.5 # Convergence criteria for adaptive refinement convergence\_threshold: float = 0.1 # 10% of initial ensemble spread max\_iterations: int = 3 @dataclass class ClimateMetrics: timestamp: str iteration: int convergence\_rate: float energy\_conservation\_error: float max\_wind\_speed\_mps: float min\_pressure\_mb: float detected\_phenomena: list\[str\] compute\_time\_seconds: float ensemble\_spread: float @dataclass class SimulationSummary: total\_iterations: int final\_resolution\_km: float avg\_convergence\_rate: float total\_compute\_time\_seconds: float hurricanes\_detected: int heatwaves\_detected: int converged: bool region: str output\_files: list\[File\] date\_range: list\[str, str\] \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* \`SimulationParams\` defines the core behavior of the simulation, including grid resolution, physics schemes, and ensemble size. The default configuration runs 800 ensemble members, which is sufficient to produce statistically meaningful uncertainty estimates. > \[!NOTE\] > Decreasing the grid spacing via \`grid\_resolution\_km\` (for example, from 10 km to 5 km) increases grid resolution and significantly increases memory usage because it introduces more data points and intermediate state. Even with 141 GB of H200 GPU memory, high-resolution or adaptively refined simulations may exceed available VRAM, especially when running large ensembles. > > To mitigate this, consider reducing the ensemble size, limiting the refined region, running fewer physics variables, or scaling the simulation across more GPUs so memory is distributed more evenly. \`ClimateMetrics\` collects diagnostics at each iteration, such as convergence rate, energy conservation, and detected phenomena. These metrics are streamed to the real-time dashboard so you can monitor how the simulation evolves as it runs. ### Task environments Different stages need different resources. Flyte's \`TaskEnvironment\` declares exactly what each task requires: \`\`\` gpu\_env = flyte.TaskEnvironment( name="climate\_modeling\_gpu", resources=flyte.Resources( cpu=5, memory="130Gi", gpu="H200:1", ), image=climate\_image, cache="auto", ) dask\_env = flyte.TaskEnvironment( name="climate\_modeling\_dask", plugin\_config=Dask( scheduler=Scheduler(resources=flyte.Resources(cpu=2, memory="6Gi")), workers=WorkerGroup( number\_of\_workers=2, resources=flyte.Resources(cpu=2, memory="12Gi"), ), ), image=climate\_image, resources=flyte.Resources(cpu=2, memory="12Gi"), # Head node cache="auto", ) cpu\_env = flyte.TaskEnvironment( name="climate\_modeling\_cpu", resources=flyte.Resources(cpu=8, memory="64Gi"), image=climate\_image, cache="auto", secrets=\[\ flyte.Secret(key="cds\_api\_key", as\_env\_var="ECMWF\_DATASTORES\_KEY"),\ flyte.Secret(key="cds\_api\_url", as\_env\_var="ECMWF\_DATASTORES\_URL"),\ \], depends\_on=\[gpu\_env, dask\_env\], ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* Here’s what each environment is responsible for: - \*\*\`gpu\_env\`\*\*: Runs the atmospheric simulations on H200 GPUs. The 130 GB of GPU memory is used to hold the ensemble members in VRAM during execution. - \*\*\`dask\_env\`\*\*: Provides a distributed Dask cluster for preprocessing. A scheduler and multiple workers handle data fusion and transformation in parallel. - \*\*\`cpu\_env\`\*\*: Handles data ingestion and orchestration. This environment also includes the secrets required to access the ERA5 API. The \`depends\_on\` setting on \`cpu\_env\` ensures that Flyte builds the GPU and Dask images first. Once those environments are ready, the orchestration task can launch the specialized simulation and preprocessing tasks. ### Data ingestion: multiple sources in parallel Climate models need data from multiple sources. Each source has different formats, APIs, and failure modes. We handle them as separate Flyte tasks that run concurrently. \*\*Satellite imagery from NOAA GOES\*\* \`\`\` @cpu\_env.task async def ingest\_satellite\_data(region: str, date\_range: list\[str, str\]) -> File: """Ingest GOES satellite imagery from NOAA's public S3 buckets.""" \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* This task fetches cloud imagery and precipitable water products from NOAA's public S3 buckets. GOES-16 covers the Atlantic; GOES-17 covers the Pacific. The task selects the appropriate satellite based on region, fetches multiple days in parallel using \`asyncio.gather\`, and combines everything into a single xarray Dataset. \*\*ERA5 reanalysis from Copernicus\*\* \`\`\` @cpu\_env.task async def ingest\_reanalysis\_data(region: str, date\_range: list\[str, str\]) -> File: """Fetch ERA5 reanalysis from Copernicus Climate Data Store.""" \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* ERA5 provides 3D atmospheric fields such as temperature, wind, humidity at multiple pressure levels from surface to stratosphere. The ECMWF datastores client handles authentication via Flyte secrets. Each day fetches in parallel, then gets concatenated. \*\*Surface observations from weather stations:\*\* \`\`\` @cpu\_env.task async def ingest\_station\_data( region: str, date\_range: list\[str, str\], max\_stations: int = 100 ) -> File: """Fetch ground observations from NOAA's Integrated Surface Database.""" \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* Ground truth comes from NOAA's Integrated Surface Database. The task filters stations by geographic bounds, fetches hourly observations, and returns a Parquet file for efficient downstream processing. All three tasks return Flyte \`File\` objects that hold references to data in blob storage. No data moves until a downstream task actually needs it. ### Preprocessing with Dask The three data sources need to be combined into a unified atmospheric state. This means: - Interpolating to a common grid - Handling missing values - Merging variables from different sources - Quality control This is a perfect fit for Dask to handle lazy evaluation over chunked arrays: \`\`\`python @dask\_env.task async def preprocess\_atmospheric\_data( satellite\_data: File, reanalysis\_data: File, station\_data: File, target\_resolution\_km: float, ) -> File: \`\`\` This task connects to the Dask cluster provisioned by Flyte, loads the datasets with appropriate chunking, merges satellite and reanalysis grids, fills in missing values, and persists the result. Flyte caches the output, so preprocessing only runs when the inputs change. ### GPU-accelerated atmospheric simulation Now the core: running atmospheric physics on the GPU. Each ensemble member is an independent forecast with slightly perturbed initial conditions. \`\`\` @gpu\_env.task async def run\_atmospheric\_simulation( input\_data: File, params: SimulationParams, partition\_id: int = 0, ensemble\_start: int | None = None, ensemble\_end: int | None = None, ) -> tuple\[File, ClimateMetrics\]: """Run GPU-accelerated atmospheric simulation with ensemble forecasting.""" \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* The task accepts a subset of ensemble members (\`ensemble\_start\` to \`ensemble\_end\`). This enables distributing 800 members across multiple GPUs. The physics step is the computational kernel. It runs advection (wind transport), pressure gradients, Coriolis forces, turbulent diffusion, and moisture condensation: \`\`\` @torch.compile(mode="reduce-overhead") def physics\_step(state\_tensor, dt\_val, dx\_val): """Compiled atmospheric physics - 3-4x faster with torch.compile.""" # Advection: transport by wind temp\_grad\_x = torch.roll(state\_tensor\[:, 0\], -1, dims=2) - torch.roll( state\_tensor\[:, 0\], 1, dims=2 ) temp\_grad\_y = torch.roll(state\_tensor\[:, 0\], -1, dims=3) - torch.roll( state\_tensor\[:, 0\], 1, dims=3 ) advection = -( state\_tensor\[:, 3\] \* temp\_grad\_x + state\_tensor\[:, 4\] \* temp\_grad\_y ) / (2 \* dx\_val) state\_tensor\[:, 0\] = state\_tensor\[:, 0\] + advection \* dt\_val # Pressure gradient with Coriolis pressure\_grad\_x = ( torch.roll(state\_tensor\[:, 1\], -1, dims=2) - torch.roll(state\_tensor\[:, 1\], 1, dims=2) ) / (2 \* dx\_val) pressure\_grad\_y = ( torch.roll(state\_tensor\[:, 1\], -1, dims=3) - torch.roll(state\_tensor\[:, 1\], 1, dims=3) ) / (2 \* dx\_val) coriolis\_param = 1e-4 # ~45°N latitude coriolis\_u = coriolis\_param \* state\_tensor\[:, 4\] coriolis\_v = -coriolis\_param \* state\_tensor\[:, 3\] state\_tensor\[:, 3\] = ( state\_tensor\[:, 3\] - pressure\_grad\_x \* dt\_val \* 0.01 + coriolis\_u \* dt\_val ) state\_tensor\[:, 4\] = ( state\_tensor\[:, 4\] - pressure\_grad\_y \* dt\_val \* 0.01 + coriolis\_v \* dt\_val ) # Turbulent diffusion diffusion\_coeff = 10.0 laplacian\_temp = ( torch.roll(state\_tensor\[:, 0\], 1, dims=2) + torch.roll(state\_tensor\[:, 0\], -1, dims=2) + torch.roll(state\_tensor\[:, 0\], 1, dims=3) + torch.roll(state\_tensor\[:, 0\], -1, dims=3) - 4 \* state\_tensor\[:, 0\] ) / (dx\_val \* dx\_val) state\_tensor\[:, 0\] = ( state\_tensor\[:, 0\] + diffusion\_coeff \* laplacian\_temp \* dt\_val ) # Moisture condensation sat\_vapor\_pressure = 611.2 \* torch.exp( 17.67 \* state\_tensor\[:, 0\] / (state\_tensor\[:, 0\] + 243.5) ) condensation = torch.clamp( state\_tensor\[:, 2\] - sat\_vapor\_pressure \* 0.001, min=0 ) state\_tensor\[:, 2\] = state\_tensor\[:, 2\] - condensation \* 0.1 state\_tensor\[:, 0\] = state\_tensor\[:, 0\] + condensation \* 2.5e6 / 1005 \* dt\_val return state\_tensor \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* \`@torch.compile(mode="reduce-overhead")\` compiles this function into optimized CUDA kernels. Combined with mixed precision (\`torch.cuda.amp.autocast\`), this runs 3-4x faster than eager PyTorch. Every 10 timesteps, the simulation checks for extreme events: - \*\*Hurricanes\*\*: Wind speed > 33 m/s with low pressure - \*\*Heatwaves\*\*: Temperature anomalies exceeding thresholds Detected phenomena get logged to the metrics, which flow to the live dashboard. ### Distributing across multiple GPUs 800 ensemble members is a lot for one GPU, so we distribute them: \`\`\` @cpu\_env.task async def run\_distributed\_simulation\_ensemble( preprocessed\_data: File, params: SimulationParams, n\_gpus: int ) -> tuple\[list\[File\], list\[ClimateMetrics\]\]: total\_members = params.ensemble\_size members\_per\_gpu = total\_members // n\_gpus # Distribute ensemble members across GPUs tasks = \[\] for gpu\_id in range(n\_gpus): # Calculate ensemble range for this GPU ensemble\_start = gpu\_id \* members\_per\_gpu # Last GPU gets any remainder members if gpu\_id == n\_gpus - 1: ensemble\_end = total\_members else: ensemble\_end = ensemble\_start + members\_per\_gpu # Launch GPU task with ensemble subset gpu\_task = run\_atmospheric\_simulation( preprocessed\_data, params, gpu\_id, ensemble\_start=ensemble\_start, ensemble\_end=ensemble\_end, ) tasks.append(gpu\_task) # Execute all GPUs in parallel results = await asyncio.gather(\*tasks) output\_files = \[r\[0\] for r in results\] metrics = \[r\[1\] for r in results\] return output\_files, metrics \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* The task splits the ensemble members evenly across the available GPUs, launches the simulation runs in parallel using \`asyncio.gather\`, and then aggregates the results. With five GPUs, each GPU runs 160 ensemble members. Flyte takes care of scheduling, so GPU tasks start automatically as soon as resources become available. ### The main workflow Everything comes together in the orchestration task: \`\`\` @cpu\_env.task(report=True) async def adaptive\_climate\_modeling\_workflow( region: str = "atlantic", date\_range: list\[str, str\] = \["2024-09-01", "2024-09-10"\], current\_params: SimulationParams = SimulationParams(), enable\_multi\_gpu: bool = True, n\_gpus: int = 5, ) -> SimulationSummary: """Orchestrates multi-source ingestion, GPU simulation, and adaptive refinement.""" \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* \`report=True\` enables Flyte Reports for live monitoring. \`\`\` # Parallel data ingestion from three sources with flyte.group("data-ingestion"): satellite\_task = ingest\_satellite\_data(region, date\_range) reanalysis\_task = ingest\_reanalysis\_data(region, date\_range) station\_task = ingest\_station\_data(region, date\_range) satellite\_data, reanalysis\_data, station\_data = await asyncio.gather( satellite\_task, reanalysis\_task, station\_task, ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* \`flyte.group("data-ingestion")\` visually groups the ingestion tasks in the Flyte UI. Inside the group, three tasks launch concurrently. \`asyncio.gather\` waits for all three to complete before preprocessing begins. The workflow then enters an iterative loop: 1. Run GPU simulation (single or multi-GPU) 2. Check convergence by comparing forecasts across iterations 3. Detect extreme events 4. If a hurricane is detected and we haven't refined yet, double the grid resolution 5. Stream metrics to the live dashboard 6. Repeat until converged or max iterations reached Adaptive mesh refinement is the key feature here. When the simulation detects a hurricane forming, it automatically increases resolution to capture the fine-scale dynamics. This is expensive, so we limit it to one refinement per run. ### Running the pipeline \`\`\` if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run\_multi\_gpu = flyte.run(adaptive\_climate\_modeling\_workflow) print(f"Run URL: {run\_multi\_gpu.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/climate\_modeling/simulation.py\* Before running, set up ERA5 API credentials: \`\`\`bash flyte create secret cds\_api\_key flyte create secret cds\_api\_url https://cds.climate.copernicus.eu/api \`\`\` Then launch: \`\`\`bash flyte create config --endpoint --project --domain --builder remote uv run simulation.py \`\`\` The default configuration uses the Atlantic region for September 2024, which is hurricane season. ## Key concepts ### Ensemble forecasting Weather prediction is inherently uncertain. Small errors in the initial conditions grow over time due to chaotic dynamics, which means a single forecast can only ever be one possible outcome. Ensemble forecasting addresses this uncertainty by: - Perturbing the initial conditions within known observational error bounds - Running many independent forecasts - Computing the ensemble mean as the most likely outcome and the ensemble spread as a measure of uncertainty ### Adaptive mesh refinement When a hurricane begins to form, coarse spatial grids are not sufficient to resolve critical features like eyewall dynamics. Adaptive mesh refinement allows the simulation to focus compute where it matters most by: - Increasing grid resolution, for example from 10 km to 5 km - Reducing the timestep to maintain numerical stability - Refining only the regions of interest instead of the entire domain This approach is computationally expensive, but it is essential for producing accurate intensity forecasts. ### Real-time event detection Rather than analyzing results after a simulation completes, this pipeline detects significant events as the simulation runs. The system monitors for conditions such as: - \*\*Hurricanes\*\*: Wind speeds exceeding 33 m/s (Category 1 threshold) combined with central pressure below 980 mb - \*\*Heatwaves\*\*: Sustained temperature anomalies over a defined period Detecting these events in real time enables adaptive responses, such as refining the simulation or triggering alerts, and supports earlier warnings for extreme weather. ## Where to go next This example is intentionally scoped to keep the ideas clear, but there are several natural ways to extend it for more realistic workloads. To model different ocean basins, change the \`region\` parameter to values like \`"pacific"\` or \`"indian"\`. The ingestion tasks automatically adjust to pull the appropriate satellite coverage for each region. To run longer forecasts, increase \`simulation\_hours\` in \`SimulationParams\`. The default of 240 hours, or 10 days, is typical for medium-range forecasting, but you can run longer simulations if you have the compute budget. Finally, the physics step here is deliberately simplified. Production systems usually incorporate additional components such as radiation schemes, boundary layer parameterizations, and land surface models. These can be added incrementally as separate steps without changing the overall structure of the pipeline. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/financial-services === # Financial Services & Fintech Tutorials for financial research, trading, and other fintech workloads. ### \*\*Financial Services & Fintech > Financial research agent\*\* Prep equity briefings for the earnings cycle with grounded You.com Research synthesis and fresh news from the Search API. ### \*\*Financial Services & Fintech > Multi-agent trading simulation\*\* A multi-agent trading simulation, modeling how agents within a firm might interact, strategize, and make trades collaboratively. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/financial-services/trading-agents === # Multi-agent trading simulation > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/trading\_agents); based on work by \[TauricResearch\](https://github.com/TauricResearch/TradingAgents). This example walks you through building a multi-agent trading simulation, modeling how agents within a firm might interact, strategize, and make trades collaboratively. !\[Trading agents execution visualization\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/tutorials/trading-agents/execution.png) \_Trading agents execution visualization\_ ## TL;DR - You'll build a trading firm made up of agents that analyze, argue, and act, modeled with Python functions. - You'll use the Flyte SDK to orchestrate this world — giving you visibility, retries, caching, and durability. - You'll learn how to plug in tools, structure conversations, and track decisions across agents. - You'll see how agents debate, use context, generate reports, and retain memory via vector DBs. ## What is an agent, anyway? Agentic workflows are a rising pattern for complex problem-solving with LLMs. Think of agents as: - An LLM (like GPT-4 or Mistral) - A loop that keeps them thinking until a goal is met - A set of optional tools they can call (APIs, search, calculators, etc.) - Enough tokens to reason about the problem at hand That's it. You define tools, bind them to an agent, and let it run, reasoning step-by-step, optionally using those tools, until it finishes. ## What's different here? We're not building yet another agent framework. You're free to use LangChain, custom code, or whatever setup you like. What we're giving you is the missing piece: a way to run these workflows \*\*reliably, observably, and at scale, with zero rewrites.\*\* With Flyte, you get: - Prompt + tool traceability and full state retention - Built-in retries, caching, and failure recovery - A native way to plug in your agents; no magic syntax required ## How it works: step-by-step walkthrough This simulation is powered by a Flyte task that orchestrates multiple intelligent agents working together to analyze a company's stock and make informed trading decisions. !\[Trading agents schema\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/tutorials/trading-agents/schema.png) \_Trading agents schema\_ ### Entry point Everything begins with a top-level Flyte task called \`main\`, which serves as the entry point to the workflow. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "akshare==1.16.98",\ # "backtrader==1.9.78.123",\ # "boto3==1.39.9",\ # "chainlit==2.5.5",\ # "eodhd==1.0.32",\ # "feedparser==6.0.11",\ # "finnhub-python==2.4.23",\ # "langchain-experimental==0.3.4",\ # "langchain-openai==0.3.23",\ # "pandas==2.3.0",\ # "parsel==1.10.0",\ # "praw==7.8.1",\ # "pytz==2025.2",\ # "questionary==2.1.0",\ # "redis==6.2.0",\ # "requests==2.32.4",\ # "stockstats==0.6.5",\ # "tqdm==4.67.1",\ # "tushare==1.4.21",\ # "typing-extensions==4.14.0",\ # "yfinance==0.2.63",\ # \] # main = "main" # params = "" # /// import asyncio from copy import deepcopy import agents import agents.analysts from agents.managers import create\_research\_manager, create\_risk\_manager from agents.researchers import create\_bear\_researcher, create\_bull\_researcher from agents.risk\_debators import ( create\_neutral\_debator, create\_risky\_debator, create\_safe\_debator, ) from agents.trader import create\_trader from agents.utils.utils import AgentState from flyte\_env import DEEP\_THINKING\_LLM, QUICK\_THINKING\_LLM, env, flyte from langchain\_openai import ChatOpenAI from reflection import ( reflect\_bear\_researcher, reflect\_bull\_researcher, reflect\_research\_manager, reflect\_risk\_manager, reflect\_trader, ) @env.task async def process\_signal(full\_signal: str, QUICK\_THINKING\_LLM: str) -> str: """Process a full trading signal to extract the core decision.""" messages = \[\ {\ "role": "system",\ "content": """You are an efficient assistant designed to analyze paragraphs or\ financial reports provided by a group of analysts.\ Your task is to extract the investment decision: SELL, BUY, or HOLD.\ Provide only the extracted decision (SELL, BUY, or HOLD) as your output,\ without adding any additional text or information.""",\ },\ {"role": "human", "content": full\_signal},\ \] return ChatOpenAI(model=QUICK\_THINKING\_LLM).invoke(messages).content async def run\_analyst(analyst\_name, state, online\_tools): # Create a copy of the state for isolation run\_fn = getattr(agents.analysts, f"create\_{analyst\_name}\_analyst") # Run the analyst's chain result\_state = await run\_fn(QUICK\_THINKING\_LLM, state, online\_tools) # Determine the report key report\_key = ( "sentiment\_report" if analyst\_name == "social\_media" else f"{analyst\_name}\_report" ) report\_value = getattr(result\_state, report\_key) return result\_state.messages\[1:\], report\_key, report\_value # {{docs-fragment main}} @env.task async def main( selected\_analysts: list\[str\] = \[\ "market",\ "fundamentals",\ "news",\ "social\_media",\ \], max\_debate\_rounds: int = 1, max\_risk\_discuss\_rounds: int = 1, online\_tools: bool = True, company\_name: str = "NVDA", trade\_date: str = "2024-05-12", ) -> tuple\[str, AgentState\]: if not selected\_analysts: raise ValueError( "No analysts selected. Please select at least one analyst from market, fundamentals, news, or social\_media." ) state = AgentState( messages=\[{"role": "human", "content": company\_name}\], company\_of\_interest=company\_name, trade\_date=str(trade\_date), ) # Run all analysts concurrently results = await asyncio.gather( \*\[\ run\_analyst(analyst, deepcopy(state), online\_tools)\ for analyst in selected\_analysts\ \] ) # Flatten and append all resulting messages into the shared state for messages, report\_attr, report in results: state.messages.extend(messages) setattr(state, report\_attr, report) # Bull/Bear debate loop state = await create\_bull\_researcher(QUICK\_THINKING\_LLM, state) # Start with bull while state.investment\_debate\_state.count < 2 \* max\_debate\_rounds: current = state.investment\_debate\_state.current\_response if current.startswith("Bull"): state = await create\_bear\_researcher(QUICK\_THINKING\_LLM, state) else: state = await create\_bull\_researcher(QUICK\_THINKING\_LLM, state) state = await create\_research\_manager(DEEP\_THINKING\_LLM, state) state = await create\_trader(QUICK\_THINKING\_LLM, state) # Risk debate loop state = await create\_risky\_debator(QUICK\_THINKING\_LLM, state) # Start with risky while state.risk\_debate\_state.count < 3 \* max\_risk\_discuss\_rounds: speaker = state.risk\_debate\_state.latest\_speaker if speaker == "Risky": state = await create\_safe\_debator(QUICK\_THINKING\_LLM, state) elif speaker == "Safe": state = await create\_neutral\_debator(QUICK\_THINKING\_LLM, state) else: state = await create\_risky\_debator(QUICK\_THINKING\_LLM, state) state = await create\_risk\_manager(DEEP\_THINKING\_LLM, state) decision = await process\_signal(state.final\_trade\_decision, QUICK\_THINKING\_LLM) return decision, state # {{/docs-fragment main}} # {{docs-fragment reflect\_on\_decisions}} @env.task async def reflect\_and\_store(state: AgentState, returns: str) -> str: await asyncio.gather( reflect\_bear\_researcher(state, returns), reflect\_bull\_researcher(state, returns), reflect\_trader(state, returns), reflect\_risk\_manager(state, returns), reflect\_research\_manager(state, returns), ) return "Reflection completed." # Run the reflection task after the main function @env.task(cache="disable") async def reflect\_on\_decisions( returns: str, selected\_analysts: list\[str\] = \[\ "market",\ "fundamentals",\ "news",\ "social\_media",\ \], max\_debate\_rounds: int = 1, max\_risk\_discuss\_rounds: int = 1, online\_tools: bool = True, company\_name: str = "NVDA", trade\_date: str = "2024-05-12", ) -> str: \_, state = await main( selected\_analysts, max\_debate\_rounds, max\_risk\_discuss\_rounds, online\_tools, company\_name, trade\_date, ) return await reflect\_and\_store(state, returns) # {{/docs-fragment reflect\_on\_decisions}} # {{docs-fragment execute\_main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # run = flyte.run(reflect\_on\_decisions, "+3.2% gain over 5 days") # print(run.url) # {{/docs-fragment execute\_main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/main.py\* This task accepts several inputs: - the list of analysts to run, - the number of debate and risk discussion rounds, - a flag to enable online tools, - the company you're evaluating, - and the target trading date. The most interesting parameter here is the list of analysts to run. It determines which analyst agents will be invoked and shapes the overall structure of the simulation. Based on this input, the task dynamically launches agent tasks, running them in parallel. The \`main\` task is written as a regular asynchronous Python function wrapped with Flyte's task decorator. No domain-specific language or orchestration glue is needed — just idiomatic Python, optionally using async for better performance. The task environment is configured once and shared across all tasks for consistency. \`\`\` # {{docs-fragment env}} import flyte QUICK\_THINKING\_LLM = "gpt-4o-mini" DEEP\_THINKING\_LLM = "o4-mini" env = flyte.TaskEnvironment( name="trading-agents", secrets=\[\ flyte.Secret(key="finnhub\_api\_key", as\_env\_var="FINNHUB\_API\_KEY"),\ flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script("main.py", name="trading-agents", pre=True), resources=flyte.Resources(cpu="1"), cache="auto", ) # {{/docs-fragment env}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/flyte\_env.py\* ### Analyst agents Each analyst agent comes equipped with a set of tools and a carefully designed prompt tailored to its specific domain. These tools are modular Flyte tasks — for example, downloading financial reports or computing technical indicators — and benefit from Flyte's built-in caching to avoid redundant computation. \`\`\` from datetime import datetime import pandas as pd import tools.interface as interface import yfinance as yf from flyte\_env import env from flyte.io import File @env.task async def get\_reddit\_news( curr\_date: str, # Date you want to get news for in yyyy-mm-dd format ) -> str: """ Retrieve global news from Reddit within a specified time frame. Args: curr\_date (str): Date you want to get news for in yyyy-mm-dd format Returns: str: A formatted dataframe containing the latest global news from Reddit in the specified time frame. """ global\_news\_result = interface.get\_reddit\_global\_news(curr\_date, 7, 5) return global\_news\_result @env.task async def get\_finnhub\_news( ticker: str, # Search query of a company, e.g. 'AAPL, TSM, etc. start\_date: str, # Start date in yyyy-mm-dd format end\_date: str, # End date in yyyy-mm-dd format ) -> str: """ Retrieve the latest news about a given stock from Finnhub within a date range Args: ticker (str): Ticker of a company. e.g. AAPL, TSM start\_date (str): Start date in yyyy-mm-dd format end\_date (str): End date in yyyy-mm-dd format Returns: str: A formatted dataframe containing news about the company within the date range from start\_date to end\_date """ end\_date\_str = end\_date end\_date = datetime.strptime(end\_date, "%Y-%m-%d") start\_date = datetime.strptime(start\_date, "%Y-%m-%d") look\_back\_days = (end\_date - start\_date).days finnhub\_news\_result = interface.get\_finnhub\_news( ticker, end\_date\_str, look\_back\_days ) return finnhub\_news\_result @env.task async def get\_reddit\_stock\_info( ticker: str, # Ticker of a company. e.g. AAPL, TSM curr\_date: str, # Current date you want to get news for ) -> str: """ Retrieve the latest news about a given stock from Reddit, given the current date. Args: ticker (str): Ticker of a company. e.g. AAPL, TSM curr\_date (str): current date in yyyy-mm-dd format to get news for Returns: str: A formatted dataframe containing the latest news about the company on the given date """ stock\_news\_results = interface.get\_reddit\_company\_news(ticker, curr\_date, 7, 5) return stock\_news\_results @env.task async def get\_YFin\_data( symbol: str, # ticker symbol of the company start\_date: str, # Start date in yyyy-mm-dd format end\_date: str, # End date in yyyy-mm-dd format ) -> str: """ Retrieve the stock price data for a given ticker symbol from Yahoo Finance. Args: symbol (str): Ticker symbol of the company, e.g. AAPL, TSM start\_date (str): Start date in yyyy-mm-dd format end\_date (str): End date in yyyy-mm-dd format Returns: str: A formatted dataframe containing the stock price data for the specified ticker symbol in the specified date range. """ result\_data = interface.get\_YFin\_data(symbol, start\_date, end\_date) return result\_data @env.task async def get\_YFin\_data\_online( symbol: str, # ticker symbol of the company start\_date: str, # Start date in yyyy-mm-dd format end\_date: str, # End date in yyyy-mm-dd format ) -> str: """ Retrieve the stock price data for a given ticker symbol from Yahoo Finance. Args: symbol (str): Ticker symbol of the company, e.g. AAPL, TSM start\_date (str): Start date in yyyy-mm-dd format end\_date (str): End date in yyyy-mm-dd format Returns: str: A formatted dataframe containing the stock price data for the specified ticker symbol in the specified date range. """ result\_data = interface.get\_YFin\_data\_online(symbol, start\_date, end\_date) return result\_data @env.task async def cache\_market\_data(symbol: str, start\_date: str, end\_date: str) -> File: data\_file = f"{symbol}-YFin-data-{start\_date}-{end\_date}.csv" data = yf.download( symbol, start=start\_date, end=end\_date, multi\_level\_index=False, progress=False, auto\_adjust=True, ) data = data.reset\_index() data.to\_csv(data\_file, index=False) return await File.from\_local(data\_file) @env.task async def get\_stockstats\_indicators\_report( symbol: str, # ticker symbol of the company indicator: str, # technical indicator to get the analysis and report of curr\_date: str, # The current trading date you are trading on, YYYY-mm-dd look\_back\_days: int = 30, # how many days to look back ) -> str: """ Retrieve stock stats indicators for a given ticker symbol and indicator. Args: symbol (str): Ticker symbol of the company, e.g. AAPL, TSM indicator (str): Technical indicator to get the analysis and report of curr\_date (str): The current trading date you are trading on, YYYY-mm-dd look\_back\_days (int): How many days to look back, default is 30 Returns: str: A formatted dataframe containing the stock stats indicators for the specified ticker symbol and indicator. """ today\_date = pd.Timestamp.today() end\_date = today\_date start\_date = today\_date - pd.DateOffset(years=15) start\_date = start\_date.strftime("%Y-%m-%d") end\_date = end\_date.strftime("%Y-%m-%d") data\_file = await cache\_market\_data(symbol, start\_date, end\_date) local\_data\_file = await data\_file.download() result\_stockstats = interface.get\_stock\_stats\_indicators\_window( symbol, indicator, curr\_date, look\_back\_days, False, local\_data\_file ) return result\_stockstats # {{docs-fragment get\_stockstats\_indicators\_report\_online}} @env.task async def get\_stockstats\_indicators\_report\_online( symbol: str, # ticker symbol of the company indicator: str, # technical indicator to get the analysis and report of curr\_date: str, # The current trading date you are trading on, YYYY-mm-dd" look\_back\_days: int = 30, # "how many days to look back" ) -> str: """ Retrieve stock stats indicators for a given ticker symbol and indicator. Args: symbol (str): Ticker symbol of the company, e.g. AAPL, TSM indicator (str): Technical indicator to get the analysis and report of curr\_date (str): The current trading date you are trading on, YYYY-mm-dd look\_back\_days (int): How many days to look back, default is 30 Returns: str: A formatted dataframe containing the stock stats indicators for the specified ticker symbol and indicator. """ today\_date = pd.Timestamp.today() end\_date = today\_date start\_date = today\_date - pd.DateOffset(years=15) start\_date = start\_date.strftime("%Y-%m-%d") end\_date = end\_date.strftime("%Y-%m-%d") data\_file = await cache\_market\_data(symbol, start\_date, end\_date) local\_data\_file = await data\_file.download() result\_stockstats = interface.get\_stock\_stats\_indicators\_window( symbol, indicator, curr\_date, look\_back\_days, True, local\_data\_file ) return result\_stockstats # {{/docs-fragment get\_stockstats\_indicators\_report\_online}} @env.task async def get\_finnhub\_company\_insider\_sentiment( ticker: str, # ticker symbol for the company curr\_date: str, # current date of you are trading at, yyyy-mm-dd ) -> str: """ Retrieve insider sentiment information about a company (retrieved from public SEC information) for the past 30 days Args: ticker (str): ticker symbol of the company curr\_date (str): current date you are trading at, yyyy-mm-dd Returns: str: a report of the sentiment in the past 30 days starting at curr\_date """ data\_sentiment = interface.get\_finnhub\_company\_insider\_sentiment( ticker, curr\_date, 30 ) return data\_sentiment @env.task async def get\_finnhub\_company\_insider\_transactions( ticker: str, # ticker symbol curr\_date: str, # current date you are trading at, yyyy-mm-dd ) -> str: """ Retrieve insider transaction information about a company (retrieved from public SEC information) for the past 30 days Args: ticker (str): ticker symbol of the company curr\_date (str): current date you are trading at, yyyy-mm-dd Returns: str: a report of the company's insider transactions/trading information in the past 30 days """ data\_trans = interface.get\_finnhub\_company\_insider\_transactions( ticker, curr\_date, 30 ) return data\_trans @env.task async def get\_simfin\_balance\_sheet( ticker: str, # ticker symbol freq: str, # reporting frequency of the company's financial history: annual/quarterly curr\_date: str, # current date you are trading at, yyyy-mm-dd ): """ Retrieve the most recent balance sheet of a company Args: ticker (str): ticker symbol of the company freq (str): reporting frequency of the company's financial history: annual / quarterly curr\_date (str): current date you are trading at, yyyy-mm-dd Returns: str: a report of the company's most recent balance sheet """ data\_balance\_sheet = interface.get\_simfin\_balance\_sheet(ticker, freq, curr\_date) return data\_balance\_sheet @env.task async def get\_simfin\_cashflow( ticker: str, # ticker symbol freq: str, # reporting frequency of the company's financial history: annual/quarterly curr\_date: str, # current date you are trading at, yyyy-mm-dd ) -> str: """ Retrieve the most recent cash flow statement of a company Args: ticker (str): ticker symbol of the company freq (str): reporting frequency of the company's financial history: annual / quarterly curr\_date (str): current date you are trading at, yyyy-mm-dd Returns: str: a report of the company's most recent cash flow statement """ data\_cashflow = interface.get\_simfin\_cashflow(ticker, freq, curr\_date) return data\_cashflow @env.task async def get\_simfin\_income\_stmt( ticker: str, # ticker symbol freq: str, # reporting frequency of the company's financial history: annual/quarterly curr\_date: str, # current date you are trading at, yyyy-mm-dd ) -> str: """ Retrieve the most recent income statement of a company Args: ticker (str): ticker symbol of the company freq (str): reporting frequency of the company's financial history: annual / quarterly curr\_date (str): current date you are trading at, yyyy-mm-dd Returns: str: a report of the company's most recent income statement """ data\_income\_stmt = interface.get\_simfin\_income\_statements(ticker, freq, curr\_date) return data\_income\_stmt @env.task async def get\_google\_news( query: str, # Query to search with curr\_date: str, # Curr date in yyyy-mm-dd format ) -> str: """ Retrieve the latest news from Google News based on a query and date range. Args: query (str): Query to search with curr\_date (str): Current date in yyyy-mm-dd format look\_back\_days (int): How many days to look back Returns: str: A formatted string containing the latest news from Google News based on the query and date range. """ google\_news\_results = interface.get\_google\_news(query, curr\_date, 7) return google\_news\_results @env.task async def get\_stock\_news\_openai( ticker: str, # the company's ticker curr\_date: str, # Current date in yyyy-mm-dd format ) -> str: """ Retrieve the latest news about a given stock by using OpenAI's news API. Args: ticker (str): Ticker of a company. e.g. AAPL, TSM curr\_date (str): Current date in yyyy-mm-dd format Returns: str: A formatted string containing the latest news about the company on the given date. """ openai\_news\_results = interface.get\_stock\_news\_openai(ticker, curr\_date) return openai\_news\_results @env.task async def get\_global\_news\_openai( curr\_date: str, # Current date in yyyy-mm-dd format ) -> str: """ Retrieve the latest macroeconomics news on a given date using OpenAI's macroeconomics news API. Args: curr\_date (str): Current date in yyyy-mm-dd format Returns: str: A formatted string containing the latest macroeconomic news on the given date. """ openai\_news\_results = interface.get\_global\_news\_openai(curr\_date) return openai\_news\_results @env.task async def get\_fundamentals\_openai( ticker: str, # the company's ticker curr\_date: str, # Current date in yyyy-mm-dd format ) -> str: """ Retrieve the latest fundamental information about a given stock on a given date by using OpenAI's news API. Args: ticker (str): Ticker of a company. e.g. AAPL, TSM curr\_date (str): Current date in yyyy-mm-dd format Returns: str: A formatted string containing the latest fundamental information about the company on the given date. """ openai\_fundamentals\_results = interface.get\_fundamentals\_openai(ticker, curr\_date) return openai\_fundamentals\_results \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/tools/toolkit.py\* When initialized, an analyst enters a structured reasoning loop (via LangChain), where it can call tools, observe outputs, and refine its internal state before generating a final report. These reports are later consumed by downstream agents. Here's an example of a news analyst that interprets global events and macroeconomic signals. We specify the tools accessible to the analyst, and the LLM selects which ones to use based on context. \`\`\` import asyncio from agents.utils.utils import AgentState from flyte\_env import env from langchain\_core.messages import ToolMessage, convert\_to\_openai\_messages from langchain\_core.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain\_openai import ChatOpenAI from tools import toolkit import flyte MAX\_ITERATIONS = 5 # {{docs-fragment agent\_helper}} async def run\_chain\_with\_tools( type: str, state: AgentState, llm: str, system\_message: str, tool\_names: list\[str\] ) -> AgentState: prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ "You are a helpful AI assistant, collaborating with other assistants."\ " Use the provided tools to progress towards answering the question."\ " If you are unable to fully answer, that's OK; another assistant with different tools"\ " will help where you left off. Execute what you can to make progress."\ " If you or any other assistant has the FINAL TRANSACTION PROPOSAL: \*\*BUY/HOLD/SELL\*\* or deliverable,"\ " prefix your response with FINAL TRANSACTION PROPOSAL: \*\*BUY/HOLD/SELL\*\* so the team knows to stop."\ " You have access to the following tools: {tool\_names}.\\n{system\_message}"\ " For your reference, the current date is {current\_date}. The company we want to look at is {ticker}.",\ ),\ MessagesPlaceholder(variable\_name="messages"),\ \] ) prompt = prompt.partial(system\_message=system\_message) prompt = prompt.partial(tool\_names=", ".join(tool\_names)) prompt = prompt.partial(current\_date=state.trade\_date) prompt = prompt.partial(ticker=state.company\_of\_interest) chain = prompt | ChatOpenAI(model=llm).bind\_tools( \[getattr(toolkit, tool\_name).func for tool\_name in tool\_names\] ) iteration = 0 while iteration < MAX\_ITERATIONS: result = await chain.ainvoke(state.messages) state.messages.append(convert\_to\_openai\_messages(result)) if not result.tool\_calls: # Final response — no tools required setattr(state, f"{type}\_report", result.content or "") break # Run all tool calls in parallel async def run\_single\_tool(tool\_call): tool\_name = tool\_call\["name"\] tool\_args = tool\_call\["args"\] tool = getattr(toolkit, tool\_name, None) if not tool: return None content = await tool(\*\*tool\_args) return ToolMessage( tool\_call\_id=tool\_call\["id"\], name=tool\_name, content=content ) with flyte.group(f"tool\_calls\_iteration\_{iteration}"): tool\_messages = await asyncio.gather( \*\[run\_single\_tool(tc) for tc in result.tool\_calls\] ) # Add valid tool results to state tool\_messages = \[msg for msg in tool\_messages if msg\] state.messages.extend(convert\_to\_openai\_messages(tool\_messages)) iteration += 1 else: # Reached iteration cap — optionally raise or log print(f"Max iterations ({MAX\_ITERATIONS}) reached for {type}") return state # {{/docs-fragment agent\_helper}} @env.task async def create\_fundamentals\_analyst( llm: str, state: AgentState, online\_tools: bool ) -> AgentState: if online\_tools: tools = \[toolkit.get\_fundamentals\_openai\] else: tools = \[\ toolkit.get\_finnhub\_company\_insider\_sentiment,\ toolkit.get\_finnhub\_company\_insider\_transactions,\ toolkit.get\_simfin\_balance\_sheet,\ toolkit.get\_simfin\_cashflow,\ toolkit.get\_simfin\_income\_stmt,\ \] system\_message = ( "You are a researcher tasked with analyzing fundamental information over the past week about a company. " "Please write a comprehensive report of the company's fundamental information such as financial documents, " "company profile, basic company financials, company financial history, insider sentiment, and insider " "transactions to gain a full view of the company's " "fundamental information to inform traders. Make sure to include as much detail as possible. " "Do not simply state the trends are mixed, " "provide detailed and finegrained analysis and insights that may help traders make decisions. " "Make sure to append a Markdown table at the end of the report to organize key points in the report, " "organized and easy to read." ) tool\_names = \[tool.func.\_\_name\_\_ for tool in tools\] return await run\_chain\_with\_tools( "fundamentals", state, llm, system\_message, tool\_names ) @env.task async def create\_market\_analyst( llm: str, state: AgentState, online\_tools: bool ) -> AgentState: if online\_tools: tools = \[\ toolkit.get\_YFin\_data\_online,\ toolkit.get\_stockstats\_indicators\_report\_online,\ \] else: tools = \[\ toolkit.get\_YFin\_data,\ toolkit.get\_stockstats\_indicators\_report,\ \] system\_message = ( """You are a trading assistant tasked with analyzing financial markets. Your role is to select the \*\*most relevant indicators\*\* for a given market condition or trading strategy from the following list. The goal is to choose up to \*\*8 indicators\*\* that provide complementary insights without redundancy. Categories and each category's indicators are: Moving Averages: - close\_50\_sma: 50 SMA: A medium-term trend indicator. Usage: Identify trend direction and serve as dynamic support/resistance. Tips: It lags price; combine with faster indicators for timely signals. - close\_200\_sma: 200 SMA: A long-term trend benchmark. Usage: Confirm overall market trend and identify golden/death cross setups. Tips: It reacts slowly; best for strategic trend confirmation rather than frequent trading entries. - close\_10\_ema: 10 EMA: A responsive short-term average. Usage: Capture quick shifts in momentum and potential entry points. Tips: Prone to noise in choppy markets; use alongside longer averages for filtering false signals. MACD Related: - macd: MACD: Computes momentum via differences of EMAs. Usage: Look for crossovers and divergence as signals of trend changes. Tips: Confirm with other indicators in low-volatility or sideways markets. - macds: MACD Signal: An EMA smoothing of the MACD line. Usage: Use crossovers with the MACD line to trigger trades. Tips: Should be part of a broader strategy to avoid false positives. - macdh: MACD Histogram: Shows the gap between the MACD line and its signal. Usage: Visualize momentum strength and spot divergence early. Tips: Can be volatile; complement with additional filters in fast-moving markets. Momentum Indicators: - rsi: RSI: Measures momentum to flag overbought/oversold conditions. Usage: Apply 70/30 thresholds and watch for divergence to signal reversals. Tips: In strong trends, RSI may remain extreme; always cross-check with trend analysis. Volatility Indicators: - boll: Bollinger Middle: A 20 SMA serving as the basis for Bollinger Bands. Usage: Acts as a dynamic benchmark for price movement. Tips: Combine with the upper and lower bands to effectively spot breakouts or reversals. - boll\_ub: Bollinger Upper Band: Typically 2 standard deviations above the middle line. Usage: Signals potential overbought conditions and breakout zones. Tips: Confirm signals with other tools; prices may ride the band in strong trends. - boll\_lb: Bollinger Lower Band: Typically 2 standard deviations below the middle line. Usage: Indicates potential oversold conditions. Tips: Use additional analysis to avoid false reversal signals. - atr: ATR: Averages true range to measure volatility. Usage: Set stop-loss levels and adjust position sizes based on current market volatility. Tips: It's a reactive measure, so use it as part of a broader risk management strategy. Volume-Based Indicators: - vwma: VWMA: A moving average weighted by volume. Usage: Confirm trends by integrating price action with volume data. Tips: Watch for skewed results from volume spikes; use in combination with other volume analyses. - Select indicators that provide diverse and complementary information. Avoid redundancy (e.g., do not select both rsi and stochrsi). Also briefly explain why they are suitable for the given market context. When you tool call, please use the exact name of the indicators provided above as they are defined parameters, otherwise your call will fail. Please make sure to call get\_YFin\_data first to retrieve the CSV that is needed to generate indicators. Write a very detailed and nuanced report of the trends you observe. Do not simply state the trends are mixed, provide detailed and finegrained analysis and insights that may help traders make decisions.""" """ Make sure to append a Markdown table at the end of the report to organize key points in the report, organized and easy to read.""" ) tool\_names = \[tool.func.\_\_name\_\_ for tool in tools\] return await run\_chain\_with\_tools("market", state, llm, system\_message, tool\_names) # {{docs-fragment news\_analyst}} @env.task async def create\_news\_analyst( llm: str, state: AgentState, online\_tools: bool ) -> AgentState: if online\_tools: tools = \[\ toolkit.get\_global\_news\_openai,\ toolkit.get\_google\_news,\ \] else: tools = \[\ toolkit.get\_finnhub\_news,\ toolkit.get\_reddit\_news,\ toolkit.get\_google\_news,\ \] system\_message = ( "You are a news researcher tasked with analyzing recent news and trends over the past week. " "Please write a comprehensive report of the current state of the world that is relevant for " "trading and macroeconomics. " "Look at news from EODHD, and finnhub to be comprehensive. Do not simply state the trends are mixed, " "provide detailed and finegrained analysis and insights that may help traders make decisions." """ Make sure to append a Markdown table at the end of the report to organize key points in the report, organized and easy to read.""" ) tool\_names = \[tool.func.\_\_name\_\_ for tool in tools\] return await run\_chain\_with\_tools("news", state, llm, system\_message, tool\_names) # {{/docs-fragment news\_analyst}} @env.task async def create\_social\_media\_analyst( llm: str, state: AgentState, online\_tools: bool ) -> AgentState: if online\_tools: tools = \[toolkit.get\_stock\_news\_openai\] else: tools = \[toolkit.get\_reddit\_stock\_info\] system\_message = ( "You are a social media and company specific news researcher/analyst tasked with analyzing social media posts, " "recent company news, and public sentiment for a specific company over the past week. " "You will be given a company's name your objective is to write a comprehensive long report " "detailing your analysis, insights, and implications for traders and investors on this company's current state " "after looking at social media and what people are saying about that company, " "analyzing sentiment data of what people feel each day about the company, and looking at recent company news. " "Try to look at all sources possible from social media to sentiment to news. Do not simply state the trends " "are mixed, provide detailed and finegrained analysis and insights that may help traders make decisions." """ Make sure to append a Makrdown table at the end of the report to organize key points in the report, organized and easy to read.""" ) tool\_names = \[tool.func.\_\_name\_\_ for tool in tools\] return await run\_chain\_with\_tools( "sentiment", state, llm, system\_message, tool\_names ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/agents/analysts.py\* Each analyst agent uses a helper function to bind tools, iterate through reasoning steps (up to a configurable maximum), and produce an answer. Setting a max iteration count is crucial to prevent runaway loops. As agents reason, their message history is preserved in their internal state and passed along to the next agent in the chain. \`\`\` import asyncio from agents.utils.utils import AgentState from flyte\_env import env from langchain\_core.messages import ToolMessage, convert\_to\_openai\_messages from langchain\_core.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain\_openai import ChatOpenAI from tools import toolkit import flyte MAX\_ITERATIONS = 5 # {{docs-fragment agent\_helper}} async def run\_chain\_with\_tools( type: str, state: AgentState, llm: str, system\_message: str, tool\_names: list\[str\] ) -> AgentState: prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ "You are a helpful AI assistant, collaborating with other assistants."\ " Use the provided tools to progress towards answering the question."\ " If you are unable to fully answer, that's OK; another assistant with different tools"\ " will help where you left off. Execute what you can to make progress."\ " If you or any other assistant has the FINAL TRANSACTION PROPOSAL: \*\*BUY/HOLD/SELL\*\* or deliverable,"\ " prefix your response with FINAL TRANSACTION PROPOSAL: \*\*BUY/HOLD/SELL\*\* so the team knows to stop."\ " You have access to the following tools: {tool\_names}.\\n{system\_message}"\ " For your reference, the current date is {current\_date}. The company we want to look at is {ticker}.",\ ),\ MessagesPlaceholder(variable\_name="messages"),\ \] ) prompt = prompt.partial(system\_message=system\_message) prompt = prompt.partial(tool\_names=", ".join(tool\_names)) prompt = prompt.partial(current\_date=state.trade\_date) prompt = prompt.partial(ticker=state.company\_of\_interest) chain = prompt | ChatOpenAI(model=llm).bind\_tools( \[getattr(toolkit, tool\_name).func for tool\_name in tool\_names\] ) iteration = 0 while iteration < MAX\_ITERATIONS: result = await chain.ainvoke(state.messages) state.messages.append(convert\_to\_openai\_messages(result)) if not result.tool\_calls: # Final response — no tools required setattr(state, f"{type}\_report", result.content or "") break # Run all tool calls in parallel async def run\_single\_tool(tool\_call): tool\_name = tool\_call\["name"\] tool\_args = tool\_call\["args"\] tool = getattr(toolkit, tool\_name, None) if not tool: return None content = await tool(\*\*tool\_args) return ToolMessage( tool\_call\_id=tool\_call\["id"\], name=tool\_name, content=content ) with flyte.group(f"tool\_calls\_iteration\_{iteration}"): tool\_messages = await asyncio.gather( \*\[run\_single\_tool(tc) for tc in result.tool\_calls\] ) # Add valid tool results to state tool\_messages = \[msg for msg in tool\_messages if msg\] state.messages.extend(convert\_to\_openai\_messages(tool\_messages)) iteration += 1 else: # Reached iteration cap — optionally raise or log print(f"Max iterations ({MAX\_ITERATIONS}) reached for {type}") return state # {{/docs-fragment agent\_helper}} @env.task async def create\_fundamentals\_analyst( llm: str, state: AgentState, online\_tools: bool ) -> AgentState: if online\_tools: tools = \[toolkit.get\_fundamentals\_openai\] else: tools = \[\ toolkit.get\_finnhub\_company\_insider\_sentiment,\ toolkit.get\_finnhub\_company\_insider\_transactions,\ toolkit.get\_simfin\_balance\_sheet,\ toolkit.get\_simfin\_cashflow,\ toolkit.get\_simfin\_income\_stmt,\ \] system\_message = ( "You are a researcher tasked with analyzing fundamental information over the past week about a company. " "Please write a comprehensive report of the company's fundamental information such as financial documents, " "company profile, basic company financials, company financial history, insider sentiment, and insider " "transactions to gain a full view of the company's " "fundamental information to inform traders. Make sure to include as much detail as possible. " "Do not simply state the trends are mixed, " "provide detailed and finegrained analysis and insights that may help traders make decisions. " "Make sure to append a Markdown table at the end of the report to organize key points in the report, " "organized and easy to read." ) tool\_names = \[tool.func.\_\_name\_\_ for tool in tools\] return await run\_chain\_with\_tools( "fundamentals", state, llm, system\_message, tool\_names ) @env.task async def create\_market\_analyst( llm: str, state: AgentState, online\_tools: bool ) -> AgentState: if online\_tools: tools = \[\ toolkit.get\_YFin\_data\_online,\ toolkit.get\_stockstats\_indicators\_report\_online,\ \] else: tools = \[\ toolkit.get\_YFin\_data,\ toolkit.get\_stockstats\_indicators\_report,\ \] system\_message = ( """You are a trading assistant tasked with analyzing financial markets. Your role is to select the \*\*most relevant indicators\*\* for a given market condition or trading strategy from the following list. The goal is to choose up to \*\*8 indicators\*\* that provide complementary insights without redundancy. Categories and each category's indicators are: Moving Averages: - close\_50\_sma: 50 SMA: A medium-term trend indicator. Usage: Identify trend direction and serve as dynamic support/resistance. Tips: It lags price; combine with faster indicators for timely signals. - close\_200\_sma: 200 SMA: A long-term trend benchmark. Usage: Confirm overall market trend and identify golden/death cross setups. Tips: It reacts slowly; best for strategic trend confirmation rather than frequent trading entries. - close\_10\_ema: 10 EMA: A responsive short-term average. Usage: Capture quick shifts in momentum and potential entry points. Tips: Prone to noise in choppy markets; use alongside longer averages for filtering false signals. MACD Related: - macd: MACD: Computes momentum via differences of EMAs. Usage: Look for crossovers and divergence as signals of trend changes. Tips: Confirm with other indicators in low-volatility or sideways markets. - macds: MACD Signal: An EMA smoothing of the MACD line. Usage: Use crossovers with the MACD line to trigger trades. Tips: Should be part of a broader strategy to avoid false positives. - macdh: MACD Histogram: Shows the gap between the MACD line and its signal. Usage: Visualize momentum strength and spot divergence early. Tips: Can be volatile; complement with additional filters in fast-moving markets. Momentum Indicators: - rsi: RSI: Measures momentum to flag overbought/oversold conditions. Usage: Apply 70/30 thresholds and watch for divergence to signal reversals. Tips: In strong trends, RSI may remain extreme; always cross-check with trend analysis. Volatility Indicators: - boll: Bollinger Middle: A 20 SMA serving as the basis for Bollinger Bands. Usage: Acts as a dynamic benchmark for price movement. Tips: Combine with the upper and lower bands to effectively spot breakouts or reversals. - boll\_ub: Bollinger Upper Band: Typically 2 standard deviations above the middle line. Usage: Signals potential overbought conditions and breakout zones. Tips: Confirm signals with other tools; prices may ride the band in strong trends. - boll\_lb: Bollinger Lower Band: Typically 2 standard deviations below the middle line. Usage: Indicates potential oversold conditions. Tips: Use additional analysis to avoid false reversal signals. - atr: ATR: Averages true range to measure volatility. Usage: Set stop-loss levels and adjust position sizes based on current market volatility. Tips: It's a reactive measure, so use it as part of a broader risk management strategy. Volume-Based Indicators: - vwma: VWMA: A moving average weighted by volume. Usage: Confirm trends by integrating price action with volume data. Tips: Watch for skewed results from volume spikes; use in combination with other volume analyses. - Select indicators that provide diverse and complementary information. Avoid redundancy (e.g., do not select both rsi and stochrsi). Also briefly explain why they are suitable for the given market context. When you tool call, please use the exact name of the indicators provided above as they are defined parameters, otherwise your call will fail. Please make sure to call get\_YFin\_data first to retrieve the CSV that is needed to generate indicators. Write a very detailed and nuanced report of the trends you observe. Do not simply state the trends are mixed, provide detailed and finegrained analysis and insights that may help traders make decisions.""" """ Make sure to append a Markdown table at the end of the report to organize key points in the report, organized and easy to read.""" ) tool\_names = \[tool.func.\_\_name\_\_ for tool in tools\] return await run\_chain\_with\_tools("market", state, llm, system\_message, tool\_names) # {{docs-fragment news\_analyst}} @env.task async def create\_news\_analyst( llm: str, state: AgentState, online\_tools: bool ) -> AgentState: if online\_tools: tools = \[\ toolkit.get\_global\_news\_openai,\ toolkit.get\_google\_news,\ \] else: tools = \[\ toolkit.get\_finnhub\_news,\ toolkit.get\_reddit\_news,\ toolkit.get\_google\_news,\ \] system\_message = ( "You are a news researcher tasked with analyzing recent news and trends over the past week. " "Please write a comprehensive report of the current state of the world that is relevant for " "trading and macroeconomics. " "Look at news from EODHD, and finnhub to be comprehensive. Do not simply state the trends are mixed, " "provide detailed and finegrained analysis and insights that may help traders make decisions." """ Make sure to append a Markdown table at the end of the report to organize key points in the report, organized and easy to read.""" ) tool\_names = \[tool.func.\_\_name\_\_ for tool in tools\] return await run\_chain\_with\_tools("news", state, llm, system\_message, tool\_names) # {{/docs-fragment news\_analyst}} @env.task async def create\_social\_media\_analyst( llm: str, state: AgentState, online\_tools: bool ) -> AgentState: if online\_tools: tools = \[toolkit.get\_stock\_news\_openai\] else: tools = \[toolkit.get\_reddit\_stock\_info\] system\_message = ( "You are a social media and company specific news researcher/analyst tasked with analyzing social media posts, " "recent company news, and public sentiment for a specific company over the past week. " "You will be given a company's name your objective is to write a comprehensive long report " "detailing your analysis, insights, and implications for traders and investors on this company's current state " "after looking at social media and what people are saying about that company, " "analyzing sentiment data of what people feel each day about the company, and looking at recent company news. " "Try to look at all sources possible from social media to sentiment to news. Do not simply state the trends " "are mixed, provide detailed and finegrained analysis and insights that may help traders make decisions." """ Make sure to append a Makrdown table at the end of the report to organize key points in the report, organized and easy to read.""" ) tool\_names = \[tool.func.\_\_name\_\_ for tool in tools\] return await run\_chain\_with\_tools( "sentiment", state, llm, system\_message, tool\_names ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/agents/analysts.py\* Once all analyst reports are complete, their outputs are collected and passed to the next stage of the workflow. ### Research agents The research phase consists of two agents: a bullish researcher and a bearish one. They evaluate the company from opposing viewpoints, drawing on the analysts' reports. Unlike analysts, they don't use tools. Their role is to interpret, critique, and develop positions based on the evidence. \`\`\` from agents.utils.utils import AgentState, InvestmentDebateState, memory\_init from flyte\_env import env from langchain\_openai import ChatOpenAI # {{docs-fragment bear\_researcher}} @env.task async def create\_bear\_researcher(llm: str, state: AgentState) -> AgentState: investment\_debate\_state = state.investment\_debate\_state history = investment\_debate\_state.history bear\_history = investment\_debate\_state.bear\_history current\_response = investment\_debate\_state.current\_response market\_research\_report = state.market\_report sentiment\_report = state.sentiment\_report news\_report = state.news\_report fundamentals\_report = state.fundamentals\_report memory = await memory\_init(name="bear-researcher") curr\_situation = f"{market\_research\_report}\\n\\n{sentiment\_report}\\n\\n{news\_report}\\n\\n{fundamentals\_report}" past\_memories = memory.get\_memories(curr\_situation, n\_matches=2) past\_memory\_str = "" for rec in past\_memories: past\_memory\_str += rec\["recommendation"\] + "\\n\\n" prompt = f"""You are a Bear Analyst making the case against investing in the stock. Your goal is to present a well-reasoned argument emphasizing risks, challenges, and negative indicators. Leverage the provided research and data to highlight potential downsides and counter bullish arguments effectively. Key points to focus on: - Risks and Challenges: Highlight factors like market saturation, financial instability, or macroeconomic threats that could hinder the stock's performance. - Competitive Weaknesses: Emphasize vulnerabilities such as weaker market positioning, declining innovation, or threats from competitors. - Negative Indicators: Use evidence from financial data, market trends, or recent adverse news to support your position. - Bull Counterpoints: Critically analyze the bull argument with specific data and sound reasoning, exposing weaknesses or over-optimistic assumptions. - Engagement: Present your argument in a conversational style, directly engaging with the bull analyst's points and debating effectively rather than simply listing facts. Resources available: Market research report: {market\_research\_report} Social media sentiment report: {sentiment\_report} Latest world affairs news: {news\_report} Company fundamentals report: {fundamentals\_report} Conversation history of the debate: {history} Last bull argument: {current\_response} Reflections from similar situations and lessons learned: {past\_memory\_str} Use this information to deliver a compelling bear argument, refute the bull's claims, and engage in a dynamic debate that demonstrates the risks and weaknesses of investing in the stock. You must also address reflections and learn from lessons and mistakes you made in the past. """ response = ChatOpenAI(model=llm).invoke(prompt) argument = f"Bear Analyst: {response.content}" new\_investment\_debate\_state = InvestmentDebateState( history=history + "\\n" + argument, bear\_history=bear\_history + "\\n" + argument, bull\_history=investment\_debate\_state.bull\_history, current\_response=argument, count=investment\_debate\_state.count + 1, ) state.investment\_debate\_state = new\_investment\_debate\_state return state # {{/docs-fragment bear\_researcher}} @env.task async def create\_bull\_researcher(llm: str, state: AgentState) -> AgentState: investment\_debate\_state = state.investment\_debate\_state history = investment\_debate\_state.history bull\_history = investment\_debate\_state.bull\_history current\_response = investment\_debate\_state.current\_response market\_research\_report = state.market\_report sentiment\_report = state.sentiment\_report news\_report = state.news\_report fundamentals\_report = state.fundamentals\_report memory = await memory\_init(name="bull-researcher") curr\_situation = f"{market\_research\_report}\\n\\n{sentiment\_report}\\n\\n{news\_report}\\n\\n{fundamentals\_report}" past\_memories = memory.get\_memories(curr\_situation, n\_matches=2) past\_memory\_str = "" for rec in past\_memories: past\_memory\_str += rec\["recommendation"\] + "\\n\\n" prompt = f"""You are a Bull Analyst advocating for investing in the stock. Your task is to build a strong, evidence-based case emphasizing growth potential, competitive advantages, and positive market indicators. Leverage the provided research and data to address concerns and counter bearish arguments effectively. Key points to focus on: - Growth Potential: Highlight the company's market opportunities, revenue projections, and scalability. - Competitive Advantages: Emphasize factors like unique products, strong branding, or dominant market positioning. - Positive Indicators: Use financial health, industry trends, and recent positive news as evidence. - Bear Counterpoints: Critically analyze the bear argument with specific data and sound reasoning, addressing concerns thoroughly and showing why the bull perspective holds stronger merit. - Engagement: Present your argument in a conversational style, engaging directly with the bear analyst's points and debating effectively rather than just listing data. Resources available: Market research report: {market\_research\_report} Social media sentiment report: {sentiment\_report} Latest world affairs news: {news\_report} Company fundamentals report: {fundamentals\_report} Conversation history of the debate: {history} Last bear argument: {current\_response} Reflections from similar situations and lessons learned: {past\_memory\_str} Use this information to deliver a compelling bull argument, refute the bear's concerns, and engage in a dynamic debate that demonstrates the strengths of the bull position. You must also address reflections and learn from lessons and mistakes you made in the past. """ response = ChatOpenAI(model=llm).invoke(prompt) argument = f"Bull Analyst: {response.content}" new\_investment\_debate\_state = InvestmentDebateState( history=history + "\\n" + argument, bull\_history=bull\_history + "\\n" + argument, bear\_history=investment\_debate\_state.bear\_history, current\_response=argument, count=investment\_debate\_state.count + 1, ) state.investment\_debate\_state = new\_investment\_debate\_state return state \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/agents/researchers.py\* To aid reasoning, the agents can also retrieve relevant "memories" from a vector database, giving them richer historical context. The number of debate rounds is configurable, and after a few iterations of back-and-forth between the bull and bear, a research manager agent reviews their arguments and makes a final investment decision. \`\`\` from agents.utils.utils import ( AgentState, InvestmentDebateState, RiskDebateState, memory\_init, ) from flyte\_env import env from langchain\_openai import ChatOpenAI # {{docs-fragment research\_manager}} @env.task async def create\_research\_manager(llm: str, state: AgentState) -> AgentState: history = state.investment\_debate\_state.history investment\_debate\_state = state.investment\_debate\_state market\_research\_report = state.market\_report sentiment\_report = state.sentiment\_report news\_report = state.news\_report fundamentals\_report = state.fundamentals\_report memory = await memory\_init(name="research-manager") curr\_situation = f"{market\_research\_report}\\n\\n{sentiment\_report}\\n\\n{news\_report}\\n\\n{fundamentals\_report}" past\_memories = memory.get\_memories(curr\_situation, n\_matches=2) past\_memory\_str = "" for rec in past\_memories: past\_memory\_str += rec\["recommendation"\] + "\\n\\n" prompt = f"""As the portfolio manager and debate facilitator, your role is to critically evaluate this round of debate and make a definitive decision: align with the bear analyst, the bull analyst, or choose Hold only if it is strongly justified based on the arguments presented. Summarize the key points from both sides concisely, focusing on the most compelling evidence or reasoning. Your recommendation—Buy, Sell, or Hold—must be clear and actionable. Avoid defaulting to Hold simply because both sides have valid points; commit to a stance grounded in the debate's strongest arguments. Additionally, develop a detailed investment plan for the trader. This should include: Your Recommendation: A decisive stance supported by the most convincing arguments. Rationale: An explanation of why these arguments lead to your conclusion. Strategic Actions: Concrete steps for implementing the recommendation. Take into account your past mistakes on similar situations. Use these insights to refine your decision-making and ensure you are learning and improving. Present your analysis conversationally, as if speaking naturally, without special formatting. Here are your past reflections on mistakes: \\"{past\_memory\_str}\\" Here is the debate: Debate History: {history}""" response = ChatOpenAI(model=llm).invoke(prompt) new\_investment\_debate\_state = InvestmentDebateState( judge\_decision=response.content, history=investment\_debate\_state.history, bear\_history=investment\_debate\_state.bear\_history, bull\_history=investment\_debate\_state.bull\_history, current\_response=response.content, count=investment\_debate\_state.count, ) state.investment\_debate\_state = new\_investment\_debate\_state state.investment\_plan = response.content return state # {{/docs-fragment research\_manager}} @env.task async def create\_risk\_manager(llm: str, state: AgentState) -> AgentState: history = state.risk\_debate\_state.history risk\_debate\_state = state.risk\_debate\_state trader\_plan = state.investment\_plan market\_research\_report = state.market\_report sentiment\_report = state.sentiment\_report news\_report = state.news\_report fundamentals\_report = state.fundamentals\_report memory = await memory\_init(name="risk-manager") curr\_situation = f"{market\_research\_report}\\n\\n{sentiment\_report}\\n\\n{news\_report}\\n\\n{fundamentals\_report}" past\_memories = memory.get\_memories(curr\_situation, n\_matches=2) past\_memory\_str = "" for rec in past\_memories: past\_memory\_str += rec\["recommendation"\] + "\\n\\n" prompt = f"""As the Risk Management Judge and Debate Facilitator, your goal is to evaluate the debate between three risk analysts—Risky, Neutral, and Safe/Conservative—and determine the best course of action for the trader. Your decision must result in a clear recommendation: Buy, Sell, or Hold. Choose Hold only if strongly justified by specific arguments, not as a fallback when all sides seem valid. Strive for clarity and decisiveness. Guidelines for Decision-Making: 1. \*\*Summarize Key Arguments\*\*: Extract the strongest points from each analyst, focusing on relevance to the context. 2. \*\*Provide Rationale\*\*: Support your recommendation with direct quotes and counterarguments from the debate. 3. \*\*Refine the Trader's Plan\*\*: Start with the trader's original plan, \*\*{trader\_plan}\*\*, and adjust it based on the analysts' insights. 4. \*\*Learn from Past Mistakes\*\*: Use lessons from \*\*{past\_memory\_str}\*\* to address prior misjudgments and improve the decision you are making now to make sure you don't make a wrong BUY/SELL/HOLD call that loses money. Deliverables: - A clear and actionable recommendation: Buy, Sell, or Hold. - Detailed reasoning anchored in the debate and past reflections. --- \*\*Analysts Debate History:\*\* {history} --- Focus on actionable insights and continuous improvement. Build on past lessons, critically evaluate all perspectives, and ensure each decision advances better outcomes.""" response = ChatOpenAI(model=llm).invoke(prompt) new\_risk\_debate\_state = RiskDebateState( judge\_decision=response.content, history=risk\_debate\_state.history, risky\_history=risk\_debate\_state.risky\_history, safe\_history=risk\_debate\_state.safe\_history, neutral\_history=risk\_debate\_state.neutral\_history, latest\_speaker="Judge", current\_risky\_response=risk\_debate\_state.current\_risky\_response, current\_safe\_response=risk\_debate\_state.current\_safe\_response, current\_neutral\_response=risk\_debate\_state.current\_neutral\_response, count=risk\_debate\_state.count, ) state.risk\_debate\_state = new\_risk\_debate\_state state.final\_trade\_decision = response.content return state \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/agents/managers.py\* ### Trading agent The trader agent consolidates the insights from analysts and researchers to generate a final recommendation. It synthesizes competing signals and produces a conclusion such as \_Buy for long-term growth despite short-term volatility\_. \`\`\` from agents.utils.utils import AgentState, memory\_init from flyte\_env import env from langchain\_core.messages import convert\_to\_openai\_messages from langchain\_openai import ChatOpenAI # {{docs-fragment trader}} @env.task async def create\_trader(llm: str, state: AgentState) -> AgentState: company\_name = state.company\_of\_interest investment\_plan = state.investment\_plan market\_research\_report = state.market\_report sentiment\_report = state.sentiment\_report news\_report = state.news\_report fundamentals\_report = state.fundamentals\_report memory = await memory\_init(name="trader") curr\_situation = f"{market\_research\_report}\\n\\n{sentiment\_report}\\n\\n{news\_report}\\n\\n{fundamentals\_report}" past\_memories = memory.get\_memories(curr\_situation, n\_matches=2) past\_memory\_str = "" for rec in past\_memories: past\_memory\_str += rec\["recommendation"\] + "\\n\\n" context = { "role": "user", "content": f"Based on a comprehensive analysis by a team of analysts, " f"here is an investment plan tailored for {company\_name}. " "This plan incorporates insights from current technical market trends, " "macroeconomic indicators, and social media sentiment. " "Use this plan as a foundation for evaluating your next trading decision.\\n\\n" f"Proposed Investment Plan: {investment\_plan}\\n\\n" "Leverage these insights to make an informed and strategic decision.", } messages = \[\ {\ "role": "system",\ "content": f"""You are a trading agent analyzing market data to make investment decisions.\ Based on your analysis, provide a specific recommendation to buy, sell, or hold.\ End with a firm decision and always conclude your response with 'FINAL TRANSACTION PROPOSAL: \*\*BUY/HOLD/SELL\*\*'\ to confirm your recommendation.\ Do not forget to utilize lessons from past decisions to learn from your mistakes.\ Here is some reflections from similar situatiosn you traded in and the lessons learned: {past\_memory\_str}""",\ },\ context,\ \] result = ChatOpenAI(model=llm).invoke(messages) state.messages.append(convert\_to\_openai\_messages(result)) state.trader\_investment\_plan = result.content state.sender = "Trader" return state # {{/docs-fragment trader}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/agents/trader.py\* ### Risk agents Risk agents comprise agents with different risk tolerances: a risky debater, a neutral one, and a conservative one. They assess the portfolio through lenses like market volatility, liquidity, and systemic risk. Similar to the bull-bear debate, these agents engage in internal discussion, after which a risk manager makes the final call. \`\`\` from agents.utils.utils import AgentState, RiskDebateState from flyte\_env import env from langchain\_openai import ChatOpenAI # {{docs-fragment risk\_debator}} @env.task async def create\_risky\_debator(llm: str, state: AgentState) -> AgentState: risk\_debate\_state = state.risk\_debate\_state history = risk\_debate\_state.history risky\_history = risk\_debate\_state.risky\_history current\_safe\_response = risk\_debate\_state.current\_safe\_response current\_neutral\_response = risk\_debate\_state.current\_neutral\_response market\_research\_report = state.market\_report sentiment\_report = state.sentiment\_report news\_report = state.news\_report fundamentals\_report = state.fundamentals\_report trader\_decision = state.trader\_investment\_plan prompt = f"""As the Risky Risk Analyst, your role is to actively champion high-reward, high-risk opportunities, emphasizing bold strategies and competitive advantages. When evaluating the trader's decision or plan, focus intently on the potential upside, growth potential, and innovative benefits—even when these come with elevated risk. Use the provided market data and sentiment analysis to strengthen your arguments and challenge the opposing views. Specifically, respond directly to each point made by the conservative and neutral analysts, countering with data-driven rebuttals and persuasive reasoning. Highlight where their caution might miss critical opportunities or where their assumptions may be overly conservative. Here is the trader's decision: {trader\_decision} Your task is to create a compelling case for the trader's decision by questioning and critiquing the conservative and neutral stances to demonstrate why your high-reward perspective offers the best path forward. Incorporate insights from the following sources into your arguments: Market Research Report: {market\_research\_report} Social Media Sentiment Report: {sentiment\_report} Latest World Affairs Report: {news\_report} Company Fundamentals Report: {fundamentals\_report} Here is the current conversation history: {history} Here are the last arguments from the conservative analyst: {current\_safe\_response} Here are the last arguments from the neutral analyst: {current\_neutral\_response}. If there are no responses from the other viewpoints, do not halluncinate and just present your point. Engage actively by addressing any specific concerns raised, refuting the weaknesses in their logic, and asserting the benefits of risk-taking to outpace market norms. Maintain a focus on debating and persuading, not just presenting data. Challenge each counterpoint to underscore why a high-risk approach is optimal. Output conversationally as if you are speaking without any special formatting.""" response = ChatOpenAI(model=llm).invoke(prompt) argument = f"Risky Analyst: {response.content}" new\_risk\_debate\_state = RiskDebateState( history=history + "\\n" + argument, risky\_history=risky\_history + "\\n" + argument, safe\_history=risk\_debate\_state.safe\_history, neutral\_history=risk\_debate\_state.neutral\_history, latest\_speaker="Risky", current\_risky\_response=argument, current\_safe\_response=current\_safe\_response, current\_neutral\_response=current\_neutral\_response, count=risk\_debate\_state.count + 1, ) state.risk\_debate\_state = new\_risk\_debate\_state return state # {{/docs-fragment risk\_debator}} @env.task async def create\_safe\_debator(llm: str, state: AgentState) -> AgentState: risk\_debate\_state = state.risk\_debate\_state history = risk\_debate\_state.history safe\_history = risk\_debate\_state.safe\_history current\_risky\_response = risk\_debate\_state.current\_risky\_response current\_neutral\_response = risk\_debate\_state.current\_neutral\_response market\_research\_report = state.market\_report sentiment\_report = state.sentiment\_report news\_report = state.news\_report fundamentals\_report = state.fundamentals\_report trader\_decision = state.trader\_investment\_plan prompt = f"""As the Safe/Conservative Risk Analyst, your primary objective is to protect assets, minimize volatility, and ensure steady, reliable growth. You prioritize stability, security, and risk mitigation, carefully assessing potential losses, economic downturns, and market volatility. When evaluating the trader's decision or plan, critically examine high-risk elements, pointing out where the decision may expose the firm to undue risk and where more cautious alternatives could secure long-term gains. Here is the trader's decision: {trader\_decision} Your task is to actively counter the arguments of the Risky and Neutral Analysts, highlighting where their views may overlook potential threats or fail to prioritize sustainability. Respond directly to their points, drawing from the following data sources to build a convincing case for a low-risk approach adjustment to the trader's decision: Market Research Report: {market\_research\_report} Social Media Sentiment Report: {sentiment\_report} Latest World Affairs Report: {news\_report} Company Fundamentals Report: {fundamentals\_report} Here is the current conversation history: {history} Here is the last response from the risky analyst: {current\_risky\_response} Here is the last response from the neutral analyst: {current\_neutral\_response}. If there are no responses from the other viewpoints, do not halluncinate and just present your point. Engage by questioning their optimism and emphasizing the potential downsides they may have overlooked. Address each of their counterpoints to showcase why a conservative stance is ultimately the safest path for the firm's assets. Focus on debating and critiquing their arguments to demonstrate the strength of a low-risk strategy over their approaches. Output conversationally as if you are speaking without any special formatting.""" response = ChatOpenAI(model=llm).invoke(prompt) argument = f"Safe Analyst: {response.content}" new\_risk\_debate\_state = RiskDebateState( history=history + "\\n" + argument, risky\_history=risk\_debate\_state.risky\_history, safe\_history=safe\_history + "\\n" + argument, neutral\_history=risk\_debate\_state.neutral\_history, latest\_speaker="Safe", current\_risky\_response=current\_risky\_response, current\_safe\_response=argument, current\_neutral\_response=current\_neutral\_response, count=risk\_debate\_state.count + 1, ) state.risk\_debate\_state = new\_risk\_debate\_state return state @env.task async def create\_neutral\_debator(llm: str, state: AgentState) -> AgentState: risk\_debate\_state = state.risk\_debate\_state history = risk\_debate\_state.history neutral\_history = risk\_debate\_state.neutral\_history current\_risky\_response = risk\_debate\_state.current\_risky\_response current\_safe\_response = risk\_debate\_state.current\_safe\_response market\_research\_report = state.market\_report sentiment\_report = state.sentiment\_report news\_report = state.news\_report fundamentals\_report = state.fundamentals\_report trader\_decision = state.trader\_investment\_plan prompt = f"""As the Neutral Risk Analyst, your role is to provide a balanced perspective, weighing both the potential benefits and risks of the trader's decision or plan. You prioritize a well-rounded approach, evaluating the upsides and downsides while factoring in broader market trends, potential economic shifts, and diversification strategies.Here is the trader's decision: {trader\_decision} Your task is to challenge both the Risky and Safe Analysts, pointing out where each perspective may be overly optimistic or overly cautious. Use insights from the following data sources to support a moderate, sustainable strategy to adjust the trader's decision: Market Research Report: {market\_research\_report} Social Media Sentiment Report: {sentiment\_report} Latest World Affairs Report: {news\_report} Company Fundamentals Report: {fundamentals\_report} Here is the current conversation history: {history} Here is the last response from the risky analyst: {current\_risky\_response} Here is the last response from the safe analyst: {current\_safe\_response}. If there are no responses from the other viewpoints, do not halluncinate and just present your point. Engage actively by analyzing both sides critically, addressing weaknesses in the risky and conservative arguments to advocate for a more balanced approach. Challenge each of their points to illustrate why a moderate risk strategy might offer the best of both worlds, providing growth potential while safeguarding against extreme volatility. Focus on debating rather than simply presenting data, aiming to show that a balanced view can lead to the most reliable outcomes. Output conversationally as if you are speaking without any special formatting.""" response = ChatOpenAI(model=llm).invoke(prompt) argument = f"Neutral Analyst: {response.content}" new\_risk\_debate\_state = RiskDebateState( history=history + "\\n" + argument, risky\_history=risk\_debate\_state.risky\_history, safe\_history=risk\_debate\_state.safe\_history, neutral\_history=neutral\_history + "\\n" + argument, latest\_speaker="Neutral", current\_risky\_response=current\_risky\_response, current\_safe\_response=current\_safe\_response, current\_neutral\_response=argument, count=risk\_debate\_state.count + 1, ) state.risk\_debate\_state = new\_risk\_debate\_state return state \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/agents/risk\_debators.py\* The outcome of the risk manager — whether to proceed with the trade or not — is considered the final decision of the trading simulation. You can visualize this full pipeline in the Flyte/Union UI, where every step is logged. You’ll see input/output metadata for each tool and agent task. Thanks to Flyte's caching, repeated steps are skipped unless inputs change, saving time and compute resources. ### Retaining agent memory with S3 vectors To help agents learn from past decisions, we persist their memory in a vector store. In this example, we use an \[S3 vector\](https://aws.amazon.com/s3/features/vectors/) bucket for their simplicity and tight integration with Flyte and Union, but any vector database can be used. Note: To use the S3 vector store, make sure your IAM role has the following permissions configured: \`\`\` s3vectors:CreateVectorBucket s3vectors:CreateIndex s3vectors:PutVectors s3vectors:GetIndex s3vectors:GetVectors s3vectors:QueryVectors s3vectors:GetVectorBucket \`\`\` After each trade decision, you can run a \`reflect\_on\_decisions\` task. This evaluates whether the final outcome aligned with the agent's recommendation and stores that reflection in the vector store. These stored insights can later be retrieved to provide historical context and improve future decision-making. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "akshare==1.16.98",\ # "backtrader==1.9.78.123",\ # "boto3==1.39.9",\ # "chainlit==2.5.5",\ # "eodhd==1.0.32",\ # "feedparser==6.0.11",\ # "finnhub-python==2.4.23",\ # "langchain-experimental==0.3.4",\ # "langchain-openai==0.3.23",\ # "pandas==2.3.0",\ # "parsel==1.10.0",\ # "praw==7.8.1",\ # "pytz==2025.2",\ # "questionary==2.1.0",\ # "redis==6.2.0",\ # "requests==2.32.4",\ # "stockstats==0.6.5",\ # "tqdm==4.67.1",\ # "tushare==1.4.21",\ # "typing-extensions==4.14.0",\ # "yfinance==0.2.63",\ # \] # main = "main" # params = "" # /// import asyncio from copy import deepcopy import agents import agents.analysts from agents.managers import create\_research\_manager, create\_risk\_manager from agents.researchers import create\_bear\_researcher, create\_bull\_researcher from agents.risk\_debators import ( create\_neutral\_debator, create\_risky\_debator, create\_safe\_debator, ) from agents.trader import create\_trader from agents.utils.utils import AgentState from flyte\_env import DEEP\_THINKING\_LLM, QUICK\_THINKING\_LLM, env, flyte from langchain\_openai import ChatOpenAI from reflection import ( reflect\_bear\_researcher, reflect\_bull\_researcher, reflect\_research\_manager, reflect\_risk\_manager, reflect\_trader, ) @env.task async def process\_signal(full\_signal: str, QUICK\_THINKING\_LLM: str) -> str: """Process a full trading signal to extract the core decision.""" messages = \[\ {\ "role": "system",\ "content": """You are an efficient assistant designed to analyze paragraphs or\ financial reports provided by a group of analysts.\ Your task is to extract the investment decision: SELL, BUY, or HOLD.\ Provide only the extracted decision (SELL, BUY, or HOLD) as your output,\ without adding any additional text or information.""",\ },\ {"role": "human", "content": full\_signal},\ \] return ChatOpenAI(model=QUICK\_THINKING\_LLM).invoke(messages).content async def run\_analyst(analyst\_name, state, online\_tools): # Create a copy of the state for isolation run\_fn = getattr(agents.analysts, f"create\_{analyst\_name}\_analyst") # Run the analyst's chain result\_state = await run\_fn(QUICK\_THINKING\_LLM, state, online\_tools) # Determine the report key report\_key = ( "sentiment\_report" if analyst\_name == "social\_media" else f"{analyst\_name}\_report" ) report\_value = getattr(result\_state, report\_key) return result\_state.messages\[1:\], report\_key, report\_value # {{docs-fragment main}} @env.task async def main( selected\_analysts: list\[str\] = \[\ "market",\ "fundamentals",\ "news",\ "social\_media",\ \], max\_debate\_rounds: int = 1, max\_risk\_discuss\_rounds: int = 1, online\_tools: bool = True, company\_name: str = "NVDA", trade\_date: str = "2024-05-12", ) -> tuple\[str, AgentState\]: if not selected\_analysts: raise ValueError( "No analysts selected. Please select at least one analyst from market, fundamentals, news, or social\_media." ) state = AgentState( messages=\[{"role": "human", "content": company\_name}\], company\_of\_interest=company\_name, trade\_date=str(trade\_date), ) # Run all analysts concurrently results = await asyncio.gather( \*\[\ run\_analyst(analyst, deepcopy(state), online\_tools)\ for analyst in selected\_analysts\ \] ) # Flatten and append all resulting messages into the shared state for messages, report\_attr, report in results: state.messages.extend(messages) setattr(state, report\_attr, report) # Bull/Bear debate loop state = await create\_bull\_researcher(QUICK\_THINKING\_LLM, state) # Start with bull while state.investment\_debate\_state.count < 2 \* max\_debate\_rounds: current = state.investment\_debate\_state.current\_response if current.startswith("Bull"): state = await create\_bear\_researcher(QUICK\_THINKING\_LLM, state) else: state = await create\_bull\_researcher(QUICK\_THINKING\_LLM, state) state = await create\_research\_manager(DEEP\_THINKING\_LLM, state) state = await create\_trader(QUICK\_THINKING\_LLM, state) # Risk debate loop state = await create\_risky\_debator(QUICK\_THINKING\_LLM, state) # Start with risky while state.risk\_debate\_state.count < 3 \* max\_risk\_discuss\_rounds: speaker = state.risk\_debate\_state.latest\_speaker if speaker == "Risky": state = await create\_safe\_debator(QUICK\_THINKING\_LLM, state) elif speaker == "Safe": state = await create\_neutral\_debator(QUICK\_THINKING\_LLM, state) else: state = await create\_risky\_debator(QUICK\_THINKING\_LLM, state) state = await create\_risk\_manager(DEEP\_THINKING\_LLM, state) decision = await process\_signal(state.final\_trade\_decision, QUICK\_THINKING\_LLM) return decision, state # {{/docs-fragment main}} # {{docs-fragment reflect\_on\_decisions}} @env.task async def reflect\_and\_store(state: AgentState, returns: str) -> str: await asyncio.gather( reflect\_bear\_researcher(state, returns), reflect\_bull\_researcher(state, returns), reflect\_trader(state, returns), reflect\_risk\_manager(state, returns), reflect\_research\_manager(state, returns), ) return "Reflection completed." # Run the reflection task after the main function @env.task(cache="disable") async def reflect\_on\_decisions( returns: str, selected\_analysts: list\[str\] = \[\ "market",\ "fundamentals",\ "news",\ "social\_media",\ \], max\_debate\_rounds: int = 1, max\_risk\_discuss\_rounds: int = 1, online\_tools: bool = True, company\_name: str = "NVDA", trade\_date: str = "2024-05-12", ) -> str: \_, state = await main( selected\_analysts, max\_debate\_rounds, max\_risk\_discuss\_rounds, online\_tools, company\_name, trade\_date, ) return await reflect\_and\_store(state, returns) # {{/docs-fragment reflect\_on\_decisions}} # {{docs-fragment execute\_main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # run = flyte.run(reflect\_on\_decisions, "+3.2% gain over 5 days") # print(run.url) # {{/docs-fragment execute\_main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/main.py\* ### Running the simulation First, set up your OpenAI secret (from \[openai.com\](https://platform.openai.com/api-keys)) and Finnhub API key (from \[finnhub.io\](https://finnhub.io/)): \`\`\` flyte create secret openai\_api\_key flyte create secret finnhub\_api\_key \`\`\` Then \[clone the repo\](https://github.com/unionai/unionai-examples), navigate to the \`tutorials-v2/trading\_agents\` directory, and run the following commands: \`\`\` flyte create config --endpoint --project --domain --builder remote uv run main.py \`\`\` If you'd like to run the \`reflect\_on\_decisions\` task instead, comment out the \`main\` function call and uncomment the \`reflect\_on\_decisions\` call in the \`\_\_main\_\_\` block: \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "akshare==1.16.98",\ # "backtrader==1.9.78.123",\ # "boto3==1.39.9",\ # "chainlit==2.5.5",\ # "eodhd==1.0.32",\ # "feedparser==6.0.11",\ # "finnhub-python==2.4.23",\ # "langchain-experimental==0.3.4",\ # "langchain-openai==0.3.23",\ # "pandas==2.3.0",\ # "parsel==1.10.0",\ # "praw==7.8.1",\ # "pytz==2025.2",\ # "questionary==2.1.0",\ # "redis==6.2.0",\ # "requests==2.32.4",\ # "stockstats==0.6.5",\ # "tqdm==4.67.1",\ # "tushare==1.4.21",\ # "typing-extensions==4.14.0",\ # "yfinance==0.2.63",\ # \] # main = "main" # params = "" # /// import asyncio from copy import deepcopy import agents import agents.analysts from agents.managers import create\_research\_manager, create\_risk\_manager from agents.researchers import create\_bear\_researcher, create\_bull\_researcher from agents.risk\_debators import ( create\_neutral\_debator, create\_risky\_debator, create\_safe\_debator, ) from agents.trader import create\_trader from agents.utils.utils import AgentState from flyte\_env import DEEP\_THINKING\_LLM, QUICK\_THINKING\_LLM, env, flyte from langchain\_openai import ChatOpenAI from reflection import ( reflect\_bear\_researcher, reflect\_bull\_researcher, reflect\_research\_manager, reflect\_risk\_manager, reflect\_trader, ) @env.task async def process\_signal(full\_signal: str, QUICK\_THINKING\_LLM: str) -> str: """Process a full trading signal to extract the core decision.""" messages = \[\ {\ "role": "system",\ "content": """You are an efficient assistant designed to analyze paragraphs or\ financial reports provided by a group of analysts.\ Your task is to extract the investment decision: SELL, BUY, or HOLD.\ Provide only the extracted decision (SELL, BUY, or HOLD) as your output,\ without adding any additional text or information.""",\ },\ {"role": "human", "content": full\_signal},\ \] return ChatOpenAI(model=QUICK\_THINKING\_LLM).invoke(messages).content async def run\_analyst(analyst\_name, state, online\_tools): # Create a copy of the state for isolation run\_fn = getattr(agents.analysts, f"create\_{analyst\_name}\_analyst") # Run the analyst's chain result\_state = await run\_fn(QUICK\_THINKING\_LLM, state, online\_tools) # Determine the report key report\_key = ( "sentiment\_report" if analyst\_name == "social\_media" else f"{analyst\_name}\_report" ) report\_value = getattr(result\_state, report\_key) return result\_state.messages\[1:\], report\_key, report\_value # {{docs-fragment main}} @env.task async def main( selected\_analysts: list\[str\] = \[\ "market",\ "fundamentals",\ "news",\ "social\_media",\ \], max\_debate\_rounds: int = 1, max\_risk\_discuss\_rounds: int = 1, online\_tools: bool = True, company\_name: str = "NVDA", trade\_date: str = "2024-05-12", ) -> tuple\[str, AgentState\]: if not selected\_analysts: raise ValueError( "No analysts selected. Please select at least one analyst from market, fundamentals, news, or social\_media." ) state = AgentState( messages=\[{"role": "human", "content": company\_name}\], company\_of\_interest=company\_name, trade\_date=str(trade\_date), ) # Run all analysts concurrently results = await asyncio.gather( \*\[\ run\_analyst(analyst, deepcopy(state), online\_tools)\ for analyst in selected\_analysts\ \] ) # Flatten and append all resulting messages into the shared state for messages, report\_attr, report in results: state.messages.extend(messages) setattr(state, report\_attr, report) # Bull/Bear debate loop state = await create\_bull\_researcher(QUICK\_THINKING\_LLM, state) # Start with bull while state.investment\_debate\_state.count < 2 \* max\_debate\_rounds: current = state.investment\_debate\_state.current\_response if current.startswith("Bull"): state = await create\_bear\_researcher(QUICK\_THINKING\_LLM, state) else: state = await create\_bull\_researcher(QUICK\_THINKING\_LLM, state) state = await create\_research\_manager(DEEP\_THINKING\_LLM, state) state = await create\_trader(QUICK\_THINKING\_LLM, state) # Risk debate loop state = await create\_risky\_debator(QUICK\_THINKING\_LLM, state) # Start with risky while state.risk\_debate\_state.count < 3 \* max\_risk\_discuss\_rounds: speaker = state.risk\_debate\_state.latest\_speaker if speaker == "Risky": state = await create\_safe\_debator(QUICK\_THINKING\_LLM, state) elif speaker == "Safe": state = await create\_neutral\_debator(QUICK\_THINKING\_LLM, state) else: state = await create\_risky\_debator(QUICK\_THINKING\_LLM, state) state = await create\_risk\_manager(DEEP\_THINKING\_LLM, state) decision = await process\_signal(state.final\_trade\_decision, QUICK\_THINKING\_LLM) return decision, state # {{/docs-fragment main}} # {{docs-fragment reflect\_on\_decisions}} @env.task async def reflect\_and\_store(state: AgentState, returns: str) -> str: await asyncio.gather( reflect\_bear\_researcher(state, returns), reflect\_bull\_researcher(state, returns), reflect\_trader(state, returns), reflect\_risk\_manager(state, returns), reflect\_research\_manager(state, returns), ) return "Reflection completed." # Run the reflection task after the main function @env.task(cache="disable") async def reflect\_on\_decisions( returns: str, selected\_analysts: list\[str\] = \[\ "market",\ "fundamentals",\ "news",\ "social\_media",\ \], max\_debate\_rounds: int = 1, max\_risk\_discuss\_rounds: int = 1, online\_tools: bool = True, company\_name: str = "NVDA", trade\_date: str = "2024-05-12", ) -> str: \_, state = await main( selected\_analysts, max\_debate\_rounds, max\_risk\_discuss\_rounds, online\_tools, company\_name, trade\_date, ) return await reflect\_and\_store(state, returns) # {{/docs-fragment reflect\_on\_decisions}} # {{docs-fragment execute\_main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # run = flyte.run(reflect\_on\_decisions, "+3.2% gain over 5 days") # print(run.url) # {{/docs-fragment execute\_main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/trading\_agents/main.py\* Then run: \`\`\` uv run main.py \`\`\` ## Why Flyte? \_(A quick note before you go)\_ You might now be wondering: can't I just build all this with Python and LangChain? Absolutely. But as your project grows, you'll likely run into these challenges: 1. \*\*Observability\*\*: Agent workflows can feel opaque. You send a prompt, get a response, but what happened in between? - Were the right tools used? - Were correct arguments passed? - How did the LLM reason through intermediate steps? - Why did it fail? Flyte gives you a window into each of these stages. 2. \*\*Multi-agent coordination\*\*: Real-world applications often require multiple agents with distinct roles and responsibilities. In such cases, you'll need: - Isolated state per agent, - Shared context where needed, - And coordination — sequential or parallel. Managing this manually gets fragile, fast. Flyte handles it for you. 3. \*\*Scalability\*\*: Agents and tools might need to run in isolated or containerized environments. Whether you're scaling out to more agents or more powerful hardware, Flyte lets you scale without taxing your local machine or racking up unnecessary cloud bills. 4. \*\*Durability & recovery\*\*: LLM-based workflows are often long-running and expensive. If something fails halfway: - Do you lose all progress? - Replay everything from scratch? With Flyte, you get built-in caching, checkpointing, and recovery, so you can resume where you left off. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/financial-services/financial-research-agent === # Financial research agent > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/financial\_research\_agent). This example demonstrates how to build a financial research and earnings-cycle agent on Flyte. For each company, the agent runs grounded, source-cited research and fresh news, then synthesizes an analyst-ready equity briefing. Financial research benefits from \*\*low-latency, ranked, source-cited results\*\* across both the general web and news streams. The \[You.com Research API\](https://you.com/docs/research/overview) produces a grounded, citation-backed synthesis, and the \[You.com Search API\](https://you.com/docs/search/overview) adds a fresh-news layer. \[Claude\](https://docs.anthropic.com/) via \[LiteLLM\](https://docs.litellm.ai/) turns that evidence into an analyst-ready briefing. Flyte's \`cache="auto"\` reuses prior results when runs converge on the same companies. Flyte provides: - \*\*Fan-out parallelism\*\* across companies - \*\*\`cache="auto"\`\*\* to reuse prior You.com and LLM results across converging runs - \*\*\`@flyte.trace\`\*\* on every external call for full prompt → citation lineage - \*\*Flyte reports\*\* with thesis, risks, watch items, and source citations per company !\[Financial research agent report\](https://www.union.ai/docs/v2/flyte/\_static/images/tutorials/financial\_research\_agent/financial-research-agent.png) ## Setting up the environment The agent runs in a \`TaskEnvironment\` with secrets for the You.com and Anthropic API keys, automatic caching, and a container image built from the \`uv\` script dependencies. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "financial\_research" # params = "" # /// """Financial research & earnings-cycle agent. For each company, runs grounded, source-cited research via the You.com Research API plus a fresh-news layer via the Search API, then uses Claude to synthesize an analyst-ready equity briefing that preserves citations. Flyte caching cuts duplicate spend when runs converge. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="financial-research", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="financial-research", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str domain: str = "" snippet: str = "" published: str = "" favicon: str = "" section: str = "research" # "research", "news", or "web" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Briefing: company: str thesis: str recent\_developments: list\[str\] = field(default\_factory=list) risks: list\[str\] = field(default\_factory=list) watch\_items: list\[str\] = field(default\_factory=list) sources: list\[Source\] = field(default\_factory=list) @dataclass class ResearchReport: briefings: list\[Briefing\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_apis}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_request(method: str, url: str, timeout: float, \*\*kwargs) -> dict: """HTTP wrapper with exponential backoff + jitter on 429 rate limits. Fanned-out tasks run in separate pods, so we retry on the client side to smooth out bursts against the You.com API rate limit. """ import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} if method == "POST": headers\["Content-Type"\] = "application/json" async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.request(method, url, headers=headers, \*\*kwargs) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str, freshness: str) -> dict: """Grounded, citation-backed research answer.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": {"freshness": freshness}, } return await \_you\_request("POST", YOU\_RESEARCH\_URL, 300.0, json=body) @flyte.trace async def you\_news(query: str, count: int = 6, freshness: str = "week") -> list\[dict\]: """Fresh news headlines for a company.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_request("GET", YOU\_SEARCH\_URL, 60.0, params=params) results = data.get("results", {}) out: list\[dict\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") out.append( { "title": item.get("title", ""), "url": url, "domain": \_domain(url), "snippet": snippets\[0\] if snippets else item.get("description", ""), "published": item.get("page\_age", "") or "", "favicon": item.get("favicon\_url") or \_favicon\_for(url), "section": section, } ) return out # {{/docs-fragment you\_apis}} # {{docs-fragment llm}} @flyte.trace async def synthesize\_briefing(company: str, focus: str, research: str, news: str) -> dict: """Use Claude to synthesize a structured equity briefing.""" from litellm import acompletion system = ( "You are an equity research analyst. Using ONLY the grounded research " "and news provided, write a concise briefing. Respond ONLY with JSON: " '{"thesis": str, "recent\_developments": \[str\], "risks": \[str\], ' '"watch\_items": \[str\]}. Keep each list to 3-5 short, specific bullets.' ) user = ( f"Company: {company}\\nFocus: {focus}\\n\\n" f"Grounded research:\\n{research}\\n\\nRecent news:\\n{news}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment research\_company}} @env.task(retries=3) async def research\_company( company: str, focus: str, research\_effort: str, freshness: str, ) -> Briefing: """Research one company and synthesize a cited briefing.""" question = ( f"Provide a grounded analysis of {company} with respect to: {focus}. " f"Cover recent financial performance, strategic moves, competitive " f"positioning, and risks." ) research\_result, news = await asyncio.gather( you\_research(question, research\_effort, freshness), you\_news(f"{company} earnings news", freshness=freshness), ) output = research\_result.get("output", {}) research\_text = output.get("content", "") if not isinstance(research\_text, str): research\_text = json.dumps(research\_text) sources: list\[Source\] = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, domain=\_domain(url), snippet=str((s.get("snippets") or \[""\])\[0\]), favicon=\_favicon\_for(url), section="research", ) ) for n in news: sources.append( Source( title=str(n.get("title", "")), url=str(n.get("url", "")), domain=str(n.get("domain", "")), snippet=str(n.get("snippet", "")), published=str(n.get("published", "")), favicon=str(n.get("favicon", "")), section=str(n.get("section", "web")), ) ) news\_text = "\\n".join( f"- {n\['title'\]} ({n\['published'\]}) {n\['domain'\]}: {n\['snippet'\]\[:120\]}" for n in news ) parsed = await synthesize\_briefing(company, focus, research\_text, news\_text) def \_list(key: str) -> list\[str\]: return \[str(x) for x in (parsed.get(key) or \[\])\] return Briefing( company=company, thesis=str(parsed.get("thesis", "")), recent\_developments=\_list("recent\_developments"), risks=\_list("risks"), watch\_items=\_list("watch\_items"), sources=sources, ) # {{/docs-fragment research\_company}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com citation (Research or Search source).""" if not s.url: return "" tag\_cls = s.section if s.section in ("research", "news") else "web" meta\_bits = \[\] if s.published: meta\_bits.append(s.published\[:10\]) if s.title: meta\_bits.append(s.title) meta = " · ".join(meta\_bits) snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"{s.section}" f"
{meta}
{snip}
" ) def \_render\_report(report: ResearchReport) -> str: def \_ul(items: list\[str\]) -> str: if not items: return "

None reported.

" return "
    " + "".join(f"
  • {x}
    • " for x in items) + "
" cards = \[\] for b in report.briefings: src = "".join(\_cite(s) for s in b.sources\[:10\]) cards.append( f"

{b.company}

" f"
{b.thesis or 'No thesis generated.'}
" f"
" f"

Recent developments

{\_ul(b.recent\_developments)}
" f"

Risks

{\_ul(b.risks)}
" f"

Watch items

{\_ul(b.watch\_items)}
" f"
" + (f"

You.com sources ({len(b.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(b.sources) for b in report.briefings) return f""" {REPORT\_CSS}

Financial Research Briefings

Grounded, citation-backed equity briefings — each company backed by You.com Research synthesis plus fresh Search news.

{len(report.briefings)} companies {total\_sources} You.com sources cited
{''.join(cards) or "

No briefings generated.

"}

Research answers from the You.com Research API (grounded synthesis with inline citations) plus fresh headlines from the You.com Search API (web + auto-classified news with timestamps and snippets).

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def financial\_research( companies: list\[str\] = \[\ "NVIDIA",\ "Advanced Micro Devices",\ "Microsoft",\ "Alphabet",\ "Amazon",\ "Meta Platforms",\ "Broadcom",\ "Taiwan Semiconductor Manufacturing",\ \], focus: str = "Q4 earnings preview and competitive positioning", research\_effort: str = "standard", freshness: str = "month", ) -> ResearchReport: """Fan out across companies and aggregate cited equity briefings.""" with flyte.group("research-companies"): briefings = await asyncio.gather( \*\[\ research\_company(c, focus, research\_effort, freshness)\ for c in companies\ \] ) report = ResearchReport(briefings=list(briefings)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(financial\_research) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial\_research\_agent/main.py\* The Python packages are declared at the top of the file using the \`uv\` script style: \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # /// \`\`\` ## Data types Each \`Briefing\` carries a thesis, recent developments, risks, watch items, and a list of \`Source\` objects from both the Research and Search APIs. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "financial\_research" # params = "" # /// """Financial research & earnings-cycle agent. For each company, runs grounded, source-cited research via the You.com Research API plus a fresh-news layer via the Search API, then uses Claude to synthesize an analyst-ready equity briefing that preserves citations. Flyte caching cuts duplicate spend when runs converge. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="financial-research", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="financial-research", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str domain: str = "" snippet: str = "" published: str = "" favicon: str = "" section: str = "research" # "research", "news", or "web" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Briefing: company: str thesis: str recent\_developments: list\[str\] = field(default\_factory=list) risks: list\[str\] = field(default\_factory=list) watch\_items: list\[str\] = field(default\_factory=list) sources: list\[Source\] = field(default\_factory=list) @dataclass class ResearchReport: briefings: list\[Briefing\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_apis}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_request(method: str, url: str, timeout: float, \*\*kwargs) -> dict: """HTTP wrapper with exponential backoff + jitter on 429 rate limits. Fanned-out tasks run in separate pods, so we retry on the client side to smooth out bursts against the You.com API rate limit. """ import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} if method == "POST": headers\["Content-Type"\] = "application/json" async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.request(method, url, headers=headers, \*\*kwargs) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str, freshness: str) -> dict: """Grounded, citation-backed research answer.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": {"freshness": freshness}, } return await \_you\_request("POST", YOU\_RESEARCH\_URL, 300.0, json=body) @flyte.trace async def you\_news(query: str, count: int = 6, freshness: str = "week") -> list\[dict\]: """Fresh news headlines for a company.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_request("GET", YOU\_SEARCH\_URL, 60.0, params=params) results = data.get("results", {}) out: list\[dict\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") out.append( { "title": item.get("title", ""), "url": url, "domain": \_domain(url), "snippet": snippets\[0\] if snippets else item.get("description", ""), "published": item.get("page\_age", "") or "", "favicon": item.get("favicon\_url") or \_favicon\_for(url), "section": section, } ) return out # {{/docs-fragment you\_apis}} # {{docs-fragment llm}} @flyte.trace async def synthesize\_briefing(company: str, focus: str, research: str, news: str) -> dict: """Use Claude to synthesize a structured equity briefing.""" from litellm import acompletion system = ( "You are an equity research analyst. Using ONLY the grounded research " "and news provided, write a concise briefing. Respond ONLY with JSON: " '{"thesis": str, "recent\_developments": \[str\], "risks": \[str\], ' '"watch\_items": \[str\]}. Keep each list to 3-5 short, specific bullets.' ) user = ( f"Company: {company}\\nFocus: {focus}\\n\\n" f"Grounded research:\\n{research}\\n\\nRecent news:\\n{news}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment research\_company}} @env.task(retries=3) async def research\_company( company: str, focus: str, research\_effort: str, freshness: str, ) -> Briefing: """Research one company and synthesize a cited briefing.""" question = ( f"Provide a grounded analysis of {company} with respect to: {focus}. " f"Cover recent financial performance, strategic moves, competitive " f"positioning, and risks." ) research\_result, news = await asyncio.gather( you\_research(question, research\_effort, freshness), you\_news(f"{company} earnings news", freshness=freshness), ) output = research\_result.get("output", {}) research\_text = output.get("content", "") if not isinstance(research\_text, str): research\_text = json.dumps(research\_text) sources: list\[Source\] = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, domain=\_domain(url), snippet=str((s.get("snippets") or \[""\])\[0\]), favicon=\_favicon\_for(url), section="research", ) ) for n in news: sources.append( Source( title=str(n.get("title", "")), url=str(n.get("url", "")), domain=str(n.get("domain", "")), snippet=str(n.get("snippet", "")), published=str(n.get("published", "")), favicon=str(n.get("favicon", "")), section=str(n.get("section", "web")), ) ) news\_text = "\\n".join( f"- {n\['title'\]} ({n\['published'\]}) {n\['domain'\]}: {n\['snippet'\]\[:120\]}" for n in news ) parsed = await synthesize\_briefing(company, focus, research\_text, news\_text) def \_list(key: str) -> list\[str\]: return \[str(x) for x in (parsed.get(key) or \[\])\] return Briefing( company=company, thesis=str(parsed.get("thesis", "")), recent\_developments=\_list("recent\_developments"), risks=\_list("risks"), watch\_items=\_list("watch\_items"), sources=sources, ) # {{/docs-fragment research\_company}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com citation (Research or Search source).""" if not s.url: return "" tag\_cls = s.section if s.section in ("research", "news") else "web" meta\_bits = \[\] if s.published: meta\_bits.append(s.published\[:10\]) if s.title: meta\_bits.append(s.title) meta = " · ".join(meta\_bits) snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"{s.section}" f"
{meta}
{snip}
" ) def \_render\_report(report: ResearchReport) -> str: def \_ul(items: list\[str\]) -> str: if not items: return "

None reported.

" return "
    " + "".join(f"
  • {x}
    • " for x in items) + "
" cards = \[\] for b in report.briefings: src = "".join(\_cite(s) for s in b.sources\[:10\]) cards.append( f"

{b.company}

" f"
{b.thesis or 'No thesis generated.'}
" f"
" f"

Recent developments

{\_ul(b.recent\_developments)}
" f"

Risks

{\_ul(b.risks)}
" f"

Watch items

{\_ul(b.watch\_items)}
" f"
" + (f"

You.com sources ({len(b.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(b.sources) for b in report.briefings) return f""" {REPORT\_CSS}

Financial Research Briefings

Grounded, citation-backed equity briefings — each company backed by You.com Research synthesis plus fresh Search news.

{len(report.briefings)} companies {total\_sources} You.com sources cited
{''.join(cards) or "

No briefings generated.

"}

Research answers from the You.com Research API (grounded synthesis with inline citations) plus fresh headlines from the You.com Search API (web + auto-classified news with timestamps and snippets).

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def financial\_research( companies: list\[str\] = \[\ "NVIDIA",\ "Advanced Micro Devices",\ "Microsoft",\ "Alphabet",\ "Amazon",\ "Meta Platforms",\ "Broadcom",\ "Taiwan Semiconductor Manufacturing",\ \], focus: str = "Q4 earnings preview and competitive positioning", research\_effort: str = "standard", freshness: str = "month", ) -> ResearchReport: """Fan out across companies and aggregate cited equity briefings.""" with flyte.group("research-companies"): briefings = await asyncio.gather( \*\[\ research\_company(c, focus, research\_effort, freshness)\ for c in companies\ \] ) report = ResearchReport(briefings=list(briefings)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(financial\_research) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial\_research\_agent/main.py\* ## You.com Research and Search APIs The agent uses both You.com APIs in parallel for each company: - \*\*Research API\*\* (\`https://api.you.com/v1/research\`) — grounded, citation-backed analysis with configurable \`research\_effort\` (\`lite\`, \`standard\`, \`deep\`, \`exhaustive\`). See the \[Research API reference\](https://you.com/docs/api-reference/research/v1-research). - \*\*Search API\*\* (\`https://ydc-index.io/v1/search\`) — fresh news headlines with \`freshness\` filtering. See the \[Search API reference\](https://you.com/docs/api-reference/search/v1-search). \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "financial\_research" # params = "" # /// """Financial research & earnings-cycle agent. For each company, runs grounded, source-cited research via the You.com Research API plus a fresh-news layer via the Search API, then uses Claude to synthesize an analyst-ready equity briefing that preserves citations. Flyte caching cuts duplicate spend when runs converge. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="financial-research", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="financial-research", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str domain: str = "" snippet: str = "" published: str = "" favicon: str = "" section: str = "research" # "research", "news", or "web" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Briefing: company: str thesis: str recent\_developments: list\[str\] = field(default\_factory=list) risks: list\[str\] = field(default\_factory=list) watch\_items: list\[str\] = field(default\_factory=list) sources: list\[Source\] = field(default\_factory=list) @dataclass class ResearchReport: briefings: list\[Briefing\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_apis}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_request(method: str, url: str, timeout: float, \*\*kwargs) -> dict: """HTTP wrapper with exponential backoff + jitter on 429 rate limits. Fanned-out tasks run in separate pods, so we retry on the client side to smooth out bursts against the You.com API rate limit. """ import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} if method == "POST": headers\["Content-Type"\] = "application/json" async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.request(method, url, headers=headers, \*\*kwargs) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str, freshness: str) -> dict: """Grounded, citation-backed research answer.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": {"freshness": freshness}, } return await \_you\_request("POST", YOU\_RESEARCH\_URL, 300.0, json=body) @flyte.trace async def you\_news(query: str, count: int = 6, freshness: str = "week") -> list\[dict\]: """Fresh news headlines for a company.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_request("GET", YOU\_SEARCH\_URL, 60.0, params=params) results = data.get("results", {}) out: list\[dict\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") out.append( { "title": item.get("title", ""), "url": url, "domain": \_domain(url), "snippet": snippets\[0\] if snippets else item.get("description", ""), "published": item.get("page\_age", "") or "", "favicon": item.get("favicon\_url") or \_favicon\_for(url), "section": section, } ) return out # {{/docs-fragment you\_apis}} # {{docs-fragment llm}} @flyte.trace async def synthesize\_briefing(company: str, focus: str, research: str, news: str) -> dict: """Use Claude to synthesize a structured equity briefing.""" from litellm import acompletion system = ( "You are an equity research analyst. Using ONLY the grounded research " "and news provided, write a concise briefing. Respond ONLY with JSON: " '{"thesis": str, "recent\_developments": \[str\], "risks": \[str\], ' '"watch\_items": \[str\]}. Keep each list to 3-5 short, specific bullets.' ) user = ( f"Company: {company}\\nFocus: {focus}\\n\\n" f"Grounded research:\\n{research}\\n\\nRecent news:\\n{news}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment research\_company}} @env.task(retries=3) async def research\_company( company: str, focus: str, research\_effort: str, freshness: str, ) -> Briefing: """Research one company and synthesize a cited briefing.""" question = ( f"Provide a grounded analysis of {company} with respect to: {focus}. " f"Cover recent financial performance, strategic moves, competitive " f"positioning, and risks." ) research\_result, news = await asyncio.gather( you\_research(question, research\_effort, freshness), you\_news(f"{company} earnings news", freshness=freshness), ) output = research\_result.get("output", {}) research\_text = output.get("content", "") if not isinstance(research\_text, str): research\_text = json.dumps(research\_text) sources: list\[Source\] = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, domain=\_domain(url), snippet=str((s.get("snippets") or \[""\])\[0\]), favicon=\_favicon\_for(url), section="research", ) ) for n in news: sources.append( Source( title=str(n.get("title", "")), url=str(n.get("url", "")), domain=str(n.get("domain", "")), snippet=str(n.get("snippet", "")), published=str(n.get("published", "")), favicon=str(n.get("favicon", "")), section=str(n.get("section", "web")), ) ) news\_text = "\\n".join( f"- {n\['title'\]} ({n\['published'\]}) {n\['domain'\]}: {n\['snippet'\]\[:120\]}" for n in news ) parsed = await synthesize\_briefing(company, focus, research\_text, news\_text) def \_list(key: str) -> list\[str\]: return \[str(x) for x in (parsed.get(key) or \[\])\] return Briefing( company=company, thesis=str(parsed.get("thesis", "")), recent\_developments=\_list("recent\_developments"), risks=\_list("risks"), watch\_items=\_list("watch\_items"), sources=sources, ) # {{/docs-fragment research\_company}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com citation (Research or Search source).""" if not s.url: return "" tag\_cls = s.section if s.section in ("research", "news") else "web" meta\_bits = \[\] if s.published: meta\_bits.append(s.published\[:10\]) if s.title: meta\_bits.append(s.title) meta = " · ".join(meta\_bits) snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"{s.section}" f"
{meta}
{snip}
" ) def \_render\_report(report: ResearchReport) -> str: def \_ul(items: list\[str\]) -> str: if not items: return "

None reported.

" return "
    " + "".join(f"
  • {x}
    • " for x in items) + "
" cards = \[\] for b in report.briefings: src = "".join(\_cite(s) for s in b.sources\[:10\]) cards.append( f"

{b.company}

" f"
{b.thesis or 'No thesis generated.'}
" f"
" f"

Recent developments

{\_ul(b.recent\_developments)}
" f"

Risks

{\_ul(b.risks)}
" f"

Watch items

{\_ul(b.watch\_items)}
" f"
" + (f"

You.com sources ({len(b.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(b.sources) for b in report.briefings) return f""" {REPORT\_CSS}

Financial Research Briefings

Grounded, citation-backed equity briefings — each company backed by You.com Research synthesis plus fresh Search news.

{len(report.briefings)} companies {total\_sources} You.com sources cited
{''.join(cards) or "

No briefings generated.

"}

Research answers from the You.com Research API (grounded synthesis with inline citations) plus fresh headlines from the You.com Search API (web + auto-classified news with timestamps and snippets).

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def financial\_research( companies: list\[str\] = \[\ "NVIDIA",\ "Advanced Micro Devices",\ "Microsoft",\ "Alphabet",\ "Amazon",\ "Meta Platforms",\ "Broadcom",\ "Taiwan Semiconductor Manufacturing",\ \], focus: str = "Q4 earnings preview and competitive positioning", research\_effort: str = "standard", freshness: str = "month", ) -> ResearchReport: """Fan out across companies and aggregate cited equity briefings.""" with flyte.group("research-companies"): briefings = await asyncio.gather( \*\[\ research\_company(c, focus, research\_effort, freshness)\ for c in companies\ \] ) report = ResearchReport(briefings=list(briefings)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(financial\_research) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial\_research\_agent/main.py\* ## Synthesize briefings with Claude Claude, routed through LiteLLM, turns the grounded research answer and news headlines into a structured equity briefing grounded in the evidence provided. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "financial\_research" # params = "" # /// """Financial research & earnings-cycle agent. For each company, runs grounded, source-cited research via the You.com Research API plus a fresh-news layer via the Search API, then uses Claude to synthesize an analyst-ready equity briefing that preserves citations. Flyte caching cuts duplicate spend when runs converge. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="financial-research", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="financial-research", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str domain: str = "" snippet: str = "" published: str = "" favicon: str = "" section: str = "research" # "research", "news", or "web" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Briefing: company: str thesis: str recent\_developments: list\[str\] = field(default\_factory=list) risks: list\[str\] = field(default\_factory=list) watch\_items: list\[str\] = field(default\_factory=list) sources: list\[Source\] = field(default\_factory=list) @dataclass class ResearchReport: briefings: list\[Briefing\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_apis}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_request(method: str, url: str, timeout: float, \*\*kwargs) -> dict: """HTTP wrapper with exponential backoff + jitter on 429 rate limits. Fanned-out tasks run in separate pods, so we retry on the client side to smooth out bursts against the You.com API rate limit. """ import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} if method == "POST": headers\["Content-Type"\] = "application/json" async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.request(method, url, headers=headers, \*\*kwargs) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str, freshness: str) -> dict: """Grounded, citation-backed research answer.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": {"freshness": freshness}, } return await \_you\_request("POST", YOU\_RESEARCH\_URL, 300.0, json=body) @flyte.trace async def you\_news(query: str, count: int = 6, freshness: str = "week") -> list\[dict\]: """Fresh news headlines for a company.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_request("GET", YOU\_SEARCH\_URL, 60.0, params=params) results = data.get("results", {}) out: list\[dict\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") out.append( { "title": item.get("title", ""), "url": url, "domain": \_domain(url), "snippet": snippets\[0\] if snippets else item.get("description", ""), "published": item.get("page\_age", "") or "", "favicon": item.get("favicon\_url") or \_favicon\_for(url), "section": section, } ) return out # {{/docs-fragment you\_apis}} # {{docs-fragment llm}} @flyte.trace async def synthesize\_briefing(company: str, focus: str, research: str, news: str) -> dict: """Use Claude to synthesize a structured equity briefing.""" from litellm import acompletion system = ( "You are an equity research analyst. Using ONLY the grounded research " "and news provided, write a concise briefing. Respond ONLY with JSON: " '{"thesis": str, "recent\_developments": \[str\], "risks": \[str\], ' '"watch\_items": \[str\]}. Keep each list to 3-5 short, specific bullets.' ) user = ( f"Company: {company}\\nFocus: {focus}\\n\\n" f"Grounded research:\\n{research}\\n\\nRecent news:\\n{news}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment research\_company}} @env.task(retries=3) async def research\_company( company: str, focus: str, research\_effort: str, freshness: str, ) -> Briefing: """Research one company and synthesize a cited briefing.""" question = ( f"Provide a grounded analysis of {company} with respect to: {focus}. " f"Cover recent financial performance, strategic moves, competitive " f"positioning, and risks." ) research\_result, news = await asyncio.gather( you\_research(question, research\_effort, freshness), you\_news(f"{company} earnings news", freshness=freshness), ) output = research\_result.get("output", {}) research\_text = output.get("content", "") if not isinstance(research\_text, str): research\_text = json.dumps(research\_text) sources: list\[Source\] = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, domain=\_domain(url), snippet=str((s.get("snippets") or \[""\])\[0\]), favicon=\_favicon\_for(url), section="research", ) ) for n in news: sources.append( Source( title=str(n.get("title", "")), url=str(n.get("url", "")), domain=str(n.get("domain", "")), snippet=str(n.get("snippet", "")), published=str(n.get("published", "")), favicon=str(n.get("favicon", "")), section=str(n.get("section", "web")), ) ) news\_text = "\\n".join( f"- {n\['title'\]} ({n\['published'\]}) {n\['domain'\]}: {n\['snippet'\]\[:120\]}" for n in news ) parsed = await synthesize\_briefing(company, focus, research\_text, news\_text) def \_list(key: str) -> list\[str\]: return \[str(x) for x in (parsed.get(key) or \[\])\] return Briefing( company=company, thesis=str(parsed.get("thesis", "")), recent\_developments=\_list("recent\_developments"), risks=\_list("risks"), watch\_items=\_list("watch\_items"), sources=sources, ) # {{/docs-fragment research\_company}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com citation (Research or Search source).""" if not s.url: return "" tag\_cls = s.section if s.section in ("research", "news") else "web" meta\_bits = \[\] if s.published: meta\_bits.append(s.published\[:10\]) if s.title: meta\_bits.append(s.title) meta = " · ".join(meta\_bits) snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"{s.section}" f"
{meta}
{snip}
" ) def \_render\_report(report: ResearchReport) -> str: def \_ul(items: list\[str\]) -> str: if not items: return "

None reported.

" return "
    " + "".join(f"
  • {x}
    • " for x in items) + "
" cards = \[\] for b in report.briefings: src = "".join(\_cite(s) for s in b.sources\[:10\]) cards.append( f"

{b.company}

" f"
{b.thesis or 'No thesis generated.'}
" f"
" f"

Recent developments

{\_ul(b.recent\_developments)}
" f"

Risks

{\_ul(b.risks)}
" f"

Watch items

{\_ul(b.watch\_items)}
" f"
" + (f"

You.com sources ({len(b.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(b.sources) for b in report.briefings) return f""" {REPORT\_CSS}

Financial Research Briefings

Grounded, citation-backed equity briefings — each company backed by You.com Research synthesis plus fresh Search news.

{len(report.briefings)} companies {total\_sources} You.com sources cited
{''.join(cards) or "

No briefings generated.

"}

Research answers from the You.com Research API (grounded synthesis with inline citations) plus fresh headlines from the You.com Search API (web + auto-classified news with timestamps and snippets).

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def financial\_research( companies: list\[str\] = \[\ "NVIDIA",\ "Advanced Micro Devices",\ "Microsoft",\ "Alphabet",\ "Amazon",\ "Meta Platforms",\ "Broadcom",\ "Taiwan Semiconductor Manufacturing",\ \], focus: str = "Q4 earnings preview and competitive positioning", research\_effort: str = "standard", freshness: str = "month", ) -> ResearchReport: """Fan out across companies and aggregate cited equity briefings.""" with flyte.group("research-companies"): briefings = await asyncio.gather( \*\[\ research\_company(c, focus, research\_effort, freshness)\ for c in companies\ \] ) report = ResearchReport(briefings=list(briefings)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(financial\_research) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial\_research\_agent/main.py\* ## Research one company The \`research\_company\` task calls both You.com APIs in parallel, collects sources, and synthesizes a structured briefing. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "financial\_research" # params = "" # /// """Financial research & earnings-cycle agent. For each company, runs grounded, source-cited research via the You.com Research API plus a fresh-news layer via the Search API, then uses Claude to synthesize an analyst-ready equity briefing that preserves citations. Flyte caching cuts duplicate spend when runs converge. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="financial-research", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="financial-research", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str domain: str = "" snippet: str = "" published: str = "" favicon: str = "" section: str = "research" # "research", "news", or "web" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Briefing: company: str thesis: str recent\_developments: list\[str\] = field(default\_factory=list) risks: list\[str\] = field(default\_factory=list) watch\_items: list\[str\] = field(default\_factory=list) sources: list\[Source\] = field(default\_factory=list) @dataclass class ResearchReport: briefings: list\[Briefing\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_apis}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_request(method: str, url: str, timeout: float, \*\*kwargs) -> dict: """HTTP wrapper with exponential backoff + jitter on 429 rate limits. Fanned-out tasks run in separate pods, so we retry on the client side to smooth out bursts against the You.com API rate limit. """ import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} if method == "POST": headers\["Content-Type"\] = "application/json" async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.request(method, url, headers=headers, \*\*kwargs) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str, freshness: str) -> dict: """Grounded, citation-backed research answer.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": {"freshness": freshness}, } return await \_you\_request("POST", YOU\_RESEARCH\_URL, 300.0, json=body) @flyte.trace async def you\_news(query: str, count: int = 6, freshness: str = "week") -> list\[dict\]: """Fresh news headlines for a company.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_request("GET", YOU\_SEARCH\_URL, 60.0, params=params) results = data.get("results", {}) out: list\[dict\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") out.append( { "title": item.get("title", ""), "url": url, "domain": \_domain(url), "snippet": snippets\[0\] if snippets else item.get("description", ""), "published": item.get("page\_age", "") or "", "favicon": item.get("favicon\_url") or \_favicon\_for(url), "section": section, } ) return out # {{/docs-fragment you\_apis}} # {{docs-fragment llm}} @flyte.trace async def synthesize\_briefing(company: str, focus: str, research: str, news: str) -> dict: """Use Claude to synthesize a structured equity briefing.""" from litellm import acompletion system = ( "You are an equity research analyst. Using ONLY the grounded research " "and news provided, write a concise briefing. Respond ONLY with JSON: " '{"thesis": str, "recent\_developments": \[str\], "risks": \[str\], ' '"watch\_items": \[str\]}. Keep each list to 3-5 short, specific bullets.' ) user = ( f"Company: {company}\\nFocus: {focus}\\n\\n" f"Grounded research:\\n{research}\\n\\nRecent news:\\n{news}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment research\_company}} @env.task(retries=3) async def research\_company( company: str, focus: str, research\_effort: str, freshness: str, ) -> Briefing: """Research one company and synthesize a cited briefing.""" question = ( f"Provide a grounded analysis of {company} with respect to: {focus}. " f"Cover recent financial performance, strategic moves, competitive " f"positioning, and risks." ) research\_result, news = await asyncio.gather( you\_research(question, research\_effort, freshness), you\_news(f"{company} earnings news", freshness=freshness), ) output = research\_result.get("output", {}) research\_text = output.get("content", "") if not isinstance(research\_text, str): research\_text = json.dumps(research\_text) sources: list\[Source\] = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, domain=\_domain(url), snippet=str((s.get("snippets") or \[""\])\[0\]), favicon=\_favicon\_for(url), section="research", ) ) for n in news: sources.append( Source( title=str(n.get("title", "")), url=str(n.get("url", "")), domain=str(n.get("domain", "")), snippet=str(n.get("snippet", "")), published=str(n.get("published", "")), favicon=str(n.get("favicon", "")), section=str(n.get("section", "web")), ) ) news\_text = "\\n".join( f"- {n\['title'\]} ({n\['published'\]}) {n\['domain'\]}: {n\['snippet'\]\[:120\]}" for n in news ) parsed = await synthesize\_briefing(company, focus, research\_text, news\_text) def \_list(key: str) -> list\[str\]: return \[str(x) for x in (parsed.get(key) or \[\])\] return Briefing( company=company, thesis=str(parsed.get("thesis", "")), recent\_developments=\_list("recent\_developments"), risks=\_list("risks"), watch\_items=\_list("watch\_items"), sources=sources, ) # {{/docs-fragment research\_company}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com citation (Research or Search source).""" if not s.url: return "" tag\_cls = s.section if s.section in ("research", "news") else "web" meta\_bits = \[\] if s.published: meta\_bits.append(s.published\[:10\]) if s.title: meta\_bits.append(s.title) meta = " · ".join(meta\_bits) snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"{s.section}" f"
{meta}
{snip}
" ) def \_render\_report(report: ResearchReport) -> str: def \_ul(items: list\[str\]) -> str: if not items: return "

None reported.

" return "
    " + "".join(f"
  • {x}
    • " for x in items) + "
" cards = \[\] for b in report.briefings: src = "".join(\_cite(s) for s in b.sources\[:10\]) cards.append( f"

{b.company}

" f"
{b.thesis or 'No thesis generated.'}
" f"
" f"

Recent developments

{\_ul(b.recent\_developments)}
" f"

Risks

{\_ul(b.risks)}
" f"

Watch items

{\_ul(b.watch\_items)}
" f"
" + (f"

You.com sources ({len(b.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(b.sources) for b in report.briefings) return f""" {REPORT\_CSS}

Financial Research Briefings

Grounded, citation-backed equity briefings — each company backed by You.com Research synthesis plus fresh Search news.

{len(report.briefings)} companies {total\_sources} You.com sources cited
{''.join(cards) or "

No briefings generated.

"}

Research answers from the You.com Research API (grounded synthesis with inline citations) plus fresh headlines from the You.com Search API (web + auto-classified news with timestamps and snippets).

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def financial\_research( companies: list\[str\] = \[\ "NVIDIA",\ "Advanced Micro Devices",\ "Microsoft",\ "Alphabet",\ "Amazon",\ "Meta Platforms",\ "Broadcom",\ "Taiwan Semiconductor Manufacturing",\ \], focus: str = "Q4 earnings preview and competitive positioning", research\_effort: str = "standard", freshness: str = "month", ) -> ResearchReport: """Fan out across companies and aggregate cited equity briefings.""" with flyte.group("research-companies"): briefings = await asyncio.gather( \*\[\ research\_company(c, focus, research\_effort, freshness)\ for c in companies\ \] ) report = ResearchReport(briefings=list(briefings)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(financial\_research) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial\_research\_agent/main.py\* ## Orchestration The \`financial\_research\` driver task fans out across all companies and renders a Flyte report with per-company briefings and citations. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "financial\_research" # params = "" # /// """Financial research & earnings-cycle agent. For each company, runs grounded, source-cited research via the You.com Research API plus a fresh-news layer via the Search API, then uses Claude to synthesize an analyst-ready equity briefing that preserves citations. Flyte caching cuts duplicate spend when runs converge. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="financial-research", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="financial-research", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str domain: str = "" snippet: str = "" published: str = "" favicon: str = "" section: str = "research" # "research", "news", or "web" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Briefing: company: str thesis: str recent\_developments: list\[str\] = field(default\_factory=list) risks: list\[str\] = field(default\_factory=list) watch\_items: list\[str\] = field(default\_factory=list) sources: list\[Source\] = field(default\_factory=list) @dataclass class ResearchReport: briefings: list\[Briefing\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_apis}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_request(method: str, url: str, timeout: float, \*\*kwargs) -> dict: """HTTP wrapper with exponential backoff + jitter on 429 rate limits. Fanned-out tasks run in separate pods, so we retry on the client side to smooth out bursts against the You.com API rate limit. """ import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} if method == "POST": headers\["Content-Type"\] = "application/json" async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.request(method, url, headers=headers, \*\*kwargs) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str, freshness: str) -> dict: """Grounded, citation-backed research answer.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": {"freshness": freshness}, } return await \_you\_request("POST", YOU\_RESEARCH\_URL, 300.0, json=body) @flyte.trace async def you\_news(query: str, count: int = 6, freshness: str = "week") -> list\[dict\]: """Fresh news headlines for a company.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_request("GET", YOU\_SEARCH\_URL, 60.0, params=params) results = data.get("results", {}) out: list\[dict\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") out.append( { "title": item.get("title", ""), "url": url, "domain": \_domain(url), "snippet": snippets\[0\] if snippets else item.get("description", ""), "published": item.get("page\_age", "") or "", "favicon": item.get("favicon\_url") or \_favicon\_for(url), "section": section, } ) return out # {{/docs-fragment you\_apis}} # {{docs-fragment llm}} @flyte.trace async def synthesize\_briefing(company: str, focus: str, research: str, news: str) -> dict: """Use Claude to synthesize a structured equity briefing.""" from litellm import acompletion system = ( "You are an equity research analyst. Using ONLY the grounded research " "and news provided, write a concise briefing. Respond ONLY with JSON: " '{"thesis": str, "recent\_developments": \[str\], "risks": \[str\], ' '"watch\_items": \[str\]}. Keep each list to 3-5 short, specific bullets.' ) user = ( f"Company: {company}\\nFocus: {focus}\\n\\n" f"Grounded research:\\n{research}\\n\\nRecent news:\\n{news}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment research\_company}} @env.task(retries=3) async def research\_company( company: str, focus: str, research\_effort: str, freshness: str, ) -> Briefing: """Research one company and synthesize a cited briefing.""" question = ( f"Provide a grounded analysis of {company} with respect to: {focus}. " f"Cover recent financial performance, strategic moves, competitive " f"positioning, and risks." ) research\_result, news = await asyncio.gather( you\_research(question, research\_effort, freshness), you\_news(f"{company} earnings news", freshness=freshness), ) output = research\_result.get("output", {}) research\_text = output.get("content", "") if not isinstance(research\_text, str): research\_text = json.dumps(research\_text) sources: list\[Source\] = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, domain=\_domain(url), snippet=str((s.get("snippets") or \[""\])\[0\]), favicon=\_favicon\_for(url), section="research", ) ) for n in news: sources.append( Source( title=str(n.get("title", "")), url=str(n.get("url", "")), domain=str(n.get("domain", "")), snippet=str(n.get("snippet", "")), published=str(n.get("published", "")), favicon=str(n.get("favicon", "")), section=str(n.get("section", "web")), ) ) news\_text = "\\n".join( f"- {n\['title'\]} ({n\['published'\]}) {n\['domain'\]}: {n\['snippet'\]\[:120\]}" for n in news ) parsed = await synthesize\_briefing(company, focus, research\_text, news\_text) def \_list(key: str) -> list\[str\]: return \[str(x) for x in (parsed.get(key) or \[\])\] return Briefing( company=company, thesis=str(parsed.get("thesis", "")), recent\_developments=\_list("recent\_developments"), risks=\_list("risks"), watch\_items=\_list("watch\_items"), sources=sources, ) # {{/docs-fragment research\_company}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com citation (Research or Search source).""" if not s.url: return "" tag\_cls = s.section if s.section in ("research", "news") else "web" meta\_bits = \[\] if s.published: meta\_bits.append(s.published\[:10\]) if s.title: meta\_bits.append(s.title) meta = " · ".join(meta\_bits) snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"{s.section}" f"
{meta}
{snip}
" ) def \_render\_report(report: ResearchReport) -> str: def \_ul(items: list\[str\]) -> str: if not items: return "

None reported.

" return "
    " + "".join(f"
  • {x}
    • " for x in items) + "
" cards = \[\] for b in report.briefings: src = "".join(\_cite(s) for s in b.sources\[:10\]) cards.append( f"

{b.company}

" f"
{b.thesis or 'No thesis generated.'}
" f"
" f"

Recent developments

{\_ul(b.recent\_developments)}
" f"

Risks

{\_ul(b.risks)}
" f"

Watch items

{\_ul(b.watch\_items)}
" f"
" + (f"

You.com sources ({len(b.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(b.sources) for b in report.briefings) return f""" {REPORT\_CSS}

Financial Research Briefings

Grounded, citation-backed equity briefings — each company backed by You.com Research synthesis plus fresh Search news.

{len(report.briefings)} companies {total\_sources} You.com sources cited
{''.join(cards) or "

No briefings generated.

"}

Research answers from the You.com Research API (grounded synthesis with inline citations) plus fresh headlines from the You.com Search API (web + auto-classified news with timestamps and snippets).

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def financial\_research( companies: list\[str\] = \[\ "NVIDIA",\ "Advanced Micro Devices",\ "Microsoft",\ "Alphabet",\ "Amazon",\ "Meta Platforms",\ "Broadcom",\ "Taiwan Semiconductor Manufacturing",\ \], focus: str = "Q4 earnings preview and competitive positioning", research\_effort: str = "standard", freshness: str = "month", ) -> ResearchReport: """Fan out across companies and aggregate cited equity briefings.""" with flyte.group("research-companies"): briefings = await asyncio.gather( \*\[\ research\_company(c, focus, research\_effort, freshness)\ for c in companies\ \] ) report = ResearchReport(briefings=list(briefings)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(financial\_research) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/financial\_research\_agent/main.py\* ## Run the agent ### Create secrets Get a You.com API key from the \[You.com platform\](https://you.com/platform) (see the \[quickstart guide\](https://you.com/docs/quickstart)). Get an Anthropic API key from the \[Anthropic console\](https://console.anthropic.com/). Register both keys as Flyte secrets. The secret key names must match those declared in the \`TaskEnvironment\`: \`\`\` flyte create secret youdotcom-api-key flyte create secret internal-anthropic-api-key \`\`\` See \[Secrets\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets/page.md) for scoping and file-based secrets. ### Run locally or remotely From the \[example directory\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/financial\_research\_agent): \`\`\` cd v2/tutorials/financial\_research\_agent uv run --script main.py \`\`\` To test locally without Flyte secrets: \`\`\` export YOU\_API\_KEY= export ANTHROPIC\_API\_KEY= uv run --script main.py \`\`\` When the run completes, open the Flyte report to review equity briefings with thesis, risks, and You.com source citations for each company. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/frontier-ai === # Frontier AI Tutorials for frontier-model pretraining, automated experimentation, and large-scale AI workloads. ### \*\*Frontier AI > Distributed LLM pretraining\*\* Pretrain large language models at scale with PyTorch Lightning, FSDP, and H200 GPUs, featuring streaming data and real-time metrics. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/frontier-ai/distributed-pretraining === # Distributed LLM pretraining When training large models, infrastructure should not be the hardest part. The real work is in the model architecture, the data, and the hyperparameters. In practice, though, teams often spend weeks just trying to get distributed training to run reliably. And when it breaks, it usually breaks in familiar ways: out-of-memory crashes, corrupted checkpoints, data loaders that silently fail, or runs that hang with no obvious explanation. Most distributed training tutorials focus on PyTorch primitives. This one focuses on getting something that actually ships. We go into the technical details, such as how FSDP shards parameters, why gradient clipping behaves differently at scale, and how streaming datasets reduce memory pressure, but always with the goal of building a system that works in production. Real training jobs need more than a training loop. They need checkpointing, fault tolerance, data streaming, visibility into what’s happening, and the ability to recover from failures. In this tutorial, we build all of that using Flyte, without having to stand up or manage any additional infrastructure. > \[!NOTE\] > Full code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/pretraining/train.py). ## Overview We're going to pretrain a GPT-2 style language model from scratch. This involves training on raw text data starting from randomly initialized weights, rather than fine-tuning or adapting a pretrained model. This is the same process used to train the original GPT-2, LLaMA, and most other foundation models. The model learns by predicting the next token. Given "The cat sat on the", it learns to predict "mat". Do this billions of times across terabytes of text, and the model develops surprisingly sophisticated language understanding. That's pretraining. The challenge is scale. A 30B parameter model doesn't fit on a single GPU. The training dataset, \[SlimPajama\](https://huggingface.co/datasets/cerebras/SlimPajama-627B) in our case, is 627 billion tokens. Training runs last for days or even weeks. To make this work, you need: - \*\*Distributed training\*\*: Split the model across multiple GPUs using \[FSDP (Fully Sharded Data Parallel)\](https://docs.pytorch.org/tutorials/intermediate/FSDP\_tutorial.html) - \*\*Data streaming\*\*: Pull training data on-demand instead of downloading terabytes upfront - \*\*Checkpointing\*\*: Save progress regularly so a failure doesn’t wipe out days of compute - \*\*Observability\*\*: See what's happening inside a multi-day training run We’ll build a Flyte pipeline that takes care of all of this, using three tasks with clearly defined responsibilities: 1. \*\*Data preparation\*\*: Tokenizes your dataset and converts it to MDS (MosaicML Data Shard) format for streaming. This Flyte task is cached, so it only needs to be run once and can be reused across runs. 2. \*\*Distributed training\*\*: Runs FSDP across 8 H200 GPUs. Flyte's \`Elastic\` plugin handles the distributed setup. Checkpoints upload to S3 automatically via Flyte's \`File\` abstraction. 3. \*\*Real-time reporting\*\*: Streams loss curves and training metrics to Flyte Reports, a live dashboard integrated into the Flyte UI. Why three separate tasks? Flyte makes this separation efficient: - \*\*Caching\*\*: The data preparation step runs once. On subsequent runs, Flyte skips it entirely. - \*\*Resource isolation\*\*: Training uses expensive H200 GPUs only while actively training, while the driver runs on inexpensive CPU instances. - \*\*Fault boundaries\*\*: If training fails, the data preparation step does not re-run. Training can resume directly from the most recent checkpoint. ## Implementation Let's walk through the code. We'll start with the infrastructure setup, build the model, then wire everything together into a pipeline. ### Setting up the environment Every distributed training job needs a consistent environment across all nodes. Flyte handles this with container images: \`\`\` import logging import math import os from pathlib import Path from typing import Optional import flyte import flyte.report import lightning as L import numpy as np import torch import torch.nn as nn from flyte.io import Dir, File from flyteplugins.pytorch.task import Elastic \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* The imports tell the story: \`flyte\` for orchestration, \`flyte.report\` for live dashboards, \`lightning\` for training loop management, and \`Elastic\` from Flyte's PyTorch plugin. This last one is key as it configures PyTorch's distributed launch without you writing any distributed setup code. \`\`\` NUM\_NODES = 1 DEVICES\_PER\_NODE = 8 VOCAB\_SIZE = ( 50257 # GPT-2 BPE tokenizer vocabulary size (constant across all model sizes) ) N\_POSITIONS = 2048 # Maximum sequence length (constant across all model sizes) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* These constants define the distributed topology. We're using 1 node with 8 GPUs, but you can scale this up by changing \`NUM\_NODES\`. The vocabulary size (50,257 tokens) and sequence length (2,048 tokens) match GPT-2's \[Byte Pair Encoding (BPE) tokenizer\](https://huggingface.co/learn/llm-course/en/chapter6/5). \`\`\` image = flyte.Image.from\_debian\_base( name="distributed\_training\_h200" ).with\_pip\_packages( "transformers==4.57.3", "datasets==4.4.1", "tokenizers==0.22.1", "huggingface-hub==0.34.0", "mosaicml-streaming>=0.7.0", "pyarrow==22.0.0", "flyteplugins-pytorch>=2.0.0b33", "torch==2.9.1", "lightning==2.5.6", "tensorboard==2.20.0", "sentencepiece==0.2.1", ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* Flyte builds this container automatically when the pipeline is run. All dependencies required for distributed training, including PyTorch, Lightning, the streaming library, and NCCL for GPU communication, are baked in. There's no Dockerfile to maintain and no "works on my machine" debugging. ### Declaring resource requirements Different parts of the pipeline need different resources. Data tokenization needs CPU and memory. Training needs GPUs. The driver just coordinates. Flyte's \`TaskEnvironment\` lets you declare exactly what each task needs: \`\`\` data\_loading\_env = flyte.TaskEnvironment( name="data\_loading\_h200", image=image, resources=flyte.Resources(cpu=5, memory="28Gi", disk="100Gi"), env\_vars={ "HF\_DATASETS\_CACHE": "/tmp/hf\_cache", # Cache directory for datasets "TOKENIZERS\_PARALLELISM": "true", # Enable parallel tokenization }, cache="auto", ) distributed\_llm\_training\_env = flyte.TaskEnvironment( name="distributed\_llm\_training\_h200", image=image, resources=flyte.Resources( cpu=64, memory="512Gi", gpu=f"H200:{DEVICES\_PER\_NODE}", disk="1Ti", shm="16Gi", # Explicit shared memory for NCCL communication ), plugin\_config=Elastic(nnodes=NUM\_NODES, nproc\_per\_node=DEVICES\_PER\_NODE), env\_vars={ "TORCH\_DISTRIBUTED\_DEBUG": "INFO", "NCCL\_DEBUG": "WARN", }, cache="auto", ) driver\_env = flyte.TaskEnvironment( name="llm\_training\_driver", image=image, resources=flyte.Resources(cpu=2, memory="4Gi"), cache="auto", depends\_on=\[data\_loading\_env, distributed\_llm\_training\_env\], ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* Let's break down the training environment, since this is where most of the complexity lives: - \*\*\`gpu=f"H200:{DEVICES\_PER\_NODE}"\`\*\*: Flyte provisions exactly 8 H200 GPUs. These have 141GB of memory each, enough to train 30B+ parameter models with FSDP. - \*\*\`shm="16Gi"\`\*\*: This allocates explicit shared memory. NCCL (NVIDIA's communication library) uses shared memory for inter-GPU communication on the same node. Without this, you'll see cryptic errors like "NCCL error: unhandled system error", which can be difficult to debug. - \*\*\`Elastic(nnodes=NUM\_NODES, nproc\_per\_node=DEVICES\_PER\_NODE)\`\*\*: This is Flyte's integration with PyTorch's elastic launch. It handles process spawning (one process per GPU), rank assignment (each process knows its ID), and environment setup (master address, world size). This replaces the boilerplate typically written in shell scripts. The \`driver\_env\` is intentionally lightweight, using 2 CPUs and 4 GB of memory. Its role is limited to orchestrating tasks and passing data between them, so allocating GPUs here would be unnecessary. ### Model configurations Training a 1.5B model uses different hyperparameters than training a 65B model. Rather than hardcoding values, we define presets: \`\`\` MODEL\_CONFIGS = { "1.5B": { "n\_embd": 2048, "n\_layer": 24, "n\_head": 16, "batch\_size": 8, "learning\_rate": 6e-4, "checkpoint\_every\_n\_steps": 10, "report\_every\_n\_steps": 5, "val\_check\_interval": 100, }, # Good for testing and debugging "30B": { "n\_embd": 6656, "n\_layer": 48, "n\_head": 52, "batch\_size": 1, "learning\_rate": 1.6e-4, "checkpoint\_every\_n\_steps": 7500, "report\_every\_n\_steps": 200, "val\_check\_interval": 1000, }, "65B": { "n\_embd": 8192, "n\_layer": 80, "n\_head": 64, "batch\_size": 1, "learning\_rate": 1.5e-4, "checkpoint\_every\_n\_steps": 10000, "report\_every\_n\_steps": 250, "val\_check\_interval": 2000, }, } def get\_model\_config(model\_size: str) -> dict: if model\_size not in MODEL\_CONFIGS: available = ", ".join(MODEL\_CONFIGS.keys()) raise ValueError(f"Unknown model size: {model\_size}. Available: {available}") return MODEL\_CONFIGS\[model\_size\] \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* A few things to notice: - \*\*Batch size decreases with model size\*\*: For a fixed GPU memory budget, larger models consume more memory for parameters, optimizer state, and activations, leaving less room for per-GPU batch size. For example, a 1.5B parameter model may fit a batch size of 8 per GPU, while a 65B model may only fit a batch size of 1. This is typically compensated for using gradient accumulation to maintain a larger effective batch size. - \*\*Learning rate decreases with model size\*\*: Larger models are more sensitive to optimization instability and typically require lower learning rates. The values here follow empirical best practices used in large-scale language model training, informed by work such as the \[Chinchilla study\](https://arxiv.org/pdf/2203.15556) on compute-optimal scaling. - \*\*Checkpoint frequency increases with model size\*\*: Checkpointing a 65B model is expensive (the checkpoint is huge). We do it less often but make sure we don't lose too much progress if something fails. The 1.5B config is good for testing your setup before committing to a serious training run. ### Building the GPT model Now for the model itself. We're building a GPT-2 style decoder-only transformer from scratch. First, the configuration class: \`\`\` class GPTConfig: """Configuration for GPT model.""" def \_\_init\_\_( self, vocab\_size: int = VOCAB\_SIZE, n\_positions: int = N\_POSITIONS, n\_embd: int = 2048, n\_layer: int = 24, n\_head: int = 16, n\_inner: Optional\[int\] = None, activation\_function: str = "gelu\_new", dropout: float = 0.1, layer\_norm\_epsilon: float = 1e-5, ): self.vocab\_size = vocab\_size self.n\_positions = n\_positions self.n\_embd = n\_embd self.n\_layer = n\_layer self.n\_head = n\_head self.n\_inner = n\_inner if n\_inner is not None else 4 \* n\_embd self.activation\_function = activation\_function self.dropout = dropout self.layer\_norm\_epsilon = layer\_norm\_epsilon \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* The key architectural parameters: - \*\*\`n\_embd\`\*\*: The hidden (embedding) dimension. Larger values increase model capacity but also increase memory and compute requirements. - \*\*\`n\_layer\`\*\*: The number of transformer blocks. Model depth strongly influences expressiveness and performance. - \*\*\`n\_head\`\*\*: The number of attention heads. Each head can attend to different patterns or relationships in the input. - \*\*\`n\_inner\`\*\*: The hidden dimension of the feed-forward network (MLP), typically set to 4x the embedding dimension. Next, we define a single transformer block: \`\`\` class GPTBlock(nn.Module): """Transformer block with causal self-attention.""" def \_\_init\_\_(self, config: GPTConfig): super().\_\_init\_\_() self.ln\_1 = nn.LayerNorm(config.n\_embd, eps=config.layer\_norm\_epsilon) self.attn = nn.MultiheadAttention( config.n\_embd, config.n\_head, dropout=config.dropout, batch\_first=True, ) self.ln\_2 = nn.LayerNorm(config.n\_embd, eps=config.layer\_norm\_epsilon) # Get activation function from config ACT\_FNS = { "gelu": nn.GELU(), "gelu\_new": nn.GELU(approximate="tanh"), # GPT-2 uses approximate GELU "relu": nn.ReLU(), "silu": nn.SiLU(), "swish": nn.SiLU(), # SiLU = Swish } act\_fn = ACT\_FNS.get(config.activation\_function, nn.GELU()) self.mlp = nn.Sequential( nn.Linear(config.n\_embd, config.n\_inner), act\_fn, nn.Linear(config.n\_inner, config.n\_embd), nn.Dropout(config.dropout), ) def forward(self, x, causal\_mask, key\_padding\_mask=None): x\_normed = self.ln\_1(x) # Self-attention with causal and padding masks attn\_output, \_ = self.attn( x\_normed, # query x\_normed, # key x\_normed, # value attn\_mask=causal\_mask, # Causal mask: (seq\_len, seq\_len) key\_padding\_mask=key\_padding\_mask, # Padding mask: (batch, seq\_len) need\_weights=False, ) x = x + attn\_output # MLP with residual x = x + self.mlp(self.ln\_2(x)) return x \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* Each block has two sub-layers: causal self-attention and a feed-forward MLP. The causal mask ensures the model can only attend to previous tokens in the sequence, so it can't "cheat" by looking at the answer. This is what makes it \*autoregressive\*. The full \`GPTModel\` class (see the complete code) stacks these blocks and adds token and positional embeddings. One important detail is that the input token embedding matrix is shared with the output projection layer (often called \[weight tying\](https://mbrenndoerfer.com/writing/weight-tying-shared-embeddings-transformers)). This reduces the number of parameters by roughly 50 million for typical vocabulary sizes and often leads to better generalization and more stable training. ### The Lightning training module PyTorch Lightning handles the training loop boilerplate. We wrap our model in a \`LightningModule\` that defines how to train it: \`\`\` class GPTPreTrainingModule(L.LightningModule): """PyTorch Lightning module for GPT pre-training.""" def \_\_init\_\_( self, vocab\_size: int = 50257, n\_positions: int = 2048, n\_embd: int = 2048, n\_layer: int = 24, n\_head: int = 16, learning\_rate: float = 6e-4, weight\_decay: float = 0.1, warmup\_steps: int = 2000, max\_steps: int = 100000, ): super().\_\_init\_\_() self.save\_hyperparameters() config = GPTConfig( vocab\_size=vocab\_size, n\_positions=n\_positions, n\_embd=n\_embd, n\_layer=n\_layer, n\_head=n\_head, ) self.model = GPTModel(config) def forward(self, input\_ids, attention\_mask=None): return self.model(input\_ids, attention\_mask) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* The \`save\_hyperparameters()\` call is important because it stores all constructor arguments in the checkpoint. This allows the model to be reloaded later without having to manually reconstruct the original configuration. The training and validation steps implement standard causal language modeling, where the model is trained to predict the next token given all previous tokens in the sequence. \`\`\` def training\_step(self, batch, \_batch\_idx): # Convert int32 to int64 (long) - MDS stores as int32 but PyTorch expects long input\_ids = batch\["input\_ids"\].long() labels = batch\["labels"\].long() # Get attention mask if present (optional, for padded sequences) # attention\_mask: 1 = real token, 0 = padding # Note: Current data pipeline creates fixed-length sequences without padding, # so attention\_mask is not present. If using padded sequences, ensure: # - Padded positions in labels are set to -100 (ignored by cross\_entropy) # - attention\_mask marks real tokens (1) vs padding (0) attention\_mask = batch.get("attention\_mask", None) # Forward pass (causal mask is created internally in GPTModel) logits = self(input\_ids, attention\_mask=attention\_mask) # Shift logits and labels for causal language modeling # Before shift: labels\[i\] = input\_ids\[i\] # After shift: predict input\_ids\[i+1\] from input\_ids\[:i+1\] shift\_logits = logits\[..., :-1, :\].contiguous() shift\_labels = labels\[..., 1:\].contiguous() # Calculate loss loss = nn.functional.cross\_entropy( shift\_logits.view(-1, shift\_logits.size(-1)), shift\_labels.view(-1), ignore\_index=-100, ) # Log loss self.log( "train/loss", loss, on\_step=True, on\_epoch=True, prog\_bar=True, sync\_dist=True, ) # Calculate and log perplexity only on epoch (exp is costly, less frequent is fine) perplexity = torch.exp(torch.clamp(loss, max=20.0)) self.log( "train/perplexity", perplexity, on\_step=False, on\_epoch=True, prog\_bar=True, sync\_dist=True, ) return loss def validation\_step(self, batch, \_batch\_idx): # Convert int32 to int64 (long) - MDS stores as int32 but PyTorch expects long input\_ids = batch\["input\_ids"\].long() labels = batch\["labels"\].long() # Get attention mask if present (optional, for padded sequences) attention\_mask = batch.get("attention\_mask", None) # Forward pass (causal mask is created internally in GPTModel) logits = self(input\_ids, attention\_mask=attention\_mask) # Shift logits and labels shift\_logits = logits\[..., :-1, :\].contiguous() shift\_labels = labels\[..., 1:\].contiguous() # Calculate loss loss = nn.functional.cross\_entropy( shift\_logits.view(-1, shift\_logits.size(-1)), shift\_labels.view(-1), ignore\_index=-100, ) # Log loss self.log("val/loss", loss, prog\_bar=True, sync\_dist=True) # Calculate and log perplexity (exp is costly, but validation is infrequent so OK) perplexity = torch.exp(torch.clamp(loss, max=20.0)) self.log("val/perplexity", perplexity, prog\_bar=True, sync\_dist=True) return loss \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* The model performs a forward pass with a causal (autoregressive) mask created internally, ensuring each token can only attend to earlier positions. To align predictions with targets, the logits and labels are shifted so that the representation at position \`i\` is used to predict token \`i + 1\`. Loss is computed using cross-entropy over the shifted logits and labels. Training loss and perplexity are logged during execution, with metrics synchronized across distributed workers. The optimizer setup is where a lot of training stability comes from: \`\`\` def configure\_optimizers(self): # Separate parameters into weight decay and no weight decay groups decay\_params = \[\] no\_decay\_params = \[\] for param in self.model.parameters(): if param.requires\_grad: # 1D parameters (biases, LayerNorm) don't get weight decay # 2D+ parameters (weight matrices) get weight decay if param.ndim == 1: no\_decay\_params.append(param) else: decay\_params.append(param) optimizer\_grouped\_parameters = \[\ {"params": decay\_params, "weight\_decay": self.hparams.weight\_decay},\ {"params": no\_decay\_params, "weight\_decay": 0.0},\ \] # AdamW optimizer optimizer = torch.optim.AdamW( optimizer\_grouped\_parameters, lr=self.hparams.learning\_rate, betas=(0.9, 0.95), eps=1e-8, ) # Learning rate scheduler: warmup + cosine decay # Warmup: linear increase from 0 to 1.0 over warmup\_steps # Decay: cosine decay from 1.0 to 0.0 over remaining steps def lr\_lambda(current\_step): if current\_step < self.hparams.warmup\_steps: # Linear warmup return float(current\_step) / float(max(1, self.hparams.warmup\_steps)) # Cosine decay after warmup progress = (current\_step - self.hparams.warmup\_steps) / max( 1, self.hparams.max\_steps - self.hparams.warmup\_steps ) # Cosine annealing from 1.0 to 0.0 (returns float, not tensor) return 0.5 \* (1.0 + math.cos(progress \* math.pi)) scheduler = torch.optim.lr\_scheduler.LambdaLR(optimizer, lr\_lambda) return { "optimizer": optimizer, "lr\_scheduler": { "scheduler": scheduler, "interval": "step", }, } \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* Two important choices here: 1. \*\*Separate weight decay groups\*\*: We only apply weight decay to the weight matrices, not to biases or LayerNorm parameters. This follows the original BERT paper and is now standard practice, as regularizing biases and normalization parameters does not improve performance and can be harmful. 2. \*\*Cosine learning rate schedule with warmup\*\*: We start with a low learning rate, ramp up linearly during warmup (helps stabilize early training when gradients are noisy), then decay following a cosine curve. This schedule outperforms constant or step decay for transformer training. ### Checkpointing for fault tolerance Training a 30B-parameter model for 15,000 steps can take days. Hardware failures and spot instance preemptions are inevitable, which makes checkpointing essential. \`\`\` class S3CheckpointCallback(L.Callback): """ Periodically upload checkpoints to S3 for durability and resumption. This ensures checkpoints are safely stored in remote storage even if the training job is interrupted or the instance fails. """ def \_\_init\_\_(self, checkpoint\_dir: Path, upload\_every\_n\_steps: int): super().\_\_init\_\_() self.checkpoint\_dir = checkpoint\_dir self.upload\_every\_n\_steps = upload\_every\_n\_steps self.last\_uploaded\_step = -1 def on\_train\_batch\_end(self, trainer, pl\_module, outputs, batch, batch\_idx): """Upload checkpoint to S3 every N steps.""" if trainer.global\_rank != 0: return # Only upload from rank 0 current\_step = trainer.global\_step # Upload every N steps (aligns with ModelCheckpoint's every\_n\_train\_steps) if ( current\_step % self.upload\_every\_n\_steps == 0 and current\_step > self.last\_uploaded\_step and current\_step > 0 ): try: # Find the most recent checkpoint file checkpoint\_files = list(self.checkpoint\_dir.glob("\*.ckpt")) if not checkpoint\_files: print("No checkpoint files found to upload") return # Get the latest checkpoint (by modification time) latest\_checkpoint = max( checkpoint\_files, key=lambda p: p.stat().st\_mtime ) # Upload the checkpoint file directly to S3 using File.from\_local\_sync checkpoint\_file = File.from\_local\_sync(str(latest\_checkpoint)) print(f"Checkpoint uploaded to S3 at: {checkpoint\_file.path}") self.last\_uploaded\_step = current\_step except Exception as e: print(f"Warning: Failed to upload checkpoint to S3: {e}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* This callback runs every \`N\` training steps and uploads the checkpoint to durable storage. The key line is \`File.from\_local\_sync()\` which is a Flyte abstraction for uploading files. There are no blob store credentials to manage and no bucket paths to hardcode. Flyte automatically uses the storage backend configured for your cluster. The callback only runs on rank 0. In distributed training, all 8 GPUs have identical model states (that's the point of data parallelism). Having all of them upload the same checkpoint would be wasteful and could cause race conditions. When you restart a failed run, pass the checkpoint via \`resume\_checkpoint\` so training resumes exactly where it left off, including the same step count, optimizer state, and learning rate schedule position. ### Real-time metrics with Flyte Reports Multi-day training runs need observability. Is the loss decreasing? Did training diverge? Is the learning rate schedule behaving correctly? Flyte Reports let you build live dashboards directly in the UI: \`\`\` class FlyteReportingCallback(L.Callback): """Custom Lightning callback to report training metrics to Flyte Report.""" def \_\_init\_\_(self, report\_every\_n\_steps: int = 100): super().\_\_init\_\_() self.report\_every\_n\_steps = report\_every\_n\_steps self.metrics\_history = { "step": \[\], "train\_loss": \[\], "learning\_rate": \[\], "val\_loss": \[\], "val\_perplexity": \[\], } self.initialized\_report = False self.last\_logged\_step = -1 def on\_train\_start(self, trainer, pl\_module): """Initialize the live dashboard on training start.""" if trainer.global\_rank == 0 and not self.initialized\_report: self.\_initialize\_report() self.initialized\_report = True \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* The \`\_initialize\_report\` method (see complete code) creates an HTML/JavaScript dashboard with interactive charts. The callback then calls \`flyte.report.log()\` every \`N\` steps to push new metrics. The charts update in real-time so you can watch your loss curve descend while training runs. There is no need to deploy Grafana, configure Prometheus, or keep a TensorBoard server running. Using \`flyte.report.log()\` is sufficient to get live training metrics directly in the Flyte UI. !\[Metrics viz\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/tutorials/distributed-llm-pretraining/metrics.png) ### Streaming data at scale Training datasets are massive. SlimPajama contains 627 billion tokens and spans hundreds of gigabytes even when compressed. Downloading the entire dataset to each training node before starting would take hours and waste storage. Instead, we convert the data to MDS (MosaicML Data Shard) format and stream it during training: \`\`\` @data\_loading\_env.task async def load\_and\_prepare\_streaming\_dataset( dataset\_name: str, dataset\_config: Optional\[str\], max\_length: int, train\_split: str, val\_split: Optional\[str\], max\_train\_samples: Optional\[int\], max\_val\_samples: Optional\[int\], shard\_size\_mb: int, ) -> Dir: """Tokenize dataset and convert to MDS format for streaming.""" from datasets import load\_dataset from streaming import MDSWriter from transformers import GPT2TokenizerFast output\_dir = Path("/tmp/streaming\_dataset") output\_dir.mkdir(parents=True, exist\_ok=True) tokenizer = GPT2TokenizerFast.from\_pretrained("gpt2") tokenizer.pad\_token = tokenizer.eos\_token # MDS schema: what each sample contains columns = { "input\_ids": "ndarray:int32", "labels": "ndarray:int32", } \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* This task does three things: 1. \*\*Tokenizes the text\*\* using GPT-2's BPE tokenizer 2. \*\*Concatenates documents\*\* into fixed-length sequences (no padding waste) 3. \*\*Writes shards\*\* to storage in a format optimized for streaming The task returns a Flyte \`Dir\` object, which is a reference to the output location. It's not the data itself, just a pointer. When the training task receives this \`Dir\`, it streams shards on-demand rather than downloading everything upfront. Flyte caches this task automatically. Run the pipeline twice with the same dataset config, and Flyte skips tokenization entirely on the second run. Change the dataset or sequence length, and it re-runs. ### Distributed training with FSDP Now we get to the core: actually training the model across multiple GPUs. FSDP is what makes this possible for large models. \`\`\` @distributed\_llm\_training\_env.task(report=True) def train\_distributed\_llm( prepared\_dataset: Dir, resume\_checkpoint: Optional\[Dir\], vocab\_size: int, n\_positions: int, n\_embd: int, n\_layer: int, n\_head: int, batch\_size: int, num\_workers: int, max\_steps: int, learning\_rate: float, weight\_decay: float, warmup\_steps: int, use\_fsdp: bool, checkpoint\_upload\_steps: int, checkpoint\_every\_n\_steps: int, report\_every\_n\_steps: int, val\_check\_interval: int, grad\_accumulation\_steps: int = 1, ) -> Optional\[Dir\]: # ... setup code ... \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* Notice \`report=True\` on the task decorator. It enables Flyte Reports for this specific task. The training task receives the prepared dataset as a \`Dir\` and streams data directly from storage: \`\`\` # StreamingDataset streams shards from the remote Flyte storage on-demand # It automatically detects torch.distributed context # and shards data across GPUs - each rank gets different data automatically train\_dataset = StreamingDataset( remote=f"{remote\_path}/train", # Remote MDS shard location local=str(local\_cache / "train"), # Local cache for downloaded shards shuffle=True, # Shuffle samples shuffle\_algo="naive", # Shuffling algorithm batch\_size=batch\_size, # Used for shuffle buffer sizing ) # Create validation StreamingDataset if it exists val\_dataset = None try: val\_dataset = StreamingDataset( remote=f"{remote\_path}/validation", local=str(local\_cache / "validation"), shuffle=False, # No shuffling for validation batch\_size=batch\_size, ) print( f"Validation dataset initialized with streaming from: {remote\_path}/validation" ) except Exception as e: print(f"No validation dataset found: {e}") # Create data loaders # StreamingDataset handles distributed sampling internally by detecting # torch.distributed.get\_rank() and torch.distributed.get\_world\_size() train\_loader = DataLoader( train\_dataset, batch\_size=batch\_size, num\_workers=num\_workers, pin\_memory=True, persistent\_workers=True, drop\_last=True, # Drop incomplete batches for distributed training collate\_fn=mds\_collate\_fn, # Handle read-only arrays ) # Create validation loader if validation dataset exists val\_loader = None if val\_dataset is not None: val\_loader = DataLoader( val\_dataset, batch\_size=batch\_size, num\_workers=num\_workers, pin\_memory=True, persistent\_workers=True, drop\_last=False, collate\_fn=mds\_collate\_fn, ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* \`prepared\_dataset.path\` provides the remote storage path for the dataset. MosaicML's \`StreamingDataset\` automatically shards data across GPUs so that each rank sees different samples, without requiring a manual distributed sampler. The credentials are already in the environment because Flyte set them up. FSDP is where the memory magic happens. Instead of each GPU holding a full copy of the model (like Distributed Data Parallel (DDP)), FSDP shards the parameters, gradients, and optimizer states across all GPUs. Each GPU only holds 1/8th of the model. When a layer needs to run, FSDP all-gathers the full parameters, runs the computation, then discards them. \`\`\` # Configure distributed strategy if use\_fsdp: from torch.distributed.fsdp.wrap import ModuleWrapPolicy strategy = FSDPStrategy( auto\_wrap\_policy=ModuleWrapPolicy(\[GPTBlock\]), activation\_checkpointing\_policy=None, cpu\_offload=False, # H200 has 141GB - no CPU offload needed state\_dict\_type="full", sharding\_strategy="FULL\_SHARD", process\_group\_backend="nccl", ) else: from lightning.pytorch.strategies import DDPStrategy strategy = DDPStrategy(process\_group\_backend="nccl") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* We wrap at the \`GPTBlock\` level because each transformer block becomes an FSDP unit. This balances communication overhead (more units = more all-gathers) against memory savings (smaller units = more granular sharding). One subtle detail: gradient clipping. With FSDP, gradients are sharded across ranks, so computing a global gradient norm would require an expensive all-reduce operation. Instead of norm-based clipping, we use value-based gradient clipping, which clamps each individual gradient element to a fixed range. This can be done independently on each rank with no coordination overhead and is commonly used for large-scale FSDP training. \`\`\` # Initialize trainer trainer = L.Trainer( strategy=strategy, accelerator="gpu", devices=DEVICES\_PER\_NODE, num\_nodes=NUM\_NODES, # Training configuration max\_steps=max\_steps, precision="bf16-mixed", # BFloat16 for better numerical stability # Optimization gradient\_clip\_val=1.0, gradient\_clip\_algorithm=( "value" if use\_fsdp else "norm" ), # FSDP requires 'value', DDP can use 'norm' accumulate\_grad\_batches=grad\_accumulation\_steps, # Logging and checkpointing callbacks=callbacks, log\_every\_n\_steps=report\_every\_n\_steps, val\_check\_interval=val\_check\_interval, # Performance benchmark=True, deterministic=False, # Enable gradient checkpointing for memory efficiency enable\_checkpointing=True, use\_distributed\_sampler=False, # StreamingDataset handles distributed sampling ) # Train the model (resume from checkpoint if provided) trainer.fit(model, train\_loader, val\_loader, ckpt\_path=ckpt\_path) # Print final results if trainer.global\_rank == 0: if val\_loader is not None: print( f"Final validation loss: {trainer.callback\_metrics.get('val/loss', 0.0):.4f}" ) print( f"Final validation perplexity: {trainer.callback\_metrics.get('val/perplexity', 0.0):.4f}" ) print(f"Checkpoints saved to: {checkpoint\_dir}") return Dir.from\_local\_sync(output\_dir) return None \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* The trainer configuration brings together all the pieces we've discussed: - \*\*\`precision="bf16-mixed"\`\*\*: BFloat16 mixed precision training. BF16 has the same dynamic range as FP32 (unlike FP16), so you don't need loss scaling. This is the standard choice for modern GPU training. - \*\*\`gradient\_clip\_val=1.0\`\*\*: Clips gradients to prevent exploding gradients during training. Combined with value-based clipping for FSDP compatibility. - \*\*\`accumulate\_grad\_batches\`\*\*: Accumulates gradients over multiple forward passes before updating weights. This lets us hit a larger effective batch size than what fits in GPU memory. - \*\*\`val\_check\_interval\`\*\*: How often to run validation. For long training runs, you don't want to validate every epoch — that would be too infrequent. Instead, validate every \`N\` training steps. - \*\*\`use\_distributed\_sampler=False\`\*\*: We disable Lightning's built-in distributed sampler because \`StreamingDataset\` handles data sharding internally. Using both would cause conflicts. - \*\*\`benchmark=True\`\*\*: Enables cuDNN autotuning. PyTorch will benchmark different convolution algorithms on the first batch and pick the fastest one for your specific input sizes. The trainer then calls \`fit()\` with the model, data loaders, and optionally a checkpoint path to resume from. ### Tying it together The pipeline task orchestrates everything: \`\`\` @driver\_env.task async def distributed\_llm\_pipeline( model\_size: str, dataset\_name: str = "Salesforce/wikitext", dataset\_config: str = "wikitext-103-raw-v1", max\_length: int = 2048, max\_train\_samples: Optional\[int\] = 10000, max\_val\_samples: Optional\[int\] = 1000, max\_steps: int = 100000, resume\_checkpoint: Optional\[Dir\] = None, checkpoint\_upload\_steps: int = 1000, # Optional overrides (if None, uses model preset defaults) batch\_size: Optional\[int\] = None, learning\_rate: Optional\[float\] = None, use\_fsdp: bool = True, ) -> Optional\[Dir\]: # Get model configuration model\_config = get\_model\_config(model\_size) # Use preset values if not overridden actual\_batch\_size = ( batch\_size if batch\_size is not None else model\_config\["batch\_size"\] ) actual\_learning\_rate = ( learning\_rate if learning\_rate is not None else model\_config\["learning\_rate"\] ) # Step 1: Load and prepare streaming dataset prepared\_dataset = await load\_and\_prepare\_streaming\_dataset( dataset\_name=dataset\_name, dataset\_config=dataset\_config, max\_length=max\_length, train\_split="train", val\_split="validation", max\_train\_samples=max\_train\_samples, max\_val\_samples=max\_val\_samples, shard\_size\_mb=64, # 64MB shards ) # Step 2: Run distributed training if resume\_checkpoint is not None: print("\\nStep 2: Resuming distributed training from checkpoint...") else: print("\\nStep 2: Starting distributed training from scratch...") target\_global\_batch = 256 world\_size = NUM\_NODES \* DEVICES\_PER\_NODE effective\_per\_step = world\_size \* actual\_batch\_size grad\_accumulation\_steps = max( 1, math.ceil(target\_global\_batch / max(1, effective\_per\_step)) ) checkpoint\_dir = train\_distributed\_llm( prepared\_dataset=prepared\_dataset, resume\_checkpoint=resume\_checkpoint, vocab\_size=VOCAB\_SIZE, n\_positions=N\_POSITIONS, n\_embd=model\_config\["n\_embd"\], n\_layer=model\_config\["n\_layer"\], n\_head=model\_config\["n\_head"\], batch\_size=actual\_batch\_size, num\_workers=8, max\_steps=max\_steps, learning\_rate=actual\_learning\_rate, weight\_decay=0.1, warmup\_steps=500, use\_fsdp=use\_fsdp, checkpoint\_upload\_steps=checkpoint\_upload\_steps, checkpoint\_every\_n\_steps=model\_config\["checkpoint\_every\_n\_steps"\], report\_every\_n\_steps=model\_config\["report\_every\_n\_steps"\], val\_check\_interval=model\_config\["val\_check\_interval"\], grad\_accumulation\_steps=grad\_accumulation\_steps, ) return checkpoint\_dir \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* The flow is straightforward: load the configuration, prepare the data, and run training. Flyte automatically manages the execution graph so data preparation runs first and training waits until it completes. If data preparation is cached from a previous run, training starts immediately. The gradient accumulation calculation is worth noting. We want a global batch size of 256 (this affects training dynamics), but each GPU can only fit a small batch. With 8 GPUs and batch size 1 each, we need 32 accumulation steps to hit 256. ## Running the pipeline With everything defined, running is simple: \`\`\` if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run( distributed\_llm\_pipeline, model\_size="30B", dataset\_name="cerebras/SlimPajama-627B", dataset\_config=None, max\_length=2048, max\_train\_samples=5\_000\_000, max\_val\_samples=50\_000, max\_steps=15\_000, use\_fsdp=True, checkpoint\_upload\_steps=1000, ) print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/pretraining/train.py\* This configuration is designed for testing and demonstration. Notice \`max\_train\_samples=5\_000\_000\` — that's 5 million samples from a dataset with 627 billion tokens. A tiny fraction, enough to verify everything works without burning through compute. For a real pretraining run, you would remove this limit by setting \`max\_train\_samples=None\`, or increase it significantly. You would also increase \`max\_steps\` to match your compute budget, likely scale to multiple nodes by setting \`NUM\_NODES=4\` or higher, and allocate more resources. The rest of the pipeline remains unchanged. \`\`\`bash flyte create config --endpoint --project --domain --builder remote uv run train.py \`\`\` When you run this, Flyte: 1. \*\*Builds the container\*\* (cached after first run) 2. \*\*Schedules data prep\*\* on CPU nodes 3. \*\*Waits for data prep\*\* (or skips if cached) 4. \*\*Provisions H200 nodes\*\* and launches distributed training 5. \*\*Streams logs and metrics\*\* to the UI in real-time Open the Flyte UI to observe the workflow execution. The data preparation task completes first, followed by the training task spinning up. As training begins, the Flyte Reports dashboard starts plotting loss curves. If anything goes wrong, the logs are immediately available in the UI. !\[Training Log\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/tutorials/distributed-llm-pretraining/logs.png) If training fails due to an out-of-memory error, a GPU driver error, or a hardware issue, check the logs, fix the problem, and restart the run with \`resume\_checkpoint\` pointing to the most recent checkpoint. Training resumes from where it left off. Flyte tracks the full execution history, so it is easy to see exactly what happened. ## Going further If you've run through this tutorial, here's where to go next depending on what you're trying to do: \*\*You want to train on your own data.\*\* The data prep task accepts any HuggingFace dataset with a \`text\` column. If your data isn't on HuggingFace, you can modify \`load\_and\_prepare\_streaming\_dataset\` to read from S3, local files, or any other source. The key is getting your data into MDS format. Once it's there, the streaming and sharding just works. For production training, look at SlimPajama, \[RedPajama\](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T), or \[The Pile\](https://huggingface.co/datasets/EleutherAI/pile) as starting points. \*\*You want to scale to more GPUs.\*\* Bump \`NUM\_NODES\` and Flyte handles the rest. The main thing to watch is the effective batch size. As you add more GPUs, you may want to reduce gradient accumulation steps to keep the same global batch size, or increase them if you want to experiment with larger batches. \*\*Your training keeps failing.\*\* Add \`retries=3\` to your task decorator for automatic retry on transient failures. This handles spot instance preemption, temporary network issues, and the occasional GPU that decides to stop working. Combined with checkpointing, you get fault-tolerant training that can survive most infrastructure hiccups. For persistent failures, the Flyte UI logs are your friend as they capture stdout/stderr from all ranks. \*\*You want better visibility into what's happening.\*\* We're actively working on surfacing GPU driver logs (xid/sxid errors), memory utilization breakdowns, and NCCL communication metrics directly in the Flyte UI. If you're hitting issues that the current logs don't explain, reach out. Your feedback helps us prioritize what observability features to build next! === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/computer-vision === # Computer Vision Tutorials for image and vision-language model workloads. ### \*\*Computer Vision > Fine-tuning a VLM\*\* Adapt Qwen2.5-VL to occluded image classification by training a 10K-parameter adapter with multi-node DeepSpeed, automatic recovery, and live training dashboards. ### \*\*Computer Vision > Multimodal retrieval evaluation\*\* Benchmark ColPali, SigLIP, and OCR+BM25 visual document retrieval on ViDoRe with warm GPU containers, dynamic batching, and an interactive report. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/computer-vision/qwen-vl-finetuning === # Fine-tuning a VLM Large vision-language models like Qwen2.5-VL are remarkably capable out of the box. But adapting one to a specialized task raises an immediate question: do you really need to update 3 billion parameters? Usually, no. The \*\*frozen backbone pattern\*\* is a practical alternative: keep all pretrained weights frozen and train only a small, task-specific adapter inserted before the vision encoder. The adapter learns to transform its input in a way that makes the frozen model perform well on your task without touching the underlying billions of parameters. The result is faster training, lower memory pressure, and a much smaller set of weights to store and version. This tutorial makes that pattern concrete. We take a partially-occluded image classification task — CIFAR-10 images with random black rectangles covering 22–45% of the frame — and train a tiny Conv2d adapter to "see through" the occlusion before the frozen VLM processes it. The adapter has approximately \*\*10,500 trainable parameters\*\*. The backbone has 3 billion. The machine learning is interesting, but the real focus here is on shipping a production-grade training pipeline: - \*\*Multi-node distributed training\*\* across 2 nodes × 4 GPUs using PyTorch Elastic and DeepSpeed Stage 2 - \*\*Automatic fault tolerance\*\*: checkpoints upload to object storage after every validation epoch; if training fails, the pipeline returns the last known-good checkpoint instead of crashing - \*\*Live observability\*\*: a streaming HTML dashboard in the Flyte UI updates in real-time as training runs, no separate monitoring infrastructure required - \*\*Cached data preparation\*\*: dataset processing runs once and is reused across all reruns - \*\*Clean task isolation\*\*: each stage runs with exactly the resources it needs, nothing more > \[!NOTE\] > Full code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning). ## Overview The pipeline has four tasks with clearly defined responsibilities: 1. \*\*Dataset preparation\*\* (\`prepare\_occlusion\_dataset\`): Downloads CIFAR-10, applies random occlusions, and writes image manifests as streaming JSONL files to object storage. Runs on CPU and is cached, so it only runs once regardless of how many times you rerun the pipeline with the same config. 2. \*\*Multi-node training\*\* (\`train\_qwen\_adapter\_multinode\`): Runs PyTorch Lightning with DeepSpeed Stage 2 across 2 nodes × 4 L40s GPUs. Only the adapter trains; the 3B backbone stays frozen. 3. \*\*Evaluation\*\* (\`evaluate\_qwen\_adapter\`): Loads the saved adapter, runs inference on validation examples, and produces a predictions report. Runs on a single GPU. 4. \*\*Driver\*\* (\`qwen\_vl\_multinode\_deepspeed\`): The pipeline entry point. Orchestrates the three tasks above, manages WandB initialization, handles recovery from training failures, and produces a final HTML report in the Flyte UI. Why this separation? It mirrors how production pipelines should be structured. Data prep is cheap and deterministic so we cache it. Training is expensive and failure-prone so we isolate it with fault tolerance. Evaluation needs different hardware than training. The driver is pure coordination, so it gets minimal resources. ## Implementation ### Setting up the environment Different tasks need different compute. Flyte's \`TaskEnvironment\` is how you declare exactly what each task needs. First, define the container images. Training needs a full CUDA stack with ML libraries, driver compatibility, and DeepSpeed's build tools: \`\`\` gpu\_image = ( flyte.Image.from\_base("nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04") .clone(name="qwen\_vl\_multinode\_deepspeed", python\_version=(3, 13), extendable=True) .with\_apt\_packages("build-essential") .with\_pip\_packages( "torch==2.9.1", "torchvision==0.24.1", "lightning==2.6.1", "transformers==4.57.3", "deepspeed==0.18.8", "datasets==4.4.1", "pillow==11.3.0", "flyteplugins-pytorch>=2.0.11", "flyteplugins-jsonl>=2.0.11", "flyteplugins-wandb>=2.0.11", ) ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/config.py\* \`from\_base\` starts from the official NVIDIA CUDA image, giving you NCCL, cuDNN, and the right driver headers out of the box. \`with\_apt\_packages("build-essential")\` is required because DeepSpeed compiles CUDA kernels at first use and without build tools, it silently falls back to slower CPU implementations. The non-GPU image for data preparation and orchestration is much lighter: \`\`\` non\_gpu\_image = flyte.Image.from\_debian\_base( name="qwen\_vl\_multinode\_deepspeed\_non\_gpu" ).with\_pip\_packages( "flyteplugins-pytorch>=2.0.11", "flyteplugins-jsonl>=2.0.11", "flyteplugins-wandb>=2.0.11", "lightning==2.6.1", "datasets==4.4.1", "pillow==11.3.0", "torchvision==0.24.1", ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/config.py\* With images defined, each task gets its own resource declaration: \`\`\` dataset\_env = flyte.TaskEnvironment( name="qwen\_vl\_dataset\_prep", image=non\_gpu\_image, resources=flyte.Resources(cpu=5, memory="15Gi"), cache="auto", ) training\_env = flyte.TaskEnvironment( name="qwen\_vl\_multinode\_training", image=gpu\_image, resources=flyte.Resources( cpu=42, memory="256Gi", gpu=f"L40s:{DEVICES\_PER\_NODE}", shm="16Gi", ), plugin\_config=Elastic(nnodes=NUM\_NODES, nproc\_per\_node=DEVICES\_PER\_NODE), secrets=\[\ flyte.Secret(key="wandb\_api\_key", as\_env\_var="WANDB\_API\_KEY")\ \], # TODO: update with your own secret key env\_vars={ "TORCH\_DISTRIBUTED\_DEBUG": "INFO", "NCCL\_DEBUG": "WARN", "TOKENIZERS\_PARALLELISM": "false", "CUDA\_HOME": "/usr/local/cuda", "DS\_SKIP\_CUDA\_CHECK": "1", }, ) evaluation\_env = flyte.TaskEnvironment( name="qwen\_vl\_adapter\_eval", image=gpu\_image, resources=flyte.Resources(cpu=16, memory="64Gi", gpu="L40s:1"), cache="auto", ) driver\_env = flyte.TaskEnvironment( name="qwen\_vl\_multinode\_driver", image=non\_gpu\_image, resources=flyte.Resources(cpu=2, memory="4Gi"), depends\_on=\[dataset\_env, training\_env, evaluation\_env\], ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/config.py\* A few things worth noting here: - \*\*\`Elastic(nnodes=2, nproc\_per\_node=4)\`\*\*: Flyte's integration with PyTorch's elastic launch. It handles process spawning (one process per GPU), rank assignment, and distributed environment setup — master address, world size, rendezvous — without any shell scripting or manual \`torchrun\` invocations. - \*\*\`shm="16Gi"\`\*\*: Shared memory is required for NCCL inter-GPU communication on the same node. Without it, you'll see cryptic errors from the communication library when training starts. - \*\*\`cache="auto"\`\*\*: The dataset preparation task is cached by input hash. Running the pipeline twice with the same hyperparameters skips it entirely on the second run. - \*\*\`depends\_on\`\*\*: The driver task declares that each worker image must finish building before it starts, ensuring containers are ready before the driver begins orchestrating. - \*\*\`secrets\`\*\*: The WandB API key is injected from Flyte's secret store as an environment variable. No credentials in code. All training hyperparameters flow through a single typed dataclass: \`\`\` @dataclass class Config: model\_name: str = DEFAULT\_MODEL\_NAME image\_size: int = IMAGE\_SIZE max\_train\_samples: int = 1024 max\_val\_samples: int = 256 epochs: int = 8 per\_device\_batch\_size: int = 1 target\_global\_batch\_size: int = 16 learning\_rate: float = 2e-4 weight\_decay: float = 1e-2 reconstruction\_loss\_weight: float = 0.35 report\_every\_n\_steps: int = 10 num\_workers: int = 4 max\_length: int = 512 eval\_examples: int = 16 train\_occlusion\_min: float = 0.22 train\_occlusion\_max: float = 0.42 eval\_occlusion\_min: float = 0.28 eval\_occlusion\_max: float = 0.45 seed: int = 7 def to\_dict(self) -> dict: return asdict(self) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/config.py\* Using a dataclass rather than scattered constants or argparse arguments means the full config is serializable, can be stored in artifact metadata alongside the model checkpoint, and flows cleanly as a typed input between tasks. The \`to\_dict()\` method serializes it for WandB logging. ### Preparing the dataset The dataset task handles everything: downloading CIFAR-10, generating occlusions, and writing the manifests. \`\`\` @dataset\_env.task async def prepare\_occlusion\_dataset(config: Config) -> DatasetArtifacts: from PIL import Image from torchvision.datasets import CIFAR10 from flyte.io import Dir from flyteplugins.jsonl import JsonlFile import random rng = random.Random(config.seed) images\_dir = Path("/tmp/qwen\_vl\_occlusion\_images") train\_images\_dir = images\_dir / "train" / "images" val\_images\_dir = images\_dir / "validation" / "images" train\_images\_dir.mkdir(parents=True, exist\_ok=True) val\_images\_dir.mkdir(parents=True, exist\_ok=True) prompt = ( "The image may be partially occluded. " "Answer with exactly one CIFAR-10 class label: " + ", ".join(CLASS\_NAMES) + ". What is the main object?" ) async def export\_split( dataset, split\_name: str, limit: int, local\_image\_dir: Path, occ\_min: float, occ\_max: float, ): out = JsonlFile.new\_remote(f"{split\_name}\_manifest.jsonl") async with out.writer() as writer: for idx in range(limit): pil\_image, label\_idx = dataset\[idx\] resized = pil\_image.resize( (config.image\_size, config.image\_size), resample=Image.Resampling.BICUBIC, ) rel\_path = f"{split\_name}/images/{split\_name}-{idx:05d}.png" resized.save(local\_image\_dir / f"{split\_name}-{idx:05d}.png") occlusion = build\_occlusion\_box( width=config.image\_size, height=config.image\_size, rng=rng, min\_fraction=occ\_min, max\_fraction=occ\_max, ) await writer.write( { "image\_path": rel\_path, "label": CLASS\_NAMES\[label\_idx\], "label\_index": int(label\_idx), "prompt": prompt, "occlusion": occlusion, } ) return out train\_dataset = CIFAR10(root="/tmp/cifar10", train=True, download=True) val\_dataset = CIFAR10(root="/tmp/cifar10", train=False, download=True) train\_manifest = await export\_split( train\_dataset, "train", config.max\_train\_samples, train\_images\_dir, config.train\_occlusion\_min, config.train\_occlusion\_max, ) val\_manifest = await export\_split( val\_dataset, "validation", config.max\_val\_samples, val\_images\_dir, config.eval\_occlusion\_min, config.eval\_occlusion\_max, ) return DatasetArtifacts( train\_manifest=train\_manifest, val\_manifest=val\_manifest, images=await Dir.from\_local(str(images\_dir)), ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/data.py\* Each image gets a randomly-placed black rectangle. The occlusion covers 22–42% of the image area during training and 28–45% during evaluation. The occlusion is deliberately harder at eval time to test how robust the adapter is. The bounding box coordinates are written into each manifest record alongside the image path and ground-truth label, so the training task can reconstruct the binary occlusion mask as the adapter's fourth input channel. Two Flyte primitives handle data persistence without any manual storage management: - \*\*\`JsonlFile.new\_remote()\`\*\* opens a streaming writer that writes directly to remote object storage. The training task reads records back via \`jf.iter\_records\_sync()\`, so no local file paths and S3 credentials to manage. - \*\*\`Dir.from\_local()\`\*\* uploads the local images directory to object storage and returns a typed handle. The training task downloads it to a local path via \`Dir.download\_sync()\`. Because \`cache="auto"\` is set on this task, dataset preparation runs once. Subsequent reruns with the same config skip it entirely. ### The adapter Here's the entire trainable component of the model with \`~10,500\` parameters: \`\`\` class ResidualOcclusionAdapter(nn.Module): def \_\_init\_\_(self, hidden\_channels: int = 32): super().\_\_init\_\_() self.net = nn.Sequential( nn.Conv2d(4, hidden\_channels, kernel\_size=3, padding=1), nn.GELU(), nn.Conv2d(hidden\_channels, hidden\_channels, kernel\_size=3, padding=1), nn.GELU(), nn.Conv2d(hidden\_channels, 3, kernel\_size=1), nn.Tanh(), ) self.gate = nn.Parameter(torch.tensor(0.10)) def forward( self, pixel\_values: torch.Tensor, occlusion\_mask: torch.Tensor ) -> torch.Tensor: if pixel\_values.ndim != 4: raise ValueError( "ResidualOcclusionAdapter expects dense image tensors with shape " f"(B, C, H, W), but received {tuple(pixel\_values.shape)}." ) if occlusion\_mask.ndim == 3: occlusion\_mask = occlusion\_mask.unsqueeze(1) adapter\_input = torch.cat( \[pixel\_values, occlusion\_mask.to(pixel\_values.dtype)\], dim=1, ) residual = self.net(adapter\_input) return pixel\_values + torch.tanh(self.gate) \* residual \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/model.py\* The adapter takes the occluded image (3 channels) concatenated with the binary occlusion mask (1 channel) as a 4-channel input. It predicts a residual correction through a small convolutional network, then adds that correction back to the original pixels. The learnable \`gate\` scalar, initialized to \`0.10\`, controls how strongly the adapter modifies the image. It starts as a near-identity transformation and gradually grows during training as the adapter gains confidence. The adapter is plugged into Qwen2.5-VL via a Lightning module: \`\`\` class QwenVLAdapterModule(L.LightningModule): def \_\_init\_\_( self, model\_name: str, learning\_rate: float, weight\_decay: float, reconstruction\_loss\_weight: float, ): super().\_\_init\_\_() from transformers import Qwen2\_5\_VLForConditionalGeneration self.save\_hyperparameters() self.adapter = ResidualOcclusionAdapter() self.backbone = Qwen2\_5\_VLForConditionalGeneration.from\_pretrained( model\_name, torch\_dtype=torch.bfloat16, attn\_implementation="sdpa", ) self.backbone.requires\_grad\_(False) self.backbone.gradient\_checkpointing\_enable() # DeepSpeed checkpoints only persist the trainable adapter weights when # \`exclude\_frozen\_parameters=True\`. On resume we rebuild the frozen # backbone from Hugging Face and load the checkpoint non-strictly. self.strict\_loading = False self.total\_params, self.trainable\_params = count\_parameters(self) self.example\_input\_array = None self.vision\_patch\_size = int(self.backbone.config.vision\_config.patch\_size) self.temporal\_patch\_size = int( getattr(self.backbone.config.vision\_config, "temporal\_patch\_size", 1) ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/model.py\* The key line is \`self.backbone.requires\_grad\_(False)\`. This freezes all 3 billion backbone parameters which means only the adapter's ~10,500 weights receive gradients. \`gradient\_checkpointing\_enable()\` trades compute for memory: instead of keeping the frozen backbone's intermediate activations in GPU memory during the backward pass, they're recomputed on the fly. This is critical when a 3B model is sitting in GPU memory alongside your optimizer state. \`strict\_loading = False\` handles an important DeepSpeed checkpoint detail. When \`exclude\_frozen\_parameters=True\` is set on the strategy, DeepSpeed only saves the adapter weights in checkpoints, not the 3B frozen backbone. On resume, the checkpoint won't contain backbone weights, so loading must be non-strict. The \`on\_load\_checkpoint\` hook fills in the missing backbone weights from the freshly-loaded HuggingFace model, combining the best of both worlds: small checkpoints and a fully initialized model. The training loss combines two objectives: \`\`\` def \_forward\_losses( self, batch: dict\[str, torch.Tensor\] ) -> dict\[str, torch.Tensor\]: backbone\_dtype = next(self.backbone.parameters()).dtype if batch\["pixel\_values"\].ndim == 2: if "image\_grid\_thw" not in batch: raise ValueError( "Packed Qwen pixel values require \`image\_grid\_thw\` to reconstruct " "dense images for the Conv2d adapter." ) grid\_thw = batch\["image\_grid\_thw"\] dense\_pixels = packed\_pixels\_to\_dense\_images( batch\["pixel\_values"\].to(dtype=backbone\_dtype), grid\_thw, patch\_size=self.vision\_patch\_size, temporal\_patch\_size=self.temporal\_patch\_size, ) clean\_pixels = packed\_pixels\_to\_dense\_images( batch\["clean\_pixel\_values"\].to(dtype=backbone\_dtype), grid\_thw, patch\_size=self.vision\_patch\_size, temporal\_patch\_size=self.temporal\_patch\_size, ) adapted\_dense = self.adapter(dense\_pixels, batch\["occlusion\_mask"\]) adapted\_pixels = dense\_images\_to\_packed\_pixels( adapted\_dense, grid\_thw, patch\_size=self.vision\_patch\_size, temporal\_patch\_size=self.temporal\_patch\_size, ) else: clean\_pixels = batch\["clean\_pixel\_values"\].to(dtype=backbone\_dtype) adapted\_dense = self.adapter( batch\["pixel\_values"\].to(dtype=backbone\_dtype), batch\["occlusion\_mask"\], ) adapted\_pixels = adapted\_dense forward\_kwargs = { "input\_ids": batch\["input\_ids"\], "attention\_mask": batch\["attention\_mask"\], "pixel\_values": adapted\_pixels, "labels": batch\["labels"\], } if "image\_grid\_thw" in batch: forward\_kwargs\["image\_grid\_thw"\] = batch\["image\_grid\_thw"\] outputs = self.backbone(\*\*forward\_kwargs) clean\_pixels = clean\_pixels.to( device=adapted\_pixels.device, dtype=backbone\_dtype ) occlusion\_mask = batch\["occlusion\_mask"\].to( device=adapted\_pixels.device, dtype=backbone\_dtype, ) if occlusion\_mask.ndim == 3: occlusion\_mask = occlusion\_mask.unsqueeze(1) if occlusion\_mask.shape\[-2:\] != adapted\_dense.shape\[-2:\]: occlusion\_mask = F.interpolate( occlusion\_mask, size=adapted\_dense.shape\[-2:\], mode="nearest", ) reconstruction\_error = (adapted\_dense - clean\_pixels).abs() \* occlusion\_mask mask\_denominator = (occlusion\_mask.sum() \* adapted\_dense.shape\[1\]).clamp\_min( 1.0 ) reconstruction\_loss = reconstruction\_error.sum() / mask\_denominator total\_loss = ( outputs.loss + self.hparams.reconstruction\_loss\_weight \* reconstruction\_loss ) return { "total\_loss": total\_loss, "lm\_loss": outputs.loss, "reconstruction\_loss": reconstruction\_loss, } \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/model.py\* The \*\*language modeling loss\*\* (cross-entropy on the predicted class label tokens) drives the model to produce correct answers. The \*\*reconstruction loss\*\* (mean absolute error between the adapter's output and the clean image, computed only in the occluded region) pushes the adapter to actually restore the missing pixels rather than finding a representation shortcut. Without it, the adapter could overfit the frozen backbone's quirks and produce correct tokens while generating noise in the masked region. The \`reconstruction\_loss\_weight\` (default \`0.35\`) balances these two objectives. Because Qwen2.5-VL's preprocessor packs image patches into a flat \`(num\_patches, patch\_dim)\` tensor, the adapter must unpack this into a spatial \`(B, C, H, W)\` tensor, apply the convolutions, then repack. The \`packed\_pixels\_to\_dense\_images\` and \`dense\_images\_to\_packed\_pixels\` utilities in \`model.py\` handle this format conversion transparently. ### Multi-node training with DeepSpeed The training task is a standard PyTorch Lightning training loop with distributed infrastructure handled by Flyte and DeepSpeed: \`\`\` @wandb\_init @training\_env.task(report=True) def train\_qwen\_adapter\_multinode( train\_manifest: JsonlFile, val\_manifest: JsonlFile, images\_dir: Dir, config: Config, resume\_from: Optional\[Dir\] = None, recovery\_uri: Optional\[str\] = None, ) -> Optional\[Dir\]: \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/tasks.py\* The \`@wandb\_init\` decorator integrates with the \`wandb\_config\` context created in the driver task. It retrieves the initialized WandB run and attaches a \`WandbLogger\` to the trainer. The \`report=True\` flag on the task decorator enables Flyte Reports for live dashboard streaming from this task. !\[Live Training\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/tutorials/qwen-vl-finetuning/live\_training\_graph.png) !\[Live Training Contd\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/tutorials/qwen-vl-finetuning/losses.png) DeepSpeed Stage 2 shards optimizer states and gradients across GPUs, reducing per-GPU memory usage significantly. The critical configuration flag here is \`exclude\_frozen\_parameters=True\`: \`\`\` strategy = DeepSpeedStrategy( stage=2, offload\_optimizer=False, offload\_parameters=False, process\_group\_backend="nccl", exclude\_frozen\_parameters=True, ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/tasks.py\* Without \`exclude\_frozen\_parameters=True\`, DeepSpeed would shard and checkpoint the frozen backbone weights too, producing enormous checkpoint files, slow checkpoint saves, and unnecessary communication overhead. With it, only the adapter participates in sharding and checkpointing. The backbone is loaded independently on each worker from HuggingFace. Gradient accumulation is computed automatically to hit the target global batch size regardless of how many GPUs are actually running: \`\`\` world\_size = NUM\_NODES \* DEVICES\_PER\_NODE per\_step\_batch = world\_size \* config.per\_device\_batch\_size grad\_accum\_steps = max( 1, math.ceil(config.target\_global\_batch\_size / max(1, per\_step\_batch)), ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/tasks.py\* With 2 nodes × 4 GPUs × per-device batch size 1, the effective per-step batch is 8. To reach the default target of 16, the trainer accumulates over 2 steps. Change \`NUM\_NODES\` or \`per\_device\_batch\_size\` and the calculation adjusts automatically. The trainer brings everything together: \`\`\` trainer = L.Trainer( accelerator="gpu", devices=DEVICES\_PER\_NODE, num\_nodes=NUM\_NODES, strategy=strategy, logger=wandb\_logger, precision="bf16-mixed", max\_epochs=config.epochs, accumulate\_grad\_batches=grad\_accum\_steps, callbacks=\[\ checkpoint\_callback,\ metrics\_callback,\ recovery\_callback,\ live\_report\_callback,\ \], gradient\_clip\_val=1.0, benchmark=True, log\_every\_n\_steps=1, ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/tasks.py\* \`precision="bf16-mixed"\` uses BFloat16, which matches FP32's dynamic range (unlike FP16), so you don't need loss scaling. This is the standard choice for modern VLM training. \`benchmark=True\` runs cuDNN autotuning on the first batch to select the fastest kernels for your specific input sizes. ### Fault tolerance and recovery Multi-node GPU jobs fail. Hardware hiccups, spot instance preemptions, NCCL timeouts, memory spikes, etc. and the question is when, not if. This pipeline handles it with a two-part system. After every validation epoch, the \`RecoveryArtifactCallback\` calls \`trainer.save\_checkpoint()\` to write a DeepSpeed checkpoint directory, then uploads all shard files to the recovery URI. Each node's local rank 0 uploads its own shards; global rank 0 uploads the metadata files (\`metrics.json\`, \`summary.json\`). A distributed barrier between save and upload ensures all workers finish before training continues. If training fails, the driver task catches the error and returns the last recovery artifact instead of propagating the failure: \`\`\` try: with wandb\_config( project=wandb\_project, entity=wandb\_entity, ): training\_artifacts = train\_qwen\_adapter\_multinode( train\_manifest=train\_manifest, val\_manifest=val\_manifest, images\_dir=images, config=config, resume\_from=resume\_training\_artifacts, recovery\_uri=recovery\_uri, ) except flyte.errors.RuntimeUserError as e: if recovery\_uri is None: raise e print(f"Training failed - recovering latest checkpoint bundle: {recovery\_uri}") try: recovered\_artifacts = Dir(path=recovery\_uri) recovered\_root = await download\_dir\_async(recovered\_artifacts) flyte.report.log( build\_qwen\_adapter\_report\_html(recovered\_root, None), do\_flush=True, ) return recovered\_artifacts except Exception: raise e \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/tasks.py\* A failed run still produces useful output: the best checkpoint reached before the failure, along with a partial training report. To resume from that point, pass the recovery artifact as \`resume\_training\_artifacts\` on the next run. The training task downloads it, finds the most recent \`.ckpt\` file, and passes it to \`trainer.fit()\` as \`ckpt\_path\`. Training picks up at the last saved epoch with optimizer state and metrics history intact. The recovery URI is constructed from the configurable base path and the run name: \`\`\` s3://your-bucket/qwen-vl-multinode-deepspeed//qwen\_vl\_training\_recovery/ \`\`\` This means each run gets its own recovery location, so you can identify exactly which run a checkpoint came from. ### Live observability \`flyte.report\` lets you push HTML content directly into the Flyte UI during task execution, with no separate monitoring infrastructure. The \`LiveTrainingReportCallback\` uses this to stream training metrics in real-time: \`\`\` def \_push\_update( self, \*, trainer, pl\_module, status: str, phase: str, train\_total=None, train\_lm=None, train\_recon=None, val\_total=None, note: str, ) -> None: adapter\_gate = float(torch.tanh(pl\_module.adapter.gate).detach().cpu()) def fmt(value): return f"{float(value):.4f}" if value is not None else "-" payload = { "step": trainer.global\_step, "phase": phase, "train\_total": fmt(train\_total), "train\_lm": fmt(train\_lm), "train\_recon": fmt(train\_recon), "val\_total": fmt(val\_total), "train\_total\_value": ( float(train\_total) if train\_total is not None else None ), "val\_total\_value": float(val\_total) if val\_total is not None else None, "adapter\_gate": f"{adapter\_gate:.4f}", "status": status, "resumed\_from": self.resumed\_from or "fresh run", "recovery\_path": self.recovery\_callback.latest\_path or "pending first checkpoint upload", "note": note, } flyte.report.log( f""" """, do\_flush=True, ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/qwen\_vl\_frozen\_backbone\_finetuning/callbacks.py\* \`on\_train\_start\` (see the full code) initializes the dashboard with an SVG loss chart and an HTML metrics table. Every \`report\_every\_n\_steps\` training steps, \`\_push\_update\` serializes the latest metrics into a \`

Visual Document Retrieval — Experiment Comparison

ViDoRe benchmark · {len(report.results)} experiment(s)

{len(report.results)}
Experiments
{best.config.name}
Best by Recall@K
{best.metrics.recall\_at\_k:.3f}
Best Recall@{best.metrics.k}
{best.metrics.ndcg\_at\_k:.3f}
Best NDCG@{best.metrics.k}
{best.metrics.mrr:.3f}
Best MRR

Metrics by Experiment

Ranked Results

{table\_rows}
ExperimentModel Recall@KNDCG@KMRRK
""" await flyte.report.replace.aio(html) await flyte.report.flush.aio() # ───────────────────────────────────────────────────────────────────────────── # Experiment orchestration # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment run\_experiment}} @driver.task async def run\_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult: """ End-to-end retrieval pipeline for a single ExperimentConfig. Flyte v2's dynamic execution means this driver task can call GPU tasks (index\_colpali, search\_colpali) based on the runtime value of config.model — no static DAG wiring required. The if/elif is plain Python; Flyte schedules the selected sub-tasks on the appropriate environment. Caching: two experiments that share the same model and corpus (e.g. ColPali at top\_k=5 and top\_k=10) will hit the same cached index. GPU work is paid at most once per (model, corpus) pair across all experiments. Search queries are sharded into chunks of SEARCH\_SHARD\_SIZE and dispatched as concurrent task invocations. All shards land on the single warm container (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the GPU saturated throughout search rather than processing one large sequential batch from a single caller. flyte.group wraps each experiment in a named span in the Flyte UI, making it easy to compare latencies and drill into individual runs. """ SEARCH\_SHARD\_SIZE = 256 with flyte.group(config.name): if config.model == RetrievalModel.COLPALI: index\_file = await index\_colpali(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_colpali(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] elif config.model == RetrievalModel.SIGLIP: index\_file = await index\_siglip(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_siglip(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] else: # RetrievalModel.OCR\_BM25 page\_texts = await extract\_page\_texts(dataset.page\_files) results = await search\_bm25(page\_texts, dataset.page\_ids, dataset.queries, config.top\_k) metrics = await evaluate(results, dataset.queries, config.top\_k) return ExperimentResult(config=config, metrics=metrics) # {{/docs-fragment run\_experiment}} # {{docs-fragment compare\_experiments}} @driver.task async def compare\_experiments( configs: list\[ExperimentConfig\], subset: str = "docvqa", max\_pages: int = 200, ) -> ComparisonReport: """ Fan out over all experiment configs and return a ranked comparison table. The dataset is loaded once and shared across all experiments. Each config runs as a concurrent Flyte task via asyncio.gather. Experiments that share a model reuse the cached index — you only pay GPU time for new work. On completion, generate\_report emits an interactive Chart.js HTML report visible directly in the Flyte execution detail page. Default dataset: vidore\_v3\_finance\_en (~2 942 corpus pages, 1 854 queries) with max\_pages=2 000 to exercise the GPU pipeline at scale. """ dataset = await load\_vidore\_pages(subset=subset, max\_pages=max\_pages) # All experiments launch concurrently. Shared cached outputs (same model, # same corpus) are served from cache rather than recomputed. experiment\_coros = \[run\_experiment(config=cfg, dataset=dataset) for cfg in configs\] results: list\[ExperimentResult\] = list(await asyncio.gather(\*experiment\_coros)) report = ComparisonReport(results=results) print(report.summary()) best = report.best\_by("recall\_at\_k") print(f"\\nBest by Recall@{best.metrics.k}: {best.config.name}") # Emit the interactive HTML report in the Flyte UI. await generate\_report(report) return report # {{/docs-fragment compare\_experiments}} # ───────────────────────────────────────────────────────────────────────────── # Entry point # ───────────────────────────────────────────────────────────────────────────── if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Define the experiment grid. Each ExperimentConfig is one point in the # design space. Adding a new model or varying top\_k is one line here — # no task code changes required. # # ColPali appears twice with different top\_k values. The cache ensures # index\_colpali runs only once and both experiments share that result. # {{docs-fragment grid}} configs = \[\ ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top\_k=5),\ ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top\_k=10),\ ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top\_k=5),\ ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR\_BM25, top\_k=5),\ \] # {{/docs-fragment grid}} run = flyte.with\_runcontext().run( compare\_experiments, configs=configs, subset="vidore\_v3\_finance\_en", max\_pages=2000, ) print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval\_eval.py\* The Python dependencies (ColPali, transformers, docTR, etc.) are declared in the \`uv\` script header at the top of the file. ## Define the task environments Each model gets its own GPU environment so their warm-container pools scale independently. The ColPali and SigLIP environments use \`ReusePolicy\` to keep model weights resident; the driver coordinates orchestration, BM25, evaluation, and reporting. \`\`\` # /// script # requires-python = ">=3.12" # dependencies = \[\ # "colpali-engine>=0.3.1",\ # "transformers>=4.41",\ # "sentencepiece>=0.2",\ # "torch>=2.0",\ # "pillow>=10",\ # "datasets>=2.18",\ # "rank-bm25>=0.2",\ # "numpy>=1.26",\ # "python-doctr\[torch\]>=0.8",\ # "pydantic>=2.0",\ # "flyte>=2.0.0",\ # \] # /// """ Multimodal Retrieval Evaluation Pipeline This tutorial is an experiment framework for benchmarking visual document retrieval approaches on the ViDoRe benchmark. Each experiment is defined by an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and returns a ranked comparison table with an interactive HTML report. The corpus is a set of PDF page images; queries are plain-text questions. Each retrieval method must find the page that answers each question — no text is provided to the model, only the raw image. ColPali-v1.2 — patch-level multi-vector embeddings from a VLM (PaliGemma). No OCR. The model produces one vector per image patch (~1024 per page). MaxSim late-interaction scoring finds the best matching patch for each query token. SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP successor. One matrix multiply per query; fast and effective but a single vector cannot localise fine-grained regions. OCR + BM25 — text-only baseline. doctr (GPU OCR) extracts text in batches, BM25 matches keywords. Strong on text-dense pages; fails on charts, tables, and figures where content is visual. """ import asyncio import enum import json import math import os import tempfile from functools import lru\_cache from io import BytesIO from itertools import islice import numpy as np from PIL import Image as PILImage from pydantic import BaseModel from rank\_bm25 import BM25Okapi from extras import DynamicBatcher import flyte import flyte.report from flyte.io import File # ───────────────────────────────────────────────────────────────────────────── # Environments # ───────────────────────────────────────────────────────────────────────────── # One Docker image for all tasks. The PEP 723 header defines Python deps. # ca-certificates is required for HTTPS calls to HuggingFace and blob stores. # {{docs-fragment image}} image = ( flyte.Image.from\_uv\_script(\_\_file\_\_, name="vidore-eval-v2") .with\_apt\_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0") # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy. # Without it every reusable container exits with StartError (exit code 128). .with\_pip\_packages("unionai-reuse>=0.1.11") ) # {{/docs-fragment image}} # GPU environment for ColPali image encoding and search. # # ReusePolicy keeps up to 3 warm GPU containers alive between task calls. # Without it, every task invocation cold-starts a new container and downloads # ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model # weights already loaded into VRAM — is reused for the next task dispatch. # # replicas=1 single warm container — all concurrent shard calls land # here so they share one DynamicBatcher process # concurrency=8 up to 8 query-shard tasks run simultaneously on the # container, all feeding the same DynamicBatcher queue # idle\_ttl=120 keep alive 2 min after the last task finishes # scaledown\_ttl=60 scale to zero after 1 min of complete inactivity # {{docs-fragment envs}} colpali\_indexer = flyte.TaskEnvironment( name="vidore-colpali-indexer", image=image, resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for SigLIP image encoding and search. # # Separate from the ColPali environment so each model's warm containers # are managed independently — ColPali and SigLIP experiments can scale # without contending for the same pool of reusable containers. siglip\_indexer = flyte.TaskEnvironment( name="vidore-siglip-indexer", image=image, resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition) # in batches on GPU — much faster than CPU Tesseract. # No ReusePolicy needed: the result is cached, so this task runs at most once. ocr\_engine = flyte.TaskEnvironment( name="vidore-ocr-engine", image=image, resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1), ) # Driver: orchestration, BM25 search, evaluation, and reporting. # depends\_on ensures the shared Docker image is built before all environments # try to schedule tasks. driver = flyte.TaskEnvironment( name="vidore-driver", image=image, resources=flyte.Resources(cpu=2, memory="12Gi"), depends\_on=\[colpali\_indexer, siglip\_indexer, ocr\_engine\], ) # {{/docs-fragment envs}} # ───────────────────────────────────────────────────────────────────────────── # Configuration types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment config\_types}} class RetrievalModel(str, enum.Enum): """Retrieval backend to evaluate.""" COLPALI = "colpali-v1.2" # multi-vector patch embeddings, MaxSim SIGLIP = "siglip-so400m" # single-vector global embedding, cosine sim OCR\_BM25 = "ocr+bm25" # text extracted by Tesseract, ranked by BM25 class ExperimentConfig(BaseModel): """ All knobs for one retrieval experiment. Passed as a typed Flyte input. Because ExperimentConfig is a Pydantic model, Flyte serialises it alongside every task output — so you can always reconstruct which config produced which metric without maintaining a separate log. """ name: str # human-readable label shown in the comparison table model: RetrievalModel top\_k: int = 5 # number of pages to retrieve per query # {{/docs-fragment config\_types}} # ───────────────────────────────────────────────────────────────────────────── # Data types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment data\_types}} class PageQuery(BaseModel): """One retrieval query with its ground-truth page.""" query\_id: str text: str # e.g. "What was revenue growth in Q3?" relevant\_page\_id: str # one correct page per query class PageDataset(BaseModel): """ A corpus of document page images paired with text queries. page\_ids: unique page identifiers (derived from ViDoRe image filenames). page\_files: the same pages stored in Flyte's blob store as JPEG File handles. Tasks read images directly from here; no live HTTP. queries: text questions with ground-truth page IDs for evaluation. """ page\_ids: list\[str\] page\_files: list\[File\] queries: list\[PageQuery\] class Config: arbitrary\_types\_allowed = True class RetrievalResult(BaseModel): query\_id: str ranked\_page\_ids: list\[str\] # ordered best → worst class Metrics(BaseModel): recall\_at\_k: float ndcg\_at\_k: float mrr: float k: int class ExperimentResult(BaseModel): config: ExperimentConfig metrics: Metrics # {{/docs-fragment data\_types}} class ComparisonReport(BaseModel): results: list\[ExperimentResult\] def best\_by(self, metric: str = "recall\_at\_k") -> ExperimentResult: return max(self.results, key=lambda r: getattr(r.metrics, metric)) def summary(self) -> str: header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}" sep = "─" \* len(header) rows = \[header, sep\] for r in sorted(self.results, key=lambda x: -x.metrics.recall\_at\_k): rows.append( f"{r.config.name:<30} " f"{r.config.model.value:<18} " f"{r.metrics.recall\_at\_k:>10.3f} " f"{r.metrics.ndcg\_at\_k:>8.3f} " f"{r.metrics.mrr:>7.3f}" ) return "\\n".join(rows) # ───────────────────────────────────────────────────────────────────────────── # Cached model loaders # ───────────────────────────────────────────────────────────────────────────── # These functions are at module level so they are shared across all tasks that # run on the same warm container (via ReusePolicy). lru\_cache(maxsize=1) means # the model is loaded from disk/HuggingFace exactly once per container process # and kept in GPU memory for every subsequent task dispatch to that container. @lru\_cache(maxsize=1) def \_colpali\_model(): """Load ColPali-v1.2 into GPU memory and cache the result. device\_map= is the correct loading pattern for ColPali's PaliGemma backbone; it handles weight placement via accelerate. torch.compile is skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross- invocation batching is the primary GPU utilisation mechanism. """ import torch from colpali\_engine.models import ColPali, ColPaliProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = ColPali.from\_pretrained( "vidore/colpali-v1.2", torch\_dtype=torch.bfloat16, device\_map=device, ) processor = ColPaliProcessor.from\_pretrained("vidore/colpali-v1.2") return model, processor, device @lru\_cache(maxsize=1) def \_siglip\_model(): """Load SigLIP SO400M into GPU memory, compile it, and cache the result. torch.compile (mode="reduce-overhead") fuses the vision and text encoder transformer layers into optimised CUDA kernels. As with ColPali, the compilation overhead is paid once per warm container lifetime. """ import torch from transformers import AutoModel, AutoProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = AutoModel.from\_pretrained("google/siglip-so400m-patch14-224").to(device) if device == "cuda": model = torch.compile(model, mode="reduce-overhead") processor = AutoProcessor.from\_pretrained("google/siglip-so400m-patch14-224") return model, processor, device @lru\_cache(maxsize=1) def \_ocr\_model(): """Load the doctr OCR predictor onto GPU and cache it. doctr's ocr\_predictor bundles a detection model (DBNet) and a recognition model (CRNN/SAR) into a single callable. pretrained=True downloads both model weights from doctr's model zoo on first use. """ import torch from doctr.models import ocr\_predictor predictor = ocr\_predictor(pretrained=True) if torch.cuda.is\_available(): predictor = predictor.cuda() return predictor # ───────────────────────────────────────────────────────────────────────────── # Search batcher singletons # ───────────────────────────────────────────────────────────────────────────── # One DynamicBatcher per model, shared across all concurrent search task # invocations on the same warm container (concurrency=3). Queries from every # concurrent caller are aggregated into a single GPU batch, maximizing # throughput compared to each invocation running its own forward pass. # # Initialised lazily on the first search call via double-checked locking and # lives for the container's lifetime. The process\_fn runs GPU work via # asyncio.to\_thread so the aggregation loop can continue collecting queries # from other callers while the GPU processes the current batch. # # File is not hashable so alru\_cache cannot be used here; module-level state # with asyncio.Lock is the correct pattern. # # Assumption: index\_colpali/index\_siglip use cache="auto", so the same corpus # always produces the same index File across all callers on this container. If # the index file ever changed between calls, the batcher would silently continue # using the corpus embeddings loaded from the first call. \_colpali\_batcher: DynamicBatcher | None = None \_colpali\_batcher\_lock = asyncio.Lock() \_siglip\_batcher: DynamicBatcher | None = None \_siglip\_batcher\_lock = asyncio.Lock() async def \_get\_colpali\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level ColPali search batcher, creating it on first call.""" global \_colpali\_batcher if \_colpali\_batcher is not None: return \_colpali\_batcher async with \_colpali\_batcher\_lock: if \_colpali\_batcher is not None: return \_colpali\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, n\_patches, dim) index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_colpali\_model() corpus\_emb = corpus\_emb.to(device, dtype=torch.float32) async def colpali\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: query\_inputs = processor.process\_queries(\[q.text for q in batch\]) query\_inputs = {k: v.to(device) for k, v in query\_inputs.items()} with torch.no\_grad(): query\_embs = model(\*\*query\_inputs).float() # (B, T, D) query\_chunk = 8 n\_pages = corpus\_emb.shape\[0\] all\_scores = torch.empty(len(batch), n\_pages, device=device) for start in range(0, len(batch), query\_chunk): chunk = query\_embs\[start : start + query\_chunk\] all\_scores\[start : start + query\_chunk\] = ( torch.einsum("ctd,pjd->ctpj", chunk, corpus\_emb) .max(dim=3).values .sum(dim=1) ) sorted\_indices = all\_scores.argsort(dim=1, descending=True).cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] # Run GPU work in a thread so the event loop — and the batcher's # aggregation loop — remain unblocked while the GPU is busy. return await asyncio.to\_thread(\_gpu\_work) batcher: DynamicBatcher\[PageQuery, list\[str\]\] = DynamicBatcher( process\_fn=colpali\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_colpali\_batcher = batcher return \_colpali\_batcher async def \_get\_siglip\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level SigLIP search batcher, creating it on first call.""" global \_siglip\_batcher if \_siglip\_batcher is not None: return \_siglip\_batcher async with \_siglip\_batcher\_lock: if \_siglip\_batcher is not None: return \_siglip\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, dim), L2-normalised index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_siglip\_model() corpus\_emb = corpus\_emb.to(device) async def siglip\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: text\_inputs = processor( text=\[q.text for q in batch\], return\_tensors="pt", padding=True, truncation=True, ).to(device) with torch.no\_grad(): text\_out = model.text\_model(\*\*text\_inputs) query\_embs = text\_out.pooler\_output # (B, dim) query\_embs = query\_embs / query\_embs.norm(dim=-1, keepdim=True) scores\_matrix = corpus\_emb @ query\_embs.T # (n\_pages, B) sorted\_indices = scores\_matrix.argsort(dim=0, descending=True).T.cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] return await asyncio.to\_thread(\_gpu\_work) batcher = DynamicBatcher( process\_fn=siglip\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_siglip\_batcher = batcher return \_siglip\_batcher # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── def \_batches(items: list, batch\_size: int): """Yield successive fixed-size batches from a list.""" for start in range(0, len(items), batch\_size): yield items\[start : start + batch\_size\] def \_load\_image\_sync(f: File) -> PILImage.Image: """Blocking download + decode. Intended to be called from a thread pool.""" with f.open\_sync("rb") as fh: data = fh.read() return PILImage.open(BytesIO(data)).convert("RGB") async def \_load\_image(f: File) -> PILImage.Image: """Download and decode a page image in a thread-pool worker. asyncio.to\_thread runs \_load\_image\_sync in a real OS thread so that blocking network I/O can overlap with GPU-bound forward passes when images are pre-submitted via loop.run\_in\_executor before the GPU kernel. """ return await asyncio.to\_thread(\_load\_image\_sync, f) async def \_load\_npz(index\_file: File) -> np.lib.npyio.NpzFile: """Download an index File to a local temp path and open with np.load.""" with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp: async with index\_file.open("rb") as fh: tmp.write(bytes(await fh.read())) return np.load(tmp.name) def \_dcg(relevances: list\[int\]) -> float: return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances)) # ───────────────────────────────────────────────────────────────────────────── # Tasks — data loading # ───────────────────────────────────────────────────────────────────────────── @driver.task(cache="auto", retries=3) async def load\_vidore\_pages(subset: str = "docvqa", max\_pages: int = 200) -> PageDataset: """ Load a ViDoRe benchmark subset and store page images in Flyte's blob store. Supports two dataset formats: Legacy (subsampled) — single 'test' split with one row per (query, page) pair; fields: image, query, image\_filename. streaming=True reads only the rows requested via islice — no full-shard download. Datasets: vidore/docvqa\_test\_subsampled, vidore/infovqa\_test\_subsampled V3 — separate corpus / queries / qrels splits following the BEIR retrieval benchmark format. corpus contains page images; queries contains question text; qrels maps query IDs to relevant corpus page IDs (many-to-many). Datasets: vidore/vidore\_v3\_finance\_en (~2 942 pages, 1 854 queries) The first call uploads page images to Flyte's blob store and caches the PageDataset; every subsequent call with the same arguments returns the cached result instantly. retries=3 guards against transient HuggingFace network failures. Available subsets: "docvqa", "infovqa", "vidore\_v3\_finance\_en" """ from datasets import load\_dataset subset\_map = { "docvqa": "vidore/docvqa\_test\_subsampled", "infovqa": "vidore/infovqa\_test\_subsampled", "vidore\_v3\_finance\_en": "vidore/vidore\_v3\_finance\_en", } dataset\_name = subset\_map.get(subset, f"vidore/{subset}\_test\_subsampled") # V3 datasets ship with separate corpus / queries / qrels splits. \_V3\_SUBSETS = {"vidore\_v3\_finance\_en"} if subset in \_V3\_SUBSETS: # ── V3 format ───────────────────────────────────────────────────────── # corpus / queries / qrels are HuggingFace configs (name=), not splits. # corpus uses streaming=True so images are decoded one at a time — # loading all 2 942 rows eagerly would hold gigabytes of PIL images in # the driver's RAM simultaneously. qrels and queries are text-only and # small enough to load fully into memory. corpus\_ds = load\_dataset(dataset\_name, name="corpus", split="test", streaming=True) qrels\_ds = load\_dataset(dataset\_name, name="qrels", split="test") queries\_ds = load\_dataset(dataset\_name, name="queries", split="test") # Normalise field names — V3 follows BEIR convention (hyphenated ids). def \_col(ds, \*candidates): cols = set(ds.column\_names) for c in candidates: if c in cols: return c raise KeyError(f"None of {candidates} found in columns {cols}") corpus\_id\_col = \_col(corpus\_ds, "corpus-id", "corpus\_id", "id", "\_id") query\_id\_col = \_col(queries\_ds, "query-id", "query\_id", "id", "\_id") query\_text\_col = \_col(queries\_ds, "query", "text") qrel\_qid\_col = \_col(qrels\_ds, "query-id", "query\_id") qrel\_cid\_col = \_col(qrels\_ds, "corpus-id", "corpus\_id") # Slice corpus to max\_pages, upload each image to Flyte blob store. page\_ids: list\[str\] = \[\] page\_files: list\[File\] = \[\] corpus\_id\_to\_page\_id: dict\[str, str\] = {} for i, row in enumerate(islice(corpus\_ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue cid = str(row\[corpus\_id\_col\]) page\_id = f"{subset}\_{i:04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) corpus\_id\_to\_page\_id\[cid\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) # Build query\_id → relevant page\_id from qrels (first match wins). # Only keep relevance judgements whose corpus page is in our slice. qrel\_map: dict\[str, str\] = {} for row in qrels\_ds: qid = str(row\[qrel\_qid\_col\]) cid = str(row\[qrel\_cid\_col\]) if cid in corpus\_id\_to\_page\_id and qid not in qrel\_map: qrel\_map\[qid\] = corpus\_id\_to\_page\_id\[cid\] # Collect queries that have at least one relevant page in our slice. queries: list\[PageQuery\] = \[\] for row in queries\_ds: qid = str(row\[query\_id\_col\]) if qid not in qrel\_map: continue queries.append( PageQuery( query\_id=qid, text=str(row\[query\_text\_col\]), relevant\_page\_id=qrel\_map\[qid\], ) ) else: # ── Legacy format ───────────────────────────────────────────────────── # Single 'test' split with one row per (query, page) pair. ds = load\_dataset(dataset\_name, split="test", streaming=True) page\_ids = \[\] page\_files = \[\] queries = \[\] seen\_pages: dict\[str, str\] = {} # image\_filename → page\_id for i, row in enumerate(islice(ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue filename: str = row.get("image\_filename") or f"page\_{i}" query\_text: str = row.get("query", "") if not query\_text: continue # Each unique page is uploaded exactly once; multiple queries may # share the same page (same image\_filename). if filename not in seen\_pages: page\_id = f"{subset}\_{len(page\_ids):04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) seen\_pages\[filename\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) else: page\_id = seen\_pages\[filename\] queries.append( PageQuery( query\_id=f"q{i:04d}", text=query\_text, relevant\_page\_id=page\_id, ) ) print(f"Loaded {len(page\_ids)} unique pages, {len(queries)} queries", flush=True) return PageDataset(page\_ids=page\_ids, page\_files=page\_files, queries=queries) # ───────────────────────────────────────────────────────────────────────────── # Tasks — indexing # ───────────────────────────────────────────────────────────────────────────── @colpali\_indexer.task(cache="auto", retries=2) async def index\_colpali(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with ColPali-v1.2 and save the multi-vector index. ColPali skips OCR entirely. It feeds the raw page image into PaliGemma (a vision-language model) and produces one embedding vector per image patch — roughly 1,024 patches per page, each of dimension 128. \_colpali\_model() is an lru\_cache'd loader. On a cold container, it downloads and loads the model once. On a warm container (kept alive by ReusePolicy), it returns the already-loaded model instantly from cache — no repeated ~7 GB download. The index is stored as a .npz file in Flyte's blob store: embeddings — float32, shape (n\_pages, n\_patches, dim) page\_ids — matching page ID strings cache="auto" + retries=2: the result is stored permanently on success; transient failures (e.g. HuggingFace rate limits) are retried twice. """ import torch model, processor, device = \_colpali\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 4)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor.process\_images(images) inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no\_grad(): emb = model(\*\*inputs) # (batch, n\_patches, dim) all\_embeddings.append(emb.cpu().float().numpy()) print(f"ColPali: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, n\_patches, dim) out\_path = os.path.join(tempfile.gettempdir(), "colpali\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @siglip\_indexer.task(cache="auto", retries=2) async def index\_siglip(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with SigLIP SO400M and save the single-vector index. SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss instead of softmax — avoiding the normalisation bottleneck that limits CLIP's scalability. Produces one global embedding per page. \_siglip\_model() caches the model across warm container reuses. The index is stored as a .npz file: embeddings — float32, shape (n\_pages, dim), L2-normalised page\_ids — matching page ID strings """ import torch model, processor, device = \_siglip\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 8)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor(images=images, return\_tensors="pt", padding=True).to(device) with torch.no\_grad(): outputs = model.vision\_model(\*\*inputs) emb = outputs.pooler\_output # (batch, dim) emb = emb / emb.norm(dim=-1, keepdim=True) # L2 normalise all\_embeddings.append(emb.cpu().float().numpy()) print(f"SigLIP: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, dim) out\_path = os.path.join(tempfile.gettempdir(), "siglip\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @ocr\_engine.task(cache="auto") async def extract\_page\_texts(page\_files: list\[File\]) -> list\[str\]: """ OCR every page with doctr on GPU to produce a text-only baseline. doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single callable predictor. Pages are downloaded in parallel then fed in batches of ocr\_batch\_size. asyncio.to\_thread keeps the event loop unblocked during GPU inference. Result structure: result.pages\[i\].blocks\[j\].lines\[k\].words\[l\].value Cached: the same corpus is OCR'd at most once across all experiments that use the OCR+BM25 backend. """ import gc predictor = \_ocr\_model() # Process in batches: download each batch just-in-time so only # ocr\_batch\_size images are in memory at once instead of all 2 000. ocr\_batch\_size = 8 total = len(page\_files) texts: list\[str\] = \[\] for start in range(0, total, ocr\_batch\_size): batch\_files = page\_files\[start : start + ocr\_batch\_size\] batch\_images = list( await asyncio.gather(\*\[asyncio.to\_thread(\_load\_image\_sync, f) for f in batch\_files\]) ) batch\_np = \[np.array(img) for img in batch\_images\] del batch\_images result = await asyncio.to\_thread(predictor, batch\_np) del batch\_np for page\_output in result.pages: texts.append( "\\n".join( " ".join(word.value for word in line.words) for block in page\_output.blocks for line in block.lines ) ) del result gc.collect() print(f"OCR: processed {min(start + ocr\_batch\_size, total)}/{total} pages", flush=True) return texts # ───────────────────────────────────────────────────────────────────────────── # Tasks — search # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment search\_colpali}} @colpali\_indexer.task async def search\_colpali( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher. MaxSim score for page p given query q: score(q, p) = Σ\_{t ∈ query tokens} max\_{j ∈ page patches} (q\_t · p\_j) Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_colpali invocations on the same warm container (concurrency=8) into a single GPU batch. This keeps the GPU saturated rather than running one small batch per caller. The batcher's process\_fn runs GPU work in asyncio.to\_thread, so the aggregation loop stays live while the GPU encodes and scores. """ batcher = await \_get\_colpali\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] # {{/docs-fragment search\_colpali}} @siglip\_indexer.task async def search\_siglip( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using SigLIP cosine similarity via DynamicBatcher. Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_siglip invocations on the same warm container (concurrency=3) into a single GPU batch. SigLIP's single-vector embeddings make full vectorisation safe — the scores matrix (n\_pages x n\_queries) is small enough to materialise in one GPU call regardless of batch size. """ batcher = await \_get\_siglip\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] @driver.task async def search\_bm25( page\_texts: list\[str\], page\_ids: list\[str\], queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using BM25 over OCR'd text. The standard keyword-based baseline. No GPU required; strong on text-dense pages, weak on visual content that Tesseract cannot read. """ tokenized = \[text.lower().split() for text in page\_texts\] bm25 = BM25Okapi(tokenized) results: list\[RetrievalResult\] = \[\] for q in queries: scores = bm25.get\_scores(q.text.lower().split()) ranked = sorted(range(len(page\_ids)), key=lambda i: -scores\[i\])\[:top\_k\] results.append( RetrievalResult( query\_id=q.query\_id, ranked\_page\_ids=\[page\_ids\[i\] for i in ranked\], ) ) return results # ───────────────────────────────────────────────────────────────────────────── # Tasks — evaluation # ───────────────────────────────────────────────────────────────────────────── @driver.task async def evaluate( results: list\[RetrievalResult\], ground\_truth: list\[PageQuery\], k: int, ) -> Metrics: """ Compute Recall@K, NDCG@K, and MRR for a single retrieval model. Recall@K — was the correct page in the top-K results? NDCG@K — normalised discounted cumulative gain; rewards earlier hits. MRR — mean reciprocal rank of the first correct result. All three are averaged over all queries. Higher is better. """ gt\_map = {q.query\_id: q.relevant\_page\_id for q in ground\_truth} recall\_vals, ndcg\_vals, mrr\_vals = \[\], \[\], \[\] for r in results: relevant = gt\_map.get(r.query\_id, "") top = r.ranked\_page\_ids\[:k\] recall\_vals.append(1.0 if relevant in top else 0.0) rels = \[1 if pid == relevant else 0 for pid in top\] idcg = \_dcg(\[1\]) # ideal: correct page at rank 1 ndcg\_vals.append(\_dcg(rels) / idcg if idcg > 0 else 0.0) rr = 0.0 for rank, pid in enumerate(r.ranked\_page\_ids, start=1): if pid == relevant: rr = 1.0 / rank break mrr\_vals.append(rr) return Metrics( recall\_at\_k=float(np.mean(recall\_vals)), ndcg\_at\_k=float(np.mean(ndcg\_vals)), mrr=float(np.mean(mrr\_vals)), k=k, ) # ───────────────────────────────────────────────────────────────────────────── # Tasks — report # ───────────────────────────────────────────────────────────────────────────── @driver.task(report=True) async def generate\_report(report: ComparisonReport) -> None: """ Emit an interactive HTML report visible in the Flyte UI. report=True marks this task as a reporting task. Flyte renders the HTML returned via flyte.report.replace.aio() directly in the execution detail page — no separate dashboard or export step required. The report contains: - Summary cards: experiment count, best model, best Recall@K. - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment. - Ranked results table with all three metrics. """ sorted\_results = sorted(report.results, key=lambda r: -r.metrics.recall\_at\_k) best = sorted\_results\[0\] labels = \[r.config.name for r in sorted\_results\] recall\_vals = \[r.metrics.recall\_at\_k for r in sorted\_results\] ndcg\_vals = \[r.metrics.ndcg\_at\_k for r in sorted\_results\] mrr\_vals = \[r.metrics.mrr for r in sorted\_results\] table\_rows = "".join( f""" {r.config.name} {r.config.model.value} {r.metrics.recall\_at\_k:.3f} {r.metrics.ndcg\_at\_k:.3f} {r.metrics.mrr:.3f} {r.metrics.k} """ for r in sorted\_results ) html = f""" Visual Document Retrieval — Results

Visual Document Retrieval — Experiment Comparison

ViDoRe benchmark · {len(report.results)} experiment(s)

{len(report.results)}
Experiments
{best.config.name}
Best by Recall@K
{best.metrics.recall\_at\_k:.3f}
Best Recall@{best.metrics.k}
{best.metrics.ndcg\_at\_k:.3f}
Best NDCG@{best.metrics.k}
{best.metrics.mrr:.3f}
Best MRR

Metrics by Experiment

Ranked Results

{table\_rows}
ExperimentModel Recall@KNDCG@KMRRK
""" await flyte.report.replace.aio(html) await flyte.report.flush.aio() # ───────────────────────────────────────────────────────────────────────────── # Experiment orchestration # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment run\_experiment}} @driver.task async def run\_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult: """ End-to-end retrieval pipeline for a single ExperimentConfig. Flyte v2's dynamic execution means this driver task can call GPU tasks (index\_colpali, search\_colpali) based on the runtime value of config.model — no static DAG wiring required. The if/elif is plain Python; Flyte schedules the selected sub-tasks on the appropriate environment. Caching: two experiments that share the same model and corpus (e.g. ColPali at top\_k=5 and top\_k=10) will hit the same cached index. GPU work is paid at most once per (model, corpus) pair across all experiments. Search queries are sharded into chunks of SEARCH\_SHARD\_SIZE and dispatched as concurrent task invocations. All shards land on the single warm container (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the GPU saturated throughout search rather than processing one large sequential batch from a single caller. flyte.group wraps each experiment in a named span in the Flyte UI, making it easy to compare latencies and drill into individual runs. """ SEARCH\_SHARD\_SIZE = 256 with flyte.group(config.name): if config.model == RetrievalModel.COLPALI: index\_file = await index\_colpali(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_colpali(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] elif config.model == RetrievalModel.SIGLIP: index\_file = await index\_siglip(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_siglip(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] else: # RetrievalModel.OCR\_BM25 page\_texts = await extract\_page\_texts(dataset.page\_files) results = await search\_bm25(page\_texts, dataset.page\_ids, dataset.queries, config.top\_k) metrics = await evaluate(results, dataset.queries, config.top\_k) return ExperimentResult(config=config, metrics=metrics) # {{/docs-fragment run\_experiment}} # {{docs-fragment compare\_experiments}} @driver.task async def compare\_experiments( configs: list\[ExperimentConfig\], subset: str = "docvqa", max\_pages: int = 200, ) -> ComparisonReport: """ Fan out over all experiment configs and return a ranked comparison table. The dataset is loaded once and shared across all experiments. Each config runs as a concurrent Flyte task via asyncio.gather. Experiments that share a model reuse the cached index — you only pay GPU time for new work. On completion, generate\_report emits an interactive Chart.js HTML report visible directly in the Flyte execution detail page. Default dataset: vidore\_v3\_finance\_en (~2 942 corpus pages, 1 854 queries) with max\_pages=2 000 to exercise the GPU pipeline at scale. """ dataset = await load\_vidore\_pages(subset=subset, max\_pages=max\_pages) # All experiments launch concurrently. Shared cached outputs (same model, # same corpus) are served from cache rather than recomputed. experiment\_coros = \[run\_experiment(config=cfg, dataset=dataset) for cfg in configs\] results: list\[ExperimentResult\] = list(await asyncio.gather(\*experiment\_coros)) report = ComparisonReport(results=results) print(report.summary()) best = report.best\_by("recall\_at\_k") print(f"\\nBest by Recall@{best.metrics.k}: {best.config.name}") # Emit the interactive HTML report in the Flyte UI. await generate\_report(report) return report # {{/docs-fragment compare\_experiments}} # ───────────────────────────────────────────────────────────────────────────── # Entry point # ───────────────────────────────────────────────────────────────────────────── if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Define the experiment grid. Each ExperimentConfig is one point in the # design space. Adding a new model or varying top\_k is one line here — # no task code changes required. # # ColPali appears twice with different top\_k values. The cache ensures # index\_colpali runs only once and both experiments share that result. # {{docs-fragment grid}} configs = \[\ ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top\_k=5),\ ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top\_k=10),\ ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top\_k=5),\ ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR\_BM25, top\_k=5),\ \] # {{/docs-fragment grid}} run = flyte.with\_runcontext().run( compare\_experiments, configs=configs, subset="vidore\_v3\_finance\_en", max\_pages=2000, ) print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval\_eval.py\* ## Configuration and data types An experiment is fully described by an \`ExperimentConfig\`. Because it's a Pydantic model, Flyte serializes it alongside every output. \`\`\` # /// script # requires-python = ">=3.12" # dependencies = \[\ # "colpali-engine>=0.3.1",\ # "transformers>=4.41",\ # "sentencepiece>=0.2",\ # "torch>=2.0",\ # "pillow>=10",\ # "datasets>=2.18",\ # "rank-bm25>=0.2",\ # "numpy>=1.26",\ # "python-doctr\[torch\]>=0.8",\ # "pydantic>=2.0",\ # "flyte>=2.0.0",\ # \] # /// """ Multimodal Retrieval Evaluation Pipeline This tutorial is an experiment framework for benchmarking visual document retrieval approaches on the ViDoRe benchmark. Each experiment is defined by an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and returns a ranked comparison table with an interactive HTML report. The corpus is a set of PDF page images; queries are plain-text questions. Each retrieval method must find the page that answers each question — no text is provided to the model, only the raw image. ColPali-v1.2 — patch-level multi-vector embeddings from a VLM (PaliGemma). No OCR. The model produces one vector per image patch (~1024 per page). MaxSim late-interaction scoring finds the best matching patch for each query token. SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP successor. One matrix multiply per query; fast and effective but a single vector cannot localise fine-grained regions. OCR + BM25 — text-only baseline. doctr (GPU OCR) extracts text in batches, BM25 matches keywords. Strong on text-dense pages; fails on charts, tables, and figures where content is visual. """ import asyncio import enum import json import math import os import tempfile from functools import lru\_cache from io import BytesIO from itertools import islice import numpy as np from PIL import Image as PILImage from pydantic import BaseModel from rank\_bm25 import BM25Okapi from extras import DynamicBatcher import flyte import flyte.report from flyte.io import File # ───────────────────────────────────────────────────────────────────────────── # Environments # ───────────────────────────────────────────────────────────────────────────── # One Docker image for all tasks. The PEP 723 header defines Python deps. # ca-certificates is required for HTTPS calls to HuggingFace and blob stores. # {{docs-fragment image}} image = ( flyte.Image.from\_uv\_script(\_\_file\_\_, name="vidore-eval-v2") .with\_apt\_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0") # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy. # Without it every reusable container exits with StartError (exit code 128). .with\_pip\_packages("unionai-reuse>=0.1.11") ) # {{/docs-fragment image}} # GPU environment for ColPali image encoding and search. # # ReusePolicy keeps up to 3 warm GPU containers alive between task calls. # Without it, every task invocation cold-starts a new container and downloads # ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model # weights already loaded into VRAM — is reused for the next task dispatch. # # replicas=1 single warm container — all concurrent shard calls land # here so they share one DynamicBatcher process # concurrency=8 up to 8 query-shard tasks run simultaneously on the # container, all feeding the same DynamicBatcher queue # idle\_ttl=120 keep alive 2 min after the last task finishes # scaledown\_ttl=60 scale to zero after 1 min of complete inactivity # {{docs-fragment envs}} colpali\_indexer = flyte.TaskEnvironment( name="vidore-colpali-indexer", image=image, resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for SigLIP image encoding and search. # # Separate from the ColPali environment so each model's warm containers # are managed independently — ColPali and SigLIP experiments can scale # without contending for the same pool of reusable containers. siglip\_indexer = flyte.TaskEnvironment( name="vidore-siglip-indexer", image=image, resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition) # in batches on GPU — much faster than CPU Tesseract. # No ReusePolicy needed: the result is cached, so this task runs at most once. ocr\_engine = flyte.TaskEnvironment( name="vidore-ocr-engine", image=image, resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1), ) # Driver: orchestration, BM25 search, evaluation, and reporting. # depends\_on ensures the shared Docker image is built before all environments # try to schedule tasks. driver = flyte.TaskEnvironment( name="vidore-driver", image=image, resources=flyte.Resources(cpu=2, memory="12Gi"), depends\_on=\[colpali\_indexer, siglip\_indexer, ocr\_engine\], ) # {{/docs-fragment envs}} # ───────────────────────────────────────────────────────────────────────────── # Configuration types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment config\_types}} class RetrievalModel(str, enum.Enum): """Retrieval backend to evaluate.""" COLPALI = "colpali-v1.2" # multi-vector patch embeddings, MaxSim SIGLIP = "siglip-so400m" # single-vector global embedding, cosine sim OCR\_BM25 = "ocr+bm25" # text extracted by Tesseract, ranked by BM25 class ExperimentConfig(BaseModel): """ All knobs for one retrieval experiment. Passed as a typed Flyte input. Because ExperimentConfig is a Pydantic model, Flyte serialises it alongside every task output — so you can always reconstruct which config produced which metric without maintaining a separate log. """ name: str # human-readable label shown in the comparison table model: RetrievalModel top\_k: int = 5 # number of pages to retrieve per query # {{/docs-fragment config\_types}} # ───────────────────────────────────────────────────────────────────────────── # Data types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment data\_types}} class PageQuery(BaseModel): """One retrieval query with its ground-truth page.""" query\_id: str text: str # e.g. "What was revenue growth in Q3?" relevant\_page\_id: str # one correct page per query class PageDataset(BaseModel): """ A corpus of document page images paired with text queries. page\_ids: unique page identifiers (derived from ViDoRe image filenames). page\_files: the same pages stored in Flyte's blob store as JPEG File handles. Tasks read images directly from here; no live HTTP. queries: text questions with ground-truth page IDs for evaluation. """ page\_ids: list\[str\] page\_files: list\[File\] queries: list\[PageQuery\] class Config: arbitrary\_types\_allowed = True class RetrievalResult(BaseModel): query\_id: str ranked\_page\_ids: list\[str\] # ordered best → worst class Metrics(BaseModel): recall\_at\_k: float ndcg\_at\_k: float mrr: float k: int class ExperimentResult(BaseModel): config: ExperimentConfig metrics: Metrics # {{/docs-fragment data\_types}} class ComparisonReport(BaseModel): results: list\[ExperimentResult\] def best\_by(self, metric: str = "recall\_at\_k") -> ExperimentResult: return max(self.results, key=lambda r: getattr(r.metrics, metric)) def summary(self) -> str: header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}" sep = "─" \* len(header) rows = \[header, sep\] for r in sorted(self.results, key=lambda x: -x.metrics.recall\_at\_k): rows.append( f"{r.config.name:<30} " f"{r.config.model.value:<18} " f"{r.metrics.recall\_at\_k:>10.3f} " f"{r.metrics.ndcg\_at\_k:>8.3f} " f"{r.metrics.mrr:>7.3f}" ) return "\\n".join(rows) # ───────────────────────────────────────────────────────────────────────────── # Cached model loaders # ───────────────────────────────────────────────────────────────────────────── # These functions are at module level so they are shared across all tasks that # run on the same warm container (via ReusePolicy). lru\_cache(maxsize=1) means # the model is loaded from disk/HuggingFace exactly once per container process # and kept in GPU memory for every subsequent task dispatch to that container. @lru\_cache(maxsize=1) def \_colpali\_model(): """Load ColPali-v1.2 into GPU memory and cache the result. device\_map= is the correct loading pattern for ColPali's PaliGemma backbone; it handles weight placement via accelerate. torch.compile is skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross- invocation batching is the primary GPU utilisation mechanism. """ import torch from colpali\_engine.models import ColPali, ColPaliProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = ColPali.from\_pretrained( "vidore/colpali-v1.2", torch\_dtype=torch.bfloat16, device\_map=device, ) processor = ColPaliProcessor.from\_pretrained("vidore/colpali-v1.2") return model, processor, device @lru\_cache(maxsize=1) def \_siglip\_model(): """Load SigLIP SO400M into GPU memory, compile it, and cache the result. torch.compile (mode="reduce-overhead") fuses the vision and text encoder transformer layers into optimised CUDA kernels. As with ColPali, the compilation overhead is paid once per warm container lifetime. """ import torch from transformers import AutoModel, AutoProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = AutoModel.from\_pretrained("google/siglip-so400m-patch14-224").to(device) if device == "cuda": model = torch.compile(model, mode="reduce-overhead") processor = AutoProcessor.from\_pretrained("google/siglip-so400m-patch14-224") return model, processor, device @lru\_cache(maxsize=1) def \_ocr\_model(): """Load the doctr OCR predictor onto GPU and cache it. doctr's ocr\_predictor bundles a detection model (DBNet) and a recognition model (CRNN/SAR) into a single callable. pretrained=True downloads both model weights from doctr's model zoo on first use. """ import torch from doctr.models import ocr\_predictor predictor = ocr\_predictor(pretrained=True) if torch.cuda.is\_available(): predictor = predictor.cuda() return predictor # ───────────────────────────────────────────────────────────────────────────── # Search batcher singletons # ───────────────────────────────────────────────────────────────────────────── # One DynamicBatcher per model, shared across all concurrent search task # invocations on the same warm container (concurrency=3). Queries from every # concurrent caller are aggregated into a single GPU batch, maximizing # throughput compared to each invocation running its own forward pass. # # Initialised lazily on the first search call via double-checked locking and # lives for the container's lifetime. The process\_fn runs GPU work via # asyncio.to\_thread so the aggregation loop can continue collecting queries # from other callers while the GPU processes the current batch. # # File is not hashable so alru\_cache cannot be used here; module-level state # with asyncio.Lock is the correct pattern. # # Assumption: index\_colpali/index\_siglip use cache="auto", so the same corpus # always produces the same index File across all callers on this container. If # the index file ever changed between calls, the batcher would silently continue # using the corpus embeddings loaded from the first call. \_colpali\_batcher: DynamicBatcher | None = None \_colpali\_batcher\_lock = asyncio.Lock() \_siglip\_batcher: DynamicBatcher | None = None \_siglip\_batcher\_lock = asyncio.Lock() async def \_get\_colpali\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level ColPali search batcher, creating it on first call.""" global \_colpali\_batcher if \_colpali\_batcher is not None: return \_colpali\_batcher async with \_colpali\_batcher\_lock: if \_colpali\_batcher is not None: return \_colpali\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, n\_patches, dim) index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_colpali\_model() corpus\_emb = corpus\_emb.to(device, dtype=torch.float32) async def colpali\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: query\_inputs = processor.process\_queries(\[q.text for q in batch\]) query\_inputs = {k: v.to(device) for k, v in query\_inputs.items()} with torch.no\_grad(): query\_embs = model(\*\*query\_inputs).float() # (B, T, D) query\_chunk = 8 n\_pages = corpus\_emb.shape\[0\] all\_scores = torch.empty(len(batch), n\_pages, device=device) for start in range(0, len(batch), query\_chunk): chunk = query\_embs\[start : start + query\_chunk\] all\_scores\[start : start + query\_chunk\] = ( torch.einsum("ctd,pjd->ctpj", chunk, corpus\_emb) .max(dim=3).values .sum(dim=1) ) sorted\_indices = all\_scores.argsort(dim=1, descending=True).cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] # Run GPU work in a thread so the event loop — and the batcher's # aggregation loop — remain unblocked while the GPU is busy. return await asyncio.to\_thread(\_gpu\_work) batcher: DynamicBatcher\[PageQuery, list\[str\]\] = DynamicBatcher( process\_fn=colpali\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_colpali\_batcher = batcher return \_colpali\_batcher async def \_get\_siglip\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level SigLIP search batcher, creating it on first call.""" global \_siglip\_batcher if \_siglip\_batcher is not None: return \_siglip\_batcher async with \_siglip\_batcher\_lock: if \_siglip\_batcher is not None: return \_siglip\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, dim), L2-normalised index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_siglip\_model() corpus\_emb = corpus\_emb.to(device) async def siglip\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: text\_inputs = processor( text=\[q.text for q in batch\], return\_tensors="pt", padding=True, truncation=True, ).to(device) with torch.no\_grad(): text\_out = model.text\_model(\*\*text\_inputs) query\_embs = text\_out.pooler\_output # (B, dim) query\_embs = query\_embs / query\_embs.norm(dim=-1, keepdim=True) scores\_matrix = corpus\_emb @ query\_embs.T # (n\_pages, B) sorted\_indices = scores\_matrix.argsort(dim=0, descending=True).T.cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] return await asyncio.to\_thread(\_gpu\_work) batcher = DynamicBatcher( process\_fn=siglip\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_siglip\_batcher = batcher return \_siglip\_batcher # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── def \_batches(items: list, batch\_size: int): """Yield successive fixed-size batches from a list.""" for start in range(0, len(items), batch\_size): yield items\[start : start + batch\_size\] def \_load\_image\_sync(f: File) -> PILImage.Image: """Blocking download + decode. Intended to be called from a thread pool.""" with f.open\_sync("rb") as fh: data = fh.read() return PILImage.open(BytesIO(data)).convert("RGB") async def \_load\_image(f: File) -> PILImage.Image: """Download and decode a page image in a thread-pool worker. asyncio.to\_thread runs \_load\_image\_sync in a real OS thread so that blocking network I/O can overlap with GPU-bound forward passes when images are pre-submitted via loop.run\_in\_executor before the GPU kernel. """ return await asyncio.to\_thread(\_load\_image\_sync, f) async def \_load\_npz(index\_file: File) -> np.lib.npyio.NpzFile: """Download an index File to a local temp path and open with np.load.""" with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp: async with index\_file.open("rb") as fh: tmp.write(bytes(await fh.read())) return np.load(tmp.name) def \_dcg(relevances: list\[int\]) -> float: return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances)) # ───────────────────────────────────────────────────────────────────────────── # Tasks — data loading # ───────────────────────────────────────────────────────────────────────────── @driver.task(cache="auto", retries=3) async def load\_vidore\_pages(subset: str = "docvqa", max\_pages: int = 200) -> PageDataset: """ Load a ViDoRe benchmark subset and store page images in Flyte's blob store. Supports two dataset formats: Legacy (subsampled) — single 'test' split with one row per (query, page) pair; fields: image, query, image\_filename. streaming=True reads only the rows requested via islice — no full-shard download. Datasets: vidore/docvqa\_test\_subsampled, vidore/infovqa\_test\_subsampled V3 — separate corpus / queries / qrels splits following the BEIR retrieval benchmark format. corpus contains page images; queries contains question text; qrels maps query IDs to relevant corpus page IDs (many-to-many). Datasets: vidore/vidore\_v3\_finance\_en (~2 942 pages, 1 854 queries) The first call uploads page images to Flyte's blob store and caches the PageDataset; every subsequent call with the same arguments returns the cached result instantly. retries=3 guards against transient HuggingFace network failures. Available subsets: "docvqa", "infovqa", "vidore\_v3\_finance\_en" """ from datasets import load\_dataset subset\_map = { "docvqa": "vidore/docvqa\_test\_subsampled", "infovqa": "vidore/infovqa\_test\_subsampled", "vidore\_v3\_finance\_en": "vidore/vidore\_v3\_finance\_en", } dataset\_name = subset\_map.get(subset, f"vidore/{subset}\_test\_subsampled") # V3 datasets ship with separate corpus / queries / qrels splits. \_V3\_SUBSETS = {"vidore\_v3\_finance\_en"} if subset in \_V3\_SUBSETS: # ── V3 format ───────────────────────────────────────────────────────── # corpus / queries / qrels are HuggingFace configs (name=), not splits. # corpus uses streaming=True so images are decoded one at a time — # loading all 2 942 rows eagerly would hold gigabytes of PIL images in # the driver's RAM simultaneously. qrels and queries are text-only and # small enough to load fully into memory. corpus\_ds = load\_dataset(dataset\_name, name="corpus", split="test", streaming=True) qrels\_ds = load\_dataset(dataset\_name, name="qrels", split="test") queries\_ds = load\_dataset(dataset\_name, name="queries", split="test") # Normalise field names — V3 follows BEIR convention (hyphenated ids). def \_col(ds, \*candidates): cols = set(ds.column\_names) for c in candidates: if c in cols: return c raise KeyError(f"None of {candidates} found in columns {cols}") corpus\_id\_col = \_col(corpus\_ds, "corpus-id", "corpus\_id", "id", "\_id") query\_id\_col = \_col(queries\_ds, "query-id", "query\_id", "id", "\_id") query\_text\_col = \_col(queries\_ds, "query", "text") qrel\_qid\_col = \_col(qrels\_ds, "query-id", "query\_id") qrel\_cid\_col = \_col(qrels\_ds, "corpus-id", "corpus\_id") # Slice corpus to max\_pages, upload each image to Flyte blob store. page\_ids: list\[str\] = \[\] page\_files: list\[File\] = \[\] corpus\_id\_to\_page\_id: dict\[str, str\] = {} for i, row in enumerate(islice(corpus\_ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue cid = str(row\[corpus\_id\_col\]) page\_id = f"{subset}\_{i:04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) corpus\_id\_to\_page\_id\[cid\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) # Build query\_id → relevant page\_id from qrels (first match wins). # Only keep relevance judgements whose corpus page is in our slice. qrel\_map: dict\[str, str\] = {} for row in qrels\_ds: qid = str(row\[qrel\_qid\_col\]) cid = str(row\[qrel\_cid\_col\]) if cid in corpus\_id\_to\_page\_id and qid not in qrel\_map: qrel\_map\[qid\] = corpus\_id\_to\_page\_id\[cid\] # Collect queries that have at least one relevant page in our slice. queries: list\[PageQuery\] = \[\] for row in queries\_ds: qid = str(row\[query\_id\_col\]) if qid not in qrel\_map: continue queries.append( PageQuery( query\_id=qid, text=str(row\[query\_text\_col\]), relevant\_page\_id=qrel\_map\[qid\], ) ) else: # ── Legacy format ───────────────────────────────────────────────────── # Single 'test' split with one row per (query, page) pair. ds = load\_dataset(dataset\_name, split="test", streaming=True) page\_ids = \[\] page\_files = \[\] queries = \[\] seen\_pages: dict\[str, str\] = {} # image\_filename → page\_id for i, row in enumerate(islice(ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue filename: str = row.get("image\_filename") or f"page\_{i}" query\_text: str = row.get("query", "") if not query\_text: continue # Each unique page is uploaded exactly once; multiple queries may # share the same page (same image\_filename). if filename not in seen\_pages: page\_id = f"{subset}\_{len(page\_ids):04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) seen\_pages\[filename\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) else: page\_id = seen\_pages\[filename\] queries.append( PageQuery( query\_id=f"q{i:04d}", text=query\_text, relevant\_page\_id=page\_id, ) ) print(f"Loaded {len(page\_ids)} unique pages, {len(queries)} queries", flush=True) return PageDataset(page\_ids=page\_ids, page\_files=page\_files, queries=queries) # ───────────────────────────────────────────────────────────────────────────── # Tasks — indexing # ───────────────────────────────────────────────────────────────────────────── @colpali\_indexer.task(cache="auto", retries=2) async def index\_colpali(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with ColPali-v1.2 and save the multi-vector index. ColPali skips OCR entirely. It feeds the raw page image into PaliGemma (a vision-language model) and produces one embedding vector per image patch — roughly 1,024 patches per page, each of dimension 128. \_colpali\_model() is an lru\_cache'd loader. On a cold container, it downloads and loads the model once. On a warm container (kept alive by ReusePolicy), it returns the already-loaded model instantly from cache — no repeated ~7 GB download. The index is stored as a .npz file in Flyte's blob store: embeddings — float32, shape (n\_pages, n\_patches, dim) page\_ids — matching page ID strings cache="auto" + retries=2: the result is stored permanently on success; transient failures (e.g. HuggingFace rate limits) are retried twice. """ import torch model, processor, device = \_colpali\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 4)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor.process\_images(images) inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no\_grad(): emb = model(\*\*inputs) # (batch, n\_patches, dim) all\_embeddings.append(emb.cpu().float().numpy()) print(f"ColPali: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, n\_patches, dim) out\_path = os.path.join(tempfile.gettempdir(), "colpali\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @siglip\_indexer.task(cache="auto", retries=2) async def index\_siglip(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with SigLIP SO400M and save the single-vector index. SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss instead of softmax — avoiding the normalisation bottleneck that limits CLIP's scalability. Produces one global embedding per page. \_siglip\_model() caches the model across warm container reuses. The index is stored as a .npz file: embeddings — float32, shape (n\_pages, dim), L2-normalised page\_ids — matching page ID strings """ import torch model, processor, device = \_siglip\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 8)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor(images=images, return\_tensors="pt", padding=True).to(device) with torch.no\_grad(): outputs = model.vision\_model(\*\*inputs) emb = outputs.pooler\_output # (batch, dim) emb = emb / emb.norm(dim=-1, keepdim=True) # L2 normalise all\_embeddings.append(emb.cpu().float().numpy()) print(f"SigLIP: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, dim) out\_path = os.path.join(tempfile.gettempdir(), "siglip\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @ocr\_engine.task(cache="auto") async def extract\_page\_texts(page\_files: list\[File\]) -> list\[str\]: """ OCR every page with doctr on GPU to produce a text-only baseline. doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single callable predictor. Pages are downloaded in parallel then fed in batches of ocr\_batch\_size. asyncio.to\_thread keeps the event loop unblocked during GPU inference. Result structure: result.pages\[i\].blocks\[j\].lines\[k\].words\[l\].value Cached: the same corpus is OCR'd at most once across all experiments that use the OCR+BM25 backend. """ import gc predictor = \_ocr\_model() # Process in batches: download each batch just-in-time so only # ocr\_batch\_size images are in memory at once instead of all 2 000. ocr\_batch\_size = 8 total = len(page\_files) texts: list\[str\] = \[\] for start in range(0, total, ocr\_batch\_size): batch\_files = page\_files\[start : start + ocr\_batch\_size\] batch\_images = list( await asyncio.gather(\*\[asyncio.to\_thread(\_load\_image\_sync, f) for f in batch\_files\]) ) batch\_np = \[np.array(img) for img in batch\_images\] del batch\_images result = await asyncio.to\_thread(predictor, batch\_np) del batch\_np for page\_output in result.pages: texts.append( "\\n".join( " ".join(word.value for word in line.words) for block in page\_output.blocks for line in block.lines ) ) del result gc.collect() print(f"OCR: processed {min(start + ocr\_batch\_size, total)}/{total} pages", flush=True) return texts # ───────────────────────────────────────────────────────────────────────────── # Tasks — search # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment search\_colpali}} @colpali\_indexer.task async def search\_colpali( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher. MaxSim score for page p given query q: score(q, p) = Σ\_{t ∈ query tokens} max\_{j ∈ page patches} (q\_t · p\_j) Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_colpali invocations on the same warm container (concurrency=8) into a single GPU batch. This keeps the GPU saturated rather than running one small batch per caller. The batcher's process\_fn runs GPU work in asyncio.to\_thread, so the aggregation loop stays live while the GPU encodes and scores. """ batcher = await \_get\_colpali\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] # {{/docs-fragment search\_colpali}} @siglip\_indexer.task async def search\_siglip( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using SigLIP cosine similarity via DynamicBatcher. Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_siglip invocations on the same warm container (concurrency=3) into a single GPU batch. SigLIP's single-vector embeddings make full vectorisation safe — the scores matrix (n\_pages x n\_queries) is small enough to materialise in one GPU call regardless of batch size. """ batcher = await \_get\_siglip\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] @driver.task async def search\_bm25( page\_texts: list\[str\], page\_ids: list\[str\], queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using BM25 over OCR'd text. The standard keyword-based baseline. No GPU required; strong on text-dense pages, weak on visual content that Tesseract cannot read. """ tokenized = \[text.lower().split() for text in page\_texts\] bm25 = BM25Okapi(tokenized) results: list\[RetrievalResult\] = \[\] for q in queries: scores = bm25.get\_scores(q.text.lower().split()) ranked = sorted(range(len(page\_ids)), key=lambda i: -scores\[i\])\[:top\_k\] results.append( RetrievalResult( query\_id=q.query\_id, ranked\_page\_ids=\[page\_ids\[i\] for i in ranked\], ) ) return results # ───────────────────────────────────────────────────────────────────────────── # Tasks — evaluation # ───────────────────────────────────────────────────────────────────────────── @driver.task async def evaluate( results: list\[RetrievalResult\], ground\_truth: list\[PageQuery\], k: int, ) -> Metrics: """ Compute Recall@K, NDCG@K, and MRR for a single retrieval model. Recall@K — was the correct page in the top-K results? NDCG@K — normalised discounted cumulative gain; rewards earlier hits. MRR — mean reciprocal rank of the first correct result. All three are averaged over all queries. Higher is better. """ gt\_map = {q.query\_id: q.relevant\_page\_id for q in ground\_truth} recall\_vals, ndcg\_vals, mrr\_vals = \[\], \[\], \[\] for r in results: relevant = gt\_map.get(r.query\_id, "") top = r.ranked\_page\_ids\[:k\] recall\_vals.append(1.0 if relevant in top else 0.0) rels = \[1 if pid == relevant else 0 for pid in top\] idcg = \_dcg(\[1\]) # ideal: correct page at rank 1 ndcg\_vals.append(\_dcg(rels) / idcg if idcg > 0 else 0.0) rr = 0.0 for rank, pid in enumerate(r.ranked\_page\_ids, start=1): if pid == relevant: rr = 1.0 / rank break mrr\_vals.append(rr) return Metrics( recall\_at\_k=float(np.mean(recall\_vals)), ndcg\_at\_k=float(np.mean(ndcg\_vals)), mrr=float(np.mean(mrr\_vals)), k=k, ) # ───────────────────────────────────────────────────────────────────────────── # Tasks — report # ───────────────────────────────────────────────────────────────────────────── @driver.task(report=True) async def generate\_report(report: ComparisonReport) -> None: """ Emit an interactive HTML report visible in the Flyte UI. report=True marks this task as a reporting task. Flyte renders the HTML returned via flyte.report.replace.aio() directly in the execution detail page — no separate dashboard or export step required. The report contains: - Summary cards: experiment count, best model, best Recall@K. - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment. - Ranked results table with all three metrics. """ sorted\_results = sorted(report.results, key=lambda r: -r.metrics.recall\_at\_k) best = sorted\_results\[0\] labels = \[r.config.name for r in sorted\_results\] recall\_vals = \[r.metrics.recall\_at\_k for r in sorted\_results\] ndcg\_vals = \[r.metrics.ndcg\_at\_k for r in sorted\_results\] mrr\_vals = \[r.metrics.mrr for r in sorted\_results\] table\_rows = "".join( f""" {r.config.name} {r.config.model.value} {r.metrics.recall\_at\_k:.3f} {r.metrics.ndcg\_at\_k:.3f} {r.metrics.mrr:.3f} {r.metrics.k} """ for r in sorted\_results ) html = f""" Visual Document Retrieval — Results

Visual Document Retrieval — Experiment Comparison

ViDoRe benchmark · {len(report.results)} experiment(s)

{len(report.results)}
Experiments
{best.config.name}
Best by Recall@K
{best.metrics.recall\_at\_k:.3f}
Best Recall@{best.metrics.k}
{best.metrics.ndcg\_at\_k:.3f}
Best NDCG@{best.metrics.k}
{best.metrics.mrr:.3f}
Best MRR

Metrics by Experiment

Ranked Results

{table\_rows}
ExperimentModel Recall@KNDCG@KMRRK
""" await flyte.report.replace.aio(html) await flyte.report.flush.aio() # ───────────────────────────────────────────────────────────────────────────── # Experiment orchestration # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment run\_experiment}} @driver.task async def run\_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult: """ End-to-end retrieval pipeline for a single ExperimentConfig. Flyte v2's dynamic execution means this driver task can call GPU tasks (index\_colpali, search\_colpali) based on the runtime value of config.model — no static DAG wiring required. The if/elif is plain Python; Flyte schedules the selected sub-tasks on the appropriate environment. Caching: two experiments that share the same model and corpus (e.g. ColPali at top\_k=5 and top\_k=10) will hit the same cached index. GPU work is paid at most once per (model, corpus) pair across all experiments. Search queries are sharded into chunks of SEARCH\_SHARD\_SIZE and dispatched as concurrent task invocations. All shards land on the single warm container (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the GPU saturated throughout search rather than processing one large sequential batch from a single caller. flyte.group wraps each experiment in a named span in the Flyte UI, making it easy to compare latencies and drill into individual runs. """ SEARCH\_SHARD\_SIZE = 256 with flyte.group(config.name): if config.model == RetrievalModel.COLPALI: index\_file = await index\_colpali(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_colpali(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] elif config.model == RetrievalModel.SIGLIP: index\_file = await index\_siglip(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_siglip(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] else: # RetrievalModel.OCR\_BM25 page\_texts = await extract\_page\_texts(dataset.page\_files) results = await search\_bm25(page\_texts, dataset.page\_ids, dataset.queries, config.top\_k) metrics = await evaluate(results, dataset.queries, config.top\_k) return ExperimentResult(config=config, metrics=metrics) # {{/docs-fragment run\_experiment}} # {{docs-fragment compare\_experiments}} @driver.task async def compare\_experiments( configs: list\[ExperimentConfig\], subset: str = "docvqa", max\_pages: int = 200, ) -> ComparisonReport: """ Fan out over all experiment configs and return a ranked comparison table. The dataset is loaded once and shared across all experiments. Each config runs as a concurrent Flyte task via asyncio.gather. Experiments that share a model reuse the cached index — you only pay GPU time for new work. On completion, generate\_report emits an interactive Chart.js HTML report visible directly in the Flyte execution detail page. Default dataset: vidore\_v3\_finance\_en (~2 942 corpus pages, 1 854 queries) with max\_pages=2 000 to exercise the GPU pipeline at scale. """ dataset = await load\_vidore\_pages(subset=subset, max\_pages=max\_pages) # All experiments launch concurrently. Shared cached outputs (same model, # same corpus) are served from cache rather than recomputed. experiment\_coros = \[run\_experiment(config=cfg, dataset=dataset) for cfg in configs\] results: list\[ExperimentResult\] = list(await asyncio.gather(\*experiment\_coros)) report = ComparisonReport(results=results) print(report.summary()) best = report.best\_by("recall\_at\_k") print(f"\\nBest by Recall@{best.metrics.k}: {best.config.name}") # Emit the interactive HTML report in the Flyte UI. await generate\_report(report) return report # {{/docs-fragment compare\_experiments}} # ───────────────────────────────────────────────────────────────────────────── # Entry point # ───────────────────────────────────────────────────────────────────────────── if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Define the experiment grid. Each ExperimentConfig is one point in the # design space. Adding a new model or varying top\_k is one line here — # no task code changes required. # # ColPali appears twice with different top\_k values. The cache ensures # index\_colpali runs only once and both experiments share that result. # {{docs-fragment grid}} configs = \[\ ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top\_k=5),\ ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top\_k=10),\ ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top\_k=5),\ ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR\_BM25, top\_k=5),\ \] # {{/docs-fragment grid}} run = flyte.with\_runcontext().run( compare\_experiments, configs=configs, subset="vidore\_v3\_finance\_en", max\_pages=2000, ) print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval\_eval.py\* The corpus, queries, retrieval results, and metrics are likewise typed. Page images are stored as \`flyte.io.File\` handles in blob storage, so tasks read images directly rather than re-fetching over HTTP. \`\`\` # /// script # requires-python = ">=3.12" # dependencies = \[\ # "colpali-engine>=0.3.1",\ # "transformers>=4.41",\ # "sentencepiece>=0.2",\ # "torch>=2.0",\ # "pillow>=10",\ # "datasets>=2.18",\ # "rank-bm25>=0.2",\ # "numpy>=1.26",\ # "python-doctr\[torch\]>=0.8",\ # "pydantic>=2.0",\ # "flyte>=2.0.0",\ # \] # /// """ Multimodal Retrieval Evaluation Pipeline This tutorial is an experiment framework for benchmarking visual document retrieval approaches on the ViDoRe benchmark. Each experiment is defined by an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and returns a ranked comparison table with an interactive HTML report. The corpus is a set of PDF page images; queries are plain-text questions. Each retrieval method must find the page that answers each question — no text is provided to the model, only the raw image. ColPali-v1.2 — patch-level multi-vector embeddings from a VLM (PaliGemma). No OCR. The model produces one vector per image patch (~1024 per page). MaxSim late-interaction scoring finds the best matching patch for each query token. SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP successor. One matrix multiply per query; fast and effective but a single vector cannot localise fine-grained regions. OCR + BM25 — text-only baseline. doctr (GPU OCR) extracts text in batches, BM25 matches keywords. Strong on text-dense pages; fails on charts, tables, and figures where content is visual. """ import asyncio import enum import json import math import os import tempfile from functools import lru\_cache from io import BytesIO from itertools import islice import numpy as np from PIL import Image as PILImage from pydantic import BaseModel from rank\_bm25 import BM25Okapi from extras import DynamicBatcher import flyte import flyte.report from flyte.io import File # ───────────────────────────────────────────────────────────────────────────── # Environments # ───────────────────────────────────────────────────────────────────────────── # One Docker image for all tasks. The PEP 723 header defines Python deps. # ca-certificates is required for HTTPS calls to HuggingFace and blob stores. # {{docs-fragment image}} image = ( flyte.Image.from\_uv\_script(\_\_file\_\_, name="vidore-eval-v2") .with\_apt\_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0") # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy. # Without it every reusable container exits with StartError (exit code 128). .with\_pip\_packages("unionai-reuse>=0.1.11") ) # {{/docs-fragment image}} # GPU environment for ColPali image encoding and search. # # ReusePolicy keeps up to 3 warm GPU containers alive between task calls. # Without it, every task invocation cold-starts a new container and downloads # ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model # weights already loaded into VRAM — is reused for the next task dispatch. # # replicas=1 single warm container — all concurrent shard calls land # here so they share one DynamicBatcher process # concurrency=8 up to 8 query-shard tasks run simultaneously on the # container, all feeding the same DynamicBatcher queue # idle\_ttl=120 keep alive 2 min after the last task finishes # scaledown\_ttl=60 scale to zero after 1 min of complete inactivity # {{docs-fragment envs}} colpali\_indexer = flyte.TaskEnvironment( name="vidore-colpali-indexer", image=image, resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for SigLIP image encoding and search. # # Separate from the ColPali environment so each model's warm containers # are managed independently — ColPali and SigLIP experiments can scale # without contending for the same pool of reusable containers. siglip\_indexer = flyte.TaskEnvironment( name="vidore-siglip-indexer", image=image, resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition) # in batches on GPU — much faster than CPU Tesseract. # No ReusePolicy needed: the result is cached, so this task runs at most once. ocr\_engine = flyte.TaskEnvironment( name="vidore-ocr-engine", image=image, resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1), ) # Driver: orchestration, BM25 search, evaluation, and reporting. # depends\_on ensures the shared Docker image is built before all environments # try to schedule tasks. driver = flyte.TaskEnvironment( name="vidore-driver", image=image, resources=flyte.Resources(cpu=2, memory="12Gi"), depends\_on=\[colpali\_indexer, siglip\_indexer, ocr\_engine\], ) # {{/docs-fragment envs}} # ───────────────────────────────────────────────────────────────────────────── # Configuration types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment config\_types}} class RetrievalModel(str, enum.Enum): """Retrieval backend to evaluate.""" COLPALI = "colpali-v1.2" # multi-vector patch embeddings, MaxSim SIGLIP = "siglip-so400m" # single-vector global embedding, cosine sim OCR\_BM25 = "ocr+bm25" # text extracted by Tesseract, ranked by BM25 class ExperimentConfig(BaseModel): """ All knobs for one retrieval experiment. Passed as a typed Flyte input. Because ExperimentConfig is a Pydantic model, Flyte serialises it alongside every task output — so you can always reconstruct which config produced which metric without maintaining a separate log. """ name: str # human-readable label shown in the comparison table model: RetrievalModel top\_k: int = 5 # number of pages to retrieve per query # {{/docs-fragment config\_types}} # ───────────────────────────────────────────────────────────────────────────── # Data types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment data\_types}} class PageQuery(BaseModel): """One retrieval query with its ground-truth page.""" query\_id: str text: str # e.g. "What was revenue growth in Q3?" relevant\_page\_id: str # one correct page per query class PageDataset(BaseModel): """ A corpus of document page images paired with text queries. page\_ids: unique page identifiers (derived from ViDoRe image filenames). page\_files: the same pages stored in Flyte's blob store as JPEG File handles. Tasks read images directly from here; no live HTTP. queries: text questions with ground-truth page IDs for evaluation. """ page\_ids: list\[str\] page\_files: list\[File\] queries: list\[PageQuery\] class Config: arbitrary\_types\_allowed = True class RetrievalResult(BaseModel): query\_id: str ranked\_page\_ids: list\[str\] # ordered best → worst class Metrics(BaseModel): recall\_at\_k: float ndcg\_at\_k: float mrr: float k: int class ExperimentResult(BaseModel): config: ExperimentConfig metrics: Metrics # {{/docs-fragment data\_types}} class ComparisonReport(BaseModel): results: list\[ExperimentResult\] def best\_by(self, metric: str = "recall\_at\_k") -> ExperimentResult: return max(self.results, key=lambda r: getattr(r.metrics, metric)) def summary(self) -> str: header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}" sep = "─" \* len(header) rows = \[header, sep\] for r in sorted(self.results, key=lambda x: -x.metrics.recall\_at\_k): rows.append( f"{r.config.name:<30} " f"{r.config.model.value:<18} " f"{r.metrics.recall\_at\_k:>10.3f} " f"{r.metrics.ndcg\_at\_k:>8.3f} " f"{r.metrics.mrr:>7.3f}" ) return "\\n".join(rows) # ───────────────────────────────────────────────────────────────────────────── # Cached model loaders # ───────────────────────────────────────────────────────────────────────────── # These functions are at module level so they are shared across all tasks that # run on the same warm container (via ReusePolicy). lru\_cache(maxsize=1) means # the model is loaded from disk/HuggingFace exactly once per container process # and kept in GPU memory for every subsequent task dispatch to that container. @lru\_cache(maxsize=1) def \_colpali\_model(): """Load ColPali-v1.2 into GPU memory and cache the result. device\_map= is the correct loading pattern for ColPali's PaliGemma backbone; it handles weight placement via accelerate. torch.compile is skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross- invocation batching is the primary GPU utilisation mechanism. """ import torch from colpali\_engine.models import ColPali, ColPaliProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = ColPali.from\_pretrained( "vidore/colpali-v1.2", torch\_dtype=torch.bfloat16, device\_map=device, ) processor = ColPaliProcessor.from\_pretrained("vidore/colpali-v1.2") return model, processor, device @lru\_cache(maxsize=1) def \_siglip\_model(): """Load SigLIP SO400M into GPU memory, compile it, and cache the result. torch.compile (mode="reduce-overhead") fuses the vision and text encoder transformer layers into optimised CUDA kernels. As with ColPali, the compilation overhead is paid once per warm container lifetime. """ import torch from transformers import AutoModel, AutoProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = AutoModel.from\_pretrained("google/siglip-so400m-patch14-224").to(device) if device == "cuda": model = torch.compile(model, mode="reduce-overhead") processor = AutoProcessor.from\_pretrained("google/siglip-so400m-patch14-224") return model, processor, device @lru\_cache(maxsize=1) def \_ocr\_model(): """Load the doctr OCR predictor onto GPU and cache it. doctr's ocr\_predictor bundles a detection model (DBNet) and a recognition model (CRNN/SAR) into a single callable. pretrained=True downloads both model weights from doctr's model zoo on first use. """ import torch from doctr.models import ocr\_predictor predictor = ocr\_predictor(pretrained=True) if torch.cuda.is\_available(): predictor = predictor.cuda() return predictor # ───────────────────────────────────────────────────────────────────────────── # Search batcher singletons # ───────────────────────────────────────────────────────────────────────────── # One DynamicBatcher per model, shared across all concurrent search task # invocations on the same warm container (concurrency=3). Queries from every # concurrent caller are aggregated into a single GPU batch, maximizing # throughput compared to each invocation running its own forward pass. # # Initialised lazily on the first search call via double-checked locking and # lives for the container's lifetime. The process\_fn runs GPU work via # asyncio.to\_thread so the aggregation loop can continue collecting queries # from other callers while the GPU processes the current batch. # # File is not hashable so alru\_cache cannot be used here; module-level state # with asyncio.Lock is the correct pattern. # # Assumption: index\_colpali/index\_siglip use cache="auto", so the same corpus # always produces the same index File across all callers on this container. If # the index file ever changed between calls, the batcher would silently continue # using the corpus embeddings loaded from the first call. \_colpali\_batcher: DynamicBatcher | None = None \_colpali\_batcher\_lock = asyncio.Lock() \_siglip\_batcher: DynamicBatcher | None = None \_siglip\_batcher\_lock = asyncio.Lock() async def \_get\_colpali\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level ColPali search batcher, creating it on first call.""" global \_colpali\_batcher if \_colpali\_batcher is not None: return \_colpali\_batcher async with \_colpali\_batcher\_lock: if \_colpali\_batcher is not None: return \_colpali\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, n\_patches, dim) index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_colpali\_model() corpus\_emb = corpus\_emb.to(device, dtype=torch.float32) async def colpali\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: query\_inputs = processor.process\_queries(\[q.text for q in batch\]) query\_inputs = {k: v.to(device) for k, v in query\_inputs.items()} with torch.no\_grad(): query\_embs = model(\*\*query\_inputs).float() # (B, T, D) query\_chunk = 8 n\_pages = corpus\_emb.shape\[0\] all\_scores = torch.empty(len(batch), n\_pages, device=device) for start in range(0, len(batch), query\_chunk): chunk = query\_embs\[start : start + query\_chunk\] all\_scores\[start : start + query\_chunk\] = ( torch.einsum("ctd,pjd->ctpj", chunk, corpus\_emb) .max(dim=3).values .sum(dim=1) ) sorted\_indices = all\_scores.argsort(dim=1, descending=True).cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] # Run GPU work in a thread so the event loop — and the batcher's # aggregation loop — remain unblocked while the GPU is busy. return await asyncio.to\_thread(\_gpu\_work) batcher: DynamicBatcher\[PageQuery, list\[str\]\] = DynamicBatcher( process\_fn=colpali\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_colpali\_batcher = batcher return \_colpali\_batcher async def \_get\_siglip\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level SigLIP search batcher, creating it on first call.""" global \_siglip\_batcher if \_siglip\_batcher is not None: return \_siglip\_batcher async with \_siglip\_batcher\_lock: if \_siglip\_batcher is not None: return \_siglip\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, dim), L2-normalised index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_siglip\_model() corpus\_emb = corpus\_emb.to(device) async def siglip\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: text\_inputs = processor( text=\[q.text for q in batch\], return\_tensors="pt", padding=True, truncation=True, ).to(device) with torch.no\_grad(): text\_out = model.text\_model(\*\*text\_inputs) query\_embs = text\_out.pooler\_output # (B, dim) query\_embs = query\_embs / query\_embs.norm(dim=-1, keepdim=True) scores\_matrix = corpus\_emb @ query\_embs.T # (n\_pages, B) sorted\_indices = scores\_matrix.argsort(dim=0, descending=True).T.cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] return await asyncio.to\_thread(\_gpu\_work) batcher = DynamicBatcher( process\_fn=siglip\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_siglip\_batcher = batcher return \_siglip\_batcher # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── def \_batches(items: list, batch\_size: int): """Yield successive fixed-size batches from a list.""" for start in range(0, len(items), batch\_size): yield items\[start : start + batch\_size\] def \_load\_image\_sync(f: File) -> PILImage.Image: """Blocking download + decode. Intended to be called from a thread pool.""" with f.open\_sync("rb") as fh: data = fh.read() return PILImage.open(BytesIO(data)).convert("RGB") async def \_load\_image(f: File) -> PILImage.Image: """Download and decode a page image in a thread-pool worker. asyncio.to\_thread runs \_load\_image\_sync in a real OS thread so that blocking network I/O can overlap with GPU-bound forward passes when images are pre-submitted via loop.run\_in\_executor before the GPU kernel. """ return await asyncio.to\_thread(\_load\_image\_sync, f) async def \_load\_npz(index\_file: File) -> np.lib.npyio.NpzFile: """Download an index File to a local temp path and open with np.load.""" with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp: async with index\_file.open("rb") as fh: tmp.write(bytes(await fh.read())) return np.load(tmp.name) def \_dcg(relevances: list\[int\]) -> float: return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances)) # ───────────────────────────────────────────────────────────────────────────── # Tasks — data loading # ───────────────────────────────────────────────────────────────────────────── @driver.task(cache="auto", retries=3) async def load\_vidore\_pages(subset: str = "docvqa", max\_pages: int = 200) -> PageDataset: """ Load a ViDoRe benchmark subset and store page images in Flyte's blob store. Supports two dataset formats: Legacy (subsampled) — single 'test' split with one row per (query, page) pair; fields: image, query, image\_filename. streaming=True reads only the rows requested via islice — no full-shard download. Datasets: vidore/docvqa\_test\_subsampled, vidore/infovqa\_test\_subsampled V3 — separate corpus / queries / qrels splits following the BEIR retrieval benchmark format. corpus contains page images; queries contains question text; qrels maps query IDs to relevant corpus page IDs (many-to-many). Datasets: vidore/vidore\_v3\_finance\_en (~2 942 pages, 1 854 queries) The first call uploads page images to Flyte's blob store and caches the PageDataset; every subsequent call with the same arguments returns the cached result instantly. retries=3 guards against transient HuggingFace network failures. Available subsets: "docvqa", "infovqa", "vidore\_v3\_finance\_en" """ from datasets import load\_dataset subset\_map = { "docvqa": "vidore/docvqa\_test\_subsampled", "infovqa": "vidore/infovqa\_test\_subsampled", "vidore\_v3\_finance\_en": "vidore/vidore\_v3\_finance\_en", } dataset\_name = subset\_map.get(subset, f"vidore/{subset}\_test\_subsampled") # V3 datasets ship with separate corpus / queries / qrels splits. \_V3\_SUBSETS = {"vidore\_v3\_finance\_en"} if subset in \_V3\_SUBSETS: # ── V3 format ───────────────────────────────────────────────────────── # corpus / queries / qrels are HuggingFace configs (name=), not splits. # corpus uses streaming=True so images are decoded one at a time — # loading all 2 942 rows eagerly would hold gigabytes of PIL images in # the driver's RAM simultaneously. qrels and queries are text-only and # small enough to load fully into memory. corpus\_ds = load\_dataset(dataset\_name, name="corpus", split="test", streaming=True) qrels\_ds = load\_dataset(dataset\_name, name="qrels", split="test") queries\_ds = load\_dataset(dataset\_name, name="queries", split="test") # Normalise field names — V3 follows BEIR convention (hyphenated ids). def \_col(ds, \*candidates): cols = set(ds.column\_names) for c in candidates: if c in cols: return c raise KeyError(f"None of {candidates} found in columns {cols}") corpus\_id\_col = \_col(corpus\_ds, "corpus-id", "corpus\_id", "id", "\_id") query\_id\_col = \_col(queries\_ds, "query-id", "query\_id", "id", "\_id") query\_text\_col = \_col(queries\_ds, "query", "text") qrel\_qid\_col = \_col(qrels\_ds, "query-id", "query\_id") qrel\_cid\_col = \_col(qrels\_ds, "corpus-id", "corpus\_id") # Slice corpus to max\_pages, upload each image to Flyte blob store. page\_ids: list\[str\] = \[\] page\_files: list\[File\] = \[\] corpus\_id\_to\_page\_id: dict\[str, str\] = {} for i, row in enumerate(islice(corpus\_ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue cid = str(row\[corpus\_id\_col\]) page\_id = f"{subset}\_{i:04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) corpus\_id\_to\_page\_id\[cid\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) # Build query\_id → relevant page\_id from qrels (first match wins). # Only keep relevance judgements whose corpus page is in our slice. qrel\_map: dict\[str, str\] = {} for row in qrels\_ds: qid = str(row\[qrel\_qid\_col\]) cid = str(row\[qrel\_cid\_col\]) if cid in corpus\_id\_to\_page\_id and qid not in qrel\_map: qrel\_map\[qid\] = corpus\_id\_to\_page\_id\[cid\] # Collect queries that have at least one relevant page in our slice. queries: list\[PageQuery\] = \[\] for row in queries\_ds: qid = str(row\[query\_id\_col\]) if qid not in qrel\_map: continue queries.append( PageQuery( query\_id=qid, text=str(row\[query\_text\_col\]), relevant\_page\_id=qrel\_map\[qid\], ) ) else: # ── Legacy format ───────────────────────────────────────────────────── # Single 'test' split with one row per (query, page) pair. ds = load\_dataset(dataset\_name, split="test", streaming=True) page\_ids = \[\] page\_files = \[\] queries = \[\] seen\_pages: dict\[str, str\] = {} # image\_filename → page\_id for i, row in enumerate(islice(ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue filename: str = row.get("image\_filename") or f"page\_{i}" query\_text: str = row.get("query", "") if not query\_text: continue # Each unique page is uploaded exactly once; multiple queries may # share the same page (same image\_filename). if filename not in seen\_pages: page\_id = f"{subset}\_{len(page\_ids):04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) seen\_pages\[filename\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) else: page\_id = seen\_pages\[filename\] queries.append( PageQuery( query\_id=f"q{i:04d}", text=query\_text, relevant\_page\_id=page\_id, ) ) print(f"Loaded {len(page\_ids)} unique pages, {len(queries)} queries", flush=True) return PageDataset(page\_ids=page\_ids, page\_files=page\_files, queries=queries) # ───────────────────────────────────────────────────────────────────────────── # Tasks — indexing # ───────────────────────────────────────────────────────────────────────────── @colpali\_indexer.task(cache="auto", retries=2) async def index\_colpali(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with ColPali-v1.2 and save the multi-vector index. ColPali skips OCR entirely. It feeds the raw page image into PaliGemma (a vision-language model) and produces one embedding vector per image patch — roughly 1,024 patches per page, each of dimension 128. \_colpali\_model() is an lru\_cache'd loader. On a cold container, it downloads and loads the model once. On a warm container (kept alive by ReusePolicy), it returns the already-loaded model instantly from cache — no repeated ~7 GB download. The index is stored as a .npz file in Flyte's blob store: embeddings — float32, shape (n\_pages, n\_patches, dim) page\_ids — matching page ID strings cache="auto" + retries=2: the result is stored permanently on success; transient failures (e.g. HuggingFace rate limits) are retried twice. """ import torch model, processor, device = \_colpali\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 4)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor.process\_images(images) inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no\_grad(): emb = model(\*\*inputs) # (batch, n\_patches, dim) all\_embeddings.append(emb.cpu().float().numpy()) print(f"ColPali: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, n\_patches, dim) out\_path = os.path.join(tempfile.gettempdir(), "colpali\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @siglip\_indexer.task(cache="auto", retries=2) async def index\_siglip(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with SigLIP SO400M and save the single-vector index. SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss instead of softmax — avoiding the normalisation bottleneck that limits CLIP's scalability. Produces one global embedding per page. \_siglip\_model() caches the model across warm container reuses. The index is stored as a .npz file: embeddings — float32, shape (n\_pages, dim), L2-normalised page\_ids — matching page ID strings """ import torch model, processor, device = \_siglip\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 8)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor(images=images, return\_tensors="pt", padding=True).to(device) with torch.no\_grad(): outputs = model.vision\_model(\*\*inputs) emb = outputs.pooler\_output # (batch, dim) emb = emb / emb.norm(dim=-1, keepdim=True) # L2 normalise all\_embeddings.append(emb.cpu().float().numpy()) print(f"SigLIP: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, dim) out\_path = os.path.join(tempfile.gettempdir(), "siglip\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @ocr\_engine.task(cache="auto") async def extract\_page\_texts(page\_files: list\[File\]) -> list\[str\]: """ OCR every page with doctr on GPU to produce a text-only baseline. doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single callable predictor. Pages are downloaded in parallel then fed in batches of ocr\_batch\_size. asyncio.to\_thread keeps the event loop unblocked during GPU inference. Result structure: result.pages\[i\].blocks\[j\].lines\[k\].words\[l\].value Cached: the same corpus is OCR'd at most once across all experiments that use the OCR+BM25 backend. """ import gc predictor = \_ocr\_model() # Process in batches: download each batch just-in-time so only # ocr\_batch\_size images are in memory at once instead of all 2 000. ocr\_batch\_size = 8 total = len(page\_files) texts: list\[str\] = \[\] for start in range(0, total, ocr\_batch\_size): batch\_files = page\_files\[start : start + ocr\_batch\_size\] batch\_images = list( await asyncio.gather(\*\[asyncio.to\_thread(\_load\_image\_sync, f) for f in batch\_files\]) ) batch\_np = \[np.array(img) for img in batch\_images\] del batch\_images result = await asyncio.to\_thread(predictor, batch\_np) del batch\_np for page\_output in result.pages: texts.append( "\\n".join( " ".join(word.value for word in line.words) for block in page\_output.blocks for line in block.lines ) ) del result gc.collect() print(f"OCR: processed {min(start + ocr\_batch\_size, total)}/{total} pages", flush=True) return texts # ───────────────────────────────────────────────────────────────────────────── # Tasks — search # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment search\_colpali}} @colpali\_indexer.task async def search\_colpali( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher. MaxSim score for page p given query q: score(q, p) = Σ\_{t ∈ query tokens} max\_{j ∈ page patches} (q\_t · p\_j) Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_colpali invocations on the same warm container (concurrency=8) into a single GPU batch. This keeps the GPU saturated rather than running one small batch per caller. The batcher's process\_fn runs GPU work in asyncio.to\_thread, so the aggregation loop stays live while the GPU encodes and scores. """ batcher = await \_get\_colpali\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] # {{/docs-fragment search\_colpali}} @siglip\_indexer.task async def search\_siglip( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using SigLIP cosine similarity via DynamicBatcher. Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_siglip invocations on the same warm container (concurrency=3) into a single GPU batch. SigLIP's single-vector embeddings make full vectorisation safe — the scores matrix (n\_pages x n\_queries) is small enough to materialise in one GPU call regardless of batch size. """ batcher = await \_get\_siglip\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] @driver.task async def search\_bm25( page\_texts: list\[str\], page\_ids: list\[str\], queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using BM25 over OCR'd text. The standard keyword-based baseline. No GPU required; strong on text-dense pages, weak on visual content that Tesseract cannot read. """ tokenized = \[text.lower().split() for text in page\_texts\] bm25 = BM25Okapi(tokenized) results: list\[RetrievalResult\] = \[\] for q in queries: scores = bm25.get\_scores(q.text.lower().split()) ranked = sorted(range(len(page\_ids)), key=lambda i: -scores\[i\])\[:top\_k\] results.append( RetrievalResult( query\_id=q.query\_id, ranked\_page\_ids=\[page\_ids\[i\] for i in ranked\], ) ) return results # ───────────────────────────────────────────────────────────────────────────── # Tasks — evaluation # ───────────────────────────────────────────────────────────────────────────── @driver.task async def evaluate( results: list\[RetrievalResult\], ground\_truth: list\[PageQuery\], k: int, ) -> Metrics: """ Compute Recall@K, NDCG@K, and MRR for a single retrieval model. Recall@K — was the correct page in the top-K results? NDCG@K — normalised discounted cumulative gain; rewards earlier hits. MRR — mean reciprocal rank of the first correct result. All three are averaged over all queries. Higher is better. """ gt\_map = {q.query\_id: q.relevant\_page\_id for q in ground\_truth} recall\_vals, ndcg\_vals, mrr\_vals = \[\], \[\], \[\] for r in results: relevant = gt\_map.get(r.query\_id, "") top = r.ranked\_page\_ids\[:k\] recall\_vals.append(1.0 if relevant in top else 0.0) rels = \[1 if pid == relevant else 0 for pid in top\] idcg = \_dcg(\[1\]) # ideal: correct page at rank 1 ndcg\_vals.append(\_dcg(rels) / idcg if idcg > 0 else 0.0) rr = 0.0 for rank, pid in enumerate(r.ranked\_page\_ids, start=1): if pid == relevant: rr = 1.0 / rank break mrr\_vals.append(rr) return Metrics( recall\_at\_k=float(np.mean(recall\_vals)), ndcg\_at\_k=float(np.mean(ndcg\_vals)), mrr=float(np.mean(mrr\_vals)), k=k, ) # ───────────────────────────────────────────────────────────────────────────── # Tasks — report # ───────────────────────────────────────────────────────────────────────────── @driver.task(report=True) async def generate\_report(report: ComparisonReport) -> None: """ Emit an interactive HTML report visible in the Flyte UI. report=True marks this task as a reporting task. Flyte renders the HTML returned via flyte.report.replace.aio() directly in the execution detail page — no separate dashboard or export step required. The report contains: - Summary cards: experiment count, best model, best Recall@K. - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment. - Ranked results table with all three metrics. """ sorted\_results = sorted(report.results, key=lambda r: -r.metrics.recall\_at\_k) best = sorted\_results\[0\] labels = \[r.config.name for r in sorted\_results\] recall\_vals = \[r.metrics.recall\_at\_k for r in sorted\_results\] ndcg\_vals = \[r.metrics.ndcg\_at\_k for r in sorted\_results\] mrr\_vals = \[r.metrics.mrr for r in sorted\_results\] table\_rows = "".join( f""" {r.config.name} {r.config.model.value} {r.metrics.recall\_at\_k:.3f} {r.metrics.ndcg\_at\_k:.3f} {r.metrics.mrr:.3f} {r.metrics.k} """ for r in sorted\_results ) html = f""" Visual Document Retrieval — Results

Visual Document Retrieval — Experiment Comparison

ViDoRe benchmark · {len(report.results)} experiment(s)

{len(report.results)}
Experiments
{best.config.name}
Best by Recall@K
{best.metrics.recall\_at\_k:.3f}
Best Recall@{best.metrics.k}
{best.metrics.ndcg\_at\_k:.3f}
Best NDCG@{best.metrics.k}
{best.metrics.mrr:.3f}
Best MRR

Metrics by Experiment

Ranked Results

{table\_rows}
ExperimentModel Recall@KNDCG@KMRRK
""" await flyte.report.replace.aio(html) await flyte.report.flush.aio() # ───────────────────────────────────────────────────────────────────────────── # Experiment orchestration # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment run\_experiment}} @driver.task async def run\_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult: """ End-to-end retrieval pipeline for a single ExperimentConfig. Flyte v2's dynamic execution means this driver task can call GPU tasks (index\_colpali, search\_colpali) based on the runtime value of config.model — no static DAG wiring required. The if/elif is plain Python; Flyte schedules the selected sub-tasks on the appropriate environment. Caching: two experiments that share the same model and corpus (e.g. ColPali at top\_k=5 and top\_k=10) will hit the same cached index. GPU work is paid at most once per (model, corpus) pair across all experiments. Search queries are sharded into chunks of SEARCH\_SHARD\_SIZE and dispatched as concurrent task invocations. All shards land on the single warm container (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the GPU saturated throughout search rather than processing one large sequential batch from a single caller. flyte.group wraps each experiment in a named span in the Flyte UI, making it easy to compare latencies and drill into individual runs. """ SEARCH\_SHARD\_SIZE = 256 with flyte.group(config.name): if config.model == RetrievalModel.COLPALI: index\_file = await index\_colpali(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_colpali(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] elif config.model == RetrievalModel.SIGLIP: index\_file = await index\_siglip(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_siglip(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] else: # RetrievalModel.OCR\_BM25 page\_texts = await extract\_page\_texts(dataset.page\_files) results = await search\_bm25(page\_texts, dataset.page\_ids, dataset.queries, config.top\_k) metrics = await evaluate(results, dataset.queries, config.top\_k) return ExperimentResult(config=config, metrics=metrics) # {{/docs-fragment run\_experiment}} # {{docs-fragment compare\_experiments}} @driver.task async def compare\_experiments( configs: list\[ExperimentConfig\], subset: str = "docvqa", max\_pages: int = 200, ) -> ComparisonReport: """ Fan out over all experiment configs and return a ranked comparison table. The dataset is loaded once and shared across all experiments. Each config runs as a concurrent Flyte task via asyncio.gather. Experiments that share a model reuse the cached index — you only pay GPU time for new work. On completion, generate\_report emits an interactive Chart.js HTML report visible directly in the Flyte execution detail page. Default dataset: vidore\_v3\_finance\_en (~2 942 corpus pages, 1 854 queries) with max\_pages=2 000 to exercise the GPU pipeline at scale. """ dataset = await load\_vidore\_pages(subset=subset, max\_pages=max\_pages) # All experiments launch concurrently. Shared cached outputs (same model, # same corpus) are served from cache rather than recomputed. experiment\_coros = \[run\_experiment(config=cfg, dataset=dataset) for cfg in configs\] results: list\[ExperimentResult\] = list(await asyncio.gather(\*experiment\_coros)) report = ComparisonReport(results=results) print(report.summary()) best = report.best\_by("recall\_at\_k") print(f"\\nBest by Recall@{best.metrics.k}: {best.config.name}") # Emit the interactive HTML report in the Flyte UI. await generate\_report(report) return report # {{/docs-fragment compare\_experiments}} # ───────────────────────────────────────────────────────────────────────────── # Entry point # ───────────────────────────────────────────────────────────────────────────── if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Define the experiment grid. Each ExperimentConfig is one point in the # design space. Adding a new model or varying top\_k is one line here — # no task code changes required. # # ColPali appears twice with different top\_k values. The cache ensures # index\_colpali runs only once and both experiments share that result. # {{docs-fragment grid}} configs = \[\ ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top\_k=5),\ ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top\_k=10),\ ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top\_k=5),\ ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR\_BM25, top\_k=5),\ \] # {{/docs-fragment grid}} run = flyte.with\_runcontext().run( compare\_experiments, configs=configs, subset="vidore\_v3\_finance\_en", max\_pages=2000, ) print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval\_eval.py\* ## Loading, indexing, and search \`load\_vidore\_pages\` downloads a ViDoRe subset and uploads each page image to blob storage (cached, with retries). Indexing tasks (\`index\_colpali\`, \`index\_siglip\`) encode every page into a \`.npz\` index, and the OCR task (\`extract\_page\_texts\`) produces the text baseline. These run on the GPU environments and are cached per corpus. Search uses the \`DynamicBatcher\` so queries from all concurrent search-task invocations on a warm container are merged into a single GPU batch: \`\`\` # /// script # requires-python = ">=3.12" # dependencies = \[\ # "colpali-engine>=0.3.1",\ # "transformers>=4.41",\ # "sentencepiece>=0.2",\ # "torch>=2.0",\ # "pillow>=10",\ # "datasets>=2.18",\ # "rank-bm25>=0.2",\ # "numpy>=1.26",\ # "python-doctr\[torch\]>=0.8",\ # "pydantic>=2.0",\ # "flyte>=2.0.0",\ # \] # /// """ Multimodal Retrieval Evaluation Pipeline This tutorial is an experiment framework for benchmarking visual document retrieval approaches on the ViDoRe benchmark. Each experiment is defined by an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and returns a ranked comparison table with an interactive HTML report. The corpus is a set of PDF page images; queries are plain-text questions. Each retrieval method must find the page that answers each question — no text is provided to the model, only the raw image. ColPali-v1.2 — patch-level multi-vector embeddings from a VLM (PaliGemma). No OCR. The model produces one vector per image patch (~1024 per page). MaxSim late-interaction scoring finds the best matching patch for each query token. SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP successor. One matrix multiply per query; fast and effective but a single vector cannot localise fine-grained regions. OCR + BM25 — text-only baseline. doctr (GPU OCR) extracts text in batches, BM25 matches keywords. Strong on text-dense pages; fails on charts, tables, and figures where content is visual. """ import asyncio import enum import json import math import os import tempfile from functools import lru\_cache from io import BytesIO from itertools import islice import numpy as np from PIL import Image as PILImage from pydantic import BaseModel from rank\_bm25 import BM25Okapi from extras import DynamicBatcher import flyte import flyte.report from flyte.io import File # ───────────────────────────────────────────────────────────────────────────── # Environments # ───────────────────────────────────────────────────────────────────────────── # One Docker image for all tasks. The PEP 723 header defines Python deps. # ca-certificates is required for HTTPS calls to HuggingFace and blob stores. # {{docs-fragment image}} image = ( flyte.Image.from\_uv\_script(\_\_file\_\_, name="vidore-eval-v2") .with\_apt\_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0") # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy. # Without it every reusable container exits with StartError (exit code 128). .with\_pip\_packages("unionai-reuse>=0.1.11") ) # {{/docs-fragment image}} # GPU environment for ColPali image encoding and search. # # ReusePolicy keeps up to 3 warm GPU containers alive between task calls. # Without it, every task invocation cold-starts a new container and downloads # ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model # weights already loaded into VRAM — is reused for the next task dispatch. # # replicas=1 single warm container — all concurrent shard calls land # here so they share one DynamicBatcher process # concurrency=8 up to 8 query-shard tasks run simultaneously on the # container, all feeding the same DynamicBatcher queue # idle\_ttl=120 keep alive 2 min after the last task finishes # scaledown\_ttl=60 scale to zero after 1 min of complete inactivity # {{docs-fragment envs}} colpali\_indexer = flyte.TaskEnvironment( name="vidore-colpali-indexer", image=image, resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for SigLIP image encoding and search. # # Separate from the ColPali environment so each model's warm containers # are managed independently — ColPali and SigLIP experiments can scale # without contending for the same pool of reusable containers. siglip\_indexer = flyte.TaskEnvironment( name="vidore-siglip-indexer", image=image, resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition) # in batches on GPU — much faster than CPU Tesseract. # No ReusePolicy needed: the result is cached, so this task runs at most once. ocr\_engine = flyte.TaskEnvironment( name="vidore-ocr-engine", image=image, resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1), ) # Driver: orchestration, BM25 search, evaluation, and reporting. # depends\_on ensures the shared Docker image is built before all environments # try to schedule tasks. driver = flyte.TaskEnvironment( name="vidore-driver", image=image, resources=flyte.Resources(cpu=2, memory="12Gi"), depends\_on=\[colpali\_indexer, siglip\_indexer, ocr\_engine\], ) # {{/docs-fragment envs}} # ───────────────────────────────────────────────────────────────────────────── # Configuration types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment config\_types}} class RetrievalModel(str, enum.Enum): """Retrieval backend to evaluate.""" COLPALI = "colpali-v1.2" # multi-vector patch embeddings, MaxSim SIGLIP = "siglip-so400m" # single-vector global embedding, cosine sim OCR\_BM25 = "ocr+bm25" # text extracted by Tesseract, ranked by BM25 class ExperimentConfig(BaseModel): """ All knobs for one retrieval experiment. Passed as a typed Flyte input. Because ExperimentConfig is a Pydantic model, Flyte serialises it alongside every task output — so you can always reconstruct which config produced which metric without maintaining a separate log. """ name: str # human-readable label shown in the comparison table model: RetrievalModel top\_k: int = 5 # number of pages to retrieve per query # {{/docs-fragment config\_types}} # ───────────────────────────────────────────────────────────────────────────── # Data types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment data\_types}} class PageQuery(BaseModel): """One retrieval query with its ground-truth page.""" query\_id: str text: str # e.g. "What was revenue growth in Q3?" relevant\_page\_id: str # one correct page per query class PageDataset(BaseModel): """ A corpus of document page images paired with text queries. page\_ids: unique page identifiers (derived from ViDoRe image filenames). page\_files: the same pages stored in Flyte's blob store as JPEG File handles. Tasks read images directly from here; no live HTTP. queries: text questions with ground-truth page IDs for evaluation. """ page\_ids: list\[str\] page\_files: list\[File\] queries: list\[PageQuery\] class Config: arbitrary\_types\_allowed = True class RetrievalResult(BaseModel): query\_id: str ranked\_page\_ids: list\[str\] # ordered best → worst class Metrics(BaseModel): recall\_at\_k: float ndcg\_at\_k: float mrr: float k: int class ExperimentResult(BaseModel): config: ExperimentConfig metrics: Metrics # {{/docs-fragment data\_types}} class ComparisonReport(BaseModel): results: list\[ExperimentResult\] def best\_by(self, metric: str = "recall\_at\_k") -> ExperimentResult: return max(self.results, key=lambda r: getattr(r.metrics, metric)) def summary(self) -> str: header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}" sep = "─" \* len(header) rows = \[header, sep\] for r in sorted(self.results, key=lambda x: -x.metrics.recall\_at\_k): rows.append( f"{r.config.name:<30} " f"{r.config.model.value:<18} " f"{r.metrics.recall\_at\_k:>10.3f} " f"{r.metrics.ndcg\_at\_k:>8.3f} " f"{r.metrics.mrr:>7.3f}" ) return "\\n".join(rows) # ───────────────────────────────────────────────────────────────────────────── # Cached model loaders # ───────────────────────────────────────────────────────────────────────────── # These functions are at module level so they are shared across all tasks that # run on the same warm container (via ReusePolicy). lru\_cache(maxsize=1) means # the model is loaded from disk/HuggingFace exactly once per container process # and kept in GPU memory for every subsequent task dispatch to that container. @lru\_cache(maxsize=1) def \_colpali\_model(): """Load ColPali-v1.2 into GPU memory and cache the result. device\_map= is the correct loading pattern for ColPali's PaliGemma backbone; it handles weight placement via accelerate. torch.compile is skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross- invocation batching is the primary GPU utilisation mechanism. """ import torch from colpali\_engine.models import ColPali, ColPaliProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = ColPali.from\_pretrained( "vidore/colpali-v1.2", torch\_dtype=torch.bfloat16, device\_map=device, ) processor = ColPaliProcessor.from\_pretrained("vidore/colpali-v1.2") return model, processor, device @lru\_cache(maxsize=1) def \_siglip\_model(): """Load SigLIP SO400M into GPU memory, compile it, and cache the result. torch.compile (mode="reduce-overhead") fuses the vision and text encoder transformer layers into optimised CUDA kernels. As with ColPali, the compilation overhead is paid once per warm container lifetime. """ import torch from transformers import AutoModel, AutoProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = AutoModel.from\_pretrained("google/siglip-so400m-patch14-224").to(device) if device == "cuda": model = torch.compile(model, mode="reduce-overhead") processor = AutoProcessor.from\_pretrained("google/siglip-so400m-patch14-224") return model, processor, device @lru\_cache(maxsize=1) def \_ocr\_model(): """Load the doctr OCR predictor onto GPU and cache it. doctr's ocr\_predictor bundles a detection model (DBNet) and a recognition model (CRNN/SAR) into a single callable. pretrained=True downloads both model weights from doctr's model zoo on first use. """ import torch from doctr.models import ocr\_predictor predictor = ocr\_predictor(pretrained=True) if torch.cuda.is\_available(): predictor = predictor.cuda() return predictor # ───────────────────────────────────────────────────────────────────────────── # Search batcher singletons # ───────────────────────────────────────────────────────────────────────────── # One DynamicBatcher per model, shared across all concurrent search task # invocations on the same warm container (concurrency=3). Queries from every # concurrent caller are aggregated into a single GPU batch, maximizing # throughput compared to each invocation running its own forward pass. # # Initialised lazily on the first search call via double-checked locking and # lives for the container's lifetime. The process\_fn runs GPU work via # asyncio.to\_thread so the aggregation loop can continue collecting queries # from other callers while the GPU processes the current batch. # # File is not hashable so alru\_cache cannot be used here; module-level state # with asyncio.Lock is the correct pattern. # # Assumption: index\_colpali/index\_siglip use cache="auto", so the same corpus # always produces the same index File across all callers on this container. If # the index file ever changed between calls, the batcher would silently continue # using the corpus embeddings loaded from the first call. \_colpali\_batcher: DynamicBatcher | None = None \_colpali\_batcher\_lock = asyncio.Lock() \_siglip\_batcher: DynamicBatcher | None = None \_siglip\_batcher\_lock = asyncio.Lock() async def \_get\_colpali\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level ColPali search batcher, creating it on first call.""" global \_colpali\_batcher if \_colpali\_batcher is not None: return \_colpali\_batcher async with \_colpali\_batcher\_lock: if \_colpali\_batcher is not None: return \_colpali\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, n\_patches, dim) index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_colpali\_model() corpus\_emb = corpus\_emb.to(device, dtype=torch.float32) async def colpali\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: query\_inputs = processor.process\_queries(\[q.text for q in batch\]) query\_inputs = {k: v.to(device) for k, v in query\_inputs.items()} with torch.no\_grad(): query\_embs = model(\*\*query\_inputs).float() # (B, T, D) query\_chunk = 8 n\_pages = corpus\_emb.shape\[0\] all\_scores = torch.empty(len(batch), n\_pages, device=device) for start in range(0, len(batch), query\_chunk): chunk = query\_embs\[start : start + query\_chunk\] all\_scores\[start : start + query\_chunk\] = ( torch.einsum("ctd,pjd->ctpj", chunk, corpus\_emb) .max(dim=3).values .sum(dim=1) ) sorted\_indices = all\_scores.argsort(dim=1, descending=True).cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] # Run GPU work in a thread so the event loop — and the batcher's # aggregation loop — remain unblocked while the GPU is busy. return await asyncio.to\_thread(\_gpu\_work) batcher: DynamicBatcher\[PageQuery, list\[str\]\] = DynamicBatcher( process\_fn=colpali\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_colpali\_batcher = batcher return \_colpali\_batcher async def \_get\_siglip\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level SigLIP search batcher, creating it on first call.""" global \_siglip\_batcher if \_siglip\_batcher is not None: return \_siglip\_batcher async with \_siglip\_batcher\_lock: if \_siglip\_batcher is not None: return \_siglip\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, dim), L2-normalised index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_siglip\_model() corpus\_emb = corpus\_emb.to(device) async def siglip\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: text\_inputs = processor( text=\[q.text for q in batch\], return\_tensors="pt", padding=True, truncation=True, ).to(device) with torch.no\_grad(): text\_out = model.text\_model(\*\*text\_inputs) query\_embs = text\_out.pooler\_output # (B, dim) query\_embs = query\_embs / query\_embs.norm(dim=-1, keepdim=True) scores\_matrix = corpus\_emb @ query\_embs.T # (n\_pages, B) sorted\_indices = scores\_matrix.argsort(dim=0, descending=True).T.cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] return await asyncio.to\_thread(\_gpu\_work) batcher = DynamicBatcher( process\_fn=siglip\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_siglip\_batcher = batcher return \_siglip\_batcher # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── def \_batches(items: list, batch\_size: int): """Yield successive fixed-size batches from a list.""" for start in range(0, len(items), batch\_size): yield items\[start : start + batch\_size\] def \_load\_image\_sync(f: File) -> PILImage.Image: """Blocking download + decode. Intended to be called from a thread pool.""" with f.open\_sync("rb") as fh: data = fh.read() return PILImage.open(BytesIO(data)).convert("RGB") async def \_load\_image(f: File) -> PILImage.Image: """Download and decode a page image in a thread-pool worker. asyncio.to\_thread runs \_load\_image\_sync in a real OS thread so that blocking network I/O can overlap with GPU-bound forward passes when images are pre-submitted via loop.run\_in\_executor before the GPU kernel. """ return await asyncio.to\_thread(\_load\_image\_sync, f) async def \_load\_npz(index\_file: File) -> np.lib.npyio.NpzFile: """Download an index File to a local temp path and open with np.load.""" with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp: async with index\_file.open("rb") as fh: tmp.write(bytes(await fh.read())) return np.load(tmp.name) def \_dcg(relevances: list\[int\]) -> float: return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances)) # ───────────────────────────────────────────────────────────────────────────── # Tasks — data loading # ───────────────────────────────────────────────────────────────────────────── @driver.task(cache="auto", retries=3) async def load\_vidore\_pages(subset: str = "docvqa", max\_pages: int = 200) -> PageDataset: """ Load a ViDoRe benchmark subset and store page images in Flyte's blob store. Supports two dataset formats: Legacy (subsampled) — single 'test' split with one row per (query, page) pair; fields: image, query, image\_filename. streaming=True reads only the rows requested via islice — no full-shard download. Datasets: vidore/docvqa\_test\_subsampled, vidore/infovqa\_test\_subsampled V3 — separate corpus / queries / qrels splits following the BEIR retrieval benchmark format. corpus contains page images; queries contains question text; qrels maps query IDs to relevant corpus page IDs (many-to-many). Datasets: vidore/vidore\_v3\_finance\_en (~2 942 pages, 1 854 queries) The first call uploads page images to Flyte's blob store and caches the PageDataset; every subsequent call with the same arguments returns the cached result instantly. retries=3 guards against transient HuggingFace network failures. Available subsets: "docvqa", "infovqa", "vidore\_v3\_finance\_en" """ from datasets import load\_dataset subset\_map = { "docvqa": "vidore/docvqa\_test\_subsampled", "infovqa": "vidore/infovqa\_test\_subsampled", "vidore\_v3\_finance\_en": "vidore/vidore\_v3\_finance\_en", } dataset\_name = subset\_map.get(subset, f"vidore/{subset}\_test\_subsampled") # V3 datasets ship with separate corpus / queries / qrels splits. \_V3\_SUBSETS = {"vidore\_v3\_finance\_en"} if subset in \_V3\_SUBSETS: # ── V3 format ───────────────────────────────────────────────────────── # corpus / queries / qrels are HuggingFace configs (name=), not splits. # corpus uses streaming=True so images are decoded one at a time — # loading all 2 942 rows eagerly would hold gigabytes of PIL images in # the driver's RAM simultaneously. qrels and queries are text-only and # small enough to load fully into memory. corpus\_ds = load\_dataset(dataset\_name, name="corpus", split="test", streaming=True) qrels\_ds = load\_dataset(dataset\_name, name="qrels", split="test") queries\_ds = load\_dataset(dataset\_name, name="queries", split="test") # Normalise field names — V3 follows BEIR convention (hyphenated ids). def \_col(ds, \*candidates): cols = set(ds.column\_names) for c in candidates: if c in cols: return c raise KeyError(f"None of {candidates} found in columns {cols}") corpus\_id\_col = \_col(corpus\_ds, "corpus-id", "corpus\_id", "id", "\_id") query\_id\_col = \_col(queries\_ds, "query-id", "query\_id", "id", "\_id") query\_text\_col = \_col(queries\_ds, "query", "text") qrel\_qid\_col = \_col(qrels\_ds, "query-id", "query\_id") qrel\_cid\_col = \_col(qrels\_ds, "corpus-id", "corpus\_id") # Slice corpus to max\_pages, upload each image to Flyte blob store. page\_ids: list\[str\] = \[\] page\_files: list\[File\] = \[\] corpus\_id\_to\_page\_id: dict\[str, str\] = {} for i, row in enumerate(islice(corpus\_ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue cid = str(row\[corpus\_id\_col\]) page\_id = f"{subset}\_{i:04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) corpus\_id\_to\_page\_id\[cid\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) # Build query\_id → relevant page\_id from qrels (first match wins). # Only keep relevance judgements whose corpus page is in our slice. qrel\_map: dict\[str, str\] = {} for row in qrels\_ds: qid = str(row\[qrel\_qid\_col\]) cid = str(row\[qrel\_cid\_col\]) if cid in corpus\_id\_to\_page\_id and qid not in qrel\_map: qrel\_map\[qid\] = corpus\_id\_to\_page\_id\[cid\] # Collect queries that have at least one relevant page in our slice. queries: list\[PageQuery\] = \[\] for row in queries\_ds: qid = str(row\[query\_id\_col\]) if qid not in qrel\_map: continue queries.append( PageQuery( query\_id=qid, text=str(row\[query\_text\_col\]), relevant\_page\_id=qrel\_map\[qid\], ) ) else: # ── Legacy format ───────────────────────────────────────────────────── # Single 'test' split with one row per (query, page) pair. ds = load\_dataset(dataset\_name, split="test", streaming=True) page\_ids = \[\] page\_files = \[\] queries = \[\] seen\_pages: dict\[str, str\] = {} # image\_filename → page\_id for i, row in enumerate(islice(ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue filename: str = row.get("image\_filename") or f"page\_{i}" query\_text: str = row.get("query", "") if not query\_text: continue # Each unique page is uploaded exactly once; multiple queries may # share the same page (same image\_filename). if filename not in seen\_pages: page\_id = f"{subset}\_{len(page\_ids):04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) seen\_pages\[filename\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) else: page\_id = seen\_pages\[filename\] queries.append( PageQuery( query\_id=f"q{i:04d}", text=query\_text, relevant\_page\_id=page\_id, ) ) print(f"Loaded {len(page\_ids)} unique pages, {len(queries)} queries", flush=True) return PageDataset(page\_ids=page\_ids, page\_files=page\_files, queries=queries) # ───────────────────────────────────────────────────────────────────────────── # Tasks — indexing # ───────────────────────────────────────────────────────────────────────────── @colpali\_indexer.task(cache="auto", retries=2) async def index\_colpali(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with ColPali-v1.2 and save the multi-vector index. ColPali skips OCR entirely. It feeds the raw page image into PaliGemma (a vision-language model) and produces one embedding vector per image patch — roughly 1,024 patches per page, each of dimension 128. \_colpali\_model() is an lru\_cache'd loader. On a cold container, it downloads and loads the model once. On a warm container (kept alive by ReusePolicy), it returns the already-loaded model instantly from cache — no repeated ~7 GB download. The index is stored as a .npz file in Flyte's blob store: embeddings — float32, shape (n\_pages, n\_patches, dim) page\_ids — matching page ID strings cache="auto" + retries=2: the result is stored permanently on success; transient failures (e.g. HuggingFace rate limits) are retried twice. """ import torch model, processor, device = \_colpali\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 4)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor.process\_images(images) inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no\_grad(): emb = model(\*\*inputs) # (batch, n\_patches, dim) all\_embeddings.append(emb.cpu().float().numpy()) print(f"ColPali: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, n\_patches, dim) out\_path = os.path.join(tempfile.gettempdir(), "colpali\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @siglip\_indexer.task(cache="auto", retries=2) async def index\_siglip(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with SigLIP SO400M and save the single-vector index. SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss instead of softmax — avoiding the normalisation bottleneck that limits CLIP's scalability. Produces one global embedding per page. \_siglip\_model() caches the model across warm container reuses. The index is stored as a .npz file: embeddings — float32, shape (n\_pages, dim), L2-normalised page\_ids — matching page ID strings """ import torch model, processor, device = \_siglip\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 8)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor(images=images, return\_tensors="pt", padding=True).to(device) with torch.no\_grad(): outputs = model.vision\_model(\*\*inputs) emb = outputs.pooler\_output # (batch, dim) emb = emb / emb.norm(dim=-1, keepdim=True) # L2 normalise all\_embeddings.append(emb.cpu().float().numpy()) print(f"SigLIP: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, dim) out\_path = os.path.join(tempfile.gettempdir(), "siglip\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @ocr\_engine.task(cache="auto") async def extract\_page\_texts(page\_files: list\[File\]) -> list\[str\]: """ OCR every page with doctr on GPU to produce a text-only baseline. doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single callable predictor. Pages are downloaded in parallel then fed in batches of ocr\_batch\_size. asyncio.to\_thread keeps the event loop unblocked during GPU inference. Result structure: result.pages\[i\].blocks\[j\].lines\[k\].words\[l\].value Cached: the same corpus is OCR'd at most once across all experiments that use the OCR+BM25 backend. """ import gc predictor = \_ocr\_model() # Process in batches: download each batch just-in-time so only # ocr\_batch\_size images are in memory at once instead of all 2 000. ocr\_batch\_size = 8 total = len(page\_files) texts: list\[str\] = \[\] for start in range(0, total, ocr\_batch\_size): batch\_files = page\_files\[start : start + ocr\_batch\_size\] batch\_images = list( await asyncio.gather(\*\[asyncio.to\_thread(\_load\_image\_sync, f) for f in batch\_files\]) ) batch\_np = \[np.array(img) for img in batch\_images\] del batch\_images result = await asyncio.to\_thread(predictor, batch\_np) del batch\_np for page\_output in result.pages: texts.append( "\\n".join( " ".join(word.value for word in line.words) for block in page\_output.blocks for line in block.lines ) ) del result gc.collect() print(f"OCR: processed {min(start + ocr\_batch\_size, total)}/{total} pages", flush=True) return texts # ───────────────────────────────────────────────────────────────────────────── # Tasks — search # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment search\_colpali}} @colpali\_indexer.task async def search\_colpali( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher. MaxSim score for page p given query q: score(q, p) = Σ\_{t ∈ query tokens} max\_{j ∈ page patches} (q\_t · p\_j) Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_colpali invocations on the same warm container (concurrency=8) into a single GPU batch. This keeps the GPU saturated rather than running one small batch per caller. The batcher's process\_fn runs GPU work in asyncio.to\_thread, so the aggregation loop stays live while the GPU encodes and scores. """ batcher = await \_get\_colpali\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] # {{/docs-fragment search\_colpali}} @siglip\_indexer.task async def search\_siglip( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using SigLIP cosine similarity via DynamicBatcher. Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_siglip invocations on the same warm container (concurrency=3) into a single GPU batch. SigLIP's single-vector embeddings make full vectorisation safe — the scores matrix (n\_pages x n\_queries) is small enough to materialise in one GPU call regardless of batch size. """ batcher = await \_get\_siglip\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] @driver.task async def search\_bm25( page\_texts: list\[str\], page\_ids: list\[str\], queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using BM25 over OCR'd text. The standard keyword-based baseline. No GPU required; strong on text-dense pages, weak on visual content that Tesseract cannot read. """ tokenized = \[text.lower().split() for text in page\_texts\] bm25 = BM25Okapi(tokenized) results: list\[RetrievalResult\] = \[\] for q in queries: scores = bm25.get\_scores(q.text.lower().split()) ranked = sorted(range(len(page\_ids)), key=lambda i: -scores\[i\])\[:top\_k\] results.append( RetrievalResult( query\_id=q.query\_id, ranked\_page\_ids=\[page\_ids\[i\] for i in ranked\], ) ) return results # ───────────────────────────────────────────────────────────────────────────── # Tasks — evaluation # ───────────────────────────────────────────────────────────────────────────── @driver.task async def evaluate( results: list\[RetrievalResult\], ground\_truth: list\[PageQuery\], k: int, ) -> Metrics: """ Compute Recall@K, NDCG@K, and MRR for a single retrieval model. Recall@K — was the correct page in the top-K results? NDCG@K — normalised discounted cumulative gain; rewards earlier hits. MRR — mean reciprocal rank of the first correct result. All three are averaged over all queries. Higher is better. """ gt\_map = {q.query\_id: q.relevant\_page\_id for q in ground\_truth} recall\_vals, ndcg\_vals, mrr\_vals = \[\], \[\], \[\] for r in results: relevant = gt\_map.get(r.query\_id, "") top = r.ranked\_page\_ids\[:k\] recall\_vals.append(1.0 if relevant in top else 0.0) rels = \[1 if pid == relevant else 0 for pid in top\] idcg = \_dcg(\[1\]) # ideal: correct page at rank 1 ndcg\_vals.append(\_dcg(rels) / idcg if idcg > 0 else 0.0) rr = 0.0 for rank, pid in enumerate(r.ranked\_page\_ids, start=1): if pid == relevant: rr = 1.0 / rank break mrr\_vals.append(rr) return Metrics( recall\_at\_k=float(np.mean(recall\_vals)), ndcg\_at\_k=float(np.mean(ndcg\_vals)), mrr=float(np.mean(mrr\_vals)), k=k, ) # ───────────────────────────────────────────────────────────────────────────── # Tasks — report # ───────────────────────────────────────────────────────────────────────────── @driver.task(report=True) async def generate\_report(report: ComparisonReport) -> None: """ Emit an interactive HTML report visible in the Flyte UI. report=True marks this task as a reporting task. Flyte renders the HTML returned via flyte.report.replace.aio() directly in the execution detail page — no separate dashboard or export step required. The report contains: - Summary cards: experiment count, best model, best Recall@K. - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment. - Ranked results table with all three metrics. """ sorted\_results = sorted(report.results, key=lambda r: -r.metrics.recall\_at\_k) best = sorted\_results\[0\] labels = \[r.config.name for r in sorted\_results\] recall\_vals = \[r.metrics.recall\_at\_k for r in sorted\_results\] ndcg\_vals = \[r.metrics.ndcg\_at\_k for r in sorted\_results\] mrr\_vals = \[r.metrics.mrr for r in sorted\_results\] table\_rows = "".join( f""" {r.config.name} {r.config.model.value} {r.metrics.recall\_at\_k:.3f} {r.metrics.ndcg\_at\_k:.3f} {r.metrics.mrr:.3f} {r.metrics.k} """ for r in sorted\_results ) html = f""" Visual Document Retrieval — Results

Visual Document Retrieval — Experiment Comparison

ViDoRe benchmark · {len(report.results)} experiment(s)

{len(report.results)}
Experiments
{best.config.name}
Best by Recall@K
{best.metrics.recall\_at\_k:.3f}
Best Recall@{best.metrics.k}
{best.metrics.ndcg\_at\_k:.3f}
Best NDCG@{best.metrics.k}
{best.metrics.mrr:.3f}
Best MRR

Metrics by Experiment

Ranked Results

{table\_rows}
ExperimentModel Recall@KNDCG@KMRRK
""" await flyte.report.replace.aio(html) await flyte.report.flush.aio() # ───────────────────────────────────────────────────────────────────────────── # Experiment orchestration # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment run\_experiment}} @driver.task async def run\_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult: """ End-to-end retrieval pipeline for a single ExperimentConfig. Flyte v2's dynamic execution means this driver task can call GPU tasks (index\_colpali, search\_colpali) based on the runtime value of config.model — no static DAG wiring required. The if/elif is plain Python; Flyte schedules the selected sub-tasks on the appropriate environment. Caching: two experiments that share the same model and corpus (e.g. ColPali at top\_k=5 and top\_k=10) will hit the same cached index. GPU work is paid at most once per (model, corpus) pair across all experiments. Search queries are sharded into chunks of SEARCH\_SHARD\_SIZE and dispatched as concurrent task invocations. All shards land on the single warm container (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the GPU saturated throughout search rather than processing one large sequential batch from a single caller. flyte.group wraps each experiment in a named span in the Flyte UI, making it easy to compare latencies and drill into individual runs. """ SEARCH\_SHARD\_SIZE = 256 with flyte.group(config.name): if config.model == RetrievalModel.COLPALI: index\_file = await index\_colpali(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_colpali(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] elif config.model == RetrievalModel.SIGLIP: index\_file = await index\_siglip(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_siglip(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] else: # RetrievalModel.OCR\_BM25 page\_texts = await extract\_page\_texts(dataset.page\_files) results = await search\_bm25(page\_texts, dataset.page\_ids, dataset.queries, config.top\_k) metrics = await evaluate(results, dataset.queries, config.top\_k) return ExperimentResult(config=config, metrics=metrics) # {{/docs-fragment run\_experiment}} # {{docs-fragment compare\_experiments}} @driver.task async def compare\_experiments( configs: list\[ExperimentConfig\], subset: str = "docvqa", max\_pages: int = 200, ) -> ComparisonReport: """ Fan out over all experiment configs and return a ranked comparison table. The dataset is loaded once and shared across all experiments. Each config runs as a concurrent Flyte task via asyncio.gather. Experiments that share a model reuse the cached index — you only pay GPU time for new work. On completion, generate\_report emits an interactive Chart.js HTML report visible directly in the Flyte execution detail page. Default dataset: vidore\_v3\_finance\_en (~2 942 corpus pages, 1 854 queries) with max\_pages=2 000 to exercise the GPU pipeline at scale. """ dataset = await load\_vidore\_pages(subset=subset, max\_pages=max\_pages) # All experiments launch concurrently. Shared cached outputs (same model, # same corpus) are served from cache rather than recomputed. experiment\_coros = \[run\_experiment(config=cfg, dataset=dataset) for cfg in configs\] results: list\[ExperimentResult\] = list(await asyncio.gather(\*experiment\_coros)) report = ComparisonReport(results=results) print(report.summary()) best = report.best\_by("recall\_at\_k") print(f"\\nBest by Recall@{best.metrics.k}: {best.config.name}") # Emit the interactive HTML report in the Flyte UI. await generate\_report(report) return report # {{/docs-fragment compare\_experiments}} # ───────────────────────────────────────────────────────────────────────────── # Entry point # ───────────────────────────────────────────────────────────────────────────── if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Define the experiment grid. Each ExperimentConfig is one point in the # design space. Adding a new model or varying top\_k is one line here — # no task code changes required. # # ColPali appears twice with different top\_k values. The cache ensures # index\_colpali runs only once and both experiments share that result. # {{docs-fragment grid}} configs = \[\ ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top\_k=5),\ ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top\_k=10),\ ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top\_k=5),\ ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR\_BM25, top\_k=5),\ \] # {{/docs-fragment grid}} run = flyte.with\_runcontext().run( compare\_experiments, configs=configs, subset="vidore\_v3\_finance\_en", max\_pages=2000, ) print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval\_eval.py\* > \[!NOTE\] > The \`DynamicBatcher\` implementation lives in the \`extras/\` package next to the example. Run the script from the example directory so the import resolves. ## Run one experiment \`run\_experiment\` selects the right index/search path based on the runtime value of \`config.model\` — Flyte v2's dynamic execution means there's no static DAG to wire up. \`flyte.group\` wraps each experiment in a named span in the UI. \`\`\` # /// script # requires-python = ">=3.12" # dependencies = \[\ # "colpali-engine>=0.3.1",\ # "transformers>=4.41",\ # "sentencepiece>=0.2",\ # "torch>=2.0",\ # "pillow>=10",\ # "datasets>=2.18",\ # "rank-bm25>=0.2",\ # "numpy>=1.26",\ # "python-doctr\[torch\]>=0.8",\ # "pydantic>=2.0",\ # "flyte>=2.0.0",\ # \] # /// """ Multimodal Retrieval Evaluation Pipeline This tutorial is an experiment framework for benchmarking visual document retrieval approaches on the ViDoRe benchmark. Each experiment is defined by an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and returns a ranked comparison table with an interactive HTML report. The corpus is a set of PDF page images; queries are plain-text questions. Each retrieval method must find the page that answers each question — no text is provided to the model, only the raw image. ColPali-v1.2 — patch-level multi-vector embeddings from a VLM (PaliGemma). No OCR. The model produces one vector per image patch (~1024 per page). MaxSim late-interaction scoring finds the best matching patch for each query token. SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP successor. One matrix multiply per query; fast and effective but a single vector cannot localise fine-grained regions. OCR + BM25 — text-only baseline. doctr (GPU OCR) extracts text in batches, BM25 matches keywords. Strong on text-dense pages; fails on charts, tables, and figures where content is visual. """ import asyncio import enum import json import math import os import tempfile from functools import lru\_cache from io import BytesIO from itertools import islice import numpy as np from PIL import Image as PILImage from pydantic import BaseModel from rank\_bm25 import BM25Okapi from extras import DynamicBatcher import flyte import flyte.report from flyte.io import File # ───────────────────────────────────────────────────────────────────────────── # Environments # ───────────────────────────────────────────────────────────────────────────── # One Docker image for all tasks. The PEP 723 header defines Python deps. # ca-certificates is required for HTTPS calls to HuggingFace and blob stores. # {{docs-fragment image}} image = ( flyte.Image.from\_uv\_script(\_\_file\_\_, name="vidore-eval-v2") .with\_apt\_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0") # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy. # Without it every reusable container exits with StartError (exit code 128). .with\_pip\_packages("unionai-reuse>=0.1.11") ) # {{/docs-fragment image}} # GPU environment for ColPali image encoding and search. # # ReusePolicy keeps up to 3 warm GPU containers alive between task calls. # Without it, every task invocation cold-starts a new container and downloads # ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model # weights already loaded into VRAM — is reused for the next task dispatch. # # replicas=1 single warm container — all concurrent shard calls land # here so they share one DynamicBatcher process # concurrency=8 up to 8 query-shard tasks run simultaneously on the # container, all feeding the same DynamicBatcher queue # idle\_ttl=120 keep alive 2 min after the last task finishes # scaledown\_ttl=60 scale to zero after 1 min of complete inactivity # {{docs-fragment envs}} colpali\_indexer = flyte.TaskEnvironment( name="vidore-colpali-indexer", image=image, resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for SigLIP image encoding and search. # # Separate from the ColPali environment so each model's warm containers # are managed independently — ColPali and SigLIP experiments can scale # without contending for the same pool of reusable containers. siglip\_indexer = flyte.TaskEnvironment( name="vidore-siglip-indexer", image=image, resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition) # in batches on GPU — much faster than CPU Tesseract. # No ReusePolicy needed: the result is cached, so this task runs at most once. ocr\_engine = flyte.TaskEnvironment( name="vidore-ocr-engine", image=image, resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1), ) # Driver: orchestration, BM25 search, evaluation, and reporting. # depends\_on ensures the shared Docker image is built before all environments # try to schedule tasks. driver = flyte.TaskEnvironment( name="vidore-driver", image=image, resources=flyte.Resources(cpu=2, memory="12Gi"), depends\_on=\[colpali\_indexer, siglip\_indexer, ocr\_engine\], ) # {{/docs-fragment envs}} # ───────────────────────────────────────────────────────────────────────────── # Configuration types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment config\_types}} class RetrievalModel(str, enum.Enum): """Retrieval backend to evaluate.""" COLPALI = "colpali-v1.2" # multi-vector patch embeddings, MaxSim SIGLIP = "siglip-so400m" # single-vector global embedding, cosine sim OCR\_BM25 = "ocr+bm25" # text extracted by Tesseract, ranked by BM25 class ExperimentConfig(BaseModel): """ All knobs for one retrieval experiment. Passed as a typed Flyte input. Because ExperimentConfig is a Pydantic model, Flyte serialises it alongside every task output — so you can always reconstruct which config produced which metric without maintaining a separate log. """ name: str # human-readable label shown in the comparison table model: RetrievalModel top\_k: int = 5 # number of pages to retrieve per query # {{/docs-fragment config\_types}} # ───────────────────────────────────────────────────────────────────────────── # Data types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment data\_types}} class PageQuery(BaseModel): """One retrieval query with its ground-truth page.""" query\_id: str text: str # e.g. "What was revenue growth in Q3?" relevant\_page\_id: str # one correct page per query class PageDataset(BaseModel): """ A corpus of document page images paired with text queries. page\_ids: unique page identifiers (derived from ViDoRe image filenames). page\_files: the same pages stored in Flyte's blob store as JPEG File handles. Tasks read images directly from here; no live HTTP. queries: text questions with ground-truth page IDs for evaluation. """ page\_ids: list\[str\] page\_files: list\[File\] queries: list\[PageQuery\] class Config: arbitrary\_types\_allowed = True class RetrievalResult(BaseModel): query\_id: str ranked\_page\_ids: list\[str\] # ordered best → worst class Metrics(BaseModel): recall\_at\_k: float ndcg\_at\_k: float mrr: float k: int class ExperimentResult(BaseModel): config: ExperimentConfig metrics: Metrics # {{/docs-fragment data\_types}} class ComparisonReport(BaseModel): results: list\[ExperimentResult\] def best\_by(self, metric: str = "recall\_at\_k") -> ExperimentResult: return max(self.results, key=lambda r: getattr(r.metrics, metric)) def summary(self) -> str: header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}" sep = "─" \* len(header) rows = \[header, sep\] for r in sorted(self.results, key=lambda x: -x.metrics.recall\_at\_k): rows.append( f"{r.config.name:<30} " f"{r.config.model.value:<18} " f"{r.metrics.recall\_at\_k:>10.3f} " f"{r.metrics.ndcg\_at\_k:>8.3f} " f"{r.metrics.mrr:>7.3f}" ) return "\\n".join(rows) # ───────────────────────────────────────────────────────────────────────────── # Cached model loaders # ───────────────────────────────────────────────────────────────────────────── # These functions are at module level so they are shared across all tasks that # run on the same warm container (via ReusePolicy). lru\_cache(maxsize=1) means # the model is loaded from disk/HuggingFace exactly once per container process # and kept in GPU memory for every subsequent task dispatch to that container. @lru\_cache(maxsize=1) def \_colpali\_model(): """Load ColPali-v1.2 into GPU memory and cache the result. device\_map= is the correct loading pattern for ColPali's PaliGemma backbone; it handles weight placement via accelerate. torch.compile is skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross- invocation batching is the primary GPU utilisation mechanism. """ import torch from colpali\_engine.models import ColPali, ColPaliProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = ColPali.from\_pretrained( "vidore/colpali-v1.2", torch\_dtype=torch.bfloat16, device\_map=device, ) processor = ColPaliProcessor.from\_pretrained("vidore/colpali-v1.2") return model, processor, device @lru\_cache(maxsize=1) def \_siglip\_model(): """Load SigLIP SO400M into GPU memory, compile it, and cache the result. torch.compile (mode="reduce-overhead") fuses the vision and text encoder transformer layers into optimised CUDA kernels. As with ColPali, the compilation overhead is paid once per warm container lifetime. """ import torch from transformers import AutoModel, AutoProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = AutoModel.from\_pretrained("google/siglip-so400m-patch14-224").to(device) if device == "cuda": model = torch.compile(model, mode="reduce-overhead") processor = AutoProcessor.from\_pretrained("google/siglip-so400m-patch14-224") return model, processor, device @lru\_cache(maxsize=1) def \_ocr\_model(): """Load the doctr OCR predictor onto GPU and cache it. doctr's ocr\_predictor bundles a detection model (DBNet) and a recognition model (CRNN/SAR) into a single callable. pretrained=True downloads both model weights from doctr's model zoo on first use. """ import torch from doctr.models import ocr\_predictor predictor = ocr\_predictor(pretrained=True) if torch.cuda.is\_available(): predictor = predictor.cuda() return predictor # ───────────────────────────────────────────────────────────────────────────── # Search batcher singletons # ───────────────────────────────────────────────────────────────────────────── # One DynamicBatcher per model, shared across all concurrent search task # invocations on the same warm container (concurrency=3). Queries from every # concurrent caller are aggregated into a single GPU batch, maximizing # throughput compared to each invocation running its own forward pass. # # Initialised lazily on the first search call via double-checked locking and # lives for the container's lifetime. The process\_fn runs GPU work via # asyncio.to\_thread so the aggregation loop can continue collecting queries # from other callers while the GPU processes the current batch. # # File is not hashable so alru\_cache cannot be used here; module-level state # with asyncio.Lock is the correct pattern. # # Assumption: index\_colpali/index\_siglip use cache="auto", so the same corpus # always produces the same index File across all callers on this container. If # the index file ever changed between calls, the batcher would silently continue # using the corpus embeddings loaded from the first call. \_colpali\_batcher: DynamicBatcher | None = None \_colpali\_batcher\_lock = asyncio.Lock() \_siglip\_batcher: DynamicBatcher | None = None \_siglip\_batcher\_lock = asyncio.Lock() async def \_get\_colpali\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level ColPali search batcher, creating it on first call.""" global \_colpali\_batcher if \_colpali\_batcher is not None: return \_colpali\_batcher async with \_colpali\_batcher\_lock: if \_colpali\_batcher is not None: return \_colpali\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, n\_patches, dim) index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_colpali\_model() corpus\_emb = corpus\_emb.to(device, dtype=torch.float32) async def colpali\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: query\_inputs = processor.process\_queries(\[q.text for q in batch\]) query\_inputs = {k: v.to(device) for k, v in query\_inputs.items()} with torch.no\_grad(): query\_embs = model(\*\*query\_inputs).float() # (B, T, D) query\_chunk = 8 n\_pages = corpus\_emb.shape\[0\] all\_scores = torch.empty(len(batch), n\_pages, device=device) for start in range(0, len(batch), query\_chunk): chunk = query\_embs\[start : start + query\_chunk\] all\_scores\[start : start + query\_chunk\] = ( torch.einsum("ctd,pjd->ctpj", chunk, corpus\_emb) .max(dim=3).values .sum(dim=1) ) sorted\_indices = all\_scores.argsort(dim=1, descending=True).cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] # Run GPU work in a thread so the event loop — and the batcher's # aggregation loop — remain unblocked while the GPU is busy. return await asyncio.to\_thread(\_gpu\_work) batcher: DynamicBatcher\[PageQuery, list\[str\]\] = DynamicBatcher( process\_fn=colpali\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_colpali\_batcher = batcher return \_colpali\_batcher async def \_get\_siglip\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level SigLIP search batcher, creating it on first call.""" global \_siglip\_batcher if \_siglip\_batcher is not None: return \_siglip\_batcher async with \_siglip\_batcher\_lock: if \_siglip\_batcher is not None: return \_siglip\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, dim), L2-normalised index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_siglip\_model() corpus\_emb = corpus\_emb.to(device) async def siglip\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: text\_inputs = processor( text=\[q.text for q in batch\], return\_tensors="pt", padding=True, truncation=True, ).to(device) with torch.no\_grad(): text\_out = model.text\_model(\*\*text\_inputs) query\_embs = text\_out.pooler\_output # (B, dim) query\_embs = query\_embs / query\_embs.norm(dim=-1, keepdim=True) scores\_matrix = corpus\_emb @ query\_embs.T # (n\_pages, B) sorted\_indices = scores\_matrix.argsort(dim=0, descending=True).T.cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] return await asyncio.to\_thread(\_gpu\_work) batcher = DynamicBatcher( process\_fn=siglip\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_siglip\_batcher = batcher return \_siglip\_batcher # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── def \_batches(items: list, batch\_size: int): """Yield successive fixed-size batches from a list.""" for start in range(0, len(items), batch\_size): yield items\[start : start + batch\_size\] def \_load\_image\_sync(f: File) -> PILImage.Image: """Blocking download + decode. Intended to be called from a thread pool.""" with f.open\_sync("rb") as fh: data = fh.read() return PILImage.open(BytesIO(data)).convert("RGB") async def \_load\_image(f: File) -> PILImage.Image: """Download and decode a page image in a thread-pool worker. asyncio.to\_thread runs \_load\_image\_sync in a real OS thread so that blocking network I/O can overlap with GPU-bound forward passes when images are pre-submitted via loop.run\_in\_executor before the GPU kernel. """ return await asyncio.to\_thread(\_load\_image\_sync, f) async def \_load\_npz(index\_file: File) -> np.lib.npyio.NpzFile: """Download an index File to a local temp path and open with np.load.""" with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp: async with index\_file.open("rb") as fh: tmp.write(bytes(await fh.read())) return np.load(tmp.name) def \_dcg(relevances: list\[int\]) -> float: return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances)) # ───────────────────────────────────────────────────────────────────────────── # Tasks — data loading # ───────────────────────────────────────────────────────────────────────────── @driver.task(cache="auto", retries=3) async def load\_vidore\_pages(subset: str = "docvqa", max\_pages: int = 200) -> PageDataset: """ Load a ViDoRe benchmark subset and store page images in Flyte's blob store. Supports two dataset formats: Legacy (subsampled) — single 'test' split with one row per (query, page) pair; fields: image, query, image\_filename. streaming=True reads only the rows requested via islice — no full-shard download. Datasets: vidore/docvqa\_test\_subsampled, vidore/infovqa\_test\_subsampled V3 — separate corpus / queries / qrels splits following the BEIR retrieval benchmark format. corpus contains page images; queries contains question text; qrels maps query IDs to relevant corpus page IDs (many-to-many). Datasets: vidore/vidore\_v3\_finance\_en (~2 942 pages, 1 854 queries) The first call uploads page images to Flyte's blob store and caches the PageDataset; every subsequent call with the same arguments returns the cached result instantly. retries=3 guards against transient HuggingFace network failures. Available subsets: "docvqa", "infovqa", "vidore\_v3\_finance\_en" """ from datasets import load\_dataset subset\_map = { "docvqa": "vidore/docvqa\_test\_subsampled", "infovqa": "vidore/infovqa\_test\_subsampled", "vidore\_v3\_finance\_en": "vidore/vidore\_v3\_finance\_en", } dataset\_name = subset\_map.get(subset, f"vidore/{subset}\_test\_subsampled") # V3 datasets ship with separate corpus / queries / qrels splits. \_V3\_SUBSETS = {"vidore\_v3\_finance\_en"} if subset in \_V3\_SUBSETS: # ── V3 format ───────────────────────────────────────────────────────── # corpus / queries / qrels are HuggingFace configs (name=), not splits. # corpus uses streaming=True so images are decoded one at a time — # loading all 2 942 rows eagerly would hold gigabytes of PIL images in # the driver's RAM simultaneously. qrels and queries are text-only and # small enough to load fully into memory. corpus\_ds = load\_dataset(dataset\_name, name="corpus", split="test", streaming=True) qrels\_ds = load\_dataset(dataset\_name, name="qrels", split="test") queries\_ds = load\_dataset(dataset\_name, name="queries", split="test") # Normalise field names — V3 follows BEIR convention (hyphenated ids). def \_col(ds, \*candidates): cols = set(ds.column\_names) for c in candidates: if c in cols: return c raise KeyError(f"None of {candidates} found in columns {cols}") corpus\_id\_col = \_col(corpus\_ds, "corpus-id", "corpus\_id", "id", "\_id") query\_id\_col = \_col(queries\_ds, "query-id", "query\_id", "id", "\_id") query\_text\_col = \_col(queries\_ds, "query", "text") qrel\_qid\_col = \_col(qrels\_ds, "query-id", "query\_id") qrel\_cid\_col = \_col(qrels\_ds, "corpus-id", "corpus\_id") # Slice corpus to max\_pages, upload each image to Flyte blob store. page\_ids: list\[str\] = \[\] page\_files: list\[File\] = \[\] corpus\_id\_to\_page\_id: dict\[str, str\] = {} for i, row in enumerate(islice(corpus\_ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue cid = str(row\[corpus\_id\_col\]) page\_id = f"{subset}\_{i:04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) corpus\_id\_to\_page\_id\[cid\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) # Build query\_id → relevant page\_id from qrels (first match wins). # Only keep relevance judgements whose corpus page is in our slice. qrel\_map: dict\[str, str\] = {} for row in qrels\_ds: qid = str(row\[qrel\_qid\_col\]) cid = str(row\[qrel\_cid\_col\]) if cid in corpus\_id\_to\_page\_id and qid not in qrel\_map: qrel\_map\[qid\] = corpus\_id\_to\_page\_id\[cid\] # Collect queries that have at least one relevant page in our slice. queries: list\[PageQuery\] = \[\] for row in queries\_ds: qid = str(row\[query\_id\_col\]) if qid not in qrel\_map: continue queries.append( PageQuery( query\_id=qid, text=str(row\[query\_text\_col\]), relevant\_page\_id=qrel\_map\[qid\], ) ) else: # ── Legacy format ───────────────────────────────────────────────────── # Single 'test' split with one row per (query, page) pair. ds = load\_dataset(dataset\_name, split="test", streaming=True) page\_ids = \[\] page\_files = \[\] queries = \[\] seen\_pages: dict\[str, str\] = {} # image\_filename → page\_id for i, row in enumerate(islice(ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue filename: str = row.get("image\_filename") or f"page\_{i}" query\_text: str = row.get("query", "") if not query\_text: continue # Each unique page is uploaded exactly once; multiple queries may # share the same page (same image\_filename). if filename not in seen\_pages: page\_id = f"{subset}\_{len(page\_ids):04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) seen\_pages\[filename\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) else: page\_id = seen\_pages\[filename\] queries.append( PageQuery( query\_id=f"q{i:04d}", text=query\_text, relevant\_page\_id=page\_id, ) ) print(f"Loaded {len(page\_ids)} unique pages, {len(queries)} queries", flush=True) return PageDataset(page\_ids=page\_ids, page\_files=page\_files, queries=queries) # ───────────────────────────────────────────────────────────────────────────── # Tasks — indexing # ───────────────────────────────────────────────────────────────────────────── @colpali\_indexer.task(cache="auto", retries=2) async def index\_colpali(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with ColPali-v1.2 and save the multi-vector index. ColPali skips OCR entirely. It feeds the raw page image into PaliGemma (a vision-language model) and produces one embedding vector per image patch — roughly 1,024 patches per page, each of dimension 128. \_colpali\_model() is an lru\_cache'd loader. On a cold container, it downloads and loads the model once. On a warm container (kept alive by ReusePolicy), it returns the already-loaded model instantly from cache — no repeated ~7 GB download. The index is stored as a .npz file in Flyte's blob store: embeddings — float32, shape (n\_pages, n\_patches, dim) page\_ids — matching page ID strings cache="auto" + retries=2: the result is stored permanently on success; transient failures (e.g. HuggingFace rate limits) are retried twice. """ import torch model, processor, device = \_colpali\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 4)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor.process\_images(images) inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no\_grad(): emb = model(\*\*inputs) # (batch, n\_patches, dim) all\_embeddings.append(emb.cpu().float().numpy()) print(f"ColPali: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, n\_patches, dim) out\_path = os.path.join(tempfile.gettempdir(), "colpali\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @siglip\_indexer.task(cache="auto", retries=2) async def index\_siglip(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with SigLIP SO400M and save the single-vector index. SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss instead of softmax — avoiding the normalisation bottleneck that limits CLIP's scalability. Produces one global embedding per page. \_siglip\_model() caches the model across warm container reuses. The index is stored as a .npz file: embeddings — float32, shape (n\_pages, dim), L2-normalised page\_ids — matching page ID strings """ import torch model, processor, device = \_siglip\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 8)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor(images=images, return\_tensors="pt", padding=True).to(device) with torch.no\_grad(): outputs = model.vision\_model(\*\*inputs) emb = outputs.pooler\_output # (batch, dim) emb = emb / emb.norm(dim=-1, keepdim=True) # L2 normalise all\_embeddings.append(emb.cpu().float().numpy()) print(f"SigLIP: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, dim) out\_path = os.path.join(tempfile.gettempdir(), "siglip\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @ocr\_engine.task(cache="auto") async def extract\_page\_texts(page\_files: list\[File\]) -> list\[str\]: """ OCR every page with doctr on GPU to produce a text-only baseline. doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single callable predictor. Pages are downloaded in parallel then fed in batches of ocr\_batch\_size. asyncio.to\_thread keeps the event loop unblocked during GPU inference. Result structure: result.pages\[i\].blocks\[j\].lines\[k\].words\[l\].value Cached: the same corpus is OCR'd at most once across all experiments that use the OCR+BM25 backend. """ import gc predictor = \_ocr\_model() # Process in batches: download each batch just-in-time so only # ocr\_batch\_size images are in memory at once instead of all 2 000. ocr\_batch\_size = 8 total = len(page\_files) texts: list\[str\] = \[\] for start in range(0, total, ocr\_batch\_size): batch\_files = page\_files\[start : start + ocr\_batch\_size\] batch\_images = list( await asyncio.gather(\*\[asyncio.to\_thread(\_load\_image\_sync, f) for f in batch\_files\]) ) batch\_np = \[np.array(img) for img in batch\_images\] del batch\_images result = await asyncio.to\_thread(predictor, batch\_np) del batch\_np for page\_output in result.pages: texts.append( "\\n".join( " ".join(word.value for word in line.words) for block in page\_output.blocks for line in block.lines ) ) del result gc.collect() print(f"OCR: processed {min(start + ocr\_batch\_size, total)}/{total} pages", flush=True) return texts # ───────────────────────────────────────────────────────────────────────────── # Tasks — search # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment search\_colpali}} @colpali\_indexer.task async def search\_colpali( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher. MaxSim score for page p given query q: score(q, p) = Σ\_{t ∈ query tokens} max\_{j ∈ page patches} (q\_t · p\_j) Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_colpali invocations on the same warm container (concurrency=8) into a single GPU batch. This keeps the GPU saturated rather than running one small batch per caller. The batcher's process\_fn runs GPU work in asyncio.to\_thread, so the aggregation loop stays live while the GPU encodes and scores. """ batcher = await \_get\_colpali\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] # {{/docs-fragment search\_colpali}} @siglip\_indexer.task async def search\_siglip( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using SigLIP cosine similarity via DynamicBatcher. Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_siglip invocations on the same warm container (concurrency=3) into a single GPU batch. SigLIP's single-vector embeddings make full vectorisation safe — the scores matrix (n\_pages x n\_queries) is small enough to materialise in one GPU call regardless of batch size. """ batcher = await \_get\_siglip\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] @driver.task async def search\_bm25( page\_texts: list\[str\], page\_ids: list\[str\], queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using BM25 over OCR'd text. The standard keyword-based baseline. No GPU required; strong on text-dense pages, weak on visual content that Tesseract cannot read. """ tokenized = \[text.lower().split() for text in page\_texts\] bm25 = BM25Okapi(tokenized) results: list\[RetrievalResult\] = \[\] for q in queries: scores = bm25.get\_scores(q.text.lower().split()) ranked = sorted(range(len(page\_ids)), key=lambda i: -scores\[i\])\[:top\_k\] results.append( RetrievalResult( query\_id=q.query\_id, ranked\_page\_ids=\[page\_ids\[i\] for i in ranked\], ) ) return results # ───────────────────────────────────────────────────────────────────────────── # Tasks — evaluation # ───────────────────────────────────────────────────────────────────────────── @driver.task async def evaluate( results: list\[RetrievalResult\], ground\_truth: list\[PageQuery\], k: int, ) -> Metrics: """ Compute Recall@K, NDCG@K, and MRR for a single retrieval model. Recall@K — was the correct page in the top-K results? NDCG@K — normalised discounted cumulative gain; rewards earlier hits. MRR — mean reciprocal rank of the first correct result. All three are averaged over all queries. Higher is better. """ gt\_map = {q.query\_id: q.relevant\_page\_id for q in ground\_truth} recall\_vals, ndcg\_vals, mrr\_vals = \[\], \[\], \[\] for r in results: relevant = gt\_map.get(r.query\_id, "") top = r.ranked\_page\_ids\[:k\] recall\_vals.append(1.0 if relevant in top else 0.0) rels = \[1 if pid == relevant else 0 for pid in top\] idcg = \_dcg(\[1\]) # ideal: correct page at rank 1 ndcg\_vals.append(\_dcg(rels) / idcg if idcg > 0 else 0.0) rr = 0.0 for rank, pid in enumerate(r.ranked\_page\_ids, start=1): if pid == relevant: rr = 1.0 / rank break mrr\_vals.append(rr) return Metrics( recall\_at\_k=float(np.mean(recall\_vals)), ndcg\_at\_k=float(np.mean(ndcg\_vals)), mrr=float(np.mean(mrr\_vals)), k=k, ) # ───────────────────────────────────────────────────────────────────────────── # Tasks — report # ───────────────────────────────────────────────────────────────────────────── @driver.task(report=True) async def generate\_report(report: ComparisonReport) -> None: """ Emit an interactive HTML report visible in the Flyte UI. report=True marks this task as a reporting task. Flyte renders the HTML returned via flyte.report.replace.aio() directly in the execution detail page — no separate dashboard or export step required. The report contains: - Summary cards: experiment count, best model, best Recall@K. - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment. - Ranked results table with all three metrics. """ sorted\_results = sorted(report.results, key=lambda r: -r.metrics.recall\_at\_k) best = sorted\_results\[0\] labels = \[r.config.name for r in sorted\_results\] recall\_vals = \[r.metrics.recall\_at\_k for r in sorted\_results\] ndcg\_vals = \[r.metrics.ndcg\_at\_k for r in sorted\_results\] mrr\_vals = \[r.metrics.mrr for r in sorted\_results\] table\_rows = "".join( f""" {r.config.name} {r.config.model.value} {r.metrics.recall\_at\_k:.3f} {r.metrics.ndcg\_at\_k:.3f} {r.metrics.mrr:.3f} {r.metrics.k} """ for r in sorted\_results ) html = f""" Visual Document Retrieval — Results

Visual Document Retrieval — Experiment Comparison

ViDoRe benchmark · {len(report.results)} experiment(s)

{len(report.results)}
Experiments
{best.config.name}
Best by Recall@K
{best.metrics.recall\_at\_k:.3f}
Best Recall@{best.metrics.k}
{best.metrics.ndcg\_at\_k:.3f}
Best NDCG@{best.metrics.k}
{best.metrics.mrr:.3f}
Best MRR

Metrics by Experiment

Ranked Results

{table\_rows}
ExperimentModel Recall@KNDCG@KMRRK
""" await flyte.report.replace.aio(html) await flyte.report.flush.aio() # ───────────────────────────────────────────────────────────────────────────── # Experiment orchestration # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment run\_experiment}} @driver.task async def run\_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult: """ End-to-end retrieval pipeline for a single ExperimentConfig. Flyte v2's dynamic execution means this driver task can call GPU tasks (index\_colpali, search\_colpali) based on the runtime value of config.model — no static DAG wiring required. The if/elif is plain Python; Flyte schedules the selected sub-tasks on the appropriate environment. Caching: two experiments that share the same model and corpus (e.g. ColPali at top\_k=5 and top\_k=10) will hit the same cached index. GPU work is paid at most once per (model, corpus) pair across all experiments. Search queries are sharded into chunks of SEARCH\_SHARD\_SIZE and dispatched as concurrent task invocations. All shards land on the single warm container (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the GPU saturated throughout search rather than processing one large sequential batch from a single caller. flyte.group wraps each experiment in a named span in the Flyte UI, making it easy to compare latencies and drill into individual runs. """ SEARCH\_SHARD\_SIZE = 256 with flyte.group(config.name): if config.model == RetrievalModel.COLPALI: index\_file = await index\_colpali(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_colpali(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] elif config.model == RetrievalModel.SIGLIP: index\_file = await index\_siglip(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_siglip(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] else: # RetrievalModel.OCR\_BM25 page\_texts = await extract\_page\_texts(dataset.page\_files) results = await search\_bm25(page\_texts, dataset.page\_ids, dataset.queries, config.top\_k) metrics = await evaluate(results, dataset.queries, config.top\_k) return ExperimentResult(config=config, metrics=metrics) # {{/docs-fragment run\_experiment}} # {{docs-fragment compare\_experiments}} @driver.task async def compare\_experiments( configs: list\[ExperimentConfig\], subset: str = "docvqa", max\_pages: int = 200, ) -> ComparisonReport: """ Fan out over all experiment configs and return a ranked comparison table. The dataset is loaded once and shared across all experiments. Each config runs as a concurrent Flyte task via asyncio.gather. Experiments that share a model reuse the cached index — you only pay GPU time for new work. On completion, generate\_report emits an interactive Chart.js HTML report visible directly in the Flyte execution detail page. Default dataset: vidore\_v3\_finance\_en (~2 942 corpus pages, 1 854 queries) with max\_pages=2 000 to exercise the GPU pipeline at scale. """ dataset = await load\_vidore\_pages(subset=subset, max\_pages=max\_pages) # All experiments launch concurrently. Shared cached outputs (same model, # same corpus) are served from cache rather than recomputed. experiment\_coros = \[run\_experiment(config=cfg, dataset=dataset) for cfg in configs\] results: list\[ExperimentResult\] = list(await asyncio.gather(\*experiment\_coros)) report = ComparisonReport(results=results) print(report.summary()) best = report.best\_by("recall\_at\_k") print(f"\\nBest by Recall@{best.metrics.k}: {best.config.name}") # Emit the interactive HTML report in the Flyte UI. await generate\_report(report) return report # {{/docs-fragment compare\_experiments}} # ───────────────────────────────────────────────────────────────────────────── # Entry point # ───────────────────────────────────────────────────────────────────────────── if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Define the experiment grid. Each ExperimentConfig is one point in the # design space. Adding a new model or varying top\_k is one line here — # no task code changes required. # # ColPali appears twice with different top\_k values. The cache ensures # index\_colpali runs only once and both experiments share that result. # {{docs-fragment grid}} configs = \[\ ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top\_k=5),\ ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top\_k=10),\ ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top\_k=5),\ ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR\_BM25, top\_k=5),\ \] # {{/docs-fragment grid}} run = flyte.with\_runcontext().run( compare\_experiments, configs=configs, subset="vidore\_v3\_finance\_en", max\_pages=2000, ) print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval\_eval.py\* ## Compare experiments The driver loads the dataset once, fans out across all configs with \`asyncio.gather\`, and emits an interactive Chart.js report in the Flyte UI. Experiments sharing a model reuse the cached index, so you only pay GPU time for new work. \`\`\` # /// script # requires-python = ">=3.12" # dependencies = \[\ # "colpali-engine>=0.3.1",\ # "transformers>=4.41",\ # "sentencepiece>=0.2",\ # "torch>=2.0",\ # "pillow>=10",\ # "datasets>=2.18",\ # "rank-bm25>=0.2",\ # "numpy>=1.26",\ # "python-doctr\[torch\]>=0.8",\ # "pydantic>=2.0",\ # "flyte>=2.0.0",\ # \] # /// """ Multimodal Retrieval Evaluation Pipeline This tutorial is an experiment framework for benchmarking visual document retrieval approaches on the ViDoRe benchmark. Each experiment is defined by an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and returns a ranked comparison table with an interactive HTML report. The corpus is a set of PDF page images; queries are plain-text questions. Each retrieval method must find the page that answers each question — no text is provided to the model, only the raw image. ColPali-v1.2 — patch-level multi-vector embeddings from a VLM (PaliGemma). No OCR. The model produces one vector per image patch (~1024 per page). MaxSim late-interaction scoring finds the best matching patch for each query token. SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP successor. One matrix multiply per query; fast and effective but a single vector cannot localise fine-grained regions. OCR + BM25 — text-only baseline. doctr (GPU OCR) extracts text in batches, BM25 matches keywords. Strong on text-dense pages; fails on charts, tables, and figures where content is visual. """ import asyncio import enum import json import math import os import tempfile from functools import lru\_cache from io import BytesIO from itertools import islice import numpy as np from PIL import Image as PILImage from pydantic import BaseModel from rank\_bm25 import BM25Okapi from extras import DynamicBatcher import flyte import flyte.report from flyte.io import File # ───────────────────────────────────────────────────────────────────────────── # Environments # ───────────────────────────────────────────────────────────────────────────── # One Docker image for all tasks. The PEP 723 header defines Python deps. # ca-certificates is required for HTTPS calls to HuggingFace and blob stores. # {{docs-fragment image}} image = ( flyte.Image.from\_uv\_script(\_\_file\_\_, name="vidore-eval-v2") .with\_apt\_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0") # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy. # Without it every reusable container exits with StartError (exit code 128). .with\_pip\_packages("unionai-reuse>=0.1.11") ) # {{/docs-fragment image}} # GPU environment for ColPali image encoding and search. # # ReusePolicy keeps up to 3 warm GPU containers alive between task calls. # Without it, every task invocation cold-starts a new container and downloads # ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model # weights already loaded into VRAM — is reused for the next task dispatch. # # replicas=1 single warm container — all concurrent shard calls land # here so they share one DynamicBatcher process # concurrency=8 up to 8 query-shard tasks run simultaneously on the # container, all feeding the same DynamicBatcher queue # idle\_ttl=120 keep alive 2 min after the last task finishes # scaledown\_ttl=60 scale to zero after 1 min of complete inactivity # {{docs-fragment envs}} colpali\_indexer = flyte.TaskEnvironment( name="vidore-colpali-indexer", image=image, resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for SigLIP image encoding and search. # # Separate from the ColPali environment so each model's warm containers # are managed independently — ColPali and SigLIP experiments can scale # without contending for the same pool of reusable containers. siglip\_indexer = flyte.TaskEnvironment( name="vidore-siglip-indexer", image=image, resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition) # in batches on GPU — much faster than CPU Tesseract. # No ReusePolicy needed: the result is cached, so this task runs at most once. ocr\_engine = flyte.TaskEnvironment( name="vidore-ocr-engine", image=image, resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1), ) # Driver: orchestration, BM25 search, evaluation, and reporting. # depends\_on ensures the shared Docker image is built before all environments # try to schedule tasks. driver = flyte.TaskEnvironment( name="vidore-driver", image=image, resources=flyte.Resources(cpu=2, memory="12Gi"), depends\_on=\[colpali\_indexer, siglip\_indexer, ocr\_engine\], ) # {{/docs-fragment envs}} # ───────────────────────────────────────────────────────────────────────────── # Configuration types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment config\_types}} class RetrievalModel(str, enum.Enum): """Retrieval backend to evaluate.""" COLPALI = "colpali-v1.2" # multi-vector patch embeddings, MaxSim SIGLIP = "siglip-so400m" # single-vector global embedding, cosine sim OCR\_BM25 = "ocr+bm25" # text extracted by Tesseract, ranked by BM25 class ExperimentConfig(BaseModel): """ All knobs for one retrieval experiment. Passed as a typed Flyte input. Because ExperimentConfig is a Pydantic model, Flyte serialises it alongside every task output — so you can always reconstruct which config produced which metric without maintaining a separate log. """ name: str # human-readable label shown in the comparison table model: RetrievalModel top\_k: int = 5 # number of pages to retrieve per query # {{/docs-fragment config\_types}} # ───────────────────────────────────────────────────────────────────────────── # Data types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment data\_types}} class PageQuery(BaseModel): """One retrieval query with its ground-truth page.""" query\_id: str text: str # e.g. "What was revenue growth in Q3?" relevant\_page\_id: str # one correct page per query class PageDataset(BaseModel): """ A corpus of document page images paired with text queries. page\_ids: unique page identifiers (derived from ViDoRe image filenames). page\_files: the same pages stored in Flyte's blob store as JPEG File handles. Tasks read images directly from here; no live HTTP. queries: text questions with ground-truth page IDs for evaluation. """ page\_ids: list\[str\] page\_files: list\[File\] queries: list\[PageQuery\] class Config: arbitrary\_types\_allowed = True class RetrievalResult(BaseModel): query\_id: str ranked\_page\_ids: list\[str\] # ordered best → worst class Metrics(BaseModel): recall\_at\_k: float ndcg\_at\_k: float mrr: float k: int class ExperimentResult(BaseModel): config: ExperimentConfig metrics: Metrics # {{/docs-fragment data\_types}} class ComparisonReport(BaseModel): results: list\[ExperimentResult\] def best\_by(self, metric: str = "recall\_at\_k") -> ExperimentResult: return max(self.results, key=lambda r: getattr(r.metrics, metric)) def summary(self) -> str: header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}" sep = "─" \* len(header) rows = \[header, sep\] for r in sorted(self.results, key=lambda x: -x.metrics.recall\_at\_k): rows.append( f"{r.config.name:<30} " f"{r.config.model.value:<18} " f"{r.metrics.recall\_at\_k:>10.3f} " f"{r.metrics.ndcg\_at\_k:>8.3f} " f"{r.metrics.mrr:>7.3f}" ) return "\\n".join(rows) # ───────────────────────────────────────────────────────────────────────────── # Cached model loaders # ───────────────────────────────────────────────────────────────────────────── # These functions are at module level so they are shared across all tasks that # run on the same warm container (via ReusePolicy). lru\_cache(maxsize=1) means # the model is loaded from disk/HuggingFace exactly once per container process # and kept in GPU memory for every subsequent task dispatch to that container. @lru\_cache(maxsize=1) def \_colpali\_model(): """Load ColPali-v1.2 into GPU memory and cache the result. device\_map= is the correct loading pattern for ColPali's PaliGemma backbone; it handles weight placement via accelerate. torch.compile is skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross- invocation batching is the primary GPU utilisation mechanism. """ import torch from colpali\_engine.models import ColPali, ColPaliProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = ColPali.from\_pretrained( "vidore/colpali-v1.2", torch\_dtype=torch.bfloat16, device\_map=device, ) processor = ColPaliProcessor.from\_pretrained("vidore/colpali-v1.2") return model, processor, device @lru\_cache(maxsize=1) def \_siglip\_model(): """Load SigLIP SO400M into GPU memory, compile it, and cache the result. torch.compile (mode="reduce-overhead") fuses the vision and text encoder transformer layers into optimised CUDA kernels. As with ColPali, the compilation overhead is paid once per warm container lifetime. """ import torch from transformers import AutoModel, AutoProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = AutoModel.from\_pretrained("google/siglip-so400m-patch14-224").to(device) if device == "cuda": model = torch.compile(model, mode="reduce-overhead") processor = AutoProcessor.from\_pretrained("google/siglip-so400m-patch14-224") return model, processor, device @lru\_cache(maxsize=1) def \_ocr\_model(): """Load the doctr OCR predictor onto GPU and cache it. doctr's ocr\_predictor bundles a detection model (DBNet) and a recognition model (CRNN/SAR) into a single callable. pretrained=True downloads both model weights from doctr's model zoo on first use. """ import torch from doctr.models import ocr\_predictor predictor = ocr\_predictor(pretrained=True) if torch.cuda.is\_available(): predictor = predictor.cuda() return predictor # ───────────────────────────────────────────────────────────────────────────── # Search batcher singletons # ───────────────────────────────────────────────────────────────────────────── # One DynamicBatcher per model, shared across all concurrent search task # invocations on the same warm container (concurrency=3). Queries from every # concurrent caller are aggregated into a single GPU batch, maximizing # throughput compared to each invocation running its own forward pass. # # Initialised lazily on the first search call via double-checked locking and # lives for the container's lifetime. The process\_fn runs GPU work via # asyncio.to\_thread so the aggregation loop can continue collecting queries # from other callers while the GPU processes the current batch. # # File is not hashable so alru\_cache cannot be used here; module-level state # with asyncio.Lock is the correct pattern. # # Assumption: index\_colpali/index\_siglip use cache="auto", so the same corpus # always produces the same index File across all callers on this container. If # the index file ever changed between calls, the batcher would silently continue # using the corpus embeddings loaded from the first call. \_colpali\_batcher: DynamicBatcher | None = None \_colpali\_batcher\_lock = asyncio.Lock() \_siglip\_batcher: DynamicBatcher | None = None \_siglip\_batcher\_lock = asyncio.Lock() async def \_get\_colpali\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level ColPali search batcher, creating it on first call.""" global \_colpali\_batcher if \_colpali\_batcher is not None: return \_colpali\_batcher async with \_colpali\_batcher\_lock: if \_colpali\_batcher is not None: return \_colpali\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, n\_patches, dim) index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_colpali\_model() corpus\_emb = corpus\_emb.to(device, dtype=torch.float32) async def colpali\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: query\_inputs = processor.process\_queries(\[q.text for q in batch\]) query\_inputs = {k: v.to(device) for k, v in query\_inputs.items()} with torch.no\_grad(): query\_embs = model(\*\*query\_inputs).float() # (B, T, D) query\_chunk = 8 n\_pages = corpus\_emb.shape\[0\] all\_scores = torch.empty(len(batch), n\_pages, device=device) for start in range(0, len(batch), query\_chunk): chunk = query\_embs\[start : start + query\_chunk\] all\_scores\[start : start + query\_chunk\] = ( torch.einsum("ctd,pjd->ctpj", chunk, corpus\_emb) .max(dim=3).values .sum(dim=1) ) sorted\_indices = all\_scores.argsort(dim=1, descending=True).cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] # Run GPU work in a thread so the event loop — and the batcher's # aggregation loop — remain unblocked while the GPU is busy. return await asyncio.to\_thread(\_gpu\_work) batcher: DynamicBatcher\[PageQuery, list\[str\]\] = DynamicBatcher( process\_fn=colpali\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_colpali\_batcher = batcher return \_colpali\_batcher async def \_get\_siglip\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level SigLIP search batcher, creating it on first call.""" global \_siglip\_batcher if \_siglip\_batcher is not None: return \_siglip\_batcher async with \_siglip\_batcher\_lock: if \_siglip\_batcher is not None: return \_siglip\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, dim), L2-normalised index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_siglip\_model() corpus\_emb = corpus\_emb.to(device) async def siglip\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: text\_inputs = processor( text=\[q.text for q in batch\], return\_tensors="pt", padding=True, truncation=True, ).to(device) with torch.no\_grad(): text\_out = model.text\_model(\*\*text\_inputs) query\_embs = text\_out.pooler\_output # (B, dim) query\_embs = query\_embs / query\_embs.norm(dim=-1, keepdim=True) scores\_matrix = corpus\_emb @ query\_embs.T # (n\_pages, B) sorted\_indices = scores\_matrix.argsort(dim=0, descending=True).T.cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] return await asyncio.to\_thread(\_gpu\_work) batcher = DynamicBatcher( process\_fn=siglip\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_siglip\_batcher = batcher return \_siglip\_batcher # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── def \_batches(items: list, batch\_size: int): """Yield successive fixed-size batches from a list.""" for start in range(0, len(items), batch\_size): yield items\[start : start + batch\_size\] def \_load\_image\_sync(f: File) -> PILImage.Image: """Blocking download + decode. Intended to be called from a thread pool.""" with f.open\_sync("rb") as fh: data = fh.read() return PILImage.open(BytesIO(data)).convert("RGB") async def \_load\_image(f: File) -> PILImage.Image: """Download and decode a page image in a thread-pool worker. asyncio.to\_thread runs \_load\_image\_sync in a real OS thread so that blocking network I/O can overlap with GPU-bound forward passes when images are pre-submitted via loop.run\_in\_executor before the GPU kernel. """ return await asyncio.to\_thread(\_load\_image\_sync, f) async def \_load\_npz(index\_file: File) -> np.lib.npyio.NpzFile: """Download an index File to a local temp path and open with np.load.""" with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp: async with index\_file.open("rb") as fh: tmp.write(bytes(await fh.read())) return np.load(tmp.name) def \_dcg(relevances: list\[int\]) -> float: return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances)) # ───────────────────────────────────────────────────────────────────────────── # Tasks — data loading # ───────────────────────────────────────────────────────────────────────────── @driver.task(cache="auto", retries=3) async def load\_vidore\_pages(subset: str = "docvqa", max\_pages: int = 200) -> PageDataset: """ Load a ViDoRe benchmark subset and store page images in Flyte's blob store. Supports two dataset formats: Legacy (subsampled) — single 'test' split with one row per (query, page) pair; fields: image, query, image\_filename. streaming=True reads only the rows requested via islice — no full-shard download. Datasets: vidore/docvqa\_test\_subsampled, vidore/infovqa\_test\_subsampled V3 — separate corpus / queries / qrels splits following the BEIR retrieval benchmark format. corpus contains page images; queries contains question text; qrels maps query IDs to relevant corpus page IDs (many-to-many). Datasets: vidore/vidore\_v3\_finance\_en (~2 942 pages, 1 854 queries) The first call uploads page images to Flyte's blob store and caches the PageDataset; every subsequent call with the same arguments returns the cached result instantly. retries=3 guards against transient HuggingFace network failures. Available subsets: "docvqa", "infovqa", "vidore\_v3\_finance\_en" """ from datasets import load\_dataset subset\_map = { "docvqa": "vidore/docvqa\_test\_subsampled", "infovqa": "vidore/infovqa\_test\_subsampled", "vidore\_v3\_finance\_en": "vidore/vidore\_v3\_finance\_en", } dataset\_name = subset\_map.get(subset, f"vidore/{subset}\_test\_subsampled") # V3 datasets ship with separate corpus / queries / qrels splits. \_V3\_SUBSETS = {"vidore\_v3\_finance\_en"} if subset in \_V3\_SUBSETS: # ── V3 format ───────────────────────────────────────────────────────── # corpus / queries / qrels are HuggingFace configs (name=), not splits. # corpus uses streaming=True so images are decoded one at a time — # loading all 2 942 rows eagerly would hold gigabytes of PIL images in # the driver's RAM simultaneously. qrels and queries are text-only and # small enough to load fully into memory. corpus\_ds = load\_dataset(dataset\_name, name="corpus", split="test", streaming=True) qrels\_ds = load\_dataset(dataset\_name, name="qrels", split="test") queries\_ds = load\_dataset(dataset\_name, name="queries", split="test") # Normalise field names — V3 follows BEIR convention (hyphenated ids). def \_col(ds, \*candidates): cols = set(ds.column\_names) for c in candidates: if c in cols: return c raise KeyError(f"None of {candidates} found in columns {cols}") corpus\_id\_col = \_col(corpus\_ds, "corpus-id", "corpus\_id", "id", "\_id") query\_id\_col = \_col(queries\_ds, "query-id", "query\_id", "id", "\_id") query\_text\_col = \_col(queries\_ds, "query", "text") qrel\_qid\_col = \_col(qrels\_ds, "query-id", "query\_id") qrel\_cid\_col = \_col(qrels\_ds, "corpus-id", "corpus\_id") # Slice corpus to max\_pages, upload each image to Flyte blob store. page\_ids: list\[str\] = \[\] page\_files: list\[File\] = \[\] corpus\_id\_to\_page\_id: dict\[str, str\] = {} for i, row in enumerate(islice(corpus\_ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue cid = str(row\[corpus\_id\_col\]) page\_id = f"{subset}\_{i:04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) corpus\_id\_to\_page\_id\[cid\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) # Build query\_id → relevant page\_id from qrels (first match wins). # Only keep relevance judgements whose corpus page is in our slice. qrel\_map: dict\[str, str\] = {} for row in qrels\_ds: qid = str(row\[qrel\_qid\_col\]) cid = str(row\[qrel\_cid\_col\]) if cid in corpus\_id\_to\_page\_id and qid not in qrel\_map: qrel\_map\[qid\] = corpus\_id\_to\_page\_id\[cid\] # Collect queries that have at least one relevant page in our slice. queries: list\[PageQuery\] = \[\] for row in queries\_ds: qid = str(row\[query\_id\_col\]) if qid not in qrel\_map: continue queries.append( PageQuery( query\_id=qid, text=str(row\[query\_text\_col\]), relevant\_page\_id=qrel\_map\[qid\], ) ) else: # ── Legacy format ───────────────────────────────────────────────────── # Single 'test' split with one row per (query, page) pair. ds = load\_dataset(dataset\_name, split="test", streaming=True) page\_ids = \[\] page\_files = \[\] queries = \[\] seen\_pages: dict\[str, str\] = {} # image\_filename → page\_id for i, row in enumerate(islice(ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue filename: str = row.get("image\_filename") or f"page\_{i}" query\_text: str = row.get("query", "") if not query\_text: continue # Each unique page is uploaded exactly once; multiple queries may # share the same page (same image\_filename). if filename not in seen\_pages: page\_id = f"{subset}\_{len(page\_ids):04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) seen\_pages\[filename\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) else: page\_id = seen\_pages\[filename\] queries.append( PageQuery( query\_id=f"q{i:04d}", text=query\_text, relevant\_page\_id=page\_id, ) ) print(f"Loaded {len(page\_ids)} unique pages, {len(queries)} queries", flush=True) return PageDataset(page\_ids=page\_ids, page\_files=page\_files, queries=queries) # ───────────────────────────────────────────────────────────────────────────── # Tasks — indexing # ───────────────────────────────────────────────────────────────────────────── @colpali\_indexer.task(cache="auto", retries=2) async def index\_colpali(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with ColPali-v1.2 and save the multi-vector index. ColPali skips OCR entirely. It feeds the raw page image into PaliGemma (a vision-language model) and produces one embedding vector per image patch — roughly 1,024 patches per page, each of dimension 128. \_colpali\_model() is an lru\_cache'd loader. On a cold container, it downloads and loads the model once. On a warm container (kept alive by ReusePolicy), it returns the already-loaded model instantly from cache — no repeated ~7 GB download. The index is stored as a .npz file in Flyte's blob store: embeddings — float32, shape (n\_pages, n\_patches, dim) page\_ids — matching page ID strings cache="auto" + retries=2: the result is stored permanently on success; transient failures (e.g. HuggingFace rate limits) are retried twice. """ import torch model, processor, device = \_colpali\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 4)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor.process\_images(images) inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no\_grad(): emb = model(\*\*inputs) # (batch, n\_patches, dim) all\_embeddings.append(emb.cpu().float().numpy()) print(f"ColPali: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, n\_patches, dim) out\_path = os.path.join(tempfile.gettempdir(), "colpali\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @siglip\_indexer.task(cache="auto", retries=2) async def index\_siglip(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with SigLIP SO400M and save the single-vector index. SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss instead of softmax — avoiding the normalisation bottleneck that limits CLIP's scalability. Produces one global embedding per page. \_siglip\_model() caches the model across warm container reuses. The index is stored as a .npz file: embeddings — float32, shape (n\_pages, dim), L2-normalised page\_ids — matching page ID strings """ import torch model, processor, device = \_siglip\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 8)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor(images=images, return\_tensors="pt", padding=True).to(device) with torch.no\_grad(): outputs = model.vision\_model(\*\*inputs) emb = outputs.pooler\_output # (batch, dim) emb = emb / emb.norm(dim=-1, keepdim=True) # L2 normalise all\_embeddings.append(emb.cpu().float().numpy()) print(f"SigLIP: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, dim) out\_path = os.path.join(tempfile.gettempdir(), "siglip\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @ocr\_engine.task(cache="auto") async def extract\_page\_texts(page\_files: list\[File\]) -> list\[str\]: """ OCR every page with doctr on GPU to produce a text-only baseline. doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single callable predictor. Pages are downloaded in parallel then fed in batches of ocr\_batch\_size. asyncio.to\_thread keeps the event loop unblocked during GPU inference. Result structure: result.pages\[i\].blocks\[j\].lines\[k\].words\[l\].value Cached: the same corpus is OCR'd at most once across all experiments that use the OCR+BM25 backend. """ import gc predictor = \_ocr\_model() # Process in batches: download each batch just-in-time so only # ocr\_batch\_size images are in memory at once instead of all 2 000. ocr\_batch\_size = 8 total = len(page\_files) texts: list\[str\] = \[\] for start in range(0, total, ocr\_batch\_size): batch\_files = page\_files\[start : start + ocr\_batch\_size\] batch\_images = list( await asyncio.gather(\*\[asyncio.to\_thread(\_load\_image\_sync, f) for f in batch\_files\]) ) batch\_np = \[np.array(img) for img in batch\_images\] del batch\_images result = await asyncio.to\_thread(predictor, batch\_np) del batch\_np for page\_output in result.pages: texts.append( "\\n".join( " ".join(word.value for word in line.words) for block in page\_output.blocks for line in block.lines ) ) del result gc.collect() print(f"OCR: processed {min(start + ocr\_batch\_size, total)}/{total} pages", flush=True) return texts # ───────────────────────────────────────────────────────────────────────────── # Tasks — search # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment search\_colpali}} @colpali\_indexer.task async def search\_colpali( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher. MaxSim score for page p given query q: score(q, p) = Σ\_{t ∈ query tokens} max\_{j ∈ page patches} (q\_t · p\_j) Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_colpali invocations on the same warm container (concurrency=8) into a single GPU batch. This keeps the GPU saturated rather than running one small batch per caller. The batcher's process\_fn runs GPU work in asyncio.to\_thread, so the aggregation loop stays live while the GPU encodes and scores. """ batcher = await \_get\_colpali\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] # {{/docs-fragment search\_colpali}} @siglip\_indexer.task async def search\_siglip( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using SigLIP cosine similarity via DynamicBatcher. Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_siglip invocations on the same warm container (concurrency=3) into a single GPU batch. SigLIP's single-vector embeddings make full vectorisation safe — the scores matrix (n\_pages x n\_queries) is small enough to materialise in one GPU call regardless of batch size. """ batcher = await \_get\_siglip\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] @driver.task async def search\_bm25( page\_texts: list\[str\], page\_ids: list\[str\], queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using BM25 over OCR'd text. The standard keyword-based baseline. No GPU required; strong on text-dense pages, weak on visual content that Tesseract cannot read. """ tokenized = \[text.lower().split() for text in page\_texts\] bm25 = BM25Okapi(tokenized) results: list\[RetrievalResult\] = \[\] for q in queries: scores = bm25.get\_scores(q.text.lower().split()) ranked = sorted(range(len(page\_ids)), key=lambda i: -scores\[i\])\[:top\_k\] results.append( RetrievalResult( query\_id=q.query\_id, ranked\_page\_ids=\[page\_ids\[i\] for i in ranked\], ) ) return results # ───────────────────────────────────────────────────────────────────────────── # Tasks — evaluation # ───────────────────────────────────────────────────────────────────────────── @driver.task async def evaluate( results: list\[RetrievalResult\], ground\_truth: list\[PageQuery\], k: int, ) -> Metrics: """ Compute Recall@K, NDCG@K, and MRR for a single retrieval model. Recall@K — was the correct page in the top-K results? NDCG@K — normalised discounted cumulative gain; rewards earlier hits. MRR — mean reciprocal rank of the first correct result. All three are averaged over all queries. Higher is better. """ gt\_map = {q.query\_id: q.relevant\_page\_id for q in ground\_truth} recall\_vals, ndcg\_vals, mrr\_vals = \[\], \[\], \[\] for r in results: relevant = gt\_map.get(r.query\_id, "") top = r.ranked\_page\_ids\[:k\] recall\_vals.append(1.0 if relevant in top else 0.0) rels = \[1 if pid == relevant else 0 for pid in top\] idcg = \_dcg(\[1\]) # ideal: correct page at rank 1 ndcg\_vals.append(\_dcg(rels) / idcg if idcg > 0 else 0.0) rr = 0.0 for rank, pid in enumerate(r.ranked\_page\_ids, start=1): if pid == relevant: rr = 1.0 / rank break mrr\_vals.append(rr) return Metrics( recall\_at\_k=float(np.mean(recall\_vals)), ndcg\_at\_k=float(np.mean(ndcg\_vals)), mrr=float(np.mean(mrr\_vals)), k=k, ) # ───────────────────────────────────────────────────────────────────────────── # Tasks — report # ───────────────────────────────────────────────────────────────────────────── @driver.task(report=True) async def generate\_report(report: ComparisonReport) -> None: """ Emit an interactive HTML report visible in the Flyte UI. report=True marks this task as a reporting task. Flyte renders the HTML returned via flyte.report.replace.aio() directly in the execution detail page — no separate dashboard or export step required. The report contains: - Summary cards: experiment count, best model, best Recall@K. - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment. - Ranked results table with all three metrics. """ sorted\_results = sorted(report.results, key=lambda r: -r.metrics.recall\_at\_k) best = sorted\_results\[0\] labels = \[r.config.name for r in sorted\_results\] recall\_vals = \[r.metrics.recall\_at\_k for r in sorted\_results\] ndcg\_vals = \[r.metrics.ndcg\_at\_k for r in sorted\_results\] mrr\_vals = \[r.metrics.mrr for r in sorted\_results\] table\_rows = "".join( f""" {r.config.name} {r.config.model.value} {r.metrics.recall\_at\_k:.3f} {r.metrics.ndcg\_at\_k:.3f} {r.metrics.mrr:.3f} {r.metrics.k} """ for r in sorted\_results ) html = f""" Visual Document Retrieval — Results

Visual Document Retrieval — Experiment Comparison

ViDoRe benchmark · {len(report.results)} experiment(s)

{len(report.results)}
Experiments
{best.config.name}
Best by Recall@K
{best.metrics.recall\_at\_k:.3f}
Best Recall@{best.metrics.k}
{best.metrics.ndcg\_at\_k:.3f}
Best NDCG@{best.metrics.k}
{best.metrics.mrr:.3f}
Best MRR

Metrics by Experiment

Ranked Results

{table\_rows}
ExperimentModel Recall@KNDCG@KMRRK
""" await flyte.report.replace.aio(html) await flyte.report.flush.aio() # ───────────────────────────────────────────────────────────────────────────── # Experiment orchestration # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment run\_experiment}} @driver.task async def run\_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult: """ End-to-end retrieval pipeline for a single ExperimentConfig. Flyte v2's dynamic execution means this driver task can call GPU tasks (index\_colpali, search\_colpali) based on the runtime value of config.model — no static DAG wiring required. The if/elif is plain Python; Flyte schedules the selected sub-tasks on the appropriate environment. Caching: two experiments that share the same model and corpus (e.g. ColPali at top\_k=5 and top\_k=10) will hit the same cached index. GPU work is paid at most once per (model, corpus) pair across all experiments. Search queries are sharded into chunks of SEARCH\_SHARD\_SIZE and dispatched as concurrent task invocations. All shards land on the single warm container (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the GPU saturated throughout search rather than processing one large sequential batch from a single caller. flyte.group wraps each experiment in a named span in the Flyte UI, making it easy to compare latencies and drill into individual runs. """ SEARCH\_SHARD\_SIZE = 256 with flyte.group(config.name): if config.model == RetrievalModel.COLPALI: index\_file = await index\_colpali(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_colpali(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] elif config.model == RetrievalModel.SIGLIP: index\_file = await index\_siglip(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_siglip(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] else: # RetrievalModel.OCR\_BM25 page\_texts = await extract\_page\_texts(dataset.page\_files) results = await search\_bm25(page\_texts, dataset.page\_ids, dataset.queries, config.top\_k) metrics = await evaluate(results, dataset.queries, config.top\_k) return ExperimentResult(config=config, metrics=metrics) # {{/docs-fragment run\_experiment}} # {{docs-fragment compare\_experiments}} @driver.task async def compare\_experiments( configs: list\[ExperimentConfig\], subset: str = "docvqa", max\_pages: int = 200, ) -> ComparisonReport: """ Fan out over all experiment configs and return a ranked comparison table. The dataset is loaded once and shared across all experiments. Each config runs as a concurrent Flyte task via asyncio.gather. Experiments that share a model reuse the cached index — you only pay GPU time for new work. On completion, generate\_report emits an interactive Chart.js HTML report visible directly in the Flyte execution detail page. Default dataset: vidore\_v3\_finance\_en (~2 942 corpus pages, 1 854 queries) with max\_pages=2 000 to exercise the GPU pipeline at scale. """ dataset = await load\_vidore\_pages(subset=subset, max\_pages=max\_pages) # All experiments launch concurrently. Shared cached outputs (same model, # same corpus) are served from cache rather than recomputed. experiment\_coros = \[run\_experiment(config=cfg, dataset=dataset) for cfg in configs\] results: list\[ExperimentResult\] = list(await asyncio.gather(\*experiment\_coros)) report = ComparisonReport(results=results) print(report.summary()) best = report.best\_by("recall\_at\_k") print(f"\\nBest by Recall@{best.metrics.k}: {best.config.name}") # Emit the interactive HTML report in the Flyte UI. await generate\_report(report) return report # {{/docs-fragment compare\_experiments}} # ───────────────────────────────────────────────────────────────────────────── # Entry point # ───────────────────────────────────────────────────────────────────────────── if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Define the experiment grid. Each ExperimentConfig is one point in the # design space. Adding a new model or varying top\_k is one line here — # no task code changes required. # # ColPali appears twice with different top\_k values. The cache ensures # index\_colpali runs only once and both experiments share that result. # {{docs-fragment grid}} configs = \[\ ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top\_k=5),\ ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top\_k=10),\ ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top\_k=5),\ ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR\_BM25, top\_k=5),\ \] # {{/docs-fragment grid}} run = flyte.with\_runcontext().run( compare\_experiments, configs=configs, subset="vidore\_v3\_finance\_en", max\_pages=2000, ) print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval\_eval.py\* ## Run the evaluation This example has no secrets — datasets and model weights are pulled from public Hugging Face repositories. It does require GPUs, so run it remotely. The experiment grid is defined in the entry point; adding a model or varying \`top\_k\` is a one-line change: \`\`\` # /// script # requires-python = ">=3.12" # dependencies = \[\ # "colpali-engine>=0.3.1",\ # "transformers>=4.41",\ # "sentencepiece>=0.2",\ # "torch>=2.0",\ # "pillow>=10",\ # "datasets>=2.18",\ # "rank-bm25>=0.2",\ # "numpy>=1.26",\ # "python-doctr\[torch\]>=0.8",\ # "pydantic>=2.0",\ # "flyte>=2.0.0",\ # \] # /// """ Multimodal Retrieval Evaluation Pipeline This tutorial is an experiment framework for benchmarking visual document retrieval approaches on the ViDoRe benchmark. Each experiment is defined by an ExperimentConfig; the pipeline fans them out as concurrent Flyte tasks and returns a ranked comparison table with an interactive HTML report. The corpus is a set of PDF page images; queries are plain-text questions. Each retrieval method must find the page that answers each question — no text is provided to the model, only the raw image. ColPali-v1.2 — patch-level multi-vector embeddings from a VLM (PaliGemma). No OCR. The model produces one vector per image patch (~1024 per page). MaxSim late-interaction scoring finds the best matching patch for each query token. SigLIP-SO400M — single global embedding per page from Google's 2023 CLIP successor. One matrix multiply per query; fast and effective but a single vector cannot localise fine-grained regions. OCR + BM25 — text-only baseline. doctr (GPU OCR) extracts text in batches, BM25 matches keywords. Strong on text-dense pages; fails on charts, tables, and figures where content is visual. """ import asyncio import enum import json import math import os import tempfile from functools import lru\_cache from io import BytesIO from itertools import islice import numpy as np from PIL import Image as PILImage from pydantic import BaseModel from rank\_bm25 import BM25Okapi from extras import DynamicBatcher import flyte import flyte.report from flyte.io import File # ───────────────────────────────────────────────────────────────────────────── # Environments # ───────────────────────────────────────────────────────────────────────────── # One Docker image for all tasks. The PEP 723 header defines Python deps. # ca-certificates is required for HTTPS calls to HuggingFace and blob stores. # {{docs-fragment image}} image = ( flyte.Image.from\_uv\_script(\_\_file\_\_, name="vidore-eval-v2") .with\_apt\_packages("ca-certificates", "libxcb1", "libgl1", "libglib2.0-0") # unionai-reuse installs the unionai-actor-bridge binary required by ReusePolicy. # Without it every reusable container exits with StartError (exit code 128). .with\_pip\_packages("unionai-reuse>=0.1.11") ) # {{/docs-fragment image}} # GPU environment for ColPali image encoding and search. # # ReusePolicy keeps up to 3 warm GPU containers alive between task calls. # Without it, every task invocation cold-starts a new container and downloads # ColPali-v1.2 (~7 GB) from scratch. With it, the container — and the model # weights already loaded into VRAM — is reused for the next task dispatch. # # replicas=1 single warm container — all concurrent shard calls land # here so they share one DynamicBatcher process # concurrency=8 up to 8 query-shard tasks run simultaneously on the # container, all feeding the same DynamicBatcher queue # idle\_ttl=120 keep alive 2 min after the last task finishes # scaledown\_ttl=60 scale to zero after 1 min of complete inactivity # {{docs-fragment envs}} colpali\_indexer = flyte.TaskEnvironment( name="vidore-colpali-indexer", image=image, resources=flyte.Resources(cpu=4, memory="16Gi", gpu="A10G:1"), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for SigLIP image encoding and search. # # Separate from the ColPali environment so each model's warm containers # are managed independently — ColPali and SigLIP experiments can scale # without contending for the same pool of reusable containers. siglip\_indexer = flyte.TaskEnvironment( name="vidore-siglip-indexer", image=image, resources=flyte.Resources(cpu=4, memory="8Gi", gpu=1), reusable=flyte.ReusePolicy( replicas=1, concurrency=8, idle\_ttl=120, scaledown\_ttl=60, ), ) # GPU environment for doctr OCR. doctr runs DBNet (detection) + CRNN (recognition) # in batches on GPU — much faster than CPU Tesseract. # No ReusePolicy needed: the result is cached, so this task runs at most once. ocr\_engine = flyte.TaskEnvironment( name="vidore-ocr-engine", image=image, resources=flyte.Resources(cpu=4, memory="20Gi", gpu=1), ) # Driver: orchestration, BM25 search, evaluation, and reporting. # depends\_on ensures the shared Docker image is built before all environments # try to schedule tasks. driver = flyte.TaskEnvironment( name="vidore-driver", image=image, resources=flyte.Resources(cpu=2, memory="12Gi"), depends\_on=\[colpali\_indexer, siglip\_indexer, ocr\_engine\], ) # {{/docs-fragment envs}} # ───────────────────────────────────────────────────────────────────────────── # Configuration types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment config\_types}} class RetrievalModel(str, enum.Enum): """Retrieval backend to evaluate.""" COLPALI = "colpali-v1.2" # multi-vector patch embeddings, MaxSim SIGLIP = "siglip-so400m" # single-vector global embedding, cosine sim OCR\_BM25 = "ocr+bm25" # text extracted by Tesseract, ranked by BM25 class ExperimentConfig(BaseModel): """ All knobs for one retrieval experiment. Passed as a typed Flyte input. Because ExperimentConfig is a Pydantic model, Flyte serialises it alongside every task output — so you can always reconstruct which config produced which metric without maintaining a separate log. """ name: str # human-readable label shown in the comparison table model: RetrievalModel top\_k: int = 5 # number of pages to retrieve per query # {{/docs-fragment config\_types}} # ───────────────────────────────────────────────────────────────────────────── # Data types # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment data\_types}} class PageQuery(BaseModel): """One retrieval query with its ground-truth page.""" query\_id: str text: str # e.g. "What was revenue growth in Q3?" relevant\_page\_id: str # one correct page per query class PageDataset(BaseModel): """ A corpus of document page images paired with text queries. page\_ids: unique page identifiers (derived from ViDoRe image filenames). page\_files: the same pages stored in Flyte's blob store as JPEG File handles. Tasks read images directly from here; no live HTTP. queries: text questions with ground-truth page IDs for evaluation. """ page\_ids: list\[str\] page\_files: list\[File\] queries: list\[PageQuery\] class Config: arbitrary\_types\_allowed = True class RetrievalResult(BaseModel): query\_id: str ranked\_page\_ids: list\[str\] # ordered best → worst class Metrics(BaseModel): recall\_at\_k: float ndcg\_at\_k: float mrr: float k: int class ExperimentResult(BaseModel): config: ExperimentConfig metrics: Metrics # {{/docs-fragment data\_types}} class ComparisonReport(BaseModel): results: list\[ExperimentResult\] def best\_by(self, metric: str = "recall\_at\_k") -> ExperimentResult: return max(self.results, key=lambda r: getattr(r.metrics, metric)) def summary(self) -> str: header = f"{'Experiment':<30} {'Model':<18} {'Recall@K':>10} {'NDCG@K':>8} {'MRR':>7}" sep = "─" \* len(header) rows = \[header, sep\] for r in sorted(self.results, key=lambda x: -x.metrics.recall\_at\_k): rows.append( f"{r.config.name:<30} " f"{r.config.model.value:<18} " f"{r.metrics.recall\_at\_k:>10.3f} " f"{r.metrics.ndcg\_at\_k:>8.3f} " f"{r.metrics.mrr:>7.3f}" ) return "\\n".join(rows) # ───────────────────────────────────────────────────────────────────────────── # Cached model loaders # ───────────────────────────────────────────────────────────────────────────── # These functions are at module level so they are shared across all tasks that # run on the same warm container (via ReusePolicy). lru\_cache(maxsize=1) means # the model is loaded from disk/HuggingFace exactly once per container process # and kept in GPU memory for every subsequent task dispatch to that container. @lru\_cache(maxsize=1) def \_colpali\_model(): """Load ColPali-v1.2 into GPU memory and cache the result. device\_map= is the correct loading pattern for ColPali's PaliGemma backbone; it handles weight placement via accelerate. torch.compile is skipped — ColPali is GPU-compute-bound and the DynamicBatcher's cross- invocation batching is the primary GPU utilisation mechanism. """ import torch from colpali\_engine.models import ColPali, ColPaliProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = ColPali.from\_pretrained( "vidore/colpali-v1.2", torch\_dtype=torch.bfloat16, device\_map=device, ) processor = ColPaliProcessor.from\_pretrained("vidore/colpali-v1.2") return model, processor, device @lru\_cache(maxsize=1) def \_siglip\_model(): """Load SigLIP SO400M into GPU memory, compile it, and cache the result. torch.compile (mode="reduce-overhead") fuses the vision and text encoder transformer layers into optimised CUDA kernels. As with ColPali, the compilation overhead is paid once per warm container lifetime. """ import torch from transformers import AutoModel, AutoProcessor device = "cuda" if torch.cuda.is\_available() else "cpu" model = AutoModel.from\_pretrained("google/siglip-so400m-patch14-224").to(device) if device == "cuda": model = torch.compile(model, mode="reduce-overhead") processor = AutoProcessor.from\_pretrained("google/siglip-so400m-patch14-224") return model, processor, device @lru\_cache(maxsize=1) def \_ocr\_model(): """Load the doctr OCR predictor onto GPU and cache it. doctr's ocr\_predictor bundles a detection model (DBNet) and a recognition model (CRNN/SAR) into a single callable. pretrained=True downloads both model weights from doctr's model zoo on first use. """ import torch from doctr.models import ocr\_predictor predictor = ocr\_predictor(pretrained=True) if torch.cuda.is\_available(): predictor = predictor.cuda() return predictor # ───────────────────────────────────────────────────────────────────────────── # Search batcher singletons # ───────────────────────────────────────────────────────────────────────────── # One DynamicBatcher per model, shared across all concurrent search task # invocations on the same warm container (concurrency=3). Queries from every # concurrent caller are aggregated into a single GPU batch, maximizing # throughput compared to each invocation running its own forward pass. # # Initialised lazily on the first search call via double-checked locking and # lives for the container's lifetime. The process\_fn runs GPU work via # asyncio.to\_thread so the aggregation loop can continue collecting queries # from other callers while the GPU processes the current batch. # # File is not hashable so alru\_cache cannot be used here; module-level state # with asyncio.Lock is the correct pattern. # # Assumption: index\_colpali/index\_siglip use cache="auto", so the same corpus # always produces the same index File across all callers on this container. If # the index file ever changed between calls, the batcher would silently continue # using the corpus embeddings loaded from the first call. \_colpali\_batcher: DynamicBatcher | None = None \_colpali\_batcher\_lock = asyncio.Lock() \_siglip\_batcher: DynamicBatcher | None = None \_siglip\_batcher\_lock = asyncio.Lock() async def \_get\_colpali\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level ColPali search batcher, creating it on first call.""" global \_colpali\_batcher if \_colpali\_batcher is not None: return \_colpali\_batcher async with \_colpali\_batcher\_lock: if \_colpali\_batcher is not None: return \_colpali\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, n\_patches, dim) index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_colpali\_model() corpus\_emb = corpus\_emb.to(device, dtype=torch.float32) async def colpali\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: query\_inputs = processor.process\_queries(\[q.text for q in batch\]) query\_inputs = {k: v.to(device) for k, v in query\_inputs.items()} with torch.no\_grad(): query\_embs = model(\*\*query\_inputs).float() # (B, T, D) query\_chunk = 8 n\_pages = corpus\_emb.shape\[0\] all\_scores = torch.empty(len(batch), n\_pages, device=device) for start in range(0, len(batch), query\_chunk): chunk = query\_embs\[start : start + query\_chunk\] all\_scores\[start : start + query\_chunk\] = ( torch.einsum("ctd,pjd->ctpj", chunk, corpus\_emb) .max(dim=3).values .sum(dim=1) ) sorted\_indices = all\_scores.argsort(dim=1, descending=True).cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] # Run GPU work in a thread so the event loop — and the batcher's # aggregation loop — remain unblocked while the GPU is busy. return await asyncio.to\_thread(\_gpu\_work) batcher: DynamicBatcher\[PageQuery, list\[str\]\] = DynamicBatcher( process\_fn=colpali\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_colpali\_batcher = batcher return \_colpali\_batcher async def \_get\_siglip\_search\_batcher(index\_file: File) -> DynamicBatcher: """Return the process-level SigLIP search batcher, creating it on first call.""" global \_siglip\_batcher if \_siglip\_batcher is not None: return \_siglip\_batcher async with \_siglip\_batcher\_lock: if \_siglip\_batcher is not None: return \_siglip\_batcher import torch data = await \_load\_npz(index\_file) corpus\_emb = torch.from\_numpy(data\["embeddings"\]) # (n\_pages, dim), L2-normalised index\_page\_ids: list\[str\] = list(data\["page\_ids"\]) model, processor, device = \_siglip\_model() corpus\_emb = corpus\_emb.to(device) async def siglip\_process\_fn(batch: list\[PageQuery\]) -> list\[list\[str\]\]: def \_gpu\_work() -> list\[list\[str\]\]: text\_inputs = processor( text=\[q.text for q in batch\], return\_tensors="pt", padding=True, truncation=True, ).to(device) with torch.no\_grad(): text\_out = model.text\_model(\*\*text\_inputs) query\_embs = text\_out.pooler\_output # (B, dim) query\_embs = query\_embs / query\_embs.norm(dim=-1, keepdim=True) scores\_matrix = corpus\_emb @ query\_embs.T # (n\_pages, B) sorted\_indices = scores\_matrix.argsort(dim=0, descending=True).T.cpu().tolist() return \[\[index\_page\_ids\[j\] for j in ranked\] for ranked in sorted\_indices\] return await asyncio.to\_thread(\_gpu\_work) batcher = DynamicBatcher( process\_fn=siglip\_process\_fn, target\_batch\_cost=128, max\_batch\_size=128, batch\_timeout\_s=0.05, default\_cost=1, prefetch\_batches=2, ) await batcher.start() \_siglip\_batcher = batcher return \_siglip\_batcher # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── def \_batches(items: list, batch\_size: int): """Yield successive fixed-size batches from a list.""" for start in range(0, len(items), batch\_size): yield items\[start : start + batch\_size\] def \_load\_image\_sync(f: File) -> PILImage.Image: """Blocking download + decode. Intended to be called from a thread pool.""" with f.open\_sync("rb") as fh: data = fh.read() return PILImage.open(BytesIO(data)).convert("RGB") async def \_load\_image(f: File) -> PILImage.Image: """Download and decode a page image in a thread-pool worker. asyncio.to\_thread runs \_load\_image\_sync in a real OS thread so that blocking network I/O can overlap with GPU-bound forward passes when images are pre-submitted via loop.run\_in\_executor before the GPU kernel. """ return await asyncio.to\_thread(\_load\_image\_sync, f) async def \_load\_npz(index\_file: File) -> np.lib.npyio.NpzFile: """Download an index File to a local temp path and open with np.load.""" with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp: async with index\_file.open("rb") as fh: tmp.write(bytes(await fh.read())) return np.load(tmp.name) def \_dcg(relevances: list\[int\]) -> float: return sum(rel / math.log2(rank + 2) for rank, rel in enumerate(relevances)) # ───────────────────────────────────────────────────────────────────────────── # Tasks — data loading # ───────────────────────────────────────────────────────────────────────────── @driver.task(cache="auto", retries=3) async def load\_vidore\_pages(subset: str = "docvqa", max\_pages: int = 200) -> PageDataset: """ Load a ViDoRe benchmark subset and store page images in Flyte's blob store. Supports two dataset formats: Legacy (subsampled) — single 'test' split with one row per (query, page) pair; fields: image, query, image\_filename. streaming=True reads only the rows requested via islice — no full-shard download. Datasets: vidore/docvqa\_test\_subsampled, vidore/infovqa\_test\_subsampled V3 — separate corpus / queries / qrels splits following the BEIR retrieval benchmark format. corpus contains page images; queries contains question text; qrels maps query IDs to relevant corpus page IDs (many-to-many). Datasets: vidore/vidore\_v3\_finance\_en (~2 942 pages, 1 854 queries) The first call uploads page images to Flyte's blob store and caches the PageDataset; every subsequent call with the same arguments returns the cached result instantly. retries=3 guards against transient HuggingFace network failures. Available subsets: "docvqa", "infovqa", "vidore\_v3\_finance\_en" """ from datasets import load\_dataset subset\_map = { "docvqa": "vidore/docvqa\_test\_subsampled", "infovqa": "vidore/infovqa\_test\_subsampled", "vidore\_v3\_finance\_en": "vidore/vidore\_v3\_finance\_en", } dataset\_name = subset\_map.get(subset, f"vidore/{subset}\_test\_subsampled") # V3 datasets ship with separate corpus / queries / qrels splits. \_V3\_SUBSETS = {"vidore\_v3\_finance\_en"} if subset in \_V3\_SUBSETS: # ── V3 format ───────────────────────────────────────────────────────── # corpus / queries / qrels are HuggingFace configs (name=), not splits. # corpus uses streaming=True so images are decoded one at a time — # loading all 2 942 rows eagerly would hold gigabytes of PIL images in # the driver's RAM simultaneously. qrels and queries are text-only and # small enough to load fully into memory. corpus\_ds = load\_dataset(dataset\_name, name="corpus", split="test", streaming=True) qrels\_ds = load\_dataset(dataset\_name, name="qrels", split="test") queries\_ds = load\_dataset(dataset\_name, name="queries", split="test") # Normalise field names — V3 follows BEIR convention (hyphenated ids). def \_col(ds, \*candidates): cols = set(ds.column\_names) for c in candidates: if c in cols: return c raise KeyError(f"None of {candidates} found in columns {cols}") corpus\_id\_col = \_col(corpus\_ds, "corpus-id", "corpus\_id", "id", "\_id") query\_id\_col = \_col(queries\_ds, "query-id", "query\_id", "id", "\_id") query\_text\_col = \_col(queries\_ds, "query", "text") qrel\_qid\_col = \_col(qrels\_ds, "query-id", "query\_id") qrel\_cid\_col = \_col(qrels\_ds, "corpus-id", "corpus\_id") # Slice corpus to max\_pages, upload each image to Flyte blob store. page\_ids: list\[str\] = \[\] page\_files: list\[File\] = \[\] corpus\_id\_to\_page\_id: dict\[str, str\] = {} for i, row in enumerate(islice(corpus\_ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue cid = str(row\[corpus\_id\_col\]) page\_id = f"{subset}\_{i:04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) corpus\_id\_to\_page\_id\[cid\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) # Build query\_id → relevant page\_id from qrels (first match wins). # Only keep relevance judgements whose corpus page is in our slice. qrel\_map: dict\[str, str\] = {} for row in qrels\_ds: qid = str(row\[qrel\_qid\_col\]) cid = str(row\[qrel\_cid\_col\]) if cid in corpus\_id\_to\_page\_id and qid not in qrel\_map: qrel\_map\[qid\] = corpus\_id\_to\_page\_id\[cid\] # Collect queries that have at least one relevant page in our slice. queries: list\[PageQuery\] = \[\] for row in queries\_ds: qid = str(row\[query\_id\_col\]) if qid not in qrel\_map: continue queries.append( PageQuery( query\_id=qid, text=str(row\[query\_text\_col\]), relevant\_page\_id=qrel\_map\[qid\], ) ) else: # ── Legacy format ───────────────────────────────────────────────────── # Single 'test' split with one row per (query, page) pair. ds = load\_dataset(dataset\_name, split="test", streaming=True) page\_ids = \[\] page\_files = \[\] queries = \[\] seen\_pages: dict\[str, str\] = {} # image\_filename → page\_id for i, row in enumerate(islice(ds, max\_pages)): img = row.get("image") if not isinstance(img, PILImage.Image): continue filename: str = row.get("image\_filename") or f"page\_{i}" query\_text: str = row.get("query", "") if not query\_text: continue # Each unique page is uploaded exactly once; multiple queries may # share the same page (same image\_filename). if filename not in seen\_pages: page\_id = f"{subset}\_{len(page\_ids):04d}" with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f: tmp\_path = f.name img.convert("RGB").save(tmp\_path, format="JPEG") del img # free PIL memory before upload page\_file = await File.from\_local(tmp\_path) os.unlink(tmp\_path) seen\_pages\[filename\] = page\_id page\_ids.append(page\_id) page\_files.append(page\_file) else: page\_id = seen\_pages\[filename\] queries.append( PageQuery( query\_id=f"q{i:04d}", text=query\_text, relevant\_page\_id=page\_id, ) ) print(f"Loaded {len(page\_ids)} unique pages, {len(queries)} queries", flush=True) return PageDataset(page\_ids=page\_ids, page\_files=page\_files, queries=queries) # ───────────────────────────────────────────────────────────────────────────── # Tasks — indexing # ───────────────────────────────────────────────────────────────────────────── @colpali\_indexer.task(cache="auto", retries=2) async def index\_colpali(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with ColPali-v1.2 and save the multi-vector index. ColPali skips OCR entirely. It feeds the raw page image into PaliGemma (a vision-language model) and produces one embedding vector per image patch — roughly 1,024 patches per page, each of dimension 128. \_colpali\_model() is an lru\_cache'd loader. On a cold container, it downloads and loads the model once. On a warm container (kept alive by ReusePolicy), it returns the already-loaded model instantly from cache — no repeated ~7 GB download. The index is stored as a .npz file in Flyte's blob store: embeddings — float32, shape (n\_pages, n\_patches, dim) page\_ids — matching page ID strings cache="auto" + retries=2: the result is stored permanently on success; transient failures (e.g. HuggingFace rate limits) are retried twice. """ import torch model, processor, device = \_colpali\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 4)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor.process\_images(images) inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no\_grad(): emb = model(\*\*inputs) # (batch, n\_patches, dim) all\_embeddings.append(emb.cpu().float().numpy()) print(f"ColPali: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, n\_patches, dim) out\_path = os.path.join(tempfile.gettempdir(), "colpali\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @siglip\_indexer.task(cache="auto", retries=2) async def index\_siglip(page\_ids: list\[str\], page\_files: list\[File\]) -> File: """ Encode every page with SigLIP SO400M and save the single-vector index. SigLIP (2023) is Google's successor to CLIP, trained with sigmoid loss instead of softmax — avoiding the normalisation bottleneck that limits CLIP's scalability. Produces one global embedding per page. \_siglip\_model() caches the model across warm container reuses. The index is stored as a .npz file: embeddings — float32, shape (n\_pages, dim), L2-normalised page\_ids — matching page ID strings """ import torch model, processor, device = \_siglip\_model() loop = asyncio.get\_running\_loop() batches = list(\_batches(page\_files, 8)) n\_batches = len(batches) # Submit the first batch to the thread pool before entering the loop so # that downloads are already in flight when we first await them. prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[0\]\] all\_embeddings: list\[np.ndarray\] = \[\] for batch\_idx in range(n\_batches): images = list(await asyncio.gather(\*prefetch)) # Submit next batch downloads immediately — OS threads run these in # parallel with the GPU forward pass below. if batch\_idx + 1 < n\_batches: prefetch = \[loop.run\_in\_executor(None, \_load\_image\_sync, f) for f in batches\[batch\_idx + 1\]\] inputs = processor(images=images, return\_tensors="pt", padding=True).to(device) with torch.no\_grad(): outputs = model.vision\_model(\*\*inputs) emb = outputs.pooler\_output # (batch, dim) emb = emb / emb.norm(dim=-1, keepdim=True) # L2 normalise all\_embeddings.append(emb.cpu().float().numpy()) print(f"SigLIP: indexed batch {batch\_idx + 1}/{n\_batches}", flush=True) embeddings = np.concatenate(all\_embeddings, axis=0) # (n\_pages, dim) out\_path = os.path.join(tempfile.gettempdir(), "siglip\_index.npz") np.savez(out\_path, embeddings=embeddings, page\_ids=np.array(page\_ids)) return await File.from\_local(out\_path) @ocr\_engine.task(cache="auto") async def extract\_page\_texts(page\_files: list\[File\]) -> list\[str\]: """ OCR every page with doctr on GPU to produce a text-only baseline. doctr bundles DBNet (detection) + CRNN/SAR (recognition) into a single callable predictor. Pages are downloaded in parallel then fed in batches of ocr\_batch\_size. asyncio.to\_thread keeps the event loop unblocked during GPU inference. Result structure: result.pages\[i\].blocks\[j\].lines\[k\].words\[l\].value Cached: the same corpus is OCR'd at most once across all experiments that use the OCR+BM25 backend. """ import gc predictor = \_ocr\_model() # Process in batches: download each batch just-in-time so only # ocr\_batch\_size images are in memory at once instead of all 2 000. ocr\_batch\_size = 8 total = len(page\_files) texts: list\[str\] = \[\] for start in range(0, total, ocr\_batch\_size): batch\_files = page\_files\[start : start + ocr\_batch\_size\] batch\_images = list( await asyncio.gather(\*\[asyncio.to\_thread(\_load\_image\_sync, f) for f in batch\_files\]) ) batch\_np = \[np.array(img) for img in batch\_images\] del batch\_images result = await asyncio.to\_thread(predictor, batch\_np) del batch\_np for page\_output in result.pages: texts.append( "\\n".join( " ".join(word.value for word in line.words) for block in page\_output.blocks for line in block.lines ) ) del result gc.collect() print(f"OCR: processed {min(start + ocr\_batch\_size, total)}/{total} pages", flush=True) return texts # ───────────────────────────────────────────────────────────────────────────── # Tasks — search # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment search\_colpali}} @colpali\_indexer.task async def search\_colpali( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using ColPali MaxSim late interaction via DynamicBatcher. MaxSim score for page p given query q: score(q, p) = Σ\_{t ∈ query tokens} max\_{j ∈ page patches} (q\_t · p\_j) Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_colpali invocations on the same warm container (concurrency=8) into a single GPU batch. This keeps the GPU saturated rather than running one small batch per caller. The batcher's process\_fn runs GPU work in asyncio.to\_thread, so the aggregation loop stays live while the GPU encodes and scores. """ batcher = await \_get\_colpali\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] # {{/docs-fragment search\_colpali}} @siglip\_indexer.task async def search\_siglip( index\_file: File, queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using SigLIP cosine similarity via DynamicBatcher. Each query is submitted to the process-level DynamicBatcher, which aggregates queries from all concurrent search\_siglip invocations on the same warm container (concurrency=3) into a single GPU batch. SigLIP's single-vector embeddings make full vectorisation safe — the scores matrix (n\_pages x n\_queries) is small enough to materialise in one GPU call regardless of batch size. """ batcher = await \_get\_siglip\_search\_batcher(index\_file) futures = await batcher.submit\_batch(queries) all\_ranked: list\[list\[str\]\] = list(await asyncio.gather(\*futures)) return \[\ RetrievalResult(query\_id=q.query\_id, ranked\_page\_ids=ranked\[:top\_k\])\ for q, ranked in zip(queries, all\_ranked)\ \] @driver.task async def search\_bm25( page\_texts: list\[str\], page\_ids: list\[str\], queries: list\[PageQuery\], top\_k: int, ) -> list\[RetrievalResult\]: """ Retrieve pages using BM25 over OCR'd text. The standard keyword-based baseline. No GPU required; strong on text-dense pages, weak on visual content that Tesseract cannot read. """ tokenized = \[text.lower().split() for text in page\_texts\] bm25 = BM25Okapi(tokenized) results: list\[RetrievalResult\] = \[\] for q in queries: scores = bm25.get\_scores(q.text.lower().split()) ranked = sorted(range(len(page\_ids)), key=lambda i: -scores\[i\])\[:top\_k\] results.append( RetrievalResult( query\_id=q.query\_id, ranked\_page\_ids=\[page\_ids\[i\] for i in ranked\], ) ) return results # ───────────────────────────────────────────────────────────────────────────── # Tasks — evaluation # ───────────────────────────────────────────────────────────────────────────── @driver.task async def evaluate( results: list\[RetrievalResult\], ground\_truth: list\[PageQuery\], k: int, ) -> Metrics: """ Compute Recall@K, NDCG@K, and MRR for a single retrieval model. Recall@K — was the correct page in the top-K results? NDCG@K — normalised discounted cumulative gain; rewards earlier hits. MRR — mean reciprocal rank of the first correct result. All three are averaged over all queries. Higher is better. """ gt\_map = {q.query\_id: q.relevant\_page\_id for q in ground\_truth} recall\_vals, ndcg\_vals, mrr\_vals = \[\], \[\], \[\] for r in results: relevant = gt\_map.get(r.query\_id, "") top = r.ranked\_page\_ids\[:k\] recall\_vals.append(1.0 if relevant in top else 0.0) rels = \[1 if pid == relevant else 0 for pid in top\] idcg = \_dcg(\[1\]) # ideal: correct page at rank 1 ndcg\_vals.append(\_dcg(rels) / idcg if idcg > 0 else 0.0) rr = 0.0 for rank, pid in enumerate(r.ranked\_page\_ids, start=1): if pid == relevant: rr = 1.0 / rank break mrr\_vals.append(rr) return Metrics( recall\_at\_k=float(np.mean(recall\_vals)), ndcg\_at\_k=float(np.mean(ndcg\_vals)), mrr=float(np.mean(mrr\_vals)), k=k, ) # ───────────────────────────────────────────────────────────────────────────── # Tasks — report # ───────────────────────────────────────────────────────────────────────────── @driver.task(report=True) async def generate\_report(report: ComparisonReport) -> None: """ Emit an interactive HTML report visible in the Flyte UI. report=True marks this task as a reporting task. Flyte renders the HTML returned via flyte.report.replace.aio() directly in the execution detail page — no separate dashboard or export step required. The report contains: - Summary cards: experiment count, best model, best Recall@K. - Grouped bar chart: Recall@K, NDCG@K, MRR side-by-side per experiment. - Ranked results table with all three metrics. """ sorted\_results = sorted(report.results, key=lambda r: -r.metrics.recall\_at\_k) best = sorted\_results\[0\] labels = \[r.config.name for r in sorted\_results\] recall\_vals = \[r.metrics.recall\_at\_k for r in sorted\_results\] ndcg\_vals = \[r.metrics.ndcg\_at\_k for r in sorted\_results\] mrr\_vals = \[r.metrics.mrr for r in sorted\_results\] table\_rows = "".join( f""" {r.config.name} {r.config.model.value} {r.metrics.recall\_at\_k:.3f} {r.metrics.ndcg\_at\_k:.3f} {r.metrics.mrr:.3f} {r.metrics.k} """ for r in sorted\_results ) html = f""" Visual Document Retrieval — Results

Visual Document Retrieval — Experiment Comparison

ViDoRe benchmark · {len(report.results)} experiment(s)

{len(report.results)}
Experiments
{best.config.name}
Best by Recall@K
{best.metrics.recall\_at\_k:.3f}
Best Recall@{best.metrics.k}
{best.metrics.ndcg\_at\_k:.3f}
Best NDCG@{best.metrics.k}
{best.metrics.mrr:.3f}
Best MRR

Metrics by Experiment

Ranked Results

{table\_rows}
ExperimentModel Recall@KNDCG@KMRRK
""" await flyte.report.replace.aio(html) await flyte.report.flush.aio() # ───────────────────────────────────────────────────────────────────────────── # Experiment orchestration # ───────────────────────────────────────────────────────────────────────────── # {{docs-fragment run\_experiment}} @driver.task async def run\_experiment(config: ExperimentConfig, dataset: PageDataset) -> ExperimentResult: """ End-to-end retrieval pipeline for a single ExperimentConfig. Flyte v2's dynamic execution means this driver task can call GPU tasks (index\_colpali, search\_colpali) based on the runtime value of config.model — no static DAG wiring required. The if/elif is plain Python; Flyte schedules the selected sub-tasks on the appropriate environment. Caching: two experiments that share the same model and corpus (e.g. ColPali at top\_k=5 and top\_k=10) will hit the same cached index. GPU work is paid at most once per (model, corpus) pair across all experiments. Search queries are sharded into chunks of SEARCH\_SHARD\_SIZE and dispatched as concurrent task invocations. All shards land on the single warm container (replicas=1) and feed the same DynamicBatcher simultaneously, keeping the GPU saturated throughout search rather than processing one large sequential batch from a single caller. flyte.group wraps each experiment in a named span in the Flyte UI, making it easy to compare latencies and drill into individual runs. """ SEARCH\_SHARD\_SIZE = 256 with flyte.group(config.name): if config.model == RetrievalModel.COLPALI: index\_file = await index\_colpali(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_colpali(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] elif config.model == RetrievalModel.SIGLIP: index\_file = await index\_siglip(dataset.page\_ids, dataset.page\_files) shards = list(\_batches(dataset.queries, SEARCH\_SHARD\_SIZE)) shard\_results = await asyncio.gather( \*\[search\_siglip(index\_file, shard, config.top\_k) for shard in shards\] ) results = \[r for shard in shard\_results for r in shard\] else: # RetrievalModel.OCR\_BM25 page\_texts = await extract\_page\_texts(dataset.page\_files) results = await search\_bm25(page\_texts, dataset.page\_ids, dataset.queries, config.top\_k) metrics = await evaluate(results, dataset.queries, config.top\_k) return ExperimentResult(config=config, metrics=metrics) # {{/docs-fragment run\_experiment}} # {{docs-fragment compare\_experiments}} @driver.task async def compare\_experiments( configs: list\[ExperimentConfig\], subset: str = "docvqa", max\_pages: int = 200, ) -> ComparisonReport: """ Fan out over all experiment configs and return a ranked comparison table. The dataset is loaded once and shared across all experiments. Each config runs as a concurrent Flyte task via asyncio.gather. Experiments that share a model reuse the cached index — you only pay GPU time for new work. On completion, generate\_report emits an interactive Chart.js HTML report visible directly in the Flyte execution detail page. Default dataset: vidore\_v3\_finance\_en (~2 942 corpus pages, 1 854 queries) with max\_pages=2 000 to exercise the GPU pipeline at scale. """ dataset = await load\_vidore\_pages(subset=subset, max\_pages=max\_pages) # All experiments launch concurrently. Shared cached outputs (same model, # same corpus) are served from cache rather than recomputed. experiment\_coros = \[run\_experiment(config=cfg, dataset=dataset) for cfg in configs\] results: list\[ExperimentResult\] = list(await asyncio.gather(\*experiment\_coros)) report = ComparisonReport(results=results) print(report.summary()) best = report.best\_by("recall\_at\_k") print(f"\\nBest by Recall@{best.metrics.k}: {best.config.name}") # Emit the interactive HTML report in the Flyte UI. await generate\_report(report) return report # {{/docs-fragment compare\_experiments}} # ───────────────────────────────────────────────────────────────────────────── # Entry point # ───────────────────────────────────────────────────────────────────────────── if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Define the experiment grid. Each ExperimentConfig is one point in the # design space. Adding a new model or varying top\_k is one line here — # no task code changes required. # # ColPali appears twice with different top\_k values. The cache ensures # index\_colpali runs only once and both experiments share that result. # {{docs-fragment grid}} configs = \[\ ExperimentConfig(name="colpali-top5", model=RetrievalModel.COLPALI, top\_k=5),\ ExperimentConfig(name="colpali-top10", model=RetrievalModel.COLPALI, top\_k=10),\ ExperimentConfig(name="siglip-top5", model=RetrievalModel.SIGLIP, top\_k=5),\ ExperimentConfig(name="ocr-bm25-top5", model=RetrievalModel.OCR\_BM25, top\_k=5),\ \] # {{/docs-fragment grid}} run = flyte.with\_runcontext().run( compare\_experiments, configs=configs, subset="vidore\_v3\_finance\_en", max\_pages=2000, ) print(f"Run URL: {run.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/multimodal-retrieval-evaluation/retrieval\_eval.py\* From the \[example directory\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/multimodal-retrieval-evaluation): \`\`\` cd v2/tutorials/multimodal-retrieval-evaluation python retrieval\_eval.py \`\`\` When the run completes, open the \`generate\_report\` task in the UI to see the summary cards, the grouped Recall@K / NDCG@K / MRR bar chart, and the ranked results table. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents === # Agents Tutorials for building agentic workflows and autonomous LLM-powered systems. ### \*\*Agents > Autoresearch agent\*\* Run an autonomous research loop that drives Claude Code in a GPU container to run experiments, then commits results and opens a pull request. ### \*\*Agents > Coding agent\*\* Securely execute and iterate on LLM-generated code using a code agent with error reflection and retry logic. ### \*\*Agents > Competitive intelligence agent\*\* Fan out across competitors, extract source-cited market deltas with the You.com Search API, and build a knowledge-graph-ready intelligence table. ### \*\*Agents > Compliance monitoring agent\*\* Monitor trusted regulatory sources with the You.com Research API and route citation-precise findings to the right team. ### \*\*Agents > Deep research\*\* Build an agentic workflow for deep research with multi-step reasoning and evaluation. ### \*\*Agents > Field data enrichment agent\*\* Enrich geo-tagged operational events with real-world public context using the You.com Search API with country and freshness targeting. ### \*\*Agents > MLE Bot: an autonomous ML engineer\*\* An autonomous ML agent that designs, runs, and iterates on experiments using Flyte's durable sandbox for safe LLM-generated code execution. ### \*\*Agents > Support resolution agent\*\* Ground support tickets in fresh public sources via the You.com Research API and draft cited, customer-ready replies for human review. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/autoresearch === # Autoresearch agent > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/autoresearch). This tutorial wraps an autonomous AI research loop in a single Flyte task. The task spins up a GPU container, installs the \[Claude Code\](https://docs.anthropic.com/en/docs/claude-code/overview) CLI, clones a research repository, and points Claude Code at a \`program.md\` brief. The agent runs experiments to improve a model, writes results to disk, and the task then commits the changes and opens a pull request — with a progress plot rendered both in the PR and in the Flyte UI. It's an example of using Flyte as durable infrastructure for long-running, autonomous agent work: - \*\*A GPU \`TaskEnvironment\`\*\* with the API-key and GitHub secrets the agent needs. - \*\*\`report=True\`\*\* to stream a progress plot into the Flyte UI. - \*\*A reconnecting \`run.wait()\`\*\* loop in the driver so a dropped client connection doesn't lose track of a multi-hour run. > \[!WARNING\] > This example drives a coding agent that executes arbitrary code and pushes commits to a GitHub repository. Run it against a repository you control, and review the constants described below before launching. ## Define the container image The image is kept in its own \`\_image.py\` module so edits to the agent logic in \`run.py\` don't invalidate the image cache. Node.js and the Claude Code CLI are installed at run time (see below) to keep the image small. \`\`\` # /// script # requires-python = ">=3.11" # dependencies = \[\ # "flyte>=2.0.0b22",\ # "PyGithub>=2.5.0",\ # "matplotlib>=3.7.0",\ # "pandas>=2.0.0",\ # \] # /// # # Stable image definition — kept separate from run.py so edits to run.py # don't invalidate the image cache. Only touch this file when the image itself needs to change. import flyte # {{docs-fragment image}} image = ( flyte.Image.from\_uv\_script(\_\_file\_\_, name="autoresearch-agent", pre=True) .with\_apt\_packages("git") ) # {{/docs-fragment image}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autoresearch/\_image.py\* ## Define the task environment The task needs a GPU, a generous disk for the cloned repo and model weights, and two secrets: a GitHub token (to clone and push) and an Anthropic API key (for Claude Code). \`\`\` # /// script # requires-python = ">=3.11" # dependencies = \[\ # "flyte>=2.0.0b22",\ # "PyGithub>=2.5.0",\ # "matplotlib>=3.7.0",\ # \] # /// """ AutoResearch Agent - Runs the autoresearch workflow using Claude Code CLI in a GPU environment. This agent: 1. Starts a GPU-enabled container 2. Installs Claude Code CLI 3. Clones the autoresearch repository 4. Points Claude Code at program.md as the prompt and lets it run 5. Commits the result (CSV + code changes in train/) and creates a PR """ import os import shlex import subprocess from dataclasses import dataclass from pathlib import Path from typing import Optional from github import Auth, Github import flyte import flyte.report from \_image import image as autoresearch\_image GITHUB\_USERNAME = "parnianz" GITHUB\_EMAIL = "parnianzargham@gmail.com" AUTORESEARCH\_REPO\_URL = "https://github.com/unionai-oss/autoresearch.git" AUTORESEARCH\_REPO\_FULL\_NAME = "unionai-oss/autoresearch" # {{docs-fragment env}} autoresearch\_env = flyte.TaskEnvironment( name="autoresearch-agent", resources=flyte.Resources( cpu=8, memory="32Gi", gpu="T4:1", disk="100Gi", ), secrets=\[\ flyte.Secret(key="github\_token", as\_env\_var="GITHUB\_TOKEN"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=autoresearch\_image, ) # {{/docs-fragment env}} # {{docs-fragment result}} @dataclass class AutoResearchResult: """Result of the autoresearch run.""" pr\_url: str pr\_number: int branch\_name: str files\_changed: list\[str\] success: bool error\_message: Optional\[str\] = None # {{/docs-fragment result}} def clone\_repository(repo\_url: str, work\_dir: Path, github\_token: str) -> Path: """Clone the autoresearch repository with authentication.""" repo\_name = repo\_url.rstrip("/").split("/")\[-1\].replace(".git", "") repo\_path = work\_dir / repo\_name # Inject token into HTTPS URL for authentication authenticated\_url = repo\_url.replace( "https://", f"https://{GITHUB\_USERNAME}:{github\_token}@" ) if repo\_path.exists(): subprocess.run(\["git", "pull"\], cwd=repo\_path, check=True) else: subprocess.run(\["git", "clone", authenticated\_url, str(repo\_path)\], check=True) return repo\_path # {{docs-fragment task}} @autoresearch\_env.task(report=True) async def run\_autoresearch() -> AutoResearchResult: """ Run the autoresearch workflow end-to-end. Steps: - Clone https://github.com/unionai-oss/autoresearch - Configure git identity - Create a new branch - Run Claude Code CLI with program.md as the prompt - Commit results (CSV + train/ changes) - Push and open a PR against the autoresearch repo """ github\_token = os.environ\["GITHUB\_TOKEN"\] anthropic\_api\_key = os.environ\["ANTHROPIC\_API\_KEY"\] # --- Install Node.js + Claude Code at runtime (keeps image small and submission fast) --- import tarfile import urllib.request as \_urllib subprocess.run(\["apt-get", "update", "-y"\], check=False) subprocess.run(\["apt-get", "install", "-y", "git"\], check=False) node\_url = "https://nodejs.org/dist/v20.19.0/node-v20.19.0-linux-x64.tar.gz" node\_tar = Path("/tmp/node.tar.gz") print(f"Downloading Node.js from {node\_url}...", flush=True) \_urllib.urlretrieve(node\_url, node\_tar) size\_mb = node\_tar.stat().st\_size / 1024 / 1024 print(f"Downloaded {size\_mb:.1f} MB to {node\_tar}", flush=True) if size\_mb < 1: raise RuntimeError(f"Node.js download appears empty/corrupt ({size\_mb:.2f} MB) — network may be restricted") node\_dir = Path("/tmp/node") node\_dir.mkdir(exist\_ok=True) print("Extracting Node.js...", flush=True) with tarfile.open(node\_tar, "r:gz") as tar: members = \[m for m in tar.getmembers() if m.name.split("/", 1)\[-1\]\] for m in members: m.name = m.name.split("/", 1)\[-1\] tar.extractall(str(node\_dir), members=\[m for m in members if m.name\]) # Add node/npm to PATH for this process and all subprocesses node\_bin = str(node\_dir / "bin") os.environ\["PATH"\] = node\_bin + ":" + os.environ.get("PATH", "") print(f"Node version: {subprocess.run(\['node', '--version'\], capture\_output=True, text=True).stdout.strip()}", flush=True) npm\_prefix = "/tmp/npm-global" Path(npm\_prefix).mkdir(exist\_ok=True) subprocess.run(\["npm", "install", "-g", "--prefix", npm\_prefix, "@anthropic-ai/claude-code"\], check=True) os.environ\["PATH"\] = str(Path(npm\_prefix) / "bin") + ":" + os.environ\["PATH"\] print("Node.js + Claude Code installed.", flush=True) # --- Clone repo --- work\_dir = Path("/tmp/autoresearch\_workspace") work\_dir.mkdir(exist\_ok=True, parents=True) repo\_path = clone\_repository(AUTORESEARCH\_REPO\_URL, work\_dir, github\_token) # --- Git identity --- subprocess.run( \["git", "config", "--global", "user.email", GITHUB\_EMAIL\], check=True ) subprocess.run( \["git", "config", "--global", "user.name", GITHUB\_USERNAME\], check=True ) # --- Create branch --- import time as \_time branch\_name = f"autoresearch/claude-run-{int(\_time.time())}" try: subprocess.run( \["git", "checkout", "-b", branch\_name\], cwd=repo\_path, check=True, ) except subprocess.CalledProcessError: subprocess.run( \["git", "checkout", branch\_name\], cwd=repo\_path, check=True, ) # --- Read program.md to use as the Claude Code prompt --- program\_md = repo\_path / "program.md" if not program\_md.exists(): raise FileNotFoundError( f"program.md not found in {repo\_path}. " "Make sure the autoresearch repo has a program.md at its root." ) program\_md\_content = program\_md.read\_text() print(f"Loaded prompt from program.md ({len(program\_md\_content)} chars)") # {{/docs-fragment task}} # Install repo dependencies before handing off to Claude for pip\_cmd in \[\ \["pip", "install", "-e", "."\],\ \["pip", "install", "-r", "requirements.txt"\],\ \]: req\_file = repo\_path / pip\_cmd\[-1\] if pip\_cmd\[-1\].startswith("req") else None if req\_file is None or req\_file.exists(): dep\_result = subprocess.run( pip\_cmd, cwd=repo\_path, capture\_output=True, text=True ) print(f"{' '.join(pip\_cmd)}:\\n{dep\_result.stdout}", flush=True) if dep\_result.returncode != 0: print(f"(non-fatal) {dep\_result.stderr}", flush=True) # Wrap the program.md content with explicit instructions to write outputs to disk prompt = f"""You are running inside an automated GPU pipeline. You MUST write all outputs to disk as actual files. Here are your instructions from program.md: {program\_md\_content} LOGGING INSTRUCTIONS (follow exactly): - Before you start any training, print this exact line: \[AUTORESEARCH\] Training started - Before training, print what change you are testing: \[AUTORESEARCH\] Change: - When training finishes, print this exact line: \[AUTORESEARCH\] Training finished - After training, print the key metric value: \[AUTORESEARCH\] Metric: = - When writing results to CSV, print this exact line: \[AUTORESEARCH\] Writing results to CSV IMPORTANT: After completing the above instructions, make sure you have: 1. Written the final results to a CSV file in this repository (e.g. results/results.csv or similar) 2. Saved all code changes you made to the train/ directory (or wherever the training code lives) 3. All files must be written to the current working directory so they appear in git status If any command fails, debug and fix it rather than stopping. Do not just print results — write them to files on disk.""" # --- Pre-flight: verify claude is installed and API key is reachable --- version\_check = subprocess.run( \["claude", "--version"\], capture\_output=True, text=True ) print(f"claude version: {version\_check.stdout.strip()} | stderr: {version\_check.stderr.strip()}", flush=True) if version\_check.returncode != 0: raise RuntimeError(f"claude CLI not found or broken: {version\_check.stderr}") # --- Disable Claude Code sandbox --- # In Kubernetes/Flyte pods, Claude Code's sandbox tries to spin up a nested container # which fails silently and causes file writes to go to an ephemeral space instead of # the real working directory. Disabling it makes writes land in the actual filesystem. claude\_config\_dir = Path("/root/.claude") claude\_config\_dir.mkdir(parents=True, exist\_ok=True) settings = claude\_config\_dir / "settings.json" import json as \_json existing = \_json.loads(settings.read\_text()) if settings.exists() else {} existing\["sandbox"\] = False settings.write\_text(\_json.dumps(existing, indent=2)) print(f"Wrote Claude Code settings: {settings.read\_text()}", flush=True) # --- Run Claude Code CLI --- # Matches swe\_agent.py exactly: prompt as positional arg, CI=true enables non-interactive mode cmd = \[\ "claude",\ "--dangerously-skip-permissions",\ "--max-turns", "100",\ "--model", "claude-haiku-4-5-20251001",\ prompt,\ \] print(f"Running: {shlex.join(cmd\[:3\])} ", flush=True) claude\_env = { \*\*os.environ, "ANTHROPIC\_API\_KEY": anthropic\_api\_key, "CLAUDE\_SKIP\_PERMISSIONS": "true", "CI": "true", # Enables non-interactive mode (no TTY required) } # Stream output line by line so logs appear in real time instead of buffering until done proc = subprocess.Popen( cmd, cwd=repo\_path, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, # merge stderr into stdout stream text=True, env=claude\_env, ) stdout\_lines = \[\] for line in proc.stdout: line = line.rstrip("\\n") print(line, flush=True) stdout\_lines.append(line) proc.wait() full\_output = "\\n".join(stdout\_lines) print(f"Claude Code exit code: {proc.returncode}", flush=True) if proc.returncode != 0: raise RuntimeError( f"Claude Code CLI exited with code {proc.returncode}\\n" f"output: {full\_output\[-2000:\]}" ) # --- Collect changed files --- git\_status = subprocess.run( \["git", "status", "--porcelain"\], cwd=repo\_path, capture\_output=True, text=True, check=True, ) print(f"Git status:\\n{git\_status.stdout}", flush=True) files\_changed = \[\] for line in git\_status.stdout.strip().splitlines(): if line: # git status --porcelain: first two chars are XY status flags file\_path = line\[3:\].strip() files\_changed.append(file\_path) # Also list all files in repo dir for debugging all\_files = subprocess.run( \["find", ".", "-type", "f", "-not", "-path", "./.git/\*"\], cwd=repo\_path, capture\_output=True, text=True, ) print(f"All files in repo:\\n{all\_files.stdout}", flush=True) if not files\_changed: raise RuntimeError( "Claude Code ran successfully but produced no file changes.\\n" f"output: {full\_output\[-2000:\]}" ) # --- Commit --- subprocess.run(\["git", "add", "."\], cwd=repo\_path, check=True) subprocess.run(\["git", "add", "-f", "results.tsv"\], cwd=repo\_path, check=False) subprocess.run(\["git", "add", "-f", "results/"\], cwd=repo\_path, check=False) commit\_message = ( "feat: autoresearch run via Claude Code\\n\\n" "Added research results (CSV) and updated train/ code changes.\\n" "Generated by the autoresearch Flyte agent." ) subprocess.run( \["git", "commit", "-m", commit\_message\], cwd=repo\_path, check=True, ) # --- Push --- print(f"GitHub token present: {bool(github\_token)}, length: {len(github\_token) if github\_token else 0}", flush=True) authenticated\_url = AUTORESEARCH\_REPO\_URL.replace( "https://", f"https://{GITHUB\_USERNAME}:{github\_token}@" ) subprocess.run( \["git", "remote", "set-url", "origin", authenticated\_url\], cwd=repo\_path, check=True, ) push\_result = subprocess.run( \["git", "push", "-u", "origin", branch\_name, "--force"\], cwd=repo\_path, capture\_output=True, text=True, ) print(f"Push stdout: {push\_result.stdout}", flush=True) print(f"Push stderr: {push\_result.stderr}", flush=True) if push\_result.returncode != 0: raise RuntimeError(f"git push failed (exit {push\_result.returncode}):\\n{push\_result.stderr}") # --- Create PR via PyGithub --- auth = Auth.Token(github\_token) gh = Github(auth=auth) repo = gh.get\_repo(AUTORESEARCH\_REPO\_FULL\_NAME) csv\_files = \[f for f in files\_changed if f.endswith(".csv")\] train\_files = \[f for f in files\_changed if "train" in f\] pr\_body = f"""## AutoResearch Run This PR was automatically generated by the autoresearch Flyte agent using Claude Code CLI. ### What changed - \*\*Result CSV files\*\*: {', '.join(f'\`{f}\`' for f in csv\_files) or 'none detected'} - \*\*Train code changes\*\*: {', '.join(f'\`{f}\`' for f in train\_files) or 'none detected'} ### All changed files {chr(10).join(f'- \`{f}\`' for f in files\_changed)} --- 🤖 Generated by \[autoresearch Flyte agent\](https://github.com/unionai-oss/autoresearch) """ existing\_prs = list(repo.get\_pulls(state="open", head=f"unionai-oss:{branch\_name}")) if existing\_prs: pr = existing\_prs\[0\] print(f"PR already exists: {pr.html\_url}", flush=True) else: pr = repo.create\_pull( title="feat: autoresearch results + train changes", body=pr\_body, head=branch\_name, base="master", ) print(f"PR created: {pr.html\_url}", flush=True) # --- Generate progress plot from results.tsv --- plot\_path = repo\_path / "progress.png" results\_tsv = repo\_path / "results.tsv" if results\_tsv.exists(): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import pandas as pd df = pd.read\_csv(str(results\_tsv), sep="\\t") df\["val\_bpb"\] = pd.to\_numeric(df\["val\_bpb"\], errors="coerce") df\["memory\_gb"\] = pd.to\_numeric(df\["memory\_gb"\], errors="coerce") df\["status"\] = df\["status"\].str.strip().str.upper() # Filter out crashes for plotting valid = df\[df\["status"\] != "CRASH"\].copy() valid = valid.reset\_index(drop=True) if len(valid) > 0 and valid\["val\_bpb"\].notna().any(): baseline\_bpb = valid.loc\[0, "val\_bpb"\] best = valid\["val\_bpb"\].min() # Only plot points at or below baseline (the interesting region) below = valid\[valid\["val\_bpb"\] <= baseline\_bpb + 0.0005\] fig, ax = plt.subplots(figsize=(16, 8)) # Plot discarded as faint background dots disc = below\[below\["status"\] == "DISCARD"\] ax.scatter(disc.index, disc\["val\_bpb"\], c="#cccccc", s=12, alpha=0.5, zorder=2, label="Discarded") # Plot kept experiments as prominent green dots kept\_v = below\[below\["status"\] == "KEEP"\] ax.scatter(kept\_v.index, kept\_v\["val\_bpb"\], c="#2ecc71", s=50, zorder=4, label="Kept", edgecolors="black", linewidths=0.5) # Running minimum step line kept\_mask = valid\["status"\] == "KEEP" kept\_idx = valid.index\[kept\_mask\] kept\_bpb = valid.loc\[kept\_mask, "val\_bpb"\] running\_min = kept\_bpb.cummin() ax.step(kept\_idx, running\_min, where="post", color="#27ae60", linewidth=2, alpha=0.7, zorder=3, label="Running best") # Label each kept experiment with its description for idx, bpb in zip(kept\_idx, kept\_bpb): desc = str(valid.loc\[idx, "description"\]).strip() if len(desc) > 45: desc = desc\[:42\] + "..." ax.annotate(desc, (idx, bpb), textcoords="offset points", xytext=(6, 6), fontsize=8.0, color="#1a7a3a", alpha=0.9, rotation=30, ha="left", va="bottom") n\_total = len(df) n\_kept = len(df\[df\["status"\] == "KEEP"\]) ax.set\_xlabel("Experiment #", fontsize=12) ax.set\_ylabel("Validation BPB (lower is better)", fontsize=12) ax.set\_title(f"Autoresearch Progress: {n\_total} Experiments, {n\_kept} Kept Improvements", fontsize=14) ax.legend(loc="upper right", fontsize=9) ax.grid(True, alpha=0.2) margin = (baseline\_bpb - best) \* 0.15 ax.set\_ylim(best - margin, baseline\_bpb + margin) plt.tight\_layout() plt.savefig(str(plot\_path), dpi=150, bbox\_inches="tight") plt.close(fig) print(f"Saved plot to {plot\_path}", flush=True) # Upload plot to PR as a comment with base64 inline image import base64 img\_b64 = base64.b64encode(plot\_path.read\_bytes()).decode() pr\_comment = ( "## Autoresearch Progress\\n\\n" f"!\[Autoresearch Progress\](data:image/png;base64,{img\_b64})" ) pr.create\_issue\_comment(pr\_comment) print("Posted plot as PR comment.", flush=True) # Force-add plot to git and amend commit subprocess.run(\["git", "add", "-f", str(plot\_path)\], cwd=repo\_path, check=False) subprocess.run( \["git", "commit", "--amend", "--no-edit"\], cwd=repo\_path, check=False, ) subprocess.run( \["git", "push", "-u", "origin", branch\_name, "--force"\], cwd=repo\_path, check=False, ) # Show plot in Flyte UI via report await flyte.report.replace.aio( f"

Autoresearch Progress

" f'' f'

View PR

' ) await flyte.report.flush.aio() else: print("results.tsv found but no valid val\_bpb rows — skipping plot.", flush=True) else: print("results.tsv not found — skipping plot.", flush=True) return AutoResearchResult( pr\_url=pr.html\_url, pr\_number=pr.number, branch\_name=branch\_name, files\_changed=files\_changed, success=True, ) # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": import time flyte.init\_from\_config() run = flyte.with\_runcontext(mode="remote").run(run\_autoresearch) print(f"AutoResearch run started: {run.url}") print("Waiting for completion...") while True: try: run.wait() break except Exception as e: print(f"Connection dropped ({e}), reconnecting in 30s...") time.sleep(30) print(f"Done! See run at: {run.url}") # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autoresearch/run.py\* The agent targets a specific repository, identity, and branch via module-level constants. Update these to point at your own fork before running: GITHUB\_USERNAME = "" GITHUB\_EMAIL = "you@example.com" AUTORESEARCH\_REPO\_URL = "https://github.com//.git" AUTORESEARCH\_REPO\_FULL\_NAME = "/" \`\`\` ## Model the result The task returns a typed result describing the pull request it created. \`\`\` # /// script # requires-python = ">=3.11" # dependencies = \[\ # "flyte>=2.0.0b22",\ # "PyGithub>=2.5.0",\ # "matplotlib>=3.7.0",\ # \] # /// """ AutoResearch Agent - Runs the autoresearch workflow using Claude Code CLI in a GPU environment. This agent: 1. Starts a GPU-enabled container 2. Installs Claude Code CLI 3. Clones the autoresearch repository 4. Points Claude Code at program.md as the prompt and lets it run 5. Commits the result (CSV + code changes in train/) and creates a PR """ import os import shlex import subprocess from dataclasses import dataclass from pathlib import Path from typing import Optional from github import Auth, Github import flyte import flyte.report from \_image import image as autoresearch\_image GITHUB\_USERNAME = "parnianz" GITHUB\_EMAIL = "parnianzargham@gmail.com" AUTORESEARCH\_REPO\_URL = "https://github.com/unionai-oss/autoresearch.git" AUTORESEARCH\_REPO\_FULL\_NAME = "unionai-oss/autoresearch" # {{docs-fragment env}} autoresearch\_env = flyte.TaskEnvironment( name="autoresearch-agent", resources=flyte.Resources( cpu=8, memory="32Gi", gpu="T4:1", disk="100Gi", ), secrets=\[\ flyte.Secret(key="github\_token", as\_env\_var="GITHUB\_TOKEN"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=autoresearch\_image, ) # {{/docs-fragment env}} # {{docs-fragment result}} @dataclass class AutoResearchResult: """Result of the autoresearch run.""" pr\_url: str pr\_number: int branch\_name: str files\_changed: list\[str\] success: bool error\_message: Optional\[str\] = None # {{/docs-fragment result}} def clone\_repository(repo\_url: str, work\_dir: Path, github\_token: str) -> Path: """Clone the autoresearch repository with authentication.""" repo\_name = repo\_url.rstrip("/").split("/")\[-1\].replace(".git", "") repo\_path = work\_dir / repo\_name # Inject token into HTTPS URL for authentication authenticated\_url = repo\_url.replace( "https://", f"https://{GITHUB\_USERNAME}:{github\_token}@" ) if repo\_path.exists(): subprocess.run(\["git", "pull"\], cwd=repo\_path, check=True) else: subprocess.run(\["git", "clone", authenticated\_url, str(repo\_path)\], check=True) return repo\_path # {{docs-fragment task}} @autoresearch\_env.task(report=True) async def run\_autoresearch() -> AutoResearchResult: """ Run the autoresearch workflow end-to-end. Steps: - Clone https://github.com/unionai-oss/autoresearch - Configure git identity - Create a new branch - Run Claude Code CLI with program.md as the prompt - Commit results (CSV + train/ changes) - Push and open a PR against the autoresearch repo """ github\_token = os.environ\["GITHUB\_TOKEN"\] anthropic\_api\_key = os.environ\["ANTHROPIC\_API\_KEY"\] # --- Install Node.js + Claude Code at runtime (keeps image small and submission fast) --- import tarfile import urllib.request as \_urllib subprocess.run(\["apt-get", "update", "-y"\], check=False) subprocess.run(\["apt-get", "install", "-y", "git"\], check=False) node\_url = "https://nodejs.org/dist/v20.19.0/node-v20.19.0-linux-x64.tar.gz" node\_tar = Path("/tmp/node.tar.gz") print(f"Downloading Node.js from {node\_url}...", flush=True) \_urllib.urlretrieve(node\_url, node\_tar) size\_mb = node\_tar.stat().st\_size / 1024 / 1024 print(f"Downloaded {size\_mb:.1f} MB to {node\_tar}", flush=True) if size\_mb < 1: raise RuntimeError(f"Node.js download appears empty/corrupt ({size\_mb:.2f} MB) — network may be restricted") node\_dir = Path("/tmp/node") node\_dir.mkdir(exist\_ok=True) print("Extracting Node.js...", flush=True) with tarfile.open(node\_tar, "r:gz") as tar: members = \[m for m in tar.getmembers() if m.name.split("/", 1)\[-1\]\] for m in members: m.name = m.name.split("/", 1)\[-1\] tar.extractall(str(node\_dir), members=\[m for m in members if m.name\]) # Add node/npm to PATH for this process and all subprocesses node\_bin = str(node\_dir / "bin") os.environ\["PATH"\] = node\_bin + ":" + os.environ.get("PATH", "") print(f"Node version: {subprocess.run(\['node', '--version'\], capture\_output=True, text=True).stdout.strip()}", flush=True) npm\_prefix = "/tmp/npm-global" Path(npm\_prefix).mkdir(exist\_ok=True) subprocess.run(\["npm", "install", "-g", "--prefix", npm\_prefix, "@anthropic-ai/claude-code"\], check=True) os.environ\["PATH"\] = str(Path(npm\_prefix) / "bin") + ":" + os.environ\["PATH"\] print("Node.js + Claude Code installed.", flush=True) # --- Clone repo --- work\_dir = Path("/tmp/autoresearch\_workspace") work\_dir.mkdir(exist\_ok=True, parents=True) repo\_path = clone\_repository(AUTORESEARCH\_REPO\_URL, work\_dir, github\_token) # --- Git identity --- subprocess.run( \["git", "config", "--global", "user.email", GITHUB\_EMAIL\], check=True ) subprocess.run( \["git", "config", "--global", "user.name", GITHUB\_USERNAME\], check=True ) # --- Create branch --- import time as \_time branch\_name = f"autoresearch/claude-run-{int(\_time.time())}" try: subprocess.run( \["git", "checkout", "-b", branch\_name\], cwd=repo\_path, check=True, ) except subprocess.CalledProcessError: subprocess.run( \["git", "checkout", branch\_name\], cwd=repo\_path, check=True, ) # --- Read program.md to use as the Claude Code prompt --- program\_md = repo\_path / "program.md" if not program\_md.exists(): raise FileNotFoundError( f"program.md not found in {repo\_path}. " "Make sure the autoresearch repo has a program.md at its root." ) program\_md\_content = program\_md.read\_text() print(f"Loaded prompt from program.md ({len(program\_md\_content)} chars)") # {{/docs-fragment task}} # Install repo dependencies before handing off to Claude for pip\_cmd in \[\ \["pip", "install", "-e", "."\],\ \["pip", "install", "-r", "requirements.txt"\],\ \]: req\_file = repo\_path / pip\_cmd\[-1\] if pip\_cmd\[-1\].startswith("req") else None if req\_file is None or req\_file.exists(): dep\_result = subprocess.run( pip\_cmd, cwd=repo\_path, capture\_output=True, text=True ) print(f"{' '.join(pip\_cmd)}:\\n{dep\_result.stdout}", flush=True) if dep\_result.returncode != 0: print(f"(non-fatal) {dep\_result.stderr}", flush=True) # Wrap the program.md content with explicit instructions to write outputs to disk prompt = f"""You are running inside an automated GPU pipeline. You MUST write all outputs to disk as actual files. Here are your instructions from program.md: {program\_md\_content} LOGGING INSTRUCTIONS (follow exactly): - Before you start any training, print this exact line: \[AUTORESEARCH\] Training started - Before training, print what change you are testing: \[AUTORESEARCH\] Change: - When training finishes, print this exact line: \[AUTORESEARCH\] Training finished - After training, print the key metric value: \[AUTORESEARCH\] Metric: = - When writing results to CSV, print this exact line: \[AUTORESEARCH\] Writing results to CSV IMPORTANT: After completing the above instructions, make sure you have: 1. Written the final results to a CSV file in this repository (e.g. results/results.csv or similar) 2. Saved all code changes you made to the train/ directory (or wherever the training code lives) 3. All files must be written to the current working directory so they appear in git status If any command fails, debug and fix it rather than stopping. Do not just print results — write them to files on disk.""" # --- Pre-flight: verify claude is installed and API key is reachable --- version\_check = subprocess.run( \["claude", "--version"\], capture\_output=True, text=True ) print(f"claude version: {version\_check.stdout.strip()} | stderr: {version\_check.stderr.strip()}", flush=True) if version\_check.returncode != 0: raise RuntimeError(f"claude CLI not found or broken: {version\_check.stderr}") # --- Disable Claude Code sandbox --- # In Kubernetes/Flyte pods, Claude Code's sandbox tries to spin up a nested container # which fails silently and causes file writes to go to an ephemeral space instead of # the real working directory. Disabling it makes writes land in the actual filesystem. claude\_config\_dir = Path("/root/.claude") claude\_config\_dir.mkdir(parents=True, exist\_ok=True) settings = claude\_config\_dir / "settings.json" import json as \_json existing = \_json.loads(settings.read\_text()) if settings.exists() else {} existing\["sandbox"\] = False settings.write\_text(\_json.dumps(existing, indent=2)) print(f"Wrote Claude Code settings: {settings.read\_text()}", flush=True) # --- Run Claude Code CLI --- # Matches swe\_agent.py exactly: prompt as positional arg, CI=true enables non-interactive mode cmd = \[\ "claude",\ "--dangerously-skip-permissions",\ "--max-turns", "100",\ "--model", "claude-haiku-4-5-20251001",\ prompt,\ \] print(f"Running: {shlex.join(cmd\[:3\])} ", flush=True) claude\_env = { \*\*os.environ, "ANTHROPIC\_API\_KEY": anthropic\_api\_key, "CLAUDE\_SKIP\_PERMISSIONS": "true", "CI": "true", # Enables non-interactive mode (no TTY required) } # Stream output line by line so logs appear in real time instead of buffering until done proc = subprocess.Popen( cmd, cwd=repo\_path, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, # merge stderr into stdout stream text=True, env=claude\_env, ) stdout\_lines = \[\] for line in proc.stdout: line = line.rstrip("\\n") print(line, flush=True) stdout\_lines.append(line) proc.wait() full\_output = "\\n".join(stdout\_lines) print(f"Claude Code exit code: {proc.returncode}", flush=True) if proc.returncode != 0: raise RuntimeError( f"Claude Code CLI exited with code {proc.returncode}\\n" f"output: {full\_output\[-2000:\]}" ) # --- Collect changed files --- git\_status = subprocess.run( \["git", "status", "--porcelain"\], cwd=repo\_path, capture\_output=True, text=True, check=True, ) print(f"Git status:\\n{git\_status.stdout}", flush=True) files\_changed = \[\] for line in git\_status.stdout.strip().splitlines(): if line: # git status --porcelain: first two chars are XY status flags file\_path = line\[3:\].strip() files\_changed.append(file\_path) # Also list all files in repo dir for debugging all\_files = subprocess.run( \["find", ".", "-type", "f", "-not", "-path", "./.git/\*"\], cwd=repo\_path, capture\_output=True, text=True, ) print(f"All files in repo:\\n{all\_files.stdout}", flush=True) if not files\_changed: raise RuntimeError( "Claude Code ran successfully but produced no file changes.\\n" f"output: {full\_output\[-2000:\]}" ) # --- Commit --- subprocess.run(\["git", "add", "."\], cwd=repo\_path, check=True) subprocess.run(\["git", "add", "-f", "results.tsv"\], cwd=repo\_path, check=False) subprocess.run(\["git", "add", "-f", "results/"\], cwd=repo\_path, check=False) commit\_message = ( "feat: autoresearch run via Claude Code\\n\\n" "Added research results (CSV) and updated train/ code changes.\\n" "Generated by the autoresearch Flyte agent." ) subprocess.run( \["git", "commit", "-m", commit\_message\], cwd=repo\_path, check=True, ) # --- Push --- print(f"GitHub token present: {bool(github\_token)}, length: {len(github\_token) if github\_token else 0}", flush=True) authenticated\_url = AUTORESEARCH\_REPO\_URL.replace( "https://", f"https://{GITHUB\_USERNAME}:{github\_token}@" ) subprocess.run( \["git", "remote", "set-url", "origin", authenticated\_url\], cwd=repo\_path, check=True, ) push\_result = subprocess.run( \["git", "push", "-u", "origin", branch\_name, "--force"\], cwd=repo\_path, capture\_output=True, text=True, ) print(f"Push stdout: {push\_result.stdout}", flush=True) print(f"Push stderr: {push\_result.stderr}", flush=True) if push\_result.returncode != 0: raise RuntimeError(f"git push failed (exit {push\_result.returncode}):\\n{push\_result.stderr}") # --- Create PR via PyGithub --- auth = Auth.Token(github\_token) gh = Github(auth=auth) repo = gh.get\_repo(AUTORESEARCH\_REPO\_FULL\_NAME) csv\_files = \[f for f in files\_changed if f.endswith(".csv")\] train\_files = \[f for f in files\_changed if "train" in f\] pr\_body = f"""## AutoResearch Run This PR was automatically generated by the autoresearch Flyte agent using Claude Code CLI. ### What changed - \*\*Result CSV files\*\*: {', '.join(f'\`{f}\`' for f in csv\_files) or 'none detected'} - \*\*Train code changes\*\*: {', '.join(f'\`{f}\`' for f in train\_files) or 'none detected'} ### All changed files {chr(10).join(f'- \`{f}\`' for f in files\_changed)} --- 🤖 Generated by \[autoresearch Flyte agent\](https://github.com/unionai-oss/autoresearch) """ existing\_prs = list(repo.get\_pulls(state="open", head=f"unionai-oss:{branch\_name}")) if existing\_prs: pr = existing\_prs\[0\] print(f"PR already exists: {pr.html\_url}", flush=True) else: pr = repo.create\_pull( title="feat: autoresearch results + train changes", body=pr\_body, head=branch\_name, base="master", ) print(f"PR created: {pr.html\_url}", flush=True) # --- Generate progress plot from results.tsv --- plot\_path = repo\_path / "progress.png" results\_tsv = repo\_path / "results.tsv" if results\_tsv.exists(): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import pandas as pd df = pd.read\_csv(str(results\_tsv), sep="\\t") df\["val\_bpb"\] = pd.to\_numeric(df\["val\_bpb"\], errors="coerce") df\["memory\_gb"\] = pd.to\_numeric(df\["memory\_gb"\], errors="coerce") df\["status"\] = df\["status"\].str.strip().str.upper() # Filter out crashes for plotting valid = df\[df\["status"\] != "CRASH"\].copy() valid = valid.reset\_index(drop=True) if len(valid) > 0 and valid\["val\_bpb"\].notna().any(): baseline\_bpb = valid.loc\[0, "val\_bpb"\] best = valid\["val\_bpb"\].min() # Only plot points at or below baseline (the interesting region) below = valid\[valid\["val\_bpb"\] <= baseline\_bpb + 0.0005\] fig, ax = plt.subplots(figsize=(16, 8)) # Plot discarded as faint background dots disc = below\[below\["status"\] == "DISCARD"\] ax.scatter(disc.index, disc\["val\_bpb"\], c="#cccccc", s=12, alpha=0.5, zorder=2, label="Discarded") # Plot kept experiments as prominent green dots kept\_v = below\[below\["status"\] == "KEEP"\] ax.scatter(kept\_v.index, kept\_v\["val\_bpb"\], c="#2ecc71", s=50, zorder=4, label="Kept", edgecolors="black", linewidths=0.5) # Running minimum step line kept\_mask = valid\["status"\] == "KEEP" kept\_idx = valid.index\[kept\_mask\] kept\_bpb = valid.loc\[kept\_mask, "val\_bpb"\] running\_min = kept\_bpb.cummin() ax.step(kept\_idx, running\_min, where="post", color="#27ae60", linewidth=2, alpha=0.7, zorder=3, label="Running best") # Label each kept experiment with its description for idx, bpb in zip(kept\_idx, kept\_bpb): desc = str(valid.loc\[idx, "description"\]).strip() if len(desc) > 45: desc = desc\[:42\] + "..." ax.annotate(desc, (idx, bpb), textcoords="offset points", xytext=(6, 6), fontsize=8.0, color="#1a7a3a", alpha=0.9, rotation=30, ha="left", va="bottom") n\_total = len(df) n\_kept = len(df\[df\["status"\] == "KEEP"\]) ax.set\_xlabel("Experiment #", fontsize=12) ax.set\_ylabel("Validation BPB (lower is better)", fontsize=12) ax.set\_title(f"Autoresearch Progress: {n\_total} Experiments, {n\_kept} Kept Improvements", fontsize=14) ax.legend(loc="upper right", fontsize=9) ax.grid(True, alpha=0.2) margin = (baseline\_bpb - best) \* 0.15 ax.set\_ylim(best - margin, baseline\_bpb + margin) plt.tight\_layout() plt.savefig(str(plot\_path), dpi=150, bbox\_inches="tight") plt.close(fig) print(f"Saved plot to {plot\_path}", flush=True) # Upload plot to PR as a comment with base64 inline image import base64 img\_b64 = base64.b64encode(plot\_path.read\_bytes()).decode() pr\_comment = ( "## Autoresearch Progress\\n\\n" f"!\[Autoresearch Progress\](data:image/png;base64,{img\_b64})" ) pr.create\_issue\_comment(pr\_comment) print("Posted plot as PR comment.", flush=True) # Force-add plot to git and amend commit subprocess.run(\["git", "add", "-f", str(plot\_path)\], cwd=repo\_path, check=False) subprocess.run( \["git", "commit", "--amend", "--no-edit"\], cwd=repo\_path, check=False, ) subprocess.run( \["git", "push", "-u", "origin", branch\_name, "--force"\], cwd=repo\_path, check=False, ) # Show plot in Flyte UI via report await flyte.report.replace.aio( f"

Autoresearch Progress

" f'' f'

View PR

' ) await flyte.report.flush.aio() else: print("results.tsv found but no valid val\_bpb rows — skipping plot.", flush=True) else: print("results.tsv not found — skipping plot.", flush=True) return AutoResearchResult( pr\_url=pr.html\_url, pr\_number=pr.number, branch\_name=branch\_name, files\_changed=files\_changed, success=True, ) # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": import time flyte.init\_from\_config() run = flyte.with\_runcontext(mode="remote").run(run\_autoresearch) print(f"AutoResearch run started: {run.url}") print("Waiting for completion...") while True: try: run.wait() break except Exception as e: print(f"Connection dropped ({e}), reconnecting in 30s...") time.sleep(30) print(f"Done! See run at: {run.url}") # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autoresearch/run.py\* ## The autoresearch task The task is a long, sequential procedure. It starts by installing Node.js and Claude Code at run time, cloning the repo, configuring git, creating a branch, and loading \`program.md\` as the prompt: \`\`\` # /// script # requires-python = ">=3.11" # dependencies = \[\ # "flyte>=2.0.0b22",\ # "PyGithub>=2.5.0",\ # "matplotlib>=3.7.0",\ # \] # /// """ AutoResearch Agent - Runs the autoresearch workflow using Claude Code CLI in a GPU environment. This agent: 1. Starts a GPU-enabled container 2. Installs Claude Code CLI 3. Clones the autoresearch repository 4. Points Claude Code at program.md as the prompt and lets it run 5. Commits the result (CSV + code changes in train/) and creates a PR """ import os import shlex import subprocess from dataclasses import dataclass from pathlib import Path from typing import Optional from github import Auth, Github import flyte import flyte.report from \_image import image as autoresearch\_image GITHUB\_USERNAME = "parnianz" GITHUB\_EMAIL = "parnianzargham@gmail.com" AUTORESEARCH\_REPO\_URL = "https://github.com/unionai-oss/autoresearch.git" AUTORESEARCH\_REPO\_FULL\_NAME = "unionai-oss/autoresearch" # {{docs-fragment env}} autoresearch\_env = flyte.TaskEnvironment( name="autoresearch-agent", resources=flyte.Resources( cpu=8, memory="32Gi", gpu="T4:1", disk="100Gi", ), secrets=\[\ flyte.Secret(key="github\_token", as\_env\_var="GITHUB\_TOKEN"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=autoresearch\_image, ) # {{/docs-fragment env}} # {{docs-fragment result}} @dataclass class AutoResearchResult: """Result of the autoresearch run.""" pr\_url: str pr\_number: int branch\_name: str files\_changed: list\[str\] success: bool error\_message: Optional\[str\] = None # {{/docs-fragment result}} def clone\_repository(repo\_url: str, work\_dir: Path, github\_token: str) -> Path: """Clone the autoresearch repository with authentication.""" repo\_name = repo\_url.rstrip("/").split("/")\[-1\].replace(".git", "") repo\_path = work\_dir / repo\_name # Inject token into HTTPS URL for authentication authenticated\_url = repo\_url.replace( "https://", f"https://{GITHUB\_USERNAME}:{github\_token}@" ) if repo\_path.exists(): subprocess.run(\["git", "pull"\], cwd=repo\_path, check=True) else: subprocess.run(\["git", "clone", authenticated\_url, str(repo\_path)\], check=True) return repo\_path # {{docs-fragment task}} @autoresearch\_env.task(report=True) async def run\_autoresearch() -> AutoResearchResult: """ Run the autoresearch workflow end-to-end. Steps: - Clone https://github.com/unionai-oss/autoresearch - Configure git identity - Create a new branch - Run Claude Code CLI with program.md as the prompt - Commit results (CSV + train/ changes) - Push and open a PR against the autoresearch repo """ github\_token = os.environ\["GITHUB\_TOKEN"\] anthropic\_api\_key = os.environ\["ANTHROPIC\_API\_KEY"\] # --- Install Node.js + Claude Code at runtime (keeps image small and submission fast) --- import tarfile import urllib.request as \_urllib subprocess.run(\["apt-get", "update", "-y"\], check=False) subprocess.run(\["apt-get", "install", "-y", "git"\], check=False) node\_url = "https://nodejs.org/dist/v20.19.0/node-v20.19.0-linux-x64.tar.gz" node\_tar = Path("/tmp/node.tar.gz") print(f"Downloading Node.js from {node\_url}...", flush=True) \_urllib.urlretrieve(node\_url, node\_tar) size\_mb = node\_tar.stat().st\_size / 1024 / 1024 print(f"Downloaded {size\_mb:.1f} MB to {node\_tar}", flush=True) if size\_mb < 1: raise RuntimeError(f"Node.js download appears empty/corrupt ({size\_mb:.2f} MB) — network may be restricted") node\_dir = Path("/tmp/node") node\_dir.mkdir(exist\_ok=True) print("Extracting Node.js...", flush=True) with tarfile.open(node\_tar, "r:gz") as tar: members = \[m for m in tar.getmembers() if m.name.split("/", 1)\[-1\]\] for m in members: m.name = m.name.split("/", 1)\[-1\] tar.extractall(str(node\_dir), members=\[m for m in members if m.name\]) # Add node/npm to PATH for this process and all subprocesses node\_bin = str(node\_dir / "bin") os.environ\["PATH"\] = node\_bin + ":" + os.environ.get("PATH", "") print(f"Node version: {subprocess.run(\['node', '--version'\], capture\_output=True, text=True).stdout.strip()}", flush=True) npm\_prefix = "/tmp/npm-global" Path(npm\_prefix).mkdir(exist\_ok=True) subprocess.run(\["npm", "install", "-g", "--prefix", npm\_prefix, "@anthropic-ai/claude-code"\], check=True) os.environ\["PATH"\] = str(Path(npm\_prefix) / "bin") + ":" + os.environ\["PATH"\] print("Node.js + Claude Code installed.", flush=True) # --- Clone repo --- work\_dir = Path("/tmp/autoresearch\_workspace") work\_dir.mkdir(exist\_ok=True, parents=True) repo\_path = clone\_repository(AUTORESEARCH\_REPO\_URL, work\_dir, github\_token) # --- Git identity --- subprocess.run( \["git", "config", "--global", "user.email", GITHUB\_EMAIL\], check=True ) subprocess.run( \["git", "config", "--global", "user.name", GITHUB\_USERNAME\], check=True ) # --- Create branch --- import time as \_time branch\_name = f"autoresearch/claude-run-{int(\_time.time())}" try: subprocess.run( \["git", "checkout", "-b", branch\_name\], cwd=repo\_path, check=True, ) except subprocess.CalledProcessError: subprocess.run( \["git", "checkout", branch\_name\], cwd=repo\_path, check=True, ) # --- Read program.md to use as the Claude Code prompt --- program\_md = repo\_path / "program.md" if not program\_md.exists(): raise FileNotFoundError( f"program.md not found in {repo\_path}. " "Make sure the autoresearch repo has a program.md at its root." ) program\_md\_content = program\_md.read\_text() print(f"Loaded prompt from program.md ({len(program\_md\_content)} chars)") # {{/docs-fragment task}} # Install repo dependencies before handing off to Claude for pip\_cmd in \[\ \["pip", "install", "-e", "."\],\ \["pip", "install", "-r", "requirements.txt"\],\ \]: req\_file = repo\_path / pip\_cmd\[-1\] if pip\_cmd\[-1\].startswith("req") else None if req\_file is None or req\_file.exists(): dep\_result = subprocess.run( pip\_cmd, cwd=repo\_path, capture\_output=True, text=True ) print(f"{' '.join(pip\_cmd)}:\\n{dep\_result.stdout}", flush=True) if dep\_result.returncode != 0: print(f"(non-fatal) {dep\_result.stderr}", flush=True) # Wrap the program.md content with explicit instructions to write outputs to disk prompt = f"""You are running inside an automated GPU pipeline. You MUST write all outputs to disk as actual files. Here are your instructions from program.md: {program\_md\_content} LOGGING INSTRUCTIONS (follow exactly): - Before you start any training, print this exact line: \[AUTORESEARCH\] Training started - Before training, print what change you are testing: \[AUTORESEARCH\] Change: - When training finishes, print this exact line: \[AUTORESEARCH\] Training finished - After training, print the key metric value: \[AUTORESEARCH\] Metric: = - When writing results to CSV, print this exact line: \[AUTORESEARCH\] Writing results to CSV IMPORTANT: After completing the above instructions, make sure you have: 1. Written the final results to a CSV file in this repository (e.g. results/results.csv or similar) 2. Saved all code changes you made to the train/ directory (or wherever the training code lives) 3. All files must be written to the current working directory so they appear in git status If any command fails, debug and fix it rather than stopping. Do not just print results — write them to files on disk.""" # --- Pre-flight: verify claude is installed and API key is reachable --- version\_check = subprocess.run( \["claude", "--version"\], capture\_output=True, text=True ) print(f"claude version: {version\_check.stdout.strip()} | stderr: {version\_check.stderr.strip()}", flush=True) if version\_check.returncode != 0: raise RuntimeError(f"claude CLI not found or broken: {version\_check.stderr}") # --- Disable Claude Code sandbox --- # In Kubernetes/Flyte pods, Claude Code's sandbox tries to spin up a nested container # which fails silently and causes file writes to go to an ephemeral space instead of # the real working directory. Disabling it makes writes land in the actual filesystem. claude\_config\_dir = Path("/root/.claude") claude\_config\_dir.mkdir(parents=True, exist\_ok=True) settings = claude\_config\_dir / "settings.json" import json as \_json existing = \_json.loads(settings.read\_text()) if settings.exists() else {} existing\["sandbox"\] = False settings.write\_text(\_json.dumps(existing, indent=2)) print(f"Wrote Claude Code settings: {settings.read\_text()}", flush=True) # --- Run Claude Code CLI --- # Matches swe\_agent.py exactly: prompt as positional arg, CI=true enables non-interactive mode cmd = \[\ "claude",\ "--dangerously-skip-permissions",\ "--max-turns", "100",\ "--model", "claude-haiku-4-5-20251001",\ prompt,\ \] print(f"Running: {shlex.join(cmd\[:3\])} ", flush=True) claude\_env = { \*\*os.environ, "ANTHROPIC\_API\_KEY": anthropic\_api\_key, "CLAUDE\_SKIP\_PERMISSIONS": "true", "CI": "true", # Enables non-interactive mode (no TTY required) } # Stream output line by line so logs appear in real time instead of buffering until done proc = subprocess.Popen( cmd, cwd=repo\_path, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, # merge stderr into stdout stream text=True, env=claude\_env, ) stdout\_lines = \[\] for line in proc.stdout: line = line.rstrip("\\n") print(line, flush=True) stdout\_lines.append(line) proc.wait() full\_output = "\\n".join(stdout\_lines) print(f"Claude Code exit code: {proc.returncode}", flush=True) if proc.returncode != 0: raise RuntimeError( f"Claude Code CLI exited with code {proc.returncode}\\n" f"output: {full\_output\[-2000:\]}" ) # --- Collect changed files --- git\_status = subprocess.run( \["git", "status", "--porcelain"\], cwd=repo\_path, capture\_output=True, text=True, check=True, ) print(f"Git status:\\n{git\_status.stdout}", flush=True) files\_changed = \[\] for line in git\_status.stdout.strip().splitlines(): if line: # git status --porcelain: first two chars are XY status flags file\_path = line\[3:\].strip() files\_changed.append(file\_path) # Also list all files in repo dir for debugging all\_files = subprocess.run( \["find", ".", "-type", "f", "-not", "-path", "./.git/\*"\], cwd=repo\_path, capture\_output=True, text=True, ) print(f"All files in repo:\\n{all\_files.stdout}", flush=True) if not files\_changed: raise RuntimeError( "Claude Code ran successfully but produced no file changes.\\n" f"output: {full\_output\[-2000:\]}" ) # --- Commit --- subprocess.run(\["git", "add", "."\], cwd=repo\_path, check=True) subprocess.run(\["git", "add", "-f", "results.tsv"\], cwd=repo\_path, check=False) subprocess.run(\["git", "add", "-f", "results/"\], cwd=repo\_path, check=False) commit\_message = ( "feat: autoresearch run via Claude Code\\n\\n" "Added research results (CSV) and updated train/ code changes.\\n" "Generated by the autoresearch Flyte agent." ) subprocess.run( \["git", "commit", "-m", commit\_message\], cwd=repo\_path, check=True, ) # --- Push --- print(f"GitHub token present: {bool(github\_token)}, length: {len(github\_token) if github\_token else 0}", flush=True) authenticated\_url = AUTORESEARCH\_REPO\_URL.replace( "https://", f"https://{GITHUB\_USERNAME}:{github\_token}@" ) subprocess.run( \["git", "remote", "set-url", "origin", authenticated\_url\], cwd=repo\_path, check=True, ) push\_result = subprocess.run( \["git", "push", "-u", "origin", branch\_name, "--force"\], cwd=repo\_path, capture\_output=True, text=True, ) print(f"Push stdout: {push\_result.stdout}", flush=True) print(f"Push stderr: {push\_result.stderr}", flush=True) if push\_result.returncode != 0: raise RuntimeError(f"git push failed (exit {push\_result.returncode}):\\n{push\_result.stderr}") # --- Create PR via PyGithub --- auth = Auth.Token(github\_token) gh = Github(auth=auth) repo = gh.get\_repo(AUTORESEARCH\_REPO\_FULL\_NAME) csv\_files = \[f for f in files\_changed if f.endswith(".csv")\] train\_files = \[f for f in files\_changed if "train" in f\] pr\_body = f"""## AutoResearch Run This PR was automatically generated by the autoresearch Flyte agent using Claude Code CLI. ### What changed - \*\*Result CSV files\*\*: {', '.join(f'\`{f}\`' for f in csv\_files) or 'none detected'} - \*\*Train code changes\*\*: {', '.join(f'\`{f}\`' for f in train\_files) or 'none detected'} ### All changed files {chr(10).join(f'- \`{f}\`' for f in files\_changed)} --- 🤖 Generated by \[autoresearch Flyte agent\](https://github.com/unionai-oss/autoresearch) """ existing\_prs = list(repo.get\_pulls(state="open", head=f"unionai-oss:{branch\_name}")) if existing\_prs: pr = existing\_prs\[0\] print(f"PR already exists: {pr.html\_url}", flush=True) else: pr = repo.create\_pull( title="feat: autoresearch results + train changes", body=pr\_body, head=branch\_name, base="master", ) print(f"PR created: {pr.html\_url}", flush=True) # --- Generate progress plot from results.tsv --- plot\_path = repo\_path / "progress.png" results\_tsv = repo\_path / "results.tsv" if results\_tsv.exists(): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import pandas as pd df = pd.read\_csv(str(results\_tsv), sep="\\t") df\["val\_bpb"\] = pd.to\_numeric(df\["val\_bpb"\], errors="coerce") df\["memory\_gb"\] = pd.to\_numeric(df\["memory\_gb"\], errors="coerce") df\["status"\] = df\["status"\].str.strip().str.upper() # Filter out crashes for plotting valid = df\[df\["status"\] != "CRASH"\].copy() valid = valid.reset\_index(drop=True) if len(valid) > 0 and valid\["val\_bpb"\].notna().any(): baseline\_bpb = valid.loc\[0, "val\_bpb"\] best = valid\["val\_bpb"\].min() # Only plot points at or below baseline (the interesting region) below = valid\[valid\["val\_bpb"\] <= baseline\_bpb + 0.0005\] fig, ax = plt.subplots(figsize=(16, 8)) # Plot discarded as faint background dots disc = below\[below\["status"\] == "DISCARD"\] ax.scatter(disc.index, disc\["val\_bpb"\], c="#cccccc", s=12, alpha=0.5, zorder=2, label="Discarded") # Plot kept experiments as prominent green dots kept\_v = below\[below\["status"\] == "KEEP"\] ax.scatter(kept\_v.index, kept\_v\["val\_bpb"\], c="#2ecc71", s=50, zorder=4, label="Kept", edgecolors="black", linewidths=0.5) # Running minimum step line kept\_mask = valid\["status"\] == "KEEP" kept\_idx = valid.index\[kept\_mask\] kept\_bpb = valid.loc\[kept\_mask, "val\_bpb"\] running\_min = kept\_bpb.cummin() ax.step(kept\_idx, running\_min, where="post", color="#27ae60", linewidth=2, alpha=0.7, zorder=3, label="Running best") # Label each kept experiment with its description for idx, bpb in zip(kept\_idx, kept\_bpb): desc = str(valid.loc\[idx, "description"\]).strip() if len(desc) > 45: desc = desc\[:42\] + "..." ax.annotate(desc, (idx, bpb), textcoords="offset points", xytext=(6, 6), fontsize=8.0, color="#1a7a3a", alpha=0.9, rotation=30, ha="left", va="bottom") n\_total = len(df) n\_kept = len(df\[df\["status"\] == "KEEP"\]) ax.set\_xlabel("Experiment #", fontsize=12) ax.set\_ylabel("Validation BPB (lower is better)", fontsize=12) ax.set\_title(f"Autoresearch Progress: {n\_total} Experiments, {n\_kept} Kept Improvements", fontsize=14) ax.legend(loc="upper right", fontsize=9) ax.grid(True, alpha=0.2) margin = (baseline\_bpb - best) \* 0.15 ax.set\_ylim(best - margin, baseline\_bpb + margin) plt.tight\_layout() plt.savefig(str(plot\_path), dpi=150, bbox\_inches="tight") plt.close(fig) print(f"Saved plot to {plot\_path}", flush=True) # Upload plot to PR as a comment with base64 inline image import base64 img\_b64 = base64.b64encode(plot\_path.read\_bytes()).decode() pr\_comment = ( "## Autoresearch Progress\\n\\n" f"!\[Autoresearch Progress\](data:image/png;base64,{img\_b64})" ) pr.create\_issue\_comment(pr\_comment) print("Posted plot as PR comment.", flush=True) # Force-add plot to git and amend commit subprocess.run(\["git", "add", "-f", str(plot\_path)\], cwd=repo\_path, check=False) subprocess.run( \["git", "commit", "--amend", "--no-edit"\], cwd=repo\_path, check=False, ) subprocess.run( \["git", "push", "-u", "origin", branch\_name, "--force"\], cwd=repo\_path, check=False, ) # Show plot in Flyte UI via report await flyte.report.replace.aio( f"

Autoresearch Progress

" f'' f'

View PR

' ) await flyte.report.flush.aio() else: print("results.tsv found but no valid val\_bpb rows — skipping plot.", flush=True) else: print("results.tsv not found — skipping plot.", flush=True) return AutoResearchResult( pr\_url=pr.html\_url, pr\_number=pr.number, branch\_name=branch\_name, files\_changed=files\_changed, success=True, ) # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": import time flyte.init\_from\_config() run = flyte.with\_runcontext(mode="remote").run(run\_autoresearch) print(f"AutoResearch run started: {run.url}") print("Waiting for completion...") while True: try: run.wait() break except Exception as e: print(f"Connection dropped ({e}), reconnecting in 30s...") time.sleep(30) print(f"Done! See run at: {run.url}") # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autoresearch/run.py\* From there the task: 1. Wraps the \`program.md\` brief with explicit logging and "write outputs to disk" instructions. 2. Disables the Claude Code sandbox (it conflicts with the Flyte pod's container) and runs the CLI non-interactively, streaming its output to the Flyte logs in real time. 3. Collects the files the agent changed via \`git status\`, commits them, and force-pushes the branch. 4. Opens (or reuses) a pull request with \[PyGithub\](https://pygithub.readthedocs.io/). 5. If the agent produced a \`results.tsv\`, renders a progress plot of validation bits-per-byte, attaches it to the PR, and streams it into the Flyte UI: \`\`\` # /// script # requires-python = ">=3.11" # dependencies = \[\ # "flyte>=2.0.0b22",\ # "PyGithub>=2.5.0",\ # "matplotlib>=3.7.0",\ # \] # /// """ AutoResearch Agent - Runs the autoresearch workflow using Claude Code CLI in a GPU environment. This agent: 1. Starts a GPU-enabled container 2. Installs Claude Code CLI 3. Clones the autoresearch repository 4. Points Claude Code at program.md as the prompt and lets it run 5. Commits the result (CSV + code changes in train/) and creates a PR """ import os import shlex import subprocess from dataclasses import dataclass from pathlib import Path from typing import Optional from github import Auth, Github import flyte import flyte.report from \_image import image as autoresearch\_image GITHUB\_USERNAME = "parnianz" GITHUB\_EMAIL = "parnianzargham@gmail.com" AUTORESEARCH\_REPO\_URL = "https://github.com/unionai-oss/autoresearch.git" AUTORESEARCH\_REPO\_FULL\_NAME = "unionai-oss/autoresearch" # {{docs-fragment env}} autoresearch\_env = flyte.TaskEnvironment( name="autoresearch-agent", resources=flyte.Resources( cpu=8, memory="32Gi", gpu="T4:1", disk="100Gi", ), secrets=\[\ flyte.Secret(key="github\_token", as\_env\_var="GITHUB\_TOKEN"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=autoresearch\_image, ) # {{/docs-fragment env}} # {{docs-fragment result}} @dataclass class AutoResearchResult: """Result of the autoresearch run.""" pr\_url: str pr\_number: int branch\_name: str files\_changed: list\[str\] success: bool error\_message: Optional\[str\] = None # {{/docs-fragment result}} def clone\_repository(repo\_url: str, work\_dir: Path, github\_token: str) -> Path: """Clone the autoresearch repository with authentication.""" repo\_name = repo\_url.rstrip("/").split("/")\[-1\].replace(".git", "") repo\_path = work\_dir / repo\_name # Inject token into HTTPS URL for authentication authenticated\_url = repo\_url.replace( "https://", f"https://{GITHUB\_USERNAME}:{github\_token}@" ) if repo\_path.exists(): subprocess.run(\["git", "pull"\], cwd=repo\_path, check=True) else: subprocess.run(\["git", "clone", authenticated\_url, str(repo\_path)\], check=True) return repo\_path # {{docs-fragment task}} @autoresearch\_env.task(report=True) async def run\_autoresearch() -> AutoResearchResult: """ Run the autoresearch workflow end-to-end. Steps: - Clone https://github.com/unionai-oss/autoresearch - Configure git identity - Create a new branch - Run Claude Code CLI with program.md as the prompt - Commit results (CSV + train/ changes) - Push and open a PR against the autoresearch repo """ github\_token = os.environ\["GITHUB\_TOKEN"\] anthropic\_api\_key = os.environ\["ANTHROPIC\_API\_KEY"\] # --- Install Node.js + Claude Code at runtime (keeps image small and submission fast) --- import tarfile import urllib.request as \_urllib subprocess.run(\["apt-get", "update", "-y"\], check=False) subprocess.run(\["apt-get", "install", "-y", "git"\], check=False) node\_url = "https://nodejs.org/dist/v20.19.0/node-v20.19.0-linux-x64.tar.gz" node\_tar = Path("/tmp/node.tar.gz") print(f"Downloading Node.js from {node\_url}...", flush=True) \_urllib.urlretrieve(node\_url, node\_tar) size\_mb = node\_tar.stat().st\_size / 1024 / 1024 print(f"Downloaded {size\_mb:.1f} MB to {node\_tar}", flush=True) if size\_mb < 1: raise RuntimeError(f"Node.js download appears empty/corrupt ({size\_mb:.2f} MB) — network may be restricted") node\_dir = Path("/tmp/node") node\_dir.mkdir(exist\_ok=True) print("Extracting Node.js...", flush=True) with tarfile.open(node\_tar, "r:gz") as tar: members = \[m for m in tar.getmembers() if m.name.split("/", 1)\[-1\]\] for m in members: m.name = m.name.split("/", 1)\[-1\] tar.extractall(str(node\_dir), members=\[m for m in members if m.name\]) # Add node/npm to PATH for this process and all subprocesses node\_bin = str(node\_dir / "bin") os.environ\["PATH"\] = node\_bin + ":" + os.environ.get("PATH", "") print(f"Node version: {subprocess.run(\['node', '--version'\], capture\_output=True, text=True).stdout.strip()}", flush=True) npm\_prefix = "/tmp/npm-global" Path(npm\_prefix).mkdir(exist\_ok=True) subprocess.run(\["npm", "install", "-g", "--prefix", npm\_prefix, "@anthropic-ai/claude-code"\], check=True) os.environ\["PATH"\] = str(Path(npm\_prefix) / "bin") + ":" + os.environ\["PATH"\] print("Node.js + Claude Code installed.", flush=True) # --- Clone repo --- work\_dir = Path("/tmp/autoresearch\_workspace") work\_dir.mkdir(exist\_ok=True, parents=True) repo\_path = clone\_repository(AUTORESEARCH\_REPO\_URL, work\_dir, github\_token) # --- Git identity --- subprocess.run( \["git", "config", "--global", "user.email", GITHUB\_EMAIL\], check=True ) subprocess.run( \["git", "config", "--global", "user.name", GITHUB\_USERNAME\], check=True ) # --- Create branch --- import time as \_time branch\_name = f"autoresearch/claude-run-{int(\_time.time())}" try: subprocess.run( \["git", "checkout", "-b", branch\_name\], cwd=repo\_path, check=True, ) except subprocess.CalledProcessError: subprocess.run( \["git", "checkout", branch\_name\], cwd=repo\_path, check=True, ) # --- Read program.md to use as the Claude Code prompt --- program\_md = repo\_path / "program.md" if not program\_md.exists(): raise FileNotFoundError( f"program.md not found in {repo\_path}. " "Make sure the autoresearch repo has a program.md at its root." ) program\_md\_content = program\_md.read\_text() print(f"Loaded prompt from program.md ({len(program\_md\_content)} chars)") # {{/docs-fragment task}} # Install repo dependencies before handing off to Claude for pip\_cmd in \[\ \["pip", "install", "-e", "."\],\ \["pip", "install", "-r", "requirements.txt"\],\ \]: req\_file = repo\_path / pip\_cmd\[-1\] if pip\_cmd\[-1\].startswith("req") else None if req\_file is None or req\_file.exists(): dep\_result = subprocess.run( pip\_cmd, cwd=repo\_path, capture\_output=True, text=True ) print(f"{' '.join(pip\_cmd)}:\\n{dep\_result.stdout}", flush=True) if dep\_result.returncode != 0: print(f"(non-fatal) {dep\_result.stderr}", flush=True) # Wrap the program.md content with explicit instructions to write outputs to disk prompt = f"""You are running inside an automated GPU pipeline. You MUST write all outputs to disk as actual files. Here are your instructions from program.md: {program\_md\_content} LOGGING INSTRUCTIONS (follow exactly): - Before you start any training, print this exact line: \[AUTORESEARCH\] Training started - Before training, print what change you are testing: \[AUTORESEARCH\] Change: - When training finishes, print this exact line: \[AUTORESEARCH\] Training finished - After training, print the key metric value: \[AUTORESEARCH\] Metric: = - When writing results to CSV, print this exact line: \[AUTORESEARCH\] Writing results to CSV IMPORTANT: After completing the above instructions, make sure you have: 1. Written the final results to a CSV file in this repository (e.g. results/results.csv or similar) 2. Saved all code changes you made to the train/ directory (or wherever the training code lives) 3. All files must be written to the current working directory so they appear in git status If any command fails, debug and fix it rather than stopping. Do not just print results — write them to files on disk.""" # --- Pre-flight: verify claude is installed and API key is reachable --- version\_check = subprocess.run( \["claude", "--version"\], capture\_output=True, text=True ) print(f"claude version: {version\_check.stdout.strip()} | stderr: {version\_check.stderr.strip()}", flush=True) if version\_check.returncode != 0: raise RuntimeError(f"claude CLI not found or broken: {version\_check.stderr}") # --- Disable Claude Code sandbox --- # In Kubernetes/Flyte pods, Claude Code's sandbox tries to spin up a nested container # which fails silently and causes file writes to go to an ephemeral space instead of # the real working directory. Disabling it makes writes land in the actual filesystem. claude\_config\_dir = Path("/root/.claude") claude\_config\_dir.mkdir(parents=True, exist\_ok=True) settings = claude\_config\_dir / "settings.json" import json as \_json existing = \_json.loads(settings.read\_text()) if settings.exists() else {} existing\["sandbox"\] = False settings.write\_text(\_json.dumps(existing, indent=2)) print(f"Wrote Claude Code settings: {settings.read\_text()}", flush=True) # --- Run Claude Code CLI --- # Matches swe\_agent.py exactly: prompt as positional arg, CI=true enables non-interactive mode cmd = \[\ "claude",\ "--dangerously-skip-permissions",\ "--max-turns", "100",\ "--model", "claude-haiku-4-5-20251001",\ prompt,\ \] print(f"Running: {shlex.join(cmd\[:3\])} ", flush=True) claude\_env = { \*\*os.environ, "ANTHROPIC\_API\_KEY": anthropic\_api\_key, "CLAUDE\_SKIP\_PERMISSIONS": "true", "CI": "true", # Enables non-interactive mode (no TTY required) } # Stream output line by line so logs appear in real time instead of buffering until done proc = subprocess.Popen( cmd, cwd=repo\_path, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, # merge stderr into stdout stream text=True, env=claude\_env, ) stdout\_lines = \[\] for line in proc.stdout: line = line.rstrip("\\n") print(line, flush=True) stdout\_lines.append(line) proc.wait() full\_output = "\\n".join(stdout\_lines) print(f"Claude Code exit code: {proc.returncode}", flush=True) if proc.returncode != 0: raise RuntimeError( f"Claude Code CLI exited with code {proc.returncode}\\n" f"output: {full\_output\[-2000:\]}" ) # --- Collect changed files --- git\_status = subprocess.run( \["git", "status", "--porcelain"\], cwd=repo\_path, capture\_output=True, text=True, check=True, ) print(f"Git status:\\n{git\_status.stdout}", flush=True) files\_changed = \[\] for line in git\_status.stdout.strip().splitlines(): if line: # git status --porcelain: first two chars are XY status flags file\_path = line\[3:\].strip() files\_changed.append(file\_path) # Also list all files in repo dir for debugging all\_files = subprocess.run( \["find", ".", "-type", "f", "-not", "-path", "./.git/\*"\], cwd=repo\_path, capture\_output=True, text=True, ) print(f"All files in repo:\\n{all\_files.stdout}", flush=True) if not files\_changed: raise RuntimeError( "Claude Code ran successfully but produced no file changes.\\n" f"output: {full\_output\[-2000:\]}" ) # --- Commit --- subprocess.run(\["git", "add", "."\], cwd=repo\_path, check=True) subprocess.run(\["git", "add", "-f", "results.tsv"\], cwd=repo\_path, check=False) subprocess.run(\["git", "add", "-f", "results/"\], cwd=repo\_path, check=False) commit\_message = ( "feat: autoresearch run via Claude Code\\n\\n" "Added research results (CSV) and updated train/ code changes.\\n" "Generated by the autoresearch Flyte agent." ) subprocess.run( \["git", "commit", "-m", commit\_message\], cwd=repo\_path, check=True, ) # --- Push --- print(f"GitHub token present: {bool(github\_token)}, length: {len(github\_token) if github\_token else 0}", flush=True) authenticated\_url = AUTORESEARCH\_REPO\_URL.replace( "https://", f"https://{GITHUB\_USERNAME}:{github\_token}@" ) subprocess.run( \["git", "remote", "set-url", "origin", authenticated\_url\], cwd=repo\_path, check=True, ) push\_result = subprocess.run( \["git", "push", "-u", "origin", branch\_name, "--force"\], cwd=repo\_path, capture\_output=True, text=True, ) print(f"Push stdout: {push\_result.stdout}", flush=True) print(f"Push stderr: {push\_result.stderr}", flush=True) if push\_result.returncode != 0: raise RuntimeError(f"git push failed (exit {push\_result.returncode}):\\n{push\_result.stderr}") # --- Create PR via PyGithub --- auth = Auth.Token(github\_token) gh = Github(auth=auth) repo = gh.get\_repo(AUTORESEARCH\_REPO\_FULL\_NAME) csv\_files = \[f for f in files\_changed if f.endswith(".csv")\] train\_files = \[f for f in files\_changed if "train" in f\] pr\_body = f"""## AutoResearch Run This PR was automatically generated by the autoresearch Flyte agent using Claude Code CLI. ### What changed - \*\*Result CSV files\*\*: {', '.join(f'\`{f}\`' for f in csv\_files) or 'none detected'} - \*\*Train code changes\*\*: {', '.join(f'\`{f}\`' for f in train\_files) or 'none detected'} ### All changed files {chr(10).join(f'- \`{f}\`' for f in files\_changed)} --- 🤖 Generated by \[autoresearch Flyte agent\](https://github.com/unionai-oss/autoresearch) """ existing\_prs = list(repo.get\_pulls(state="open", head=f"unionai-oss:{branch\_name}")) if existing\_prs: pr = existing\_prs\[0\] print(f"PR already exists: {pr.html\_url}", flush=True) else: pr = repo.create\_pull( title="feat: autoresearch results + train changes", body=pr\_body, head=branch\_name, base="master", ) print(f"PR created: {pr.html\_url}", flush=True) # --- Generate progress plot from results.tsv --- plot\_path = repo\_path / "progress.png" results\_tsv = repo\_path / "results.tsv" if results\_tsv.exists(): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import pandas as pd df = pd.read\_csv(str(results\_tsv), sep="\\t") df\["val\_bpb"\] = pd.to\_numeric(df\["val\_bpb"\], errors="coerce") df\["memory\_gb"\] = pd.to\_numeric(df\["memory\_gb"\], errors="coerce") df\["status"\] = df\["status"\].str.strip().str.upper() # Filter out crashes for plotting valid = df\[df\["status"\] != "CRASH"\].copy() valid = valid.reset\_index(drop=True) if len(valid) > 0 and valid\["val\_bpb"\].notna().any(): baseline\_bpb = valid.loc\[0, "val\_bpb"\] best = valid\["val\_bpb"\].min() # Only plot points at or below baseline (the interesting region) below = valid\[valid\["val\_bpb"\] <= baseline\_bpb + 0.0005\] fig, ax = plt.subplots(figsize=(16, 8)) # Plot discarded as faint background dots disc = below\[below\["status"\] == "DISCARD"\] ax.scatter(disc.index, disc\["val\_bpb"\], c="#cccccc", s=12, alpha=0.5, zorder=2, label="Discarded") # Plot kept experiments as prominent green dots kept\_v = below\[below\["status"\] == "KEEP"\] ax.scatter(kept\_v.index, kept\_v\["val\_bpb"\], c="#2ecc71", s=50, zorder=4, label="Kept", edgecolors="black", linewidths=0.5) # Running minimum step line kept\_mask = valid\["status"\] == "KEEP" kept\_idx = valid.index\[kept\_mask\] kept\_bpb = valid.loc\[kept\_mask, "val\_bpb"\] running\_min = kept\_bpb.cummin() ax.step(kept\_idx, running\_min, where="post", color="#27ae60", linewidth=2, alpha=0.7, zorder=3, label="Running best") # Label each kept experiment with its description for idx, bpb in zip(kept\_idx, kept\_bpb): desc = str(valid.loc\[idx, "description"\]).strip() if len(desc) > 45: desc = desc\[:42\] + "..." ax.annotate(desc, (idx, bpb), textcoords="offset points", xytext=(6, 6), fontsize=8.0, color="#1a7a3a", alpha=0.9, rotation=30, ha="left", va="bottom") n\_total = len(df) n\_kept = len(df\[df\["status"\] == "KEEP"\]) ax.set\_xlabel("Experiment #", fontsize=12) ax.set\_ylabel("Validation BPB (lower is better)", fontsize=12) ax.set\_title(f"Autoresearch Progress: {n\_total} Experiments, {n\_kept} Kept Improvements", fontsize=14) ax.legend(loc="upper right", fontsize=9) ax.grid(True, alpha=0.2) margin = (baseline\_bpb - best) \* 0.15 ax.set\_ylim(best - margin, baseline\_bpb + margin) plt.tight\_layout() plt.savefig(str(plot\_path), dpi=150, bbox\_inches="tight") plt.close(fig) print(f"Saved plot to {plot\_path}", flush=True) # Upload plot to PR as a comment with base64 inline image import base64 img\_b64 = base64.b64encode(plot\_path.read\_bytes()).decode() pr\_comment = ( "## Autoresearch Progress\\n\\n" f"!\[Autoresearch Progress\](data:image/png;base64,{img\_b64})" ) pr.create\_issue\_comment(pr\_comment) print("Posted plot as PR comment.", flush=True) # Force-add plot to git and amend commit subprocess.run(\["git", "add", "-f", str(plot\_path)\], cwd=repo\_path, check=False) subprocess.run( \["git", "commit", "--amend", "--no-edit"\], cwd=repo\_path, check=False, ) subprocess.run( \["git", "push", "-u", "origin", branch\_name, "--force"\], cwd=repo\_path, check=False, ) # Show plot in Flyte UI via report await flyte.report.replace.aio( f"

Autoresearch Progress

" f'' f'

View PR

' ) await flyte.report.flush.aio() else: print("results.tsv found but no valid val\_bpb rows — skipping plot.", flush=True) else: print("results.tsv not found — skipping plot.", flush=True) return AutoResearchResult( pr\_url=pr.html\_url, pr\_number=pr.number, branch\_name=branch\_name, files\_changed=files\_changed, success=True, ) # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": import time flyte.init\_from\_config() run = flyte.with\_runcontext(mode="remote").run(run\_autoresearch) print(f"AutoResearch run started: {run.url}") print("Waiting for completion...") while True: try: run.wait() break except Exception as e: print(f"Connection dropped ({e}), reconnecting in 30s...") time.sleep(30) print(f"Done! See run at: {run.url}") # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/autoresearch/run.py\* The entry point submits the task in \`remote\` mode and reconnects automatically if the client connection drops during the long run. ## Run the agent ### Create secrets Get an Anthropic API key from the \[Anthropic console\](https://console.anthropic.com/) and a \[GitHub personal access token\](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens) with permission to push and open PRs on the target repository. Register both as Flyte secrets. The key names must match those declared in the \`TaskEnvironment\`: \`\`\` flyte create secret github\_token flyte create secret internal-anthropic-api-key \`\`\` See \[Secrets\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets/page.md) for scoping and file-based secrets. ### Prepare the research repository The target repository must contain a \`program.md\` at its root describing the research task for the agent. Point \`AUTORESEARCH\_REPO\_URL\` / \`AUTORESEARCH\_REPO\_FULL\_NAME\` (and the git identity constants) at a repo you control. ### Run remotely From the \[example directory\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/autoresearch): \`\`\` cd v2/tutorials/autoresearch python run.py \`\`\` This task runs remotely (it needs a GPU and network access). Follow the printed run URL to watch the agent's logs stream in, and open the run's report panel to see the progress plot once results are available. When the task finishes, the returned \`AutoResearchResult\` contains the pull request URL. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/code-agent === # Coding agent > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/code\_runner). This example demonstrates how to run code generated by a large language model (LLM) using a \`ContainerTask\`. The agent takes a user’s question, generates Flyte 2 code using the Flyte 2 documentation as context, and runs it in an isolated container. If the execution fails, the agent reflects on the error and retries up to a configurable limit until it succeeds. Using \`ContainerTask\` ensures that all generated code runs in a secure environment. This gives you full flexibility to execute arbitrary logic safely and reliably. ## What this example demonstrates - How to combine LLM generation with programmatic execution. - How to run untrusted or dynamically generated code securely. - How to iteratively improve code using agent-like behavior. ## Setting up the agent environment Let's start by importing the necessary libraries and setting up two environments: one for the container task and another for the agent task. This example follows the \`uv\` script format to declare dependencies. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b23",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # /// \`\`\` \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # main = "main" # params = "" # /// # {{docs-fragment code\_runner\_task}} import flyte from flyte.extras import ContainerTask from flyte.io import File code\_runner\_task = ContainerTask( name="run\_flyte\_v2", image=flyte.Image.from\_debian\_base(), input\_data\_dir="/var/inputs", output\_data\_dir="/var/outputs", inputs={"script": File}, outputs={"result": str, "exit\_code": str}, command=\[\ "/bin/bash",\ "-c",\ (\ "set -o pipefail && "\ "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "\ "echo $? > /var/outputs/exit\_code"\ ),\ \], resources=flyte.Resources(cpu=1, memory="1Gi"), ) # {{/docs-fragment code\_runner\_task}} # {{docs-fragment env}} import tempfile from typing import Optional from langchain\_core.runnables import Runnable from pydantic import BaseModel, Field container\_env = flyte.TaskEnvironment.from\_task( "code-runner-container", code\_runner\_task ) env = flyte.TaskEnvironment( name="code\_runner", secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="code-runner-agent"), resources=flyte.Resources(cpu=1), depends\_on=\[container\_env\], ) # {{/docs-fragment env}} # {{docs-fragment code\_base\_model}} class Code(BaseModel): """Schema for code solutions to questions about Flyte v2.""" prefix: str = Field( default="", description="Description of the problem and approach" ) imports: str = Field( default="", description="Code block with just import statements" ) code: str = Field( default="", description="Code block not including import statements" ) # {{/docs-fragment code\_base\_model}} # {{docs-fragment agent\_state}} class AgentState(BaseModel): messages: list\[dict\[str, str\]\] = Field(default\_factory=list) generation: Code = Field(default\_factory=Code) iterations: int = 0 error: str = "no" output: Optional\[str\] = None # {{/docs-fragment agent\_state}} # {{docs-fragment generate\_code\_gen\_chain}} async def generate\_code\_gen\_chain(debug: bool) -> Runnable: from langchain\_core.prompts import ChatPromptTemplate from langchain\_openai import ChatOpenAI # Grader prompt code\_gen\_prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ """\ You are a coding assistant with expertise in Python.\ You are able to execute the Flyte v2 code locally in a sandbox environment.\ \ Use the following pattern to execute the code:\ \ \ if \_\_name\_\_ == "\_\_main\_\_":\ flyte.init\_from\_config()\ print(flyte.run(...))\ \ \ Your response will be shown to the user.\ Here is a full set of documentation:\ \ -------\ {context}\ -------\ \ Answer the user question based on the above provided documentation.\ Ensure any code you provide can be executed with all required imports and variables defined.\ Structure your answer with a description of the code solution.\ Then list the imports. And finally list the functioning code block.\ Here is the user question:""",\ ),\ ("placeholder", "{messages}"),\ \] ) expt\_llm = "gpt-4o" if not debug else "gpt-4o-mini" llm = ChatOpenAI(temperature=0, model=expt\_llm) code\_gen\_chain = code\_gen\_prompt | llm.with\_structured\_output(Code) return code\_gen\_chain # {{/docs-fragment generate\_code\_gen\_chain}} # {{docs-fragment docs\_retriever}} @env.task async def docs\_retriever(url: str) -> str: from bs4 import BeautifulSoup from langchain\_community.document\_loaders.recursive\_url\_loader import ( RecursiveUrlLoader, ) loader = RecursiveUrlLoader( url=url, max\_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text ) docs = loader.load() # Sort the list based on the URLs and get the text d\_sorted = sorted(docs, key=lambda x: x.metadata\["source"\]) d\_reversed = list(reversed(d\_sorted)) concatenated\_content = "\\n\\n\\n --- \\n\\n\\n".join( \[doc.page\_content for doc in d\_reversed\] ) return concatenated\_content # {{/docs-fragment docs\_retriever}} # {{docs-fragment generate}} @env.task async def generate( question: str, state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Generate a code solution Args: question (str): The user question state (dict): The current graph state concatenated\_content (str): The concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, generation """ print("---GENERATING CODE SOLUTION---") messages = state.messages iterations = state.iterations error = state.error # We have been routed back to generation with an error if error == "yes": messages += \[\ {\ "role": "user",\ "content": (\ "Now, try again. Invoke the code tool to structure the output "\ "with a prefix, imports, and code block:"\ ),\ }\ \] code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Solution code\_solution = code\_gen\_chain.invoke( { "context": concatenated\_content, "messages": ( messages if messages else \[{"role": "user", "content": question}\] ), } ) messages += \[\ {\ "role": "assistant",\ "content": f"{code\_solution.prefix} \\n Imports: {code\_solution.imports} \\n Code: {code\_solution.code}",\ }\ \] return AgentState( messages=messages, generation=code\_solution, iterations=iterations + 1, error=error, output=state.output, ) # {{/docs-fragment generate}} # {{docs-fragment code\_check}} @env.task async def code\_check(state: AgentState) -> AgentState: """ Check code Args: state (dict): The current graph state Returns: state (dict): New key added to state, error """ print("---CHECKING CODE---") # State messages = state.messages code\_solution = state.generation iterations = state.iterations # Get solution components imports = code\_solution.imports.strip() code = code\_solution.code.strip() # Create temp file for imports with tempfile.NamedTemporaryFile( mode="w", suffix=".py", delete=False ) as imports\_file: imports\_file.write(imports + "\\n") imports\_path = imports\_file.name # Create temp file for code body with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code\_file: code\_file.write(imports + "\\n" + code + "\\n") code\_path = code\_file.name # Check imports import\_output, import\_exit\_code = await code\_runner\_task( script=await File.from\_local(imports\_path) ) if import\_exit\_code.strip() != "0": print("---CODE IMPORT CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the import test: {import\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=import\_output, ) else: print("---CODE IMPORT CHECK: PASSED---") # Check execution code\_output, code\_exit\_code = await code\_runner\_task( script=await File.from\_local(code\_path) ) if code\_exit\_code.strip() != "0": print("---CODE BLOCK CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the code execution test: {code\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=code\_output, ) else: print("---CODE BLOCK CHECK: PASSED---") # No errors print("---NO CODE TEST FAILURES---") return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="no", output=code\_output, ) # {{/docs-fragment code\_check}} # {{docs-fragment reflect}} @env.task async def reflect( state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Reflect on errors Args: state (dict): The current graph state concatenated\_content (str): Concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, reflection """ print("---REFLECTING---") # State messages = state.messages iterations = state.iterations code\_solution = state.generation # Prompt reflection code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Add reflection reflections = code\_gen\_chain.invoke( {"context": concatenated\_content, "messages": messages} ) messages += \[\ {\ "role": "assistant",\ "content": f"Here are reflections on the error: {reflections}",\ }\ \] return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error=state.error, output=state.output, ) # {{/docs-fragment reflect}} # {{docs-fragment main}} @env.task async def main( question: str = ( "Define a two-task pattern where the second catches OOM from the first and retries with more memory." ), url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/", max\_iterations: int = 3, debug: bool = False, ) -> str: concatenated\_content = await docs\_retriever(url=url) state: AgentState = AgentState() iterations = 0 while True: with flyte.group(f"code-generation-pass-{iterations + 1}"): state = await generate(question, state, concatenated\_content, debug) state = await code\_check(state) error = state.error iterations = state.iterations if error == "no" or iterations >= max\_iterations: print("---DECISION: FINISH---") code\_solution = state.generation prefix = code\_solution.prefix imports = code\_solution.imports code = code\_solution.code code\_output = state.output return f"""{prefix} {imports} {code} Result of code execution: {code\_output} """ else: print("---DECISION: RE-TRY SOLUTION---") state = await reflect(state, concatenated\_content, debug) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code\_runner/agent.py\* > \[!NOTE\] > You can set up access to the OpenAI API using a Flyte secret. > > \`\`\` > flyte create secret openai\_api\_key > \`\`\` We store the LLM-generated code in a structured format. This allows us to: - Enforce consistent formatting - Make debugging easier - Log and analyze generations systematically By capturing metadata alongside the raw code, we maintain transparency and make it easier to iterate or trace issues over time. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # main = "main" # params = "" # /// # {{docs-fragment code\_runner\_task}} import flyte from flyte.extras import ContainerTask from flyte.io import File code\_runner\_task = ContainerTask( name="run\_flyte\_v2", image=flyte.Image.from\_debian\_base(), input\_data\_dir="/var/inputs", output\_data\_dir="/var/outputs", inputs={"script": File}, outputs={"result": str, "exit\_code": str}, command=\[\ "/bin/bash",\ "-c",\ (\ "set -o pipefail && "\ "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "\ "echo $? > /var/outputs/exit\_code"\ ),\ \], resources=flyte.Resources(cpu=1, memory="1Gi"), ) # {{/docs-fragment code\_runner\_task}} # {{docs-fragment env}} import tempfile from typing import Optional from langchain\_core.runnables import Runnable from pydantic import BaseModel, Field container\_env = flyte.TaskEnvironment.from\_task( "code-runner-container", code\_runner\_task ) env = flyte.TaskEnvironment( name="code\_runner", secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="code-runner-agent"), resources=flyte.Resources(cpu=1), depends\_on=\[container\_env\], ) # {{/docs-fragment env}} # {{docs-fragment code\_base\_model}} class Code(BaseModel): """Schema for code solutions to questions about Flyte v2.""" prefix: str = Field( default="", description="Description of the problem and approach" ) imports: str = Field( default="", description="Code block with just import statements" ) code: str = Field( default="", description="Code block not including import statements" ) # {{/docs-fragment code\_base\_model}} # {{docs-fragment agent\_state}} class AgentState(BaseModel): messages: list\[dict\[str, str\]\] = Field(default\_factory=list) generation: Code = Field(default\_factory=Code) iterations: int = 0 error: str = "no" output: Optional\[str\] = None # {{/docs-fragment agent\_state}} # {{docs-fragment generate\_code\_gen\_chain}} async def generate\_code\_gen\_chain(debug: bool) -> Runnable: from langchain\_core.prompts import ChatPromptTemplate from langchain\_openai import ChatOpenAI # Grader prompt code\_gen\_prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ """\ You are a coding assistant with expertise in Python.\ You are able to execute the Flyte v2 code locally in a sandbox environment.\ \ Use the following pattern to execute the code:\ \ \ if \_\_name\_\_ == "\_\_main\_\_":\ flyte.init\_from\_config()\ print(flyte.run(...))\ \ \ Your response will be shown to the user.\ Here is a full set of documentation:\ \ -------\ {context}\ -------\ \ Answer the user question based on the above provided documentation.\ Ensure any code you provide can be executed with all required imports and variables defined.\ Structure your answer with a description of the code solution.\ Then list the imports. And finally list the functioning code block.\ Here is the user question:""",\ ),\ ("placeholder", "{messages}"),\ \] ) expt\_llm = "gpt-4o" if not debug else "gpt-4o-mini" llm = ChatOpenAI(temperature=0, model=expt\_llm) code\_gen\_chain = code\_gen\_prompt | llm.with\_structured\_output(Code) return code\_gen\_chain # {{/docs-fragment generate\_code\_gen\_chain}} # {{docs-fragment docs\_retriever}} @env.task async def docs\_retriever(url: str) -> str: from bs4 import BeautifulSoup from langchain\_community.document\_loaders.recursive\_url\_loader import ( RecursiveUrlLoader, ) loader = RecursiveUrlLoader( url=url, max\_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text ) docs = loader.load() # Sort the list based on the URLs and get the text d\_sorted = sorted(docs, key=lambda x: x.metadata\["source"\]) d\_reversed = list(reversed(d\_sorted)) concatenated\_content = "\\n\\n\\n --- \\n\\n\\n".join( \[doc.page\_content for doc in d\_reversed\] ) return concatenated\_content # {{/docs-fragment docs\_retriever}} # {{docs-fragment generate}} @env.task async def generate( question: str, state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Generate a code solution Args: question (str): The user question state (dict): The current graph state concatenated\_content (str): The concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, generation """ print("---GENERATING CODE SOLUTION---") messages = state.messages iterations = state.iterations error = state.error # We have been routed back to generation with an error if error == "yes": messages += \[\ {\ "role": "user",\ "content": (\ "Now, try again. Invoke the code tool to structure the output "\ "with a prefix, imports, and code block:"\ ),\ }\ \] code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Solution code\_solution = code\_gen\_chain.invoke( { "context": concatenated\_content, "messages": ( messages if messages else \[{"role": "user", "content": question}\] ), } ) messages += \[\ {\ "role": "assistant",\ "content": f"{code\_solution.prefix} \\n Imports: {code\_solution.imports} \\n Code: {code\_solution.code}",\ }\ \] return AgentState( messages=messages, generation=code\_solution, iterations=iterations + 1, error=error, output=state.output, ) # {{/docs-fragment generate}} # {{docs-fragment code\_check}} @env.task async def code\_check(state: AgentState) -> AgentState: """ Check code Args: state (dict): The current graph state Returns: state (dict): New key added to state, error """ print("---CHECKING CODE---") # State messages = state.messages code\_solution = state.generation iterations = state.iterations # Get solution components imports = code\_solution.imports.strip() code = code\_solution.code.strip() # Create temp file for imports with tempfile.NamedTemporaryFile( mode="w", suffix=".py", delete=False ) as imports\_file: imports\_file.write(imports + "\\n") imports\_path = imports\_file.name # Create temp file for code body with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code\_file: code\_file.write(imports + "\\n" + code + "\\n") code\_path = code\_file.name # Check imports import\_output, import\_exit\_code = await code\_runner\_task( script=await File.from\_local(imports\_path) ) if import\_exit\_code.strip() != "0": print("---CODE IMPORT CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the import test: {import\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=import\_output, ) else: print("---CODE IMPORT CHECK: PASSED---") # Check execution code\_output, code\_exit\_code = await code\_runner\_task( script=await File.from\_local(code\_path) ) if code\_exit\_code.strip() != "0": print("---CODE BLOCK CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the code execution test: {code\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=code\_output, ) else: print("---CODE BLOCK CHECK: PASSED---") # No errors print("---NO CODE TEST FAILURES---") return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="no", output=code\_output, ) # {{/docs-fragment code\_check}} # {{docs-fragment reflect}} @env.task async def reflect( state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Reflect on errors Args: state (dict): The current graph state concatenated\_content (str): Concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, reflection """ print("---REFLECTING---") # State messages = state.messages iterations = state.iterations code\_solution = state.generation # Prompt reflection code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Add reflection reflections = code\_gen\_chain.invoke( {"context": concatenated\_content, "messages": messages} ) messages += \[\ {\ "role": "assistant",\ "content": f"Here are reflections on the error: {reflections}",\ }\ \] return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error=state.error, output=state.output, ) # {{/docs-fragment reflect}} # {{docs-fragment main}} @env.task async def main( question: str = ( "Define a two-task pattern where the second catches OOM from the first and retries with more memory." ), url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/", max\_iterations: int = 3, debug: bool = False, ) -> str: concatenated\_content = await docs\_retriever(url=url) state: AgentState = AgentState() iterations = 0 while True: with flyte.group(f"code-generation-pass-{iterations + 1}"): state = await generate(question, state, concatenated\_content, debug) state = await code\_check(state) error = state.error iterations = state.iterations if error == "no" or iterations >= max\_iterations: print("---DECISION: FINISH---") code\_solution = state.generation prefix = code\_solution.prefix imports = code\_solution.imports code = code\_solution.code code\_output = state.output return f"""{prefix} {imports} {code} Result of code execution: {code\_output} """ else: print("---DECISION: RE-TRY SOLUTION---") state = await reflect(state, concatenated\_content, debug) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code\_runner/agent.py\* We then define a state model to persist the agent's history across iterations. This includes previous messages, generated code, and any errors encountered. Maintaining this history allows the agent to reflect on past attempts, avoid repeating mistakes, and iteratively improve the generated code. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # main = "main" # params = "" # /// # {{docs-fragment code\_runner\_task}} import flyte from flyte.extras import ContainerTask from flyte.io import File code\_runner\_task = ContainerTask( name="run\_flyte\_v2", image=flyte.Image.from\_debian\_base(), input\_data\_dir="/var/inputs", output\_data\_dir="/var/outputs", inputs={"script": File}, outputs={"result": str, "exit\_code": str}, command=\[\ "/bin/bash",\ "-c",\ (\ "set -o pipefail && "\ "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "\ "echo $? > /var/outputs/exit\_code"\ ),\ \], resources=flyte.Resources(cpu=1, memory="1Gi"), ) # {{/docs-fragment code\_runner\_task}} # {{docs-fragment env}} import tempfile from typing import Optional from langchain\_core.runnables import Runnable from pydantic import BaseModel, Field container\_env = flyte.TaskEnvironment.from\_task( "code-runner-container", code\_runner\_task ) env = flyte.TaskEnvironment( name="code\_runner", secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="code-runner-agent"), resources=flyte.Resources(cpu=1), depends\_on=\[container\_env\], ) # {{/docs-fragment env}} # {{docs-fragment code\_base\_model}} class Code(BaseModel): """Schema for code solutions to questions about Flyte v2.""" prefix: str = Field( default="", description="Description of the problem and approach" ) imports: str = Field( default="", description="Code block with just import statements" ) code: str = Field( default="", description="Code block not including import statements" ) # {{/docs-fragment code\_base\_model}} # {{docs-fragment agent\_state}} class AgentState(BaseModel): messages: list\[dict\[str, str\]\] = Field(default\_factory=list) generation: Code = Field(default\_factory=Code) iterations: int = 0 error: str = "no" output: Optional\[str\] = None # {{/docs-fragment agent\_state}} # {{docs-fragment generate\_code\_gen\_chain}} async def generate\_code\_gen\_chain(debug: bool) -> Runnable: from langchain\_core.prompts import ChatPromptTemplate from langchain\_openai import ChatOpenAI # Grader prompt code\_gen\_prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ """\ You are a coding assistant with expertise in Python.\ You are able to execute the Flyte v2 code locally in a sandbox environment.\ \ Use the following pattern to execute the code:\ \ \ if \_\_name\_\_ == "\_\_main\_\_":\ flyte.init\_from\_config()\ print(flyte.run(...))\ \ \ Your response will be shown to the user.\ Here is a full set of documentation:\ \ -------\ {context}\ -------\ \ Answer the user question based on the above provided documentation.\ Ensure any code you provide can be executed with all required imports and variables defined.\ Structure your answer with a description of the code solution.\ Then list the imports. And finally list the functioning code block.\ Here is the user question:""",\ ),\ ("placeholder", "{messages}"),\ \] ) expt\_llm = "gpt-4o" if not debug else "gpt-4o-mini" llm = ChatOpenAI(temperature=0, model=expt\_llm) code\_gen\_chain = code\_gen\_prompt | llm.with\_structured\_output(Code) return code\_gen\_chain # {{/docs-fragment generate\_code\_gen\_chain}} # {{docs-fragment docs\_retriever}} @env.task async def docs\_retriever(url: str) -> str: from bs4 import BeautifulSoup from langchain\_community.document\_loaders.recursive\_url\_loader import ( RecursiveUrlLoader, ) loader = RecursiveUrlLoader( url=url, max\_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text ) docs = loader.load() # Sort the list based on the URLs and get the text d\_sorted = sorted(docs, key=lambda x: x.metadata\["source"\]) d\_reversed = list(reversed(d\_sorted)) concatenated\_content = "\\n\\n\\n --- \\n\\n\\n".join( \[doc.page\_content for doc in d\_reversed\] ) return concatenated\_content # {{/docs-fragment docs\_retriever}} # {{docs-fragment generate}} @env.task async def generate( question: str, state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Generate a code solution Args: question (str): The user question state (dict): The current graph state concatenated\_content (str): The concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, generation """ print("---GENERATING CODE SOLUTION---") messages = state.messages iterations = state.iterations error = state.error # We have been routed back to generation with an error if error == "yes": messages += \[\ {\ "role": "user",\ "content": (\ "Now, try again. Invoke the code tool to structure the output "\ "with a prefix, imports, and code block:"\ ),\ }\ \] code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Solution code\_solution = code\_gen\_chain.invoke( { "context": concatenated\_content, "messages": ( messages if messages else \[{"role": "user", "content": question}\] ), } ) messages += \[\ {\ "role": "assistant",\ "content": f"{code\_solution.prefix} \\n Imports: {code\_solution.imports} \\n Code: {code\_solution.code}",\ }\ \] return AgentState( messages=messages, generation=code\_solution, iterations=iterations + 1, error=error, output=state.output, ) # {{/docs-fragment generate}} # {{docs-fragment code\_check}} @env.task async def code\_check(state: AgentState) -> AgentState: """ Check code Args: state (dict): The current graph state Returns: state (dict): New key added to state, error """ print("---CHECKING CODE---") # State messages = state.messages code\_solution = state.generation iterations = state.iterations # Get solution components imports = code\_solution.imports.strip() code = code\_solution.code.strip() # Create temp file for imports with tempfile.NamedTemporaryFile( mode="w", suffix=".py", delete=False ) as imports\_file: imports\_file.write(imports + "\\n") imports\_path = imports\_file.name # Create temp file for code body with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code\_file: code\_file.write(imports + "\\n" + code + "\\n") code\_path = code\_file.name # Check imports import\_output, import\_exit\_code = await code\_runner\_task( script=await File.from\_local(imports\_path) ) if import\_exit\_code.strip() != "0": print("---CODE IMPORT CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the import test: {import\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=import\_output, ) else: print("---CODE IMPORT CHECK: PASSED---") # Check execution code\_output, code\_exit\_code = await code\_runner\_task( script=await File.from\_local(code\_path) ) if code\_exit\_code.strip() != "0": print("---CODE BLOCK CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the code execution test: {code\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=code\_output, ) else: print("---CODE BLOCK CHECK: PASSED---") # No errors print("---NO CODE TEST FAILURES---") return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="no", output=code\_output, ) # {{/docs-fragment code\_check}} # {{docs-fragment reflect}} @env.task async def reflect( state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Reflect on errors Args: state (dict): The current graph state concatenated\_content (str): Concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, reflection """ print("---REFLECTING---") # State messages = state.messages iterations = state.iterations code\_solution = state.generation # Prompt reflection code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Add reflection reflections = code\_gen\_chain.invoke( {"context": concatenated\_content, "messages": messages} ) messages += \[\ {\ "role": "assistant",\ "content": f"Here are reflections on the error: {reflections}",\ }\ \] return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error=state.error, output=state.output, ) # {{/docs-fragment reflect}} # {{docs-fragment main}} @env.task async def main( question: str = ( "Define a two-task pattern where the second catches OOM from the first and retries with more memory." ), url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/", max\_iterations: int = 3, debug: bool = False, ) -> str: concatenated\_content = await docs\_retriever(url=url) state: AgentState = AgentState() iterations = 0 while True: with flyte.group(f"code-generation-pass-{iterations + 1}"): state = await generate(question, state, concatenated\_content, debug) state = await code\_check(state) error = state.error iterations = state.iterations if error == "no" or iterations >= max\_iterations: print("---DECISION: FINISH---") code\_solution = state.generation prefix = code\_solution.prefix imports = code\_solution.imports code = code\_solution.code code\_output = state.output return f"""{prefix} {imports} {code} Result of code execution: {code\_output} """ else: print("---DECISION: RE-TRY SOLUTION---") state = await reflect(state, concatenated\_content, debug) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code\_runner/agent.py\* ## Retrieve docs We define a task to load documents from a given URL and concatenate them into a single string. This string is then used as part of the LLM prompt. We set \`max\_depth = 20\` to avoid loading an excessive number of documents. However, even with this limit, the resulting context can still be quite large. To handle this, we use an LLM (GPT-4 in this case) that supports extended context windows. > \[!NOTE\] > Appending all documents into a single string can result in extremely large contexts, potentially exceeding the LLM’s token limit. > If your dataset grows beyond what a single prompt can handle, there are a couple of strategies you can use. > One option is to apply Retrieval-Augmented Generation (RAG), where you chunk the documents, embed them using a model, > store the vectors in a vector database, and retrieve only the most relevant pieces at inference time. > > An alternative approach is to pass references to full files into the prompt, allowing the LLM to decide which files are most relevant based > on natural-language search over file paths, summaries, or even contents. This method assumes that only a subset of files > will be necessary for a given task, and the LLM is responsible for navigating the structure and identifying what to read. > While this can be a lighter-weight solution for smaller datasets, its effectiveness depends on how well the LLM can > reason over file references and the reliability of its internal search heuristics. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # main = "main" # params = "" # /// # {{docs-fragment code\_runner\_task}} import flyte from flyte.extras import ContainerTask from flyte.io import File code\_runner\_task = ContainerTask( name="run\_flyte\_v2", image=flyte.Image.from\_debian\_base(), input\_data\_dir="/var/inputs", output\_data\_dir="/var/outputs", inputs={"script": File}, outputs={"result": str, "exit\_code": str}, command=\[\ "/bin/bash",\ "-c",\ (\ "set -o pipefail && "\ "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "\ "echo $? > /var/outputs/exit\_code"\ ),\ \], resources=flyte.Resources(cpu=1, memory="1Gi"), ) # {{/docs-fragment code\_runner\_task}} # {{docs-fragment env}} import tempfile from typing import Optional from langchain\_core.runnables import Runnable from pydantic import BaseModel, Field container\_env = flyte.TaskEnvironment.from\_task( "code-runner-container", code\_runner\_task ) env = flyte.TaskEnvironment( name="code\_runner", secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="code-runner-agent"), resources=flyte.Resources(cpu=1), depends\_on=\[container\_env\], ) # {{/docs-fragment env}} # {{docs-fragment code\_base\_model}} class Code(BaseModel): """Schema for code solutions to questions about Flyte v2.""" prefix: str = Field( default="", description="Description of the problem and approach" ) imports: str = Field( default="", description="Code block with just import statements" ) code: str = Field( default="", description="Code block not including import statements" ) # {{/docs-fragment code\_base\_model}} # {{docs-fragment agent\_state}} class AgentState(BaseModel): messages: list\[dict\[str, str\]\] = Field(default\_factory=list) generation: Code = Field(default\_factory=Code) iterations: int = 0 error: str = "no" output: Optional\[str\] = None # {{/docs-fragment agent\_state}} # {{docs-fragment generate\_code\_gen\_chain}} async def generate\_code\_gen\_chain(debug: bool) -> Runnable: from langchain\_core.prompts import ChatPromptTemplate from langchain\_openai import ChatOpenAI # Grader prompt code\_gen\_prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ """\ You are a coding assistant with expertise in Python.\ You are able to execute the Flyte v2 code locally in a sandbox environment.\ \ Use the following pattern to execute the code:\ \ \ if \_\_name\_\_ == "\_\_main\_\_":\ flyte.init\_from\_config()\ print(flyte.run(...))\ \ \ Your response will be shown to the user.\ Here is a full set of documentation:\ \ -------\ {context}\ -------\ \ Answer the user question based on the above provided documentation.\ Ensure any code you provide can be executed with all required imports and variables defined.\ Structure your answer with a description of the code solution.\ Then list the imports. And finally list the functioning code block.\ Here is the user question:""",\ ),\ ("placeholder", "{messages}"),\ \] ) expt\_llm = "gpt-4o" if not debug else "gpt-4o-mini" llm = ChatOpenAI(temperature=0, model=expt\_llm) code\_gen\_chain = code\_gen\_prompt | llm.with\_structured\_output(Code) return code\_gen\_chain # {{/docs-fragment generate\_code\_gen\_chain}} # {{docs-fragment docs\_retriever}} @env.task async def docs\_retriever(url: str) -> str: from bs4 import BeautifulSoup from langchain\_community.document\_loaders.recursive\_url\_loader import ( RecursiveUrlLoader, ) loader = RecursiveUrlLoader( url=url, max\_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text ) docs = loader.load() # Sort the list based on the URLs and get the text d\_sorted = sorted(docs, key=lambda x: x.metadata\["source"\]) d\_reversed = list(reversed(d\_sorted)) concatenated\_content = "\\n\\n\\n --- \\n\\n\\n".join( \[doc.page\_content for doc in d\_reversed\] ) return concatenated\_content # {{/docs-fragment docs\_retriever}} # {{docs-fragment generate}} @env.task async def generate( question: str, state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Generate a code solution Args: question (str): The user question state (dict): The current graph state concatenated\_content (str): The concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, generation """ print("---GENERATING CODE SOLUTION---") messages = state.messages iterations = state.iterations error = state.error # We have been routed back to generation with an error if error == "yes": messages += \[\ {\ "role": "user",\ "content": (\ "Now, try again. Invoke the code tool to structure the output "\ "with a prefix, imports, and code block:"\ ),\ }\ \] code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Solution code\_solution = code\_gen\_chain.invoke( { "context": concatenated\_content, "messages": ( messages if messages else \[{"role": "user", "content": question}\] ), } ) messages += \[\ {\ "role": "assistant",\ "content": f"{code\_solution.prefix} \\n Imports: {code\_solution.imports} \\n Code: {code\_solution.code}",\ }\ \] return AgentState( messages=messages, generation=code\_solution, iterations=iterations + 1, error=error, output=state.output, ) # {{/docs-fragment generate}} # {{docs-fragment code\_check}} @env.task async def code\_check(state: AgentState) -> AgentState: """ Check code Args: state (dict): The current graph state Returns: state (dict): New key added to state, error """ print("---CHECKING CODE---") # State messages = state.messages code\_solution = state.generation iterations = state.iterations # Get solution components imports = code\_solution.imports.strip() code = code\_solution.code.strip() # Create temp file for imports with tempfile.NamedTemporaryFile( mode="w", suffix=".py", delete=False ) as imports\_file: imports\_file.write(imports + "\\n") imports\_path = imports\_file.name # Create temp file for code body with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code\_file: code\_file.write(imports + "\\n" + code + "\\n") code\_path = code\_file.name # Check imports import\_output, import\_exit\_code = await code\_runner\_task( script=await File.from\_local(imports\_path) ) if import\_exit\_code.strip() != "0": print("---CODE IMPORT CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the import test: {import\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=import\_output, ) else: print("---CODE IMPORT CHECK: PASSED---") # Check execution code\_output, code\_exit\_code = await code\_runner\_task( script=await File.from\_local(code\_path) ) if code\_exit\_code.strip() != "0": print("---CODE BLOCK CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the code execution test: {code\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=code\_output, ) else: print("---CODE BLOCK CHECK: PASSED---") # No errors print("---NO CODE TEST FAILURES---") return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="no", output=code\_output, ) # {{/docs-fragment code\_check}} # {{docs-fragment reflect}} @env.task async def reflect( state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Reflect on errors Args: state (dict): The current graph state concatenated\_content (str): Concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, reflection """ print("---REFLECTING---") # State messages = state.messages iterations = state.iterations code\_solution = state.generation # Prompt reflection code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Add reflection reflections = code\_gen\_chain.invoke( {"context": concatenated\_content, "messages": messages} ) messages += \[\ {\ "role": "assistant",\ "content": f"Here are reflections on the error: {reflections}",\ }\ \] return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error=state.error, output=state.output, ) # {{/docs-fragment reflect}} # {{docs-fragment main}} @env.task async def main( question: str = ( "Define a two-task pattern where the second catches OOM from the first and retries with more memory." ), url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/", max\_iterations: int = 3, debug: bool = False, ) -> str: concatenated\_content = await docs\_retriever(url=url) state: AgentState = AgentState() iterations = 0 while True: with flyte.group(f"code-generation-pass-{iterations + 1}"): state = await generate(question, state, concatenated\_content, debug) state = await code\_check(state) error = state.error iterations = state.iterations if error == "no" or iterations >= max\_iterations: print("---DECISION: FINISH---") code\_solution = state.generation prefix = code\_solution.prefix imports = code\_solution.imports code = code\_solution.code code\_output = state.output return f"""{prefix} {imports} {code} Result of code execution: {code\_output} """ else: print("---DECISION: RE-TRY SOLUTION---") state = await reflect(state, concatenated\_content, debug) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code\_runner/agent.py\* ## Code generation Next, we define a utility function to construct the LLM chain responsible for generating Python code from user input. This chain leverages a LangChain \`PromptTemplate\` to structure the input and an OpenAI chat model to generate well-formed, Flyte 2-compatible Python scripts. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # main = "main" # params = "" # /// # {{docs-fragment code\_runner\_task}} import flyte from flyte.extras import ContainerTask from flyte.io import File code\_runner\_task = ContainerTask( name="run\_flyte\_v2", image=flyte.Image.from\_debian\_base(), input\_data\_dir="/var/inputs", output\_data\_dir="/var/outputs", inputs={"script": File}, outputs={"result": str, "exit\_code": str}, command=\[\ "/bin/bash",\ "-c",\ (\ "set -o pipefail && "\ "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "\ "echo $? > /var/outputs/exit\_code"\ ),\ \], resources=flyte.Resources(cpu=1, memory="1Gi"), ) # {{/docs-fragment code\_runner\_task}} # {{docs-fragment env}} import tempfile from typing import Optional from langchain\_core.runnables import Runnable from pydantic import BaseModel, Field container\_env = flyte.TaskEnvironment.from\_task( "code-runner-container", code\_runner\_task ) env = flyte.TaskEnvironment( name="code\_runner", secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="code-runner-agent"), resources=flyte.Resources(cpu=1), depends\_on=\[container\_env\], ) # {{/docs-fragment env}} # {{docs-fragment code\_base\_model}} class Code(BaseModel): """Schema for code solutions to questions about Flyte v2.""" prefix: str = Field( default="", description="Description of the problem and approach" ) imports: str = Field( default="", description="Code block with just import statements" ) code: str = Field( default="", description="Code block not including import statements" ) # {{/docs-fragment code\_base\_model}} # {{docs-fragment agent\_state}} class AgentState(BaseModel): messages: list\[dict\[str, str\]\] = Field(default\_factory=list) generation: Code = Field(default\_factory=Code) iterations: int = 0 error: str = "no" output: Optional\[str\] = None # {{/docs-fragment agent\_state}} # {{docs-fragment generate\_code\_gen\_chain}} async def generate\_code\_gen\_chain(debug: bool) -> Runnable: from langchain\_core.prompts import ChatPromptTemplate from langchain\_openai import ChatOpenAI # Grader prompt code\_gen\_prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ """\ You are a coding assistant with expertise in Python.\ You are able to execute the Flyte v2 code locally in a sandbox environment.\ \ Use the following pattern to execute the code:\ \ \ if \_\_name\_\_ == "\_\_main\_\_":\ flyte.init\_from\_config()\ print(flyte.run(...))\ \ \ Your response will be shown to the user.\ Here is a full set of documentation:\ \ -------\ {context}\ -------\ \ Answer the user question based on the above provided documentation.\ Ensure any code you provide can be executed with all required imports and variables defined.\ Structure your answer with a description of the code solution.\ Then list the imports. And finally list the functioning code block.\ Here is the user question:""",\ ),\ ("placeholder", "{messages}"),\ \] ) expt\_llm = "gpt-4o" if not debug else "gpt-4o-mini" llm = ChatOpenAI(temperature=0, model=expt\_llm) code\_gen\_chain = code\_gen\_prompt | llm.with\_structured\_output(Code) return code\_gen\_chain # {{/docs-fragment generate\_code\_gen\_chain}} # {{docs-fragment docs\_retriever}} @env.task async def docs\_retriever(url: str) -> str: from bs4 import BeautifulSoup from langchain\_community.document\_loaders.recursive\_url\_loader import ( RecursiveUrlLoader, ) loader = RecursiveUrlLoader( url=url, max\_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text ) docs = loader.load() # Sort the list based on the URLs and get the text d\_sorted = sorted(docs, key=lambda x: x.metadata\["source"\]) d\_reversed = list(reversed(d\_sorted)) concatenated\_content = "\\n\\n\\n --- \\n\\n\\n".join( \[doc.page\_content for doc in d\_reversed\] ) return concatenated\_content # {{/docs-fragment docs\_retriever}} # {{docs-fragment generate}} @env.task async def generate( question: str, state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Generate a code solution Args: question (str): The user question state (dict): The current graph state concatenated\_content (str): The concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, generation """ print("---GENERATING CODE SOLUTION---") messages = state.messages iterations = state.iterations error = state.error # We have been routed back to generation with an error if error == "yes": messages += \[\ {\ "role": "user",\ "content": (\ "Now, try again. Invoke the code tool to structure the output "\ "with a prefix, imports, and code block:"\ ),\ }\ \] code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Solution code\_solution = code\_gen\_chain.invoke( { "context": concatenated\_content, "messages": ( messages if messages else \[{"role": "user", "content": question}\] ), } ) messages += \[\ {\ "role": "assistant",\ "content": f"{code\_solution.prefix} \\n Imports: {code\_solution.imports} \\n Code: {code\_solution.code}",\ }\ \] return AgentState( messages=messages, generation=code\_solution, iterations=iterations + 1, error=error, output=state.output, ) # {{/docs-fragment generate}} # {{docs-fragment code\_check}} @env.task async def code\_check(state: AgentState) -> AgentState: """ Check code Args: state (dict): The current graph state Returns: state (dict): New key added to state, error """ print("---CHECKING CODE---") # State messages = state.messages code\_solution = state.generation iterations = state.iterations # Get solution components imports = code\_solution.imports.strip() code = code\_solution.code.strip() # Create temp file for imports with tempfile.NamedTemporaryFile( mode="w", suffix=".py", delete=False ) as imports\_file: imports\_file.write(imports + "\\n") imports\_path = imports\_file.name # Create temp file for code body with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code\_file: code\_file.write(imports + "\\n" + code + "\\n") code\_path = code\_file.name # Check imports import\_output, import\_exit\_code = await code\_runner\_task( script=await File.from\_local(imports\_path) ) if import\_exit\_code.strip() != "0": print("---CODE IMPORT CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the import test: {import\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=import\_output, ) else: print("---CODE IMPORT CHECK: PASSED---") # Check execution code\_output, code\_exit\_code = await code\_runner\_task( script=await File.from\_local(code\_path) ) if code\_exit\_code.strip() != "0": print("---CODE BLOCK CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the code execution test: {code\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=code\_output, ) else: print("---CODE BLOCK CHECK: PASSED---") # No errors print("---NO CODE TEST FAILURES---") return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="no", output=code\_output, ) # {{/docs-fragment code\_check}} # {{docs-fragment reflect}} @env.task async def reflect( state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Reflect on errors Args: state (dict): The current graph state concatenated\_content (str): Concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, reflection """ print("---REFLECTING---") # State messages = state.messages iterations = state.iterations code\_solution = state.generation # Prompt reflection code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Add reflection reflections = code\_gen\_chain.invoke( {"context": concatenated\_content, "messages": messages} ) messages += \[\ {\ "role": "assistant",\ "content": f"Here are reflections on the error: {reflections}",\ }\ \] return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error=state.error, output=state.output, ) # {{/docs-fragment reflect}} # {{docs-fragment main}} @env.task async def main( question: str = ( "Define a two-task pattern where the second catches OOM from the first and retries with more memory." ), url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/", max\_iterations: int = 3, debug: bool = False, ) -> str: concatenated\_content = await docs\_retriever(url=url) state: AgentState = AgentState() iterations = 0 while True: with flyte.group(f"code-generation-pass-{iterations + 1}"): state = await generate(question, state, concatenated\_content, debug) state = await code\_check(state) error = state.error iterations = state.iterations if error == "no" or iterations >= max\_iterations: print("---DECISION: FINISH---") code\_solution = state.generation prefix = code\_solution.prefix imports = code\_solution.imports code = code\_solution.code code\_output = state.output return f"""{prefix} {imports} {code} Result of code execution: {code\_output} """ else: print("---DECISION: RE-TRY SOLUTION---") state = await reflect(state, concatenated\_content, debug) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code\_runner/agent.py\* We then define a \`generate\` task responsible for producing the code solution. To improve clarity and testability, the output is structured in three parts: a short summary of the generated solution, a list of necessary imports, and the main body of executable code. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # main = "main" # params = "" # /// # {{docs-fragment code\_runner\_task}} import flyte from flyte.extras import ContainerTask from flyte.io import File code\_runner\_task = ContainerTask( name="run\_flyte\_v2", image=flyte.Image.from\_debian\_base(), input\_data\_dir="/var/inputs", output\_data\_dir="/var/outputs", inputs={"script": File}, outputs={"result": str, "exit\_code": str}, command=\[\ "/bin/bash",\ "-c",\ (\ "set -o pipefail && "\ "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "\ "echo $? > /var/outputs/exit\_code"\ ),\ \], resources=flyte.Resources(cpu=1, memory="1Gi"), ) # {{/docs-fragment code\_runner\_task}} # {{docs-fragment env}} import tempfile from typing import Optional from langchain\_core.runnables import Runnable from pydantic import BaseModel, Field container\_env = flyte.TaskEnvironment.from\_task( "code-runner-container", code\_runner\_task ) env = flyte.TaskEnvironment( name="code\_runner", secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="code-runner-agent"), resources=flyte.Resources(cpu=1), depends\_on=\[container\_env\], ) # {{/docs-fragment env}} # {{docs-fragment code\_base\_model}} class Code(BaseModel): """Schema for code solutions to questions about Flyte v2.""" prefix: str = Field( default="", description="Description of the problem and approach" ) imports: str = Field( default="", description="Code block with just import statements" ) code: str = Field( default="", description="Code block not including import statements" ) # {{/docs-fragment code\_base\_model}} # {{docs-fragment agent\_state}} class AgentState(BaseModel): messages: list\[dict\[str, str\]\] = Field(default\_factory=list) generation: Code = Field(default\_factory=Code) iterations: int = 0 error: str = "no" output: Optional\[str\] = None # {{/docs-fragment agent\_state}} # {{docs-fragment generate\_code\_gen\_chain}} async def generate\_code\_gen\_chain(debug: bool) -> Runnable: from langchain\_core.prompts import ChatPromptTemplate from langchain\_openai import ChatOpenAI # Grader prompt code\_gen\_prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ """\ You are a coding assistant with expertise in Python.\ You are able to execute the Flyte v2 code locally in a sandbox environment.\ \ Use the following pattern to execute the code:\ \ \ if \_\_name\_\_ == "\_\_main\_\_":\ flyte.init\_from\_config()\ print(flyte.run(...))\ \ \ Your response will be shown to the user.\ Here is a full set of documentation:\ \ -------\ {context}\ -------\ \ Answer the user question based on the above provided documentation.\ Ensure any code you provide can be executed with all required imports and variables defined.\ Structure your answer with a description of the code solution.\ Then list the imports. And finally list the functioning code block.\ Here is the user question:""",\ ),\ ("placeholder", "{messages}"),\ \] ) expt\_llm = "gpt-4o" if not debug else "gpt-4o-mini" llm = ChatOpenAI(temperature=0, model=expt\_llm) code\_gen\_chain = code\_gen\_prompt | llm.with\_structured\_output(Code) return code\_gen\_chain # {{/docs-fragment generate\_code\_gen\_chain}} # {{docs-fragment docs\_retriever}} @env.task async def docs\_retriever(url: str) -> str: from bs4 import BeautifulSoup from langchain\_community.document\_loaders.recursive\_url\_loader import ( RecursiveUrlLoader, ) loader = RecursiveUrlLoader( url=url, max\_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text ) docs = loader.load() # Sort the list based on the URLs and get the text d\_sorted = sorted(docs, key=lambda x: x.metadata\["source"\]) d\_reversed = list(reversed(d\_sorted)) concatenated\_content = "\\n\\n\\n --- \\n\\n\\n".join( \[doc.page\_content for doc in d\_reversed\] ) return concatenated\_content # {{/docs-fragment docs\_retriever}} # {{docs-fragment generate}} @env.task async def generate( question: str, state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Generate a code solution Args: question (str): The user question state (dict): The current graph state concatenated\_content (str): The concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, generation """ print("---GENERATING CODE SOLUTION---") messages = state.messages iterations = state.iterations error = state.error # We have been routed back to generation with an error if error == "yes": messages += \[\ {\ "role": "user",\ "content": (\ "Now, try again. Invoke the code tool to structure the output "\ "with a prefix, imports, and code block:"\ ),\ }\ \] code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Solution code\_solution = code\_gen\_chain.invoke( { "context": concatenated\_content, "messages": ( messages if messages else \[{"role": "user", "content": question}\] ), } ) messages += \[\ {\ "role": "assistant",\ "content": f"{code\_solution.prefix} \\n Imports: {code\_solution.imports} \\n Code: {code\_solution.code}",\ }\ \] return AgentState( messages=messages, generation=code\_solution, iterations=iterations + 1, error=error, output=state.output, ) # {{/docs-fragment generate}} # {{docs-fragment code\_check}} @env.task async def code\_check(state: AgentState) -> AgentState: """ Check code Args: state (dict): The current graph state Returns: state (dict): New key added to state, error """ print("---CHECKING CODE---") # State messages = state.messages code\_solution = state.generation iterations = state.iterations # Get solution components imports = code\_solution.imports.strip() code = code\_solution.code.strip() # Create temp file for imports with tempfile.NamedTemporaryFile( mode="w", suffix=".py", delete=False ) as imports\_file: imports\_file.write(imports + "\\n") imports\_path = imports\_file.name # Create temp file for code body with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code\_file: code\_file.write(imports + "\\n" + code + "\\n") code\_path = code\_file.name # Check imports import\_output, import\_exit\_code = await code\_runner\_task( script=await File.from\_local(imports\_path) ) if import\_exit\_code.strip() != "0": print("---CODE IMPORT CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the import test: {import\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=import\_output, ) else: print("---CODE IMPORT CHECK: PASSED---") # Check execution code\_output, code\_exit\_code = await code\_runner\_task( script=await File.from\_local(code\_path) ) if code\_exit\_code.strip() != "0": print("---CODE BLOCK CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the code execution test: {code\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=code\_output, ) else: print("---CODE BLOCK CHECK: PASSED---") # No errors print("---NO CODE TEST FAILURES---") return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="no", output=code\_output, ) # {{/docs-fragment code\_check}} # {{docs-fragment reflect}} @env.task async def reflect( state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Reflect on errors Args: state (dict): The current graph state concatenated\_content (str): Concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, reflection """ print("---REFLECTING---") # State messages = state.messages iterations = state.iterations code\_solution = state.generation # Prompt reflection code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Add reflection reflections = code\_gen\_chain.invoke( {"context": concatenated\_content, "messages": messages} ) messages += \[\ {\ "role": "assistant",\ "content": f"Here are reflections on the error: {reflections}",\ }\ \] return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error=state.error, output=state.output, ) # {{/docs-fragment reflect}} # {{docs-fragment main}} @env.task async def main( question: str = ( "Define a two-task pattern where the second catches OOM from the first and retries with more memory." ), url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/", max\_iterations: int = 3, debug: bool = False, ) -> str: concatenated\_content = await docs\_retriever(url=url) state: AgentState = AgentState() iterations = 0 while True: with flyte.group(f"code-generation-pass-{iterations + 1}"): state = await generate(question, state, concatenated\_content, debug) state = await code\_check(state) error = state.error iterations = state.iterations if error == "no" or iterations >= max\_iterations: print("---DECISION: FINISH---") code\_solution = state.generation prefix = code\_solution.prefix imports = code\_solution.imports code = code\_solution.code code\_output = state.output return f"""{prefix} {imports} {code} Result of code execution: {code\_output} """ else: print("---DECISION: RE-TRY SOLUTION---") state = await reflect(state, concatenated\_content, debug) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code\_runner/agent.py\* A \`ContainerTask\` then executes this code in an isolated container environment. It takes the code as input, runs it safely, and returns the program’s output and exit code. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # main = "main" # params = "" # /// # {{docs-fragment code\_runner\_task}} import flyte from flyte.extras import ContainerTask from flyte.io import File code\_runner\_task = ContainerTask( name="run\_flyte\_v2", image=flyte.Image.from\_debian\_base(), input\_data\_dir="/var/inputs", output\_data\_dir="/var/outputs", inputs={"script": File}, outputs={"result": str, "exit\_code": str}, command=\[\ "/bin/bash",\ "-c",\ (\ "set -o pipefail && "\ "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "\ "echo $? > /var/outputs/exit\_code"\ ),\ \], resources=flyte.Resources(cpu=1, memory="1Gi"), ) # {{/docs-fragment code\_runner\_task}} # {{docs-fragment env}} import tempfile from typing import Optional from langchain\_core.runnables import Runnable from pydantic import BaseModel, Field container\_env = flyte.TaskEnvironment.from\_task( "code-runner-container", code\_runner\_task ) env = flyte.TaskEnvironment( name="code\_runner", secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="code-runner-agent"), resources=flyte.Resources(cpu=1), depends\_on=\[container\_env\], ) # {{/docs-fragment env}} # {{docs-fragment code\_base\_model}} class Code(BaseModel): """Schema for code solutions to questions about Flyte v2.""" prefix: str = Field( default="", description="Description of the problem and approach" ) imports: str = Field( default="", description="Code block with just import statements" ) code: str = Field( default="", description="Code block not including import statements" ) # {{/docs-fragment code\_base\_model}} # {{docs-fragment agent\_state}} class AgentState(BaseModel): messages: list\[dict\[str, str\]\] = Field(default\_factory=list) generation: Code = Field(default\_factory=Code) iterations: int = 0 error: str = "no" output: Optional\[str\] = None # {{/docs-fragment agent\_state}} # {{docs-fragment generate\_code\_gen\_chain}} async def generate\_code\_gen\_chain(debug: bool) -> Runnable: from langchain\_core.prompts import ChatPromptTemplate from langchain\_openai import ChatOpenAI # Grader prompt code\_gen\_prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ """\ You are a coding assistant with expertise in Python.\ You are able to execute the Flyte v2 code locally in a sandbox environment.\ \ Use the following pattern to execute the code:\ \ \ if \_\_name\_\_ == "\_\_main\_\_":\ flyte.init\_from\_config()\ print(flyte.run(...))\ \ \ Your response will be shown to the user.\ Here is a full set of documentation:\ \ -------\ {context}\ -------\ \ Answer the user question based on the above provided documentation.\ Ensure any code you provide can be executed with all required imports and variables defined.\ Structure your answer with a description of the code solution.\ Then list the imports. And finally list the functioning code block.\ Here is the user question:""",\ ),\ ("placeholder", "{messages}"),\ \] ) expt\_llm = "gpt-4o" if not debug else "gpt-4o-mini" llm = ChatOpenAI(temperature=0, model=expt\_llm) code\_gen\_chain = code\_gen\_prompt | llm.with\_structured\_output(Code) return code\_gen\_chain # {{/docs-fragment generate\_code\_gen\_chain}} # {{docs-fragment docs\_retriever}} @env.task async def docs\_retriever(url: str) -> str: from bs4 import BeautifulSoup from langchain\_community.document\_loaders.recursive\_url\_loader import ( RecursiveUrlLoader, ) loader = RecursiveUrlLoader( url=url, max\_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text ) docs = loader.load() # Sort the list based on the URLs and get the text d\_sorted = sorted(docs, key=lambda x: x.metadata\["source"\]) d\_reversed = list(reversed(d\_sorted)) concatenated\_content = "\\n\\n\\n --- \\n\\n\\n".join( \[doc.page\_content for doc in d\_reversed\] ) return concatenated\_content # {{/docs-fragment docs\_retriever}} # {{docs-fragment generate}} @env.task async def generate( question: str, state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Generate a code solution Args: question (str): The user question state (dict): The current graph state concatenated\_content (str): The concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, generation """ print("---GENERATING CODE SOLUTION---") messages = state.messages iterations = state.iterations error = state.error # We have been routed back to generation with an error if error == "yes": messages += \[\ {\ "role": "user",\ "content": (\ "Now, try again. Invoke the code tool to structure the output "\ "with a prefix, imports, and code block:"\ ),\ }\ \] code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Solution code\_solution = code\_gen\_chain.invoke( { "context": concatenated\_content, "messages": ( messages if messages else \[{"role": "user", "content": question}\] ), } ) messages += \[\ {\ "role": "assistant",\ "content": f"{code\_solution.prefix} \\n Imports: {code\_solution.imports} \\n Code: {code\_solution.code}",\ }\ \] return AgentState( messages=messages, generation=code\_solution, iterations=iterations + 1, error=error, output=state.output, ) # {{/docs-fragment generate}} # {{docs-fragment code\_check}} @env.task async def code\_check(state: AgentState) -> AgentState: """ Check code Args: state (dict): The current graph state Returns: state (dict): New key added to state, error """ print("---CHECKING CODE---") # State messages = state.messages code\_solution = state.generation iterations = state.iterations # Get solution components imports = code\_solution.imports.strip() code = code\_solution.code.strip() # Create temp file for imports with tempfile.NamedTemporaryFile( mode="w", suffix=".py", delete=False ) as imports\_file: imports\_file.write(imports + "\\n") imports\_path = imports\_file.name # Create temp file for code body with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code\_file: code\_file.write(imports + "\\n" + code + "\\n") code\_path = code\_file.name # Check imports import\_output, import\_exit\_code = await code\_runner\_task( script=await File.from\_local(imports\_path) ) if import\_exit\_code.strip() != "0": print("---CODE IMPORT CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the import test: {import\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=import\_output, ) else: print("---CODE IMPORT CHECK: PASSED---") # Check execution code\_output, code\_exit\_code = await code\_runner\_task( script=await File.from\_local(code\_path) ) if code\_exit\_code.strip() != "0": print("---CODE BLOCK CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the code execution test: {code\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=code\_output, ) else: print("---CODE BLOCK CHECK: PASSED---") # No errors print("---NO CODE TEST FAILURES---") return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="no", output=code\_output, ) # {{/docs-fragment code\_check}} # {{docs-fragment reflect}} @env.task async def reflect( state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Reflect on errors Args: state (dict): The current graph state concatenated\_content (str): Concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, reflection """ print("---REFLECTING---") # State messages = state.messages iterations = state.iterations code\_solution = state.generation # Prompt reflection code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Add reflection reflections = code\_gen\_chain.invoke( {"context": concatenated\_content, "messages": messages} ) messages += \[\ {\ "role": "assistant",\ "content": f"Here are reflections on the error: {reflections}",\ }\ \] return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error=state.error, output=state.output, ) # {{/docs-fragment reflect}} # {{docs-fragment main}} @env.task async def main( question: str = ( "Define a two-task pattern where the second catches OOM from the first and retries with more memory." ), url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/", max\_iterations: int = 3, debug: bool = False, ) -> str: concatenated\_content = await docs\_retriever(url=url) state: AgentState = AgentState() iterations = 0 while True: with flyte.group(f"code-generation-pass-{iterations + 1}"): state = await generate(question, state, concatenated\_content, debug) state = await code\_check(state) error = state.error iterations = state.iterations if error == "no" or iterations >= max\_iterations: print("---DECISION: FINISH---") code\_solution = state.generation prefix = code\_solution.prefix imports = code\_solution.imports code = code\_solution.code code\_output = state.output return f"""{prefix} {imports} {code} Result of code execution: {code\_output} """ else: print("---DECISION: RE-TRY SOLUTION---") state = await reflect(state, concatenated\_content, debug) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code\_runner/agent.py\* This task verifies that the generated code runs as expected. It tests the import statements first, then executes the full code. It records the output and any error messages in the agent state for further analysis. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # main = "main" # params = "" # /// # {{docs-fragment code\_runner\_task}} import flyte from flyte.extras import ContainerTask from flyte.io import File code\_runner\_task = ContainerTask( name="run\_flyte\_v2", image=flyte.Image.from\_debian\_base(), input\_data\_dir="/var/inputs", output\_data\_dir="/var/outputs", inputs={"script": File}, outputs={"result": str, "exit\_code": str}, command=\[\ "/bin/bash",\ "-c",\ (\ "set -o pipefail && "\ "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "\ "echo $? > /var/outputs/exit\_code"\ ),\ \], resources=flyte.Resources(cpu=1, memory="1Gi"), ) # {{/docs-fragment code\_runner\_task}} # {{docs-fragment env}} import tempfile from typing import Optional from langchain\_core.runnables import Runnable from pydantic import BaseModel, Field container\_env = flyte.TaskEnvironment.from\_task( "code-runner-container", code\_runner\_task ) env = flyte.TaskEnvironment( name="code\_runner", secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="code-runner-agent"), resources=flyte.Resources(cpu=1), depends\_on=\[container\_env\], ) # {{/docs-fragment env}} # {{docs-fragment code\_base\_model}} class Code(BaseModel): """Schema for code solutions to questions about Flyte v2.""" prefix: str = Field( default="", description="Description of the problem and approach" ) imports: str = Field( default="", description="Code block with just import statements" ) code: str = Field( default="", description="Code block not including import statements" ) # {{/docs-fragment code\_base\_model}} # {{docs-fragment agent\_state}} class AgentState(BaseModel): messages: list\[dict\[str, str\]\] = Field(default\_factory=list) generation: Code = Field(default\_factory=Code) iterations: int = 0 error: str = "no" output: Optional\[str\] = None # {{/docs-fragment agent\_state}} # {{docs-fragment generate\_code\_gen\_chain}} async def generate\_code\_gen\_chain(debug: bool) -> Runnable: from langchain\_core.prompts import ChatPromptTemplate from langchain\_openai import ChatOpenAI # Grader prompt code\_gen\_prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ """\ You are a coding assistant with expertise in Python.\ You are able to execute the Flyte v2 code locally in a sandbox environment.\ \ Use the following pattern to execute the code:\ \ \ if \_\_name\_\_ == "\_\_main\_\_":\ flyte.init\_from\_config()\ print(flyte.run(...))\ \ \ Your response will be shown to the user.\ Here is a full set of documentation:\ \ -------\ {context}\ -------\ \ Answer the user question based on the above provided documentation.\ Ensure any code you provide can be executed with all required imports and variables defined.\ Structure your answer with a description of the code solution.\ Then list the imports. And finally list the functioning code block.\ Here is the user question:""",\ ),\ ("placeholder", "{messages}"),\ \] ) expt\_llm = "gpt-4o" if not debug else "gpt-4o-mini" llm = ChatOpenAI(temperature=0, model=expt\_llm) code\_gen\_chain = code\_gen\_prompt | llm.with\_structured\_output(Code) return code\_gen\_chain # {{/docs-fragment generate\_code\_gen\_chain}} # {{docs-fragment docs\_retriever}} @env.task async def docs\_retriever(url: str) -> str: from bs4 import BeautifulSoup from langchain\_community.document\_loaders.recursive\_url\_loader import ( RecursiveUrlLoader, ) loader = RecursiveUrlLoader( url=url, max\_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text ) docs = loader.load() # Sort the list based on the URLs and get the text d\_sorted = sorted(docs, key=lambda x: x.metadata\["source"\]) d\_reversed = list(reversed(d\_sorted)) concatenated\_content = "\\n\\n\\n --- \\n\\n\\n".join( \[doc.page\_content for doc in d\_reversed\] ) return concatenated\_content # {{/docs-fragment docs\_retriever}} # {{docs-fragment generate}} @env.task async def generate( question: str, state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Generate a code solution Args: question (str): The user question state (dict): The current graph state concatenated\_content (str): The concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, generation """ print("---GENERATING CODE SOLUTION---") messages = state.messages iterations = state.iterations error = state.error # We have been routed back to generation with an error if error == "yes": messages += \[\ {\ "role": "user",\ "content": (\ "Now, try again. Invoke the code tool to structure the output "\ "with a prefix, imports, and code block:"\ ),\ }\ \] code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Solution code\_solution = code\_gen\_chain.invoke( { "context": concatenated\_content, "messages": ( messages if messages else \[{"role": "user", "content": question}\] ), } ) messages += \[\ {\ "role": "assistant",\ "content": f"{code\_solution.prefix} \\n Imports: {code\_solution.imports} \\n Code: {code\_solution.code}",\ }\ \] return AgentState( messages=messages, generation=code\_solution, iterations=iterations + 1, error=error, output=state.output, ) # {{/docs-fragment generate}} # {{docs-fragment code\_check}} @env.task async def code\_check(state: AgentState) -> AgentState: """ Check code Args: state (dict): The current graph state Returns: state (dict): New key added to state, error """ print("---CHECKING CODE---") # State messages = state.messages code\_solution = state.generation iterations = state.iterations # Get solution components imports = code\_solution.imports.strip() code = code\_solution.code.strip() # Create temp file for imports with tempfile.NamedTemporaryFile( mode="w", suffix=".py", delete=False ) as imports\_file: imports\_file.write(imports + "\\n") imports\_path = imports\_file.name # Create temp file for code body with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code\_file: code\_file.write(imports + "\\n" + code + "\\n") code\_path = code\_file.name # Check imports import\_output, import\_exit\_code = await code\_runner\_task( script=await File.from\_local(imports\_path) ) if import\_exit\_code.strip() != "0": print("---CODE IMPORT CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the import test: {import\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=import\_output, ) else: print("---CODE IMPORT CHECK: PASSED---") # Check execution code\_output, code\_exit\_code = await code\_runner\_task( script=await File.from\_local(code\_path) ) if code\_exit\_code.strip() != "0": print("---CODE BLOCK CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the code execution test: {code\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=code\_output, ) else: print("---CODE BLOCK CHECK: PASSED---") # No errors print("---NO CODE TEST FAILURES---") return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="no", output=code\_output, ) # {{/docs-fragment code\_check}} # {{docs-fragment reflect}} @env.task async def reflect( state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Reflect on errors Args: state (dict): The current graph state concatenated\_content (str): Concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, reflection """ print("---REFLECTING---") # State messages = state.messages iterations = state.iterations code\_solution = state.generation # Prompt reflection code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Add reflection reflections = code\_gen\_chain.invoke( {"context": concatenated\_content, "messages": messages} ) messages += \[\ {\ "role": "assistant",\ "content": f"Here are reflections on the error: {reflections}",\ }\ \] return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error=state.error, output=state.output, ) # {{/docs-fragment reflect}} # {{docs-fragment main}} @env.task async def main( question: str = ( "Define a two-task pattern where the second catches OOM from the first and retries with more memory." ), url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/", max\_iterations: int = 3, debug: bool = False, ) -> str: concatenated\_content = await docs\_retriever(url=url) state: AgentState = AgentState() iterations = 0 while True: with flyte.group(f"code-generation-pass-{iterations + 1}"): state = await generate(question, state, concatenated\_content, debug) state = await code\_check(state) error = state.error iterations = state.iterations if error == "no" or iterations >= max\_iterations: print("---DECISION: FINISH---") code\_solution = state.generation prefix = code\_solution.prefix imports = code\_solution.imports code = code\_solution.code code\_output = state.output return f"""{prefix} {imports} {code} Result of code execution: {code\_output} """ else: print("---DECISION: RE-TRY SOLUTION---") state = await reflect(state, concatenated\_content, debug) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code\_runner/agent.py\* If an error occurs, a separate task reflects on the failure and generates a response. This reflection is added to the agent state to guide future iterations. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # main = "main" # params = "" # /// # {{docs-fragment code\_runner\_task}} import flyte from flyte.extras import ContainerTask from flyte.io import File code\_runner\_task = ContainerTask( name="run\_flyte\_v2", image=flyte.Image.from\_debian\_base(), input\_data\_dir="/var/inputs", output\_data\_dir="/var/outputs", inputs={"script": File}, outputs={"result": str, "exit\_code": str}, command=\[\ "/bin/bash",\ "-c",\ (\ "set -o pipefail && "\ "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "\ "echo $? > /var/outputs/exit\_code"\ ),\ \], resources=flyte.Resources(cpu=1, memory="1Gi"), ) # {{/docs-fragment code\_runner\_task}} # {{docs-fragment env}} import tempfile from typing import Optional from langchain\_core.runnables import Runnable from pydantic import BaseModel, Field container\_env = flyte.TaskEnvironment.from\_task( "code-runner-container", code\_runner\_task ) env = flyte.TaskEnvironment( name="code\_runner", secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="code-runner-agent"), resources=flyte.Resources(cpu=1), depends\_on=\[container\_env\], ) # {{/docs-fragment env}} # {{docs-fragment code\_base\_model}} class Code(BaseModel): """Schema for code solutions to questions about Flyte v2.""" prefix: str = Field( default="", description="Description of the problem and approach" ) imports: str = Field( default="", description="Code block with just import statements" ) code: str = Field( default="", description="Code block not including import statements" ) # {{/docs-fragment code\_base\_model}} # {{docs-fragment agent\_state}} class AgentState(BaseModel): messages: list\[dict\[str, str\]\] = Field(default\_factory=list) generation: Code = Field(default\_factory=Code) iterations: int = 0 error: str = "no" output: Optional\[str\] = None # {{/docs-fragment agent\_state}} # {{docs-fragment generate\_code\_gen\_chain}} async def generate\_code\_gen\_chain(debug: bool) -> Runnable: from langchain\_core.prompts import ChatPromptTemplate from langchain\_openai import ChatOpenAI # Grader prompt code\_gen\_prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ """\ You are a coding assistant with expertise in Python.\ You are able to execute the Flyte v2 code locally in a sandbox environment.\ \ Use the following pattern to execute the code:\ \ \ if \_\_name\_\_ == "\_\_main\_\_":\ flyte.init\_from\_config()\ print(flyte.run(...))\ \ \ Your response will be shown to the user.\ Here is a full set of documentation:\ \ -------\ {context}\ -------\ \ Answer the user question based on the above provided documentation.\ Ensure any code you provide can be executed with all required imports and variables defined.\ Structure your answer with a description of the code solution.\ Then list the imports. And finally list the functioning code block.\ Here is the user question:""",\ ),\ ("placeholder", "{messages}"),\ \] ) expt\_llm = "gpt-4o" if not debug else "gpt-4o-mini" llm = ChatOpenAI(temperature=0, model=expt\_llm) code\_gen\_chain = code\_gen\_prompt | llm.with\_structured\_output(Code) return code\_gen\_chain # {{/docs-fragment generate\_code\_gen\_chain}} # {{docs-fragment docs\_retriever}} @env.task async def docs\_retriever(url: str) -> str: from bs4 import BeautifulSoup from langchain\_community.document\_loaders.recursive\_url\_loader import ( RecursiveUrlLoader, ) loader = RecursiveUrlLoader( url=url, max\_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text ) docs = loader.load() # Sort the list based on the URLs and get the text d\_sorted = sorted(docs, key=lambda x: x.metadata\["source"\]) d\_reversed = list(reversed(d\_sorted)) concatenated\_content = "\\n\\n\\n --- \\n\\n\\n".join( \[doc.page\_content for doc in d\_reversed\] ) return concatenated\_content # {{/docs-fragment docs\_retriever}} # {{docs-fragment generate}} @env.task async def generate( question: str, state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Generate a code solution Args: question (str): The user question state (dict): The current graph state concatenated\_content (str): The concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, generation """ print("---GENERATING CODE SOLUTION---") messages = state.messages iterations = state.iterations error = state.error # We have been routed back to generation with an error if error == "yes": messages += \[\ {\ "role": "user",\ "content": (\ "Now, try again. Invoke the code tool to structure the output "\ "with a prefix, imports, and code block:"\ ),\ }\ \] code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Solution code\_solution = code\_gen\_chain.invoke( { "context": concatenated\_content, "messages": ( messages if messages else \[{"role": "user", "content": question}\] ), } ) messages += \[\ {\ "role": "assistant",\ "content": f"{code\_solution.prefix} \\n Imports: {code\_solution.imports} \\n Code: {code\_solution.code}",\ }\ \] return AgentState( messages=messages, generation=code\_solution, iterations=iterations + 1, error=error, output=state.output, ) # {{/docs-fragment generate}} # {{docs-fragment code\_check}} @env.task async def code\_check(state: AgentState) -> AgentState: """ Check code Args: state (dict): The current graph state Returns: state (dict): New key added to state, error """ print("---CHECKING CODE---") # State messages = state.messages code\_solution = state.generation iterations = state.iterations # Get solution components imports = code\_solution.imports.strip() code = code\_solution.code.strip() # Create temp file for imports with tempfile.NamedTemporaryFile( mode="w", suffix=".py", delete=False ) as imports\_file: imports\_file.write(imports + "\\n") imports\_path = imports\_file.name # Create temp file for code body with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code\_file: code\_file.write(imports + "\\n" + code + "\\n") code\_path = code\_file.name # Check imports import\_output, import\_exit\_code = await code\_runner\_task( script=await File.from\_local(imports\_path) ) if import\_exit\_code.strip() != "0": print("---CODE IMPORT CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the import test: {import\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=import\_output, ) else: print("---CODE IMPORT CHECK: PASSED---") # Check execution code\_output, code\_exit\_code = await code\_runner\_task( script=await File.from\_local(code\_path) ) if code\_exit\_code.strip() != "0": print("---CODE BLOCK CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the code execution test: {code\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=code\_output, ) else: print("---CODE BLOCK CHECK: PASSED---") # No errors print("---NO CODE TEST FAILURES---") return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="no", output=code\_output, ) # {{/docs-fragment code\_check}} # {{docs-fragment reflect}} @env.task async def reflect( state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Reflect on errors Args: state (dict): The current graph state concatenated\_content (str): Concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, reflection """ print("---REFLECTING---") # State messages = state.messages iterations = state.iterations code\_solution = state.generation # Prompt reflection code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Add reflection reflections = code\_gen\_chain.invoke( {"context": concatenated\_content, "messages": messages} ) messages += \[\ {\ "role": "assistant",\ "content": f"Here are reflections on the error: {reflections}",\ }\ \] return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error=state.error, output=state.output, ) # {{/docs-fragment reflect}} # {{docs-fragment main}} @env.task async def main( question: str = ( "Define a two-task pattern where the second catches OOM from the first and retries with more memory." ), url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/", max\_iterations: int = 3, debug: bool = False, ) -> str: concatenated\_content = await docs\_retriever(url=url) state: AgentState = AgentState() iterations = 0 while True: with flyte.group(f"code-generation-pass-{iterations + 1}"): state = await generate(question, state, concatenated\_content, debug) state = await code\_check(state) error = state.error iterations = state.iterations if error == "no" or iterations >= max\_iterations: print("---DECISION: FINISH---") code\_solution = state.generation prefix = code\_solution.prefix imports = code\_solution.imports code = code\_solution.code code\_output = state.output return f"""{prefix} {imports} {code} Result of code execution: {code\_output} """ else: print("---DECISION: RE-TRY SOLUTION---") state = await reflect(state, concatenated\_content, debug) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code\_runner/agent.py\* Finally, we define a \`main\` task that runs the code agent and orchestrates the steps above. If the code execution fails, we reflect on the error and retry until we reach the maximum number of iterations. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "langchain-core==0.3.66",\ # "langchain-openai==0.3.24",\ # "langchain-community==0.3.26",\ # "beautifulsoup4==4.13.4",\ # "docker==7.1.0",\ # \] # main = "main" # params = "" # /// # {{docs-fragment code\_runner\_task}} import flyte from flyte.extras import ContainerTask from flyte.io import File code\_runner\_task = ContainerTask( name="run\_flyte\_v2", image=flyte.Image.from\_debian\_base(), input\_data\_dir="/var/inputs", output\_data\_dir="/var/outputs", inputs={"script": File}, outputs={"result": str, "exit\_code": str}, command=\[\ "/bin/bash",\ "-c",\ (\ "set -o pipefail && "\ "uv run --script /var/inputs/script > /var/outputs/result 2>&1; "\ "echo $? > /var/outputs/exit\_code"\ ),\ \], resources=flyte.Resources(cpu=1, memory="1Gi"), ) # {{/docs-fragment code\_runner\_task}} # {{docs-fragment env}} import tempfile from typing import Optional from langchain\_core.runnables import Runnable from pydantic import BaseModel, Field container\_env = flyte.TaskEnvironment.from\_task( "code-runner-container", code\_runner\_task ) env = flyte.TaskEnvironment( name="code\_runner", secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="code-runner-agent"), resources=flyte.Resources(cpu=1), depends\_on=\[container\_env\], ) # {{/docs-fragment env}} # {{docs-fragment code\_base\_model}} class Code(BaseModel): """Schema for code solutions to questions about Flyte v2.""" prefix: str = Field( default="", description="Description of the problem and approach" ) imports: str = Field( default="", description="Code block with just import statements" ) code: str = Field( default="", description="Code block not including import statements" ) # {{/docs-fragment code\_base\_model}} # {{docs-fragment agent\_state}} class AgentState(BaseModel): messages: list\[dict\[str, str\]\] = Field(default\_factory=list) generation: Code = Field(default\_factory=Code) iterations: int = 0 error: str = "no" output: Optional\[str\] = None # {{/docs-fragment agent\_state}} # {{docs-fragment generate\_code\_gen\_chain}} async def generate\_code\_gen\_chain(debug: bool) -> Runnable: from langchain\_core.prompts import ChatPromptTemplate from langchain\_openai import ChatOpenAI # Grader prompt code\_gen\_prompt = ChatPromptTemplate.from\_messages( \[\ (\ "system",\ """\ You are a coding assistant with expertise in Python.\ You are able to execute the Flyte v2 code locally in a sandbox environment.\ \ Use the following pattern to execute the code:\ \ \ if \_\_name\_\_ == "\_\_main\_\_":\ flyte.init\_from\_config()\ print(flyte.run(...))\ \ \ Your response will be shown to the user.\ Here is a full set of documentation:\ \ -------\ {context}\ -------\ \ Answer the user question based on the above provided documentation.\ Ensure any code you provide can be executed with all required imports and variables defined.\ Structure your answer with a description of the code solution.\ Then list the imports. And finally list the functioning code block.\ Here is the user question:""",\ ),\ ("placeholder", "{messages}"),\ \] ) expt\_llm = "gpt-4o" if not debug else "gpt-4o-mini" llm = ChatOpenAI(temperature=0, model=expt\_llm) code\_gen\_chain = code\_gen\_prompt | llm.with\_structured\_output(Code) return code\_gen\_chain # {{/docs-fragment generate\_code\_gen\_chain}} # {{docs-fragment docs\_retriever}} @env.task async def docs\_retriever(url: str) -> str: from bs4 import BeautifulSoup from langchain\_community.document\_loaders.recursive\_url\_loader import ( RecursiveUrlLoader, ) loader = RecursiveUrlLoader( url=url, max\_depth=20, extractor=lambda x: BeautifulSoup(x, "html.parser").text ) docs = loader.load() # Sort the list based on the URLs and get the text d\_sorted = sorted(docs, key=lambda x: x.metadata\["source"\]) d\_reversed = list(reversed(d\_sorted)) concatenated\_content = "\\n\\n\\n --- \\n\\n\\n".join( \[doc.page\_content for doc in d\_reversed\] ) return concatenated\_content # {{/docs-fragment docs\_retriever}} # {{docs-fragment generate}} @env.task async def generate( question: str, state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Generate a code solution Args: question (str): The user question state (dict): The current graph state concatenated\_content (str): The concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, generation """ print("---GENERATING CODE SOLUTION---") messages = state.messages iterations = state.iterations error = state.error # We have been routed back to generation with an error if error == "yes": messages += \[\ {\ "role": "user",\ "content": (\ "Now, try again. Invoke the code tool to structure the output "\ "with a prefix, imports, and code block:"\ ),\ }\ \] code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Solution code\_solution = code\_gen\_chain.invoke( { "context": concatenated\_content, "messages": ( messages if messages else \[{"role": "user", "content": question}\] ), } ) messages += \[\ {\ "role": "assistant",\ "content": f"{code\_solution.prefix} \\n Imports: {code\_solution.imports} \\n Code: {code\_solution.code}",\ }\ \] return AgentState( messages=messages, generation=code\_solution, iterations=iterations + 1, error=error, output=state.output, ) # {{/docs-fragment generate}} # {{docs-fragment code\_check}} @env.task async def code\_check(state: AgentState) -> AgentState: """ Check code Args: state (dict): The current graph state Returns: state (dict): New key added to state, error """ print("---CHECKING CODE---") # State messages = state.messages code\_solution = state.generation iterations = state.iterations # Get solution components imports = code\_solution.imports.strip() code = code\_solution.code.strip() # Create temp file for imports with tempfile.NamedTemporaryFile( mode="w", suffix=".py", delete=False ) as imports\_file: imports\_file.write(imports + "\\n") imports\_path = imports\_file.name # Create temp file for code body with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as code\_file: code\_file.write(imports + "\\n" + code + "\\n") code\_path = code\_file.name # Check imports import\_output, import\_exit\_code = await code\_runner\_task( script=await File.from\_local(imports\_path) ) if import\_exit\_code.strip() != "0": print("---CODE IMPORT CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the import test: {import\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=import\_output, ) else: print("---CODE IMPORT CHECK: PASSED---") # Check execution code\_output, code\_exit\_code = await code\_runner\_task( script=await File.from\_local(code\_path) ) if code\_exit\_code.strip() != "0": print("---CODE BLOCK CHECK: FAILED---") error\_message = \[\ {\ "role": "user",\ "content": f"Your solution failed the code execution test: {code\_output}",\ }\ \] messages += error\_message return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="yes", output=code\_output, ) else: print("---CODE BLOCK CHECK: PASSED---") # No errors print("---NO CODE TEST FAILURES---") return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error="no", output=code\_output, ) # {{/docs-fragment code\_check}} # {{docs-fragment reflect}} @env.task async def reflect( state: AgentState, concatenated\_content: str, debug: bool ) -> AgentState: """ Reflect on errors Args: state (dict): The current graph state concatenated\_content (str): Concatenated docs content debug (bool): Debug mode Returns: state (dict): New key added to state, reflection """ print("---REFLECTING---") # State messages = state.messages iterations = state.iterations code\_solution = state.generation # Prompt reflection code\_gen\_chain = await generate\_code\_gen\_chain(debug) # Add reflection reflections = code\_gen\_chain.invoke( {"context": concatenated\_content, "messages": messages} ) messages += \[\ {\ "role": "assistant",\ "content": f"Here are reflections on the error: {reflections}",\ }\ \] return AgentState( generation=code\_solution, messages=messages, iterations=iterations, error=state.error, output=state.output, ) # {{/docs-fragment reflect}} # {{docs-fragment main}} @env.task async def main( question: str = ( "Define a two-task pattern where the second catches OOM from the first and retries with more memory." ), url: str = "https://pre-release-v2.docs-builder.pages.dev/docs/byoc/user-guide/", max\_iterations: int = 3, debug: bool = False, ) -> str: concatenated\_content = await docs\_retriever(url=url) state: AgentState = AgentState() iterations = 0 while True: with flyte.group(f"code-generation-pass-{iterations + 1}"): state = await generate(question, state, concatenated\_content, debug) state = await code\_check(state) error = state.error iterations = state.iterations if error == "no" or iterations >= max\_iterations: print("---DECISION: FINISH---") code\_solution = state.generation prefix = code\_solution.prefix imports = code\_solution.imports code = code\_solution.code code\_output = state.output return f"""{prefix} {imports} {code} Result of code execution: {code\_output} """ else: print("---DECISION: RE-TRY SOLUTION---") state = await reflect(state, concatenated\_content, debug) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/code\_runner/agent.py\* ## Running the code agent If things are working properly, you should see output similar to the following: \`\`\` ---GENERATING CODE SOLUTION--- ---CHECKING CODE--- ---CODE BLOCK CHECK: PASSED--- ---NO CODE TEST FAILURES--- ---DECISION: FINISH--- In this solution, we define two tasks using Flyte v2. The first task, \`oomer\`, is designed to simulate an out-of-memory (OOM) error by attempting to allocate a large list. The second task, \`failure\_recovery\`, attempts to execute \`oomer\` and catches any OOM errors. If an OOM error is caught, it retries the \`oomer\` task with increased memory resources. This pattern demonstrates how to handle resource-related exceptions and dynamically adjust task configurations in Flyte workflows. import asyncio import flyte import flyte.errors env = flyte.TaskEnvironment(name="oom\_example", resources=flyte.Resources(cpu=1, memory="250Mi")) @env.task async def oomer(x: int): large\_list = \[0\] \* 100000000 # Simulate OOM print(len(large\_list)) @env.task async def always\_succeeds() -> int: await asyncio.sleep(1) return 42 ... \`\`\` You can run the code agent on a Flyte/Union cluster using the following command: \`\`\` uv run agent.py \`\`\` === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/mle-bot === # MLE Bot: an autonomous ML engineer > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/mle\_bot). You have a dataset and a business question. Today, going from a raw CSV to a trained, evaluated model with a written report takes an ML engineer hours of experimentation: profiling the data, picking algorithms, engineering features, tuning hyperparameters, analyzing results, and iterating. What if you could describe the problem in plain English and let an agent handle the rest? This tutorial walks you through building exactly that. You'll construct an autonomous ML engineer that takes a problem description and a dataset, designs experiments, runs them on cloud infrastructure, analyzes results, iterates, and produces a report summarizing the best model it found. ## TL;DR - You'll build an agent that takes a natural language problem description and a CSV file, then produces a trained model and a detailed report comparing the results. - The LLM reasons over dataset statistics, never raw data. Trusted tools compute statistics in the cloud, and only those statistics reach the LLM. - LLM-generated orchestration code runs inside Flyte's sandbox: no imports, no network access, no filesystem. It can only call pre-approved tool functions. - Each tool function runs as a durable Flyte task in the cloud, with retries, observability, and full traceability. ## The problem with LLMs and ML pipelines If you ask an LLM to "train a model on this dataset," you run into a few issues fast. The LLM might hallucinate sklearn APIs that don't exist. It has no access to real compute, so it can't actually train anything. It runs everything in a single context with no way to handle large datasets. And if something goes wrong, there's no structured way to iterate. The core tension is that LLMs are genuinely good at reasoning about \*what\* to try. Given a dataset profile showing class imbalance and temporal structure, a capable model will suggest rolling window features and appropriate resampling strategies. But LLMs are unreliable at \*executing\* those plans. They generate buggy code, lose track of variable names, and have no way to dispatch real compute. The solution is to separate the two concerns. Let the LLM handle the planning: which algorithms to try, what feature engineering to apply, which hyperparameters to tune. Then hand the execution to trusted tool functions that run on real infrastructure. The LLM controls \*what\* happens. The tools control \*how\*. Think of it like giving a junior engineer access to a curated set of approved tools and reviewing their work. They can compose those tools in creative ways, but they can't go off-script and install random packages or hit arbitrary endpoints. ## How it works The agent runs in five phases: 1. \*\*Profile\*\* the dataset using a trusted tool. The tool returns statistics (shape, class balance, feature correlations, missing values). The LLM never touches the raw data. 2. \*\*Design\*\* a batch of experiments. The LLM reads the profile and proposes 2 to 3 experiments, each with an algorithm, hyperparameters, and a feature engineering strategy. 3. \*\*Execute\*\* each experiment in parallel. For each one, the LLM generates Python orchestration code that chains together pre-approved tool functions. That code runs inside a restricted sandbox, and each tool call dispatches as a durable Flyte task on cloud compute. 4. \*\*Analyze\*\* the results. The LLM reviews metrics across experiments, optionally requests targeted data explorations (e.g., "are failures concentrated on specific machines?"), and decides whether to iterate with new experiments. 5. \*\*Produce a report\*\* summarizing the winning model, the experiment journey, and deployment recommendations. Two things make this work. First, the LLM never sees raw data. The profiling tool runs in the cloud on managed compute and returns only aggregated statistics. This keeps prompt sizes manageable and avoids leaking sensitive data into LLM API calls. Second, the LLM-generated code runs inside Flyte's sandbox where the only thing it can do is call your pre-approved tool functions. More on that shortly. ### What to expect Here's what an actual run looks like on a synthetic predictive maintenance dataset (175k rows of sensor data from 20 industrial pumps, ~3% failure rate). In the first iteration, the agent designed three experiments: a logistic regression baseline, an XGBoost model with rolling window features, and a random forest with lag features. After reviewing the results, it decided to continue. It requested two targeted explorations ("do failure cases show meaningfully higher vibration?" and "how do feature-target correlations vary by pump?"), then used those findings to design a second round of experiments with tuned feature engineering and class weighting. After two iterations and five total experiments, the final rankings looked like this: | Rank | Experiment | ROC-AUC | F1 | Recall | Precision | |------|-----------|---------|------|--------|-----------| | 1 | \*\*Random Forest with Balanced Class Weights\*\* | 0.7983 | 0.4284 | 0.4561 | 0.4038 | | 2 | XGBoost with Feature Engineering | 0.7847 | 0.4568 | 0.4722 | 0.4425 | | 3 | Enhanced XGBoost with Focused Feature Engineering | 0.7821 | 0.3565 | 0.4973 | 0.2778 | | 4 | Random Forest with Lag Features | 0.7651 | 0.5206 | 0.4104 | 0.7116 | | 5 | Baseline Logistic Regression | 0.7528 | 0.118 | 0.6496 | 0.0649 | The agent autonomously explored different algorithms, feature strategies, and class imbalance techniques, then ranked everything by ROC-AUC. The full report includes the LLM's reasoning and generated code for every experiment, so you can trace exactly why it chose each approach and what code it wrote to implement it. Since the LLM makes different decisions each run, your results will vary, but the overall pattern (profile, design, execute, analyze, iterate) stays the same. ## Declaring task environments Before writing any tasks, you need to declare \*where\* and \*how\* they run. In Flyte v2, a \`TaskEnvironment\` bundles together the container image, resource requirements, secrets, and dependencies for a group of tasks. The MLE Bot uses two environments. One for the ML tools (pandas, sklearn, xgboost) and one for the agent itself (the OpenAI client and the sandbox runtime): \`\`\` """Flyte TaskEnvironment definitions for mle-bot. Two environments: - tool\_env: Runs the ML tools (data loading, feature engineering, training, evaluation). Has sklearn, xgboost, pandas, numpy, joblib. - agent\_env: Runs the orchestrating agent (OpenAI calls, sandbox orchestration). Has openai, pydantic-monty. Depends on tool\_env. """ # {{docs-fragment environments}} import flyte tool\_env = flyte.TaskEnvironment( "mle-tools", resources=flyte.Resources(cpu=2, memory="4Gi"), image=( flyte.Image.from\_debian\_base(name="mle-tools-image").with\_pip\_packages( "pandas>=2.0.0", "scikit-learn>=1.3.0", "xgboost>=2.0.0", "numpy>=1.24.0", "joblib>=1.3.0", ) ), ) agent\_env = flyte.TaskEnvironment( "mle-agent", resources=flyte.Resources(cpu=1, memory="2Gi"), secrets=\[flyte.Secret(key="OPENAI\_API\_KEY", as\_env\_var="OPENAI\_API\_KEY")\], env\_vars={"PYTHONUNBUFFERED": "1"}, image=( flyte.Image.from\_debian\_base(name="mle-agent-image") .with\_apt\_packages("git") .with\_pip\_packages( "openai>=1.0.0", "flyte\[sandbox\]", ) ), depends\_on=\[tool\_env\], ) # {{/docs-fragment environments}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle\_bot/mle\_bot/environments.py\* A few things to note. \`flyte.Resources\` sets the CPU and memory for every task in that environment. \`flyte.Image.from\_debian\_base()\` builds a container image on the fly with the packages you declare, so you never need to manage Dockerfiles. \`flyte.Secret\` injects a secret from your cluster's secret store as an environment variable. And \`depends\_on=\[tool\_env\]\` tells Flyte that the agent environment needs to be able to dispatch tasks in the tool environment. This is what enables the sandbox to call tool functions that run on separate, appropriately-resourced compute. ## Building durable tool functions Each tool is a regular Python async function decorated with \`@env.task\`. That decorator turns it into a durable Flyte task: it runs in its own container with the resources declared on the environment, it's automatically retried on transient failures, and every invocation is tracked in the Flyte UI. Data flows between tasks as \`flyte.io.File\` objects. A \`File\` is a reference to data in cloud storage. When a task needs the actual bytes, it calls \`await data.download()\` to pull them into the container's local filesystem. When it produces output, it creates a \`File\` from a local path and returns it. Flyte handles the upload to cloud storage when the task completes. The data itself never passes through the agent or the LLM. Here's what the training tool looks like: \`\`\` """Model training tools. A single unified interface for training classifiers with different algorithms. The tool handles serialization, class imbalance, and basic hyperparameter passing. """ from flyte.io import File from mle\_bot.environments import tool\_env from mle\_bot.schemas import ( GradientBoostingParams, LogisticRegressionParams, RandomForestParams, XGBoostParams, ) # {{docs-fragment train\_model}} @tool\_env.task async def train\_model( data: File, target\_column: str, algorithm: str, hyperparams: dict, ) -> File: """Train a classification model and return the serialized model and training metrics. Supports multiple algorithms through a single interface so the agent can dispatch different approaches without knowing implementation details. Args: data: CSV file with training data (features + target column). target\_column: Name of the column to predict. algorithm: One of: "xgboost" — Gradient boosted trees. Good default for tabular data. Handles missing values and class imbalance natively. "random\_forest" — Ensemble of decision trees. More robust to outliers. "logistic\_regression"— Linear model. Fast baseline, good for linearly separable problems. "gradient\_boosting" — Sklearn GradientBoostingClassifier. Slower than xgboost but sometimes better on small datasets. hyperparams: Dict of algorithm-specific hyperparameters. Common keys: For xgboost / gradient\_boosting: n\_estimators (int, default 100): Number of trees. max\_depth (int, default 6): Maximum tree depth. learning\_rate (float, default 0.1): Step size shrinkage. scale\_pos\_weight (float): Ratio negative/positive — use for imbalanced data. Set to (n\_negative / n\_positive) to upweight minority class. subsample (float, default 1.0): Fraction of samples used per tree. colsample\_bytree (float, default 1.0): Fraction of features per tree. For random\_forest: n\_estimators (int, default 100): Number of trees. max\_depth (int or null, default null): Maximum tree depth (null = unlimited). min\_samples\_leaf (int, default 1): Minimum samples at a leaf node. class\_weight (str, default "balanced"): "balanced" reweights by class frequency. For logistic\_regression: C (float, default 1.0): Inverse regularization strength (higher = less regularization). max\_iter (int, default 1000): Maximum iterations for solver. class\_weight (str, default "balanced"): "balanced" reweights by class frequency. Returns: File — serialized model (joblib format, contains model + feature columns + target column). """ # {{/docs-fragment train\_model}} import tempfile import joblib import numpy as np import pandas as pd from flyte.io import File as FlyteFile from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier from sklearn.linear\_model import LogisticRegression from sklearn.metrics import accuracy\_score, f1\_score, precision\_score, recall\_score, roc\_auc\_score path = await data.download() df = pd.read\_csv(path) # Only use numeric columns — drop strings like machine\_id automatically feature\_cols = \[c for c in df.select\_dtypes(include=\[np.number\]).columns if c != target\_column\] X = df\[feature\_cols\].values y = df\[target\_column\].values class\_dist = {str(k): int(v) for k, v in zip(\*np.unique(y, return\_counts=True))} n\_positive = int((y == 1).sum()) n\_negative = int((y == 0).sum()) default\_scale = max(1.0, n\_negative / n\_positive) if n\_positive > 0 else 1.0 if algorithm == "xgboost": from xgboost import XGBClassifier p = XGBoostParams.model\_validate({\*\*hyperparams, "scale\_pos\_weight": hyperparams.get("scale\_pos\_weight", default\_scale)}) params = {\*\*p.model\_dump(), "eval\_metric": "logloss", "random\_state": 42} model = XGBClassifier(\*\*params) elif algorithm == "random\_forest": p = RandomForestParams.model\_validate(hyperparams) params = {\*\*p.model\_dump(), "random\_state": 42, "n\_jobs": -1} model = RandomForestClassifier(\*\*params) elif algorithm == "gradient\_boosting": p = GradientBoostingParams.model\_validate(hyperparams) params = {\*\*p.model\_dump(), "random\_state": 42} model = GradientBoostingClassifier(\*\*params) elif algorithm == "logistic\_regression": p = LogisticRegressionParams.model\_validate(hyperparams) params = {\*\*p.model\_dump(), "random\_state": 42} model = LogisticRegression(\*\*params) else: raise ValueError(f"Unknown algorithm: {algorithm!r}. Choose from: xgboost, random\_forest, gradient\_boosting, logistic\_regression") model.fit(X, y) y\_pred = model.predict(X) y\_prob = model.predict\_proba(X)\[:, 1\] if hasattr(model, "predict\_proba") else y\_pred train\_metrics = { "accuracy": round(float(accuracy\_score(y, y\_pred)), 4), "f1": round(float(f1\_score(y, y\_pred, average="binary", zero\_division=0)), 4), "precision": round(float(precision\_score(y, y\_pred, average="binary", zero\_division=0)), 4), "recall": round(float(recall\_score(y, y\_pred, average="binary", zero\_division=0)), 4), "roc\_auc": round(float(roc\_auc\_score(y, y\_prob)), 4), } # Feature importance (top 20) if hasattr(model, "feature\_importances\_"): importances = model.feature\_importances\_ importance\_dict = {feature\_cols\[i\]: round(float(importances\[i\]), 4) for i in range(len(feature\_cols))} importance\_dict = dict(sorted(importance\_dict.items(), key=lambda x: x\[1\], reverse=True)\[:20\]) elif hasattr(model, "coef\_"): importances = abs(model.coef\_\[0\]) importance\_dict = {feature\_cols\[i\]: round(float(importances\[i\]), 4) for i in range(len(feature\_cols))} importance\_dict = dict(sorted(importance\_dict.items(), key=lambda x: x\[1\], reverse=True)\[:20\]) else: importance\_dict = {} model\_file = tempfile.NamedTemporaryFile(suffix=".joblib", delete=False) joblib.dump({"model": model, "feature\_columns": feature\_cols, "target\_column": target\_column}, model\_file.name) model\_file.close() return await FlyteFile.from\_local(model\_file.name) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle\_bot/mle\_bot/tools/training.py\* And here's the profiling tool, which is the first thing the agent calls. It computes dataset statistics that the LLM will use to design experiments: \`\`\` """Data loading, profiling, and splitting tools. These tools are safe, general-purpose, and side-effect free. They run as durable Flyte tasks so they execute in the cloud on managed compute. """ from flyte.io import File from mle\_bot.environments import tool\_env # {{docs-fragment profile\_dataset}} @tool\_env.task async def profile\_dataset(data: File, target\_column: str) -> dict: """Profile a dataset and return statistics that inform ML problem design. Call this first before designing any experiments. The returned profile tells you the shape, column types, class balance, missing values, and numeric statistics — everything needed to choose algorithms and feature strategies. Args: data: CSV file to profile. target\_column: Name of the column to predict. Returns a dict with keys: - shape: \[n\_rows, n\_cols\] - columns: list of all column names - dtypes: {col: dtype\_string, ...} - numeric\_columns: list of numeric column names (excluding target) - categorical\_columns: list of non-numeric column names (excluding target) - target\_distribution: {class\_value: count, ...} - class\_balance: {class\_value: pct, ...} (proportions, sum to 100) - missing\_pct: {col: pct\_missing, ...} - numeric\_stats: {col: {mean, std, min, max, median}, ...} - n\_classes: int — number of unique target values - is\_imbalanced: bool — True if minority class < 20% of data - sample: list of 3 example rows as dicts """ import numpy as np import pandas as pd path = await data.download() df = pd.read\_csv(path) target\_counts = df\[target\_column\].value\_counts() class\_balance = (df\[target\_column\].value\_counts(normalize=True) \* 100).round(2).to\_dict() minority\_pct = float(min(class\_balance.values())) numeric\_cols = df.select\_dtypes(include=\[np.number\]).columns.tolist() categorical\_cols = df.select\_dtypes(exclude=\[np.number\]).columns.tolist() numeric\_stats = {} for col in numeric\_cols: if col == target\_column: continue numeric\_stats\[col\] = { "mean": round(float(df\[col\].mean()), 4), "std": round(float(df\[col\].std()), 4), "min": round(float(df\[col\].min()), 4), "max": round(float(df\[col\].max()), 4), "median": round(float(df\[col\].median()), 4), } # Point-biserial correlation between each numeric feature and the target feature\_target\_corr = {} for col in numeric\_cols: if col == target\_column: continue corr = float(df\[col\].corr(df\[target\_column\])) if not np.isnan(corr): feature\_target\_corr\[col\] = round(corr, 4) # Sort by absolute correlation descending feature\_target\_corr = dict( sorted(feature\_target\_corr.items(), key=lambda x: abs(x\[1\]), reverse=True) ) return { "shape": list(df.shape), "columns": list(df.columns), "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()}, "numeric\_columns": \[c for c in numeric\_cols if c != target\_column\], "categorical\_columns": \[c for c in categorical\_cols if c != target\_column\], "target\_distribution": {str(k): int(v) for k, v in target\_counts.items()}, "class\_balance": {str(k): float(v) for k, v in class\_balance.items()}, "missing\_pct": {col: round(float(pct \* 100), 2) for col, pct in df.isnull().mean().items()}, "numeric\_stats": numeric\_stats, "feature\_target\_corr": feature\_target\_corr, "n\_classes": int(df\[target\_column\].nunique()), "is\_imbalanced": minority\_pct < 20.0, "sample": df.head(3).fillna("").to\_dict(orient="records"), } # {{/docs-fragment profile\_dataset}} @tool\_env.task async def split\_dataset( data: File, target\_column: str, test\_size: float, time\_column: str, split\_type: str, ) -> File: """Split a dataset and return either the train or test half. Call this twice — once with split\_type="train" and once with split\_type="test" — to get both halves. Always split before feature engineering to prevent data leakage. Args: data: CSV file to split. target\_column: Name of the column to predict. test\_size: Fraction of data to use for testing (e.g. 0.2 for 20%). time\_column: If non-empty, sort by this column and take the last \`test\_size\` fraction as test (time-based split, no shuffling). If empty string "", use stratified random split. split\_type: Which half to return — "train" or "test". Returns: File — CSV file containing the requested split (train or test rows). """ import tempfile import pandas as pd from flyte.io import File as FlyteFile from sklearn.model\_selection import train\_test\_split path = await data.download() df = pd.read\_csv(path) if time\_column: df = df.sort\_values(time\_column).reset\_index(drop=True) split\_idx = int(len(df) \* (1 - test\_size)) train\_df = df.iloc\[:split\_idx\] test\_df = df.iloc\[split\_idx:\] else: train\_df, test\_df = train\_test\_split( df, test\_size=test\_size, stratify=df\[target\_column\], random\_state=42, ) selected\_df = train\_df if split\_type == "train" else test\_df out = tempfile.NamedTemporaryFile(suffix=".csv", delete=False) selected\_df.to\_csv(out.name, index=False) out.close() return await FlyteFile.from\_local(out.name) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle\_bot/mle\_bot/tools/data.py\* The full tool inventory includes ten functions: \`profile\_dataset\`, \`split\_dataset\`, \`explore\_dataset\`, \`engineer\_features\`, \`select\_features\`, \`resample\_dataset\`, \`train\_model\`, \`get\_predictions\`, \`evaluate\_model\`, and \`rank\_experiments\`. Each one does exactly one thing. The LLM composes them into pipelines, but each tool enforces its own correctness guarantees internally. For example, \`resample\_dataset\` only applies resampling to training data, never test data, regardless of what the LLM asks for. ## Guiding the LLM with domain knowledge The quality of the agent's experiments depends heavily on what you tell it. The MLE Bot bakes ML best practices directly into its system prompts, so the LLM starts from a solid foundation rather than relying on whatever it picked up during pretraining. The orchestration prompt, for example, includes guidance on feature engineering strategies, class imbalance handling, and algorithm selection. It's dynamically built from the dataset profile, so the LLM sees concrete context alongside the general advice: \`\`\`python def \_build\_orchestration\_system\_prompt(profile: dict) -> str: return f"""\\ You are an expert ML engineer. Your job is to design and write the best possible pipeline for a machine learning experiment. ## Dataset context Shape: {shape\[0\]:,} rows × {shape\[1\]} columns Numeric features: {numeric\_cols} Class balance: {class\_balance}, imbalanced: {is\_imbalanced} Feature-target correlations (raw): {corr\_str} ## General ML best practices \*\*Feature engineering\*\*: - Sequential/time-series data: rolling window features capture trends that point-in-time readings miss. Choose window sizes relative to the prediction horizon and temporal resolution of the data. - Consider skipping feature engineering entirely for a baseline. \*\*Class imbalance\*\* (when is\_imbalanced=true): - Tree ensembles: use class\_weight="balanced" or scale\_pos\_weight. - The default 0.5 decision threshold may not be optimal. \*\*Algorithm selection\*\*: - XGBoost: strong default for tabular data. Start here. - RandomForest: more robust to outliers, good for noisy data. - LogisticRegression: fast linear baseline. ... """ \`\`\` This means the LLM doesn't just get a blank canvas. It gets a structured briefing that combines the actual dataset characteristics with best practices for handling them. When the profile shows class imbalance, the prompt tells it which hyperparameters to adjust and which resampling strategies to consider. When there's a timestamp column, the prompt suggests rolling window features with guidance on window sizing. The user's problem description also has a significant impact on the agent's behavior. A query like "Predict pump failures 24 hours before they happen based on sensor readings" tells the LLM that this is a time-series classification problem with a specific prediction horizon. That shapes everything: the LLM will favor temporal feature engineering (rolling windows sized relative to that 24-hour horizon), pick algorithms that handle imbalanced binary classification well, and focus on recall as a key metric because missing a failure is worse than a false alarm. Change the query to something like "Classify machine health status from the latest sensor snapshot" and the same dataset would produce a completely different set of experiments, with less emphasis on temporal features and more on cross-sectional patterns. ## The agent loop: profile, design, execute, iterate The agent's main function orchestrates five phases. Let's walk through each one. \*\*Phase 1: Profile.\*\* The agent calls \`profile\_dataset\` directly as a trusted tool. This isn't sandboxed because there's nothing to protect against here: the function is your code, running on your compute. The \`flyte.group\` call organizes this step in the Flyte UI so you can inspect it later. \`\`\`python with flyte.group("profile"): profile = await profile\_dataset(data, target\_column) \`\`\` \*\*Phase 2: Design.\*\* The profile dict goes to the LLM along with the problem description. The LLM returns a structured response matching the \`InitialDesign\` schema: \`\`\` """Pydantic schemas for tool inputs and agent data structures. These models define the expected shape of configs and results throughout the agent. Important: Tool functions that are called from the Monty sandbox must accept plain \`dict\` at the boundary (Monty can't import or instantiate classes). Each tool parses its incoming dict into the appropriate model internally for validation. In agent.py, use \`.model\_dump()\` to convert models back to dicts before passing to the sandbox. """ from typing import Literal from pydantic import BaseModel, Field # --------------------------------------------------------------------------- # Feature engineering # --------------------------------------------------------------------------- class FeatureConfig(BaseModel): """Configuration for the engineer\_features tool.""" group\_column: str = Field( default="", description="Column to group by for rolling/lag features (e.g. 'machine\_id'). " "Required when rolling\_columns or lag\_columns is specified.", ) time\_column: str = Field( default="", description="Timestamp column to sort by before computing rolling/lag features.", ) rolling\_columns: list\[str\] = Field( default\_factory=list, description="Numeric columns to compute rolling statistics for (mean, std, min, max).", ) windows: list\[int\] = Field( default\_factory=list, description="Rolling window sizes in rows (e.g. \[6, 12, 24\]).", ) lag\_columns: list\[str\] = Field( default\_factory=list, description="Numeric columns to create lag features for.", ) lags: list\[int\] = Field( default\_factory=list, description="Lag steps in rows (e.g. \[1, 3, 6\]).", ) normalize: bool = Field( default=False, description="If true, z-score normalize all numeric columns except target\_column.", ) target\_column: str = Field( default="", description="Column to exclude from normalization. Required when normalize=True.", ) drop\_columns: list\[str\] = Field( default\_factory=list, description="Columns to remove from output (e.g. raw timestamp after rolling).", ) fillna\_method: Literal\["forward", "zero", "drop"\] = Field( default="forward", description="How to fill NaN values introduced by rolling/lag. " "'forward' forward-fills then fills remaining with 0. " "'zero' fills all NaN with 0. 'drop' drops rows with NaN.", ) # --------------------------------------------------------------------------- # Training hyperparameters (per algorithm) # --------------------------------------------------------------------------- class XGBoostParams(BaseModel): n\_estimators: int = Field(default=100, ge=1) max\_depth: int = Field(default=6, ge=1, le=20) learning\_rate: float = Field(default=0.1, gt=0, le=1) scale\_pos\_weight: float = Field( default=1.0, ge=0, description="Set to n\_negative/n\_positive for imbalanced datasets.", ) subsample: float = Field(default=1.0, gt=0, le=1) colsample\_bytree: float = Field(default=1.0, gt=0, le=1) class RandomForestParams(BaseModel): n\_estimators: int = Field(default=100, ge=1) max\_depth: int | None = Field( default=None, description="Maximum tree depth. None means unlimited.", ) min\_samples\_leaf: int = Field(default=1, ge=1) class\_weight: Literal\["balanced", "balanced\_subsample"\] | None = Field(default="balanced") class GradientBoostingParams(BaseModel): n\_estimators: int = Field(default=100, ge=1) max\_depth: int = Field(default=3, ge=1, le=10) learning\_rate: float = Field(default=0.1, gt=0, le=1) subsample: float = Field(default=1.0, gt=0, le=1) class LogisticRegressionParams(BaseModel): C: float = Field(default=1.0, gt=0, description="Inverse regularization strength.") max\_iter: int = Field(default=1000, ge=100) class\_weight: Literal\["balanced"\] | None = Field(default="balanced") # --------------------------------------------------------------------------- # Experiment design (used by agent.py, validated when parsing LLM JSON) # --------------------------------------------------------------------------- Algorithm = Literal\["xgboost", "random\_forest", "gradient\_boosting", "logistic\_regression"\] # {{docs-fragment schemas}} class ExperimentConfig(BaseModel): """One experiment to run — produced by the LLM and executed by the agent.""" name: str = Field(description="Short descriptive name for this experiment.") algorithm: Algorithm hyperparams: dict = Field( default\_factory=dict, description="Algorithm-specific hyperparameters. Will be validated inside train\_model.", ) feature\_config: FeatureConfig = Field(default\_factory=FeatureConfig) rationale: str = Field(default="", description="Why this experiment is worth running.") class InitialDesign(BaseModel): """LLM response for initial experiment design.""" problem\_type: str = Field(default="binary\_classification") primary\_metric: Literal\["roc\_auc", "f1", "recall"\] = Field(default="roc\_auc") reasoning: str experiments: list\[ExperimentConfig\] class IterationDecision(BaseModel): """LLM response after analyzing experiment results.""" should\_continue: bool reasoning: str exploration\_requests: list\[dict\] = Field( default\_factory=list, description="Optional list of explore\_dataset config dicts to run before designing " "the next batch. Each dict is passed directly to explore\_dataset.", ) next\_experiments: list\[ExperimentConfig\] = Field(default\_factory=list) # {{/docs-fragment schemas}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle\_bot/mle\_bot/schemas.py\* The LLM typically proposes 2 to 3 experiments: a baseline with minimal feature engineering, an experiment with rolling window features for temporal data, and perhaps one testing a different algorithm or resampling strategy. \*\*Phase 3: Execute in parallel.\*\* All experiments in a batch run simultaneously using \`asyncio.gather()\`. Each experiment dispatches its own set of durable Flyte tasks: \`\`\` """MLE Agent — orchestrates ML experiments using Flyte's durable sandbox. The agent: 1. Profiles the dataset using a trusted tool (data never touches the LLM). 2. Asks OpenAI to design a set of experiments (algorithms, hyperparams, feature strategy). 3. For each experiment, generates Monty orchestration code and executes it via flyte.sandbox.orchestrate\_local(), which dispatches the heavy compute as durable tasks. 4. Analyzes results, iterates if needed. 5. Produces a model card summarizing the winning model. The Monty sandbox ensures the LLM-generated orchestration code is safe — it can only call the pre-approved tool functions and has no access to imports, network, or filesystem. """ import asyncio import inspect import json import os import textwrap from dataclasses import dataclass import flyte import flyte.sandbox from flyte.io import File from mle\_bot.schemas import ExperimentConfig, InitialDesign, IterationDecision from mle\_bot.environments import agent\_env from mle\_bot.tools.data import profile\_dataset, split\_dataset from mle\_bot.tools.evaluation import evaluate\_model, rank\_experiments from mle\_bot.tools.exploration import explore\_dataset from mle\_bot.tools.features import engineer\_features from mle\_bot.tools.predictions import get\_predictions from mle\_bot.tools.resampling import resample\_dataset from mle\_bot.tools.selection import select\_features from mle\_bot.tools.training import train\_model # {{docs-fragment tools}} # All tools exposed to the sandbox. # Keys must match the function names used in LLM-generated orchestration code. TOOLS = \[\ profile\_dataset, split\_dataset, explore\_dataset,\ engineer\_features, resample\_dataset, select\_features,\ train\_model, get\_predictions, evaluate\_model, rank\_experiments,\ \] TOOLS\_BY\_NAME = {t.func.\_\_name\_\_ if hasattr(t, "func") else t.\_\_name\_\_: t for t in TOOLS} # {{/docs-fragment tools}} # --------------------------------------------------------------------------- # Prompt builders # --------------------------------------------------------------------------- def \_tool\_signatures() -> str: """Build a summary of available tool signatures and docstrings for the system prompt.""" parts = \[\] for t in TOOLS: func = t.func if hasattr(t, "func") else t sig = inspect.signature(func) doc = inspect.getdoc(func) or "" # Trim docstring to first 40 lines so prompt stays manageable doc\_lines = doc.splitlines()\[:40\] short\_doc = "\\n ".join(doc\_lines) parts.append(f"async def {func.\_\_name\_\_}{sig}:\\n \\"\\"\\"{short\_doc}\\"\\"\\"\\n ...") return "\\n\\n".join(parts) # {{docs-fragment orchestration\_prompt}} def \_build\_orchestration\_system\_prompt(profile: dict) -> str: monty\_rules = flyte.sandbox.ORCHESTRATOR\_SYNTAX\_PROMPT tools\_section = \_tool\_signatures() is\_imbalanced = profile.get("is\_imbalanced", False) class\_balance = profile.get("class\_balance", {}) columns = profile.get("columns", \[\]) numeric\_cols = profile.get("numeric\_columns", \[\]) categorical\_cols = profile.get("categorical\_columns", \[\]) corr = profile.get("feature\_target\_corr", {}) corr\_str = ", ".join(f"{k}: {v:+.3f}" for k, v in list(corr.items())\[:8\]) if corr else "n/a" shape = profile.get("shape", \[0, 0\]) return f"""\\ You are an expert ML engineer. Your job is to design and write the best possible pipeline for a machine learning experiment, then generate the Python orchestration code to execute it. The code runs inside a restricted sandbox. The last expression in your code is returned as the result. All tool calls are made like regular function calls — you do NOT need to await them. ## Dataset context Shape: {shape\[0\]:,} rows × {shape\[1\]} columns Numeric features: {numeric\_cols} Categorical features (excluded from model — not supported): {categorical\_cols} Class balance: {class\_balance}, imbalanced: {is\_imbalanced} Feature-target correlations (raw, point-biserial): {corr\_str} ## General ML best practices — apply these based on the dataset context above \*\*Feature engineering\*\* (engineer\_features tool): - Sequential/time-series data (timestamp column present, rows ordered over time): rolling window features (means, stds, min/max) capture trends that point-in-time readings miss. Lag features capture recent history. Choose window sizes relative to the prediction horizon and temporal resolution of the data. - Tabular cross-sectional data: normalization helps linear models and distance-based methods. Interaction terms can help if correlations are weak individually. - Consider skipping feature engineering entirely for a baseline — it establishes whether raw features already carry enough signal. \*\*Class imbalance\*\* (when is\_imbalanced=true): - Tree ensembles: use class\_weight="balanced" or scale\_pos\_weight=n\_neg/n\_pos. - Threshold: the default 0.5 decision threshold may not be optimal — the model's probability output is what matters, threshold is tuned at deployment time. - Metric: ROC-AUC is robust to imbalance; avg\_precision is better when positives are very rare. \*\*Algorithm selection\*\*: - XGBoost / GradientBoosting: strong default for tabular data, handles missing values, built-in imbalance handling. Start here unless data is very small. - RandomForest: more robust to outliers, good for noisy data, parallelizes well. - LogisticRegression: fast linear baseline. Useful to establish whether the problem is linearly separable before adding complexity. - Prefer simpler models when n\_samples < 5,000 to avoid overfitting. \*\*Resampling\*\* (resample\_dataset tool) — data-level imbalance handling: - Use when class\_weight/scale\_pos\_weight alone isn't improving recall adequately, or when you want to test whether data-level vs algorithm-level imbalance handling works better for this dataset. - ONLY resample the TRAIN split — never test. Resampling test data gives misleading metrics. - "oversample": fast, no new information, good first try. - "smote": synthetic samples via interpolation — more diverse than random oversampling, better for high-dimensional or sparse minority classes. - "undersample": loses majority data — only useful when majority class is very large and training speed is a concern. \*\*Feature selection\*\* (select\_features tool) — prune after feature engineering: - Use after engineer\_features when the feature count is large (20+) and you suspect many features are redundant or noisy (e.g., rolling stats at many window sizes). - "mutual\_info": ranks by non-linear dependence with target — best general choice. - "variance\_threshold": drops near-constant features — cheap first pass. - "correlation\_filter": drops redundant features that are highly correlated with each other — useful when many rolling windows capture the same trend. - Can be applied before or after splitting. Apply the same selection to both train and test to ensure the model sees the same features at evaluation time. \*\*Prediction output\*\* (get\_predictions tool) — enables two advanced patterns: 1. Error analysis: train a model → get\_predictions(model, test\_file, target) → explore\_dataset(predictions\_file, {{"class\_distributions": \["feature\_x"\], "target\_column": "correct"}}) to see which examples the model gets wrong. Use this to inform feature engineering for the next iteration. 2. Stacking: train base\_model → get\_predictions(base\_model, train\_file, target) → train a meta\_model on the predictions CSV (use "predicted\_prob" as a feature alongside original features) → evaluate meta\_model on test. get\_predictions returns a CSV with columns: all originals + predicted\_prob, predicted\_class, correct. \*\*Pipeline structure\*\* — you are not required to follow a fixed sequence. Design what makes sense for this specific experiment. ## Available tools {tools\_section} ## Monty sandbox restrictions {monty\_rules} ## Critical patterns for using tool results split\_dataset returns a File — call it twice: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") engineer\_features returns a File — chain calls freely: eng = engineer\_features(train\_file, {{"rolling\_columns": \[...\], "windows": \[...\]}}) eng2 = engineer\_features(eng, {{"normalize": true, "target\_column": target\_column}}) train\_model returns a File — pass directly to evaluate\_model: model\_file = train\_model(train\_file, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, test\_file, target\_column) evaluate\_model returns a dict — subscript reads are allowed: roc = eval\_result\["metrics"\]\["roc\_auc"\] Do NOT use augmented assignment (+=), subscript assignment (d\["k"\]=v), or try/except. Build dicts as literals only. The last expression (no assignment) is the return value. ## When fixing a previous error Read the error and the failing code carefully before writing a fix. Identify the root cause — do not just change variable names or add no-ops. Trace what each tool returns, what each subsequent call expects, and where the mismatch is. Then fix the underlying logic, not just the surface symptom. ## Pipeline design — you decide the structure You are NOT required to follow a fixed sequence. Design the pipeline that makes most sense for the experiment. Examples of valid approaches: Baseline (no feature engineering): train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") model\_file = train\_model(train\_file, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, test\_file, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Two-stage feature engineering (rolling then normalize separately): train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") rolled\_train = engineer\_features(train\_file, {{"rolling\_columns": \["vibration"\], "windows": \[6, 24\]}}) rolled\_test = engineer\_features(test\_file, {{"rolling\_columns": \["vibration"\], "windows": \[6, 24\]}}) eng\_train = engineer\_features(rolled\_train, {{"normalize": true, "target\_column": target\_column}}) eng\_test = engineer\_features(rolled\_test, {{"normalize": true, "target\_column": target\_column}}) model\_file = train\_model(eng\_train, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, eng\_test, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Compare two class weightings and return the better model: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") model\_a = train\_model(train\_file, target\_column, "xgboost", {{"n\_estimators": 100, "scale\_pos\_weight": 10}}) model\_b = train\_model(train\_file, target\_column, "xgboost", {{"n\_estimators": 100, "scale\_pos\_weight": 33}}) eval\_a = evaluate\_model(model\_a, test\_file, target\_column) eval\_b = evaluate\_model(model\_b, test\_file, target\_column) best\_eval = eval\_a if eval\_a\["metrics"\]\["roc\_auc"\] > eval\_b\["metrics"\]\["roc\_auc"\] else eval\_b {{"experiment\_name": experiment\_name, "algorithm": "xgboost", "metrics": best\_eval\["metrics"\], "confusion\_matrix": best\_eval\["confusion\_matrix"\], "threshold\_analysis": best\_eval\["threshold\_analysis"\], "n\_samples": best\_eval\["n\_samples"\]}} SMOTE oversampling before training: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") eng\_train = engineer\_features(train\_file, {{"rolling\_columns": \["vibration\_mms"\], "windows": \[6, 12\]}}) eng\_test = engineer\_features(test\_file, {{"rolling\_columns": \["vibration\_mms"\], "windows": \[6, 12\]}}) resampled\_train = resample\_dataset(eng\_train, target\_column, {{"strategy": "smote", "target\_ratio": 0.2}}) model\_file = train\_model(resampled\_train, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, eng\_test, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Feature engineering followed by feature selection: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") eng\_train = engineer\_features(train\_file, {{"rolling\_columns": \["vibration\_mms", "temperature\_c"\], "windows": \[6, 12, 24\]}}) eng\_test = engineer\_features(test\_file, {{"rolling\_columns": \["vibration\_mms", "temperature\_c"\], "windows": \[6, 12, 24\]}}) sel\_train = select\_features(eng\_train, target\_column, {{"method": "mutual\_info", "k": 15}}) sel\_test = select\_features(eng\_test, target\_column, {{"method": "mutual\_info", "k": 15}}) model\_file = train\_model(sel\_train, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, sel\_test, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Error analysis — explore what the model gets wrong, then return that as insight: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") model\_file = train\_model(train\_file, target\_column, algorithm, hyperparams) pred\_file = get\_predictions(model\_file, test\_file, target\_column) error\_analysis = explore\_dataset(pred\_file, {{"target\_column": "correct", "class\_distributions": \["vibration\_mms", "temperature\_c"\]}}) eval\_result = evaluate\_model(model\_file, test\_file, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\], "error\_analysis": error\_analysis}} The last expression MUST be a dict with at minimum these keys: experiment\_name, algorithm, metrics, confusion\_matrix, threshold\_analysis, n\_samples Additional keys (e.g. error\_analysis) are allowed and will appear in the report. ## Response format Respond in exactly this format: ## Reasoning \[Your thinking: what pipeline makes sense for this experiment and why. Consider whether\ feature engineering helps, whether class imbalance needs special treatment, whether\ chaining multiple steps adds value, etc.\] ## Code \`\`\`python \[your orchestration code\] \`\`\` """ # {{/docs-fragment orchestration\_prompt}} def \_build\_analysis\_system\_prompt(max\_iterations: int, current\_iteration: int) -> str: remaining = max\_iterations - current\_iteration - 1 return f"""\\ You are an expert ML engineer analyzing experiment results to guide the next iteration of model development. You must respond with valid JSON only — no markdown, no explanation outside the JSON. Response format: {{ "should\_continue": true | false, "reasoning": "What you observed, what it tells you, and what to try next", "exploration\_requests": \[\ {{\ "question": "The specific hypothesis you are testing, e.g. 'Do failure cases show meaningfully higher vibration than healthy cases?'",\ "analysis\_type": "class\_distributions",\ "target\_column": "failure\_24h",\ "class\_distributions": \["vibration\_mms", "temperature\_c"\]\ }}\ \], "next\_experiments": \[\ {{\ "name": "descriptive experiment name",\ "algorithm": "xgboost" | "random\_forest" | "gradient\_boosting" | "logistic\_regression",\ "hyperparams": {{ ... algorithm-specific hyperparams ... }},\ "feature\_config": {{\ "group\_column": "...",\ "time\_column": "...",\ "rolling\_columns": \[...\],\ "windows": \[...\],\ "lag\_columns": \[...\],\ "lags": \[...\],\ "normalize": true | false,\ "drop\_columns": \[...\],\ "fillna\_method": "forward"\ }},\ "rationale": "Why this experiment is worth trying"\ }}\ \] }} exploration\_requests rules: - Max 2 requests per iteration. - Each request targets EXACTLY ONE analysis\_type. Do not mix multiple types in one request. - Supported analysis\_type values and their required config fields: "class\_distributions" → requires: target\_column, class\_distributions (list of columns) "correlation\_matrix" → requires: correlation\_matrix: true "temporal\_trend" → requires: temporal\_trend: {{time\_column, target\_column, freq}} "group\_stats" → requires: group\_stats: {{group\_column, target\_column}} "outlier\_summary" → requires: outlier\_summary (list of columns) "feature\_target\_corr\_by\_group" → requires: feature\_target\_corr\_by\_group: {{group\_column, target\_column, feature\_columns}} - The "question" field is required. It must be a specific testable hypothesis, not a general request. Bad: "explore the data". Good: "Is vibration\_mms higher for failures?" - Set exploration\_requests to \[\] if the current results already tell you enough to design the next experiments. Only explore when you have a concrete unanswered question. When deciding next experiments, reason about WHAT WAS TRIED vs what hasn't been explored. Each result includes used\_feature\_engineering, used\_rolling\_features, used\_lag\_features. Think systematically: if no feature engineering was tried yet, does the data profile suggest it would help (weak raw correlations, temporal/sequential structure)? If feature engineering helped, can it be improved? Avoid experiments identical to ones tried. Iteration context: this is iteration {current\_iteration + 1} of {max\_iterations} requested. Remaining iterations allowed: {remaining}. Set should\_continue=false only if: - Best ROC-AUC >= 0.97, OR - No remaining iterations (remaining == 0), OR - Results have genuinely plateaued (< 0.005 ROC-AUC improvement over last iteration AND you have already tried the most promising directions) Otherwise keep exploring — the user asked for {max\_iterations} iterations for a reason. """ def \_build\_initial\_design\_system\_prompt() -> str: return """\\ You are an expert ML engineer. Given a dataset profile and a problem description, design the first batch of experiments to run. You must respond with valid JSON only — no markdown, no explanation outside the JSON. Response format: { "problem\_type": "binary\_classification", "primary\_metric": "roc\_auc" | "f1" | "recall", "reasoning": "Brief description of your strategy", "experiments": \[\ {\ "name": "descriptive experiment name",\ "algorithm": "xgboost" | "random\_forest" | "gradient\_boosting" | "logistic\_regression",\ "hyperparams": { ... algorithm-specific hyperparams ... },\ "feature\_config": {\ "group\_column": "",\ "time\_column": "",\ "rolling\_columns": \[\],\ "windows": \[\],\ "lag\_columns": \[\],\ "lags": \[\],\ "normalize": false,\ "drop\_columns": \[\],\ "fillna\_method": "forward"\ },\ "rationale": "Why this experiment makes sense given the data profile"\ }\ \] } Design 2-3 experiments for the first batch. Good first batches typically include: - A fast baseline to establish a floor (e.g. logistic\_regression with default settings) - Your best initial hypothesis given the data profile - Optionally one variant that tests a specific idea suggested by the profile Use the dataset profile to guide your choices: - feature\_target\_corr: weak raw correlations suggest feature engineering may help - categorical\_columns: note these are excluded from the model automatically - is\_imbalanced: handle with class\_weight or scale\_pos\_weight - Shape and column types should inform algorithm complexity (simpler models for small datasets) - A time column suggests sequential structure; lag/rolling features may capture temporal patterns The feature\_config in each experiment describes what engineer\_features should apply. Leave all fields empty/false if no feature engineering is needed for that experiment. The orchestration code generator will decide the exact pipeline — your job here is to specify what the experiment is trying to learn, not to prescribe every implementation detail. """ # --------------------------------------------------------------------------- # LLM client # --------------------------------------------------------------------------- def \_openai\_client(): from openai import OpenAI return OpenAI(api\_key=os.environ\["OPENAI\_API\_KEY"\]) async def \_call\_llm(system: str, messages: list\[dict\], model: str = "gpt-4o") -> str: """Call OpenAI and return the response text.""" client = \_openai\_client() response = await asyncio.to\_thread( client.chat.completions.create, model=model, messages=\[{"role": "system", "content": system}, \*messages\], temperature=0.2, ) return response.choices\[0\].message.content def \_extract\_code(text: str) -> str: """Pull Python code out of a markdown code block.""" if "\`\`\`python" in text: start = text.index("\`\`\`python") + len("\`\`\`python") end = text.index("\`\`\`", start) return text\[start:end\].strip() if "\`\`\`" in text: start = text.index("\`\`\`") + 3 end = text.index("\`\`\`", start) return text\[start:end\].strip() return text.strip() def \_extract\_reasoning(text: str) -> str: """Extract the ## Reasoning section from LLM response.""" if "## Reasoning" in text: start = text.index("## Reasoning") + len("## Reasoning") if "## Code" in text: end = text.index("## Code") return text\[start:end\].strip() return text\[start:\].strip() return "" def \_parse\_json(text: str) -> dict: """Extract and parse JSON from LLM response.""" text = text.strip() if "\`\`\`json" in text: start = text.index("\`\`\`json") + 7 end = text.index("\`\`\`", start) text = text\[start:end\].strip() elif "\`\`\`" in text: start = text.index("\`\`\`") + 3 end = text.index("\`\`\`", start) text = text\[start:end\].strip() return json.loads(text) # --------------------------------------------------------------------------- # Display helpers # --------------------------------------------------------------------------- def \_recommend\_threshold(threshold\_analysis: list, min\_precision: float = 0.70) -> dict | None: """Find the threshold that maximises recall subject to precision >= min\_precision.""" candidates = \[t for t in threshold\_analysis if t\["precision"\] >= min\_precision\] if not candidates: return None return max(candidates, key=lambda t: t\["recall"\]) def \_print\_experiment\_table(results: list\["ExperimentResult"\], best\_name: str) -> None: """Print a ranked comparison table of all experiments.""" sorted\_results = sorted(results, key=lambda r: r.metrics.get("roc\_auc", 0), reverse=True) print("\\n" + "─" \* 78) print(f" {'Rank':<5} {'Experiment':<32} {'ROC-AUC':<9} {'F1':<7} {'Recall':<8} {'Note'}") print("─" \* 78) for rank, r in enumerate(sorted\_results, 1): note = "◀ winner" if r.name == best\_name else "" roc = r.metrics.get("roc\_auc", 0) f1 = r.metrics.get("f1", 0) recall = r.metrics.get("recall", 0) print(f" {rank:<5} {r.name:<32} {roc:<9.4f} {f1:<7.4f} {recall:<8.4f} {note}") print("─" \* 78) def \_print\_threshold\_recommendation(threshold\_analysis: list, default\_metrics: dict) -> None: """Print the operational threshold recommendation.""" rec = \_recommend\_threshold(threshold\_analysis) if not rec: return default\_recall = default\_metrics.get("recall", 0) default\_precision = default\_metrics.get("precision", 0) missed\_pct = round((1 - rec\["recall"\]) \* 100, 1) false\_alarm\_pct = round((1 - rec\["precision"\]) \* 100, 1) print(f"\\n Recommended decision threshold: {rec\['threshold'\]}") print(f" ├─ Precision : {rec\['precision'\]:.0%} ({false\_alarm\_pct}% of alerts are false alarms)") print(f" ├─ Recall : {rec\['recall'\]:.0%} (catches {rec\['recall'\]\*100:.0f}% of actual failures)") print(f" └─ F1 : {rec\['f1'\]:.4f}") print(f" Default threshold (0.5): Precision={default\_precision:.0%}, Recall={default\_recall:.0%}") if rec\["recall"\] > default\_recall: extra = round((rec\["recall"\] - default\_recall) \* 100, 1) print(f" → Lowering threshold catches {extra}% more failures at cost of more alerts") # --------------------------------------------------------------------------- # Orchestration code generation (durable Flyte task with Flyte report) # --------------------------------------------------------------------------- @agent\_env.task async def plan\_experiment( experiment\_json: str, profile\_json: str, target\_column: str, time\_column: str, previous\_error: str = "", previous\_code: str = "", llm\_model: str = "gpt-4o", ) -> str: """LLM plans a single experiment: reasons about the pipeline and generates Monty code. Runs as a durable Flyte task so each experiment's planning step is traceable. Returns a JSON string: {"code": "...", "reasoning": "..."}. Args: experiment\_json: JSON string of the experiment spec (name, algorithm, hyperparams, ...). profile\_json: JSON string of the dataset profile from profile\_dataset. target\_column: Name of the target column. time\_column: Time column for temporal splitting, or empty string. previous\_error: Error message from the previous attempt (empty on first try). previous\_code: Code that failed on the previous attempt (empty on first try). llm\_model: OpenAI model identifier. Returns: str — JSON string with keys "code" and "reasoning". """ experiment = json.loads(experiment\_json) profile = json.loads(profile\_json) exp\_name = experiment.get("name", "experiment") # Strip rationale — it was written by the design LLM to explain \*why\* this # experiment was chosen. Passing it here causes plan\_experiment to parrot it # back as "reasoning" instead of independently thinking about \*how\* to build # the best pipeline. Keep only the technical spec. pipeline\_spec = { k: v for k, v in experiment.items() if k not in ("rationale",) } system = \_build\_orchestration\_system\_prompt(profile) user\_content = textwrap.dedent(f""" Design and implement the best pipeline for this experiment: Name: {exp\_name} Algorithm: {pipeline\_spec.get("algorithm")} Hyperparams: {json.dumps(pipeline\_spec.get("hyperparams", {}), indent=2)} Feature config hint: {json.dumps(pipeline\_spec.get("feature\_config", {}), indent=2)} Available sandbox inputs: - data: File — the full dataset CSV - target\_column: str = "{target\_column}" - time\_column: str = "{time\_column}" (empty string means no time ordering) - experiment\_name: str = "{exp\_name}" The feature config hint is a suggestion from the experiment designer — you can follow it, improve on it, or override it if the dataset context and your ML judgment suggest a better approach. In your ## Reasoning, explain your actual pipeline decisions: what you chose to do (or not do) and why, based on the dataset profile above. Do not restate the experiment name or why it was chosen. """).strip() messages = \[{"role": "user", "content": user\_content}\] if previous\_code and previous\_error: messages = \[\ {"role": "user", "content": user\_content},\ {"role": "assistant", "content": f"\`\`\`python\\n{previous\_code}\\n\`\`\`"},\ {"role": "user", "content": f"That code failed with this error:\\n\\n{previous\_error}\\n\\nPlease fix it."},\ \] response = await \_call\_llm(system, messages, llm\_model) reasoning = \_extract\_reasoning(response) code = \_extract\_code(response) return json.dumps({"code": code, "reasoning": reasoning}) @flyte.trace async def design\_experiments( problem\_description: str, profile\_json: str, llm\_model: str = "gpt-4o", ) -> str: """LLM designs the initial batch of experiments given problem + dataset profile. Traced so the prompt/response is visible in the Flyte UI and results are cached for deterministic replay on crash/retry. Returns raw LLM response (JSON string matching InitialDesign schema). """ design\_prompt = textwrap.dedent(f""" Problem description: {problem\_description} Dataset profile: {profile\_json} Design the first batch of experiments. """).strip() return await \_call\_llm( \_build\_initial\_design\_system\_prompt(), \[{"role": "user", "content": design\_prompt}\], llm\_model, ) @flyte.trace async def analyze\_iteration( analysis\_prompt: str, max\_iterations: int, current\_iteration: int, llm\_model: str = "gpt-4o", ) -> str: """LLM analyzes experiment results and decides whether/how to continue. Traced so the prompt/response is visible in the Flyte UI and results are cached for deterministic replay on crash/retry. Returns raw LLM response (JSON string matching IterationDecision schema). """ return await \_call\_llm( \_build\_analysis\_system\_prompt(max\_iterations, current\_iteration), \[{"role": "user", "content": analysis\_prompt}\], llm\_model, ) @flyte.trace async def plan\_followup( analysis\_prompt: str, analysis\_response: str, followup\_prompt: str, max\_iterations: int, current\_iteration: int, llm\_model: str = "gpt-4o", ) -> str: """LLM designs next experiments after targeted data explorations. Traced so the prompt/response is visible in the Flyte UI and results are cached for deterministic replay on crash/retry. Returns raw LLM response (JSON string with {"next\_experiments": \[...\]}). """ return await \_call\_llm( \_build\_analysis\_system\_prompt(max\_iterations, current\_iteration), \[\ {"role": "user", "content": analysis\_prompt},\ {"role": "assistant", "content": analysis\_response},\ {"role": "user", "content": followup\_prompt},\ \], llm\_model, ) def \_corrupt\_experiment\_for\_demo(exp\_dict: dict) -> dict: """Introduce a deliberate error into the first experiment for demo purposes. Corrupts the algorithm name so the LLM must recover from a known-bad value. The retry loop will catch this, regenerate with the error message, and fix it. """ corrupted = dict(exp\_dict) corrupted\["algorithm"\] = corrupted\["algorithm"\] + "\_INVALID" return corrupted # --------------------------------------------------------------------------- # Main agent loop # --------------------------------------------------------------------------- @dataclass class ExperimentResult: name: str algorithm: str metrics: dict confusion\_matrix: dict threshold\_analysis: list n\_samples: int code: str attempts: int reasoning: str = "" error: str = "" @dataclass class AgentResult: model\_card: str best\_experiment: str best\_metrics: dict all\_results: list\[ExperimentResult\] iterations: int total\_experiments: int async def \_run\_experiment( exp: "ExperimentConfig", exp\_dict: dict, inject\_failure: bool, data: File, target\_column: str, time\_column: str, profile: dict, llm\_model: str, max\_retries: int, ) -> "ExperimentResult | None": """Run a single experiment with retries. Returns None on total failure.""" exp\_name = exp.name profile\_json = json.dumps(profile) print(f"\\n ┌─ {exp\_name} \[{exp.algorithm}\]") if exp.rationale: for line in textwrap.wrap(exp.rationale, width=58): print(f" │ {line}") if inject\_failure: print(f" │ \[injecting failure for demo: algorithm='{exp\_dict\['algorithm'\]}'\]") code = "" error = "" result = None attempt = 0 reasoning = "" # {{docs-fragment retry\_loop}} for attempt in range(max\_retries): try: with flyte.group(exp\_name): plan\_json = await plan\_experiment.aio( experiment\_json=json.dumps(exp\_dict), profile\_json=profile\_json, target\_column=target\_column, time\_column=time\_column, previous\_error=error, previous\_code=code, llm\_model=llm\_model, ) plan = json.loads(plan\_json) code = plan\["code"\] reasoning = plan.get("reasoning", "") result = await flyte.sandbox.orchestrate\_local( code, inputs={"data": data, "target\_column": target\_column, "time\_column": time\_column, "experiment\_name": exp\_name}, tasks=TOOLS, ) error = "" break except Exception as exc: error = str(exc) short\_error = error\[:100\] + "..." if len(error) > 100 else error print(f" │ attempt {attempt + 1} failed: {short\_error}") print(f" │ → asking LLM to fix and retry...") if inject\_failure and attempt == 0: exp\_dict = exp.model\_dump() # {{/docs-fragment retry\_loop}} if result and not error: exp\_result = ExperimentResult( name=exp\_name, algorithm=exp.algorithm, metrics=result.get("metrics", {}), confusion\_matrix=result.get("confusion\_matrix", {}), threshold\_analysis=result.get("threshold\_analysis", \[\]), n\_samples=result.get("n\_samples", 0), code=code, reasoning=reasoning, attempts=attempt + 1, ) m = exp\_result.metrics attempts\_note = f" (recovered after {attempt + 1} attempts)" if attempt > 0 else "" print(f" └─ ROC-AUC={m.get('roc\_auc')}, F1={m.get('f1')}, Recall={m.get('recall')}{attempts\_note}") return exp\_result print(f" └─ FAILED after {max\_retries} attempts — skipping.") return None async def run\_agent( data: File, problem\_description: str, target\_column: str, time\_column: str = "", max\_iterations: int = 3, max\_retries\_per\_experiment: int = 3, llm\_model: str = "gpt-4o", inject\_failure: bool = False, ) -> AgentResult: """Run the MLE agent end-to-end. Args: data: CSV file containing the dataset. problem\_description: Natural language description of the ML problem. target\_column: Name of the target column to predict. time\_column: Optional column to use for time-based train/test split. max\_iterations: Maximum number of experiment iterations to run. max\_retries\_per\_experiment: Max times to retry a failed sandbox execution. llm\_model: OpenAI model to use (default: gpt-4o). inject\_failure: If True, corrupts the first experiment to demonstrate self-healing. """ print(f"\\n{'='\*60}") print(f"MLE Agent starting") print(f"Problem: {problem\_description}") print(f"Target: {target\_column}") if inject\_failure: print(f"\[demo mode: failure injection enabled\]") print(f"{'='\*60}\\n") # {{docs-fragment phase1\_profile}} # --- Phase 1: Profile the dataset (trusted tool, LLM never sees raw data) --- print(">> Phase 1: Profiling dataset...") with flyte.group("profile"): profile = await profile\_dataset(data, target\_column) # {{/docs-fragment phase1\_profile}} print(f" Shape: {profile\['shape'\]}, Classes: {profile\['target\_distribution'\]}") print(f" Imbalanced: {profile\['is\_imbalanced'\]}, Columns: {len(profile\['columns'\])}") corr = profile.get("feature\_target\_corr", {}) top\_corr = list(corr.items())\[:5\] print(f" Top correlations: {', '.join(f'{k}={v:+.3f}' for k,v in top\_corr)}") # Stream report: dataset summary await flyte.report.log.aio( f"

MLE Agent Run

" f"

Problem: {problem\_description}

" f"

Dataset: {profile\['shape'\]\[0\]:,} rows × {profile\['shape'\]\[1\]} cols  |  " f"Class balance: {profile\['class\_balance'\]}  |  Imbalanced: {profile\['is\_imbalanced'\]}

" f"

Top feature-target correlations (raw): " + ", ".join(f"{k}: {v:+.3f}" for k, v in top\_corr) + f"

", do\_flush=True, ) # --- Phase 2: LLM designs initial experiments --- print("\\n>> Phase 2: Designing initial experiments...") design\_response = await design\_experiments( problem\_description=problem\_description, profile\_json=json.dumps(profile), llm\_model=llm\_model, ) design = InitialDesign.model\_validate(\_parse\_json(design\_response)) print(f" Primary metric: {design.primary\_metric}") print(f" Strategy: {design.reasoning}") print(f" Experiments planned: {len(design.experiments)}") all\_results: list\[ExperimentResult\] = \[\] iteration\_log: list\[dict\] = \[\] # tracks per-iteration decisions + explorations for summary current\_experiments: list\[ExperimentConfig\] = design.experiments first\_experiment = True # --- Phase 3: Iterative experiment loop --- for iteration in range(max\_iterations): experiments = current\_experiments if not experiments: print(f"\\n>> No experiments to run in iteration {iteration + 1}. Stopping.") break print(f"\\n>> Phase 3.{iteration + 1}: Running {len(experiments)} experiment(s) in parallel...") # Assign names and prepare dicts before launching in parallel exp\_batch = \[\] for i, exp in enumerate(experiments): if not exp.name: exp.name = f"experiment\_{len(all\_results) + i + 1}" exp\_dict = exp.model\_dump() inject\_this = inject\_failure and first\_experiment and i == 0 if inject\_this: exp\_dict = \_corrupt\_experiment\_for\_demo(exp\_dict) first\_experiment = False exp\_batch.append((exp, exp\_dict, inject\_this)) # {{docs-fragment parallel\_execute}} batch\_results = await asyncio.gather(\*\[\ \_run\_experiment(\ exp=exp,\ exp\_dict=exp\_dict,\ inject\_failure=inject\_this,\ data=data,\ target\_column=target\_column,\ time\_column=time\_column,\ profile=profile,\ llm\_model=llm\_model,\ max\_retries=max\_retries\_per\_experiment,\ )\ for exp, exp\_dict, inject\_this in exp\_batch\ \]) # {{/docs-fragment parallel\_execute}} for exp\_result in batch\_results: if exp\_result is not None: all\_results.append(exp\_result) # Stream report: each experiment as it completes m = exp\_result.metrics html = ( f"

Iteration {iteration + 1} — {exp\_result.name}

" f"

Algorithm: {exp\_result.algorithm}  |  " f"ROC-AUC: {m.get('roc\_auc')}  |  " f"F1: {m.get('f1')}  |  " f"Recall: {m.get('recall')}  |  " f"Attempts: {exp\_result.attempts}

" ) if exp\_result.reasoning: html += f"
Reasoning
{exp\_result.reasoning}
" html += f"
Generated Code
{exp\_result.code}
" await flyte.report.log.aio(html, do\_flush=True) # --- Phase 4: Analyze results, decide whether to iterate --- if all\_results and iteration < max\_iterations - 1: print(f"\\n>> Phase 4.{iteration + 1}: Analyzing results, deciding next steps...") results\_summary = \[\ {\ "experiment\_name": r.name,\ "algorithm": r.algorithm,\ "metrics": r.metrics,\ "confusion\_matrix": r.confusion\_matrix,\ "used\_feature\_engineering": "engineer\_features" in r.code,\ "used\_rolling\_features": "rolling\_columns" in r.code,\ "used\_lag\_features": "lag\_columns" in r.code,\ }\ for r in all\_results\ \] analysis\_prompt = textwrap.dedent(f""" Problem: {problem\_description} Dataset profile: shape={profile\['shape'\]}, imbalanced={profile\['is\_imbalanced'\]} Feature-target correlations (raw): {json.dumps(profile.get('feature\_target\_corr', {}), indent=2)} Experiment results so far (iteration {iteration + 1}): {json.dumps(results\_summary, indent=2)} Should we run more experiments? If yes, request any data explorations you need, then specify what experiments to run next. """).strip() analysis\_response = await analyze\_iteration( analysis\_prompt=analysis\_prompt, max\_iterations=max\_iterations, current\_iteration=iteration, llm\_model=llm\_model, ) decision = IterationDecision.model\_validate(\_parse\_json(analysis\_response)) verdict = "continuing" if decision.should\_continue else "stopping" print(f" Decision: {verdict}") print(f" Reasoning: {decision.reasoning}") # Stream report: analysis decision await flyte.report.log.aio( f"

Analysis — Iteration {iteration + 1}

" f"

Decision: {verdict}

" f"

Reasoning: {decision.reasoning}

", do\_flush=True, ) # Track this iteration for the experiment journey summary iter\_entry = { "iteration": iteration + 1, "experiments": \[r.name for r in batch\_results if r is not None\], "best\_roc\_auc": max( (r.metrics.get("roc\_auc", 0) for r in all\_results), default=0 ), "reasoning": decision.reasoning, "explorations": \[\], } # --- Targeted exploration before next iteration --- if decision.should\_continue and decision.exploration\_requests: print(f" Running {len(decision.exploration\_requests)} exploration request(s)...") exploration\_questions = \[\] exploration\_results = \[\] for i, req in enumerate(decision.exploration\_requests): question = req.get("question", f"Exploration {i + 1}") # Strip agent-level metadata — tool only needs the analysis config tool\_config = {k: v for k, v in req.items() if k not in ("question", "analysis\_type")} print(f" Q: {question}") with flyte.group(f"explore\_{iteration + 1}\_{i + 1}"): result = await explore\_dataset(data, tool\_config) exploration\_questions.append(question) exploration\_results.append(result) iter\_entry\["explorations"\].append({"question": question}) await flyte.report.log.aio( f"

Exploration {i + 1}

" f"

Question: {question}

" f"
Results
{json.dumps(result, indent=2)}
", do\_flush=True, ) # Build follow-up that explicitly connects each question to its answer qa\_pairs = "\\n\\n".join( f'Question {i + 1}: "{q}"\\nResult:\\n{json.dumps(r, indent=2)}' for i, (q, r) in enumerate(zip(exploration\_questions, exploration\_results)) ) followup\_prompt = textwrap.dedent(f""" You requested {len(exploration\_results)} targeted exploration(s). Here is what you asked and what you learned: {qa\_pairs} Given what you learned and your earlier reasoning: "{decision.reasoning}" Now specify the next experiments. For each experiment, briefly state which exploration insight informed your choice. Respond with valid JSON: {{"next\_experiments": \[...same schema as before...\]}} """).strip() followup\_response = await plan\_followup( analysis\_prompt=analysis\_prompt, analysis\_response=analysis\_response, followup\_prompt=followup\_prompt, max\_iterations=max\_iterations, current\_iteration=iteration, llm\_model=llm\_model, ) followup = \_parse\_json(followup\_response) current\_experiments = IterationDecision.model\_validate({ "should\_continue": True, "reasoning": decision.reasoning, "next\_experiments": followup.get("next\_experiments", \[\]), }).next\_experiments print(f" Post-exploration: {len(current\_experiments)} experiment(s) planned") else: current\_experiments = decision.next\_experiments iteration\_log.append(iter\_entry) if not decision.should\_continue: break # --- Phase 5: Rank all results and generate model card --- print(f"\\n>> Phase 5: Ranking {len(all\_results)} experiment(s) and generating model card...") if not all\_results: return AgentResult( model\_card="No experiments completed successfully.", best\_experiment="", best\_metrics={}, all\_results=\[\], iterations=iteration + 1, total\_experiments=0, ) ranking\_input = \[\ {\ "experiment\_name": r.name,\ "metrics": r.metrics,\ "confusion\_matrix": r.confusion\_matrix,\ }\ for r in all\_results\ \] with flyte.group("rank"): ranking = await rank\_experiments(json.dumps(ranking\_input)) best\_name = ranking\["best\_experiment"\] best\_result = next(r for r in all\_results if r.name == best\_name) \_print\_experiment\_table(all\_results, best\_name) \_print\_threshold\_recommendation(best\_result.threshold\_analysis, best\_result.metrics) # Stream report: final rankings table rows = "".join( f"{row\['rank'\]}" f"{'' if row\['experiment\_name'\] == best\_name else ''}" f"{row\['experiment\_name'\]}" f"{'' if row\['experiment\_name'\] == best\_name else ''}" f"{row\['roc\_auc'\]}{row\['f1'\]}" f"{row\['recall'\]}{row\['precision'\]}" for row in ranking.get("ranking", \[\]) ) await flyte.report.log.aio( f"

Final Rankings

" f"" f"" f"{rows}
RankExperimentROC-AUCF1RecallPrecision
" f"

{ranking.get('summary', '')}

", do\_flush=True, ) # Stream report: experiment journey summary journey\_rows = "" for entry in iteration\_log: exps = ", ".join(entry\["experiments"\]) if entry\["experiments"\] else "—" explorations = "; ".join(e\["question"\] for e in entry\["explorations"\]) if entry\["explorations"\] else "—" short\_reasoning = (entry\["reasoning"\]\[:120\] + "…") if len(entry\["reasoning"\]) > 120 else entry\["reasoning"\] journey\_rows += ( f"" f"{entry\['iteration'\]}" f"{exps}" f"{entry\['best\_roc\_auc'\]:.4f}" f"{short\_reasoning}" f"{explorations}" f"" ) await flyte.report.log.aio( f"

Experiment Journey

" f"" f"" f"{journey\_rows}" f"
IterExperimentsBest ROC-AUCKey insightExplorations
", do\_flush=True, ) model\_card = await \_generate\_model\_card( problem\_description=problem\_description, profile=profile, all\_results=all\_results, best\_result=best\_result, ranking=ranking, iteration\_log=iteration\_log, llm\_model=llm\_model, ) print(f"\\n{'='\*60}") print(f"DONE — Best model: {best\_name}") print(f" ROC-AUC={best\_result.metrics.get('roc\_auc')}, F1={best\_result.metrics.get('f1')}") print(f"{'='\*60}\\n") return AgentResult( model\_card=model\_card, best\_experiment=best\_name, best\_metrics=best\_result.metrics, all\_results=all\_results, iterations=iteration + 1, total\_experiments=len(all\_results), ) async def \_generate\_model\_card( problem\_description: str, profile: dict, all\_results: list\[ExperimentResult\], best\_result: ExperimentResult, ranking: dict, iteration\_log: list\[dict\], llm\_model: str, ) -> str: """Generate a markdown model card summarizing the winning model.""" system = textwrap.dedent(""" You are an ML engineer writing a model card for a trained model. Write in markdown. Be concise but informative. Include: - Problem statement - Dataset summary - Experiment journey (brief per-iteration narrative: what was tried, what was learned, what changed) - Experiment summary (table of all experiments with metrics) - Winning model details (algorithm, key hyperparams, metrics, threshold analysis) - Recommendations for deployment (decision threshold, monitoring) """).strip() results\_text = "\\n".join( f"- {r.name} ({r.algorithm}): ROC-AUC={r.metrics.get('roc\_auc')}, " f"F1={r.metrics.get('f1')}, Recall={r.metrics.get('recall')}" for r in all\_results ) journey\_text = "" if iteration\_log: journey\_text = "\\n\\nIteration log:\\n" + "\\n".join( f" Iteration {e\['iteration'\]}: ran \[{', '.join(e\['experiments'\])}\], " f"best ROC-AUC so far={e\['best\_roc\_auc'\]:.4f}. " f"Key insight: {e\['reasoning'\]\[:200\]}. " + (f"Explorations: {'; '.join(x\['question'\] for x in e\['explorations'\])}" if e\['explorations'\] else "") for e in iteration\_log ) user\_content = textwrap.dedent(f""" Problem: {problem\_description} Dataset: {profile\['shape'\]\[0\]} rows × {profile\['shape'\]\[1\]} cols. Class balance: {profile\['class\_balance'\]} Imbalanced: {profile\['is\_imbalanced'\]} {journey\_text} All experiments: {results\_text} Best model: {best\_result.name} ({best\_result.algorithm}) Metrics: {json.dumps(best\_result.metrics, indent=2)} Confusion matrix: {json.dumps(best\_result.confusion\_matrix, indent=2)} Threshold analysis: {json.dumps(best\_result.threshold\_analysis, indent=2)} Ranking summary: {ranking\['summary'\]} """).strip() response = await \_call\_llm(system, \[{"role": "user", "content": user\_content}\], llm\_model) return response # --------------------------------------------------------------------------- # Durable entrypoint (runs the agent as a Flyte task in the cloud) # --------------------------------------------------------------------------- # {{docs-fragment entrypoint}} @agent\_env.task(retries=1, report=True) async def mle\_agent\_task( data: File, problem\_description: str, target\_column: str, time\_column: str = "", max\_iterations: int = 3, ) -> str: """Durable Flyte task entrypoint for the MLE agent.""" result = await run\_agent( data=data, problem\_description=problem\_description, target\_column=target\_column, time\_column=time\_column, max\_iterations=max\_iterations, ) return result.model\_card # {{/docs-fragment entrypoint}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle\_bot/mle\_bot/agent.py\* \*\*Phase 4: Analyze and iterate.\*\* After each batch completes, the LLM reviews the results and decides whether to continue. It can optionally request targeted data explorations before designing the next round. If the LLM requests explorations (e.g., "do failure cases show higher vibration readings?"), the agent runs \`explore\_dataset\` with those configurations, feeds the results back to the LLM, and lets it refine the next batch of experiments based on what it learned. The loop continues until the LLM decides to stop, the target metric threshold is reached, or the maximum number of iterations is exhausted. ## Running LLM-generated code in Flyte's sandbox This is where it gets interesting. The LLM doesn't just pick parameters from a dropdown. For each experiment, it writes actual Python code that decides how to compose the tool functions into a pipeline. Maybe it splits the data, engineers rolling window features, applies SMOTE resampling on the training split, trains an XGBoost model, and evaluates it. Or maybe it skips feature engineering entirely for a baseline. The LLM decides the structure. That code runs inside Flyte's sandbox, a restricted execution environment that enforces strict constraints: - No \`import\` statements. The only callable functions are the ones you explicitly provide. - No network access and no filesystem access. - No \`try\`/\`except\`, no \`class\` definitions, no augmented assignment (\`+=\`). - No \`with\` statements, no generators, no \`global\`/\`nonlocal\`. The sandbox sees your pre-approved tool functions as plain function calls. When the code calls \`train\_model(...)\`, the sandbox pauses execution, dispatches the call to Flyte (which runs it as a durable task on cloud compute with the resources declared on \`tool\_env\`), waits for the result, and resumes. The LLM-generated code looks like synchronous Python, but under the hood each tool call is a full Flyte task execution. Here's how the sandbox is invoked: \`\`\`python result = await flyte.sandbox.orchestrate\_local( code, inputs={ "data": data, "target\_column": target\_column, "time\_column": time\_column, "experiment\_name": exp\_name, }, tasks=TOOLS, ) \`\`\` The \`code\` parameter is a string of Python generated by the LLM. \`inputs\` provides the variables that the code can reference. \`tasks\` is the allowlist: a list of Flyte task functions that the code is permitted to call. Nothing else is available. Here's an example of what the LLM might generate for a single experiment: \`\`\`python train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") eng\_train = engineer\_features(train\_file, { "rolling\_columns": \["vibration\_mms", "temperature\_c"\], "windows": \[6, 12, 24\], "group\_column": "machine\_id", "time\_column": "timestamp" }) eng\_test = engineer\_features(test\_file, { "rolling\_columns": \["vibration\_mms", "temperature\_c"\], "windows": \[6, 12, 24\], "group\_column": "machine\_id", "time\_column": "timestamp" }) model\_file = train\_model(eng\_train, target\_column, "xgboost", { "n\_estimators": 200, "max\_depth": 8, "scale\_pos\_weight": 33 }) eval\_result = evaluate\_model(model\_file, eng\_test, target\_column) {"experiment\_name": experiment\_name, "algorithm": "xgboost", "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]} \`\`\` Each function call in that snippet dispatches a separate Flyte task. The \`split\_dataset\` calls run on the tool environment's compute (2 CPU, 4Gi memory). The \`train\_model\` call trains an actual XGBoost model. The last expression (a dict literal) is returned as the sandbox result. Sometimes the LLM generates code with bugs, like a wrong variable name or a missing argument. The agent handles this with a retry loop. If the sandbox raises an exception, the error message and the failing code are fed back to the LLM, which gets a chance to fix the issue: \`\`\` """MLE Agent — orchestrates ML experiments using Flyte's durable sandbox. The agent: 1. Profiles the dataset using a trusted tool (data never touches the LLM). 2. Asks OpenAI to design a set of experiments (algorithms, hyperparams, feature strategy). 3. For each experiment, generates Monty orchestration code and executes it via flyte.sandbox.orchestrate\_local(), which dispatches the heavy compute as durable tasks. 4. Analyzes results, iterates if needed. 5. Produces a model card summarizing the winning model. The Monty sandbox ensures the LLM-generated orchestration code is safe — it can only call the pre-approved tool functions and has no access to imports, network, or filesystem. """ import asyncio import inspect import json import os import textwrap from dataclasses import dataclass import flyte import flyte.sandbox from flyte.io import File from mle\_bot.schemas import ExperimentConfig, InitialDesign, IterationDecision from mle\_bot.environments import agent\_env from mle\_bot.tools.data import profile\_dataset, split\_dataset from mle\_bot.tools.evaluation import evaluate\_model, rank\_experiments from mle\_bot.tools.exploration import explore\_dataset from mle\_bot.tools.features import engineer\_features from mle\_bot.tools.predictions import get\_predictions from mle\_bot.tools.resampling import resample\_dataset from mle\_bot.tools.selection import select\_features from mle\_bot.tools.training import train\_model # {{docs-fragment tools}} # All tools exposed to the sandbox. # Keys must match the function names used in LLM-generated orchestration code. TOOLS = \[\ profile\_dataset, split\_dataset, explore\_dataset,\ engineer\_features, resample\_dataset, select\_features,\ train\_model, get\_predictions, evaluate\_model, rank\_experiments,\ \] TOOLS\_BY\_NAME = {t.func.\_\_name\_\_ if hasattr(t, "func") else t.\_\_name\_\_: t for t in TOOLS} # {{/docs-fragment tools}} # --------------------------------------------------------------------------- # Prompt builders # --------------------------------------------------------------------------- def \_tool\_signatures() -> str: """Build a summary of available tool signatures and docstrings for the system prompt.""" parts = \[\] for t in TOOLS: func = t.func if hasattr(t, "func") else t sig = inspect.signature(func) doc = inspect.getdoc(func) or "" # Trim docstring to first 40 lines so prompt stays manageable doc\_lines = doc.splitlines()\[:40\] short\_doc = "\\n ".join(doc\_lines) parts.append(f"async def {func.\_\_name\_\_}{sig}:\\n \\"\\"\\"{short\_doc}\\"\\"\\"\\n ...") return "\\n\\n".join(parts) # {{docs-fragment orchestration\_prompt}} def \_build\_orchestration\_system\_prompt(profile: dict) -> str: monty\_rules = flyte.sandbox.ORCHESTRATOR\_SYNTAX\_PROMPT tools\_section = \_tool\_signatures() is\_imbalanced = profile.get("is\_imbalanced", False) class\_balance = profile.get("class\_balance", {}) columns = profile.get("columns", \[\]) numeric\_cols = profile.get("numeric\_columns", \[\]) categorical\_cols = profile.get("categorical\_columns", \[\]) corr = profile.get("feature\_target\_corr", {}) corr\_str = ", ".join(f"{k}: {v:+.3f}" for k, v in list(corr.items())\[:8\]) if corr else "n/a" shape = profile.get("shape", \[0, 0\]) return f"""\\ You are an expert ML engineer. Your job is to design and write the best possible pipeline for a machine learning experiment, then generate the Python orchestration code to execute it. The code runs inside a restricted sandbox. The last expression in your code is returned as the result. All tool calls are made like regular function calls — you do NOT need to await them. ## Dataset context Shape: {shape\[0\]:,} rows × {shape\[1\]} columns Numeric features: {numeric\_cols} Categorical features (excluded from model — not supported): {categorical\_cols} Class balance: {class\_balance}, imbalanced: {is\_imbalanced} Feature-target correlations (raw, point-biserial): {corr\_str} ## General ML best practices — apply these based on the dataset context above \*\*Feature engineering\*\* (engineer\_features tool): - Sequential/time-series data (timestamp column present, rows ordered over time): rolling window features (means, stds, min/max) capture trends that point-in-time readings miss. Lag features capture recent history. Choose window sizes relative to the prediction horizon and temporal resolution of the data. - Tabular cross-sectional data: normalization helps linear models and distance-based methods. Interaction terms can help if correlations are weak individually. - Consider skipping feature engineering entirely for a baseline — it establishes whether raw features already carry enough signal. \*\*Class imbalance\*\* (when is\_imbalanced=true): - Tree ensembles: use class\_weight="balanced" or scale\_pos\_weight=n\_neg/n\_pos. - Threshold: the default 0.5 decision threshold may not be optimal — the model's probability output is what matters, threshold is tuned at deployment time. - Metric: ROC-AUC is robust to imbalance; avg\_precision is better when positives are very rare. \*\*Algorithm selection\*\*: - XGBoost / GradientBoosting: strong default for tabular data, handles missing values, built-in imbalance handling. Start here unless data is very small. - RandomForest: more robust to outliers, good for noisy data, parallelizes well. - LogisticRegression: fast linear baseline. Useful to establish whether the problem is linearly separable before adding complexity. - Prefer simpler models when n\_samples < 5,000 to avoid overfitting. \*\*Resampling\*\* (resample\_dataset tool) — data-level imbalance handling: - Use when class\_weight/scale\_pos\_weight alone isn't improving recall adequately, or when you want to test whether data-level vs algorithm-level imbalance handling works better for this dataset. - ONLY resample the TRAIN split — never test. Resampling test data gives misleading metrics. - "oversample": fast, no new information, good first try. - "smote": synthetic samples via interpolation — more diverse than random oversampling, better for high-dimensional or sparse minority classes. - "undersample": loses majority data — only useful when majority class is very large and training speed is a concern. \*\*Feature selection\*\* (select\_features tool) — prune after feature engineering: - Use after engineer\_features when the feature count is large (20+) and you suspect many features are redundant or noisy (e.g., rolling stats at many window sizes). - "mutual\_info": ranks by non-linear dependence with target — best general choice. - "variance\_threshold": drops near-constant features — cheap first pass. - "correlation\_filter": drops redundant features that are highly correlated with each other — useful when many rolling windows capture the same trend. - Can be applied before or after splitting. Apply the same selection to both train and test to ensure the model sees the same features at evaluation time. \*\*Prediction output\*\* (get\_predictions tool) — enables two advanced patterns: 1. Error analysis: train a model → get\_predictions(model, test\_file, target) → explore\_dataset(predictions\_file, {{"class\_distributions": \["feature\_x"\], "target\_column": "correct"}}) to see which examples the model gets wrong. Use this to inform feature engineering for the next iteration. 2. Stacking: train base\_model → get\_predictions(base\_model, train\_file, target) → train a meta\_model on the predictions CSV (use "predicted\_prob" as a feature alongside original features) → evaluate meta\_model on test. get\_predictions returns a CSV with columns: all originals + predicted\_prob, predicted\_class, correct. \*\*Pipeline structure\*\* — you are not required to follow a fixed sequence. Design what makes sense for this specific experiment. ## Available tools {tools\_section} ## Monty sandbox restrictions {monty\_rules} ## Critical patterns for using tool results split\_dataset returns a File — call it twice: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") engineer\_features returns a File — chain calls freely: eng = engineer\_features(train\_file, {{"rolling\_columns": \[...\], "windows": \[...\]}}) eng2 = engineer\_features(eng, {{"normalize": true, "target\_column": target\_column}}) train\_model returns a File — pass directly to evaluate\_model: model\_file = train\_model(train\_file, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, test\_file, target\_column) evaluate\_model returns a dict — subscript reads are allowed: roc = eval\_result\["metrics"\]\["roc\_auc"\] Do NOT use augmented assignment (+=), subscript assignment (d\["k"\]=v), or try/except. Build dicts as literals only. The last expression (no assignment) is the return value. ## When fixing a previous error Read the error and the failing code carefully before writing a fix. Identify the root cause — do not just change variable names or add no-ops. Trace what each tool returns, what each subsequent call expects, and where the mismatch is. Then fix the underlying logic, not just the surface symptom. ## Pipeline design — you decide the structure You are NOT required to follow a fixed sequence. Design the pipeline that makes most sense for the experiment. Examples of valid approaches: Baseline (no feature engineering): train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") model\_file = train\_model(train\_file, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, test\_file, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Two-stage feature engineering (rolling then normalize separately): train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") rolled\_train = engineer\_features(train\_file, {{"rolling\_columns": \["vibration"\], "windows": \[6, 24\]}}) rolled\_test = engineer\_features(test\_file, {{"rolling\_columns": \["vibration"\], "windows": \[6, 24\]}}) eng\_train = engineer\_features(rolled\_train, {{"normalize": true, "target\_column": target\_column}}) eng\_test = engineer\_features(rolled\_test, {{"normalize": true, "target\_column": target\_column}}) model\_file = train\_model(eng\_train, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, eng\_test, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Compare two class weightings and return the better model: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") model\_a = train\_model(train\_file, target\_column, "xgboost", {{"n\_estimators": 100, "scale\_pos\_weight": 10}}) model\_b = train\_model(train\_file, target\_column, "xgboost", {{"n\_estimators": 100, "scale\_pos\_weight": 33}}) eval\_a = evaluate\_model(model\_a, test\_file, target\_column) eval\_b = evaluate\_model(model\_b, test\_file, target\_column) best\_eval = eval\_a if eval\_a\["metrics"\]\["roc\_auc"\] > eval\_b\["metrics"\]\["roc\_auc"\] else eval\_b {{"experiment\_name": experiment\_name, "algorithm": "xgboost", "metrics": best\_eval\["metrics"\], "confusion\_matrix": best\_eval\["confusion\_matrix"\], "threshold\_analysis": best\_eval\["threshold\_analysis"\], "n\_samples": best\_eval\["n\_samples"\]}} SMOTE oversampling before training: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") eng\_train = engineer\_features(train\_file, {{"rolling\_columns": \["vibration\_mms"\], "windows": \[6, 12\]}}) eng\_test = engineer\_features(test\_file, {{"rolling\_columns": \["vibration\_mms"\], "windows": \[6, 12\]}}) resampled\_train = resample\_dataset(eng\_train, target\_column, {{"strategy": "smote", "target\_ratio": 0.2}}) model\_file = train\_model(resampled\_train, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, eng\_test, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Feature engineering followed by feature selection: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") eng\_train = engineer\_features(train\_file, {{"rolling\_columns": \["vibration\_mms", "temperature\_c"\], "windows": \[6, 12, 24\]}}) eng\_test = engineer\_features(test\_file, {{"rolling\_columns": \["vibration\_mms", "temperature\_c"\], "windows": \[6, 12, 24\]}}) sel\_train = select\_features(eng\_train, target\_column, {{"method": "mutual\_info", "k": 15}}) sel\_test = select\_features(eng\_test, target\_column, {{"method": "mutual\_info", "k": 15}}) model\_file = train\_model(sel\_train, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, sel\_test, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Error analysis — explore what the model gets wrong, then return that as insight: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") model\_file = train\_model(train\_file, target\_column, algorithm, hyperparams) pred\_file = get\_predictions(model\_file, test\_file, target\_column) error\_analysis = explore\_dataset(pred\_file, {{"target\_column": "correct", "class\_distributions": \["vibration\_mms", "temperature\_c"\]}}) eval\_result = evaluate\_model(model\_file, test\_file, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\], "error\_analysis": error\_analysis}} The last expression MUST be a dict with at minimum these keys: experiment\_name, algorithm, metrics, confusion\_matrix, threshold\_analysis, n\_samples Additional keys (e.g. error\_analysis) are allowed and will appear in the report. ## Response format Respond in exactly this format: ## Reasoning \[Your thinking: what pipeline makes sense for this experiment and why. Consider whether\ feature engineering helps, whether class imbalance needs special treatment, whether\ chaining multiple steps adds value, etc.\] ## Code \`\`\`python \[your orchestration code\] \`\`\` """ # {{/docs-fragment orchestration\_prompt}} def \_build\_analysis\_system\_prompt(max\_iterations: int, current\_iteration: int) -> str: remaining = max\_iterations - current\_iteration - 1 return f"""\\ You are an expert ML engineer analyzing experiment results to guide the next iteration of model development. You must respond with valid JSON only — no markdown, no explanation outside the JSON. Response format: {{ "should\_continue": true | false, "reasoning": "What you observed, what it tells you, and what to try next", "exploration\_requests": \[\ {{\ "question": "The specific hypothesis you are testing, e.g. 'Do failure cases show meaningfully higher vibration than healthy cases?'",\ "analysis\_type": "class\_distributions",\ "target\_column": "failure\_24h",\ "class\_distributions": \["vibration\_mms", "temperature\_c"\]\ }}\ \], "next\_experiments": \[\ {{\ "name": "descriptive experiment name",\ "algorithm": "xgboost" | "random\_forest" | "gradient\_boosting" | "logistic\_regression",\ "hyperparams": {{ ... algorithm-specific hyperparams ... }},\ "feature\_config": {{\ "group\_column": "...",\ "time\_column": "...",\ "rolling\_columns": \[...\],\ "windows": \[...\],\ "lag\_columns": \[...\],\ "lags": \[...\],\ "normalize": true | false,\ "drop\_columns": \[...\],\ "fillna\_method": "forward"\ }},\ "rationale": "Why this experiment is worth trying"\ }}\ \] }} exploration\_requests rules: - Max 2 requests per iteration. - Each request targets EXACTLY ONE analysis\_type. Do not mix multiple types in one request. - Supported analysis\_type values and their required config fields: "class\_distributions" → requires: target\_column, class\_distributions (list of columns) "correlation\_matrix" → requires: correlation\_matrix: true "temporal\_trend" → requires: temporal\_trend: {{time\_column, target\_column, freq}} "group\_stats" → requires: group\_stats: {{group\_column, target\_column}} "outlier\_summary" → requires: outlier\_summary (list of columns) "feature\_target\_corr\_by\_group" → requires: feature\_target\_corr\_by\_group: {{group\_column, target\_column, feature\_columns}} - The "question" field is required. It must be a specific testable hypothesis, not a general request. Bad: "explore the data". Good: "Is vibration\_mms higher for failures?" - Set exploration\_requests to \[\] if the current results already tell you enough to design the next experiments. Only explore when you have a concrete unanswered question. When deciding next experiments, reason about WHAT WAS TRIED vs what hasn't been explored. Each result includes used\_feature\_engineering, used\_rolling\_features, used\_lag\_features. Think systematically: if no feature engineering was tried yet, does the data profile suggest it would help (weak raw correlations, temporal/sequential structure)? If feature engineering helped, can it be improved? Avoid experiments identical to ones tried. Iteration context: this is iteration {current\_iteration + 1} of {max\_iterations} requested. Remaining iterations allowed: {remaining}. Set should\_continue=false only if: - Best ROC-AUC >= 0.97, OR - No remaining iterations (remaining == 0), OR - Results have genuinely plateaued (< 0.005 ROC-AUC improvement over last iteration AND you have already tried the most promising directions) Otherwise keep exploring — the user asked for {max\_iterations} iterations for a reason. """ def \_build\_initial\_design\_system\_prompt() -> str: return """\\ You are an expert ML engineer. Given a dataset profile and a problem description, design the first batch of experiments to run. You must respond with valid JSON only — no markdown, no explanation outside the JSON. Response format: { "problem\_type": "binary\_classification", "primary\_metric": "roc\_auc" | "f1" | "recall", "reasoning": "Brief description of your strategy", "experiments": \[\ {\ "name": "descriptive experiment name",\ "algorithm": "xgboost" | "random\_forest" | "gradient\_boosting" | "logistic\_regression",\ "hyperparams": { ... algorithm-specific hyperparams ... },\ "feature\_config": {\ "group\_column": "",\ "time\_column": "",\ "rolling\_columns": \[\],\ "windows": \[\],\ "lag\_columns": \[\],\ "lags": \[\],\ "normalize": false,\ "drop\_columns": \[\],\ "fillna\_method": "forward"\ },\ "rationale": "Why this experiment makes sense given the data profile"\ }\ \] } Design 2-3 experiments for the first batch. Good first batches typically include: - A fast baseline to establish a floor (e.g. logistic\_regression with default settings) - Your best initial hypothesis given the data profile - Optionally one variant that tests a specific idea suggested by the profile Use the dataset profile to guide your choices: - feature\_target\_corr: weak raw correlations suggest feature engineering may help - categorical\_columns: note these are excluded from the model automatically - is\_imbalanced: handle with class\_weight or scale\_pos\_weight - Shape and column types should inform algorithm complexity (simpler models for small datasets) - A time column suggests sequential structure; lag/rolling features may capture temporal patterns The feature\_config in each experiment describes what engineer\_features should apply. Leave all fields empty/false if no feature engineering is needed for that experiment. The orchestration code generator will decide the exact pipeline — your job here is to specify what the experiment is trying to learn, not to prescribe every implementation detail. """ # --------------------------------------------------------------------------- # LLM client # --------------------------------------------------------------------------- def \_openai\_client(): from openai import OpenAI return OpenAI(api\_key=os.environ\["OPENAI\_API\_KEY"\]) async def \_call\_llm(system: str, messages: list\[dict\], model: str = "gpt-4o") -> str: """Call OpenAI and return the response text.""" client = \_openai\_client() response = await asyncio.to\_thread( client.chat.completions.create, model=model, messages=\[{"role": "system", "content": system}, \*messages\], temperature=0.2, ) return response.choices\[0\].message.content def \_extract\_code(text: str) -> str: """Pull Python code out of a markdown code block.""" if "\`\`\`python" in text: start = text.index("\`\`\`python") + len("\`\`\`python") end = text.index("\`\`\`", start) return text\[start:end\].strip() if "\`\`\`" in text: start = text.index("\`\`\`") + 3 end = text.index("\`\`\`", start) return text\[start:end\].strip() return text.strip() def \_extract\_reasoning(text: str) -> str: """Extract the ## Reasoning section from LLM response.""" if "## Reasoning" in text: start = text.index("## Reasoning") + len("## Reasoning") if "## Code" in text: end = text.index("## Code") return text\[start:end\].strip() return text\[start:\].strip() return "" def \_parse\_json(text: str) -> dict: """Extract and parse JSON from LLM response.""" text = text.strip() if "\`\`\`json" in text: start = text.index("\`\`\`json") + 7 end = text.index("\`\`\`", start) text = text\[start:end\].strip() elif "\`\`\`" in text: start = text.index("\`\`\`") + 3 end = text.index("\`\`\`", start) text = text\[start:end\].strip() return json.loads(text) # --------------------------------------------------------------------------- # Display helpers # --------------------------------------------------------------------------- def \_recommend\_threshold(threshold\_analysis: list, min\_precision: float = 0.70) -> dict | None: """Find the threshold that maximises recall subject to precision >= min\_precision.""" candidates = \[t for t in threshold\_analysis if t\["precision"\] >= min\_precision\] if not candidates: return None return max(candidates, key=lambda t: t\["recall"\]) def \_print\_experiment\_table(results: list\["ExperimentResult"\], best\_name: str) -> None: """Print a ranked comparison table of all experiments.""" sorted\_results = sorted(results, key=lambda r: r.metrics.get("roc\_auc", 0), reverse=True) print("\\n" + "─" \* 78) print(f" {'Rank':<5} {'Experiment':<32} {'ROC-AUC':<9} {'F1':<7} {'Recall':<8} {'Note'}") print("─" \* 78) for rank, r in enumerate(sorted\_results, 1): note = "◀ winner" if r.name == best\_name else "" roc = r.metrics.get("roc\_auc", 0) f1 = r.metrics.get("f1", 0) recall = r.metrics.get("recall", 0) print(f" {rank:<5} {r.name:<32} {roc:<9.4f} {f1:<7.4f} {recall:<8.4f} {note}") print("─" \* 78) def \_print\_threshold\_recommendation(threshold\_analysis: list, default\_metrics: dict) -> None: """Print the operational threshold recommendation.""" rec = \_recommend\_threshold(threshold\_analysis) if not rec: return default\_recall = default\_metrics.get("recall", 0) default\_precision = default\_metrics.get("precision", 0) missed\_pct = round((1 - rec\["recall"\]) \* 100, 1) false\_alarm\_pct = round((1 - rec\["precision"\]) \* 100, 1) print(f"\\n Recommended decision threshold: {rec\['threshold'\]}") print(f" ├─ Precision : {rec\['precision'\]:.0%} ({false\_alarm\_pct}% of alerts are false alarms)") print(f" ├─ Recall : {rec\['recall'\]:.0%} (catches {rec\['recall'\]\*100:.0f}% of actual failures)") print(f" └─ F1 : {rec\['f1'\]:.4f}") print(f" Default threshold (0.5): Precision={default\_precision:.0%}, Recall={default\_recall:.0%}") if rec\["recall"\] > default\_recall: extra = round((rec\["recall"\] - default\_recall) \* 100, 1) print(f" → Lowering threshold catches {extra}% more failures at cost of more alerts") # --------------------------------------------------------------------------- # Orchestration code generation (durable Flyte task with Flyte report) # --------------------------------------------------------------------------- @agent\_env.task async def plan\_experiment( experiment\_json: str, profile\_json: str, target\_column: str, time\_column: str, previous\_error: str = "", previous\_code: str = "", llm\_model: str = "gpt-4o", ) -> str: """LLM plans a single experiment: reasons about the pipeline and generates Monty code. Runs as a durable Flyte task so each experiment's planning step is traceable. Returns a JSON string: {"code": "...", "reasoning": "..."}. Args: experiment\_json: JSON string of the experiment spec (name, algorithm, hyperparams, ...). profile\_json: JSON string of the dataset profile from profile\_dataset. target\_column: Name of the target column. time\_column: Time column for temporal splitting, or empty string. previous\_error: Error message from the previous attempt (empty on first try). previous\_code: Code that failed on the previous attempt (empty on first try). llm\_model: OpenAI model identifier. Returns: str — JSON string with keys "code" and "reasoning". """ experiment = json.loads(experiment\_json) profile = json.loads(profile\_json) exp\_name = experiment.get("name", "experiment") # Strip rationale — it was written by the design LLM to explain \*why\* this # experiment was chosen. Passing it here causes plan\_experiment to parrot it # back as "reasoning" instead of independently thinking about \*how\* to build # the best pipeline. Keep only the technical spec. pipeline\_spec = { k: v for k, v in experiment.items() if k not in ("rationale",) } system = \_build\_orchestration\_system\_prompt(profile) user\_content = textwrap.dedent(f""" Design and implement the best pipeline for this experiment: Name: {exp\_name} Algorithm: {pipeline\_spec.get("algorithm")} Hyperparams: {json.dumps(pipeline\_spec.get("hyperparams", {}), indent=2)} Feature config hint: {json.dumps(pipeline\_spec.get("feature\_config", {}), indent=2)} Available sandbox inputs: - data: File — the full dataset CSV - target\_column: str = "{target\_column}" - time\_column: str = "{time\_column}" (empty string means no time ordering) - experiment\_name: str = "{exp\_name}" The feature config hint is a suggestion from the experiment designer — you can follow it, improve on it, or override it if the dataset context and your ML judgment suggest a better approach. In your ## Reasoning, explain your actual pipeline decisions: what you chose to do (or not do) and why, based on the dataset profile above. Do not restate the experiment name or why it was chosen. """).strip() messages = \[{"role": "user", "content": user\_content}\] if previous\_code and previous\_error: messages = \[\ {"role": "user", "content": user\_content},\ {"role": "assistant", "content": f"\`\`\`python\\n{previous\_code}\\n\`\`\`"},\ {"role": "user", "content": f"That code failed with this error:\\n\\n{previous\_error}\\n\\nPlease fix it."},\ \] response = await \_call\_llm(system, messages, llm\_model) reasoning = \_extract\_reasoning(response) code = \_extract\_code(response) return json.dumps({"code": code, "reasoning": reasoning}) @flyte.trace async def design\_experiments( problem\_description: str, profile\_json: str, llm\_model: str = "gpt-4o", ) -> str: """LLM designs the initial batch of experiments given problem + dataset profile. Traced so the prompt/response is visible in the Flyte UI and results are cached for deterministic replay on crash/retry. Returns raw LLM response (JSON string matching InitialDesign schema). """ design\_prompt = textwrap.dedent(f""" Problem description: {problem\_description} Dataset profile: {profile\_json} Design the first batch of experiments. """).strip() return await \_call\_llm( \_build\_initial\_design\_system\_prompt(), \[{"role": "user", "content": design\_prompt}\], llm\_model, ) @flyte.trace async def analyze\_iteration( analysis\_prompt: str, max\_iterations: int, current\_iteration: int, llm\_model: str = "gpt-4o", ) -> str: """LLM analyzes experiment results and decides whether/how to continue. Traced so the prompt/response is visible in the Flyte UI and results are cached for deterministic replay on crash/retry. Returns raw LLM response (JSON string matching IterationDecision schema). """ return await \_call\_llm( \_build\_analysis\_system\_prompt(max\_iterations, current\_iteration), \[{"role": "user", "content": analysis\_prompt}\], llm\_model, ) @flyte.trace async def plan\_followup( analysis\_prompt: str, analysis\_response: str, followup\_prompt: str, max\_iterations: int, current\_iteration: int, llm\_model: str = "gpt-4o", ) -> str: """LLM designs next experiments after targeted data explorations. Traced so the prompt/response is visible in the Flyte UI and results are cached for deterministic replay on crash/retry. Returns raw LLM response (JSON string with {"next\_experiments": \[...\]}). """ return await \_call\_llm( \_build\_analysis\_system\_prompt(max\_iterations, current\_iteration), \[\ {"role": "user", "content": analysis\_prompt},\ {"role": "assistant", "content": analysis\_response},\ {"role": "user", "content": followup\_prompt},\ \], llm\_model, ) def \_corrupt\_experiment\_for\_demo(exp\_dict: dict) -> dict: """Introduce a deliberate error into the first experiment for demo purposes. Corrupts the algorithm name so the LLM must recover from a known-bad value. The retry loop will catch this, regenerate with the error message, and fix it. """ corrupted = dict(exp\_dict) corrupted\["algorithm"\] = corrupted\["algorithm"\] + "\_INVALID" return corrupted # --------------------------------------------------------------------------- # Main agent loop # --------------------------------------------------------------------------- @dataclass class ExperimentResult: name: str algorithm: str metrics: dict confusion\_matrix: dict threshold\_analysis: list n\_samples: int code: str attempts: int reasoning: str = "" error: str = "" @dataclass class AgentResult: model\_card: str best\_experiment: str best\_metrics: dict all\_results: list\[ExperimentResult\] iterations: int total\_experiments: int async def \_run\_experiment( exp: "ExperimentConfig", exp\_dict: dict, inject\_failure: bool, data: File, target\_column: str, time\_column: str, profile: dict, llm\_model: str, max\_retries: int, ) -> "ExperimentResult | None": """Run a single experiment with retries. Returns None on total failure.""" exp\_name = exp.name profile\_json = json.dumps(profile) print(f"\\n ┌─ {exp\_name} \[{exp.algorithm}\]") if exp.rationale: for line in textwrap.wrap(exp.rationale, width=58): print(f" │ {line}") if inject\_failure: print(f" │ \[injecting failure for demo: algorithm='{exp\_dict\['algorithm'\]}'\]") code = "" error = "" result = None attempt = 0 reasoning = "" # {{docs-fragment retry\_loop}} for attempt in range(max\_retries): try: with flyte.group(exp\_name): plan\_json = await plan\_experiment.aio( experiment\_json=json.dumps(exp\_dict), profile\_json=profile\_json, target\_column=target\_column, time\_column=time\_column, previous\_error=error, previous\_code=code, llm\_model=llm\_model, ) plan = json.loads(plan\_json) code = plan\["code"\] reasoning = plan.get("reasoning", "") result = await flyte.sandbox.orchestrate\_local( code, inputs={"data": data, "target\_column": target\_column, "time\_column": time\_column, "experiment\_name": exp\_name}, tasks=TOOLS, ) error = "" break except Exception as exc: error = str(exc) short\_error = error\[:100\] + "..." if len(error) > 100 else error print(f" │ attempt {attempt + 1} failed: {short\_error}") print(f" │ → asking LLM to fix and retry...") if inject\_failure and attempt == 0: exp\_dict = exp.model\_dump() # {{/docs-fragment retry\_loop}} if result and not error: exp\_result = ExperimentResult( name=exp\_name, algorithm=exp.algorithm, metrics=result.get("metrics", {}), confusion\_matrix=result.get("confusion\_matrix", {}), threshold\_analysis=result.get("threshold\_analysis", \[\]), n\_samples=result.get("n\_samples", 0), code=code, reasoning=reasoning, attempts=attempt + 1, ) m = exp\_result.metrics attempts\_note = f" (recovered after {attempt + 1} attempts)" if attempt > 0 else "" print(f" └─ ROC-AUC={m.get('roc\_auc')}, F1={m.get('f1')}, Recall={m.get('recall')}{attempts\_note}") return exp\_result print(f" └─ FAILED after {max\_retries} attempts — skipping.") return None async def run\_agent( data: File, problem\_description: str, target\_column: str, time\_column: str = "", max\_iterations: int = 3, max\_retries\_per\_experiment: int = 3, llm\_model: str = "gpt-4o", inject\_failure: bool = False, ) -> AgentResult: """Run the MLE agent end-to-end. Args: data: CSV file containing the dataset. problem\_description: Natural language description of the ML problem. target\_column: Name of the target column to predict. time\_column: Optional column to use for time-based train/test split. max\_iterations: Maximum number of experiment iterations to run. max\_retries\_per\_experiment: Max times to retry a failed sandbox execution. llm\_model: OpenAI model to use (default: gpt-4o). inject\_failure: If True, corrupts the first experiment to demonstrate self-healing. """ print(f"\\n{'='\*60}") print(f"MLE Agent starting") print(f"Problem: {problem\_description}") print(f"Target: {target\_column}") if inject\_failure: print(f"\[demo mode: failure injection enabled\]") print(f"{'='\*60}\\n") # {{docs-fragment phase1\_profile}} # --- Phase 1: Profile the dataset (trusted tool, LLM never sees raw data) --- print(">> Phase 1: Profiling dataset...") with flyte.group("profile"): profile = await profile\_dataset(data, target\_column) # {{/docs-fragment phase1\_profile}} print(f" Shape: {profile\['shape'\]}, Classes: {profile\['target\_distribution'\]}") print(f" Imbalanced: {profile\['is\_imbalanced'\]}, Columns: {len(profile\['columns'\])}") corr = profile.get("feature\_target\_corr", {}) top\_corr = list(corr.items())\[:5\] print(f" Top correlations: {', '.join(f'{k}={v:+.3f}' for k,v in top\_corr)}") # Stream report: dataset summary await flyte.report.log.aio( f"

MLE Agent Run

" f"

Problem: {problem\_description}

" f"

Dataset: {profile\['shape'\]\[0\]:,} rows × {profile\['shape'\]\[1\]} cols  |  " f"Class balance: {profile\['class\_balance'\]}  |  Imbalanced: {profile\['is\_imbalanced'\]}

" f"

Top feature-target correlations (raw): " + ", ".join(f"{k}: {v:+.3f}" for k, v in top\_corr) + f"

", do\_flush=True, ) # --- Phase 2: LLM designs initial experiments --- print("\\n>> Phase 2: Designing initial experiments...") design\_response = await design\_experiments( problem\_description=problem\_description, profile\_json=json.dumps(profile), llm\_model=llm\_model, ) design = InitialDesign.model\_validate(\_parse\_json(design\_response)) print(f" Primary metric: {design.primary\_metric}") print(f" Strategy: {design.reasoning}") print(f" Experiments planned: {len(design.experiments)}") all\_results: list\[ExperimentResult\] = \[\] iteration\_log: list\[dict\] = \[\] # tracks per-iteration decisions + explorations for summary current\_experiments: list\[ExperimentConfig\] = design.experiments first\_experiment = True # --- Phase 3: Iterative experiment loop --- for iteration in range(max\_iterations): experiments = current\_experiments if not experiments: print(f"\\n>> No experiments to run in iteration {iteration + 1}. Stopping.") break print(f"\\n>> Phase 3.{iteration + 1}: Running {len(experiments)} experiment(s) in parallel...") # Assign names and prepare dicts before launching in parallel exp\_batch = \[\] for i, exp in enumerate(experiments): if not exp.name: exp.name = f"experiment\_{len(all\_results) + i + 1}" exp\_dict = exp.model\_dump() inject\_this = inject\_failure and first\_experiment and i == 0 if inject\_this: exp\_dict = \_corrupt\_experiment\_for\_demo(exp\_dict) first\_experiment = False exp\_batch.append((exp, exp\_dict, inject\_this)) # {{docs-fragment parallel\_execute}} batch\_results = await asyncio.gather(\*\[\ \_run\_experiment(\ exp=exp,\ exp\_dict=exp\_dict,\ inject\_failure=inject\_this,\ data=data,\ target\_column=target\_column,\ time\_column=time\_column,\ profile=profile,\ llm\_model=llm\_model,\ max\_retries=max\_retries\_per\_experiment,\ )\ for exp, exp\_dict, inject\_this in exp\_batch\ \]) # {{/docs-fragment parallel\_execute}} for exp\_result in batch\_results: if exp\_result is not None: all\_results.append(exp\_result) # Stream report: each experiment as it completes m = exp\_result.metrics html = ( f"

Iteration {iteration + 1} — {exp\_result.name}

" f"

Algorithm: {exp\_result.algorithm}  |  " f"ROC-AUC: {m.get('roc\_auc')}  |  " f"F1: {m.get('f1')}  |  " f"Recall: {m.get('recall')}  |  " f"Attempts: {exp\_result.attempts}

" ) if exp\_result.reasoning: html += f"
Reasoning
{exp\_result.reasoning}
" html += f"
Generated Code
{exp\_result.code}
" await flyte.report.log.aio(html, do\_flush=True) # --- Phase 4: Analyze results, decide whether to iterate --- if all\_results and iteration < max\_iterations - 1: print(f"\\n>> Phase 4.{iteration + 1}: Analyzing results, deciding next steps...") results\_summary = \[\ {\ "experiment\_name": r.name,\ "algorithm": r.algorithm,\ "metrics": r.metrics,\ "confusion\_matrix": r.confusion\_matrix,\ "used\_feature\_engineering": "engineer\_features" in r.code,\ "used\_rolling\_features": "rolling\_columns" in r.code,\ "used\_lag\_features": "lag\_columns" in r.code,\ }\ for r in all\_results\ \] analysis\_prompt = textwrap.dedent(f""" Problem: {problem\_description} Dataset profile: shape={profile\['shape'\]}, imbalanced={profile\['is\_imbalanced'\]} Feature-target correlations (raw): {json.dumps(profile.get('feature\_target\_corr', {}), indent=2)} Experiment results so far (iteration {iteration + 1}): {json.dumps(results\_summary, indent=2)} Should we run more experiments? If yes, request any data explorations you need, then specify what experiments to run next. """).strip() analysis\_response = await analyze\_iteration( analysis\_prompt=analysis\_prompt, max\_iterations=max\_iterations, current\_iteration=iteration, llm\_model=llm\_model, ) decision = IterationDecision.model\_validate(\_parse\_json(analysis\_response)) verdict = "continuing" if decision.should\_continue else "stopping" print(f" Decision: {verdict}") print(f" Reasoning: {decision.reasoning}") # Stream report: analysis decision await flyte.report.log.aio( f"

Analysis — Iteration {iteration + 1}

" f"

Decision: {verdict}

" f"

Reasoning: {decision.reasoning}

", do\_flush=True, ) # Track this iteration for the experiment journey summary iter\_entry = { "iteration": iteration + 1, "experiments": \[r.name for r in batch\_results if r is not None\], "best\_roc\_auc": max( (r.metrics.get("roc\_auc", 0) for r in all\_results), default=0 ), "reasoning": decision.reasoning, "explorations": \[\], } # --- Targeted exploration before next iteration --- if decision.should\_continue and decision.exploration\_requests: print(f" Running {len(decision.exploration\_requests)} exploration request(s)...") exploration\_questions = \[\] exploration\_results = \[\] for i, req in enumerate(decision.exploration\_requests): question = req.get("question", f"Exploration {i + 1}") # Strip agent-level metadata — tool only needs the analysis config tool\_config = {k: v for k, v in req.items() if k not in ("question", "analysis\_type")} print(f" Q: {question}") with flyte.group(f"explore\_{iteration + 1}\_{i + 1}"): result = await explore\_dataset(data, tool\_config) exploration\_questions.append(question) exploration\_results.append(result) iter\_entry\["explorations"\].append({"question": question}) await flyte.report.log.aio( f"

Exploration {i + 1}

" f"

Question: {question}

" f"
Results
{json.dumps(result, indent=2)}
", do\_flush=True, ) # Build follow-up that explicitly connects each question to its answer qa\_pairs = "\\n\\n".join( f'Question {i + 1}: "{q}"\\nResult:\\n{json.dumps(r, indent=2)}' for i, (q, r) in enumerate(zip(exploration\_questions, exploration\_results)) ) followup\_prompt = textwrap.dedent(f""" You requested {len(exploration\_results)} targeted exploration(s). Here is what you asked and what you learned: {qa\_pairs} Given what you learned and your earlier reasoning: "{decision.reasoning}" Now specify the next experiments. For each experiment, briefly state which exploration insight informed your choice. Respond with valid JSON: {{"next\_experiments": \[...same schema as before...\]}} """).strip() followup\_response = await plan\_followup( analysis\_prompt=analysis\_prompt, analysis\_response=analysis\_response, followup\_prompt=followup\_prompt, max\_iterations=max\_iterations, current\_iteration=iteration, llm\_model=llm\_model, ) followup = \_parse\_json(followup\_response) current\_experiments = IterationDecision.model\_validate({ "should\_continue": True, "reasoning": decision.reasoning, "next\_experiments": followup.get("next\_experiments", \[\]), }).next\_experiments print(f" Post-exploration: {len(current\_experiments)} experiment(s) planned") else: current\_experiments = decision.next\_experiments iteration\_log.append(iter\_entry) if not decision.should\_continue: break # --- Phase 5: Rank all results and generate model card --- print(f"\\n>> Phase 5: Ranking {len(all\_results)} experiment(s) and generating model card...") if not all\_results: return AgentResult( model\_card="No experiments completed successfully.", best\_experiment="", best\_metrics={}, all\_results=\[\], iterations=iteration + 1, total\_experiments=0, ) ranking\_input = \[\ {\ "experiment\_name": r.name,\ "metrics": r.metrics,\ "confusion\_matrix": r.confusion\_matrix,\ }\ for r in all\_results\ \] with flyte.group("rank"): ranking = await rank\_experiments(json.dumps(ranking\_input)) best\_name = ranking\["best\_experiment"\] best\_result = next(r for r in all\_results if r.name == best\_name) \_print\_experiment\_table(all\_results, best\_name) \_print\_threshold\_recommendation(best\_result.threshold\_analysis, best\_result.metrics) # Stream report: final rankings table rows = "".join( f"{row\['rank'\]}" f"{'' if row\['experiment\_name'\] == best\_name else ''}" f"{row\['experiment\_name'\]}" f"{'' if row\['experiment\_name'\] == best\_name else ''}" f"{row\['roc\_auc'\]}{row\['f1'\]}" f"{row\['recall'\]}{row\['precision'\]}" for row in ranking.get("ranking", \[\]) ) await flyte.report.log.aio( f"

Final Rankings

" f"" f"" f"{rows}
RankExperimentROC-AUCF1RecallPrecision
" f"

{ranking.get('summary', '')}

", do\_flush=True, ) # Stream report: experiment journey summary journey\_rows = "" for entry in iteration\_log: exps = ", ".join(entry\["experiments"\]) if entry\["experiments"\] else "—" explorations = "; ".join(e\["question"\] for e in entry\["explorations"\]) if entry\["explorations"\] else "—" short\_reasoning = (entry\["reasoning"\]\[:120\] + "…") if len(entry\["reasoning"\]) > 120 else entry\["reasoning"\] journey\_rows += ( f"" f"{entry\['iteration'\]}" f"{exps}" f"{entry\['best\_roc\_auc'\]:.4f}" f"{short\_reasoning}" f"{explorations}" f"" ) await flyte.report.log.aio( f"

Experiment Journey

" f"" f"" f"{journey\_rows}" f"
IterExperimentsBest ROC-AUCKey insightExplorations
", do\_flush=True, ) model\_card = await \_generate\_model\_card( problem\_description=problem\_description, profile=profile, all\_results=all\_results, best\_result=best\_result, ranking=ranking, iteration\_log=iteration\_log, llm\_model=llm\_model, ) print(f"\\n{'='\*60}") print(f"DONE — Best model: {best\_name}") print(f" ROC-AUC={best\_result.metrics.get('roc\_auc')}, F1={best\_result.metrics.get('f1')}") print(f"{'='\*60}\\n") return AgentResult( model\_card=model\_card, best\_experiment=best\_name, best\_metrics=best\_result.metrics, all\_results=all\_results, iterations=iteration + 1, total\_experiments=len(all\_results), ) async def \_generate\_model\_card( problem\_description: str, profile: dict, all\_results: list\[ExperimentResult\], best\_result: ExperimentResult, ranking: dict, iteration\_log: list\[dict\], llm\_model: str, ) -> str: """Generate a markdown model card summarizing the winning model.""" system = textwrap.dedent(""" You are an ML engineer writing a model card for a trained model. Write in markdown. Be concise but informative. Include: - Problem statement - Dataset summary - Experiment journey (brief per-iteration narrative: what was tried, what was learned, what changed) - Experiment summary (table of all experiments with metrics) - Winning model details (algorithm, key hyperparams, metrics, threshold analysis) - Recommendations for deployment (decision threshold, monitoring) """).strip() results\_text = "\\n".join( f"- {r.name} ({r.algorithm}): ROC-AUC={r.metrics.get('roc\_auc')}, " f"F1={r.metrics.get('f1')}, Recall={r.metrics.get('recall')}" for r in all\_results ) journey\_text = "" if iteration\_log: journey\_text = "\\n\\nIteration log:\\n" + "\\n".join( f" Iteration {e\['iteration'\]}: ran \[{', '.join(e\['experiments'\])}\], " f"best ROC-AUC so far={e\['best\_roc\_auc'\]:.4f}. " f"Key insight: {e\['reasoning'\]\[:200\]}. " + (f"Explorations: {'; '.join(x\['question'\] for x in e\['explorations'\])}" if e\['explorations'\] else "") for e in iteration\_log ) user\_content = textwrap.dedent(f""" Problem: {problem\_description} Dataset: {profile\['shape'\]\[0\]} rows × {profile\['shape'\]\[1\]} cols. Class balance: {profile\['class\_balance'\]} Imbalanced: {profile\['is\_imbalanced'\]} {journey\_text} All experiments: {results\_text} Best model: {best\_result.name} ({best\_result.algorithm}) Metrics: {json.dumps(best\_result.metrics, indent=2)} Confusion matrix: {json.dumps(best\_result.confusion\_matrix, indent=2)} Threshold analysis: {json.dumps(best\_result.threshold\_analysis, indent=2)} Ranking summary: {ranking\['summary'\]} """).strip() response = await \_call\_llm(system, \[{"role": "user", "content": user\_content}\], llm\_model) return response # --------------------------------------------------------------------------- # Durable entrypoint (runs the agent as a Flyte task in the cloud) # --------------------------------------------------------------------------- # {{docs-fragment entrypoint}} @agent\_env.task(retries=1, report=True) async def mle\_agent\_task( data: File, problem\_description: str, target\_column: str, time\_column: str = "", max\_iterations: int = 3, ) -> str: """Durable Flyte task entrypoint for the MLE agent.""" result = await run\_agent( data=data, problem\_description=problem\_description, target\_column=target\_column, time\_column=time\_column, max\_iterations=max\_iterations, ) return result.model\_card # {{/docs-fragment entrypoint}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle\_bot/mle\_bot/agent.py\* On the first attempt, \`previous\_error\` and \`previous\_code\` are empty. On subsequent attempts, the LLM sees exactly what went wrong and can fix it. In practice, most experiments succeed on the first try, with occasional recoveries on the second. ## Streaming results to a live report While the agent runs, it streams results to the Flyte UI in real time using \`flyte.report.log.aio()\`. You don't have to wait for the full run to finish to see how experiments are performing. The entrypoint task enables this with \`report=True\`: \`\`\` """MLE Agent — orchestrates ML experiments using Flyte's durable sandbox. The agent: 1. Profiles the dataset using a trusted tool (data never touches the LLM). 2. Asks OpenAI to design a set of experiments (algorithms, hyperparams, feature strategy). 3. For each experiment, generates Monty orchestration code and executes it via flyte.sandbox.orchestrate\_local(), which dispatches the heavy compute as durable tasks. 4. Analyzes results, iterates if needed. 5. Produces a model card summarizing the winning model. The Monty sandbox ensures the LLM-generated orchestration code is safe — it can only call the pre-approved tool functions and has no access to imports, network, or filesystem. """ import asyncio import inspect import json import os import textwrap from dataclasses import dataclass import flyte import flyte.sandbox from flyte.io import File from mle\_bot.schemas import ExperimentConfig, InitialDesign, IterationDecision from mle\_bot.environments import agent\_env from mle\_bot.tools.data import profile\_dataset, split\_dataset from mle\_bot.tools.evaluation import evaluate\_model, rank\_experiments from mle\_bot.tools.exploration import explore\_dataset from mle\_bot.tools.features import engineer\_features from mle\_bot.tools.predictions import get\_predictions from mle\_bot.tools.resampling import resample\_dataset from mle\_bot.tools.selection import select\_features from mle\_bot.tools.training import train\_model # {{docs-fragment tools}} # All tools exposed to the sandbox. # Keys must match the function names used in LLM-generated orchestration code. TOOLS = \[\ profile\_dataset, split\_dataset, explore\_dataset,\ engineer\_features, resample\_dataset, select\_features,\ train\_model, get\_predictions, evaluate\_model, rank\_experiments,\ \] TOOLS\_BY\_NAME = {t.func.\_\_name\_\_ if hasattr(t, "func") else t.\_\_name\_\_: t for t in TOOLS} # {{/docs-fragment tools}} # --------------------------------------------------------------------------- # Prompt builders # --------------------------------------------------------------------------- def \_tool\_signatures() -> str: """Build a summary of available tool signatures and docstrings for the system prompt.""" parts = \[\] for t in TOOLS: func = t.func if hasattr(t, "func") else t sig = inspect.signature(func) doc = inspect.getdoc(func) or "" # Trim docstring to first 40 lines so prompt stays manageable doc\_lines = doc.splitlines()\[:40\] short\_doc = "\\n ".join(doc\_lines) parts.append(f"async def {func.\_\_name\_\_}{sig}:\\n \\"\\"\\"{short\_doc}\\"\\"\\"\\n ...") return "\\n\\n".join(parts) # {{docs-fragment orchestration\_prompt}} def \_build\_orchestration\_system\_prompt(profile: dict) -> str: monty\_rules = flyte.sandbox.ORCHESTRATOR\_SYNTAX\_PROMPT tools\_section = \_tool\_signatures() is\_imbalanced = profile.get("is\_imbalanced", False) class\_balance = profile.get("class\_balance", {}) columns = profile.get("columns", \[\]) numeric\_cols = profile.get("numeric\_columns", \[\]) categorical\_cols = profile.get("categorical\_columns", \[\]) corr = profile.get("feature\_target\_corr", {}) corr\_str = ", ".join(f"{k}: {v:+.3f}" for k, v in list(corr.items())\[:8\]) if corr else "n/a" shape = profile.get("shape", \[0, 0\]) return f"""\\ You are an expert ML engineer. Your job is to design and write the best possible pipeline for a machine learning experiment, then generate the Python orchestration code to execute it. The code runs inside a restricted sandbox. The last expression in your code is returned as the result. All tool calls are made like regular function calls — you do NOT need to await them. ## Dataset context Shape: {shape\[0\]:,} rows × {shape\[1\]} columns Numeric features: {numeric\_cols} Categorical features (excluded from model — not supported): {categorical\_cols} Class balance: {class\_balance}, imbalanced: {is\_imbalanced} Feature-target correlations (raw, point-biserial): {corr\_str} ## General ML best practices — apply these based on the dataset context above \*\*Feature engineering\*\* (engineer\_features tool): - Sequential/time-series data (timestamp column present, rows ordered over time): rolling window features (means, stds, min/max) capture trends that point-in-time readings miss. Lag features capture recent history. Choose window sizes relative to the prediction horizon and temporal resolution of the data. - Tabular cross-sectional data: normalization helps linear models and distance-based methods. Interaction terms can help if correlations are weak individually. - Consider skipping feature engineering entirely for a baseline — it establishes whether raw features already carry enough signal. \*\*Class imbalance\*\* (when is\_imbalanced=true): - Tree ensembles: use class\_weight="balanced" or scale\_pos\_weight=n\_neg/n\_pos. - Threshold: the default 0.5 decision threshold may not be optimal — the model's probability output is what matters, threshold is tuned at deployment time. - Metric: ROC-AUC is robust to imbalance; avg\_precision is better when positives are very rare. \*\*Algorithm selection\*\*: - XGBoost / GradientBoosting: strong default for tabular data, handles missing values, built-in imbalance handling. Start here unless data is very small. - RandomForest: more robust to outliers, good for noisy data, parallelizes well. - LogisticRegression: fast linear baseline. Useful to establish whether the problem is linearly separable before adding complexity. - Prefer simpler models when n\_samples < 5,000 to avoid overfitting. \*\*Resampling\*\* (resample\_dataset tool) — data-level imbalance handling: - Use when class\_weight/scale\_pos\_weight alone isn't improving recall adequately, or when you want to test whether data-level vs algorithm-level imbalance handling works better for this dataset. - ONLY resample the TRAIN split — never test. Resampling test data gives misleading metrics. - "oversample": fast, no new information, good first try. - "smote": synthetic samples via interpolation — more diverse than random oversampling, better for high-dimensional or sparse minority classes. - "undersample": loses majority data — only useful when majority class is very large and training speed is a concern. \*\*Feature selection\*\* (select\_features tool) — prune after feature engineering: - Use after engineer\_features when the feature count is large (20+) and you suspect many features are redundant or noisy (e.g., rolling stats at many window sizes). - "mutual\_info": ranks by non-linear dependence with target — best general choice. - "variance\_threshold": drops near-constant features — cheap first pass. - "correlation\_filter": drops redundant features that are highly correlated with each other — useful when many rolling windows capture the same trend. - Can be applied before or after splitting. Apply the same selection to both train and test to ensure the model sees the same features at evaluation time. \*\*Prediction output\*\* (get\_predictions tool) — enables two advanced patterns: 1. Error analysis: train a model → get\_predictions(model, test\_file, target) → explore\_dataset(predictions\_file, {{"class\_distributions": \["feature\_x"\], "target\_column": "correct"}}) to see which examples the model gets wrong. Use this to inform feature engineering for the next iteration. 2. Stacking: train base\_model → get\_predictions(base\_model, train\_file, target) → train a meta\_model on the predictions CSV (use "predicted\_prob" as a feature alongside original features) → evaluate meta\_model on test. get\_predictions returns a CSV with columns: all originals + predicted\_prob, predicted\_class, correct. \*\*Pipeline structure\*\* — you are not required to follow a fixed sequence. Design what makes sense for this specific experiment. ## Available tools {tools\_section} ## Monty sandbox restrictions {monty\_rules} ## Critical patterns for using tool results split\_dataset returns a File — call it twice: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") engineer\_features returns a File — chain calls freely: eng = engineer\_features(train\_file, {{"rolling\_columns": \[...\], "windows": \[...\]}}) eng2 = engineer\_features(eng, {{"normalize": true, "target\_column": target\_column}}) train\_model returns a File — pass directly to evaluate\_model: model\_file = train\_model(train\_file, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, test\_file, target\_column) evaluate\_model returns a dict — subscript reads are allowed: roc = eval\_result\["metrics"\]\["roc\_auc"\] Do NOT use augmented assignment (+=), subscript assignment (d\["k"\]=v), or try/except. Build dicts as literals only. The last expression (no assignment) is the return value. ## When fixing a previous error Read the error and the failing code carefully before writing a fix. Identify the root cause — do not just change variable names or add no-ops. Trace what each tool returns, what each subsequent call expects, and where the mismatch is. Then fix the underlying logic, not just the surface symptom. ## Pipeline design — you decide the structure You are NOT required to follow a fixed sequence. Design the pipeline that makes most sense for the experiment. Examples of valid approaches: Baseline (no feature engineering): train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") model\_file = train\_model(train\_file, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, test\_file, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Two-stage feature engineering (rolling then normalize separately): train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") rolled\_train = engineer\_features(train\_file, {{"rolling\_columns": \["vibration"\], "windows": \[6, 24\]}}) rolled\_test = engineer\_features(test\_file, {{"rolling\_columns": \["vibration"\], "windows": \[6, 24\]}}) eng\_train = engineer\_features(rolled\_train, {{"normalize": true, "target\_column": target\_column}}) eng\_test = engineer\_features(rolled\_test, {{"normalize": true, "target\_column": target\_column}}) model\_file = train\_model(eng\_train, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, eng\_test, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Compare two class weightings and return the better model: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") model\_a = train\_model(train\_file, target\_column, "xgboost", {{"n\_estimators": 100, "scale\_pos\_weight": 10}}) model\_b = train\_model(train\_file, target\_column, "xgboost", {{"n\_estimators": 100, "scale\_pos\_weight": 33}}) eval\_a = evaluate\_model(model\_a, test\_file, target\_column) eval\_b = evaluate\_model(model\_b, test\_file, target\_column) best\_eval = eval\_a if eval\_a\["metrics"\]\["roc\_auc"\] > eval\_b\["metrics"\]\["roc\_auc"\] else eval\_b {{"experiment\_name": experiment\_name, "algorithm": "xgboost", "metrics": best\_eval\["metrics"\], "confusion\_matrix": best\_eval\["confusion\_matrix"\], "threshold\_analysis": best\_eval\["threshold\_analysis"\], "n\_samples": best\_eval\["n\_samples"\]}} SMOTE oversampling before training: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") eng\_train = engineer\_features(train\_file, {{"rolling\_columns": \["vibration\_mms"\], "windows": \[6, 12\]}}) eng\_test = engineer\_features(test\_file, {{"rolling\_columns": \["vibration\_mms"\], "windows": \[6, 12\]}}) resampled\_train = resample\_dataset(eng\_train, target\_column, {{"strategy": "smote", "target\_ratio": 0.2}}) model\_file = train\_model(resampled\_train, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, eng\_test, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Feature engineering followed by feature selection: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") eng\_train = engineer\_features(train\_file, {{"rolling\_columns": \["vibration\_mms", "temperature\_c"\], "windows": \[6, 12, 24\]}}) eng\_test = engineer\_features(test\_file, {{"rolling\_columns": \["vibration\_mms", "temperature\_c"\], "windows": \[6, 12, 24\]}}) sel\_train = select\_features(eng\_train, target\_column, {{"method": "mutual\_info", "k": 15}}) sel\_test = select\_features(eng\_test, target\_column, {{"method": "mutual\_info", "k": 15}}) model\_file = train\_model(sel\_train, target\_column, algorithm, hyperparams) eval\_result = evaluate\_model(model\_file, sel\_test, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\]}} Error analysis — explore what the model gets wrong, then return that as insight: train\_file = split\_dataset(data, target\_column, 0.2, time\_column, "train") test\_file = split\_dataset(data, target\_column, 0.2, time\_column, "test") model\_file = train\_model(train\_file, target\_column, algorithm, hyperparams) pred\_file = get\_predictions(model\_file, test\_file, target\_column) error\_analysis = explore\_dataset(pred\_file, {{"target\_column": "correct", "class\_distributions": \["vibration\_mms", "temperature\_c"\]}}) eval\_result = evaluate\_model(model\_file, test\_file, target\_column) {{"experiment\_name": experiment\_name, "algorithm": algorithm, "metrics": eval\_result\["metrics"\], "confusion\_matrix": eval\_result\["confusion\_matrix"\], "threshold\_analysis": eval\_result\["threshold\_analysis"\], "n\_samples": eval\_result\["n\_samples"\], "error\_analysis": error\_analysis}} The last expression MUST be a dict with at minimum these keys: experiment\_name, algorithm, metrics, confusion\_matrix, threshold\_analysis, n\_samples Additional keys (e.g. error\_analysis) are allowed and will appear in the report. ## Response format Respond in exactly this format: ## Reasoning \[Your thinking: what pipeline makes sense for this experiment and why. Consider whether\ feature engineering helps, whether class imbalance needs special treatment, whether\ chaining multiple steps adds value, etc.\] ## Code \`\`\`python \[your orchestration code\] \`\`\` """ # {{/docs-fragment orchestration\_prompt}} def \_build\_analysis\_system\_prompt(max\_iterations: int, current\_iteration: int) -> str: remaining = max\_iterations - current\_iteration - 1 return f"""\\ You are an expert ML engineer analyzing experiment results to guide the next iteration of model development. You must respond with valid JSON only — no markdown, no explanation outside the JSON. Response format: {{ "should\_continue": true | false, "reasoning": "What you observed, what it tells you, and what to try next", "exploration\_requests": \[\ {{\ "question": "The specific hypothesis you are testing, e.g. 'Do failure cases show meaningfully higher vibration than healthy cases?'",\ "analysis\_type": "class\_distributions",\ "target\_column": "failure\_24h",\ "class\_distributions": \["vibration\_mms", "temperature\_c"\]\ }}\ \], "next\_experiments": \[\ {{\ "name": "descriptive experiment name",\ "algorithm": "xgboost" | "random\_forest" | "gradient\_boosting" | "logistic\_regression",\ "hyperparams": {{ ... algorithm-specific hyperparams ... }},\ "feature\_config": {{\ "group\_column": "...",\ "time\_column": "...",\ "rolling\_columns": \[...\],\ "windows": \[...\],\ "lag\_columns": \[...\],\ "lags": \[...\],\ "normalize": true | false,\ "drop\_columns": \[...\],\ "fillna\_method": "forward"\ }},\ "rationale": "Why this experiment is worth trying"\ }}\ \] }} exploration\_requests rules: - Max 2 requests per iteration. - Each request targets EXACTLY ONE analysis\_type. Do not mix multiple types in one request. - Supported analysis\_type values and their required config fields: "class\_distributions" → requires: target\_column, class\_distributions (list of columns) "correlation\_matrix" → requires: correlation\_matrix: true "temporal\_trend" → requires: temporal\_trend: {{time\_column, target\_column, freq}} "group\_stats" → requires: group\_stats: {{group\_column, target\_column}} "outlier\_summary" → requires: outlier\_summary (list of columns) "feature\_target\_corr\_by\_group" → requires: feature\_target\_corr\_by\_group: {{group\_column, target\_column, feature\_columns}} - The "question" field is required. It must be a specific testable hypothesis, not a general request. Bad: "explore the data". Good: "Is vibration\_mms higher for failures?" - Set exploration\_requests to \[\] if the current results already tell you enough to design the next experiments. Only explore when you have a concrete unanswered question. When deciding next experiments, reason about WHAT WAS TRIED vs what hasn't been explored. Each result includes used\_feature\_engineering, used\_rolling\_features, used\_lag\_features. Think systematically: if no feature engineering was tried yet, does the data profile suggest it would help (weak raw correlations, temporal/sequential structure)? If feature engineering helped, can it be improved? Avoid experiments identical to ones tried. Iteration context: this is iteration {current\_iteration + 1} of {max\_iterations} requested. Remaining iterations allowed: {remaining}. Set should\_continue=false only if: - Best ROC-AUC >= 0.97, OR - No remaining iterations (remaining == 0), OR - Results have genuinely plateaued (< 0.005 ROC-AUC improvement over last iteration AND you have already tried the most promising directions) Otherwise keep exploring — the user asked for {max\_iterations} iterations for a reason. """ def \_build\_initial\_design\_system\_prompt() -> str: return """\\ You are an expert ML engineer. Given a dataset profile and a problem description, design the first batch of experiments to run. You must respond with valid JSON only — no markdown, no explanation outside the JSON. Response format: { "problem\_type": "binary\_classification", "primary\_metric": "roc\_auc" | "f1" | "recall", "reasoning": "Brief description of your strategy", "experiments": \[\ {\ "name": "descriptive experiment name",\ "algorithm": "xgboost" | "random\_forest" | "gradient\_boosting" | "logistic\_regression",\ "hyperparams": { ... algorithm-specific hyperparams ... },\ "feature\_config": {\ "group\_column": "",\ "time\_column": "",\ "rolling\_columns": \[\],\ "windows": \[\],\ "lag\_columns": \[\],\ "lags": \[\],\ "normalize": false,\ "drop\_columns": \[\],\ "fillna\_method": "forward"\ },\ "rationale": "Why this experiment makes sense given the data profile"\ }\ \] } Design 2-3 experiments for the first batch. Good first batches typically include: - A fast baseline to establish a floor (e.g. logistic\_regression with default settings) - Your best initial hypothesis given the data profile - Optionally one variant that tests a specific idea suggested by the profile Use the dataset profile to guide your choices: - feature\_target\_corr: weak raw correlations suggest feature engineering may help - categorical\_columns: note these are excluded from the model automatically - is\_imbalanced: handle with class\_weight or scale\_pos\_weight - Shape and column types should inform algorithm complexity (simpler models for small datasets) - A time column suggests sequential structure; lag/rolling features may capture temporal patterns The feature\_config in each experiment describes what engineer\_features should apply. Leave all fields empty/false if no feature engineering is needed for that experiment. The orchestration code generator will decide the exact pipeline — your job here is to specify what the experiment is trying to learn, not to prescribe every implementation detail. """ # --------------------------------------------------------------------------- # LLM client # --------------------------------------------------------------------------- def \_openai\_client(): from openai import OpenAI return OpenAI(api\_key=os.environ\["OPENAI\_API\_KEY"\]) async def \_call\_llm(system: str, messages: list\[dict\], model: str = "gpt-4o") -> str: """Call OpenAI and return the response text.""" client = \_openai\_client() response = await asyncio.to\_thread( client.chat.completions.create, model=model, messages=\[{"role": "system", "content": system}, \*messages\], temperature=0.2, ) return response.choices\[0\].message.content def \_extract\_code(text: str) -> str: """Pull Python code out of a markdown code block.""" if "\`\`\`python" in text: start = text.index("\`\`\`python") + len("\`\`\`python") end = text.index("\`\`\`", start) return text\[start:end\].strip() if "\`\`\`" in text: start = text.index("\`\`\`") + 3 end = text.index("\`\`\`", start) return text\[start:end\].strip() return text.strip() def \_extract\_reasoning(text: str) -> str: """Extract the ## Reasoning section from LLM response.""" if "## Reasoning" in text: start = text.index("## Reasoning") + len("## Reasoning") if "## Code" in text: end = text.index("## Code") return text\[start:end\].strip() return text\[start:\].strip() return "" def \_parse\_json(text: str) -> dict: """Extract and parse JSON from LLM response.""" text = text.strip() if "\`\`\`json" in text: start = text.index("\`\`\`json") + 7 end = text.index("\`\`\`", start) text = text\[start:end\].strip() elif "\`\`\`" in text: start = text.index("\`\`\`") + 3 end = text.index("\`\`\`", start) text = text\[start:end\].strip() return json.loads(text) # --------------------------------------------------------------------------- # Display helpers # --------------------------------------------------------------------------- def \_recommend\_threshold(threshold\_analysis: list, min\_precision: float = 0.70) -> dict | None: """Find the threshold that maximises recall subject to precision >= min\_precision.""" candidates = \[t for t in threshold\_analysis if t\["precision"\] >= min\_precision\] if not candidates: return None return max(candidates, key=lambda t: t\["recall"\]) def \_print\_experiment\_table(results: list\["ExperimentResult"\], best\_name: str) -> None: """Print a ranked comparison table of all experiments.""" sorted\_results = sorted(results, key=lambda r: r.metrics.get("roc\_auc", 0), reverse=True) print("\\n" + "─" \* 78) print(f" {'Rank':<5} {'Experiment':<32} {'ROC-AUC':<9} {'F1':<7} {'Recall':<8} {'Note'}") print("─" \* 78) for rank, r in enumerate(sorted\_results, 1): note = "◀ winner" if r.name == best\_name else "" roc = r.metrics.get("roc\_auc", 0) f1 = r.metrics.get("f1", 0) recall = r.metrics.get("recall", 0) print(f" {rank:<5} {r.name:<32} {roc:<9.4f} {f1:<7.4f} {recall:<8.4f} {note}") print("─" \* 78) def \_print\_threshold\_recommendation(threshold\_analysis: list, default\_metrics: dict) -> None: """Print the operational threshold recommendation.""" rec = \_recommend\_threshold(threshold\_analysis) if not rec: return default\_recall = default\_metrics.get("recall", 0) default\_precision = default\_metrics.get("precision", 0) missed\_pct = round((1 - rec\["recall"\]) \* 100, 1) false\_alarm\_pct = round((1 - rec\["precision"\]) \* 100, 1) print(f"\\n Recommended decision threshold: {rec\['threshold'\]}") print(f" ├─ Precision : {rec\['precision'\]:.0%} ({false\_alarm\_pct}% of alerts are false alarms)") print(f" ├─ Recall : {rec\['recall'\]:.0%} (catches {rec\['recall'\]\*100:.0f}% of actual failures)") print(f" └─ F1 : {rec\['f1'\]:.4f}") print(f" Default threshold (0.5): Precision={default\_precision:.0%}, Recall={default\_recall:.0%}") if rec\["recall"\] > default\_recall: extra = round((rec\["recall"\] - default\_recall) \* 100, 1) print(f" → Lowering threshold catches {extra}% more failures at cost of more alerts") # --------------------------------------------------------------------------- # Orchestration code generation (durable Flyte task with Flyte report) # --------------------------------------------------------------------------- @agent\_env.task async def plan\_experiment( experiment\_json: str, profile\_json: str, target\_column: str, time\_column: str, previous\_error: str = "", previous\_code: str = "", llm\_model: str = "gpt-4o", ) -> str: """LLM plans a single experiment: reasons about the pipeline and generates Monty code. Runs as a durable Flyte task so each experiment's planning step is traceable. Returns a JSON string: {"code": "...", "reasoning": "..."}. Args: experiment\_json: JSON string of the experiment spec (name, algorithm, hyperparams, ...). profile\_json: JSON string of the dataset profile from profile\_dataset. target\_column: Name of the target column. time\_column: Time column for temporal splitting, or empty string. previous\_error: Error message from the previous attempt (empty on first try). previous\_code: Code that failed on the previous attempt (empty on first try). llm\_model: OpenAI model identifier. Returns: str — JSON string with keys "code" and "reasoning". """ experiment = json.loads(experiment\_json) profile = json.loads(profile\_json) exp\_name = experiment.get("name", "experiment") # Strip rationale — it was written by the design LLM to explain \*why\* this # experiment was chosen. Passing it here causes plan\_experiment to parrot it # back as "reasoning" instead of independently thinking about \*how\* to build # the best pipeline. Keep only the technical spec. pipeline\_spec = { k: v for k, v in experiment.items() if k not in ("rationale",) } system = \_build\_orchestration\_system\_prompt(profile) user\_content = textwrap.dedent(f""" Design and implement the best pipeline for this experiment: Name: {exp\_name} Algorithm: {pipeline\_spec.get("algorithm")} Hyperparams: {json.dumps(pipeline\_spec.get("hyperparams", {}), indent=2)} Feature config hint: {json.dumps(pipeline\_spec.get("feature\_config", {}), indent=2)} Available sandbox inputs: - data: File — the full dataset CSV - target\_column: str = "{target\_column}" - time\_column: str = "{time\_column}" (empty string means no time ordering) - experiment\_name: str = "{exp\_name}" The feature config hint is a suggestion from the experiment designer — you can follow it, improve on it, or override it if the dataset context and your ML judgment suggest a better approach. In your ## Reasoning, explain your actual pipeline decisions: what you chose to do (or not do) and why, based on the dataset profile above. Do not restate the experiment name or why it was chosen. """).strip() messages = \[{"role": "user", "content": user\_content}\] if previous\_code and previous\_error: messages = \[\ {"role": "user", "content": user\_content},\ {"role": "assistant", "content": f"\`\`\`python\\n{previous\_code}\\n\`\`\`"},\ {"role": "user", "content": f"That code failed with this error:\\n\\n{previous\_error}\\n\\nPlease fix it."},\ \] response = await \_call\_llm(system, messages, llm\_model) reasoning = \_extract\_reasoning(response) code = \_extract\_code(response) return json.dumps({"code": code, "reasoning": reasoning}) @flyte.trace async def design\_experiments( problem\_description: str, profile\_json: str, llm\_model: str = "gpt-4o", ) -> str: """LLM designs the initial batch of experiments given problem + dataset profile. Traced so the prompt/response is visible in the Flyte UI and results are cached for deterministic replay on crash/retry. Returns raw LLM response (JSON string matching InitialDesign schema). """ design\_prompt = textwrap.dedent(f""" Problem description: {problem\_description} Dataset profile: {profile\_json} Design the first batch of experiments. """).strip() return await \_call\_llm( \_build\_initial\_design\_system\_prompt(), \[{"role": "user", "content": design\_prompt}\], llm\_model, ) @flyte.trace async def analyze\_iteration( analysis\_prompt: str, max\_iterations: int, current\_iteration: int, llm\_model: str = "gpt-4o", ) -> str: """LLM analyzes experiment results and decides whether/how to continue. Traced so the prompt/response is visible in the Flyte UI and results are cached for deterministic replay on crash/retry. Returns raw LLM response (JSON string matching IterationDecision schema). """ return await \_call\_llm( \_build\_analysis\_system\_prompt(max\_iterations, current\_iteration), \[{"role": "user", "content": analysis\_prompt}\], llm\_model, ) @flyte.trace async def plan\_followup( analysis\_prompt: str, analysis\_response: str, followup\_prompt: str, max\_iterations: int, current\_iteration: int, llm\_model: str = "gpt-4o", ) -> str: """LLM designs next experiments after targeted data explorations. Traced so the prompt/response is visible in the Flyte UI and results are cached for deterministic replay on crash/retry. Returns raw LLM response (JSON string with {"next\_experiments": \[...\]}). """ return await \_call\_llm( \_build\_analysis\_system\_prompt(max\_iterations, current\_iteration), \[\ {"role": "user", "content": analysis\_prompt},\ {"role": "assistant", "content": analysis\_response},\ {"role": "user", "content": followup\_prompt},\ \], llm\_model, ) def \_corrupt\_experiment\_for\_demo(exp\_dict: dict) -> dict: """Introduce a deliberate error into the first experiment for demo purposes. Corrupts the algorithm name so the LLM must recover from a known-bad value. The retry loop will catch this, regenerate with the error message, and fix it. """ corrupted = dict(exp\_dict) corrupted\["algorithm"\] = corrupted\["algorithm"\] + "\_INVALID" return corrupted # --------------------------------------------------------------------------- # Main agent loop # --------------------------------------------------------------------------- @dataclass class ExperimentResult: name: str algorithm: str metrics: dict confusion\_matrix: dict threshold\_analysis: list n\_samples: int code: str attempts: int reasoning: str = "" error: str = "" @dataclass class AgentResult: model\_card: str best\_experiment: str best\_metrics: dict all\_results: list\[ExperimentResult\] iterations: int total\_experiments: int async def \_run\_experiment( exp: "ExperimentConfig", exp\_dict: dict, inject\_failure: bool, data: File, target\_column: str, time\_column: str, profile: dict, llm\_model: str, max\_retries: int, ) -> "ExperimentResult | None": """Run a single experiment with retries. Returns None on total failure.""" exp\_name = exp.name profile\_json = json.dumps(profile) print(f"\\n ┌─ {exp\_name} \[{exp.algorithm}\]") if exp.rationale: for line in textwrap.wrap(exp.rationale, width=58): print(f" │ {line}") if inject\_failure: print(f" │ \[injecting failure for demo: algorithm='{exp\_dict\['algorithm'\]}'\]") code = "" error = "" result = None attempt = 0 reasoning = "" # {{docs-fragment retry\_loop}} for attempt in range(max\_retries): try: with flyte.group(exp\_name): plan\_json = await plan\_experiment.aio( experiment\_json=json.dumps(exp\_dict), profile\_json=profile\_json, target\_column=target\_column, time\_column=time\_column, previous\_error=error, previous\_code=code, llm\_model=llm\_model, ) plan = json.loads(plan\_json) code = plan\["code"\] reasoning = plan.get("reasoning", "") result = await flyte.sandbox.orchestrate\_local( code, inputs={"data": data, "target\_column": target\_column, "time\_column": time\_column, "experiment\_name": exp\_name}, tasks=TOOLS, ) error = "" break except Exception as exc: error = str(exc) short\_error = error\[:100\] + "..." if len(error) > 100 else error print(f" │ attempt {attempt + 1} failed: {short\_error}") print(f" │ → asking LLM to fix and retry...") if inject\_failure and attempt == 0: exp\_dict = exp.model\_dump() # {{/docs-fragment retry\_loop}} if result and not error: exp\_result = ExperimentResult( name=exp\_name, algorithm=exp.algorithm, metrics=result.get("metrics", {}), confusion\_matrix=result.get("confusion\_matrix", {}), threshold\_analysis=result.get("threshold\_analysis", \[\]), n\_samples=result.get("n\_samples", 0), code=code, reasoning=reasoning, attempts=attempt + 1, ) m = exp\_result.metrics attempts\_note = f" (recovered after {attempt + 1} attempts)" if attempt > 0 else "" print(f" └─ ROC-AUC={m.get('roc\_auc')}, F1={m.get('f1')}, Recall={m.get('recall')}{attempts\_note}") return exp\_result print(f" └─ FAILED after {max\_retries} attempts — skipping.") return None async def run\_agent( data: File, problem\_description: str, target\_column: str, time\_column: str = "", max\_iterations: int = 3, max\_retries\_per\_experiment: int = 3, llm\_model: str = "gpt-4o", inject\_failure: bool = False, ) -> AgentResult: """Run the MLE agent end-to-end. Args: data: CSV file containing the dataset. problem\_description: Natural language description of the ML problem. target\_column: Name of the target column to predict. time\_column: Optional column to use for time-based train/test split. max\_iterations: Maximum number of experiment iterations to run. max\_retries\_per\_experiment: Max times to retry a failed sandbox execution. llm\_model: OpenAI model to use (default: gpt-4o). inject\_failure: If True, corrupts the first experiment to demonstrate self-healing. """ print(f"\\n{'='\*60}") print(f"MLE Agent starting") print(f"Problem: {problem\_description}") print(f"Target: {target\_column}") if inject\_failure: print(f"\[demo mode: failure injection enabled\]") print(f"{'='\*60}\\n") # {{docs-fragment phase1\_profile}} # --- Phase 1: Profile the dataset (trusted tool, LLM never sees raw data) --- print(">> Phase 1: Profiling dataset...") with flyte.group("profile"): profile = await profile\_dataset(data, target\_column) # {{/docs-fragment phase1\_profile}} print(f" Shape: {profile\['shape'\]}, Classes: {profile\['target\_distribution'\]}") print(f" Imbalanced: {profile\['is\_imbalanced'\]}, Columns: {len(profile\['columns'\])}") corr = profile.get("feature\_target\_corr", {}) top\_corr = list(corr.items())\[:5\] print(f" Top correlations: {', '.join(f'{k}={v:+.3f}' for k,v in top\_corr)}") # Stream report: dataset summary await flyte.report.log.aio( f"

MLE Agent Run

" f"

Problem: {problem\_description}

" f"

Dataset: {profile\['shape'\]\[0\]:,} rows × {profile\['shape'\]\[1\]} cols  |  " f"Class balance: {profile\['class\_balance'\]}  |  Imbalanced: {profile\['is\_imbalanced'\]}

" f"

Top feature-target correlations (raw): " + ", ".join(f"{k}: {v:+.3f}" for k, v in top\_corr) + f"

", do\_flush=True, ) # --- Phase 2: LLM designs initial experiments --- print("\\n>> Phase 2: Designing initial experiments...") design\_response = await design\_experiments( problem\_description=problem\_description, profile\_json=json.dumps(profile), llm\_model=llm\_model, ) design = InitialDesign.model\_validate(\_parse\_json(design\_response)) print(f" Primary metric: {design.primary\_metric}") print(f" Strategy: {design.reasoning}") print(f" Experiments planned: {len(design.experiments)}") all\_results: list\[ExperimentResult\] = \[\] iteration\_log: list\[dict\] = \[\] # tracks per-iteration decisions + explorations for summary current\_experiments: list\[ExperimentConfig\] = design.experiments first\_experiment = True # --- Phase 3: Iterative experiment loop --- for iteration in range(max\_iterations): experiments = current\_experiments if not experiments: print(f"\\n>> No experiments to run in iteration {iteration + 1}. Stopping.") break print(f"\\n>> Phase 3.{iteration + 1}: Running {len(experiments)} experiment(s) in parallel...") # Assign names and prepare dicts before launching in parallel exp\_batch = \[\] for i, exp in enumerate(experiments): if not exp.name: exp.name = f"experiment\_{len(all\_results) + i + 1}" exp\_dict = exp.model\_dump() inject\_this = inject\_failure and first\_experiment and i == 0 if inject\_this: exp\_dict = \_corrupt\_experiment\_for\_demo(exp\_dict) first\_experiment = False exp\_batch.append((exp, exp\_dict, inject\_this)) # {{docs-fragment parallel\_execute}} batch\_results = await asyncio.gather(\*\[\ \_run\_experiment(\ exp=exp,\ exp\_dict=exp\_dict,\ inject\_failure=inject\_this,\ data=data,\ target\_column=target\_column,\ time\_column=time\_column,\ profile=profile,\ llm\_model=llm\_model,\ max\_retries=max\_retries\_per\_experiment,\ )\ for exp, exp\_dict, inject\_this in exp\_batch\ \]) # {{/docs-fragment parallel\_execute}} for exp\_result in batch\_results: if exp\_result is not None: all\_results.append(exp\_result) # Stream report: each experiment as it completes m = exp\_result.metrics html = ( f"

Iteration {iteration + 1} — {exp\_result.name}

" f"

Algorithm: {exp\_result.algorithm}  |  " f"ROC-AUC: {m.get('roc\_auc')}  |  " f"F1: {m.get('f1')}  |  " f"Recall: {m.get('recall')}  |  " f"Attempts: {exp\_result.attempts}

" ) if exp\_result.reasoning: html += f"
Reasoning
{exp\_result.reasoning}
" html += f"
Generated Code
{exp\_result.code}
" await flyte.report.log.aio(html, do\_flush=True) # --- Phase 4: Analyze results, decide whether to iterate --- if all\_results and iteration < max\_iterations - 1: print(f"\\n>> Phase 4.{iteration + 1}: Analyzing results, deciding next steps...") results\_summary = \[\ {\ "experiment\_name": r.name,\ "algorithm": r.algorithm,\ "metrics": r.metrics,\ "confusion\_matrix": r.confusion\_matrix,\ "used\_feature\_engineering": "engineer\_features" in r.code,\ "used\_rolling\_features": "rolling\_columns" in r.code,\ "used\_lag\_features": "lag\_columns" in r.code,\ }\ for r in all\_results\ \] analysis\_prompt = textwrap.dedent(f""" Problem: {problem\_description} Dataset profile: shape={profile\['shape'\]}, imbalanced={profile\['is\_imbalanced'\]} Feature-target correlations (raw): {json.dumps(profile.get('feature\_target\_corr', {}), indent=2)} Experiment results so far (iteration {iteration + 1}): {json.dumps(results\_summary, indent=2)} Should we run more experiments? If yes, request any data explorations you need, then specify what experiments to run next. """).strip() analysis\_response = await analyze\_iteration( analysis\_prompt=analysis\_prompt, max\_iterations=max\_iterations, current\_iteration=iteration, llm\_model=llm\_model, ) decision = IterationDecision.model\_validate(\_parse\_json(analysis\_response)) verdict = "continuing" if decision.should\_continue else "stopping" print(f" Decision: {verdict}") print(f" Reasoning: {decision.reasoning}") # Stream report: analysis decision await flyte.report.log.aio( f"

Analysis — Iteration {iteration + 1}

" f"

Decision: {verdict}

" f"

Reasoning: {decision.reasoning}

", do\_flush=True, ) # Track this iteration for the experiment journey summary iter\_entry = { "iteration": iteration + 1, "experiments": \[r.name for r in batch\_results if r is not None\], "best\_roc\_auc": max( (r.metrics.get("roc\_auc", 0) for r in all\_results), default=0 ), "reasoning": decision.reasoning, "explorations": \[\], } # --- Targeted exploration before next iteration --- if decision.should\_continue and decision.exploration\_requests: print(f" Running {len(decision.exploration\_requests)} exploration request(s)...") exploration\_questions = \[\] exploration\_results = \[\] for i, req in enumerate(decision.exploration\_requests): question = req.get("question", f"Exploration {i + 1}") # Strip agent-level metadata — tool only needs the analysis config tool\_config = {k: v for k, v in req.items() if k not in ("question", "analysis\_type")} print(f" Q: {question}") with flyte.group(f"explore\_{iteration + 1}\_{i + 1}"): result = await explore\_dataset(data, tool\_config) exploration\_questions.append(question) exploration\_results.append(result) iter\_entry\["explorations"\].append({"question": question}) await flyte.report.log.aio( f"

Exploration {i + 1}

" f"

Question: {question}

" f"
Results
{json.dumps(result, indent=2)}
", do\_flush=True, ) # Build follow-up that explicitly connects each question to its answer qa\_pairs = "\\n\\n".join( f'Question {i + 1}: "{q}"\\nResult:\\n{json.dumps(r, indent=2)}' for i, (q, r) in enumerate(zip(exploration\_questions, exploration\_results)) ) followup\_prompt = textwrap.dedent(f""" You requested {len(exploration\_results)} targeted exploration(s). Here is what you asked and what you learned: {qa\_pairs} Given what you learned and your earlier reasoning: "{decision.reasoning}" Now specify the next experiments. For each experiment, briefly state which exploration insight informed your choice. Respond with valid JSON: {{"next\_experiments": \[...same schema as before...\]}} """).strip() followup\_response = await plan\_followup( analysis\_prompt=analysis\_prompt, analysis\_response=analysis\_response, followup\_prompt=followup\_prompt, max\_iterations=max\_iterations, current\_iteration=iteration, llm\_model=llm\_model, ) followup = \_parse\_json(followup\_response) current\_experiments = IterationDecision.model\_validate({ "should\_continue": True, "reasoning": decision.reasoning, "next\_experiments": followup.get("next\_experiments", \[\]), }).next\_experiments print(f" Post-exploration: {len(current\_experiments)} experiment(s) planned") else: current\_experiments = decision.next\_experiments iteration\_log.append(iter\_entry) if not decision.should\_continue: break # --- Phase 5: Rank all results and generate model card --- print(f"\\n>> Phase 5: Ranking {len(all\_results)} experiment(s) and generating model card...") if not all\_results: return AgentResult( model\_card="No experiments completed successfully.", best\_experiment="", best\_metrics={}, all\_results=\[\], iterations=iteration + 1, total\_experiments=0, ) ranking\_input = \[\ {\ "experiment\_name": r.name,\ "metrics": r.metrics,\ "confusion\_matrix": r.confusion\_matrix,\ }\ for r in all\_results\ \] with flyte.group("rank"): ranking = await rank\_experiments(json.dumps(ranking\_input)) best\_name = ranking\["best\_experiment"\] best\_result = next(r for r in all\_results if r.name == best\_name) \_print\_experiment\_table(all\_results, best\_name) \_print\_threshold\_recommendation(best\_result.threshold\_analysis, best\_result.metrics) # Stream report: final rankings table rows = "".join( f"{row\['rank'\]}" f"{'' if row\['experiment\_name'\] == best\_name else ''}" f"{row\['experiment\_name'\]}" f"{'' if row\['experiment\_name'\] == best\_name else ''}" f"{row\['roc\_auc'\]}{row\['f1'\]}" f"{row\['recall'\]}{row\['precision'\]}" for row in ranking.get("ranking", \[\]) ) await flyte.report.log.aio( f"

Final Rankings

" f"" f"" f"{rows}
RankExperimentROC-AUCF1RecallPrecision
" f"

{ranking.get('summary', '')}

", do\_flush=True, ) # Stream report: experiment journey summary journey\_rows = "" for entry in iteration\_log: exps = ", ".join(entry\["experiments"\]) if entry\["experiments"\] else "—" explorations = "; ".join(e\["question"\] for e in entry\["explorations"\]) if entry\["explorations"\] else "—" short\_reasoning = (entry\["reasoning"\]\[:120\] + "…") if len(entry\["reasoning"\]) > 120 else entry\["reasoning"\] journey\_rows += ( f"" f"{entry\['iteration'\]}" f"{exps}" f"{entry\['best\_roc\_auc'\]:.4f}" f"{short\_reasoning}" f"{explorations}" f"" ) await flyte.report.log.aio( f"

Experiment Journey

" f"" f"" f"{journey\_rows}" f"
IterExperimentsBest ROC-AUCKey insightExplorations
", do\_flush=True, ) model\_card = await \_generate\_model\_card( problem\_description=problem\_description, profile=profile, all\_results=all\_results, best\_result=best\_result, ranking=ranking, iteration\_log=iteration\_log, llm\_model=llm\_model, ) print(f"\\n{'='\*60}") print(f"DONE — Best model: {best\_name}") print(f" ROC-AUC={best\_result.metrics.get('roc\_auc')}, F1={best\_result.metrics.get('f1')}") print(f"{'='\*60}\\n") return AgentResult( model\_card=model\_card, best\_experiment=best\_name, best\_metrics=best\_result.metrics, all\_results=all\_results, iterations=iteration + 1, total\_experiments=len(all\_results), ) async def \_generate\_model\_card( problem\_description: str, profile: dict, all\_results: list\[ExperimentResult\], best\_result: ExperimentResult, ranking: dict, iteration\_log: list\[dict\], llm\_model: str, ) -> str: """Generate a markdown model card summarizing the winning model.""" system = textwrap.dedent(""" You are an ML engineer writing a model card for a trained model. Write in markdown. Be concise but informative. Include: - Problem statement - Dataset summary - Experiment journey (brief per-iteration narrative: what was tried, what was learned, what changed) - Experiment summary (table of all experiments with metrics) - Winning model details (algorithm, key hyperparams, metrics, threshold analysis) - Recommendations for deployment (decision threshold, monitoring) """).strip() results\_text = "\\n".join( f"- {r.name} ({r.algorithm}): ROC-AUC={r.metrics.get('roc\_auc')}, " f"F1={r.metrics.get('f1')}, Recall={r.metrics.get('recall')}" for r in all\_results ) journey\_text = "" if iteration\_log: journey\_text = "\\n\\nIteration log:\\n" + "\\n".join( f" Iteration {e\['iteration'\]}: ran \[{', '.join(e\['experiments'\])}\], " f"best ROC-AUC so far={e\['best\_roc\_auc'\]:.4f}. " f"Key insight: {e\['reasoning'\]\[:200\]}. " + (f"Explorations: {'; '.join(x\['question'\] for x in e\['explorations'\])}" if e\['explorations'\] else "") for e in iteration\_log ) user\_content = textwrap.dedent(f""" Problem: {problem\_description} Dataset: {profile\['shape'\]\[0\]} rows × {profile\['shape'\]\[1\]} cols. Class balance: {profile\['class\_balance'\]} Imbalanced: {profile\['is\_imbalanced'\]} {journey\_text} All experiments: {results\_text} Best model: {best\_result.name} ({best\_result.algorithm}) Metrics: {json.dumps(best\_result.metrics, indent=2)} Confusion matrix: {json.dumps(best\_result.confusion\_matrix, indent=2)} Threshold analysis: {json.dumps(best\_result.threshold\_analysis, indent=2)} Ranking summary: {ranking\['summary'\]} """).strip() response = await \_call\_llm(system, \[{"role": "user", "content": user\_content}\], llm\_model) return response # --------------------------------------------------------------------------- # Durable entrypoint (runs the agent as a Flyte task in the cloud) # --------------------------------------------------------------------------- # {{docs-fragment entrypoint}} @agent\_env.task(retries=1, report=True) async def mle\_agent\_task( data: File, problem\_description: str, target\_column: str, time\_column: str = "", max\_iterations: int = 3, ) -> str: """Durable Flyte task entrypoint for the MLE agent.""" result = await run\_agent( data=data, problem\_description=problem\_description, target\_column=target\_column, time\_column=time\_column, max\_iterations=max\_iterations, ) return result.model\_card # {{/docs-fragment entrypoint}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/mle\_bot/mle\_bot/agent.py\* As each experiment completes, the agent streams its metrics to the report: \`\`\`python await flyte.report.log.aio( f"

Iteration {iteration + 1}: {exp\_result.name}

" f"

Algorithm: {exp\_result.algorithm}  |  " f"ROC-AUC: {m.get('roc\_auc')}  |  " f"F1: {m.get('f1')}

", do\_flush=True, ) \`\`\` The final report includes a dataset summary, per-experiment metrics with expandable reasoning and generated code, the analysis decisions at each iteration, a final rankings table, and an experiment journey summary showing how the agent's strategy evolved. ## Running it First, generate the synthetic demo dataset (a predictive maintenance scenario with 175k+ rows of simulated sensor data from 20 industrial pumps): \`\`\`bash uv run main.py generate-data \`\`\` Then submit the agent to your Flyte cluster: \`\`\`bash uv run main.py run \\ --data data/predictive\_maintenance.csv \\ --problem "Predict pump failures 24 hours before they happen" \\ --target failure\_24h \\ --time-column timestamp \\ --max-iterations 3 \\ --output results/report.md \`\`\` The agent connects to your cluster via \`~/.flyte/config.yaml\`, uploads the CSV, and submits the agent task. You'll see a URL to track the execution in the Flyte UI, and logs will stream to your terminal. > \[!NOTE\] > You'll need to register your OpenAI API key as a cluster secret before running: > \`flyte create secret openai-api-key \` If you want to see the self-healing retry loop in action, add the \`--inject-failure\` flag. This deliberately corrupts the first experiment so the agent has to detect the error and recover, which makes for a nice demo of the durability guarantees. ## Why Flyte? You could build something similar with plain Python and \`exec()\`. But there are a few things you'd lose. \*\*Safety.\*\* Flyte's sandbox restricts LLM-generated code to calling your pre-approved functions and nothing else. No imports, no network, no filesystem. If you wouldn't give an intern root access to your production cluster, you probably shouldn't give an LLM unrestricted code execution either. \*\*Durability.\*\* Every tool call is a Flyte task. If the agent process crashes halfway through iteration 3, the experiments that already completed are cached. You restart and pick up where you left off instead of retraining models from scratch. For long-running ML experiments, this matters. \*\*Observability.\*\* You can see every LLM prompt, every generated code snippet, every tool invocation, and every result in the Flyte UI. When the agent makes a questionable decision (like skipping feature engineering on temporal data), you can trace exactly why: the prompt it received, the profile it read, the reasoning it generated. \*\*Compute isolation.\*\* The ML tools run on cloud instances with the CPU and memory they need. The agent itself runs on a small 1-CPU instance since all it does is call the LLM and dispatch tool tasks. You're not bottlenecked by your laptop, and you're not paying for GPU-class compute to run an orchestration loop. \*\*Parallelism.\*\* Multiple experiments run simultaneously via \`asyncio.gather()\`, each dispatching its own durable tasks. Flyte handles the scheduling. If you have three experiments in a batch and each involves training + evaluation, that's six tasks running concurrently on cloud compute. The MLE Bot is a specific example of a more general pattern: giving an LLM the ability to reason about \*what\* work should be done, while Flyte handles \*how\* that work gets executed safely, durably, and at scale. The sandbox is the boundary between the two. Everything above the boundary is LLM-generated and untrusted. Everything below it is your code, running on your infrastructure, with all the guarantees you'd expect from a production orchestrator. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/deep-research === # Deep research > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/deep\_research\_agent); based on work by \[Together AI\](https://github.com/togethercomputer/open\_deep\_research). This example demonstrates how to build an agentic workflow for deep research—a multi-step reasoning system that mirrors how a human researcher explores, analyzes, and synthesizes information from the web. Deep research refers to the iterative process of thoroughly investigating a topic: identifying relevant sources, evaluating their usefulness, refining the research direction, and ultimately producing a well-structured summary or report. It's a long-running task that requires the agent to reason over time, adapt its strategy, and chain multiple steps together, making it an ideal fit for an agentic architecture. In this example, we use: - \[Tavily\](https://www.tavily.com/) to search for and retrieve high-quality online resources. - \[LiteLLM\](https://litellm.ai/) to route LLM calls that perform reasoning, evaluation, and synthesis. The agent executes a multi-step trajectory: - Parallel search across multiple queries. - Evaluation of retrieved results. - Adaptive iteration: If results are insufficient, it formulates new research queries and repeats the search-evaluate cycle. - Synthesis: After a fixed number of iterations, it produces a comprehensive research report. What makes this workflow compelling is its dynamic, evolving nature. The agent isn't just following a fixed plan; it's making decisions in context, using multiple prompts and reasoning steps to steer the process. Flyte is uniquely well-suited for this kind of system. It provides: - Structured composition of dynamic reasoning steps - Built-in parallelism for faster search and evaluation - Traceability and observability into each step and iteration - Scalability for long-running or compute-intensive workloads !\[Result\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/gifs/tutorials/deep-research/result.gif) Throughout this guide, we'll show how to design this workflow using the Flyte SDK, and how to unlock the full potential of agentic development with tools you already know and trust. ## Setting up the environment Let's begin by setting up the task environment. We define the following components: - Secrets for Together and Tavily API keys - A custom image with required Python packages and apt dependencies (\`pandoc\`, \`texlive-xetex\`) - External YAML file with all LLM prompts baked into the container \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pydantic==2.11.5",\ # "litellm==1.72.2",\ # "tavily-python==0.7.5",\ # "together==1.5.24",\ # "markdown==3.8.2",\ # "pymdown-extensions==10.16.1",\ # \] # main = "main" # params = "" # /// # {{docs-fragment env}} import asyncio import json from pathlib import Path import flyte import yaml from flyte.io.\_file import File from libs.utils.data\_types import ( DeepResearchResult, DeepResearchResults, ResearchPlan, SourceList, ) from libs.utils.generation import generate\_html, generate\_toc\_image from libs.utils.llms import asingle\_shot\_llm\_call from libs.utils.log import AgentLogger from libs.utils.tavily\_search import atavily\_search\_results TIME\_LIMIT\_MULTIPLIER = 5 MAX\_COMPLETION\_TOKENS = 4096 logging = AgentLogger("together.open\_deep\_research") env = flyte.TaskEnvironment( name="deep-researcher", secrets=\[\ flyte.Secret(key="together\_api\_key", as\_env\_var="TOGETHER\_API\_KEY"),\ flyte.Secret(key="tavily\_api\_key", as\_env\_var="TAVILY\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="deep-research-agent", pre=True) .with\_apt\_packages("pandoc", "texlive-xetex") .with\_source\_file(Path("prompts.yaml"), "/root"), resources=flyte.Resources(cpu=1), ) # {{/docs-fragment env}} # {{docs-fragment generate\_research\_queries}} @env.task async def generate\_research\_queries( topic: str, planning\_model: str, json\_model: str, prompts\_file: File, ) -> list\[str\]: async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) PLANNING\_PROMPT = prompts\["planning\_prompt"\] plan = "" logging.info(f"\\n\\nGenerated deep research plan for topic: {topic}\\n\\nPlan:") async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=PLANNING\_PROMPT, message=f"Research Topic: {topic}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): plan += chunk print(chunk, end="", flush=True) SEARCH\_PROMPT = prompts\["plan\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=SEARCH\_PROMPT, message=f"Plan to be parsed: {plan}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk plan = json.loads(response\_json) return plan\["queries"\] # {{/docs-fragment generate\_research\_queries}} async def \_summarize\_content\_async( raw\_content: str, query: str, prompt: str, summarization\_model: str, ) -> str: """Summarize content asynchronously using the LLM""" logging.info("Summarizing content asynchronously using the LLM") result = "" async for chunk in asingle\_shot\_llm\_call( model=summarization\_model, system\_prompt=prompt, message=f"{raw\_content}\\n\\n{query}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): result += chunk return result # {{docs-fragment search\_and\_summarize}} @env.task async def search\_and\_summarize( query: str, prompts\_file: File, summarization\_model: str, ) -> DeepResearchResults: """Perform search for a single query""" if len(query) > 400: # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues query = query\[:400\] logging.info(f"Truncated query to 400 characters: {query}") response = await atavily\_search\_results(query) logging.info("Tavily Search Called.") async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) RAW\_CONTENT\_SUMMARIZER\_PROMPT = prompts\["raw\_content\_summarizer\_prompt"\] with flyte.group("summarize-content"): # Create tasks for summarization summarization\_tasks = \[\] result\_info = \[\] for result in response.results: if result.raw\_content is None: continue task = \_summarize\_content\_async( result.raw\_content, query, RAW\_CONTENT\_SUMMARIZER\_PROMPT, summarization\_model, ) summarization\_tasks.append(task) result\_info.append(result) # Use return\_exceptions=True to prevent exceptions from propagating summarized\_contents = await asyncio.gather( \*summarization\_tasks, return\_exceptions=True ) # Filter out exceptions summarized\_contents = \[\ result for result in summarized\_contents if not isinstance(result, Exception)\ \] formatted\_results = \[\] for result, summarized\_content in zip(result\_info, summarized\_contents): formatted\_results.append( DeepResearchResult( title=result.title, link=result.link, content=result.content, raw\_content=result.raw\_content, filtered\_raw\_content=summarized\_content, ) ) return DeepResearchResults(results=formatted\_results) # {{/docs-fragment search\_and\_summarize}} @env.task async def search\_all\_queries( queries: list\[str\], summarization\_model: str, prompts\_file: File ) -> DeepResearchResults: """Execute searches for all queries in parallel""" tasks = \[\] results\_list = \[\] tasks = \[\ search\_and\_summarize(query, prompts\_file, summarization\_model)\ for query in queries\ \] if tasks: res\_list = await asyncio.gather(\*tasks) results\_list.extend(res\_list) # Combine all results combined\_results = DeepResearchResults(results=\[\]) for results in results\_list: combined\_results = combined\_results + results return combined\_results # {{docs-fragment evaluate\_research\_completeness}} @env.task async def evaluate\_research\_completeness( topic: str, results: DeepResearchResults, queries: list\[str\], prompts\_file: File, planning\_model: str, json\_model: str, ) -> list\[str\]: """ Evaluate if the current search results are sufficient or if more research is needed. Returns an empty list if research is complete, or a list of additional queries if more research is needed. """ # Format the search results for the LLM formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) EVALUATION\_PROMPT = prompts\["evaluation\_prompt"\] logging.info("\\nEvaluation: ") evaluation = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=EVALUATION\_PROMPT, message=( f"{topic}\\n\\n" f"{queries}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=None, ): evaluation += chunk print(chunk, end="", flush=True) EVALUATION\_PARSING\_PROMPT = prompts\["evaluation\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=EVALUATION\_PARSING\_PROMPT, message=f"Evaluation to be parsed: {evaluation}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk evaluation = json.loads(response\_json) return evaluation\["queries"\] # {{/docs-fragment evaluate\_research\_completeness}} # {{docs-fragment filter\_results}} @env.task async def filter\_results( topic: str, results: DeepResearchResults, prompts\_file: File, planning\_model: str, json\_model: str, max\_sources: int, ) -> DeepResearchResults: """Filter the search results based on the research plan""" # Format the search results for the LLM, without the raw content formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) FILTER\_PROMPT = prompts\["filter\_prompt"\] logging.info("\\nFilter response: ") filter\_response = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=FILTER\_PROMPT, message=( f"{topic}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): filter\_response += chunk print(chunk, end="", flush=True) logging.info(f"Filter response: {filter\_response}") FILTER\_PARSING\_PROMPT = prompts\["filter\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=FILTER\_PARSING\_PROMPT, message=f"Filter response to be parsed: {filter\_response}", response\_format={ "type": "json\_object", "schema": SourceList.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk sources = json.loads(response\_json)\["sources"\] logging.info(f"Filtered sources: {sources}") if max\_sources != -1: sources = sources\[:max\_sources\] # Filter the results based on the source list filtered\_results = \[\ results.results\[i - 1\] for i in sources if i - 1 < len(results.results)\ \] return DeepResearchResults(results=filtered\_results) # {{/docs-fragment filter\_results}} def \_remove\_thinking\_tags(answer: str) -> str: """Remove content within tags""" while "" in answer and "" in answer: start = answer.find("") end = answer.find("") + len("") answer = answer\[:start\] + answer\[end:\] return answer # {{docs-fragment generate\_research\_answer}} @env.task async def generate\_research\_answer( topic: str, results: DeepResearchResults, remove\_thinking\_tags: bool, prompts\_file: File, answer\_model: str, ) -> str: """ Generate a comprehensive answer to the research topic based on the search results. Returns a detailed response that synthesizes information from all search results. """ formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) ANSWER\_PROMPT = prompts\["answer\_prompt"\] answer = "" async for chunk in asingle\_shot\_llm\_call( model=answer\_model, system\_prompt=ANSWER\_PROMPT, message=f"Research Topic: {topic}\\n\\nSearch Results:\\n{formatted\_results}", response\_format=None, # NOTE: This is the max\_token parameter for the LLM call on Together AI, # may need to be changed for other providers max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): answer += chunk # this is just to avoid typing complaints if answer is None or not isinstance(answer, str): logging.error("No answer generated") return "No answer generated" if remove\_thinking\_tags: # Remove content within tags answer = \_remove\_thinking\_tags(answer) # Remove markdown code block markers if they exist at the beginning if answer.lstrip().startswith("\`\`\`"): # Find the first line break after the opening backticks first\_linebreak = answer.find("\\n", answer.find("\`\`\`")) if first\_linebreak != -1: # Remove everything up to and including the first line break answer = answer\[first\_linebreak + 1 :\] # Remove closing code block if it exists if answer.rstrip().endswith("\`\`\`"): answer = answer.rstrip()\[:-3\].rstrip() return answer.strip() # {{/docs-fragment generate\_research\_answer}} # {{docs-fragment research\_topic}} @env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff\_factor=2)) async def research\_topic( topic: str, budget: int = 3, remove\_thinking\_tags: bool = True, max\_queries: int = 5, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 40, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", prompts\_file: File | str = "prompts.yaml", ) -> str: """Main method to conduct research on a topic. Will be used for weave evals.""" if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) # Step 1: Generate initial queries queries = await generate\_research\_queries( topic=topic, planning\_model=planning\_model, json\_model=json\_model, prompts\_file=prompts\_file, ) queries = \[topic, \*queries\[: max\_queries - 1\]\] all\_queries = queries.copy() logging.info(f"Initial queries: {queries}") if len(queries) == 0: logging.error("No initial queries generated") return "No initial queries generated" # Step 2: Perform initial search results = await search\_all\_queries(queries, summarization\_model, prompts\_file) logging.info(f"Initial search complete, found {len(results.results)} results") # Step 3: Conduct iterative research within budget for iteration in range(budget): with flyte.group(f"eval\_iteration\_{iteration}"): # Evaluate if more research is needed additional\_queries = await evaluate\_research\_completeness( topic=topic, results=results, queries=all\_queries, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, ) # Filter out empty strings and check if any queries remain additional\_queries = \[q for q in additional\_queries if q\] if not additional\_queries: logging.info("No need for additional research") break # for debugging purposes we limit the number of queries additional\_queries = additional\_queries\[:max\_queries\] logging.info(f"Additional queries: {additional\_queries}") # Expand research with new queries new\_results = await search\_all\_queries( additional\_queries, summarization\_model, prompts\_file ) logging.info( f"Follow-up search complete, found {len(new\_results.results)} results" ) results = results + new\_results all\_queries.extend(additional\_queries) # Step 4: Generate final answer logging.info(f"Generating final answer for topic: {topic}") results = results.dedup() logging.info(f"Deduplication complete, kept {len(results.results)} results") filtered\_results = await filter\_results( topic=topic, results=results, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, ) logging.info( f"LLM Filtering complete, kept {len(filtered\_results.results)} results" ) # Generate final answer answer = await generate\_research\_answer( topic=topic, results=filtered\_results, remove\_thinking\_tags=remove\_thinking\_tags, prompts\_file=prompts\_file, answer\_model=answer\_model, ) return answer # {{/docs-fragment research\_topic}} # {{docs-fragment main}} @env.task(report=True) async def main( topic: str = ( "List the essential requirements for a developer-focused agent orchestration system." ), prompts\_file: File | str = "/root/prompts.yaml", budget: int = 2, remove\_thinking\_tags: bool = True, max\_queries: int = 3, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 10, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", ) -> str: if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) answer = await research\_topic( topic=topic, budget=budget, remove\_thinking\_tags=remove\_thinking\_tags, max\_queries=max\_queries, answer\_model=answer\_model, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, summarization\_model=summarization\_model, prompts\_file=prompts\_file, ) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") toc\_image\_url = await generate\_toc\_image( yaml.safe\_load(yaml\_contents)\["data\_visualization\_prompt"\], planning\_model, topic, ) html\_content = await generate\_html(answer, toc\_image\_url) await flyte.report.replace.aio(html\_content, do\_flush=True) await flyte.report.flush.aio() return html\_content # {{/docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep\_research\_agent/agent.py\* The Python packages are declared at the top of the file using the \`uv\` script style: \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b6",\ # "pydantic==2.11.5",\ # "litellm==1.72.2",\ # "tavily-python==0.7.5",\ # "together==1.5.24",\ # "markdown==3.8.2",\ # "pymdown-extensions==10.16.1",\ # \] # /// \`\`\` ## Generate research queries This task converts a user prompt into a list of focused queries. It makes two LLM calls to generate a high-level research plan and parse that plan into atomic search queries. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pydantic==2.11.5",\ # "litellm==1.72.2",\ # "tavily-python==0.7.5",\ # "together==1.5.24",\ # "markdown==3.8.2",\ # "pymdown-extensions==10.16.1",\ # \] # main = "main" # params = "" # /// # {{docs-fragment env}} import asyncio import json from pathlib import Path import flyte import yaml from flyte.io.\_file import File from libs.utils.data\_types import ( DeepResearchResult, DeepResearchResults, ResearchPlan, SourceList, ) from libs.utils.generation import generate\_html, generate\_toc\_image from libs.utils.llms import asingle\_shot\_llm\_call from libs.utils.log import AgentLogger from libs.utils.tavily\_search import atavily\_search\_results TIME\_LIMIT\_MULTIPLIER = 5 MAX\_COMPLETION\_TOKENS = 4096 logging = AgentLogger("together.open\_deep\_research") env = flyte.TaskEnvironment( name="deep-researcher", secrets=\[\ flyte.Secret(key="together\_api\_key", as\_env\_var="TOGETHER\_API\_KEY"),\ flyte.Secret(key="tavily\_api\_key", as\_env\_var="TAVILY\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="deep-research-agent", pre=True) .with\_apt\_packages("pandoc", "texlive-xetex") .with\_source\_file(Path("prompts.yaml"), "/root"), resources=flyte.Resources(cpu=1), ) # {{/docs-fragment env}} # {{docs-fragment generate\_research\_queries}} @env.task async def generate\_research\_queries( topic: str, planning\_model: str, json\_model: str, prompts\_file: File, ) -> list\[str\]: async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) PLANNING\_PROMPT = prompts\["planning\_prompt"\] plan = "" logging.info(f"\\n\\nGenerated deep research plan for topic: {topic}\\n\\nPlan:") async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=PLANNING\_PROMPT, message=f"Research Topic: {topic}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): plan += chunk print(chunk, end="", flush=True) SEARCH\_PROMPT = prompts\["plan\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=SEARCH\_PROMPT, message=f"Plan to be parsed: {plan}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk plan = json.loads(response\_json) return plan\["queries"\] # {{/docs-fragment generate\_research\_queries}} async def \_summarize\_content\_async( raw\_content: str, query: str, prompt: str, summarization\_model: str, ) -> str: """Summarize content asynchronously using the LLM""" logging.info("Summarizing content asynchronously using the LLM") result = "" async for chunk in asingle\_shot\_llm\_call( model=summarization\_model, system\_prompt=prompt, message=f"{raw\_content}\\n\\n{query}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): result += chunk return result # {{docs-fragment search\_and\_summarize}} @env.task async def search\_and\_summarize( query: str, prompts\_file: File, summarization\_model: str, ) -> DeepResearchResults: """Perform search for a single query""" if len(query) > 400: # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues query = query\[:400\] logging.info(f"Truncated query to 400 characters: {query}") response = await atavily\_search\_results(query) logging.info("Tavily Search Called.") async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) RAW\_CONTENT\_SUMMARIZER\_PROMPT = prompts\["raw\_content\_summarizer\_prompt"\] with flyte.group("summarize-content"): # Create tasks for summarization summarization\_tasks = \[\] result\_info = \[\] for result in response.results: if result.raw\_content is None: continue task = \_summarize\_content\_async( result.raw\_content, query, RAW\_CONTENT\_SUMMARIZER\_PROMPT, summarization\_model, ) summarization\_tasks.append(task) result\_info.append(result) # Use return\_exceptions=True to prevent exceptions from propagating summarized\_contents = await asyncio.gather( \*summarization\_tasks, return\_exceptions=True ) # Filter out exceptions summarized\_contents = \[\ result for result in summarized\_contents if not isinstance(result, Exception)\ \] formatted\_results = \[\] for result, summarized\_content in zip(result\_info, summarized\_contents): formatted\_results.append( DeepResearchResult( title=result.title, link=result.link, content=result.content, raw\_content=result.raw\_content, filtered\_raw\_content=summarized\_content, ) ) return DeepResearchResults(results=formatted\_results) # {{/docs-fragment search\_and\_summarize}} @env.task async def search\_all\_queries( queries: list\[str\], summarization\_model: str, prompts\_file: File ) -> DeepResearchResults: """Execute searches for all queries in parallel""" tasks = \[\] results\_list = \[\] tasks = \[\ search\_and\_summarize(query, prompts\_file, summarization\_model)\ for query in queries\ \] if tasks: res\_list = await asyncio.gather(\*tasks) results\_list.extend(res\_list) # Combine all results combined\_results = DeepResearchResults(results=\[\]) for results in results\_list: combined\_results = combined\_results + results return combined\_results # {{docs-fragment evaluate\_research\_completeness}} @env.task async def evaluate\_research\_completeness( topic: str, results: DeepResearchResults, queries: list\[str\], prompts\_file: File, planning\_model: str, json\_model: str, ) -> list\[str\]: """ Evaluate if the current search results are sufficient or if more research is needed. Returns an empty list if research is complete, or a list of additional queries if more research is needed. """ # Format the search results for the LLM formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) EVALUATION\_PROMPT = prompts\["evaluation\_prompt"\] logging.info("\\nEvaluation: ") evaluation = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=EVALUATION\_PROMPT, message=( f"{topic}\\n\\n" f"{queries}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=None, ): evaluation += chunk print(chunk, end="", flush=True) EVALUATION\_PARSING\_PROMPT = prompts\["evaluation\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=EVALUATION\_PARSING\_PROMPT, message=f"Evaluation to be parsed: {evaluation}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk evaluation = json.loads(response\_json) return evaluation\["queries"\] # {{/docs-fragment evaluate\_research\_completeness}} # {{docs-fragment filter\_results}} @env.task async def filter\_results( topic: str, results: DeepResearchResults, prompts\_file: File, planning\_model: str, json\_model: str, max\_sources: int, ) -> DeepResearchResults: """Filter the search results based on the research plan""" # Format the search results for the LLM, without the raw content formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) FILTER\_PROMPT = prompts\["filter\_prompt"\] logging.info("\\nFilter response: ") filter\_response = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=FILTER\_PROMPT, message=( f"{topic}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): filter\_response += chunk print(chunk, end="", flush=True) logging.info(f"Filter response: {filter\_response}") FILTER\_PARSING\_PROMPT = prompts\["filter\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=FILTER\_PARSING\_PROMPT, message=f"Filter response to be parsed: {filter\_response}", response\_format={ "type": "json\_object", "schema": SourceList.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk sources = json.loads(response\_json)\["sources"\] logging.info(f"Filtered sources: {sources}") if max\_sources != -1: sources = sources\[:max\_sources\] # Filter the results based on the source list filtered\_results = \[\ results.results\[i - 1\] for i in sources if i - 1 < len(results.results)\ \] return DeepResearchResults(results=filtered\_results) # {{/docs-fragment filter\_results}} def \_remove\_thinking\_tags(answer: str) -> str: """Remove content within tags""" while "" in answer and "" in answer: start = answer.find("") end = answer.find("") + len("") answer = answer\[:start\] + answer\[end:\] return answer # {{docs-fragment generate\_research\_answer}} @env.task async def generate\_research\_answer( topic: str, results: DeepResearchResults, remove\_thinking\_tags: bool, prompts\_file: File, answer\_model: str, ) -> str: """ Generate a comprehensive answer to the research topic based on the search results. Returns a detailed response that synthesizes information from all search results. """ formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) ANSWER\_PROMPT = prompts\["answer\_prompt"\] answer = "" async for chunk in asingle\_shot\_llm\_call( model=answer\_model, system\_prompt=ANSWER\_PROMPT, message=f"Research Topic: {topic}\\n\\nSearch Results:\\n{formatted\_results}", response\_format=None, # NOTE: This is the max\_token parameter for the LLM call on Together AI, # may need to be changed for other providers max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): answer += chunk # this is just to avoid typing complaints if answer is None or not isinstance(answer, str): logging.error("No answer generated") return "No answer generated" if remove\_thinking\_tags: # Remove content within tags answer = \_remove\_thinking\_tags(answer) # Remove markdown code block markers if they exist at the beginning if answer.lstrip().startswith("\`\`\`"): # Find the first line break after the opening backticks first\_linebreak = answer.find("\\n", answer.find("\`\`\`")) if first\_linebreak != -1: # Remove everything up to and including the first line break answer = answer\[first\_linebreak + 1 :\] # Remove closing code block if it exists if answer.rstrip().endswith("\`\`\`"): answer = answer.rstrip()\[:-3\].rstrip() return answer.strip() # {{/docs-fragment generate\_research\_answer}} # {{docs-fragment research\_topic}} @env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff\_factor=2)) async def research\_topic( topic: str, budget: int = 3, remove\_thinking\_tags: bool = True, max\_queries: int = 5, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 40, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", prompts\_file: File | str = "prompts.yaml", ) -> str: """Main method to conduct research on a topic. Will be used for weave evals.""" if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) # Step 1: Generate initial queries queries = await generate\_research\_queries( topic=topic, planning\_model=planning\_model, json\_model=json\_model, prompts\_file=prompts\_file, ) queries = \[topic, \*queries\[: max\_queries - 1\]\] all\_queries = queries.copy() logging.info(f"Initial queries: {queries}") if len(queries) == 0: logging.error("No initial queries generated") return "No initial queries generated" # Step 2: Perform initial search results = await search\_all\_queries(queries, summarization\_model, prompts\_file) logging.info(f"Initial search complete, found {len(results.results)} results") # Step 3: Conduct iterative research within budget for iteration in range(budget): with flyte.group(f"eval\_iteration\_{iteration}"): # Evaluate if more research is needed additional\_queries = await evaluate\_research\_completeness( topic=topic, results=results, queries=all\_queries, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, ) # Filter out empty strings and check if any queries remain additional\_queries = \[q for q in additional\_queries if q\] if not additional\_queries: logging.info("No need for additional research") break # for debugging purposes we limit the number of queries additional\_queries = additional\_queries\[:max\_queries\] logging.info(f"Additional queries: {additional\_queries}") # Expand research with new queries new\_results = await search\_all\_queries( additional\_queries, summarization\_model, prompts\_file ) logging.info( f"Follow-up search complete, found {len(new\_results.results)} results" ) results = results + new\_results all\_queries.extend(additional\_queries) # Step 4: Generate final answer logging.info(f"Generating final answer for topic: {topic}") results = results.dedup() logging.info(f"Deduplication complete, kept {len(results.results)} results") filtered\_results = await filter\_results( topic=topic, results=results, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, ) logging.info( f"LLM Filtering complete, kept {len(filtered\_results.results)} results" ) # Generate final answer answer = await generate\_research\_answer( topic=topic, results=filtered\_results, remove\_thinking\_tags=remove\_thinking\_tags, prompts\_file=prompts\_file, answer\_model=answer\_model, ) return answer # {{/docs-fragment research\_topic}} # {{docs-fragment main}} @env.task(report=True) async def main( topic: str = ( "List the essential requirements for a developer-focused agent orchestration system." ), prompts\_file: File | str = "/root/prompts.yaml", budget: int = 2, remove\_thinking\_tags: bool = True, max\_queries: int = 3, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 10, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", ) -> str: if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) answer = await research\_topic( topic=topic, budget=budget, remove\_thinking\_tags=remove\_thinking\_tags, max\_queries=max\_queries, answer\_model=answer\_model, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, summarization\_model=summarization\_model, prompts\_file=prompts\_file, ) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") toc\_image\_url = await generate\_toc\_image( yaml.safe\_load(yaml\_contents)\["data\_visualization\_prompt"\], planning\_model, topic, ) html\_content = await generate\_html(answer, toc\_image\_url) await flyte.report.replace.aio(html\_content, do\_flush=True) await flyte.report.flush.aio() return html\_content # {{/docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep\_research\_agent/agent.py\* LLM calls use LiteLLM, and each is wrapped with \`flyte.trace\` for observability: \`\`\` from typing import Any, AsyncIterator, Optional from litellm import acompletion, completion import flyte # {{docs-fragment asingle\_shot\_llm\_call}} @flyte.trace async def asingle\_shot\_llm\_call( model: str, system\_prompt: str, message: str, response\_format: Optional\[dict\[str, str | dict\[str, Any\]\]\] = None, max\_completion\_tokens: int | None = None, ) -> AsyncIterator\[str\]: stream = await acompletion( model=model, messages=\[\ {"role": "system", "content": system\_prompt},\ {"role": "user", "content": message},\ \], temperature=0.0, response\_format=response\_format, # NOTE: max\_token is deprecated per OpenAI API docs, use max\_completion\_tokens instead if possible # NOTE: max\_completion\_tokens is not currently supported by Together AI, so we use max\_tokens instead max\_tokens=max\_completion\_tokens, timeout=600, stream=True, ) async for chunk in stream: content = chunk.choices\[0\].delta.get("content", "") if content: yield content # {{/docs-fragment asingle\_shot\_llm\_call}} def single\_shot\_llm\_call( model: str, system\_prompt: str, message: str, response\_format: Optional\[dict\[str, str | dict\[str, Any\]\]\] = None, max\_completion\_tokens: int | None = None, ) -> str: response = completion( model=model, messages=\[\ {"role": "system", "content": system\_prompt},\ {"role": "user", "content": message},\ \], temperature=0.0, response\_format=response\_format, # NOTE: max\_token is deprecated per OpenAI API docs, use max\_completion\_tokens instead if possible # NOTE: max\_completion\_tokens is not currently supported by Together AI, so we use max\_tokens instead max\_tokens=max\_completion\_tokens, timeout=600, ) return response.choices\[0\].message\["content"\] # type: ignore \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep\_research\_agent/libs/utils/llms.py\* > \[!NOTE\] > We use \`flyte.trace\` to track intermediate steps within a task, like LLM calls or specific function executions. This lightweight decorator adds observability with minimal overhead and is especially useful for inspecting reasoning chains during task execution. ## Search and summarize We submit each research query to Tavily and summarize the results using an LLM. We run all summarization tasks with \`asyncio.gather\`, which signals to Flyte that these tasks can be distributed across separate compute resources. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pydantic==2.11.5",\ # "litellm==1.72.2",\ # "tavily-python==0.7.5",\ # "together==1.5.24",\ # "markdown==3.8.2",\ # "pymdown-extensions==10.16.1",\ # \] # main = "main" # params = "" # /// # {{docs-fragment env}} import asyncio import json from pathlib import Path import flyte import yaml from flyte.io.\_file import File from libs.utils.data\_types import ( DeepResearchResult, DeepResearchResults, ResearchPlan, SourceList, ) from libs.utils.generation import generate\_html, generate\_toc\_image from libs.utils.llms import asingle\_shot\_llm\_call from libs.utils.log import AgentLogger from libs.utils.tavily\_search import atavily\_search\_results TIME\_LIMIT\_MULTIPLIER = 5 MAX\_COMPLETION\_TOKENS = 4096 logging = AgentLogger("together.open\_deep\_research") env = flyte.TaskEnvironment( name="deep-researcher", secrets=\[\ flyte.Secret(key="together\_api\_key", as\_env\_var="TOGETHER\_API\_KEY"),\ flyte.Secret(key="tavily\_api\_key", as\_env\_var="TAVILY\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="deep-research-agent", pre=True) .with\_apt\_packages("pandoc", "texlive-xetex") .with\_source\_file(Path("prompts.yaml"), "/root"), resources=flyte.Resources(cpu=1), ) # {{/docs-fragment env}} # {{docs-fragment generate\_research\_queries}} @env.task async def generate\_research\_queries( topic: str, planning\_model: str, json\_model: str, prompts\_file: File, ) -> list\[str\]: async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) PLANNING\_PROMPT = prompts\["planning\_prompt"\] plan = "" logging.info(f"\\n\\nGenerated deep research plan for topic: {topic}\\n\\nPlan:") async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=PLANNING\_PROMPT, message=f"Research Topic: {topic}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): plan += chunk print(chunk, end="", flush=True) SEARCH\_PROMPT = prompts\["plan\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=SEARCH\_PROMPT, message=f"Plan to be parsed: {plan}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk plan = json.loads(response\_json) return plan\["queries"\] # {{/docs-fragment generate\_research\_queries}} async def \_summarize\_content\_async( raw\_content: str, query: str, prompt: str, summarization\_model: str, ) -> str: """Summarize content asynchronously using the LLM""" logging.info("Summarizing content asynchronously using the LLM") result = "" async for chunk in asingle\_shot\_llm\_call( model=summarization\_model, system\_prompt=prompt, message=f"{raw\_content}\\n\\n{query}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): result += chunk return result # {{docs-fragment search\_and\_summarize}} @env.task async def search\_and\_summarize( query: str, prompts\_file: File, summarization\_model: str, ) -> DeepResearchResults: """Perform search for a single query""" if len(query) > 400: # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues query = query\[:400\] logging.info(f"Truncated query to 400 characters: {query}") response = await atavily\_search\_results(query) logging.info("Tavily Search Called.") async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) RAW\_CONTENT\_SUMMARIZER\_PROMPT = prompts\["raw\_content\_summarizer\_prompt"\] with flyte.group("summarize-content"): # Create tasks for summarization summarization\_tasks = \[\] result\_info = \[\] for result in response.results: if result.raw\_content is None: continue task = \_summarize\_content\_async( result.raw\_content, query, RAW\_CONTENT\_SUMMARIZER\_PROMPT, summarization\_model, ) summarization\_tasks.append(task) result\_info.append(result) # Use return\_exceptions=True to prevent exceptions from propagating summarized\_contents = await asyncio.gather( \*summarization\_tasks, return\_exceptions=True ) # Filter out exceptions summarized\_contents = \[\ result for result in summarized\_contents if not isinstance(result, Exception)\ \] formatted\_results = \[\] for result, summarized\_content in zip(result\_info, summarized\_contents): formatted\_results.append( DeepResearchResult( title=result.title, link=result.link, content=result.content, raw\_content=result.raw\_content, filtered\_raw\_content=summarized\_content, ) ) return DeepResearchResults(results=formatted\_results) # {{/docs-fragment search\_and\_summarize}} @env.task async def search\_all\_queries( queries: list\[str\], summarization\_model: str, prompts\_file: File ) -> DeepResearchResults: """Execute searches for all queries in parallel""" tasks = \[\] results\_list = \[\] tasks = \[\ search\_and\_summarize(query, prompts\_file, summarization\_model)\ for query in queries\ \] if tasks: res\_list = await asyncio.gather(\*tasks) results\_list.extend(res\_list) # Combine all results combined\_results = DeepResearchResults(results=\[\]) for results in results\_list: combined\_results = combined\_results + results return combined\_results # {{docs-fragment evaluate\_research\_completeness}} @env.task async def evaluate\_research\_completeness( topic: str, results: DeepResearchResults, queries: list\[str\], prompts\_file: File, planning\_model: str, json\_model: str, ) -> list\[str\]: """ Evaluate if the current search results are sufficient or if more research is needed. Returns an empty list if research is complete, or a list of additional queries if more research is needed. """ # Format the search results for the LLM formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) EVALUATION\_PROMPT = prompts\["evaluation\_prompt"\] logging.info("\\nEvaluation: ") evaluation = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=EVALUATION\_PROMPT, message=( f"{topic}\\n\\n" f"{queries}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=None, ): evaluation += chunk print(chunk, end="", flush=True) EVALUATION\_PARSING\_PROMPT = prompts\["evaluation\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=EVALUATION\_PARSING\_PROMPT, message=f"Evaluation to be parsed: {evaluation}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk evaluation = json.loads(response\_json) return evaluation\["queries"\] # {{/docs-fragment evaluate\_research\_completeness}} # {{docs-fragment filter\_results}} @env.task async def filter\_results( topic: str, results: DeepResearchResults, prompts\_file: File, planning\_model: str, json\_model: str, max\_sources: int, ) -> DeepResearchResults: """Filter the search results based on the research plan""" # Format the search results for the LLM, without the raw content formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) FILTER\_PROMPT = prompts\["filter\_prompt"\] logging.info("\\nFilter response: ") filter\_response = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=FILTER\_PROMPT, message=( f"{topic}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): filter\_response += chunk print(chunk, end="", flush=True) logging.info(f"Filter response: {filter\_response}") FILTER\_PARSING\_PROMPT = prompts\["filter\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=FILTER\_PARSING\_PROMPT, message=f"Filter response to be parsed: {filter\_response}", response\_format={ "type": "json\_object", "schema": SourceList.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk sources = json.loads(response\_json)\["sources"\] logging.info(f"Filtered sources: {sources}") if max\_sources != -1: sources = sources\[:max\_sources\] # Filter the results based on the source list filtered\_results = \[\ results.results\[i - 1\] for i in sources if i - 1 < len(results.results)\ \] return DeepResearchResults(results=filtered\_results) # {{/docs-fragment filter\_results}} def \_remove\_thinking\_tags(answer: str) -> str: """Remove content within tags""" while "" in answer and "" in answer: start = answer.find("") end = answer.find("") + len("") answer = answer\[:start\] + answer\[end:\] return answer # {{docs-fragment generate\_research\_answer}} @env.task async def generate\_research\_answer( topic: str, results: DeepResearchResults, remove\_thinking\_tags: bool, prompts\_file: File, answer\_model: str, ) -> str: """ Generate a comprehensive answer to the research topic based on the search results. Returns a detailed response that synthesizes information from all search results. """ formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) ANSWER\_PROMPT = prompts\["answer\_prompt"\] answer = "" async for chunk in asingle\_shot\_llm\_call( model=answer\_model, system\_prompt=ANSWER\_PROMPT, message=f"Research Topic: {topic}\\n\\nSearch Results:\\n{formatted\_results}", response\_format=None, # NOTE: This is the max\_token parameter for the LLM call on Together AI, # may need to be changed for other providers max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): answer += chunk # this is just to avoid typing complaints if answer is None or not isinstance(answer, str): logging.error("No answer generated") return "No answer generated" if remove\_thinking\_tags: # Remove content within tags answer = \_remove\_thinking\_tags(answer) # Remove markdown code block markers if they exist at the beginning if answer.lstrip().startswith("\`\`\`"): # Find the first line break after the opening backticks first\_linebreak = answer.find("\\n", answer.find("\`\`\`")) if first\_linebreak != -1: # Remove everything up to and including the first line break answer = answer\[first\_linebreak + 1 :\] # Remove closing code block if it exists if answer.rstrip().endswith("\`\`\`"): answer = answer.rstrip()\[:-3\].rstrip() return answer.strip() # {{/docs-fragment generate\_research\_answer}} # {{docs-fragment research\_topic}} @env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff\_factor=2)) async def research\_topic( topic: str, budget: int = 3, remove\_thinking\_tags: bool = True, max\_queries: int = 5, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 40, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", prompts\_file: File | str = "prompts.yaml", ) -> str: """Main method to conduct research on a topic. Will be used for weave evals.""" if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) # Step 1: Generate initial queries queries = await generate\_research\_queries( topic=topic, planning\_model=planning\_model, json\_model=json\_model, prompts\_file=prompts\_file, ) queries = \[topic, \*queries\[: max\_queries - 1\]\] all\_queries = queries.copy() logging.info(f"Initial queries: {queries}") if len(queries) == 0: logging.error("No initial queries generated") return "No initial queries generated" # Step 2: Perform initial search results = await search\_all\_queries(queries, summarization\_model, prompts\_file) logging.info(f"Initial search complete, found {len(results.results)} results") # Step 3: Conduct iterative research within budget for iteration in range(budget): with flyte.group(f"eval\_iteration\_{iteration}"): # Evaluate if more research is needed additional\_queries = await evaluate\_research\_completeness( topic=topic, results=results, queries=all\_queries, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, ) # Filter out empty strings and check if any queries remain additional\_queries = \[q for q in additional\_queries if q\] if not additional\_queries: logging.info("No need for additional research") break # for debugging purposes we limit the number of queries additional\_queries = additional\_queries\[:max\_queries\] logging.info(f"Additional queries: {additional\_queries}") # Expand research with new queries new\_results = await search\_all\_queries( additional\_queries, summarization\_model, prompts\_file ) logging.info( f"Follow-up search complete, found {len(new\_results.results)} results" ) results = results + new\_results all\_queries.extend(additional\_queries) # Step 4: Generate final answer logging.info(f"Generating final answer for topic: {topic}") results = results.dedup() logging.info(f"Deduplication complete, kept {len(results.results)} results") filtered\_results = await filter\_results( topic=topic, results=results, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, ) logging.info( f"LLM Filtering complete, kept {len(filtered\_results.results)} results" ) # Generate final answer answer = await generate\_research\_answer( topic=topic, results=filtered\_results, remove\_thinking\_tags=remove\_thinking\_tags, prompts\_file=prompts\_file, answer\_model=answer\_model, ) return answer # {{/docs-fragment research\_topic}} # {{docs-fragment main}} @env.task(report=True) async def main( topic: str = ( "List the essential requirements for a developer-focused agent orchestration system." ), prompts\_file: File | str = "/root/prompts.yaml", budget: int = 2, remove\_thinking\_tags: bool = True, max\_queries: int = 3, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 10, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", ) -> str: if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) answer = await research\_topic( topic=topic, budget=budget, remove\_thinking\_tags=remove\_thinking\_tags, max\_queries=max\_queries, answer\_model=answer\_model, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, summarization\_model=summarization\_model, prompts\_file=prompts\_file, ) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") toc\_image\_url = await generate\_toc\_image( yaml.safe\_load(yaml\_contents)\["data\_visualization\_prompt"\], planning\_model, topic, ) html\_content = await generate\_html(answer, toc\_image\_url) await flyte.report.replace.aio(html\_content, do\_flush=True) await flyte.report.flush.aio() return html\_content # {{/docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep\_research\_agent/agent.py\* ## Evaluate research completeness Now we assess whether the gathered research is sufficient. Again, the task uses two LLM calls to evaluate the completeness of the results and propose additional queries if necessary. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pydantic==2.11.5",\ # "litellm==1.72.2",\ # "tavily-python==0.7.5",\ # "together==1.5.24",\ # "markdown==3.8.2",\ # "pymdown-extensions==10.16.1",\ # \] # main = "main" # params = "" # /// # {{docs-fragment env}} import asyncio import json from pathlib import Path import flyte import yaml from flyte.io.\_file import File from libs.utils.data\_types import ( DeepResearchResult, DeepResearchResults, ResearchPlan, SourceList, ) from libs.utils.generation import generate\_html, generate\_toc\_image from libs.utils.llms import asingle\_shot\_llm\_call from libs.utils.log import AgentLogger from libs.utils.tavily\_search import atavily\_search\_results TIME\_LIMIT\_MULTIPLIER = 5 MAX\_COMPLETION\_TOKENS = 4096 logging = AgentLogger("together.open\_deep\_research") env = flyte.TaskEnvironment( name="deep-researcher", secrets=\[\ flyte.Secret(key="together\_api\_key", as\_env\_var="TOGETHER\_API\_KEY"),\ flyte.Secret(key="tavily\_api\_key", as\_env\_var="TAVILY\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="deep-research-agent", pre=True) .with\_apt\_packages("pandoc", "texlive-xetex") .with\_source\_file(Path("prompts.yaml"), "/root"), resources=flyte.Resources(cpu=1), ) # {{/docs-fragment env}} # {{docs-fragment generate\_research\_queries}} @env.task async def generate\_research\_queries( topic: str, planning\_model: str, json\_model: str, prompts\_file: File, ) -> list\[str\]: async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) PLANNING\_PROMPT = prompts\["planning\_prompt"\] plan = "" logging.info(f"\\n\\nGenerated deep research plan for topic: {topic}\\n\\nPlan:") async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=PLANNING\_PROMPT, message=f"Research Topic: {topic}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): plan += chunk print(chunk, end="", flush=True) SEARCH\_PROMPT = prompts\["plan\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=SEARCH\_PROMPT, message=f"Plan to be parsed: {plan}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk plan = json.loads(response\_json) return plan\["queries"\] # {{/docs-fragment generate\_research\_queries}} async def \_summarize\_content\_async( raw\_content: str, query: str, prompt: str, summarization\_model: str, ) -> str: """Summarize content asynchronously using the LLM""" logging.info("Summarizing content asynchronously using the LLM") result = "" async for chunk in asingle\_shot\_llm\_call( model=summarization\_model, system\_prompt=prompt, message=f"{raw\_content}\\n\\n{query}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): result += chunk return result # {{docs-fragment search\_and\_summarize}} @env.task async def search\_and\_summarize( query: str, prompts\_file: File, summarization\_model: str, ) -> DeepResearchResults: """Perform search for a single query""" if len(query) > 400: # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues query = query\[:400\] logging.info(f"Truncated query to 400 characters: {query}") response = await atavily\_search\_results(query) logging.info("Tavily Search Called.") async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) RAW\_CONTENT\_SUMMARIZER\_PROMPT = prompts\["raw\_content\_summarizer\_prompt"\] with flyte.group("summarize-content"): # Create tasks for summarization summarization\_tasks = \[\] result\_info = \[\] for result in response.results: if result.raw\_content is None: continue task = \_summarize\_content\_async( result.raw\_content, query, RAW\_CONTENT\_SUMMARIZER\_PROMPT, summarization\_model, ) summarization\_tasks.append(task) result\_info.append(result) # Use return\_exceptions=True to prevent exceptions from propagating summarized\_contents = await asyncio.gather( \*summarization\_tasks, return\_exceptions=True ) # Filter out exceptions summarized\_contents = \[\ result for result in summarized\_contents if not isinstance(result, Exception)\ \] formatted\_results = \[\] for result, summarized\_content in zip(result\_info, summarized\_contents): formatted\_results.append( DeepResearchResult( title=result.title, link=result.link, content=result.content, raw\_content=result.raw\_content, filtered\_raw\_content=summarized\_content, ) ) return DeepResearchResults(results=formatted\_results) # {{/docs-fragment search\_and\_summarize}} @env.task async def search\_all\_queries( queries: list\[str\], summarization\_model: str, prompts\_file: File ) -> DeepResearchResults: """Execute searches for all queries in parallel""" tasks = \[\] results\_list = \[\] tasks = \[\ search\_and\_summarize(query, prompts\_file, summarization\_model)\ for query in queries\ \] if tasks: res\_list = await asyncio.gather(\*tasks) results\_list.extend(res\_list) # Combine all results combined\_results = DeepResearchResults(results=\[\]) for results in results\_list: combined\_results = combined\_results + results return combined\_results # {{docs-fragment evaluate\_research\_completeness}} @env.task async def evaluate\_research\_completeness( topic: str, results: DeepResearchResults, queries: list\[str\], prompts\_file: File, planning\_model: str, json\_model: str, ) -> list\[str\]: """ Evaluate if the current search results are sufficient or if more research is needed. Returns an empty list if research is complete, or a list of additional queries if more research is needed. """ # Format the search results for the LLM formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) EVALUATION\_PROMPT = prompts\["evaluation\_prompt"\] logging.info("\\nEvaluation: ") evaluation = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=EVALUATION\_PROMPT, message=( f"{topic}\\n\\n" f"{queries}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=None, ): evaluation += chunk print(chunk, end="", flush=True) EVALUATION\_PARSING\_PROMPT = prompts\["evaluation\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=EVALUATION\_PARSING\_PROMPT, message=f"Evaluation to be parsed: {evaluation}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk evaluation = json.loads(response\_json) return evaluation\["queries"\] # {{/docs-fragment evaluate\_research\_completeness}} # {{docs-fragment filter\_results}} @env.task async def filter\_results( topic: str, results: DeepResearchResults, prompts\_file: File, planning\_model: str, json\_model: str, max\_sources: int, ) -> DeepResearchResults: """Filter the search results based on the research plan""" # Format the search results for the LLM, without the raw content formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) FILTER\_PROMPT = prompts\["filter\_prompt"\] logging.info("\\nFilter response: ") filter\_response = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=FILTER\_PROMPT, message=( f"{topic}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): filter\_response += chunk print(chunk, end="", flush=True) logging.info(f"Filter response: {filter\_response}") FILTER\_PARSING\_PROMPT = prompts\["filter\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=FILTER\_PARSING\_PROMPT, message=f"Filter response to be parsed: {filter\_response}", response\_format={ "type": "json\_object", "schema": SourceList.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk sources = json.loads(response\_json)\["sources"\] logging.info(f"Filtered sources: {sources}") if max\_sources != -1: sources = sources\[:max\_sources\] # Filter the results based on the source list filtered\_results = \[\ results.results\[i - 1\] for i in sources if i - 1 < len(results.results)\ \] return DeepResearchResults(results=filtered\_results) # {{/docs-fragment filter\_results}} def \_remove\_thinking\_tags(answer: str) -> str: """Remove content within tags""" while "" in answer and "" in answer: start = answer.find("") end = answer.find("") + len("") answer = answer\[:start\] + answer\[end:\] return answer # {{docs-fragment generate\_research\_answer}} @env.task async def generate\_research\_answer( topic: str, results: DeepResearchResults, remove\_thinking\_tags: bool, prompts\_file: File, answer\_model: str, ) -> str: """ Generate a comprehensive answer to the research topic based on the search results. Returns a detailed response that synthesizes information from all search results. """ formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) ANSWER\_PROMPT = prompts\["answer\_prompt"\] answer = "" async for chunk in asingle\_shot\_llm\_call( model=answer\_model, system\_prompt=ANSWER\_PROMPT, message=f"Research Topic: {topic}\\n\\nSearch Results:\\n{formatted\_results}", response\_format=None, # NOTE: This is the max\_token parameter for the LLM call on Together AI, # may need to be changed for other providers max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): answer += chunk # this is just to avoid typing complaints if answer is None or not isinstance(answer, str): logging.error("No answer generated") return "No answer generated" if remove\_thinking\_tags: # Remove content within tags answer = \_remove\_thinking\_tags(answer) # Remove markdown code block markers if they exist at the beginning if answer.lstrip().startswith("\`\`\`"): # Find the first line break after the opening backticks first\_linebreak = answer.find("\\n", answer.find("\`\`\`")) if first\_linebreak != -1: # Remove everything up to and including the first line break answer = answer\[first\_linebreak + 1 :\] # Remove closing code block if it exists if answer.rstrip().endswith("\`\`\`"): answer = answer.rstrip()\[:-3\].rstrip() return answer.strip() # {{/docs-fragment generate\_research\_answer}} # {{docs-fragment research\_topic}} @env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff\_factor=2)) async def research\_topic( topic: str, budget: int = 3, remove\_thinking\_tags: bool = True, max\_queries: int = 5, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 40, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", prompts\_file: File | str = "prompts.yaml", ) -> str: """Main method to conduct research on a topic. Will be used for weave evals.""" if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) # Step 1: Generate initial queries queries = await generate\_research\_queries( topic=topic, planning\_model=planning\_model, json\_model=json\_model, prompts\_file=prompts\_file, ) queries = \[topic, \*queries\[: max\_queries - 1\]\] all\_queries = queries.copy() logging.info(f"Initial queries: {queries}") if len(queries) == 0: logging.error("No initial queries generated") return "No initial queries generated" # Step 2: Perform initial search results = await search\_all\_queries(queries, summarization\_model, prompts\_file) logging.info(f"Initial search complete, found {len(results.results)} results") # Step 3: Conduct iterative research within budget for iteration in range(budget): with flyte.group(f"eval\_iteration\_{iteration}"): # Evaluate if more research is needed additional\_queries = await evaluate\_research\_completeness( topic=topic, results=results, queries=all\_queries, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, ) # Filter out empty strings and check if any queries remain additional\_queries = \[q for q in additional\_queries if q\] if not additional\_queries: logging.info("No need for additional research") break # for debugging purposes we limit the number of queries additional\_queries = additional\_queries\[:max\_queries\] logging.info(f"Additional queries: {additional\_queries}") # Expand research with new queries new\_results = await search\_all\_queries( additional\_queries, summarization\_model, prompts\_file ) logging.info( f"Follow-up search complete, found {len(new\_results.results)} results" ) results = results + new\_results all\_queries.extend(additional\_queries) # Step 4: Generate final answer logging.info(f"Generating final answer for topic: {topic}") results = results.dedup() logging.info(f"Deduplication complete, kept {len(results.results)} results") filtered\_results = await filter\_results( topic=topic, results=results, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, ) logging.info( f"LLM Filtering complete, kept {len(filtered\_results.results)} results" ) # Generate final answer answer = await generate\_research\_answer( topic=topic, results=filtered\_results, remove\_thinking\_tags=remove\_thinking\_tags, prompts\_file=prompts\_file, answer\_model=answer\_model, ) return answer # {{/docs-fragment research\_topic}} # {{docs-fragment main}} @env.task(report=True) async def main( topic: str = ( "List the essential requirements for a developer-focused agent orchestration system." ), prompts\_file: File | str = "/root/prompts.yaml", budget: int = 2, remove\_thinking\_tags: bool = True, max\_queries: int = 3, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 10, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", ) -> str: if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) answer = await research\_topic( topic=topic, budget=budget, remove\_thinking\_tags=remove\_thinking\_tags, max\_queries=max\_queries, answer\_model=answer\_model, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, summarization\_model=summarization\_model, prompts\_file=prompts\_file, ) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") toc\_image\_url = await generate\_toc\_image( yaml.safe\_load(yaml\_contents)\["data\_visualization\_prompt"\], planning\_model, topic, ) html\_content = await generate\_html(answer, toc\_image\_url) await flyte.report.replace.aio(html\_content, do\_flush=True) await flyte.report.flush.aio() return html\_content # {{/docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep\_research\_agent/agent.py\* ## Filter results In this step, we evaluate the relevance of search results and rank them. This task returns the most useful sources for the final synthesis. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pydantic==2.11.5",\ # "litellm==1.72.2",\ # "tavily-python==0.7.5",\ # "together==1.5.24",\ # "markdown==3.8.2",\ # "pymdown-extensions==10.16.1",\ # \] # main = "main" # params = "" # /// # {{docs-fragment env}} import asyncio import json from pathlib import Path import flyte import yaml from flyte.io.\_file import File from libs.utils.data\_types import ( DeepResearchResult, DeepResearchResults, ResearchPlan, SourceList, ) from libs.utils.generation import generate\_html, generate\_toc\_image from libs.utils.llms import asingle\_shot\_llm\_call from libs.utils.log import AgentLogger from libs.utils.tavily\_search import atavily\_search\_results TIME\_LIMIT\_MULTIPLIER = 5 MAX\_COMPLETION\_TOKENS = 4096 logging = AgentLogger("together.open\_deep\_research") env = flyte.TaskEnvironment( name="deep-researcher", secrets=\[\ flyte.Secret(key="together\_api\_key", as\_env\_var="TOGETHER\_API\_KEY"),\ flyte.Secret(key="tavily\_api\_key", as\_env\_var="TAVILY\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="deep-research-agent", pre=True) .with\_apt\_packages("pandoc", "texlive-xetex") .with\_source\_file(Path("prompts.yaml"), "/root"), resources=flyte.Resources(cpu=1), ) # {{/docs-fragment env}} # {{docs-fragment generate\_research\_queries}} @env.task async def generate\_research\_queries( topic: str, planning\_model: str, json\_model: str, prompts\_file: File, ) -> list\[str\]: async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) PLANNING\_PROMPT = prompts\["planning\_prompt"\] plan = "" logging.info(f"\\n\\nGenerated deep research plan for topic: {topic}\\n\\nPlan:") async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=PLANNING\_PROMPT, message=f"Research Topic: {topic}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): plan += chunk print(chunk, end="", flush=True) SEARCH\_PROMPT = prompts\["plan\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=SEARCH\_PROMPT, message=f"Plan to be parsed: {plan}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk plan = json.loads(response\_json) return plan\["queries"\] # {{/docs-fragment generate\_research\_queries}} async def \_summarize\_content\_async( raw\_content: str, query: str, prompt: str, summarization\_model: str, ) -> str: """Summarize content asynchronously using the LLM""" logging.info("Summarizing content asynchronously using the LLM") result = "" async for chunk in asingle\_shot\_llm\_call( model=summarization\_model, system\_prompt=prompt, message=f"{raw\_content}\\n\\n{query}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): result += chunk return result # {{docs-fragment search\_and\_summarize}} @env.task async def search\_and\_summarize( query: str, prompts\_file: File, summarization\_model: str, ) -> DeepResearchResults: """Perform search for a single query""" if len(query) > 400: # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues query = query\[:400\] logging.info(f"Truncated query to 400 characters: {query}") response = await atavily\_search\_results(query) logging.info("Tavily Search Called.") async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) RAW\_CONTENT\_SUMMARIZER\_PROMPT = prompts\["raw\_content\_summarizer\_prompt"\] with flyte.group("summarize-content"): # Create tasks for summarization summarization\_tasks = \[\] result\_info = \[\] for result in response.results: if result.raw\_content is None: continue task = \_summarize\_content\_async( result.raw\_content, query, RAW\_CONTENT\_SUMMARIZER\_PROMPT, summarization\_model, ) summarization\_tasks.append(task) result\_info.append(result) # Use return\_exceptions=True to prevent exceptions from propagating summarized\_contents = await asyncio.gather( \*summarization\_tasks, return\_exceptions=True ) # Filter out exceptions summarized\_contents = \[\ result for result in summarized\_contents if not isinstance(result, Exception)\ \] formatted\_results = \[\] for result, summarized\_content in zip(result\_info, summarized\_contents): formatted\_results.append( DeepResearchResult( title=result.title, link=result.link, content=result.content, raw\_content=result.raw\_content, filtered\_raw\_content=summarized\_content, ) ) return DeepResearchResults(results=formatted\_results) # {{/docs-fragment search\_and\_summarize}} @env.task async def search\_all\_queries( queries: list\[str\], summarization\_model: str, prompts\_file: File ) -> DeepResearchResults: """Execute searches for all queries in parallel""" tasks = \[\] results\_list = \[\] tasks = \[\ search\_and\_summarize(query, prompts\_file, summarization\_model)\ for query in queries\ \] if tasks: res\_list = await asyncio.gather(\*tasks) results\_list.extend(res\_list) # Combine all results combined\_results = DeepResearchResults(results=\[\]) for results in results\_list: combined\_results = combined\_results + results return combined\_results # {{docs-fragment evaluate\_research\_completeness}} @env.task async def evaluate\_research\_completeness( topic: str, results: DeepResearchResults, queries: list\[str\], prompts\_file: File, planning\_model: str, json\_model: str, ) -> list\[str\]: """ Evaluate if the current search results are sufficient or if more research is needed. Returns an empty list if research is complete, or a list of additional queries if more research is needed. """ # Format the search results for the LLM formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) EVALUATION\_PROMPT = prompts\["evaluation\_prompt"\] logging.info("\\nEvaluation: ") evaluation = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=EVALUATION\_PROMPT, message=( f"{topic}\\n\\n" f"{queries}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=None, ): evaluation += chunk print(chunk, end="", flush=True) EVALUATION\_PARSING\_PROMPT = prompts\["evaluation\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=EVALUATION\_PARSING\_PROMPT, message=f"Evaluation to be parsed: {evaluation}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk evaluation = json.loads(response\_json) return evaluation\["queries"\] # {{/docs-fragment evaluate\_research\_completeness}} # {{docs-fragment filter\_results}} @env.task async def filter\_results( topic: str, results: DeepResearchResults, prompts\_file: File, planning\_model: str, json\_model: str, max\_sources: int, ) -> DeepResearchResults: """Filter the search results based on the research plan""" # Format the search results for the LLM, without the raw content formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) FILTER\_PROMPT = prompts\["filter\_prompt"\] logging.info("\\nFilter response: ") filter\_response = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=FILTER\_PROMPT, message=( f"{topic}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): filter\_response += chunk print(chunk, end="", flush=True) logging.info(f"Filter response: {filter\_response}") FILTER\_PARSING\_PROMPT = prompts\["filter\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=FILTER\_PARSING\_PROMPT, message=f"Filter response to be parsed: {filter\_response}", response\_format={ "type": "json\_object", "schema": SourceList.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk sources = json.loads(response\_json)\["sources"\] logging.info(f"Filtered sources: {sources}") if max\_sources != -1: sources = sources\[:max\_sources\] # Filter the results based on the source list filtered\_results = \[\ results.results\[i - 1\] for i in sources if i - 1 < len(results.results)\ \] return DeepResearchResults(results=filtered\_results) # {{/docs-fragment filter\_results}} def \_remove\_thinking\_tags(answer: str) -> str: """Remove content within tags""" while "" in answer and "" in answer: start = answer.find("") end = answer.find("") + len("") answer = answer\[:start\] + answer\[end:\] return answer # {{docs-fragment generate\_research\_answer}} @env.task async def generate\_research\_answer( topic: str, results: DeepResearchResults, remove\_thinking\_tags: bool, prompts\_file: File, answer\_model: str, ) -> str: """ Generate a comprehensive answer to the research topic based on the search results. Returns a detailed response that synthesizes information from all search results. """ formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) ANSWER\_PROMPT = prompts\["answer\_prompt"\] answer = "" async for chunk in asingle\_shot\_llm\_call( model=answer\_model, system\_prompt=ANSWER\_PROMPT, message=f"Research Topic: {topic}\\n\\nSearch Results:\\n{formatted\_results}", response\_format=None, # NOTE: This is the max\_token parameter for the LLM call on Together AI, # may need to be changed for other providers max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): answer += chunk # this is just to avoid typing complaints if answer is None or not isinstance(answer, str): logging.error("No answer generated") return "No answer generated" if remove\_thinking\_tags: # Remove content within tags answer = \_remove\_thinking\_tags(answer) # Remove markdown code block markers if they exist at the beginning if answer.lstrip().startswith("\`\`\`"): # Find the first line break after the opening backticks first\_linebreak = answer.find("\\n", answer.find("\`\`\`")) if first\_linebreak != -1: # Remove everything up to and including the first line break answer = answer\[first\_linebreak + 1 :\] # Remove closing code block if it exists if answer.rstrip().endswith("\`\`\`"): answer = answer.rstrip()\[:-3\].rstrip() return answer.strip() # {{/docs-fragment generate\_research\_answer}} # {{docs-fragment research\_topic}} @env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff\_factor=2)) async def research\_topic( topic: str, budget: int = 3, remove\_thinking\_tags: bool = True, max\_queries: int = 5, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 40, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", prompts\_file: File | str = "prompts.yaml", ) -> str: """Main method to conduct research on a topic. Will be used for weave evals.""" if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) # Step 1: Generate initial queries queries = await generate\_research\_queries( topic=topic, planning\_model=planning\_model, json\_model=json\_model, prompts\_file=prompts\_file, ) queries = \[topic, \*queries\[: max\_queries - 1\]\] all\_queries = queries.copy() logging.info(f"Initial queries: {queries}") if len(queries) == 0: logging.error("No initial queries generated") return "No initial queries generated" # Step 2: Perform initial search results = await search\_all\_queries(queries, summarization\_model, prompts\_file) logging.info(f"Initial search complete, found {len(results.results)} results") # Step 3: Conduct iterative research within budget for iteration in range(budget): with flyte.group(f"eval\_iteration\_{iteration}"): # Evaluate if more research is needed additional\_queries = await evaluate\_research\_completeness( topic=topic, results=results, queries=all\_queries, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, ) # Filter out empty strings and check if any queries remain additional\_queries = \[q for q in additional\_queries if q\] if not additional\_queries: logging.info("No need for additional research") break # for debugging purposes we limit the number of queries additional\_queries = additional\_queries\[:max\_queries\] logging.info(f"Additional queries: {additional\_queries}") # Expand research with new queries new\_results = await search\_all\_queries( additional\_queries, summarization\_model, prompts\_file ) logging.info( f"Follow-up search complete, found {len(new\_results.results)} results" ) results = results + new\_results all\_queries.extend(additional\_queries) # Step 4: Generate final answer logging.info(f"Generating final answer for topic: {topic}") results = results.dedup() logging.info(f"Deduplication complete, kept {len(results.results)} results") filtered\_results = await filter\_results( topic=topic, results=results, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, ) logging.info( f"LLM Filtering complete, kept {len(filtered\_results.results)} results" ) # Generate final answer answer = await generate\_research\_answer( topic=topic, results=filtered\_results, remove\_thinking\_tags=remove\_thinking\_tags, prompts\_file=prompts\_file, answer\_model=answer\_model, ) return answer # {{/docs-fragment research\_topic}} # {{docs-fragment main}} @env.task(report=True) async def main( topic: str = ( "List the essential requirements for a developer-focused agent orchestration system." ), prompts\_file: File | str = "/root/prompts.yaml", budget: int = 2, remove\_thinking\_tags: bool = True, max\_queries: int = 3, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 10, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", ) -> str: if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) answer = await research\_topic( topic=topic, budget=budget, remove\_thinking\_tags=remove\_thinking\_tags, max\_queries=max\_queries, answer\_model=answer\_model, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, summarization\_model=summarization\_model, prompts\_file=prompts\_file, ) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") toc\_image\_url = await generate\_toc\_image( yaml.safe\_load(yaml\_contents)\["data\_visualization\_prompt"\], planning\_model, topic, ) html\_content = await generate\_html(answer, toc\_image\_url) await flyte.report.replace.aio(html\_content, do\_flush=True) await flyte.report.flush.aio() return html\_content # {{/docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep\_research\_agent/agent.py\* ## Generate the final answer Finally, we generate a detailed research report by synthesizing the top-ranked results. This is the output returned to the user. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pydantic==2.11.5",\ # "litellm==1.72.2",\ # "tavily-python==0.7.5",\ # "together==1.5.24",\ # "markdown==3.8.2",\ # "pymdown-extensions==10.16.1",\ # \] # main = "main" # params = "" # /// # {{docs-fragment env}} import asyncio import json from pathlib import Path import flyte import yaml from flyte.io.\_file import File from libs.utils.data\_types import ( DeepResearchResult, DeepResearchResults, ResearchPlan, SourceList, ) from libs.utils.generation import generate\_html, generate\_toc\_image from libs.utils.llms import asingle\_shot\_llm\_call from libs.utils.log import AgentLogger from libs.utils.tavily\_search import atavily\_search\_results TIME\_LIMIT\_MULTIPLIER = 5 MAX\_COMPLETION\_TOKENS = 4096 logging = AgentLogger("together.open\_deep\_research") env = flyte.TaskEnvironment( name="deep-researcher", secrets=\[\ flyte.Secret(key="together\_api\_key", as\_env\_var="TOGETHER\_API\_KEY"),\ flyte.Secret(key="tavily\_api\_key", as\_env\_var="TAVILY\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="deep-research-agent", pre=True) .with\_apt\_packages("pandoc", "texlive-xetex") .with\_source\_file(Path("prompts.yaml"), "/root"), resources=flyte.Resources(cpu=1), ) # {{/docs-fragment env}} # {{docs-fragment generate\_research\_queries}} @env.task async def generate\_research\_queries( topic: str, planning\_model: str, json\_model: str, prompts\_file: File, ) -> list\[str\]: async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) PLANNING\_PROMPT = prompts\["planning\_prompt"\] plan = "" logging.info(f"\\n\\nGenerated deep research plan for topic: {topic}\\n\\nPlan:") async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=PLANNING\_PROMPT, message=f"Research Topic: {topic}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): plan += chunk print(chunk, end="", flush=True) SEARCH\_PROMPT = prompts\["plan\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=SEARCH\_PROMPT, message=f"Plan to be parsed: {plan}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk plan = json.loads(response\_json) return plan\["queries"\] # {{/docs-fragment generate\_research\_queries}} async def \_summarize\_content\_async( raw\_content: str, query: str, prompt: str, summarization\_model: str, ) -> str: """Summarize content asynchronously using the LLM""" logging.info("Summarizing content asynchronously using the LLM") result = "" async for chunk in asingle\_shot\_llm\_call( model=summarization\_model, system\_prompt=prompt, message=f"{raw\_content}\\n\\n{query}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): result += chunk return result # {{docs-fragment search\_and\_summarize}} @env.task async def search\_and\_summarize( query: str, prompts\_file: File, summarization\_model: str, ) -> DeepResearchResults: """Perform search for a single query""" if len(query) > 400: # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues query = query\[:400\] logging.info(f"Truncated query to 400 characters: {query}") response = await atavily\_search\_results(query) logging.info("Tavily Search Called.") async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) RAW\_CONTENT\_SUMMARIZER\_PROMPT = prompts\["raw\_content\_summarizer\_prompt"\] with flyte.group("summarize-content"): # Create tasks for summarization summarization\_tasks = \[\] result\_info = \[\] for result in response.results: if result.raw\_content is None: continue task = \_summarize\_content\_async( result.raw\_content, query, RAW\_CONTENT\_SUMMARIZER\_PROMPT, summarization\_model, ) summarization\_tasks.append(task) result\_info.append(result) # Use return\_exceptions=True to prevent exceptions from propagating summarized\_contents = await asyncio.gather( \*summarization\_tasks, return\_exceptions=True ) # Filter out exceptions summarized\_contents = \[\ result for result in summarized\_contents if not isinstance(result, Exception)\ \] formatted\_results = \[\] for result, summarized\_content in zip(result\_info, summarized\_contents): formatted\_results.append( DeepResearchResult( title=result.title, link=result.link, content=result.content, raw\_content=result.raw\_content, filtered\_raw\_content=summarized\_content, ) ) return DeepResearchResults(results=formatted\_results) # {{/docs-fragment search\_and\_summarize}} @env.task async def search\_all\_queries( queries: list\[str\], summarization\_model: str, prompts\_file: File ) -> DeepResearchResults: """Execute searches for all queries in parallel""" tasks = \[\] results\_list = \[\] tasks = \[\ search\_and\_summarize(query, prompts\_file, summarization\_model)\ for query in queries\ \] if tasks: res\_list = await asyncio.gather(\*tasks) results\_list.extend(res\_list) # Combine all results combined\_results = DeepResearchResults(results=\[\]) for results in results\_list: combined\_results = combined\_results + results return combined\_results # {{docs-fragment evaluate\_research\_completeness}} @env.task async def evaluate\_research\_completeness( topic: str, results: DeepResearchResults, queries: list\[str\], prompts\_file: File, planning\_model: str, json\_model: str, ) -> list\[str\]: """ Evaluate if the current search results are sufficient or if more research is needed. Returns an empty list if research is complete, or a list of additional queries if more research is needed. """ # Format the search results for the LLM formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) EVALUATION\_PROMPT = prompts\["evaluation\_prompt"\] logging.info("\\nEvaluation: ") evaluation = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=EVALUATION\_PROMPT, message=( f"{topic}\\n\\n" f"{queries}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=None, ): evaluation += chunk print(chunk, end="", flush=True) EVALUATION\_PARSING\_PROMPT = prompts\["evaluation\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=EVALUATION\_PARSING\_PROMPT, message=f"Evaluation to be parsed: {evaluation}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk evaluation = json.loads(response\_json) return evaluation\["queries"\] # {{/docs-fragment evaluate\_research\_completeness}} # {{docs-fragment filter\_results}} @env.task async def filter\_results( topic: str, results: DeepResearchResults, prompts\_file: File, planning\_model: str, json\_model: str, max\_sources: int, ) -> DeepResearchResults: """Filter the search results based on the research plan""" # Format the search results for the LLM, without the raw content formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) FILTER\_PROMPT = prompts\["filter\_prompt"\] logging.info("\\nFilter response: ") filter\_response = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=FILTER\_PROMPT, message=( f"{topic}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): filter\_response += chunk print(chunk, end="", flush=True) logging.info(f"Filter response: {filter\_response}") FILTER\_PARSING\_PROMPT = prompts\["filter\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=FILTER\_PARSING\_PROMPT, message=f"Filter response to be parsed: {filter\_response}", response\_format={ "type": "json\_object", "schema": SourceList.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk sources = json.loads(response\_json)\["sources"\] logging.info(f"Filtered sources: {sources}") if max\_sources != -1: sources = sources\[:max\_sources\] # Filter the results based on the source list filtered\_results = \[\ results.results\[i - 1\] for i in sources if i - 1 < len(results.results)\ \] return DeepResearchResults(results=filtered\_results) # {{/docs-fragment filter\_results}} def \_remove\_thinking\_tags(answer: str) -> str: """Remove content within tags""" while "" in answer and "" in answer: start = answer.find("") end = answer.find("") + len("") answer = answer\[:start\] + answer\[end:\] return answer # {{docs-fragment generate\_research\_answer}} @env.task async def generate\_research\_answer( topic: str, results: DeepResearchResults, remove\_thinking\_tags: bool, prompts\_file: File, answer\_model: str, ) -> str: """ Generate a comprehensive answer to the research topic based on the search results. Returns a detailed response that synthesizes information from all search results. """ formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) ANSWER\_PROMPT = prompts\["answer\_prompt"\] answer = "" async for chunk in asingle\_shot\_llm\_call( model=answer\_model, system\_prompt=ANSWER\_PROMPT, message=f"Research Topic: {topic}\\n\\nSearch Results:\\n{formatted\_results}", response\_format=None, # NOTE: This is the max\_token parameter for the LLM call on Together AI, # may need to be changed for other providers max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): answer += chunk # this is just to avoid typing complaints if answer is None or not isinstance(answer, str): logging.error("No answer generated") return "No answer generated" if remove\_thinking\_tags: # Remove content within tags answer = \_remove\_thinking\_tags(answer) # Remove markdown code block markers if they exist at the beginning if answer.lstrip().startswith("\`\`\`"): # Find the first line break after the opening backticks first\_linebreak = answer.find("\\n", answer.find("\`\`\`")) if first\_linebreak != -1: # Remove everything up to and including the first line break answer = answer\[first\_linebreak + 1 :\] # Remove closing code block if it exists if answer.rstrip().endswith("\`\`\`"): answer = answer.rstrip()\[:-3\].rstrip() return answer.strip() # {{/docs-fragment generate\_research\_answer}} # {{docs-fragment research\_topic}} @env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff\_factor=2)) async def research\_topic( topic: str, budget: int = 3, remove\_thinking\_tags: bool = True, max\_queries: int = 5, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 40, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", prompts\_file: File | str = "prompts.yaml", ) -> str: """Main method to conduct research on a topic. Will be used for weave evals.""" if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) # Step 1: Generate initial queries queries = await generate\_research\_queries( topic=topic, planning\_model=planning\_model, json\_model=json\_model, prompts\_file=prompts\_file, ) queries = \[topic, \*queries\[: max\_queries - 1\]\] all\_queries = queries.copy() logging.info(f"Initial queries: {queries}") if len(queries) == 0: logging.error("No initial queries generated") return "No initial queries generated" # Step 2: Perform initial search results = await search\_all\_queries(queries, summarization\_model, prompts\_file) logging.info(f"Initial search complete, found {len(results.results)} results") # Step 3: Conduct iterative research within budget for iteration in range(budget): with flyte.group(f"eval\_iteration\_{iteration}"): # Evaluate if more research is needed additional\_queries = await evaluate\_research\_completeness( topic=topic, results=results, queries=all\_queries, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, ) # Filter out empty strings and check if any queries remain additional\_queries = \[q for q in additional\_queries if q\] if not additional\_queries: logging.info("No need for additional research") break # for debugging purposes we limit the number of queries additional\_queries = additional\_queries\[:max\_queries\] logging.info(f"Additional queries: {additional\_queries}") # Expand research with new queries new\_results = await search\_all\_queries( additional\_queries, summarization\_model, prompts\_file ) logging.info( f"Follow-up search complete, found {len(new\_results.results)} results" ) results = results + new\_results all\_queries.extend(additional\_queries) # Step 4: Generate final answer logging.info(f"Generating final answer for topic: {topic}") results = results.dedup() logging.info(f"Deduplication complete, kept {len(results.results)} results") filtered\_results = await filter\_results( topic=topic, results=results, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, ) logging.info( f"LLM Filtering complete, kept {len(filtered\_results.results)} results" ) # Generate final answer answer = await generate\_research\_answer( topic=topic, results=filtered\_results, remove\_thinking\_tags=remove\_thinking\_tags, prompts\_file=prompts\_file, answer\_model=answer\_model, ) return answer # {{/docs-fragment research\_topic}} # {{docs-fragment main}} @env.task(report=True) async def main( topic: str = ( "List the essential requirements for a developer-focused agent orchestration system." ), prompts\_file: File | str = "/root/prompts.yaml", budget: int = 2, remove\_thinking\_tags: bool = True, max\_queries: int = 3, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 10, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", ) -> str: if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) answer = await research\_topic( topic=topic, budget=budget, remove\_thinking\_tags=remove\_thinking\_tags, max\_queries=max\_queries, answer\_model=answer\_model, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, summarization\_model=summarization\_model, prompts\_file=prompts\_file, ) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") toc\_image\_url = await generate\_toc\_image( yaml.safe\_load(yaml\_contents)\["data\_visualization\_prompt"\], planning\_model, topic, ) html\_content = await generate\_html(answer, toc\_image\_url) await flyte.report.replace.aio(html\_content, do\_flush=True) await flyte.report.flush.aio() return html\_content # {{/docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep\_research\_agent/agent.py\* ## Orchestration Next, we define a \`research\_topic\` task to orchestrate the entire deep research workflow. It runs the core stages in sequence: generating research queries, performing search and summarization, evaluating the completeness of results, and producing the final report. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pydantic==2.11.5",\ # "litellm==1.72.2",\ # "tavily-python==0.7.5",\ # "together==1.5.24",\ # "markdown==3.8.2",\ # "pymdown-extensions==10.16.1",\ # \] # main = "main" # params = "" # /// # {{docs-fragment env}} import asyncio import json from pathlib import Path import flyte import yaml from flyte.io.\_file import File from libs.utils.data\_types import ( DeepResearchResult, DeepResearchResults, ResearchPlan, SourceList, ) from libs.utils.generation import generate\_html, generate\_toc\_image from libs.utils.llms import asingle\_shot\_llm\_call from libs.utils.log import AgentLogger from libs.utils.tavily\_search import atavily\_search\_results TIME\_LIMIT\_MULTIPLIER = 5 MAX\_COMPLETION\_TOKENS = 4096 logging = AgentLogger("together.open\_deep\_research") env = flyte.TaskEnvironment( name="deep-researcher", secrets=\[\ flyte.Secret(key="together\_api\_key", as\_env\_var="TOGETHER\_API\_KEY"),\ flyte.Secret(key="tavily\_api\_key", as\_env\_var="TAVILY\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="deep-research-agent", pre=True) .with\_apt\_packages("pandoc", "texlive-xetex") .with\_source\_file(Path("prompts.yaml"), "/root"), resources=flyte.Resources(cpu=1), ) # {{/docs-fragment env}} # {{docs-fragment generate\_research\_queries}} @env.task async def generate\_research\_queries( topic: str, planning\_model: str, json\_model: str, prompts\_file: File, ) -> list\[str\]: async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) PLANNING\_PROMPT = prompts\["planning\_prompt"\] plan = "" logging.info(f"\\n\\nGenerated deep research plan for topic: {topic}\\n\\nPlan:") async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=PLANNING\_PROMPT, message=f"Research Topic: {topic}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): plan += chunk print(chunk, end="", flush=True) SEARCH\_PROMPT = prompts\["plan\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=SEARCH\_PROMPT, message=f"Plan to be parsed: {plan}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk plan = json.loads(response\_json) return plan\["queries"\] # {{/docs-fragment generate\_research\_queries}} async def \_summarize\_content\_async( raw\_content: str, query: str, prompt: str, summarization\_model: str, ) -> str: """Summarize content asynchronously using the LLM""" logging.info("Summarizing content asynchronously using the LLM") result = "" async for chunk in asingle\_shot\_llm\_call( model=summarization\_model, system\_prompt=prompt, message=f"{raw\_content}\\n\\n{query}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): result += chunk return result # {{docs-fragment search\_and\_summarize}} @env.task async def search\_and\_summarize( query: str, prompts\_file: File, summarization\_model: str, ) -> DeepResearchResults: """Perform search for a single query""" if len(query) > 400: # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues query = query\[:400\] logging.info(f"Truncated query to 400 characters: {query}") response = await atavily\_search\_results(query) logging.info("Tavily Search Called.") async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) RAW\_CONTENT\_SUMMARIZER\_PROMPT = prompts\["raw\_content\_summarizer\_prompt"\] with flyte.group("summarize-content"): # Create tasks for summarization summarization\_tasks = \[\] result\_info = \[\] for result in response.results: if result.raw\_content is None: continue task = \_summarize\_content\_async( result.raw\_content, query, RAW\_CONTENT\_SUMMARIZER\_PROMPT, summarization\_model, ) summarization\_tasks.append(task) result\_info.append(result) # Use return\_exceptions=True to prevent exceptions from propagating summarized\_contents = await asyncio.gather( \*summarization\_tasks, return\_exceptions=True ) # Filter out exceptions summarized\_contents = \[\ result for result in summarized\_contents if not isinstance(result, Exception)\ \] formatted\_results = \[\] for result, summarized\_content in zip(result\_info, summarized\_contents): formatted\_results.append( DeepResearchResult( title=result.title, link=result.link, content=result.content, raw\_content=result.raw\_content, filtered\_raw\_content=summarized\_content, ) ) return DeepResearchResults(results=formatted\_results) # {{/docs-fragment search\_and\_summarize}} @env.task async def search\_all\_queries( queries: list\[str\], summarization\_model: str, prompts\_file: File ) -> DeepResearchResults: """Execute searches for all queries in parallel""" tasks = \[\] results\_list = \[\] tasks = \[\ search\_and\_summarize(query, prompts\_file, summarization\_model)\ for query in queries\ \] if tasks: res\_list = await asyncio.gather(\*tasks) results\_list.extend(res\_list) # Combine all results combined\_results = DeepResearchResults(results=\[\]) for results in results\_list: combined\_results = combined\_results + results return combined\_results # {{docs-fragment evaluate\_research\_completeness}} @env.task async def evaluate\_research\_completeness( topic: str, results: DeepResearchResults, queries: list\[str\], prompts\_file: File, planning\_model: str, json\_model: str, ) -> list\[str\]: """ Evaluate if the current search results are sufficient or if more research is needed. Returns an empty list if research is complete, or a list of additional queries if more research is needed. """ # Format the search results for the LLM formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) EVALUATION\_PROMPT = prompts\["evaluation\_prompt"\] logging.info("\\nEvaluation: ") evaluation = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=EVALUATION\_PROMPT, message=( f"{topic}\\n\\n" f"{queries}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=None, ): evaluation += chunk print(chunk, end="", flush=True) EVALUATION\_PARSING\_PROMPT = prompts\["evaluation\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=EVALUATION\_PARSING\_PROMPT, message=f"Evaluation to be parsed: {evaluation}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk evaluation = json.loads(response\_json) return evaluation\["queries"\] # {{/docs-fragment evaluate\_research\_completeness}} # {{docs-fragment filter\_results}} @env.task async def filter\_results( topic: str, results: DeepResearchResults, prompts\_file: File, planning\_model: str, json\_model: str, max\_sources: int, ) -> DeepResearchResults: """Filter the search results based on the research plan""" # Format the search results for the LLM, without the raw content formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) FILTER\_PROMPT = prompts\["filter\_prompt"\] logging.info("\\nFilter response: ") filter\_response = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=FILTER\_PROMPT, message=( f"{topic}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): filter\_response += chunk print(chunk, end="", flush=True) logging.info(f"Filter response: {filter\_response}") FILTER\_PARSING\_PROMPT = prompts\["filter\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=FILTER\_PARSING\_PROMPT, message=f"Filter response to be parsed: {filter\_response}", response\_format={ "type": "json\_object", "schema": SourceList.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk sources = json.loads(response\_json)\["sources"\] logging.info(f"Filtered sources: {sources}") if max\_sources != -1: sources = sources\[:max\_sources\] # Filter the results based on the source list filtered\_results = \[\ results.results\[i - 1\] for i in sources if i - 1 < len(results.results)\ \] return DeepResearchResults(results=filtered\_results) # {{/docs-fragment filter\_results}} def \_remove\_thinking\_tags(answer: str) -> str: """Remove content within tags""" while "" in answer and "" in answer: start = answer.find("") end = answer.find("") + len("") answer = answer\[:start\] + answer\[end:\] return answer # {{docs-fragment generate\_research\_answer}} @env.task async def generate\_research\_answer( topic: str, results: DeepResearchResults, remove\_thinking\_tags: bool, prompts\_file: File, answer\_model: str, ) -> str: """ Generate a comprehensive answer to the research topic based on the search results. Returns a detailed response that synthesizes information from all search results. """ formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) ANSWER\_PROMPT = prompts\["answer\_prompt"\] answer = "" async for chunk in asingle\_shot\_llm\_call( model=answer\_model, system\_prompt=ANSWER\_PROMPT, message=f"Research Topic: {topic}\\n\\nSearch Results:\\n{formatted\_results}", response\_format=None, # NOTE: This is the max\_token parameter for the LLM call on Together AI, # may need to be changed for other providers max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): answer += chunk # this is just to avoid typing complaints if answer is None or not isinstance(answer, str): logging.error("No answer generated") return "No answer generated" if remove\_thinking\_tags: # Remove content within tags answer = \_remove\_thinking\_tags(answer) # Remove markdown code block markers if they exist at the beginning if answer.lstrip().startswith("\`\`\`"): # Find the first line break after the opening backticks first\_linebreak = answer.find("\\n", answer.find("\`\`\`")) if first\_linebreak != -1: # Remove everything up to and including the first line break answer = answer\[first\_linebreak + 1 :\] # Remove closing code block if it exists if answer.rstrip().endswith("\`\`\`"): answer = answer.rstrip()\[:-3\].rstrip() return answer.strip() # {{/docs-fragment generate\_research\_answer}} # {{docs-fragment research\_topic}} @env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff\_factor=2)) async def research\_topic( topic: str, budget: int = 3, remove\_thinking\_tags: bool = True, max\_queries: int = 5, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 40, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", prompts\_file: File | str = "prompts.yaml", ) -> str: """Main method to conduct research on a topic. Will be used for weave evals.""" if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) # Step 1: Generate initial queries queries = await generate\_research\_queries( topic=topic, planning\_model=planning\_model, json\_model=json\_model, prompts\_file=prompts\_file, ) queries = \[topic, \*queries\[: max\_queries - 1\]\] all\_queries = queries.copy() logging.info(f"Initial queries: {queries}") if len(queries) == 0: logging.error("No initial queries generated") return "No initial queries generated" # Step 2: Perform initial search results = await search\_all\_queries(queries, summarization\_model, prompts\_file) logging.info(f"Initial search complete, found {len(results.results)} results") # Step 3: Conduct iterative research within budget for iteration in range(budget): with flyte.group(f"eval\_iteration\_{iteration}"): # Evaluate if more research is needed additional\_queries = await evaluate\_research\_completeness( topic=topic, results=results, queries=all\_queries, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, ) # Filter out empty strings and check if any queries remain additional\_queries = \[q for q in additional\_queries if q\] if not additional\_queries: logging.info("No need for additional research") break # for debugging purposes we limit the number of queries additional\_queries = additional\_queries\[:max\_queries\] logging.info(f"Additional queries: {additional\_queries}") # Expand research with new queries new\_results = await search\_all\_queries( additional\_queries, summarization\_model, prompts\_file ) logging.info( f"Follow-up search complete, found {len(new\_results.results)} results" ) results = results + new\_results all\_queries.extend(additional\_queries) # Step 4: Generate final answer logging.info(f"Generating final answer for topic: {topic}") results = results.dedup() logging.info(f"Deduplication complete, kept {len(results.results)} results") filtered\_results = await filter\_results( topic=topic, results=results, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, ) logging.info( f"LLM Filtering complete, kept {len(filtered\_results.results)} results" ) # Generate final answer answer = await generate\_research\_answer( topic=topic, results=filtered\_results, remove\_thinking\_tags=remove\_thinking\_tags, prompts\_file=prompts\_file, answer\_model=answer\_model, ) return answer # {{/docs-fragment research\_topic}} # {{docs-fragment main}} @env.task(report=True) async def main( topic: str = ( "List the essential requirements for a developer-focused agent orchestration system." ), prompts\_file: File | str = "/root/prompts.yaml", budget: int = 2, remove\_thinking\_tags: bool = True, max\_queries: int = 3, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 10, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", ) -> str: if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) answer = await research\_topic( topic=topic, budget=budget, remove\_thinking\_tags=remove\_thinking\_tags, max\_queries=max\_queries, answer\_model=answer\_model, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, summarization\_model=summarization\_model, prompts\_file=prompts\_file, ) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") toc\_image\_url = await generate\_toc\_image( yaml.safe\_load(yaml\_contents)\["data\_visualization\_prompt"\], planning\_model, topic, ) html\_content = await generate\_html(answer, toc\_image\_url) await flyte.report.replace.aio(html\_content, do\_flush=True) await flyte.report.flush.aio() return html\_content # {{/docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep\_research\_agent/agent.py\* The \`main\` task wraps this entire pipeline and adds report generation in HTML format as the final step. It also serves as the main entry point to the workflow, allowing us to pass in all configuration parameters, including which LLMs to use at each stage. This flexibility lets us mix and match models for planning, summarization, and final synthesis, helping us optimize for both cost and quality. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pydantic==2.11.5",\ # "litellm==1.72.2",\ # "tavily-python==0.7.5",\ # "together==1.5.24",\ # "markdown==3.8.2",\ # "pymdown-extensions==10.16.1",\ # \] # main = "main" # params = "" # /// # {{docs-fragment env}} import asyncio import json from pathlib import Path import flyte import yaml from flyte.io.\_file import File from libs.utils.data\_types import ( DeepResearchResult, DeepResearchResults, ResearchPlan, SourceList, ) from libs.utils.generation import generate\_html, generate\_toc\_image from libs.utils.llms import asingle\_shot\_llm\_call from libs.utils.log import AgentLogger from libs.utils.tavily\_search import atavily\_search\_results TIME\_LIMIT\_MULTIPLIER = 5 MAX\_COMPLETION\_TOKENS = 4096 logging = AgentLogger("together.open\_deep\_research") env = flyte.TaskEnvironment( name="deep-researcher", secrets=\[\ flyte.Secret(key="together\_api\_key", as\_env\_var="TOGETHER\_API\_KEY"),\ flyte.Secret(key="tavily\_api\_key", as\_env\_var="TAVILY\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="deep-research-agent", pre=True) .with\_apt\_packages("pandoc", "texlive-xetex") .with\_source\_file(Path("prompts.yaml"), "/root"), resources=flyte.Resources(cpu=1), ) # {{/docs-fragment env}} # {{docs-fragment generate\_research\_queries}} @env.task async def generate\_research\_queries( topic: str, planning\_model: str, json\_model: str, prompts\_file: File, ) -> list\[str\]: async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) PLANNING\_PROMPT = prompts\["planning\_prompt"\] plan = "" logging.info(f"\\n\\nGenerated deep research plan for topic: {topic}\\n\\nPlan:") async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=PLANNING\_PROMPT, message=f"Research Topic: {topic}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): plan += chunk print(chunk, end="", flush=True) SEARCH\_PROMPT = prompts\["plan\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=SEARCH\_PROMPT, message=f"Plan to be parsed: {plan}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk plan = json.loads(response\_json) return plan\["queries"\] # {{/docs-fragment generate\_research\_queries}} async def \_summarize\_content\_async( raw\_content: str, query: str, prompt: str, summarization\_model: str, ) -> str: """Summarize content asynchronously using the LLM""" logging.info("Summarizing content asynchronously using the LLM") result = "" async for chunk in asingle\_shot\_llm\_call( model=summarization\_model, system\_prompt=prompt, message=f"{raw\_content}\\n\\n{query}", response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): result += chunk return result # {{docs-fragment search\_and\_summarize}} @env.task async def search\_and\_summarize( query: str, prompts\_file: File, summarization\_model: str, ) -> DeepResearchResults: """Perform search for a single query""" if len(query) > 400: # NOTE: we are truncating the query to 400 characters to avoid Tavily Search issues query = query\[:400\] logging.info(f"Truncated query to 400 characters: {query}") response = await atavily\_search\_results(query) logging.info("Tavily Search Called.") async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) RAW\_CONTENT\_SUMMARIZER\_PROMPT = prompts\["raw\_content\_summarizer\_prompt"\] with flyte.group("summarize-content"): # Create tasks for summarization summarization\_tasks = \[\] result\_info = \[\] for result in response.results: if result.raw\_content is None: continue task = \_summarize\_content\_async( result.raw\_content, query, RAW\_CONTENT\_SUMMARIZER\_PROMPT, summarization\_model, ) summarization\_tasks.append(task) result\_info.append(result) # Use return\_exceptions=True to prevent exceptions from propagating summarized\_contents = await asyncio.gather( \*summarization\_tasks, return\_exceptions=True ) # Filter out exceptions summarized\_contents = \[\ result for result in summarized\_contents if not isinstance(result, Exception)\ \] formatted\_results = \[\] for result, summarized\_content in zip(result\_info, summarized\_contents): formatted\_results.append( DeepResearchResult( title=result.title, link=result.link, content=result.content, raw\_content=result.raw\_content, filtered\_raw\_content=summarized\_content, ) ) return DeepResearchResults(results=formatted\_results) # {{/docs-fragment search\_and\_summarize}} @env.task async def search\_all\_queries( queries: list\[str\], summarization\_model: str, prompts\_file: File ) -> DeepResearchResults: """Execute searches for all queries in parallel""" tasks = \[\] results\_list = \[\] tasks = \[\ search\_and\_summarize(query, prompts\_file, summarization\_model)\ for query in queries\ \] if tasks: res\_list = await asyncio.gather(\*tasks) results\_list.extend(res\_list) # Combine all results combined\_results = DeepResearchResults(results=\[\]) for results in results\_list: combined\_results = combined\_results + results return combined\_results # {{docs-fragment evaluate\_research\_completeness}} @env.task async def evaluate\_research\_completeness( topic: str, results: DeepResearchResults, queries: list\[str\], prompts\_file: File, planning\_model: str, json\_model: str, ) -> list\[str\]: """ Evaluate if the current search results are sufficient or if more research is needed. Returns an empty list if research is complete, or a list of additional queries if more research is needed. """ # Format the search results for the LLM formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) EVALUATION\_PROMPT = prompts\["evaluation\_prompt"\] logging.info("\\nEvaluation: ") evaluation = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=EVALUATION\_PROMPT, message=( f"{topic}\\n\\n" f"{queries}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=None, ): evaluation += chunk print(chunk, end="", flush=True) EVALUATION\_PARSING\_PROMPT = prompts\["evaluation\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=EVALUATION\_PARSING\_PROMPT, message=f"Evaluation to be parsed: {evaluation}", response\_format={ "type": "json\_object", "schema": ResearchPlan.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk evaluation = json.loads(response\_json) return evaluation\["queries"\] # {{/docs-fragment evaluate\_research\_completeness}} # {{docs-fragment filter\_results}} @env.task async def filter\_results( topic: str, results: DeepResearchResults, prompts\_file: File, planning\_model: str, json\_model: str, max\_sources: int, ) -> DeepResearchResults: """Filter the search results based on the research plan""" # Format the search results for the LLM, without the raw content formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) FILTER\_PROMPT = prompts\["filter\_prompt"\] logging.info("\\nFilter response: ") filter\_response = "" async for chunk in asingle\_shot\_llm\_call( model=planning\_model, system\_prompt=FILTER\_PROMPT, message=( f"{topic}\\n\\n" f"{formatted\_results}" ), response\_format=None, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): filter\_response += chunk print(chunk, end="", flush=True) logging.info(f"Filter response: {filter\_response}") FILTER\_PARSING\_PROMPT = prompts\["filter\_parsing\_prompt"\] response\_json = "" async for chunk in asingle\_shot\_llm\_call( model=json\_model, system\_prompt=FILTER\_PARSING\_PROMPT, message=f"Filter response to be parsed: {filter\_response}", response\_format={ "type": "json\_object", "schema": SourceList.model\_json\_schema(), }, max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): response\_json += chunk sources = json.loads(response\_json)\["sources"\] logging.info(f"Filtered sources: {sources}") if max\_sources != -1: sources = sources\[:max\_sources\] # Filter the results based on the source list filtered\_results = \[\ results.results\[i - 1\] for i in sources if i - 1 < len(results.results)\ \] return DeepResearchResults(results=filtered\_results) # {{/docs-fragment filter\_results}} def \_remove\_thinking\_tags(answer: str) -> str: """Remove content within tags""" while "" in answer and "" in answer: start = answer.find("") end = answer.find("") + len("") answer = answer\[:start\] + answer\[end:\] return answer # {{docs-fragment generate\_research\_answer}} @env.task async def generate\_research\_answer( topic: str, results: DeepResearchResults, remove\_thinking\_tags: bool, prompts\_file: File, answer\_model: str, ) -> str: """ Generate a comprehensive answer to the research topic based on the search results. Returns a detailed response that synthesizes information from all search results. """ formatted\_results = str(results) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") prompts = yaml.safe\_load(yaml\_contents) ANSWER\_PROMPT = prompts\["answer\_prompt"\] answer = "" async for chunk in asingle\_shot\_llm\_call( model=answer\_model, system\_prompt=ANSWER\_PROMPT, message=f"Research Topic: {topic}\\n\\nSearch Results:\\n{formatted\_results}", response\_format=None, # NOTE: This is the max\_token parameter for the LLM call on Together AI, # may need to be changed for other providers max\_completion\_tokens=MAX\_COMPLETION\_TOKENS, ): answer += chunk # this is just to avoid typing complaints if answer is None or not isinstance(answer, str): logging.error("No answer generated") return "No answer generated" if remove\_thinking\_tags: # Remove content within tags answer = \_remove\_thinking\_tags(answer) # Remove markdown code block markers if they exist at the beginning if answer.lstrip().startswith("\`\`\`"): # Find the first line break after the opening backticks first\_linebreak = answer.find("\\n", answer.find("\`\`\`")) if first\_linebreak != -1: # Remove everything up to and including the first line break answer = answer\[first\_linebreak + 1 :\] # Remove closing code block if it exists if answer.rstrip().endswith("\`\`\`"): answer = answer.rstrip()\[:-3\].rstrip() return answer.strip() # {{/docs-fragment generate\_research\_answer}} # {{docs-fragment research\_topic}} @env.task(retries=flyte.RetryStrategy(count=3, backoff=10, backoff\_factor=2)) async def research\_topic( topic: str, budget: int = 3, remove\_thinking\_tags: bool = True, max\_queries: int = 5, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 40, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", prompts\_file: File | str = "prompts.yaml", ) -> str: """Main method to conduct research on a topic. Will be used for weave evals.""" if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) # Step 1: Generate initial queries queries = await generate\_research\_queries( topic=topic, planning\_model=planning\_model, json\_model=json\_model, prompts\_file=prompts\_file, ) queries = \[topic, \*queries\[: max\_queries - 1\]\] all\_queries = queries.copy() logging.info(f"Initial queries: {queries}") if len(queries) == 0: logging.error("No initial queries generated") return "No initial queries generated" # Step 2: Perform initial search results = await search\_all\_queries(queries, summarization\_model, prompts\_file) logging.info(f"Initial search complete, found {len(results.results)} results") # Step 3: Conduct iterative research within budget for iteration in range(budget): with flyte.group(f"eval\_iteration\_{iteration}"): # Evaluate if more research is needed additional\_queries = await evaluate\_research\_completeness( topic=topic, results=results, queries=all\_queries, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, ) # Filter out empty strings and check if any queries remain additional\_queries = \[q for q in additional\_queries if q\] if not additional\_queries: logging.info("No need for additional research") break # for debugging purposes we limit the number of queries additional\_queries = additional\_queries\[:max\_queries\] logging.info(f"Additional queries: {additional\_queries}") # Expand research with new queries new\_results = await search\_all\_queries( additional\_queries, summarization\_model, prompts\_file ) logging.info( f"Follow-up search complete, found {len(new\_results.results)} results" ) results = results + new\_results all\_queries.extend(additional\_queries) # Step 4: Generate final answer logging.info(f"Generating final answer for topic: {topic}") results = results.dedup() logging.info(f"Deduplication complete, kept {len(results.results)} results") filtered\_results = await filter\_results( topic=topic, results=results, prompts\_file=prompts\_file, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, ) logging.info( f"LLM Filtering complete, kept {len(filtered\_results.results)} results" ) # Generate final answer answer = await generate\_research\_answer( topic=topic, results=filtered\_results, remove\_thinking\_tags=remove\_thinking\_tags, prompts\_file=prompts\_file, answer\_model=answer\_model, ) return answer # {{/docs-fragment research\_topic}} # {{docs-fragment main}} @env.task(report=True) async def main( topic: str = ( "List the essential requirements for a developer-focused agent orchestration system." ), prompts\_file: File | str = "/root/prompts.yaml", budget: int = 2, remove\_thinking\_tags: bool = True, max\_queries: int = 3, answer\_model: str = "together\_ai/deepseek-ai/DeepSeek-V3", planning\_model: str = "together\_ai/Qwen/Qwen2.5-72B-Instruct-Turbo", json\_model: str = "together\_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max\_sources: int = 10, summarization\_model: str = "together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", ) -> str: if isinstance(prompts\_file, str): prompts\_file = await File.from\_local(prompts\_file) answer = await research\_topic( topic=topic, budget=budget, remove\_thinking\_tags=remove\_thinking\_tags, max\_queries=max\_queries, answer\_model=answer\_model, planning\_model=planning\_model, json\_model=json\_model, max\_sources=max\_sources, summarization\_model=summarization\_model, prompts\_file=prompts\_file, ) async with prompts\_file.open() as fh: data = await fh.read() yaml\_contents = str(data, "utf-8") toc\_image\_url = await generate\_toc\_image( yaml.safe\_load(yaml\_contents)\["data\_visualization\_prompt"\], planning\_model, topic, ) html\_content = await generate\_html(answer, toc\_image\_url) await flyte.report.replace.aio(html\_content, do\_flush=True) await flyte.report.flush.aio() return html\_content # {{/docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep\_research\_agent/agent.py\* ## Run the deep research agent First, create the required secrets: \`\`\` flyte create secret TOGETHER\_API\_KEY <> flyte create secret TAVILY\_API\_KEY <> \`\`\` Run the agent: \`\`\` uv run agent.py \`\`\` If you want to test it locally first, run the following commands: \`\`\` brew install pandoc brew install basictex # restart your terminal after install export TOGETHER\_API\_KEY=<> export TAVILY\_API\_KEY=<> uv run agent.py \`\`\` ## Evaluate with Weights & Biases Weave We use W&B Weave to evaluate the full agent pipeline and analyze LLM-generated responses. The evaluation runs as a Flyte pipeline and uses an LLM-as-a-judge scorer to measure the quality of LLM-generated responses. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "weave==0.51.51",\ # "datasets==3.6.0",\ # "huggingface-hub==0.32.6",\ # "litellm==1.72.2",\ # "tavily-python==0.7.5",\ # \] # /// import os import weave from agent import research\_topic from datasets import load\_dataset from huggingface\_hub import login from libs.utils.log import AgentLogger from litellm import completion import flyte logging = AgentLogger() weave.init(project\_name="deep-researcher") env = flyte.TaskEnvironment(name="deep-researcher-eval") @weave.op def llm\_as\_a\_judge\_scoring(answer: str, output: str, question: str) -> bool: prompt = f""" Given the following question and answer, evaluate the answer against the correct answer: {question} {output} {answer} Note that the agent answer might be a long text containing a lot of information or it might be a short answer. You should read the entire text and think if the agent answers the question somewhere in the text. You should try to be flexible with the answer but careful. For example, answering with names instead of name and surname is fine. The important thing is that the answer of the agent either contains the correct answer or is equal to the correct answer. The agent answer is correct because I can read that .... 1 Otherwise, return The agent answer is incorrect because there is ... 0 """ messages = \[\ {\ "role": "system",\ "content": "You are an helpful assistant that returns a number between 0 and 1.",\ },\ {"role": "user", "content": prompt},\ \] answer = ( completion( model="together\_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", messages=messages, max\_tokens=1000, temperature=0.0, ) .choices\[0\] # type: ignore .message\["content"\] # type: ignore ) return bool(int(answer.split("")\[1\].split("")\[0\].strip())) def authenticate\_huggingface(): """Authenticate with Hugging Face Hub using token from environment variable.""" token = os.getenv("HUGGINGFACE\_TOKEN") if not token: raise ValueError( "HUGGINGFACE\_TOKEN environment variable not set. " "Please set it with your token from https://huggingface.co/settings/tokens" ) try: login(token=token) print("Successfully authenticated with Hugging Face Hub") except Exception as e: raise RuntimeError(f"Failed to authenticate with Hugging Face Hub: {e!s}") @env.task async def load\_questions( dataset\_names: list\[str\] | None = None, ) -> list\[dict\[str, str\]\]: """ Load questions from the specified Hugging Face dataset configurations. Args: dataset\_names: List of dataset configurations to load Options: "smolagents:simpleqa", "hotpotqa", "simpleqa", "together-search-bench" If None, all available configurations except hotpotqa will be loaded Returns: List of question-answer pairs """ if dataset\_names is None: dataset\_names = \["smolagents:simpleqa"\] all\_questions = \[\] # Authenticate with Hugging Face Hub (once and for all) authenticate\_huggingface() for dataset\_name in dataset\_names: print(f"Loading dataset: {dataset\_name}") try: if dataset\_name == "together-search-bench": # Load Together-Search-Bench dataset dataset\_path = "togethercomputer/together-search-bench" ds = load\_dataset(dataset\_path) if "test" in ds: split\_data = ds\["test"\] else: print(f"No 'test' split found in dataset at {dataset\_path}") continue for i in range(len(split\_data)): item = split\_data\[i\] question\_data = { "question": item\["question"\], "answer": item\["answer"\], "dataset": item.get("dataset", "together-search-bench"), } all\_questions.append(question\_data) print(f"Loaded {len(split\_data)} questions from together-search-bench dataset") continue elif dataset\_name == "hotpotqa": # Load HotpotQA dataset (using distractor version for validation) ds = load\_dataset("hotpotqa/hotpot\_qa", "distractor", trust\_remote\_code=True) split\_name = "validation" elif dataset\_name == "simpleqa": ds = load\_dataset("basicv8vc/SimpleQA") split\_name = "test" else: # Strip "smolagents:" prefix when loading the dataset actual\_dataset = dataset\_name.split(":")\[-1\] ds = load\_dataset("smolagents/benchmark-v1", actual\_dataset) split\_name = "test" except Exception as e: print(f"Failed to load dataset {dataset\_name}: {e!s}") continue # Skip this dataset if it fails to load print(f"Dataset structure for {dataset\_name}: {ds}") print(f"Available splits: {list(ds)}") split\_data = ds\[split\_name\] # type: ignore for i in range(len(split\_data)): item = split\_data\[i\] if dataset\_name == "hotpotqa": # we remove questions that are easy or medium (if any) just to reduce the number of questions if item\["level"\] != "hard": continue question\_data = { "question": item\["question"\], "answer": item\["answer"\], "dataset": dataset\_name, } elif dataset\_name == "simpleqa": # Handle SimpleQA dataset format question\_data = { "question": item\["problem"\], "answer": item\["answer"\], "dataset": dataset\_name, } else: question\_data = { "question": item\["question"\], "answer": item\["true\_answer"\], "dataset": dataset\_name, } all\_questions.append(question\_data) print(f"Loaded {len(all\_questions)} questions in total") return all\_questions @weave.op async def predict(question: str): return await research\_topic(topic=str(question)) @env.task async def main(datasets: list\[str\] = \["together-search-bench"\], limit: int | None = 1): questions = await load\_questions(datasets) if limit is not None: questions = questions\[:limit\] print(f"Limited to {len(questions)} question(s)") evaluation = weave.Evaluation(dataset=questions, scorers=\[llm\_as\_a\_judge\_scoring\]) await evaluation.evaluate(predict) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() flyte.with\_runcontext(raw\_data\_path="data").run(main) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/deep\_research\_agent/weave\_evals.py\* You can run this pipeline locally as follows: \`\`\` export HUGGINGFACE\_TOKEN=<> # https://huggingface.co/settings/tokens export WANDB\_API\_KEY=<> # https://wandb.ai/settings uv run weave\_evals.py \`\`\` The script will run all tasks in the pipeline and log the evaluation results to Weights & Biases. While you can also evaluate individual tasks, this script focuses on end-to-end evaluation of the end-to-end deep research workflow. !\[Weave evaluations\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/tutorials/deep-research/weave\_evals.png) === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/competitive-intelligence-agent === # Competitive intelligence agent > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/competitive\_intelligence\_agent). This example demonstrates how to build a continuous competitive and market intelligence agent on Flyte. The agent fans out across a list of competitors, pulls fresh, source-cited web and news results from the \[You.com Search API\](https://you.com/docs/search/overview), and uses \[Claude\](https://docs.anthropic.com/) via \[LiteLLM\](https://docs.litellm.ai/) to extract structured \*\*deltas\*\* — pricing changes, product launches, funding events, leadership moves, and more — into a knowledge-graph-ready table. You.com returns ranked web and news results with snippets and publication timestamps, giving the LLM attributable sources to cite. Flyte orchestrates the rest: - \*\*Fan-out parallelism\*\* across competitors with \`asyncio.gather\` - \*\*\`cache="auto"\`\*\* so converging parallel or repeat runs reuse prior You.com and LLM results when queries overlap - \*\*\`@flyte.trace\`\*\* on every You.com and LLM call for full prompt → query → source lineage - \*\*Flyte reports\*\* that render an HTML dashboard grouping deltas by competitor and category !\[Competitive intelligence agent report\](https://www.union.ai/docs/v2/flyte/\_static/images/tutorials/competitive\_intelligence\_agent/competitive-intelligence-agent.png) ## Setting up the environment The agent runs in a single \`TaskEnvironment\` with secrets for the You.com and Anthropic API keys, automatic caching, and a container image built from the \`uv\` script dependencies. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "competitive\_intelligence" # params = "" # /// """Continuous competitive & market intelligence agent. A Dragonfly-style agent that fans out across competitors, pulls fresh, source-cited web + news results from the You.com Search API, and uses Claude to extract structured "deltas" (pricing, features, funding, leadership, etc.) into a knowledge-graph-ready table. """ # {{docs-fragment env}} import asyncio import json from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="competitive-intelligence", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="competitive-intelligence", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class SearchHit: """A You.com Search result with its full structured metadata.""" title: str url: str domain: str snippet: str published: str # You.com page\_age timestamp author: str favicon: str # You.com favicon\_url thumbnail: str section: str # "news" or "web" — You.com's auto classification @dataclass class Delta: competitor: str category: str summary: str confidence: float source: SearchHit | None = None @dataclass class CompetitorWatch: competitor: str deltas: list\[Delta\] = field(default\_factory=list) sources: list\[SearchHit\] = field(default\_factory=list) @dataclass class IntelReport: watches: list\[CompetitorWatch\] = field(default\_factory=list) @property def deltas(self) -> list\[Delta\]: return \[d for w in self.watches for d in w.deltas\] # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import os import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) @flyte.trace async def you\_search(query: str, count: int = 8, freshness: str = "week") -> list\[SearchHit\]: """Call the You.com Search API and return unified web + news hits.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), thumbnail=item.get("thumbnail\_url", "") or "", section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict | list: """Call Claude via LiteLLM and parse a JSON response.""" from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=2048, ) content = resp.choices\[0\].message.content return \_parse\_json(content) def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min( (i for i in (text.find("{"), text.find("\[")) if i != -1),\ default=0,\ )\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} EXTRACT\_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \\ search results about a competitor, extract concrete, recently-changed signals \\ ("deltas") in the requested categories. Only report changes that are supported \\ by a specific search result. Respond with a JSON object of the form: {"deltas": \[{"category": str, "summary": str, "source\_index": int (the \[n\] of \\\ the supporting search result), "confidence": float between 0 and 1}\]} If there are no clear changes, return {"deltas": \[\]}.""" # {{docs-fragment watch\_competitor}} @env.task(retries=3) async def watch\_competitor( competitor: str, categories: list\[str\], freshness: str, ) -> CompetitorWatch: """Search for fresh signals on one competitor and extract structured deltas.""" query = ( f"{competitor} " + " OR ".join(categories) + " announcement OR news OR update" ) hits = await you\_search(query, count=8, freshness=freshness) if not hits: return CompetitorWatch(competitor=competitor) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Competitor: {competitor}\\n" f"Categories to watch: {', '.join(categories)}\\n\\n" f"Search results:\\n{evidence}" ) parsed = await llm\_json(EXTRACT\_SYSTEM, user) raw\_deltas = parsed.get("deltas", \[\]) if isinstance(parsed, dict) else \[\] deltas: list\[Delta\] = \[\] cited: list\[SearchHit\] = \[\] for d in raw\_deltas: idx = int(d.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None if src is not None and src not in cited: cited.append(src) deltas.append( Delta( competitor=competitor, category=str(d.get("category", "unknown")), summary=str(d.get("summary", "")), confidence=float(d.get("confidence", 0.0) or 0.0), source=src, ) ) return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited) # {{/docs-fragment watch\_competitor}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_conf\_bar(conf: float) -> str: pct = max(0, min(100, int(conf \* 100))) return ( f"" f"{conf:.0%} confidence" ) def \_cite(src: SearchHit) -> str: """Render a rich You.com citation: favicon, domain, date, author, snippet.""" if src is None: return "" tag = ( f"news" if src.section == "news" else "web" ) meta\_bits = \[\] if src.published: meta\_bits.append(src.published\[:10\]) if src.author: meta\_bits.append(f"by {src.author}") meta = " · ".join(meta\_bits) snip = f"
“{src.snippet}”
" if src.snippet else "" return ( f"
" f"" f"
" f"{src.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: IntelReport) -> str: watches = sorted(report.watches, key=lambda w: w.competitor) total\_sources = sum(len(w.sources) for w in watches) cards = \[\] for w in watches: deltas = sorted(w.deltas, key=lambda d: -d.confidence) rows = "".join( f"
{d.category}" f"
{d.summary}
" f"{\_conf\_bar(d.confidence)}" f"{\_cite(d.source)}" "
" for d in deltas ) cards.append( f"

{w.competitor}

" f"{len(deltas)} signal(s) · " f"{len(w.sources)} You.com source(s){rows or ''}
" ) return f""" {REPORT\_CSS}

Competitive Intelligence Deltas

Fresh, source-cited market signals — every delta links back to a ranked, timestamped You.com Search result.

{len(report.deltas)} signals {len(watches)} competitors tracked {total\_sources} cited You.com sources
{''.join(cards) or "

No signals detected in this window.

"}

Sources retrieved and ranked by the You.com Search API (web + auto-classified news), with publication timestamps, authors, and snippet provenance preserved for full prompt → citation lineage.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def competitive\_intelligence( competitors: list\[str\] = \[\ "Anthropic",\ "OpenAI",\ "Mistral AI",\ "Google DeepMind",\ "Cohere",\ "Perplexity AI",\ "xAI",\ "Hugging Face",\ "Databricks",\ "Together AI",\ \], categories: list\[str\] = \[\ "pricing",\ "product launch",\ "model release",\ "funding",\ "leadership",\ "partnership",\ \], freshness: str = "week", ) -> IntelReport: """Fan out across competitors and aggregate structured deltas.""" with flyte.group("watch-competitors"): results = await asyncio.gather( \*\[watch\_competitor(c, categories, freshness) for c in competitors\] ) report = IntelReport(watches=list(results)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(competitive\_intelligence) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive\_intelligence\_agent/main.py\* The Python packages are declared at the top of the file using the \`uv\` script style: \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # /// \`\`\` ## Data types The agent models search hits, deltas, and the final report as dataclasses. Each \`Delta\` links back to a \`SearchHit\` that preserves You.com metadata — domain, publication date, author, and snippet. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "competitive\_intelligence" # params = "" # /// """Continuous competitive & market intelligence agent. A Dragonfly-style agent that fans out across competitors, pulls fresh, source-cited web + news results from the You.com Search API, and uses Claude to extract structured "deltas" (pricing, features, funding, leadership, etc.) into a knowledge-graph-ready table. """ # {{docs-fragment env}} import asyncio import json from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="competitive-intelligence", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="competitive-intelligence", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class SearchHit: """A You.com Search result with its full structured metadata.""" title: str url: str domain: str snippet: str published: str # You.com page\_age timestamp author: str favicon: str # You.com favicon\_url thumbnail: str section: str # "news" or "web" — You.com's auto classification @dataclass class Delta: competitor: str category: str summary: str confidence: float source: SearchHit | None = None @dataclass class CompetitorWatch: competitor: str deltas: list\[Delta\] = field(default\_factory=list) sources: list\[SearchHit\] = field(default\_factory=list) @dataclass class IntelReport: watches: list\[CompetitorWatch\] = field(default\_factory=list) @property def deltas(self) -> list\[Delta\]: return \[d for w in self.watches for d in w.deltas\] # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import os import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) @flyte.trace async def you\_search(query: str, count: int = 8, freshness: str = "week") -> list\[SearchHit\]: """Call the You.com Search API and return unified web + news hits.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), thumbnail=item.get("thumbnail\_url", "") or "", section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict | list: """Call Claude via LiteLLM and parse a JSON response.""" from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=2048, ) content = resp.choices\[0\].message.content return \_parse\_json(content) def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min( (i for i in (text.find("{"), text.find("\[")) if i != -1),\ default=0,\ )\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} EXTRACT\_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \\ search results about a competitor, extract concrete, recently-changed signals \\ ("deltas") in the requested categories. Only report changes that are supported \\ by a specific search result. Respond with a JSON object of the form: {"deltas": \[{"category": str, "summary": str, "source\_index": int (the \[n\] of \\\ the supporting search result), "confidence": float between 0 and 1}\]} If there are no clear changes, return {"deltas": \[\]}.""" # {{docs-fragment watch\_competitor}} @env.task(retries=3) async def watch\_competitor( competitor: str, categories: list\[str\], freshness: str, ) -> CompetitorWatch: """Search for fresh signals on one competitor and extract structured deltas.""" query = ( f"{competitor} " + " OR ".join(categories) + " announcement OR news OR update" ) hits = await you\_search(query, count=8, freshness=freshness) if not hits: return CompetitorWatch(competitor=competitor) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Competitor: {competitor}\\n" f"Categories to watch: {', '.join(categories)}\\n\\n" f"Search results:\\n{evidence}" ) parsed = await llm\_json(EXTRACT\_SYSTEM, user) raw\_deltas = parsed.get("deltas", \[\]) if isinstance(parsed, dict) else \[\] deltas: list\[Delta\] = \[\] cited: list\[SearchHit\] = \[\] for d in raw\_deltas: idx = int(d.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None if src is not None and src not in cited: cited.append(src) deltas.append( Delta( competitor=competitor, category=str(d.get("category", "unknown")), summary=str(d.get("summary", "")), confidence=float(d.get("confidence", 0.0) or 0.0), source=src, ) ) return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited) # {{/docs-fragment watch\_competitor}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_conf\_bar(conf: float) -> str: pct = max(0, min(100, int(conf \* 100))) return ( f"" f"{conf:.0%} confidence" ) def \_cite(src: SearchHit) -> str: """Render a rich You.com citation: favicon, domain, date, author, snippet.""" if src is None: return "" tag = ( f"news" if src.section == "news" else "web" ) meta\_bits = \[\] if src.published: meta\_bits.append(src.published\[:10\]) if src.author: meta\_bits.append(f"by {src.author}") meta = " · ".join(meta\_bits) snip = f"
“{src.snippet}”
" if src.snippet else "" return ( f"
" f"" f"
" f"{src.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: IntelReport) -> str: watches = sorted(report.watches, key=lambda w: w.competitor) total\_sources = sum(len(w.sources) for w in watches) cards = \[\] for w in watches: deltas = sorted(w.deltas, key=lambda d: -d.confidence) rows = "".join( f"
{d.category}" f"
{d.summary}
" f"{\_conf\_bar(d.confidence)}" f"{\_cite(d.source)}" "
" for d in deltas ) cards.append( f"

{w.competitor}

" f"{len(deltas)} signal(s) · " f"{len(w.sources)} You.com source(s){rows or ''}
" ) return f""" {REPORT\_CSS}

Competitive Intelligence Deltas

Fresh, source-cited market signals — every delta links back to a ranked, timestamped You.com Search result.

{len(report.deltas)} signals {len(watches)} competitors tracked {total\_sources} cited You.com sources
{''.join(cards) or "

No signals detected in this window.

"}

Sources retrieved and ranked by the You.com Search API (web + auto-classified news), with publication timestamps, authors, and snippet provenance preserved for full prompt → citation lineage.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def competitive\_intelligence( competitors: list\[str\] = \[\ "Anthropic",\ "OpenAI",\ "Mistral AI",\ "Google DeepMind",\ "Cohere",\ "Perplexity AI",\ "xAI",\ "Hugging Face",\ "Databricks",\ "Together AI",\ \], categories: list\[str\] = \[\ "pricing",\ "product launch",\ "model release",\ "funding",\ "leadership",\ "partnership",\ \], freshness: str = "week", ) -> IntelReport: """Fan out across competitors and aggregate structured deltas.""" with flyte.group("watch-competitors"): results = await asyncio.gather( \*\[watch\_competitor(c, categories, freshness) for c in competitors\] ) report = IntelReport(watches=list(results)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(competitive\_intelligence) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive\_intelligence\_agent/main.py\* ## Search with the You.com Search API The \`you\_search\` helper calls the \[You.com Search API\](https://you.com/docs/search/overview) at \`https://ydc-index.io/v1/search\`. It requests unified web and news results with a \`freshness\` filter (\`day\`, \`week\`, \`month\`, or \`year\`) and returns structured hits the LLM can cite by index. See the \[Search API reference\](https://you.com/docs/api-reference/search/v1-search) for all supported parameters, including \`count\`, \`country\`, and search operators. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "competitive\_intelligence" # params = "" # /// """Continuous competitive & market intelligence agent. A Dragonfly-style agent that fans out across competitors, pulls fresh, source-cited web + news results from the You.com Search API, and uses Claude to extract structured "deltas" (pricing, features, funding, leadership, etc.) into a knowledge-graph-ready table. """ # {{docs-fragment env}} import asyncio import json from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="competitive-intelligence", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="competitive-intelligence", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class SearchHit: """A You.com Search result with its full structured metadata.""" title: str url: str domain: str snippet: str published: str # You.com page\_age timestamp author: str favicon: str # You.com favicon\_url thumbnail: str section: str # "news" or "web" — You.com's auto classification @dataclass class Delta: competitor: str category: str summary: str confidence: float source: SearchHit | None = None @dataclass class CompetitorWatch: competitor: str deltas: list\[Delta\] = field(default\_factory=list) sources: list\[SearchHit\] = field(default\_factory=list) @dataclass class IntelReport: watches: list\[CompetitorWatch\] = field(default\_factory=list) @property def deltas(self) -> list\[Delta\]: return \[d for w in self.watches for d in w.deltas\] # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import os import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) @flyte.trace async def you\_search(query: str, count: int = 8, freshness: str = "week") -> list\[SearchHit\]: """Call the You.com Search API and return unified web + news hits.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), thumbnail=item.get("thumbnail\_url", "") or "", section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict | list: """Call Claude via LiteLLM and parse a JSON response.""" from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=2048, ) content = resp.choices\[0\].message.content return \_parse\_json(content) def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min( (i for i in (text.find("{"), text.find("\[")) if i != -1),\ default=0,\ )\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} EXTRACT\_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \\ search results about a competitor, extract concrete, recently-changed signals \\ ("deltas") in the requested categories. Only report changes that are supported \\ by a specific search result. Respond with a JSON object of the form: {"deltas": \[{"category": str, "summary": str, "source\_index": int (the \[n\] of \\\ the supporting search result), "confidence": float between 0 and 1}\]} If there are no clear changes, return {"deltas": \[\]}.""" # {{docs-fragment watch\_competitor}} @env.task(retries=3) async def watch\_competitor( competitor: str, categories: list\[str\], freshness: str, ) -> CompetitorWatch: """Search for fresh signals on one competitor and extract structured deltas.""" query = ( f"{competitor} " + " OR ".join(categories) + " announcement OR news OR update" ) hits = await you\_search(query, count=8, freshness=freshness) if not hits: return CompetitorWatch(competitor=competitor) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Competitor: {competitor}\\n" f"Categories to watch: {', '.join(categories)}\\n\\n" f"Search results:\\n{evidence}" ) parsed = await llm\_json(EXTRACT\_SYSTEM, user) raw\_deltas = parsed.get("deltas", \[\]) if isinstance(parsed, dict) else \[\] deltas: list\[Delta\] = \[\] cited: list\[SearchHit\] = \[\] for d in raw\_deltas: idx = int(d.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None if src is not None and src not in cited: cited.append(src) deltas.append( Delta( competitor=competitor, category=str(d.get("category", "unknown")), summary=str(d.get("summary", "")), confidence=float(d.get("confidence", 0.0) or 0.0), source=src, ) ) return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited) # {{/docs-fragment watch\_competitor}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_conf\_bar(conf: float) -> str: pct = max(0, min(100, int(conf \* 100))) return ( f"" f"{conf:.0%} confidence" ) def \_cite(src: SearchHit) -> str: """Render a rich You.com citation: favicon, domain, date, author, snippet.""" if src is None: return "" tag = ( f"news" if src.section == "news" else "web" ) meta\_bits = \[\] if src.published: meta\_bits.append(src.published\[:10\]) if src.author: meta\_bits.append(f"by {src.author}") meta = " · ".join(meta\_bits) snip = f"
“{src.snippet}”
" if src.snippet else "" return ( f"
" f"" f"
" f"{src.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: IntelReport) -> str: watches = sorted(report.watches, key=lambda w: w.competitor) total\_sources = sum(len(w.sources) for w in watches) cards = \[\] for w in watches: deltas = sorted(w.deltas, key=lambda d: -d.confidence) rows = "".join( f"
{d.category}" f"
{d.summary}
" f"{\_conf\_bar(d.confidence)}" f"{\_cite(d.source)}" "
" for d in deltas ) cards.append( f"

{w.competitor}

" f"{len(deltas)} signal(s) · " f"{len(w.sources)} You.com source(s){rows or ''}
" ) return f""" {REPORT\_CSS}

Competitive Intelligence Deltas

Fresh, source-cited market signals — every delta links back to a ranked, timestamped You.com Search result.

{len(report.deltas)} signals {len(watches)} competitors tracked {total\_sources} cited You.com sources
{''.join(cards) or "

No signals detected in this window.

"}

Sources retrieved and ranked by the You.com Search API (web + auto-classified news), with publication timestamps, authors, and snippet provenance preserved for full prompt → citation lineage.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def competitive\_intelligence( competitors: list\[str\] = \[\ "Anthropic",\ "OpenAI",\ "Mistral AI",\ "Google DeepMind",\ "Cohere",\ "Perplexity AI",\ "xAI",\ "Hugging Face",\ "Databricks",\ "Together AI",\ \], categories: list\[str\] = \[\ "pricing",\ "product launch",\ "model release",\ "funding",\ "leadership",\ "partnership",\ \], freshness: str = "week", ) -> IntelReport: """Fan out across competitors and aggregate structured deltas.""" with flyte.group("watch-competitors"): results = await asyncio.gather( \*\[watch\_competitor(c, categories, freshness) for c in competitors\] ) report = IntelReport(watches=list(results)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(competitive\_intelligence) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive\_intelligence\_agent/main.py\* > \[!NOTE\] > We use \`@flyte.trace\` to track intermediate steps within a task, like You.com API calls and LLM invocations. Each traced call appears as a span in the Flyte dashboard with its inputs and outputs captured. ## Extract deltas with Claude A shared \`llm\_json\` helper routes to Claude through LiteLLM and parses structured JSON from the response. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "competitive\_intelligence" # params = "" # /// """Continuous competitive & market intelligence agent. A Dragonfly-style agent that fans out across competitors, pulls fresh, source-cited web + news results from the You.com Search API, and uses Claude to extract structured "deltas" (pricing, features, funding, leadership, etc.) into a knowledge-graph-ready table. """ # {{docs-fragment env}} import asyncio import json from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="competitive-intelligence", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="competitive-intelligence", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class SearchHit: """A You.com Search result with its full structured metadata.""" title: str url: str domain: str snippet: str published: str # You.com page\_age timestamp author: str favicon: str # You.com favicon\_url thumbnail: str section: str # "news" or "web" — You.com's auto classification @dataclass class Delta: competitor: str category: str summary: str confidence: float source: SearchHit | None = None @dataclass class CompetitorWatch: competitor: str deltas: list\[Delta\] = field(default\_factory=list) sources: list\[SearchHit\] = field(default\_factory=list) @dataclass class IntelReport: watches: list\[CompetitorWatch\] = field(default\_factory=list) @property def deltas(self) -> list\[Delta\]: return \[d for w in self.watches for d in w.deltas\] # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import os import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) @flyte.trace async def you\_search(query: str, count: int = 8, freshness: str = "week") -> list\[SearchHit\]: """Call the You.com Search API and return unified web + news hits.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), thumbnail=item.get("thumbnail\_url", "") or "", section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict | list: """Call Claude via LiteLLM and parse a JSON response.""" from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=2048, ) content = resp.choices\[0\].message.content return \_parse\_json(content) def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min( (i for i in (text.find("{"), text.find("\[")) if i != -1),\ default=0,\ )\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} EXTRACT\_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \\ search results about a competitor, extract concrete, recently-changed signals \\ ("deltas") in the requested categories. Only report changes that are supported \\ by a specific search result. Respond with a JSON object of the form: {"deltas": \[{"category": str, "summary": str, "source\_index": int (the \[n\] of \\\ the supporting search result), "confidence": float between 0 and 1}\]} If there are no clear changes, return {"deltas": \[\]}.""" # {{docs-fragment watch\_competitor}} @env.task(retries=3) async def watch\_competitor( competitor: str, categories: list\[str\], freshness: str, ) -> CompetitorWatch: """Search for fresh signals on one competitor and extract structured deltas.""" query = ( f"{competitor} " + " OR ".join(categories) + " announcement OR news OR update" ) hits = await you\_search(query, count=8, freshness=freshness) if not hits: return CompetitorWatch(competitor=competitor) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Competitor: {competitor}\\n" f"Categories to watch: {', '.join(categories)}\\n\\n" f"Search results:\\n{evidence}" ) parsed = await llm\_json(EXTRACT\_SYSTEM, user) raw\_deltas = parsed.get("deltas", \[\]) if isinstance(parsed, dict) else \[\] deltas: list\[Delta\] = \[\] cited: list\[SearchHit\] = \[\] for d in raw\_deltas: idx = int(d.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None if src is not None and src not in cited: cited.append(src) deltas.append( Delta( competitor=competitor, category=str(d.get("category", "unknown")), summary=str(d.get("summary", "")), confidence=float(d.get("confidence", 0.0) or 0.0), source=src, ) ) return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited) # {{/docs-fragment watch\_competitor}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_conf\_bar(conf: float) -> str: pct = max(0, min(100, int(conf \* 100))) return ( f"" f"{conf:.0%} confidence" ) def \_cite(src: SearchHit) -> str: """Render a rich You.com citation: favicon, domain, date, author, snippet.""" if src is None: return "" tag = ( f"news" if src.section == "news" else "web" ) meta\_bits = \[\] if src.published: meta\_bits.append(src.published\[:10\]) if src.author: meta\_bits.append(f"by {src.author}") meta = " · ".join(meta\_bits) snip = f"
“{src.snippet}”
" if src.snippet else "" return ( f"
" f"" f"
" f"{src.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: IntelReport) -> str: watches = sorted(report.watches, key=lambda w: w.competitor) total\_sources = sum(len(w.sources) for w in watches) cards = \[\] for w in watches: deltas = sorted(w.deltas, key=lambda d: -d.confidence) rows = "".join( f"
{d.category}" f"
{d.summary}
" f"{\_conf\_bar(d.confidence)}" f"{\_cite(d.source)}" "
" for d in deltas ) cards.append( f"

{w.competitor}

" f"{len(deltas)} signal(s) · " f"{len(w.sources)} You.com source(s){rows or ''}
" ) return f""" {REPORT\_CSS}

Competitive Intelligence Deltas

Fresh, source-cited market signals — every delta links back to a ranked, timestamped You.com Search result.

{len(report.deltas)} signals {len(watches)} competitors tracked {total\_sources} cited You.com sources
{''.join(cards) or "

No signals detected in this window.

"}

Sources retrieved and ranked by the You.com Search API (web + auto-classified news), with publication timestamps, authors, and snippet provenance preserved for full prompt → citation lineage.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def competitive\_intelligence( competitors: list\[str\] = \[\ "Anthropic",\ "OpenAI",\ "Mistral AI",\ "Google DeepMind",\ "Cohere",\ "Perplexity AI",\ "xAI",\ "Hugging Face",\ "Databricks",\ "Together AI",\ \], categories: list\[str\] = \[\ "pricing",\ "product launch",\ "model release",\ "funding",\ "leadership",\ "partnership",\ \], freshness: str = "week", ) -> IntelReport: """Fan out across competitors and aggregate structured deltas.""" with flyte.group("watch-competitors"): results = await asyncio.gather( \*\[watch\_competitor(c, categories, freshness) for c in competitors\] ) report = IntelReport(watches=list(results)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(competitive\_intelligence) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive\_intelligence\_agent/main.py\* ## Watch one competitor The \`watch\_competitor\` task builds a category-scoped search query, calls the You.com Search API, and asks Claude to extract only changes that are supported by a specific search result. Each delta carries a confidence score and a link to its source hit. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "competitive\_intelligence" # params = "" # /// """Continuous competitive & market intelligence agent. A Dragonfly-style agent that fans out across competitors, pulls fresh, source-cited web + news results from the You.com Search API, and uses Claude to extract structured "deltas" (pricing, features, funding, leadership, etc.) into a knowledge-graph-ready table. """ # {{docs-fragment env}} import asyncio import json from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="competitive-intelligence", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="competitive-intelligence", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class SearchHit: """A You.com Search result with its full structured metadata.""" title: str url: str domain: str snippet: str published: str # You.com page\_age timestamp author: str favicon: str # You.com favicon\_url thumbnail: str section: str # "news" or "web" — You.com's auto classification @dataclass class Delta: competitor: str category: str summary: str confidence: float source: SearchHit | None = None @dataclass class CompetitorWatch: competitor: str deltas: list\[Delta\] = field(default\_factory=list) sources: list\[SearchHit\] = field(default\_factory=list) @dataclass class IntelReport: watches: list\[CompetitorWatch\] = field(default\_factory=list) @property def deltas(self) -> list\[Delta\]: return \[d for w in self.watches for d in w.deltas\] # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import os import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) @flyte.trace async def you\_search(query: str, count: int = 8, freshness: str = "week") -> list\[SearchHit\]: """Call the You.com Search API and return unified web + news hits.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), thumbnail=item.get("thumbnail\_url", "") or "", section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict | list: """Call Claude via LiteLLM and parse a JSON response.""" from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=2048, ) content = resp.choices\[0\].message.content return \_parse\_json(content) def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min( (i for i in (text.find("{"), text.find("\[")) if i != -1),\ default=0,\ )\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} EXTRACT\_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \\ search results about a competitor, extract concrete, recently-changed signals \\ ("deltas") in the requested categories. Only report changes that are supported \\ by a specific search result. Respond with a JSON object of the form: {"deltas": \[{"category": str, "summary": str, "source\_index": int (the \[n\] of \\\ the supporting search result), "confidence": float between 0 and 1}\]} If there are no clear changes, return {"deltas": \[\]}.""" # {{docs-fragment watch\_competitor}} @env.task(retries=3) async def watch\_competitor( competitor: str, categories: list\[str\], freshness: str, ) -> CompetitorWatch: """Search for fresh signals on one competitor and extract structured deltas.""" query = ( f"{competitor} " + " OR ".join(categories) + " announcement OR news OR update" ) hits = await you\_search(query, count=8, freshness=freshness) if not hits: return CompetitorWatch(competitor=competitor) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Competitor: {competitor}\\n" f"Categories to watch: {', '.join(categories)}\\n\\n" f"Search results:\\n{evidence}" ) parsed = await llm\_json(EXTRACT\_SYSTEM, user) raw\_deltas = parsed.get("deltas", \[\]) if isinstance(parsed, dict) else \[\] deltas: list\[Delta\] = \[\] cited: list\[SearchHit\] = \[\] for d in raw\_deltas: idx = int(d.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None if src is not None and src not in cited: cited.append(src) deltas.append( Delta( competitor=competitor, category=str(d.get("category", "unknown")), summary=str(d.get("summary", "")), confidence=float(d.get("confidence", 0.0) or 0.0), source=src, ) ) return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited) # {{/docs-fragment watch\_competitor}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_conf\_bar(conf: float) -> str: pct = max(0, min(100, int(conf \* 100))) return ( f"" f"{conf:.0%} confidence" ) def \_cite(src: SearchHit) -> str: """Render a rich You.com citation: favicon, domain, date, author, snippet.""" if src is None: return "" tag = ( f"news" if src.section == "news" else "web" ) meta\_bits = \[\] if src.published: meta\_bits.append(src.published\[:10\]) if src.author: meta\_bits.append(f"by {src.author}") meta = " · ".join(meta\_bits) snip = f"
“{src.snippet}”
" if src.snippet else "" return ( f"
" f"" f"
" f"{src.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: IntelReport) -> str: watches = sorted(report.watches, key=lambda w: w.competitor) total\_sources = sum(len(w.sources) for w in watches) cards = \[\] for w in watches: deltas = sorted(w.deltas, key=lambda d: -d.confidence) rows = "".join( f"
{d.category}" f"
{d.summary}
" f"{\_conf\_bar(d.confidence)}" f"{\_cite(d.source)}" "
" for d in deltas ) cards.append( f"

{w.competitor}

" f"{len(deltas)} signal(s) · " f"{len(w.sources)} You.com source(s){rows or ''}
" ) return f""" {REPORT\_CSS}

Competitive Intelligence Deltas

Fresh, source-cited market signals — every delta links back to a ranked, timestamped You.com Search result.

{len(report.deltas)} signals {len(watches)} competitors tracked {total\_sources} cited You.com sources
{''.join(cards) or "

No signals detected in this window.

"}

Sources retrieved and ranked by the You.com Search API (web + auto-classified news), with publication timestamps, authors, and snippet provenance preserved for full prompt → citation lineage.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def competitive\_intelligence( competitors: list\[str\] = \[\ "Anthropic",\ "OpenAI",\ "Mistral AI",\ "Google DeepMind",\ "Cohere",\ "Perplexity AI",\ "xAI",\ "Hugging Face",\ "Databricks",\ "Together AI",\ \], categories: list\[str\] = \[\ "pricing",\ "product launch",\ "model release",\ "funding",\ "leadership",\ "partnership",\ \], freshness: str = "week", ) -> IntelReport: """Fan out across competitors and aggregate structured deltas.""" with flyte.group("watch-competitors"): results = await asyncio.gather( \*\[watch\_competitor(c, categories, freshness) for c in competitors\] ) report = IntelReport(watches=list(results)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(competitive\_intelligence) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive\_intelligence\_agent/main.py\* ## Orchestration The \`competitive\_intelligence\` driver task fans out across all competitors with \`asyncio.gather\`, aggregates the results, and renders a Flyte report. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "competitive\_intelligence" # params = "" # /// """Continuous competitive & market intelligence agent. A Dragonfly-style agent that fans out across competitors, pulls fresh, source-cited web + news results from the You.com Search API, and uses Claude to extract structured "deltas" (pricing, features, funding, leadership, etc.) into a knowledge-graph-ready table. """ # {{docs-fragment env}} import asyncio import json from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="competitive-intelligence", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="competitive-intelligence", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class SearchHit: """A You.com Search result with its full structured metadata.""" title: str url: str domain: str snippet: str published: str # You.com page\_age timestamp author: str favicon: str # You.com favicon\_url thumbnail: str section: str # "news" or "web" — You.com's auto classification @dataclass class Delta: competitor: str category: str summary: str confidence: float source: SearchHit | None = None @dataclass class CompetitorWatch: competitor: str deltas: list\[Delta\] = field(default\_factory=list) sources: list\[SearchHit\] = field(default\_factory=list) @dataclass class IntelReport: watches: list\[CompetitorWatch\] = field(default\_factory=list) @property def deltas(self) -> list\[Delta\]: return \[d for w in self.watches for d in w.deltas\] # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import os import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) @flyte.trace async def you\_search(query: str, count: int = 8, freshness: str = "week") -> list\[SearchHit\]: """Call the You.com Search API and return unified web + news hits.""" params = {"query": query, "count": count, "freshness": freshness} data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), thumbnail=item.get("thumbnail\_url", "") or "", section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict | list: """Call Claude via LiteLLM and parse a JSON response.""" from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=2048, ) content = resp.choices\[0\].message.content return \_parse\_json(content) def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min( (i for i in (text.find("{"), text.find("\[")) if i != -1),\ default=0,\ )\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} EXTRACT\_SYSTEM = """You are a competitive-intelligence analyst. Given fresh \\ search results about a competitor, extract concrete, recently-changed signals \\ ("deltas") in the requested categories. Only report changes that are supported \\ by a specific search result. Respond with a JSON object of the form: {"deltas": \[{"category": str, "summary": str, "source\_index": int (the \[n\] of \\\ the supporting search result), "confidence": float between 0 and 1}\]} If there are no clear changes, return {"deltas": \[\]}.""" # {{docs-fragment watch\_competitor}} @env.task(retries=3) async def watch\_competitor( competitor: str, categories: list\[str\], freshness: str, ) -> CompetitorWatch: """Search for fresh signals on one competitor and extract structured deltas.""" query = ( f"{competitor} " + " OR ".join(categories) + " announcement OR news OR update" ) hits = await you\_search(query, count=8, freshness=freshness) if not hits: return CompetitorWatch(competitor=competitor) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Competitor: {competitor}\\n" f"Categories to watch: {', '.join(categories)}\\n\\n" f"Search results:\\n{evidence}" ) parsed = await llm\_json(EXTRACT\_SYSTEM, user) raw\_deltas = parsed.get("deltas", \[\]) if isinstance(parsed, dict) else \[\] deltas: list\[Delta\] = \[\] cited: list\[SearchHit\] = \[\] for d in raw\_deltas: idx = int(d.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None if src is not None and src not in cited: cited.append(src) deltas.append( Delta( competitor=competitor, category=str(d.get("category", "unknown")), summary=str(d.get("summary", "")), confidence=float(d.get("confidence", 0.0) or 0.0), source=src, ) ) return CompetitorWatch(competitor=competitor, deltas=deltas, sources=cited) # {{/docs-fragment watch\_competitor}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_conf\_bar(conf: float) -> str: pct = max(0, min(100, int(conf \* 100))) return ( f"" f"{conf:.0%} confidence" ) def \_cite(src: SearchHit) -> str: """Render a rich You.com citation: favicon, domain, date, author, snippet.""" if src is None: return "" tag = ( f"news" if src.section == "news" else "web" ) meta\_bits = \[\] if src.published: meta\_bits.append(src.published\[:10\]) if src.author: meta\_bits.append(f"by {src.author}") meta = " · ".join(meta\_bits) snip = f"
“{src.snippet}”
" if src.snippet else "" return ( f"
" f"" f"
" f"{src.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: IntelReport) -> str: watches = sorted(report.watches, key=lambda w: w.competitor) total\_sources = sum(len(w.sources) for w in watches) cards = \[\] for w in watches: deltas = sorted(w.deltas, key=lambda d: -d.confidence) rows = "".join( f"
{d.category}" f"
{d.summary}
" f"{\_conf\_bar(d.confidence)}" f"{\_cite(d.source)}" "
" for d in deltas ) cards.append( f"

{w.competitor}

" f"{len(deltas)} signal(s) · " f"{len(w.sources)} You.com source(s){rows or ''}
" ) return f""" {REPORT\_CSS}

Competitive Intelligence Deltas

Fresh, source-cited market signals — every delta links back to a ranked, timestamped You.com Search result.

{len(report.deltas)} signals {len(watches)} competitors tracked {total\_sources} cited You.com sources
{''.join(cards) or "

No signals detected in this window.

"}

Sources retrieved and ranked by the You.com Search API (web + auto-classified news), with publication timestamps, authors, and snippet provenance preserved for full prompt → citation lineage.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} @env.task(report=True) async def competitive\_intelligence( competitors: list\[str\] = \[\ "Anthropic",\ "OpenAI",\ "Mistral AI",\ "Google DeepMind",\ "Cohere",\ "Perplexity AI",\ "xAI",\ "Hugging Face",\ "Databricks",\ "Together AI",\ \], categories: list\[str\] = \[\ "pricing",\ "product launch",\ "model release",\ "funding",\ "leadership",\ "partnership",\ \], freshness: str = "week", ) -> IntelReport: """Fan out across competitors and aggregate structured deltas.""" with flyte.group("watch-competitors"): results = await asyncio.gather( \*\[watch\_competitor(c, categories, freshness) for c in competitors\] ) report = IntelReport(watches=list(results)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(competitive\_intelligence) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/competitive\_intelligence\_agent/main.py\* ## Run the agent ### Create secrets Get a You.com API key from the \[You.com platform\](https://you.com/platform) (see the \[quickstart guide\](https://you.com/docs/quickstart)). Get an Anthropic API key from the \[Anthropic console\](https://console.anthropic.com/). Register both keys as Flyte secrets. The secret key names must match those declared in the \`TaskEnvironment\`: \`\`\` flyte create secret youdotcom-api-key flyte create secret internal-anthropic-api-key \`\`\` See \[Secrets\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets/page.md) for scoping and file-based secrets. ### Run locally or remotely From the \[example directory\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/competitive\_intelligence\_agent): \`\`\` cd v2/tutorials/competitive\_intelligence\_agent uv run --script main.py \`\`\` Or pass custom competitors with the Flyte CLI: \`\`\` flyte run main.py competitive\_intelligence \\ --competitors '\["Anthropic", "OpenAI"\]' \`\`\` To test locally without Flyte secrets, export the environment variables directly: \`\`\` export YOU\_API\_KEY= export ANTHROPIC\_API\_KEY= uv run --script main.py \`\`\` When the run completes, open the Flyte report in the UI to review deltas grouped by competitor, each with a clickable You.com source citation. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/compliance-monitoring-agent === # Compliance monitoring agent > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/compliance\_monitoring\_agent). This example demonstrates how to build a regulatory and compliance monitoring agent on Flyte. The agent watches trusted regulatory sources — FDA guidance, SEC filings, sanctions lists, state-level privacy laws — and routes structured, \*\*citation-precise\*\* findings to the right downstream team (compliance, legal, or clinical ops). Compliance monitoring requires \*\*citation precision and recency\*\* so every finding can be verified. The \[You.com Research API\](https://you.com/docs/research/overview) returns a grounded, synthesized answer plus structured sources (URL, title, snippet). Use \`source\_control\` to restrict research to trusted government and regulator domains within a recency window, and \`output\_schema\` when you need machine-readable findings. \[Claude\](https://docs.anthropic.com/) via \[LiteLLM\](https://docs.litellm.ai/) triages each finding for severity and routing. Combined with Flyte's audit lineage, you get end-to-end traceability from query to citation. Flyte provides: - \*\*Fan-out parallelism\*\* across watch items - \*\*\`@flyte.trace\`\*\* on every You.com Research and LLM call - \*\*Retries\*\* on monitoring tasks for robustness - \*\*Flyte reports\*\* grouped by team and severity !\[Compliance monitoring agent report\](https://www.union.ai/docs/v2/flyte/\_static/images/tutorials/compliance\_monitoring\_agent/compliance-monitoring-agent.png) ## Setting up the environment The agent runs in a \`TaskEnvironment\` with secrets for the You.com and Anthropic API keys and a container image built from the \`uv\` script dependencies. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "compliance\_monitoring" # params = "" # /// """Regulatory & compliance monitoring agent. Watches trusted regulatory sources via the You.com Research API (with domain/freshness source controls and a structured output schema), then uses Claude to assign severity and route citation-precise findings to the right team. Every external call is traced so Flyte's audit lineage extends to the web layer. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="compliance-monitoring", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="compliance-monitoring", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class WatchItem: topic: str trusted\_domains: list\[str\] team: str @dataclass class Finding: topic: str team: str title: str summary: str source\_url: str published\_date: str snippet: str domain: str = "" favicon: str = "" severity: str = "info" rationale: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class ComplianceReport: findings: list\[Finding\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" FINDINGS\_SCHEMA = { "type": "object", "properties": { "findings": { "type": "array", "items": { "type": "object", "properties": { "title": {"type": "string"}, "summary": {"type": "string"}, "source\_url": {"type": "string"}, "published\_date": {"type": "string"}, "snippet": {"type": "string"}, }, "required": \[\ "title",\ "summary",\ "source\_url",\ "published\_date",\ "snippet",\ \], "additionalProperties": False, }, } }, "required": \["findings"\], "additionalProperties": False, } async def \_you\_post(url: str, body: dict, timeout: float = 300.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research( question: str, include\_domains: list\[str\], freshness: str, research\_effort: str = "standard", ) -> dict: """Call the You.com Research API with domain + freshness source controls.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": { "include\_domains": include\_domains, "freshness": freshness, }, "output\_schema": FINDINGS\_SCHEMA, } return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment llm}} @flyte.trace async def triage(topic: str, findings: list\[dict\]) -> list\[dict\]: """Use Claude to assign a severity + rationale to each finding.""" from litellm import acompletion if not findings: return \[\] system = ( "You are a regulatory-compliance triage analyst. For each finding, " "assign a severity of 'info' (FYI), 'watch' (monitor closely), or " "'action' (requires a concrete compliance/legal response), and a one-" "sentence rationale. Respond ONLY with JSON: " '{"triage": \[{"severity": str, "rationale": str}\]} with one entry per ' "finding, in order." ) listing = "\\n".join( f"\[{i + 1}\] {f.get('title', '')}: {f.get('summary', '')}" for i, f in enumerate(findings) ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": f"Topic: {topic}\\n\\nFindings:\\n{listing}"},\ \], temperature=0.0, max\_tokens=1024, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed.get("triage", \[\]) if isinstance(parsed, dict) else \[\] def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment monitor\_watch\_item}} @env.task(retries=3) async def monitor\_watch\_item(item: WatchItem, freshness: str) -> list\[Finding\]: """Research one regulatory topic and produce triaged, cited findings.""" question = ( f"What are the most recent changes, updates, or new guidance regarding " f"'{item.topic}'? Report concrete, dated changes with their sources." ) result = await you\_research(question, item.trusted\_domains, freshness) output = result.get("output", {}) # Build a lookup from the Research API's full source list (url -> metadata). src\_by\_url: dict\[str, dict\] = {} for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) if url: src\_by\_url\[url\] = s content = output.get("content", {}) if isinstance(content, str): content = \_parse\_json(content) if content.strip() else {} raw\_findings = content.get("findings", \[\]) if isinstance(content, dict) else \[\] triage\_results = await triage(item.topic, raw\_findings) findings: list\[Finding\] = \[\] for i, f in enumerate(raw\_findings): t = triage\_results\[i\] if i < len(triage\_results) else {} url = str(f.get("source\_url", "")) meta = src\_by\_url.get(url, {}) snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or \[""\])\[0\]) findings.append( Finding( topic=item.topic, team=item.team, title=str(f.get("title", "") or meta.get("title", "")), summary=str(f.get("summary", "")), source\_url=url, published\_date=str(f.get("published\_date", "")), snippet=snippet, domain=\_domain(url), favicon=\_favicon\_for(url), severity=str(t.get("severity", "info")), rationale=str(t.get("rationale", "")), ) ) return findings # {{/docs-fragment monitor\_watch\_item}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"action": 0, "watch": 1, "info": 2} \_SEVERITY\_STYLE = { "action": ("#fdecea", "#c0392b"), "watch": ("#fdf3e1", "#b7791f"), "info": ("#e3f1fb", "#2b6cb0"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#edf0f3", "#52606d")) return f"{sev}" def \_cite(f: Finding) -> str: """Render a rich You.com Research citation with domain, date, and snippet.""" if not f.source\_url: return "" meta = f.published\_date\[:10\] if f.published\_date else "" snip = f"
“{f.snippet}”
" if f.snippet else "" return ( f"
" f"
" f"{f.domain or 'source'}" f"research" f"
{meta} · {f.title}
{snip}
" ) def \_render\_report(report: ComplianceReport) -> str: findings = sorted( report.findings, key=lambda f: (\_SEVERITY\_ORDER.get(f.severity, 3), f.team), ) counts = {s: sum(1 for f in findings if f.severity == s) for s in \_SEVERITY\_ORDER} cited = sum(1 for f in findings if f.source\_url) cards = \[\] for f in findings: cards.append( f"
" f"
{\_sev\_badge(f.severity)}{f.team}
" f"

{f.title or f.topic}

" f"
{f.summary}
" f"
{f.rationale}
" f"
{f.topic}
" f"{\_cite(f)}
" ) return f""" {REPORT\_CSS}

Compliance Monitoring Findings

Citation-precise regulatory changes from trusted domains — every finding links to a You.com Research source with snippet provenance.

{len(findings)} findings {cited} cited You.com sources {counts\['action'\]} action {counts\['watch'\]} watch {counts\['info'\]} info
{''.join(cards) or "

No findings in this window.

"}

Findings retrieved via the You.com Research API with source\_control domain allowlists and freshness filters. Flyte logs which agent called which query and got which document — full prompt → citation lineage for audit.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_watch\_items() -> list\[WatchItem\]: return \[\ WatchItem(\ topic="FDA guidance on AI/ML-enabled medical device software",\ trusted\_domains=\["fda.gov", "federalregister.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="SEC climate-related disclosure rules for public companies",\ trusted\_domains=\["sec.gov", "federalregister.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="OFAC sanctions list additions and updates",\ trusted\_domains=\["treasury.gov", "ofac.treasury.gov"\],\ team="compliance",\ ),\ WatchItem(\ topic="State-level consumer data privacy laws and amendments",\ trusted\_domains=\["iapp.org", "oag.ca.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="FDA drug recalls and safety communications",\ trusted\_domains=\["fda.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="HIPAA enforcement actions and guidance updates",\ trusted\_domains=\["hhs.gov"\],\ team="compliance",\ ),\ \] @env.task(report=True) async def compliance\_monitoring( watch\_items: list\[WatchItem\] | None = None, freshness: str = "month", ) -> ComplianceReport: """Fan out across regulatory watch items and aggregate triaged findings.""" if watch\_items is None: watch\_items = \_default\_watch\_items() with flyte.group("monitor-watch-items"): results = await asyncio.gather( \*\[monitor\_watch\_item(item, freshness) for item in watch\_items\] ) report = ComplianceReport(findings=\[f for fs in results for f in fs\]) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(compliance\_monitoring) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance\_monitoring\_agent/main.py\* The Python packages are declared at the top of the file using the \`uv\` script style: \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # /// \`\`\` ## Data types Each \`WatchItem\` specifies a regulatory topic, a list of trusted domains for \`source\_control\`, and a routing destination team. Findings carry citation metadata — source URL, published date, and snippet — so every claim can be verified. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "compliance\_monitoring" # params = "" # /// """Regulatory & compliance monitoring agent. Watches trusted regulatory sources via the You.com Research API (with domain/freshness source controls and a structured output schema), then uses Claude to assign severity and route citation-precise findings to the right team. Every external call is traced so Flyte's audit lineage extends to the web layer. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="compliance-monitoring", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="compliance-monitoring", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class WatchItem: topic: str trusted\_domains: list\[str\] team: str @dataclass class Finding: topic: str team: str title: str summary: str source\_url: str published\_date: str snippet: str domain: str = "" favicon: str = "" severity: str = "info" rationale: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class ComplianceReport: findings: list\[Finding\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" FINDINGS\_SCHEMA = { "type": "object", "properties": { "findings": { "type": "array", "items": { "type": "object", "properties": { "title": {"type": "string"}, "summary": {"type": "string"}, "source\_url": {"type": "string"}, "published\_date": {"type": "string"}, "snippet": {"type": "string"}, }, "required": \[\ "title",\ "summary",\ "source\_url",\ "published\_date",\ "snippet",\ \], "additionalProperties": False, }, } }, "required": \["findings"\], "additionalProperties": False, } async def \_you\_post(url: str, body: dict, timeout: float = 300.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research( question: str, include\_domains: list\[str\], freshness: str, research\_effort: str = "standard", ) -> dict: """Call the You.com Research API with domain + freshness source controls.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": { "include\_domains": include\_domains, "freshness": freshness, }, "output\_schema": FINDINGS\_SCHEMA, } return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment llm}} @flyte.trace async def triage(topic: str, findings: list\[dict\]) -> list\[dict\]: """Use Claude to assign a severity + rationale to each finding.""" from litellm import acompletion if not findings: return \[\] system = ( "You are a regulatory-compliance triage analyst. For each finding, " "assign a severity of 'info' (FYI), 'watch' (monitor closely), or " "'action' (requires a concrete compliance/legal response), and a one-" "sentence rationale. Respond ONLY with JSON: " '{"triage": \[{"severity": str, "rationale": str}\]} with one entry per ' "finding, in order." ) listing = "\\n".join( f"\[{i + 1}\] {f.get('title', '')}: {f.get('summary', '')}" for i, f in enumerate(findings) ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": f"Topic: {topic}\\n\\nFindings:\\n{listing}"},\ \], temperature=0.0, max\_tokens=1024, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed.get("triage", \[\]) if isinstance(parsed, dict) else \[\] def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment monitor\_watch\_item}} @env.task(retries=3) async def monitor\_watch\_item(item: WatchItem, freshness: str) -> list\[Finding\]: """Research one regulatory topic and produce triaged, cited findings.""" question = ( f"What are the most recent changes, updates, or new guidance regarding " f"'{item.topic}'? Report concrete, dated changes with their sources." ) result = await you\_research(question, item.trusted\_domains, freshness) output = result.get("output", {}) # Build a lookup from the Research API's full source list (url -> metadata). src\_by\_url: dict\[str, dict\] = {} for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) if url: src\_by\_url\[url\] = s content = output.get("content", {}) if isinstance(content, str): content = \_parse\_json(content) if content.strip() else {} raw\_findings = content.get("findings", \[\]) if isinstance(content, dict) else \[\] triage\_results = await triage(item.topic, raw\_findings) findings: list\[Finding\] = \[\] for i, f in enumerate(raw\_findings): t = triage\_results\[i\] if i < len(triage\_results) else {} url = str(f.get("source\_url", "")) meta = src\_by\_url.get(url, {}) snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or \[""\])\[0\]) findings.append( Finding( topic=item.topic, team=item.team, title=str(f.get("title", "") or meta.get("title", "")), summary=str(f.get("summary", "")), source\_url=url, published\_date=str(f.get("published\_date", "")), snippet=snippet, domain=\_domain(url), favicon=\_favicon\_for(url), severity=str(t.get("severity", "info")), rationale=str(t.get("rationale", "")), ) ) return findings # {{/docs-fragment monitor\_watch\_item}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"action": 0, "watch": 1, "info": 2} \_SEVERITY\_STYLE = { "action": ("#fdecea", "#c0392b"), "watch": ("#fdf3e1", "#b7791f"), "info": ("#e3f1fb", "#2b6cb0"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#edf0f3", "#52606d")) return f"{sev}" def \_cite(f: Finding) -> str: """Render a rich You.com Research citation with domain, date, and snippet.""" if not f.source\_url: return "" meta = f.published\_date\[:10\] if f.published\_date else "" snip = f"
“{f.snippet}”
" if f.snippet else "" return ( f"
" f"
" f"{f.domain or 'source'}" f"research" f"
{meta} · {f.title}
{snip}
" ) def \_render\_report(report: ComplianceReport) -> str: findings = sorted( report.findings, key=lambda f: (\_SEVERITY\_ORDER.get(f.severity, 3), f.team), ) counts = {s: sum(1 for f in findings if f.severity == s) for s in \_SEVERITY\_ORDER} cited = sum(1 for f in findings if f.source\_url) cards = \[\] for f in findings: cards.append( f"
" f"
{\_sev\_badge(f.severity)}{f.team}
" f"

{f.title or f.topic}

" f"
{f.summary}
" f"
{f.rationale}
" f"
{f.topic}
" f"{\_cite(f)}
" ) return f""" {REPORT\_CSS}

Compliance Monitoring Findings

Citation-precise regulatory changes from trusted domains — every finding links to a You.com Research source with snippet provenance.

{len(findings)} findings {cited} cited You.com sources {counts\['action'\]} action {counts\['watch'\]} watch {counts\['info'\]} info
{''.join(cards) or "

No findings in this window.

"}

Findings retrieved via the You.com Research API with source\_control domain allowlists and freshness filters. Flyte logs which agent called which query and got which document — full prompt → citation lineage for audit.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_watch\_items() -> list\[WatchItem\]: return \[\ WatchItem(\ topic="FDA guidance on AI/ML-enabled medical device software",\ trusted\_domains=\["fda.gov", "federalregister.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="SEC climate-related disclosure rules for public companies",\ trusted\_domains=\["sec.gov", "federalregister.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="OFAC sanctions list additions and updates",\ trusted\_domains=\["treasury.gov", "ofac.treasury.gov"\],\ team="compliance",\ ),\ WatchItem(\ topic="State-level consumer data privacy laws and amendments",\ trusted\_domains=\["iapp.org", "oag.ca.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="FDA drug recalls and safety communications",\ trusted\_domains=\["fda.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="HIPAA enforcement actions and guidance updates",\ trusted\_domains=\["hhs.gov"\],\ team="compliance",\ ),\ \] @env.task(report=True) async def compliance\_monitoring( watch\_items: list\[WatchItem\] | None = None, freshness: str = "month", ) -> ComplianceReport: """Fan out across regulatory watch items and aggregate triaged findings.""" if watch\_items is None: watch\_items = \_default\_watch\_items() with flyte.group("monitor-watch-items"): results = await asyncio.gather( \*\[monitor\_watch\_item(item, freshness) for item in watch\_items\] ) report = ComplianceReport(findings=\[f for fs in results for f in fs\]) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(compliance\_monitoring) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance\_monitoring\_agent/main.py\* ## Research with the You.com Research API The \`you\_research\` helper calls the \[You.com Research API\](https://you.com/docs/research/overview) at \`https://api.you.com/v1/research\`. It passes \`source\_control\` with an \`include\_domains\` allowlist and a \`freshness\` filter, and requests structured output via \`output\_schema\`. See the \[Research API reference\](https://you.com/docs/api-reference/research/v1-research) for \`research\_effort\` levels (\`lite\`, \`standard\`, \`deep\`, \`exhaustive\`), \`source\_control\`, and \`output\_schema\` parameters. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "compliance\_monitoring" # params = "" # /// """Regulatory & compliance monitoring agent. Watches trusted regulatory sources via the You.com Research API (with domain/freshness source controls and a structured output schema), then uses Claude to assign severity and route citation-precise findings to the right team. Every external call is traced so Flyte's audit lineage extends to the web layer. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="compliance-monitoring", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="compliance-monitoring", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class WatchItem: topic: str trusted\_domains: list\[str\] team: str @dataclass class Finding: topic: str team: str title: str summary: str source\_url: str published\_date: str snippet: str domain: str = "" favicon: str = "" severity: str = "info" rationale: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class ComplianceReport: findings: list\[Finding\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" FINDINGS\_SCHEMA = { "type": "object", "properties": { "findings": { "type": "array", "items": { "type": "object", "properties": { "title": {"type": "string"}, "summary": {"type": "string"}, "source\_url": {"type": "string"}, "published\_date": {"type": "string"}, "snippet": {"type": "string"}, }, "required": \[\ "title",\ "summary",\ "source\_url",\ "published\_date",\ "snippet",\ \], "additionalProperties": False, }, } }, "required": \["findings"\], "additionalProperties": False, } async def \_you\_post(url: str, body: dict, timeout: float = 300.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research( question: str, include\_domains: list\[str\], freshness: str, research\_effort: str = "standard", ) -> dict: """Call the You.com Research API with domain + freshness source controls.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": { "include\_domains": include\_domains, "freshness": freshness, }, "output\_schema": FINDINGS\_SCHEMA, } return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment llm}} @flyte.trace async def triage(topic: str, findings: list\[dict\]) -> list\[dict\]: """Use Claude to assign a severity + rationale to each finding.""" from litellm import acompletion if not findings: return \[\] system = ( "You are a regulatory-compliance triage analyst. For each finding, " "assign a severity of 'info' (FYI), 'watch' (monitor closely), or " "'action' (requires a concrete compliance/legal response), and a one-" "sentence rationale. Respond ONLY with JSON: " '{"triage": \[{"severity": str, "rationale": str}\]} with one entry per ' "finding, in order." ) listing = "\\n".join( f"\[{i + 1}\] {f.get('title', '')}: {f.get('summary', '')}" for i, f in enumerate(findings) ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": f"Topic: {topic}\\n\\nFindings:\\n{listing}"},\ \], temperature=0.0, max\_tokens=1024, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed.get("triage", \[\]) if isinstance(parsed, dict) else \[\] def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment monitor\_watch\_item}} @env.task(retries=3) async def monitor\_watch\_item(item: WatchItem, freshness: str) -> list\[Finding\]: """Research one regulatory topic and produce triaged, cited findings.""" question = ( f"What are the most recent changes, updates, or new guidance regarding " f"'{item.topic}'? Report concrete, dated changes with their sources." ) result = await you\_research(question, item.trusted\_domains, freshness) output = result.get("output", {}) # Build a lookup from the Research API's full source list (url -> metadata). src\_by\_url: dict\[str, dict\] = {} for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) if url: src\_by\_url\[url\] = s content = output.get("content", {}) if isinstance(content, str): content = \_parse\_json(content) if content.strip() else {} raw\_findings = content.get("findings", \[\]) if isinstance(content, dict) else \[\] triage\_results = await triage(item.topic, raw\_findings) findings: list\[Finding\] = \[\] for i, f in enumerate(raw\_findings): t = triage\_results\[i\] if i < len(triage\_results) else {} url = str(f.get("source\_url", "")) meta = src\_by\_url.get(url, {}) snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or \[""\])\[0\]) findings.append( Finding( topic=item.topic, team=item.team, title=str(f.get("title", "") or meta.get("title", "")), summary=str(f.get("summary", "")), source\_url=url, published\_date=str(f.get("published\_date", "")), snippet=snippet, domain=\_domain(url), favicon=\_favicon\_for(url), severity=str(t.get("severity", "info")), rationale=str(t.get("rationale", "")), ) ) return findings # {{/docs-fragment monitor\_watch\_item}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"action": 0, "watch": 1, "info": 2} \_SEVERITY\_STYLE = { "action": ("#fdecea", "#c0392b"), "watch": ("#fdf3e1", "#b7791f"), "info": ("#e3f1fb", "#2b6cb0"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#edf0f3", "#52606d")) return f"{sev}" def \_cite(f: Finding) -> str: """Render a rich You.com Research citation with domain, date, and snippet.""" if not f.source\_url: return "" meta = f.published\_date\[:10\] if f.published\_date else "" snip = f"
“{f.snippet}”
" if f.snippet else "" return ( f"
" f"
" f"{f.domain or 'source'}" f"research" f"
{meta} · {f.title}
{snip}
" ) def \_render\_report(report: ComplianceReport) -> str: findings = sorted( report.findings, key=lambda f: (\_SEVERITY\_ORDER.get(f.severity, 3), f.team), ) counts = {s: sum(1 for f in findings if f.severity == s) for s in \_SEVERITY\_ORDER} cited = sum(1 for f in findings if f.source\_url) cards = \[\] for f in findings: cards.append( f"
" f"
{\_sev\_badge(f.severity)}{f.team}
" f"

{f.title or f.topic}

" f"
{f.summary}
" f"
{f.rationale}
" f"
{f.topic}
" f"{\_cite(f)}
" ) return f""" {REPORT\_CSS}

Compliance Monitoring Findings

Citation-precise regulatory changes from trusted domains — every finding links to a You.com Research source with snippet provenance.

{len(findings)} findings {cited} cited You.com sources {counts\['action'\]} action {counts\['watch'\]} watch {counts\['info'\]} info
{''.join(cards) or "

No findings in this window.

"}

Findings retrieved via the You.com Research API with source\_control domain allowlists and freshness filters. Flyte logs which agent called which query and got which document — full prompt → citation lineage for audit.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_watch\_items() -> list\[WatchItem\]: return \[\ WatchItem(\ topic="FDA guidance on AI/ML-enabled medical device software",\ trusted\_domains=\["fda.gov", "federalregister.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="SEC climate-related disclosure rules for public companies",\ trusted\_domains=\["sec.gov", "federalregister.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="OFAC sanctions list additions and updates",\ trusted\_domains=\["treasury.gov", "ofac.treasury.gov"\],\ team="compliance",\ ),\ WatchItem(\ topic="State-level consumer data privacy laws and amendments",\ trusted\_domains=\["iapp.org", "oag.ca.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="FDA drug recalls and safety communications",\ trusted\_domains=\["fda.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="HIPAA enforcement actions and guidance updates",\ trusted\_domains=\["hhs.gov"\],\ team="compliance",\ ),\ \] @env.task(report=True) async def compliance\_monitoring( watch\_items: list\[WatchItem\] | None = None, freshness: str = "month", ) -> ComplianceReport: """Fan out across regulatory watch items and aggregate triaged findings.""" if watch\_items is None: watch\_items = \_default\_watch\_items() with flyte.group("monitor-watch-items"): results = await asyncio.gather( \*\[monitor\_watch\_item(item, freshness) for item in watch\_items\] ) report = ComplianceReport(findings=\[f for fs in results for f in fs\]) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(compliance\_monitoring) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance\_monitoring\_agent/main.py\* ## Triage findings with Claude After the Research API returns structured findings, Claude assigns a severity (\`info\`, \`watch\`, or \`action\`) and a routing rationale for each one. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "compliance\_monitoring" # params = "" # /// """Regulatory & compliance monitoring agent. Watches trusted regulatory sources via the You.com Research API (with domain/freshness source controls and a structured output schema), then uses Claude to assign severity and route citation-precise findings to the right team. Every external call is traced so Flyte's audit lineage extends to the web layer. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="compliance-monitoring", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="compliance-monitoring", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class WatchItem: topic: str trusted\_domains: list\[str\] team: str @dataclass class Finding: topic: str team: str title: str summary: str source\_url: str published\_date: str snippet: str domain: str = "" favicon: str = "" severity: str = "info" rationale: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class ComplianceReport: findings: list\[Finding\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" FINDINGS\_SCHEMA = { "type": "object", "properties": { "findings": { "type": "array", "items": { "type": "object", "properties": { "title": {"type": "string"}, "summary": {"type": "string"}, "source\_url": {"type": "string"}, "published\_date": {"type": "string"}, "snippet": {"type": "string"}, }, "required": \[\ "title",\ "summary",\ "source\_url",\ "published\_date",\ "snippet",\ \], "additionalProperties": False, }, } }, "required": \["findings"\], "additionalProperties": False, } async def \_you\_post(url: str, body: dict, timeout: float = 300.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research( question: str, include\_domains: list\[str\], freshness: str, research\_effort: str = "standard", ) -> dict: """Call the You.com Research API with domain + freshness source controls.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": { "include\_domains": include\_domains, "freshness": freshness, }, "output\_schema": FINDINGS\_SCHEMA, } return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment llm}} @flyte.trace async def triage(topic: str, findings: list\[dict\]) -> list\[dict\]: """Use Claude to assign a severity + rationale to each finding.""" from litellm import acompletion if not findings: return \[\] system = ( "You are a regulatory-compliance triage analyst. For each finding, " "assign a severity of 'info' (FYI), 'watch' (monitor closely), or " "'action' (requires a concrete compliance/legal response), and a one-" "sentence rationale. Respond ONLY with JSON: " '{"triage": \[{"severity": str, "rationale": str}\]} with one entry per ' "finding, in order." ) listing = "\\n".join( f"\[{i + 1}\] {f.get('title', '')}: {f.get('summary', '')}" for i, f in enumerate(findings) ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": f"Topic: {topic}\\n\\nFindings:\\n{listing}"},\ \], temperature=0.0, max\_tokens=1024, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed.get("triage", \[\]) if isinstance(parsed, dict) else \[\] def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment monitor\_watch\_item}} @env.task(retries=3) async def monitor\_watch\_item(item: WatchItem, freshness: str) -> list\[Finding\]: """Research one regulatory topic and produce triaged, cited findings.""" question = ( f"What are the most recent changes, updates, or new guidance regarding " f"'{item.topic}'? Report concrete, dated changes with their sources." ) result = await you\_research(question, item.trusted\_domains, freshness) output = result.get("output", {}) # Build a lookup from the Research API's full source list (url -> metadata). src\_by\_url: dict\[str, dict\] = {} for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) if url: src\_by\_url\[url\] = s content = output.get("content", {}) if isinstance(content, str): content = \_parse\_json(content) if content.strip() else {} raw\_findings = content.get("findings", \[\]) if isinstance(content, dict) else \[\] triage\_results = await triage(item.topic, raw\_findings) findings: list\[Finding\] = \[\] for i, f in enumerate(raw\_findings): t = triage\_results\[i\] if i < len(triage\_results) else {} url = str(f.get("source\_url", "")) meta = src\_by\_url.get(url, {}) snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or \[""\])\[0\]) findings.append( Finding( topic=item.topic, team=item.team, title=str(f.get("title", "") or meta.get("title", "")), summary=str(f.get("summary", "")), source\_url=url, published\_date=str(f.get("published\_date", "")), snippet=snippet, domain=\_domain(url), favicon=\_favicon\_for(url), severity=str(t.get("severity", "info")), rationale=str(t.get("rationale", "")), ) ) return findings # {{/docs-fragment monitor\_watch\_item}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"action": 0, "watch": 1, "info": 2} \_SEVERITY\_STYLE = { "action": ("#fdecea", "#c0392b"), "watch": ("#fdf3e1", "#b7791f"), "info": ("#e3f1fb", "#2b6cb0"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#edf0f3", "#52606d")) return f"{sev}" def \_cite(f: Finding) -> str: """Render a rich You.com Research citation with domain, date, and snippet.""" if not f.source\_url: return "" meta = f.published\_date\[:10\] if f.published\_date else "" snip = f"
“{f.snippet}”
" if f.snippet else "" return ( f"
" f"
" f"{f.domain or 'source'}" f"research" f"
{meta} · {f.title}
{snip}
" ) def \_render\_report(report: ComplianceReport) -> str: findings = sorted( report.findings, key=lambda f: (\_SEVERITY\_ORDER.get(f.severity, 3), f.team), ) counts = {s: sum(1 for f in findings if f.severity == s) for s in \_SEVERITY\_ORDER} cited = sum(1 for f in findings if f.source\_url) cards = \[\] for f in findings: cards.append( f"
" f"
{\_sev\_badge(f.severity)}{f.team}
" f"

{f.title or f.topic}

" f"
{f.summary}
" f"
{f.rationale}
" f"
{f.topic}
" f"{\_cite(f)}
" ) return f""" {REPORT\_CSS}

Compliance Monitoring Findings

Citation-precise regulatory changes from trusted domains — every finding links to a You.com Research source with snippet provenance.

{len(findings)} findings {cited} cited You.com sources {counts\['action'\]} action {counts\['watch'\]} watch {counts\['info'\]} info
{''.join(cards) or "

No findings in this window.

"}

Findings retrieved via the You.com Research API with source\_control domain allowlists and freshness filters. Flyte logs which agent called which query and got which document — full prompt → citation lineage for audit.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_watch\_items() -> list\[WatchItem\]: return \[\ WatchItem(\ topic="FDA guidance on AI/ML-enabled medical device software",\ trusted\_domains=\["fda.gov", "federalregister.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="SEC climate-related disclosure rules for public companies",\ trusted\_domains=\["sec.gov", "federalregister.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="OFAC sanctions list additions and updates",\ trusted\_domains=\["treasury.gov", "ofac.treasury.gov"\],\ team="compliance",\ ),\ WatchItem(\ topic="State-level consumer data privacy laws and amendments",\ trusted\_domains=\["iapp.org", "oag.ca.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="FDA drug recalls and safety communications",\ trusted\_domains=\["fda.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="HIPAA enforcement actions and guidance updates",\ trusted\_domains=\["hhs.gov"\],\ team="compliance",\ ),\ \] @env.task(report=True) async def compliance\_monitoring( watch\_items: list\[WatchItem\] | None = None, freshness: str = "month", ) -> ComplianceReport: """Fan out across regulatory watch items and aggregate triaged findings.""" if watch\_items is None: watch\_items = \_default\_watch\_items() with flyte.group("monitor-watch-items"): results = await asyncio.gather( \*\[monitor\_watch\_item(item, freshness) for item in watch\_items\] ) report = ComplianceReport(findings=\[f for fs in results for f in fs\]) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(compliance\_monitoring) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance\_monitoring\_agent/main.py\* ## Monitor one watch item The \`monitor\_watch\_item\` task researches a single regulatory topic, enriches findings with source metadata from the Research API response, and triages each finding for severity and routing. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "compliance\_monitoring" # params = "" # /// """Regulatory & compliance monitoring agent. Watches trusted regulatory sources via the You.com Research API (with domain/freshness source controls and a structured output schema), then uses Claude to assign severity and route citation-precise findings to the right team. Every external call is traced so Flyte's audit lineage extends to the web layer. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="compliance-monitoring", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="compliance-monitoring", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class WatchItem: topic: str trusted\_domains: list\[str\] team: str @dataclass class Finding: topic: str team: str title: str summary: str source\_url: str published\_date: str snippet: str domain: str = "" favicon: str = "" severity: str = "info" rationale: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class ComplianceReport: findings: list\[Finding\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" FINDINGS\_SCHEMA = { "type": "object", "properties": { "findings": { "type": "array", "items": { "type": "object", "properties": { "title": {"type": "string"}, "summary": {"type": "string"}, "source\_url": {"type": "string"}, "published\_date": {"type": "string"}, "snippet": {"type": "string"}, }, "required": \[\ "title",\ "summary",\ "source\_url",\ "published\_date",\ "snippet",\ \], "additionalProperties": False, }, } }, "required": \["findings"\], "additionalProperties": False, } async def \_you\_post(url: str, body: dict, timeout: float = 300.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research( question: str, include\_domains: list\[str\], freshness: str, research\_effort: str = "standard", ) -> dict: """Call the You.com Research API with domain + freshness source controls.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": { "include\_domains": include\_domains, "freshness": freshness, }, "output\_schema": FINDINGS\_SCHEMA, } return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment llm}} @flyte.trace async def triage(topic: str, findings: list\[dict\]) -> list\[dict\]: """Use Claude to assign a severity + rationale to each finding.""" from litellm import acompletion if not findings: return \[\] system = ( "You are a regulatory-compliance triage analyst. For each finding, " "assign a severity of 'info' (FYI), 'watch' (monitor closely), or " "'action' (requires a concrete compliance/legal response), and a one-" "sentence rationale. Respond ONLY with JSON: " '{"triage": \[{"severity": str, "rationale": str}\]} with one entry per ' "finding, in order." ) listing = "\\n".join( f"\[{i + 1}\] {f.get('title', '')}: {f.get('summary', '')}" for i, f in enumerate(findings) ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": f"Topic: {topic}\\n\\nFindings:\\n{listing}"},\ \], temperature=0.0, max\_tokens=1024, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed.get("triage", \[\]) if isinstance(parsed, dict) else \[\] def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment monitor\_watch\_item}} @env.task(retries=3) async def monitor\_watch\_item(item: WatchItem, freshness: str) -> list\[Finding\]: """Research one regulatory topic and produce triaged, cited findings.""" question = ( f"What are the most recent changes, updates, or new guidance regarding " f"'{item.topic}'? Report concrete, dated changes with their sources." ) result = await you\_research(question, item.trusted\_domains, freshness) output = result.get("output", {}) # Build a lookup from the Research API's full source list (url -> metadata). src\_by\_url: dict\[str, dict\] = {} for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) if url: src\_by\_url\[url\] = s content = output.get("content", {}) if isinstance(content, str): content = \_parse\_json(content) if content.strip() else {} raw\_findings = content.get("findings", \[\]) if isinstance(content, dict) else \[\] triage\_results = await triage(item.topic, raw\_findings) findings: list\[Finding\] = \[\] for i, f in enumerate(raw\_findings): t = triage\_results\[i\] if i < len(triage\_results) else {} url = str(f.get("source\_url", "")) meta = src\_by\_url.get(url, {}) snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or \[""\])\[0\]) findings.append( Finding( topic=item.topic, team=item.team, title=str(f.get("title", "") or meta.get("title", "")), summary=str(f.get("summary", "")), source\_url=url, published\_date=str(f.get("published\_date", "")), snippet=snippet, domain=\_domain(url), favicon=\_favicon\_for(url), severity=str(t.get("severity", "info")), rationale=str(t.get("rationale", "")), ) ) return findings # {{/docs-fragment monitor\_watch\_item}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"action": 0, "watch": 1, "info": 2} \_SEVERITY\_STYLE = { "action": ("#fdecea", "#c0392b"), "watch": ("#fdf3e1", "#b7791f"), "info": ("#e3f1fb", "#2b6cb0"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#edf0f3", "#52606d")) return f"{sev}" def \_cite(f: Finding) -> str: """Render a rich You.com Research citation with domain, date, and snippet.""" if not f.source\_url: return "" meta = f.published\_date\[:10\] if f.published\_date else "" snip = f"
“{f.snippet}”
" if f.snippet else "" return ( f"
" f"
" f"{f.domain or 'source'}" f"research" f"
{meta} · {f.title}
{snip}
" ) def \_render\_report(report: ComplianceReport) -> str: findings = sorted( report.findings, key=lambda f: (\_SEVERITY\_ORDER.get(f.severity, 3), f.team), ) counts = {s: sum(1 for f in findings if f.severity == s) for s in \_SEVERITY\_ORDER} cited = sum(1 for f in findings if f.source\_url) cards = \[\] for f in findings: cards.append( f"
" f"
{\_sev\_badge(f.severity)}{f.team}
" f"

{f.title or f.topic}

" f"
{f.summary}
" f"
{f.rationale}
" f"
{f.topic}
" f"{\_cite(f)}
" ) return f""" {REPORT\_CSS}

Compliance Monitoring Findings

Citation-precise regulatory changes from trusted domains — every finding links to a You.com Research source with snippet provenance.

{len(findings)} findings {cited} cited You.com sources {counts\['action'\]} action {counts\['watch'\]} watch {counts\['info'\]} info
{''.join(cards) or "

No findings in this window.

"}

Findings retrieved via the You.com Research API with source\_control domain allowlists and freshness filters. Flyte logs which agent called which query and got which document — full prompt → citation lineage for audit.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_watch\_items() -> list\[WatchItem\]: return \[\ WatchItem(\ topic="FDA guidance on AI/ML-enabled medical device software",\ trusted\_domains=\["fda.gov", "federalregister.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="SEC climate-related disclosure rules for public companies",\ trusted\_domains=\["sec.gov", "federalregister.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="OFAC sanctions list additions and updates",\ trusted\_domains=\["treasury.gov", "ofac.treasury.gov"\],\ team="compliance",\ ),\ WatchItem(\ topic="State-level consumer data privacy laws and amendments",\ trusted\_domains=\["iapp.org", "oag.ca.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="FDA drug recalls and safety communications",\ trusted\_domains=\["fda.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="HIPAA enforcement actions and guidance updates",\ trusted\_domains=\["hhs.gov"\],\ team="compliance",\ ),\ \] @env.task(report=True) async def compliance\_monitoring( watch\_items: list\[WatchItem\] | None = None, freshness: str = "month", ) -> ComplianceReport: """Fan out across regulatory watch items and aggregate triaged findings.""" if watch\_items is None: watch\_items = \_default\_watch\_items() with flyte.group("monitor-watch-items"): results = await asyncio.gather( \*\[monitor\_watch\_item(item, freshness) for item in watch\_items\] ) report = ComplianceReport(findings=\[f for fs in results for f in fs\]) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(compliance\_monitoring) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance\_monitoring\_agent/main.py\* ## Orchestration The \`compliance\_monitoring\` driver task fans out across all watch items, aggregates findings, and renders a Flyte report sorted by severity and team. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "compliance\_monitoring" # params = "" # /// """Regulatory & compliance monitoring agent. Watches trusted regulatory sources via the You.com Research API (with domain/freshness source controls and a structured output schema), then uses Claude to assign severity and route citation-precise findings to the right team. Every external call is traced so Flyte's audit lineage extends to the web layer. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="compliance-monitoring", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="compliance-monitoring", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class WatchItem: topic: str trusted\_domains: list\[str\] team: str @dataclass class Finding: topic: str team: str title: str summary: str source\_url: str published\_date: str snippet: str domain: str = "" favicon: str = "" severity: str = "info" rationale: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class ComplianceReport: findings: list\[Finding\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" FINDINGS\_SCHEMA = { "type": "object", "properties": { "findings": { "type": "array", "items": { "type": "object", "properties": { "title": {"type": "string"}, "summary": {"type": "string"}, "source\_url": {"type": "string"}, "published\_date": {"type": "string"}, "snippet": {"type": "string"}, }, "required": \[\ "title",\ "summary",\ "source\_url",\ "published\_date",\ "snippet",\ \], "additionalProperties": False, }, } }, "required": \["findings"\], "additionalProperties": False, } async def \_you\_post(url: str, body: dict, timeout: float = 300.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research( question: str, include\_domains: list\[str\], freshness: str, research\_effort: str = "standard", ) -> dict: """Call the You.com Research API with domain + freshness source controls.""" body = { "input": question, "research\_effort": research\_effort, "source\_control": { "include\_domains": include\_domains, "freshness": freshness, }, "output\_schema": FINDINGS\_SCHEMA, } return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment llm}} @flyte.trace async def triage(topic: str, findings: list\[dict\]) -> list\[dict\]: """Use Claude to assign a severity + rationale to each finding.""" from litellm import acompletion if not findings: return \[\] system = ( "You are a regulatory-compliance triage analyst. For each finding, " "assign a severity of 'info' (FYI), 'watch' (monitor closely), or " "'action' (requires a concrete compliance/legal response), and a one-" "sentence rationale. Respond ONLY with JSON: " '{"triage": \[{"severity": str, "rationale": str}\]} with one entry per ' "finding, in order." ) listing = "\\n".join( f"\[{i + 1}\] {f.get('title', '')}: {f.get('summary', '')}" for i, f in enumerate(findings) ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": f"Topic: {topic}\\n\\nFindings:\\n{listing}"},\ \], temperature=0.0, max\_tokens=1024, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed.get("triage", \[\]) if isinstance(parsed, dict) else \[\] def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} # {{docs-fragment monitor\_watch\_item}} @env.task(retries=3) async def monitor\_watch\_item(item: WatchItem, freshness: str) -> list\[Finding\]: """Research one regulatory topic and produce triaged, cited findings.""" question = ( f"What are the most recent changes, updates, or new guidance regarding " f"'{item.topic}'? Report concrete, dated changes with their sources." ) result = await you\_research(question, item.trusted\_domains, freshness) output = result.get("output", {}) # Build a lookup from the Research API's full source list (url -> metadata). src\_by\_url: dict\[str, dict\] = {} for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) if url: src\_by\_url\[url\] = s content = output.get("content", {}) if isinstance(content, str): content = \_parse\_json(content) if content.strip() else {} raw\_findings = content.get("findings", \[\]) if isinstance(content, dict) else \[\] triage\_results = await triage(item.topic, raw\_findings) findings: list\[Finding\] = \[\] for i, f in enumerate(raw\_findings): t = triage\_results\[i\] if i < len(triage\_results) else {} url = str(f.get("source\_url", "")) meta = src\_by\_url.get(url, {}) snippet = str(f.get("snippet", "")) or str((meta.get("snippets") or \[""\])\[0\]) findings.append( Finding( topic=item.topic, team=item.team, title=str(f.get("title", "") or meta.get("title", "")), summary=str(f.get("summary", "")), source\_url=url, published\_date=str(f.get("published\_date", "")), snippet=snippet, domain=\_domain(url), favicon=\_favicon\_for(url), severity=str(t.get("severity", "info")), rationale=str(t.get("rationale", "")), ) ) return findings # {{/docs-fragment monitor\_watch\_item}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"action": 0, "watch": 1, "info": 2} \_SEVERITY\_STYLE = { "action": ("#fdecea", "#c0392b"), "watch": ("#fdf3e1", "#b7791f"), "info": ("#e3f1fb", "#2b6cb0"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#edf0f3", "#52606d")) return f"{sev}" def \_cite(f: Finding) -> str: """Render a rich You.com Research citation with domain, date, and snippet.""" if not f.source\_url: return "" meta = f.published\_date\[:10\] if f.published\_date else "" snip = f"
“{f.snippet}”
" if f.snippet else "" return ( f"
" f"
" f"{f.domain or 'source'}" f"research" f"
{meta} · {f.title}
{snip}
" ) def \_render\_report(report: ComplianceReport) -> str: findings = sorted( report.findings, key=lambda f: (\_SEVERITY\_ORDER.get(f.severity, 3), f.team), ) counts = {s: sum(1 for f in findings if f.severity == s) for s in \_SEVERITY\_ORDER} cited = sum(1 for f in findings if f.source\_url) cards = \[\] for f in findings: cards.append( f"
" f"
{\_sev\_badge(f.severity)}{f.team}
" f"

{f.title or f.topic}

" f"
{f.summary}
" f"
{f.rationale}
" f"
{f.topic}
" f"{\_cite(f)}
" ) return f""" {REPORT\_CSS}

Compliance Monitoring Findings

Citation-precise regulatory changes from trusted domains — every finding links to a You.com Research source with snippet provenance.

{len(findings)} findings {cited} cited You.com sources {counts\['action'\]} action {counts\['watch'\]} watch {counts\['info'\]} info
{''.join(cards) or "

No findings in this window.

"}

Findings retrieved via the You.com Research API with source\_control domain allowlists and freshness filters. Flyte logs which agent called which query and got which document — full prompt → citation lineage for audit.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_watch\_items() -> list\[WatchItem\]: return \[\ WatchItem(\ topic="FDA guidance on AI/ML-enabled medical device software",\ trusted\_domains=\["fda.gov", "federalregister.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="SEC climate-related disclosure rules for public companies",\ trusted\_domains=\["sec.gov", "federalregister.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="OFAC sanctions list additions and updates",\ trusted\_domains=\["treasury.gov", "ofac.treasury.gov"\],\ team="compliance",\ ),\ WatchItem(\ topic="State-level consumer data privacy laws and amendments",\ trusted\_domains=\["iapp.org", "oag.ca.gov"\],\ team="legal",\ ),\ WatchItem(\ topic="FDA drug recalls and safety communications",\ trusted\_domains=\["fda.gov"\],\ team="clinical",\ ),\ WatchItem(\ topic="HIPAA enforcement actions and guidance updates",\ trusted\_domains=\["hhs.gov"\],\ team="compliance",\ ),\ \] @env.task(report=True) async def compliance\_monitoring( watch\_items: list\[WatchItem\] | None = None, freshness: str = "month", ) -> ComplianceReport: """Fan out across regulatory watch items and aggregate triaged findings.""" if watch\_items is None: watch\_items = \_default\_watch\_items() with flyte.group("monitor-watch-items"): results = await asyncio.gather( \*\[monitor\_watch\_item(item, freshness) for item in watch\_items\] ) report = ComplianceReport(findings=\[f for fs in results for f in fs\]) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(compliance\_monitoring) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/compliance\_monitoring\_agent/main.py\* ## Run the agent ### Create secrets Get a You.com API key from the \[You.com platform\](https://you.com/platform) (see the \[quickstart guide\](https://you.com/docs/quickstart)). Get an Anthropic API key from the \[Anthropic console\](https://console.anthropic.com/). Register both keys as Flyte secrets. The secret key names must match those declared in the \`TaskEnvironment\`: \`\`\` flyte create secret youdotcom-api-key flyte create secret internal-anthropic-api-key \`\`\` See \[Secrets\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets/page.md) for scoping and file-based secrets. ### Run locally or remotely From the \[example directory\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/compliance\_monitoring\_agent): \`\`\` cd v2/tutorials/compliance\_monitoring\_agent uv run --script main.py \`\`\` To test locally without Flyte secrets: \`\`\` export YOU\_API\_KEY= export ANTHROPIC\_API\_KEY= uv run --script main.py \`\`\` When the run completes, open the Flyte report to review findings grouped by severity, each with a verifiable You.com Research citation. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/field-data-enrichment-agent === # Field data enrichment agent > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/field\_data\_enrichment\_agent). This example demonstrates how to build an autonomous systems and field-data enrichment agent on Flyte. The agent enriches geo-tagged operational events — from autonomous vehicles, aircraft, satellites, or field sensors — with \*\*real-world public context\*\*: road closures, weather events, airspace changes, or local incidents tied to a geofence. Operational data stays in your environment while public-web grounding queries go to the \[You.com Search API\](https://you.com/docs/search/overview). The API provides unified web and news results with \`freshness\` and \`country\` targeting, and \[Claude\](https://docs.anthropic.com/) via \[LiteLLM\](https://docs.litellm.ai/) summarizes the relevant context for each geo-tagged event. Flyte provides: - \*\*Fan-out parallelism\*\* across geo-tagged events - \*\*\`cache="auto"\`\*\* so repeated geofence checks within the cache window reuse prior results - \*\*\`@flyte.trace\`\*\* on every external call for lineage - \*\*Flyte reports\*\* with operational severity and per-incident citations !\[Field data enrichment agent report\](https://www.union.ai/docs/v2/flyte/\_static/images/tutorials/field\_data\_enrichment\_agent/field-data-enrichment-data.png) ## Setting up the environment The agent runs in a \`TaskEnvironment\` with secrets for the You.com and Anthropic API keys, automatic caching, and a container image built from the \`uv\` script dependencies. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "field\_data\_enrichment" # params = "" # /// """Autonomous systems & field-data enrichment agent. Enriches geo-tagged operational events with real-world public context (road closures, weather, incidents) using the You.com Search API with country + freshness targeting, then uses Claude to summarize the relevant context. Only public-web grounding queries leave the customer's cloud, never operational data. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="field-data-enrichment", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="field-data-enrichment", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class GeoEvent: event\_id: str location: str country: str event\_type: str @dataclass class Incident: description: str source\_url: str published: str domain: str = "" author: str = "" favicon: str = "" snippet: str = "" section: str = "web" @dataclass class EnrichedEvent: event\_id: str location: str context\_summary: str severity: str incidents: list\[Incident\] = field(default\_factory=list) @dataclass class EnrichmentReport: events: list\[EnrichedEvent\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" @dataclass class SearchHit: title: str url: str domain: str snippet: str published: str author: str favicon: str section: str def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_search( query: str, country: str, freshness: str = "day", count: int = 8 ) -> list\[SearchHit\]: """Search the public web + news for context near a geofenced location.""" params = { "query": query, "count": count, "freshness": freshness, "country": country, } data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict: from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} ENRICH\_SYSTEM = """You are an operational-context analyst for autonomous and \\ field systems. Given fresh local search results near a geofenced location, \\ summarize the real-world context relevant to operations, extract discrete \\ incidents (road closures, weather events, regulatory/airspace changes, local \\ incidents), and assign an operational severity of 'none', 'low', 'medium', or \\ 'high'. Each incident must reference the supporting search result by its index. \\ Respond ONLY with JSON: {"context\_summary": str, "severity": str, "incidents": \[{"description": str, \\\ "source\_index": int (the \[n\] of the supporting search result)}\]}""" # {{docs-fragment enrich\_event}} @env.task(retries=3) async def enrich\_event(event: GeoEvent, freshness: str) -> EnrichedEvent: """Ground one geo-tagged event in fresh public context.""" query = f"{event.location} {event.event\_type.replace('\_', ' ')} road closure weather incident" hits = await you\_search(query, country=event.country, freshness=freshness) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Location: {event.location}\\n" f"Event type: {event.event\_type}\\n\\n" f"Search results:\\n{evidence or 'No results.'}" ) parsed = await llm\_json(ENRICH\_SYSTEM, user) def \_incident(it: dict) -> Incident: idx = int(it.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None return Incident( description=str(it.get("description", "")), source\_url=src.url if src else "", published=src.published if src else "", domain=src.domain if src else "", author=src.author if src else "", favicon=src.favicon if src else "", snippet=src.snippet if src else "", section=src.section if src else "web", ) incidents = \[\_incident(it) for it in (parsed.get("incidents", \[\]) or \[\])\] return EnrichedEvent( event\_id=event.event\_id, location=event.location, context\_summary=str(parsed.get("context\_summary", "")), severity=str(parsed.get("severity", "none")), incidents=incidents, ) # {{/docs-fragment enrich\_event}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"high": 0, "medium": 1, "low": 2, "none": 3} \_SEVERITY\_STYLE = { "high": ("#fdecea", "#c0392b"), "medium": ("#fdf3e1", "#b7791f"), "low": ("#e3f1fb", "#2b6cb0"), "none": ("#eef1f4", "#627d98"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#eef1f4", "#627d98")) return f"{sev}" def \_cite(it: Incident) -> str: """Render a rich You.com citation for an incident's supporting source.""" if not it.source\_url: return "" tag = ( "news" if it.section == "news" else "web" ) meta\_bits = \[\] if it.published: meta\_bits.append(it.published\[:10\]) if it.author: meta\_bits.append(f"by {it.author}") meta = " · ".join(meta\_bits) snip = f"
“{it.snippet}”
" if it.snippet else "" return ( f"
" f"
" f"{it.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: EnrichmentReport) -> str: events = sorted(report.events, key=lambda e: \_SEVERITY\_ORDER.get(e.severity, 4)) flagged = sum(1 for e in events if e.severity in ("high", "medium")) total\_sources = sum(len(e.incidents) for e in events) cards = \[\] for e in events: incidents = "".join( f"
• {it.description}{\_cite(it)}
" for it in e.incidents ) cards.append( f"
" f"
{\_sev\_badge(e.severity)}" f"{e.event\_id} · {e.location}
" f"
{e.context\_summary or 'No relevant public context found.'}
" f"{incidents}
" ) return f""" {REPORT\_CSS}

Field-Data Enrichment

Geo-tagged events grounded in fresh public context — each incident cites a timestamped You.com Search result.

{len(events)} events {flagged} flagged (high/medium) {total\_sources} cited You.com sources
{''.join(cards) or "

No events processed.

"}

Public context retrieved via the You.com Search API with country + freshness targeting. Operational data never leaves the BYOC boundary — only public-web queries go out.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} DEFAULT\_EVENTS = \[\ GeoEvent("evt-1", "Mountain View, CA", "US", "road\_closure\_check"),\ GeoEvent("evt-2", "Tokyo, Japan", "JP", "weather"),\ GeoEvent("evt-3", "Austin, TX", "US", "road\_closure\_check"),\ GeoEvent("evt-4", "Phoenix, AZ", "US", "weather"),\ GeoEvent("evt-5", "London, UK", "GB", "incident"),\ GeoEvent("evt-6", "San Francisco, CA", "US", "incident"),\ GeoEvent("evt-7", "Seattle, WA", "US", "weather"),\ GeoEvent("evt-8", "Miami, FL", "US", "weather"),\ GeoEvent("evt-9", "Denver, CO", "US", "road\_closure\_check"),\ GeoEvent("evt-10", "Berlin, Germany", "DE", "incident"),\ \] @env.task(report=True) async def field\_data\_enrichment( events: list\[GeoEvent\] = DEFAULT\_EVENTS, freshness: str = "day", ) -> EnrichmentReport: """Fan out across geo-tagged events and enrich each with public context.""" with flyte.group("enrich-events"): enriched = await asyncio.gather( \*\[enrich\_event(e, freshness) for e in events\] ) report = EnrichmentReport(events=list(enriched)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(field\_data\_enrichment) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/field\_data\_enrichment\_agent/main.py\* The Python packages are declared at the top of the file using the \`uv\` script style: \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # /// \`\`\` ## Data types Each \`GeoEvent\` carries an event ID, location, ISO country code for geo-targeting, and an event type. Enriched events include a context summary, operational severity, and discrete incidents with source citations. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "field\_data\_enrichment" # params = "" # /// """Autonomous systems & field-data enrichment agent. Enriches geo-tagged operational events with real-world public context (road closures, weather, incidents) using the You.com Search API with country + freshness targeting, then uses Claude to summarize the relevant context. Only public-web grounding queries leave the customer's cloud, never operational data. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="field-data-enrichment", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="field-data-enrichment", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class GeoEvent: event\_id: str location: str country: str event\_type: str @dataclass class Incident: description: str source\_url: str published: str domain: str = "" author: str = "" favicon: str = "" snippet: str = "" section: str = "web" @dataclass class EnrichedEvent: event\_id: str location: str context\_summary: str severity: str incidents: list\[Incident\] = field(default\_factory=list) @dataclass class EnrichmentReport: events: list\[EnrichedEvent\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" @dataclass class SearchHit: title: str url: str domain: str snippet: str published: str author: str favicon: str section: str def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_search( query: str, country: str, freshness: str = "day", count: int = 8 ) -> list\[SearchHit\]: """Search the public web + news for context near a geofenced location.""" params = { "query": query, "count": count, "freshness": freshness, "country": country, } data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict: from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} ENRICH\_SYSTEM = """You are an operational-context analyst for autonomous and \\ field systems. Given fresh local search results near a geofenced location, \\ summarize the real-world context relevant to operations, extract discrete \\ incidents (road closures, weather events, regulatory/airspace changes, local \\ incidents), and assign an operational severity of 'none', 'low', 'medium', or \\ 'high'. Each incident must reference the supporting search result by its index. \\ Respond ONLY with JSON: {"context\_summary": str, "severity": str, "incidents": \[{"description": str, \\\ "source\_index": int (the \[n\] of the supporting search result)}\]}""" # {{docs-fragment enrich\_event}} @env.task(retries=3) async def enrich\_event(event: GeoEvent, freshness: str) -> EnrichedEvent: """Ground one geo-tagged event in fresh public context.""" query = f"{event.location} {event.event\_type.replace('\_', ' ')} road closure weather incident" hits = await you\_search(query, country=event.country, freshness=freshness) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Location: {event.location}\\n" f"Event type: {event.event\_type}\\n\\n" f"Search results:\\n{evidence or 'No results.'}" ) parsed = await llm\_json(ENRICH\_SYSTEM, user) def \_incident(it: dict) -> Incident: idx = int(it.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None return Incident( description=str(it.get("description", "")), source\_url=src.url if src else "", published=src.published if src else "", domain=src.domain if src else "", author=src.author if src else "", favicon=src.favicon if src else "", snippet=src.snippet if src else "", section=src.section if src else "web", ) incidents = \[\_incident(it) for it in (parsed.get("incidents", \[\]) or \[\])\] return EnrichedEvent( event\_id=event.event\_id, location=event.location, context\_summary=str(parsed.get("context\_summary", "")), severity=str(parsed.get("severity", "none")), incidents=incidents, ) # {{/docs-fragment enrich\_event}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"high": 0, "medium": 1, "low": 2, "none": 3} \_SEVERITY\_STYLE = { "high": ("#fdecea", "#c0392b"), "medium": ("#fdf3e1", "#b7791f"), "low": ("#e3f1fb", "#2b6cb0"), "none": ("#eef1f4", "#627d98"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#eef1f4", "#627d98")) return f"{sev}" def \_cite(it: Incident) -> str: """Render a rich You.com citation for an incident's supporting source.""" if not it.source\_url: return "" tag = ( "news" if it.section == "news" else "web" ) meta\_bits = \[\] if it.published: meta\_bits.append(it.published\[:10\]) if it.author: meta\_bits.append(f"by {it.author}") meta = " · ".join(meta\_bits) snip = f"
“{it.snippet}”
" if it.snippet else "" return ( f"
" f"
" f"{it.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: EnrichmentReport) -> str: events = sorted(report.events, key=lambda e: \_SEVERITY\_ORDER.get(e.severity, 4)) flagged = sum(1 for e in events if e.severity in ("high", "medium")) total\_sources = sum(len(e.incidents) for e in events) cards = \[\] for e in events: incidents = "".join( f"
• {it.description}{\_cite(it)}
" for it in e.incidents ) cards.append( f"
" f"
{\_sev\_badge(e.severity)}" f"{e.event\_id} · {e.location}
" f"
{e.context\_summary or 'No relevant public context found.'}
" f"{incidents}
" ) return f""" {REPORT\_CSS}

Field-Data Enrichment

Geo-tagged events grounded in fresh public context — each incident cites a timestamped You.com Search result.

{len(events)} events {flagged} flagged (high/medium) {total\_sources} cited You.com sources
{''.join(cards) or "

No events processed.

"}

Public context retrieved via the You.com Search API with country + freshness targeting. Operational data never leaves the BYOC boundary — only public-web queries go out.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} DEFAULT\_EVENTS = \[\ GeoEvent("evt-1", "Mountain View, CA", "US", "road\_closure\_check"),\ GeoEvent("evt-2", "Tokyo, Japan", "JP", "weather"),\ GeoEvent("evt-3", "Austin, TX", "US", "road\_closure\_check"),\ GeoEvent("evt-4", "Phoenix, AZ", "US", "weather"),\ GeoEvent("evt-5", "London, UK", "GB", "incident"),\ GeoEvent("evt-6", "San Francisco, CA", "US", "incident"),\ GeoEvent("evt-7", "Seattle, WA", "US", "weather"),\ GeoEvent("evt-8", "Miami, FL", "US", "weather"),\ GeoEvent("evt-9", "Denver, CO", "US", "road\_closure\_check"),\ GeoEvent("evt-10", "Berlin, Germany", "DE", "incident"),\ \] @env.task(report=True) async def field\_data\_enrichment( events: list\[GeoEvent\] = DEFAULT\_EVENTS, freshness: str = "day", ) -> EnrichmentReport: """Fan out across geo-tagged events and enrich each with public context.""" with flyte.group("enrich-events"): enriched = await asyncio.gather( \*\[enrich\_event(e, freshness) for e in events\] ) report = EnrichmentReport(events=list(enriched)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(field\_data\_enrichment) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/field\_data\_enrichment\_agent/main.py\* ## Search with the You.com Search API The \`you\_search\` helper calls the \[You.com Search API\](https://you.com/docs/search/overview) with \`freshness\` and \`country\` parameters to retrieve location-relevant web and news results. See the \[Search API reference\](https://you.com/docs/api-reference/search/v1-search) for supported country codes and freshness values. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "field\_data\_enrichment" # params = "" # /// """Autonomous systems & field-data enrichment agent. Enriches geo-tagged operational events with real-world public context (road closures, weather, incidents) using the You.com Search API with country + freshness targeting, then uses Claude to summarize the relevant context. Only public-web grounding queries leave the customer's cloud, never operational data. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="field-data-enrichment", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="field-data-enrichment", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class GeoEvent: event\_id: str location: str country: str event\_type: str @dataclass class Incident: description: str source\_url: str published: str domain: str = "" author: str = "" favicon: str = "" snippet: str = "" section: str = "web" @dataclass class EnrichedEvent: event\_id: str location: str context\_summary: str severity: str incidents: list\[Incident\] = field(default\_factory=list) @dataclass class EnrichmentReport: events: list\[EnrichedEvent\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" @dataclass class SearchHit: title: str url: str domain: str snippet: str published: str author: str favicon: str section: str def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_search( query: str, country: str, freshness: str = "day", count: int = 8 ) -> list\[SearchHit\]: """Search the public web + news for context near a geofenced location.""" params = { "query": query, "count": count, "freshness": freshness, "country": country, } data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict: from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} ENRICH\_SYSTEM = """You are an operational-context analyst for autonomous and \\ field systems. Given fresh local search results near a geofenced location, \\ summarize the real-world context relevant to operations, extract discrete \\ incidents (road closures, weather events, regulatory/airspace changes, local \\ incidents), and assign an operational severity of 'none', 'low', 'medium', or \\ 'high'. Each incident must reference the supporting search result by its index. \\ Respond ONLY with JSON: {"context\_summary": str, "severity": str, "incidents": \[{"description": str, \\\ "source\_index": int (the \[n\] of the supporting search result)}\]}""" # {{docs-fragment enrich\_event}} @env.task(retries=3) async def enrich\_event(event: GeoEvent, freshness: str) -> EnrichedEvent: """Ground one geo-tagged event in fresh public context.""" query = f"{event.location} {event.event\_type.replace('\_', ' ')} road closure weather incident" hits = await you\_search(query, country=event.country, freshness=freshness) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Location: {event.location}\\n" f"Event type: {event.event\_type}\\n\\n" f"Search results:\\n{evidence or 'No results.'}" ) parsed = await llm\_json(ENRICH\_SYSTEM, user) def \_incident(it: dict) -> Incident: idx = int(it.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None return Incident( description=str(it.get("description", "")), source\_url=src.url if src else "", published=src.published if src else "", domain=src.domain if src else "", author=src.author if src else "", favicon=src.favicon if src else "", snippet=src.snippet if src else "", section=src.section if src else "web", ) incidents = \[\_incident(it) for it in (parsed.get("incidents", \[\]) or \[\])\] return EnrichedEvent( event\_id=event.event\_id, location=event.location, context\_summary=str(parsed.get("context\_summary", "")), severity=str(parsed.get("severity", "none")), incidents=incidents, ) # {{/docs-fragment enrich\_event}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"high": 0, "medium": 1, "low": 2, "none": 3} \_SEVERITY\_STYLE = { "high": ("#fdecea", "#c0392b"), "medium": ("#fdf3e1", "#b7791f"), "low": ("#e3f1fb", "#2b6cb0"), "none": ("#eef1f4", "#627d98"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#eef1f4", "#627d98")) return f"{sev}" def \_cite(it: Incident) -> str: """Render a rich You.com citation for an incident's supporting source.""" if not it.source\_url: return "" tag = ( "news" if it.section == "news" else "web" ) meta\_bits = \[\] if it.published: meta\_bits.append(it.published\[:10\]) if it.author: meta\_bits.append(f"by {it.author}") meta = " · ".join(meta\_bits) snip = f"
“{it.snippet}”
" if it.snippet else "" return ( f"
" f"
" f"{it.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: EnrichmentReport) -> str: events = sorted(report.events, key=lambda e: \_SEVERITY\_ORDER.get(e.severity, 4)) flagged = sum(1 for e in events if e.severity in ("high", "medium")) total\_sources = sum(len(e.incidents) for e in events) cards = \[\] for e in events: incidents = "".join( f"
• {it.description}{\_cite(it)}
" for it in e.incidents ) cards.append( f"
" f"
{\_sev\_badge(e.severity)}" f"{e.event\_id} · {e.location}
" f"
{e.context\_summary or 'No relevant public context found.'}
" f"{incidents}
" ) return f""" {REPORT\_CSS}

Field-Data Enrichment

Geo-tagged events grounded in fresh public context — each incident cites a timestamped You.com Search result.

{len(events)} events {flagged} flagged (high/medium) {total\_sources} cited You.com sources
{''.join(cards) or "

No events processed.

"}

Public context retrieved via the You.com Search API with country + freshness targeting. Operational data never leaves the BYOC boundary — only public-web queries go out.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} DEFAULT\_EVENTS = \[\ GeoEvent("evt-1", "Mountain View, CA", "US", "road\_closure\_check"),\ GeoEvent("evt-2", "Tokyo, Japan", "JP", "weather"),\ GeoEvent("evt-3", "Austin, TX", "US", "road\_closure\_check"),\ GeoEvent("evt-4", "Phoenix, AZ", "US", "weather"),\ GeoEvent("evt-5", "London, UK", "GB", "incident"),\ GeoEvent("evt-6", "San Francisco, CA", "US", "incident"),\ GeoEvent("evt-7", "Seattle, WA", "US", "weather"),\ GeoEvent("evt-8", "Miami, FL", "US", "weather"),\ GeoEvent("evt-9", "Denver, CO", "US", "road\_closure\_check"),\ GeoEvent("evt-10", "Berlin, Germany", "DE", "incident"),\ \] @env.task(report=True) async def field\_data\_enrichment( events: list\[GeoEvent\] = DEFAULT\_EVENTS, freshness: str = "day", ) -> EnrichmentReport: """Fan out across geo-tagged events and enrich each with public context.""" with flyte.group("enrich-events"): enriched = await asyncio.gather( \*\[enrich\_event(e, freshness) for e in events\] ) report = EnrichmentReport(events=list(enriched)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(field\_data\_enrichment) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/field\_data\_enrichment\_agent/main.py\* ## Enrich one event The \`enrich\_event\` task builds a location- and type-scoped query, calls the You.com Search API, and asks Claude to summarize relevant real-world context, extract discrete incidents, and assign an operational severity — all grounded in the returned sources. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "field\_data\_enrichment" # params = "" # /// """Autonomous systems & field-data enrichment agent. Enriches geo-tagged operational events with real-world public context (road closures, weather, incidents) using the You.com Search API with country + freshness targeting, then uses Claude to summarize the relevant context. Only public-web grounding queries leave the customer's cloud, never operational data. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="field-data-enrichment", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="field-data-enrichment", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class GeoEvent: event\_id: str location: str country: str event\_type: str @dataclass class Incident: description: str source\_url: str published: str domain: str = "" author: str = "" favicon: str = "" snippet: str = "" section: str = "web" @dataclass class EnrichedEvent: event\_id: str location: str context\_summary: str severity: str incidents: list\[Incident\] = field(default\_factory=list) @dataclass class EnrichmentReport: events: list\[EnrichedEvent\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" @dataclass class SearchHit: title: str url: str domain: str snippet: str published: str author: str favicon: str section: str def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_search( query: str, country: str, freshness: str = "day", count: int = 8 ) -> list\[SearchHit\]: """Search the public web + news for context near a geofenced location.""" params = { "query": query, "count": count, "freshness": freshness, "country": country, } data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict: from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} ENRICH\_SYSTEM = """You are an operational-context analyst for autonomous and \\ field systems. Given fresh local search results near a geofenced location, \\ summarize the real-world context relevant to operations, extract discrete \\ incidents (road closures, weather events, regulatory/airspace changes, local \\ incidents), and assign an operational severity of 'none', 'low', 'medium', or \\ 'high'. Each incident must reference the supporting search result by its index. \\ Respond ONLY with JSON: {"context\_summary": str, "severity": str, "incidents": \[{"description": str, \\\ "source\_index": int (the \[n\] of the supporting search result)}\]}""" # {{docs-fragment enrich\_event}} @env.task(retries=3) async def enrich\_event(event: GeoEvent, freshness: str) -> EnrichedEvent: """Ground one geo-tagged event in fresh public context.""" query = f"{event.location} {event.event\_type.replace('\_', ' ')} road closure weather incident" hits = await you\_search(query, country=event.country, freshness=freshness) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Location: {event.location}\\n" f"Event type: {event.event\_type}\\n\\n" f"Search results:\\n{evidence or 'No results.'}" ) parsed = await llm\_json(ENRICH\_SYSTEM, user) def \_incident(it: dict) -> Incident: idx = int(it.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None return Incident( description=str(it.get("description", "")), source\_url=src.url if src else "", published=src.published if src else "", domain=src.domain if src else "", author=src.author if src else "", favicon=src.favicon if src else "", snippet=src.snippet if src else "", section=src.section if src else "web", ) incidents = \[\_incident(it) for it in (parsed.get("incidents", \[\]) or \[\])\] return EnrichedEvent( event\_id=event.event\_id, location=event.location, context\_summary=str(parsed.get("context\_summary", "")), severity=str(parsed.get("severity", "none")), incidents=incidents, ) # {{/docs-fragment enrich\_event}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"high": 0, "medium": 1, "low": 2, "none": 3} \_SEVERITY\_STYLE = { "high": ("#fdecea", "#c0392b"), "medium": ("#fdf3e1", "#b7791f"), "low": ("#e3f1fb", "#2b6cb0"), "none": ("#eef1f4", "#627d98"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#eef1f4", "#627d98")) return f"{sev}" def \_cite(it: Incident) -> str: """Render a rich You.com citation for an incident's supporting source.""" if not it.source\_url: return "" tag = ( "news" if it.section == "news" else "web" ) meta\_bits = \[\] if it.published: meta\_bits.append(it.published\[:10\]) if it.author: meta\_bits.append(f"by {it.author}") meta = " · ".join(meta\_bits) snip = f"
“{it.snippet}”
" if it.snippet else "" return ( f"
" f"
" f"{it.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: EnrichmentReport) -> str: events = sorted(report.events, key=lambda e: \_SEVERITY\_ORDER.get(e.severity, 4)) flagged = sum(1 for e in events if e.severity in ("high", "medium")) total\_sources = sum(len(e.incidents) for e in events) cards = \[\] for e in events: incidents = "".join( f"
• {it.description}{\_cite(it)}
" for it in e.incidents ) cards.append( f"
" f"
{\_sev\_badge(e.severity)}" f"{e.event\_id} · {e.location}
" f"
{e.context\_summary or 'No relevant public context found.'}
" f"{incidents}
" ) return f""" {REPORT\_CSS}

Field-Data Enrichment

Geo-tagged events grounded in fresh public context — each incident cites a timestamped You.com Search result.

{len(events)} events {flagged} flagged (high/medium) {total\_sources} cited You.com sources
{''.join(cards) or "

No events processed.

"}

Public context retrieved via the You.com Search API with country + freshness targeting. Operational data never leaves the BYOC boundary — only public-web queries go out.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} DEFAULT\_EVENTS = \[\ GeoEvent("evt-1", "Mountain View, CA", "US", "road\_closure\_check"),\ GeoEvent("evt-2", "Tokyo, Japan", "JP", "weather"),\ GeoEvent("evt-3", "Austin, TX", "US", "road\_closure\_check"),\ GeoEvent("evt-4", "Phoenix, AZ", "US", "weather"),\ GeoEvent("evt-5", "London, UK", "GB", "incident"),\ GeoEvent("evt-6", "San Francisco, CA", "US", "incident"),\ GeoEvent("evt-7", "Seattle, WA", "US", "weather"),\ GeoEvent("evt-8", "Miami, FL", "US", "weather"),\ GeoEvent("evt-9", "Denver, CO", "US", "road\_closure\_check"),\ GeoEvent("evt-10", "Berlin, Germany", "DE", "incident"),\ \] @env.task(report=True) async def field\_data\_enrichment( events: list\[GeoEvent\] = DEFAULT\_EVENTS, freshness: str = "day", ) -> EnrichmentReport: """Fan out across geo-tagged events and enrich each with public context.""" with flyte.group("enrich-events"): enriched = await asyncio.gather( \*\[enrich\_event(e, freshness) for e in events\] ) report = EnrichmentReport(events=list(enriched)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(field\_data\_enrichment) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/field\_data\_enrichment\_agent/main.py\* ## Orchestration The \`field\_data\_enrichment\` driver task fans out across all events and renders a Flyte report sorted by severity. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "field\_data\_enrichment" # params = "" # /// """Autonomous systems & field-data enrichment agent. Enriches geo-tagged operational events with real-world public context (road closures, weather, incidents) using the You.com Search API with country + freshness targeting, then uses Claude to summarize the relevant context. Only public-web grounding queries leave the customer's cloud, never operational data. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="field-data-enrichment", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="field-data-enrichment", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), cache="auto", ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class GeoEvent: event\_id: str location: str country: str event\_type: str @dataclass class Incident: description: str source\_url: str published: str domain: str = "" author: str = "" favicon: str = "" snippet: str = "" section: str = "web" @dataclass class EnrichedEvent: event\_id: str location: str context\_summary: str severity: str incidents: list\[Incident\] = field(default\_factory=list) @dataclass class EnrichmentReport: events: list\[EnrichedEvent\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_search}} YOU\_SEARCH\_URL = "https://ydc-index.io/v1/search" @dataclass class SearchHit: title: str url: str domain: str snippet: str published: str author: str favicon: str section: str def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon(item: dict, url: str) -> str: return item.get("favicon\_url") or ( f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" ) async def \_you\_get(url: str, params: dict, timeout: float = 60.0) -> dict: """GET with exponential backoff + jitter on 429 rate limits.""" import asyncio import random import httpx headers = {"X-API-Key": os.environ\["YOU\_API\_KEY"\]} async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.get(url, headers=headers, params=params) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_search( query: str, country: str, freshness: str = "day", count: int = 8 ) -> list\[SearchHit\]: """Search the public web + news for context near a geofenced location.""" params = { "query": query, "count": count, "freshness": freshness, "country": country, } data = await \_you\_get(YOU\_SEARCH\_URL, params) results = data.get("results", {}) hits: list\[SearchHit\] = \[\] for section in ("news", "web"): for item in results.get(section, \[\]) or \[\]: snippets = item.get("snippets") or \[\] url = item.get("url", "") hits.append( SearchHit( title=item.get("title", ""), url=url, domain=\_domain(url), snippet=(snippets\[0\] if snippets else item.get("description", "")), published=item.get("page\_age", "") or "", author=", ".join(item.get("authors") or \[\]), favicon=\_favicon(item, url), section=section, ) ) return hits # {{/docs-fragment you\_search}} # {{docs-fragment llm}} @flyte.trace async def llm\_json(system: str, user: str) -> dict: from litellm import acompletion resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.0, max\_tokens=1536, ) parsed = \_parse\_json(resp.choices\[0\].message.content) return parsed if isinstance(parsed, dict) else {} def \_parse\_json(text: str) -> dict | list: text = text.strip() if text.startswith("\`\`\`"): text = text.split("\`\`\`", 2)\[1\] if text.lstrip().startswith("json"): text = text.lstrip()\[4:\] start = min((i for i in (text.find("{"), text.find("\[")) if i != -1), default=0)\ end = max(text.rfind("}"), text.rfind("\]")) + 1 return json.loads(text\[start:end\]) # {{/docs-fragment llm}} ENRICH\_SYSTEM = """You are an operational-context analyst for autonomous and \\ field systems. Given fresh local search results near a geofenced location, \\ summarize the real-world context relevant to operations, extract discrete \\ incidents (road closures, weather events, regulatory/airspace changes, local \\ incidents), and assign an operational severity of 'none', 'low', 'medium', or \\ 'high'. Each incident must reference the supporting search result by its index. \\ Respond ONLY with JSON: {"context\_summary": str, "severity": str, "incidents": \[{"description": str, \\\ "source\_index": int (the \[n\] of the supporting search result)}\]}""" # {{docs-fragment enrich\_event}} @env.task(retries=3) async def enrich\_event(event: GeoEvent, freshness: str) -> EnrichedEvent: """Ground one geo-tagged event in fresh public context.""" query = f"{event.location} {event.event\_type.replace('\_', ' ')} road closure weather incident" hits = await you\_search(query, country=event.country, freshness=freshness) evidence = "\\n\\n".join( f"\[{i + 1}\] {h.title} ({h.published}) — {h.domain}\\n{h.url}\\n{h.snippet}" for i, h in enumerate(hits) ) user = ( f"Location: {event.location}\\n" f"Event type: {event.event\_type}\\n\\n" f"Search results:\\n{evidence or 'No results.'}" ) parsed = await llm\_json(ENRICH\_SYSTEM, user) def \_incident(it: dict) -> Incident: idx = int(it.get("source\_index", 0) or 0) src = hits\[idx - 1\] if 1 <= idx <= len(hits) else None return Incident( description=str(it.get("description", "")), source\_url=src.url if src else "", published=src.published if src else "", domain=src.domain if src else "", author=src.author if src else "", favicon=src.favicon if src else "", snippet=src.snippet if src else "", section=src.section if src else "web", ) incidents = \[\_incident(it) for it in (parsed.get("incidents", \[\]) or \[\])\] return EnrichedEvent( event\_id=event.event\_id, location=event.location, context\_summary=str(parsed.get("context\_summary", "")), severity=str(parsed.get("severity", "none")), incidents=incidents, ) # {{/docs-fragment enrich\_event}} # {{docs-fragment report}} \_SEVERITY\_ORDER = {"high": 0, "medium": 1, "low": 2, "none": 3} \_SEVERITY\_STYLE = { "high": ("#fdecea", "#c0392b"), "medium": ("#fdf3e1", "#b7791f"), "low": ("#e3f1fb", "#2b6cb0"), "none": ("#eef1f4", "#627d98"), } REPORT\_CSS = """ """ def \_sev\_badge(sev: str) -> str: bg, fg = \_SEVERITY\_STYLE.get(sev, ("#eef1f4", "#627d98")) return f"{sev}" def \_cite(it: Incident) -> str: """Render a rich You.com citation for an incident's supporting source.""" if not it.source\_url: return "" tag = ( "news" if it.section == "news" else "web" ) meta\_bits = \[\] if it.published: meta\_bits.append(it.published\[:10\]) if it.author: meta\_bits.append(f"by {it.author}") meta = " · ".join(meta\_bits) snip = f"
“{it.snippet}”
" if it.snippet else "" return ( f"
" f"
" f"{it.domain or 'source'}{tag}" f"
{meta}
{snip}
" ) def \_render\_report(report: EnrichmentReport) -> str: events = sorted(report.events, key=lambda e: \_SEVERITY\_ORDER.get(e.severity, 4)) flagged = sum(1 for e in events if e.severity in ("high", "medium")) total\_sources = sum(len(e.incidents) for e in events) cards = \[\] for e in events: incidents = "".join( f"
• {it.description}{\_cite(it)}
" for it in e.incidents ) cards.append( f"
" f"
{\_sev\_badge(e.severity)}" f"{e.event\_id} · {e.location}
" f"
{e.context\_summary or 'No relevant public context found.'}
" f"{incidents}
" ) return f""" {REPORT\_CSS}

Field-Data Enrichment

Geo-tagged events grounded in fresh public context — each incident cites a timestamped You.com Search result.

{len(events)} events {flagged} flagged (high/medium) {total\_sources} cited You.com sources
{''.join(cards) or "

No events processed.

"}

Public context retrieved via the You.com Search API with country + freshness targeting. Operational data never leaves the BYOC boundary — only public-web queries go out.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} DEFAULT\_EVENTS = \[\ GeoEvent("evt-1", "Mountain View, CA", "US", "road\_closure\_check"),\ GeoEvent("evt-2", "Tokyo, Japan", "JP", "weather"),\ GeoEvent("evt-3", "Austin, TX", "US", "road\_closure\_check"),\ GeoEvent("evt-4", "Phoenix, AZ", "US", "weather"),\ GeoEvent("evt-5", "London, UK", "GB", "incident"),\ GeoEvent("evt-6", "San Francisco, CA", "US", "incident"),\ GeoEvent("evt-7", "Seattle, WA", "US", "weather"),\ GeoEvent("evt-8", "Miami, FL", "US", "weather"),\ GeoEvent("evt-9", "Denver, CO", "US", "road\_closure\_check"),\ GeoEvent("evt-10", "Berlin, Germany", "DE", "incident"),\ \] @env.task(report=True) async def field\_data\_enrichment( events: list\[GeoEvent\] = DEFAULT\_EVENTS, freshness: str = "day", ) -> EnrichmentReport: """Fan out across geo-tagged events and enrich each with public context.""" with flyte.group("enrich-events"): enriched = await asyncio.gather( \*\[enrich\_event(e, freshness) for e in events\] ) report = EnrichmentReport(events=list(enriched)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(field\_data\_enrichment) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/field\_data\_enrichment\_agent/main.py\* ## Run the agent ### Create secrets Get a You.com API key from the \[You.com platform\](https://you.com/platform) (see the \[quickstart guide\](https://you.com/docs/quickstart)). Get an Anthropic API key from the \[Anthropic console\](https://console.anthropic.com/). Register both keys as Flyte secrets. The secret key names must match those declared in the \`TaskEnvironment\`: \`\`\` flyte create secret youdotcom-api-key flyte create secret internal-anthropic-api-key \`\`\` See \[Secrets\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets/page.md) for scoping and file-based secrets. ### Run locally or remotely From the \[example directory\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/field\_data\_enrichment\_agent): \`\`\` cd v2/tutorials/field\_data\_enrichment\_agent uv run --script main.py \`\`\` To test locally without Flyte secrets: \`\`\` export YOU\_API\_KEY= export ANTHROPIC\_API\_KEY= uv run --script main.py \`\`\` When the run completes, open the Flyte report to review enriched events with operational severity and timestamped You.com source citations for each incident. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/agents/support-resolution-agent === # Support resolution agent > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/support\_resolution\_agent). This example demonstrates how to build a customer-support and field-service resolution agent on Flyte. The agent resolves tickets that need current public information — return policies, weather advisories, product recalls, manufacturer specs — and drafts a customer-ready reply with sources a human agent can verify before sending. The \[You.com Research API\](https://you.com/docs/research/overview) grounds each ticket in fresh, citable sources. \[Claude\](https://docs.anthropic.com/) via \[LiteLLM\](https://docs.litellm.ai/) turns that research into a reply draft. With \`research\_effort="lite"\`, the research step stays fast enough for human-in-the-loop support flows. Flyte provides: - \*\*Fan-out parallelism\*\* across support tickets - \*\*\`@flyte.trace\`\*\* on every external call for lineage - A \*\*two-step pipeline\*\* per ticket: ground the answer, then draft the reply - \*\*Flyte reports\*\* with draft replies and verifiable source citations !\[Support resolution agent report\](https://www.union.ai/docs/v2/flyte/\_static/images/tutorials/support\_resolution\_agent/support-resolutions-agent.png) ## Setting up the environment The agent runs in a \`TaskEnvironment\` with secrets for the You.com and Anthropic API keys and a container image built from the \`uv\` script dependencies. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "support\_resolution" # params = "" # /// """Customer-support & field-service resolution agent. Grounds a support ticket in fresh, public, citable sources via the You.com Research API (low effort for low latency, human-in-the-loop use), then uses Claude to draft a customer-ready reply that cites its sources inline so a human agent can verify before sending. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="support-resolution", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="support-resolution", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str snippet: str domain: str = "" favicon: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Ticket: ticket\_id: str question: str context: str = "" @dataclass class Grounding: answer: str sources: list\[Source\] = field(default\_factory=list) @dataclass class Resolution: ticket\_id: str ticket: str grounded\_answer: str draft\_reply: str sources: list\[Source\] = field(default\_factory=list) @dataclass class ResolutionReport: resolutions: list\[Resolution\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" async def \_you\_post(url: str, body: dict, timeout: float = 120.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str = "lite") -> dict: """Fast, citation-backed grounding for a support question.""" body = {"input": question, "research\_effort": research\_effort} return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment ground\_answer}} @env.task(retries=3) async def ground\_answer(ticket: str, context: str, research\_effort: str) -> Grounding: """Ground the ticket in fresh public sources via the Research API.""" question = ticket if not context else f"{ticket}\\n\\nContext: {context}" result = await you\_research(question, research\_effort) output = result.get("output", {}) answer = output.get("content", "") if not isinstance(answer, str): answer = json.dumps(answer) sources = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, snippet=str((s.get("snippets") or \[""\])\[0\]), domain=\_domain(url), favicon=\_favicon\_for(url), ) ) return Grounding(answer=answer, sources=sources) # {{/docs-fragment ground\_answer}} # {{docs-fragment draft\_reply}} @flyte.trace async def \_draft(ticket: str, answer: str, sources\_text: str) -> str: from litellm import acompletion system = ( "You are a senior customer-support agent. Using ONLY the grounded " "answer and sources provided, draft a concise, friendly, customer-ready " "reply. Cite the relevant source URL inline in parentheses after any " "factual claim so a human agent can verify before sending. If the " "sources do not answer the question, say so plainly." ) user = ( f"Customer ticket: {ticket}\\n\\n" f"Grounded answer:\\n{answer}\\n\\nSources:\\n{sources\_text}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.2, max\_tokens=1024, ) return resp.choices\[0\].message.content @env.task async def draft\_reply(ticket: Ticket, grounding: Grounding) -> Resolution: """Turn the grounded answer into a cited, customer-ready reply.""" sources\_text = "\\n".join( f"- {s.title} ({s.domain}): {s.url}\\n \\"{s.snippet}\\"" for s in grounding.sources ) reply = await \_draft(ticket.question, grounding.answer, sources\_text) return Resolution( ticket\_id=ticket.ticket\_id, ticket=ticket.question, grounded\_answer=grounding.answer, draft\_reply=reply, sources=grounding.sources, ) # {{/docs-fragment draft\_reply}} # {{docs-fragment resolve\_ticket}} async def resolve\_ticket(ticket: Ticket, research\_effort: str) -> Resolution: """Ground one ticket then draft its reply.""" grounding = await ground\_answer(ticket.question, ticket.context, research\_effort) return await draft\_reply(ticket, grounding) # {{/docs-fragment resolve\_ticket}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com Research citation for a support source.""" if not s.url: return "" snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"research" f"
{s.title}
{snip}
" ) def \_render\_report(report: ResolutionReport) -> str: cards = \[\] for res in report.resolutions: src = "".join(\_cite(s) for s in res.sources\[:8\]) reply\_html = res.draft\_reply.replace("\\n", "
") cards.append( f"
" f"
{res.ticket\_id}
" f"
{res.ticket}
" f"

Draft reply (for human review)

{reply\_html}
" + (f"

You.com sources ({len(res.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(r.sources) for r in report.resolutions) return f""" {REPORT\_CSS}

Support Resolutions

Tickets grounded in fresh public sources via the You.com Research API — draft replies cite sources a human agent can verify.

{len(report.resolutions)} tickets {total\_sources} You.com sources cited
{''.join(cards) or "

No tickets processed.

"}

Each ticket grounded by the You.com Research API (lite effort for low-latency, human-in-the-loop use). Sources include domain, title, and snippet provenance — ready to paste into a customer reply with verification links.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_tickets() -> list\[Ticket\]: return \[\ Ticket(\ "tkt-1",\ "Is there a recall on the DeWalt DCD777 cordless drill, and what should "\ "the customer do if there is?",\ "Customer purchased the drill recently and is asking about safety recalls.",\ ),\ Ticket(\ "tkt-2",\ "What is Sony's current return policy for the WH-1000XM5 headphones?",\ "Customer wants to return an opened pair bought 20 days ago.",\ ),\ Ticket(\ "tkt-3",\ "Are there any current weather advisories that could delay flights out of "\ "Denver International Airport today?",\ "Customer is worried about a connecting flight.",\ ),\ Ticket(\ "tkt-4",\ "What are the dimensions and weight capacity of the IKEA BEKANT desk?",\ "Customer is checking if it fits their space before resolving a complaint.",\ ),\ Ticket(\ "tkt-5",\ "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",\ "Customer reports overheating and wants to know about known issues.",\ ),\ Ticket(\ "tkt-6",\ "What is the warranty period for a Dyson V15 Detect vacuum in the US?",\ "Customer's vacuum stopped working and asks about coverage.",\ ),\ \] @env.task(report=True) async def support\_resolution( tickets: list\[Ticket\] | None = None, research\_effort: str = "lite", ) -> ResolutionReport: """Fan out across support tickets, grounding and drafting cited replies.""" if tickets is None: tickets = \_default\_tickets() with flyte.group("resolve-tickets"): resolutions = await asyncio.gather( \*\[resolve\_ticket(t, research\_effort) for t in tickets\] ) report = ResolutionReport(resolutions=list(resolutions)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(support\_resolution) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support\_resolution\_agent/main.py\* The Python packages are declared at the top of the file using the \`uv\` script style: \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # /// \`\`\` ## Data types Each \`Ticket\` carries a ticket ID, a customer question, and optional product or vendor context. The final \`Resolution\` includes the grounded answer, a draft reply, and the list of You.com sources. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "support\_resolution" # params = "" # /// """Customer-support & field-service resolution agent. Grounds a support ticket in fresh, public, citable sources via the You.com Research API (low effort for low latency, human-in-the-loop use), then uses Claude to draft a customer-ready reply that cites its sources inline so a human agent can verify before sending. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="support-resolution", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="support-resolution", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str snippet: str domain: str = "" favicon: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Ticket: ticket\_id: str question: str context: str = "" @dataclass class Grounding: answer: str sources: list\[Source\] = field(default\_factory=list) @dataclass class Resolution: ticket\_id: str ticket: str grounded\_answer: str draft\_reply: str sources: list\[Source\] = field(default\_factory=list) @dataclass class ResolutionReport: resolutions: list\[Resolution\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" async def \_you\_post(url: str, body: dict, timeout: float = 120.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str = "lite") -> dict: """Fast, citation-backed grounding for a support question.""" body = {"input": question, "research\_effort": research\_effort} return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment ground\_answer}} @env.task(retries=3) async def ground\_answer(ticket: str, context: str, research\_effort: str) -> Grounding: """Ground the ticket in fresh public sources via the Research API.""" question = ticket if not context else f"{ticket}\\n\\nContext: {context}" result = await you\_research(question, research\_effort) output = result.get("output", {}) answer = output.get("content", "") if not isinstance(answer, str): answer = json.dumps(answer) sources = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, snippet=str((s.get("snippets") or \[""\])\[0\]), domain=\_domain(url), favicon=\_favicon\_for(url), ) ) return Grounding(answer=answer, sources=sources) # {{/docs-fragment ground\_answer}} # {{docs-fragment draft\_reply}} @flyte.trace async def \_draft(ticket: str, answer: str, sources\_text: str) -> str: from litellm import acompletion system = ( "You are a senior customer-support agent. Using ONLY the grounded " "answer and sources provided, draft a concise, friendly, customer-ready " "reply. Cite the relevant source URL inline in parentheses after any " "factual claim so a human agent can verify before sending. If the " "sources do not answer the question, say so plainly." ) user = ( f"Customer ticket: {ticket}\\n\\n" f"Grounded answer:\\n{answer}\\n\\nSources:\\n{sources\_text}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.2, max\_tokens=1024, ) return resp.choices\[0\].message.content @env.task async def draft\_reply(ticket: Ticket, grounding: Grounding) -> Resolution: """Turn the grounded answer into a cited, customer-ready reply.""" sources\_text = "\\n".join( f"- {s.title} ({s.domain}): {s.url}\\n \\"{s.snippet}\\"" for s in grounding.sources ) reply = await \_draft(ticket.question, grounding.answer, sources\_text) return Resolution( ticket\_id=ticket.ticket\_id, ticket=ticket.question, grounded\_answer=grounding.answer, draft\_reply=reply, sources=grounding.sources, ) # {{/docs-fragment draft\_reply}} # {{docs-fragment resolve\_ticket}} async def resolve\_ticket(ticket: Ticket, research\_effort: str) -> Resolution: """Ground one ticket then draft its reply.""" grounding = await ground\_answer(ticket.question, ticket.context, research\_effort) return await draft\_reply(ticket, grounding) # {{/docs-fragment resolve\_ticket}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com Research citation for a support source.""" if not s.url: return "" snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"research" f"
{s.title}
{snip}
" ) def \_render\_report(report: ResolutionReport) -> str: cards = \[\] for res in report.resolutions: src = "".join(\_cite(s) for s in res.sources\[:8\]) reply\_html = res.draft\_reply.replace("\\n", "
") cards.append( f"
" f"
{res.ticket\_id}
" f"
{res.ticket}
" f"

Draft reply (for human review)

{reply\_html}
" + (f"

You.com sources ({len(res.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(r.sources) for r in report.resolutions) return f""" {REPORT\_CSS}

Support Resolutions

Tickets grounded in fresh public sources via the You.com Research API — draft replies cite sources a human agent can verify.

{len(report.resolutions)} tickets {total\_sources} You.com sources cited
{''.join(cards) or "

No tickets processed.

"}

Each ticket grounded by the You.com Research API (lite effort for low-latency, human-in-the-loop use). Sources include domain, title, and snippet provenance — ready to paste into a customer reply with verification links.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_tickets() -> list\[Ticket\]: return \[\ Ticket(\ "tkt-1",\ "Is there a recall on the DeWalt DCD777 cordless drill, and what should "\ "the customer do if there is?",\ "Customer purchased the drill recently and is asking about safety recalls.",\ ),\ Ticket(\ "tkt-2",\ "What is Sony's current return policy for the WH-1000XM5 headphones?",\ "Customer wants to return an opened pair bought 20 days ago.",\ ),\ Ticket(\ "tkt-3",\ "Are there any current weather advisories that could delay flights out of "\ "Denver International Airport today?",\ "Customer is worried about a connecting flight.",\ ),\ Ticket(\ "tkt-4",\ "What are the dimensions and weight capacity of the IKEA BEKANT desk?",\ "Customer is checking if it fits their space before resolving a complaint.",\ ),\ Ticket(\ "tkt-5",\ "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",\ "Customer reports overheating and wants to know about known issues.",\ ),\ Ticket(\ "tkt-6",\ "What is the warranty period for a Dyson V15 Detect vacuum in the US?",\ "Customer's vacuum stopped working and asks about coverage.",\ ),\ \] @env.task(report=True) async def support\_resolution( tickets: list\[Ticket\] | None = None, research\_effort: str = "lite", ) -> ResolutionReport: """Fan out across support tickets, grounding and drafting cited replies.""" if tickets is None: tickets = \_default\_tickets() with flyte.group("resolve-tickets"): resolutions = await asyncio.gather( \*\[resolve\_ticket(t, research\_effort) for t in tickets\] ) report = ResolutionReport(resolutions=list(resolutions)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(support\_resolution) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support\_resolution\_agent/main.py\* ## Ground answers with the You.com Research API The \`you\_research\` helper calls the \[You.com Research API\](https://you.com/docs/research/overview) with a configurable \`research\_effort\`. For support use cases, \`lite\` provides a fast, citation-backed answer suitable for real-time, human-in-the-loop flows. See the \[Research API reference\](https://you.com/docs/api-reference/research/v1-research) for effort levels and parameters. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "support\_resolution" # params = "" # /// """Customer-support & field-service resolution agent. Grounds a support ticket in fresh, public, citable sources via the You.com Research API (low effort for low latency, human-in-the-loop use), then uses Claude to draft a customer-ready reply that cites its sources inline so a human agent can verify before sending. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="support-resolution", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="support-resolution", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str snippet: str domain: str = "" favicon: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Ticket: ticket\_id: str question: str context: str = "" @dataclass class Grounding: answer: str sources: list\[Source\] = field(default\_factory=list) @dataclass class Resolution: ticket\_id: str ticket: str grounded\_answer: str draft\_reply: str sources: list\[Source\] = field(default\_factory=list) @dataclass class ResolutionReport: resolutions: list\[Resolution\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" async def \_you\_post(url: str, body: dict, timeout: float = 120.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str = "lite") -> dict: """Fast, citation-backed grounding for a support question.""" body = {"input": question, "research\_effort": research\_effort} return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment ground\_answer}} @env.task(retries=3) async def ground\_answer(ticket: str, context: str, research\_effort: str) -> Grounding: """Ground the ticket in fresh public sources via the Research API.""" question = ticket if not context else f"{ticket}\\n\\nContext: {context}" result = await you\_research(question, research\_effort) output = result.get("output", {}) answer = output.get("content", "") if not isinstance(answer, str): answer = json.dumps(answer) sources = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, snippet=str((s.get("snippets") or \[""\])\[0\]), domain=\_domain(url), favicon=\_favicon\_for(url), ) ) return Grounding(answer=answer, sources=sources) # {{/docs-fragment ground\_answer}} # {{docs-fragment draft\_reply}} @flyte.trace async def \_draft(ticket: str, answer: str, sources\_text: str) -> str: from litellm import acompletion system = ( "You are a senior customer-support agent. Using ONLY the grounded " "answer and sources provided, draft a concise, friendly, customer-ready " "reply. Cite the relevant source URL inline in parentheses after any " "factual claim so a human agent can verify before sending. If the " "sources do not answer the question, say so plainly." ) user = ( f"Customer ticket: {ticket}\\n\\n" f"Grounded answer:\\n{answer}\\n\\nSources:\\n{sources\_text}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.2, max\_tokens=1024, ) return resp.choices\[0\].message.content @env.task async def draft\_reply(ticket: Ticket, grounding: Grounding) -> Resolution: """Turn the grounded answer into a cited, customer-ready reply.""" sources\_text = "\\n".join( f"- {s.title} ({s.domain}): {s.url}\\n \\"{s.snippet}\\"" for s in grounding.sources ) reply = await \_draft(ticket.question, grounding.answer, sources\_text) return Resolution( ticket\_id=ticket.ticket\_id, ticket=ticket.question, grounded\_answer=grounding.answer, draft\_reply=reply, sources=grounding.sources, ) # {{/docs-fragment draft\_reply}} # {{docs-fragment resolve\_ticket}} async def resolve\_ticket(ticket: Ticket, research\_effort: str) -> Resolution: """Ground one ticket then draft its reply.""" grounding = await ground\_answer(ticket.question, ticket.context, research\_effort) return await draft\_reply(ticket, grounding) # {{/docs-fragment resolve\_ticket}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com Research citation for a support source.""" if not s.url: return "" snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"research" f"
{s.title}
{snip}
" ) def \_render\_report(report: ResolutionReport) -> str: cards = \[\] for res in report.resolutions: src = "".join(\_cite(s) for s in res.sources\[:8\]) reply\_html = res.draft\_reply.replace("\\n", "
") cards.append( f"
" f"
{res.ticket\_id}
" f"
{res.ticket}
" f"

Draft reply (for human review)

{reply\_html}
" + (f"

You.com sources ({len(res.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(r.sources) for r in report.resolutions) return f""" {REPORT\_CSS}

Support Resolutions

Tickets grounded in fresh public sources via the You.com Research API — draft replies cite sources a human agent can verify.

{len(report.resolutions)} tickets {total\_sources} You.com sources cited
{''.join(cards) or "

No tickets processed.

"}

Each ticket grounded by the You.com Research API (lite effort for low-latency, human-in-the-loop use). Sources include domain, title, and snippet provenance — ready to paste into a customer reply with verification links.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_tickets() -> list\[Ticket\]: return \[\ Ticket(\ "tkt-1",\ "Is there a recall on the DeWalt DCD777 cordless drill, and what should "\ "the customer do if there is?",\ "Customer purchased the drill recently and is asking about safety recalls.",\ ),\ Ticket(\ "tkt-2",\ "What is Sony's current return policy for the WH-1000XM5 headphones?",\ "Customer wants to return an opened pair bought 20 days ago.",\ ),\ Ticket(\ "tkt-3",\ "Are there any current weather advisories that could delay flights out of "\ "Denver International Airport today?",\ "Customer is worried about a connecting flight.",\ ),\ Ticket(\ "tkt-4",\ "What are the dimensions and weight capacity of the IKEA BEKANT desk?",\ "Customer is checking if it fits their space before resolving a complaint.",\ ),\ Ticket(\ "tkt-5",\ "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",\ "Customer reports overheating and wants to know about known issues.",\ ),\ Ticket(\ "tkt-6",\ "What is the warranty period for a Dyson V15 Detect vacuum in the US?",\ "Customer's vacuum stopped working and asks about coverage.",\ ),\ \] @env.task(report=True) async def support\_resolution( tickets: list\[Ticket\] | None = None, research\_effort: str = "lite", ) -> ResolutionReport: """Fan out across support tickets, grounding and drafting cited replies.""" if tickets is None: tickets = \_default\_tickets() with flyte.group("resolve-tickets"): resolutions = await asyncio.gather( \*\[resolve\_ticket(t, research\_effort) for t in tickets\] ) report = ResolutionReport(resolutions=list(resolutions)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(support\_resolution) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support\_resolution\_agent/main.py\* ## Ground one ticket The \`ground\_answer\` task combines the ticket question and context into a research query and collects the grounded answer plus structured sources from the Research API response. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "support\_resolution" # params = "" # /// """Customer-support & field-service resolution agent. Grounds a support ticket in fresh, public, citable sources via the You.com Research API (low effort for low latency, human-in-the-loop use), then uses Claude to draft a customer-ready reply that cites its sources inline so a human agent can verify before sending. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="support-resolution", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="support-resolution", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str snippet: str domain: str = "" favicon: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Ticket: ticket\_id: str question: str context: str = "" @dataclass class Grounding: answer: str sources: list\[Source\] = field(default\_factory=list) @dataclass class Resolution: ticket\_id: str ticket: str grounded\_answer: str draft\_reply: str sources: list\[Source\] = field(default\_factory=list) @dataclass class ResolutionReport: resolutions: list\[Resolution\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" async def \_you\_post(url: str, body: dict, timeout: float = 120.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str = "lite") -> dict: """Fast, citation-backed grounding for a support question.""" body = {"input": question, "research\_effort": research\_effort} return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment ground\_answer}} @env.task(retries=3) async def ground\_answer(ticket: str, context: str, research\_effort: str) -> Grounding: """Ground the ticket in fresh public sources via the Research API.""" question = ticket if not context else f"{ticket}\\n\\nContext: {context}" result = await you\_research(question, research\_effort) output = result.get("output", {}) answer = output.get("content", "") if not isinstance(answer, str): answer = json.dumps(answer) sources = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, snippet=str((s.get("snippets") or \[""\])\[0\]), domain=\_domain(url), favicon=\_favicon\_for(url), ) ) return Grounding(answer=answer, sources=sources) # {{/docs-fragment ground\_answer}} # {{docs-fragment draft\_reply}} @flyte.trace async def \_draft(ticket: str, answer: str, sources\_text: str) -> str: from litellm import acompletion system = ( "You are a senior customer-support agent. Using ONLY the grounded " "answer and sources provided, draft a concise, friendly, customer-ready " "reply. Cite the relevant source URL inline in parentheses after any " "factual claim so a human agent can verify before sending. If the " "sources do not answer the question, say so plainly." ) user = ( f"Customer ticket: {ticket}\\n\\n" f"Grounded answer:\\n{answer}\\n\\nSources:\\n{sources\_text}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.2, max\_tokens=1024, ) return resp.choices\[0\].message.content @env.task async def draft\_reply(ticket: Ticket, grounding: Grounding) -> Resolution: """Turn the grounded answer into a cited, customer-ready reply.""" sources\_text = "\\n".join( f"- {s.title} ({s.domain}): {s.url}\\n \\"{s.snippet}\\"" for s in grounding.sources ) reply = await \_draft(ticket.question, grounding.answer, sources\_text) return Resolution( ticket\_id=ticket.ticket\_id, ticket=ticket.question, grounded\_answer=grounding.answer, draft\_reply=reply, sources=grounding.sources, ) # {{/docs-fragment draft\_reply}} # {{docs-fragment resolve\_ticket}} async def resolve\_ticket(ticket: Ticket, research\_effort: str) -> Resolution: """Ground one ticket then draft its reply.""" grounding = await ground\_answer(ticket.question, ticket.context, research\_effort) return await draft\_reply(ticket, grounding) # {{/docs-fragment resolve\_ticket}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com Research citation for a support source.""" if not s.url: return "" snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"research" f"
{s.title}
{snip}
" ) def \_render\_report(report: ResolutionReport) -> str: cards = \[\] for res in report.resolutions: src = "".join(\_cite(s) for s in res.sources\[:8\]) reply\_html = res.draft\_reply.replace("\\n", "
") cards.append( f"
" f"
{res.ticket\_id}
" f"
{res.ticket}
" f"

Draft reply (for human review)

{reply\_html}
" + (f"

You.com sources ({len(res.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(r.sources) for r in report.resolutions) return f""" {REPORT\_CSS}

Support Resolutions

Tickets grounded in fresh public sources via the You.com Research API — draft replies cite sources a human agent can verify.

{len(report.resolutions)} tickets {total\_sources} You.com sources cited
{''.join(cards) or "

No tickets processed.

"}

Each ticket grounded by the You.com Research API (lite effort for low-latency, human-in-the-loop use). Sources include domain, title, and snippet provenance — ready to paste into a customer reply with verification links.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_tickets() -> list\[Ticket\]: return \[\ Ticket(\ "tkt-1",\ "Is there a recall on the DeWalt DCD777 cordless drill, and what should "\ "the customer do if there is?",\ "Customer purchased the drill recently and is asking about safety recalls.",\ ),\ Ticket(\ "tkt-2",\ "What is Sony's current return policy for the WH-1000XM5 headphones?",\ "Customer wants to return an opened pair bought 20 days ago.",\ ),\ Ticket(\ "tkt-3",\ "Are there any current weather advisories that could delay flights out of "\ "Denver International Airport today?",\ "Customer is worried about a connecting flight.",\ ),\ Ticket(\ "tkt-4",\ "What are the dimensions and weight capacity of the IKEA BEKANT desk?",\ "Customer is checking if it fits their space before resolving a complaint.",\ ),\ Ticket(\ "tkt-5",\ "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",\ "Customer reports overheating and wants to know about known issues.",\ ),\ Ticket(\ "tkt-6",\ "What is the warranty period for a Dyson V15 Detect vacuum in the US?",\ "Customer's vacuum stopped working and asks about coverage.",\ ),\ \] @env.task(report=True) async def support\_resolution( tickets: list\[Ticket\] | None = None, research\_effort: str = "lite", ) -> ResolutionReport: """Fan out across support tickets, grounding and drafting cited replies.""" if tickets is None: tickets = \_default\_tickets() with flyte.group("resolve-tickets"): resolutions = await asyncio.gather( \*\[resolve\_ticket(t, research\_effort) for t in tickets\] ) report = ResolutionReport(resolutions=list(resolutions)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(support\_resolution) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support\_resolution\_agent/main.py\* ## Draft a customer-ready reply The \`draft\_reply\` task turns the grounded answer into a concise, friendly reply that cites source URLs inline so a human agent can verify before sending. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "support\_resolution" # params = "" # /// """Customer-support & field-service resolution agent. Grounds a support ticket in fresh, public, citable sources via the You.com Research API (low effort for low latency, human-in-the-loop use), then uses Claude to draft a customer-ready reply that cites its sources inline so a human agent can verify before sending. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="support-resolution", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="support-resolution", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str snippet: str domain: str = "" favicon: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Ticket: ticket\_id: str question: str context: str = "" @dataclass class Grounding: answer: str sources: list\[Source\] = field(default\_factory=list) @dataclass class Resolution: ticket\_id: str ticket: str grounded\_answer: str draft\_reply: str sources: list\[Source\] = field(default\_factory=list) @dataclass class ResolutionReport: resolutions: list\[Resolution\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" async def \_you\_post(url: str, body: dict, timeout: float = 120.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str = "lite") -> dict: """Fast, citation-backed grounding for a support question.""" body = {"input": question, "research\_effort": research\_effort} return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment ground\_answer}} @env.task(retries=3) async def ground\_answer(ticket: str, context: str, research\_effort: str) -> Grounding: """Ground the ticket in fresh public sources via the Research API.""" question = ticket if not context else f"{ticket}\\n\\nContext: {context}" result = await you\_research(question, research\_effort) output = result.get("output", {}) answer = output.get("content", "") if not isinstance(answer, str): answer = json.dumps(answer) sources = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, snippet=str((s.get("snippets") or \[""\])\[0\]), domain=\_domain(url), favicon=\_favicon\_for(url), ) ) return Grounding(answer=answer, sources=sources) # {{/docs-fragment ground\_answer}} # {{docs-fragment draft\_reply}} @flyte.trace async def \_draft(ticket: str, answer: str, sources\_text: str) -> str: from litellm import acompletion system = ( "You are a senior customer-support agent. Using ONLY the grounded " "answer and sources provided, draft a concise, friendly, customer-ready " "reply. Cite the relevant source URL inline in parentheses after any " "factual claim so a human agent can verify before sending. If the " "sources do not answer the question, say so plainly." ) user = ( f"Customer ticket: {ticket}\\n\\n" f"Grounded answer:\\n{answer}\\n\\nSources:\\n{sources\_text}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.2, max\_tokens=1024, ) return resp.choices\[0\].message.content @env.task async def draft\_reply(ticket: Ticket, grounding: Grounding) -> Resolution: """Turn the grounded answer into a cited, customer-ready reply.""" sources\_text = "\\n".join( f"- {s.title} ({s.domain}): {s.url}\\n \\"{s.snippet}\\"" for s in grounding.sources ) reply = await \_draft(ticket.question, grounding.answer, sources\_text) return Resolution( ticket\_id=ticket.ticket\_id, ticket=ticket.question, grounded\_answer=grounding.answer, draft\_reply=reply, sources=grounding.sources, ) # {{/docs-fragment draft\_reply}} # {{docs-fragment resolve\_ticket}} async def resolve\_ticket(ticket: Ticket, research\_effort: str) -> Resolution: """Ground one ticket then draft its reply.""" grounding = await ground\_answer(ticket.question, ticket.context, research\_effort) return await draft\_reply(ticket, grounding) # {{/docs-fragment resolve\_ticket}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com Research citation for a support source.""" if not s.url: return "" snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"research" f"
{s.title}
{snip}
" ) def \_render\_report(report: ResolutionReport) -> str: cards = \[\] for res in report.resolutions: src = "".join(\_cite(s) for s in res.sources\[:8\]) reply\_html = res.draft\_reply.replace("\\n", "
") cards.append( f"
" f"
{res.ticket\_id}
" f"
{res.ticket}
" f"

Draft reply (for human review)

{reply\_html}
" + (f"

You.com sources ({len(res.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(r.sources) for r in report.resolutions) return f""" {REPORT\_CSS}

Support Resolutions

Tickets grounded in fresh public sources via the You.com Research API — draft replies cite sources a human agent can verify.

{len(report.resolutions)} tickets {total\_sources} You.com sources cited
{''.join(cards) or "

No tickets processed.

"}

Each ticket grounded by the You.com Research API (lite effort for low-latency, human-in-the-loop use). Sources include domain, title, and snippet provenance — ready to paste into a customer reply with verification links.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_tickets() -> list\[Ticket\]: return \[\ Ticket(\ "tkt-1",\ "Is there a recall on the DeWalt DCD777 cordless drill, and what should "\ "the customer do if there is?",\ "Customer purchased the drill recently and is asking about safety recalls.",\ ),\ Ticket(\ "tkt-2",\ "What is Sony's current return policy for the WH-1000XM5 headphones?",\ "Customer wants to return an opened pair bought 20 days ago.",\ ),\ Ticket(\ "tkt-3",\ "Are there any current weather advisories that could delay flights out of "\ "Denver International Airport today?",\ "Customer is worried about a connecting flight.",\ ),\ Ticket(\ "tkt-4",\ "What are the dimensions and weight capacity of the IKEA BEKANT desk?",\ "Customer is checking if it fits their space before resolving a complaint.",\ ),\ Ticket(\ "tkt-5",\ "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",\ "Customer reports overheating and wants to know about known issues.",\ ),\ Ticket(\ "tkt-6",\ "What is the warranty period for a Dyson V15 Detect vacuum in the US?",\ "Customer's vacuum stopped working and asks about coverage.",\ ),\ \] @env.task(report=True) async def support\_resolution( tickets: list\[Ticket\] | None = None, research\_effort: str = "lite", ) -> ResolutionReport: """Fan out across support tickets, grounding and drafting cited replies.""" if tickets is None: tickets = \_default\_tickets() with flyte.group("resolve-tickets"): resolutions = await asyncio.gather( \*\[resolve\_ticket(t, research\_effort) for t in tickets\] ) report = ResolutionReport(resolutions=list(resolutions)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(support\_resolution) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support\_resolution\_agent/main.py\* ## Resolve one ticket Each ticket runs \`ground\_answer\` followed by \`draft\_reply\` in sequence. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "support\_resolution" # params = "" # /// """Customer-support & field-service resolution agent. Grounds a support ticket in fresh, public, citable sources via the You.com Research API (low effort for low latency, human-in-the-loop use), then uses Claude to draft a customer-ready reply that cites its sources inline so a human agent can verify before sending. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="support-resolution", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="support-resolution", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str snippet: str domain: str = "" favicon: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Ticket: ticket\_id: str question: str context: str = "" @dataclass class Grounding: answer: str sources: list\[Source\] = field(default\_factory=list) @dataclass class Resolution: ticket\_id: str ticket: str grounded\_answer: str draft\_reply: str sources: list\[Source\] = field(default\_factory=list) @dataclass class ResolutionReport: resolutions: list\[Resolution\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" async def \_you\_post(url: str, body: dict, timeout: float = 120.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str = "lite") -> dict: """Fast, citation-backed grounding for a support question.""" body = {"input": question, "research\_effort": research\_effort} return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment ground\_answer}} @env.task(retries=3) async def ground\_answer(ticket: str, context: str, research\_effort: str) -> Grounding: """Ground the ticket in fresh public sources via the Research API.""" question = ticket if not context else f"{ticket}\\n\\nContext: {context}" result = await you\_research(question, research\_effort) output = result.get("output", {}) answer = output.get("content", "") if not isinstance(answer, str): answer = json.dumps(answer) sources = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, snippet=str((s.get("snippets") or \[""\])\[0\]), domain=\_domain(url), favicon=\_favicon\_for(url), ) ) return Grounding(answer=answer, sources=sources) # {{/docs-fragment ground\_answer}} # {{docs-fragment draft\_reply}} @flyte.trace async def \_draft(ticket: str, answer: str, sources\_text: str) -> str: from litellm import acompletion system = ( "You are a senior customer-support agent. Using ONLY the grounded " "answer and sources provided, draft a concise, friendly, customer-ready " "reply. Cite the relevant source URL inline in parentheses after any " "factual claim so a human agent can verify before sending. If the " "sources do not answer the question, say so plainly." ) user = ( f"Customer ticket: {ticket}\\n\\n" f"Grounded answer:\\n{answer}\\n\\nSources:\\n{sources\_text}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.2, max\_tokens=1024, ) return resp.choices\[0\].message.content @env.task async def draft\_reply(ticket: Ticket, grounding: Grounding) -> Resolution: """Turn the grounded answer into a cited, customer-ready reply.""" sources\_text = "\\n".join( f"- {s.title} ({s.domain}): {s.url}\\n \\"{s.snippet}\\"" for s in grounding.sources ) reply = await \_draft(ticket.question, grounding.answer, sources\_text) return Resolution( ticket\_id=ticket.ticket\_id, ticket=ticket.question, grounded\_answer=grounding.answer, draft\_reply=reply, sources=grounding.sources, ) # {{/docs-fragment draft\_reply}} # {{docs-fragment resolve\_ticket}} async def resolve\_ticket(ticket: Ticket, research\_effort: str) -> Resolution: """Ground one ticket then draft its reply.""" grounding = await ground\_answer(ticket.question, ticket.context, research\_effort) return await draft\_reply(ticket, grounding) # {{/docs-fragment resolve\_ticket}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com Research citation for a support source.""" if not s.url: return "" snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"research" f"
{s.title}
{snip}
" ) def \_render\_report(report: ResolutionReport) -> str: cards = \[\] for res in report.resolutions: src = "".join(\_cite(s) for s in res.sources\[:8\]) reply\_html = res.draft\_reply.replace("\\n", "
") cards.append( f"
" f"
{res.ticket\_id}
" f"
{res.ticket}
" f"

Draft reply (for human review)

{reply\_html}
" + (f"

You.com sources ({len(res.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(r.sources) for r in report.resolutions) return f""" {REPORT\_CSS}

Support Resolutions

Tickets grounded in fresh public sources via the You.com Research API — draft replies cite sources a human agent can verify.

{len(report.resolutions)} tickets {total\_sources} You.com sources cited
{''.join(cards) or "

No tickets processed.

"}

Each ticket grounded by the You.com Research API (lite effort for low-latency, human-in-the-loop use). Sources include domain, title, and snippet provenance — ready to paste into a customer reply with verification links.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_tickets() -> list\[Ticket\]: return \[\ Ticket(\ "tkt-1",\ "Is there a recall on the DeWalt DCD777 cordless drill, and what should "\ "the customer do if there is?",\ "Customer purchased the drill recently and is asking about safety recalls.",\ ),\ Ticket(\ "tkt-2",\ "What is Sony's current return policy for the WH-1000XM5 headphones?",\ "Customer wants to return an opened pair bought 20 days ago.",\ ),\ Ticket(\ "tkt-3",\ "Are there any current weather advisories that could delay flights out of "\ "Denver International Airport today?",\ "Customer is worried about a connecting flight.",\ ),\ Ticket(\ "tkt-4",\ "What are the dimensions and weight capacity of the IKEA BEKANT desk?",\ "Customer is checking if it fits their space before resolving a complaint.",\ ),\ Ticket(\ "tkt-5",\ "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",\ "Customer reports overheating and wants to know about known issues.",\ ),\ Ticket(\ "tkt-6",\ "What is the warranty period for a Dyson V15 Detect vacuum in the US?",\ "Customer's vacuum stopped working and asks about coverage.",\ ),\ \] @env.task(report=True) async def support\_resolution( tickets: list\[Ticket\] | None = None, research\_effort: str = "lite", ) -> ResolutionReport: """Fan out across support tickets, grounding and drafting cited replies.""" if tickets is None: tickets = \_default\_tickets() with flyte.group("resolve-tickets"): resolutions = await asyncio.gather( \*\[resolve\_ticket(t, research\_effort) for t in tickets\] ) report = ResolutionReport(resolutions=list(resolutions)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(support\_resolution) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support\_resolution\_agent/main.py\* ## Orchestration The \`support\_resolution\` driver task fans out across all tickets and renders a Flyte report with every draft reply and its sources. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.4.0",\ # "httpx>=0.27.0",\ # "litellm>=1.72.0",\ # \] # main = "support\_resolution" # params = "" # /// """Customer-support & field-service resolution agent. Grounds a support ticket in fresh, public, citable sources via the You.com Research API (low effort for low latency, human-in-the-loop use), then uses Claude to draft a customer-ready reply that cites its sources inline so a human agent can verify before sending. """ # {{docs-fragment env}} import asyncio import json import os from dataclasses import dataclass, field import flyte MODEL = "anthropic/claude-haiku-4-5" env = flyte.TaskEnvironment( name="support-resolution", secrets=\[\ flyte.Secret(key="youdotcom-api-key", as\_env\_var="YOU\_API\_KEY"),\ flyte.Secret(key="internal-anthropic-api-key", as\_env\_var="ANTHROPIC\_API\_KEY"),\ \], image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="support-resolution", pre=True), resources=flyte.Resources(cpu="1", memory="1Gi"), ) # {{/docs-fragment env}} # {{docs-fragment data\_types}} @dataclass class Source: title: str url: str snippet: str domain: str = "" favicon: str = "" def \_domain(url: str) -> str: from urllib.parse import urlparse try: return urlparse(url).netloc.replace("www.", "") except Exception: return "" def \_favicon\_for(url: str) -> str: return f"https://ydc-index.io/favicon?domain={\_domain(url)}&size=128" @dataclass class Ticket: ticket\_id: str question: str context: str = "" @dataclass class Grounding: answer: str sources: list\[Source\] = field(default\_factory=list) @dataclass class Resolution: ticket\_id: str ticket: str grounded\_answer: str draft\_reply: str sources: list\[Source\] = field(default\_factory=list) @dataclass class ResolutionReport: resolutions: list\[Resolution\] = field(default\_factory=list) # {{/docs-fragment data\_types}} # {{docs-fragment you\_research}} YOU\_RESEARCH\_URL = "https://api.you.com/v1/research" async def \_you\_post(url: str, body: dict, timeout: float = 120.0) -> dict: """POST with exponential backoff + jitter on 429 rate limits.""" import random import httpx headers = { "X-API-Key": os.environ\["YOU\_API\_KEY"\], "Content-Type": "application/json", } async with httpx.AsyncClient(timeout=timeout) as client: for attempt in range(7): resp = await client.post(url, headers=headers, json=body) if resp.status\_code == 429 and attempt < 6: wait = float(resp.headers.get("retry-after") or 0) or min(2\*\*attempt, 30) await asyncio.sleep(wait + random.uniform(0, 2)) continue resp.raise\_for\_status() return resp.json() resp.raise\_for\_status() return resp.json() @flyte.trace async def you\_research(question: str, research\_effort: str = "lite") -> dict: """Fast, citation-backed grounding for a support question.""" body = {"input": question, "research\_effort": research\_effort} return await \_you\_post(YOU\_RESEARCH\_URL, body) # {{/docs-fragment you\_research}} # {{docs-fragment ground\_answer}} @env.task(retries=3) async def ground\_answer(ticket: str, context: str, research\_effort: str) -> Grounding: """Ground the ticket in fresh public sources via the Research API.""" question = ticket if not context else f"{ticket}\\n\\nContext: {context}" result = await you\_research(question, research\_effort) output = result.get("output", {}) answer = output.get("content", "") if not isinstance(answer, str): answer = json.dumps(answer) sources = \[\] for s in output.get("sources", \[\]) or \[\]: url = str(s.get("url", "")) sources.append( Source( title=str(s.get("title", "") or url), url=url, snippet=str((s.get("snippets") or \[""\])\[0\]), domain=\_domain(url), favicon=\_favicon\_for(url), ) ) return Grounding(answer=answer, sources=sources) # {{/docs-fragment ground\_answer}} # {{docs-fragment draft\_reply}} @flyte.trace async def \_draft(ticket: str, answer: str, sources\_text: str) -> str: from litellm import acompletion system = ( "You are a senior customer-support agent. Using ONLY the grounded " "answer and sources provided, draft a concise, friendly, customer-ready " "reply. Cite the relevant source URL inline in parentheses after any " "factual claim so a human agent can verify before sending. If the " "sources do not answer the question, say so plainly." ) user = ( f"Customer ticket: {ticket}\\n\\n" f"Grounded answer:\\n{answer}\\n\\nSources:\\n{sources\_text}" ) resp = await acompletion( model=MODEL, messages=\[\ {"role": "system", "content": system},\ {"role": "user", "content": user},\ \], temperature=0.2, max\_tokens=1024, ) return resp.choices\[0\].message.content @env.task async def draft\_reply(ticket: Ticket, grounding: Grounding) -> Resolution: """Turn the grounded answer into a cited, customer-ready reply.""" sources\_text = "\\n".join( f"- {s.title} ({s.domain}): {s.url}\\n \\"{s.snippet}\\"" for s in grounding.sources ) reply = await \_draft(ticket.question, grounding.answer, sources\_text) return Resolution( ticket\_id=ticket.ticket\_id, ticket=ticket.question, grounded\_answer=grounding.answer, draft\_reply=reply, sources=grounding.sources, ) # {{/docs-fragment draft\_reply}} # {{docs-fragment resolve\_ticket}} async def resolve\_ticket(ticket: Ticket, research\_effort: str) -> Resolution: """Ground one ticket then draft its reply.""" grounding = await ground\_answer(ticket.question, ticket.context, research\_effort) return await draft\_reply(ticket, grounding) # {{/docs-fragment resolve\_ticket}} # {{docs-fragment report}} REPORT\_CSS = """ """ def \_cite(s: Source) -> str: """Render a rich You.com Research citation for a support source.""" if not s.url: return "" snip = f"
“{s.snippet}”
" if s.snippet else "" return ( f"
" f"
" f"{s.domain or 'source'}" f"research" f"
{s.title}
{snip}
" ) def \_render\_report(report: ResolutionReport) -> str: cards = \[\] for res in report.resolutions: src = "".join(\_cite(s) for s in res.sources\[:8\]) reply\_html = res.draft\_reply.replace("\\n", "
") cards.append( f"
" f"
{res.ticket\_id}
" f"
{res.ticket}
" f"

Draft reply (for human review)

{reply\_html}
" + (f"

You.com sources ({len(res.sources)})

{src}
" if src else "") + "
" ) total\_sources = sum(len(r.sources) for r in report.resolutions) return f""" {REPORT\_CSS}

Support Resolutions

Tickets grounded in fresh public sources via the You.com Research API — draft replies cite sources a human agent can verify.

{len(report.resolutions)} tickets {total\_sources} You.com sources cited
{''.join(cards) or "

No tickets processed.

"}

Each ticket grounded by the You.com Research API (lite effort for low-latency, human-in-the-loop use). Sources include domain, title, and snippet provenance — ready to paste into a customer reply with verification links.

""" # {{/docs-fragment report}} # {{docs-fragment driver}} def \_default\_tickets() -> list\[Ticket\]: return \[\ Ticket(\ "tkt-1",\ "Is there a recall on the DeWalt DCD777 cordless drill, and what should "\ "the customer do if there is?",\ "Customer purchased the drill recently and is asking about safety recalls.",\ ),\ Ticket(\ "tkt-2",\ "What is Sony's current return policy for the WH-1000XM5 headphones?",\ "Customer wants to return an opened pair bought 20 days ago.",\ ),\ Ticket(\ "tkt-3",\ "Are there any current weather advisories that could delay flights out of "\ "Denver International Airport today?",\ "Customer is worried about a connecting flight.",\ ),\ Ticket(\ "tkt-4",\ "What are the dimensions and weight capacity of the IKEA BEKANT desk?",\ "Customer is checking if it fits their space before resolving a complaint.",\ ),\ Ticket(\ "tkt-5",\ "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",\ "Customer reports overheating and wants to know about known issues.",\ ),\ Ticket(\ "tkt-6",\ "What is the warranty period for a Dyson V15 Detect vacuum in the US?",\ "Customer's vacuum stopped working and asks about coverage.",\ ),\ \] @env.task(report=True) async def support\_resolution( tickets: list\[Ticket\] | None = None, research\_effort: str = "lite", ) -> ResolutionReport: """Fan out across support tickets, grounding and drafting cited replies.""" if tickets is None: tickets = \_default\_tickets() with flyte.group("resolve-tickets"): resolutions = await asyncio.gather( \*\[resolve\_ticket(t, research\_effort) for t in tickets\] ) report = ResolutionReport(resolutions=list(resolutions)) await flyte.report.replace.aio(\_render\_report(report), do\_flush=True) await flyte.report.flush.aio() return report # {{/docs-fragment driver}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(support\_resolution) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/support\_resolution\_agent/main.py\* ## Run the agent ### Create secrets Get a You.com API key from the \[You.com platform\](https://you.com/platform) (see the \[quickstart guide\](https://you.com/docs/quickstart)). Get an Anthropic API key from the \[Anthropic console\](https://console.anthropic.com/). Register both keys as Flyte secrets. The secret key names must match those declared in the \`TaskEnvironment\`: \`\`\` flyte create secret youdotcom-api-key flyte create secret internal-anthropic-api-key \`\`\` See \[Secrets\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets/page.md) for scoping and file-based secrets. ### Run locally or remotely From the \[example directory\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/support\_resolution\_agent): \`\`\` cd v2/tutorials/support\_resolution\_agent uv run --script main.py \`\`\` To test locally without Flyte secrets: \`\`\` export YOU\_API\_KEY= export ANTHROPIC\_API\_KEY= uv run --script main.py \`\`\` When the run completes, open the Flyte report to review draft replies for each ticket, with You.com source citations ready for a human agent to verify and paste into a customer response. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/context-engineering === # Context Engineering Tutorials for prompt engineering, prompt optimization, and context construction. ### \*\*Context Engineering > Automatic prompt engineering\*\* Easily run prompt optimization with real-time observability, traceability, and automatic recovery. ### \*\*Context Engineering > Text-to-SQL prompt optimization\*\* Learn how to turn natural language questions into SQL queries with Flyte and LlamaIndex, and explore prompt optimization in practice. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/context-engineering/text\_to\_sql === # Text-to-SQL prompt optimization > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/text\_to\_sql); based on work by \[LlamaIndex\](https://docs.llamaindex.ai/en/stable/examples/workflow/advanced\_text\_to\_sql/). Data analytics drives modern decision-making, but SQL often creates a bottleneck. Writing queries requires technical expertise, so non-technical stakeholders must rely on data teams. That translation layer slows everyone down. Text-to-SQL narrows this gap by turning natural language into executable SQL queries. It lowers the barrier to structured data and makes databases accessible to more people. In this tutorial, we build a Text-to-SQL workflow using LlamaIndex and evaluate it on the \[WikiTableQuestions dataset\](https://ppasupat.github.io/WikiTableQuestions/) (a benchmark of natural language questions over semi-structured tables). We then explore prompt optimization to see whether it improves accuracy and show how to track prompts and results over time. Along the way, we'll see what worked, what didn't, and what we learned about building durable evaluation pipelines. The pattern here can be adapted to your own datasets and workflows. !\[Evaluation\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/tutorials/text-to-sql/evaluation.png) ## Ingesting data We start by ingesting the WikiTableQuestions dataset, which comes as CSV files, into a SQLite database. This database serves as the source of truth for our Text-to-SQL pipeline. \`\`\` import asyncio import fnmatch import os import re import zipfile import flyte import pandas as pd import requests from flyte.io import Dir, File from llama\_index.core.llms import ChatMessage from llama\_index.core.prompts import ChatPromptTemplate from llama\_index.llms.openai import OpenAI from pydantic import BaseModel, Field from sqlalchemy import Column, Integer, MetaData, String, Table, create\_engine from utils import env # {{docs-fragment table\_info}} class TableInfo(BaseModel): """Information regarding a structured table.""" table\_name: str = Field(..., description="table name (underscores only, no spaces)") table\_summary: str = Field( ..., description="short, concise summary/caption of the table" ) # {{/docs-fragment table\_info}} @env.task async def download\_and\_extract(zip\_path: str, search\_glob: str) -> Dir: """Download and extract the dataset zip file if not already available.""" output\_zip = "data.zip" extract\_dir = "wiki\_table\_questions" if not os.path.exists(zip\_path): response = requests.get(zip\_path, stream=True) with open(output\_zip, "wb") as f: for chunk in response.iter\_content(chunk\_size=8192): f.write(chunk) else: output\_zip = zip\_path print(f"Using existing file {output\_zip}") os.makedirs(extract\_dir, exist\_ok=True) with zipfile.ZipFile(output\_zip, "r") as zip\_ref: for member in zip\_ref.namelist(): if fnmatch.fnmatch(member, search\_glob): zip\_ref.extract(member, extract\_dir) remote\_dir = await Dir.from\_local(extract\_dir) return remote\_dir async def read\_csv\_file( csv\_file: File, nrows: int | None = None ) -> pd.DataFrame | None: """Safely download and parse a CSV file into a DataFrame.""" try: local\_csv\_file = await csv\_file.download() return pd.read\_csv(local\_csv\_file, nrows=nrows) except Exception as e: print(f"Error parsing {csv\_file.path}: {e}") return None def sanitize\_column\_name(col\_name: str) -> str: """Sanitize column names by replacing spaces/special chars with underscores.""" return re.sub(r"\\W+", "\_", col\_name) async def create\_table\_from\_dataframe( df: pd.DataFrame, table\_name: str, engine, metadata\_obj ): """Create a SQL table from a Pandas DataFrame.""" # Sanitize column names sanitized\_columns = {col: sanitize\_column\_name(col) for col in df.columns} df = df.rename(columns=sanitized\_columns) # Define table columns based on DataFrame dtypes columns = \[\ Column(col, String if dtype == "object" else Integer)\ for col, dtype in zip(df.columns, df.dtypes)\ \] table = Table(table\_name, metadata\_obj, \*columns) # Create table in database metadata\_obj.create\_all(engine) # Insert data into table with engine.begin() as conn: for \_, row in df.iterrows(): conn.execute(table.insert().values(\*\*row.to\_dict())) @flyte.trace async def create\_table( csv\_file: File, table\_info: TableInfo, database\_path: str ) -> str: """Safely create a table from CSV if parsing succeeds.""" df = await read\_csv\_file(csv\_file) if df is None: return "false" print(f"Creating table: {table\_info.table\_name}") engine = create\_engine(f"sqlite:///{database\_path}") metadata\_obj = MetaData() await create\_table\_from\_dataframe(df, table\_info.table\_name, engine, metadata\_obj) return "true" @flyte.trace async def llm\_structured\_predict( df\_str: str, table\_names: list\[str\], prompt\_tmpl: ChatPromptTemplate, feedback: str, llm: OpenAI, ) -> TableInfo: return llm.structured\_predict( TableInfo, prompt\_tmpl, feedback=feedback, table\_str=df\_str, exclude\_table\_name\_list=str(list(table\_names)), ) async def generate\_unique\_table\_info( df\_str: str, table\_names: list\[str\], prompt\_tmpl: ChatPromptTemplate, llm: OpenAI, tablename\_lock: asyncio.Lock, retries: int = 3, ) -> TableInfo | None: """Process a single CSV file to generate a unique TableInfo.""" last\_table\_name = None for attempt in range(retries): feedback = "" if attempt > 0: feedback = f"Note: '{last\_table\_name}' already exists. Please pick a new name not in {table\_names}." table\_info = await llm\_structured\_predict( df\_str, table\_names, prompt\_tmpl, feedback, llm ) last\_table\_name = table\_info.table\_name async with tablename\_lock: if table\_info.table\_name not in table\_names: table\_names.append(table\_info.table\_name) return table\_info print(f"Table name {table\_info.table\_name} already exists, retrying...") return None async def process\_csv\_file( csv\_file: File, table\_names: list\[str\], semaphore: asyncio.Semaphore, tablename\_lock: asyncio.Lock, llm: OpenAI, prompt\_tmpl: ChatPromptTemplate, ) -> TableInfo | None: """Process a single CSV file to generate a unique TableInfo.""" async with semaphore: df = await read\_csv\_file(csv\_file, nrows=10) if df is None: return None return await generate\_unique\_table\_info( df.to\_csv(), table\_names, prompt\_tmpl, llm, tablename\_lock ) @env.task async def extract\_table\_info( data\_dir: Dir, model: str, concurrency: int ) -> list\[TableInfo | None\]: """Extract structured table information from CSV files.""" table\_names: list\[str\] = \[\] semaphore = asyncio.Semaphore(concurrency) tablename\_lock = asyncio.Lock() llm = OpenAI(model=model) prompt\_str = """\\ Provide a JSON object with the following fields: - \`table\_name\`: must be unique and descriptive (underscores only, no generic names). - \`table\_summary\`: short and concise summary of the table. Do NOT use any of these table names: {exclude\_table\_name\_list} Table: {table\_str} {feedback} """ prompt\_tmpl = ChatPromptTemplate( message\_templates=\[ChatMessage.from\_str(prompt\_str, role="user")\] ) tasks = \[\ process\_csv\_file(\ csv\_file, table\_names, semaphore, tablename\_lock, llm, prompt\_tmpl\ )\ async for csv\_file in data\_dir.walk()\ \] return await asyncio.gather(\*tasks) # {{docs-fragment data\_ingestion}} @env.task async def data\_ingestion( csv\_zip\_path: str = "https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip", search\_glob: str = "WikiTableQuestions/csv/200-csv/\*.csv", concurrency: int = 5, model: str = "gpt-4o-mini", ) -> tuple\[File, list\[TableInfo | None\]\]: """Main data ingestion pipeline: download → extract → analyze → create DB.""" data\_dir = await download\_and\_extract(csv\_zip\_path, search\_glob) table\_infos = await extract\_table\_info(data\_dir, model, concurrency) database\_path = "wiki\_table\_questions.db" i = 0 async for csv\_file in data\_dir.walk(): table\_info = table\_infos\[i\] if table\_info: ok = await create\_table(csv\_file, table\_info, database\_path) if ok == "false": table\_infos\[i\] = None else: print(f"Skipping table creation for {csv\_file} due to missing TableInfo.") i += 1 db\_file = await File.from\_local(database\_path) return db\_file, table\_infos # {{/docs-fragment data\_ingestion}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/data\_ingestion.py\* The ingestion step: 1. Downloads the dataset (a zip archive from GitHub). 2. Extracts the CSV files locally. 3. Generates table metadata (names and descriptions). 4. Creates corresponding tables in SQLite. The Flyte task returns both the path to the database and the generated table metadata. \`\`\` import asyncio import fnmatch import os import re import zipfile import flyte import pandas as pd import requests from flyte.io import Dir, File from llama\_index.core.llms import ChatMessage from llama\_index.core.prompts import ChatPromptTemplate from llama\_index.llms.openai import OpenAI from pydantic import BaseModel, Field from sqlalchemy import Column, Integer, MetaData, String, Table, create\_engine from utils import env # {{docs-fragment table\_info}} class TableInfo(BaseModel): """Information regarding a structured table.""" table\_name: str = Field(..., description="table name (underscores only, no spaces)") table\_summary: str = Field( ..., description="short, concise summary/caption of the table" ) # {{/docs-fragment table\_info}} @env.task async def download\_and\_extract(zip\_path: str, search\_glob: str) -> Dir: """Download and extract the dataset zip file if not already available.""" output\_zip = "data.zip" extract\_dir = "wiki\_table\_questions" if not os.path.exists(zip\_path): response = requests.get(zip\_path, stream=True) with open(output\_zip, "wb") as f: for chunk in response.iter\_content(chunk\_size=8192): f.write(chunk) else: output\_zip = zip\_path print(f"Using existing file {output\_zip}") os.makedirs(extract\_dir, exist\_ok=True) with zipfile.ZipFile(output\_zip, "r") as zip\_ref: for member in zip\_ref.namelist(): if fnmatch.fnmatch(member, search\_glob): zip\_ref.extract(member, extract\_dir) remote\_dir = await Dir.from\_local(extract\_dir) return remote\_dir async def read\_csv\_file( csv\_file: File, nrows: int | None = None ) -> pd.DataFrame | None: """Safely download and parse a CSV file into a DataFrame.""" try: local\_csv\_file = await csv\_file.download() return pd.read\_csv(local\_csv\_file, nrows=nrows) except Exception as e: print(f"Error parsing {csv\_file.path}: {e}") return None def sanitize\_column\_name(col\_name: str) -> str: """Sanitize column names by replacing spaces/special chars with underscores.""" return re.sub(r"\\W+", "\_", col\_name) async def create\_table\_from\_dataframe( df: pd.DataFrame, table\_name: str, engine, metadata\_obj ): """Create a SQL table from a Pandas DataFrame.""" # Sanitize column names sanitized\_columns = {col: sanitize\_column\_name(col) for col in df.columns} df = df.rename(columns=sanitized\_columns) # Define table columns based on DataFrame dtypes columns = \[\ Column(col, String if dtype == "object" else Integer)\ for col, dtype in zip(df.columns, df.dtypes)\ \] table = Table(table\_name, metadata\_obj, \*columns) # Create table in database metadata\_obj.create\_all(engine) # Insert data into table with engine.begin() as conn: for \_, row in df.iterrows(): conn.execute(table.insert().values(\*\*row.to\_dict())) @flyte.trace async def create\_table( csv\_file: File, table\_info: TableInfo, database\_path: str ) -> str: """Safely create a table from CSV if parsing succeeds.""" df = await read\_csv\_file(csv\_file) if df is None: return "false" print(f"Creating table: {table\_info.table\_name}") engine = create\_engine(f"sqlite:///{database\_path}") metadata\_obj = MetaData() await create\_table\_from\_dataframe(df, table\_info.table\_name, engine, metadata\_obj) return "true" @flyte.trace async def llm\_structured\_predict( df\_str: str, table\_names: list\[str\], prompt\_tmpl: ChatPromptTemplate, feedback: str, llm: OpenAI, ) -> TableInfo: return llm.structured\_predict( TableInfo, prompt\_tmpl, feedback=feedback, table\_str=df\_str, exclude\_table\_name\_list=str(list(table\_names)), ) async def generate\_unique\_table\_info( df\_str: str, table\_names: list\[str\], prompt\_tmpl: ChatPromptTemplate, llm: OpenAI, tablename\_lock: asyncio.Lock, retries: int = 3, ) -> TableInfo | None: """Process a single CSV file to generate a unique TableInfo.""" last\_table\_name = None for attempt in range(retries): feedback = "" if attempt > 0: feedback = f"Note: '{last\_table\_name}' already exists. Please pick a new name not in {table\_names}." table\_info = await llm\_structured\_predict( df\_str, table\_names, prompt\_tmpl, feedback, llm ) last\_table\_name = table\_info.table\_name async with tablename\_lock: if table\_info.table\_name not in table\_names: table\_names.append(table\_info.table\_name) return table\_info print(f"Table name {table\_info.table\_name} already exists, retrying...") return None async def process\_csv\_file( csv\_file: File, table\_names: list\[str\], semaphore: asyncio.Semaphore, tablename\_lock: asyncio.Lock, llm: OpenAI, prompt\_tmpl: ChatPromptTemplate, ) -> TableInfo | None: """Process a single CSV file to generate a unique TableInfo.""" async with semaphore: df = await read\_csv\_file(csv\_file, nrows=10) if df is None: return None return await generate\_unique\_table\_info( df.to\_csv(), table\_names, prompt\_tmpl, llm, tablename\_lock ) @env.task async def extract\_table\_info( data\_dir: Dir, model: str, concurrency: int ) -> list\[TableInfo | None\]: """Extract structured table information from CSV files.""" table\_names: list\[str\] = \[\] semaphore = asyncio.Semaphore(concurrency) tablename\_lock = asyncio.Lock() llm = OpenAI(model=model) prompt\_str = """\\ Provide a JSON object with the following fields: - \`table\_name\`: must be unique and descriptive (underscores only, no generic names). - \`table\_summary\`: short and concise summary of the table. Do NOT use any of these table names: {exclude\_table\_name\_list} Table: {table\_str} {feedback} """ prompt\_tmpl = ChatPromptTemplate( message\_templates=\[ChatMessage.from\_str(prompt\_str, role="user")\] ) tasks = \[\ process\_csv\_file(\ csv\_file, table\_names, semaphore, tablename\_lock, llm, prompt\_tmpl\ )\ async for csv\_file in data\_dir.walk()\ \] return await asyncio.gather(\*tasks) # {{docs-fragment data\_ingestion}} @env.task async def data\_ingestion( csv\_zip\_path: str = "https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip", search\_glob: str = "WikiTableQuestions/csv/200-csv/\*.csv", concurrency: int = 5, model: str = "gpt-4o-mini", ) -> tuple\[File, list\[TableInfo | None\]\]: """Main data ingestion pipeline: download → extract → analyze → create DB.""" data\_dir = await download\_and\_extract(csv\_zip\_path, search\_glob) table\_infos = await extract\_table\_info(data\_dir, model, concurrency) database\_path = "wiki\_table\_questions.db" i = 0 async for csv\_file in data\_dir.walk(): table\_info = table\_infos\[i\] if table\_info: ok = await create\_table(csv\_file, table\_info, database\_path) if ok == "false": table\_infos\[i\] = None else: print(f"Skipping table creation for {csv\_file} due to missing TableInfo.") i += 1 db\_file = await File.from\_local(database\_path) return db\_file, table\_infos # {{/docs-fragment data\_ingestion}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/data\_ingestion.py\* ## From question to SQL Next, we define a workflow that converts natural language into executable SQL using a retrieval-augmented generation (RAG) approach. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # "sqlalchemy>=2.0.0",\ # "pandas>=2.0.0",\ # "requests>=2.25.0",\ # "pydantic>=2.0.0",\ # \] # main = "text\_to\_sql" # params = "" # /// import asyncio from pathlib import Path import flyte from data\_ingestion import TableInfo, data\_ingestion from flyte.io import Dir, File from llama\_index.core import ( PromptTemplate, SQLDatabase, StorageContext, VectorStoreIndex, load\_index\_from\_storage, ) from llama\_index.core.llms import ChatResponse from llama\_index.core.objects import ObjectIndex, SQLTableNodeMapping, SQLTableSchema from llama\_index.core.prompts.prompt\_type import PromptType from llama\_index.core.retrievers import SQLRetriever from llama\_index.core.schema import TextNode from llama\_index.llms.openai import OpenAI from sqlalchemy import create\_engine, text from utils import env # {{docs-fragment index\_tables}} @flyte.trace async def index\_table(table\_name: str, table\_index\_dir: str, database\_uri: str) -> str: """Index a single table into vector store.""" path = f"{table\_index\_dir}/{table\_name}" engine = create\_engine(database\_uri) def \_fetch\_rows(): with engine.connect() as conn: cursor = conn.execute(text(f'SELECT \* FROM "{table\_name}"')) return cursor.fetchall() result = await asyncio.to\_thread(\_fetch\_rows) nodes = \[TextNode(text=str(tuple(row))) for row in result\] index = VectorStoreIndex(nodes) index.set\_index\_id("vector\_index") index.storage\_context.persist(path) return path @env.task async def index\_all\_tables(db\_file: File) -> Dir: """Index all tables concurrently.""" table\_index\_dir = "table\_indices" Path(table\_index\_dir).mkdir(exist\_ok=True) await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) tasks = \[\ index\_table(t, table\_index\_dir, "sqlite:///local\_db.sqlite")\ for t in sql\_database.get\_usable\_table\_names()\ \] await asyncio.gather(\*tasks) remote\_dir = await Dir.from\_local(table\_index\_dir) return remote\_dir # {{/docs-fragment index\_tables}} @flyte.trace async def get\_table\_schema\_context( table\_schema\_obj: SQLTableSchema, database\_uri: str, ) -> str: """Retrieve schema + optional description context for a single table.""" engine = create\_engine(database\_uri) sql\_database = SQLDatabase(engine) table\_info = sql\_database.get\_single\_table\_info(table\_schema\_obj.table\_name) if table\_schema\_obj.context\_str: table\_info += f" The table description is: {table\_schema\_obj.context\_str}" return table\_info @flyte.trace async def get\_table\_row\_context( table\_schema\_obj: SQLTableSchema, local\_vector\_index\_dir: str, query: str, ) -> str: """Retrieve row-level context examples using vector search.""" storage\_context = StorageContext.from\_defaults( persist\_dir=str(f"{local\_vector\_index\_dir}/{table\_schema\_obj.table\_name}") ) vector\_index = load\_index\_from\_storage(storage\_context, index\_id="vector\_index") vector\_retriever = vector\_index.as\_retriever(similarity\_top\_k=2) relevant\_nodes = vector\_retriever.retrieve(query) if not relevant\_nodes: return "" row\_context = "\\nHere are some relevant example rows (values in the same order as columns above)\\n" for node in relevant\_nodes: row\_context += str(node.get\_content()) + "\\n" return row\_context async def process\_table( table\_schema\_obj: SQLTableSchema, database\_uri: str, local\_vector\_index\_dir: str, query: str, ) -> str: """Combine schema + row context for one table.""" table\_info = await get\_table\_schema\_context(table\_schema\_obj, database\_uri) row\_context = await get\_table\_row\_context( table\_schema\_obj, local\_vector\_index\_dir, query ) full\_context = table\_info if row\_context: full\_context += "\\n" + row\_context print(f"Table Info: {full\_context}") return full\_context async def get\_table\_context\_and\_rows\_str( query: str, database\_uri: str, table\_schema\_objs: list\[SQLTableSchema\], vector\_index\_dir: Dir, ): """Get combined schema + row context for all tables.""" local\_vector\_index\_dir = await vector\_index\_dir.download() # run per-table work concurrently context\_strs = await asyncio.gather( \*\[\ process\_table(t, database\_uri, local\_vector\_index\_dir, query)\ for t in table\_schema\_objs\ \] ) return "\\n\\n".join(context\_strs) # {{docs-fragment retrieve\_tables}} @env.task async def retrieve\_tables( query: str, table\_infos: list\[TableInfo | None\], db\_file: File, vector\_index\_dir: Dir, ) -> str: """Retrieve relevant tables and return schema context string.""" await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) table\_node\_mapping = SQLTableNodeMapping(sql\_database) table\_schema\_objs = \[\ SQLTableSchema(table\_name=t.table\_name, context\_str=t.table\_summary)\ for t in table\_infos\ if t is not None\ \] obj\_index = ObjectIndex.from\_objects( table\_schema\_objs, table\_node\_mapping, VectorStoreIndex, ) obj\_retriever = obj\_index.as\_retriever(similarity\_top\_k=3) retrieved\_schemas = obj\_retriever.retrieve(query) return await get\_table\_context\_and\_rows\_str( query, "sqlite:///local\_db.sqlite", retrieved\_schemas, vector\_index\_dir ) # {{/docs-fragment retrieve\_tables}} def parse\_response\_to\_sql(chat\_response: ChatResponse) -> str: """Extract SQL query from LLM response.""" response = chat\_response.message.content sql\_query\_start = response.find("SQLQuery:") if sql\_query\_start != -1: response = response\[sql\_query\_start:\] if response.startswith("SQLQuery:"): response = response\[len("SQLQuery:") :\] sql\_result\_start = response.find("SQLResult:") if sql\_result\_start != -1: response = response\[:sql\_result\_start\] return response.strip().strip("\`\`\`").strip() # {{docs-fragment sql\_and\_response}} @env.task async def generate\_sql(query: str, table\_context: str, model: str, prompt: str) -> str: """Generate SQL query from natural language question and table context.""" llm = OpenAI(model=model) fmt\_messages = ( PromptTemplate( prompt, prompt\_type=PromptType.TEXT\_TO\_SQL, ) .partial\_format(dialect="sqlite") .format\_messages(query\_str=query, schema=table\_context) ) chat\_response = await llm.achat(fmt\_messages) return parse\_response\_to\_sql(chat\_response) @env.task async def generate\_response(query: str, sql: str, db\_file: File, model: str) -> str: """Run SQL query on database and synthesize final response.""" await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) sql\_retriever = SQLRetriever(sql\_database) retrieved\_rows = sql\_retriever.retrieve(sql) response\_synthesis\_prompt = PromptTemplate( "Given an input question, synthesize a response from the query results.\\n" "Query: {query\_str}\\n" "SQL: {sql\_query}\\n" "SQL Response: {context\_str}\\n" "Response: " ) llm = OpenAI(model=model) fmt\_messages = response\_synthesis\_prompt.format\_messages( sql\_query=sql, context\_str=str(retrieved\_rows), query\_str=query, ) chat\_response = await llm.achat(fmt\_messages) return chat\_response.message.content # {{/docs-fragment sql\_and\_response}} # {{docs-fragment text\_to\_sql}} @env.task async def text\_to\_sql( system\_prompt: str = ( "Given an input question, first create a syntactically correct {dialect} " "query to run, then look at the results of the query and return the answer. " "You can order the results by a relevant column to return the most " "interesting examples in the database.\\n\\n" "Never query for all the columns from a specific table, only ask for a " "few relevant columns given the question.\\n\\n" "Pay attention to use only the column names that you can see in the schema " "description. " "Be careful to not query for columns that do not exist. " "Pay attention to which column is in which table. " "Also, qualify column names with the table name when needed. " "You are required to use the following format, each taking one line:\\n\\n" "Question: Question here\\n" "SQLQuery: SQL Query to run\\n" "SQLResult: Result of the SQLQuery\\n" "Answer: Final answer here\\n\\n" "Only use tables listed below.\\n" "{schema}\\n\\n" "Question: {query\_str}\\n" "SQLQuery: " ), query: str = "What was the year that The Notorious BIG was signed to Bad Boy?", model: str = "gpt-4o-mini", ) -> str: db\_file, table\_infos = await data\_ingestion() vector\_index\_dir = await index\_all\_tables(db\_file) table\_context = await retrieve\_tables(query, table\_infos, db\_file, vector\_index\_dir) sql = await generate\_sql(query, table\_context, model, system\_prompt) return await generate\_response(query, sql, db\_file, model) # {{/docs-fragment text\_to\_sql}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(text\_to\_sql) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/text\_to\_sql.py\* The main \`text\_to\_sql\` task orchestrates the pipeline: - Ingest data - Build vector indices for each table - Retrieve relevant tables and rows - Generate SQL queries with an LLM - Execute queries and synthesize answers We use OpenAI GPT models with carefully structured prompts to maximize SQL correctness. ### Vector indexing We index each table's rows semantically so the model can retrieve relevant examples during SQL generation. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # "sqlalchemy>=2.0.0",\ # "pandas>=2.0.0",\ # "requests>=2.25.0",\ # "pydantic>=2.0.0",\ # \] # main = "text\_to\_sql" # params = "" # /// import asyncio from pathlib import Path import flyte from data\_ingestion import TableInfo, data\_ingestion from flyte.io import Dir, File from llama\_index.core import ( PromptTemplate, SQLDatabase, StorageContext, VectorStoreIndex, load\_index\_from\_storage, ) from llama\_index.core.llms import ChatResponse from llama\_index.core.objects import ObjectIndex, SQLTableNodeMapping, SQLTableSchema from llama\_index.core.prompts.prompt\_type import PromptType from llama\_index.core.retrievers import SQLRetriever from llama\_index.core.schema import TextNode from llama\_index.llms.openai import OpenAI from sqlalchemy import create\_engine, text from utils import env # {{docs-fragment index\_tables}} @flyte.trace async def index\_table(table\_name: str, table\_index\_dir: str, database\_uri: str) -> str: """Index a single table into vector store.""" path = f"{table\_index\_dir}/{table\_name}" engine = create\_engine(database\_uri) def \_fetch\_rows(): with engine.connect() as conn: cursor = conn.execute(text(f'SELECT \* FROM "{table\_name}"')) return cursor.fetchall() result = await asyncio.to\_thread(\_fetch\_rows) nodes = \[TextNode(text=str(tuple(row))) for row in result\] index = VectorStoreIndex(nodes) index.set\_index\_id("vector\_index") index.storage\_context.persist(path) return path @env.task async def index\_all\_tables(db\_file: File) -> Dir: """Index all tables concurrently.""" table\_index\_dir = "table\_indices" Path(table\_index\_dir).mkdir(exist\_ok=True) await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) tasks = \[\ index\_table(t, table\_index\_dir, "sqlite:///local\_db.sqlite")\ for t in sql\_database.get\_usable\_table\_names()\ \] await asyncio.gather(\*tasks) remote\_dir = await Dir.from\_local(table\_index\_dir) return remote\_dir # {{/docs-fragment index\_tables}} @flyte.trace async def get\_table\_schema\_context( table\_schema\_obj: SQLTableSchema, database\_uri: str, ) -> str: """Retrieve schema + optional description context for a single table.""" engine = create\_engine(database\_uri) sql\_database = SQLDatabase(engine) table\_info = sql\_database.get\_single\_table\_info(table\_schema\_obj.table\_name) if table\_schema\_obj.context\_str: table\_info += f" The table description is: {table\_schema\_obj.context\_str}" return table\_info @flyte.trace async def get\_table\_row\_context( table\_schema\_obj: SQLTableSchema, local\_vector\_index\_dir: str, query: str, ) -> str: """Retrieve row-level context examples using vector search.""" storage\_context = StorageContext.from\_defaults( persist\_dir=str(f"{local\_vector\_index\_dir}/{table\_schema\_obj.table\_name}") ) vector\_index = load\_index\_from\_storage(storage\_context, index\_id="vector\_index") vector\_retriever = vector\_index.as\_retriever(similarity\_top\_k=2) relevant\_nodes = vector\_retriever.retrieve(query) if not relevant\_nodes: return "" row\_context = "\\nHere are some relevant example rows (values in the same order as columns above)\\n" for node in relevant\_nodes: row\_context += str(node.get\_content()) + "\\n" return row\_context async def process\_table( table\_schema\_obj: SQLTableSchema, database\_uri: str, local\_vector\_index\_dir: str, query: str, ) -> str: """Combine schema + row context for one table.""" table\_info = await get\_table\_schema\_context(table\_schema\_obj, database\_uri) row\_context = await get\_table\_row\_context( table\_schema\_obj, local\_vector\_index\_dir, query ) full\_context = table\_info if row\_context: full\_context += "\\n" + row\_context print(f"Table Info: {full\_context}") return full\_context async def get\_table\_context\_and\_rows\_str( query: str, database\_uri: str, table\_schema\_objs: list\[SQLTableSchema\], vector\_index\_dir: Dir, ): """Get combined schema + row context for all tables.""" local\_vector\_index\_dir = await vector\_index\_dir.download() # run per-table work concurrently context\_strs = await asyncio.gather( \*\[\ process\_table(t, database\_uri, local\_vector\_index\_dir, query)\ for t in table\_schema\_objs\ \] ) return "\\n\\n".join(context\_strs) # {{docs-fragment retrieve\_tables}} @env.task async def retrieve\_tables( query: str, table\_infos: list\[TableInfo | None\], db\_file: File, vector\_index\_dir: Dir, ) -> str: """Retrieve relevant tables and return schema context string.""" await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) table\_node\_mapping = SQLTableNodeMapping(sql\_database) table\_schema\_objs = \[\ SQLTableSchema(table\_name=t.table\_name, context\_str=t.table\_summary)\ for t in table\_infos\ if t is not None\ \] obj\_index = ObjectIndex.from\_objects( table\_schema\_objs, table\_node\_mapping, VectorStoreIndex, ) obj\_retriever = obj\_index.as\_retriever(similarity\_top\_k=3) retrieved\_schemas = obj\_retriever.retrieve(query) return await get\_table\_context\_and\_rows\_str( query, "sqlite:///local\_db.sqlite", retrieved\_schemas, vector\_index\_dir ) # {{/docs-fragment retrieve\_tables}} def parse\_response\_to\_sql(chat\_response: ChatResponse) -> str: """Extract SQL query from LLM response.""" response = chat\_response.message.content sql\_query\_start = response.find("SQLQuery:") if sql\_query\_start != -1: response = response\[sql\_query\_start:\] if response.startswith("SQLQuery:"): response = response\[len("SQLQuery:") :\] sql\_result\_start = response.find("SQLResult:") if sql\_result\_start != -1: response = response\[:sql\_result\_start\] return response.strip().strip("\`\`\`").strip() # {{docs-fragment sql\_and\_response}} @env.task async def generate\_sql(query: str, table\_context: str, model: str, prompt: str) -> str: """Generate SQL query from natural language question and table context.""" llm = OpenAI(model=model) fmt\_messages = ( PromptTemplate( prompt, prompt\_type=PromptType.TEXT\_TO\_SQL, ) .partial\_format(dialect="sqlite") .format\_messages(query\_str=query, schema=table\_context) ) chat\_response = await llm.achat(fmt\_messages) return parse\_response\_to\_sql(chat\_response) @env.task async def generate\_response(query: str, sql: str, db\_file: File, model: str) -> str: """Run SQL query on database and synthesize final response.""" await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) sql\_retriever = SQLRetriever(sql\_database) retrieved\_rows = sql\_retriever.retrieve(sql) response\_synthesis\_prompt = PromptTemplate( "Given an input question, synthesize a response from the query results.\\n" "Query: {query\_str}\\n" "SQL: {sql\_query}\\n" "SQL Response: {context\_str}\\n" "Response: " ) llm = OpenAI(model=model) fmt\_messages = response\_synthesis\_prompt.format\_messages( sql\_query=sql, context\_str=str(retrieved\_rows), query\_str=query, ) chat\_response = await llm.achat(fmt\_messages) return chat\_response.message.content # {{/docs-fragment sql\_and\_response}} # {{docs-fragment text\_to\_sql}} @env.task async def text\_to\_sql( system\_prompt: str = ( "Given an input question, first create a syntactically correct {dialect} " "query to run, then look at the results of the query and return the answer. " "You can order the results by a relevant column to return the most " "interesting examples in the database.\\n\\n" "Never query for all the columns from a specific table, only ask for a " "few relevant columns given the question.\\n\\n" "Pay attention to use only the column names that you can see in the schema " "description. " "Be careful to not query for columns that do not exist. " "Pay attention to which column is in which table. " "Also, qualify column names with the table name when needed. " "You are required to use the following format, each taking one line:\\n\\n" "Question: Question here\\n" "SQLQuery: SQL Query to run\\n" "SQLResult: Result of the SQLQuery\\n" "Answer: Final answer here\\n\\n" "Only use tables listed below.\\n" "{schema}\\n\\n" "Question: {query\_str}\\n" "SQLQuery: " ), query: str = "What was the year that The Notorious BIG was signed to Bad Boy?", model: str = "gpt-4o-mini", ) -> str: db\_file, table\_infos = await data\_ingestion() vector\_index\_dir = await index\_all\_tables(db\_file) table\_context = await retrieve\_tables(query, table\_infos, db\_file, vector\_index\_dir) sql = await generate\_sql(query, table\_context, model, system\_prompt) return await generate\_response(query, sql, db\_file, model) # {{/docs-fragment text\_to\_sql}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(text\_to\_sql) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/text\_to\_sql.py\* Each row becomes a text node stored in LlamaIndex’s \`VectorStoreIndex\`. This lets the system pull semantically similar rows when handling queries. ### Table retrieval and context building We then retrieve the most relevant tables for a given query and build rich context that combines schema information with sample rows. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # "sqlalchemy>=2.0.0",\ # "pandas>=2.0.0",\ # "requests>=2.25.0",\ # "pydantic>=2.0.0",\ # \] # main = "text\_to\_sql" # params = "" # /// import asyncio from pathlib import Path import flyte from data\_ingestion import TableInfo, data\_ingestion from flyte.io import Dir, File from llama\_index.core import ( PromptTemplate, SQLDatabase, StorageContext, VectorStoreIndex, load\_index\_from\_storage, ) from llama\_index.core.llms import ChatResponse from llama\_index.core.objects import ObjectIndex, SQLTableNodeMapping, SQLTableSchema from llama\_index.core.prompts.prompt\_type import PromptType from llama\_index.core.retrievers import SQLRetriever from llama\_index.core.schema import TextNode from llama\_index.llms.openai import OpenAI from sqlalchemy import create\_engine, text from utils import env # {{docs-fragment index\_tables}} @flyte.trace async def index\_table(table\_name: str, table\_index\_dir: str, database\_uri: str) -> str: """Index a single table into vector store.""" path = f"{table\_index\_dir}/{table\_name}" engine = create\_engine(database\_uri) def \_fetch\_rows(): with engine.connect() as conn: cursor = conn.execute(text(f'SELECT \* FROM "{table\_name}"')) return cursor.fetchall() result = await asyncio.to\_thread(\_fetch\_rows) nodes = \[TextNode(text=str(tuple(row))) for row in result\] index = VectorStoreIndex(nodes) index.set\_index\_id("vector\_index") index.storage\_context.persist(path) return path @env.task async def index\_all\_tables(db\_file: File) -> Dir: """Index all tables concurrently.""" table\_index\_dir = "table\_indices" Path(table\_index\_dir).mkdir(exist\_ok=True) await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) tasks = \[\ index\_table(t, table\_index\_dir, "sqlite:///local\_db.sqlite")\ for t in sql\_database.get\_usable\_table\_names()\ \] await asyncio.gather(\*tasks) remote\_dir = await Dir.from\_local(table\_index\_dir) return remote\_dir # {{/docs-fragment index\_tables}} @flyte.trace async def get\_table\_schema\_context( table\_schema\_obj: SQLTableSchema, database\_uri: str, ) -> str: """Retrieve schema + optional description context for a single table.""" engine = create\_engine(database\_uri) sql\_database = SQLDatabase(engine) table\_info = sql\_database.get\_single\_table\_info(table\_schema\_obj.table\_name) if table\_schema\_obj.context\_str: table\_info += f" The table description is: {table\_schema\_obj.context\_str}" return table\_info @flyte.trace async def get\_table\_row\_context( table\_schema\_obj: SQLTableSchema, local\_vector\_index\_dir: str, query: str, ) -> str: """Retrieve row-level context examples using vector search.""" storage\_context = StorageContext.from\_defaults( persist\_dir=str(f"{local\_vector\_index\_dir}/{table\_schema\_obj.table\_name}") ) vector\_index = load\_index\_from\_storage(storage\_context, index\_id="vector\_index") vector\_retriever = vector\_index.as\_retriever(similarity\_top\_k=2) relevant\_nodes = vector\_retriever.retrieve(query) if not relevant\_nodes: return "" row\_context = "\\nHere are some relevant example rows (values in the same order as columns above)\\n" for node in relevant\_nodes: row\_context += str(node.get\_content()) + "\\n" return row\_context async def process\_table( table\_schema\_obj: SQLTableSchema, database\_uri: str, local\_vector\_index\_dir: str, query: str, ) -> str: """Combine schema + row context for one table.""" table\_info = await get\_table\_schema\_context(table\_schema\_obj, database\_uri) row\_context = await get\_table\_row\_context( table\_schema\_obj, local\_vector\_index\_dir, query ) full\_context = table\_info if row\_context: full\_context += "\\n" + row\_context print(f"Table Info: {full\_context}") return full\_context async def get\_table\_context\_and\_rows\_str( query: str, database\_uri: str, table\_schema\_objs: list\[SQLTableSchema\], vector\_index\_dir: Dir, ): """Get combined schema + row context for all tables.""" local\_vector\_index\_dir = await vector\_index\_dir.download() # run per-table work concurrently context\_strs = await asyncio.gather( \*\[\ process\_table(t, database\_uri, local\_vector\_index\_dir, query)\ for t in table\_schema\_objs\ \] ) return "\\n\\n".join(context\_strs) # {{docs-fragment retrieve\_tables}} @env.task async def retrieve\_tables( query: str, table\_infos: list\[TableInfo | None\], db\_file: File, vector\_index\_dir: Dir, ) -> str: """Retrieve relevant tables and return schema context string.""" await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) table\_node\_mapping = SQLTableNodeMapping(sql\_database) table\_schema\_objs = \[\ SQLTableSchema(table\_name=t.table\_name, context\_str=t.table\_summary)\ for t in table\_infos\ if t is not None\ \] obj\_index = ObjectIndex.from\_objects( table\_schema\_objs, table\_node\_mapping, VectorStoreIndex, ) obj\_retriever = obj\_index.as\_retriever(similarity\_top\_k=3) retrieved\_schemas = obj\_retriever.retrieve(query) return await get\_table\_context\_and\_rows\_str( query, "sqlite:///local\_db.sqlite", retrieved\_schemas, vector\_index\_dir ) # {{/docs-fragment retrieve\_tables}} def parse\_response\_to\_sql(chat\_response: ChatResponse) -> str: """Extract SQL query from LLM response.""" response = chat\_response.message.content sql\_query\_start = response.find("SQLQuery:") if sql\_query\_start != -1: response = response\[sql\_query\_start:\] if response.startswith("SQLQuery:"): response = response\[len("SQLQuery:") :\] sql\_result\_start = response.find("SQLResult:") if sql\_result\_start != -1: response = response\[:sql\_result\_start\] return response.strip().strip("\`\`\`").strip() # {{docs-fragment sql\_and\_response}} @env.task async def generate\_sql(query: str, table\_context: str, model: str, prompt: str) -> str: """Generate SQL query from natural language question and table context.""" llm = OpenAI(model=model) fmt\_messages = ( PromptTemplate( prompt, prompt\_type=PromptType.TEXT\_TO\_SQL, ) .partial\_format(dialect="sqlite") .format\_messages(query\_str=query, schema=table\_context) ) chat\_response = await llm.achat(fmt\_messages) return parse\_response\_to\_sql(chat\_response) @env.task async def generate\_response(query: str, sql: str, db\_file: File, model: str) -> str: """Run SQL query on database and synthesize final response.""" await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) sql\_retriever = SQLRetriever(sql\_database) retrieved\_rows = sql\_retriever.retrieve(sql) response\_synthesis\_prompt = PromptTemplate( "Given an input question, synthesize a response from the query results.\\n" "Query: {query\_str}\\n" "SQL: {sql\_query}\\n" "SQL Response: {context\_str}\\n" "Response: " ) llm = OpenAI(model=model) fmt\_messages = response\_synthesis\_prompt.format\_messages( sql\_query=sql, context\_str=str(retrieved\_rows), query\_str=query, ) chat\_response = await llm.achat(fmt\_messages) return chat\_response.message.content # {{/docs-fragment sql\_and\_response}} # {{docs-fragment text\_to\_sql}} @env.task async def text\_to\_sql( system\_prompt: str = ( "Given an input question, first create a syntactically correct {dialect} " "query to run, then look at the results of the query and return the answer. " "You can order the results by a relevant column to return the most " "interesting examples in the database.\\n\\n" "Never query for all the columns from a specific table, only ask for a " "few relevant columns given the question.\\n\\n" "Pay attention to use only the column names that you can see in the schema " "description. " "Be careful to not query for columns that do not exist. " "Pay attention to which column is in which table. " "Also, qualify column names with the table name when needed. " "You are required to use the following format, each taking one line:\\n\\n" "Question: Question here\\n" "SQLQuery: SQL Query to run\\n" "SQLResult: Result of the SQLQuery\\n" "Answer: Final answer here\\n\\n" "Only use tables listed below.\\n" "{schema}\\n\\n" "Question: {query\_str}\\n" "SQLQuery: " ), query: str = "What was the year that The Notorious BIG was signed to Bad Boy?", model: str = "gpt-4o-mini", ) -> str: db\_file, table\_infos = await data\_ingestion() vector\_index\_dir = await index\_all\_tables(db\_file) table\_context = await retrieve\_tables(query, table\_infos, db\_file, vector\_index\_dir) sql = await generate\_sql(query, table\_context, model, system\_prompt) return await generate\_response(query, sql, db\_file, model) # {{/docs-fragment text\_to\_sql}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(text\_to\_sql) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/text\_to\_sql.py\* The retriever selects tables via semantic similarity, then attaches their schema and example rows. This context grounds the model's SQL generation in the database's actual structure and content. ### SQL generation and response synthesis Finally, we generate SQL queries and produce natural language answers. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # "sqlalchemy>=2.0.0",\ # "pandas>=2.0.0",\ # "requests>=2.25.0",\ # "pydantic>=2.0.0",\ # \] # main = "text\_to\_sql" # params = "" # /// import asyncio from pathlib import Path import flyte from data\_ingestion import TableInfo, data\_ingestion from flyte.io import Dir, File from llama\_index.core import ( PromptTemplate, SQLDatabase, StorageContext, VectorStoreIndex, load\_index\_from\_storage, ) from llama\_index.core.llms import ChatResponse from llama\_index.core.objects import ObjectIndex, SQLTableNodeMapping, SQLTableSchema from llama\_index.core.prompts.prompt\_type import PromptType from llama\_index.core.retrievers import SQLRetriever from llama\_index.core.schema import TextNode from llama\_index.llms.openai import OpenAI from sqlalchemy import create\_engine, text from utils import env # {{docs-fragment index\_tables}} @flyte.trace async def index\_table(table\_name: str, table\_index\_dir: str, database\_uri: str) -> str: """Index a single table into vector store.""" path = f"{table\_index\_dir}/{table\_name}" engine = create\_engine(database\_uri) def \_fetch\_rows(): with engine.connect() as conn: cursor = conn.execute(text(f'SELECT \* FROM "{table\_name}"')) return cursor.fetchall() result = await asyncio.to\_thread(\_fetch\_rows) nodes = \[TextNode(text=str(tuple(row))) for row in result\] index = VectorStoreIndex(nodes) index.set\_index\_id("vector\_index") index.storage\_context.persist(path) return path @env.task async def index\_all\_tables(db\_file: File) -> Dir: """Index all tables concurrently.""" table\_index\_dir = "table\_indices" Path(table\_index\_dir).mkdir(exist\_ok=True) await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) tasks = \[\ index\_table(t, table\_index\_dir, "sqlite:///local\_db.sqlite")\ for t in sql\_database.get\_usable\_table\_names()\ \] await asyncio.gather(\*tasks) remote\_dir = await Dir.from\_local(table\_index\_dir) return remote\_dir # {{/docs-fragment index\_tables}} @flyte.trace async def get\_table\_schema\_context( table\_schema\_obj: SQLTableSchema, database\_uri: str, ) -> str: """Retrieve schema + optional description context for a single table.""" engine = create\_engine(database\_uri) sql\_database = SQLDatabase(engine) table\_info = sql\_database.get\_single\_table\_info(table\_schema\_obj.table\_name) if table\_schema\_obj.context\_str: table\_info += f" The table description is: {table\_schema\_obj.context\_str}" return table\_info @flyte.trace async def get\_table\_row\_context( table\_schema\_obj: SQLTableSchema, local\_vector\_index\_dir: str, query: str, ) -> str: """Retrieve row-level context examples using vector search.""" storage\_context = StorageContext.from\_defaults( persist\_dir=str(f"{local\_vector\_index\_dir}/{table\_schema\_obj.table\_name}") ) vector\_index = load\_index\_from\_storage(storage\_context, index\_id="vector\_index") vector\_retriever = vector\_index.as\_retriever(similarity\_top\_k=2) relevant\_nodes = vector\_retriever.retrieve(query) if not relevant\_nodes: return "" row\_context = "\\nHere are some relevant example rows (values in the same order as columns above)\\n" for node in relevant\_nodes: row\_context += str(node.get\_content()) + "\\n" return row\_context async def process\_table( table\_schema\_obj: SQLTableSchema, database\_uri: str, local\_vector\_index\_dir: str, query: str, ) -> str: """Combine schema + row context for one table.""" table\_info = await get\_table\_schema\_context(table\_schema\_obj, database\_uri) row\_context = await get\_table\_row\_context( table\_schema\_obj, local\_vector\_index\_dir, query ) full\_context = table\_info if row\_context: full\_context += "\\n" + row\_context print(f"Table Info: {full\_context}") return full\_context async def get\_table\_context\_and\_rows\_str( query: str, database\_uri: str, table\_schema\_objs: list\[SQLTableSchema\], vector\_index\_dir: Dir, ): """Get combined schema + row context for all tables.""" local\_vector\_index\_dir = await vector\_index\_dir.download() # run per-table work concurrently context\_strs = await asyncio.gather( \*\[\ process\_table(t, database\_uri, local\_vector\_index\_dir, query)\ for t in table\_schema\_objs\ \] ) return "\\n\\n".join(context\_strs) # {{docs-fragment retrieve\_tables}} @env.task async def retrieve\_tables( query: str, table\_infos: list\[TableInfo | None\], db\_file: File, vector\_index\_dir: Dir, ) -> str: """Retrieve relevant tables and return schema context string.""" await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) table\_node\_mapping = SQLTableNodeMapping(sql\_database) table\_schema\_objs = \[\ SQLTableSchema(table\_name=t.table\_name, context\_str=t.table\_summary)\ for t in table\_infos\ if t is not None\ \] obj\_index = ObjectIndex.from\_objects( table\_schema\_objs, table\_node\_mapping, VectorStoreIndex, ) obj\_retriever = obj\_index.as\_retriever(similarity\_top\_k=3) retrieved\_schemas = obj\_retriever.retrieve(query) return await get\_table\_context\_and\_rows\_str( query, "sqlite:///local\_db.sqlite", retrieved\_schemas, vector\_index\_dir ) # {{/docs-fragment retrieve\_tables}} def parse\_response\_to\_sql(chat\_response: ChatResponse) -> str: """Extract SQL query from LLM response.""" response = chat\_response.message.content sql\_query\_start = response.find("SQLQuery:") if sql\_query\_start != -1: response = response\[sql\_query\_start:\] if response.startswith("SQLQuery:"): response = response\[len("SQLQuery:") :\] sql\_result\_start = response.find("SQLResult:") if sql\_result\_start != -1: response = response\[:sql\_result\_start\] return response.strip().strip("\`\`\`").strip() # {{docs-fragment sql\_and\_response}} @env.task async def generate\_sql(query: str, table\_context: str, model: str, prompt: str) -> str: """Generate SQL query from natural language question and table context.""" llm = OpenAI(model=model) fmt\_messages = ( PromptTemplate( prompt, prompt\_type=PromptType.TEXT\_TO\_SQL, ) .partial\_format(dialect="sqlite") .format\_messages(query\_str=query, schema=table\_context) ) chat\_response = await llm.achat(fmt\_messages) return parse\_response\_to\_sql(chat\_response) @env.task async def generate\_response(query: str, sql: str, db\_file: File, model: str) -> str: """Run SQL query on database and synthesize final response.""" await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) sql\_retriever = SQLRetriever(sql\_database) retrieved\_rows = sql\_retriever.retrieve(sql) response\_synthesis\_prompt = PromptTemplate( "Given an input question, synthesize a response from the query results.\\n" "Query: {query\_str}\\n" "SQL: {sql\_query}\\n" "SQL Response: {context\_str}\\n" "Response: " ) llm = OpenAI(model=model) fmt\_messages = response\_synthesis\_prompt.format\_messages( sql\_query=sql, context\_str=str(retrieved\_rows), query\_str=query, ) chat\_response = await llm.achat(fmt\_messages) return chat\_response.message.content # {{/docs-fragment sql\_and\_response}} # {{docs-fragment text\_to\_sql}} @env.task async def text\_to\_sql( system\_prompt: str = ( "Given an input question, first create a syntactically correct {dialect} " "query to run, then look at the results of the query and return the answer. " "You can order the results by a relevant column to return the most " "interesting examples in the database.\\n\\n" "Never query for all the columns from a specific table, only ask for a " "few relevant columns given the question.\\n\\n" "Pay attention to use only the column names that you can see in the schema " "description. " "Be careful to not query for columns that do not exist. " "Pay attention to which column is in which table. " "Also, qualify column names with the table name when needed. " "You are required to use the following format, each taking one line:\\n\\n" "Question: Question here\\n" "SQLQuery: SQL Query to run\\n" "SQLResult: Result of the SQLQuery\\n" "Answer: Final answer here\\n\\n" "Only use tables listed below.\\n" "{schema}\\n\\n" "Question: {query\_str}\\n" "SQLQuery: " ), query: str = "What was the year that The Notorious BIG was signed to Bad Boy?", model: str = "gpt-4o-mini", ) -> str: db\_file, table\_infos = await data\_ingestion() vector\_index\_dir = await index\_all\_tables(db\_file) table\_context = await retrieve\_tables(query, table\_infos, db\_file, vector\_index\_dir) sql = await generate\_sql(query, table\_context, model, system\_prompt) return await generate\_response(query, sql, db\_file, model) # {{/docs-fragment text\_to\_sql}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(text\_to\_sql) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/text\_to\_sql.py\* The SQL generation prompt includes schema, example rows, and formatting rules. After execution, the system returns a final answer. At this point, we have an end-to-end Text-to-SQL pipeline: natural language questions go in, SQL queries run, and answers come back. To make this workflow production-ready, we leveraged several Flyte 2 capabilities. Caching ensures that repeated steps, like table ingestion or vector indexing, don’t need to rerun unnecessarily, saving time and compute. Containerization provides consistent, reproducible execution across environments, making it easier to scale and deploy. Observability features let us track every step of the pipeline, monitor performance, and debug issues quickly. While the pipeline works end-to-end, to get a pulse on how it performs across multiple prompts and to gradually improve performance, we can start experimenting with prompt tuning. Two things help make this process meaningful: - \*\*A clean evaluation dataset\*\* - so we can measure accuracy against trusted ground truth. - \*\*A systematic evaluation loop\*\* - so we can see whether prompt changes or other adjustments actually help. With these in place, the next step is to build a "golden" QA dataset that will guide iterative prompt optimization. ## Building the QA dataset > \[!NOTE\] > The WikiTableQuestions dataset already includes question–answer pairs, available in its \[GitHub repository\](https://github.com/ppasupat/WikiTableQuestions/tree/master/data). To use them for this workflow, you'll need to adapt the data into the required format, but the raw material is there for you to build on. We generate a dataset of natural language questions paired with executable SQL queries. This dataset acts as the benchmark for prompt tuning and evaluation. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas>=2.0.0",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # "pydantic>=2.0.0",\ # \] # main = "build\_eval\_dataset" # params = "" # /// import sqlite3 import flyte import pandas as pd from data\_ingestion import data\_ingestion from flyte.io import File from llama\_index.core import PromptTemplate from llama\_index.llms.openai import OpenAI from utils import env from pydantic import BaseModel class QAItem(BaseModel): question: str sql: str class QAList(BaseModel): items: list\[QAItem\] # {{docs-fragment get\_and\_split\_schema}} @env.task async def get\_and\_split\_schema(db\_file: File, tables\_per\_chunk: int) -> list\[str\]: """ Download the SQLite DB, extract schema info (columns + sample rows), then split it into chunks with up to \`tables\_per\_chunk\` tables each. """ await db\_file.download(local\_path="local\_db.sqlite") conn = sqlite3.connect("local\_db.sqlite") cursor = conn.cursor() tables = cursor.execute( "SELECT name FROM sqlite\_master WHERE type='table';" ).fetchall() schema\_blocks = \[\] for table in tables: table\_name = table\[0\] # columns cursor.execute(f"PRAGMA table\_info({table\_name});") columns = \[col\[1\] for col in cursor.fetchall()\] block = f"Table: {table\_name}({', '.join(columns)})" # sample rows cursor.execute(f"SELECT \* FROM {table\_name} LIMIT 10;") rows = cursor.fetchall() if rows: block += "\\nSample rows:\\n" for row in rows: block += f"{row}\\n" schema\_blocks.append(block) conn.close() chunks = \[\] current\_chunk = \[\] for block in schema\_blocks: current\_chunk.append(block) if len(current\_chunk) >= tables\_per\_chunk: chunks.append("\\n".join(current\_chunk)) current\_chunk = \[\] if current\_chunk: chunks.append("\\n".join(current\_chunk)) return chunks # {{/docs-fragment get\_and\_split\_schema}} # {{docs-fragment generate\_questions\_and\_sql}} @flyte.trace async def generate\_questions\_and\_sql( schema: str, num\_samples: int, batch\_size: int ) -> QAList: llm = OpenAI(model="gpt-4.1") prompt\_tmpl = PromptTemplate( """Prompt: You are helping build a Text-to-SQL dataset. Here is the database schema: {schema} Generate {num} natural language questions a user might ask about this database. For each question, also provide the correct SQL query. Reasoning process (you must follow this internally): - Given an input question, first create a syntactically correct {dialect} SQL query. - Never use SELECT \*; only include the relevant columns. - Use only columns/tables from the schema. Qualify column names when ambiguous. - You may order results by a meaningful column to make the query more useful. - Be careful not to add unnecessary columns. - Use filters, aggregations, joins, grouping, and subqueries when relevant. Final Output: Return only a JSON object with one field: - "items": a list of {num} objects, each with: - "question": the natural language question - "sql": the corresponding SQL query """ ) all\_items: list\[QAItem\] = \[\] # batch generation for start in range(0, num\_samples, batch\_size): current\_num = min(batch\_size, num\_samples - start) response = llm.structured\_predict( QAList, prompt\_tmpl, schema=schema, num=current\_num, ) all\_items.extend(response.items) # deduplicate seen = set() unique\_items: list\[QAItem\] = \[\] for item in all\_items: key = (item.question.strip().lower(), item.sql.strip().lower()) if key not in seen: seen.add(key) unique\_items.append(item) return QAList(items=unique\_items\[:num\_samples\]) # {{/docs-fragment generate\_questions\_and\_sql}} @flyte.trace async def llm\_validate\_batch(pairs: list\[dict\[str, str\]\]) -> list\[str\]: """Validate a batch of question/sql/result dicts using one LLM call.""" batch\_prompt = """You are validating the correctness of SQL query results against the question. For each example, answer only "True" (correct) or "False" (incorrect). Output one answer per line, in the same order as the examples. --- """ for i, pair in enumerate(pairs, start=1): batch\_prompt += f""" Example {i}: Question: {pair\['question'\]} SQL: {pair\['sql'\]} Result: {pair\['rows'\]} --- """ llm = OpenAI(model="gpt-4.1") resp = await llm.acomplete(batch\_prompt) # Expect exactly one True/False per example results = \[\ line.strip()\ for line in resp.text.splitlines()\ if line.strip() in ("True", "False")\ \] return results # {{docs-fragment validate\_sql}} @env.task async def validate\_sql( db\_file: File, question\_sql\_pairs: QAList, batch\_size: int ) -> list\[dict\[str, str\]\]: await db\_file.download(local\_path="local\_db.sqlite") conn = sqlite3.connect("local\_db.sqlite") cursor = conn.cursor() qa\_data = \[\] batch = \[\] for pair in question\_sql\_pairs.items: q, sql = pair.question, pair.sql try: cursor.execute(sql) rows = cursor.fetchall() batch.append({"question": q, "sql": sql, "rows": str(rows)}) # process when batch is full if len(batch) == batch\_size: results = await llm\_validate\_batch(batch) for pair, is\_valid in zip(batch, results): if is\_valid == "True": qa\_data.append( { "input": pair\["question"\], "sql": pair\["sql"\], "target": pair\["rows"\], } ) else: print(f"Filtered out incorrect result for: {pair\['question'\]}") batch = \[\] except Exception as e: print(f"Skipping invalid SQL: {sql} ({e})") # process leftover batch if batch: results = await llm\_validate\_batch(batch) for pair, is\_valid in zip(batch, results): if is\_valid == "True": qa\_data.append( { "input": pair\["question"\], "sql": pair\["sql"\], "target": pair\["rows"\], } ) else: print(f"Filtered out incorrect result for: {pair\['question'\]}") conn.close() return qa\_data # {{/docs-fragment validate\_sql}} @flyte.trace async def save\_to\_csv(qa\_data: list\[dict\]) -> File: df = pd.DataFrame(qa\_data, columns=\["input", "target", "sql"\]) csv\_file = "qa\_dataset.csv" df.to\_csv(csv\_file, index=False) return await File.from\_local(csv\_file) # {{docs-fragment build\_eval\_dataset}} @env.task async def build\_eval\_dataset( num\_samples: int = 300, batch\_size: int = 30, tables\_per\_chunk: int = 3 ) -> File: db\_file, \_ = await data\_ingestion() schema\_chunks = await get\_and\_split\_schema(db\_file, tables\_per\_chunk) per\_chunk\_samples = max(1, num\_samples // len(schema\_chunks)) final\_qa\_data = \[\] for chunk in schema\_chunks: qa\_list = await generate\_questions\_and\_sql( schema=chunk, num\_samples=per\_chunk\_samples, batch\_size=batch\_size, ) qa\_data = await validate\_sql(db\_file, qa\_list, batch\_size) final\_qa\_data.extend(qa\_data) csv\_file = await save\_to\_csv(final\_qa\_data) return csv\_file # {{/docs-fragment build\_eval\_dataset}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(build\_eval\_dataset) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/create\_qa\_dataset.py\* The pipeline does the following: - Schema extraction – pull full database schemas, including table names, columns, and sample rows. - Question–SQL generation – use an LLM to produce natural language questions with matching SQL queries. - Validation – run each query against the database, filter out invalid results, and also remove results that aren't relevant. - Final export – store the clean, validated pairs in CSV format for downstream use. ### Schema extraction and chunking We break schemas into smaller chunks to cover all tables evenly. This avoids overfitting to a subset of tables and ensures broad coverage across the dataset. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas>=2.0.0",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # "pydantic>=2.0.0",\ # \] # main = "build\_eval\_dataset" # params = "" # /// import sqlite3 import flyte import pandas as pd from data\_ingestion import data\_ingestion from flyte.io import File from llama\_index.core import PromptTemplate from llama\_index.llms.openai import OpenAI from utils import env from pydantic import BaseModel class QAItem(BaseModel): question: str sql: str class QAList(BaseModel): items: list\[QAItem\] # {{docs-fragment get\_and\_split\_schema}} @env.task async def get\_and\_split\_schema(db\_file: File, tables\_per\_chunk: int) -> list\[str\]: """ Download the SQLite DB, extract schema info (columns + sample rows), then split it into chunks with up to \`tables\_per\_chunk\` tables each. """ await db\_file.download(local\_path="local\_db.sqlite") conn = sqlite3.connect("local\_db.sqlite") cursor = conn.cursor() tables = cursor.execute( "SELECT name FROM sqlite\_master WHERE type='table';" ).fetchall() schema\_blocks = \[\] for table in tables: table\_name = table\[0\] # columns cursor.execute(f"PRAGMA table\_info({table\_name});") columns = \[col\[1\] for col in cursor.fetchall()\] block = f"Table: {table\_name}({', '.join(columns)})" # sample rows cursor.execute(f"SELECT \* FROM {table\_name} LIMIT 10;") rows = cursor.fetchall() if rows: block += "\\nSample rows:\\n" for row in rows: block += f"{row}\\n" schema\_blocks.append(block) conn.close() chunks = \[\] current\_chunk = \[\] for block in schema\_blocks: current\_chunk.append(block) if len(current\_chunk) >= tables\_per\_chunk: chunks.append("\\n".join(current\_chunk)) current\_chunk = \[\] if current\_chunk: chunks.append("\\n".join(current\_chunk)) return chunks # {{/docs-fragment get\_and\_split\_schema}} # {{docs-fragment generate\_questions\_and\_sql}} @flyte.trace async def generate\_questions\_and\_sql( schema: str, num\_samples: int, batch\_size: int ) -> QAList: llm = OpenAI(model="gpt-4.1") prompt\_tmpl = PromptTemplate( """Prompt: You are helping build a Text-to-SQL dataset. Here is the database schema: {schema} Generate {num} natural language questions a user might ask about this database. For each question, also provide the correct SQL query. Reasoning process (you must follow this internally): - Given an input question, first create a syntactically correct {dialect} SQL query. - Never use SELECT \*; only include the relevant columns. - Use only columns/tables from the schema. Qualify column names when ambiguous. - You may order results by a meaningful column to make the query more useful. - Be careful not to add unnecessary columns. - Use filters, aggregations, joins, grouping, and subqueries when relevant. Final Output: Return only a JSON object with one field: - "items": a list of {num} objects, each with: - "question": the natural language question - "sql": the corresponding SQL query """ ) all\_items: list\[QAItem\] = \[\] # batch generation for start in range(0, num\_samples, batch\_size): current\_num = min(batch\_size, num\_samples - start) response = llm.structured\_predict( QAList, prompt\_tmpl, schema=schema, num=current\_num, ) all\_items.extend(response.items) # deduplicate seen = set() unique\_items: list\[QAItem\] = \[\] for item in all\_items: key = (item.question.strip().lower(), item.sql.strip().lower()) if key not in seen: seen.add(key) unique\_items.append(item) return QAList(items=unique\_items\[:num\_samples\]) # {{/docs-fragment generate\_questions\_and\_sql}} @flyte.trace async def llm\_validate\_batch(pairs: list\[dict\[str, str\]\]) -> list\[str\]: """Validate a batch of question/sql/result dicts using one LLM call.""" batch\_prompt = """You are validating the correctness of SQL query results against the question. For each example, answer only "True" (correct) or "False" (incorrect). Output one answer per line, in the same order as the examples. --- """ for i, pair in enumerate(pairs, start=1): batch\_prompt += f""" Example {i}: Question: {pair\['question'\]} SQL: {pair\['sql'\]} Result: {pair\['rows'\]} --- """ llm = OpenAI(model="gpt-4.1") resp = await llm.acomplete(batch\_prompt) # Expect exactly one True/False per example results = \[\ line.strip()\ for line in resp.text.splitlines()\ if line.strip() in ("True", "False")\ \] return results # {{docs-fragment validate\_sql}} @env.task async def validate\_sql( db\_file: File, question\_sql\_pairs: QAList, batch\_size: int ) -> list\[dict\[str, str\]\]: await db\_file.download(local\_path="local\_db.sqlite") conn = sqlite3.connect("local\_db.sqlite") cursor = conn.cursor() qa\_data = \[\] batch = \[\] for pair in question\_sql\_pairs.items: q, sql = pair.question, pair.sql try: cursor.execute(sql) rows = cursor.fetchall() batch.append({"question": q, "sql": sql, "rows": str(rows)}) # process when batch is full if len(batch) == batch\_size: results = await llm\_validate\_batch(batch) for pair, is\_valid in zip(batch, results): if is\_valid == "True": qa\_data.append( { "input": pair\["question"\], "sql": pair\["sql"\], "target": pair\["rows"\], } ) else: print(f"Filtered out incorrect result for: {pair\['question'\]}") batch = \[\] except Exception as e: print(f"Skipping invalid SQL: {sql} ({e})") # process leftover batch if batch: results = await llm\_validate\_batch(batch) for pair, is\_valid in zip(batch, results): if is\_valid == "True": qa\_data.append( { "input": pair\["question"\], "sql": pair\["sql"\], "target": pair\["rows"\], } ) else: print(f"Filtered out incorrect result for: {pair\['question'\]}") conn.close() return qa\_data # {{/docs-fragment validate\_sql}} @flyte.trace async def save\_to\_csv(qa\_data: list\[dict\]) -> File: df = pd.DataFrame(qa\_data, columns=\["input", "target", "sql"\]) csv\_file = "qa\_dataset.csv" df.to\_csv(csv\_file, index=False) return await File.from\_local(csv\_file) # {{docs-fragment build\_eval\_dataset}} @env.task async def build\_eval\_dataset( num\_samples: int = 300, batch\_size: int = 30, tables\_per\_chunk: int = 3 ) -> File: db\_file, \_ = await data\_ingestion() schema\_chunks = await get\_and\_split\_schema(db\_file, tables\_per\_chunk) per\_chunk\_samples = max(1, num\_samples // len(schema\_chunks)) final\_qa\_data = \[\] for chunk in schema\_chunks: qa\_list = await generate\_questions\_and\_sql( schema=chunk, num\_samples=per\_chunk\_samples, batch\_size=batch\_size, ) qa\_data = await validate\_sql(db\_file, qa\_list, batch\_size) final\_qa\_data.extend(qa\_data) csv\_file = await save\_to\_csv(final\_qa\_data) return csv\_file # {{/docs-fragment build\_eval\_dataset}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(build\_eval\_dataset) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/create\_qa\_dataset.py\* ### Question and SQL generation Using structured prompts, we ask an LLM to generate realistic questions users might ask, then pair them with syntactically valid SQL queries. Deduplication ensures diversity across queries. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas>=2.0.0",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # "pydantic>=2.0.0",\ # \] # main = "build\_eval\_dataset" # params = "" # /// import sqlite3 import flyte import pandas as pd from data\_ingestion import data\_ingestion from flyte.io import File from llama\_index.core import PromptTemplate from llama\_index.llms.openai import OpenAI from utils import env from pydantic import BaseModel class QAItem(BaseModel): question: str sql: str class QAList(BaseModel): items: list\[QAItem\] # {{docs-fragment get\_and\_split\_schema}} @env.task async def get\_and\_split\_schema(db\_file: File, tables\_per\_chunk: int) -> list\[str\]: """ Download the SQLite DB, extract schema info (columns + sample rows), then split it into chunks with up to \`tables\_per\_chunk\` tables each. """ await db\_file.download(local\_path="local\_db.sqlite") conn = sqlite3.connect("local\_db.sqlite") cursor = conn.cursor() tables = cursor.execute( "SELECT name FROM sqlite\_master WHERE type='table';" ).fetchall() schema\_blocks = \[\] for table in tables: table\_name = table\[0\] # columns cursor.execute(f"PRAGMA table\_info({table\_name});") columns = \[col\[1\] for col in cursor.fetchall()\] block = f"Table: {table\_name}({', '.join(columns)})" # sample rows cursor.execute(f"SELECT \* FROM {table\_name} LIMIT 10;") rows = cursor.fetchall() if rows: block += "\\nSample rows:\\n" for row in rows: block += f"{row}\\n" schema\_blocks.append(block) conn.close() chunks = \[\] current\_chunk = \[\] for block in schema\_blocks: current\_chunk.append(block) if len(current\_chunk) >= tables\_per\_chunk: chunks.append("\\n".join(current\_chunk)) current\_chunk = \[\] if current\_chunk: chunks.append("\\n".join(current\_chunk)) return chunks # {{/docs-fragment get\_and\_split\_schema}} # {{docs-fragment generate\_questions\_and\_sql}} @flyte.trace async def generate\_questions\_and\_sql( schema: str, num\_samples: int, batch\_size: int ) -> QAList: llm = OpenAI(model="gpt-4.1") prompt\_tmpl = PromptTemplate( """Prompt: You are helping build a Text-to-SQL dataset. Here is the database schema: {schema} Generate {num} natural language questions a user might ask about this database. For each question, also provide the correct SQL query. Reasoning process (you must follow this internally): - Given an input question, first create a syntactically correct {dialect} SQL query. - Never use SELECT \*; only include the relevant columns. - Use only columns/tables from the schema. Qualify column names when ambiguous. - You may order results by a meaningful column to make the query more useful. - Be careful not to add unnecessary columns. - Use filters, aggregations, joins, grouping, and subqueries when relevant. Final Output: Return only a JSON object with one field: - "items": a list of {num} objects, each with: - "question": the natural language question - "sql": the corresponding SQL query """ ) all\_items: list\[QAItem\] = \[\] # batch generation for start in range(0, num\_samples, batch\_size): current\_num = min(batch\_size, num\_samples - start) response = llm.structured\_predict( QAList, prompt\_tmpl, schema=schema, num=current\_num, ) all\_items.extend(response.items) # deduplicate seen = set() unique\_items: list\[QAItem\] = \[\] for item in all\_items: key = (item.question.strip().lower(), item.sql.strip().lower()) if key not in seen: seen.add(key) unique\_items.append(item) return QAList(items=unique\_items\[:num\_samples\]) # {{/docs-fragment generate\_questions\_and\_sql}} @flyte.trace async def llm\_validate\_batch(pairs: list\[dict\[str, str\]\]) -> list\[str\]: """Validate a batch of question/sql/result dicts using one LLM call.""" batch\_prompt = """You are validating the correctness of SQL query results against the question. For each example, answer only "True" (correct) or "False" (incorrect). Output one answer per line, in the same order as the examples. --- """ for i, pair in enumerate(pairs, start=1): batch\_prompt += f""" Example {i}: Question: {pair\['question'\]} SQL: {pair\['sql'\]} Result: {pair\['rows'\]} --- """ llm = OpenAI(model="gpt-4.1") resp = await llm.acomplete(batch\_prompt) # Expect exactly one True/False per example results = \[\ line.strip()\ for line in resp.text.splitlines()\ if line.strip() in ("True", "False")\ \] return results # {{docs-fragment validate\_sql}} @env.task async def validate\_sql( db\_file: File, question\_sql\_pairs: QAList, batch\_size: int ) -> list\[dict\[str, str\]\]: await db\_file.download(local\_path="local\_db.sqlite") conn = sqlite3.connect("local\_db.sqlite") cursor = conn.cursor() qa\_data = \[\] batch = \[\] for pair in question\_sql\_pairs.items: q, sql = pair.question, pair.sql try: cursor.execute(sql) rows = cursor.fetchall() batch.append({"question": q, "sql": sql, "rows": str(rows)}) # process when batch is full if len(batch) == batch\_size: results = await llm\_validate\_batch(batch) for pair, is\_valid in zip(batch, results): if is\_valid == "True": qa\_data.append( { "input": pair\["question"\], "sql": pair\["sql"\], "target": pair\["rows"\], } ) else: print(f"Filtered out incorrect result for: {pair\['question'\]}") batch = \[\] except Exception as e: print(f"Skipping invalid SQL: {sql} ({e})") # process leftover batch if batch: results = await llm\_validate\_batch(batch) for pair, is\_valid in zip(batch, results): if is\_valid == "True": qa\_data.append( { "input": pair\["question"\], "sql": pair\["sql"\], "target": pair\["rows"\], } ) else: print(f"Filtered out incorrect result for: {pair\['question'\]}") conn.close() return qa\_data # {{/docs-fragment validate\_sql}} @flyte.trace async def save\_to\_csv(qa\_data: list\[dict\]) -> File: df = pd.DataFrame(qa\_data, columns=\["input", "target", "sql"\]) csv\_file = "qa\_dataset.csv" df.to\_csv(csv\_file, index=False) return await File.from\_local(csv\_file) # {{docs-fragment build\_eval\_dataset}} @env.task async def build\_eval\_dataset( num\_samples: int = 300, batch\_size: int = 30, tables\_per\_chunk: int = 3 ) -> File: db\_file, \_ = await data\_ingestion() schema\_chunks = await get\_and\_split\_schema(db\_file, tables\_per\_chunk) per\_chunk\_samples = max(1, num\_samples // len(schema\_chunks)) final\_qa\_data = \[\] for chunk in schema\_chunks: qa\_list = await generate\_questions\_and\_sql( schema=chunk, num\_samples=per\_chunk\_samples, batch\_size=batch\_size, ) qa\_data = await validate\_sql(db\_file, qa\_list, batch\_size) final\_qa\_data.extend(qa\_data) csv\_file = await save\_to\_csv(final\_qa\_data) return csv\_file # {{/docs-fragment build\_eval\_dataset}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(build\_eval\_dataset) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/create\_qa\_dataset.py\* ### Validation and quality control Each generated SQL query runs against the database, and another LLM double-checks that the result matches the intent of the natural language question. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas>=2.0.0",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # "pydantic>=2.0.0",\ # \] # main = "build\_eval\_dataset" # params = "" # /// import sqlite3 import flyte import pandas as pd from data\_ingestion import data\_ingestion from flyte.io import File from llama\_index.core import PromptTemplate from llama\_index.llms.openai import OpenAI from utils import env from pydantic import BaseModel class QAItem(BaseModel): question: str sql: str class QAList(BaseModel): items: list\[QAItem\] # {{docs-fragment get\_and\_split\_schema}} @env.task async def get\_and\_split\_schema(db\_file: File, tables\_per\_chunk: int) -> list\[str\]: """ Download the SQLite DB, extract schema info (columns + sample rows), then split it into chunks with up to \`tables\_per\_chunk\` tables each. """ await db\_file.download(local\_path="local\_db.sqlite") conn = sqlite3.connect("local\_db.sqlite") cursor = conn.cursor() tables = cursor.execute( "SELECT name FROM sqlite\_master WHERE type='table';" ).fetchall() schema\_blocks = \[\] for table in tables: table\_name = table\[0\] # columns cursor.execute(f"PRAGMA table\_info({table\_name});") columns = \[col\[1\] for col in cursor.fetchall()\] block = f"Table: {table\_name}({', '.join(columns)})" # sample rows cursor.execute(f"SELECT \* FROM {table\_name} LIMIT 10;") rows = cursor.fetchall() if rows: block += "\\nSample rows:\\n" for row in rows: block += f"{row}\\n" schema\_blocks.append(block) conn.close() chunks = \[\] current\_chunk = \[\] for block in schema\_blocks: current\_chunk.append(block) if len(current\_chunk) >= tables\_per\_chunk: chunks.append("\\n".join(current\_chunk)) current\_chunk = \[\] if current\_chunk: chunks.append("\\n".join(current\_chunk)) return chunks # {{/docs-fragment get\_and\_split\_schema}} # {{docs-fragment generate\_questions\_and\_sql}} @flyte.trace async def generate\_questions\_and\_sql( schema: str, num\_samples: int, batch\_size: int ) -> QAList: llm = OpenAI(model="gpt-4.1") prompt\_tmpl = PromptTemplate( """Prompt: You are helping build a Text-to-SQL dataset. Here is the database schema: {schema} Generate {num} natural language questions a user might ask about this database. For each question, also provide the correct SQL query. Reasoning process (you must follow this internally): - Given an input question, first create a syntactically correct {dialect} SQL query. - Never use SELECT \*; only include the relevant columns. - Use only columns/tables from the schema. Qualify column names when ambiguous. - You may order results by a meaningful column to make the query more useful. - Be careful not to add unnecessary columns. - Use filters, aggregations, joins, grouping, and subqueries when relevant. Final Output: Return only a JSON object with one field: - "items": a list of {num} objects, each with: - "question": the natural language question - "sql": the corresponding SQL query """ ) all\_items: list\[QAItem\] = \[\] # batch generation for start in range(0, num\_samples, batch\_size): current\_num = min(batch\_size, num\_samples - start) response = llm.structured\_predict( QAList, prompt\_tmpl, schema=schema, num=current\_num, ) all\_items.extend(response.items) # deduplicate seen = set() unique\_items: list\[QAItem\] = \[\] for item in all\_items: key = (item.question.strip().lower(), item.sql.strip().lower()) if key not in seen: seen.add(key) unique\_items.append(item) return QAList(items=unique\_items\[:num\_samples\]) # {{/docs-fragment generate\_questions\_and\_sql}} @flyte.trace async def llm\_validate\_batch(pairs: list\[dict\[str, str\]\]) -> list\[str\]: """Validate a batch of question/sql/result dicts using one LLM call.""" batch\_prompt = """You are validating the correctness of SQL query results against the question. For each example, answer only "True" (correct) or "False" (incorrect). Output one answer per line, in the same order as the examples. --- """ for i, pair in enumerate(pairs, start=1): batch\_prompt += f""" Example {i}: Question: {pair\['question'\]} SQL: {pair\['sql'\]} Result: {pair\['rows'\]} --- """ llm = OpenAI(model="gpt-4.1") resp = await llm.acomplete(batch\_prompt) # Expect exactly one True/False per example results = \[\ line.strip()\ for line in resp.text.splitlines()\ if line.strip() in ("True", "False")\ \] return results # {{docs-fragment validate\_sql}} @env.task async def validate\_sql( db\_file: File, question\_sql\_pairs: QAList, batch\_size: int ) -> list\[dict\[str, str\]\]: await db\_file.download(local\_path="local\_db.sqlite") conn = sqlite3.connect("local\_db.sqlite") cursor = conn.cursor() qa\_data = \[\] batch = \[\] for pair in question\_sql\_pairs.items: q, sql = pair.question, pair.sql try: cursor.execute(sql) rows = cursor.fetchall() batch.append({"question": q, "sql": sql, "rows": str(rows)}) # process when batch is full if len(batch) == batch\_size: results = await llm\_validate\_batch(batch) for pair, is\_valid in zip(batch, results): if is\_valid == "True": qa\_data.append( { "input": pair\["question"\], "sql": pair\["sql"\], "target": pair\["rows"\], } ) else: print(f"Filtered out incorrect result for: {pair\['question'\]}") batch = \[\] except Exception as e: print(f"Skipping invalid SQL: {sql} ({e})") # process leftover batch if batch: results = await llm\_validate\_batch(batch) for pair, is\_valid in zip(batch, results): if is\_valid == "True": qa\_data.append( { "input": pair\["question"\], "sql": pair\["sql"\], "target": pair\["rows"\], } ) else: print(f"Filtered out incorrect result for: {pair\['question'\]}") conn.close() return qa\_data # {{/docs-fragment validate\_sql}} @flyte.trace async def save\_to\_csv(qa\_data: list\[dict\]) -> File: df = pd.DataFrame(qa\_data, columns=\["input", "target", "sql"\]) csv\_file = "qa\_dataset.csv" df.to\_csv(csv\_file, index=False) return await File.from\_local(csv\_file) # {{docs-fragment build\_eval\_dataset}} @env.task async def build\_eval\_dataset( num\_samples: int = 300, batch\_size: int = 30, tables\_per\_chunk: int = 3 ) -> File: db\_file, \_ = await data\_ingestion() schema\_chunks = await get\_and\_split\_schema(db\_file, tables\_per\_chunk) per\_chunk\_samples = max(1, num\_samples // len(schema\_chunks)) final\_qa\_data = \[\] for chunk in schema\_chunks: qa\_list = await generate\_questions\_and\_sql( schema=chunk, num\_samples=per\_chunk\_samples, batch\_size=batch\_size, ) qa\_data = await validate\_sql(db\_file, qa\_list, batch\_size) final\_qa\_data.extend(qa\_data) csv\_file = await save\_to\_csv(final\_qa\_data) return csv\_file # {{/docs-fragment build\_eval\_dataset}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(build\_eval\_dataset) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/create\_qa\_dataset.py\* Even with automated checks, human review remains critical. Since this dataset serves as the ground truth, mislabeled pairs can distort evaluation. For production use, always invest in human-in-the-loop review. ## Optimizing prompts With the QA dataset in place, we can turn to prompt optimization. The idea: start from a baseline prompt, generate new variants, and measure whether accuracy improves. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas>=2.0.0",\ # "sqlalchemy>=2.0.0",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from data\_ingestion import TableInfo from flyte.io import Dir, File from llama\_index.core import SQLDatabase from llama\_index.core.retrievers import SQLRetriever from sqlalchemy import create\_engine from text\_to\_sql import data\_ingestion, generate\_sql, index\_all\_tables, retrieve\_tables from utils import env CSS = """ """ @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into val/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Val/Test split df\_renamed = df.rename(columns={"input": "question", "target": "answer"}) n = len(df\_renamed) split = n // 2 df\_val = df\_renamed.iloc\[:split\] df\_test = df\_renamed.iloc\[split:\] return df\_val, df\_test @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] @flyte.trace async def generate\_response(db\_file: File, sql: str) -> str: await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) sql\_retriever = SQLRetriever(sql\_database) retrieved\_rows = sql\_retriever.retrieve(sql) if retrieved\_rows: # Get the structured result and stringify return str(retrieved\_rows\[0\].node.metadata\["result"\]) return "" async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, db\_file: File, table\_infos: list\[TableInfo | None\], vector\_index\_dir: Dir, ) -> dict: # Generate response from target model table\_context = await retrieve\_tables( question, table\_infos, db\_file, vector\_index\_dir ) sql = await generate\_sql( question, table\_context, target\_model\_config.model\_name, target\_model\_config.prompt, ) sql = sql.replace("sql\\n", "") try: response = await generate\_response(db\_file, sql) except Exception as e: print(f"Failed to generate response for question {question}: {e}") response = None # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ query\_str=question,\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "sql": sql, "is\_correct": verdict\_clean == "true", } async def run\_grouped\_task( i, index, question, answer, sql, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, db\_file, table\_infos, vector\_index\_dir, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, db\_file, table\_infos, vector\_index\_dir, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {html.escape(sql)} {result\['model\_response'\]} {result\['sql'\]} {correct\_html} """, do\_flush=True, ) return result @dataclass class DatabaseConfig: csv\_zip\_path: str search\_glob: str concurrency: int model: str # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, db\_config: DatabaseConfig, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) db\_file, table\_infos = await data\_ingestion( db\_config.csv\_zip\_path, db\_config.search\_glob, db\_config.concurrency, db\_config.model, ) vector\_index\_dir = await index\_all\_tables(db\_file) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ row.sql,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ db\_file,\ table\_infos,\ vector\_index\_dir,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Ground Truth Answer Ground Truth SQL Model Response Model SQL Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_val: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, db\_config: DatabaseConfig, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_val, target\_model\_config, review\_model\_config, concurrency, db\_config, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_val, target\_model\_config, review\_model\_config, concurrency, db\_config, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( ground\_truth\_csv: File | str = "/root/ground\_truth.csv", db\_config: DatabaseConfig = DatabaseConfig( csv\_zip\_path="https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip", search\_glob="WikiTableQuestions/csv/200-csv/\*.csv", concurrency=5, model="gpt-4o-mini", ), target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""Given an input question, create a syntactically correct {dialect} query to run. Schema: {schema} Question: {query\_str} SQL query to run: """, max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, prompt="""Your job is to determine whether the model's response is correct compared to the ground truth taking into account the context of the question. Both answers were generated by running SQL queries on the same database. - If the model's response contains all of the ground truth values, and any additional information is harmless (e.g., extra columns or metadata), output "True". - If it adds incorrect or unrelated rows, or omits required values, output "False". Question: {query\_str} Ground Truth: {answer} Model Response: {response} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicates better quality. {prompt\_scores\_str} Each prompt was used to translate a natural-language question into a SQL query against a provided database schema. artists(id, name) albums(id, title, artist\_id, release\_year) How many albums did The Beatles release? SELECT COUNT(\*) FROM albums a JOIN artists r ON a.artist\_id = r.id WHERE r.name = 'The Beatles'; Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past. - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't work in the past. - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc. for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating a system prompt. Always use three placeholders for each prompt: dialect, schema, query\_str. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 5, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(ground\_truth\_csv, str) and os.path.isfile(ground\_truth\_csv): ground\_truth\_csv = await File.from\_local(ground\_truth\_csv) df\_val, df\_test = await data\_prep(ground\_truth\_csv) best\_prompt, val\_accuracy = await prompt\_optimizer( df\_val, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, db\_config, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, db\_config, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, db\_config, ) return { "best\_prompt": best\_prompt, "validation\_accuracy": val\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/optimizer.py\* ### Evaluation pipeline We evaluate each prompt variant against the golden dataset, split into validation and test sets, and record accuracy metrics in real time. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas>=2.0.0",\ # "sqlalchemy>=2.0.0",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from data\_ingestion import TableInfo from flyte.io import Dir, File from llama\_index.core import SQLDatabase from llama\_index.core.retrievers import SQLRetriever from sqlalchemy import create\_engine from text\_to\_sql import data\_ingestion, generate\_sql, index\_all\_tables, retrieve\_tables from utils import env CSS = """ """ @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into val/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Val/Test split df\_renamed = df.rename(columns={"input": "question", "target": "answer"}) n = len(df\_renamed) split = n // 2 df\_val = df\_renamed.iloc\[:split\] df\_test = df\_renamed.iloc\[split:\] return df\_val, df\_test @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] @flyte.trace async def generate\_response(db\_file: File, sql: str) -> str: await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) sql\_retriever = SQLRetriever(sql\_database) retrieved\_rows = sql\_retriever.retrieve(sql) if retrieved\_rows: # Get the structured result and stringify return str(retrieved\_rows\[0\].node.metadata\["result"\]) return "" async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, db\_file: File, table\_infos: list\[TableInfo | None\], vector\_index\_dir: Dir, ) -> dict: # Generate response from target model table\_context = await retrieve\_tables( question, table\_infos, db\_file, vector\_index\_dir ) sql = await generate\_sql( question, table\_context, target\_model\_config.model\_name, target\_model\_config.prompt, ) sql = sql.replace("sql\\n", "") try: response = await generate\_response(db\_file, sql) except Exception as e: print(f"Failed to generate response for question {question}: {e}") response = None # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ query\_str=question,\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "sql": sql, "is\_correct": verdict\_clean == "true", } async def run\_grouped\_task( i, index, question, answer, sql, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, db\_file, table\_infos, vector\_index\_dir, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, db\_file, table\_infos, vector\_index\_dir, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {html.escape(sql)} {result\['model\_response'\]} {result\['sql'\]} {correct\_html} """, do\_flush=True, ) return result @dataclass class DatabaseConfig: csv\_zip\_path: str search\_glob: str concurrency: int model: str # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, db\_config: DatabaseConfig, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) db\_file, table\_infos = await data\_ingestion( db\_config.csv\_zip\_path, db\_config.search\_glob, db\_config.concurrency, db\_config.model, ) vector\_index\_dir = await index\_all\_tables(db\_file) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ row.sql,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ db\_file,\ table\_infos,\ vector\_index\_dir,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Ground Truth Answer Ground Truth SQL Model Response Model SQL Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_val: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, db\_config: DatabaseConfig, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_val, target\_model\_config, review\_model\_config, concurrency, db\_config, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_val, target\_model\_config, review\_model\_config, concurrency, db\_config, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( ground\_truth\_csv: File | str = "/root/ground\_truth.csv", db\_config: DatabaseConfig = DatabaseConfig( csv\_zip\_path="https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip", search\_glob="WikiTableQuestions/csv/200-csv/\*.csv", concurrency=5, model="gpt-4o-mini", ), target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""Given an input question, create a syntactically correct {dialect} query to run. Schema: {schema} Question: {query\_str} SQL query to run: """, max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, prompt="""Your job is to determine whether the model's response is correct compared to the ground truth taking into account the context of the question. Both answers were generated by running SQL queries on the same database. - If the model's response contains all of the ground truth values, and any additional information is harmless (e.g., extra columns or metadata), output "True". - If it adds incorrect or unrelated rows, or omits required values, output "False". Question: {query\_str} Ground Truth: {answer} Model Response: {response} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicates better quality. {prompt\_scores\_str} Each prompt was used to translate a natural-language question into a SQL query against a provided database schema. artists(id, name) albums(id, title, artist\_id, release\_year) How many albums did The Beatles release? SELECT COUNT(\*) FROM albums a JOIN artists r ON a.artist\_id = r.id WHERE r.name = 'The Beatles'; Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past. - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't work in the past. - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc. for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating a system prompt. Always use three placeholders for each prompt: dialect, schema, query\_str. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 5, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(ground\_truth\_csv, str) and os.path.isfile(ground\_truth\_csv): ground\_truth\_csv = await File.from\_local(ground\_truth\_csv) df\_val, df\_test = await data\_prep(ground\_truth\_csv) best\_prompt, val\_accuracy = await prompt\_optimizer( df\_val, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, db\_config, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, db\_config, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, db\_config, ) return { "best\_prompt": best\_prompt, "validation\_accuracy": val\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/optimizer.py\* Here's how prompt accuracy evolves over time, as shown in the UI report: !\[Prompt accuracies\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/tutorials/text-to-sql/prompt\_accuracies.png) ### Iterative optimization An optimizer LLM proposes new prompts by analyzing patterns in successful and failed generations. Each candidate runs through the evaluation loop, and we select the best performer. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas>=2.0.0",\ # "sqlalchemy>=2.0.0",\ # "llama-index-core>=0.11.0",\ # "llama-index-llms-openai>=0.2.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from data\_ingestion import TableInfo from flyte.io import Dir, File from llama\_index.core import SQLDatabase from llama\_index.core.retrievers import SQLRetriever from sqlalchemy import create\_engine from text\_to\_sql import data\_ingestion, generate\_sql, index\_all\_tables, retrieve\_tables from utils import env CSS = """ """ @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into val/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Val/Test split df\_renamed = df.rename(columns={"input": "question", "target": "answer"}) n = len(df\_renamed) split = n // 2 df\_val = df\_renamed.iloc\[:split\] df\_test = df\_renamed.iloc\[split:\] return df\_val, df\_test @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] @flyte.trace async def generate\_response(db\_file: File, sql: str) -> str: await db\_file.download(local\_path="local\_db.sqlite") engine = create\_engine("sqlite:///local\_db.sqlite") sql\_database = SQLDatabase(engine) sql\_retriever = SQLRetriever(sql\_database) retrieved\_rows = sql\_retriever.retrieve(sql) if retrieved\_rows: # Get the structured result and stringify return str(retrieved\_rows\[0\].node.metadata\["result"\]) return "" async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, db\_file: File, table\_infos: list\[TableInfo | None\], vector\_index\_dir: Dir, ) -> dict: # Generate response from target model table\_context = await retrieve\_tables( question, table\_infos, db\_file, vector\_index\_dir ) sql = await generate\_sql( question, table\_context, target\_model\_config.model\_name, target\_model\_config.prompt, ) sql = sql.replace("sql\\n", "") try: response = await generate\_response(db\_file, sql) except Exception as e: print(f"Failed to generate response for question {question}: {e}") response = None # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ query\_str=question,\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "sql": sql, "is\_correct": verdict\_clean == "true", } async def run\_grouped\_task( i, index, question, answer, sql, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, db\_file, table\_infos, vector\_index\_dir, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, db\_file, table\_infos, vector\_index\_dir, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {html.escape(sql)} {result\['model\_response'\]} {result\['sql'\]} {correct\_html} """, do\_flush=True, ) return result @dataclass class DatabaseConfig: csv\_zip\_path: str search\_glob: str concurrency: int model: str # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, db\_config: DatabaseConfig, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) db\_file, table\_infos = await data\_ingestion( db\_config.csv\_zip\_path, db\_config.search\_glob, db\_config.concurrency, db\_config.model, ) vector\_index\_dir = await index\_all\_tables(db\_file) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ row.sql,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ db\_file,\ table\_infos,\ vector\_index\_dir,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Ground Truth Answer Ground Truth SQL Model Response Model SQL Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_val: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, db\_config: DatabaseConfig, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_val, target\_model\_config, review\_model\_config, concurrency, db\_config, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_val, target\_model\_config, review\_model\_config, concurrency, db\_config, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( ground\_truth\_csv: File | str = "/root/ground\_truth.csv", db\_config: DatabaseConfig = DatabaseConfig( csv\_zip\_path="https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip", search\_glob="WikiTableQuestions/csv/200-csv/\*.csv", concurrency=5, model="gpt-4o-mini", ), target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""Given an input question, create a syntactically correct {dialect} query to run. Schema: {schema} Question: {query\_str} SQL query to run: """, max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, prompt="""Your job is to determine whether the model's response is correct compared to the ground truth taking into account the context of the question. Both answers were generated by running SQL queries on the same database. - If the model's response contains all of the ground truth values, and any additional information is harmless (e.g., extra columns or metadata), output "True". - If it adds incorrect or unrelated rows, or omits required values, output "False". Question: {query\_str} Ground Truth: {answer} Model Response: {response} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicates better quality. {prompt\_scores\_str} Each prompt was used to translate a natural-language question into a SQL query against a provided database schema. artists(id, name) albums(id, title, artist\_id, release\_year) How many albums did The Beatles release? SELECT COUNT(\*) FROM albums a JOIN artists r ON a.artist\_id = r.id WHERE r.name = 'The Beatles'; Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past. - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't work in the past. - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc. for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating a system prompt. Always use three placeholders for each prompt: dialect, schema, query\_str. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 5, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(ground\_truth\_csv, str) and os.path.isfile(ground\_truth\_csv): ground\_truth\_csv = await File.from\_local(ground\_truth\_csv) df\_val, df\_test = await data\_prep(ground\_truth\_csv) best\_prompt, val\_accuracy = await prompt\_optimizer( df\_val, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, db\_config, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, db\_config, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, db\_config, ) return { "best\_prompt": best\_prompt, "validation\_accuracy": val\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/text\_to\_sql/optimizer.py\* On paper, this creates a continuous improvement cycle: baseline → new variants → measured gains. ## Run it To create the QA dataset: \`\`\` python create\_qa\_dataset.py \`\`\` To run the prompt optimization loop: \`\`\` python optimizer.py \`\`\` ## What we observed Prompt optimization didn't consistently lift SQL accuracy in this workflow. Accuracy plateaued near the baseline. But the process surfaced valuable lessons about what matters when building LLM-powered systems on real infrastructure. - \*\*Schema clarity matters\*\*: CSV ingestion produced tables with overlapping names, creating ambiguity. This showed how schema design and metadata hygiene directly affect downstream evaluation. - \*\*Ground truth needs trust\*\*: Because the dataset came from LLM outputs, noise remained even after filtering. Human review proved essential. Golden datasets need deliberate curation, not just automation. - \*\*Optimization needs context\*\*: The optimizer couldn't “see” which examples failed, limiting its ability to improve. Feeding failures directly risks overfitting. A structured way to capture and reuse evaluation signals is the right long-term path. Sometimes prompt tweaks alone can lift accuracy, but other times the real bottleneck lives in the data, the schema, or the evaluation loop. The lesson isn't "prompt optimization doesn't work", but that its impact depends on the system around it. Accuracy improves most reliably when prompts evolve alongside clean data, trusted evaluation, and observable feedback loops. ## The bigger lesson Evaluation and optimization aren’t one-off experiments; they’re continuous processes. What makes them sustainable isn't a clever prompt, it’s the platform around it. Systems succeed when they: - \*\*Observe\*\* failures with clarity — track exactly what failed and why. - \*\*Remain durable\*\* across iterations — run pipelines that are stable, reproducible, and comparable over time. That's where Flyte 2 comes in. Prompt optimization is one lever, but it becomes powerful only when combined with: - Clean, human-validated evaluation datasets. - Systematic reporting and feedback loops. \*\*The real takeaway: improving LLM pipelines isn't about chasing the perfect prompt. It's about designing workflows with observability and durability at the core, so that every experiment compounds into long-term progress.\*\* === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/context-engineering/auto\_prompt\_engineering === # Automatic prompt engineering > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/auto\_prompt\_engineering). When building with LLMs and agents, the first prompt almost never works. We usually need several iterations before results are useful. Doing this manually is slow, inconsistent, and hard to reproduce. Flyte turns prompt engineering into a systematic process. With Flyte we can: - Generate candidate prompts automatically. - Run evaluations in parallel. - Track results in real time with built-in observability. - Recover from failures without losing progress. - Trace the lineage of every experiment for reproducibility. And we're not limited to prompts. Just like \[hyperparameter optimization\](../../model-training/hpo/\_index) in ML, we can tune model temperature, retrieval strategies, tool usage, and more. Over time, this grows into full agentic evaluations, tracking not only prompts but also how agents behave, make decisions, and interact with their environment. In this tutorial, we'll build an automated prompt engineering pipeline with Flyte, step by step. ## Set up the environment First, let's configure our task environment. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas==2.3.1",\ # "pyarrow==21.0.0",\ # "litellm==1.75.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// # {{docs-fragment env}} import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from flyte.io.\_file import File env = flyte.TaskEnvironment( name="auto-prompt-engineering", image=flyte.Image.from\_uv\_script( \_\_file\_\_, name="auto-prompt-engineering", pre=True ), secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], resources=flyte.Resources(cpu=1), ) CSS = """ """ # {{/docs-fragment env}} # {{docs-fragment data\_prep}} @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Train/Test split df\_train = df.iloc\[:150\].rename(columns={"input": "question", "target": "answer"}) df\_test = df.iloc\[150:250\].rename(columns={"input": "question", "target": "answer"}) return df\_train, df\_test # {{/docs-fragment data\_prep}} # {{docs-fragment model\_config}} @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" # {{/docs-fragment model\_config}} # {{docs-fragment call\_model}} @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] # {{/docs-fragment call\_model}} # {{docs-fragment generate\_and\_review}} async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, ) -> dict: # Generate response from target model response = await call\_model( target\_model\_config, \[\ {"role": "system", "content": target\_model\_config.prompt},\ {"role": "user", "content": question},\ \], ) # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "is\_correct": verdict\_clean == "true", } # {{/docs-fragment generate\_and\_review}} async def run\_grouped\_task( i, index, question, answer, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {result\['model\_response'\]} {correct\_html} """, do\_flush=True, ) return result # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Answer Model Response Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_train: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( csv\_file: File | str = "https://dub.sh/geometric-shapes", target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="Solve the given problem about geometric shapes. Think step by step.", max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem. The response may contain detailed steps and explanations, but the final answer is the key point. Please determine if the final answer provided in the response is correct based on the ground truth number. Respond with 'True' if the final answer is correct and 'False' if it is not. Only respond with 'True' or 'False', nothing else. Model Response: {response} Ground Truth: {answer} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality. {prompt\_scores\_str} Each prompt was used together with a problem statement around geometric shapes. This SVG path element draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle (B) Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 3, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(csv\_file, str) and os.path.isfile(csv\_file): csv\_file = await File.from\_local(csv\_file) df\_train, df\_test = await data\_prep(csv\_file) best\_prompt, training\_accuracy = await prompt\_optimizer( df\_train, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) return { "best\_prompt": best\_prompt, "training\_accuracy": training\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto\_prompt\_engineering/optimizer.py\* We need an API key to call GPT-4.1 (our optimization model). Add it as a Flyte secret: \`\`\` flyte create secret openai\_api\_key \`\`\` We also define CSS styles for live HTML reports that track prompt optimization in real time: !\[Results\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/gifs/tutorials/prompt\_engineering/results.gif) ## Prepare the evaluation dataset Next, we define our golden dataset, a set of prompts with known outputs. This dataset is used to evaluate the quality of generated prompts. For this tutorial, we use a small geometric shapes dataset. To keep it portable, the data prep task takes a CSV file (as a Flyte \`File\` or a string for files available remotely) and splits it into train and test subsets. If you already have prompts and outputs in Google Sheets, simply export them as CSV with two columns: \`input\` and \`target\`. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas==2.3.1",\ # "pyarrow==21.0.0",\ # "litellm==1.75.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// # {{docs-fragment env}} import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from flyte.io.\_file import File env = flyte.TaskEnvironment( name="auto-prompt-engineering", image=flyte.Image.from\_uv\_script( \_\_file\_\_, name="auto-prompt-engineering", pre=True ), secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], resources=flyte.Resources(cpu=1), ) CSS = """ """ # {{/docs-fragment env}} # {{docs-fragment data\_prep}} @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Train/Test split df\_train = df.iloc\[:150\].rename(columns={"input": "question", "target": "answer"}) df\_test = df.iloc\[150:250\].rename(columns={"input": "question", "target": "answer"}) return df\_train, df\_test # {{/docs-fragment data\_prep}} # {{docs-fragment model\_config}} @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" # {{/docs-fragment model\_config}} # {{docs-fragment call\_model}} @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] # {{/docs-fragment call\_model}} # {{docs-fragment generate\_and\_review}} async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, ) -> dict: # Generate response from target model response = await call\_model( target\_model\_config, \[\ {"role": "system", "content": target\_model\_config.prompt},\ {"role": "user", "content": question},\ \], ) # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "is\_correct": verdict\_clean == "true", } # {{/docs-fragment generate\_and\_review}} async def run\_grouped\_task( i, index, question, answer, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {result\['model\_response'\]} {correct\_html} """, do\_flush=True, ) return result # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Answer Model Response Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_train: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( csv\_file: File | str = "https://dub.sh/geometric-shapes", target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="Solve the given problem about geometric shapes. Think step by step.", max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem. The response may contain detailed steps and explanations, but the final answer is the key point. Please determine if the final answer provided in the response is correct based on the ground truth number. Respond with 'True' if the final answer is correct and 'False' if it is not. Only respond with 'True' or 'False', nothing else. Model Response: {response} Ground Truth: {answer} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality. {prompt\_scores\_str} Each prompt was used together with a problem statement around geometric shapes. This SVG path element draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle (B) Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 3, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(csv\_file, str) and os.path.isfile(csv\_file): csv\_file = await File.from\_local(csv\_file) df\_train, df\_test = await data\_prep(csv\_file) best\_prompt, training\_accuracy = await prompt\_optimizer( df\_train, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) return { "best\_prompt": best\_prompt, "training\_accuracy": training\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto\_prompt\_engineering/optimizer.py\* This approach works with any dataset. You can swap in your own with no extra dependencies. ## Define models We use two models: - \*\*Target model\*\* → the one we want to optimize. - \*\*Review model\*\* → the one that evaluates candidate prompts. First, we capture all model parameters in a dataclass: \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas==2.3.1",\ # "pyarrow==21.0.0",\ # "litellm==1.75.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// # {{docs-fragment env}} import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from flyte.io.\_file import File env = flyte.TaskEnvironment( name="auto-prompt-engineering", image=flyte.Image.from\_uv\_script( \_\_file\_\_, name="auto-prompt-engineering", pre=True ), secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], resources=flyte.Resources(cpu=1), ) CSS = """ """ # {{/docs-fragment env}} # {{docs-fragment data\_prep}} @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Train/Test split df\_train = df.iloc\[:150\].rename(columns={"input": "question", "target": "answer"}) df\_test = df.iloc\[150:250\].rename(columns={"input": "question", "target": "answer"}) return df\_train, df\_test # {{/docs-fragment data\_prep}} # {{docs-fragment model\_config}} @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" # {{/docs-fragment model\_config}} # {{docs-fragment call\_model}} @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] # {{/docs-fragment call\_model}} # {{docs-fragment generate\_and\_review}} async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, ) -> dict: # Generate response from target model response = await call\_model( target\_model\_config, \[\ {"role": "system", "content": target\_model\_config.prompt},\ {"role": "user", "content": question},\ \], ) # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "is\_correct": verdict\_clean == "true", } # {{/docs-fragment generate\_and\_review}} async def run\_grouped\_task( i, index, question, answer, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {result\['model\_response'\]} {correct\_html} """, do\_flush=True, ) return result # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Answer Model Response Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_train: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( csv\_file: File | str = "https://dub.sh/geometric-shapes", target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="Solve the given problem about geometric shapes. Think step by step.", max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem. The response may contain detailed steps and explanations, but the final answer is the key point. Please determine if the final answer provided in the response is correct based on the ground truth number. Respond with 'True' if the final answer is correct and 'False' if it is not. Only respond with 'True' or 'False', nothing else. Model Response: {response} Ground Truth: {answer} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality. {prompt\_scores\_str} Each prompt was used together with a problem statement around geometric shapes. This SVG path element draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle (B) Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 3, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(csv\_file, str) and os.path.isfile(csv\_file): csv\_file = await File.from\_local(csv\_file) df\_train, df\_test = await data\_prep(csv\_file) best\_prompt, training\_accuracy = await prompt\_optimizer( df\_train, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) return { "best\_prompt": best\_prompt, "training\_accuracy": training\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto\_prompt\_engineering/optimizer.py\* Then we define a Flyte \`trace\` to call the model. Unlike a task, a trace runs within the same runtime as the parent process. Since the model is hosted externally, this keeps the call lightweight but still observable. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas==2.3.1",\ # "pyarrow==21.0.0",\ # "litellm==1.75.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// # {{docs-fragment env}} import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from flyte.io.\_file import File env = flyte.TaskEnvironment( name="auto-prompt-engineering", image=flyte.Image.from\_uv\_script( \_\_file\_\_, name="auto-prompt-engineering", pre=True ), secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], resources=flyte.Resources(cpu=1), ) CSS = """ """ # {{/docs-fragment env}} # {{docs-fragment data\_prep}} @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Train/Test split df\_train = df.iloc\[:150\].rename(columns={"input": "question", "target": "answer"}) df\_test = df.iloc\[150:250\].rename(columns={"input": "question", "target": "answer"}) return df\_train, df\_test # {{/docs-fragment data\_prep}} # {{docs-fragment model\_config}} @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" # {{/docs-fragment model\_config}} # {{docs-fragment call\_model}} @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] # {{/docs-fragment call\_model}} # {{docs-fragment generate\_and\_review}} async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, ) -> dict: # Generate response from target model response = await call\_model( target\_model\_config, \[\ {"role": "system", "content": target\_model\_config.prompt},\ {"role": "user", "content": question},\ \], ) # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "is\_correct": verdict\_clean == "true", } # {{/docs-fragment generate\_and\_review}} async def run\_grouped\_task( i, index, question, answer, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {result\['model\_response'\]} {correct\_html} """, do\_flush=True, ) return result # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Answer Model Response Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_train: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( csv\_file: File | str = "https://dub.sh/geometric-shapes", target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="Solve the given problem about geometric shapes. Think step by step.", max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem. The response may contain detailed steps and explanations, but the final answer is the key point. Please determine if the final answer provided in the response is correct based on the ground truth number. Respond with 'True' if the final answer is correct and 'False' if it is not. Only respond with 'True' or 'False', nothing else. Model Response: {response} Ground Truth: {answer} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality. {prompt\_scores\_str} Each prompt was used together with a problem statement around geometric shapes. This SVG path element draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle (B) Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 3, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(csv\_file, str) and os.path.isfile(csv\_file): csv\_file = await File.from\_local(csv\_file) df\_train, df\_test = await data\_prep(csv\_file) best\_prompt, training\_accuracy = await prompt\_optimizer( df\_train, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) return { "best\_prompt": best\_prompt, "training\_accuracy": training\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto\_prompt\_engineering/optimizer.py\* Finally, we wrap the trace in a task to call both target and review models: \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas==2.3.1",\ # "pyarrow==21.0.0",\ # "litellm==1.75.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// # {{docs-fragment env}} import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from flyte.io.\_file import File env = flyte.TaskEnvironment( name="auto-prompt-engineering", image=flyte.Image.from\_uv\_script( \_\_file\_\_, name="auto-prompt-engineering", pre=True ), secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], resources=flyte.Resources(cpu=1), ) CSS = """ """ # {{/docs-fragment env}} # {{docs-fragment data\_prep}} @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Train/Test split df\_train = df.iloc\[:150\].rename(columns={"input": "question", "target": "answer"}) df\_test = df.iloc\[150:250\].rename(columns={"input": "question", "target": "answer"}) return df\_train, df\_test # {{/docs-fragment data\_prep}} # {{docs-fragment model\_config}} @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" # {{/docs-fragment model\_config}} # {{docs-fragment call\_model}} @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] # {{/docs-fragment call\_model}} # {{docs-fragment generate\_and\_review}} async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, ) -> dict: # Generate response from target model response = await call\_model( target\_model\_config, \[\ {"role": "system", "content": target\_model\_config.prompt},\ {"role": "user", "content": question},\ \], ) # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "is\_correct": verdict\_clean == "true", } # {{/docs-fragment generate\_and\_review}} async def run\_grouped\_task( i, index, question, answer, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {result\['model\_response'\]} {correct\_html} """, do\_flush=True, ) return result # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Answer Model Response Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_train: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( csv\_file: File | str = "https://dub.sh/geometric-shapes", target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="Solve the given problem about geometric shapes. Think step by step.", max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem. The response may contain detailed steps and explanations, but the final answer is the key point. Please determine if the final answer provided in the response is correct based on the ground truth number. Respond with 'True' if the final answer is correct and 'False' if it is not. Only respond with 'True' or 'False', nothing else. Model Response: {response} Ground Truth: {answer} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality. {prompt\_scores\_str} Each prompt was used together with a problem statement around geometric shapes. This SVG path element draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle (B) Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 3, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(csv\_file, str) and os.path.isfile(csv\_file): csv\_file = await File.from\_local(csv\_file) df\_train, df\_test = await data\_prep(csv\_file) best\_prompt, training\_accuracy = await prompt\_optimizer( df\_train, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) return { "best\_prompt": best\_prompt, "training\_accuracy": training\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto\_prompt\_engineering/optimizer.py\* ## Evaluate prompts We now define the evaluation process. Each prompt in the dataset is tested in parallel, but we use a semaphore to control concurrency. A helper function ties together the \`generate\_and\_review\` task with an HTML report template. Using \`asyncio.gather\`, we evaluate multiple prompts at once. The function measures accuracy as the fraction of responses that match the ground truth. Flyte streams these results to the UI, so you can watch evaluations happen live. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas==2.3.1",\ # "pyarrow==21.0.0",\ # "litellm==1.75.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// # {{docs-fragment env}} import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from flyte.io.\_file import File env = flyte.TaskEnvironment( name="auto-prompt-engineering", image=flyte.Image.from\_uv\_script( \_\_file\_\_, name="auto-prompt-engineering", pre=True ), secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], resources=flyte.Resources(cpu=1), ) CSS = """ """ # {{/docs-fragment env}} # {{docs-fragment data\_prep}} @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Train/Test split df\_train = df.iloc\[:150\].rename(columns={"input": "question", "target": "answer"}) df\_test = df.iloc\[150:250\].rename(columns={"input": "question", "target": "answer"}) return df\_train, df\_test # {{/docs-fragment data\_prep}} # {{docs-fragment model\_config}} @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" # {{/docs-fragment model\_config}} # {{docs-fragment call\_model}} @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] # {{/docs-fragment call\_model}} # {{docs-fragment generate\_and\_review}} async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, ) -> dict: # Generate response from target model response = await call\_model( target\_model\_config, \[\ {"role": "system", "content": target\_model\_config.prompt},\ {"role": "user", "content": question},\ \], ) # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "is\_correct": verdict\_clean == "true", } # {{/docs-fragment generate\_and\_review}} async def run\_grouped\_task( i, index, question, answer, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {result\['model\_response'\]} {correct\_html} """, do\_flush=True, ) return result # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Answer Model Response Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_train: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( csv\_file: File | str = "https://dub.sh/geometric-shapes", target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="Solve the given problem about geometric shapes. Think step by step.", max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem. The response may contain detailed steps and explanations, but the final answer is the key point. Please determine if the final answer provided in the response is correct based on the ground truth number. Respond with 'True' if the final answer is correct and 'False' if it is not. Only respond with 'True' or 'False', nothing else. Model Response: {response} Ground Truth: {answer} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality. {prompt\_scores\_str} Each prompt was used together with a problem statement around geometric shapes. This SVG path element draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle (B) Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 3, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(csv\_file, str) and os.path.isfile(csv\_file): csv\_file = await File.from\_local(csv\_file) df\_train, df\_test = await data\_prep(csv\_file) best\_prompt, training\_accuracy = await prompt\_optimizer( df\_train, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) return { "best\_prompt": best\_prompt, "training\_accuracy": training\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto\_prompt\_engineering/optimizer.py\* ## Optimize prompts Optimization builds on evaluation. We give the optimizer model: - the history of prompts tested so far, and - their accuracies. The model then proposes a new prompt. We start with a \_baseline\_ evaluation using the user-provided prompt. Then for each iteration, the optimizer suggests a new prompt, which we evaluate and log. We continue until we hit the iteration limit. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas==2.3.1",\ # "pyarrow==21.0.0",\ # "litellm==1.75.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// # {{docs-fragment env}} import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from flyte.io.\_file import File env = flyte.TaskEnvironment( name="auto-prompt-engineering", image=flyte.Image.from\_uv\_script( \_\_file\_\_, name="auto-prompt-engineering", pre=True ), secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], resources=flyte.Resources(cpu=1), ) CSS = """ """ # {{/docs-fragment env}} # {{docs-fragment data\_prep}} @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Train/Test split df\_train = df.iloc\[:150\].rename(columns={"input": "question", "target": "answer"}) df\_test = df.iloc\[150:250\].rename(columns={"input": "question", "target": "answer"}) return df\_train, df\_test # {{/docs-fragment data\_prep}} # {{docs-fragment model\_config}} @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" # {{/docs-fragment model\_config}} # {{docs-fragment call\_model}} @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] # {{/docs-fragment call\_model}} # {{docs-fragment generate\_and\_review}} async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, ) -> dict: # Generate response from target model response = await call\_model( target\_model\_config, \[\ {"role": "system", "content": target\_model\_config.prompt},\ {"role": "user", "content": question},\ \], ) # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "is\_correct": verdict\_clean == "true", } # {{/docs-fragment generate\_and\_review}} async def run\_grouped\_task( i, index, question, answer, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {result\['model\_response'\]} {correct\_html} """, do\_flush=True, ) return result # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Answer Model Response Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_train: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( csv\_file: File | str = "https://dub.sh/geometric-shapes", target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="Solve the given problem about geometric shapes. Think step by step.", max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem. The response may contain detailed steps and explanations, but the final answer is the key point. Please determine if the final answer provided in the response is correct based on the ground truth number. Respond with 'True' if the final answer is correct and 'False' if it is not. Only respond with 'True' or 'False', nothing else. Model Response: {response} Ground Truth: {answer} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality. {prompt\_scores\_str} Each prompt was used together with a problem statement around geometric shapes. This SVG path element draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle (B) Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 3, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(csv\_file, str) and os.path.isfile(csv\_file): csv\_file = await File.from\_local(csv\_file) df\_train, df\_test = await data\_prep(csv\_file) best\_prompt, training\_accuracy = await prompt\_optimizer( df\_train, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) return { "best\_prompt": best\_prompt, "training\_accuracy": training\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto\_prompt\_engineering/optimizer.py\* At the end, we return the best prompt and its accuracy. The report shows how accuracy improves over time and which prompts were tested. !\[Report\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/gifs/tutorials/prompt\_engineering/prompt\_accuracies.png) ## Build the full pipeline The entrypoint task wires everything together: - Accepts model configs, dataset, iteration count, and concurrency. - Runs data preparation. - Calls the optimizer. - Evaluates both baseline and best prompts on the test set. \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas==2.3.1",\ # "pyarrow==21.0.0",\ # "litellm==1.75.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// # {{docs-fragment env}} import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from flyte.io.\_file import File env = flyte.TaskEnvironment( name="auto-prompt-engineering", image=flyte.Image.from\_uv\_script( \_\_file\_\_, name="auto-prompt-engineering", pre=True ), secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], resources=flyte.Resources(cpu=1), ) CSS = """ """ # {{/docs-fragment env}} # {{docs-fragment data\_prep}} @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Train/Test split df\_train = df.iloc\[:150\].rename(columns={"input": "question", "target": "answer"}) df\_test = df.iloc\[150:250\].rename(columns={"input": "question", "target": "answer"}) return df\_train, df\_test # {{/docs-fragment data\_prep}} # {{docs-fragment model\_config}} @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" # {{/docs-fragment model\_config}} # {{docs-fragment call\_model}} @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] # {{/docs-fragment call\_model}} # {{docs-fragment generate\_and\_review}} async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, ) -> dict: # Generate response from target model response = await call\_model( target\_model\_config, \[\ {"role": "system", "content": target\_model\_config.prompt},\ {"role": "user", "content": question},\ \], ) # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "is\_correct": verdict\_clean == "true", } # {{/docs-fragment generate\_and\_review}} async def run\_grouped\_task( i, index, question, answer, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {result\['model\_response'\]} {correct\_html} """, do\_flush=True, ) return result # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Answer Model Response Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_train: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( csv\_file: File | str = "https://dub.sh/geometric-shapes", target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="Solve the given problem about geometric shapes. Think step by step.", max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem. The response may contain detailed steps and explanations, but the final answer is the key point. Please determine if the final answer provided in the response is correct based on the ground truth number. Respond with 'True' if the final answer is correct and 'False' if it is not. Only respond with 'True' or 'False', nothing else. Model Response: {response} Ground Truth: {answer} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality. {prompt\_scores\_str} Each prompt was used together with a problem statement around geometric shapes. This SVG path element draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle (B) Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 3, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(csv\_file, str) and os.path.isfile(csv\_file): csv\_file = await File.from\_local(csv\_file) df\_train, df\_test = await data\_prep(csv\_file) best\_prompt, training\_accuracy = await prompt\_optimizer( df\_train, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) return { "best\_prompt": best\_prompt, "training\_accuracy": training\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto\_prompt\_engineering/optimizer.py\* ## Run it We add a simple main block so we can run the workflow as a script: \`\`\` # /// script # requires-python = "==3.13" # dependencies = \[\ # "flyte>=2.0.0b52",\ # "pandas==2.3.1",\ # "pyarrow==21.0.0",\ # "litellm==1.75.0",\ # \] # main = "auto\_prompt\_engineering" # params = "" # /// # {{docs-fragment env}} import asyncio import html import os import re from dataclasses import dataclass from typing import Optional, Union import flyte import flyte.report import pandas as pd from flyte.io.\_file import File env = flyte.TaskEnvironment( name="auto-prompt-engineering", image=flyte.Image.from\_uv\_script( \_\_file\_\_, name="auto-prompt-engineering", pre=True ), secrets=\[flyte.Secret(key="openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY")\], resources=flyte.Resources(cpu=1), ) CSS = """ """ # {{/docs-fragment env}} # {{docs-fragment data\_prep}} @env.task async def data\_prep(csv\_file: File | str) -> tuple\[pd.DataFrame, pd.DataFrame\]: """ Load Q&A data from a public Google Sheet CSV export URL and split into train/test DataFrames. The sheet should have columns: 'input' and 'target'. """ df = pd.read\_csv( await csv\_file.download() if isinstance(csv\_file, File) else csv\_file ) if "input" not in df.columns or "target" not in df.columns: raise ValueError("Sheet must contain 'input' and 'target' columns.") # Shuffle rows df = df.sample(frac=1, random\_state=1234).reset\_index(drop=True) # Train/Test split df\_train = df.iloc\[:150\].rename(columns={"input": "question", "target": "answer"}) df\_test = df.iloc\[150:250\].rename(columns={"input": "question", "target": "answer"}) return df\_train, df\_test # {{/docs-fragment data\_prep}} # {{docs-fragment model\_config}} @dataclass class ModelConfig: model\_name: str hosted\_model\_uri: Optional\[str\] = None temperature: float = 0.0 max\_tokens: Optional\[int\] = 1000 timeout: int = 600 prompt: str = "" # {{/docs-fragment model\_config}} # {{docs-fragment call\_model}} @flyte.trace async def call\_model( model\_config: ModelConfig, messages: list\[dict\[str, str\]\], ) -> str: from litellm import acompletion response = await acompletion( model=model\_config.model\_name, api\_base=model\_config.hosted\_model\_uri, messages=messages, temperature=model\_config.temperature, timeout=model\_config.timeout, max\_tokens=model\_config.max\_tokens, ) return response.choices\[0\].message\["content"\] # {{/docs-fragment call\_model}} # {{docs-fragment generate\_and\_review}} async def generate\_and\_review( index: int, question: str, answer: str, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, ) -> dict: # Generate response from target model response = await call\_model( target\_model\_config, \[\ {"role": "system", "content": target\_model\_config.prompt},\ {"role": "user", "content": question},\ \], ) # Format review prompt with response + answer review\_messages = \[\ {\ "role": "system",\ "content": review\_model\_config.prompt.format(\ response=response,\ answer=answer,\ ),\ }\ \] verdict = await call\_model(review\_model\_config, review\_messages) # Normalize verdict verdict\_clean = verdict.strip().lower() if verdict\_clean not in {"true", "false"}: verdict\_clean = "not sure" return { "index": index, "model\_response": response, "is\_correct": verdict\_clean == "true", } # {{/docs-fragment generate\_and\_review}} async def run\_grouped\_task( i, index, question, answer, semaphore, target\_model\_config, review\_model\_config, counter, counter\_lock, ): async with semaphore: with flyte.group(name=f"row-{i}"): result = await generate\_and\_review( index, question, answer, target\_model\_config, review\_model\_config, ) async with counter\_lock: # Update counters counter\["processed"\] += 1 if result\["is\_correct"\]: counter\["correct"\] += 1 correct\_html = "✔ Yes" else: correct\_html = "✘ No" # Calculate accuracy accuracy\_pct = (counter\["correct"\] / counter\["processed"\]) \* 100 # Update chart await flyte.report.log.aio( f"", do\_flush=True, ) # Add row to table await flyte.report.log.aio( f""" {html.escape(question)} {html.escape(answer)} {result\['model\_response'\]} {correct\_html} """, do\_flush=True, ) return result # {{docs-fragment evaluate\_prompt}} @env.task(report=True) async def evaluate\_prompt( df: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, concurrency: int, ) -> float: semaphore = asyncio.Semaphore(concurrency) counter = {"correct": 0, "processed": 0} counter\_lock = asyncio.Lock() # Write initial HTML structure await flyte.report.log.aio( CSS + """

Model Evaluation Results

Live Accuracy

Accuracy: 0.0% """, do\_flush=True, ) # Launch tasks concurrently tasks = \[\ run\_grouped\_task(\ i,\ row.Index,\ row.question,\ row.answer,\ semaphore,\ target\_model\_config,\ review\_model\_config,\ counter,\ counter\_lock,\ )\ for i, row in enumerate(df.itertuples(index=True))\ \] await asyncio.gather(\*tasks) # Close table await flyte.report.log.aio("
Question Answer Model Response Correct?
", do\_flush=True) async with counter\_lock: return ( (counter\["correct"\] / counter\["processed"\]) if counter\["processed"\] else 0.0 ) # {{/docs-fragment evaluate\_prompt}} @dataclass class PromptResult: prompt: str accuracy: float # {{docs-fragment prompt\_optimizer}} @env.task(report=True) async def prompt\_optimizer( df\_train: pd.DataFrame, target\_model\_config: ModelConfig, review\_model\_config: ModelConfig, optimizer\_model\_config: ModelConfig, max\_iterations: int, concurrency: int, ) -> tuple\[str, float\]: prompt\_accuracies: list\[PromptResult\] = \[\] # Send styling + table header immediately await flyte.report.log.aio( CSS + """

📊 Prompt Accuracy Comparison

""", do\_flush=True, ) # Step 1: Evaluate starting prompt and stream row with flyte.group(name="baseline\_evaluation"): starting\_accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append( PromptResult(prompt=target\_model\_config.prompt, accuracy=starting\_accuracy) ) await \_log\_prompt\_row(target\_model\_config.prompt, starting\_accuracy) # Step 2: Optimize prompts one by one, streaming after each while len(prompt\_accuracies) <= max\_iterations: with flyte.group(name=f"prompt\_optimization\_step\_{len(prompt\_accuracies)}"): # Prepare prompt scores string for optimizer prompt\_scores\_str = "\\n".join( f"{result.prompt}: {result.accuracy:.2f}" for result in sorted(prompt\_accuracies, key=lambda x: x.accuracy) ) optimizer\_model\_prompt = optimizer\_model\_config.prompt.format( prompt\_scores\_str=prompt\_scores\_str ) response = await call\_model( optimizer\_model\_config, \[{"role": "system", "content": optimizer\_model\_prompt}\], ) response = response.strip() match = re.search(r"\\\[\\\[(.\*?)\\\]\\\]", response, re.DOTALL) if not match: print("No new prompt found. Skipping.") continue new\_prompt = match.group(1) target\_model\_config.prompt = new\_prompt accuracy = await evaluate\_prompt( df\_train, target\_model\_config, review\_model\_config, concurrency, ) prompt\_accuracies.append(PromptResult(prompt=new\_prompt, accuracy=accuracy)) # Log this new prompt row immediately await \_log\_prompt\_row(new\_prompt, accuracy) # Close table await flyte.report.log.aio("
Prompt Accuracy
", do\_flush=True) # Find best best\_result = max(prompt\_accuracies, key=lambda x: x.accuracy) improvement = best\_result.accuracy - starting\_accuracy # Summary await flyte.report.log.aio( f"""

🏆 Summary

Best Prompt: {html.escape(best\_result.prompt)}

Best Accuracy: {best\_result.accuracy\*100:.2f}%

Improvement Over Baseline: {improvement\*100:.2f}%

""", do\_flush=True, ) return best\_result.prompt, best\_result.accuracy # {{/docs-fragment prompt\_optimizer}} async def \_log\_prompt\_row(prompt: str, accuracy: float): """Helper to log a single prompt/accuracy row to Flyte report.""" pct = accuracy \* 100 if pct > 80: color = "linear-gradient(90deg, #4CAF50, #81C784)" elif pct > 60: color = "linear-gradient(90deg, #FFC107, #FFD54F)" else: color = "linear-gradient(90deg, #F44336, #E57373)" await flyte.report.log.aio( f""" {html.escape(prompt)} {pct:.1f}%
""", do\_flush=True, ) # {{docs-fragment auto\_prompt\_engineering}} @env.task async def auto\_prompt\_engineering( csv\_file: File | str = "https://dub.sh/geometric-shapes", target\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="Solve the given problem about geometric shapes. Think step by step.", max\_tokens=10000, ), review\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1-mini", hosted\_model\_uri=None, prompt="""You are a review model tasked with evaluating the correctness of a response to a navigation problem. The response may contain detailed steps and explanations, but the final answer is the key point. Please determine if the final answer provided in the response is correct based on the ground truth number. Respond with 'True' if the final answer is correct and 'False' if it is not. Only respond with 'True' or 'False', nothing else. Model Response: {response} Ground Truth: {answer} """, ), optimizer\_model\_config: ModelConfig = ModelConfig( model\_name="gpt-4.1", hosted\_model\_uri=None, temperature=0.7, max\_tokens=None, prompt=""" I have some prompts along with their corresponding accuracies. The prompts are arranged in ascending order based on their accuracy, where higher accuracy indicate better quality. {prompt\_scores\_str} Each prompt was used together with a problem statement around geometric shapes. This SVG path element draws a Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I) sector (J) triangle (B) Write a new prompt that will achieve an accuracy as high as possible and that is different from the old ones. - It is very important that the new prompt is distinct from ALL the old ones! - Ensure that you analyse the prompts with a high accuracy and reuse the patterns that worked in the past - Ensure that you analyse the prompts with a low accuracy and avoid the patterns that didn't worked in the past - Think out loud before creating the prompt. Describe what has worked in the past and what hasn't. Only then create the new prompt. - Use all available information like prompt length, formal/informal use of language, etc for your analysis. - Be creative, try out different ways of prompting the model. You may even come up with hypothetical scenarios that might improve the accuracy. - You are generating system prompts. This means that there should be no placeholders in the prompt, as they cannot be filled at runtime. Instead focus on general instructions that will help the model to solve the task. - Write your new prompt in double square brackets. Use only plain text for the prompt text and do not add any markdown (i.e. no hashtags, backticks, quotes, etc). """, ), max\_iterations: int = 3, concurrency: int = 10, ) -> dict\[str, Union\[str, float\]\]: if isinstance(csv\_file, str) and os.path.isfile(csv\_file): csv\_file = await File.from\_local(csv\_file) df\_train, df\_test = await data\_prep(csv\_file) best\_prompt, training\_accuracy = await prompt\_optimizer( df\_train, target\_model\_config, review\_model\_config, optimizer\_model\_config, max\_iterations, concurrency, ) with flyte.group(name="test\_data\_evaluation"): baseline\_test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) target\_model\_config.prompt = best\_prompt test\_accuracy = await evaluate\_prompt( df\_test, target\_model\_config, review\_model\_config, concurrency, ) return { "best\_prompt": best\_prompt, "training\_accuracy": training\_accuracy, "baseline\_test\_accuracy": baseline\_test\_accuracy, "test\_accuracy": test\_accuracy, } # {{/docs-fragment auto\_prompt\_engineering}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(auto\_prompt\_engineering) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/auto\_prompt\_engineering/optimizer.py\* Run it with: \`\`\` uv run optimizer.py \`\`\` !\[Execution\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/gifs/tutorials/prompt\_engineering/execution.gif) ## Why this matters Most prompt engineering pipelines start as quick scripts or notebooks. They're fine for experimenting, but they're difficult to scale, reproduce, or debug when things go wrong. With Flyte 2, we get a more reliable setup: - Run many evaluations in parallel with \[async Python\](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/async/page.md) or \[native DSL\](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/async/page.md). - Watch accuracy improve in real time and link results back to the exact dataset, prompt, and model config used. - Resume cleanly after failures without rerunning everything from scratch. - Reuse the same pattern to tune other parameters like temperature, retrieval depth, or agent strategies, not just prompts. ## Next steps You now have a working automated prompt engineering pipeline. Here’s how you can take it further: - \*\*Optimize beyond prompts\*\*: Tune temperature, retrieval strategies, or tool usage just like prompts. - \*\*Expand evaluation metrics\*\*: Add latency, cost, robustness, or diversity alongside accuracy. - \*\*Move toward agentic evaluation\*\*: Instead of single prompts, test how agents plan, use tools, and recover from failures in long-horizon tasks. With this foundation, prompt engineering becomes repeatable, observable, and scalable, ready for production-grade LLM and agent systems. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/model-training === # Model Training Tutorials for training, fine-tuning, and hyperparameter optimization of models at scale. ### \*\*Model Training > Hyperparameter optimization\*\* Run large-scale HPO experiments with zero manual tracking, deterministic results, and automatic recovery. === PAGE: https://www.union.ai/docs/v2/flyte/tutorials/model-training/hpo === # Hyperparameter optimization > \[!NOTE\] > Code available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/ml/optimizer.py). Hyperparameter Optimization (HPO) is a critical step in the machine learning (ML) lifecycle. Hyperparameters are the knobs and dials of a model—values such as learning rates, tree depths, or dropout rates that significantly impact performance but cannot be learned during training. Instead, we must select them manually or optimize them through guided search. Model developers often enjoy the flexibility of choosing from a wide variety of model types, whether gradient boosted machines (GBMs), generalized linear models (GLMs), deep learning architectures, or dozens of others. A common challenge across all these options is the need to systematically explore model performance across hyperparameter configurations tailored to the specific dataset and task. Thankfully, this exploration can be automated. Frameworks like \[Optuna\](https://optuna.org/), \[Hyperopt\](https://hyperopt.github.io/hyperopt/), and \[Ray Tune\](https://docs.ray.io/en/latest/tune/index.html) use advanced sampling algorithms to efficiently search the hyperparameter space and identify optimal configurations. HPO may be executed in two distinct ways: - \*\*Serial HPO\*\* runs one trial at a time, which is easy to set up but can be painfully slow. - \*\*Parallel HPO\*\* distributes trials across multiple processes. It typically follows a pattern with two parameters: \*\*\_N\_\*\*, the total number of trials to run, and \*\*\_C\_\*\*, the maximum number of trials that can run concurrently. Trials are executed asynchronously, and new ones are scheduled based on the results and status of completed or in-progress ones. However, parallel HPO introduces a new complexity: the need for a centralized state that tracks: - All past trials (successes and failures) - All ongoing trials This state is essential so that the optimization algorithm can make informed decisions about which hyperparameters to try next. ## A better way to run HPO This is where Flyte shines. - There's no need to manage a separate centralized database for state tracking, as every objective run is \*\*cached\*\*, \*\*recorded\*\*, and \*\*recoverable\*\* via Flyte's execution engine. - The entire HPO process is observable in the UI with full lineage and metadata for each trial. - Each objective is seeded for reproducibility, enabling deterministic trial results. - If the main optimization task crashes or is terminated, \*\*Flyte can resume from the last successful or failed trial, making the experiment highly fault-tolerant\*\*. - Trial functions can be strongly typed, enabling rich, flexible hyperparameter spaces while maintaining strict type safety across trials. In this example, we combine Flyte with Optuna to optimize a \`RandomForestClassifier\` on the Iris dataset. Each trial runs in an isolated task, and the optimization process is orchestrated asynchronously, with Flyte handling the underlying scheduling, retries, and caching. ## Declare dependencies We start by declaring a Python environment using Python 3.13 and specifying our runtime dependencies. \`\`\` # /// script requires-python = "==3.13" dependencies = \[\ "optuna>=4.0.0,<5.0.0",\ "flyte>=2.0.0b0",\ "scikit-learn==1.7.0",\ \] # /// \`\`\` With the environment defined, we begin by importing standard library and third-party modules necessary for both the ML task and distributed execution. \`\`\` import asyncio import typing from collections import Counter from typing import Optional, Union \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* These standard library imports are essential for asynchronous execution (\`asyncio\`), type annotations (\`typing\`, \`Optional\`, \`Union\`), and aggregating trial state counts (\`Counter\`). \`\`\` import optuna from optuna import Trial from sklearn.datasets import load\_iris from sklearn.ensemble import RandomForestClassifier from sklearn.model\_selection import cross\_val\_score from sklearn.utils import shuffle \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* We use Optuna for hyperparameter optimization and several utilities from scikit-learn to prepare data (\`load\_iris\`), define the model (\`RandomForestClassifier\`), evaluate it (\`cross\_val\_score\`), and shuffle the dataset for randomness (\`shuffle\`). \`\`\` import flyte import flyte.errors \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* Flyte is our orchestration framework. We use it to define tasks, manage resources, and recover from execution errors. ## Define the task environment We define a Flyte task environment called \`driver\`, which encapsulates metadata, compute resources, the container image context needed for remote execution, and caching behavior. \`\`\` driver = flyte.TaskEnvironment( name="driver", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="optimizer"), cache="auto", ) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* This environment specifies that the tasks will run with 1 CPU and 250Mi of memory, the image is built using the current script (\`\_\_file\_\_\`), and caching is enabled. ## Define the optimizer Next, we define an \`Optimizer\` class that handles parallel execution of Optuna trials using async coroutines. This class abstracts the full optimization loop and supports concurrent trial execution with live logging. \`\`\` class Optimizer: def \_\_init\_\_( self, objective: callable, n\_trials: int, concurrency: int = 1, delay: float = 0.1, study: Optional\[optuna.Study\] = None, log\_delay: float = 0.1, ): self.n\_trials: int = n\_trials self.concurrency: int = concurrency self.objective: typing.Callable = objective self.delay: float = delay self.log\_delay = log\_delay self.study = study if study else optuna.create\_study() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* We pass the \`objective\` function, number of trials to run (\`n\_trials\`), and maximum parallel trials (\`concurrency\`). The optional delay throttles execution between trials, while \`log\_delay\` controls how often logging runs. If no existing Optuna Study is provided, a new one is created automatically. \`\`\` async def log(self): while True: await asyncio.sleep(self.log\_delay) counter = Counter() for trial in self.study.trials: counter\[trial.state.name.lower()\] += 1 counts = dict(counter, queued=self.n\_trials - len(self)) # print items in dictionary in a readable format formatted = \[f"{name}: {count}" for name, count in counts.items()\] print(f"{' '.join(formatted)}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* This method periodically prints the number of trials in each state (e.g., running, complete, fail). It keeps users informed of ongoing optimization progress and is invoked as a background task when logging is enabled. !\[Optuna logging\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/tutorials/hpo/logging.png) \_Logs are streamed live as the execution progresses.\_ \`\`\` async def spawn(self, semaphore: asyncio.Semaphore): async with semaphore: trial: Trial = self.study.ask() try: print("Starting trial", trial.number) params = { "n\_estimators": trial.suggest\_int("n\_estimators", 10, 200), "max\_depth": trial.suggest\_int("max\_depth", 2, 20), "min\_samples\_split": trial.suggest\_float( "min\_samples\_split", 0.1, 1.0 ), } output = await self.objective(params) self.study.tell(trial, output, state=optuna.trial.TrialState.COMPLETE) except flyte.errors.RuntimeUserError as e: print(f"Trial {trial.number} failed: {e}") self.study.tell(trial, state=optuna.trial.TrialState.FAIL) await asyncio.sleep(self.delay) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* Each call to \`spawn\` runs a single Optuna trial. The \`semaphore\` ensures that only a fixed number of concurrent trials are active at once, respecting the \`concurrency\` parameter. We first ask Optuna for a new trial and generate a parameter dictionary by querying the trial object for suggested hyperparameters. The trial is then evaluated by the objective function. If successful, we mark it as \`COMPLETE\`. If the trial fails due to a \`RuntimeUserError\` from Flyte, we log and record the failure in the Optuna study. \`\`\` async def \_\_call\_\_(self): # create semaphore to manage concurrency semaphore = asyncio.Semaphore(self.concurrency) # create list of async trials trials = \[self.spawn(semaphore) for \_ in range(self.n\_trials)\] logger: Optional\[asyncio.Task\] = None if self.log\_delay: logger = asyncio.create\_task(self.log()) # await all trials to complete await asyncio.gather(\*trials) if self.log\_delay and logger: logger.cancel() try: await logger except asyncio.CancelledError: pass \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* The \`\_\_call\_\_\` method defines the overall async optimization routine. It creates the semaphore, spawns \`n\_trials\` coroutines, and optionally starts the background logging task. All trials are awaited with \`asyncio.gather\`. \`\`\` def \_\_len\_\_(self) -> int: """Return the number of trials in history.""" return len(self.study.trials) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* This method simply allows us to query the number of trials already associated with the study. ## Define the objective function The objective task defines how we evaluate a particular set of hyperparameters. It's an async task, allowing for caching, tracking, and recoverability across executions. \`\`\` @driver.task async def objective(params: dict\[str, Union\[int, float\]\]) -> float: data = load\_iris() X, y = shuffle(data.data, data.target, random\_state=42) clf = RandomForestClassifier( n\_estimators=params\["n\_estimators"\], max\_depth=params\["max\_depth"\], min\_samples\_split=params\["min\_samples\_split"\], random\_state=42, n\_jobs=-1, ) # Use cross-validation to evaluate performance score = cross\_val\_score(clf, X, y, cv=3, scoring="accuracy").mean() return score.item() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* We use the Iris dataset as a toy classification problem. The input params dictionary contains the trial's hyperparameters, which we unpack into a \`RandomForestClassifier\`. We shuffle the dataset for randomness, and compute a 3-fold cross-validation accuracy. ## Define the main optimization loop The optimize task is the main driver of our optimization experiment. It creates the \`Optimizer\` instance and invokes it. \`\`\` @driver.task async def optimize( n\_trials: int = 20, concurrency: int = 5, delay: float = 0.05, log\_delay: float = 0.1, ) -> dict\[str, Union\[int, float\]\]: optimizer = Optimizer( objective=objective, n\_trials=n\_trials, concurrency=concurrency, delay=delay, log\_delay=log\_delay, study=optuna.create\_study( direction="maximize", sampler=optuna.samplers.TPESampler(seed=42) ), ) await optimizer() best = optimizer.study.best\_trial print("✅ Best Trial") print(" Number :", best.number) print(" Params :", best.params) print(" Score :", best.value) return best.params \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* We configure a \`TPESampler\` for Optuna and \`seed\` it for determinism. After running all trials, we extract the best-performing trial and print its parameters and score. Returning the best params allows downstream tasks or clients to use the tuned model. ## Run the experiment Finally, we include an executable entry point to run this optimization using \`flyte.run\`. \`\`\` if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(optimize, 100, 10) print(run.url) run.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/ml/optimizer.py\* We load Flyte config from \`config.yaml\`, launch the optimize task with 100 trials and concurrency of 10, and print a link to view the execution in the Flyte UI. !\[HPO execution\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/tutorials/hpo/execution.png) \_Each objective run is cached, recorded, and recoverable. With concurrency set to 10, only 10 trials execute in parallel at any given time.\_ --- # Unknown \# Documentation Welcome to the documentation. ## Subpages - \[User Guide\](https://www.union.ai/docs/v1/flyte/user-guide/page.md) - \[Introduction\](https://www.union.ai/docs/v1/flyte/introduction/page.md) - \[Getting started\](https://www.union.ai/docs/v1/flyte/getting-started/page.md) - \[Core concepts\](https://www.union.ai/docs/v1/flyte/core-concepts/page.md) - \[Development cycle\](https://www.union.ai/docs/v1/flyte/development-cycle/page.md) - \[Data input/output\](https://www.union.ai/docs/v1/flyte/data-input-output/page.md) - \[Programming\](https://www.union.ai/docs/v1/flyte/programming/page.md) - \[Tutorials\](https://www.union.ai/docs/v1/flyte/tutorials/page.md) - \[Creating a RAG App with LanceDB and Google Gemini\](https://www.union.ai/docs/v1/flyte/retrieval-augmented-generation/lance-db-rag) - \[Taking NVIDIA’s Enterprise RAG Blueprint to Production\](https://www.union.ai/docs/v1/flyte/compound-ai-systems/enterprise-rag-blueprint) - \[Integrations\](https://www.union.ai/docs/v1/flyte/integrations/page.md) - Connectors - Flytekit plugins - Using Flytekit plugins - Native backend plugins - External service backend plugins - Enabling backend plugins - Flyte operators - \[API Reference\](https://www.union.ai/docs/v1/flyte/api-reference/page.md) - \[Flytekit SDK\](https://www.union.ai/docs/v1/flyte/flytekit-sdk/page.md) - \[Pyflyte CLI\](https://www.union.ai/docs/v1/flyte/pyflyte-cli/page.md) - \[Flytectl CLI\](https://www.union.ai/docs/v1/flyte/flytectl-cli/page.md) - \[Flyteidl\](https://www.union.ai/docs/v1/flyte/flyteidl/page.md) - \[Community\](https://www.union.ai/docs/v1/flyte/community/page.md) - Contributing to the codebase - Contributing to documentation - \[Architecture\](https://www.union.ai/docs/v1/flyte/architecture/page.md) - \[Deployment\](https://www.union.ai/docs/v1/flyte/deployment/page.md) --- \*\*Source\*\*: https://github.com/unionai/unionai-docs/blob/main/content/\_index.md \*\*HTML\*\*: https://www.union.ai/docs/v1/flyte/ --- # Unknown \# Documentation Welcome to the documentation. ## Subpages - \[User Guide\](https://www.union.ai/docs/v2/union/user-guide/page.md) - Basics - \[Flyte 2\](https://www.union.ai/docs/v2/union/overview/page.md) - \[Quickstart\](https://www.union.ai/docs/v2/union/quickstart/page.md) - \[Core concepts\](https://www.union.ai/docs/v2/union/core-concepts/page.md) - \[Run modes\](https://www.union.ai/docs/v2/union/run-modes/page.md) - Tasks - \[Configure tasks\](https://www.union.ai/docs/v2/union/task-configuration/page.md) - \[Build tasks\](https://www.union.ai/docs/v2/union/task-programming/page.md) - \[Run and deploy tasks\](https://www.union.ai/docs/v2/union/task-deployment/page.md) - Apps - \[Configure apps\](https://www.union.ai/docs/v2/union/configure-apps/page.md) - \[Build apps\](https://www.union.ai/docs/v2/union/build-apps/page.md) - \[Native app integrations\](https://www.union.ai/docs/v2/union/native-app-integrations/page.md) - \[Serve and deploy apps\](https://www.union.ai/docs/v2/union/serve-and-deploy-apps/page.md) - Agents - \[Build agents\](https://www.union.ai/docs/v2/union/build-agent/page.md) - \[Agent framework integrations\](https://www.union.ai/docs/v2/union/agent-framework-integrations/page.md) - \[Sandboxing\](https://www.union.ai/docs/v2/union/sandboxing/page.md) - \[Build an MCP\](https://www.union.ai/docs/v2/union/build-mcp/page.md) - Access and identity - \[Authenticating\](https://www.union.ai/docs/v2/union/authenticating/page.md) - \[User management\](https://www.union.ai/docs/v2/union/user-management/page.md) - Advanced Guides - \[Project patterns\](https://www.union.ai/docs/v2/union/project-patterns/page.md) - \[Run scaling\](https://www.union.ai/docs/v2/union/run-scaling/page.md) - \[Advanced project\](https://www.union.ai/docs/v2/union/advanced-project/page.md) - \[Migration\](https://www.union.ai/docs/v2/union/migration/page.md) - \[Tutorials\](https://www.union.ai/docs/v2/union/tutorials/page.md) - Industry verticals - \[Biotech & Healthcare\](https://www.union.ai/docs/v2/union/biotech-healthcare/page.md) - \[Geospatial\](https://www.union.ai/docs/v2/union/geospatial/page.md) - \[Financial Services & Fintech\](https://www.union.ai/docs/v2/union/financial-services/page.md) - \[Frontier AI\](https://www.union.ai/docs/v2/union/frontier-ai/page.md) - Technical topics - \[Computer Vision\](https://www.union.ai/docs/v2/union/computer-vision/page.md) - \[Agents\](https://www.union.ai/docs/v2/union/agents/page.md) - \[Context Engineering\](https://www.union.ai/docs/v2/union/context-engineering/page.md) - \[Model Training\](https://www.union.ai/docs/v2/union/model-training/page.md) - \[Data Processing\](https://www.union.ai/docs/v2/union/data-processing/page.md) - \[Integrations\](https://www.union.ai/docs/v2/union/integrations/page.md) - Integration categories - Distributed compute - Supported distributed compute integrations - How the plugin system works - Example: Using the Dask plugin - Key design principle - Agentic AI - Supported agentic AI integrations - Experiment tracking - Supported experiment tracking integrations - Configuration - Supported configuration integrations - Data validation - Supported data validation integrations - Connectors - Supported connectors - Creating a new connector - Async connector interface - Example: Batch job connector - Connector-level secrets - Deploy a custom connector - LLM Serving - Supported LLM serving integrations - Notebook execution - Supported notebook execution integrations - \[API Reference\](https://www.union.ai/docs/v2/union/api-reference/page.md) - \[Flyte SDK\](https://www.union.ai/docs/v2/union/flyte-sdk/page.md) - \[Flyte CLI\](https://www.union.ai/docs/v2/union/flyte-cli/page.md) - \[Migration from Flyte 1\](https://www.union.ai/docs/v2/union/migration/page.md) - \[Community\](https://www.union.ai/docs/v2/union/community/page.md) - Contributing to documentation - \[Release Notes\](https://www.union.ai/docs/v2/union/release-notes/page.md) - June 2026 - :rocket: Retries with Backoff and Timeout Controls - :robot: Flyte-Native Agent Construct - :memo: Multi-Pod Log Streaming - :computer: Build and Deploy MCP Servers - :sparkles: Smarter Failure Classification - :zap: Faster CLI Startup and Unified Local Caches - :sparkles: Friendlier Build and Deploy Errors - :zap: More Resilient Uploads - :wrench: SDK Reliability Improvements - :wrench: Exclude Files from Code Bundles with \`.flyteignore\` - :wrench: Settings Applied at Run Creation - :sparkles: Console Run Exploration Improvements - :sparkles: Task Environment Details Drawer - :gear: Cluster and Cluster Pool Management from the CLI - :chart\_with\_upwards\_trend: GPU Configuration for Ray Clusters - :sparkles: Pydantic Union Types with Field Annotations - :sparkles: Queues with Concurrency and Depth Control (Beta) - :sparkles: Events API (Beta) - :gear: Self-Managed Dataplane Updates - May 2026 - :robot: AI Agent Components in \`flyte.ai\` - :wrench: Hierarchical Settings - :sparkles: Interactive Triggers - :hammer: Non-Root Images by Default - :sparkles: Actionable CLI Error Messages - :hammer: Image Building Enhancements - :sparkles: Console Enhancements - :computer: TUI and Shell Task Improvements - :sparkles: New Plugins and Examples - :wrench: User-Facing Logger - :wrench: Keyring Opt-Out - :gear: Cluster-Aware Data Access - :gear: OpenShift Support for Self-Managed Deployments - March 2026 - :wrench: Extended Idle Timeout for Panel Apps - :wrench: Plugin Variants Documentation - :rocket: Google Gemini Plugin Integration - :hammer: Forced Image Build Caching - :computer: LLM-Powered Code Generation - :wrench: Updated AI Plugin Examples - :wrench: Debug Mode Integration - :sparkles: Improved CLI Enum Support - :memo: Programmatic Log Access - :zap: Simplified PyTorch Example Setup - :chart\_with\_upwards\_trend: Distributed Training Evaluation - :zap: Improved Benchmark Flexibility - :computer: CLI Project Management - :robot: Anthropic Claude Integration - :hourglass\_flowing\_sand: Panel App Enhancements - :gear: AWS Config File Support - :sparkles: Improved Task Execution Reliability - :wrench: Enhanced Action Service Integration - :rocket: Async Training with Early Stopping - :wrench: Improved Include Path Handling - :zap: Enhanced Retry Management - :zap: Improved Module Loading - :zap: Dynamic Batching for Improved GPU Utilization - :sparkles: Run Cache Disabling - :computer: Vim Key Navigation for TUI - :sparkles: Clickable Image Build URLs - :sparkles: Enhanced Run Filters - :wrench: Simplified Dependency Management - :robot: MLE Agent Enhancements - :sparkles: Improved Task Command Initialization - :zap: New Example Applications & Bug Fixes - :gear: Phase Transitions Tracking - :wrench: Multiple Source Files Support - :package: Simplified Code Bundling - :wrench: Improved Error Messaging for Deployment - :wrench: Improved Debugging for Reusable Tasks - :sparkles: JSONL Plugin Support - February 2026 - :sparkles: JSON Schema Enhancement - :calculator: Panel Calculator Example - :sparkles: Spark Plugin Update - :lock: Secure Package Specification - :zap: Enum Name Acceptance in CLI - :wrench: Enhanced Pod Template Handling - :zap: Stress Testing Example Added - :bug: Correct Serialization Field - :wrench: Improved Async Task Handling - :wrench: Sync Alignment of File Upload Methods - :hourglass: Request Timeout Configuration - :wrench: Enhanced Bundling and Error Handling - :wrench: Dynamic Pydantic Model Creation - :busts\_in\_silhouette: Human-in-the-Loop Plugin - :rocket: Stateless Code Sandbox - :wrench: Improved CLI Logging Initialization - :wrench: Enhanced Ignore Handling - :whale: CI Image Builder - :wrench: TypedDict Compatibility Fix - :globe\_with\_meridians: Cross-Platform Code Bundling - :wrench: Improved CLI JSON Formatting - :wrench: Improved Pod Image Handling - :sparkles: Flyte Webhook Environment - :repeat: Retry Interceptor for gRPC - :sparkles: Orchestration Sandbox Feature - :wrench: Task Shortname Override Fix - :sparkles: NVIDIA H100 GPU Support - :zap: Enhanced Error Handling in PyTorch Elastic Jobs - :wrench: Reverse Path Priority Fix - :globe\_with\_meridians: S3 Virtual Hosted-Style Support - November 2025 - :fast\_forward: Grouped Runs - :globe\_with\_meridians: Apps (beta) - :label: Custom context - :lock: Secrets UI - Image builds now run in the same project-domain - Support for secret mounts in Poetry and UV projects - October 2025 - :infinity: Larger fanouts - :computer: Remote debugging for Ray head nodes - :zap: Triggers and audit history - :arrow\_up: Deployed tasks and input passing - \[Support\](https://www.union.ai/docs/v2/union/support/page.md) - Severity levels - Response time targets - Shared Slack channel - Support portal - Union Cloud console - Email - \[Security\](https://www.union.ai/docs/v2/union/security/page.md) - Overview - Deployment models - \[Deployment\](https://www.union.ai/docs/v2/union/deployment/page.md) - BYOC deployment - Self-managed deployment - Data plane - Control plane --- \*\*Source\*\*: https://github.com/unionai/unionai-docs/blob/main/content/\_index.md \*\*HTML\*\*: https://www.union.ai/docs/v2/union/ --- # Unknown \# Documentation Welcome to the documentation. ## Subpages - \[User Guide\](https://www.union.ai/docs/v2/flyte/user-guide/page.md) - Basics - \[Flyte 2\](https://www.union.ai/docs/v2/flyte/overview/page.md) - \[Quickstart\](https://www.union.ai/docs/v2/flyte/quickstart/page.md) - \[Core concepts\](https://www.union.ai/docs/v2/flyte/core-concepts/page.md) - \[Run modes\](https://www.union.ai/docs/v2/flyte/run-modes/page.md) - Tasks - \[Configure tasks\](https://www.union.ai/docs/v2/flyte/task-configuration/page.md) - \[Build tasks\](https://www.union.ai/docs/v2/flyte/task-programming/page.md) - \[Run and deploy tasks\](https://www.union.ai/docs/v2/flyte/task-deployment/page.md) - Apps - \[Configure apps\](https://www.union.ai/docs/v2/flyte/configure-apps/page.md) - \[Build apps\](https://www.union.ai/docs/v2/flyte/build-apps/page.md) - \[Native app integrations\](https://www.union.ai/docs/v2/flyte/native-app-integrations/page.md) - \[Serve and deploy apps\](https://www.union.ai/docs/v2/flyte/serve-and-deploy-apps/page.md) - Agents - \[Build agents\](https://www.union.ai/docs/v2/flyte/build-agent/page.md) - \[Agent framework integrations\](https://www.union.ai/docs/v2/flyte/agent-framework-integrations/page.md) - \[Sandboxing\](https://www.union.ai/docs/v2/flyte/sandboxing/page.md) - \[Build an MCP\](https://www.union.ai/docs/v2/flyte/build-mcp/page.md) - Advanced Guides - \[Project patterns\](https://www.union.ai/docs/v2/flyte/project-patterns/page.md) - \[Run scaling\](https://www.union.ai/docs/v2/flyte/run-scaling/page.md) - \[Advanced project\](https://www.union.ai/docs/v2/flyte/advanced-project/page.md) - \[Migration\](https://www.union.ai/docs/v2/flyte/migration/page.md) - \[Tutorials\](https://www.union.ai/docs/v2/flyte/tutorials/page.md) - Industry verticals - \[Biotech & Healthcare\](https://www.union.ai/docs/v2/flyte/biotech-healthcare/page.md) - \[Geospatial\](https://www.union.ai/docs/v2/flyte/geospatial/page.md) - \[Financial Services & Fintech\](https://www.union.ai/docs/v2/flyte/financial-services/page.md) - \[Frontier AI\](https://www.union.ai/docs/v2/flyte/frontier-ai/page.md) - Technical topics - \[Computer Vision\](https://www.union.ai/docs/v2/flyte/computer-vision/page.md) - \[Agents\](https://www.union.ai/docs/v2/flyte/agents/page.md) - \[Context Engineering\](https://www.union.ai/docs/v2/flyte/context-engineering/page.md) - \[Model Training\](https://www.union.ai/docs/v2/flyte/model-training/page.md) - \[Data Processing\](https://www.union.ai/docs/v2/flyte/data-processing) - \[Integrations\](https://www.union.ai/docs/v2/flyte/integrations/page.md) - Integration categories - Distributed compute - Supported distributed compute integrations - How the plugin system works - Example: Using the Dask plugin - Key design principle - Agentic AI - Supported agentic AI integrations - Experiment tracking - Supported experiment tracking integrations - Configuration - Supported configuration integrations - Data validation - Supported data validation integrations - Connectors - Supported connectors - Creating a new connector - Async connector interface - Example: Batch job connector - Connector-level secrets - Deploy a custom connector - LLM Serving - Supported LLM serving integrations - Notebook execution - Supported notebook execution integrations - \[API Reference\](https://www.union.ai/docs/v2/flyte/api-reference/page.md) - \[Flyte SDK\](https://www.union.ai/docs/v2/flyte/flyte-sdk/page.md) - \[Flyte CLI\](https://www.union.ai/docs/v2/flyte/flyte-cli/page.md) - \[Migration from Flyte 1\](https://www.union.ai/docs/v2/flyte/migration/page.md) - \[Community\](https://www.union.ai/docs/v2/flyte/community/page.md) - Contributing to the codebase - Contributing to documentation --- \*\*Source\*\*: https://github.com/unionai/unionai-docs/blob/main/content/\_index.md \*\*HTML\*\*: https://www.union.ai/docs/v2/flyte/ --- # Unknown \# Integrations > This bundle contains all pages in the Integrations section. > Source: https://www.union.ai/docs/v2/flyte/integrations/ === PAGE: https://www.union.ai/docs/v2/flyte/integrations === # Integrations > \*\*📝 Note\*\* > > An LLM-optimized bundle of this entire section is available at \[\`section.md\`\](section.md). > This single file contains all pages in this section, optimized for AI coding agent context. Flyte 2 is designed to be extensible by default. While the core platform covers the most common orchestration needs, many production workloads require specialized infrastructure, external services or execution semantics that go beyond the core runtime. Flyte 2 exposes these capabilities through integrations. Under the hood, integrations are implemented using Flyte 2's plugin system, which provides a consistent way to extend the platform without modifying core execution logic. An integration allows you to declaratively enable new capabilities such as distributed compute frameworks or third-party services without manually managing infrastructure. You specify what you need, and Flyte takes care of how it is provisioned, used and cleaned up. This page covers: - The types of integrations Flyte 2 supports today - How integrations fit into Flyte 2's execution model - How to use integrations in your tasks - The integrations available out of the box If you need functionality that doesn't exist yet, Flyte 2's plugin system is intentionally open-ended. You can build and register your own integrations using the same architecture described here. ## Integration categories Flyte 2 integrations fall into the following categories: 1. \*\*Distributed compute\*\*: Provision transient compute clusters to run tasks across multiple nodes, with automatic lifecycle management. 2. \*\*Agentic AI\*\*: Support for various common aspects of agentic AI applications. 3. \*\*Configuration\*\*: Compose and pass hierarchical configuration objects between tasks, with type-safe schemas and CLI/YAML composition. 4. \*\*Experiment tracking\*\*: Integrate with experiment tracking platforms for logging metrics, parameters, and artifacts. 5. \*\*Data validation\*\*: Enforce schema contracts on dataframes flowing between tasks, with automatic validation reports. 6. \*\*Connectors\*\*: Stateless, long-running services that receive execution requests via gRPC and then submit work to external (or internal) systems. 7. \*\*LLM Serving\*\*: Deploy and serve large language models with an OpenAI-compatible API. 8. \*\*Notebook execution\*\*: Run parameterized Jupyter notebooks as typed Flyte tasks with cell-level reports. ## Distributed compute Distributed compute integrations allow tasks to run on dynamically provisioned clusters. These clusters are created just-in-time, scoped to the task execution and torn down automatically when the task completes. This enables large-scale parallelism without requiring users to operate or maintain long-running infrastructure. ### Supported distributed compute integrations | Plugin | Description | Common use cases | | --------------------------- | ------------------------------------------------ | ------------------------------------------------------ | | \[Ray\](./ray/\_index) | Provisions Ray clusters via KubeRay | Distributed Python, ML training, hyperparameter tuning | | \[Spark\](./spark/\_index) | Provisions Spark clusters via Spark Operator | Large-scale data processing, ETL pipelines | | \[Dask\](./dask/\_index) | Provisions Dask clusters via Dask Operator | Parallel Python workloads, dataframe operations | | \[PyTorch\](./pytorch/\_index) | Distributed PyTorch training with elastic launch | Single-node and multi-node training | Each plugin encapsulates: - Cluster provisioning - Resource configuration - Networking and service discovery - Lifecycle management and teardown From the task author's perspective, these details are abstracted away. ### How the plugin system works At a high level, Flyte 2's distributed compute plugin architecture follows a simple and consistent pattern. #### 1. Registration Each plugin registers itself with Flyte 2's core plugin registry: - \*\*\`TaskPluginRegistry\`\*\*: The central registry for all distributed compute plugins - Each plugin declares: - Its configuration schema - How that configuration maps to execution behavior This registration step makes the plugin discoverable by the runtime. #### 2. Task environments and plugin configuration Integrations are activated through a \`TaskEnvironment\`. A \`TaskEnvironment\` bundles: - A container image - Execution settings - A plugin configuration object enabled with \`plugin\_config\` The plugin configuration describes \_what\_ infrastructure or integration the task requires. #### 3. Automatic provisioning and execution When a task associated with a \`TaskEnvironment\` runs: 1. Flyte inspects the environment's plugin configuration 2. The plugin provisions the required infrastructure or integration 3. The task executes with access to that capability 4. Flyte cleans up all transient resources after completion ### Example: Using the Dask plugin Below is a complete example showing how a task gains access to a Dask cluster simply by running inside an environment configured with the Dask plugin. \`\`\`python from flyteplugins.dask import Dask, WorkerGroup import flyte # Define the Dask cluster configuration dask\_config = Dask( workers=WorkerGroup(number\_of\_workers=4) ) # Create a task environment that enables Dask env = flyte.TaskEnvironment( name="dask\_env", plugin\_config=dask\_config, image=image, ) # Any task in this environment has access to the Dask cluster @env.task async def process\_data(data: list) -> list: from distributed import Client client = Client() # Automatically connects to the provisioned cluster futures = client.map(transform, data) return client.gather(futures) \`\`\` When \`process\_data\` executes, Flyte performs the following steps: 1. Provisions a Dask cluster with 4 workers 2. Executes the task with network access to the cluster 3. Tears down the cluster once the task completes No cluster management logic appears in the task code. The task only expresses intent. ### Key design principle All distributed compute integrations follow the same mental model: - You declare the required capability via configuration - You attach that configuration to a task environment - Tasks decorated with that environment automatically gain access to the capability This makes it easy to swap execution backends or introduce distributed compute incrementally without rewriting workflows. ## Agentic AI Agentic AI integrations provide drop-in replacements for LLM provider SDKs. They let you use Flyte tasks as agent tools so that tool calls run with full Flyte observability, retries, and caching. ### Supported agentic AI integrations | Plugin | Description | Common use cases | | ----------------------------------- | -------------------------------------------------------------- | ---------------------------------------- | | \[OpenAI\](./openai/\_index) | Drop-in replacement for OpenAI Agents SDK \`function\_tool\` | Agentic workflows with OpenAI models | | \[Anthropic\](./anthropic/\_index) | Agent loop and \`function\_tool\` for the Anthropic Claude SDK | Agentic workflows with Claude | | \[Gemini\](./gemini/\_index) | Agent loop and \`function\_tool\` for the Google Gemini SDK | Agentic workflows with Gemini | | \[Code generation\](./codegen/\_index) | LLM-driven code generation with automatic testing in sandboxes | Data processing, ETL, analysis pipelines | ## Experiment tracking Experiment tracking integrations let you log metrics, parameters, and artifacts to external tracking platforms during Flyte task execution. ### Supported experiment tracking integrations | Plugin | Description | Common use cases | | ------------------------------------ | ---------------------------- | ------------------------------------------------ | | \[MLflow\](./mlflow/\_index) | MLflow experiment tracking | Experiment tracking, autologging, model registry | | \[Weights and Biases\](./wandb/\_index) | Weights & Biases integration | Experiment tracking and hyperparameter tuning | ## Configuration Configuration integrations let you compose and pass hierarchical configuration objects between Flyte tasks, with type-safe schemas and CLI/YAML composition. ### Supported configuration integrations | Plugin | Description | Common use cases | | ------------------------------- | ----------------------------------------------------------------- | --------------------------------------------------------------------------------- | | \[OmegaConf\](./omegaconf/\_index) | \`DictConfig\` / \`ListConfig\` as native task input and output types | Passing composed configs between tasks, structured configs, YAML-driven pipelines | | \[Hydra\](./hydra/\_index) | Hydra config composition and sweep submission for Flyte tasks | YAML-driven experiment composition, grid and Bayesian sweeps, hardware presets | ## Data validation Data validation integrations enforce schema contracts on the dataframes flowing between tasks. They validate data at task boundaries, catch type and constraint violations early, and produce HTML reports visible in the Flyte UI. ### Supported data validation integrations | Plugin | Description | Common use cases | | --------------------------- | ---------------------------------------------------------- | ----------------------------------------------------------- | | \[Pandera\](./pandera/\_index) | Validates dataframes with pandera \`DataFrameModel\` schemas | Schema enforcement, data quality checks, validation reports | ## Connectors Connectors are stateless, long‑running services that receive execution requests via gRPC and then submit work to external (or internal) systems. Each connector runs as its own Kubernetes deployment, and is triggered when a Flyte task of the matching type is executed. Although they normally run inside the data plane, you can also run connectors locally as long as the required secrets/credentials are present locally. This is useful because connectors are just Python services that can be spawned in‑process. Connectors are designed to scale horizontally and reduce load on the core Flyte backend because they execute \_outside\_ the core system. This decoupling makes connectors efficient, resilient, and easy to iterate on. You can even test them locally without modifying backend configuration, which reduces friction during development. ### Supported connectors | Connector | Description | Common use cases | | --------------------------------- | ------------------------------------------- | ---------------------------------------- | | \[Snowflake\](./snowflake/\_index) | Run SQL queries on Snowflake asynchronously | Data warehousing, ETL, analytics queries | | \[BigQuery\](./bigquery/\_index) | Run SQL queries on Google BigQuery | Data warehousing, ETL, analytics queries | | \[Databricks\](./databricks/\_index) | Run PySpark jobs on Databricks clusters | Large-scale data processing, Spark ETL | ### Creating a new connector If none of the existing connectors meet your needs, you can build your own. > \[!NOTE\] > Connectors communicate via Protobuf, so in theory they can be implemented in any language. > Today, only \*\*Python\*\* connectors are supported. ### Async connector interface To implement a new async connector, extend \`AsyncConnector\` and implement the following methods, all of which must be idempotent: | Method | Purpose | | ---------- | ----------------------------------------------------------- | | \`create\` | Launch the external job (via REST, gRPC, SDK, or other API) | | \`get\` | Fetch current job state (return job status or output) | | \`delete\` | Delete / cancel the external job | | \`get\_logs\` | Stream paginated log lines to the Flyte UI | To test the connector locally, the connector task should inherit from \[AsyncConnectorExecutorMixin\](https://github.com/flyteorg/flyte-sdk/blob/1d49299294cd5e15385fe8c48089b3454b7a4cd1/src/flyte/connectors/\_connector.py#L206). This mixin simulates how the Flyte 2 system executes asynchronous connector tasks, making it easier to validate your connector implementation before deploying it. ### Example: Batch job connector The following example implements a connector that simulates submitting and polling an external batch job. Replace the mock logic with real API calls for your use case. \*\*Connector\*\* (\`my\_connector/connector.py\`): \`\`\` import time import uuid from dataclasses import dataclass from typing import Any, Dict, Optional from flyteidl2.connector.connector\_pb2 import ( GetTaskLogsResponse, GetTaskLogsResponseBody, GetTaskLogsResponseHeader, ) from flyteidl2.core.execution\_pb2 import TaskExecution from flyteidl2.logs.dataplane.payload\_pb2 import LogLine, LogLineOriginator from google.protobuf.timestamp\_pb2 import Timestamp from flyte import logger from flyte.connectors import AsyncConnector, ConnectorRegistry, Resource, ResourceMeta @dataclass class BatchJobMetadata(ResourceMeta): job\_id: str created\_at: float class BatchJobConnector(AsyncConnector): name = "Batch Job Connector" task\_type\_name = "batch\_job" metadata\_type = BatchJobMetadata async def create(self, task\_template, inputs: Optional\[Dict\[str, Any\]\] = None, \*\*kwargs) -> BatchJobMetadata: job\_id = str(uuid.uuid4())\[:8\] logger.info(f"Submitted batch job {job\_id}") return BatchJobMetadata(job\_id=job\_id, created\_at=time.time()) async def get(self, resource\_meta: BatchJobMetadata, \*\*kwargs) -> Resource: elapsed = time.time() - resource\_meta.created\_at if elapsed < 5: return Resource(phase=TaskExecution.RUNNING, message="Job in progress") return Resource( phase=TaskExecution.SUCCEEDED, message="Job completed", outputs={"result": f"output-from-{resource\_meta.job\_id}"}, ) async def delete(self, resource\_meta: BatchJobMetadata, \*\*kwargs): logger.info(f"Cancelled job {resource\_meta.job\_id}") async def get\_logs(self, resource\_meta: BatchJobMetadata, token: str = "", \*\*kwargs): def line(message: str, ts: float) -> LogLine: t = Timestamp() t.FromSeconds(int(ts)) return LogLine(timestamp=t, message=message, originator=LogLineOriginator.USER) start = resource\_meta.created\_at job\_id = resource\_meta.job\_id pages = { "": GetTaskLogsResponseBody(lines=\[ line(f"\[INFO\] Job {job\_id} submitted", start), line(f"\[INFO\] Job {job\_id} started", start + 1), \]), "page-2": GetTaskLogsResponseBody(lines=\[ line(f"\[INFO\] Job {job\_id} finished", start + 5), \]), } next\_tokens = {"": "page-2", "page-2": ""} yield GetTaskLogsResponse(body=pages.get(token, GetTaskLogsResponseBody(lines=\[\]))) next\_token = next\_tokens.get(token, "") if next\_token: yield GetTaskLogsResponse(header=GetTaskLogsResponseHeader(token=next\_token)) ConnectorRegistry.register(BatchJobConnector()) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/batch\_job/connector.py\* \*\*Task plugin\*\* (\`my\_connector/task.py\`): \`\`\` from dataclasses import dataclass from typing import Any, Dict, Optional, Type from flyte.connectors import AsyncConnectorExecutorMixin from flyte.extend import TaskTemplate from flyte.models import NativeInterface, SerializationContext @dataclass class BatchJobConfig: timeout\_seconds: int = 300 class BatchJobTask(AsyncConnectorExecutorMixin, TaskTemplate): \_TASK\_TYPE = "batch\_job" def \_\_init\_\_(self, name: str, plugin\_config: BatchJobConfig, inputs: Optional\[Dict\[str, Type\]\] = None, outputs: Optional\[Dict\[str, Type\]\] = None, \*\*kwargs): super().\_\_init\_\_( name=name, interface=NativeInterface( {k: (v, None) for k, v in inputs.items()} if inputs else {}, outputs or {}, ), task\_type=self.\_TASK\_TYPE, image=None, \*\*kwargs, ) self.plugin\_config = plugin\_config def custom\_config(self, sctx: SerializationContext) -> Optional\[Dict\[str, Any\]\]: return {"timeout\_seconds": self.plugin\_config.timeout\_seconds} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/batch\_job/task.py\* \*\*Usage\*\*: \`\`\`python import flyte from my\_connector.task import BatchJobConfig, BatchJobTask batch\_job = BatchJobTask( name="my\_batch\_job", plugin\_config=BatchJobConfig(timeout\_seconds=60), inputs={"name": str}, outputs={"result": str}, ) flyte.TaskEnvironment.from\_task("batch-job-env", batch\_job) \`\`\` ### Connector-level secrets If your connector needs credentials (API keys, tokens) shared across all tasks, pass them as environment variables into the connector process. Set environment variables on the connector Kubernetes deployment: \`\`\`bash kubectl set env deployment/ MY\_API\_KEY= -n \`\`\` Inside the connector, read the secret from the environment: \`\`\`python import os api\_key = os.environ\["MY\_API\_KEY"\] \`\`\` See \[Secrets\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets/page.md) for how to store and manage secrets. ### Deploy a custom connector Deploying a connector requires two steps: building a Docker image that contains your connector code and then patching the connector Kubernetes deployment to use it. \*\*Step 1: Build the connector image\*\* \`\`\`python import asyncio from flyte import Image from flyte.extend import ImageBuildEngine async def build\_connector\_image(registry: str, name: str, builder: str = "local"): image = Image.from\_debian\_base( registry=registry, name=name ).with\_pip\_packages("flyte\[connector\]", "my-connector-package") await ImageBuildEngine.build(image, builder=builder) if \_\_name\_\_ == "\_\_main\_\_": asyncio.run( build\_connector\_image( registry="", name="my-connector", builder="local" ) ) \`\`\` \*\*Step 2: Override the connector deployment image\*\* Once the image is pushed, patch the connector Kubernetes deployment to use it: \`\`\`bash kubectl set image deployment/ \\ connector=/my-connector: \\ -n \`\`\` Replace \`\` with the name of your connector deployment (e.g. \`flyte-connector\`), and \`\` with the namespace where Flyte is installed (typically \`flyte\`). ## LLM Serving LLM serving integrations let you deploy and serve large language models as Flyte apps with an OpenAI-compatible API. They handle model loading, GPU management, and autoscaling. ### Supported LLM serving integrations | Plugin | Description | Common use cases | | --------------------------------------------- | --------------------------------------------------- | ---------------------------- | | \[SGLang\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/sglang-app/page.md) | Deploy models with SGLang's high-throughput runtime | LLM inference, model serving | | \[vLLM\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/vllm-app/page.md) | Deploy models with vLLM's PagedAttention engine | LLM inference, model serving | For full setup instructions including multi-GPU deployment, model prefetching, and autoscaling, see the \[SGLang app\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/sglang-app/page.md) and \[vLLM app\](https://www.union.ai/docs/v2/flyte/user-guide/native-app-integrations/vllm-app/page.md) pages. ## Notebook execution Notebook execution integrations let you run Jupyter notebooks as first-class Flyte tasks with typed inputs and outputs, HTML reports surfaced in the Flyte UI, and the ability to call other Flyte tasks from within the notebook. ### Supported notebook execution integrations | Plugin | Description | Common use cases | | ------------------------------- | ------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------- | | \[Papermill\](./papermill/\_index) | Parameterize and execute \`.ipynb\` files via \[papermill\](https://papermill.readthedocs.io/) | Productionizing exploratory notebooks, cell-by-cell HTML reports, notebook-driven analysis pipelines | === PAGE: https://www.union.ai/docs/v2/flyte/integrations/anthropic === # Anthropic The Anthropic plugin lets you build agentic workflows with \[Claude\](https://www.anthropic.com/) on Flyte. It provides a \`function\_tool\` decorator that wraps Flyte tasks as tools that Claude can call, and a \`run\_agent\` function that drives the agent conversation loop. When Claude calls a tool, the call executes as a Flyte task with full observability, retries, and caching. ## Installation \`\`\`bash pip install flyteplugins-anthropic \`\`\` Requires \`anthropic >= 0.40.0\`. ## Quick start \`\`\`python import flyte from flyteplugins.anthropic import function\_tool, run\_agent env = flyte.TaskEnvironment( name="claude-agent", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="anthropic\_agent"), secrets=flyte.Secret("anthropic\_api\_key", as\_env\_var="ANTHROPIC\_API\_KEY"), ) @function\_tool @env.task async def get\_weather(city: str) -> str: """Get the current weather for a city.""" return f"The weather in {city} is sunny, 72F" @env.task async def main(prompt: str) -> str: tools = \[get\_weather\] return await run\_agent(prompt=prompt, tools=tools) \`\`\` ## API ### \`function\_tool\` Converts a Flyte task, \`@flyte.trace\`-decorated function, or plain callable into a tool that Claude can invoke. \`\`\`python @function\_tool @env.task async def my\_tool(param: str) -> str: """Tool description sent to Claude.""" ... \`\`\` Can also be called with optional overrides: \`\`\`python @function\_tool(name="custom\_name", description="Custom description") @env.task async def my\_tool(param: str) -> str: ... \`\`\` Parameters: | Parameter | Type | Description | |-----------|------|-------------| | \`func\` | callable | The function to wrap | | \`name\` | \`str\` | Override the tool name (defaults to the function name) | | \`description\` | \`str\` | Override the tool description (defaults to the docstring) | > \[!NOTE\] > The docstring on each \`@function\_tool\` task is sent to Claude as the tool description. Write clear, concise docstrings. ### \`Agent\` A dataclass for bundling agent configuration: \`\`\`python from flyteplugins.anthropic import Agent agent = Agent( name="my-agent", instructions="You are a helpful assistant.", model="claude-sonnet-4-20250514", tools=\[get\_weather\], max\_tokens=4096, max\_iterations=10, ) \`\`\` | Field | Type | Default | Description | |-------|------|---------|-------------| | \`name\` | \`str\` | \`"assistant"\` | Agent name | | \`instructions\` | \`str\` | \`"You are a helpful assistant."\` | System prompt | | \`model\` | \`str\` | \`"claude-sonnet-4-20250514"\` | Claude model ID | | \`tools\` | \`list\[FunctionTool\]\` | \`\[\]\` | Tools available to the agent | | \`max\_tokens\` | \`int\` | \`4096\` | Maximum tokens per response | | \`max\_iterations\` | \`int\` | \`10\` | Maximum tool-call loop iterations | ### \`run\_agent\` Runs a Claude conversation loop, dispatching tool calls to Flyte tasks until Claude returns a final response. \`\`\`python result = await run\_agent( prompt="What's the weather in Tokyo?", tools=\[get\_weather\], model="claude-sonnet-4-20250514", ) \`\`\` You can also pass an \`Agent\` object: \`\`\`python result = await run\_agent(prompt="What's the weather?", agent=agent) \`\`\` | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | \`prompt\` | \`str\` | required | User message | | \`tools\` | \`list\[FunctionTool\]\` | \`None\` | Tools available to the agent | | \`agent\` | \`Agent\` | \`None\` | Agent config (overrides individual params) | | \`model\` | \`str\` | \`"claude-sonnet-4-20250514"\` | Claude model ID | | \`system\` | \`str\` | \`None\` | System prompt | | \`max\_tokens\` | \`int\` | \`4096\` | Maximum tokens per response | | \`max\_iterations\` | \`int\` | \`10\` | Maximum iterations (prevents infinite loops) | | \`api\_key\` | \`str\` | \`None\` | API key (falls back to \`ANTHROPIC\_API\_KEY\` env var) | ## Secrets Store your Anthropic API key as a Flyte secret and expose it as an environment variable: \`\`\`python secrets=flyte.Secret("anthropic\_api\_key", as\_env\_var="ANTHROPIC\_API\_KEY") \`\`\` ## API reference See the \[Anthropic API reference\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/anthropic/\_index) for full details. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/bigquery === # BigQuery The BigQuery connector lets you run SQL queries against \[Google BigQuery\](https://cloud.google.com/bigquery) directly from Flyte tasks. Queries are submitted asynchronously via the BigQuery Jobs API and polled for completion, so they don't block a worker while waiting for results. The connector supports: - Parameterized SQL queries with typed inputs - Google Cloud service account authentication - Returns query results as DataFrames - Query cancellation on task abort ## Installation \`\`\`bash pip install flyteplugins-bigquery \`\`\` This installs the Google Cloud BigQuery client libraries. ## Quick start Here's a minimal example that runs a SQL query on BigQuery: \`\`\`python from flyte.io import DataFrame from flyteplugins.bigquery import BigQueryConfig, BigQueryTask config = BigQueryConfig( ProjectID="my-gcp-project", Location="US", ) count\_users = BigQueryTask( name="count\_users", query\_template="SELECT COUNT(\*) FROM dataset.users", plugin\_config=config, output\_dataframe\_type=DataFrame, ) \`\`\` This defines a task called \`count\_users\` that runs the query on the configured BigQuery instance. When executed, the connector: 1. Connects to BigQuery using the provided configuration 2. Submits the query asynchronously via the Jobs API 3. Polls until the query completes or fails To run the task, create a \`TaskEnvironment\` from it and execute it locally or remotely: \`\`\`python import flyte bigquery\_env = flyte.TaskEnvironment.from\_task("bigquery\_env", count\_users) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Run locally (connector runs in-process, requires credentials locally) run = flyte.with\_runcontext(mode="local").run(count\_users) # Run remotely (connector runs as a service in your data plane) run = flyte.with\_runcontext(mode="remote").run(count\_users) print(run.url) \`\`\` > \[!NOTE\] > The \`TaskEnvironment\` created by \`from\_task\` does not need an image or pip packages. BigQuery tasks are connector tasks, which means the query executes on the connector service, not in your task container. In \`local\` mode, the connector runs in-process and requires \`flyteplugins-bigquery\` and credentials to be available on your machine. ## Configuration ### \`BigQueryConfig\` parameters | Field | Type | Required | Description | |-------|------|----------|-------------| | \`ProjectID\` | \`str\` | Yes | GCP project ID | | \`Location\` | \`str\` | No | BigQuery region (e.g., \`"US"\`, \`"EU"\`) | | \`QueryJobConfig\` | \`bigquery.QueryJobConfig\` | No | Native BigQuery \[QueryJobConfig\](https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJobConfig) object for advanced settings | ### \`BigQueryTask\` parameters | Parameter | Type | Description | |-----------|------|-------------| | \`name\` | \`str\` | Unique task name | | \`query\_template\` | \`str\` | SQL query (whitespace is normalized before execution) | | \`plugin\_config\` | \`BigQueryConfig\` | Connection configuration | | \`inputs\` | \`Dict\[str, Type\]\` | Named typed inputs bound as query parameters | | \`output\_dataframe\_type\` | \`Type\[DataFrame\]\` | If set, query results are returned as a \`DataFrame\` | | \`google\_application\_credentials\` | \`str\` | Name of the Flyte secret containing the GCP service account JSON key | ## Authentication Pass the name of a Flyte secret containing your GCP service account JSON key: \`\`\`python query = BigQueryTask( name="secure\_query", query\_template="SELECT \* FROM dataset.sensitive\_data", plugin\_config=config, google\_application\_credentials="my-gcp-sa-key", ) \`\`\` ## Query templating Use the \`inputs\` parameter to define typed inputs for your query. Input values are bound as BigQuery \`ScalarQueryParameter\` values. ### Supported input types | Python type | BigQuery type | |-------------|---------------| | \`int\` | \`INT64\` | | \`float\` | \`FLOAT64\` | | \`str\` | \`STRING\` | | \`bool\` | \`BOOL\` | | \`bytes\` | \`BYTES\` | | \`datetime\` | \`DATETIME\` | | \`list\` | \`ARRAY\` | ### Parameterized query example \`\`\`python from flyte.io import DataFrame events\_by\_region = BigQueryTask( name="events\_by\_region", query\_template="SELECT \* FROM dataset.events WHERE region = @region AND score > @min\_score", plugin\_config=config, inputs={"region": str, "min\_score": float}, output\_dataframe\_type=DataFrame, ) \`\`\` > \[!NOTE\] > The query template is normalized before execution: newlines and tabs are replaced with spaces and consecutive whitespace is collapsed. You can format your queries across multiple lines for readability without affecting execution. ## Retrieving query results Set \`output\_dataframe\_type\` to capture results as a DataFrame: \`\`\`python from flyte.io import DataFrame top\_customers = BigQueryTask( name="top\_customers", query\_template=""" SELECT customer\_id, SUM(amount) AS total\_spend FROM dataset.orders GROUP BY customer\_id ORDER BY total\_spend DESC LIMIT 100 """, plugin\_config=config, output\_dataframe\_type=DataFrame, ) \`\`\` If you don't need query results (for example, DDL statements or INSERT queries), omit \`output\_dataframe\_type\`. ## API reference See the \[BigQuery API reference\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/bigquery/\_index) for full details. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/codegen === # Code generation The code generation plugin turns natural-language prompts into tested, production-ready Python code. You describe what the code should do, along with sample data, schema definitions, constraints, and typed inputs/outputs, and the plugin handles the rest: generating code, writing tests, building an isolated \[code sandbox\](/docs/v2/union//user-guide/sandboxing/code-sandboxing) with the right dependencies, running the tests, diagnosing failures, and iterating until everything passes. The result is a validated script you can execute against real data or deploy as a reusable Flyte task. ## Installation \`\`\`bash pip install flyteplugins-codegen # For Agent mode (Claude-only) pip install flyteplugins-codegen\[agent\] \`\`\` ## Quick start \`\`\`python{hl\_lines=\[3, 4, 6, 12, 14, "20-25"\]} import flyte from flyte.io import File from flyte.sandbox import sandbox\_environment from flyteplugins.codegen import AutoCoderAgent agent = AutoCoderAgent(model="gpt-4.1", name="summarize-sales") env = flyte.TaskEnvironment( name="my-env", secrets=\[flyte.Secret(key="openai\_key", as\_env\_var="OPENAI\_API\_KEY")\], image=flyte.Image.from\_debian\_base().with\_pip\_packages( "flyteplugins-codegen", ), depends\_on=\[sandbox\_environment\], ) @env.task async def process\_data(csv\_file: File) -> tuple\[float, int, int\]: result = await agent.generate.aio( prompt="Read the CSV and compute total\_revenue, total\_units and row\_count.", samples={"sales": csv\_file}, outputs={"total\_revenue": float, "total\_units": int, "row\_count": int}, ) return await result.run.aio() \`\`\` The \`depends\_on=\[sandbox\_environment\]\` declaration is required. It ensures the sandbox runtime is available when dynamically-created sandboxes execute. !\[Sandbox\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/codegen/sandbox.png) ## Two execution backends The plugin supports two backends for generating and validating code. Both share the same \`AutoCoderAgent\` interface and produce the same \`CodeGenEvalResult\`. ### LiteLLM (default) Uses structured-output LLM calls to generate code, detect packages, build sandbox images, run tests, diagnose failures, and iterate. Works with any model that supports structured outputs (GPT-4, Claude, Gemini, etc. via LiteLLM). \`\`\`python{hl\_lines=\[1, 3\]} agent = AutoCoderAgent( name="my-task", model="gpt-4.1", max\_iterations=10, ) \`\`\` The LiteLLM backend follows a fixed pipeline: \`\`\`mermaid flowchart TD A\["prompt + samples"\] --> B\["generate\_plan"\] B --> C\["generate\_code"\] C --> D\["detect\_packages"\] D --> E\["build\_image"\] E --> F{skip\_tests?} F -- yes --> G\["return result"\] F -- no --> H\["generate\_tests"\] H --> I\["execute\_tests"\] I --> J{pass?} J -- yes --> G J -- no --> K\["diagnose\_error"\] K --> L{error type?} L -- "logic error" --> M\["regenerate code"\] L -- "environment error" --> N\["add packages, rebuild image"\] L -- "test error" --> O\["fix test expectations"\] M --> I N --> I O --> I \`\`\` The loop continues until tests pass or \`max\_iterations\` is reached. !\[LiteLLM\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/codegen/litellm.png) ### Agent (Claude) Uses the Claude Agent SDK to autonomously generate, test, and fix code. The agent has access to \`Bash\`, \`Read\`, \`Write\`, and \`Edit\` tools and decides what to do at each step. Test execution commands (\`pytest\`) are intercepted and run inside isolated sandboxes. \`\`\`python{hl\_lines=\["3-4"\]} agent = AutoCoderAgent( name="my-task", model="claude-sonnet-4-5-20250929", backend="claude", ) \`\`\` > \[!NOTE\] > Agent mode requires \`ANTHROPIC\_API\_KEY\` as a Flyte secret and is Claude-only. \*\*Key differences from LiteLLM:\*\* | | LiteLLM | Agent | | --------------------- | --------------------------------- | ---------------------------------------------- | | \*\*Execution\*\* | Fixed generate-test-fix pipeline | Autonomous agent decides actions | | \*\*Model support\*\* | Any model with structured outputs | Claude only | | \*\*Iteration control\*\* | \`max\_iterations\` | \`agent\_max\_turns\` | | \*\*Test execution\*\* | Direct sandbox execution | \`pytest\` commands intercepted via hooks | | \*\*Tool safety\*\* | N/A | Commands classified as safe/denied/intercepted | | \*\*Observability\*\* | Logs + token counts | Full tool call tracing in Flyte UI | In Agent mode, Bash commands are classified before execution: - \*\*Safe\*\* (\`ls\`, \`cat\`, \`grep\`, \`head\`, etc.) — allowed to run directly - \*\*Intercepted\*\* (\`pytest\`) — routed to sandbox execution - \*\*Denied\*\* (\`apt\`, \`pip install\`, \`curl\`, etc.) — blocked for safety ## Providing data ### Sample data Pass sample data via \`samples\` as \`File\` objects or pandas \`DataFrame\`s. The plugin automatically: 1. Converts DataFrames to CSV files 2. Infers \[Pandera\](https://pandera.readthedocs.io/) schemas from the data — column types, nullability 3. Parses natural-language \`constraints\` into Pandera checks (e.g., \`"quantity must be positive"\` becomes \`pa.Check.gt(0)\`) 4. Extracts data context — column statistics, distributions, patterns, sample rows 5. Injects all of this into the LLM prompt so the generated code is aware of the exact data structure Pandera is used purely for prompt enrichment, not runtime validation. The generated code does not import Pandera — it benefits from the LLM knowing the precise data structure. The generated schemas are stored on \`result.generated\_schemas\` for inspection. \`\`\`python{hl\_lines=\[3\]} result = await agent.generate.aio( prompt="Clean and validate the data, remove duplicates", samples={"orders": orders\_df, "products": products\_file}, constraints=\["quantity must be positive", "price between 0 and 10000"\], outputs={"cleaned\_orders": File}, ) \`\`\` ### Schema and constraints Use \`schema\` to provide free-form context about data formats or target structures (e.g., a database schema). Use \`constraints\` to declare business rules that the generated code must respect: \`\`\`python{hl\_lines=\["4-17"\]} result = await agent.generate.aio( prompt=prompt, samples={"readings": sensor\_df}, schema="""Output JSON schema for report\_json: { "sensor\_id": str, "avg\_temp": float, "min\_temp": float, "max\_temp": float, "avg\_humidity": float, } """, constraints=\[ "Temperature values must be between -40 and 60 Celsius", "Humidity values must be between 0 and 100 percent", "Output report must have one row per unique sensor\_id", \], outputs={ "report\_json": str, "total\_anomalies": int, }, ) \`\`\` !\[Pandera Constraints\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/codegen/pandera\_constraints.png) ### Inputs and outputs Declare \`inputs\` for non-sample arguments (e.g., thresholds, flags) and \`outputs\` for the expected result types. Supported output types: \`str\`, \`int\`, \`float\`, \`bool\`, \`datetime.datetime\`, \`datetime.timedelta\`, \`File\`. Sample entries are automatically added as \`File\` inputs — you do not need to redeclare them. \`\`\`python{hl\_lines=\[4, 5\]} result = await agent.generate.aio( prompt="Filter transactions above the threshold", samples={"transactions": tx\_file}, inputs={"threshold": float, "include\_pending": bool}, outputs={"filtered": File, "count": int}, ) \`\`\` ## Running generated code \`agent.generate()\` returns a \`CodeGenEvalResult\`. If \`result.success\` is \`True\`, the generated code passed all tests and you can execute it against real data. If \`max\_iterations\` (LiteLLM) or \`agent\_max\_turns\` (Agent) is reached without tests passing, \`result.success\` is \`False\` and \`result.error\` contains the failure details. Both \`run()\` and \`as\_task()\` return output values as a tuple in the order declared in \`outputs\`. If there is a single output, the value is returned directly (not wrapped in a tuple). ### One-shot execution with \`result.run()\` Runs the generated code in a sandbox. If samples were provided during \`generate()\`, they are used as default inputs. \`\`\`python # Use sample data as defaults total\_revenue, total\_units, count = await result.run.aio() # Override specific inputs total\_revenue, total\_units, count = await result.run.aio(threshold=0.5) # Sync version total\_revenue, total\_units, count = result.run() \`\`\` \`result.run()\` accepts optional configuration: \`\`\`python{hl\_lines=\["4-6"\]} total\_revenue, total\_units, count = await result.run.aio( name="execute-on-data", resources=flyte.Resources(cpu=2, memory="4Gi"), retries=2, timeout=600, cache="auto", ) \`\`\` ### Reusable task with \`result.as\_task()\` Creates a callable sandbox task from the generated code. Useful when you want to run the same generated code against different data. \`\`\`python{hl\_lines=\[1, "6-7", "9-10"\]} task = result.as\_task( name="run-sensor-analysis", resources=flyte.Resources(cpu=1, memory="512Mi"), ) # Call with sample defaults report, total\_anomalies = await task.aio() # Call with different data report, total\_anomalies = await task.aio(readings=new\_data\_file) \`\`\` ## Error diagnosis The LiteLLM backend classifies test failures into three categories and applies targeted fixes: | Error type | Meaning | Action | | ------------- | ----------------------------- | ------------------------------------------------ | | \`logic\` | Bug in the generated code | Regenerate code with specific patch instructions | | \`environment\` | Missing package or dependency | Add the package and rebuild the sandbox image | | \`test\_error\` | Bug in the generated test | Fix the test expectations | If the same error persists after a fix, the plugin reclassifies it (e.g., \`logic\` to \`test\_error\`) to try the other approach. In Agent mode, the agent diagnoses and fixes issues autonomously based on error output. ## Durable execution Code generation is expensive — it involves multiple LLM calls, image builds, and sandbox executions. Without durability, a transient failure in the pipeline (network blip, OOM, downstream service error) would force the entire process to restart from scratch: regenerating code, rebuilding images, re-running sandboxes, making additional LLM calls. Flyte solves this through two complementary mechanisms: \*\*replay logs\*\* and \*\*caching\*\*. ### Replay logs Flyte maintains a replay log that records every trace and task execution within a run. When a task crashes and retries, the system replays the log from the previous attempt rather than recomputing everything: - No additional model calls - No code regeneration - No sandbox re-execution - No container rebuilds The workflow breezes through the earlier steps and resumes from the failure point. This applies as long as the traces and tasks execute in the same order and use the same inputs as the first attempt. ### Caching Separately, Flyte can cache task results across runs. With \`cache="auto"\`, sandbox executions (image builds, test runs, code execution) are cached. This is useful when you re-run the same pipeline — not just when recovering from a crash, but across entirely separate invocations with the same inputs. Together, replay logs handle crash recovery within a run, and caching avoids redundant work across runs. ### Non-determinism in Agent mode One challenge with agents is that they are inherently non-deterministic — the sequence of actions can vary between runs, which could break replay. In practice, the codegen agent follows a predictable pattern (write code, generate tests, run tests, inspect results), which works in replay's favor. The plugin also embeds logic that instructs the agent not to regenerate or re-execute steps that already completed successfully in the first run. This acts as an additional safety check alongside the replay log to account for non-determinism. !\[Agent\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/codegen/agent.png) On the first attempt, the full pipeline runs. If a transient failure occurs, the system instantly replays the traces (which track model calls) and sandbox executions, allowing the pipeline to resume from the point of failure. !\[Durability\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/codegen/durability.png) ## Observability ### LiteLLM backend - Logs every iteration with attempt count, error type, and package changes - Tracks total input/output tokens across all LLM calls (available on \`result.total\_input\_tokens\` and \`result.total\_output\_tokens\`) - Results include full conversation history for debugging (\`result.conversation\_history\`) ### Agent backend - Traces each tool call (name + input) via \`PostToolUse\` hooks - Traces tool failures via \`PostToolUseFailure\` hooks - Traces a summary when the agent finishes (total tool calls, tool distribution, final image/packages) - Classifies Bash commands as safe, denied, or intercepted (for sandbox execution) - All traces appear in the Flyte UI ## Examples ### Processing CSVs with different schemas Generate code that handles varying CSV formats, then run on real data: \`\`\`python{hl\_lines=\[1, 3, 14, 16, 27\]} from flyteplugins.codegen import AutoCoderAgent agent = AutoCoderAgent( name="sales-processor", model="gpt-4.1", max\_iterations=5, resources=flyte.Resources(cpu=1, memory="512Mi"), litellm\_params={"temperature": 0.2, "max\_tokens": 4096}, ) @env.task async def process\_sales(csv\_file: File) -> dict\[str, float | int\]: result = await agent.generate.aio( prompt="Read the CSV and compute total\_revenue, total\_units, and transaction\_count.", samples={"csv\_data": csv\_file}, outputs={ "total\_revenue": float, "total\_units": int, "transaction\_count": int, }, ) if not result.success: raise RuntimeError(f"Code generation failed: {result.error}") total\_revenue, total\_units, transaction\_count = await result.run.aio() return { "total\_revenue": total\_revenue, "total\_units": total\_units, "transaction\_count": transaction\_count, } \`\`\` ### DataFrame analysis with constraints Pass DataFrames directly and enforce business rules with constraints: \`\`\`python{hl\_lines=\[10, "15-19"\]} agent = AutoCoderAgent( model="gpt-4.1", name="sensor-analysis", base\_packages=\["numpy"\], max\_sample\_rows=30, ) @env.task async def analyze\_sensors(sensor\_df: pd.DataFrame) -> tuple\[File, int\]: result = await agent.generate.aio( prompt="""Analyze IoT sensor data. For each sensor, calculate mean/min/max temperature, mean humidity, and count warnings. Output a summary CSV.""", samples={"readings": sensor\_df}, constraints=\[ "Temperature values must be between -40 and 60 Celsius", "Humidity values must be between 0 and 100 percent", "Output report must have one row per unique sensor\_id", \], outputs={ "report": File, "total\_anomalies": int, }, ) if not result.success: raise RuntimeError(f"Code generation failed: {result.error}") task = result.as\_task( name="run-sensor-analysis", resources=flyte.Resources(cpu=1, memory="512Mi"), ) return await task.aio(readings=result.original\_samples\["readings"\]) \`\`\` ### Agent mode The same task using Claude as an autonomous agent: \`\`\`python{hl\_lines=\[3\]} agent = AutoCoderAgent( name="sales-agent", backend="claude", model="claude-sonnet-4-5-20250929", resources=flyte.Resources(cpu=1, memory="512Mi"), ) @env.task async def process\_sales\_with\_agent(csv\_file: File) -> dict\[str, float | int\]: result = await agent.generate.aio( prompt="Read the CSV and compute total\_revenue, total\_units, and transaction\_count.", samples={"csv\_data": csv\_file}, outputs={ "total\_revenue": float, "total\_units": int, "transaction\_count": int, }, ) if not result.success: raise RuntimeError(f"Agent code generation failed: {result.error}") total\_revenue, total\_units, transaction\_count = await result.run.aio() return { "total\_revenue": total\_revenue, "total\_units": total\_units, "transaction\_count": transaction\_count, } \`\`\` ## Configuration ### LiteLLM parameters Tune model behavior with \`litellm\_params\`: \`\`\`python{hl\_lines=\["5-8"\]} agent = AutoCoderAgent( name="my-task", model="anthropic/claude-sonnet-4-20250514", api\_key="ANTHROPIC\_API\_KEY", litellm\_params={ "temperature": 0.3, "max\_tokens": 4000, }, ) \`\`\` ### Image configuration Control the registry and Python version for sandbox images: \`\`\`python{hl\_lines=\["6-10"\]} from flyte.sandbox import ImageConfig agent = AutoCoderAgent( name="my-task", model="gpt-4.1", image\_config=ImageConfig( registry="my-registry.io", registry\_secret="registry-creds", python\_version=(3, 12), ), ) \`\`\` ### Skipping tests Set \`skip\_tests=True\` to skip test generation and execution. The agent still generates code, detects packages, and builds the sandbox image, but does not generate or run tests. \`\`\`python{hl\_lines=\[4\]} agent = AutoCoderAgent( name="my-task", model="gpt-4.1", skip\_tests=True, ) \`\`\` > \[!NOTE\] > \`skip\_tests\` only applies to LiteLLM mode. In Agent mode, the agent autonomously decides when to test. ### Base packages Ensure specific packages are always installed in every sandbox: \`\`\`python{hl\_lines=\[4\]} agent = AutoCoderAgent( name="my-task", model="gpt-4.1", base\_packages=\["numpy", "pandas"\], ) \`\`\` ## Best practices - \*\*One agent per task.\*\* Each \`generate()\` call builds its own sandbox image and manages its own package state. Running multiple agents in the same task can cause resource contention and makes failures harder to diagnose. - \*\*Keep \`cache="auto"\` (the default).\*\* Caching flows to all internal sandboxes, making retries near-instant. Use \`"disable"\` during development if you want fresh executions, or \`"override"\` to force re-execution and update the cached result. - \*\*Set \`max\_iterations\` conservatively.\*\* Start with 5-10 iterations. If the model cannot produce correct code in that budget, the prompt or constraints likely need refinement. - \*\*Provide constraints for data-heavy tasks.\*\* Explicit constraints (e.g., \`"quantity must be positive"\`) produce better schemas and better generated code. - \*\*Inspect \`result.generated\_schemas\`.\*\* Review the inferred Pandera schemas to verify the model understood your data structure correctly. ## API reference ### \`AutoCoderAgent\` constructor | Parameter | Type | Default | Description | | ----------------- | ----------------- | -------------- | -------------------------------------------------------------------------------------- | | \`name\` | \`str\` | \`"auto-coder"\` | Unique name for tracking and image naming | | \`model\` | \`str\` | \`"gpt-4.1"\` | LiteLLM model identifier | | \`backend\` | \`str\` | \`"litellm"\` | Execution backend: \`"litellm"\` or \`"claude"\` | | \`system\_prompt\` | \`str\` | \`None\` | Custom system prompt override | | \`api\_key\` | \`str\` | \`None\` | Name of the environment variable containing the LLM API key (e.g., \`"OPENAI\_API\_KEY"\`) | | \`api\_base\` | \`str\` | \`None\` | Custom API base URL | | \`litellm\_params\` | \`dict\` | \`None\` | Extra LiteLLM params (temperature, max\_tokens, etc.) | | \`base\_packages\` | \`list\[str\]\` | \`None\` | Always-install pip packages | | \`resources\` | \`flyte.Resources\` | \`None\` | Resources for sandbox execution (default: 1 CPU, 1Gi) | | \`image\_config\` | \`ImageConfig\` | \`None\` | Registry, secret, and Python version | | \`max\_iterations\` | \`int\` | \`10\` | Max generate-test-fix iterations (LiteLLM mode) | | \`max\_sample\_rows\` | \`int\` | \`100\` | Rows to sample from data for LLM context | | \`skip\_tests\` | \`bool\` | \`False\` | Skip test generation and execution (LiteLLM mode) | | \`sandbox\_retries\` | \`int\` | \`0\` | Flyte task-level retries for each sandbox execution | | \`timeout\` | \`int\` | \`None\` | Timeout in seconds for sandboxes | | \`env\_vars\` | \`dict\[str, str\]\` | \`None\` | Environment variables for sandboxes | | \`secrets\` | \`list\[Secret\]\` | \`None\` | Flyte secrets for sandboxes | | \`cache\` | \`str\` | \`"auto"\` | Cache behavior: \`"auto"\`, \`"override"\`, or \`"disable"\` | | \`agent\_max\_turns\` | \`int\` | \`50\` | Max turns when \`backend="claude"\` | ### \`generate()\` parameters | Parameter | Type | Default | Description | | ------------- | ------------------------------ | -------- | --------------------------------------------------------------------------------------- | | \`prompt\` | \`str\` | required | Natural-language task description | | \`schema\` | \`str\` | \`None\` | Free-form context about data formats or target structures | | \`constraints\` | \`list\[str\]\` | \`None\` | Natural-language constraints (e.g., \`"quantity must be positive"\`) | | \`samples\` | \`dict\[str, File \\| DataFrame\]\` | \`None\` | Sample data. DataFrames are auto-converted to CSV files. | | \`inputs\` | \`dict\[str, type\]\` | \`None\` | Non-sample input types (e.g., \`{"threshold": float}\`) | | \`outputs\` | \`dict\[str, type\]\` | \`None\` | Output types. Supported: \`str\`, \`int\`, \`float\`, \`bool\`, \`datetime\`, \`timedelta\`, \`File\` | ### \`CodeGenEvalResult\` fields | Field | Type | Description | | -------------------------- | ------------------------- | --------------------------------------------------------- | | \`success\` | \`bool\` | Whether tests passed | | \`solution\` | \`CodeSolution\` | Generated code (\`.code\`, \`.language\`, \`.system\_packages\`) | | \`tests\` | \`str\` | Generated test code | | \`output\` | \`str\` | Test output | | \`exit\_code\` | \`int\` | Test exit code | | \`error\` | \`str \\| None\` | Error message if failed | | \`attempts\` | \`int\` | Number of iterations used | | \`image\` | \`str\` | Built sandbox image with all dependencies | | \`detected\_packages\` | \`list\[str\]\` | Pip packages detected | | \`detected\_system\_packages\` | \`list\[str\]\` | Apt packages detected | | \`generated\_schemas\` | \`dict\[str, str\] \\| None\` | Pandera schemas as Python code strings | | \`data\_context\` | \`str \\| None\` | Extracted data context | | \`original\_samples\` | \`dict\[str, File\] \\| None\` | Sample data as Files (defaults for \`run()\`/\`as\_task()\`) | | \`total\_input\_tokens\` | \`int\` | Total input tokens across all LLM calls | | \`total\_output\_tokens\` | \`int\` | Total output tokens across all LLM calls | | \`conversation\_history\` | \`list\[dict\]\` | Full LLM conversation history for debugging | ### \`CodeGenEvalResult\` methods | Method | Description | | ----------------------------------- | ------------------------------------------------------------------ | | \`result.run(\*\*overrides)\` | Execute generated code in a sandbox. Sample data used as defaults. | | \`await result.run.aio(\*\*overrides)\` | Async version of \`run()\`. | | \`result.as\_task(name, ...)\` | Create a reusable callable sandbox task from the generated code. | Both \`run()\` and \`as\_task()\` accept optional \`name\`, \`resources\`, \`retries\`, \`timeout\`, \`env\_vars\`, \`secrets\`, and \`cache\` parameters. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/dask === # Dask The Dask plugin lets you run \[Dask\](https://www.dask.org/) jobs natively on Kubernetes. Flyte provisions a transient Dask cluster for each task execution using the \[Dask Kubernetes Operator\](https://kubernetes.dask.org/en/latest/operator.html) and tears it down on completion. ## When to use this plugin - Parallel Python workloads that outgrow a single machine - Distributed DataFrame operations on large datasets - Workloads that use Dask's task scheduler for arbitrary computation graphs - Jobs that need to scale NumPy, pandas, or scikit-learn workflows across multiple nodes ## Installation \`\`\`bash pip install flyteplugins-dask \`\`\` Your task image must also include the Dask distributed scheduler: \`\`\`python image = flyte.Image.from\_debian\_base(name="dask").with\_pip\_packages("flyteplugins-dask") \`\`\` ## Configuration Create a \`Dask\` configuration and pass it as \`plugin\_config\` to a \`TaskEnvironment\`: \`\`\`python from flyteplugins.dask import Dask, Scheduler, WorkerGroup dask\_config = Dask( scheduler=Scheduler(), workers=WorkerGroup(number\_of\_workers=4), ) dask\_env = flyte.TaskEnvironment( name="dask\_env", plugin\_config=dask\_config, image=image, ) \`\`\` ### \`Dask\` parameters | Parameter | Type | Description | |-----------|------|-------------| | \`scheduler\` | \`Scheduler\` | Scheduler pod configuration (defaults to \`Scheduler()\`) | | \`workers\` | \`WorkerGroup\` | Worker group configuration (defaults to \`WorkerGroup()\`) | ### \`Scheduler\` parameters | Parameter | Type | Description | |-----------|------|-------------| | \`image\` | \`str\` | Custom scheduler image (must include \`dask\[distributed\]\`) | | \`resources\` | \`Resources\` | Resource requests for the scheduler pod | ### \`WorkerGroup\` parameters | Parameter | Type | Description | |-----------|------|-------------| | \`number\_of\_workers\` | \`int\` | Number of worker pods (default: \`1\`) | | \`image\` | \`str\` | Custom worker image (must include \`dask\[distributed\]\`) | | \`resources\` | \`Resources\` | Resource requests per worker pod | > \[!NOTE\] > The scheduler and all workers should use the same Python environment to avoid serialization issues. ### Accessing the Dask client Inside a Dask task, create a \`distributed.Client()\` with no arguments. It automatically connects to the provisioned cluster: \`\`\`python from distributed import Client @dask\_env.task async def my\_dask\_task(n: int) -> list: client = Client() futures = client.map(lambda x: x + 1, range(n)) return client.gather(futures) \`\`\` ## Example \`\`\`python # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-dask", # "distributed" # \] # main = "hello\_dask\_nested" # params = "" # /// import asyncio import typing from distributed import Client from flyteplugins.dask import Dask, Scheduler, WorkerGroup import flyte.remote import flyte.storage from flyte import Resources image = flyte.Image.from\_debian\_base(python\_version=(3, 12)).with\_pip\_packages("flyteplugins-dask") dask\_config = Dask( scheduler=Scheduler(), workers=WorkerGroup(number\_of\_workers=4), ) task\_env = flyte.TaskEnvironment( name="hello\_dask", resources=Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image ) dask\_env = flyte.TaskEnvironment( name="dask\_env", plugin\_config=dask\_config, image=image, resources=Resources(cpu="1", memory="1Gi"), depends\_on=\[task\_env\], ) @task\_env.task() async def hello\_dask(): await asyncio.sleep(5) print("Hello from the Dask task!") @dask\_env.task async def hello\_dask\_nested(n: int = 3) -> typing.List\[int\]: print("running dask task") t = asyncio.create\_task(hello\_dask()) client = Client() futures = client.map(lambda x: x + 1, range(n)) res = client.gather(futures) await t return res if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() r = flyte.run(hello\_dask\_nested) print(r.name) print(r.url) r.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/dask/dask\_example.py\* ## API reference See the \[Dask API reference\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/dask/\_index) for full details. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/databricks === # Databricks The Databricks plugin lets you run PySpark jobs on \[Databricks\](https://www.databricks.com/) clusters directly from Flyte tasks. You write normal PySpark code in a Flyte task, and the plugin submits it to Databricks via the \[Jobs API 2.1\](https://docs.databricks.com/api/workspace/jobs/submit). The connector handles job submission, polling, and cancellation. The plugin supports: - Running PySpark tasks on new or existing Databricks clusters - Full Spark configuration (driver/executor memory, cores, instances) - Databricks cluster auto-scaling - API token-based authentication ## Installation \`\`\`bash pip install flyteplugins-databricks \`\`\` This also installs \`flyteplugins-spark\` as a dependency, since the Databricks plugin extends the Spark plugin. ## Quick start Create a \`Databricks\` configuration and pass it as \`plugin\_config\` to a \`TaskEnvironment\`: \`\`\`python from flyteplugins.databricks import Databricks import flyte image = ( flyte.Image.from\_base("databricksruntime/standard:16.4-LTS") .clone(name="spark", registry="ghcr.io/flyteorg", extendable=True) .with\_env\_vars({"UV\_PYTHON": "/databricks/python3/bin/python"}) .with\_pip\_packages("flyteplugins-databricks", pre=True) ) databricks\_conf = Databricks( spark\_conf={ "spark.driver.memory": "2000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", }, executor\_path="/databricks/python3/bin/python", databricks\_conf={ "run\_name": "flyte databricks plugin", "new\_cluster": { "spark\_version": "13.3.x-scala2.12", "node\_type\_id": "m6i.large", "autoscale": {"min\_workers": 1, "max\_workers": 2}, }, "timeout\_seconds": 3600, "max\_retries": 1, }, databricks\_instance="myaccount.cloud.databricks.com", databricks\_token="DATABRICKS\_TOKEN", ) databricks\_env = flyte.TaskEnvironment( name="databricks\_env", resources=flyte.Resources(cpu=(1, 2), memory=("3000Mi", "5000Mi")), plugin\_config=databricks\_conf, image=image, ) \`\`\` Then use the environment to decorate your task: \`\`\`python @databricks\_env.task async def hello\_databricks() -> float: spark = flyte.ctx().data\["spark\_session"\] # Use spark as a normal SparkSession count = spark.sparkContext.parallelize(range(100)).count() return float(count) \`\`\` ## Configuration The \`Databricks\` config extends the \[Spark\](../spark/\_index) config with Databricks-specific fields. ### Spark fields (inherited) | Parameter | Type | Description | |-----------|------|-------------| | \`spark\_conf\` | \`Dict\[str, str\]\` | Spark configuration key-value pairs | | \`hadoop\_conf\` | \`Dict\[str, str\]\` | Hadoop configuration key-value pairs | | \`executor\_path\` | \`str\` | Path to the Python binary on the Databricks cluster (e.g., \`/databricks/python3/bin/python\`) | | \`applications\_path\` | \`str\` | Path to the main application file | ### Databricks-specific fields | Parameter | Type | Description | |-----------|------|-------------| | \`databricks\_conf\` | \`Dict\[str, Union\[str, dict\]\]\` | Databricks \[run-submit\](https://docs.databricks.com/api/workspace/jobs/submit) job configuration. Must contain either \`existing\_cluster\_id\` or \`new\_cluster\` | | \`databricks\_instance\` | \`str\` | Your workspace domain (e.g., \`myaccount.cloud.databricks.com\`). Can also be set via the \`FLYTE\_DATABRICKS\_INSTANCE\` env var on the connector | | \`databricks\_token\` | \`str\` | Name of the Flyte secret containing the Databricks API token | ### \`databricks\_conf\` structure The \`databricks\_conf\` dict maps to the Databricks run-submit API payload. Key fields: | Field | Description | |-------|-------------| | \`new\_cluster\` | Cluster spec with \`spark\_version\`, \`node\_type\_id\`, \`autoscale\`, etc. | | \`existing\_cluster\_id\` | ID of an existing cluster to use instead of creating a new one | | \`run\_name\` | Display name in the Databricks UI | | \`timeout\_seconds\` | Maximum job duration | | \`max\_retries\` | Number of retries before marking the job as failed | The connector automatically injects the Docker image, Spark configuration, and environment variables from the task container into the cluster spec. ## Authentication Store your Databricks API token as a Flyte secret. The \`databricks\_token\` parameter specifies the secret name: \`\`\`python databricks\_conf = Databricks( # ... databricks\_token="DATABRICKS\_TOKEN", ) \`\`\` ## Accessing the Spark session Inside a Databricks task, the \`SparkSession\` is available through the task context, just like the \[Spark plugin\](../spark/\_index): \`\`\`python @databricks\_env.task async def my\_databricks\_task() -> float: spark = flyte.ctx().data\["spark\_session"\] df = spark.read.parquet("s3://my-bucket/data.parquet") return float(df.count()) \`\`\` ## API reference See the \[Databricks API reference\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/databricks/\_index) for full details. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/gemini === # Gemini The Gemini plugin lets you build agentic workflows with \[Gemini\](https://ai.google.dev/) on Flyte. It provides a \`function\_tool\` decorator that wraps Flyte tasks as tools that Gemini can call, and a \`run\_agent\` function that drives the agent conversation loop. When Gemini calls a tool, the call executes as a Flyte task with full observability, retries, and caching. Gemini's native parallel function calling is supported: multiple tool calls in a single turn are all dispatched and their results bundled into one response. ## Installation \`\`\`bash pip install flyteplugins-gemini \`\`\` Requires \`google-genai >= 1.0.0\`. ## Quick start \`\`\`python import flyte from flyteplugins.gemini import function\_tool, run\_agent env = flyte.TaskEnvironment( name="gemini-agent", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="gemini\_agent"), secrets=flyte.Secret("google\_api\_key", as\_env\_var="GOOGLE\_API\_KEY"), ) @function\_tool @env.task async def get\_weather(city: str) -> str: """Get the current weather for a city.""" return f"The weather in {city} is sunny, 72F" @env.task async def main(prompt: str) -> str: tools = \[get\_weather\] return await run\_agent(prompt=prompt, tools=tools) \`\`\` ## API ### \`function\_tool\` Converts a Flyte task, \`@flyte.trace\`-decorated function, or plain callable into a tool that Gemini can invoke. \`\`\`python @function\_tool @env.task async def my\_tool(param: str) -> str: """Tool description sent to Gemini.""" ... \`\`\` Can also be called with optional overrides: \`\`\`python @function\_tool(name="custom\_name", description="Custom description") @env.task async def my\_tool(param: str) -> str: ... \`\`\` Parameters: | Parameter | Type | Description | |-----------|------|-------------| | \`func\` | callable | The function to wrap | | \`name\` | \`str\` | Override the tool name (defaults to the function name) | | \`description\` | \`str\` | Override the tool description (defaults to the docstring) | > \[!NOTE\] > The docstring on each \`@function\_tool\` task is sent to Gemini as the tool description. Write clear, concise docstrings. ### \`Agent\` A dataclass for bundling agent configuration: \`\`\`python from flyteplugins.gemini import Agent agent = Agent( name="my-agent", instructions="You are a helpful assistant.", model="gemini-2.5-flash", tools=\[get\_weather\], max\_output\_tokens=8192, max\_iterations=10, ) \`\`\` | Field | Type | Default | Description | |-------|------|---------|-------------| | \`name\` | \`str\` | \`"assistant"\` | Agent name | | \`instructions\` | \`str\` | \`"You are a helpful assistant."\` | System prompt | | \`model\` | \`str\` | \`"gemini-2.5-flash"\` | Gemini model ID | | \`tools\` | \`list\[FunctionTool\]\` | \`\[\]\` | Tools available to the agent | | \`max\_output\_tokens\` | \`int\` | \`8192\` | Maximum tokens per response | | \`max\_iterations\` | \`int\` | \`10\` | Maximum tool-call loop iterations | ### \`run\_agent\` Runs a Gemini conversation loop, dispatching tool calls to Flyte tasks until Gemini returns a final response. \`\`\`python result = await run\_agent( prompt="What's the weather in Tokyo?", tools=\[get\_weather\], model="gemini-2.5-flash", ) \`\`\` You can also pass an \`Agent\` object: \`\`\`python result = await run\_agent(prompt="What's the weather?", agent=agent) \`\`\` | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | \`prompt\` | \`str\` | required | User message | | \`tools\` | \`list\[FunctionTool\]\` | \`None\` | Tools available to the agent | | \`agent\` | \`Agent\` | \`None\` | Agent config (overrides individual params) | | \`model\` | \`str\` | \`"gemini-2.5-flash"\` | Gemini model ID | | \`system\` | \`str\` | \`None\` | System prompt | | \`max\_output\_tokens\` | \`int\` | \`8192\` | Maximum tokens per response | | \`max\_iterations\` | \`int\` | \`10\` | Maximum iterations (prevents infinite loops) | | \`api\_key\` | \`str\` | \`None\` | API key (falls back to \`GOOGLE\_API\_KEY\` env var) | ## Secrets Store your Google API key as a Flyte secret and expose it as an environment variable: \`\`\`python secrets=flyte.Secret("google\_api\_key", as\_env\_var="GOOGLE\_API\_KEY") \`\`\` ## API reference See the \[Gemini API reference\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/gemini/\_index) for full details. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/hydra === # Hydra \[Hydra\](https://hydra.cc) is a framework for composing and overriding configuration trees from YAML files, dataclasses and the command line. The \`flyteplugins-hydra\` plugin makes Hydra a first-class submission layer for Flyte, so you can compose a config exactly as you would in any other Hydra app and have each composed run executed as a Flyte task, locally or as a remote execution on a Flyte cluster. The plugin offers three complementary entry points that share a single launcher implementation: | Entry point | Use it when | | ---------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------- | | \`hydra/launcher=flyte\` (Hydra Launcher plugin) | You already have a \`@hydra.main\` script and want standard Hydra CLI ergonomics, including \`--multirun\` and custom sweepers. | | \`flyte hydra run\` (Flyte CLI extension) | You want a Flyte-style CLI that imports a task from a Python file and composes a Hydra config without requiring a \`@hydra.main\` wrapper. | | \`hydra\_run\` / \`hydra\_sweep\` (Python SDK) | You want to submit runs directly from Python -- notebooks, tests, examples or another orchestration script. | All three paths converge on the same \`FlyteLauncher\`. ## Installation \`\`\`bash pip install flyteplugins-hydra \`\`\` The plugin depends on \`flyteplugins-omegaconf\`, which is installed automatically and provides the \`DictConfig\`/\`ListConfig\` type transformers that allow Hydra-composed configs to flow into Flyte tasks. Both packages must be available in the same environment as \`flyte\`. If you call \`apply\_task\_env\` for child tasks (see \*\*Hydra > Task environment overrides\*\*), include \`flyteplugins-hydra\` in the task image as well. ## Requirements on tasks Every task launched through this plugin must accept an OmegaConf \`DictConfig\` input. Any other parameters are passed through as ordinary task arguments. \`\`\`python{hl\_lines=\[1, 5\]} from omegaconf import DictConfig @env.task async def pipeline(cfg: DictConfig, dataset: str) -> float: ... \`\`\` The plugin auto-detects the \`DictConfig\` parameter name. If your parameter is \`cfg\`, app-level overrides are passed through \`--cfg\` on the CLI; if it is \`config\`, they are passed through \`--config\`; and so on. ## A walkthrough config The examples in this page assume a small project layout: \`\`\` project/ ├── train.py └── conf/ ├── training.yaml ├── model/ │ ├── resnet.yaml │ └── vit.yaml ├── optimizer/ │ ├── adam.yaml │ └── sgd.yaml └── task\_env/ ├── a100.yaml └── prebuilt\_image.yaml \`\`\` \`conf/training.yaml\`: \`\`\`yaml defaults: - optimizer: adam - model: resnet - \_self\_ data: path: s3://my-bucket/imagenet dataset: imagenet training: epochs: 30 batch\_size: 64 \`\`\` \`train.py\` (abbreviated): \`\`\`python import flyte from omegaconf import DictConfig from flyteplugins.hydra import apply\_task\_env env = flyte.TaskEnvironment(name="training", image=...) @env.task async def preprocess(cfg: DictConfig) -> flyte.io.Dir: ... @env.task async def train\_model(cfg: DictConfig, data: flyte.io.Dir) -> tuple\[flyte.io.Dir, float\]: ... @env.task async def pipeline(cfg: DictConfig, dataset: str) -> float: data = await preprocess(cfg) train\_task = apply\_task\_env(train\_model, cfg) \_, val\_loss = await train\_task(cfg, data) return val\_loss \`\`\` The same \`pipeline\` task is the target of every example below. > \*\*📝 Note\*\* > > \`config\_path\` is resolved relative to the current working directory. If you submit runs from a directory other than \`project/\`, pass an absolute path (or an absolute path on the CLI via \`--config-path /abs/path/to/conf\`). For structured-config-only setups (no YAML files), omit \`config\_path\` / \`--config-path\` entirely. ## Execution mode Remote execution is the default. Every entry point exposes an explicit knob: | Surface | Local | Remote | | ---------------------- | --------------------------- | -------------------------------------- | | \`@hydra.main\` launcher | \`hydra.launcher.mode=local\` | \`hydra.launcher.mode=remote\` (default) | | \`flyte hydra run\` | \`--local\` | \`--mode remote\` (default) | | Python SDK | \`mode="local"\` | \`mode="remote"\` (default) | For the \`@hydra.main\` launcher, the default applies as soon as \`hydra/launcher=flyte\` is selected. Remote runs print the Flyte run URL immediately after submission, before any waiting. By default the plugin then waits for every submitted run to reach a terminal phase, capped at 32 worker threads. To tune or disable waiting: | Surface | Tune wait threads | Fire and forget | | ---------------------- | ------------------------------------ | --------------------------- | | \`@hydra.main\` launcher | \`hydra.launcher.wait\_max\_workers=64\` | \`hydra.launcher.wait=false\` | | \`flyte hydra run\` | \`--wait-max-workers 64\` | \`--no-wait\` | | Python SDK | \`wait\_max\_workers=64\` | \`wait=False\` | For a sweep, every job is submitted first, and then the plugin waits for all runs concurrently. Submission is not blocked by earlier runs reaching a terminal phase. ## Hydra launcher (\`@hydra.main\` scripts) Use this path when your script already has a \`@hydra.main\` entry point. Selecting \`hydra/launcher=flyte\` swaps Hydra's built-in \`BasicLauncher\` for \`FlyteLauncher\`. Single remote run: \`\`\`bash python train.py hydra/launcher=flyte hydra.launcher.mode=remote \`\`\` Single local run: \`\`\`bash python train.py hydra/launcher=flyte hydra.launcher.mode=local \`\`\` Remote grid sweep submission: Each comma-separated value expands into a separate Flyte execution; six executions in this example: \`\`\`bash{hl\_lines=\[4\]} python train.py --multirun \\ hydra/launcher=flyte hydra.launcher.mode=remote \\ hydra.launcher.wait\_max\_workers=64 \\ optimizer.lr=0.001,0.01,0.1 training.epochs=10,20 \`\`\` Fire-and-forget sweep submission: \`\`\`bash{hl\_lines=\[2\]} python train.py --multirun \\ hydra/launcher=flyte hydra.launcher.wait=false \\ optimizer.lr=0.001,0.01,0.1 \`\`\` Custom sweepers (Optuna) work exactly as they do with the BasicLauncher. Selecting \`hydra/sweeper=...\` activates the sweeper and \`FlyteLauncher\` runs each trial as a Flyte execution: \`\`\`bash{hl\_lines=\["3-5"\]} python train.py --multirun \\ hydra/launcher=flyte hydra.launcher.mode=remote \\ hydra/sweeper=optuna hydra.sweeper.n\_trials=20 \\ hydra.sweeper.n\_jobs=4 \\ "optimizer.lr=interval(1e-4,1e-1)" \`\`\` Inside \`@hydra.main\`, the standard pattern is: \`\`\`python{hl\_lines=\[7\]} import flyte import hydra from omegaconf import DictConfig from flyteplugins.hydra import apply\_task\_env @hydra.main(version\_base=None, config\_path="conf", config\_name="training") def main(cfg: DictConfig): flyte.init\_from\_config() entry\_task = apply\_task\_env(pipeline, cfg) return flyte.run(entry\_task, cfg=cfg, dataset=cfg.data.dataset) if \_\_name\_\_ == "\_\_main\_\_": main() \`\`\` ## Python SDK \`hydra\_run\` composes one config and runs the task once. \`hydra\_sweep\` expands sweep overrides and runs the task once per combination. ### Single run \`\`\`python{hl\_lines=\[1, 3, 7\]} from flyteplugins.hydra import hydra\_run run = hydra\_run( pipeline, config\_path="conf", config\_name="training", overrides=\["optimizer.lr=0.01"\], dataset="s3://my-bucket/imagenet", mode="remote", wait=True, wait\_max\_workers=64, ) \`\`\` For a remote run with \`wait=True\`, the return value is a wrapper exposing both \`run.url\` and \`run.value\` (the resolved task output). The wrapper is \`float()\`-castable so Hydra sweepers such as Optuna can consume scalar objectives directly. With \`wait=False\`, the return value is the underlying \`flyte.remote.Run\`. ### Grid sweep \`\`\`python{hl\_lines=\[7\]} from flyteplugins.hydra import hydra\_sweep runs = hydra\_sweep( pipeline, config\_path="conf", config\_name="training", overrides=\["optimizer.lr=0.001,0.01,0.1", "training.epochs=10,20"\], dataset="s3://my-bucket/imagenet", mode="remote", ) \`\`\` Six executions are submitted (3 × 2). \`runs\` is a list aligned with the Cartesian-product order Hydra's \`BasicSweeper\` produces. ### Custom sweepers Custom sweeper plugins are activated by passing their selection in \`overrides\`: \`\`\`python{hl\_lines=\["5-10"\]} runs = hydra\_sweep( pipeline, config\_path="conf", config\_name="training", overrides=\[ "hydra/sweeper=optuna", "hydra.sweeper.n\_trials=20", "hydra.sweeper.n\_jobs=4", "optimizer.lr=interval(1e-4,1e-1)", \], dataset="s3://my-bucket/imagenet", mode="remote", ) \`\`\` Whenever an override starts with \`hydra/\`, the plugin invokes the full Hydra runtime so plugin discovery (sweepers, launchers, callbacks) can run. Pure value overrides on the \`hydra.\*\` namespace (for example \`hydra.run.dir=...\`) do not need the full runtime and are applied per-job by the launcher directly. ### Forwarding \`flyte.with\_runcontext\` options Use \`run\_options\` to pass Flyte runtime options through to every job: \`\`\`python{hl\_lines=\["8-14"\]} runs = hydra\_sweep( pipeline, config\_path="conf", config\_name="training", overrides=\["optimizer.lr=0.001,0.01,0.1"\], dataset="s3://my-bucket/imagenet", mode="remote", run\_options={ "name": "my-training-sweep", "service\_account": "default", "copy\_style": "all", "raw\_data\_path": "s3://my-bucket/raw-data", "debug": True, }, ) \`\`\` ## Flyte CLI (\`flyte hydra run\`) \`flyte hydra run\` is registered through the \`flyte.plugins.cli.commands\` entry point. It loads a task from a Python file, composes a Hydra config, and runs the task without requiring the script to have its own \`@hydra.main\` function. It also inherits the relevant flags from \`flyte run\` (\`--project\`, \`--domain\`, \`--image\`, \`--name\`, \`--service-account\`, \`--raw-data-path\`, \`--copy-style\`, \`--debug\`, \`--local\`, \`--follow\`). ### Single run Remote (default): \`\`\`bash flyte hydra run --config-path conf --config-name training \\ train.py pipeline --dataset s3://my-bucket/imagenet \`\`\` Forced local: \`\`\`bash{hl\_lines=\[1\]} flyte hydra run --local --config-path conf --config-name training \\ train.py pipeline --dataset s3://my-bucket/imagenet \`\`\` ### Grid sweep \`\`\`bash{hl\_lines=\[4\]} flyte hydra run --multirun --config-path conf --config-name training \\ --wait-max-workers 64 \\ train.py pipeline --dataset s3://my-bucket/imagenet \\ --cfg "optimizer.lr=0.001,0.01,0.1" --cfg "training.epochs=10,20" \`\`\` ### App-level vs Hydra-namespace overrides The CLI keeps app-level overrides separate from Hydra runtime overrides so they do not collide with ordinary Flyte task arguments. App-level overrides target the composed config and are passed through the \*\*task's \`DictConfig\` parameter name\*\*. For \`pipeline(cfg: DictConfig, ...)\`, use \`--cfg\`. For \`pipeline\_with\_config(config: DictConfig, ...)\`, use \`--config\`: \`\`\`bash{hl\_lines=\["3-4", 8\]} flyte hydra run --config-path conf --config-name training \\ train.py pipeline \\ --cfg optimizer.lr=0.01 \\ --cfg training.epochs=20 flyte hydra run --config-path conf --config-name training \\ train.py pipeline\_with\_config \\ --config optimizer.lr=0.01 \`\`\` Hydra runtime overrides: Anything in the \`hydra.\*\` or \`hydra/\*\` namespace go through \`--hydra-override\`: \`\`\`bash{hl\_lines=\[3, 4\]} flyte hydra run --config-path conf --config-name training \\ train.py pipeline \\ --hydra-override hydra.run.dir=./outputs/exp1 \\ --hydra-override hydra/launcher=flyte \`\`\` Custom sweepers combine the two: \`\`\`bash{hl\_lines=\["3-7"\]} flyte hydra run --multirun --config-path conf --config-name training \\ train.py pipeline --dataset s3://my-bucket/imagenet \\ --hydra-override hydra/sweeper=optuna \\ --hydra-override hydra.sweeper.n\_trials=20 \\ --hydra-override hydra.sweeper.n\_jobs=4 \\ --cfg "optimizer.lr=interval(1e-4,1e-1)" \\ --cfg "training.epochs=choice(10,20,50)" \`\`\` ### \`--follow\` and \`--no-wait\` \`--follow\` streams logs from the launched run after submission; it implies waiting and cannot be combined with \`--no-wait\`. \`--no-wait\` returns immediately after submission and skips log streaming. ### Shell completion Install Click's completion hook for the \`flyte\` executable. For zsh: \`\`\`zsh eval "$(\_FLYTE\_COMPLETE=zsh\_source flyte)" \`\`\` For bash: \`\`\`bash eval "$(\_FLYTE\_COMPLETE=bash\_source flyte)" \`\`\` Once installed, \`flyte hydra run\` adds Hydra-aware completion after \`SCRIPT TASK\_NAME\`. The command imports the script, inspects the task signature, and suggests: - The app override flag matching the task's \`DictConfig\` parameter (\`--cfg\`, \`--config\`, ...). - Override values for that flag and \`--hydra-override\` via Hydra's own completion engine, including config keys, config-group selections and sweep functions. \`\`\`bash{hl\_lines=\["2-3", "6-7"\]} flyte hydra run --config-path conf --config-name training \\ train.py pipeline --cfg optimizer. # suggests optimizer.lr=, optimizer.weight\_decay=, ... flyte hydra run --config-path conf --config-name training \\ train.py pipeline --hydra-override hydra/launcher= # suggests hydra launcher choices \`\`\` Because completion has to import the target script, keep task definitions and \`ConfigStore\` registration import-safe, and avoid expensive top-level work in scripts you reach via \`flyte hydra run\`. !\[Auto Completion\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/hydra/auto\_complete.gif) ## Override grammar The override grammar is identical to standard Hydra; what differs is only how you pass the strings (positional in \`python train.py ...\`, list entries in \`overrides=\[...\]\`, repeated \`--cfg\`/\`--hydra-override\` on the Flyte CLI). | Form | Meaning | | ---------------------------------- | ---------------------------------------------------------------------------------------- | | \`optimizer.lr=0.01\` | Set an existing key. | | \`optimizer=sgd\` | Select a config group (replaces the \`optimizer\` subtree with \`conf/optimizer/sgd.yaml\`). | | \`+task\_env=a100\` | Append a config group whose key is not currently in the config. | | \`+training.grad\_clip=1.0\` | Append a key that does not exist. | | \`++optimizer.lr=0.05\` | Force-set a key, creating it if missing and overriding strict-schema errors. | | \`~training.warmup\_steps\` | Delete a key from the composed config. | | \`optimizer.lr=0.001,0.01,0.1\` | Sweep value (with \`--multirun\`); expanded into one job per element. | | \`optimizer.lr=interval(1e-4,1e-1)\` | Continuous sweep range; consumed by samplers like Optuna. | | \`optimizer=choice(adam,sgd)\` | Categorical sweep; consumed by samplers. | | \`hydra.run.dir=./outputs/exp1\` | Hydra-namespace value override (single run output dir). | | \`hydra.sweep.dir=./outputs/sweep1\` | Hydra-namespace sweep output dir. | | \`hydra/sweeper=optuna\` | Hydra-namespace config group selection (activates the Optuna sweeper plugin). | ## Sweeps ### Grid sweeps (BasicSweeper) Comma-separated overrides expand into a Cartesian product. The plugin uses Hydra's \`BasicSweeper\` to expand them, then submits one Flyte execution per combination. \`\`\`python{hl\_lines=\[1, 4, 7\]} from flyteplugins.hydra import hydra\_sweep runs = hydra\_sweep( pipeline, config\_path="conf", config\_name="training", overrides=\["model=resnet,vit", "optimizer.lr=0.001,0.01,0.1"\], dataset="s3://my-bucket/imagenet", mode="remote", ) # 6 executions \`\`\` \`\`\`bash{hl\_lines=\[3\]} flyte hydra run --multirun --config-path conf --config-name training \\ train.py pipeline --dataset s3://my-bucket/imagenet \\ --cfg "model=resnet,vit" --cfg "optimizer.lr=0.001,0.01,0.1" \`\`\` Hardware presets can sweep alongside hyperparameters: \`\`\`bash{hl\_lines=\[3\]} flyte hydra run --multirun --config-path conf --config-name training \\ train.py pipeline --dataset s3://my-bucket/imagenet \\ --cfg "+task\_env=a10g,a100" --cfg "optimizer.lr=0.001,0.01,0.1" \`\`\` ### Bayesian / TPE sweeps (Optuna) Install the sweeper, then activate it via \`hydra/sweeper=optuna\`. Continuous parameters use \`interval(...)\`; categorical parameters use \`choice(...)\`. \`\`\`bash pip install hydra-optuna-sweeper \`\`\` \`\`\`bash{hl\_lines=\["3-8"\]} flyte hydra run --multirun --config-path conf --config-name training \\ train.py pipeline --dataset s3://my-bucket/imagenet \\ --hydra-override "hydra/sweeper=optuna" \\ --hydra-override "hydra.sweeper.n\_trials=30" \\ --hydra-override "hydra.sweeper.n\_jobs=5" \\ --cfg "optimizer.lr=interval(1e-4,1e-1)" \\ --cfg "optimizer.weight\_decay=interval(1e-6,1e-2)" \\ --cfg "model=choice(resnet,vit)" \`\`\` When \`wait=True\`, each remote run's wrapped result exposes the task output as a float (via \`\_\_float\_\_\`), so Optuna can use it directly as the trial objective. With \`wait=False\`, the sweeper sees the run URL but cannot read objective values; use this only for fire-and-forget submission. Other sweepers that respect Hydra's plugin protocol are activated the same way: install the package, select \`hydra/sweeper=\`, and set the sweeper's parameters under \`hydra.sweeper.\*\`. ### Sweep output directories Hydra-namespace overrides redirect where Hydra writes per-job logs and config snapshots: \`\`\`bash{hl\_lines=\[3, 4\]} flyte hydra run --multirun --config-path conf --config-name training \\ train.py pipeline --dataset s3://my-bucket/imagenet \\ --hydra-override "hydra.sweep.dir=./outputs/sweep1" \\ --hydra-override "hydra.sweep.subdir=\\${hydra.job.num}" \\ --cfg "optimizer.lr=0.001,0.01,0.1" \`\`\` ## Task environment overrides Hydra is good at composing flat YAML; Flyte tasks need richer settings such as resources and container images. The plugin reserves a config key named \`task\_env\` by default that maps task names to \`task.override\` kwargs. \`\`\`yaml task\_env: pipeline: resources: cpu: "2" memory: 8Gi train\_model: resources: cpu: "16" memory: 64Gi gpu: "A100:1" \`\`\` When the plugin launches a task, it looks up \`task\_env\[\]\` (\`pipeline\` in this example) and applies the values via \`task.override(...)\`. Resource mappings are converted into \`flyte.Resources(\*\*values)\` automatically. ### Prebuilt images To run a task in a prebuilt container image, set \`image\` (and optionally \`primary\_container\_name\`): \`\`\`yaml{hl\_lines=\[3\]} task\_env: pipeline: image: ghcr.io/acme/flyte-training:latest primary\_container\_name: main resources: cpu: "4" memory: 16Gi \`\`\` \`task.override\` does not accept \`image\` directly. The task image is part of the task definition. Instead, the plugin lowers the override to a \`flyte.PodTemplate\` whose primary container uses the requested image: - If the task has no inline pod template, a new one is created. - If the task already has an inline \`flyte.PodTemplate\`, the plugin deep-copies it and sets only the image on the primary container. - If the task references a pod template by name (a string), the plugin raises an error. You must patch a string-named template by editing it in cluster config rather than at submission time. ### Applying overrides to child tasks The launcher only controls the entry task it submits. Child tasks called from within the entry task are not patched automatically. Use \`apply\_task\_env\` to apply the same \`resources\`/\`image\` handling to a child task before invoking it: \`\`\`python{hl\_lines=\[1, 7\]} from flyteplugins.hydra import apply\_task\_env @env.task async def pipeline(cfg: DictConfig, dataset: str) -> float: data = await preprocess(cfg) train\_task = apply\_task\_env(train\_model, cfg) \_, val\_loss = await train\_task(cfg, data) return val\_loss \`\`\` This keeps the override knobs in YAML/CLI surfaces while leaving each task in control of which children it patches. ### Renaming the task-env key If your config uses a different name for the task-env subtree, pass it explicitly: \`\`\`python hydra\_run(..., task\_env\_key="task\_environment") \`\`\` \`\`\`bash flyte hydra run --task-env-key task\_environment ... \`\`\` ### What \`task\_env\` should not model The YAML schema intentionally omits the full Kubernetes \`V1PodSpec\`. Keep advanced pod configuration (volumes, init containers, node selectors, etc.) in Python task/environment code where you have a real type. Use Hydra \`task\_env\` presets for the common knobs only: image, primary container name and resources. ## Structured configs (without YAML) Structured configs work with this plugin as long as they are registered before the launcher composes the config. \`flyte hydra run\` imports the script first, so top-level \`ConfigStore.instance().store(...)\` calls run before composition. \`\`\`python{hl\_lines=\[17\]} from dataclasses import dataclass, field from hydra.core.config\_store import ConfigStore from omegaconf import DictConfig @dataclass class TrainingConf: epochs: int = 30 batch\_size: int = 64 @dataclass class RootConf: training: TrainingConf = field(default\_factory=TrainingConf) ConfigStore.instance().store(name="structured\_training", node=RootConf) \`\`\` Run a fully-structured config without YAML: \`\`\`bash{hl\_lines=\[1\]} flyte hydra run --config-name structured\_training \\ train.py pipeline --dataset s3://my-bucket/imagenet \`\`\` The same config also works through \`@hydra.main\`: \`\`\`bash python train.py --config-name structured\_training \`\`\` If the structured config still references YAML config groups, keep \`--config-path conf\`. If everything is registered in \`ConfigStore\`, omit \`--config-path\`. > \*\*⚠️ Warning\*\* > > Do not register structured configs only inside \`if \_\_name\_\_ == "\_\_main\_\_":\` or inside the \`@hydra.main\` function body. \`flyte hydra run\` and shell completion inspect the script at import time, before either of those blocks runs, and registrations placed there will not be visible. Structured configs sweep just like YAML configs: \`\`\`python{hl\_lines=\[4, 5\]} runs = hydra\_sweep( pipeline, config\_path=None, config\_name="structured\_training", overrides=\["training.epochs=10,20", "training.batch\_size=32,64"\], dataset="s3://my-bucket/imagenet", mode="remote", ) \`\`\` === PAGE: https://www.union.ai/docs/v2/flyte/integrations/mlflow === # MLflow The MLflow plugin integrates \[MLflow\](https://mlflow.org/) experiment tracking with Flyte. It provides a \`@mlflow\_run\` decorator that automatically manages MLflow runs within Flyte tasks, with support for autologging, parent-child run sharing, distributed training, and auto-generated UI links. The decorator works with both sync and async tasks. ## Installation \`\`\`bash pip install flyteplugins-mlflow \`\`\` Requires \`mlflow\` and \`flyte\`. ## Quick start \`\`\`python{hl\_lines=\[3, 9, "13-16", 22\]} import flyte import mlflow from flyteplugins.mlflow import mlflow\_run, get\_mlflow\_run env = flyte.TaskEnvironment( name="mlflow-tracking", resources=flyte.Resources(cpu=1, memory="500Mi"), image=flyte.Image.from\_debian\_base(name="mlflow\_example").with\_pip\_packages( "flyteplugins-mlflow" ), ) @mlflow\_run( tracking\_uri="http://localhost:5000", experiment\_name="my-experiment", ) @env.task async def train\_model(learning\_rate: float) -> str: mlflow.log\_param("lr", learning\_rate) mlflow.log\_metric("loss", 0.42) run = get\_mlflow\_run() return run.info.run\_id \`\`\` !\[Link\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/link.png) !\[Mlflow UI\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/mlflow\_dashboard.png) > \[!NOTE\] > \`@mlflow\_run\` must be the outermost decorator, before \`@env.task\`: > > \`\`\`python{hl\_lines=\["1-2"\]} > @mlflow\_run # outermost > @env.task # innermost > async def my\_task(): ... > \`\`\` ## Autologging Enable MLflow's autologging to automatically capture parameters, metrics, and models without manual \`mlflow.log\_\*\` calls. ### Generic autologging \`\`\`python{hl\_lines=\[1\]} @mlflow\_run(autolog=True) @env.task async def train(): from sklearn.linear\_model import LogisticRegression model = LogisticRegression() model.fit(X, y) # Parameters, metrics, and model are logged automatically \`\`\` ### Framework-specific autologging Pass \`framework\` to use a framework-specific autolog implementation: \`\`\`python{hl\_lines=\[3\]} @mlflow\_run( autolog=True, framework="sklearn", log\_models=True, log\_datasets=False, ) @env.task async def train\_sklearn(): from sklearn.ensemble import RandomForestClassifier model = RandomForestClassifier(n\_estimators=100) model.fit(X\_train, y\_train) \`\`\` Supported frameworks include any framework with an \`mlflow.{framework}.autolog()\` function. You can find the full list of supported frameworks \[here\](https://mlflow.org/docs/latest/ml/tracking/autolog/#supported-libraries). You can pass additional autolog parameters via \`autolog\_kwargs\`: \`\`\`python{hl\_lines=\[4\]} @mlflow\_run( autolog=True, framework="pytorch", autolog\_kwargs={"log\_every\_n\_epoch": 5}, ) @env.task async def train\_pytorch(): ... \`\`\` !\[Autolog\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/autolog.png) ## Run modes The \`run\_mode\` parameter controls how MLflow runs are created and shared across tasks: | Mode | Behavior | | ------------------ | --------------------------------------------------------------------- | | \`"auto"\` (default) | Reuse the parent's run if one exists, otherwise create a new run | | \`"new"\` | Always create a new independent run | | \`"nested"\` | Create a new run nested under the parent via \`mlflow.parentRunId\` tag | ### Sharing a run across tasks With \`run\_mode="auto"\` (the default), child tasks reuse the parent's MLflow run: \`\`\`python{hl\_lines=\[1, 5, 7\]} @mlflow\_run @env.task async def parent\_task(): mlflow.log\_param("stage", "parent") await child\_task() # Shares the same MLflow run @mlflow\_run @env.task async def child\_task(): mlflow.log\_metric("child\_metric", 1.0) # Logged to the parent's run \`\`\` ### Creating independent runs Use \`run\_mode="new"\` when a task should always create its own top-level MLflow run, completely independent of any parent: \`\`\`python{hl\_lines=\[1\]} @mlflow\_run(run\_mode="new") @env.task async def standalone\_experiment(): mlflow.log\_param("experiment\_type", "baseline") mlflow.log\_metric("accuracy", 0.95) \`\`\` ### Nested runs Use \`run\_mode="nested"\` to create a child run that appears under the parent in the MLflow UI. This works across processes and containers via the \`mlflow.parentRunId\` tag. !\[Nested runs\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/mlflow\_hpo.png) This is the recommended pattern for hyperparameter optimization, where each trial should be tracked as a child of the parent study run: \`\`\`python{hl\_lines=\[1, 2, 15, "22-25"\]} from flyteplugins.mlflow import Mlflow @mlflow\_run(run\_mode="nested") @env.task(links=\[Mlflow()\]) async def run\_trial(trial\_number: int, n\_estimators: int, max\_depth: int) -> float: """Each trial creates a nested MLflow run under the parent.""" mlflow.log\_params({"n\_estimators": n\_estimators, "max\_depth": max\_depth}) mlflow.log\_param("trial\_number", trial\_number) model = RandomForestRegressor(n\_estimators=n\_estimators, max\_depth=max\_depth) model.fit(X\_train, y\_train) rmse = float(np.sqrt(mean\_squared\_error(y\_val, model.predict(X\_val)))) mlflow.log\_metric("rmse", rmse) return rmse @mlflow\_run @env.task async def hpo\_search(n\_trials: int = 30) -> str: """Parent run tracks the overall study.""" run = get\_mlflow\_run() mlflow.log\_param("n\_trials", n\_trials) # Run trials in parallel — each gets a nested MLflow run rmses = await asyncio.gather( \*(run\_trial(trial\_number=i, \*\*params) for i, params in enumerate(trial\_params)) ) mlflow.log\_metric("best\_rmse", min(rmses)) return run.info.run\_id \`\`\` !\[HPO\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/hpo.png) ## Workflow-level configuration Use \`mlflow\_config()\` with \`flyte.with\_runcontext()\` to set MLflow configuration for an entire workflow. All \`@mlflow\_run\`-decorated tasks in the workflow inherit these settings: \`\`\`python{hl\_lines=\[1, "4-8"\]} from flyteplugins.mlflow import mlflow\_config r = flyte.with\_runcontext( custom\_context=mlflow\_config( tracking\_uri="http://localhost:5000", experiment\_id="846992856162999", tags={"team": "ml"}, ) ).run(train\_model, learning\_rate=0.001) \`\`\` This eliminates the need to repeat \`tracking\_uri\` and experiment settings on every \`@mlflow\_run\` decorator. ### Per-task overrides Use \`mlflow\_config()\` as a context manager inside a task to override configuration for specific child tasks: \`\`\`python{hl\_lines=\[6\]} @mlflow\_run @env.task async def parent\_task(): await shared\_child() # Inherits parent config with mlflow\_config(run\_mode="new", tags={"role": "independent"}): await independent\_child() # Gets its own run \`\`\` ### Configuration priority Settings are resolved in priority order: 1. Explicit \`@mlflow\_run\` decorator arguments 2. \`mlflow\_config()\` context configuration 3. Environment variables (for \`tracking\_uri\`) 4. MLflow defaults ## Distributed training In distributed training, only rank 0 logs to MLflow by default. The plugin detects rank automatically from the \`RANK\` environment variable: \`\`\`python{hl\_lines=\[1, "4-6"\]} @mlflow\_run @env.task async def distributed\_train(): # Only rank 0 creates an MLflow run and logs metrics. # Other ranks execute the task function directly without # creating an MLflow run or incurring any MLflow overhead. ... \`\`\` On non-rank-0 workers, no MLflow run is created and \`get\_mlflow\_run()\` returns \`None\`. The task function still executes normally — only the MLflow instrumentation is skipped. !\[Distributed training\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/mlflow/distributed\_training.png) You can also set rank explicitly: \`\`\`python{hl\_lines=\[1\]} @mlflow\_run(rank=0) @env.task async def train(): ... \`\`\` ## MLflow UI links The \`Mlflow\` link class displays links to the MLflow UI in the Flyte UI. Since the MLflow run is created inside the task at execution time, the run URL cannot be determined before the task starts. Links are only shown when a run URL is already available from context, either because a parent task created the run, or because an explicit URL is provided. The recommended pattern is for the parent task to create the MLflow run, and child tasks that inherit the run (via \`run\_mode="auto"\`) display the link to that run. For nested runs (\`run\_mode="nested"\`), children display a link to the parent run. ### Setup Set \`link\_host\` via \`mlflow\_config()\` and attach \`Mlflow()\` links to child tasks: \`\`\`python{hl\_lines=\[4, 17\]} from flyteplugins.mlflow import Mlflow, mlflow\_config @mlflow\_run @env.task(links=\[Mlflow()\]) async def child\_task(): ... # Link points to the parent's MLflow run @mlflow\_run @env.task async def parent\_task(): await child\_task() if \_\_name\_\_ == "\_\_main\_\_": r = flyte.with\_runcontext( custom\_context=mlflow\_config( tracking\_uri="http://localhost:5000", link\_host="http://localhost:5000", ) ).run(parent\_task) \`\`\` > \[!NOTE\] > \`Mlflow()\` is instantiated without a \`link\` argument because the URL is auto-generated at runtime. When the parent task creates an MLflow run, the plugin builds the URL from \`link\_host\` and the run's experiment/run IDs, then propagates it to child tasks via the Flyte context. Passing an explicit \`link\` would bypass this auto-generation. ### Custom URL templates The default link format is: \`\`\` {host}/#/experiments/{experiment\_id}/runs/{run\_id} \`\`\` For platforms like Databricks that use a different URL structure, provide a custom template: \`\`\`python{hl\_lines=\[3\]} mlflow\_config( link\_host="https://dbc-xxx.cloud.databricks.com", link\_template="{host}/ml/experiments/{experiment\_id}/runs/{run\_id}", ) \`\`\` ### Explicit links If you know the run URL ahead of time, you can set it directly: \`\`\`python{hl\_lines=\[1\]} @env.task(links=\[Mlflow(link="https://mlflow.example.com/#/experiments/1/runs/abc123")\]) async def my\_task(): ... \`\`\` ### Link behavior by run mode | Run mode | Link behavior | | ---------- | ---------------------------------------------------------------------------------------------- | | \`"auto"\` | Parent link propagates to child tasks sharing the run | | \`"new"\` | Parent link is cleared; no link is shown until the task's own run is available to its children | | \`"nested"\` | Parent link is kept and renamed to "MLflow (parent)" | ## Automatic Flyte tags When running inside Flyte, the plugin automatically tags MLflow runs with execution metadata: | Tag | Description | | ------------------- | ---------------- | | \`flyte.action\_name\` | Task action name | | \`flyte.run\_name\` | Flyte run name | | \`flyte.project\` | Flyte project | | \`flyte.domain\` | Flyte domain | These tags are merged with any user-provided tags. ## API reference ### \`mlflow\_run\` and \`mlflow\_config\` \`mlflow\_run\` is a decorator that manages MLflow runs for Flyte tasks. \`mlflow\_config\` creates workflow-level configuration or per-task overrides. Both accept the same core parameters: | Parameter | Type | Default | Description | | ----------------- | ---------------- | -------- | ----------------------------------------------------------------------------- | | \`run\_mode\` | \`str\` | \`"auto"\` | \`"auto"\`, \`"new"\`, or \`"nested"\` | | \`tracking\_uri\` | \`str\` | \`None\` | MLflow tracking server URL | | \`experiment\_name\` | \`str\` | \`None\` | MLflow experiment name (raises \`ValueError\` if combined with \`experiment\_id\`) | | \`experiment\_id\` | \`str\` | \`None\` | MLflow experiment ID (raises \`ValueError\` if combined with \`experiment\_name\`) | | \`run\_name\` | \`str\` | \`None\` | Human-readable run name (raises \`ValueError\` if combined with \`run\_id\`) | | \`run\_id\` | \`str\` | \`None\` | Explicit MLflow run ID (raises \`ValueError\` if combined with \`run\_name\`) | | \`tags\` | \`dict\[str, str\]\` | \`None\` | Tags for the run | | \`autolog\` | \`bool\` | \`False\` | Enable MLflow autologging | | \`framework\` | \`str\` | \`None\` | Framework for autolog (e.g. \`"sklearn"\`, \`"pytorch"\`) | | \`log\_models\` | \`bool\` | \`None\` | Log models automatically (requires \`autolog\`) | | \`log\_datasets\` | \`bool\` | \`None\` | Log datasets automatically (requires \`autolog\`) | | \`autolog\_kwargs\` | \`dict\` | \`None\` | Extra parameters for \`mlflow.autolog()\` | Additional keyword arguments are passed to \`mlflow.start\_run()\`. \`mlflow\_run\` also accepts: | Parameter | Type | Default | Description | | --------- | ----- | ------- | -------------------------------------------------------- | | \`rank\` | \`int\` | \`None\` | Process rank for distributed training (only rank 0 logs) | \`mlflow\_config\` also accepts: | Parameter | Type | Default | Description | | --------------- | ----- | ------- | --------------------------------------------------------------------------- | | \`link\_host\` | \`str\` | \`None\` | MLflow UI host for auto-generating links | | \`link\_template\` | \`str\` | \`None\` | Custom URL template (placeholders: \`{host}\`, \`{experiment\_id}\`, \`{run\_id}\`) | ### \`get\_mlflow\_run\` Returns the current \`mlflow.ActiveRun\` if within a \`@mlflow\_run\`-decorated task. Returns \`None\` otherwise. \`\`\`python from flyteplugins.mlflow import get\_mlflow\_run run = get\_mlflow\_run() if run: print(run.info.run\_id) \`\`\` ### \`get\_mlflow\_context\` Returns the current \`mlflow\_config\` settings from the Flyte context, or \`None\` if no MLflow configuration is set. Useful for inspecting the inherited configuration inside a task: \`\`\`python from flyteplugins.mlflow import get\_mlflow\_context @mlflow\_run @env.task async def my\_task(): config = get\_mlflow\_context() if config: print(config.tracking\_uri, config.experiment\_id) \`\`\` ### \`Mlflow\` Link class for displaying MLflow UI links in the Flyte console. | Field | Type | Default | Description | | ------ | ----- | ---------- | --------------------------------------- | | \`name\` | \`str\` | \`"MLflow"\` | Display name for the link | | \`link\` | \`str\` | \`""\` | Explicit URL (bypasses auto-generation) | === PAGE: https://www.union.ai/docs/v2/flyte/integrations/omegaconf === # OmegaConf \[OmegaConf\](https://omegaconf.readthedocs.io/) is a hierarchical configuration system used by many ML frameworks (and the foundation of \[Hydra\](../hydra/\_index)). The \`flyteplugins-omegaconf\` plugin makes OmegaConf's \`DictConfig\` and \`ListConfig\` first-class types in Flyte tasks, so you can pass entire configs like plain dicts, YAML files or dataclass-backed structured configs between tasks without flattening them into individual scalar arguments. The plugin enables: - \`DictConfig\` and \`ListConfig\` as native task input and output types - Round-tripping of structured configs (dataclass schemas) across task boundaries - Preservation of OmegaConf-specific values: \`MISSING\` sentinels, \`Enum\`s, \`pathlib.Path\`s, \`tuple\`s, and \`bytes\` - Resolved variable interpolations on the wire - A YAML-rendered Flyte report tab for human-readable config inspection ## Installation \`\`\`bash pip install flyteplugins-omegaconf \`\`\` Installing the package automatically registers \`DictConfig\` and \`ListConfig\` with Flyte's \`TypeEngine\`. No manual setup is required. If you are using the \[Hydra plugin\](../hydra/\_index), \`flyteplugins-omegaconf\` is installed as a transitive dependency. ## Quick start \`\`\`python{hl\_lines=\[2, "8-9", "14-17"\]} import flyte from omegaconf import DictConfig, OmegaConf env = flyte.TaskEnvironment(name="training", image=...) @env.task async def train(cfg: DictConfig) -> float: return run\_experiment(cfg.optimizer.lr, cfg.training.epochs) @env.task async def pipeline() -> float: cfg = OmegaConf.create( {"optimizer": {"lr": 0.001}, "training": {"epochs": 10}} ) return await train(cfg) \`\`\` The config is serialized when \`train\` is invoked and reconstructed as a \`DictConfig\` inside the task. No type registration, manual encoding or schema declaration is required. ## When to use this plugin Use \`flyteplugins-omegaconf\` when: - You already use OmegaConf. For example, you have YAML configs, dataclass-based config trees or a Hydra app, and want to keep that representation intact across task boundaries. - You want to pass a single composed config object instead of widening task signatures with dozens of scalar arguments. - You want to enforce schema validation at the task entry point via dataclass-backed structured configs. - You want resolved interpolations (\`${other.value}\`) to be materialized at submission time rather than at task runtime. If you do not use OmegaConf elsewhere, prefer plain dataclasses, \`pydantic.BaseModel\` or \`dict\` for task inputs as they are supported by Flyte natively without an extra dependency. ## Building a DictConfig Any of the standard OmegaConf construction methods produce a value the plugin can serialize. ### From a plain dict \`\`\`python{hl\_lines=\["1-3"\]} cfg = OmegaConf.create( {"optimizer": {"lr": 0.001}, "training": {"epochs": 10}} ) flyte.run(train, cfg=cfg) \`\`\` ### From a YAML file \`\`\`python{hl\_lines=\[1\]} cfg = OmegaConf.load("configs/training.yaml") flyte.run(train, cfg=cfg) \`\`\` The file is read locally on the submitter, not on the worker. If the YAML lives in your project tree and needs to be packaged into the task image, use \`flyte.with\_runcontext(copy\_style="all").run(...)\`. ### From a dataclass (structured config) \`\`\`python{hl\_lines=\["3-6", 8\]} from dataclasses import dataclass @dataclass class TrainConf: lr: float = 0.001 epochs: int = 10 cfg = OmegaConf.structured(TrainConf()) flyte.run(train, cfg=cfg) \`\`\` Structured configs are covered in detail in \*\*OmegaConf > Structured configs\*\* below. ### From a base config plus overrides \`\`\`python{hl\_lines=\["1-3"\]} base = OmegaConf.load("configs/training.yaml") override = OmegaConf.create({"optimizer": {"lr": 0.01}}) cfg = OmegaConf.merge(base, override) flyte.run(train, cfg=cfg) \`\`\` This is the same pattern Hydra uses internally. See the \[Hydra integration\](../hydra/\_index) for a full composition layer on top of this plugin. ## Variable interpolation OmegaConf supports \`${...}\` interpolations that resolve relative to the config tree: \`\`\`python{hl\_lines=\[3, 4\]} cfg = OmegaConf.create( { "base\_lr": 0.01, "optimizer": {"lr": "${base\_lr}", "momentum": 0.9}, } ) flyte.run(train, cfg=cfg) \`\`\` Interpolations are resolved at serialization time. By the time the task runs, \`cfg.optimizer.lr\` is the concrete float \`0.01\`, not the string \`"${base\_lr}"\`. This means: - The receiving task does not need any context that only existed in the submitter's environment. - Resolved values appear in the Flyte I/O panel. - A reference that fails to resolve at submission time fails fast, before any task runs. If you need lazy resolution on the worker, resolve the reference yourself inside the task or pass the unresolved string through a normal \`str\` input. ## Nested and deeply structured configs Nested configs are supported, including deeply structured OmegaConf objects. \`\`\`python{hl\_lines=\["1-13", 18\]} cfg = OmegaConf.create( { "experiment": { "model": { "encoder": { "attention": {"num\_heads": 8, "head\_dim": 64}, "ffn": {"hidden\_dim": 2048, "activation": "gelu"}, }, "decoder": {"num\_layers": 6}, } } } ) @env.task async def extract\_leaf(cfg: DictConfig) -> int: return int(cfg.experiment.model.encoder.attention.num\_heads) \`\`\` ## DictConfigs that contain lists A \`DictConfig\` may hold list values; they are reconstructed as nested \`ListConfig\`s on the receiving side. \`\`\`python{hl\_lines=\[4, 5, 8, 9\]} cfg = OmegaConf.create( { "model": { "layer\_sizes": \[64, 128, 256, 512\], "activations": \["relu", "relu", "relu", "sigmoid"\], }, "data": { "augmentations": \["random\_flip", "random\_crop", "color\_jitter"\], "input\_size": \[224, 224\], }, } ) @env.task async def double\_layer\_sizes(cfg: DictConfig) -> DictConfig: doubled = \[size \* 2 for size in cfg.model.layer\_sizes\] return OmegaConf.merge(cfg, {"model": {"layer\_sizes": doubled}}) \`\`\` ## ListConfig as input and output \`ListConfig\` is symmetric with \`DictConfig\` and supports the same construction patterns. ### Lists of primitives \`\`\`python{hl\_lines=\[2\]} @env.task async def scale\_values(values: ListConfig, factor: float) -> ListConfig: return OmegaConf.create(\[v \* factor for v in values\]) \`\`\` ### Building a schedule from another task \`\`\`python{hl\_lines=\[3, 7, 8\]} @env.task async def build\_lr\_schedule(base\_lr: float, num\_stages: int) -> ListConfig: return OmegaConf.create(\[base\_lr \* (0.5 \*\* i) for i in range(num\_stages)\]) @env.task async def train\_with\_schedule(cfg: DictConfig, lr\_schedule: ListConfig) -> float: final\_lr = float(lr\_schedule\[-1\]) ... \`\`\` ### Nested lists (list of lists) \`\`\`python{hl\_lines=\[1, 6\]} grid = OmegaConf.create(\[\[0.001, 0.01, 0.1\], \[10, 20, 50\]\]) @env.task async def flatten\_grid(grid: ListConfig) -> ListConfig: flat = \[item for sublist in OmegaConf.to\_container(grid) for item in sublist\] return OmegaConf.create(flat) \`\`\` ### Lists of DictConfigs \`\`\`python{hl\_lines=\["2-6"\]} configs = OmegaConf.create( \[ {"optimizer": {"lr": 0.001}, "training": {"epochs": 10}}, {"optimizer": {"lr": 0.01}, "training": {"epochs": 20}}, {"optimizer": {"lr": 0.1}, "training": {"epochs": 5}}, \] ) @env.task async def select\_best\_config(configs: ListConfig) -> DictConfig: best = max(OmegaConf.to\_container(configs), key=lambda c: c\["optimizer"\]\["lr"\]) return OmegaConf.create(best) \`\`\` ### Lists of dataclass instances \`\`\`python{hl\_lines=\["9-13"\]} @dataclass class LayerConf: name: str width: int activation: str layers = OmegaConf.create( \[ LayerConf(name="encoder", width=768, activation="gelu"), LayerConf(name="bottleneck", width=128, activation="relu"), LayerConf(name="decoder", width=768, activation="linear"), \] ) \`\`\` Each element round-trips as a typed \`DictConfig\` backed by \`LayerConf\`, so the receiving task can call \`OmegaConf.get\_type(layers\[0\])\` and access fields with attribute notation. > \*\*📝 Note\*\* > > ListConfig is always plain. Even when its elements are dataclass-backed, the outer \`ListConfig\` does not carry a list-level schema as there is no structured (typed-element) \`ListConfig\` in OmegaConf. This affects only the outer container; nested elements retain their schemas. ## Structured configs A structured config is a \`DictConfig\` that is bound to a Python dataclass. The dataclass acts as a schema: assigning a value of the wrong type raises \`omegaconf.ValidationError\`, and merging unknown keys raises an error instead of silently extending the config. ### Basic structured config \`\`\`python{hl\_lines=\["5-8", "11-14", 17, 20\]} from dataclasses import dataclass, field from omegaconf import OmegaConf, DictConfig @dataclass class OptimizerConf: lr: float = 0.001 weight\_decay: float = 1e-4 @dataclass class TrainConf: optimizer: OptimizerConf = field(default\_factory=OptimizerConf) epochs: int = 10 cfg = OmegaConf.structured(TrainConf()) flyte.run(train, cfg=cfg) # cfg.optimizer.lr = "oops" # raises omegaconf.ValidationError \`\`\` ### Schema reconstruction in the receiving task When a structured \`DictConfig\` is deserialized in a downstream task, the plugin operates in \*\*Auto mode\*\*: it reads the originating dataclass name from the wire payload and tries to import it. Two outcomes are possible: - Dataclass importable in the receiving task: \`cfg\` is reconstructed as a \`TrainConf\`-backed \`DictConfig\`. \`OmegaConf.get\_type(cfg)\` returns \`TrainConf\`, and type validation is enforced. - Dataclass not importable: \`cfg\` falls back to a plain \`DictConfig\` carrying the raw values. \`OmegaConf.get\_type(cfg)\` returns \`dict\`. The values are intact but the schema is lost. To keep schemas across task hops, define dataclasses in modules that are importable from every task in the pipeline (for example, in a shared \`configs.py\` module bundled into the task image). ### Required (\`MISSING\`) fields OmegaConf's \`MISSING\` sentinel marks a required field that has no default: \`\`\`python{hl\_lines=\[1, 5, "8-9", "12-13"\]} from omegaconf import MISSING @dataclass class TrainConf: data\_path: str = MISSING epochs: int = 10 # Pass with MISSING still unset — serialization succeeds. cfg = OmegaConf.structured(TrainConf()) flyte.run(train, cfg=cfg) # Or fill it before passing. cfg = OmegaConf.structured(TrainConf(data\_path="/data/imagenet")) flyte.run(train, cfg=cfg) \`\`\` A config with an unset \`MISSING\` field serializes and deserializes successfully as the sentinel is preserved on the wire. Accessing the field on the receiving side raises \`MissingMandatoryValue\`. > \*\*📝 Note\*\* > > Type annotations are preserved only in Auto mode. When the dataclass is importable on the receiving side, an unfilled \`MISSING\` field still carries its declared type (e.g. \`StringNode\` for \`str\`). When the plugin falls back to a plain \`DictConfig\` because the dataclass is not importable, the field becomes an \`AnyNode\` where the value is preserved, but the type annotation is not. ### Advanced field types Beyond primitives and nested dataclasses, structured configs may declare fields of these types and they will round-trip with their schemas intact: - \`Enum\` subclasses - \`pathlib.Path\` - \`Optional\[T\]\` - \`bytes\` - \`dict\[str, T\]\` where \`T\` is a dataclass - \`list\[T\]\` where \`T\` is a dataclass \`\`\`python{hl\_lines=\["6-8", "20-35"\]} from enum import Enum from pathlib import Path from typing import Optional class RunMode(Enum): TRAIN = "train" EVAL = "eval" @dataclass class CallbackConf: name: str = "early\_stop" patience: int = 3 monitor: str = MISSING @dataclass class AdvancedTrainConf: mode: RunMode = RunMode.TRAIN checkpoint\_dir: Path = Path("/tmp/checkpoints") maybe\_seed: Optional\[int\] = None payload: bytes = b"default-token" callbacks\_by\_name: dict\[str, CallbackConf\] = field( default\_factory=lambda: { "early\_stop": CallbackConf(name="early\_stop", patience=3), "checkpoint": CallbackConf(name="checkpoint", monitor="val\_loss"), } ) callbacks: list\[CallbackConf\] = field( default\_factory=lambda: \[ CallbackConf(name="lr\_monitor", patience=2, monitor="lr"), CallbackConf(name="nan\_guard", patience=1, monitor="loss"), \] ) \`\`\` Inside a downstream task: \`\`\`python @env.task async def inspect(cfg: DictConfig) -> str: assert OmegaConf.get\_type(cfg) == AdvancedTrainConf assert OmegaConf.get\_type(cfg.callbacks\[0\]) == CallbackConf assert isinstance(cfg.mode, RunMode) assert isinstance(cfg.checkpoint\_dir, Path) assert isinstance(cfg.payload, bytes) return cfg.mode.value \`\`\` ### Merging overrides on top of a structured base \`\`\`python{hl\_lines=\[3, 11\]} @env.task async def structured\_merge\_pipeline() -> str: base = OmegaConf.structured(TrainConf()) overrides = OmegaConf.create( { "optimizer": {"lr": 0.05}, "training": {"epochs": 100}, "experiment\_name": "sweep-run-1", } ) cfg = OmegaConf.merge(base, overrides) return await validate\_config(cfg) \`\`\` Merging an unknown key against a structured config raises an error, so define every key the override layer might supply on the dataclass. ## Embedding rich Python values inside a plain DictConfig A plain \`DictConfig\` (one not bound to a dataclass) can still hold Python values that OmegaConf does not natively model. The plugin preserves the following types end-to-end whether they appear in plain or structured configs: - \`pathlib.Path\` and any subclass of \`pathlib.PurePath\` - \`enum.Enum\` members - \`tuple\` (round-trips as \`tuple\`, not \`list\`) - \`bytes\` \`\`\`python{hl\_lines=\[1\]} cfg = OmegaConf.create({"model\_path": Path("/opt/models/model.bin")}) @env.task async def use\_path(cfg: DictConfig) -> str: assert isinstance(cfg.model\_path, Path) return f"model\_path={cfg.model\_path}" \`\`\` If an \`Enum\`'s class cannot be imported in the receiving environment, the value is returned as the underlying primitive (\`int\`, \`str\`, ...) instead of the enum member. ## Reserved-looking keys The plugin's wire format uses an internal payload marker (\`\_\_flyte\_omegaconf\_\_\`), which means user-facing keys named \`kind\`, \`values\`, \`name\`, \`value\`, \`type\`, or \`schema\` round-trip unchanged: \`\`\`python{hl\_lines=\[1, 8\]} cfg = OmegaConf.create({"kind": "training-job", "values": {"lr": 0.001}}) @env.task async def use\_payload\_shaped\_config(cfg: DictConfig) -> str: # cfg.values resolves to DictConfig.values() — use bracket notation # to reach the user key named "values". return f"kind={cfg.kind} lr={cfg\['values'\].lr}" \`\`\` The only practical consideration is Python's normal attribute-vs-method conflict: \`cfg.values\` is the \`.values()\` method, so reach for \`cfg\["values"\]\` when your config has a key with that name. ## YAML reports The Flyte I/O panel displays the literal wire representation of a \`DictConfig\`. !\[Wire Representation\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/omegaconf/input.png) For a YAML view, enable a Flyte report on the task and log the config with \`log\_yaml\`: \`\`\`python{hl\_lines=\[1, 4, 6\]} from flyteplugins.omegaconf import log\_yaml @env.task(report=True) async def train(cfg: DictConfig) -> DictConfig: await log\_yaml.aio(cfg, title="Input config") ... \`\`\` !\[YAML Report\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/omegaconf/yaml\_repr.png) The plugin also exposes: - \`to\_yaml(cfg)\`: render an OmegaConf container as a YAML string. - \`to\_html(cfg, title=...)\`: wrap the YAML in escaped HTML for embedding in a custom report. - \`replace\_yaml(cfg, ...)\`: replace the contents of a report tab instead of appending. \`\`\`python from flyteplugins.omegaconf.report import to\_yaml, replace\_yaml text = to\_yaml(cfg) await replace\_yaml.aio(cfg, tab="Final config") \`\`\` \`MISSING\` fields appear as \`???\` in the YAML output, matching OmegaConf's own convention. ## Wire format Both \`DictConfig\` and \`ListConfig\` are serialized as MessagePack blobs with the literal representation: \`\`\` Literal(scalar=Scalar(binary=Binary(value=, tag="msgpack"))) \`\`\` The msgpack payload uses an internal tagged structure to distinguish OmegaConf-specific concepts from raw values: - A \`DictConfig\` payload includes the originating dataclass name (\`builtins.dict\` for plain configs) plus its values. - \`MISSING\`, \`Enum\`, \`Path\`, and \`tuple\` values carry tagged shapes so they can be reconstructed faithfully. You normally do not need to inspect this format. It is documented here because: - The plugin serializes with \`resolve=True\`, so the wire representation always contains concrete values for \`${...}\` interpolations. - Cache-key metadata is set via Flyte's \`MESSAGEPACK\` serialization format, so two tasks given equivalent configs hit the same cache entry. ## End-to-end example The example below ties the pieces together: a structured \`DictConfig\` is created in a parent task, flows through several child tasks that read and modify it, and a \`ListConfig\` produced midway is consumed by a later stage. Each hop serializes and deserializes the config; the dataclass schema is recovered on the receiving side because \`TrainConf\` (and friends) are importable in every task in the pipeline. \`\`\` from dataclasses import dataclass, field import flyte from omegaconf import DictConfig, ListConfig, OmegaConf env = flyte.TaskEnvironment( name="omegaconf-pipeline-example", image=flyte.Image.from\_debian\_base().with\_pip\_packages("flyteplugins-omegaconf"), ) @dataclass class OptimizerConf: lr: float = 0.001 weight\_decay: float = 1e-4 @dataclass class DataConf: path: str = "" preprocessed: bool = False @dataclass class ResultsConf: val\_loss: float = 0.0 final\_lr: float = 0.0 num\_lr\_steps: int = 0 @dataclass class TrainConf: optimizer: OptimizerConf = field(default\_factory=OptimizerConf) data: DataConf = field(default\_factory=DataConf) results: ResultsConf = field(default\_factory=ResultsConf) epochs: int = 10 batch\_size: int = 32 experiment: str = "baseline" @env.task async def preprocess(cfg: DictConfig, dataset: str) -> DictConfig: """First stage: fills in the data section of cfg.""" return OmegaConf.merge(cfg, {"data": {"path": dataset, "preprocessed": True}}) @env.task async def build\_schedule(cfg: DictConfig) -> ListConfig: """Produces an LR schedule from cfg as a ListConfig.""" lrs = \[cfg.optimizer.lr \* (0.5\*\*i) for i in range(cfg.epochs)\] return OmegaConf.create(lrs) @env.task async def train(cfg: DictConfig, lr\_schedule: ListConfig) -> tuple\[DictConfig, float\]: """Simulates training. Returns the final cfg (with results filled in) and val loss.""" final\_lr = float(lr\_schedule\[-1\]) val\_loss = final\_lr \* 10 # placeholder result\_cfg = OmegaConf.merge( cfg, { "results": { "val\_loss": val\_loss, "final\_lr": final\_lr, "num\_lr\_steps": len(lr\_schedule), } }, ) return result\_cfg, val\_loss @env.task async def evaluate(result\_cfg: DictConfig, val\_loss: float) -> str: """Final stage: formats a report from the result config.""" return ( f"experiment={result\_cfg.experiment} " f"data={result\_cfg.data.path} " f"val\_loss={val\_loss:.6f} " f"final\_lr={result\_cfg.results.final\_lr:.6f} " f"lr\_steps={result\_cfg.results.num\_lr\_steps}" ) @env.task async def training\_pipeline(dataset: str) -> str: """Full pipeline: cfg flows preprocess, build\_schedule, train and evaluate.""" cfg = OmegaConf.structured( TrainConf( optimizer=OptimizerConf(lr=0.01, weight\_decay=1e-5), epochs=5, batch\_size=64, experiment="structured-cfg-pipeline", ) ) preprocessed\_cfg = await preprocess(cfg, dataset=dataset) lr\_schedule = await build\_schedule(preprocessed\_cfg) result\_cfg, val\_loss = await train(preprocessed\_cfg, lr\_schedule=lr\_schedule) return await evaluate(result\_cfg, val\_loss=val\_loss) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(training\_pipeline, dataset="s3://my-bucket/imagenet") print(f"Run URL: {run.url}") print(f"Outputs: {run.outputs()}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/omegaconf/example.py\* For more focused examples such as plain \`DictConfig\` patterns, advanced \`ListConfig\` shapes, all \`MISSING\`/\`Enum\`/\`Path\`/\`bytes\` cases, see the \[plugin repository\](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/omegaconf/examples). === PAGE: https://www.union.ai/docs/v2/flyte/integrations/openai === # OpenAI The OpenAI plugin provides a drop-in replacement for the \[OpenAI Agents SDK\](https://openai.github.io/openai-agents-python/) \`function\_tool\` decorator. It lets you use Flyte tasks as tools in agentic workflows so that tool calls run as tracked, reproducible Flyte task executions. ## When to use this plugin - Building agentic workflows with the OpenAI Agents SDK on Flyte - You want tool calls to run as Flyte tasks with full observability, retries, and caching - You want to combine LLM agents with existing Flyte pipelines ## Installation \`\`\`bash pip install flyteplugins-openai \`\`\` Requires \`openai-agents >= 0.2.4\`. ## Usage The plugin provides a single decorator, \`function\_tool\`, that wraps Flyte tasks as OpenAI agent tools. ### \`function\_tool\` When applied to a Flyte task (a function decorated with \`@env.task\`), \`function\_tool\` makes that task available as an OpenAI \`FunctionTool\`. The agent can call it like any other tool, and the call executes as a Flyte task. When applied to a regular function or a \`@flyte.trace\`-decorated function, it delegates directly to the OpenAI Agents SDK's built-in \`function\_tool\`. ### Basic pattern 1. Define a \`TaskEnvironment\` with your image and secrets 2. Decorate your task functions with \`@function\_tool\` and \`@env.task\` 3. Pass the tools to an \`Agent\` 4. Run the agent from another Flyte task \`\`\`python from agents import Agent, Runner from flyteplugins.openai.agents import function\_tool env = flyte.TaskEnvironment( name="openai\_agents", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="openai\_agents\_image"), secrets=flyte.Secret("openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY"), ) @function\_tool @env.task async def get\_weather(city: str) -> Weather: """Get the weather for a given city.""" return Weather(city=city, temperature\_range="14-20C", conditions="Sunny") agent = Agent( name="Weather Agent", instructions="You are a helpful agent.", tools=\[get\_weather\], ) @env.task async def main() -> str: result = await Runner.run(agent, input="What's the weather in Tokyo?") return result.final\_output \`\`\` > \[!NOTE\] > The docstring on each \`@function\_tool\` task is sent to the LLM as the tool description. Write clear, concise docstrings that describe what the tool does and what its parameters mean. ### Secrets Store your OpenAI API key as a Flyte secret and expose it as an environment variable: \`\`\`python secrets=flyte.Secret("openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY") \`\`\` ## Example \`\`\`python """OpenAI Agents with Flyte, basic tool example. Usage: Create secret: \`\`\` flyte create secret openai\_api\_key uv run agents\_tools.py \`\`\` """ # {{docs-fragment uv-script}} # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-openai>=2.0.0b7", # "openai-agents>=0.2.4", # "pydantic>=2.10.6", # \] # main = "main" # params = "" # /// # {{/docs-fragment uv-script}} # {{docs-fragment imports-task-env}} from agents import Agent, Runner from pydantic import BaseModel import flyte from flyteplugins.openai.agents import function\_tool env = flyte.TaskEnvironment( name="openai\_agents\_tools", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="openai\_agents\_image"), secrets=flyte.Secret("openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY"), ) # {{/docs-fragment imports-task-env}} # {{docs-fragment tools}} class Weather(BaseModel): city: str temperature\_range: str conditions: str @function\_tool @env.task async def get\_weather(city: str) -> Weather: """Get the weather for a given city.""" return Weather(city=city, temperature\_range="14-20C", conditions="Sunny with wind.") # {{/docs-fragment tools}} # {{docs-fragment agent}} agent = Agent( name="Hello world", instructions="You are a helpful agent.", tools=\[get\_weather\], ) @env.task async def main() -> str: result = await Runner.run(agent, input="What's the weather in Tokyo?") print(result.final\_output) return result.final\_output # {{/docs-fragment agent}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/openai/openai/agents\_tools.py\* ## API reference See the \[OpenAI API reference\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/openai/\_index) for full details. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/openai/agent\_tools === # Agent tools In this example, we will use the \`openai-agents\` library to create a simple agent that can use tools to perform tasks. This example is based on the \[basic tools example\](https://github.com/openai/openai-agents-python/blob/main/examples/basic/tools.py) example from the \`openai-agents-python\` repo. First, create an OpenAI API key, which you can get from the \[OpenAI website\](https://platform.openai.com/account/api-keys). Then, create a secret on your Flyte cluster with: \`\`\` flyte create secret OPENAI\_API\_KEY --value \`\`\` Then, we'll use \`uv script\` to specify our dependencies. \`\`\` """OpenAI Agents with Flyte, basic tool example. Usage: Create secret: \`\`\` flyte create secret openai\_api\_key uv run agents\_tools.py \`\`\` """ # {{docs-fragment uv-script}} # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-openai>=2.0.0b7", # "openai-agents>=0.2.4", # "pydantic>=2.10.6", # \] # main = "main" # params = "" # /// # {{/docs-fragment uv-script}} # {{docs-fragment imports-task-env}} from agents import Agent, Runner from pydantic import BaseModel import flyte from flyteplugins.openai.agents import function\_tool env = flyte.TaskEnvironment( name="openai\_agents\_tools", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="openai\_agents\_image"), secrets=flyte.Secret("openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY"), ) # {{/docs-fragment imports-task-env}} # {{docs-fragment tools}} class Weather(BaseModel): city: str temperature\_range: str conditions: str @function\_tool @env.task async def get\_weather(city: str) -> Weather: """Get the weather for a given city.""" return Weather(city=city, temperature\_range="14-20C", conditions="Sunny with wind.") # {{/docs-fragment tools}} # {{docs-fragment agent}} agent = Agent( name="Hello world", instructions="You are a helpful agent.", tools=\[get\_weather\], ) @env.task async def main() -> str: result = await Runner.run(agent, input="What's the weather in Tokyo?") print(result.final\_output) return result.final\_output # {{/docs-fragment agent}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/openai/openai/agents\_tools.py\* Next, we'll import the libraries and create a \`TaskEnvironment\`, which we need to run the example: \`\`\` """OpenAI Agents with Flyte, basic tool example. Usage: Create secret: \`\`\` flyte create secret openai\_api\_key uv run agents\_tools.py \`\`\` """ # {{docs-fragment uv-script}} # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-openai>=2.0.0b7", # "openai-agents>=0.2.4", # "pydantic>=2.10.6", # \] # main = "main" # params = "" # /// # {{/docs-fragment uv-script}} # {{docs-fragment imports-task-env}} from agents import Agent, Runner from pydantic import BaseModel import flyte from flyteplugins.openai.agents import function\_tool env = flyte.TaskEnvironment( name="openai\_agents\_tools", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="openai\_agents\_image"), secrets=flyte.Secret("openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY"), ) # {{/docs-fragment imports-task-env}} # {{docs-fragment tools}} class Weather(BaseModel): city: str temperature\_range: str conditions: str @function\_tool @env.task async def get\_weather(city: str) -> Weather: """Get the weather for a given city.""" return Weather(city=city, temperature\_range="14-20C", conditions="Sunny with wind.") # {{/docs-fragment tools}} # {{docs-fragment agent}} agent = Agent( name="Hello world", instructions="You are a helpful agent.", tools=\[get\_weather\], ) @env.task async def main() -> str: result = await Runner.run(agent, input="What's the weather in Tokyo?") print(result.final\_output) return result.final\_output # {{/docs-fragment agent}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/openai/openai/agents\_tools.py\* ## Define the tools We'll define a tool that can get weather information for a given city. In this case, we'll use a toy function that returns a hard-coded \`Weather\` object. \`\`\` """OpenAI Agents with Flyte, basic tool example. Usage: Create secret: \`\`\` flyte create secret openai\_api\_key uv run agents\_tools.py \`\`\` """ # {{docs-fragment uv-script}} # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-openai>=2.0.0b7", # "openai-agents>=0.2.4", # "pydantic>=2.10.6", # \] # main = "main" # params = "" # /// # {{/docs-fragment uv-script}} # {{docs-fragment imports-task-env}} from agents import Agent, Runner from pydantic import BaseModel import flyte from flyteplugins.openai.agents import function\_tool env = flyte.TaskEnvironment( name="openai\_agents\_tools", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="openai\_agents\_image"), secrets=flyte.Secret("openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY"), ) # {{/docs-fragment imports-task-env}} # {{docs-fragment tools}} class Weather(BaseModel): city: str temperature\_range: str conditions: str @function\_tool @env.task async def get\_weather(city: str) -> Weather: """Get the weather for a given city.""" return Weather(city=city, temperature\_range="14-20C", conditions="Sunny with wind.") # {{/docs-fragment tools}} # {{docs-fragment agent}} agent = Agent( name="Hello world", instructions="You are a helpful agent.", tools=\[get\_weather\], ) @env.task async def main() -> str: result = await Runner.run(agent, input="What's the weather in Tokyo?") print(result.final\_output) return result.final\_output # {{/docs-fragment agent}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/openai/openai/agents\_tools.py\* In this code snippet, the \`@function\_tool\` decorator is imported from \`flyteplugins.openai.agents\`, which is a drop-in replacement for the \`@function\_tool\` decorator from \`openai-agents\` library. ## Define the agent Then, we'll define the agent, which calls the tool: \`\`\` """OpenAI Agents with Flyte, basic tool example. Usage: Create secret: \`\`\` flyte create secret openai\_api\_key uv run agents\_tools.py \`\`\` """ # {{docs-fragment uv-script}} # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-openai>=2.0.0b7", # "openai-agents>=0.2.4", # "pydantic>=2.10.6", # \] # main = "main" # params = "" # /// # {{/docs-fragment uv-script}} # {{docs-fragment imports-task-env}} from agents import Agent, Runner from pydantic import BaseModel import flyte from flyteplugins.openai.agents import function\_tool env = flyte.TaskEnvironment( name="openai\_agents\_tools", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="openai\_agents\_image"), secrets=flyte.Secret("openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY"), ) # {{/docs-fragment imports-task-env}} # {{docs-fragment tools}} class Weather(BaseModel): city: str temperature\_range: str conditions: str @function\_tool @env.task async def get\_weather(city: str) -> Weather: """Get the weather for a given city.""" return Weather(city=city, temperature\_range="14-20C", conditions="Sunny with wind.") # {{/docs-fragment tools}} # {{docs-fragment agent}} agent = Agent( name="Hello world", instructions="You are a helpful agent.", tools=\[get\_weather\], ) @env.task async def main() -> str: result = await Runner.run(agent, input="What's the weather in Tokyo?") print(result.final\_output) return result.final\_output # {{/docs-fragment agent}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/openai/openai/agents\_tools.py\* ## Run the agent Finally, we'll run the agent. Create \`config.yaml\` file, which the \`flyte.init\_from\_config()\` function will use to connect to the Flyte cluster: \`\`\`bash flyte create config \\ --output ~/.flyte/config.yaml \\ --endpoint demo.hosted.unionai.cloud/ \\ --project flytesnacks \\ --domain development \\ --builder remote \`\`\` \`\`\` """OpenAI Agents with Flyte, basic tool example. Usage: Create secret: \`\`\` flyte create secret openai\_api\_key uv run agents\_tools.py \`\`\` """ # {{docs-fragment uv-script}} # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-openai>=2.0.0b7", # "openai-agents>=0.2.4", # "pydantic>=2.10.6", # \] # main = "main" # params = "" # /// # {{/docs-fragment uv-script}} # {{docs-fragment imports-task-env}} from agents import Agent, Runner from pydantic import BaseModel import flyte from flyteplugins.openai.agents import function\_tool env = flyte.TaskEnvironment( name="openai\_agents\_tools", resources=flyte.Resources(cpu=1, memory="250Mi"), image=flyte.Image.from\_uv\_script(\_\_file\_\_, name="openai\_agents\_image"), secrets=flyte.Secret("openai\_api\_key", as\_env\_var="OPENAI\_API\_KEY"), ) # {{/docs-fragment imports-task-env}} # {{docs-fragment tools}} class Weather(BaseModel): city: str temperature\_range: str conditions: str @function\_tool @env.task async def get\_weather(city: str) -> Weather: """Get the weather for a given city.""" return Weather(city=city, temperature\_range="14-20C", conditions="Sunny with wind.") # {{/docs-fragment tools}} # {{docs-fragment agent}} agent = Agent( name="Hello world", instructions="You are a helpful agent.", tools=\[get\_weather\], ) @env.task async def main() -> str: result = await Runner.run(agent, input="What's the weather in Tokyo?") print(result.final\_output) return result.final\_output # {{/docs-fragment agent}} # {{docs-fragment main}} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() # {{/docs-fragment main}} \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/openai/openai/agents\_tools.py\* ## Conclusion In this example, we've seen how to use the \`openai-agents\` library to create a simple agent that can use tools to perform tasks. The full code is available \[here\](https://github.com/unionai/unionai-examples/tree/main/v2/integrations/flyte-plugins/openai/openai). === PAGE: https://www.union.ai/docs/v2/flyte/integrations/pandera === # Pandera The \[Pandera\](https://pandera.readthedocs.io/en/latest/) plugin validates dataframes at task boundaries using \[\`DataFrameModel\`\](https://pandera.readthedocs.io/en/latest/dataframe\_models.html) schemas. When a task receives or returns a pandera-typed dataframe, the plugin automatically validates the data, raises or warns on schema violations, and writes an HTML validation report to the Flyte deck. Pandera supports multiple dataframe backends. The \`flyteplugins-pandera\` plugin handles: | Pandera typing module | DataFrame library | Additional plugin | |-|-|-| | \`pandera.typing.pandas\` | pandas | — | | \`pandera.typing.polars\` | Polars (eager and lazy) | \`flyteplugins-polars\` | | \`pandera.typing.pyspark\_sql\` | PySpark SQL | \`flyteplugins-spark\` | ## When to use this plugin - You want compile-time-style guarantees that data flowing between tasks conforms to a declared schema - You need column-level type, constraint, and statistical checks on task inputs and outputs - You want automatic validation reports visible in the Flyte UI ## Installation Install the plugin with the pandera extras for your dataframe backend: ### pandas \`\`\`bash pip install flyteplugins-pandera 'pandera\[pandas\]' \`\`\` ### Polars \`\`\`bash pip install flyteplugins-pandera flyteplugins-polars 'pandera\[polars\]' \`\`\` ### PySpark SQL \`\`\`bash pip install flyteplugins-pandera flyteplugins-spark 'pandera\[pyspark\]' \`\`\` ## Defining schemas Schemas are defined as Python classes that inherit from pandera's \`DataFrameModel\`. Each field declares a column name, type, and optional constraints: \`\`\`python import pandera.pandas as pa class EmployeeSchema(pa.DataFrameModel): employee\_id: int = pa.Field(ge=0) name: str class EmployeeSchemaWithStatus(EmployeeSchema): status: str = pa.Field(isin=\["active", "inactive"\]) \`\`\` Schemas compose through inheritance: \`EmployeeSchemaWithStatus\` includes all columns from \`EmployeeSchema\` plus the \`status\` column. For full details on schema definition—including custom checks, regex column matching, and \`Config\` options—see the \[pandera DataFrameModel documentation\](https://pandera.readthedocs.io/en/latest/dataframe\_models.html). ## Using schemas in tasks Annotate task inputs and outputs with pandera's generic \`DataFrame\` type. The plugin validates data on every encode (output) and decode (input): \`\`\`python import pandera.typing.pandas as pt @env.task(report=True) async def build\_employees() -> pt.DataFrame\[EmployeeSchema\]: return pd.DataFrame({ "employee\_id": \[1, 2, 3\], "name": \["Ada", "Grace", "Barbara"\], }) @env.task(report=True) async def add\_status( df: pt.DataFrame\[EmployeeSchema\], ) -> pt.DataFrame\[EmployeeSchemaWithStatus\]: return df.assign(status="active") \`\`\` Setting \`report=True\` on the task makes validation reports visible as deck tabs in the Flyte UI. ## Error handling with \`ValidationConfig\` By default, a validation failure raises an exception and fails the task. To downgrade failures to warnings instead, annotate the parameter with \`ValidationConfig(on\_error="warn")\`: \`\`\`python from typing import Annotated from flyteplugins.pandera import ValidationConfig @env.task(report=True) async def lenient\_pass\_through( df: Annotated\[pt.DataFrame\[EmployeeSchema\], ValidationConfig(on\_error="warn")\], ) -> Annotated\[pt.DataFrame\[EmployeeSchemaWithStatus\], ValidationConfig(on\_error="warn")\]: ... \`\`\` | \`on\_error\` value | Behavior | |-|-| | \`"raise"\` (default) | Validation failure raises \`pandera.errors.SchemaError\` and the task fails | | \`"warn"\` | Validation failure logs a warning and writes the report, but the task continues | You can mix \`"raise"\` and \`"warn"\` across inputs and outputs of the same task. For example, use \`"warn"\` on inputs to accept best-effort data while still enforcing strict output contracts. ## Image configuration Include the plugin in your task image. The exact setup depends on your dataframe backend: ### Pandas \`\`\`python import flyte img = flyte.Image.from\_debian\_base( python\_version=(3, 12), ).with\_pip\_packages("flyteplugins-pandera") env = flyte.TaskEnvironment( "pandera\_pandas", image=img, resources=flyte.Resources(cpu="1", memory="2Gi"), ) \`\`\` ### Polars \`\`\`python import flyte img = ( flyte.Image.from\_debian\_base(python\_version=(3, 12)) .with\_pip\_packages("flyteplugins-polars", "pandera\[polars\]") ) env = flyte.TaskEnvironment( "pandera\_polars", image=img, resources=flyte.Resources(cpu="1", memory="2Gi"), ) \`\`\` ### PySpark SQL \`\`\`python import flyte from flyteplugins.spark.task import Spark image = ( flyte.Image.from\_base("apache/spark-py:v3.4.0") .clone(name="pandera-pyspark-sql", python\_version=(3, 10), extendable=True) .with\_pip\_packages("flyteplugins-spark", "pandera\[pyspark\]") ) spark\_conf = Spark( spark\_conf={ "spark.driver.memory": "1000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", }, ) env = flyte.TaskEnvironment( name="pandera\_pyspark", plugin\_config=spark\_conf, image=image, resources=flyte.Resources(cpu="1", memory="2Gi"), ) \`\`\` ## Polars lazy frames The Polars backend supports both \`pt.DataFrame\` (eager) and \`pt.LazyFrame\` (lazy). With lazy frames, pandera validates the data when the frame is materialized at task I/O boundaries: \`\`\`python import pandera.typing.polars as pt import polars as pl @env.task(report=True) async def create\_lazy() -> pt.LazyFrame\[MetricsSchema\]: return pl.LazyFrame({"item": \["x", "y"\], "value": \[3.0, 4.0\]}) @env.task(report=True) async def consume\_lazy( lf: pt.LazyFrame\[MetricsSchema\], ) -> pt.DataFrame\[MetricsSchema\]: return lf.filter(pl.col("value") > 0.0).collect() \`\`\` ## Examples ### pandas \`\`\`python # /// script # requires-python = ">=3.12" # dependencies = \[ # "flyte", # "flyteplugins-pandera", # "pandera\[pandas\]", # \] # main = "main" # /// from \_\_future\_\_ import annotations from typing import Annotated import pandas as pd import pandera.pandas as pa import pandera.typing.pandas as pt from flyteplugins.pandera import ValidationConfig import flyte img = flyte.Image.from\_debian\_base(python\_version=(3, 12)).with\_pip\_packages( "flyteplugins-pandera", "pandera\[pandas\]" ) env = flyte.TaskEnvironment( "pandera\_pandas\_schema", image=img, resources=flyte.Resources(cpu="1", memory="2Gi"), ) class EmployeeSchema(pa.DataFrameModel): employee\_id: int = pa.Field(ge=0) name: str class EmployeeSchemaWithStatus(EmployeeSchema): status: str = pa.Field(isin=\["active", "inactive"\]) # {{docs-fragment build\_valid\_employees}} @env.task(report=True) async def build\_valid\_employees() -> pt.DataFrame\[EmployeeSchema\]: return pd.DataFrame( { "employee\_id": \[1, 2, 3\], "name": \["Ada", "Grace", "Barbara"\], } ) # {{/docs-fragment}} # {{docs-fragment pass\_through}} @env.task(report=True) async def pass\_through( df: pt.DataFrame\[EmployeeSchema\], ) -> pt.DataFrame\[EmployeeSchemaWithStatus\]: return df.assign(status="active") # {{/docs-fragment}} # {{docs-fragment pass\_through\_with\_error\_warn}} @env.task(report=True) async def pass\_through\_with\_error\_warn( df: Annotated\[ pt.DataFrame\[EmployeeSchema\], ValidationConfig(on\_error="warn") \], ) -> Annotated\[ pt.DataFrame\[EmployeeSchemaWithStatus\], ValidationConfig(on\_error="warn") \]: del df\["name"\] return df # {{/docs-fragment}} # {{docs-fragment pass\_through\_with\_error\_raise}} @env.task(report=True) async def pass\_through\_with\_error\_raise( df: Annotated\[ pt.DataFrame\[EmployeeSchema\], ValidationConfig(on\_error="warn") \], ) -> Annotated\[ pt.DataFrame\[EmployeeSchemaWithStatus\], ValidationConfig(on\_error="raise") \]: del df\["name"\] return df # {{/docs-fragment}} @env.task(report=True) async def main() -> pt.DataFrame\[EmployeeSchemaWithStatus\]: df = await build\_valid\_employees() df2 = await pass\_through(df) await pass\_through\_with\_error\_warn(df.drop(\["employee\_id"\], axis="columns")) await pass\_through\_with\_error\_warn(df.assign(employee\_id=-1)) try: await pass\_through\_with\_error\_raise(df) except Exception as exc: print(exc) return df2 if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() print("pandas pandera example OK:", run.outputs()\[0\]) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pandera/pandas\_schema.py\* ### Polars \`\`\`python # /// script # requires-python = ">=3.12" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-pandera", # "flyteplugins-polars", # "pandera\[polars\]", # \] # main = "main" # /// from \_\_future\_\_ import annotations from typing import Annotated import pandera.polars as pa import pandera.typing.polars as pt import polars as pl from flyteplugins.pandera import ValidationConfig import flyte img = ( flyte.Image.from\_debian\_base(python\_version=(3, 12)) .with\_pip\_packages("flyteplugins-pandera", "flyteplugins-polars", "pandera\[polars\]") ) env = flyte.TaskEnvironment( "pandera\_polars\_schema", image=img, resources=flyte.Resources(cpu="1", memory="2Gi"), ) class EmployeeSchema(pa.DataFrameModel): employee\_id: int = pa.Field(ge=0) name: str class EmployeeSchemaWithStatus(EmployeeSchema): status: str = pa.Field(isin=\["active", "inactive"\]) class MetricsSchema(pa.DataFrameModel): item: str value: float # {{docs-fragment build\_valid\_employees}} @env.task(report=True) async def build\_valid\_employees() -> pt.DataFrame\[EmployeeSchema\]: return pl.DataFrame( { "employee\_id": \[1, 2, 3\], "name": \["Ada", "Grace", "Barbara"\], } ) # {{/docs-fragment}} # {{docs-fragment pass\_through}} @env.task(report=True) async def pass\_through( df: pt.DataFrame\[EmployeeSchema\], ) -> pt.DataFrame\[EmployeeSchemaWithStatus\]: return df.with\_columns(pl.lit("active").alias("status")) # {{/docs-fragment}} @env.task(report=True) async def pass\_through\_with\_error\_warn( df: Annotated\[ pt.DataFrame\[EmployeeSchema\], ValidationConfig(on\_error="warn") \], ) -> Annotated\[ pt.DataFrame\[EmployeeSchemaWithStatus\], ValidationConfig(on\_error="warn") \]: return df.drop("name") @env.task(report=True) async def pass\_through\_with\_error\_raise( df: Annotated\[ pt.DataFrame\[EmployeeSchema\], ValidationConfig(on\_error="warn") \], ) -> Annotated\[ pt.DataFrame\[EmployeeSchemaWithStatus\], ValidationConfig(on\_error="raise") \]: return df.drop("name") # {{docs-fragment metrics\_lazy}} @env.task(report=True) async def metrics\_eager() -> pt.DataFrame\[MetricsSchema\]: return pl.DataFrame({"item": \["a", "b"\], "value": \[1.0, 2.0\]}) @env.task(report=True) async def metrics\_lazy() -> pt.LazyFrame\[MetricsSchema\]: return pl.LazyFrame({"item": \["x", "y"\], "value": \[3.0, 4.0\]}) @env.task(report=True) async def filter\_metrics( lf: pt.LazyFrame\[MetricsSchema\], ) -> pt.DataFrame\[MetricsSchema\]: return lf.filter(pl.col("value") > 0.0).collect() # {{/docs-fragment}} @env.task(report=True) async def main() -> pt.DataFrame\[EmployeeSchemaWithStatus\]: df = await build\_valid\_employees() df2 = await pass\_through(df) await pass\_through\_with\_error\_warn(df.drop("employee\_id")) await pass\_through\_with\_error\_warn( df.with\_columns(pl.lit(-1).alias("employee\_id")) ) try: await pass\_through\_with\_error\_raise(df) except Exception as exc: print(exc) \_ = await metrics\_eager() lazy = await metrics\_lazy() \_ = await filter\_metrics(lazy) return df2 if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() print("polars pandera example OK:", run.outputs()\[0\]) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pandera/polars\_schema.py\* ### PySpark SQL \`\`\`python # /// script # requires-python = ">=3.10" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-pandera", # "flyteplugins-spark", # "pandera\[pyspark\]", # \] # main = "main" # /// from \_\_future\_\_ import annotations from typing import Annotated, cast import pandera.typing.pyspark\_sql as pt import pyspark.sql.types as T from flyteplugins.pandera import ValidationConfig from flyteplugins.spark.task import Spark from pandera.pyspark import DataFrameModel, Field from pyspark.sql import SparkSession from pyspark.sql import functions as F import flyte image = ( flyte.Image.from\_base("apache/spark-py:v3.4.0") .clone(name="pandera-pyspark-sql", python\_version=(3, 10), extendable=True) .with\_pip\_packages( "flyteplugins-pandera", "flyteplugins-spark", "pandera\[pyspark\]", ) ) spark\_conf = Spark( spark\_conf={ "spark.driver.memory": "1000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", "spark.kubernetes.file.upload.path": "/opt/spark/work-dir", "spark.jars": ( "https://storage.googleapis.com/hadoop-lib/gcs/" "gcs-connector-hadoop3-latest.jar," "https://repo1.maven.org/maven2/org/apache/hadoop/" "hadoop-aws/3.2.2/hadoop-aws-3.2.2.jar," "https://repo1.maven.org/maven2/com/amazonaws/" "aws-java-sdk-bundle/1.12.262/aws-java-sdk-bundle-1.12.262.jar" ), }, ) env = flyte.TaskEnvironment( name="pandera\_pyspark\_sql\_schema", plugin\_config=spark\_conf, image=image, resources=flyte.Resources(cpu="1", memory="2Gi"), ) # {{docs-fragment schemas}} class EmployeeSchema(DataFrameModel): employee\_id: int = Field(ge=0) name: str = Field() job\_title: str = Field() class EmployeeSchemaWithStatus(EmployeeSchema): status: str = Field(isin=\["active", "inactive"\]) # {{/docs-fragment}} # {{docs-fragment build\_valid\_employees}} @env.task(report=True) async def build\_valid\_employees() -> pt.DataFrame\[EmployeeSchema\]: spark = cast(SparkSession, flyte.ctx().data\["spark\_session"\]) data = \[ (1, "Ada", "Engineer"), (2, "Grace", "Mathematician"), (3, "Barbara", "Computer scientist"), \] schema = T.StructType( \[ T.StructField("employee\_id", T.IntegerType(), False), T.StructField("name", T.StringType(), False), T.StructField("job\_title", T.StringType(), False), \] ) return spark.createDataFrame(data, schema=schema) # {{/docs-fragment}} # {{docs-fragment pass\_through}} @env.task(report=True) async def pass\_through( df: pt.DataFrame\[EmployeeSchema\], ) -> pt.DataFrame\[EmployeeSchemaWithStatus\]: return df.withColumn("status", F.lit("active")) # {{/docs-fragment}} @env.task(report=True) async def pass\_through\_with\_error\_warn( df: Annotated\[ pt.DataFrame\[EmployeeSchema\], ValidationConfig(on\_error="warn") \], ) -> Annotated\[ pt.DataFrame\[EmployeeSchemaWithStatus\], ValidationConfig(on\_error="warn") \]: return df.drop("name") @env.task(report=True) async def pass\_through\_with\_error\_raise( df: Annotated\[ pt.DataFrame\[EmployeeSchema\], ValidationConfig(on\_error="warn") \], ) -> Annotated\[ pt.DataFrame\[EmployeeSchemaWithStatus\], ValidationConfig(on\_error="raise") \]: return df.drop("name") @env.task(report=True) async def main() -> pt.DataFrame\[EmployeeSchemaWithStatus\]: df = await build\_valid\_employees() df2 = await pass\_through(df) await pass\_through\_with\_error\_warn(df.drop("employee\_id")) await pass\_through\_with\_error\_warn(df.withColumn("employee\_id", F.lit(-1))) try: await pass\_through\_with\_error\_raise(df) except Exception as exc: print(exc) return df2 if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(main) print(run.url) run.wait() print("pyspark\_sql pandera example OK:", run.outputs()\[0\]) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pandera/pyspark\_sql\_schema.py\* === PAGE: https://www.union.ai/docs/v2/flyte/integrations/papermill === # Papermill The Papermill plugin lets you run Jupyter notebooks as Flyte tasks. It uses \[papermill\](https://papermill.readthedocs.io/) to parameterize and execute \`.ipynb\` files, capture their outputs as typed Flyte values, and render the executed notebook as an HTML report visible in the Flyte UI. A \`NotebookTask\` behaves like any other Flyte task: it has typed inputs and outputs, participates in workflows, runs remotely, integrates with the Flyte type system (including \`File\`, \`Dir\`, and \`DataFrame\`), and can call other Flyte tasks from within the notebook. ## When to use this plugin - Productionizing exploratory notebooks without rewriting them as Python modules - Generating cell-by-cell HTML reports as task artifacts (charts, tables, narrative analysis) - Letting data scientists iterate in notebooks while platform teams orchestrate them - Running notebooks on Spark or with GPU/CPU resources configured on the task environment ## Installation \`\`\`bash pip install flyteplugins-papermill \`\`\` The plugin must also be installed in the task image. For example: \`\`\`python{hl\_lines=\["3-5"\]} import flyte image = flyte.Image.from\_debian\_base(name="papermill-env").with\_pip\_packages( "flyteplugins-papermill" ) env = flyte.TaskEnvironment(name="papermill\_env", image=image) \`\`\` ## Quick start \`\`\`python{hl\_lines=\[1, 6, "9-15", 19\]} from flyteplugins.papermill import NotebookTask import flyte env = flyte.TaskEnvironment( name="my\_env", image=flyte.Image.from\_debian\_base(name="my-env").with\_pip\_packages("flyteplugins-papermill"), ) add\_numbers = NotebookTask( name="add\_numbers", notebook\_path="notebooks/basic\_math.ipynb", task\_environment=env, inputs={"x": int, "y": float}, outputs={"result": float}, ) @env.task def workflow(x: int = 5, y: float = 3.14) -> float: return add\_numbers(x=x, y=y) \`\`\` \`notebook\_path\` may be relative (resolved against the calling file's directory) or absolute. ## Notebook setup Each notebook driven by a \`NotebookTask\` needs two specially tagged cells. ### \`parameters\` cell Tag a cell with \`parameters\` and assign default values matching the names declared in \`inputs={...}\`. Papermill injects the actual values into a cell appended right after this one at execution time. \`\`\`python # tagged: parameters x = 0 y = 0.0 \`\`\` ### \`outputs\` cell Tag a cell with \`outputs\` and call \`record\_outputs(...)\` as the last expression of the cell. The function returns a serialized representation of the values, which Jupyter captures as the cell's displayed output. \`NotebookTask\` then reads that captured output from the executed notebook to recover the typed values. \`\`\`python # tagged: outputs from flyteplugins.papermill import record\_outputs record\_outputs(result=x + y) \`\`\` \`record\_outputs\` accepts any value that the Flyte type system supports such as primitives, \`File\`, \`Dir\`, \`DataFrame\`, dataclasses, etc. The output names and types must match the \`outputs={...}\` declaration on the \`NotebookTask\`. > \[!NOTE\] > Inputs and outputs have different type rules. Inputs are restricted to JSON-serializable primitives plus \`File\`/\`Dir\`/\`DataFrame\` because papermill's parameter mechanism is JSON-only. Outputs go through the full Flyte type engine inside the notebook via \`record\_outputs\`, so dataclasses and any other Flyte-supported type work there. If a notebook has no outputs, omit the \`outputs\` cell and don't pass \`outputs\` to \`NotebookTask\`. The notebook still runs and its HTML report is rendered, but no values are returned. ## Inputs and outputs ### Supported input types Notebook parameters are passed through papermill, which only accepts JSON-serializable values. The plugin allows: - Primitives: \`int\`, \`float\`, \`str\`, \`bool\`, \`list\`, \`dict\`, \`None\` - Flyte I/O types: \`flyte.io.File\`, \`flyte.io.Dir\`, \`flyte.io.DataFrame\` (serialized to their path/URI strings) Passing any other type raises \`TypeError\` at call time. Wrap unsupported values in a dataclass and serialize them to a primitive container, or write them to a \`File\`/\`Dir\` first. ### Complex types: File, Dir, DataFrame \`File\`, \`Dir\` and \`DataFrame\` are passed to the notebook as plain path/URI strings. Reconstruct them inside the notebook with the provided helpers: \`\`\`python from flyteplugins.papermill import load\_file, load\_dir, load\_dataframe # input\_file, input\_dir, input\_df were injected as strings by papermill f = load\_file(input\_file) # -> flyte.io.File d = load\_dir(input\_dir) # -> flyte.io.Dir df = load\_dataframe(input\_df) # -> flyte.io.DataFrame (parquet by default) \`\`\` \`load\_dataframe\` accepts a \`fmt\` argument (default \`"parquet"\`) for non-parquet storage formats. Jupyter supports top-level \`await\`, so use it directly for async I/O: \`\`\`python{hl\_lines=\[4, 5\]} import pandas as pd from flyte.io import DataFrame pdf = await df.open(pd.DataFrame).all() output\_df = await DataFrame.from\_local(pdf) \`\`\` To return a \`DataFrame\` from a notebook, materialize it as a \`flyte.io.DataFrame\` and pass it to \`record\_outputs\`: \`\`\`python{hl\_lines=\[6, 7, 9\]} # tagged: outputs import pandas as pd from flyte.io import DataFrame from flyteplugins.papermill import record\_outputs result\_df = pd.DataFrame({"name": \["alice", "bob"\], "score": \[90, 75\]}) output = await DataFrame.from\_local(result\_df) record\_outputs(filtered\_df=output, row\_count=len(result\_df)) \`\`\` The same pattern applies to \`File\` (\`await File.from\_local(...)\`) and \`Dir\` (\`await Dir.from\_local(...)\`). ### Outputs: single, multiple, none A \`NotebookTask\` returns: - A single value when \`outputs\` has one entry - A tuple in the order declared in \`outputs\` when there are multiple entries - \`None\` when \`outputs\` is omitted \`\`\`python{hl\_lines=\[7, 12\]} # Multiple outputs text\_analysis = NotebookTask( name="text\_analysis", notebook\_path="notebooks/text.ipynb", task\_environment=env, inputs={"text": str, "n": int}, outputs={"repeated": str, "word\_count": int, "char\_count": int}, ) @env.task def workflow(text: str, n: int) -> tuple\[str, int, int\]: repeated, word\_count, char\_count = text\_analysis(text=text, n=n) return repeated, word\_count, char\_count \`\`\` \`\`\`python{hl\_lines=\[11\]} # No outputs — useful for side-effect-only notebooks (reports, exports) printer = NotebookTask( name="printer", notebook\_path="notebooks/print\_report.ipynb", task\_environment=env, inputs={"message": str}, ) @env.task def report\_workflow(message: str = "hello"): printer(message=message) \`\`\` If a declared output is missing from \`record\_outputs(...)\`, \`NotebookTask\` raises \`TypeError\` listing the missing names. ## Calling Flyte tasks from notebooks You can call other Flyte tasks directly from inside a notebook. The plugin injects the parent task's runtime context into the notebook kernel at the start of execution, so task calls are routed through the Flyte controller automatically, so no manual setup required. When running remotely, each task call is submitted to Flyte and appears as a separate node in the run graph. When running locally, the calls execute in-process as regular Python functions. \`\`\`python{hl\_lines=\[1, 4\]} # Inside a notebook cell from my\_tasks import expensive\_task result = await expensive\_task(data=42) \`\`\` Sync tasks can be called the same way: \`\`\`python{hl\_lines=\[3\]} from my\_tasks import compute\_total total = compute\_total(values=\[1, 2, 3\]) \`\`\` > \[!NOTE\] > The setup cell that initializes the runtime context is injected automatically and stripped from the rendered HTML report and the uploaded \`.ipynb\` files, so it never shows up to users. ## Workflow patterns ### Chaining notebooks Outputs from one \`NotebookTask\` can feed directly into another: \`\`\`python{hl\_lines=\[3, 4\]} @env.task def chained\_workflow(a: int, b: float, c: float) -> float: intermediate = step1\_add(x=a, y=b) final = step2\_add(x=int(intermediate), y=c) return final \`\`\` ### Mixing notebooks with regular tasks \`NotebookTask\` composes with \`@env.task\` functions in either direction: \`\`\`python{hl\_lines=\["3-5"\]} @env.task def mixed\_workflow(n: int) -> float: doubled = double(n=n) # regular task nb\_result = notebook\_add(x=doubled, y=100.0) # notebook task return add(a=nb\_result, b=0.5) # regular task \`\`\` ### Inline definition \`NotebookTask\` can be created inside a task function rather than at module scope. The resolver bakes the notebook path and type schemas into the task spec at registration time, so no module-level reference is required at execution. \`\`\`python{hl\_lines=\[3, 5\]} @env.task def workflow(x: int = 3, y: float = 1.5) -> int: from flyteplugins.papermill import NotebookTask nb = NotebookTask( name="add\_numbers", notebook\_path="notebooks/basic\_math.ipynb", task\_environment=env, inputs={"x": int, "y": float}, outputs={"result": float}, ) return nb(x=x, y=y) \`\`\` ### Calling from sync vs. async tasks \`NotebookTask\` is internally synchronous. Papermill blocks while the notebook runs. Call it directly from a sync task or use \`.aio()\` from an async task: \`\`\`python{hl\_lines=\[2, 6, 7\]} @env.task def sync\_parent(x: int) -> float: return notebook(x=x) @env.task async def async\_parent(x: int) -> float: return await notebook.aio(x=x) \`\`\` ### Running a NotebookTask directly as the entrypoint A \`NotebookTask\` can be the workflow entrypoint without wrapping it in another task: \`\`\`python{hl\_lines=\[1, 11\]} nb = NotebookTask( name="add\_numbers", notebook\_path="notebooks/basic\_math.ipynb", task\_environment=env, inputs={"x": int, "y": float}, outputs={"result": float}, ) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.with\_runcontext(mode="remote", copy\_style="all").run(nb, x=3, y=1.5) print(run.url) \`\`\` ## Reports and notebook artifacts ### HTML report (default) Every \`NotebookTask\` execution renders the executed notebook to HTML and logs it to the Flyte Report tab for that task. This happens whether the notebook succeeds or fails — see \*\*Papermill > Reports and notebook artifacts > Failure reports\*\* below. The report is on by default and requires no configuration. !\[HTML Report\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/papermill/default\_report.png) ### Notebook artifacts By default the executed notebook lives only inside the rendered HTML report. To get the source and executed \`.ipynb\` files as typed Flyte outputs (so downstream tasks can read them or so they show up as artifacts in the run UI), set \`output\_notebooks=True\`: \`\`\`python{hl\_lines=\[7, 12\]} notebook = NotebookTask( name="analysis", notebook\_path="notebooks/analysis.ipynb", task\_environment=env, inputs={"x": int}, outputs={"result": float}, output\_notebooks=True, ) @env.task def workflow(x: int = 5) -> tuple\[float, File, File\]: result, source\_nb, executed\_nb = notebook(x=x) return result, source\_nb, executed\_nb \`\`\` When enabled, two outputs are appended to the task's interface automatically: - \`output\_notebook\`: The source \`.ipynb\` (no executed cell outputs) - \`output\_notebook\_executed\`: The executed \`.ipynb\` (with cell outputs) > \[!WARNING\] > The names \`output\_notebook\` and \`output\_notebook\_executed\` are reserved when \`output\_notebooks=True\`. Don't use them as your own user output names. ### Clean reports \`report\_mode=True\` tells papermill to mark input cells with a \`source\_hidden\` flag during execution. The plugin then strips those input cells from both the rendered HTML report and the uploaded \`.ipynb\` files, so only cell outputs (charts, tables, text) remain. This produces a clean stakeholder-facing report without exposing the underlying code. \`\`\`python{hl\_lines=\[3\]} notebook = NotebookTask( ... report\_mode=True, output\_notebooks=True, ) \`\`\` !\[Clean Report\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/papermill/clean\_report.png) ### Failure reports The HTML report is rendered even when the notebook fails. Papermill writes the output notebook cell-by-cell as it executes, so the partial notebook is on disk when an exception propagates out. The plugin renders this partial notebook to HTML and flushes it to the Flyte Report before re-raising the error, giving full visibility into which cell failed and what output the earlier cells produced. This is especially useful for long-running notebooks: you can inspect partial results without re-running the whole pipeline. !\[Failed Report\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/papermill/failed\_report.png) ## Spark notebooks Pass \`plugin\_config=Spark(...)\` to run a notebook inside a Spark driver pod managed by the Spark on Kubernetes Operator: \`\`\`python{hl\_lines=\["8-16"\]} from flyteplugins.papermill import NotebookTask from flyteplugins.spark import Spark spark\_nb = NotebookTask( name="spark\_analysis", notebook\_path="notebooks/spark\_analysis.ipynb", task\_environment=env, plugin\_config=Spark( spark\_conf={ "spark.executor.instances": "2", "spark.executor.memory": "2g", "spark.executor.cores": "1", "spark.driver.memory": "1g", "spark.driver.cores": "1", }, ), inputs={"data": list}, outputs={"total": int, "count": int}, ) \`\`\` Inside the notebook, build the \`SparkSession\` directly: \`\`\`python from pyspark.sql import SparkSession spark = SparkSession.builder.appName("FlyteSpark").getOrCreate() \`\`\` > \[!WARNING\] > \`SparkContext.addPyFile()\` is not called for notebook tasks. The notebook kernel runs in a subprocess that cannot share state with the parent task process, so dynamic code distribution via \`addPyFile\` is not supported. Executor pods use the same Docker image as the driver, so any package needed in UDFs must be installed in the image. See the \[Spark plugin\](../spark/\_index) page for the full \`Spark\` configuration reference. ## Local testing Calling a \`NotebookTask\` as a regular Python function outside any Flyte runner executes the notebook synchronously through papermill and returns Python values: \`\`\`python result = add\_numbers(x=1, y=2.5) \`\`\` In this mode: - The notebook runs in-process (no remote submission) - No HTML report is rendered (no task context) - \`File\` and \`Dir\` outputs created inside the notebook resolve to local paths - No plugin lifecycle hooks fire (so no Spark cluster is provisioned, etc.) This makes iteration on notebook logic fast. You can run the task from a script, REPL or test without going through Flyte at all. ## Execution options \`NotebookTask\` exposes the full set of papermill execution knobs. The snippet below shows example values. See \*\*Papermill > \`NotebookTask\` reference\*\* for defaults. \`\`\`python NotebookTask( name="all\_options", notebook\_path="notebooks/basic\_math.ipynb", task\_environment=env, inputs={"x": int, "y": float}, outputs={"result": float}, kernel\_name="python3", # default None - use kernel from notebook metadata language=None, # rarely needed; overrides notebook language execution\_timeout=300, # default None - no per-cell timeout start\_timeout=120, # default 60 seconds to wait for kernel startup log\_output=True, # default False; stream cell output to task log progress\_bar=True, # default True; tqdm-style progress in logs report\_mode=False, # default False; True hides input cells in report request\_save\_on\_cell\_execute=True, # default True; save after every cell (nbclient) engine\_name=None, # default None - nbclient engine\_kwargs={"autosave\_cell\_every": 30}, # extra kwargs forwarded to engine ) \`\`\` > \[!NOTE\] > \`request\_save\_on\_cell\_execute\` is largely redundant in remote execution: the plugin always renders and uploads the partial notebook on failure, so crash diagnostics don't depend on it. Leave it on its default unless using a custom engine that requires it. ## \`NotebookTask\` reference | Parameter | Default | Description | | ------------------------------ | ------- | ----------------------------------------------------------------------------------------- | | \`name\` | — | Task name | | \`notebook\_path\` | — | Path to the \`.ipynb\`, relative to the calling file or absolute | | \`task\_environment\` | — | \`TaskEnvironment\` for registration and remote execution | | \`inputs\` | \`None\` | \`{name: type}\` dict of notebook inputs | | \`outputs\` | \`None\` | \`{name: type}\` dict of notebook outputs | | \`plugin\_config\` | \`None\` | Plugin config — currently only \`Spark(...)\` is supported. Sets the task type accordingly. | | \`kernel\_name\` | \`None\` | Jupyter kernel name; \`None\` uses the kernel from notebook metadata | | \`engine\_name\` | \`None\` | Papermill engine; \`None\` uses the default \`nbclient\` engine | | \`log\_output\` | \`False\` | Stream cell output to the task log | | \`start\_timeout\` | \`60\` | Seconds to wait for kernel startup | | \`execution\_timeout\` | \`None\` | Per-cell timeout in seconds; \`None\` means no timeout | | \`report\_mode\` | \`False\` | Strip input cells from the report and uploaded \`.ipynb\` | | \`request\_save\_on\_cell\_execute\` | \`True\` | Save notebook after every cell (nbclient engine only) | | \`progress\_bar\` | \`True\` | Show a tqdm-style progress bar during execution | | \`language\` | \`None\` | Override notebook language (rarely needed) | | \`engine\_kwargs\` | \`{}\` | Extra kwargs forwarded to the papermill engine | | \`output\_notebooks\` | \`False\` | Upload source and executed \`.ipynb\` as \`File\` task outputs | ## Helper functions These are imported from \`flyteplugins.papermill\` and called from inside the notebook. | Function | Purpose | | ------------------------------------ | ------------------------------------------------------------------------------------------------------------------- | | \`record\_outputs(\*\*kwargs)\` | Records outputs from the \`outputs\`-tagged cell. Must be the cell's last expression. Accepts any Flyte-typed values. | | \`load\_file(path)\` | Reconstructs a \`flyte.io.File\` from the path string injected by papermill. | | \`load\_dir(path)\` | Reconstructs a \`flyte.io.Dir\` from the path string injected by papermill. | | \`load\_dataframe(uri, fmt="parquet")\` | Reconstructs a \`flyte.io.DataFrame\` from the URI string injected by papermill. | === PAGE: https://www.union.ai/docs/v2/flyte/integrations/pytorch === # PyTorch The PyTorch plugin lets you run distributed \[PyTorch\](https://pytorch.org/) training jobs natively on Kubernetes. It uses the \[Kubeflow Training Operator\](https://github.com/kubeflow/training-operator) to manage multi-node training with PyTorch's elastic launch (\`torchrun\`). ## When to use this plugin - Single-node or multi-node distributed training with \`DistributedDataParallel\` (DDP) - Elastic training that can scale up and down during execution - Any workload that uses \`torch.distributed\` for data-parallel or model-parallel training ## Installation \`\`\`bash pip install flyteplugins-pytorch \`\`\` ## Configuration Create an \`Elastic\` configuration and pass it as \`plugin\_config\` to a \`TaskEnvironment\`: \`\`\`python from flyteplugins.pytorch import Elastic torch\_env = flyte.TaskEnvironment( name="torch\_env", resources=flyte.Resources(cpu=(1, 2), memory=("1Gi", "2Gi")), plugin\_config=Elastic( nnodes=2, nproc\_per\_node=1, ), image=image, ) \`\`\` ### \`Elastic\` parameters | Parameter | Type | Description | |-----------|------|-------------| | \`nnodes\` | \`int\` or \`str\` | \*\*Required.\*\* Number of nodes. Use an int for a fixed count or a range string (e.g., \`"2:4"\`) for elastic training | | \`nproc\_per\_node\` | \`int\` | \*\*Required.\*\* Number of processes (workers) per node | | \`rdzv\_backend\` | \`str\` | Rendezvous backend: \`"c10d"\` (default), \`"etcd"\`, or \`"etcd-v2"\` | | \`max\_restarts\` | \`int\` | Maximum worker group restarts (default: \`3\`) | | \`monitor\_interval\` | \`int\` | Agent health check interval in seconds (default: \`3\`) | | \`run\_policy\` | \`RunPolicy\` | Job run policy (cleanup, TTL, deadlines, retries) | ### \`RunPolicy\` parameters | Parameter | Type | Description | |-----------|------|-------------| | \`clean\_pod\_policy\` | \`str\` | Pod cleanup policy: \`"None"\`, \`"all"\`, or \`"Running"\` | | \`ttl\_seconds\_after\_finished\` | \`int\` | Seconds to keep pods after job completion | | \`active\_deadline\_seconds\` | \`int\` | Maximum time the job can run (seconds) | | \`backoff\_limit\` | \`int\` | Number of retries before marking the job as failed | ### NCCL tuning parameters The plugin includes built-in NCCL timeout tuning to reduce failure-detection latency (PyTorch defaults to 1800 seconds): | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | \`nccl\_heartbeat\_timeout\_sec\` | \`int\` | \`300\` | NCCL heartbeat timeout (seconds) | | \`nccl\_async\_error\_handling\` | \`bool\` | \`False\` | Enable async NCCL error handling | | \`nccl\_collective\_timeout\_sec\` | \`int\` | \`None\` | Timeout for NCCL collective operations | | \`nccl\_enable\_monitoring\` | \`bool\` | \`True\` | Enable NCCL monitoring | ### Writing a distributed training task Tasks using this plugin do not need to be \`async\`. Initialize the process group and use \`DistributedDataParallel\` as you normally would with \`torchrun\`: \`\`\`python import torch import torch.distributed from torch.nn.parallel import DistributedDataParallel as DDP @torch\_env.task def train(epochs: int) -> float: torch.distributed.init\_process\_group("gloo") model = DDP(MyModel()) # ... training loop ... return final\_loss \`\`\` > \[!NOTE\] > When \`nnodes=1\`, the task runs as a regular Python task (no Kubernetes training job is created). Set \`nnodes >= 2\` for multi-node distributed training. ## Example \`\`\`python # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-pytorch", # "torch" # \] # main = "torch\_distributed\_train" # params = "3" # /// import typing import torch import torch.distributed import torch.nn as nn import torch.optim as optim from flyteplugins.pytorch.task import Elastic from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler, TensorDataset import flyte image = flyte.Image.from\_debian\_base(name="torch").with\_pip\_packages("flyteplugins-pytorch", pre=True) torch\_env = flyte.TaskEnvironment( name="torch\_env", resources=flyte.Resources(cpu=(1, 2), memory=("1Gi", "2Gi")), plugin\_config=Elastic( nproc\_per\_node=1, # if you want to do local testing set nnodes=1 nnodes=2, ), image=image, ) class LinearRegressionModel(nn.Module): def \_\_init\_\_(self): super().\_\_init\_\_() self.linear = nn.Linear(1, 1) def forward(self, x): return self.linear(x) def prepare\_dataloader(rank: int, world\_size: int, batch\_size: int = 2) -> DataLoader: """ Prepare a DataLoader with a DistributedSampler so each rank gets a shard of the dataset. """ # Dummy dataset x\_train = torch.tensor(\[\[1.0\], \[2.0\], \[3.0\], \[4.0\]\]) y\_train = torch.tensor(\[\[3.0\], \[5.0\], \[7.0\], \[9.0\]\]) dataset = TensorDataset(x\_train, y\_train) # Distributed-aware sampler sampler = DistributedSampler(dataset, num\_replicas=world\_size, rank=rank, shuffle=True) return DataLoader(dataset, batch\_size=batch\_size, sampler=sampler) def train\_loop(epochs: int = 3) -> float: """ A simple training loop for linear regression. """ torch.distributed.init\_process\_group("gloo") model = DDP(LinearRegressionModel()) rank = torch.distributed.get\_rank() world\_size = torch.distributed.get\_world\_size() dataloader = prepare\_dataloader( rank=rank, world\_size=world\_size, batch\_size=64, ) criterion = nn.MSELoss() optimizer = optim.SGD(model.parameters(), lr=0.01) final\_loss = 0.0 for \_ in range(epochs): for x, y in dataloader: outputs = model(x) loss = criterion(outputs, y) optimizer.zero\_grad() loss.backward() optimizer.step() final\_loss = loss.item() if torch.distributed.get\_rank() == 0: print(f"Loss: {final\_loss}") return final\_loss @torch\_env.task def torch\_distributed\_train(epochs: int) -> typing.Optional\[float\]: """ A nested task that sets up a simple distributed training job using PyTorch's """ print("starting launcher") loss = train\_loop(epochs=epochs) print("Training complete") return loss if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() r = flyte.run(torch\_distributed\_train, epochs=3) print(r.name) print(r.url) r.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/pytorch/pytorch\_example.py\* ## API reference See the \[PyTorch API reference\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/pytorch/\_index) for full details. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/ray === # Ray The Ray plugin lets you run \[Ray\](https://www.ray.io/) jobs natively on Kubernetes. Flyte provisions a transient Ray cluster for each task execution using \[KubeRay\](https://github.com/ray-project/kuberay) and tears it down on completion. ## When to use this plugin - Distributed Python workloads (parallel computation, data processing) - ML training with Ray Train or hyperparameter tuning with Ray Tune - Ray Serve inference workloads - Any workload that benefits from Ray's actor model or task parallelism ## Installation \`\`\`bash pip install flyteplugins-ray \`\`\` Your task image must also include a compatible version of Ray: \`\`\`python image = ( flyte.Image.from\_debian\_base(name="ray") .with\_pip\_packages("ray\[default\]==2.46.0", "flyteplugins-ray") ) \`\`\` ## Configuration Create a \`RayJobConfig\` and pass it as \`plugin\_config\` to a \`TaskEnvironment\`: \`\`\`python from flyteplugins.ray import HeadNodeConfig, RayJobConfig, WorkerNodeConfig ray\_config = RayJobConfig( head\_node\_config=HeadNodeConfig(ray\_start\_params={"log-color": "True"}), worker\_node\_config=\[WorkerNodeConfig(group\_name="ray-group", replicas=2)\], runtime\_env={"pip": \["numpy", "pandas"\]}, enable\_autoscaling=False, shutdown\_after\_job\_finishes=True, ttl\_seconds\_after\_finished=300, ) ray\_env = flyte.TaskEnvironment( name="ray\_env", plugin\_config=ray\_config, image=image, ) \`\`\` ### \`RayJobConfig\` parameters | Parameter | Type | Description | |-----------|------|-------------| | \`worker\_node\_config\` | \`List\[WorkerNodeConfig\]\` | \*\*Required.\*\* List of worker group configurations | | \`head\_node\_config\` | \`HeadNodeConfig\` | Head node configuration (optional) | | \`enable\_autoscaling\` | \`bool\` | Enable Ray autoscaler (default: \`False\`) | | \`runtime\_env\` | \`dict\` | Ray runtime environment (pip packages, env vars, etc.) | | \`address\` | \`str\` | Connect to an existing Ray cluster instead of provisioning one | | \`shutdown\_after\_job\_finishes\` | \`bool\` | Shut down the cluster after the job completes (default: \`False\`) | | \`ttl\_seconds\_after\_finished\` | \`int\` | Seconds to keep the cluster after completion before cleanup | ### \`WorkerNodeConfig\` parameters | Parameter | Type | Description | |-----------|------|-------------| | \`group\_name\` | \`str\` | \*\*Required.\*\* Name of this worker group | | \`replicas\` | \`int\` | \*\*Required.\*\* Number of worker replicas | | \`min\_replicas\` | \`int\` | Minimum replicas (for autoscaling) | | \`max\_replicas\` | \`int\` | Maximum replicas (for autoscaling) | | \`ray\_start\_params\` | \`Dict\[str, str\]\` | Ray start parameters for workers | | \`requests\` | \`Resources\` | Resource requests per worker | | \`limits\` | \`Resources\` | Resource limits per worker | | \`pod\_template\` | \`PodTemplate\` | Full pod template (mutually exclusive with \`requests\`/\`limits\`) | ### \`HeadNodeConfig\` parameters | Parameter | Type | Description | |-----------|------|-------------| | \`ray\_start\_params\` | \`Dict\[str, str\]\` | Ray start parameters for the head node | | \`requests\` | \`Resources\` | Resource requests for the head node | | \`limits\` | \`Resources\` | Resource limits for the head node | | \`pod\_template\` | \`PodTemplate\` | Full pod template (mutually exclusive with \`requests\`/\`limits\`) | ### Connecting to an existing cluster To connect to an existing Ray cluster instead of provisioning a new one, set the \`address\` parameter: \`\`\`python ray\_config = RayJobConfig( worker\_node\_config=\[WorkerNodeConfig(group\_name="ray-group", replicas=2)\], address="ray://existing-cluster:10001", ) \`\`\` ## Examples The following example shows how to configure Ray in a \`TaskEnvironment\`. Flyte automatically provisions a Ray cluster for each task using this configuration: \`\`\`python # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-ray", # "ray\[default\]==2.46.0" # \] # main = "hello\_ray\_nested" # params = "3" # /// import asyncio import typing import ray from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig import flyte.remote import flyte.storage @ray.remote def f(x): return x \* x ray\_config = RayJobConfig( head\_node\_config=HeadNodeConfig(ray\_start\_params={"log-color": "True"}), worker\_node\_config=\[WorkerNodeConfig(group\_name="ray-group", replicas=2)\], runtime\_env={"pip": \["numpy", "pandas"\]}, enable\_autoscaling=False, shutdown\_after\_job\_finishes=True, ttl\_seconds\_after\_finished=300, ) image = ( flyte.Image.from\_debian\_base(name="ray") .with\_apt\_packages("wget") .with\_pip\_packages("ray\[default\]==2.46.0", "flyteplugins-ray", "pip", "mypy") ) task\_env = flyte.TaskEnvironment( name="hello\_ray", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image ) ray\_env = flyte.TaskEnvironment( name="ray\_env", plugin\_config=ray\_config, image=image, resources=flyte.Resources(cpu=(3, 4), memory=("3000Mi", "5000Mi")), depends\_on=\[task\_env\], ) @task\_env.task() async def hello\_ray(): await asyncio.sleep(20) print("Hello from the Ray task!") @ray\_env.task async def hello\_ray\_nested(n: int = 3) -> typing.List\[int\]: print("running ray task") t = asyncio.create\_task(hello\_ray()) futures = \[f.remote(i) for i in range(n)\] res = ray.get(futures) await t return res if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() r = flyte.run(hello\_ray\_nested) print(r.name) print(r.url) r.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/ray/ray\_example.py\* The next example demonstrates how Flyte can create ephemeral Ray clusters and run a subtask that connects to an existing Ray cluster: \`\`\`python # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-ray", # "ray\[default\]==2.46.0" # \] # main = "create\_ray\_cluster" # params = "" # /// import os import typing import ray from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig import flyte.storage @ray.remote def f(x): return x \* x ray\_config = RayJobConfig( head\_node\_config=HeadNodeConfig(ray\_start\_params={"log-color": "True"}), worker\_node\_config=\[WorkerNodeConfig(group\_name="ray-group", replicas=2)\], enable\_autoscaling=False, shutdown\_after\_job\_finishes=True, ttl\_seconds\_after\_finished=3600, ) image = ( flyte.Image.from\_debian\_base(name="ray") .with\_apt\_packages("wget") .with\_pip\_packages("ray\[default\]==2.46.0", "flyteplugins-ray") ) task\_env = flyte.TaskEnvironment( name="ray\_client", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image ) ray\_env = flyte.TaskEnvironment( name="ray\_cluster", plugin\_config=ray\_config, image=image, resources=flyte.Resources(cpu=(2, 4), memory=("2000Mi", "4000Mi")), depends\_on=\[task\_env\], ) @task\_env.task() async def hello\_ray(cluster\_ip: str) -> typing.List\[int\]: """ Run a simple Ray task that connects to an existing Ray cluster. """ ray.init(address=f"ray://{cluster\_ip}:10001") futures = \[f.remote(i) for i in range(5)\] res = ray.get(futures) return res @ray\_env.task async def create\_ray\_cluster() -> str: """ Create a Ray cluster and return the head node IP address. """ print("creating ray cluster") cluster\_ip = os.getenv("MY\_POD\_IP") if cluster\_ip is None: raise ValueError("MY\_POD\_IP environment variable is not set") return f"{cluster\_ip}" if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() run = flyte.run(create\_ray\_cluster) run.wait() print("run url:", run.url) print("cluster created, running ray task") print("ray address:", run.outputs()\[0\]) run = flyte.run(hello\_ray, cluster\_ip=run.outputs()\[0\]) print("run url:", run.url) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/ray/ray\_existing\_example.py\* ## API reference See the \[Ray API reference\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/ray/\_index) for full details. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/snowflake === # Snowflake The Snowflake connector lets you run SQL queries against \[Snowflake\](https://www.snowflake.com/) directly from Flyte tasks. Queries are submitted asynchronously and polled for completion, so they don't block a worker while waiting for results. The connector supports: - Parameterized SQL queries with typed inputs - Key-pair and password-based authentication - Returns query results as DataFrames - Automatic links to the Snowflake query dashboard in the Flyte UI - Query cancellation on task abort ## Installation \`\`\`bash pip install flyteplugins-snowflake \`\`\` This installs the Snowflake Python connector and the \`cryptography\` library for key-pair authentication. ## Quick start Here's a minimal example that runs a SQL query on Snowflake: \`\`\`python {hl\_lines=\[2, 4, 12\]} from flyte.io import DataFrame from flyteplugins.connectors.snowflake import Snowflake, SnowflakeConfig config = SnowflakeConfig( account="myorg-myaccount", user="flyte\_user", database="ANALYTICS", schema="PUBLIC", warehouse="COMPUTE\_WH", ) count\_users = Snowflake( name="count\_users", query\_template="SELECT COUNT(\*) FROM users", plugin\_config=config, output\_dataframe\_type=DataFrame, ) \`\`\` This defines a task called \`count\_users\` that runs \`SELECT COUNT(\*) FROM users\` on the configured Snowflake instance. When executed, the connector: 1. Connects to Snowflake using the provided configuration 2. Submits the query asynchronously 3. Polls until the query completes or fails 4. Provides a link to the query in the Snowflake dashboard !\[Snowflake Link\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/snowflake/ui.png) To run the task, create a \`TaskEnvironment\` from it and execute it locally or remotely: \`\`\`python {hl\_lines=3} import flyte snowflake\_env = flyte.TaskEnvironment.from\_task("snowflake\_env", count\_users) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Run locally (connector runs in-process, requires credentials and packages locally) run = flyte.with\_runcontext(mode="local").run(count\_users) # Run remotely (connector runs as a service in your data plane) run = flyte.with\_runcontext(mode="remote").run(count\_users) print(run.url) \`\`\` > \[!NOTE\] > The \`TaskEnvironment\` created by \`from\_task\` does not need an image or pip packages. Snowflake tasks are connector tasks, which means the query executes on the connector service, not in your task container. In \`local\` mode, the connector runs in-process and requires \`flyteplugins-snowflake\` and credentials to be available on your machine. In \`remote\` mode, the connector runs as a service in your data plane. ## Configuration The \`SnowflakeConfig\` dataclass defines the connection settings for your Snowflake instance. ### Required fields | Field | Type | Description | | ----------- | ----- | ------------------------------------------------------- | | \`account\` | \`str\` | Snowflake account identifier (e.g. \`"myorg-myaccount"\`) | | \`database\` | \`str\` | Target database name | | \`schema\` | \`str\` | Target schema name (e.g. \`"PUBLIC"\`) | | \`warehouse\` | \`str\` | Compute warehouse to use for query execution | | \`user\` | \`str\` | Snowflake username | ### Additional connection parameters Use \`connection\_kwargs\` to pass any additional parameters supported by the \[Snowflake Python connector\](https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api). This is a dictionary that gets forwarded directly to \`snowflake.connector.connect()\`. Common options include: | Parameter | Type | Description | | --------------- | ----- | -------------------------------------------------------------------------- | | \`role\` | \`str\` | Snowflake role to use for the session | | \`authenticator\` | \`str\` | Authentication method (e.g. \`"snowflake"\`, \`"externalbrowser"\`, \`"oauth"\`) | | \`token\` | \`str\` | OAuth token when using \`authenticator="oauth"\` | | \`login\_timeout\` | \`int\` | Timeout in seconds for the login request | Example with a role: \`\`\`python {hl\_lines=8} config = SnowflakeConfig( account="myorg-myaccount", user="flyte\_user", database="ANALYTICS", schema="PUBLIC", warehouse="COMPUTE\_WH", connection\_kwargs={ "role": "DATA\_ANALYST", }, ) \`\`\` ## Authentication The connector supports two authentication approaches: key-pair authentication, and password-based or other authentication methods provided through \`connection\_kwargs\`. ### Key-pair authentication Key-pair authentication is the recommended approach for automated workloads. Pass the names of the Flyte secrets containing the private key and optional passphrase: \`\`\`python {hl\_lines=\[5, 6\]} query = Snowflake( name="secure\_query", query\_template="SELECT \* FROM sensitive\_data", plugin\_config=config, snowflake\_private\_key="my-snowflake-private-key", snowflake\_private\_key\_passphrase="my-snowflake-pk-passphrase", ) \`\`\` The \`snowflake\_private\_key\` parameter is the name of the secret (or secret key) that contains your PEM-encoded private key. The \`snowflake\_private\_key\_passphrase\` parameter is the name of the secret (or secret key) that contains the passphrase, if your key is encrypted. If your key is not encrypted, omit the passphrase parameter. The connector decodes the PEM key and converts it to DER format for Snowflake authentication. > \[!NOTE\] > If your credentials are stored in a secret group, you can pass \`secret\_group\` to the \`Snowflake\` task. The plugin expects \`snowflake\_private\_key\` and > \`snowflake\_private\_key\_passphrase\` to be keys within the same secret group. ### Password authentication Send the password via \`connection\_kwargs\`: \`\`\`python {hl\_lines=8} config = SnowflakeConfig( account="myorg-myaccount", user="flyte\_user", database="ANALYTICS", schema="PUBLIC", warehouse="COMPUTE\_WH", connection\_kwargs={ "password": "my-password", }, ) \`\`\` ### OAuth authentication For OAuth-based authentication, specify the authenticator and token: \`\`\`python {hl\_lines=\["8-9"\]} config = SnowflakeConfig( account="myorg-myaccount", user="flyte\_user", database="ANALYTICS", schema="PUBLIC", warehouse="COMPUTE\_WH", connection\_kwargs={ "authenticator": "oauth", "token": "", }, ) \`\`\` ## Query templating Use the \`inputs\` parameter to define typed inputs for your query. Input values are bound using the \`%(param)s\` syntax supported by the \[Snowflake Python connector\](https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api), which handles type conversion and escaping automatically. ### Supported input types The \`inputs\` dictionary maps parameter names to Python values. Supported scalar types include \`str\`, \`int\`, \`float\`, and \`bool\`. To insert multiple rows in a single query, you can also provide lists as input values. When using list inputs, be sure to set \`batch=True\` on the \`Snowflake\` task. This enables automatic batching, where the inputs are expanded and sent as a single multi-row query instead of you having to write multiple individual statements. ### Batched \`INSERT\` with list inputs When \`batch=True\` is enabled, a parameterized \`INSERT\` query with list inputs is automatically expanded into a multi-row \`VALUES\` statement. Example: \`\`\`python query = "INSERT INTO t (a, b) VALUES (%(a)s, %(b)s)" inputs = {"a": \[1, 2\], "b": \["x", "y"\]} \`\`\` This is expanded into: \`\`\`sql INSERT INTO t (a, b) VALUES (%(a\_0)s, %(b\_0)s), (%(a\_1)s, %(b\_1)s) \`\`\` with the following flattened parameters: \`\`\`python flat\_params = { "a\_0": 1, "b\_0": "x", "a\_1": 2, "b\_1": "y", } \`\`\` #### Constraints - The query must contain exactly one \`VALUES (...)\` clause. - All list inputs must have the same non-zero length. ### Parameterized \`SELECT\` \`\`\`python {hl\_lines=\[5, 7\]} from flyte.io import DataFrame events\_by\_date = Snowflake( name="events\_by\_date", query\_template="SELECT \* FROM events WHERE event\_date = %(event\_date)s", plugin\_config=config, inputs={"event\_date": str}, output\_dataframe\_type=DataFrame, ) \`\`\` You can call the task with the required inputs: \`\`\`python {hl\_lines=3} @env.task async def fetch\_events() -> DataFrame: return await events\_by\_date(event\_date="2024-01-15") \`\`\` ### Multiple inputs You can define multiple input parameters of different types: \`\`\`python {hl\_lines=\["4-8", "12-15"\]} filtered\_events = Snowflake( name="filtered\_events", query\_template=""" SELECT \* FROM events WHERE event\_date >= %(start\_date)s AND event\_date <= %(end\_date)s AND region = %(region)s AND score > %(min\_score)s """, plugin\_config=config, inputs={ "start\_date": str, "end\_date": str, "region": str, "min\_score": float, }, output\_dataframe\_type=DataFrame, ) \`\`\` > \[!NOTE\] > The query template is normalized before execution: newlines and tabs are replaced with spaces, and consecutive whitespace is collapsed. You can format your queries across multiple lines for readability without affecting execution. ## Retrieving query results If your query produces output, set \`output\_dataframe\_type\` to capture the results. \`output\_dataframe\_type\` accepts \`DataFrame\` from \`flyte.io\`. This is a meta-wrapper type that represents tabular results and can be materialized into a concrete DataFrame implementation using \`open()\` where you specify the target type and \`all()\`. \`\`\`python {hl\_lines=13} from flyte.io import DataFrame top\_customers = Snowflake( name="top\_customers", query\_template=""" SELECT customer\_id, SUM(amount) AS total\_spend FROM orders GROUP BY customer\_id ORDER BY total\_spend DESC LIMIT 100 """, plugin\_config=config, output\_dataframe\_type=DataFrame, ) \`\`\` At present, only \`pandas.DataFrame\` is supported. The results are returned directly when you call the task: \`\`\`python {hl\_lines=6} import pandas as pd @env.task async def analyze\_top\_customers() -> dict: df = await top\_customers() pandas\_df = await df.open(pd.DataFrame).all() total\_spend = pandas\_df\["total\_spend"\].sum() return {"total\_spend": float(total\_spend)} \`\`\` If you specify \`pandas.DataFrame\` as the \`output\_dataframe\_type\`, you do not need to call \`open()\` and \`all()\` to materialize the results. \`\`\`python {hl\_lines=\[1, 13, "18-19"\]} import pandas as pd top\_customers = Snowflake( name="top\_customers", query\_template=""" SELECT customer\_id, SUM(amount) AS total\_spend FROM orders GROUP BY customer\_id ORDER BY total\_spend DESC LIMIT 100 """, plugin\_config=config, output\_dataframe\_type=pd.DataFrame, ) @env.task async def analyze\_top\_customers() -> dict: df = await top\_customers() total\_spend = df\["total\_spend"\].sum() return {"total\_spend": float(total\_spend)} \`\`\` > \[!NOTE\] > Be sure to inject the \`SNOWFLAKE\_PRIVATE\_KEY\` and \`SNOWFLAKE\_PRIVATE\_KEY\_PASSPHRASE\` environment variables as secrets into your downstream tasks, as they must have access to Snowflake credentials in order to retrieve the DataFrame results. More on how you can create secrets \[here\](https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/secrets/page.md). If you don't need query results (for example, \`DDL\` statements or \`INSERT\` queries), omit \`output\_dataframe\_type\`. ## End-to-end example Here's a complete workflow that uses the Snowflake connector as part of a data pipeline. The workflow creates a staging table, inserts records, queries aggregated results and processes them in a downstream task. \`\`\` import flyte from flyte.io import DataFrame from flyteplugins.connectors.snowflake import Snowflake, SnowflakeConfig config = SnowflakeConfig( account="myorg-myaccount", user="flyte\_user", database="ANALYTICS", schema="PUBLIC", warehouse="COMPUTE\_WH", connection\_kwargs={ "role": "ETL\_ROLE", }, ) # Step 1: Create the staging table if it doesn't exist create\_staging = Snowflake( name="create\_staging", query\_template=""" CREATE TABLE IF NOT EXISTS staging.daily\_events ( event\_id STRING, event\_date DATE, user\_id STRING, event\_type STRING, payload VARIANT ) """, plugin\_config=config, snowflake\_private\_key="snowflake", snowflake\_private\_key\_passphrase="snowflake\_passphrase", ) # Step 2: Insert a record into the staging table insert\_events = Snowflake( name="insert\_event", query\_template=""" INSERT INTO staging.daily\_events (event\_id, event\_date, user\_id, event\_type) VALUES (%(event\_id)s, %(event\_date)s, %(user\_id)s, %(event\_type)s) """, plugin\_config=config, inputs={ "event\_id": list\[str\], "event\_date": list\[str\], "user\_id": list\[str\], "event\_type": list\[str\], }, snowflake\_private\_key="snowflake", snowflake\_private\_key\_passphrase="snowflake\_passphrase", batch=True, ) # Step 3: Query aggregated results for a given date daily\_summary = Snowflake( name="daily\_summary", query\_template=""" SELECT event\_type, COUNT(\*) AS event\_count, COUNT(DISTINCT user\_id) AS unique\_users FROM staging.daily\_events WHERE event\_date = %(report\_date)s GROUP BY event\_type ORDER BY event\_count DESC """, plugin\_config=config, inputs={"report\_date": str}, output\_dataframe\_type=DataFrame, snowflake\_private\_key="snowflake", snowflake\_private\_key\_passphrase="snowflake\_passphrase", ) # Create environments for all Snowflake tasks snowflake\_env = flyte.TaskEnvironment.from\_task( "snowflake\_env", create\_staging, insert\_events, daily\_summary ) # Main pipeline environment that depends on the Snowflake task environments env = flyte.TaskEnvironment( name="analytics\_env", resources=flyte.Resources(memory="512Mi"), image=flyte.Image.from\_debian\_base(name="analytics").with\_pip\_packages( "flyteplugins-snowflake", pre=True ), secrets=\[ flyte.Secret(key="snowflake", as\_env\_var="SNOWFLAKE\_PRIVATE\_KEY"), flyte.Secret( key="snowflake\_passphrase", as\_env\_var="SNOWFLAKE\_PRIVATE\_KEY\_PASSPHRASE" ), \], depends\_on=\[snowflake\_env\], ) # Step 4: Process the results in Python @env.task async def generate\_report(summary: DataFrame) -> dict: import pandas as pd df = await summary.open(pd.DataFrame).all() total\_events = df\["event\_count"\].sum() top\_event = df.iloc\[0\]\["event\_type"\] return { "total\_events": int(total\_events), "top\_event\_type": top\_event, "event\_types\_count": len(df), } # Compose the pipeline @env.task async def run\_daily\_pipeline( event\_ids: list\[str\], event\_dates: list\[str\], user\_ids: list\[str\], event\_types: list\[str\], ) -> dict: await create\_staging() await insert\_events( event\_id=event\_ids, event\_date=event\_dates, user\_id=user\_ids, event\_type=event\_types, ) summary = await daily\_summary(report\_date=event\_dates\[0\]) return await generate\_report(summary=summary) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() # Run locally run = flyte.with\_runcontext(mode="local").run( run\_daily\_pipeline, event\_ids=\["event-1", "event-2"\], event\_dates=\["2023-01-01", "2023-01-02"\], user\_ids=\["user-1", "user-2"\], event\_types=\["click", "view"\], ) # Or run remotely run = flyte.with\_runcontext(mode="remote").run( run\_daily\_pipeline, event\_ids=\["event-1", "event-2"\], event\_dates=\["2023-01-01", "2023-01-02"\], user\_ids=\["user-1", "user-2"\], event\_types=\["click", "view"\], ) print(run.url) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/connectors/snowflake/example.py\* === PAGE: https://www.union.ai/docs/v2/flyte/integrations/spark === # Spark The Spark plugin lets you run \[Apache Spark\](https://spark.apache.org/) jobs natively on Kubernetes. Flyte manages the full cluster lifecycle: provisioning a transient Spark cluster for each task execution, running the job, and tearing the cluster down on completion. Under the hood, the plugin uses the \[Spark on Kubernetes Operator\](https://github.com/GoogleCloudPlatform/spark-on-k8s-operator) to create and manage Spark applications. No external Spark service or long-running cluster is required. ## When to use this plugin - Large-scale data processing and ETL pipelines - Jobs that benefit from Spark's distributed execution engine (Spark SQL, PySpark, Spark MLlib) - Workloads that need Hadoop-compatible storage access (S3, GCS, HDFS) ## Installation \`\`\`bash pip install flyteplugins-spark \`\`\` ## Configuration Create a \`Spark\` configuration and pass it as \`plugin\_config\` to a \`TaskEnvironment\`: \`\`\`python from flyteplugins.spark import Spark spark\_config = Spark( spark\_conf={ "spark.driver.memory": "3000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", }, ) spark\_env = flyte.TaskEnvironment( name="spark\_env", plugin\_config=spark\_config, image=image, ) \`\`\` ### \`Spark\` parameters | Parameter | Type | Description | |-----------|------|-------------| | \`spark\_conf\` | \`Dict\[str, str\]\` | Spark configuration key-value pairs (e.g., executor memory, cores, instances) | | \`hadoop\_conf\` | \`Dict\[str, str\]\` | Hadoop configuration key-value pairs (e.g., S3/GCS access settings) | | \`executor\_path\` | \`str\` | Path to the Python binary for PySpark executors | | \`applications\_path\` | \`str\` | Path to the main Spark application file | | \`driver\_pod\` | \`PodTemplate\` | Pod template for the Spark driver pod | | \`executor\_pod\` | \`PodTemplate\` | Pod template for the Spark executor pods | ### Accessing the Spark session Inside a Spark task, the \`SparkSession\` is available through the task context: \`\`\`python from flyte.\_context import internal\_ctx @spark\_env.task async def my\_spark\_task() -> float: ctx = internal\_ctx() spark = ctx.data.task\_context.data\["spark\_session"\] # Use spark as a normal SparkSession df = spark.read.parquet("s3://my-bucket/data.parquet") return df.count() \`\`\` ### Overriding configuration at runtime You can override Spark configuration for individual task calls using \`.override()\`: \`\`\`python from copy import deepcopy updated\_config = deepcopy(spark\_config) updated\_config.spark\_conf\["spark.executor.instances"\] = "4" result = await my\_spark\_task.override(plugin\_config=updated\_config)() \`\`\` ## Example \`\`\`python # /// script # requires-python = "==3.13" # dependencies = \[ # "flyte>=2.0.0b52", # "flyteplugins-spark" # \] # main = "hello\_spark\_nested" # params = "3" # /// import random from copy import deepcopy from operator import add from flyteplugins.spark.task import Spark import flyte.remote from flyte.\_context import internal\_ctx image = ( flyte.Image.from\_base("apache/spark-py:v3.4.0") .clone(name="spark", python\_version=(3, 10), registry="ghcr.io/flyteorg") .with\_pip\_packages("flyteplugins-spark", pre=True) ) task\_env = flyte.TaskEnvironment( name="get\_pi", resources=flyte.Resources(cpu=(1, 2), memory=("400Mi", "1000Mi")), image=image ) spark\_conf = Spark( spark\_conf={ "spark.driver.memory": "3000M", "spark.executor.memory": "1000M", "spark.executor.cores": "1", "spark.executor.instances": "2", "spark.driver.cores": "1", "spark.kubernetes.file.upload.path": "/opt/spark/work-dir", "spark.jars": "https://storage.googleapis.com/hadoop-lib/gcs/gcs-connector-hadoop3-latest.jar,https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/3.2.2/hadoop-aws-3.2.2.jar,https://repo1.maven.org/maven2/com/amazonaws/aws-java-sdk-bundle/1.12.262/aws-java-sdk-bundle-1.12.262.jar", }, ) spark\_env = flyte.TaskEnvironment( name="spark\_env", resources=flyte.Resources(cpu=(1, 2), memory=("3000Mi", "5000Mi")), plugin\_config=spark\_conf, image=image, depends\_on=\[task\_env\], ) def f(\_): x = random.random() \* 2 - 1 y = random.random() \* 2 - 1 return 1 if x\*\*2 + y\*\*2 <= 1 else 0 @task\_env.task async def get\_pi(count: int, partitions: int) -> float: return 4.0 \* count / partitions @spark\_env.task async def hello\_spark\_nested(partitions: int = 3) -> float: n = 1 \* partitions ctx = internal\_ctx() spark = ctx.data.task\_context.data\["spark\_session"\] count = spark.sparkContext.parallelize(range(1, n + 1), partitions).map(f).reduce(add) return await get\_pi(count, partitions) @task\_env.task async def spark\_overrider(executor\_instances: int = 3, partitions: int = 4) -> float: updated\_spark\_conf = deepcopy(spark\_conf) updated\_spark\_conf.spark\_conf\["spark.executor.instances"\] = str(executor\_instances) return await hello\_spark\_nested.override(plugin\_config=updated\_spark\_conf)(partitions=partitions) if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() r = flyte.run(hello\_spark\_nested) print(r.name) print(r.url) r.wait() \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/spark/spark\_example.py\* ## API reference See the \[Spark API reference\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/spark/\_index) for full details. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb === # Weights & Biases \[Weights & Biases\](https://wandb.ai) (W&B) is a platform for tracking machine learning experiments, visualizing metrics and optimizing hyperparameters. This plugin integrates W&B with Flyte, enabling you to: - Automatically initialize W&B runs in your tasks without boilerplate - Link directly from the Flyte UI to your W&B runs and sweeps - Share W&B runs across parent and child tasks - Track distributed training jobs across multiple GPUs and nodes - Run hyperparameter sweeps with parallel agents ## Installation \`\`\`bash pip install flyteplugins-wandb \`\`\` You also need a W&B API key. Store it as a Flyte secret so your tasks can authenticate with W&B. ## Quick start Here's a minimal example that logs metrics to W&B from a Flyte task: \`\`\` import flyte from flyteplugins.wandb import get\_wandb\_run, wandb\_config, wandb\_init env = flyte.TaskEnvironment( name="wandb-example", image=flyte.Image.from\_debian\_base(name="wandb-example").with\_pip\_packages( "flyteplugins-wandb" ), secrets=\[flyte.Secret(key="wandb\_api\_key", as\_env\_var="WANDB\_API\_KEY")\], ) @wandb\_init @env.task async def train\_model() -> str: wandb\_run = get\_wandb\_run() # Your training code here for epoch in range(10): loss = 1.0 / (epoch + 1) wandb\_run.log({"epoch": epoch, "loss": loss}) return "Training complete" if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() r = flyte.with\_runcontext( custom\_context=wandb\_config( project="my-project", entity="my-team", ), ).run(train\_model) print(f"run url: {r.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/quick\_start.py\* This example demonstrates the core pattern: 1. \*\*Define a task environment\*\* with the plugin installed and your W&B API key as a secret 2. \*\*Decorate your task\*\* with \`@wandb\_init\` (must be the outermost decorator, above \`@env.task\`) 3. \*\*Access the run\*\* with \`get\_wandb\_run()\` to log metrics 4. \*\*Provide configuration\*\* via \`wandb\_config()\` when running the task The plugin handles calling \`wandb.init()\` and \`wandb.finish()\` for you, and automatically adds a link to the W&B run in the Flyte UI. !\[UI\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/wandb/ui.png) ## What's next This integration guide is split into focused sections, depending on how you want to use Weights & Biases with Flyte: - \*\*\*\*Weights & Biases > Experiments\*\*\*\*: Create and manage W&B runs from Flyte tasks. - \*\*\*\*Weights & Biases > Distributed training\*\*\*\*: Track experiments across multi-GPU and multi-node training jobs. - \*\*\*\*Weights & Biases > Sweeps\*\*\*\*: Run hyperparameter searches and manage sweep execution from Flyte tasks. - \*\*\*\*Weights & Biases > Downloading logs\*\*\*\*: Download logs and execution metadata from Weights & Biases. - \*\*\*\*Weights & Biases > Constraints and best practices\*\*\*\*: Learn about limitations, edge cases and recommended patterns. - \*\*\*\*Weights & Biases > Manual integration\*\*\*\*: Use Weights & Biases directly in Flyte tasks without decorators or helpers. > \*\*📝 Note\*\* > > We've included additional examples developed while testing edge cases of the plugin \[here\](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/wandb/examples). === PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/experiments === # Experiments The \`@wandb\_init\` decorator automatically initializes a W&B run when your task executes and finishes it when the task completes. This section covers the different ways to use it. ## Basic usage Apply \`@wandb\_init\` as the outermost decorator on your task: \`\`\`python {hl\_lines=1} @wandb\_init @env.task async def my\_task() -> str: run = get\_wandb\_run() run.log({"metric": 42}) return "done" \`\`\` The decorator: - Calls \`wandb.init()\` before your task code runs - Calls \`wandb.finish()\` after your task completes (or fails) - Adds a link to the W&B run in the Flyte UI You can also use it on synchronous tasks: \`\`\`python {hl\_lines=\[1, 3\]} @wandb\_init @env.task def my\_sync\_task() -> str: run = get\_wandb\_run() run.log({"metric": 42}) return "done" \`\`\` ## Accessing the run object Use \`get\_wandb\_run()\` to access the current W&B run object: \`\`\`python {hl\_lines=6} from flyteplugins.wandb import get\_wandb\_run @wandb\_init @env.task async def train() -> str: run = get\_wandb\_run() # Log metrics run.log({"loss": 0.5, "accuracy": 0.9}) # Access run properties print(f"Run ID: {run.id}") print(f"Run URL: {run.url}") print(f"Project: {run.project}") # Log configuration run.config.update({"learning\_rate": 0.001, "batch\_size": 32}) return run.id \`\`\` ## Parent-child task relationships When a parent task calls child tasks, the plugin can share the same W&B run across all of them. This is useful for tracking an entire workflow in a single run. \`\`\`python {hl\_lines=\[1, 9, 16\]} @wandb\_init @env.task async def child\_task(x: int) -> int: run = get\_wandb\_run() run.log({"child\_metric": x \* 2}) return x \* 2 @wandb\_init @env.task async def parent\_task() -> int: run = get\_wandb\_run() run.log({"parent\_metric": 100}) # Child task shares the parent's run by default result = await child\_task(5) return result \`\`\` By default (\`run\_mode="auto"\`), child tasks reuse their parent's W&B run. All metrics logged by the parent and children appear in the same run in the W&B UI. ## Run modes The \`run\_mode\` parameter controls how tasks create or reuse W&B runs. There are three modes: | Mode | Behavior | | ---------------- | -------------------------------------------------------------------------- | | \`auto\` (default) | Create a new run if no parent run exists, otherwise reuse the parent's run | | \`new\` | Always create a new run, even if a parent run exists | | \`shared\` | Always reuse the parent's run (fails if no parent run exists) | ### Using \`run\_mode="new"\` for independent runs \`\`\`python {hl\_lines=1} @wandb\_init(run\_mode="new") @env.task async def independent\_child(x: int) -> int: run = get\_wandb\_run() # This task gets its own separate run run.log({"independent\_metric": x}) return x @wandb\_init @env.task async def parent\_task() -> str: run = get\_wandb\_run() parent\_run\_id = run.id # This child creates its own run await independent\_child(5) # Parent's run is unchanged assert run.id == parent\_run\_id return parent\_run\_id \`\`\` ### Using \`run\_mode="shared"\` for explicit sharing \`\`\`python {hl\_lines=1} @wandb\_init(run\_mode="shared") @env.task async def must\_share\_run(x: int) -> int: # This task requires a parent run to exist # It will fail if called as a top-level task run = get\_wandb\_run() run.log({"shared\_metric": x}) return x \`\`\` ## Configuration with \`wandb\_config\` Use \`wandb\_config()\` to configure W&B runs. You can set it at the workflow level or override it for specific tasks, allowing you to provide configuration values at runtime. ### Workflow-level configuration \`\`\`python {hl\_lines=\["5-9"\]} if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() flyte.with\_runcontext( custom\_context=wandb\_config( project="my-project", entity="my-team", tags=\["experiment-1", "production"\], config={"model": "resnet50", "dataset": "imagenet"}, ), ).run(train\_task) \`\`\` ### Overriding configuration for child tasks Use \`wandb\_config()\` as a context manager to override settings for specific child task calls: \`\`\`python {hl\_lines=\[8, 12\]} @wandb\_init @env.task async def parent\_task() -> str: run = get\_wandb\_run() run.log({"parent\_metric": 100}) # Override tags and config for this child call with wandb\_config(tags=\["special-run"\], config={"learning\_rate": 0.01}): await child\_task(10) # Override run\_mode for this child call with wandb\_config(run\_mode="new"): await child\_task(20) # Gets its own run return "done" \`\`\` ## Using traces with W&B runs Flyte traces can access the parent task's W&B run without needing the \`@wandb\_init\` decorator. This is useful for helper functions that should log to the same run: \`\`\`python {hl\_lines=\[1, 3\]} @flyte.trace async def log\_validation\_metrics(accuracy: float, f1: float): run = get\_wandb\_run() if run: run.log({"val\_accuracy": accuracy, "val\_f1": f1}) @wandb\_init @env.task async def train\_and\_validate() -> str: run = get\_wandb\_run() # Training loop for epoch in range(10): run.log({"train\_loss": 1.0 / (epoch + 1)}) # Trace logs to the same run await log\_validation\_metrics(accuracy=0.95, f1=0.92) return "done" \`\`\` > \*\*📝 Note\*\* > > Do not apply \`@wandb\_init\` to traces. Traces automatically access the parent task's run via \`get\_wandb\_run()\`. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/distributed\_training === # Distributed training When running distributed training jobs, multiple processes run simultaneously across GPUs. The \`@wandb\_init\` decorator automatically detects distributed training environments and coordinates W&B logging across processes. The plugin: - Auto-detects distributed context from environment variables (set by launchers like \`torchrun\`) - Controls which processes initialize W&B runs based on the \`run\_mode\` and \`rank\_scope\` parameters - Generates unique run IDs that distinguish between workers and ranks - Adds links to W&B runs in the Flyte UI ## Quick start Here's a minimal single-node example that logs metrics from a distributed training task. By default (\`run\_mode="auto"\`, \`rank\_scope="global"\`), only rank 0 logs to W&B: \`\`\` import flyte import torch import torch.distributed from flyteplugins.pytorch.task import Elastic from flyteplugins.wandb import get\_wandb\_run, wandb\_config, wandb\_init image = flyte.Image.from\_debian\_base(name="torch-wandb").with\_pip\_packages( "flyteplugins-wandb", "flyteplugins-pytorch" ) env = flyte.TaskEnvironment( name="distributed\_env", image=image, resources=flyte.Resources(gpu="A100:2"), plugin\_config=Elastic(nproc\_per\_node=2, nnodes=1), secrets=flyte.Secret(key="wandb\_api\_key", as\_env\_var="WANDB\_API\_KEY"), ) @wandb\_init @env.task def train() -> float: torch.distributed.init\_process\_group("nccl") # Only rank 0 gets a W&B run object; others get None run = get\_wandb\_run() # Simulate training for step in range(100): loss = 1.0 / (step + 1) # Safe to call on all ranks - only rank 0 actually logs if run: run.log({"loss": loss, "step": step}) torch.distributed.destroy\_process\_group() return loss if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() flyte.with\_runcontext( custom\_context=wandb\_config(project="my-project", entity="my-team") ).run(train) \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/distributed\_training\_quick\_start.py\* A few things to note: 1. Use the \`Elastic\` plugin to configure distributed training (number of processes, nodes) 2. Apply \`@wandb\_init\` as the outermost decorator 3. Check if \`run\` is not None before logging - only the primary rank has a run object in \`auto\` mode > \*\*📝 Note\*\* > > The \`if run:\` check is always safe regardless of run mode. In \`shared\` and \`new\` modes all ranks get a run object, but the check doesn't hurt and keeps your code portable across modes. !\[Single-node auto\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/wandb/single\_node\_auto\_flyte.png) ## Run modes in distributed training The \`run\_mode\` parameter controls how W&B runs are created across distributed processes. The behavior differs between single-node (one machine, multiple GPUs) and multi-node (multiple machines) setups. ### Single-node behavior | Mode | Which ranks log | Result | | ---------------- | --------------------- | -------------------------------------- | | \`auto\` (default) | Only rank 0 | 1 W&B run | | \`shared\` | All ranks to same run | 1 W&B run with metrics labeled by rank | | \`new\` | Each rank separately | N W&B runs (grouped in UI) | ### Multi-node behavior For multi-node training, the \`rank\_scope\` parameter controls the granularity of W&B runs: - \*\*\`global\`\*\* (default): Treat all workers as one unit - \*\*\`worker\`\*\*: Treat each worker/node independently The combination of \`run\_mode\` and \`rank\_scope\` determines logging behavior: | \`run\_mode\` | \`rank\_scope\` | Who initializes W&B | W&B Runs | Grouping | | ---------- | ------------ | ---------------------- | -------- | -------- | | \`auto\` | \`global\` | Global rank 0 only | 1 | - | | \`auto\` | \`worker\` | Local rank 0 per worker | N | - | | \`shared\` | \`global\` | All ranks (shared globally) | 1 | - | | \`shared\` | \`worker\` | All ranks (shared per worker) | N | - | | \`new\` | \`global\` | All ranks | N × M | 1 group | | \`new\` | \`worker\` | All ranks | N × M | N groups | Where \`N\` = number of workers/nodes, \`M\` = processes per worker. ### Choosing run mode and rank scope - \*\*\`auto\`\*\* (recommended): Use when you want clean dashboards with minimal runs. Most metrics (loss, accuracy) are the same across ranks after gradient synchronization, so logging from one rank is sufficient. - \*\*\`shared\`\*\*: Use when you need to compare metrics across ranks in a single view. Each rank's metrics are labeled with an \`x\_label\` identifier. Useful for debugging load imbalance or per-GPU throughput. - \*\*\`new\`\*\*: Use when you need completely separate runs per GPU, for example to track GPU-specific metrics or compare training dynamics across devices. For multi-node training: - Use \*\*\`rank\_scope="global"\`\*\* (default) for most cases. A single consolidated run across all nodes is sufficient since metrics like loss and accuracy converge after gradient synchronization. - Use \*\*\`rank\_scope="worker"\`\*\* for debugging and per-node analysis. This is useful when you need to inspect data distribution across nodes, compare predictions from different workers, or track metrics on individual batches outside the main node. ## Single-node multi-GPU For single-node distributed training, configure the \`Elastic\` plugin with \`nnodes=1\` and set \`nproc\_per\_node\` to your GPU count. ### Basic example with \`auto\` mode \`\`\`python {hl\_lines=\["6-7", 13, 18, 30\]} import os import torch import torch.distributed import flyte from flyteplugins.pytorch.task import Elastic from flyteplugins.wandb import wandb\_init, get\_wandb\_run env = flyte.TaskEnvironment( name="single\_node\_env", image=image, resources=flyte.Resources(gpu="A100:4"), plugin\_config=Elastic(nproc\_per\_node=4, nnodes=1), secrets=flyte.Secret(key="wandb\_api\_key", as\_env\_var="WANDB\_API\_KEY"), ) @wandb\_init # run\_mode="auto" (default) @env.task def train\_single\_node() -> float: torch.distributed.init\_process\_group("nccl") rank = torch.distributed.get\_rank() local\_rank = int(os.environ.get("LOCAL\_RANK", 0)) device = torch.device(f"cuda:{local\_rank}") torch.cuda.set\_device(device) run = get\_wandb\_run() # Training loop - only rank 0 logs for epoch in range(10): loss = train\_epoch(model, dataloader, device) if run: run.log({"epoch": epoch, "loss": loss}) torch.distributed.destroy\_process\_group() return loss \`\`\` ### Using \`shared\` mode for per-rank metrics When you need to see metrics from all GPUs in a single run, use \`run\_mode="shared"\`: \`\`\`python {hl\_lines=\[3, 13, 19\]} import os @wandb\_init(run\_mode="shared") @env.task def train\_with\_per\_gpu\_metrics() -> float: torch.distributed.init\_process\_group("nccl") rank = torch.distributed.get\_rank() local\_rank = int(os.environ.get("LOCAL\_RANK", 0)) device = torch.device(f"cuda:{local\_rank}") torch.cuda.set\_device(device) # In shared mode, all ranks get a run object run = get\_wandb\_run() for step in range(1000): loss, throughput = train\_step(model, batch, device) # Each rank logs with automatic x\_label identification if run: run.log({ "loss": loss, "throughput\_samples\_per\_sec": throughput, "gpu\_memory\_used": torch.cuda.memory\_allocated(device), }) torch.distributed.destroy\_process\_group() return loss \`\`\` !\[Single-node shared\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/wandb/single\_node\_shared\_flyte.png) In the W&B UI, metrics from each rank appear with distinct labels, allowing you to compare GPU utilization and throughput across devices. !\[Single-node shared W&B UI\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/wandb/single\_node\_shared\_wandb.png) ### Using \`new\` mode for per-rank runs When you need completely separate W&B runs for each GPU, use \`run\_mode="new"\`. Each rank gets its own run, and runs are grouped together in the W&B UI: \`\`\`python {hl\_lines=\[1, "11-12"\]} @wandb\_init(run\_mode="new") # Each rank gets its own run @env.task def train\_per\_rank() -> float: torch.distributed.init\_process\_group("nccl") rank = torch.distributed.get\_rank() # ... # Each rank has its own W&B run run = get\_wandb\_run() # Run IDs: {base}-rank-{rank} # All runs are grouped under {base} in W&B UI run.log({"train/loss": loss.item(), "rank": rank}) # ... \`\`\` With \`run\_mode="new"\`: - Each rank creates its own W&B run - Run IDs follow the pattern \`{run\_name}-{action\_name}-rank-{rank}\` - All runs are grouped together in the W&B UI for comparison ## Multi-node training with \`Elastic\` For multi-node distributed training, set \`nnodes\` to your node count. The \`rank\_scope\` parameter controls whether you get a single W&B run across all nodes (\`global\`) or one run per node (\`worker\`). ### Global scope (default): Single run across all nodes With \`run\_mode="auto"\` and \`rank\_scope="global"\` (both defaults), only global rank 0 initializes W&B, resulting in a single run for the entire distributed job: \`\`\`python {hl\_lines=\["11-12", "27-30", "35", "59-60", "95-98"\]} import os import torch import torch.distributed import torch.nn as nn import torch.optim as optim from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler import flyte from flyteplugins.pytorch.task import Elastic from flyteplugins.wandb import wandb\_init, wandb\_config, get\_wandb\_run image = flyte.Image.from\_debian\_base(name="torch-wandb").with\_pip\_packages( "flyteplugins-wandb", "flyteplugins-pytorch", pre=True ) multi\_node\_env = flyte.TaskEnvironment( name="multi\_node\_env", image=image, resources=flyte.Resources( cpu=(1, 2), memory=("1Gi", "10Gi"), gpu="A100:4", shm="auto", ), plugin\_config=Elastic( nproc\_per\_node=4, # GPUs per node nnodes=2, # Number of nodes ), secrets=flyte.Secret(key="wandb\_api\_key", as\_env\_var="WANDB\_API\_KEY"), ) @wandb\_init # rank\_scope="global" by default → 1 run total @multi\_node\_env.task def train\_multi\_node(epochs: int, batch\_size: int) -> float: torch.distributed.init\_process\_group("nccl") rank = torch.distributed.get\_rank() world\_size = torch.distributed.get\_world\_size() local\_rank = int(os.environ.get("LOCAL\_RANK", 0)) device = torch.device(f"cuda:{local\_rank}") torch.cuda.set\_device(device) # Model with DDP model = MyModel().to(device) model = DDP(model, device\_ids=\[local\_rank\]) # Distributed data loading dataset = MyDataset() sampler = DistributedSampler(dataset, num\_replicas=world\_size, rank=rank) dataloader = DataLoader(dataset, batch\_size=batch\_size, sampler=sampler) optimizer = optim.AdamW(model.parameters(), lr=1e-3) criterion = nn.CrossEntropyLoss() # Only global rank 0 gets a W&B run run = get\_wandb\_run() for epoch in range(epochs): sampler.set\_epoch(epoch) model.train() for batch\_idx, (data, target) in enumerate(dataloader): data, target = data.to(device), target.to(device) optimizer.zero\_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() if run and batch\_idx % 100 == 0: run.log({ "train/loss": loss.item(), "train/epoch": epoch, "train/batch": batch\_idx, }) if run: run.log({"train/epoch\_complete": epoch}) # Barrier ensures all ranks finish before cleanup torch.distributed.barrier() torch.distributed.destroy\_process\_group() return loss.item() if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() flyte.with\_runcontext( custom\_context=wandb\_config( project="multi-node-training", tags=\["distributed", "multi-node"\], ) ).run(train\_multi\_node, epochs=10, batch\_size=32) \`\`\` With this configuration: - Two nodes run the task, each with 4 GPUs (8 total processes) - Only global rank 0 creates a W&B run - Run ID follows the pattern \`{run\_name}-{action\_name}\` - The Flyte UI shows a single link to the W&B run ### Worker scope: One run per node Use \`rank\_scope="worker"\` when you want each node to have its own W&B run for per-node analysis: \`\`\`python {hl\_lines=\[1, 8\]} @wandb\_init(rank\_scope="worker") # 1 run per worker/node @multi\_node\_env.task def train\_per\_worker(epochs: int, batch\_size: int) -> float: torch.distributed.init\_process\_group("nccl") local\_rank = int(os.environ.get("LOCAL\_RANK", 0)) # ... # Local rank 0 of each worker gets a W&B run run = get\_wandb\_run() if run: # Each worker logs to its own run run.log({"train/loss": loss.item()}) # ... \`\`\` With \`run\_mode="auto"\`, \`rank\_scope="worker"\`: - Each node's local rank 0 creates a W&B run - Run IDs follow the pattern \`{run\_name}-{action\_name}-worker-{worker\_index}\` - The Flyte UI shows links to each worker's W&B run !\[Multi-node\](https://raw.githubusercontent.com/unionai/unionai-docs-static/refs/heads/main/images/integrations/wandb/multi\_node.png) ### Shared mode: All ranks log to the same run Use \`run\_mode="shared"\` when you need metrics from all ranks in a single view. Each rank's metrics are labeled with an \`x\_label\` identifier. #### Shared + global scope (1 run total) \`\`\`python {hl\_lines=\[1, 7\]} @wandb\_init(run\_mode="shared") # All ranks log to 1 shared run @multi\_node\_env.task def train\_shared\_global() -> float: torch.distributed.init\_process\_group("nccl") # ... # All ranks get a run object, all log to the same run run = get\_wandb\_run() # Each rank logs with automatic x\_label identification run.log({"train/loss": loss.item(), "rank": rank}) # ... \`\`\` #### Shared + worker scope (N runs, 1 per node) \`\`\`python {hl\_lines=\[1, 7, 10\]} @wandb\_init(run\_mode="shared", rank\_scope="worker") # 1 shared run per worker @multi\_node\_env.task def train\_shared\_worker() -> float: torch.distributed.init\_process\_group("nccl") # ... # All ranks get a run object, grouped by worker run = get\_wandb\_run() # Ranks on the same worker share a run run.log({"train/loss": loss.item(), "local\_rank": local\_rank}) # ... \`\`\` ### New mode: Separate run per rank Use \`run\_mode="new"\` when you need completely separate runs per GPU. Runs are grouped in the W&B UI for easy comparison. #### New + global scope (N×M runs, 1 group) \`\`\`python {hl\_lines=\[1, 7, 10\]} @wandb\_init(run\_mode="new") # Each rank gets its own run, all in 1 group @multi\_node\_env.task def train\_new\_global() -> float: torch.distributed.init\_process\_group("nccl") # ... # Each rank has its own run run = get\_wandb\_run() # Run IDs: {base}-rank-{global\_rank} run.log({"train/loss": loss.item()}) # ... \`\`\` #### New + worker scope (N×M runs, N groups) \`\`\`python {hl\_lines=\[1, 7, 10\]} @wandb\_init(run\_mode="new", rank\_scope="worker") # Each rank gets own run, grouped per worker @multi\_node\_env.task def train\_new\_worker() -> float: torch.distributed.init\_process\_group("nccl") # ... # Each rank has its own run, grouped by worker run = get\_wandb\_run() # Run IDs: {base}-worker-{idx}-rank-{local\_rank} run.log({"train/loss": loss.item()}) # ... \`\`\` ## How it works The plugin automatically detects distributed training by checking environment variables set by distributed launchers like \`torchrun\`: | Environment variable | Description | | -------------------- | -------------------------------------------------------- | | \`RANK\` | Global rank across all processes | | \`WORLD\_SIZE\` | Total number of processes | | \`LOCAL\_RANK\` | Rank within the current node | | \`LOCAL\_WORLD\_SIZE\` | Number of processes on the current node | | \`GROUP\_RANK\` | Node/worker index (0 for first node, 1 for second, etc.) | When these variables are present, the plugin: 1. \*\*Determines which ranks should initialize W&B\*\* based on \`run\_mode\` and \`rank\_scope\` 2. \*\*Generates unique run IDs\*\* that include worker and rank information 4. \*\*Creates UI links\*\* for each W&B run (single link with \`rank\_scope="global"\`, one per worker with \`rank\_scope="worker"\`) The plugin automatically adapts to your training setup, eliminating the need for manual distributed configuration. ### Run ID patterns | Scenario | Run ID Pattern | Group | | ---------------------------- | --------------------------------------------- | ------------------------ | | Single-node auto/shared | \`{base}\` | - | | Single-node new | \`{base}-rank-{rank}\` | \`{base}\` | | Multi-node auto/shared (global) | \`{base}\` | - | | Multi-node auto/shared (worker) | \`{base}-worker-{idx}\` | - | | Multi-node new (global) | \`{base}-rank-{global\_rank}\` | \`{base}\` | | Multi-node new (worker) | \`{base}-worker-{idx}-rank-{local\_rank}\` | \`{base}-worker-{idx}\` | Where \`{base}\` = \`{run\_name}-{action\_name}\` === PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/sweeps === # Sweeps W&B sweeps automate hyperparameter optimization by running multiple trials with different parameter combinations. The \`@wandb\_sweep\` decorator creates a sweep and makes it easy to run trials in parallel using Flyte's distributed execution. ## Creating a sweep Use \`@wandb\_sweep\` to create a W&B sweep when the task executes: \`\`\` import flyte import wandb from flyteplugins.wandb import ( get\_wandb\_sweep\_id, wandb\_config, wandb\_init, wandb\_sweep, wandb\_sweep\_config, ) env = flyte.TaskEnvironment( name="wandb-example", image=flyte.Image.from\_debian\_base(name="wandb-example").with\_pip\_packages( "flyteplugins-wandb" ), secrets=\[flyte.Secret(key="wandb\_api\_key", as\_env\_var="WANDB\_API\_KEY")\], ) @wandb\_init def objective(): """Objective function that W&B calls for each trial.""" wandb\_run = wandb.run config = wandb\_run.config # Simulate training with hyperparameters from the sweep for epoch in range(config.epochs): loss = 1.0 / (config.learning\_rate \* config.batch\_size) + epoch \* 0.1 wandb\_run.log({"epoch": epoch, "loss": loss}) @wandb\_sweep @env.task async def run\_sweep() -> str: sweep\_id = get\_wandb\_sweep\_id() # Run 10 trials wandb.agent(sweep\_id, function=objective, count=10) return sweep\_id if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() r = flyte.with\_runcontext( custom\_context={ \*\*wandb\_config(project="my-project", entity="my-team"), \*\*wandb\_sweep\_config( method="random", metric={"name": "loss", "goal": "minimize"}, parameters={ "learning\_rate": {"min": 0.0001, "max": 0.1}, "batch\_size": {"values": \[16, 32, 64, 128\]}, "epochs": {"values": \[5, 10, 20\]}, }, ), }, ).run(run\_sweep) print(f"run url: {r.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/sweep.py\* The \`@wandb\_sweep\` decorator: - Creates a W&B sweep when the task starts - Makes the sweep ID available via \`get\_wandb\_sweep\_id()\` - Adds a link to the main sweeps page in the Flyte UI Use \`wandb\_sweep\_config()\` to define the sweep parameters. This is passed to W&B's sweep API. > \*\*📝 Note\*\* > > Random and Bayesian searches run indefinitely, and the sweep remains in the \`Running\` state until you stop it. > You can stop a running sweep from the Weights & Biases UI or from the command line. ## Running parallel agents Flyte's distributed execution makes it easy to run multiple sweep agents in parallel, each on its own compute resources: \`\`\` import asyncio from datetime import timedelta import flyte import wandb from flyteplugins.wandb import ( get\_wandb\_sweep\_id, wandb\_config, wandb\_init, wandb\_sweep, wandb\_sweep\_config, get\_wandb\_context, ) env = flyte.TaskEnvironment( name="wandb-parallel-sweep-example", image=flyte.Image.from\_debian\_base( name="wandb-parallel-sweep-example" ).with\_pip\_packages("flyteplugins-wandb"), secrets=\[flyte.Secret(key="wandb\_api\_key", as\_env\_var="WANDB\_API\_KEY")\], ) @wandb\_init def objective(): wandb\_run = wandb.run config = wandb\_run.config for epoch in range(config.epochs): loss = 1.0 / (config.learning\_rate \* config.batch\_size) + epoch \* 0.1 wandb\_run.log({"epoch": epoch, "loss": loss}) @wandb\_sweep @env.task async def sweep\_agent(agent\_id: int, sweep\_id: str, count: int = 5) -> int: """Single agent that runs a subset of trials.""" wandb.agent( sweep\_id, function=objective, count=count, project=get\_wandb\_context().project ) return agent\_id @wandb\_sweep @env.task async def run\_parallel\_sweep(total\_trials: int = 20, trials\_per\_agent: int = 5) -> str: """Orchestrate multiple agents running in parallel.""" sweep\_id = get\_wandb\_sweep\_id() num\_agents = (total\_trials + trials\_per\_agent - 1) // trials\_per\_agent # Launch agents in parallel, each with its own resources agent\_tasks = \[ sweep\_agent.override( resources=flyte.Resources(cpu="2", memory="4Gi"), retries=3, timeout=timedelta(minutes=30), )(agent\_id=i, sweep\_id=sweep\_id, count=trials\_per\_agent) for i in range(num\_agents) \] await asyncio.gather(\*agent\_tasks) return sweep\_id if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() r = flyte.with\_runcontext( custom\_context={ \*\*wandb\_config(project="my-project", entity="my-team"), \*\*wandb\_sweep\_config( method="random", metric={"name": "loss", "goal": "minimize"}, parameters={ "learning\_rate": {"min": 0.0001, "max": 0.1}, "batch\_size": {"values": \[16, 32, 64\]}, "epochs": {"values": \[5, 10, 20\]}, }, ), }, ).run( run\_parallel\_sweep, total\_trials=20, trials\_per\_agent=5, ) print(f"run url: {r.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/parallel\_sweep.py\* This pattern provides: - \*\*Distributed execution\*\*: Each agent runs on separate compute nodes - \*\*Resource allocation\*\*: Specify CPU, memory, and GPU per agent - \*\*Fault tolerance\*\*: Failed agents can retry without affecting others - \*\*Timeout protection\*\*: Prevent runaway trials > \*\*📝 Note\*\* > > \`run\_parallel\_sweep\` links to the main Weights & Biases sweeps page and \`sweep\_agent\` links to the specific sweep URL because we cannot determine the sweep ID at link rendering time. !\[Sweep\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/wandb/sweep.png) ## Writing objective functions The objective function is called by \`wandb.agent()\` for each trial. It must be a regular Python function decorated with \`@wandb\_init\`: \`\`\`python {hl\_lines=\["1-2", "5-6"\]} @wandb\_init def objective(): """Objective function for sweep trials.""" # Access hyperparameters from wandb.run.config run = wandb.run config = run.config # Your training code model = create\_model( learning\_rate=config.learning\_rate, hidden\_size=config.hidden\_size, ) for epoch in range(config.epochs): train\_loss = train\_epoch(model) val\_loss = validate(model) # Log metrics - W&B tracks these for the sweep run.log({ "epoch": epoch, "train\_loss": train\_loss, "val\_loss": val\_loss, }) # The final val\_loss is used by the sweep to rank trials \`\`\` Key points: - Use \`@wandb\_init\` on the objective function (not \`@env.task\`) - Access hyperparameters via \`wandb.run.config\` (not \`get\_wandb\_run()\` since this is outside Flyte context) - Log the metric specified in \`wandb\_sweep\_config(metric=...)\` so the sweep can optimize it - The function is called multiple times by \`wandb.agent()\`, once per trial === PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/downloading\_logs === # Downloading logs This integration enables downloading Weights & Biases run data, including metrics history, summary data, and synced files. ## Automatic download Set \`download\_logs=True\` to automatically download run data after your task completes: \`\`\`python {hl\_lines=1} @wandb\_init(download\_logs=True) @env.task async def train\_with\_download(): run = get\_wandb\_run() for epoch in range(10): run.log({"loss": 1.0 / (epoch + 1)}) return run.id \`\`\` The downloaded data is traced by Flyte and appears as a \`Dir\` output in the Flyte UI. Downloaded files include: - \`summary.json\`: Final summary metrics - \`metrics\_history.json\`: Step-by-step metrics history - Any files synced by W&B (\`requirements.txt\`, \`wandb\_metadata.json\`, etc.) You can also set \`download\_logs=True\` in \`wandb\_config()\`: \`\`\`python {hl\_lines=5} flyte.with\_runcontext( custom\_context=wandb\_config( project="my-project", entity="my-team", download\_logs=True, ), ).run(train\_task) \`\`\` !\[Logs\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/wandb/logs.png) For sweeps, set \`download\_logs=True\` on \`@wandb\_sweep\` or \`wandb\_sweep\_config()\` to download all trial data: \`\`\`python {hl\_lines=1} @wandb\_sweep(download\_logs=True) @env.task async def run\_sweep(): sweep\_id = get\_wandb\_sweep\_id() wandb.agent(sweep\_id, function=objective, count=10) return sweep\_id \`\`\` !\[Sweep Logs\](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/wandb/sweep\_logs.png) ## Accessing run directories during execution Use \`get\_wandb\_run\_dir()\` to access the local W&B run directory during task execution. This is useful for writing custom files that get synced to W&B: \`\`\`python {hl\_lines=\[1, 7, "18-19"\]} from flyteplugins.wandb import get\_wandb\_run\_dir @wandb\_init @env.task def train\_with\_artifacts(): run = get\_wandb\_run() local\_dir = get\_wandb\_run\_dir() # Train your model for epoch in range(10): run.log({"loss": 1.0 / (epoch + 1)}) # Save model checkpoint to the run directory model\_path = f"{local\_dir}/model\_checkpoint.pt" torch.save(model.state\_dict(), model\_path) # Save custom metrics file with open(f"{local\_dir}/custom\_metrics.json", "w") as f: json.dump({"final\_accuracy": 0.95}, f) return run.id \`\`\` Files written to the run directory are automatically synced to W&B and can be accessed later via the W&B UI or by setting \`download\_logs=True\`. > \*\*📝 Note\*\* > > \`get\_wandb\_run\_dir()\` accesses the local directory without making network calls. Files written here may have a brief delay before appearing in the W&B cloud. === PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/constraints\_and\_best\_practices === # Constraints and best practices ## Decorator ordering \`@wandb\_init\` and \`@wandb\_sweep\` must be the \*\*outermost decorators\*\*, applied after \`@env.task\`: \`\`\`python # Correct @wandb\_init @env.task async def my\_task(): ... # Incorrect - will not work @env.task @wandb\_init async def my\_task(): ... \`\`\` ## Traces cannot use decorators Do not apply \`@wandb\_init\` to traces. Traces automatically access the parent task's run via \`get\_wandb\_run()\`: \`\`\`python # Correct @flyte.trace async def my\_trace(): run = get\_wandb\_run() if run: run.log({"metric": 42}) # Incorrect - don't decorate traces @wandb\_init @flyte.trace async def my\_trace(): ... \`\`\` ## Maximum sweep agents \[W&B limits sweeps to a maximum of 20 concurrent agents\](https://docs.wandb.ai/models/sweeps/existing-project#3-launch-agents). ## Configuration priority Configuration is merged with the following priority (highest to lowest): 1. Decorator parameters (\`@wandb\_init(project="...")\`) 2. Context manager (\`with wandb\_config(...)\`) 3. Workflow-level context (\`flyte.with\_runcontext(custom\_context=wandb\_config(...))\`) 4. Auto-generated values (run ID from Flyte context) ## Run ID generation When no explicit \`id\` is provided, the plugin generates run IDs using the pattern: \`\`\` {run\_name}-{action\_name} \`\`\` This ensures unique, predictable IDs that can be matched between the \`Wandb\` link class and manual \`wandb.init()\` calls. ## Sync delay for local files Files written to the run directory (via \`get\_wandb\_run\_dir()\`) are synced to W&B asynchronously. There may be a brief delay before they appear in the W&B cloud or can be downloaded via \`download\_wandb\_run\_dir()\`. ## Shared run mode requirements When using \`run\_mode="shared"\`, the task requires a parent task to have already created a W&B run. Calling a task with \`run\_mode="shared"\` as a top-level task will fail. ## Objective functions for sweeps Objective functions passed to \`wandb.agent()\` should: - Be regular Python functions (not Flyte tasks) - Be decorated with \`@wandb\_init\` - Access hyperparameters via \`wandb.run.config\` (not \`get\_wandb\_run()\`) - Log the metric specified in \`wandb\_sweep\_config(metric=...)\` so the sweep can optimize it ## Error handling The plugin raises standard exceptions: - \`RuntimeError\`: When \`download\_wandb\_run\_dir()\` is called without a run ID and no active run exists - \`wandb.errors.AuthenticationError\`: When \`WANDB\_API\_KEY\` is not set or invalid - \`wandb.errors.CommError\`: When a run cannot be found in the W&B cloud === PAGE: https://www.union.ai/docs/v2/flyte/integrations/wandb/manual === # Manual integration If you need more control over W&B initialization, you can use the \`Wandb\` and \`WandbSweep\` link classes directly instead of the decorators. This lets you call \`wandb.init()\` and \`wandb.finish()\` yourself while still getting automatic links in the Flyte UI. ## Using the Wandb link class Add a \`Wandb\` link to your task to generate a link to the W&B run in the Flyte UI: \`\`\` import flyte import wandb from flyteplugins.wandb import Wandb env = flyte.TaskEnvironment( name="wandb-manual-init-example", image=flyte.Image.from\_debian\_base( name="wandb-manual-init-example" ).with\_pip\_packages("flyteplugins-wandb"), secrets=\[flyte.Secret(key="wandb\_api\_key", as\_env\_var="WANDB\_API\_KEY")\], ) @env.task( links=( Wandb( project="my-project", entity="my-team", run\_mode="new", # No id parameter - link will auto-generate from run\_name-action\_name ), ) ) async def train\_model(learning\_rate: float) -> str: ctx = flyte.ctx() # Generate run ID matching the link's auto-generated ID run\_id = f"{ctx.action.run\_name}-{ctx.action.name}" # Manually initialize W&B wandb\_run = wandb.init( project="my-project", entity="my-team", id=run\_id, config={"learning\_rate": learning\_rate}, ) # Your training code for epoch in range(10): loss = 1.0 / (learning\_rate \* (epoch + 1)) wandb\_run.log({"epoch": epoch, "loss": loss}) # Manually finish the run wandb\_run.finish() return wandb\_run.id if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() r = flyte.with\_runcontext().run( train\_model, learning\_rate=0.01, ) print(f"run url: {r.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/init\_manual.py\* ### With a custom run ID If you want to use your own run ID, specify it in both the link and the \`wandb.init()\` call: \`\`\`python {hl\_lines=\[6, 14\]} @env.task( links=( Wandb( project="my-project", entity="my-team", id="my-custom-run-id", ), ) ) async def train\_with\_custom\_id() -> str: run = wandb.init( project="my-project", entity="my-team", id="my-custom-run-id", # Must match the link's ID resume="allow", ) # Training code... run.finish() return run.id \`\`\` ### Adding links at runtime with override You can also add links when calling a task using \`.override()\`: \`\`\`python {hl\_lines=9} @env.task async def train\_model(learning\_rate: float) -> str: # ... training code with manual wandb.init() ... return run.id # Add link when running the task result = await train\_model.override( links=(Wandb(project="my-project", entity="my-team", run\_mode="new"),) )(learning\_rate=0.01) \`\`\` ## Using the \`WandbSweep\` link class Use \`WandbSweep\` to add a link to a W&B sweep: \`\`\` import flyte import wandb from flyteplugins.wandb import WandbSweep env = flyte.TaskEnvironment( name="wandb-manual-sweep-example", image=flyte.Image.from\_debian\_base( name="wandb-manual-sweep-example" ).with\_pip\_packages("flyteplugins-wandb"), secrets=\[flyte.Secret(key="wandb\_api\_key", as\_env\_var="WANDB\_API\_KEY")\], ) def objective(): with wandb.init(project="my-project", entity="my-team") as wandb\_run: config = wandb\_run.config for epoch in range(config.epochs): loss = 1.0 / (config.learning\_rate \* config.batch\_size) + epoch \* 0.1 wandb\_run.log({"epoch": epoch, "loss": loss}) @env.task( links=( WandbSweep( project="my-project", entity="my-team", ), ) ) async def manual\_sweep() -> str: # Manually create the sweep sweep\_config = { "method": "random", "metric": {"name": "loss", "goal": "minimize"}, "parameters": { "learning\_rate": {"min": 0.0001, "max": 0.1}, "batch\_size": {"values": \[16, 32, 64\]}, "epochs": {"value": 10}, }, } sweep\_id = wandb.sweep(sweep\_config, project="my-project", entity="my-team") # Run the sweep wandb.agent(sweep\_id, function=objective, count=10, project="my-project") return sweep\_id if \_\_name\_\_ == "\_\_main\_\_": flyte.init\_from\_config() r = flyte.with\_runcontext().run(manual\_sweep) print(f"run url: {r.url}") \`\`\` \*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/sweep\_manual.py\* The link will point to the project's sweeps page. If you have the sweep ID, you can specify it in the link: \`\`\`python {hl\_lines=6} @env.task( links=( WandbSweep( project="my-project", entity="my-team", id="known-sweep-id", ), ) ) async def resume\_sweep() -> str: # Resume an existing sweep wandb.agent("known-sweep-id", function=objective, count=10) return "known-sweep-id" \`\`\` --- # Unknown \# Community Flyte is an open source project that is built and maintained by a community of contributors. Union AI is the primary maintainer of Flyte and developer of Union.ai, a closed source commercial product that is built on top of Flyte. Since the success of Flyte is essential to the success of Union.ai, the company is dedicated to building and expanding the Flyte open source project and community. For information on how to get involved and how to keep in touch, see \[Joining the community\](https://www.union.ai/docs/v2/flyte/community/joining-the-community/page.md). ## Contributing to the codebase The full Flyte codebase is open source and available on GitHub. If you are interested, see \[Contributing code\](https://www.union.ai/docs/v2/flyte/community/contributing-code/page.md). ## Contributing to documentation Union AI maintains and hosts both Flyte and Union documentation at \[www.union.ai/docs\](https://www.union.ai/docs/v2/root/). The two sets of documentation are deeply integrated, as the Union product is built on top of Flyte. To better maintain both sets of docs, they are hosted in the same repository (but rendered so that you can choose to view one or the other). Both the Flyte and Union documentation are open source. Flyte community members and Union customers are both welcome to contribute to the documentation. If you are interested, see \[Contributing documentation and examples\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/\_index). ## Subpages - \[Joining the community\](https://www.union.ai/docs/v2/flyte/community/joining-the-community/page.md) - Community sync - Contributor's sync - Newsletter - Slack guidelines - Abide by the \[LF's Code of Conduct\](https://lfprojects.org/policies/code-of-conduct/) - Avoid using DMs and @mentions - Make use of threads - Do not post the same question across multiple channels - Do not solicit members of our Slack - \[Contributing code\](https://www.union.ai/docs/v2/flyte/community/contributing-code/page.md) - Flyte 2 - Becoming a contributor - Before submitting your PR - 🐞 File an issue - Component Reference - \`flyte\` - \`flyteidl\` - \`flytepropeller\` - \`flyteadmin\` - \`flytekit\` - \`flyteconsole\` - \`datacatalog\` - \`flyteplugins\` - \`flytestdlib\` - \`flytectl\` - Development Environment Setup Guide - Requirements - Content - How to setup dev environment for flyteidl, flyteadmin, flyteplugins, flytepropeller, datacatalog and flytestdlib? - How to setup dev environment for flytekit? - How to setup dev environment for flyteconsole? - How to access Flyte UI, minio, postgres, k3s, and endpoints? - \[Contributing docs and examples\](https://www.union.ai/docs/v2/flyte/community/contributing-docs/page.md) - The combined Flyte and Union docs site - Versions - Common build infrastructure - Variants - Both Flyte and Union docs are open source --- \*\*Source\*\*: https://github.com/unionai/unionai-docs/blob/main/content/community/\_index.md \*\*HTML\*\*: https://www.union.ai/docs/v2/flyte/community/ --- # Unknown \# Reference This section provides the reference material for the Flyte SDK and CLI. To get started, add \`flyte\` to your project \`\`\`shell $ uv pip install --no-cache --upgrade flyte \`\`\` This will install both the Flyte SDK and CLI. ### \[Flyte SDK\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/page.md) The Flyte SDK provides the core Python API for building workflows and apps on your Union instance. ### \[Flyte CLI\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/page.md) The Flyte CLI is the command-line interface for interacting with your Union instance. ### \[Migration from Flyte 1\](https://www.union.ai/docs/v2/flyte/api-reference/migration/page.md) Comprehensive reference for migrating Flyte 1 workflows to Flyte 2. ## Subpages - \[LLM-optimized documentation\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-context/page.md) - Per-page Markdown (\`page.md\`) - Section bundles (\`section.md\`) - Page index (\`llms.txt\`) - Full documentation (\`llms-full.txt\`) - \[Migration from Flyte 1\](https://www.union.ai/docs/v2/flyte/api-reference/migration/page.md) - Key API changes at a glance - Topics - \[Philosophy and imports\](https://www.union.ai/docs/v2/flyte/api-reference/overview/page.md) - \[Container images\](https://www.union.ai/docs/v2/flyte/api-reference/images/page.md) - \[Configuration and CLI\](https://www.union.ai/docs/v2/flyte/api-reference/configuration-and-cli/page.md) - \[Tasks and workflows\](https://www.union.ai/docs/v2/flyte/api-reference/tasks-and-workflows/page.md) - \[Secrets, resources, and caching\](https://www.union.ai/docs/v2/flyte/api-reference/secrets-resources-caching/page.md) - \[Parallelism and async\](https://www.union.ai/docs/v2/flyte/api-reference/parallelism-and-async/page.md) - \[Triggers and dynamic workflows\](https://www.union.ai/docs/v2/flyte/api-reference/triggers-and-dynamic/page.md) - \[Examples and common gotchas\](https://www.union.ai/docs/v2/flyte/api-reference/examples-and-gotchas/page.md) - \[Flyte CLI\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-cli/page.md) - flyte - flyte abort - flyte build - flyte create - flyte delete - flyte deploy - flyte edit - flyte gen - flyte get - flyte prefetch - flyte run - flyte serve - flyte signal - flyte start - flyte stop - flyte update - flyte whoami - \[Flyte SDK\](https://www.union.ai/docs/v2/flyte/api-reference/flyte-sdk/page.md) - \[Integrations\](https://www.union.ai/docs/v2/flyte/api-reference/integrations/page.md) --- \*\*Source\*\*: https://github.com/unionai/unionai-docs/blob/main/content/api-reference/\_index.md \*\*HTML\*\*: https://www.union.ai/docs/v2/flyte/api-reference/ --- # Unknown \# Flyte Open Source Documentation > \*\*This is legacy (v1) documentation.\*\* Do not use unless explicitly asked about this version. For current documentation, see https://www.union.ai/docs/v2/llms.txt > Full documentation (single file): https://www.union.ai/docs/v1/flyte/llms-full.txt > Site: https://www.union.ai/docs/v1/flyte Each entry below is \`- \[Page title\](URL)\` followed by the H2/H3 headings found on that page. Pages link to individual \`page.md\` files. Sections marked with a "Section bundle" link have a \`section.md\` that concatenates all pages in the section into a single file — use it to load an entire section into context at once. ## User guide - \[Introduction\](https://www.union.ai/docs/v1/flyte/user-guide/introduction/page.md) - Flyte - Trying out Flyte - Flyte in production - \[Getting started\](https://www.union.ai/docs/v1/flyte/user-guide/getting-started/page.md) - Try Flyte on your local machine > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/user-guide/getting-started/section.md - \[Getting started > Local setup\](https://www.union.ai/docs/v1/flyte/user-guide/getting-started/local-setup/page.md) - Install \`uv\` - Ensure the correct version of Python is installed - Install the \`pyflyte\` CLI - Install Docker and get access to a container registry - Install \`flytectl\` to set up a local cluster - macOS - Linux - Windows - Start Docker and the local cluster - Configure the connection to your cluster - Check your CLI configuration - \[Getting started > First project\](https://www.union.ai/docs/v1/flyte/user-guide/getting-started/first-project/page.md) - Create a new Flyte project - Initialize a local project - \[Getting started > Understanding the code\](https://www.union.ai/docs/v1/flyte/user-guide/getting-started/understanding-the-code/page.md) - Python code - ImageSpec - Tasks - Workflow - pyproject.toml - uv.lock - \[Getting started > Running your workflow\](https://www.union.ai/docs/v1/flyte/user-guide/getting-started/running-your-workflow/page.md) - Python virtual environment - Run the code locally - Running remotely on Flyte in the cloud - Register the workflow without running - Run the workflow from the Flyte interface - View the workflow execution on Flyte - \[Core concepts\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/page.md) - Defining tasks and workflows - Type annotation is required - Workflows \*are not\* full Python functions - Registering tasks and workflows - Registering on the command line with \`pyflyte\` or \`flytectl\` - Registering in Python with \`FlyteRemote\` - Results of registration - Changing tasks and workflows - Inspecting tasks and workflows - Inspecting workflows in the UI - Inspecting tasks in the UI - Inspecting workflows on the command line with \`flytectl\` - Inspecting tasks on the command line with \`flytectl\` - Inspecting tasks and workflows in Python with \`FlyteRemote\` - Running tasks and workflows - Running a task or workflow in the UI - Running a task or workflow locally on the command line with \`pyflyte\` or \`python\` - Running a task or workflow remotely on the command line with \`pyflyte\` - Running a task or workflow remotely in Python with \`FlyteRemote\` > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/section.md - \[Core concepts > Workflows\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows/page.md) > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows/section.md - \[Core concepts > Workflows > Standard workflows\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows/standard-workflows/page.md) - Evaluation of a standard workflow - Conditional construct - Chaining operator - Workflow decorator parameters - \[Core concepts > Workflows > Subworkflows and sub-launch plans\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows/subworkflows-and-sub-launch-plans/page.md) - When to use subworkflows - When to use sub-launch plans - \[Core concepts > Workflows > Dynamic workflows\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows/dynamic-workflows/page.md) - Defining a dynamic workflow - Advantages of dynamic workflows - Flexibility - Lower pressure on \`etcd\` - Dynamic workflows vs. map tasks - Using dynamic workflows to achieve recursion - \[Core concepts > Workflows > Imperative workflows\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows/imperative-workflows/page.md) - Example - \[Core concepts > Workflows > Launching workflows\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows/launching-workflows/page.md) - \[Core concepts > Workflows > Viewing workflows\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows/viewing-workflows/page.md) - Workflows list - Workflow view - Workflow versions list - Workflow and task descriptions - \[Core concepts > Workflows > Viewing workflow executions\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows/viewing-workflow-executions/page.md) - Domain Settings - All Executions in the Project - Execution view - Nodes - Graph - Timeline - \[Core concepts > Tasks\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/page.md) - Tasks are independently executable - Tasks are strongly typed - Tasks are containerized - Tasks are named, versioned, and immutable - Tasks are (usually) deterministic and cacheable - Workflows can contain many types of tasks - Mix and match task characteristics - Task configuration > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/section.md - \[Core concepts > Tasks > Page\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/map-tasks/page.md) - Map tasks - \[Core concepts > Tasks > Other task types\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-types/page.md) - PythonFunctionTask - ContainerTask - Shell tasks - Example - Specialized plugin task classes and configs - @fl.task parameters - \[Core concepts > Tasks > Task parameters\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-parameters/page.md) - Use \`partial\` to provide default arguments to tasks - Named outputs - \[Core concepts > Tasks > Launching tasks\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/launching-tasks/page.md) - \[Core concepts > Tasks > Viewing tasks\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/viewing-tasks/page.md) - Tasks list - Task view - Task versions list - \[Core concepts > Tasks > Task software environment\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-software-environment/page.md) - \[Core concepts > Tasks > Task software environment > Local image building\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-software-environment/image-spec/page.md) - Project structure - requirements.txt - imagespec-simple-example.py - Install and configure \`pyflyte\` and Docker - Set up an image registry - Authenticate to the registry - Set up your project and domain on Flyte - Understand the requirements - Set up a virtual Python environment - Run the workflow locally - Register the workflow - Ensure that the image is publicly accessible - Run the workflow on Flyte - Multi-image workflows - \[Core concepts > Tasks > Task software environment > ImageSpec with ECR\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-software-environment/image-spec-with-ecr/page.md) - Set up the image repository - Authenticate to the registry - Register your workflow to Flyte - \[Core concepts > Tasks > Task software environment > ImageSpec with GAR\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-software-environment/image-spec-with-gar/page.md) - Set up the image repository - Authenticate to the registry - Register your workflow to Flyte - \[Core concepts > Tasks > Task software environment > ImageSpec with ACR\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-software-environment/image-spec-with-acr/page.md) - Authenticate to the registry - Register your workflow to Flyte - \[Core concepts > Tasks > Task software environment > Environment variables\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-software-environment/environment-variables/page.md) - \[Core concepts > Tasks > Viewing logs\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/viewing-logs/page.md) - Cloud provider logs - \[Core concepts > Tasks > Reference tasks\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/reference-tasks/page.md) - Example - \[Core concepts > Tasks > Task hardware environment\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-hardware-environment/page.md) - Customizing task resources - Using the \`@fl.task\` decorator - Using PodTemplate - \`pod\_template\` and \`pod\_template\_name\` @fl.task parameters - Accelerators - Task-level monitoring - \[Core concepts > Tasks > Task hardware environment > Customizing task resources\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-hardware-environment/customizing-task-resources/page.md) - The \`requests\` and \`limits\` settings - The \`accelerator\` setting - The \`with\_overrides\` method - \[Core concepts > Tasks > Task hardware environment > Accelerators\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-hardware-environment/accelerators/page.md) - Using predefined accelerator constants - List of predefined accelerator constants - \[Core concepts > Tasks > Task hardware environment > Retries and timeouts\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-hardware-environment/retries-and-timeouts/page.md) - Retry types - Configuring retries - Retrying interruptible tasks - Retrying map tasks - Timeouts - \[Core concepts > Tasks > Task hardware environment > Interruptible instances\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks/task-hardware-environment/interruptible-instances/page.md) - Configuring tasks to use interruptible instances - Workflow level interruptible - Advantages and disadvantages of interruptible instances - \[Core concepts > Launch plans\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans/page.md) - Default launch plan - Launch plans are versioned - Custom launch plans - Viewing launch plans for a workflow - Registering a launch plan - Registering a launch plan on the command line - Registering a launch plan in Python with \`FlyteRemote\` - Results of registration - Changing a launch plan > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans/section.md - \[Core concepts > Launch plans > Defining launch plans\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans/defining-launch-plans/page.md) - Default and Fixed Inputs - Scheduled Execution - Labels and Annotations - Execution Parameters - Security and Authentication - Raw Output Data Configuration - Putting It All Together - \[Core concepts > Launch plans > Viewing launch plans\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans/viewing-launch-plans/page.md) - Viewing launch plans in the UI - Viewing launch plans on the command line with \`uctl\` - Viewing launch plans in Python with \`FlyteRemote\` - \[Core concepts > Launch plans > Notifications\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans/notifications/page.md) - \[Core concepts > Launch plans > Schedules\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans/schedules/page.md) - Fixed-rate schedules - Cron schedules - Cron expression format - Cron expression examples - Cron aliases - kickoff\_time\_input\_arg - \[Core concepts > Launch plans > Activating and deactivating\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans/activating-and-deactivating/page.md) - Activating and deactivating a launch plan in the UI - Activating and deactivating a launch plan on the command line with \`uctl\` - Activating and deactivating a launch plan in Python with \`FlyteRemote\` - \[Core concepts > Launch plans > Running launch plans\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans/running-launch-plans/page.md) - Running a launch plan in the UI - Running a launch plan on the command line with \`uctl\` - Running a launch plan in Python with \`FlyteRemote\` - Sub-launch plans - \[Core concepts > Launch plans > Reference launch plans\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans/reference-launch-plans/page.md) - Example - \[Core concepts > Launch plans > Concurrency control\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans/concurrency-control/page.md) - How it works - Basic usage - Scheduled workflows with concurrency control - Defining the policy - Key behaviors and considerations - Version-agnostic check, version-specific enforcement - Concurrency limit on manual trigger - Scheduled execution behavior - Limitations - "At most" enforcement - Notifications for skipped executions - Best practices - \[Core concepts > Caching\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/caching/page.md) - Inputs caching - Outputs caching - Enabling and configuring caching - The \`Cache\` object - The \`overwrite-cache\` flag - Overwrite cache on the command line - Overwrite cache in the UI - Overwrite cache programmatically - How caching works - Explicit cache version - Node signature - Caching when running locally - Cache serialization - Enabling cache serialization - How does cache serialization work? - Caching of offloaded objects - How does caching of offloaded objects work? - \[Core concepts > Named outputs\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/named-outputs/page.md) - \[Core concepts > ImageSpec\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/image-spec/page.md) - Install Python or APT packages - Install Conda packages - Use different Python versions in the image - Import modules only in a specify imageSpec environment - Install CUDA in the image - Use Nvidia docker image - Install packages from extra index - Build an image in different architecture - Install flytekit from GitHub - Customize the tag of the image - Copy additional files or directories - Define ImageSpec in a YAML File - Build the image without registering the workflow - Force push an image - Getting source files into ImageSpec - \[Development cycle\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/page.md) > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/section.md - \[Development cycle > Project structure\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/project-structure/page.md) - Recommended Directory Structure - Organizing Tasks and Workflows - Orchestration Directory for Helper Constructs - Core Logic for Workflow-Specific Functionality - Importance of \`\_\_init\_\_.py\` - Monorepo vs Multi-repo: Choosing a structure - CI/CD - Documentation and Docstrings - \[Development cycle > Projects and domains\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/projects-and-domains/page.md) - Projects - Domains - When to use different Flyte projects? - Projects and Domains: The Power of the Project-Domain Pair - Domains: Clear Environment Separation - Projects: Organizing Workflows by Teams, Business Areas, or Applications - \[Development cycle > Building workflows\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/building-workflows/page.md) - When should I decompose tasks? - Differing runtime requirements - Improved cache performance - Take advantage of interruptible tasks - When should I parallelize tasks? - Parallelization constructs - When should I use caching? - \[Development cycle > Setting up a production project\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/setting-up-a-project/page.md) - Terminology - Create a Flyte project - Creating a local production project directory using \`pyflyte init\` - Directory structure - \[Development cycle > Local dependencies\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/local-dependencies/page.md) - Define your dependencies in your \`pyproject.toml\` - Create a Python virtual environment - \[Development cycle > ImageSpec\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/image-spec/page.md) - \[Development cycle > Running your code\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/running-your-code/page.md) - Set up your development environment - CLI commands for running your code - Registration pattern summary - Running a script in local Python with \`pyflyte run\` {#running-a-script-in-local-python} - Running a script on Flyte with \`pyflyte run --remote\` - Running tasks through flytectl - Generate execution spec file - Update the input spec file for arguments to the workflow - Create execution using the exec spec file - Monitor the execution by providing the execution id from create command - Running workflows through flytectl - Running launchplans through flytectl - Generate an execution spec file - Update the input spec file for arguments to the workflow - Create execution using the exec spec file - Monitor the execution by providing the execution id from create command - Deploying your code to Flyte with \`pyflyte register\` - Fast registration - Inspecting executions - Deploying your code to production - Package your code with \`pyflyte package\` - Register the package with \`flytectl register\` - Using pyflyte register versus pyflyte package + flytectl register - Image management and registration method - Building your own images - CI/CD with Flyte and GitHub Actions - Some CI/CD best practices - \[Development cycle > Overriding parameters\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/overriding-parameters/page.md) - Task parameters - Using \`with\_overrides\` with \`name\` and \`node\_name\` - Subworkflow and sub-launch plan parameters - \[Development cycle > Run details\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/run-details/page.md) - Passing parameters - Why \`pyflyte run\` rather than \`python\`? - \[Development cycle > Debugging with interactive tasks\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/debugging-with-interactive-tasks/page.md) - Enabling interactive tasks in your code - Basic example - requirements.txt - example.py - Register and run the workflow - Access the IDE - Inspect the task code - Interactive debugging - Update your code - Resume task - Auxiliary Python files - flyteinteractive\_interactive\_entrypoint.py - flyteinteractive\_resume\_task.py - launch.json - Integrated terminal - Install extensions - example-extensions.py - Manage resources - example-manage-resources.py - Pre and post hooks - example-pre-post-hooks.py - Only initiate VSCode on task failure - example-run-task-first.py - Debugging execution issues - \[Development cycle > Task resource validation\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/task-resource-validation/page.md) - \[Development cycle > Running in a local cluster\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/running-in-a-local-cluster/page.md) - Running in a local Kubernetes cluster - Configuration - Start the workflow - Inspect the results - Local cluster with default image - Local cluster with custom image - \[Development cycle > Jupyter notebooks\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/jupyter-notebooks/page.md) - Write your workflows and tasks in cells - Enable the notebook to register workflows to Flyte - \[Development cycle > Decks\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/decks/page.md) - Deck tabs - Deck renderers - Frame profiling renderer - Top-frame renderer - Markdown renderer - Box renderer - Image renderer - Contribute to renderers - Custom renderers - Streaming Decks - Union Deck Succeed Video - Union Deck Fail Video - \[Development cycle > Migrating from Airflow to Flyte\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/migrating\_from\_airflow\_to\_flyte/page.md) - Prerequisites - Steps - 1. Define your Airflow tasks in a Flyte workflow - 2. Test your workflow locally - 3. Move your workflow to production - \[Development cycle > FlyteRemote\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/remote-management/page.md) - Creating a \`FlyteRemote\` object - Authenticating using a client secret - \[Development cycle > FlyteRemote > FlyteRemote examples\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/remote-management/remote-examples/page.md) - Registering and running a workflow - Fetching outputs - Terminating all running executions for a workflow - Rerunning all failed executions of a workflow - Filtering for executions using a \`Filter\` - Launch task via FlyteRemote with a new version - Launch workflow via FlyteRemote - Launch launchplan via FlyteRemote - Inspecting executions - \[Development cycle > Testing\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/testing/page.md) - \[Development cycle > Testing > Mocking tasks\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/testing/mocking-tasks/page.md) - \[Data input/output\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/page.md) - Mapping Python to Flyte types > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/section.md - \[Data input/output > FlyteFile and FlyteDirectory\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/flyte-file-and-flyte-directory/page.md) - FlyteFile - Mac OS - Linux - Streaming support - FlyteDirectory - Changing the data upload location - Changing the raw data prefix - Specifying \`remote\_path\` for a \`FlyteFile\` or \`FlyteDirectory\` - Remote examples - Remote file example - Remote directory example - Streaming - Downloading - Implicit downloading - Explicit downloading - Typed aliases - \[Data input/output > Downloading with FlyteFile and FlyteDirectory\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/downloading-with-ff-and-fd/page.md) - FlyteFile - FlyteDirectory - \[Data input/output > Task input and output\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/task-input-and-output/page.md) - Metadata and raw data - Metadata store - Raw data store - Changing the raw data storage location - Setting up your own object store - \[Data input/output > Accessing attributes\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/accessing-attributes/page.md) - List - Dictionary - Data class - Complex type - Failure scenario - \[Data input/output > Dataclass\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/dataclass/page.md) - Python types - Flyte types - \[Data input/output > Enum type\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/enum/page.md) - \[Data input/output > Pickle type\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/pickle/page.md) - \[Data input/output > Pydantic BaseModel\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/pydantic/page.md) - Python types - Flyte types - \[Data input/output > PyTorch type\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/pytorch/page.md) - Tensors and modules - Checkpoint - Auto GPU to CPU and CPU to GPU conversion - \[Data input/output > StructuredDataset\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/structured-dataset/page.md) - Usage - Example - Column type information - Serialized byte format - Storage driver and location - Inner workings of a structured dataset plugin - The \`uri\` argument - Note that no format was specified in the structured dataset constructor, or in the signature. So how did the BigQuery encoder get invoked? - How to return multiple DataFrames from a task? - How to define a custom structured dataset plugin? - NumPy encoder - NumPy decoder - NumPy renderer - The nested typed columns - \[Data input/output > TensorFlow types\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/tensorflow/page.md) - Import necessary libraries and modules - Tensorflow model - Transformer - Usage - TFRecord files - Transformer - Usage - TFRecord directories - Transformer - Usage - Configuration class: \`TFRecordDatasetConfig\` - Attributes - \[Programming\](https://www.union.ai/docs/v1/flyte/user-guide/programming/page.md) > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/user-guide/programming/section.md - \[Programming > Chaining Entities\](https://www.union.ai/docs/v1/flyte/user-guide/programming/chaining-entities/page.md) - Tasks - Subworkflows - \[Programming > Conditionals\](https://www.union.ai/docs/v1/flyte/user-guide/programming/conditionals/page.md) - Simple branch - Multiple branches - Consuming the output of a conditional - Using the output of a previous task in a conditional - Using boolean workflow inputs in a conditional - Nested conditionals - Using the output of a task in a conditional - Running a noop task in a conditional - Run the example on the Flyte cluster - \[Programming > Decorating tasks\](https://www.union.ai/docs/v1/flyte/user-guide/programming/decorating\_tasks/page.md) - Using a single decorator - Stacking multiple decorators - Run the example on Flyte - \[Programming > Decorating workflows\](https://www.union.ai/docs/v1/flyte/user-guide/programming/decorating\_workflows/page.md) - Setup-teardown pattern - Workflow decorator - Defining the DAG - Run the example on the Flyte cluster - \[Programming > Intratask checkpoints\](https://www.union.ai/docs/v1/flyte/user-guide/programming/intratask\_checkpoints/page.md) - Why intratask checkpoints? - Use case: Model training - Run the example on the Flyte cluster - \[Programming > Waiting for external inputs\](https://www.union.ai/docs/v1/flyte/user-guide/programming/waiting\_for\_external\_inputs/page.md) - Pause executions with the \`sleep\` node - Supply external inputs with \`wait\_for\_input\` - Continue executions with \`approve\` - Working with conditionals - Sending inputs to \`wait\_for\_input\` and \`approve\` nodes - Using the Flyte UI - Using \`FlyteRemote\` - \[Programming > Nested parallelism\](https://www.union.ai/docs/v1/flyte/user-guide/programming/nested-parallelism/page.md) - Nested dynamic workflows - Example code - Mixed parallelism - Example code - Design considerations - \[Programming > Failure node\](https://www.union.ai/docs/v1/flyte/user-guide/programming/failure-node/page.md) --- ## Tutorials - \[Bioinformatics\](https://www.union.ai/docs/v1/flyte/tutorials/bioinformatics/page.md) - \[Bioinformatics > Nucleotide Sequence Querying with BLASTX\](https://www.union.ai/docs/v1/flyte/tutorials/bioinformatics/blast/page.md) - BLAST - BLASTX - Data - Dockerfile - \[Bioinformatics > Nucleotide Sequence Querying with BLASTX > Page\](https://www.union.ai/docs/v1/flyte/tutorials/bioinformatics/blast/blastx-example/page.md) - \[Feature engineering\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/page.md) - \[Feature engineering > EDA, Feature Engineering, and Modeling With Papermill\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/exploratory-data-analysis/page.md) - Papermill - Examples - Notebook Etiquette - \[Feature engineering > EDA, Feature Engineering, and Modeling With Papermill > Page\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/exploratory-data-analysis/notebook-and-task/page.md) - \[Feature engineering > EDA, Feature Engineering, and Modeling With Papermill > Page\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/exploratory-data-analysis/notebooks-as-tasks/page.md) - \[Feature engineering > EDA, Feature Engineering, and Modeling With Papermill > Page\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/exploratory-data-analysis/notebook/page.md) - \[Feature engineering > EDA, Feature Engineering, and Modeling With Papermill > Supermarket Regression 2 Notebook\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/exploratory-data-analysis/supermarket\_regression\_1/page.md) - \[Feature engineering > EDA, Feature Engineering, and Modeling With Papermill > Supermarket Regression 2 Notebook\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/exploratory-data-analysis/supermarket\_regression\_2/page.md) - \[Feature engineering > EDA, Feature Engineering, and Modeling With Papermill > Supermarket Regression Notebook\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/exploratory-data-analysis/supermarket\_regression/page.md) - \[Feature engineering > Feast Integration\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/feast-integration/page.md) - Dataset - Takeaways - \[Feature engineering > Feast Integration > Page\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/feast-integration/feature\_eng\_tasks/page.md) - \[Feature engineering > Feast Integration > Page\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/feast-integration/feast\_workflow/page.md) - \[Feature engineering > Feast Integration > How to Trigger the Feast Workflow using FlyteRemote\](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/feast-integration/feast\_flyte\_remote/page.md) - 01. Register the code - 02: Launch an execution - 03. Sync an execution - 04. Retrieve the output - 05. Generate predictions - Load features from the online feature store - Generate a prediction - \[Flytelab\](https://www.union.ai/docs/v1/flyte/tutorials/flytelab/page.md) - \[Flytelab > Weather Forecasting\](https://www.union.ai/docs/v1/flyte/tutorials/flytelab/weather-forecasting/page.md) - \[Model training\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/page.md) - \[Model training > Diabetes Classification\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/pima-diabetes/page.md) - Why a Workflow? - Pros: - Cons: - Steps of the Pipeline - Takeaways - \[Model training > Diabetes Classification > Page\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/pima-diabetes/diabetes/page.md) - \[Model training > Forecasting Rossman Store Sales with Horovod and Spark\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/forecasting-sales/page.md) - About Horovod - About Spark - Horovod and Spark - Flyte and Spark - \[Model training > Forecasting Rossman Store Sales with Horovod and Spark > Page\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/forecasting-sales/keras-spark-rossmann-estimator/page.md) - \[Model training > House Price Regression\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/house-price-prediction/page.md) - Where Does Flyte Fit In? - Dataset - Takeaways - \[Model training > House Price Regression > Page\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/house-price-prediction/house-price-predictor/page.md) - \[Model training > House Price Regression > Page\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/house-price-prediction/multiregion-house-price-predictor/page.md) - \[Model training > MNIST Classification With PyTorch and W&B\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/mnist-classifier/page.md) - PyTorch - Model Development - Specify GPU Requirement - Distributed Data-Parallel Training - Weights & Biases Integration - PyTorch Dockerfile for Deployment - \[Model training > MNIST Classification With PyTorch and W&B > Page\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/mnist-classifier/pytorch-single-node-multi-gpu/page.md) - \[Model training > MNIST Classification With PyTorch and W&B > Page\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/mnist-classifier/pytorch-single-node-and-gpu/page.md) - \[Model training > NLP Processing\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/nlp-processing/page.md) - About Gensim - Data - Step-by-Step Process - \[Model training > NLP Processing > Page\](https://www.union.ai/docs/v1/flyte/tutorials/model-training/nlp-processing/word2vec-and-lda/page.md) --- ## Integrations - \[Connectors\](https://www.union.ai/docs/v1/flyte/integrations/connectors/page.md) - Creating a new connector - Async connector interface specification - Sync connector interface specification - Testing your connector locally - Enabling a connector in your Flyte deployment - \[Connectors > Airflow connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/airflow-connector/page.md) - Installation - Example usage - Local testing - Flyte deployment configuration - \[Connectors > Airflow connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/airflow-connector/airflow-connector-example-usage/page.md) - \[Connectors > BigQuery connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/bigquery-connector/page.md) - Installation - Example usage - Local testing - Flyte deployment configuration - \[Connectors > BigQuery connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/bigquery-connector/bigquery-connector-example-usage/page.md) - \[Connectors > ChatGPT connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/chatgpt-connector/page.md) - Installation - Example usage - Local testing - Flyte deployment configuration - \[Connectors > ChatGPT connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/chatgpt-connector/chatgpt-connector-example-usage/page.md) - \[Connectors > Databricks connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/databricks-connector/page.md) - Installation - Example usage - Local testing - Flyte deployment configuration - \[Connectors > Databricks connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/databricks-connector/databricks-connector-example-usage/page.md) - \[Connectors > Memory Machine Cloud connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/mmcloud-connector/page.md) - Installation - Example usage - Local testing - Flyte deployment configuration - \[Connectors > Memory Machine Cloud connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/mmcloud-connector/mmcloud-connector-example-usage/page.md) - \[Connectors > OpenAI Batch connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/openai-batch-connector/page.md) - Installation - Example usage - Local testing - Flyte deployment configuration - \[Connectors > OpenAI Batch connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/openai-batch-connector/openai-batch-connector-example-usage/page.md) - \[Connectors > Perian connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/perian-connector/page.md) - Example usage - Connector setup - \[Connectors > Perian connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/perian-connector/example/page.md) - \[Connectors > SageMaker connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/sagemaker-inference-connector/page.md) - Installation - Example usage - Local testing - Flyte deployment configuration - \[Connectors > SageMaker connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/sagemaker-inference-connector/sagemaker-inference-connector-example-usage/page.md) - \[Connectors > Sensor connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/sensor/page.md) - Example usage - Flyte deployment configuration - \[Connectors > Sensor connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/sensor/file-sensor-example/page.md) - \[Connectors > Slurm connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/slurm-connector/page.md) - Installation - Example usage - Local testing - Flyte deployment configuration - \[Connectors > Slurm connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/slurm-connector/slurm-connector-example-usage/page.md) - \[Connectors > Snowflake connector\](https://www.union.ai/docs/v1/flyte/integrations/connectors/snowflake-connector/page.md) - Installation - Example usage - Local testing - Flyte deployment configuration - \[Connectors > Snowflake connector > Page\](https://www.union.ai/docs/v1/flyte/integrations/connectors/snowflake-connector/snowflake-connector-example-usage/page.md) - \[Flytekit plugins\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/page.md) - \[Flytekit plugins > Comet ML\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/comet-ml-plugin/page.md) - \[Flytekit plugins > Comet ML > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/comet-ml-plugin/comet-ml-example/page.md) - \[Flytekit plugins > DBT\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/dbt-plugin/page.md) - Prerequisities - Running the Example - \[Flytekit plugins > DBT > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/dbt-plugin/dbt-example/page.md) - \[Flytekit plugins > Dolt\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/dolt-plugin/page.md) - Installation - \[Flytekit plugins > Dolt > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/dolt-plugin/dolt-branch-example/page.md) - \[Flytekit plugins > Dolt > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/dolt-plugin/dolt-quickstart-example/page.md) - \[Flytekit plugins > Great Expectations\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/greatexpectations-plugin/page.md) - How to Define the Integration - Data Validation Failure - Plugin Parameters - Optional Parameters - Plugin Installation - \[Flytekit plugins > Great Expectations > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/greatexpectations-plugin/task-example/page.md) - \[Flytekit plugins > Great Expectations > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/greatexpectations-plugin/type-example/page.md) - \[Flytekit plugins > Memray Profiling\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/memray-plugin/page.md) - \[Flytekit plugins > Memray Profiling > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/memray-plugin/memray-example/page.md) - \[Flytekit plugins > MLflow\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/mlflow-plugin/page.md) - \[Flytekit plugins > MLflow > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/mlflow-plugin/mlflow-example/page.md) - \[Flytekit plugins > Modin\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/modin-plugin/page.md) - Installation - How is Modin different? - \[Flytekit plugins > Modin > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/modin-plugin/knn-classifier/page.md) - \[Flytekit plugins > Neptune\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/neptune-plugin/page.md) - Installation - Local testing - Flyte deployment configuration - \[Flytekit plugins > Neptune > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/neptune-plugin/neptune-example/page.md) - \[Flytekit plugins > NIM\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/nim-plugin/page.md) - Installation - Example usage - \[Flytekit plugins > NIM > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/nim-plugin/serve-nim-container/page.md) - \[Flytekit plugins > Ollama\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/ollama-plugin/page.md) - Installation - Example usage - \[Flytekit plugins > Ollama > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/ollama-plugin/serve-llm/page.md) - \[Flytekit plugins > ONNX\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/onnx-plugin/page.md) - \[Flytekit plugins > ONNX > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/onnx-plugin/pytorch-onnx/page.md) - \[Flytekit plugins > ONNX > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/onnx-plugin/scikitlearn-onnx/page.md) - \[Flytekit plugins > ONNX > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/onnx-plugin/tensorflow-onnx/page.md) - \[Flytekit plugins > Pandera\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/pandera-plugin/page.md) - Installation - Quick Start - \[Flytekit plugins > Pandera > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/pandera-plugin/basic-schema-example/page.md) - \[Flytekit plugins > Pandera > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/pandera-plugin/validating-and-testing-ml-pipelines/page.md) - \[Flytekit plugins > Papermill\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/papermill-plugin/page.md) - Installation - \[Flytekit plugins > Papermill > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/papermill-plugin/simple/page.md) - \[Flytekit plugins > Papermill > Simple Papermill Notebook\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/papermill-plugin/nb-simple/page.md) - \[Flytekit plugins > DuckDB\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/duckdb-plugin/page.md) - \[Flytekit plugins > DuckDB > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/duckdb-plugin/duckdb-example/page.md) - \[Flytekit plugins > SQL\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/sql-plugin/page.md) - \[Flytekit plugins > SQL > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/sql-plugin/sql-alchemy/page.md) - \[Flytekit plugins > SQL > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/sql-plugin/sqlite3-integration/page.md) - \[Flytekit plugins > Weights and Biases\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/wandb-plugin/page.md) - \[Flytekit plugins > Weights and Biases > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/wandb-plugin/wandb-example/page.md) - \[Flytekit plugins > WhyLogs\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/whylogs-plugin/page.md) - whylogs with Flyte - whylogs Flyte Type - Renderers - Installing the plugin - \[Flytekit plugins > WhyLogs > Page\](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/whylogs-plugin/whylogs-example/page.md) - \[Native backend plugins\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/page.md) - \[Native backend plugins > Dask\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/k8s-dask-plugin/page.md) - Why use Kubernetes Dask? - Install the plugin - Implementation details - Local execution - Remote execution - Resource specification - Images - Environment variables - Labels and annotations - Interruptible tasks - Run the example on the Flyte cluster - \[Native backend plugins > Dask > Page\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/k8s-dask-plugin/dask-example/page.md) - \[Native backend plugins > MPI\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/kfmpi-plugin/page.md) - Horovod - MPI (Message Passing Interface) - Install the plugin - Build a Docker image - Run the example on the Flyte cluster - MPI Plugin Troubleshooting Guide - \[Native backend plugins > MPI > Page\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/kfmpi-plugin/mpi-mnist/page.md) - \[Native backend plugins > PyTorch Distributed\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/kfpytorch-plugin/page.md) - Install the plugin - Run the example on the Flyte cluster - \[Native backend plugins > PyTorch Distributed > Page\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/kfpytorch-plugin/pytorch-mnist/page.md) - \[Native backend plugins > PyTorch Distributed > Page\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/kfpytorch-plugin/pytorch-lightning-mnist-autoencoder/page.md) - \[Native backend plugins > Ray\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/ray-plugin/page.md) - Install the plugin - Implementation details - Submit a Ray job to existing cluster - Create a Ray cluster managed by Flyte and run a Ray Job on the cluster - Run the example on the Flyte cluster - \[Native backend plugins > Ray > Page\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/ray-plugin/ray-example/page.md) - \[Native backend plugins > Spark\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/k8s-spark-plugin/page.md) - Why use Kubernetes Spark? - Implementation details - Step 1: Deploy Spark plugin in the Flyte backend - Step 2: Environment Setup - Step 3: Optionally, set up visibility - Run the examples on the Flyte cluster - \[Native backend plugins > Spark > Page\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/k8s-spark-plugin/dataframe-passing/page.md) - \[Native backend plugins > Spark > Page\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/k8s-spark-plugin/pyspark-pi/page.md) - \[Native backend plugins > TensorFlow Distributed\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/kftensorflow-plugin/page.md) - Install the plugin - Run the example on the Flyte cluster - \[Native backend plugins > TensorFlow Distributed > Page\](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/kftensorflow-plugin/tf-mnist/page.md) - \[External service backend plugins\](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/page.md) - \[External service backend plugins > AWS Athena\](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/athena-plugin/page.md) - Executing Athena Queries - Installation - \[External service backend plugins > AWS Athena > Page\](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/athena-plugin/athena/page.md) - \[External service backend plugins > AWS Batch\](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/aws-batch-plugin/page.md) - Executing Batch Job - Installation - Configuring the backend to get AWS Batch working - Quick Start - \[External service backend plugins > AWS Batch > Page\](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/aws-batch-plugin/batch/page.md) - \[External service backend plugins > FlyteInteractive\](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/flyteinteractive-plugin/page.md) - Installation - Acknowledgement - \[External service backend plugins > FlyteInteractive > Page\](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/flyteinteractive-plugin/jupyter/page.md) - \[External service backend plugins > FlyteInteractive > Page\](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/flyteinteractive-plugin/vscode/page.md) - \[External service backend plugins > Hive\](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/hive-plugin/page.md) - Installation - No Need of a dockerfile - \[External service backend plugins > Hive > Page\](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/hive-plugin/hive/page.md) - \[Flyte operators\](https://www.union.ai/docs/v1/flyte/integrations/flyte-operators/page.md) - \[Flyte operators > Airflow Provider\](https://www.union.ai/docs/v1/flyte/integrations/flyte-operators/airflow-plugin/page.md) - Installation - \[Flyte operators > Airflow Provider > Page\](https://www.union.ai/docs/v1/flyte/integrations/flyte-operators/airflow-plugin/airflow/page.md) - \[Deprecated integrations\](https://www.union.ai/docs/v1/flyte/integrations/deprecated-integrations/page.md) - \[Deprecated integrations > BigQuery plugin\](https://www.union.ai/docs/v1/flyte/integrations/deprecated-integrations/bigquery-plugin/page.md) - \[Deprecated integrations > BigQuery plugin > Page\](https://www.union.ai/docs/v1/flyte/integrations/deprecated-integrations/bigquery-plugin/bigquery-plugin-example/page.md) - \[Deprecated integrations > Databricks plugin\](https://www.union.ai/docs/v1/flyte/integrations/deprecated-integrations/databricks-plugin/page.md) - \[Deprecated integrations > Databricks plugin > Page\](https://www.union.ai/docs/v1/flyte/integrations/deprecated-integrations/databricks-plugin/databricks-plugin-example/page.md) - \[Deprecated integrations > Kubernetes Pods\](https://www.union.ai/docs/v1/flyte/integrations/deprecated-integrations/k8s-pod-plugin/page.md) - Installation - \[Deprecated integrations > Kubernetes Pods > Page\](https://www.union.ai/docs/v1/flyte/integrations/deprecated-integrations/k8s-pod-plugin/pod/page.md) - \[Deprecated integrations > Snowflake plugin\](https://www.union.ai/docs/v1/flyte/integrations/deprecated-integrations/snowflake-plugin/page.md) - \[Deprecated integrations > Snowflake plugin > Page\](https://www.union.ai/docs/v1/flyte/integrations/deprecated-integrations/snowflake-plugin/snowflake-plugin-example/page.md) --- ## Reference - \[LLM-optimized documentation\](https://www.union.ai/docs/v1/flyte/api-reference/flyte-context/page.md) - Per-page Markdown (\`page.md\`) - Section bundles (\`section.md\`) - Page index (\`llms.txt\`) - Full documentation (\`llms-full.txt\`) - \[Pyflyte CLI\](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/page.md) - Installation - Configure the \`pyflyte\` CLI - Overriding the configuration file location - \`pyflyte\` CLI configuration search path - \`pyflyte\` CLI commands - backfill - build - execution - fetch - get - info - init - launchplan - local-cache - metrics - package - register - run - serve - \[Flytectl CLI\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/page.md) - Installation - macOS - Linux - Windows - Configuration - Configuration file location hierarchy - Options - Commands - Entities - \[Flytectl CLI > flytectl\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl/page.md) - Synopsis - Options - \[Flytectl CLI > flytectl version\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-version/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl append\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-append/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl append > flytectl append identityassignments\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-append/flytectl-append-identityassignments/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl apply\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-apply/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl apply > flytectl apply app\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-apply/flytectl-apply-app/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl apply > flytectl apply clusterconfig\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-apply/flytectl-apply-clusterconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl apply > flytectl apply clusterconfigid\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-apply/flytectl-apply-clusterconfigid/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl apply > flytectl apply clusterpoolconfig\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-apply/flytectl-apply-clusterpoolconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl config\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-config/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl config > flytectl config discover\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-config/flytectl-config-discover/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl config > flytectl config docs\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-config/flytectl-config-docs/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl config > flytectl config init\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-config/flytectl-config-init/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl config > flytectl config validate\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-config/flytectl-config-validate/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl create\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-create/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl create > flytectl create app\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-create/flytectl-create-app/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl create > flytectl create clusterpool\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-create/flytectl-create-clusterpool/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl create > flytectl create clusterpoolassignment\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-create/flytectl-create-clusterpoolassignment/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl create > flytectl create execution\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-create/flytectl-create-execution/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl create > flytectl create policy\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-create/flytectl-create-policy/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl create > flytectl create project\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-create/flytectl-create-project/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl create > flytectl create role\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-create/flytectl-create-role/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete app\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-app/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete cluster\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-cluster/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete cluster-pool-attributes\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-cluster-pool-attributes/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete cluster-resource-attribute\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-cluster-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete clusterconfig\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-clusterconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete clusterpool\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-clusterpool/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete clusterpoolassignment\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-clusterpoolassignment/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete execution\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-execution/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete execution-cluster-label\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-execution-cluster-label/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete execution-queue-attribute\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-execution-queue-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete identityassignments\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-identityassignments/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete plugin-override\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-plugin-override/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete policy\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-policy/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete role\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-role/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete task-resource-attribute\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-task-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl delete > flytectl delete workflow-execution-config\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-delete/flytectl-delete-workflow-execution-config/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl demo\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-demo/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl demo > flytectl demo exec\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-demo/flytectl-demo-exec/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl demo > flytectl demo reload\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-demo/flytectl-demo-reload/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl demo > flytectl demo start\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-demo/flytectl-demo-start/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl demo > flytectl demo status\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-demo/flytectl-demo-status/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl demo > flytectl demo teardown\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-demo/flytectl-demo-teardown/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get app\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-app/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get cluster\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-cluster/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get cluster-pool-attributes\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-cluster-pool-attributes/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get cluster-resource-attribute\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-cluster-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get clusterconfig\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-clusterconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get clusterconfigs\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-clusterconfigs/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get clusterpool\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-clusterpool/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get clusterpoolconfig\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-clusterpoolconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get clusterswithconfig\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-clusterswithconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get echo\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-echo/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get execution\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-execution/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get execution-cluster-label\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-execution-cluster-label/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get execution-queue-attribute\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-execution-queue-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get executionoperation\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-executionoperation/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get identityassignment\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-identityassignment/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get launchplan\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-launchplan/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get plugin-override\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-plugin-override/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get policy\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-policy/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get project\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-project/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get role\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-role/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get task\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-task/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get task-resource-attribute\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-task-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get workflow\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-workflow/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl get > flytectl get workflow-execution-config\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-get/flytectl-get-workflow-execution-config/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl register\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-register/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl register > flytectl register examples\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-register/flytectl-register-examples/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl register > flytectl register files\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-register/flytectl-register-files/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update cluster-pool-attributes\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-cluster-pool-attributes/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update cluster-resource-attribute\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-cluster-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update execution\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-execution/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update execution-cluster-label\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-execution-cluster-label/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update execution-queue-attribute\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-execution-queue-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update launchplan\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-launchplan/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update launchplan-meta\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-launchplan-meta/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update plugin-override\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-plugin-override/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update project\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-project/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update task-meta\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-task-meta/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update task-resource-attribute\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-task-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update workflow-execution-config\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-workflow-execution-config/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flytectl CLI > flytectl update > flytectl update workflow-meta\](https://www.union.ai/docs/v1/flyte/api-reference/flytectl-cli/flytectl-update/flytectl-update-workflow-meta/page.md) - Synopsis - Options - Options inherited from parent commands - \[Flyteidl\](https://www.union.ai/docs/v1/flyte/api-reference/flyteidl/page.md) - flyteidl/core/compiler.proto - CompiledLaunchPlan {#flyteidl-core-CompiledLaunchPlan} - CompiledTask {#flyteidl-core-CompiledTask} - CompiledWorkflow {#flyteidl-core-CompiledWorkflow} - CompiledWorkflowClosure {#flyteidl-core-CompiledWorkflowClosure} - ConnectionSet {#flyteidl-core-ConnectionSet} - ConnectionSet.DownstreamEntry {#flyteidl-core-ConnectionSet-DownstreamEntry} - ConnectionSet.IdList {#flyteidl-core-ConnectionSet-IdList} - ConnectionSet.UpstreamEntry {#flyteidl-core-ConnectionSet-UpstreamEntry} - flyteidl/core/interface.proto - Parameter {#flyteidl-core-Parameter} - ParameterMap {#flyteidl-core-ParameterMap} - ParameterMap.ParametersEntry {#flyteidl-core-ParameterMap-ParametersEntry} - TypedInterface {#flyteidl-core-TypedInterface} - Variable {#flyteidl-core-Variable} - VariableMap {#flyteidl-core-VariableMap} - VariableMap.VariablesEntry {#flyteidl-core-VariableMap-VariablesEntry} - flyteidl/core/catalog.proto - CatalogArtifactTag {#flyteidl-core-CatalogArtifactTag} - CatalogMetadata {#flyteidl-core-CatalogMetadata} - CatalogReservation {#flyteidl-core-CatalogReservation} - CatalogCacheStatus {#flyteidl-core-CatalogCacheStatus} - CatalogReservation.Status {#flyteidl-core-CatalogReservation-Status} - flyteidl/core/literals.proto - Binary {#flyteidl-core-Binary} - Binding {#flyteidl-core-Binding} - BindingData {#flyteidl-core-BindingData} - BindingDataCollection {#flyteidl-core-BindingDataCollection} - BindingDataMap {#flyteidl-core-BindingDataMap} - BindingDataMap.BindingsEntry {#flyteidl-core-BindingDataMap-BindingsEntry} - Blob {#flyteidl-core-Blob} - BlobMetadata {#flyteidl-core-BlobMetadata} - KeyValuePair {#flyteidl-core-KeyValuePair} - Literal {#flyteidl-core-Literal} - Literal.MetadataEntry {#flyteidl-core-Literal-MetadataEntry} - LiteralCollection {#flyteidl-core-LiteralCollection} - LiteralMap {#flyteidl-core-LiteralMap} - LiteralMap.LiteralsEntry {#flyteidl-core-LiteralMap-LiteralsEntry} - LiteralOffloadedMetadata {#flyteidl-core-LiteralOffloadedMetadata} - Primitive {#flyteidl-core-Primitive} - RetryStrategy {#flyteidl-core-RetryStrategy} - Scalar {#flyteidl-core-Scalar} - Schema {#flyteidl-core-Schema} - StructuredDataset {#flyteidl-core-StructuredDataset} - StructuredDatasetMetadata {#flyteidl-core-StructuredDatasetMetadata} - Union {#flyteidl-core-Union} - UnionInfo {#flyteidl-core-UnionInfo} - Void {#flyteidl-core-Void} - flyteidl/core/tasks.proto - Container {#flyteidl-core-Container} - ContainerPort {#flyteidl-core-ContainerPort} - DataLoadingConfig {#flyteidl-core-DataLoadingConfig} - ExtendedResources {#flyteidl-core-ExtendedResources} - GPUAccelerator {#flyteidl-core-GPUAccelerator} - IOStrategy {#flyteidl-core-IOStrategy} - K8sObjectMetadata {#flyteidl-core-K8sObjectMetadata} - K8sObjectMetadata.AnnotationsEntry {#flyteidl-core-K8sObjectMetadata-AnnotationsEntry} - K8sObjectMetadata.LabelsEntry {#flyteidl-core-K8sObjectMetadata-LabelsEntry} - K8sPod {#flyteidl-core-K8sPod} - Resources {#flyteidl-core-Resources} - Resources.ResourceEntry {#flyteidl-core-Resources-ResourceEntry} - RuntimeMetadata {#flyteidl-core-RuntimeMetadata} - SharedMemory {#flyteidl-core-SharedMemory} - Sql {#flyteidl-core-Sql} - TaskMetadata {#flyteidl-core-TaskMetadata} - TaskMetadata.TagsEntry {#flyteidl-core-TaskMetadata-TagsEntry} - TaskTemplate {#flyteidl-core-TaskTemplate} - TaskTemplate.ConfigEntry {#flyteidl-core-TaskTemplate-ConfigEntry} - Container.Architecture {#flyteidl-core-Container-Architecture} - DataLoadingConfig.LiteralMapFormat {#flyteidl-core-DataLoadingConfig-LiteralMapFormat} - IOStrategy.DownloadMode {#flyteidl-core-IOStrategy-DownloadMode} - IOStrategy.UploadMode {#flyteidl-core-IOStrategy-UploadMode} - Resources.ResourceName {#flyteidl-core-Resources-ResourceName} - RuntimeMetadata.RuntimeType {#flyteidl-core-RuntimeMetadata-RuntimeType} - Sql.Dialect {#flyteidl-core-Sql-Dialect} - flyteidl/core/metrics.proto - ExecutionMetricResult {#flyteidl-core-ExecutionMetricResult} - Span {#flyteidl-core-Span} - flyteidl/core/errors.proto - ContainerError {#flyteidl-core-ContainerError} - ErrorDocument {#flyteidl-core-ErrorDocument} - ContainerError.Kind {#flyteidl-core-ContainerError-Kind} - flyteidl/core/identifier.proto - Identifier {#flyteidl-core-Identifier} - NodeExecutionIdentifier {#flyteidl-core-NodeExecutionIdentifier} - SignalIdentifier {#flyteidl-core-SignalIdentifier} - TaskExecutionIdentifier {#flyteidl-core-TaskExecutionIdentifier} - WorkflowExecutionIdentifier {#flyteidl-core-WorkflowExecutionIdentifier} - ResourceType {#flyteidl-core-ResourceType} - flyteidl/core/artifact\_id.proto - ArtifactBindingData {#flyteidl-core-ArtifactBindingData} - ArtifactID {#flyteidl-core-ArtifactID} - ArtifactKey {#flyteidl-core-ArtifactKey} - ArtifactQuery {#flyteidl-core-ArtifactQuery} - ArtifactTag {#flyteidl-core-ArtifactTag} - InputBindingData {#flyteidl-core-InputBindingData} - LabelValue {#flyteidl-core-LabelValue} - Partitions {#flyteidl-core-Partitions} - Partitions.ValueEntry {#flyteidl-core-Partitions-ValueEntry} - RuntimeBinding {#flyteidl-core-RuntimeBinding} - TimePartition {#flyteidl-core-TimePartition} - TimeTransform {#flyteidl-core-TimeTransform} - Granularity {#flyteidl-core-Granularity} - Operator {#flyteidl-core-Operator} - flyteidl/core/types.proto - BlobType {#flyteidl-core-BlobType} - EnumType {#flyteidl-core-EnumType} - Error {#flyteidl-core-Error} - LiteralType {#flyteidl-core-LiteralType} - OutputReference {#flyteidl-core-OutputReference} - PromiseAttribute {#flyteidl-core-PromiseAttribute} - SchemaType {#flyteidl-core-SchemaType} - SchemaType.SchemaColumn {#flyteidl-core-SchemaType-SchemaColumn} - StructuredDatasetType {#flyteidl-core-StructuredDatasetType} - StructuredDatasetType.DatasetColumn {#flyteidl-core-StructuredDatasetType-DatasetColumn} - TypeAnnotation {#flyteidl-core-TypeAnnotation} - TypeStructure {#flyteidl-core-TypeStructure} - TypeStructure.DataclassTypeEntry {#flyteidl-core-TypeStructure-DataclassTypeEntry} - UnionType {#flyteidl-core-UnionType} - BlobType.BlobDimensionality {#flyteidl-core-BlobType-BlobDimensionality} - SchemaType.SchemaColumn.SchemaColumnType {#flyteidl-core-SchemaType-SchemaColumn-SchemaColumnType} - SimpleType {#flyteidl-core-SimpleType} - flyteidl/core/execution\_envs.proto - ExecutionEnv {#flyteidl-core-ExecutionEnv} - ExecutionEnvAssignment {#flyteidl-core-ExecutionEnvAssignment} - flyteidl/core/execution.proto - ExecutionError {#flyteidl-core-ExecutionError} - NodeExecution {#flyteidl-core-NodeExecution} - QualityOfService {#flyteidl-core-QualityOfService} - QualityOfServiceSpec {#flyteidl-core-QualityOfServiceSpec} - TaskExecution {#flyteidl-core-TaskExecution} - TaskLog {#flyteidl-core-TaskLog} - WorkflowExecution {#flyteidl-core-WorkflowExecution} - ExecutionError.ErrorKind {#flyteidl-core-ExecutionError-ErrorKind} - NodeExecution.Phase {#flyteidl-core-NodeExecution-Phase} - QualityOfService.Tier {#flyteidl-core-QualityOfService-Tier} - TaskExecution.Phase {#flyteidl-core-TaskExecution-Phase} - TaskLog.MessageFormat {#flyteidl-core-TaskLog-MessageFormat} - WorkflowExecution.Phase {#flyteidl-core-WorkflowExecution-Phase} - flyteidl/core/security.proto - Identity {#flyteidl-core-Identity} - OAuth2Client {#flyteidl-core-OAuth2Client} - OAuth2TokenRequest {#flyteidl-core-OAuth2TokenRequest} - Secret {#flyteidl-core-Secret} - SecurityContext {#flyteidl-core-SecurityContext} - OAuth2TokenRequest.Type {#flyteidl-core-OAuth2TokenRequest-Type} - Secret.MountType {#flyteidl-core-Secret-MountType} - flyteidl/core/workflow.proto - Alias {#flyteidl-core-Alias} - ApproveCondition {#flyteidl-core-ApproveCondition} - ArrayNode {#flyteidl-core-ArrayNode} - BranchNode {#flyteidl-core-BranchNode} - GateNode {#flyteidl-core-GateNode} - IfBlock {#flyteidl-core-IfBlock} - IfElseBlock {#flyteidl-core-IfElseBlock} - LaunchPlanTemplate {#flyteidl-core-LaunchPlanTemplate} - Node {#flyteidl-core-Node} - NodeMetadata {#flyteidl-core-NodeMetadata} - NodeMetadata.ConfigEntry {#flyteidl-core-NodeMetadata-ConfigEntry} - SignalCondition {#flyteidl-core-SignalCondition} - SleepCondition {#flyteidl-core-SleepCondition} - TaskNode {#flyteidl-core-TaskNode} - TaskNodeOverrides {#flyteidl-core-TaskNodeOverrides} - WorkflowMetadata {#flyteidl-core-WorkflowMetadata} - WorkflowMetadata.TagsEntry {#flyteidl-core-WorkflowMetadata-TagsEntry} - WorkflowMetadataDefaults {#flyteidl-core-WorkflowMetadataDefaults} - WorkflowNode {#flyteidl-core-WorkflowNode} - WorkflowTemplate {#flyteidl-core-WorkflowTemplate} - ArrayNode.DataMode {#flyteidl-core-ArrayNode-DataMode} - ArrayNode.ExecutionMode {#flyteidl-core-ArrayNode-ExecutionMode} - WorkflowMetadata.OnFailurePolicy {#flyteidl-core-WorkflowMetadata-OnFailurePolicy} - flyteidl/core/workflow\_closure.proto - WorkflowClosure {#flyteidl-core-WorkflowClosure} - flyteidl/core/condition.proto - BooleanExpression {#flyteidl-core-BooleanExpression} - ComparisonExpression {#flyteidl-core-ComparisonExpression} - ConjunctionExpression {#flyteidl-core-ConjunctionExpression} - Operand {#flyteidl-core-Operand} - ComparisonExpression.Operator {#flyteidl-core-ComparisonExpression-Operator} - ConjunctionExpression.LogicalOperator {#flyteidl-core-ConjunctionExpression-LogicalOperator} - flyteidl/core/dynamic\_job.proto - DynamicJobSpec {#flyteidl-core-DynamicJobSpec} - flyteidl/plugins/presto.proto - PrestoQuery {#flyteidl-plugins-PrestoQuery} - flyteidl/plugins/qubole.proto - HiveQuery {#flyteidl-plugins-HiveQuery} - HiveQueryCollection {#flyteidl-plugins-HiveQueryCollection} - QuboleHiveJob {#flyteidl-plugins-QuboleHiveJob} - flyteidl/plugins/ray.proto - HeadGroupSpec {#flyteidl-plugins-HeadGroupSpec} - HeadGroupSpec.RayStartParamsEntry {#flyteidl-plugins-HeadGroupSpec-RayStartParamsEntry} - RayCluster {#flyteidl-plugins-RayCluster} - RayJob {#flyteidl-plugins-RayJob} - WorkerGroupSpec {#flyteidl-plugins-WorkerGroupSpec} - WorkerGroupSpec.RayStartParamsEntry {#flyteidl-plugins-WorkerGroupSpec-RayStartParamsEntry} - flyteidl/plugins/spark.proto - SparkApplication {#flyteidl-plugins-SparkApplication} - SparkJob {#flyteidl-plugins-SparkJob} - SparkJob.HadoopConfEntry {#flyteidl-plugins-SparkJob-HadoopConfEntry} - SparkJob.SparkConfEntry {#flyteidl-plugins-SparkJob-SparkConfEntry} - SparkApplication.Type {#flyteidl-plugins-SparkApplication-Type} - flyteidl/plugins/array\_job.proto - ArrayJob {#flyteidl-plugins-ArrayJob} - flyteidl/plugins/waitable.proto - Waitable {#flyteidl-plugins-Waitable} - flyteidl/plugins/dask.proto - DaskJob {#flyteidl-plugins-DaskJob} - DaskScheduler {#flyteidl-plugins-DaskScheduler} - DaskWorkerGroup {#flyteidl-plugins-DaskWorkerGroup} - flyteidl/plugins/mpi.proto - DistributedMPITrainingTask {#flyteidl-plugins-DistributedMPITrainingTask} - flyteidl/plugins/pytorch.proto - DistributedPyTorchTrainingTask {#flyteidl-plugins-DistributedPyTorchTrainingTask} - ElasticConfig {#flyteidl-plugins-ElasticConfig} - flyteidl/plugins/kubeflow/mpi.proto - DistributedMPITrainingReplicaSpec {#flyteidl-plugins-kubeflow-DistributedMPITrainingReplicaSpec} - DistributedMPITrainingTask {#flyteidl-plugins-kubeflow-DistributedMPITrainingTask} - flyteidl/plugins/kubeflow/pytorch.proto - DistributedPyTorchTrainingReplicaSpec {#flyteidl-plugins-kubeflow-DistributedPyTorchTrainingReplicaSpec} - DistributedPyTorchTrainingTask {#flyteidl-plugins-kubeflow-DistributedPyTorchTrainingTask} - ElasticConfig {#flyteidl-plugins-kubeflow-ElasticConfig} - flyteidl/plugins/kubeflow/tensorflow.proto - DistributedTensorflowTrainingReplicaSpec {#flyteidl-plugins-kubeflow-DistributedTensorflowTrainingReplicaSpec} - DistributedTensorflowTrainingTask {#flyteidl-plugins-kubeflow-DistributedTensorflowTrainingTask} - flyteidl/plugins/kubeflow/common.proto - RunPolicy {#flyteidl-plugins-kubeflow-RunPolicy} - CleanPodPolicy {#flyteidl-plugins-kubeflow-CleanPodPolicy} - flyteidl/plugins/tensorflow.proto - DistributedTensorflowTrainingTask {#flyteidl-plugins-DistributedTensorflowTrainingTask} - flyteidl/plugins/common.proto - CommonReplicaSpec {#flyteidl-plugins-CommonReplicaSpec} - RestartPolicy {#flyteidl-plugins-RestartPolicy} - flyteidl/admin/schedule.proto - CronSchedule {#flyteidl-admin-CronSchedule} - FixedRate {#flyteidl-admin-FixedRate} - Schedule {#flyteidl-admin-Schedule} - FixedRateUnit {#flyteidl-admin-FixedRateUnit} - flyteidl/admin/project.proto - Domain {#flyteidl-admin-Domain} - GetDomainRequest {#flyteidl-admin-GetDomainRequest} - GetDomainsResponse {#flyteidl-admin-GetDomainsResponse} - InactiveProject {#flyteidl-admin-InactiveProject} - Project {#flyteidl-admin-Project} - ProjectGetRequest {#flyteidl-admin-ProjectGetRequest} - ProjectListRequest {#flyteidl-admin-ProjectListRequest} - ProjectRegisterRequest {#flyteidl-admin-ProjectRegisterRequest} - ProjectRegisterResponse {#flyteidl-admin-ProjectRegisterResponse} - ProjectUpdateResponse {#flyteidl-admin-ProjectUpdateResponse} - Projects {#flyteidl-admin-Projects} - Project.ProjectState {#flyteidl-admin-Project-ProjectState} - flyteidl/admin/cluster\_assignment.proto - ClusterAssignment {#flyteidl-admin-ClusterAssignment} - flyteidl/admin/notification.proto - EmailMessage {#flyteidl-admin-EmailMessage} - flyteidl/admin/task.proto - Task {#flyteidl-admin-Task} - TaskClosure {#flyteidl-admin-TaskClosure} - TaskCreateRequest {#flyteidl-admin-TaskCreateRequest} - TaskCreateResponse {#flyteidl-admin-TaskCreateResponse} - TaskList {#flyteidl-admin-TaskList} - TaskSpec {#flyteidl-admin-TaskSpec} - flyteidl/admin/launch\_plan.proto - ActiveLaunchPlanListRequest {#flyteidl-admin-ActiveLaunchPlanListRequest} - ActiveLaunchPlanRequest {#flyteidl-admin-ActiveLaunchPlanRequest} - Auth {#flyteidl-admin-Auth} - LaunchPlan {#flyteidl-admin-LaunchPlan} - LaunchPlanClosure {#flyteidl-admin-LaunchPlanClosure} - LaunchPlanCreateRequest {#flyteidl-admin-LaunchPlanCreateRequest} - LaunchPlanCreateResponse {#flyteidl-admin-LaunchPlanCreateResponse} - LaunchPlanList {#flyteidl-admin-LaunchPlanList} - LaunchPlanMetadata {#flyteidl-admin-LaunchPlanMetadata} - LaunchPlanSpec {#flyteidl-admin-LaunchPlanSpec} - LaunchPlanUpdateRequest {#flyteidl-admin-LaunchPlanUpdateRequest} - LaunchPlanUpdateResponse {#flyteidl-admin-LaunchPlanUpdateResponse} - LaunchPlanState {#flyteidl-admin-LaunchPlanState} - flyteidl/admin/signal.proto - Signal {#flyteidl-admin-Signal} - SignalGetOrCreateRequest {#flyteidl-admin-SignalGetOrCreateRequest} - SignalList {#flyteidl-admin-SignalList} - SignalListRequest {#flyteidl-admin-SignalListRequest} - SignalSetRequest {#flyteidl-admin-SignalSetRequest} - SignalSetResponse {#flyteidl-admin-SignalSetResponse} - flyteidl/admin/task\_execution.proto - Reason {#flyteidl-admin-Reason} - TaskExecution {#flyteidl-admin-TaskExecution} - TaskExecutionClosure {#flyteidl-admin-TaskExecutionClosure} - TaskExecutionGetDataRequest {#flyteidl-admin-TaskExecutionGetDataRequest} - TaskExecutionGetDataResponse {#flyteidl-admin-TaskExecutionGetDataResponse} - TaskExecutionGetRequest {#flyteidl-admin-TaskExecutionGetRequest} - TaskExecutionList {#flyteidl-admin-TaskExecutionList} - TaskExecutionListRequest {#flyteidl-admin-TaskExecutionListRequest} - flyteidl/admin/node\_execution.proto - DynamicNodeWorkflowResponse {#flyteidl-admin-DynamicNodeWorkflowResponse} - DynamicWorkflowNodeMetadata {#flyteidl-admin-DynamicWorkflowNodeMetadata} - GetDynamicNodeWorkflowRequest {#flyteidl-admin-GetDynamicNodeWorkflowRequest} - NodeExecution {#flyteidl-admin-NodeExecution} - NodeExecutionClosure {#flyteidl-admin-NodeExecutionClosure} - NodeExecutionForTaskListRequest {#flyteidl-admin-NodeExecutionForTaskListRequest} - NodeExecutionGetDataRequest {#flyteidl-admin-NodeExecutionGetDataRequest} - NodeExecutionGetDataResponse {#flyteidl-admin-NodeExecutionGetDataResponse} - NodeExecutionGetRequest {#flyteidl-admin-NodeExecutionGetRequest} - NodeExecutionList {#flyteidl-admin-NodeExecutionList} - NodeExecutionListRequest {#flyteidl-admin-NodeExecutionListRequest} - NodeExecutionMetaData {#flyteidl-admin-NodeExecutionMetaData} - TaskNodeMetadata {#flyteidl-admin-TaskNodeMetadata} - WorkflowNodeMetadata {#flyteidl-admin-WorkflowNodeMetadata} - flyteidl/admin/execution.proto - AbortMetadata {#flyteidl-admin-AbortMetadata} - Execution {#flyteidl-admin-Execution} - ExecutionClosure {#flyteidl-admin-ExecutionClosure} - ExecutionCreateRequest {#flyteidl-admin-ExecutionCreateRequest} - ExecutionCreateResponse {#flyteidl-admin-ExecutionCreateResponse} - ExecutionList {#flyteidl-admin-ExecutionList} - ExecutionMetadata {#flyteidl-admin-ExecutionMetadata} - ExecutionRecoverRequest {#flyteidl-admin-ExecutionRecoverRequest} - ExecutionRelaunchRequest {#flyteidl-admin-ExecutionRelaunchRequest} - ExecutionSpec {#flyteidl-admin-ExecutionSpec} - ExecutionStateChangeDetails {#flyteidl-admin-ExecutionStateChangeDetails} - ExecutionTerminateRequest {#flyteidl-admin-ExecutionTerminateRequest} - ExecutionTerminateResponse {#flyteidl-admin-ExecutionTerminateResponse} - ExecutionUpdateRequest {#flyteidl-admin-ExecutionUpdateRequest} - ExecutionUpdateResponse {#flyteidl-admin-ExecutionUpdateResponse} - LiteralMapBlob {#flyteidl-admin-LiteralMapBlob} - NotificationList {#flyteidl-admin-NotificationList} - SystemMetadata {#flyteidl-admin-SystemMetadata} - WorkflowExecutionGetDataRequest {#flyteidl-admin-WorkflowExecutionGetDataRequest} - WorkflowExecutionGetDataResponse {#flyteidl-admin-WorkflowExecutionGetDataResponse} - WorkflowExecutionGetMetricsRequest {#flyteidl-admin-WorkflowExecutionGetMetricsRequest} - WorkflowExecutionGetMetricsResponse {#flyteidl-admin-WorkflowExecutionGetMetricsResponse} - WorkflowExecutionGetRequest {#flyteidl-admin-WorkflowExecutionGetRequest} - ExecutionMetadata.ExecutionMode {#flyteidl-admin-ExecutionMetadata-ExecutionMode} - ExecutionState {#flyteidl-admin-ExecutionState} - flyteidl/admin/workflow\_attributes.proto - WorkflowAttributes {#flyteidl-admin-WorkflowAttributes} - WorkflowAttributesDeleteRequest {#flyteidl-admin-WorkflowAttributesDeleteRequest} - WorkflowAttributesDeleteResponse {#flyteidl-admin-WorkflowAttributesDeleteResponse} - WorkflowAttributesGetRequest {#flyteidl-admin-WorkflowAttributesGetRequest} - WorkflowAttributesGetResponse {#flyteidl-admin-WorkflowAttributesGetResponse} - WorkflowAttributesUpdateRequest {#flyteidl-admin-WorkflowAttributesUpdateRequest} - WorkflowAttributesUpdateResponse {#flyteidl-admin-WorkflowAttributesUpdateResponse} - flyteidl/admin/event.proto - EventErrorAlreadyInTerminalState {#flyteidl-admin-EventErrorAlreadyInTerminalState} - EventErrorIncompatibleCluster {#flyteidl-admin-EventErrorIncompatibleCluster} - EventFailureReason {#flyteidl-admin-EventFailureReason} - NodeExecutionEventRequest {#flyteidl-admin-NodeExecutionEventRequest} - NodeExecutionEventResponse {#flyteidl-admin-NodeExecutionEventResponse} - TaskExecutionEventRequest {#flyteidl-admin-TaskExecutionEventRequest} - TaskExecutionEventResponse {#flyteidl-admin-TaskExecutionEventResponse} - WorkflowExecutionEventRequest {#flyteidl-admin-WorkflowExecutionEventRequest} - WorkflowExecutionEventResponse {#flyteidl-admin-WorkflowExecutionEventResponse} - flyteidl/admin/matchable\_resource.proto - ClusterResourceAttributes {#flyteidl-admin-ClusterResourceAttributes} - ClusterResourceAttributes.AttributesEntry {#flyteidl-admin-ClusterResourceAttributes-AttributesEntry} - ExecutionClusterLabel {#flyteidl-admin-ExecutionClusterLabel} - ExecutionQueueAttributes {#flyteidl-admin-ExecutionQueueAttributes} - ListMatchableAttributesRequest {#flyteidl-admin-ListMatchableAttributesRequest} - ListMatchableAttributesResponse {#flyteidl-admin-ListMatchableAttributesResponse} - MatchableAttributesConfiguration {#flyteidl-admin-MatchableAttributesConfiguration} - MatchingAttributes {#flyteidl-admin-MatchingAttributes} - PluginOverride {#flyteidl-admin-PluginOverride} - PluginOverrides {#flyteidl-admin-PluginOverrides} - TaskResourceAttributes {#flyteidl-admin-TaskResourceAttributes} - TaskResourceSpec {#flyteidl-admin-TaskResourceSpec} - WorkflowExecutionConfig {#flyteidl-admin-WorkflowExecutionConfig} - MatchableResource {#flyteidl-admin-MatchableResource} - PluginOverride.MissingPluginBehavior {#flyteidl-admin-PluginOverride-MissingPluginBehavior} - flyteidl/admin/project\_attributes.proto - ProjectAttributes {#flyteidl-admin-ProjectAttributes} - ProjectAttributesDeleteRequest {#flyteidl-admin-ProjectAttributesDeleteRequest} - ProjectAttributesDeleteResponse {#flyteidl-admin-ProjectAttributesDeleteResponse} - ProjectAttributesGetRequest {#flyteidl-admin-ProjectAttributesGetRequest} - ProjectAttributesGetResponse {#flyteidl-admin-ProjectAttributesGetResponse} - ProjectAttributesUpdateRequest {#flyteidl-admin-ProjectAttributesUpdateRequest} - ProjectAttributesUpdateResponse {#flyteidl-admin-ProjectAttributesUpdateResponse} - flyteidl/admin/version.proto - GetVersionRequest {#flyteidl-admin-GetVersionRequest} - GetVersionResponse {#flyteidl-admin-GetVersionResponse} - Version {#flyteidl-admin-Version} - flyteidl/admin/workflow.proto - CreateWorkflowFailureReason {#flyteidl-admin-CreateWorkflowFailureReason} - Workflow {#flyteidl-admin-Workflow} - WorkflowClosure {#flyteidl-admin-WorkflowClosure} - WorkflowCreateRequest {#flyteidl-admin-WorkflowCreateRequest} - WorkflowCreateResponse {#flyteidl-admin-WorkflowCreateResponse} - WorkflowErrorExistsDifferentStructure {#flyteidl-admin-WorkflowErrorExistsDifferentStructure} - WorkflowErrorExistsIdenticalStructure {#flyteidl-admin-WorkflowErrorExistsIdenticalStructure} - WorkflowList {#flyteidl-admin-WorkflowList} - WorkflowSpec {#flyteidl-admin-WorkflowSpec} - flyteidl/admin/description\_entity.proto - Description {#flyteidl-admin-Description} - DescriptionEntity {#flyteidl-admin-DescriptionEntity} - DescriptionEntityList {#flyteidl-admin-DescriptionEntityList} - DescriptionEntityListRequest {#flyteidl-admin-DescriptionEntityListRequest} - SourceCode {#flyteidl-admin-SourceCode} - DescriptionFormat {#flyteidl-admin-DescriptionFormat} - flyteidl/admin/project\_domain\_attributes.proto - ProjectDomainAttributes {#flyteidl-admin-ProjectDomainAttributes} - ProjectDomainAttributesDeleteRequest {#flyteidl-admin-ProjectDomainAttributesDeleteRequest} - ProjectDomainAttributesDeleteResponse {#flyteidl-admin-ProjectDomainAttributesDeleteResponse} - ProjectDomainAttributesGetRequest {#flyteidl-admin-ProjectDomainAttributesGetRequest} - ProjectDomainAttributesGetResponse {#flyteidl-admin-ProjectDomainAttributesGetResponse} - ProjectDomainAttributesUpdateRequest {#flyteidl-admin-ProjectDomainAttributesUpdateRequest} - ProjectDomainAttributesUpdateResponse {#flyteidl-admin-ProjectDomainAttributesUpdateResponse} - flyteidl/admin/agent.proto - Agent {#flyteidl-admin-Agent} - AgentError {#flyteidl-admin-AgentError} - CreateRequestHeader {#flyteidl-admin-CreateRequestHeader} - CreateTaskRequest {#flyteidl-admin-CreateTaskRequest} - CreateTaskResponse {#flyteidl-admin-CreateTaskResponse} - DeleteTaskRequest {#flyteidl-admin-DeleteTaskRequest} - DeleteTaskResponse {#flyteidl-admin-DeleteTaskResponse} - ExecuteTaskSyncRequest {#flyteidl-admin-ExecuteTaskSyncRequest} - ExecuteTaskSyncResponse {#flyteidl-admin-ExecuteTaskSyncResponse} - ExecuteTaskSyncResponseHeader {#flyteidl-admin-ExecuteTaskSyncResponseHeader} - GetAgentRequest {#flyteidl-admin-GetAgentRequest} - GetAgentResponse {#flyteidl-admin-GetAgentResponse} - GetTaskLogsRequest {#flyteidl-admin-GetTaskLogsRequest} - GetTaskLogsResponse {#flyteidl-admin-GetTaskLogsResponse} - GetTaskLogsResponseBody {#flyteidl-admin-GetTaskLogsResponseBody} - GetTaskLogsResponseHeader {#flyteidl-admin-GetTaskLogsResponseHeader} - GetTaskMetricsRequest {#flyteidl-admin-GetTaskMetricsRequest} - GetTaskMetricsResponse {#flyteidl-admin-GetTaskMetricsResponse} - GetTaskRequest {#flyteidl-admin-GetTaskRequest} - GetTaskResponse {#flyteidl-admin-GetTaskResponse} - ListAgentsRequest {#flyteidl-admin-ListAgentsRequest} - ListAgentsResponse {#flyteidl-admin-ListAgentsResponse} - LogLine {#flyteidl-admin-LogLine} - Resource {#flyteidl-admin-Resource} - TaskCategory {#flyteidl-admin-TaskCategory} - TaskExecutionMetadata {#flyteidl-admin-TaskExecutionMetadata} - TaskExecutionMetadata.AnnotationsEntry {#flyteidl-admin-TaskExecutionMetadata-AnnotationsEntry} - TaskExecutionMetadata.EnvironmentVariablesEntry {#flyteidl-admin-TaskExecutionMetadata-EnvironmentVariablesEntry} - TaskExecutionMetadata.LabelsEntry {#flyteidl-admin-TaskExecutionMetadata-LabelsEntry} - AgentError.Kind {#flyteidl-admin-AgentError-Kind} - LogLineOriginator {#flyteidl-admin-LogLineOriginator} - State {#flyteidl-admin-State} - flyteidl/admin/common.proto - Annotations {#flyteidl-admin-Annotations} - Annotations.ValuesEntry {#flyteidl-admin-Annotations-ValuesEntry} - AuthRole {#flyteidl-admin-AuthRole} - EmailNotification {#flyteidl-admin-EmailNotification} - Envs {#flyteidl-admin-Envs} - FlyteURLs {#flyteidl-admin-FlyteURLs} - Labels {#flyteidl-admin-Labels} - Labels.ValuesEntry {#flyteidl-admin-Labels-ValuesEntry} - NamedEntity {#flyteidl-admin-NamedEntity} - NamedEntityGetRequest {#flyteidl-admin-NamedEntityGetRequest} - NamedEntityIdentifier {#flyteidl-admin-NamedEntityIdentifier} - NamedEntityIdentifierList {#flyteidl-admin-NamedEntityIdentifierList} - NamedEntityIdentifierListRequest {#flyteidl-admin-NamedEntityIdentifierListRequest} - NamedEntityList {#flyteidl-admin-NamedEntityList} - NamedEntityListRequest {#flyteidl-admin-NamedEntityListRequest} - NamedEntityMetadata {#flyteidl-admin-NamedEntityMetadata} - NamedEntityUpdateRequest {#flyteidl-admin-NamedEntityUpdateRequest} - NamedEntityUpdateResponse {#flyteidl-admin-NamedEntityUpdateResponse} - Notification {#flyteidl-admin-Notification} - ObjectGetRequest {#flyteidl-admin-ObjectGetRequest} - PagerDutyNotification {#flyteidl-admin-PagerDutyNotification} - RawOutputDataConfig {#flyteidl-admin-RawOutputDataConfig} - ResourceListRequest {#flyteidl-admin-ResourceListRequest} - SlackNotification {#flyteidl-admin-SlackNotification} - Sort {#flyteidl-admin-Sort} - UrlBlob {#flyteidl-admin-UrlBlob} - NamedEntityState {#flyteidl-admin-NamedEntityState} - Sort.Direction {#flyteidl-admin-Sort-Direction} - flyteidl/cacheservice/cacheservice.proto - CachedOutput {#flyteidl-cacheservice-CachedOutput} - DeleteCacheRequest {#flyteidl-cacheservice-DeleteCacheRequest} - DeleteCacheResponse {#flyteidl-cacheservice-DeleteCacheResponse} - GetCacheRequest {#flyteidl-cacheservice-GetCacheRequest} - GetCacheResponse {#flyteidl-cacheservice-GetCacheResponse} - GetOrExtendReservationRequest {#flyteidl-cacheservice-GetOrExtendReservationRequest} - GetOrExtendReservationResponse {#flyteidl-cacheservice-GetOrExtendReservationResponse} - KeyMapMetadata {#flyteidl-cacheservice-KeyMapMetadata} - KeyMapMetadata.ValuesEntry {#flyteidl-cacheservice-KeyMapMetadata-ValuesEntry} - Metadata {#flyteidl-cacheservice-Metadata} - PutCacheRequest {#flyteidl-cacheservice-PutCacheRequest} - PutCacheResponse {#flyteidl-cacheservice-PutCacheResponse} - ReleaseReservationRequest {#flyteidl-cacheservice-ReleaseReservationRequest} - ReleaseReservationResponse {#flyteidl-cacheservice-ReleaseReservationResponse} - Reservation {#flyteidl-cacheservice-Reservation} - CacheService {#flyteidl-cacheservice-CacheService} - flyteidl/service/signal.proto - SignalService {#flyteidl-service-SignalService} - flyteidl/service/external\_plugin\_service.proto - TaskCreateRequest {#flyteidl-service-TaskCreateRequest} - TaskCreateResponse {#flyteidl-service-TaskCreateResponse} - TaskDeleteRequest {#flyteidl-service-TaskDeleteRequest} - TaskDeleteResponse {#flyteidl-service-TaskDeleteResponse} - TaskGetRequest {#flyteidl-service-TaskGetRequest} - TaskGetResponse {#flyteidl-service-TaskGetResponse} - State {#flyteidl-service-State} - ExternalPluginService {#flyteidl-service-ExternalPluginService} - flyteidl/service/dataproxy.proto - CreateDownloadLinkRequest {#flyteidl-service-CreateDownloadLinkRequest} - CreateDownloadLinkResponse {#flyteidl-service-CreateDownloadLinkResponse} - CreateDownloadLocationRequest {#flyteidl-service-CreateDownloadLocationRequest} - CreateDownloadLocationResponse {#flyteidl-service-CreateDownloadLocationResponse} - CreateUploadLocationRequest {#flyteidl-service-CreateUploadLocationRequest} - CreateUploadLocationResponse {#flyteidl-service-CreateUploadLocationResponse} - CreateUploadLocationResponse.HeadersEntry {#flyteidl-service-CreateUploadLocationResponse-HeadersEntry} - GetDataRequest {#flyteidl-service-GetDataRequest} - GetDataResponse {#flyteidl-service-GetDataResponse} - PreSignedURLs {#flyteidl-service-PreSignedURLs} - ArtifactType {#flyteidl-service-ArtifactType} - DataProxyService {#flyteidl-service-DataProxyService} - flyteidl/service/identity.proto - UserInfoRequest {#flyteidl-service-UserInfoRequest} - UserInfoResponse {#flyteidl-service-UserInfoResponse} - IdentityService {#flyteidl-service-IdentityService} - flyteidl/service/auth.proto - OAuth2MetadataRequest {#flyteidl-service-OAuth2MetadataRequest} - OAuth2MetadataResponse {#flyteidl-service-OAuth2MetadataResponse} - PublicClientAuthConfigRequest {#flyteidl-service-PublicClientAuthConfigRequest} - PublicClientAuthConfigResponse {#flyteidl-service-PublicClientAuthConfigResponse} - AuthMetadataService {#flyteidl-service-AuthMetadataService} - flyteidl/service/agent.proto - AgentMetadataService {#flyteidl-service-AgentMetadataService} - AsyncAgentService {#flyteidl-service-AsyncAgentService} - SyncAgentService {#flyteidl-service-SyncAgentService} - flyteidl/service/admin.proto - AdminService {#flyteidl-service-AdminService} - flyteidl/event/cloudevents.proto - CloudEventExecutionStart {#flyteidl-event-CloudEventExecutionStart} - CloudEventNodeExecution {#flyteidl-event-CloudEventNodeExecution} - CloudEventNodeExecution.LabelsEntry {#flyteidl-event-CloudEventNodeExecution-LabelsEntry} - CloudEventTaskExecution {#flyteidl-event-CloudEventTaskExecution} - CloudEventTaskExecution.LabelsEntry {#flyteidl-event-CloudEventTaskExecution-LabelsEntry} - CloudEventWorkflowExecution {#flyteidl-event-CloudEventWorkflowExecution} - CloudEventWorkflowExecution.LabelsEntry {#flyteidl-event-CloudEventWorkflowExecution-LabelsEntry} - flyteidl/event/event.proto - DynamicWorkflowNodeMetadata {#flyteidl-event-DynamicWorkflowNodeMetadata} - EventReason {#flyteidl-event-EventReason} - ExternalResourceInfo {#flyteidl-event-ExternalResourceInfo} - NodeExecutionEvent {#flyteidl-event-NodeExecutionEvent} - ParentNodeExecutionMetadata {#flyteidl-event-ParentNodeExecutionMetadata} - ParentTaskExecutionMetadata {#flyteidl-event-ParentTaskExecutionMetadata} - ResourcePoolInfo {#flyteidl-event-ResourcePoolInfo} - TaskExecutionEvent {#flyteidl-event-TaskExecutionEvent} - TaskExecutionMetadata {#flyteidl-event-TaskExecutionMetadata} - TaskNodeMetadata {#flyteidl-event-TaskNodeMetadata} - WorkflowExecutionEvent {#flyteidl-event-WorkflowExecutionEvent} - WorkflowNodeMetadata {#flyteidl-event-WorkflowNodeMetadata} - TaskExecutionMetadata.InstanceClass {#flyteidl-event-TaskExecutionMetadata-InstanceClass} - flyteidl/datacatalog/datacatalog.proto - AddTagRequest {#datacatalog-AddTagRequest} - AddTagResponse {#datacatalog-AddTagResponse} - Artifact {#datacatalog-Artifact} - ArtifactData {#datacatalog-ArtifactData} - ArtifactPropertyFilter {#datacatalog-ArtifactPropertyFilter} - CreateArtifactRequest {#datacatalog-CreateArtifactRequest} - CreateArtifactResponse {#datacatalog-CreateArtifactResponse} - CreateDatasetRequest {#datacatalog-CreateDatasetRequest} - CreateDatasetResponse {#datacatalog-CreateDatasetResponse} - Dataset {#datacatalog-Dataset} - DatasetID {#datacatalog-DatasetID} - DatasetPropertyFilter {#datacatalog-DatasetPropertyFilter} - FilterExpression {#datacatalog-FilterExpression} - GetArtifactRequest {#datacatalog-GetArtifactRequest} - GetArtifactResponse {#datacatalog-GetArtifactResponse} - GetDatasetRequest {#datacatalog-GetDatasetRequest} - GetDatasetResponse {#datacatalog-GetDatasetResponse} - GetOrExtendReservationRequest {#datacatalog-GetOrExtendReservationRequest} - GetOrExtendReservationResponse {#datacatalog-GetOrExtendReservationResponse} - KeyValuePair {#datacatalog-KeyValuePair} - ListArtifactsRequest {#datacatalog-ListArtifactsRequest} - ListArtifactsResponse {#datacatalog-ListArtifactsResponse} - ListDatasetsRequest {#datacatalog-ListDatasetsRequest} - ListDatasetsResponse {#datacatalog-ListDatasetsResponse} - Metadata {#datacatalog-Metadata} - Metadata.KeyMapEntry {#datacatalog-Metadata-KeyMapEntry} - PaginationOptions {#datacatalog-PaginationOptions} - Partition {#datacatalog-Partition} - PartitionPropertyFilter {#datacatalog-PartitionPropertyFilter} - ReleaseReservationRequest {#datacatalog-ReleaseReservationRequest} - ReleaseReservationResponse {#datacatalog-ReleaseReservationResponse} - Reservation {#datacatalog-Reservation} - ReservationID {#datacatalog-ReservationID} - SinglePropertyFilter {#datacatalog-SinglePropertyFilter} - Tag {#datacatalog-Tag} - TagPropertyFilter {#datacatalog-TagPropertyFilter} - UpdateArtifactRequest {#datacatalog-UpdateArtifactRequest} - UpdateArtifactResponse {#datacatalog-UpdateArtifactResponse} - PaginationOptions.SortKey {#datacatalog-PaginationOptions-SortKey} - PaginationOptions.SortOrder {#datacatalog-PaginationOptions-SortOrder} - SinglePropertyFilter.ComparisonOperator {#datacatalog-SinglePropertyFilter-ComparisonOperator} - DataCatalog {#datacatalog-DataCatalog} - Scalar Value Types - \[FlyteKit Plugins\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/page.md) - \[FlyteKit Plugins > Apache Airflow\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/airflow/page.md) - \[FlyteKit Plugins > Apache Airflow > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/airflow/classes/page.md) - \[FlyteKit Plugins > Apache Airflow > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/airflow/packages/page.md) - \[FlyteKit Plugins > Apache Airflow > Packages > flytekitplugins.airflow.connector\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/airflow/packages/flytekitplugins.airflow.connector/page.md) - Directory - Classes - Methods - Methods - flytekitplugins.airflow.connector.AirflowConnector - Parameters - Properties - Methods - flytekitplugins.airflow.connector.AirflowMetadata - Parameters - Methods - \[FlyteKit Plugins > Apache Airflow > Packages > flytekitplugins.airflow.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/airflow/packages/flytekitplugins.airflow.task/page.md) - Directory - Classes - Variables - flytekitplugins.airflow.task.AirflowContainerTask - Parameters - Properties - Methods - flytekitplugins.airflow.task.AirflowObj - Parameters - flytekitplugins.airflow.task.AirflowTask - Parameters - Properties - Methods - flytekitplugins.airflow.task.AirflowTaskResolver - Parameters - Properties - Methods - \[FlyteKit Plugins > Async FSSpec\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/async-fsspec/page.md) - \[FlyteKit Plugins > Async FSSpec > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/async-fsspec/classes/page.md) - \[FlyteKit Plugins > Async FSSpec > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/async-fsspec/packages/page.md) - \[FlyteKit Plugins > Async FSSpec > Packages > flytekitplugins.async\_fsspec.s3fs.constants\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/async-fsspec/packages/flytekitplugins.async\_fsspec.s3fs.constants/page.md) - Directory - Variables - \[FlyteKit Plugins > Async FSSpec > Packages > flytekitplugins.async\_fsspec.s3fs.s3fs\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/async-fsspec/packages/flytekitplugins.async\_fsspec.s3fs.s3fs/page.md) - Directory - Classes - Variables - flytekitplugins.async\_fsspec.s3fs.s3fs.AsyncS3FileSystem - Parameters - Properties - Methods - \[FlyteKit Plugins > AWS Athena\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/aws-athena/page.md) - \[FlyteKit Plugins > AWS Batch\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/aws-batch/page.md) - \[FlyteKit Plugins > Page\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/aws-sagemaker/page.md) - Inference - \[FlyteKit Plugins > BigQuery\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/bigquery/page.md) - \[FlyteKit Plugins > BigQuery > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/bigquery/classes/page.md) - \[FlyteKit Plugins > BigQuery > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/bigquery/packages/page.md) - \[FlyteKit Plugins > BigQuery > Packages > flytekitplugins.bigquery.connector\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/bigquery/packages/flytekitplugins.bigquery.connector/page.md) - Directory - Classes - Variables - flytekitplugins.bigquery.connector.BigQueryConnector - Parameters - Properties - Methods - flytekitplugins.bigquery.connector.BigQueryMetadata - Parameters - Methods - \[FlyteKit Plugins > BigQuery > Packages > flytekitplugins.bigquery.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/bigquery/packages/flytekitplugins.bigquery.task/page.md) - Directory - Classes - flytekitplugins.bigquery.task.BigQueryConfig - Parameters - flytekitplugins.bigquery.task.BigQueryTask - Parameters - Properties - Methods - \[FlyteKit Plugins > Comet ML\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/comet-ml/page.md) - \[FlyteKit Plugins > Dask\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/dask/page.md) - \[FlyteKit Plugins > DBT\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/dbt/page.md) - \[FlyteKit Plugins > Page\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/deck-standard/page.md) - Installation - Renderer Requirements - Renderer Descriptions - SourceCodeRenderer - FrameProfilingRenderer - MarkdownRenderer - BoxRenderer - ImageRenderer - TableRenderer - GanttChartRenderer - \[FlyteKit Plugins > Page\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/dolt/page.md) - \[FlyteKit Plugins > DuckDB\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/duckdb/page.md) - \[FlyteKit Plugins > Envd\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/envd/page.md) - \[FlyteKit Plugins > Flyte Interactive\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/flyteinteractive/page.md) - \[FlyteKit Plugins > Page\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/data-fsspec/page.md) - \[FlyteKit Plugins > Google IAP\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/identity-aware-proxy/page.md) - \[FlyteKit Plugins > Google IAP > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/identity-aware-proxy/classes/page.md) - \[FlyteKit Plugins > Google IAP > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/identity-aware-proxy/packages/page.md) - \[FlyteKit Plugins > Google IAP > Packages > flytekitplugins.identity\_aware\_proxy.cli\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/identity-aware-proxy/packages/flytekitplugins.identity\_aware\_proxy.cli/page.md) - Directory - Classes - Methods - Variables - Methods - flytekitplugins.identity\_aware\_proxy.cli.GCPIdentityAwareProxyAuthenticator - Parameters - Methods - \[FlyteKit Plugins > Great Expectations\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/greatexpectations/page.md) - \[FlyteKit Plugins > Hive\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/hive/page.md) - \[FlyteKit Plugins > Hive > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/hive/classes/page.md) - \[FlyteKit Plugins > Hive > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/hive/packages/page.md) - \[FlyteKit Plugins > Hive > Packages > flytekitplugins.hive.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/hive/packages/flytekitplugins.hive.task/page.md) - Directory - Classes - flytekitplugins.hive.task.HiveConfig - Parameters - flytekitplugins.hive.task.HiveSelectTask - Parameters - Properties - Methods - flytekitplugins.hive.task.HiveTask - Parameters - Properties - Methods - \[FlyteKit Plugins > Hugging Face\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/huggingface/page.md) - \[FlyteKit Plugins > Hugging Face > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/huggingface/classes/page.md) - \[FlyteKit Plugins > Hugging Face > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/huggingface/packages/page.md) - \[FlyteKit Plugins > Hugging Face > Packages > flytekitplugins.huggingface.sd\_transformers\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/huggingface/packages/flytekitplugins.huggingface.sd\_transformers/page.md) - Directory - Classes - Variables - flytekitplugins.huggingface.sd\_transformers.HuggingFaceDatasetRenderer - Methods - flytekitplugins.huggingface.sd\_transformers.HuggingFaceDatasetToParquetEncodingHandler - Parameters - Properties - Methods - flytekitplugins.huggingface.sd\_transformers.ParquetToHuggingFaceDatasetDecodingHandler - Parameters - Properties - Methods - \[FlyteKit Plugins > Inference\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/inference/page.md) - \[FlyteKit Plugins > Inference > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/inference/classes/page.md) - \[FlyteKit Plugins > Inference > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/inference/packages/page.md) - \[FlyteKit Plugins > Inference > Packages > flytekitplugins.inference.nim.serve\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/inference/packages/flytekitplugins.inference.nim.serve/page.md) - Directory - Classes - flytekitplugins.inference.nim.serve.NIM - Parameters - Properties - Methods - flytekitplugins.inference.nim.serve.NIMSecrets - Parameters - \[FlyteKit Plugins > Inference > Packages > flytekitplugins.inference.ollama.serve\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/inference/packages/flytekitplugins.inference.ollama.serve/page.md) - Directory - Classes - flytekitplugins.inference.ollama.serve.Model - Parameters - flytekitplugins.inference.ollama.serve.Ollama - Parameters - Properties - Methods - \[FlyteKit Plugins > Inference > Packages > flytekitplugins.inference.sidecar\_template\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/inference/packages/flytekitplugins.inference.sidecar\_template/page.md) - Directory - Classes - flytekitplugins.inference.sidecar\_template.ModelInferenceTemplate - Parameters - Properties - \[FlyteKit Plugins > Inference > Packages > flytekitplugins.inference.vllm.serve\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/inference/packages/flytekitplugins.inference.vllm.serve/page.md) - Directory - Classes - flytekitplugins.inference.vllm.serve.HFSecret - Parameters - flytekitplugins.inference.vllm.serve.VLLM - Parameters - Properties - Methods - \[FlyteKit Plugins > Kubeflow MPI\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-mpi/page.md) - \[FlyteKit Plugins > Kubeflow MPI > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-mpi/classes/page.md) - \[FlyteKit Plugins > Kubeflow MPI > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-mpi/packages/page.md) - \[FlyteKit Plugins > Kubeflow MPI > Packages > flytekitplugins.kfmpi.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-mpi/packages/flytekitplugins.kfmpi.task/page.md) - Directory - Classes - flytekitplugins.kfmpi.task.CleanPodPolicy - flytekitplugins.kfmpi.task.HorovodFunctionTask - Parameters - Properties - Methods - flytekitplugins.kfmpi.task.HorovodJob - Parameters - flytekitplugins.kfmpi.task.Launcher - Parameters - flytekitplugins.kfmpi.task.MPIFunctionTask - Parameters - Properties - Methods - flytekitplugins.kfmpi.task.MPIJob - Parameters - flytekitplugins.kfmpi.task.RestartPolicy - flytekitplugins.kfmpi.task.RunPolicy - Parameters - flytekitplugins.kfmpi.task.Worker - Parameters - \[FlyteKit Plugins > Kubeflow PyTorch\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-pytorch/page.md) - \[FlyteKit Plugins > Kubeflow PyTorch > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-pytorch/classes/page.md) - \[FlyteKit Plugins > Kubeflow PyTorch > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-pytorch/packages/page.md) - \[FlyteKit Plugins > Kubeflow PyTorch > Packages > flytekitplugins.kfpytorch.error\_handling\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-pytorch/packages/flytekitplugins.kfpytorch.error\_handling/page.md) - Directory - Methods - Variables - Methods - \[FlyteKit Plugins > Kubeflow PyTorch > Packages > flytekitplugins.kfpytorch.pod\_template\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-pytorch/packages/flytekitplugins.kfpytorch.pod\_template/page.md) - Directory - Methods - Methods - \[FlyteKit Plugins > Kubeflow PyTorch > Packages > flytekitplugins.kfpytorch.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-pytorch/packages/flytekitplugins.kfpytorch.task/page.md) - Directory - Classes - Methods - Variables - Methods - flytekitplugins.kfpytorch.task.CleanPodPolicy - flytekitplugins.kfpytorch.task.Elastic - Parameters - flytekitplugins.kfpytorch.task.ElasticWorkerResult - flytekitplugins.kfpytorch.task.Master - Parameters - flytekitplugins.kfpytorch.task.PyTorch - Parameters - flytekitplugins.kfpytorch.task.PyTorchFunctionTask - Parameters - Properties - Methods - flytekitplugins.kfpytorch.task.PytorchElasticFunctionTask - Parameters - Properties - Methods - flytekitplugins.kfpytorch.task.RestartPolicy - flytekitplugins.kfpytorch.task.RunPolicy - Parameters - flytekitplugins.kfpytorch.task.Worker - Parameters - \[FlyteKit Plugins > Kubeflow TensorFlow\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-tensorflow/page.md) - \[FlyteKit Plugins > Kubeflow TensorFlow > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-tensorflow/classes/page.md) - \[FlyteKit Plugins > Kubeflow TensorFlow > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-tensorflow/packages/page.md) - \[FlyteKit Plugins > Kubeflow TensorFlow > Packages > flytekitplugins.kftensorflow.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-tensorflow/packages/flytekitplugins.kftensorflow.task/page.md) - Directory - Classes - flytekitplugins.kftensorflow.task.Chief - Parameters - flytekitplugins.kftensorflow.task.CleanPodPolicy - flytekitplugins.kftensorflow.task.Evaluator - Parameters - flytekitplugins.kftensorflow.task.PS - Parameters - flytekitplugins.kftensorflow.task.RestartPolicy - flytekitplugins.kftensorflow.task.RunPolicy - Parameters - flytekitplugins.kftensorflow.task.TensorflowFunctionTask - Parameters - Properties - Methods - flytekitplugins.kftensorflow.task.TfJob - Parameters - flytekitplugins.kftensorflow.task.Worker - Parameters - \[FlyteKit Plugins > Kubernetes Pod\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/k8s-pod/page.md) - \[FlyteKit Plugins > Kubernetes Pod > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/k8s-pod/classes/page.md) - \[FlyteKit Plugins > Kubernetes Pod > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/k8s-pod/packages/page.md) - \[FlyteKit Plugins > Kubernetes Pod > Packages > flytekitplugins.pod.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/k8s-pod/packages/flytekitplugins.pod.task/page.md) - Directory - Classes - Variables - flytekitplugins.pod.task.Pod - Parameters - flytekitplugins.pod.task.PodFunctionTask - Parameters - Properties - Methods - \[FlyteKit Plugins > Memory Machine Cloud\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/mmcloud/page.md) - \[FlyteKit Plugins > Memory Machine Cloud > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/mmcloud/classes/page.md) - \[FlyteKit Plugins > Memory Machine Cloud > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/mmcloud/packages/page.md) - \[FlyteKit Plugins > Memory Machine Cloud > Packages > flytekitplugins.mmcloud.connector\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/mmcloud/packages/flytekitplugins.mmcloud.connector/page.md) - Directory - Classes - flytekitplugins.mmcloud.connector.MMCloudConnector - Parameters - Properties - Methods - flytekitplugins.mmcloud.connector.MMCloudMetadata - Parameters - Methods - \[FlyteKit Plugins > Memory Machine Cloud > Packages > flytekitplugins.mmcloud.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/mmcloud/packages/flytekitplugins.mmcloud.task/page.md) - Directory - Classes - flytekitplugins.mmcloud.task.MMCloudConfig - Parameters - flytekitplugins.mmcloud.task.MMCloudTask - Parameters - Properties - Methods - \[FlyteKit Plugins > Memory Machine Cloud > Packages > flytekitplugins.mmcloud.utils\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/mmcloud/packages/flytekitplugins.mmcloud.utils/page.md) - Directory - Methods - Variables - Methods - \[FlyteKit Plugins > Memray Profiling\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/memray/page.md) - \[FlyteKit Plugins > Memray Profiling > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/memray/classes/page.md) - \[FlyteKit Plugins > Memray Profiling > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/memray/packages/page.md) - \[FlyteKit Plugins > Memray Profiling > Packages > flytekitplugins.memray\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/memray/packages/flytekitplugins.memray/page.md) - Directory - Classes - flytekitplugins.memray.memray\_profiling - Parameters - Methods - \[FlyteKit Plugins > MLflow\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/mlflow/page.md) - \[FlyteKit Plugins > Page\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/modin/page.md) - \[FlyteKit Plugins > Neptune\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/neptune/page.md) - \[FlyteKit Plugins > OmegaConf\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/omegaconf/page.md) - \[FlyteKit Plugins > OmegaConf > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/omegaconf/classes/page.md) - \[FlyteKit Plugins > OmegaConf > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/omegaconf/packages/page.md) - \[FlyteKit Plugins > OmegaConf > Packages > flytekitplugins.omegaconf\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/omegaconf/packages/flytekitplugins.omegaconf/page.md) - Directory - Classes - Methods - Methods - flytekitplugins.omegaconf.OmegaConfTransformerMode - \[FlyteKit Plugins > Page\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/onnx-pytorch/page.md) - \[FlyteKit Plugins > ONNX ScikitLearn\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/onnx-scikitlearn/page.md) - \[FlyteKit Plugins > ONNX ScikitLearn > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/onnx-scikitlearn/classes/page.md) - \[FlyteKit Plugins > ONNX ScikitLearn > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/onnx-scikitlearn/packages/page.md) - \[FlyteKit Plugins > ONNX ScikitLearn > Packages > flytekitplugins.onnxscikitlearn.schema\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/onnx-scikitlearn/packages/flytekitplugins.onnxscikitlearn.schema/page.md) - Directory - Classes - Methods - Methods - flytekitplugins.onnxscikitlearn.schema.ScikitLearn2ONNX - Parameters - Methods - flytekitplugins.onnxscikitlearn.schema.ScikitLearn2ONNXConfig - Parameters - Methods - flytekitplugins.onnxscikitlearn.schema.ScikitLearn2ONNXTransformer - Parameters - Properties - Methods - \[FlyteKit Plugins > Page\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/onnx-tensorflow/page.md) - \[FlyteKit Plugins > OpenAI\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/openai/page.md) - \[FlyteKit Plugins > OpenAI > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/openai/classes/page.md) - \[FlyteKit Plugins > OpenAI > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/openai/packages/page.md) - \[FlyteKit Plugins > OpenAI > Packages > flytekitplugins.openai.batch.connector\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/openai/packages/flytekitplugins.openai.batch.connector/page.md) - Directory - Classes - Variables - flytekitplugins.openai.batch.connector.BatchEndpointConnector - Parameters - Properties - Methods - flytekitplugins.openai.batch.connector.BatchEndpointMetadata - Parameters - Methods - flytekitplugins.openai.batch.connector.State - \[FlyteKit Plugins > OpenAI > Packages > flytekitplugins.openai.batch.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/openai/packages/flytekitplugins.openai.batch.task/page.md) - Directory - Classes - flytekitplugins.openai.batch.task.BatchEndpointTask - Parameters - Properties - Methods - flytekitplugins.openai.batch.task.BatchResult - Parameters - Methods - flytekitplugins.openai.batch.task.DownloadJSONFilesExecutor - Parameters - Properties - Methods - flytekitplugins.openai.batch.task.DownloadJSONFilesTask - Parameters - Properties - Methods - flytekitplugins.openai.batch.task.OpenAIFileConfig - Parameters - flytekitplugins.openai.batch.task.OpenAIFileDefaultImages - Methods - flytekitplugins.openai.batch.task.UploadJSONLFileExecutor - Parameters - Properties - Methods - flytekitplugins.openai.batch.task.UploadJSONLFileTask - Parameters - Properties - Methods - \[FlyteKit Plugins > OpenAI > Packages > flytekitplugins.openai.batch.workflow\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/openai/packages/flytekitplugins.openai.batch.workflow/page.md) - Directory - Methods - Methods - \[FlyteKit Plugins > OpenAI > Packages > flytekitplugins.openai.chatgpt.connector\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/openai/packages/flytekitplugins.openai.chatgpt.connector/page.md) - Directory - Classes - Variables - flytekitplugins.openai.chatgpt.connector.ChatGPTConnector - Parameters - Properties - Methods - \[FlyteKit Plugins > OpenAI > Packages > flytekitplugins.openai.chatgpt.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/openai/packages/flytekitplugins.openai.chatgpt.task/page.md) - Directory - Classes - flytekitplugins.openai.chatgpt.task.ChatGPTTask - Parameters - Properties - Methods - \[FlyteKit Plugins > Optuna\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/optuna/page.md) - \[FlyteKit Plugins > Optuna > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/optuna/classes/page.md) - \[FlyteKit Plugins > Optuna > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/optuna/packages/page.md) - \[FlyteKit Plugins > Optuna > Packages > flytekitplugins.optuna\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/optuna/packages/flytekitplugins.optuna/page.md) - Directory - Classes - Methods - Variables - Methods - flytekitplugins.optuna.Optimizer - Parameters - Properties - Methods - \[FlyteKit Plugins > Pandera\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/pandera/page.md) - \[FlyteKit Plugins > Pandera > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/pandera/classes/page.md) - \[FlyteKit Plugins > Pandera > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/pandera/packages/page.md) - \[FlyteKit Plugins > Pandera > Packages > flytekitplugins.pandera.config\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/pandera/packages/flytekitplugins.pandera.config/page.md) - Directory - Classes - flytekitplugins.pandera.config.ValidationConfig - Parameters - \[FlyteKit Plugins > Pandera > Packages > flytekitplugins.pandera.pandas\_renderer\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/pandera/packages/flytekitplugins.pandera.pandas\_renderer/page.md) - Directory - Classes - Variables - flytekitplugins.pandera.pandas\_renderer.PandasReport - Parameters - flytekitplugins.pandera.pandas\_renderer.PandasReportRenderer - Parameters - Methods - \[FlyteKit Plugins > Pandera > Packages > flytekitplugins.pandera.pandas\_transformer\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/pandera/packages/flytekitplugins.pandera.pandas\_transformer/page.md) - Directory - Classes - Variables - flytekitplugins.pandera.pandas\_transformer.PanderaPandasTransformer - Parameters - Properties - Methods - \[FlyteKit Plugins > Papermill\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/papermill/page.md) - \[FlyteKit Plugins > Papermill > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/papermill/classes/page.md) - \[FlyteKit Plugins > Papermill > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/papermill/packages/page.md) - \[FlyteKit Plugins > Papermill > Packages > flytekitplugins.papermill.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/papermill/packages/flytekitplugins.papermill.task/page.md) - Directory - Classes - Methods - Variables - Methods - flytekitplugins.papermill.task.NotebookTask - Parameters - Properties - Methods - \[FlyteKit Plugins > Page\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/perian/page.md) - \[FlyteKit Plugins > Page > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/perian/classes/page.md) - \[FlyteKit Plugins > Page > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/perian/packages/page.md) - \[FlyteKit Plugins > Page > Packages > flytekitplugins.perian\_job.agent\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/perian/packages/flytekitplugins.perian\_job.agent/page.md) - Directory - Classes - Variables - flytekitplugins.perian\_job.agent.PerianAgent - Methods - Properties - flytekitplugins.perian\_job.agent.PerianMetadata - Methods - \[FlyteKit Plugins > Page > Packages > flytekitplugins.perian\_job.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/perian/packages/flytekitplugins.perian\_job.task/page.md) - Directory - Classes - flytekitplugins.perian\_job.task.PerianConfig - flytekitplugins.perian\_job.task.PerianContainerTask - Methods - Properties - flytekitplugins.perian\_job.task.PerianTask - Methods - Properties - \[FlyteKit Plugins > Polars\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/polars/page.md) - \[FlyteKit Plugins > Polars > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/polars/classes/page.md) - \[FlyteKit Plugins > Polars > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/polars/packages/page.md) - \[FlyteKit Plugins > Polars > Packages > flytekitplugins.polars.sd\_transformers\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/polars/packages/flytekitplugins.polars.sd\_transformers/page.md) - Directory - Classes - Variables - flytekitplugins.polars.sd\_transformers.ParquetToPolarsDataFrameDecodingHandler - Parameters - Properties - Methods - flytekitplugins.polars.sd\_transformers.ParquetToPolarsLazyFrameDecodingHandler - Parameters - Properties - Methods - flytekitplugins.polars.sd\_transformers.PolarsDataFrameRenderer - Methods - flytekitplugins.polars.sd\_transformers.PolarsDataFrameToParquetEncodingHandler - Parameters - Properties - Methods - flytekitplugins.polars.sd\_transformers.PolarsLazyFrameToParquetEncodingHandler - Parameters - Properties - Methods - \[FlyteKit Plugins > Ray\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/ray/page.md) - \[FlyteKit Plugins > Ray > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/ray/classes/page.md) - \[FlyteKit Plugins > Ray > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/ray/packages/page.md) - \[FlyteKit Plugins > Ray > Packages > flytekitplugins.ray.models\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/ray/packages/flytekitplugins.ray.models/page.md) - Directory - Classes - flytekitplugins.ray.models.AutoscalerOptions - Parameters - Properties - Methods - flytekitplugins.ray.models.HeadGroupSpec - Parameters - Properties - Methods - flytekitplugins.ray.models.RayCluster - Parameters - Properties - Methods - flytekitplugins.ray.models.RayJob - Parameters - Properties - Methods - flytekitplugins.ray.models.WorkerGroupSpec - Parameters - Properties - Methods - \[FlyteKit Plugins > Ray > Packages > flytekitplugins.ray.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/ray/packages/flytekitplugins.ray.task/page.md) - Directory - Classes - flytekitplugins.ray.task.AutoscalerOptionsConfig - Parameters - flytekitplugins.ray.task.HeadNodeConfig - Parameters - flytekitplugins.ray.task.RayFunctionTask - Parameters - Properties - Methods - flytekitplugins.ray.task.RayJobConfig - Parameters - flytekitplugins.ray.task.WorkerNodeConfig - Parameters - \[FlyteKit Plugins > Slurm\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/slurm/page.md) - \[FlyteKit Plugins > Slurm > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/slurm/classes/page.md) - \[FlyteKit Plugins > Slurm > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/slurm/packages/page.md) - \[FlyteKit Plugins > Slurm > Packages > flytekitplugins.slurm.ssh\_utils\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/slurm/packages/flytekitplugins.slurm.ssh\_utils/page.md) - Directory - Classes - Methods - Variables - Methods - flytekitplugins.slurm.ssh\_utils.SSHConfig - Parameters - Methods - flytekitplugins.slurm.ssh\_utils.SlurmCluster - Parameters - \[FlyteKit Plugins > Snowflake\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/snowflake/page.md) - \[FlyteKit Plugins > Snowflake > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/snowflake/classes/page.md) - \[FlyteKit Plugins > Snowflake > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/snowflake/packages/page.md) - \[FlyteKit Plugins > Snowflake > Packages > flytekitplugins.snowflake.connector\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/snowflake/packages/flytekitplugins.snowflake.connector/page.md) - Directory - Classes - Methods - Variables - Methods - flytekitplugins.snowflake.connector.SnowflakeConnector - Parameters - Properties - Methods - flytekitplugins.snowflake.connector.SnowflakeJobMetadata - Parameters - Methods - \[FlyteKit Plugins > Snowflake > Packages > flytekitplugins.snowflake.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/snowflake/packages/flytekitplugins.snowflake.task/page.md) - Directory - Classes - flytekitplugins.snowflake.task.SnowflakeConfig - Parameters - flytekitplugins.snowflake.task.SnowflakeTask - Parameters - Properties - Methods - \[FlyteKit Plugins > Spark\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/page.md) - \[FlyteKit Plugins > Spark > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/classes/page.md) - \[FlyteKit Plugins > Spark > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/packages/page.md) - \[FlyteKit Plugins > Spark > Packages > flytekitplugins.spark.connector\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/packages/flytekitplugins.spark.connector/page.md) - Directory - Classes - Methods - Variables - Methods - flytekitplugins.spark.connector.DatabricksConnector - Parameters - Properties - Methods - flytekitplugins.spark.connector.DatabricksConnectorV2 - Parameters - Properties - Methods - flytekitplugins.spark.connector.DatabricksJobMetadata - Parameters - Methods - \[FlyteKit Plugins > Spark > Packages > flytekitplugins.spark.generic\_task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/packages/flytekitplugins.spark.generic\_task/page.md) - Directory - Classes - flytekitplugins.spark.generic\_task.GenericSparkConf - Parameters - flytekitplugins.spark.generic\_task.GenericSparkTask - Parameters - Properties - Methods - \[FlyteKit Plugins > Spark > Packages > flytekitplugins.spark.models\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/packages/flytekitplugins.spark.models/page.md) - Directory - Classes - flytekitplugins.spark.models.SparkJob - Parameters - Properties - Methods - flytekitplugins.spark.models.SparkType - \[FlyteKit Plugins > Spark > Packages > flytekitplugins.spark.pyspark\_transformers\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/packages/flytekitplugins.spark.pyspark\_transformers/page.md) - Directory - Classes - flytekitplugins.spark.pyspark\_transformers.PySparkPipelineModelTransformer - Parameters - Properties - Methods - \[FlyteKit Plugins > Spark > Packages > flytekitplugins.spark.schema\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/packages/flytekitplugins.spark.schema/page.md) - Directory - Classes - flytekitplugins.spark.schema.ClassicSparkDataFrameSchemaReader - Parameters - Properties - Methods - flytekitplugins.spark.schema.ClassicSparkDataFrameSchemaWriter - Parameters - Properties - Methods - flytekitplugins.spark.schema.ClassicSparkDataFrameTransformer - Parameters - Properties - Methods - flytekitplugins.spark.schema.SparkDataFrameSchemaReader - Parameters - Properties - Methods - flytekitplugins.spark.schema.SparkDataFrameSchemaWriter - Parameters - Properties - Methods - flytekitplugins.spark.schema.SparkDataFrameTransformer - Parameters - Properties - Methods - \[FlyteKit Plugins > Spark > Packages > flytekitplugins.spark.sd\_transformers\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/packages/flytekitplugins.spark.sd\_transformers/page.md) - Directory - Classes - Variables - flytekitplugins.spark.sd\_transformers.ClassicSparkToParquetEncodingHandler - Parameters - Properties - Methods - flytekitplugins.spark.sd\_transformers.ParquetToClassicSparkDecodingHandler - Parameters - Properties - Methods - flytekitplugins.spark.sd\_transformers.ParquetToSparkDecodingHandler - Parameters - Properties - Methods - flytekitplugins.spark.sd\_transformers.SparkDataFrameRenderer - Methods - flytekitplugins.spark.sd\_transformers.SparkToParquetEncodingHandler - Parameters - Properties - Methods - \[FlyteKit Plugins > Spark > Packages > flytekitplugins.spark.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/packages/flytekitplugins.spark.task/page.md) - Directory - Classes - Methods - Variables - Methods - flytekitplugins.spark.task.Databricks - Parameters - flytekitplugins.spark.task.DatabricksV2 - Parameters - flytekitplugins.spark.task.PysparkFunctionTask - Parameters - Properties - Methods - flytekitplugins.spark.task.Spark - Parameters - \[FlyteKit Plugins > Spark > Packages > flytekitplugins.spark.utils\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/packages/flytekitplugins.spark.utils/page.md) - Directory - Methods - Methods - \[FlyteKit Plugins > SQLAlchemy\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/sqlalchemy/page.md) - \[FlyteKit Plugins > SQLAlchemy > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/sqlalchemy/classes/page.md) - \[FlyteKit Plugins > SQLAlchemy > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/sqlalchemy/packages/page.md) - \[FlyteKit Plugins > SQLAlchemy > Packages > flytekitplugins.sqlalchemy.task\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/sqlalchemy/packages/flytekitplugins.sqlalchemy.task/page.md) - Directory - Classes - flytekitplugins.sqlalchemy.task.SQLAlchemyConfig - Parameters - Methods - flytekitplugins.sqlalchemy.task.SQLAlchemyDefaultImages - Methods - flytekitplugins.sqlalchemy.task.SQLAlchemyTask - Parameters - Properties - Methods - flytekitplugins.sqlalchemy.task.SQLAlchemyTaskExecutor - Parameters - Properties - Methods - \[FlyteKit Plugins > Page\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/vaex/page.md) - \[FlyteKit Plugins > Weights & Biases\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/wandb/page.md) - \[FlyteKit Plugins > Weights & Biases > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/wandb/classes/page.md) - \[FlyteKit Plugins > Weights & Biases > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/wandb/packages/page.md) - \[FlyteKit Plugins > Weights & Biases > Packages > flytekitplugins.wandb\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/wandb/packages/flytekitplugins.wandb/page.md) - Directory - Classes - flytekitplugins.wandb.wandb\_init - Parameters - Methods - \[FlyteKit Plugins > whylogs\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/whylogs/page.md) - \[FlyteKit Plugins > whylogs > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/whylogs/classes/page.md) - \[FlyteKit Plugins > whylogs > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/whylogs/packages/page.md) - \[FlyteKit Plugins > whylogs > Packages > flytekitplugins.whylogs\](https://www.union.ai/docs/v1/flyte/api-reference/plugins/whylogs/packages/flytekitplugins.whylogs/page.md) - Directory - Classes - flytekitplugins.whylogs.WhylogsConstraintsRenderer - Methods - flytekitplugins.whylogs.WhylogsDatasetProfileTransformer - Parameters - Properties - Methods - flytekitplugins.whylogs.WhylogsSummaryDriftRenderer - Methods - \[Flytekit SDK\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/page.md) - Developing on Flyte - Developing on Union - \[Flytekit SDK > Classes\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/classes/page.md) - \[Flytekit SDK > Packages\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/page.md) - \[Flytekit SDK > Packages > flytekit\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit/page.md) - Basic Authoring - Branching and Conditionals - Customizing Tasks & Workflows - Directory - Methods - Variables - Methods - \[Flytekit SDK > Packages > flytekit.bin.entrypoint\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.bin.entrypoint/page.md) - Directory - Methods - Variables - Methods - \[Flytekit SDK > Packages > flytekit.clients.auth\_helper\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.auth\_helper/page.md) - Directory - Classes - Methods - Methods - flytekit.clients.auth\_helper.AuthenticationHTTPAdapter - Parameters - Methods - flytekit.clients.auth\_helper.RemoteClientConfigStore - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.clients.auth.auth\_client\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.auth.auth\_client/page.md) - Directory - Classes - flytekit.clients.auth.auth\_client.AuthorizationClient - Parameters - Methods - flytekit.clients.auth.auth\_client.AuthorizationCode - Parameters - Properties - flytekit.clients.auth.auth\_client.EndpointMetadata - Parameters - flytekit.clients.auth.auth\_client.OAuthCallbackHandler - Parameters - Methods - flytekit.clients.auth.auth\_client.OAuthHTTPServer - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.clients.auth.authenticator\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.auth.authenticator/page.md) - Directory - Classes - flytekit.clients.auth.authenticator.Authenticator - Parameters - Methods - flytekit.clients.auth.authenticator.ClientConfig - Parameters - flytekit.clients.auth.authenticator.ClientConfigStore - Methods - flytekit.clients.auth.authenticator.ClientCredentialsAuthenticator - Parameters - Methods - flytekit.clients.auth.authenticator.CommandAuthenticator - Parameters - Methods - flytekit.clients.auth.authenticator.DeviceCodeAuthenticator - Parameters - Methods - flytekit.clients.auth.authenticator.PKCEAuthenticator - Parameters - Methods - flytekit.clients.auth.authenticator.StaticClientConfigStore - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.clients.auth.default\_html\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.auth.default\_html/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.clients.auth.exceptions\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.auth.exceptions/page.md) - Directory - Errors - flytekit.clients.auth.exceptions.AccessTokenNotFoundError - flytekit.clients.auth.exceptions.AuthenticationError - flytekit.clients.auth.exceptions.AuthenticationPending - \[Flytekit SDK > Packages > flytekit.clients.auth.keyring\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.auth.keyring/page.md) - Directory - Classes - flytekit.clients.auth.keyring.Credentials - Parameters - flytekit.clients.auth.keyring.KeyringStore - Methods - \[Flytekit SDK > Packages > flytekit.clients.auth.token\_client\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.auth.token\_client/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.clients.auth.token\_client.DeviceCodeResponse - Parameters - Methods - flytekit.clients.auth.token\_client.GrantType - Parameters - \[Flytekit SDK > Packages > flytekit.clients.friendly\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.friendly/page.md) - Directory - Classes - Variables - flytekit.clients.friendly.SynchronousFlyteClient - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.clients.grpc\_utils.auth\_interceptor\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.grpc\_utils.auth\_interceptor/page.md) - Directory - Classes - flytekit.clients.grpc\_utils.auth\_interceptor.AuthUnaryInterceptor - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.clients.grpc\_utils.deadline\_interceptor\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.grpc\_utils.deadline\_interceptor/page.md) - Directory - Classes - Methods - Methods - flytekit.clients.grpc\_utils.deadline\_interceptor.ScopedGrpcDeadlineInterceptor - Methods - \[Flytekit SDK > Packages > flytekit.clients.grpc\_utils.default\_metadata\_interceptor\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.grpc\_utils.default\_metadata\_interceptor/page.md) - Directory - Classes - flytekit.clients.grpc\_utils.default\_metadata\_interceptor.DefaultMetadataInterceptor - Methods - \[Flytekit SDK > Packages > flytekit.clients.grpc\_utils.wrap\_exception\_interceptor\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.grpc\_utils.wrap\_exception\_interceptor/page.md) - Directory - Classes - flytekit.clients.grpc\_utils.wrap\_exception\_interceptor.RetryExceptionWrapperInterceptor - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.clients.helpers\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.helpers/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.clients.raw\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clients.raw/page.md) - Directory - Classes - flytekit.clients.raw.RawSynchronousFlyteClient - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.clis.helpers\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.helpers/page.md) - Directory - Methods - Variables - Methods - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.backfill\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.backfill/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.build\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.build/page.md) - Directory - Classes - Methods - Methods - flytekit.clis.sdk\_in\_container.build.BuildCommand - Parameters - Properties - Methods - flytekit.clis.sdk\_in\_container.build.BuildParams - Parameters - Properties - Methods - flytekit.clis.sdk\_in\_container.build.BuildWorkflowCommand - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.constants\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.constants/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.executions\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.executions/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.helpers\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.helpers/page.md) - Directory - Methods - Variables - Methods - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.metrics\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.metrics/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.package\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.package/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.pyflyte\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.pyflyte/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.run\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.run/page.md) - Directory - Classes - Methods - Methods - flytekit.clis.sdk\_in\_container.run.DynamicEntityLaunchCommand - Parameters - Properties - Methods - flytekit.clis.sdk\_in\_container.run.Entities - Methods - flytekit.clis.sdk\_in\_container.run.RemoteEntityGroup - Parameters - Properties - Methods - flytekit.clis.sdk\_in\_container.run.RunCommand - Parameters - Properties - Methods - flytekit.clis.sdk\_in\_container.run.RunLevelComputedParams - Parameters - flytekit.clis.sdk\_in\_container.run.RunLevelParams - Parameters - Properties - Methods - flytekit.clis.sdk\_in\_container.run.WorkflowCommand - Parameters - Properties - Methods - flytekit.clis.sdk\_in\_container.run.YamlFileReadingCommand - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.serialize\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.serialize/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.clis.sdk\_in\_container.serialize.SerializationMode - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.serve\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.serve/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.clis.sdk\_in\_container.utils\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.sdk\_in\_container.utils/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.clis.sdk\_in\_container.utils.ErrorHandlingCommand - Parameters - Properties - Methods - flytekit.clis.sdk\_in\_container.utils.PyFlyteParams - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.clis.version\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.clis.version/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.configuration\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.configuration/page.md) - Flytekit Configuration Sources - Command Line Arguments - Python Config Object - Environment Variables - YAML Format Configuration File - INI Format Configuration File - How is configuration used? - Configuration Objects - Serialization Time Settings - Execution Time Settings - Directory - Classes - Variables - flytekit.configuration.AuthType - flytekit.configuration.AzureBlobStorageConfig - Parameters - Methods - flytekit.configuration.Config - Parameters - Methods - flytekit.configuration.DataConfig - Parameters - Methods - flytekit.configuration.EntrypointSettings - Parameters - Methods - flytekit.configuration.FastSerializationSettings - Parameters - Methods - flytekit.configuration.GCSConfig - Parameters - Methods - flytekit.configuration.GenericPersistenceConfig - Parameters - Methods - flytekit.configuration.Image - Parameters - Properties - Methods - flytekit.configuration.ImageConfig - Parameters - Methods - flytekit.configuration.LocalConfig - Parameters - Methods - flytekit.configuration.PlatformConfig - Parameters - Methods - flytekit.configuration.S3Config - Parameters - Methods - flytekit.configuration.SecretsConfig - Parameters - Methods - flytekit.configuration.SerializationSettings - Parameters - Properties - Methods - flytekit.configuration.StatsConfig - Parameters - Methods - flytekit.configuration.TaskConfig - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.configuration.default\_images\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.configuration.default\_images/page.md) - Directory - Classes - Variables - flytekit.configuration.default\_images.DefaultImages - Methods - flytekit.configuration.default\_images.PythonVersion - \[Flytekit SDK > Packages > flytekit.configuration.feature\_flags\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.configuration.feature\_flags/page.md) - Directory - Classes - flytekit.configuration.feature\_flags.FeatureFlags - \[Flytekit SDK > Packages > flytekit.configuration.file\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.configuration.file/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.configuration.file.ConfigEntry - Parameters - Methods - flytekit.configuration.file.ConfigFile - Parameters - Properties - Methods - flytekit.configuration.file.LegacyConfigEntry - Parameters - Methods - flytekit.configuration.file.YamlConfigEntry - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.configuration.internal\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.configuration.internal/page.md) - Directory - Classes - flytekit.configuration.internal.AWS - flytekit.configuration.internal.AZURE - flytekit.configuration.internal.Credentials - flytekit.configuration.internal.GCP - flytekit.configuration.internal.Images - Methods - flytekit.configuration.internal.Local - flytekit.configuration.internal.LocalSDK - flytekit.configuration.internal.Persistence - flytekit.configuration.internal.Platform - flytekit.configuration.internal.Secrets - flytekit.configuration.internal.StatsD - \[Flytekit SDK > Packages > flytekit.configuration.plugin\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.configuration.plugin/page.md) - Directory - Classes - Protocols - Methods - Methods - flytekit.configuration.plugin.FlytekitPlugin - Methods - flytekit.configuration.plugin.FlytekitPluginProtocol - Methods - \[Flytekit SDK > Packages > flytekit.constants\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.constants/page.md) - Directory - Classes - flytekit.constants.CopyFileDetection - \[Flytekit SDK > Packages > flytekit.core.annotation\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.annotation/page.md) - Directory - Classes - flytekit.core.annotation.FlyteAnnotation - Parameters - Properties - \[Flytekit SDK > Packages > flytekit.core.array\_node\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.array\_node/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.array\_node.ArrayNode - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.array\_node\_map\_task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.array\_node\_map\_task/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.array\_node\_map\_task.ArrayNodeMapTask - Parameters - Properties - Methods - flytekit.core.array\_node\_map\_task.ArrayNodeMapTaskResolver - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.artifact\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.artifact/page.md) - Directory - Classes - Protocols - Variables - flytekit.core.artifact.Artifact - Parameters - Properties - Methods - flytekit.core.artifact.ArtifactIDSpecification - Parameters - Methods - flytekit.core.artifact.ArtifactQuery - Parameters - Properties - Methods - flytekit.core.artifact.ArtifactSerializationHandler - Methods - flytekit.core.artifact.DefaultArtifactSerializationHandler - Methods - flytekit.core.artifact.InputsBase - flytekit.core.artifact.Partition - Parameters - flytekit.core.artifact.Partitions - Parameters - Properties - Methods - flytekit.core.artifact.Serializer - Methods - flytekit.core.artifact.TimePartition - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.core.artifact\_utils\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.artifact\_utils/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.core.base\_sql\_task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.base\_sql\_task/page.md) - Directory - Classes - Variables - flytekit.core.base\_sql\_task.SQLTask - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.base\_task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.base\_task/page.md) - Core Components - kwtypes - PythonTask - Task - TaskResolverMixin - IgnoreOutputs - Directory - Classes - Errors - Methods - Variables - Methods - flytekit.core.base\_task.IgnoreOutputs - flytekit.core.base\_task.PythonTask - Parameters - Properties - Methods - flytekit.core.base\_task.Task - Parameters - Properties - Methods - flytekit.core.base\_task.TaskMetadata - Parameters - Properties - Methods - flytekit.core.base\_task.TaskResolverMixin - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.cache\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.cache/page.md) - Directory - Classes - Protocols - Variables - flytekit.core.cache.Cache - Parameters - Methods - flytekit.core.cache.CachePolicy - Methods - flytekit.core.cache.VersionParameters - Parameters - \[Flytekit SDK > Packages > flytekit.core.checkpointer\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.checkpointer/page.md) - Directory - Classes - flytekit.core.checkpointer.Checkpoint - Methods - flytekit.core.checkpointer.SyncCheckpoint - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.core.class\_based\_resolver\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.class\_based\_resolver/page.md) - Directory - Classes - flytekit.core.class\_based\_resolver.ClassStorageTaskResolver - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.condition\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.condition/page.md) - Directory - Classes - Methods - Methods - flytekit.core.condition.BranchNode - Parameters - Properties - flytekit.core.condition.Case - Parameters - Properties - Methods - flytekit.core.condition.Condition - Parameters - Methods - flytekit.core.condition.ConditionalSection - Parameters - Properties - Methods - flytekit.core.condition.LocalExecutedConditionalSection - Parameters - Properties - Methods - flytekit.core.condition.SkippedConditionalSection - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.constants\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.constants/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.core.container\_task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.container\_task/page.md) - Directory - Classes - Variables - flytekit.core.container\_task.ContainerTask - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.context\_manager\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.context\_manager/page.md) - Directory - Classes - Protocols - Variables - flytekit.core.context\_manager.BranchEvalMode - flytekit.core.context\_manager.CompilationState - Parameters - Methods - flytekit.core.context\_manager.ExecutionParameters - Parameters - Properties - Methods - flytekit.core.context\_manager.ExecutionState - Parameters - Methods - flytekit.core.context\_manager.FlyteContext - Parameters - Properties - Methods - flytekit.core.context\_manager.FlyteContextManager - Methods - flytekit.core.context\_manager.FlyteEntities - flytekit.core.context\_manager.OutputMetadata - Parameters - flytekit.core.context\_manager.OutputMetadataTracker - Parameters - Methods - flytekit.core.context\_manager.SecretsManager - Parameters - Methods - flytekit.core.context\_manager.SerializableToString - Methods - \[Flytekit SDK > Packages > flytekit.core.data\_persistence\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.data\_persistence/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.data\_persistence.FileAccessProvider - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.docstring\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.docstring/page.md) - Directory - Classes - flytekit.core.docstring.Docstring - Parameters - Properties - \[Flytekit SDK > Packages > flytekit.core.environment\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.environment/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.environment.Environment - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.core.gate\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.gate/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.gate.Gate - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.hash\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.hash/page.md) - Directory - Classes - Variables - flytekit.core.hash.HashMethod - Parameters - Methods - flytekit.core.hash.HashOnReferenceMixin - \[Flytekit SDK > Packages > flytekit.core.interface\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.interface/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.interface.Interface - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.launch\_plan\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.launch\_plan/page.md) - Directory - Classes - Methods - Methods - flytekit.core.launch\_plan.LaunchPlan - Parameters - Properties - Methods - flytekit.core.launch\_plan.ReferenceLaunchPlan - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.legacy\_map\_task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.legacy\_map\_task/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.legacy\_map\_task.MapPythonTask - Parameters - Properties - Methods - flytekit.core.legacy\_map\_task.MapTaskResolver - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.local\_cache\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.local\_cache/page.md) - Directory - Classes - Variables - flytekit.core.local\_cache.LocalTaskCache - Methods - \[Flytekit SDK > Packages > flytekit.core.local\_fsspec\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.local\_fsspec/page.md) - Directory - Classes - flytekit.core.local\_fsspec.FlyteLocalFileSystem - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.mock\_stats\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.mock\_stats/page.md) - Directory - Classes - flytekit.core.mock\_stats.MockStats - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.core.node\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.node/page.md) - Directory - Classes - Methods - Methods - flytekit.core.node.Node - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.node\_creation\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.node\_creation/page.md) - Directory - Methods - Variables - Methods - \[Flytekit SDK > Packages > flytekit.core.notification\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.notification/page.md) - Directory - Classes - flytekit.core.notification.Email - Parameters - Properties - Methods - flytekit.core.notification.Notification - Parameters - Properties - Methods - flytekit.core.notification.PagerDuty - Parameters - Properties - Methods - flytekit.core.notification.Slack - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.options\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.options/page.md) - Directory - Classes - flytekit.core.options.Options - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.core.pod\_template\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.pod\_template/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.pod\_template.PodTemplate - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.core.promise\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.promise/page.md) - Directory - Classes - Protocols - Methods - Variables - Methods - flytekit.core.promise.ComparisonExpression - Parameters - Properties - Methods - flytekit.core.promise.ComparisonOps - flytekit.core.promise.ConjunctionExpression - Parameters - Properties - Methods - flytekit.core.promise.ConjunctionOps - flytekit.core.promise.HasFlyteInterface - Properties - Methods - flytekit.core.promise.LocallyExecutable - Methods - flytekit.core.promise.NodeOutput - Parameters - Properties - Methods - flytekit.core.promise.Promise - Parameters - Properties - Methods - flytekit.core.promise.SupportsNodeCreation - Properties - Methods - flytekit.core.promise.VoidPromise - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.python\_auto\_container\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.python\_auto\_container/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.python\_auto\_container.DefaultNotebookTaskResolver - Parameters - Properties - Methods - flytekit.core.python\_auto\_container.DefaultTaskResolver - Parameters - Properties - Methods - flytekit.core.python\_auto\_container.PickledEntity - Parameters - flytekit.core.python\_auto\_container.PickledEntityMetadata - Parameters - flytekit.core.python\_auto\_container.PythonAutoContainerTask - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.python\_customized\_container\_task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.python\_customized\_container\_task/page.md) - Directory - Classes - Variables - flytekit.core.python\_customized\_container\_task.PythonCustomizedContainerTask - Parameters - Properties - Methods - flytekit.core.python\_customized\_container\_task.TaskTemplateResolver - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.python\_function\_task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.python\_function\_task/page.md) - Directory - Classes - Variables - flytekit.core.python\_function\_task.AsyncPythonFunctionTask - Parameters - Properties - Methods - flytekit.core.python\_function\_task.EagerAsyncPythonFunctionTask - Parameters - Properties - Methods - flytekit.core.python\_function\_task.EagerFailureHandlerTask - Parameters - Properties - Methods - flytekit.core.python\_function\_task.EagerFailureTaskResolver - Properties - Methods - flytekit.core.python\_function\_task.PythonFunctionTask - Parameters - Properties - Methods - flytekit.core.python\_function\_task.PythonInstanceTask - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.reference\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.reference/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.core.reference\_entity\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.reference\_entity/page.md) - Directory - Classes - flytekit.core.reference\_entity.LaunchPlanReference - Parameters - Properties - flytekit.core.reference\_entity.Reference - Parameters - Properties - flytekit.core.reference\_entity.ReferenceEntity - Parameters - Properties - Methods - flytekit.core.reference\_entity.ReferenceSpec - Parameters - Properties - flytekit.core.reference\_entity.ReferenceTemplate - Parameters - Properties - flytekit.core.reference\_entity.TaskReference - Parameters - Properties - flytekit.core.reference\_entity.WorkflowReference - Parameters - Properties - \[Flytekit SDK > Packages > flytekit.core.resources\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.resources/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.resources.ResourceSpec - Parameters - Methods - flytekit.core.resources.Resources - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.core.schedule\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.schedule/page.md) - Directory - Classes - Protocols - flytekit.core.schedule.CronSchedule - Parameters - Properties - Methods - flytekit.core.schedule.FixedRate - Parameters - Properties - Methods - flytekit.core.schedule.LaunchPlanTriggerBase - Methods - flytekit.core.schedule.OnSchedule - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.core.shim\_task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.shim\_task/page.md) - Directory - Classes - Variables - flytekit.core.shim\_task.ExecutableTemplateShimTask - Parameters - Properties - Methods - flytekit.core.shim\_task.ShimTaskExecutor - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.task/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.task.Echo - Parameters - Properties - Methods - flytekit.core.task.ReferenceTask - Parameters - Properties - Methods - flytekit.core.task.TaskPlugins - Methods - \[Flytekit SDK > Packages > flytekit.core.testing\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.testing/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.core.tracked\_abc\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.tracked\_abc/page.md) - Directory - Classes - flytekit.core.tracked\_abc.FlyteTrackedABC - Methods - \[Flytekit SDK > Packages > flytekit.core.tracker\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.tracker/page.md) - Directory - Classes - Methods - Methods - flytekit.core.tracker.InstanceTrackingMeta - flytekit.core.tracker.TrackedInstance - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.type\_engine\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.type\_engine/page.md) - Directory - Classes - Errors - Methods - Variables - Methods - flytekit.core.type\_engine.AsyncTypeTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.BatchSize - Parameters - Properties - flytekit.core.type\_engine.BinaryIOTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.DataclassTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.DictTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.EnumTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.ListTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.LiteralTypeTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.LiteralsResolver - Parameters - Properties - Methods - flytekit.core.type\_engine.ProtobufTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.RestrictedTypeError - flytekit.core.type\_engine.RestrictedTypeTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.SimpleTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.TextIOTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.TypeEngine - Methods - flytekit.core.type\_engine.TypeTransformer - Parameters - Properties - Methods - flytekit.core.type\_engine.TypeTransformerFailedError - flytekit.core.type\_engine.UnionTransformer - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.core.type\_helpers\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.type\_helpers/page.md) - Directory - Methods - Variables - Methods - \[Flytekit SDK > Packages > flytekit.core.type\_match\_checking\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.type\_match\_checking/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.core.utils\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.utils/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.utils.AutoDeletingTempDir - Parameters - Properties - Methods - flytekit.core.utils.ClassDecorator - Parameters - Methods - flytekit.core.utils.Directory - Parameters - Properties - Methods - flytekit.core.utils.timeit - Parameters - \[Flytekit SDK > Packages > flytekit.core.worker\_queue\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.worker\_queue/page.md) - Directory - Classes - Variables - flytekit.core.worker\_queue.Controller - Parameters - Methods - flytekit.core.worker\_queue.ItemStatus - flytekit.core.worker\_queue.Update - Parameters - flytekit.core.worker\_queue.WorkItem - Parameters - Properties - \[Flytekit SDK > Packages > flytekit.core.workflow\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.core.workflow/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.core.workflow.ImperativeWorkflow - Parameters - Properties - Methods - flytekit.core.workflow.PythonFunctionWorkflow - Parameters - Properties - Methods - flytekit.core.workflow.ReferenceWorkflow - Parameters - Properties - Methods - flytekit.core.workflow.WorkflowBase - Parameters - Properties - Methods - flytekit.core.workflow.WorkflowFailurePolicy - flytekit.core.workflow.WorkflowMetadata - Parameters - Methods - flytekit.core.workflow.WorkflowMetadataDefaults - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.deck.deck\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.deck.deck/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.deck.deck.Deck - Parameters - Properties - Methods - flytekit.deck.deck.DeckField - Parameters - flytekit.deck.deck.TimeLineDeck - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.deck.renderer\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.deck.renderer/page.md) - Directory - Classes - Protocols - Variables - flytekit.deck.renderer.ArrowRenderer - Methods - flytekit.deck.renderer.MarkdownRenderer - Methods - flytekit.deck.renderer.PythonDependencyRenderer - Parameters - Methods - flytekit.deck.renderer.Renderable - Methods - flytekit.deck.renderer.SourceCodeRenderer - Parameters - Methods - flytekit.deck.renderer.TopFrameRenderer - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.exceptions.base\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.exceptions.base/page.md) - Directory - Errors - flytekit.exceptions.base.FlyteException - Parameters - Properties - flytekit.exceptions.base.FlyteRecoverableException - Parameters - Properties - \[Flytekit SDK > Packages > flytekit.exceptions.eager\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.exceptions.eager/page.md) - Directory - Errors - flytekit.exceptions.eager.EagerException - \[Flytekit SDK > Packages > flytekit.exceptions.scopes\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.exceptions.scopes/page.md) - Directory - Errors - Methods - Methods - flytekit.exceptions.scopes.FlyteScopedException - Parameters - Properties - flytekit.exceptions.scopes.FlyteScopedSystemException - Parameters - Properties - flytekit.exceptions.scopes.FlyteScopedUserException - Parameters - Properties - \[Flytekit SDK > Packages > flytekit.exceptions.system\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.exceptions.system/page.md) - Directory - Errors - flytekit.exceptions.system.FlyteAgentNotFound - Parameters - Properties - flytekit.exceptions.system.FlyteConnectorNotFound - Parameters - Properties - flytekit.exceptions.system.FlyteDownloadDataException - Parameters - Properties - flytekit.exceptions.system.FlyteEntrypointNotLoadable - Parameters - Properties - flytekit.exceptions.system.FlyteNonRecoverableSystemException - Parameters - Properties - flytekit.exceptions.system.FlyteNotImplementedException - Parameters - Properties - flytekit.exceptions.system.FlyteSystemAssertion - Parameters - Properties - flytekit.exceptions.system.FlyteSystemException - Parameters - Properties - flytekit.exceptions.system.FlyteSystemUnavailableException - Parameters - Properties - flytekit.exceptions.system.FlyteUploadDataException - Parameters - Properties - \[Flytekit SDK > Packages > flytekit.exceptions.user\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.exceptions.user/page.md) - Directory - Errors - flytekit.exceptions.user.FlyteAssertion - Parameters - Properties - flytekit.exceptions.user.FlyteAuthenticationException - Parameters - Properties - flytekit.exceptions.user.FlyteCompilationException - Parameters - Properties - flytekit.exceptions.user.FlyteDataNotFoundException - Parameters - Properties - flytekit.exceptions.user.FlyteDisapprovalException - Parameters - Properties - flytekit.exceptions.user.FlyteEntityAlreadyExistsException - Parameters - Properties - flytekit.exceptions.user.FlyteEntityNotExistException - Parameters - Properties - flytekit.exceptions.user.FlyteEntityNotFoundException - Parameters - Properties - flytekit.exceptions.user.FlyteFailureNodeInputMismatchException - Parameters - Properties - flytekit.exceptions.user.FlyteInvalidInputException - Parameters - Properties - flytekit.exceptions.user.FlyteMissingReturnValueException - Parameters - Properties - flytekit.exceptions.user.FlyteMissingTypeException - Parameters - Properties - flytekit.exceptions.user.FlytePromiseAttributeResolveException - Parameters - Properties - flytekit.exceptions.user.FlyteRecoverableException - Parameters - Properties - flytekit.exceptions.user.FlyteTimeout - Parameters - Properties - flytekit.exceptions.user.FlyteTypeException - Parameters - Properties - flytekit.exceptions.user.FlyteUserException - Parameters - Properties - flytekit.exceptions.user.FlyteUserRuntimeException - Parameters - Properties - flytekit.exceptions.user.FlyteValidationException - Parameters - Properties - flytekit.exceptions.user.FlyteValueException - Parameters - Properties - \[Flytekit SDK > Packages > flytekit.exceptions.utils\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.exceptions.utils/page.md) - Directory - Methods - Variables - Methods - \[Flytekit SDK > Packages > flytekit.experimental.eager\_function\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.experimental.eager\_function/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.extend.backend.base\_connector\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extend.backend.base\_connector/page.md) - Directory - Classes - flytekit.extend.backend.base\_connector.AsyncConnectorBase - Parameters - Properties - Methods - flytekit.extend.backend.base\_connector.AsyncConnectorExecutorMixin - Methods - flytekit.extend.backend.base\_connector.ConnectorBase - Parameters - Properties - flytekit.extend.backend.base\_connector.ConnectorRegistry - Methods - flytekit.extend.backend.base\_connector.Resource - Parameters - Methods - flytekit.extend.backend.base\_connector.ResourceMeta - Parameters - Methods - flytekit.extend.backend.base\_connector.SyncConnectorBase - Parameters - Properties - Methods - flytekit.extend.backend.base\_connector.SyncConnectorExecutorMixin - Methods - flytekit.extend.backend.base\_connector.TaskCategory - Parameters - Properties - \[Flytekit SDK > Packages > flytekit.extend.backend.connector\_service\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extend.backend.connector\_service/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.extend.backend.connector\_service.AsyncConnectorService - Methods - flytekit.extend.backend.connector\_service.ConnectorMetadataService - Methods - flytekit.extend.backend.connector\_service.SyncConnectorService - Methods - \[Flytekit SDK > Packages > flytekit.extend.backend.utils\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extend.backend.utils/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.extras.accelerators\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extras.accelerators/page.md) - Specifying Accelerators - Base Classes - Predefined Accelerator Constants - Directory - Classes - Variables - flytekit.extras.accelerators.BaseAccelerator - Methods - flytekit.extras.accelerators.GPUAccelerator - Parameters - Methods - flytekit.extras.accelerators.MultiInstanceGPUAccelerator - Properties - Methods - \[Flytekit SDK > Packages > flytekit.extras.cloud\_pickle\_resolver\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extras.cloud\_pickle\_resolver/page.md) - Directory - Classes - Variables - flytekit.extras.cloud\_pickle\_resolver.ExperimentalNaiveCloudPickleResolver - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.extras.pydantic\_transformer.transformer\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extras.pydantic\_transformer.transformer/page.md) - Directory - Classes - Variables - flytekit.extras.pydantic\_transformer.transformer.PydanticTransformer - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.extras.sklearn.native\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extras.sklearn.native/page.md) - Directory - Classes - Variables - flytekit.extras.sklearn.native.SklearnEstimatorTransformer - Parameters - Properties - Methods - flytekit.extras.sklearn.native.SklearnTypeTransformer - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.extras.sqlite3.task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extras.sqlite3.task/page.md) - Directory - Classes - Methods - Methods - flytekit.extras.sqlite3.task.SQLite3Config - Parameters - flytekit.extras.sqlite3.task.SQLite3Task - Parameters - Properties - Methods - flytekit.extras.sqlite3.task.SQLite3TaskExecutor - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.extras.tasks.shell\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extras.tasks.shell/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.extras.tasks.shell.AttrDict - Parameters - flytekit.extras.tasks.shell.OutputLocation - Parameters - flytekit.extras.tasks.shell.ProcessResult - Parameters - flytekit.extras.tasks.shell.RawShellTask - Parameters - Properties - Methods - flytekit.extras.tasks.shell.ShellTask - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.extras.webhook\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extras.webhook/page.md) - Directory - Classes - flytekit.extras.webhook.WebhookConnector - Parameters - Properties - Methods - flytekit.extras.webhook.WebhookTask - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.extras.webhook.connector\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extras.webhook.connector/page.md) - Directory - Classes - Variables - flytekit.extras.webhook.connector.WebhookConnector - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.extras.webhook.constants\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extras.webhook.constants/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.extras.webhook.task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.extras.webhook.task/page.md) - Directory - Classes - Variables - flytekit.extras.webhook.task.WebhookTask - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.image\_spec.default\_builder\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.image\_spec.default\_builder/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.image\_spec.default\_builder.DefaultImageBuilder - Methods - \[Flytekit SDK > Packages > flytekit.image\_spec.image\_spec\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.image\_spec.image\_spec/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.image\_spec.image\_spec.ImageBuildEngine - Methods - flytekit.image\_spec.image\_spec.ImageSpec - Parameters - Properties - Methods - flytekit.image\_spec.image\_spec.ImageSpecBuilder - Methods - \[Flytekit SDK > Packages > flytekit.image\_spec.noop\_builder\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.image\_spec.noop\_builder/page.md) - Directory - Classes - flytekit.image\_spec.noop\_builder.NoOpBuilder - Methods - \[Flytekit SDK > Packages > flytekit.interaction.click\_types\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interaction.click\_types/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.interaction.click\_types.DateTimeType - Parameters - Methods - flytekit.interaction.click\_types.DirParamType - Methods - flytekit.interaction.click\_types.DurationParamType - Methods - flytekit.interaction.click\_types.EnumParamType - Parameters - Methods - flytekit.interaction.click\_types.FileParamType - Methods - flytekit.interaction.click\_types.FlyteLiteralConverter - Parameters - Properties - Methods - flytekit.interaction.click\_types.JSONIteratorParamType - Methods - flytekit.interaction.click\_types.JsonParamType - Parameters - Methods - flytekit.interaction.click\_types.PickleParamType - Methods - flytekit.interaction.click\_types.StructuredDatasetParamType - Methods - flytekit.interaction.click\_types.UnionParamType - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.interaction.parse\_stdin\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interaction.parse\_stdin/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.interaction.rich\_utils\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interaction.rich\_utils/page.md) - Directory - Classes - flytekit.interaction.rich\_utils.RichCallback - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.interaction.string\_literals\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interaction.string\_literals/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.interactive\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interactive/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.interactive.constants\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interactive.constants/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.interactive.utils\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interactive.utils/page.md) - Directory - Methods - Variables - Methods - \[Flytekit SDK > Packages > flytekit.interactive.vscode\_lib.config\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interactive.vscode\_lib.config/page.md) - Directory - Classes - Variables - flytekit.interactive.vscode\_lib.config.VscodeConfig - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.interactive.vscode\_lib.decorator\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interactive.vscode\_lib.decorator/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.interactive.vscode\_lib.decorator.vscode - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.interactive.vscode\_lib.vscode\_constants\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interactive.vscode\_lib.vscode\_constants/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.interfaces.cli\_identifiers\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interfaces.cli\_identifiers/page.md) - Directory - Classes - flytekit.interfaces.cli\_identifiers.Identifier - Parameters - Properties - Methods - flytekit.interfaces.cli\_identifiers.TaskExecutionIdentifier - Parameters - Properties - Methods - flytekit.interfaces.cli\_identifiers.WorkflowExecutionIdentifier - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.interfaces.random\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interfaces.random/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.interfaces.stats.client\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interfaces.stats.client/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.interfaces.stats.client.DummyStatsClient - Parameters - Methods - flytekit.interfaces.stats.client.ScopeableStatsProxy - Parameters - Methods - flytekit.interfaces.stats.client.StatsClientProxy - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.interfaces.stats.taggable\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.interfaces.stats.taggable/page.md) - Directory - Classes - Methods - Methods - flytekit.interfaces.stats.taggable.TaggableStats - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.lazy\_import.lazy\_module\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.lazy\_import.lazy\_module/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.loggers\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.loggers/page.md) - Directory - Methods - Variables - Methods - \[Flytekit SDK > Packages > flytekit.models.admin.common\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.admin.common/page.md) - Directory - Classes - flytekit.models.admin.common.Sort - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.admin.task\_execution\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.admin.task\_execution/page.md) - Directory - Classes - flytekit.models.admin.task\_execution.TaskExecution - Parameters - Properties - Methods - flytekit.models.admin.task\_execution.TaskExecutionClosure - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.admin.workflow\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.admin.workflow/page.md) - Directory - Classes - flytekit.models.admin.workflow.Workflow - Parameters - Properties - Methods - flytekit.models.admin.workflow.WorkflowClosure - Parameters - Properties - Methods - flytekit.models.admin.workflow.WorkflowSpec - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.annotation\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.annotation/page.md) - Directory - Classes - flytekit.models.annotation.TypeAnnotation - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.array\_job\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.array\_job/page.md) - Directory - Classes - flytekit.models.array\_job.ArrayJob - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.common\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.common/page.md) - Directory - Classes - flytekit.models.common.Annotations - Parameters - Properties - Methods - flytekit.models.common.AuthRole - Parameters - Properties - Methods - flytekit.models.common.EmailNotification - Parameters - Properties - Methods - flytekit.models.common.Envs - Parameters - Properties - Methods - flytekit.models.common.FlyteABCMeta - Methods - flytekit.models.common.FlyteCustomIdlEntity - Properties - Methods - flytekit.models.common.FlyteIdlEntity - Properties - Methods - flytekit.models.common.FlyteType - Methods - flytekit.models.common.Labels - Parameters - Properties - Methods - flytekit.models.common.NamedEntityIdentifier - Parameters - Properties - Methods - flytekit.models.common.Notification - Parameters - Properties - Methods - flytekit.models.common.PagerDutyNotification - Parameters - Properties - Methods - flytekit.models.common.RawOutputDataConfig - Parameters - Properties - Methods - flytekit.models.common.SlackNotification - Parameters - Properties - Methods - flytekit.models.common.UrlBlob - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.concurrency\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.concurrency/page.md) - Directory - Classes - flytekit.models.concurrency.ConcurrencyLimitBehavior - Methods - flytekit.models.concurrency.ConcurrencyPolicy - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.core.catalog\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.core.catalog/page.md) - Directory - Classes - flytekit.models.core.catalog.CatalogArtifactTag - Parameters - Properties - Methods - flytekit.models.core.catalog.CatalogMetadata - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.core.compiler\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.core.compiler/page.md) - Directory - Classes - flytekit.models.core.compiler.CompiledTask - Parameters - Properties - Methods - flytekit.models.core.compiler.CompiledWorkflow - Parameters - Properties - Methods - flytekit.models.core.compiler.CompiledWorkflowClosure - Parameters - Properties - Methods - flytekit.models.core.compiler.ConnectionSet - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.core.condition\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.core.condition/page.md) - Directory - Classes - flytekit.models.core.condition.BooleanExpression - Parameters - Properties - Methods - flytekit.models.core.condition.ComparisonExpression - Parameters - Properties - Methods - flytekit.models.core.condition.ConjunctionExpression - Parameters - Properties - Methods - flytekit.models.core.condition.Operand - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.core.errors\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.core.errors/page.md) - Directory - Classes - flytekit.models.core.errors.ContainerError - Parameters - Properties - Methods - flytekit.models.core.errors.ErrorDocument - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.core.execution\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.core.execution/page.md) - Directory - Classes - flytekit.models.core.execution.ExecutionError - Parameters - Properties - Methods - flytekit.models.core.execution.NodeExecutionPhase - Methods - flytekit.models.core.execution.TaskExecutionPhase - Methods - flytekit.models.core.execution.TaskLog - Parameters - Properties - Methods - flytekit.models.core.execution.WorkflowExecutionPhase - Methods - \[Flytekit SDK > Packages > flytekit.models.core.identifier\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.core.identifier/page.md) - Directory - Classes - flytekit.models.core.identifier.Identifier - Parameters - Properties - Methods - flytekit.models.core.identifier.NodeExecutionIdentifier - Parameters - Properties - Methods - flytekit.models.core.identifier.ResourceType - flytekit.models.core.identifier.SignalIdentifier - Parameters - Properties - Methods - flytekit.models.core.identifier.TaskExecutionIdentifier - Parameters - Properties - Methods - flytekit.models.core.identifier.WorkflowExecutionIdentifier - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.core.types\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.core.types/page.md) - Directory - Classes - flytekit.models.core.types.BlobType - Parameters - Properties - Methods - flytekit.models.core.types.EnumType - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.core.workflow\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.core.workflow/page.md) - Directory - Classes - flytekit.models.core.workflow.Alias - Parameters - Properties - Methods - flytekit.models.core.workflow.ApproveCondition - Parameters - Properties - Methods - flytekit.models.core.workflow.ArrayNode - Parameters - Properties - Methods - flytekit.models.core.workflow.BranchNode - Parameters - Properties - Methods - flytekit.models.core.workflow.GateNode - Parameters - Properties - Methods - flytekit.models.core.workflow.IfBlock - Parameters - Properties - Methods - flytekit.models.core.workflow.IfElseBlock - Parameters - Properties - Methods - flytekit.models.core.workflow.Node - Parameters - Properties - Methods - flytekit.models.core.workflow.NodeMetadata - Parameters - Properties - Methods - flytekit.models.core.workflow.SignalCondition - Parameters - Properties - Methods - flytekit.models.core.workflow.SleepCondition - Parameters - Properties - Methods - flytekit.models.core.workflow.TaskNode - Parameters - Properties - Methods - flytekit.models.core.workflow.TaskNodeOverrides - Parameters - Properties - Methods - flytekit.models.core.workflow.WorkflowMetadata - Parameters - Properties - Methods - flytekit.models.core.workflow.WorkflowMetadataDefaults - Parameters - Properties - Methods - flytekit.models.core.workflow.WorkflowNode - Parameters - Properties - Methods - flytekit.models.core.workflow.WorkflowTemplate - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.documentation\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.documentation/page.md) - Directory - Classes - flytekit.models.documentation.Description - Parameters - Properties - Methods - flytekit.models.documentation.Documentation - Parameters - Properties - Methods - flytekit.models.documentation.SourceCode - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.domain\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.domain/page.md) - Directory - Classes - flytekit.models.domain.Domain - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.dynamic\_job\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.dynamic\_job/page.md) - Directory - Classes - flytekit.models.dynamic\_job.DynamicJobSpec - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.event\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.event/page.md) - Directory - Classes - flytekit.models.event.TaskExecutionMetadata - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.execution\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.execution/page.md) - Directory - Classes - flytekit.models.execution.AbortMetadata - Parameters - Properties - Methods - flytekit.models.execution.ClusterAssignment - Parameters - Properties - Methods - flytekit.models.execution.Execution - Parameters - Properties - Methods - flytekit.models.execution.ExecutionClosure - Parameters - Properties - Methods - flytekit.models.execution.ExecutionMetadata - Parameters - Properties - Methods - flytekit.models.execution.ExecutionSpec - Parameters - Properties - Methods - flytekit.models.execution.LiteralMapBlob - Parameters - Properties - Methods - flytekit.models.execution.NodeExecutionGetDataResponse - Parameters - Properties - Methods - flytekit.models.execution.NotificationList - Parameters - Properties - Methods - flytekit.models.execution.SystemMetadata - Parameters - Properties - Methods - flytekit.models.execution.TaskExecutionGetDataResponse - Parameters - Properties - Methods - flytekit.models.execution.WorkflowExecutionGetDataResponse - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.filters\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.filters/page.md) - Directory - Classes - flytekit.models.filters.Contains - Parameters - Properties - Methods - flytekit.models.filters.Equal - Parameters - Properties - Methods - flytekit.models.filters.Filter - Parameters - Properties - Methods - flytekit.models.filters.FilterList - Parameters - Properties - Methods - flytekit.models.filters.GreaterThan - Parameters - Properties - Methods - flytekit.models.filters.GreaterThanOrEqual - Parameters - Properties - Methods - flytekit.models.filters.LessThan - Parameters - Properties - Methods - flytekit.models.filters.LessThanOrEqual - Parameters - Properties - Methods - flytekit.models.filters.NotEqual - Parameters - Properties - Methods - flytekit.models.filters.SetFilter - Parameters - Properties - Methods - flytekit.models.filters.ValueIn - Parameters - Properties - Methods - flytekit.models.filters.ValueNotIn - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.interface\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.interface/page.md) - Directory - Classes - flytekit.models.interface.Parameter - Parameters - Properties - Methods - flytekit.models.interface.ParameterMap - Parameters - Properties - Methods - flytekit.models.interface.TypedInterface - Parameters - Properties - Methods - flytekit.models.interface.Variable - Parameters - Properties - Methods - flytekit.models.interface.VariableMap - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.launch\_plan\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.launch\_plan/page.md) - Directory - Classes - flytekit.models.launch\_plan.Auth - Parameters - Properties - Methods - flytekit.models.launch\_plan.LaunchPlan - Parameters - Properties - Methods - flytekit.models.launch\_plan.LaunchPlanClosure - Parameters - Properties - Methods - flytekit.models.launch\_plan.LaunchPlanMetadata - Parameters - Properties - Methods - flytekit.models.launch\_plan.LaunchPlanSpec - Parameters - Properties - Methods - flytekit.models.launch\_plan.LaunchPlanState - Methods - \[Flytekit SDK > Packages > flytekit.models.literals\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.literals/page.md) - Directory - Classes - flytekit.models.literals.Binary - Parameters - Properties - Methods - flytekit.models.literals.Binding - Parameters - Properties - Methods - flytekit.models.literals.BindingData - Parameters - Properties - Methods - flytekit.models.literals.BindingDataCollection - Parameters - Properties - Methods - flytekit.models.literals.BindingDataMap - Parameters - Properties - Methods - flytekit.models.literals.Blob - Parameters - Properties - Methods - flytekit.models.literals.BlobMetadata - Parameters - Properties - Methods - flytekit.models.literals.Literal - Parameters - Properties - Methods - flytekit.models.literals.LiteralCollection - Parameters - Properties - Methods - flytekit.models.literals.LiteralMap - Parameters - Properties - Methods - flytekit.models.literals.LiteralOffloadedMetadata - Parameters - Properties - Methods - flytekit.models.literals.Primitive - Parameters - Properties - Methods - flytekit.models.literals.RetryStrategy - Parameters - Properties - Methods - flytekit.models.literals.Scalar - Parameters - Properties - Methods - flytekit.models.literals.Schema - Parameters - Properties - Methods - flytekit.models.literals.StructuredDataset - Parameters - Properties - Methods - flytekit.models.literals.StructuredDatasetMetadata - Parameters - Properties - Methods - flytekit.models.literals.Union - Parameters - Properties - Methods - flytekit.models.literals.Void - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.matchable\_resource\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.matchable\_resource/page.md) - Directory - Classes - flytekit.models.matchable\_resource.ClusterResourceAttributes - Parameters - Properties - Methods - flytekit.models.matchable\_resource.ExecutionClusterLabel - Parameters - Properties - Methods - flytekit.models.matchable\_resource.ExecutionQueueAttributes - Parameters - Properties - Methods - flytekit.models.matchable\_resource.MatchableResource - Methods - flytekit.models.matchable\_resource.MatchingAttributes - Parameters - Properties - Methods - flytekit.models.matchable\_resource.PluginOverride - Parameters - Properties - Methods - flytekit.models.matchable\_resource.PluginOverrides - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.named\_entity\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.named\_entity/page.md) - Directory - Classes - flytekit.models.named\_entity.NamedEntityIdentifier - Parameters - Properties - Methods - flytekit.models.named\_entity.NamedEntityMetadata - Parameters - Properties - Methods - flytekit.models.named\_entity.NamedEntityState - Methods - \[Flytekit SDK > Packages > flytekit.models.node\_execution\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.node\_execution/page.md) - Directory - Classes - flytekit.models.node\_execution.DynamicWorkflowNodeMetadata - Parameters - Properties - Methods - flytekit.models.node\_execution.NodeExecution - Parameters - Properties - Methods - flytekit.models.node\_execution.NodeExecutionClosure - Parameters - Properties - Methods - flytekit.models.node\_execution.TaskNodeMetadata - Parameters - Properties - Methods - flytekit.models.node\_execution.WorkflowNodeMetadata - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.presto\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.presto/page.md) - Directory - Classes - flytekit.models.presto.PrestoQuery - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.project\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.project/page.md) - Directory - Classes - flytekit.models.project.Project - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.qubole\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.qubole/page.md) - Directory - Classes - flytekit.models.qubole.HiveQuery - Parameters - Properties - Methods - flytekit.models.qubole.HiveQueryCollection - Parameters - Properties - Methods - flytekit.models.qubole.QuboleHiveJob - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.schedule\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.schedule/page.md) - Directory - Classes - flytekit.models.schedule.Schedule - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.security\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.security/page.md) - Directory - Classes - flytekit.models.security.Identity - Parameters - Properties - Methods - flytekit.models.security.OAuth2Client - Parameters - Properties - Methods - flytekit.models.security.OAuth2TokenRequest - Parameters - Properties - Methods - flytekit.models.security.Secret - Parameters - Properties - Methods - flytekit.models.security.SecurityContext - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.task\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.task/page.md) - Directory - Classes - flytekit.models.task.CompiledTask - Parameters - Properties - Methods - flytekit.models.task.Container - Parameters - Properties - Methods - flytekit.models.task.DataLoadingConfig - Parameters - Properties - Methods - flytekit.models.task.IOStrategy - Parameters - Properties - Methods - flytekit.models.task.K8sObjectMetadata - Parameters - Properties - Methods - flytekit.models.task.K8sPod - Parameters - Properties - Methods - flytekit.models.task.Resources - Parameters - Properties - Methods - flytekit.models.task.RuntimeMetadata - Parameters - Properties - Methods - flytekit.models.task.Sql - Parameters - Properties - Methods - flytekit.models.task.Task - Parameters - Properties - Methods - flytekit.models.task.TaskClosure - Parameters - Properties - Methods - flytekit.models.task.TaskExecutionMetadata - Parameters - Properties - Methods - flytekit.models.task.TaskMetadata - Parameters - Properties - Methods - flytekit.models.task.TaskSpec - Parameters - Properties - Methods - flytekit.models.task.TaskTemplate - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.types\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.types/page.md) - Directory - Classes - flytekit.models.types.Error - Parameters - Properties - Methods - flytekit.models.types.LiteralType - Parameters - Properties - Methods - flytekit.models.types.OutputReference - Parameters - Properties - Methods - flytekit.models.types.SchemaType - Parameters - Properties - Methods - flytekit.models.types.SimpleType - flytekit.models.types.StructuredDatasetType - Parameters - Properties - Methods - flytekit.models.types.TypeStructure - Parameters - Properties - Methods - flytekit.models.types.UnionType - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.models.workflow\_closure\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.models.workflow\_closure/page.md) - Directory - Classes - flytekit.models.workflow\_closure.WorkflowClosure - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.remote.backfill\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.remote.backfill/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.remote.data\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.remote.data/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.remote.entities\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.remote.entities/page.md) - Directory - Classes - flytekit.remote.entities.FlyteArrayNode - Parameters - Properties - Methods - flytekit.remote.entities.FlyteBranchNode - Parameters - Properties - Methods - flytekit.remote.entities.FlyteGateNode - Parameters - Properties - Methods - flytekit.remote.entities.FlyteLaunchPlan - Parameters - Properties - Methods - flytekit.remote.entities.FlyteNode - Parameters - Properties - Methods - flytekit.remote.entities.FlyteTask - Parameters - Properties - Methods - flytekit.remote.entities.FlyteTaskNode - Parameters - Properties - Methods - flytekit.remote.entities.FlyteWorkflow - Parameters - Properties - Methods - flytekit.remote.entities.FlyteWorkflowNode - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.remote.executions\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.remote.executions/page.md) - Directory - Classes - flytekit.remote.executions.FlyteNodeExecution - Parameters - Properties - Methods - flytekit.remote.executions.FlyteTaskExecution - Parameters - Properties - Methods - flytekit.remote.executions.FlyteWorkflowExecution - Parameters - Properties - Methods - flytekit.remote.executions.RemoteExecutionBase - Parameters - Properties - \[Flytekit SDK > Packages > flytekit.remote.interface\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.remote.interface/page.md) - Directory - Classes - flytekit.remote.interface.TypedInterface - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.remote.lazy\_entity\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.remote.lazy\_entity/page.md) - Directory - Classes - Variables - flytekit.remote.lazy\_entity.LazyEntity - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.remote.metrics\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.remote.metrics/page.md) - Directory - Classes - Methods - Methods - flytekit.remote.metrics.FlyteExecutionSpan - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.remote.remote\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.remote.remote/page.md) - Directory - Classes - Errors - Variables - flytekit.remote.remote.FlyteRemote - Parameters - Properties - Methods - flytekit.remote.remote.RegistrationSkipped - flytekit.remote.remote.ResolvedIdentifiers - Parameters - \[Flytekit SDK > Packages > flytekit.remote.remote\_callable\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.remote.remote\_callable/page.md) - Directory - Classes - flytekit.remote.remote\_callable.RemoteEntity - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.remote.remote\_fs\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.remote.remote\_fs/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.remote.remote\_fs.FlyteFS - Parameters - Properties - Methods - flytekit.remote.remote\_fs.FlytePathResolver - Methods - flytekit.remote.remote\_fs.HttpFileWriter - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.sensor.base\_sensor\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.sensor.base\_sensor/page.md) - Directory - Classes - Protocols - Variables - flytekit.sensor.base\_sensor.BaseSensor - Parameters - Properties - Methods - flytekit.sensor.base\_sensor.SensorConfig - Methods - flytekit.sensor.base\_sensor.SensorMetadata - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.sensor.file\_sensor\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.sensor.file\_sensor/page.md) - Directory - Classes - flytekit.sensor.file\_sensor.FileSensor - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.sensor.sensor\_engine\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.sensor.sensor\_engine/page.md) - Directory - Classes - flytekit.sensor.sensor\_engine.SensorEngine - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.tools.fast\_registration\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.tools.fast\_registration/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.tools.fast\_registration.FastPackageOptions - Parameters - \[Flytekit SDK > Packages > flytekit.tools.ignore\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.tools.ignore/page.md) - Directory - Classes - Variables - flytekit.tools.ignore.DockerIgnore - Parameters - Methods - flytekit.tools.ignore.FlyteIgnore - Parameters - Methods - flytekit.tools.ignore.GitIgnore - Parameters - Methods - flytekit.tools.ignore.Ignore - Parameters - Methods - flytekit.tools.ignore.IgnoreGroup - Parameters - Methods - flytekit.tools.ignore.StandardIgnore - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.tools.interactive\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.tools.interactive/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.tools.module\_loader\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.tools.module\_loader/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.tools.repo\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.tools.repo/page.md) - Directory - Errors - Methods - Methods - flytekit.tools.repo.NoSerializableEntitiesError - \[Flytekit SDK > Packages > flytekit.tools.script\_mode\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.tools.script\_mode/page.md) - Directory - Methods - Variables - Methods - \[Flytekit SDK > Packages > flytekit.tools.serialize\_helpers\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.tools.serialize\_helpers/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.tools.subprocess\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.tools.subprocess/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.tools.translator\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.tools.translator/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.types.directory\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.directory/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.types.directory.types\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.directory.types/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.types.directory.types.FlyteDirToMultipartBlobTransformer - Parameters - Properties - Methods - flytekit.types.directory.types.FlyteDirectory - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.error.error\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.error.error/page.md) - Directory - Classes - Variables - flytekit.types.error.error.ErrorTransformer - Parameters - Properties - Methods - flytekit.types.error.error.FlyteError - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.types.file\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.file/page.md) - Directory - Classes - flytekit.types.file.FileExt - Parameters - Methods - \[Flytekit SDK > Packages > flytekit.types.file.file\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.file.file/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.types.file.file.FlyteFile - Parameters - Properties - Methods - flytekit.types.file.file.FlyteFilePathTransformer - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.file.image\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.file.image/page.md) - Directory - Classes - Variables - flytekit.types.file.image.PILImageTransformer - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.iterator.iterator\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.iterator.iterator/page.md) - Directory - Classes - Variables - flytekit.types.iterator.iterator.FlyteIterator - Parameters - flytekit.types.iterator.iterator.IteratorTransformer - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.iterator.json\_iterator\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.iterator.json\_iterator/page.md) - Directory - Classes - flytekit.types.iterator.json\_iterator.JSONIterator - Parameters - flytekit.types.iterator.json\_iterator.JSONIteratorTransformer - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.numpy.ndarray\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.numpy.ndarray/page.md) - Directory - Classes - Methods - Methods - flytekit.types.numpy.ndarray.NumpyArrayTransformer - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.pickle.pickle\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.pickle.pickle/page.md) - Directory - Classes - Variables - flytekit.types.pickle.pickle.FlytePickle - Methods - flytekit.types.pickle.pickle.FlytePickleTransformer - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.schema.types\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.schema.types/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.types.schema.types.FlyteSchema - Parameters - Properties - Methods - flytekit.types.schema.types.FlyteSchemaTransformer - Parameters - Properties - Methods - flytekit.types.schema.types.LocalIOSchemaReader - Parameters - Properties - Methods - flytekit.types.schema.types.LocalIOSchemaWriter - Parameters - Properties - Methods - flytekit.types.schema.types.SchemaEngine - Methods - flytekit.types.schema.types.SchemaFormat - flytekit.types.schema.types.SchemaHandler - Parameters - flytekit.types.schema.types.SchemaOpenMode - flytekit.types.schema.types.SchemaReader - Parameters - Properties - Methods - flytekit.types.schema.types.SchemaWriter - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.schema.types\_pandas\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.schema.types\_pandas/page.md) - Directory - Classes - flytekit.types.schema.types\_pandas.PandasDataFrameTransformer - Parameters - Properties - Methods - flytekit.types.schema.types\_pandas.PandasSchemaReader - Parameters - Properties - Methods - flytekit.types.schema.types\_pandas.PandasSchemaWriter - Parameters - Properties - Methods - flytekit.types.schema.types\_pandas.ParquetIO - Methods - \[Flytekit SDK > Packages > flytekit.types.structured\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.structured/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.types.structured.basic\_dfs\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.structured.basic\_dfs/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.types.structured.basic\_dfs.ArrowToParquetEncodingHandler - Parameters - Properties - Methods - flytekit.types.structured.basic\_dfs.CSVToPandasDecodingHandler - Parameters - Properties - Methods - flytekit.types.structured.basic\_dfs.PandasToCSVEncodingHandler - Parameters - Properties - Methods - flytekit.types.structured.basic\_dfs.PandasToParquetEncodingHandler - Parameters - Properties - Methods - flytekit.types.structured.basic\_dfs.ParquetToArrowDecodingHandler - Parameters - Properties - Methods - flytekit.types.structured.basic\_dfs.ParquetToPandasDecodingHandler - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.structured.bigquery\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.structured.bigquery/page.md) - Directory - Classes - Variables - flytekit.types.structured.bigquery.ArrowToBQEncodingHandlers - Parameters - Properties - Methods - flytekit.types.structured.bigquery.BQToArrowDecodingHandler - Parameters - Properties - Methods - flytekit.types.structured.bigquery.BQToPandasDecodingHandler - Parameters - Properties - Methods - flytekit.types.structured.bigquery.PandasToBQEncodingHandlers - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.structured.snowflake\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.structured.snowflake/page.md) - Directory - Classes - Methods - Variables - Methods - flytekit.types.structured.snowflake.PandasToSnowflakeEncodingHandlers - Parameters - Properties - Methods - flytekit.types.structured.snowflake.SnowflakeToPandasDecodingHandler - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.types.structured.structured\_dataset\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.types.structured.structured\_dataset/page.md) - Directory - Classes - Errors - Methods - Variables - Methods - flytekit.types.structured.structured\_dataset.DuplicateHandlerError - flytekit.types.structured.structured\_dataset.StructuredDataset - Parameters - Properties - Methods - flytekit.types.structured.structured\_dataset.StructuredDatasetDecoder - Parameters - Properties - Methods - flytekit.types.structured.structured\_dataset.StructuredDatasetEncoder - Parameters - Properties - Methods - flytekit.types.structured.structured\_dataset.StructuredDatasetTransformerEngine - Parameters - Properties - Methods - \[Flytekit SDK > Packages > flytekit.utils.asyn\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.utils.asyn/page.md) - Directory - Variables - \[Flytekit SDK > Packages > flytekit.utils.dict\_formatter\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.utils.dict\_formatter/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.utils.pbhash\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.utils.pbhash/page.md) - Directory - Methods - Methods - \[Flytekit SDK > Packages > flytekit.utils.rate\_limiter\](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/packages/flytekit.utils.rate\_limiter/page.md) - Directory - Classes - flytekit.utils.rate\_limiter.RateLimiter - Parameters - Methods --- ## Community - \[Joining the community\](https://www.union.ai/docs/v1/flyte/community/joining-the-community/page.md) - Community sync - Contributor's sync - Newsletter - Slack guidelines - Abide by the \[LF's Code of Conduct\](https://lfprojects.org/policies/code-of-conduct/) - Avoid using DMs and @mentions - Make use of threads - Do not post the same question across multiple channels - Do not solicit members of our Slack - \[Contributing code\](https://www.union.ai/docs/v1/flyte/community/contributing-code/page.md) - Becoming a contributor - Before submitting your PR - 🐞 File an issue - Component Reference - \`flyte\` - \`flyteidl\` - \`flytepropeller\` - \`flyteadmin\` - \`flytekit\` - \`flyteconsole\` - \`datacatalog\` - \`flyteplugins\` - \`flytestdlib\` - \`flytectl\` - Development Environment Setup Guide - Requirements - Content - How to setup dev environment for flyteidl, flyteadmin, flyteplugins, flytepropeller, datacatalog and flytestdlib? - How to setup dev environment for flytekit? - How to debug flytekit remote workflow? - How to setup dev environment for flyteconsole? - How to access Flyte UI, minio, postgres, k3s, and endpoints? - \[Contributing docs and examples\](https://www.union.ai/docs/v1/flyte/community/contributing-docs/page.md) - The combined Flyte and Union docs site - Versions - Common build infrastructure - Variants - Both Flyte and Union docs are open source > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/community/contributing-docs/section.md - \[Contributing docs and examples > Quick start\](https://www.union.ai/docs/v1/flyte/community/contributing-docs/quick-start/page.md) - Prerequisites - Clone the repository - Live preview - Distribution build - \[Contributing docs and examples > Variants\](https://www.union.ai/docs/v1/flyte/community/contributing-docs/variants/page.md) - Variants at the whole-page level - Conditional rendering within a page - {{}} - {{}} - Full example - Adding a new variant - Location - Creating a new variant - Testing the new variant - Building (just) the variant - \[Contributing docs and examples > Versions\](https://www.union.ai/docs/v1/flyte/community/contributing-docs/versions/page.md) - Versions are branches - How to create an archive version - How to create an archive version - Publishing an archive version - \[Contributing docs and examples > Authoring\](https://www.union.ai/docs/v1/flyte/community/contributing-docs/authoring/page.md) - Getting started - Target the right branch - Live preview - Pull Requests + Site Preview - Page Visibility - Page order - Page settings - Conditional Content - Linking to the API reference - Warnings and Notices - Special Content Generation - Python Generated Content - Run on Union Instructions - Jupyter Notebooks - Mapped Keys (\`{{}}\`) - Mermaid Graphs - \[Contributing docs and examples > Shortcodes\](https://www.union.ai/docs/v1/flyte/community/contributing-docs/shortcodes/page.md) - How to specify a "shortcode" - Variants - Component Library - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` and \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\`, \`{{}}\`, and \`{{}}\` - \`{{}}\` - \`{{}}\` - \[Contributing docs and examples > Redirects\](https://www.union.ai/docs/v1/flyte/community/contributing-docs/redirects/page.md) - \`docs.union.ai\` redirects - \`docs.flyte.org\` redirects - \[Contributing docs and examples > API docs\](https://www.union.ai/docs/v1/flyte/community/contributing-docs/api-docs/page.md) - API naming convention - Package Resource Resolution - Tips and Tricks - Auto-linking - Short vs. fully-qualified names - How auto-linking works - Magic-marker syntax for inline code - \[Contributing docs and examples > Publishing\](https://www.union.ai/docs/v1/flyte/community/contributing-docs/publishing/page.md) - Requirements - Managing the Tutorial Pages - Building and running locally - Developer Experience - Controlling Development Environment - Changing 'variants' - Troubleshootting - Identifying Problems: Missing Content - Identifying Problems: Page Visibility - Building Production - Testing Production Build - \[Roadmap\](https://www.union.ai/docs/v1/flyte/community/roadmap/page.md) - How the Community Works - Milestones and Release Processes - Release Cadence - Versioning Scheme - Release Branches and Patching - Documentation Versioning - Planning Process - Quarterly Planning - Change Management - Issue Lifecycle - Browse Features and Issues - Issues by Theme - Issues by Components - \[Troubleshooting guide\](https://www.union.ai/docs/v1/flyte/community/troubleshooting/page.md) - Debugging Common Execution Errors - Error: Cannot connect to the Docker daemon at unix:///var/run/docker.sock. Is the docker daemon running? - message: '0/1 nodes are available: 1 Insufficient cpu. preemption: 0/1 nodes are available: 1 No preemption victims found for incoming pod.' - terminated with exit code (137). Reason \[OOMKilled\] - Error: ImagePullBackOff - Issues Running Workloads - OPENSSL\_internal:WRONG\_VERSION\_NUMBER - ModuleNotFoundError - An error occurred (AccessDenied) when calling the PutObject operation in an EKS deployment - FlyteScopedUserException: 'JavaPackage' object is not callable when running a Spark task - authentication handshake failed: x509: "Kubernetes Ingress Controller Fake Certificate" certificate is not trusted when deploying flyte-core to your own Kubernetes cluster --- ## Architecture - \[Registration\](https://www.union.ai/docs/v1/flyte/architecture/registration/page.md) - Typical Flow - Registration in the Backend - \[Executions\](https://www.union.ai/docs/v1/flyte/architecture/executions/page.md) - Typical Flow Using Flytectl - \[Workflow state transitions\](https://www.union.ai/docs/v1/flyte/architecture/workflow-state-transitions/page.md) - Workflow States - Node States - Task States - \[Workflow timeline\](https://www.union.ai/docs/v1/flyte/architecture/workflow-timeline/page.md) - Acceptance Latency - Transition Latency - Queuing Latency - Completion Latency - Overview of Various Latencies in FlytePropeller - \[Data handling\](https://www.union.ai/docs/v1/flyte/architecture/data-handling/page.md) - Example - Raw data path - \`LiteralType\` and \`Literal\` - Serialization time - Runtime - Data movement - Between FlytePropeller and tasks - Between tasks - Practical example - Bringing in your own datastores for raw data - Deleting raw data in your own datastores - \[Data catalog\](https://www.union.ai/docs/v1/flyte/architecture/data-catalog/page.md) - How Flyte memoizes task executions on data catalog - \[Versions\](https://www.union.ai/docs/v1/flyte/architecture/versions/page.md) - Why do you need versioning? - Operational benefits of completely versioned workflows/pipelines - Why is versioning hard? - How is versioning tied to reproducibility? - What is the cost of versioning and reproducibility? - What is the best way to version your tasks and workflows? - \[Workflow lifecycle\](https://www.union.ai/docs/v1/flyte/architecture/workflow-lifecycle/page.md) - Recap - \[Component Architecture\](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/page.md) - FlyteIDL - Planes - User Plane - Control Plane - Data Plane - Component Code Architecture - Component Code References > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/architecture/component-architecture/section.md - \[Component Architecture > FlytePropeller architecture\](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/flytepropeller\_architecture/page.md) - Introduction - Components - FlyteAdmin - FlyteWorkflow CRD / K8s integration - WorkQueue/WorkerPool - WorkflowExecutor - NodeExecutor - NodeHandlers - FlyteAdmin events - FlytePlugins - \[Component Architecture > Flyte native scheduler architecture\](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/native\_scheduler\_architecture/page.md) - Introduction - Characteristics - Components - Schedule management - Scheduler - Snapshoter - CatchupAll-system - GOCronWrapper - Job executor - Monitoring - \[Control Plane\](https://www.union.ai/docs/v1/flyte/architecture/control-plane/page.md) > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/architecture/control-plane/section.md - \[Control Plane > Projects\](https://www.union.ai/docs/v1/flyte/architecture/control-plane/projects/page.md) - \[Control Plane > Domains\](https://www.union.ai/docs/v1/flyte/architecture/control-plane/domains/page.md) - \[Control Plane > Flyte Admin\](https://www.union.ai/docs/v1/flyte/architecture/control-plane/admin/page.md) - Admin Structure - RPC - Manager {#manager} - Additional Components - Repository {#repository} - Models - Component Details - Async Processes {#async-processes} - Common - Data {#data} - Errors - Runtime {#runtime} - Workflow engine {#workflow-engine} - FlyteAdmin Service Background {#flyteadmin-service-background} - Entities - Execution entities - Platform entities - Using the Admin Service - Adding request filters - Adding sorting to requests - Sorting syntax - \[Control Plane > FlyteConsole\](https://www.union.ai/docs/v1/flyte/architecture/control-plane/console/page.md) - Running FlyteConsole - Install Dependencies - Environment Variables - Run the Server - Development - Storybook - Protobuf and the Network tab - Debug Output - CORS Proxying {#cors-request-proxying} - \[Control Plane > Dynamic Job Spec\](https://www.union.ai/docs/v1/flyte/architecture/control-plane/dynamic-job-spec/page.md) - Tasks - Subworkflows - Nodes - Outputs - \[Extending Flyte\](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/page.md) - Define a Custom Type - Add a New Task Plugin - Flytekit-only task plugin - User Container vs. Pre-built Container Task Plugin - Backend Plugin - Flyte Connector Service - Summary > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/section.md - \[Extending Flyte > Custom types\](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/custom-types/page.md) - \[Extending Flyte > Prebuilt container task plugins\](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/prebuilt-container-task-plugins/page.md) - Usage - How to write a task - Python library - Defining a task - Executor - Image - \[Extending Flyte > User container task plugins\](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/user-container-task-plugins/page.md) - Sensor plugin - Plugin API - Plugin structure - Config objects - Actual usage - \[Extending Flyte > Backend plugins\](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/backend-plugins/page.md) - Basics - Interface specification - Flytekit plugin implementation - FlytePropeller backend plugin - \[Extending Flyte > Container interface\](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/container-interface/page.md) - Command templating --- ## Platform deployment - \[Flyte deployment\](https://www.union.ai/docs/v1/flyte/deployment/flyte-deployment/page.md) > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/deployment/flyte-deployment/section.md - \[Flyte deployment > Components of a Flyte deployment\](https://www.union.ai/docs/v1/flyte/deployment/flyte-deployment/planning/page.md) - Kubernetes cluster - Relational Database - Object store - Optional dependencies - Ingress controller - DNS - SSL/TLS - Helm chart variants - Sandbox - flyte-binary - flyte-core - Additional resources - Terraform reference implementations - Support - \[Flyte deployment > Installing Flyte\](https://www.union.ai/docs/v1/flyte/deployment/flyte-deployment/installing/page.md) - Verify the Installation - Port Forward Flyte Service - Connect to your Flyte instance - \[Flyte deployment > Multi-cluster\](https://www.union.ai/docs/v1/flyte/deployment/flyte-deployment/multicluster/page.md) - Scaling Beyond Kubernetes - Prerequisites - Data Plane Deployment - Control Plane configuration - Configure Execution Cluster Labels - Project-domain execution labels - Configure a Specific Workflow mapping - Day 2 Operations - Add another Kubernetes cluster - \[Platform configuration\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/page.md) > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/section.md - \[Platform configuration > Configuring authentication\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/configuring-authentication/page.md) - Configuring the identity layer - Prerequisites - Configuring your IdP for OIDC - Apply the OIDC configuration to the Flyte backend - Configuring your IdP as an External Authorization Server - Okta - Keycloak - Microsoft Entra ID - Apply the external auth server configuration to Flyte - Configuring supported authorization flows - PKCE - Client Credentials - Device Code - Disable Helm secret management - Continuous Integration - CI - Flytekit / pyflyte - \[Platform configuration > Monitoring a Flyte deployment\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/monitoring/page.md) - Metrics for Executions - Flyte statistics schema - User Stats With Flyte - Use Published Dashboards to Monitor Flyte Deployment - Setup instructions - \[Platform configuration > Configuring logging links in the UI\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/configuring-logging-links-in-the-ui/page.md) - How to configure? - Example configurations - AWS Cloudwatch - Stackdriver (Google Cloud Logging) - Datadog - Kubernetes dashboard - Configure lifetime of logging links - Configure dynamic log links - \[Platform configuration > Configuring Access to GPUs\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/configuring-access-to-gpus/page.md) - Requesting a GPU with no device preference - How it works - Requesting a specific GPU device - How it works - Configuring the nodeSelector - Requesting a GPU partition - How it works - Additional use cases - Request an A100 device with no preference for partition configuration - Request an unpartitioned A100 device - \[Platform configuration > Configuring task pods with K8s PodTemplates\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/configuring-podtemplates/page.md) - A note about containers kinds - Compile-time PodTemplates - Runtime PodTemplates - Set the \`\`default-pod-template-name\`\` in FlytePropeller - Create a PodTemplate resource - Using \`\`pod\_template\_name\`\` in a Task - Flyte's K8s Plugin Configuration - Evaluation Order in PodTemplates - Example 1: Runtime PodTemplate and K8s Plugin Configuration - Example 2: A Runtime and Compile-time PodTemplates - Example 3: Runtime and Compile-time PodTemplates and K8s Plugin Configuration - \[Platform configuration > Cloud Events\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/cloud\_events/page.md) - Use cases - Supported Implementations - Configuration - Helm values configuration - Usage - CloudEvent Spec - \[Platform configuration > Customizing project, domain, and workflow resources with flytectl\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/customizable\_resources/page.md) - Configuring existing resources - About the resource hierarchy - Task resources - Customizing task resource configuration - Cluster resources - Workflow execution configuration - Execution queues - Adding new customizable resources - \[Platform configuration > Platform Events\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/configuring-platform-events/page.md) - Use cases - Supported Implementations - Configuration - AWS SNS - GCP Pub/Sub - Helm configuration - Usage - \[Platform configuration > Workflow notifications\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/configuring-notifications/page.md) - Usage - Setting up workflow notifications - AWS configuration - GCP configuration - Email notifications - Helm configuration - \[Platform configuration > Optimizing Performance\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/performance/page.md) - Introduction - Summarized steps of a workflow execution - Performance tuning at each stage - 1. Workers, the WorkQueue, and the evaluation loop - 2. Querying observed state - 3. Evaluating the DAG and reconciling state - 4. Recording execution status - 5. Report status to the control plane - Concurrency vs parallelism - Scaling out FlyteAdmin - Scaling out Datacatalog - Scaling out FlytePropeller - Sharded scale-out - Multi-Cluster mode - Improving etcd Performance - Offloading Static Workflow Information from CRD - \[Platform configuration > Flyte ResourceManager\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/resource\_manager/page.md) - How Flyte plugins request resources - Plugin resource allocation - Plugin resource deallocation - Configuring ResourceManager to force runtime quota allocation constraints - \[Platform configuration > Secrets\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/secrets/page.md) - Creating secrets with a secrets manager - Prerequisites - Using secrets in tasks - Multiple keys grouped into one secret - Mounting secrets as files or environment variables - Testing with mock secrets - Using secrets in task templates - How secrets injection works - Secret discovery - Configuring a secret management system plugin - Helm Chart Config - AWS secrets manager - GCP secrets manager - Vault secrets manager - Scaling the webhook - Vertical scaling - Horizontal scaling - \[Platform configuration > Security Overview\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/security-overview/page.md) - Changes - Overriding base configuration - Running flyteadmin and flyteconsole on different domains - Modify FlyteAdmin Config - \[Platform configuration > Flyte API Playground: Swagger\](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/swagger/page.md) - \[Connector setup\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/page.md) > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/section.md - \[Connector setup > Airflow connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/airflow/page.md) - Specify connector configuration - Flyte binary - flyte-core - Upgrade the Helm release - \[Connector setup > Google BigQuery connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/bigquery/page.md) - Set up the GCP Flyte cluster - Specify connector configuration - flyte-binary - flyte-core - Upgrade the Helm release - \[Connector setup > ChatGPT connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/chatgpt/page.md) - Specify connector configuration - flyte-binary - flyte-core - Add the OpenAI API token - Upgrade the Helm release - \[Connector setup > Databricks connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/databricks/page.md) - Spin up a cluster - Flyte binary - Flyte core - Databricks workspace - Create an instance profile using the AWS console (For AWS Users) - Locate the IAM role that created the Databricks deployment - Edit the IAM role that created the Databricks deployment - Specify connector configuration - Flyte binary - Add the Databricks access token - Upgrade the deployment - Flyte binary - \[Connector setup > Page\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/k8sservice/page.md) - Kubernetes (K8s) Data Service Connector - Spin up a cluster - Specify connector configuration - Setup the RBAC - Upgrade the deployment - \[Connector setup > MMCloud Connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/mmcloud/page.md) - Set up MMCloud - Spin up a cluster - flyte-binary - flyte-core - Specify connector configuration - \[Connector setup > OpenAI Batch Connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/openai-batch/page.md) - Specify connector configuration - flyte-binary - flyte-core - Add the OpenAI API token - Upgrade the Flyte Helm release - flyte-binary - flyte-core - \[Connector setup > SageMaker Inference Connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/sagemaker-inference/page.md) - Specify connector configuration - flyte-binary - flyte-core - AWS credentials - Upgrade the Flyte Helm release - flyte-binary - flyte-core - \[Connector setup > Sensor connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/sensor/page.md) - Spin up a cluster - flyte-binary - flyte-core - Specify connector configuration - flyte-binary - flyte-core - Upgrade the deployment - Demo cluster - flyte-binary - flyte-core - \[Connector setup > Slurm connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/slurm/page.md) - Spin up a Slurm cluster - Install MUNGE - Create a dedicated Slurm user - Run the Slurm cluster - Test your Slurm connector locally - Overview - Set up a local test environment - Specify connector configuration - Flyte binary - Flyte core - Add the Slurm Private Key - 1. Install flyteconnector pod - 2. Set Private Key as a Secret - 3. Restart development - \[Connector setup > Snowflake connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/snowflake/page.md) - Specify connector configuration - Flyte binary - flyte-core - Upgrade the Helm release - \[Connector setup > DGXC Lepton Connector\](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/dgxc-lepton/page.md) - Prerequisites - Specify connector configuration - flyte-binary - flyte-core - Configure DGXC Lepton connector service - Deploy DGXC Lepton connector service - Configure DGXC Lepton API credentials - Required secrets - Setup instructions - Upgrade the Flyte Helm release - flyte-binary - flyte-core - Verify the setup - Supported task types - Configuration options - \[Plugins\](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/page.md) > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/section.md - \[Plugins > Kubernetes Plugins\](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/kubernetes-plugins/page.md) - Install the Kubernetes operator - Specify plugin configuration - flyte-binary - flyte-core - flyte-binary - flyte-core - flyte-binary - flyte-binary - flyte-core - flyte-binary on AWS - flyte-binary on GCP - flyte-core on AWS - flyte-core on GCP - flyte-sandbox - flyte-binary - flyte-core - flyte-binary - flyte-core - Upgrade the deployment - \[Plugins > Athena Plugin\](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/athena/page.md) - Set up the AWS Flyte cluster - Specify plugin configuration - flyte-binary - flyte-core - Upgrade the Flyte Helm release - \[Plugins > AWS Batch\](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/batch/page.md) - Set up AWS Batch - Modify AWS IAM role trust policy document - Modify system's AWS IAM role policies - Update FlyteAdmin configuration - Update FlytePropeller's configuration - \[Plugins > Sagemaker Plugin Setup\](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/sagemaker/page.md) - Prerequisites - Specify Plugin Configuration - Upgrade the Flyte Helm release - Register the Sagemaker plugin example - Launch an execution - \[Plugins > Google BigQuery Plugin\](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/bigquery/page.md) - Set up the GCP Flyte cluster - Specify plugin configuration - flyte-binary - flyte-core - Upgrade the Flyte Helm release - \[Plugins > Databricks Plugin\](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/databricks/page.md) - Databricks workspace - Specify plugin configuration - flyte-binary - flyte-core - Add the Databricks access token - flyte-binary - flyte-core - Upgrade the deployment - \[Plugins > Snowflake Plugin\](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/snowflake/page.md) - Specify plugin configuration - flyte-binary - flyte-core - Obtain and add the Snowflake JWT token - flyte-binary - flyte-core - Upgrade the deployment - \[Configuration reference\](https://www.union.ai/docs/v1/flyte/deployment/configuration-reference/page.md) > Section bundle (all pages): https://www.union.ai/docs/v1/flyte/deployment/configuration-reference/section.md - \[Configuration reference > Flyte Datacatalog Configuration\](https://www.union.ai/docs/v1/flyte/deployment/configuration-reference/datacatalog-config/page.md) - Section: application - grpcPort (int) - grpcServerReflection (bool) - grpcMaxRecvMsgSizeMBs (int) - httpPort (int) - secure (bool) - readHeaderTimeoutSeconds (int) - Section: database - host (string) - port (int) - dbname (string) - username (string) - password (string) - passwordPath (string) - options (string) - debug (bool) - enableForeignKeyConstraintWhenMigrating (bool) - maxIdleConnections (int) - maxOpenConnections (int) - connMaxLifeTime (config.Duration) - postgres (database.PostgresConfig) - sqlite (database.SQLiteConfig) - Section: datacatalog - storage-prefix (string) - metrics-scope (string) - profiler-port (int) - heartbeat-grace-period-multiplier (int) - max-reservation-heartbeat (config.Duration) - Section: logger - show-source (bool) - mute (bool) - level (int) - formatter (logger.FormatterConfig) - Section: otel - type (string) - file (otelutils.FileConfig) - jaeger (otelutils.JaegerConfig) - otlpgrpc (otelutils.OtlpGrpcConfig) - otlphttp (otelutils.OtlpHttpConfig) - sampler (otelutils.SamplerConfig) - Section: storage - type (string) - connection (storage.ConnectionConfig) - stow (storage.StowConfig) - container (string) - enable-multicontainer (bool) - cache (storage.CachingConfig) - limits (storage.LimitsConfig) - defaultHttpClient (storage.HTTPClientConfig) - signedUrl (storage.SignedURLConfig) - \[Configuration reference > Flyte Admin Configuration\](https://www.union.ai/docs/v1/flyte/deployment/configuration-reference/flyteadmin-config/page.md) - Section: admin - endpoint (config.URL) - insecure (bool) - insecureSkipVerify (bool) - caCertFilePath (string) - maxBackoffDelay (config.Duration) - perRetryTimeout (config.Duration) - maxRetries (int) - maxMessageSizeBytes (int) - authType (uint8) - tokenRefreshWindow (config.Duration) - useAuth (bool) - clientId (string) - clientSecretLocation (string) - clientSecretEnvVar (string) - scopes (\\\[\\\]string) - useAudienceFromAdmin (bool) - audience (string) - authorizationServerUrl (string) - tokenUrl (string) - authorizationHeader (string) - pkceConfig (pkce.Config) - deviceFlowConfig (deviceflow.Config) - command (\\\[\\\]string) - proxyCommand (\\\[\\\]string) - defaultServiceConfig (string) - httpProxyURL (config.URL) - Section: auth - httpAuthorizationHeader (string) - grpcAuthorizationHeader (string) - disableForHttp (bool) - disableForGrpc (bool) - authorizedUris (\\\[\\\]config.URL) - httpProxyURL (config.URL) - userAuth (config.UserAuthConfig) - appAuth (config.OAuth2Options) - Section: catalog-cache - type (string) - endpoint (string) - insecure (bool) - max-cache-age (config.Duration) - use-admin-auth (bool) - max-retries (int) - base-scalar (int) - backoff-jitter (string) - default-service-config (string) - Section: cloudevents - enable (bool) - type (string) - aws (interfaces.AWSConfig) - gcp (interfaces.GCPConfig) - kafka (interfaces.KafkaConfig) - eventsPublisher (interfaces.EventsPublisherConfig) - reconnectAttempts (int) - reconnectDelaySeconds (int) - cloudEventVersion (uint8) - Section: cluster\_resources - templatePath (string) - templateData (map\\\[string\\\]interfaces.DataSource) - refreshInterval (config.Duration) - customData (map\\\[string\\\]map\\\[string\\\]interfaces.DataSource) - standaloneDeployment (bool) - Section: clusterpools - clusterPoolAssignments (map\\\[string\\\]interfaces.ClusterPoolAssignment) - Section: clusters - clusterConfigs (\\\[\\\]interfaces.ClusterConfig) - labelClusterMap (map\\\[string\\\]\\\[\\\]interfaces.ClusterEntity) - defaultExecutionLabel (string) - Section: database - host (string) - port (int) - dbname (string) - username (string) - password (string) - passwordPath (string) - options (string) - debug (bool) - enableForeignKeyConstraintWhenMigrating (bool) - maxIdleConnections (int) - maxOpenConnections (int) - connMaxLifeTime (config.Duration) - postgres (database.PostgresConfig) - sqlite (database.SQLiteConfig) - Section: domains - id (string) - name (string) - Section: event - type (string) - file-path (string) - rate (int64) - capacity (int) - max-retries (int) - base-scalar (int) - backoff-jitter (string) - Section: externalevents - enable (bool) - type (string) - aws (interfaces.AWSConfig) - gcp (interfaces.GCPConfig) - eventsPublisher (interfaces.EventsPublisherConfig) - reconnectAttempts (int) - reconnectDelaySeconds (int) - Section: flyteadmin - roleNameKey (string) - metricsScope (string) - metricsKeys (\\\[\\\]string) - profilerPort (int) - metadataStoragePrefix (\\\[\\\]string) - eventVersion (int) - asyncEventsBufferSize (int) - maxParallelism (int32) - labels (map\\\[string\\\]string) - annotations (map\\\[string\\\]string) - interruptible (bool) - overwriteCache (bool) - assumableIamRole (string) - k8sServiceAccount (string) - outputLocationPrefix (string) - useOffloadedWorkflowClosure (bool) - envs (map\\\[string\\\]string) - featureGates (interfaces.FeatureGates) - consoleUrl (string) - useOffloadedInputs (bool) - Section: logger - show-source (bool) - mute (bool) - level (int) - formatter (logger.FormatterConfig) - Section: namespace\_mapping - mapping (string) - template (string) - templateData (map\\\[string\\\]interfaces.DataSource) - Section: notifications - type (string) - region (string) - aws (interfaces.AWSConfig) - gcp (interfaces.GCPConfig) - publisher (interfaces.NotificationsPublisherConfig) - processor (interfaces.NotificationsProcessorConfig) - emailer (interfaces.NotificationsEmailerConfig) - reconnectAttempts (int) - reconnectDelaySeconds (int) - Section: otel - type (string) - file (otelutils.FileConfig) - jaeger (otelutils.JaegerConfig) - otlpgrpc (otelutils.OtlpGrpcConfig) - otlphttp (otelutils.OtlpHttpConfig) - sampler (otelutils.SamplerConfig) - Section: plugins - connector-service (connector.Config) - catalogcache (catalog.Config) - connector-service (connector.Config) - k8s (config.K8sPluginConfig) - k8s-array (k8s.Config) - logs (logs.LogConfig) - Section: propeller - kube-config (string) - master (string) - workers (int) - workflow-reeval-duration (config.Duration) - downstream-eval-duration (config.Duration) - limit-namespace (string) - prof-port (config.Port) - metadata-prefix (string) - rawoutput-prefix (string) - queue (config.CompositeQueueConfig) - metrics-prefix (string) - metrics-keys (\\\[\\\]string) - enable-admin-launcher (bool) - max-workflow-retries (int) - max-ttl-hours (int) - gc-interval (config.Duration) - leader-election (config.LeaderElectionConfig) - publish-k8s-events (bool) - max-output-size-bytes (int64) - enable-grpc-latency-metrics (bool) - kube-client-config (config.KubeClientConfig) - node-config (config.NodeConfig) - max-streak-length (int) - event-config (config.EventConfig) - include-shard-key-label (\\\[\\\]string) - exclude-shard-key-label (\\\[\\\]string) - include-project-label (\\\[\\\]string) - exclude-project-label (\\\[\\\]string) - include-domain-label (\\\[\\\]string) - exclude-domain-label (\\\[\\\]string) - cluster-id (string) - create-flyteworkflow-crd (bool) - node-execution-worker-count (int) - array-node-config (config.ArrayNodeConfig) - literal-offloading-config (config.LiteralOffloadingConfig) - admin-launcher (launchplan.AdminConfig) - resourcemanager (config.Config) - workflowstore (workflowstore.Config) - Section: qualityofservice - tierExecutionValues (map\\\[string\\\]interfaces.QualityOfServiceSpec) - defaultTiers (map\\\[string\\\]string) - Section: queues - executionQueues (interfaces.ExecutionQueues) - workflowConfigs (interfaces.WorkflowConfigs) - Section: registration - maxWorkflowNodes (int) - maxLabelEntries (int) - maxAnnotationEntries (int) - workflowSizeLimit (string) - Section: remotedata - scheme (string) - region (string) - signedUrls (interfaces.SignedURL) - maxSizeInBytes (int64) - inlineEventDataPolicy (int) - Section: scheduler - profilerPort (config.Port) - eventScheduler (interfaces.EventSchedulerConfig) - workflowExecutor (interfaces.WorkflowExecutorConfig) - reconnectAttempts (int) - reconnectDelaySeconds (int) - Section: secrets - secrets-prefix (string) - env-prefix (string) - Section: server - httpPort (int) - grpcPort (int) - grpcServerReflection (bool) - kube-config (string) - master (string) - security (config.ServerSecurityOptions) - grpc (config.GrpcConfig) - thirdPartyConfig (config.ThirdPartyConfigOptions) - dataProxy (config.DataProxyConfig) - readHeaderTimeoutSeconds (int) - kubeClientConfig (config.KubeClientConfig (kubeClientConfig)) - Section: storage - type (string) - connection (storage.ConnectionConfig) - stow (storage.StowConfig) - container (string) - enable-multicontainer (bool) - cache (storage.CachingConfig) - limits (storage.LimitsConfig) - defaultHttpClient (storage.HTTPClientConfig) - signedUrl (storage.SignedURLConfig) - Section: task\_resources - defaults (interfaces.TaskResourceSet) - limits (interfaces.TaskResourceSet) - Section: tasks - task-plugins (config.TaskPluginConfig) - max-plugin-phase-versions (int32) - backoff (config.BackOffConfig) - maxLogMessageLength (int) - \[Configuration reference > Flyte Propeller Configuration\](https://www.union.ai/docs/v1/flyte/deployment/configuration-reference/flytepropeller-config/page.md) - Section: admin - endpoint (config.URL) - insecure (bool) - insecureSkipVerify (bool) - caCertFilePath (string) - maxBackoffDelay (config.Duration) - perRetryTimeout (config.Duration) - maxRetries (int) - maxMessageSizeBytes (int) - authType (uint8) - tokenRefreshWindow (config.Duration) - useAuth (bool) - clientId (string) - clientSecretLocation (string) - clientSecretEnvVar (string) - scopes (\\\[\\\]string) - useAudienceFromAdmin (bool) - audience (string) - authorizationServerUrl (string) - tokenUrl (string) - authorizationHeader (string) - pkceConfig (pkce.Config) - deviceFlowConfig (deviceflow.Config) - command (\\\[\\\]string) - proxyCommand (\\\[\\\]string) - defaultServiceConfig (string) - httpProxyURL (config.URL) - Section: catalog-cache - type (string) - endpoint (string) - insecure (bool) - max-cache-age (config.Duration) - use-admin-auth (bool) - max-retries (int) - base-scalar (int) - backoff-jitter (string) - default-service-config (string) - Section: event - type (string) - file-path (string) - rate (int64) - capacity (int) - max-retries (int) - base-scalar (int) - backoff-jitter (string) - Section: logger - show-source (bool) - mute (bool) - level (int) - formatter (logger.FormatterConfig) - Section: otel - type (string) - file (otelutils.FileConfig) - jaeger (otelutils.JaegerConfig) - otlpgrpc (otelutils.OtlpGrpcConfig) - otlphttp (otelutils.OtlpHttpConfig) - sampler (otelutils.SamplerConfig) - Section: plugins - connector-service (connector.Config) - athena (athena.Config) - aws (aws.Config) - bigquery (bigquery.Config) - catalogcache (catalog.Config) - connector-service (connector.Config) - dask (dask.Config) - databricks (databricks.Config) - echo (testing.Config) - k8s (config.K8sPluginConfig) - k8s-array (k8s.Config) - kf-operator (common.Config) - logs (logs.LogConfig) - qubole (config.Config) - ray (ray.Config) - snowflake (snowflake.Config) - spark (spark.Config) - Section: propeller - kube-config (string) - master (string) - workers (int) - workflow-reeval-duration (config.Duration) - downstream-eval-duration (config.Duration) - limit-namespace (string) - prof-port (config.Port) - metadata-prefix (string) - rawoutput-prefix (string) - queue (config.CompositeQueueConfig) - metrics-prefix (string) - metrics-keys (\\\[\\\]string) - enable-admin-launcher (bool) - max-workflow-retries (int) - max-ttl-hours (int) - gc-interval (config.Duration) - leader-election (config.LeaderElectionConfig) - publish-k8s-events (bool) - max-output-size-bytes (int64) - enable-grpc-latency-metrics (bool) - kube-client-config (config.KubeClientConfig) - node-config (config.NodeConfig) - max-streak-length (int) - event-config (config.EventConfig) - include-shard-key-label (\\\[\\\]string) - exclude-shard-key-label (\\\[\\\]string) - include-project-label (\\\[\\\]string) - exclude-project-label (\\\[\\\]string) - include-domain-label (\\\[\\\]string) - exclude-domain-label (\\\[\\\]string) - cluster-id (string) - create-flyteworkflow-crd (bool) - node-execution-worker-count (int) - array-node-config (config.ArrayNodeConfig) - literal-offloading-config (config.LiteralOffloadingConfig) - admin-launcher (launchplan.AdminConfig) - resourcemanager (config.Config (resourcemanager)) - workflowstore (workflowstore.Config) - Section: secrets - secrets-prefix (string) - env-prefix (string) - Section: storage - type (string) - connection (storage.ConnectionConfig) - stow (storage.StowConfig) - container (string) - enable-multicontainer (bool) - cache (storage.CachingConfig) - limits (storage.LimitsConfig) - defaultHttpClient (storage.HTTPClientConfig) - signedUrl (storage.SignedURLConfig) - Section: tasks - task-plugins (config.TaskPluginConfig) - max-plugin-phase-versions (int32) - backoff (config.BackOffConfig) - maxLogMessageLength (int) - Section: webhook - metrics-prefix (string) - certDir (string) - localCert (bool) - listenPort (int) - serviceName (string) - servicePort (int32) - secretName (string) - secretManagerType (int) - awsSecretManager (config.AWSSecretManagerConfig) - gcpSecretManager (config.GCPSecretManagerConfig) - vaultSecretManager (config.VaultSecretManagerConfig) - \[Configuration reference > Flyte Scheduler Configuration\](https://www.union.ai/docs/v1/flyte/deployment/configuration-reference/scheduler-config/page.md) - Section: admin - endpoint (config.URL) - insecure (bool) - insecureSkipVerify (bool) - caCertFilePath (string) - maxBackoffDelay (config.Duration) - perRetryTimeout (config.Duration) - maxRetries (int) - maxMessageSizeBytes (int) - authType (uint8) - tokenRefreshWindow (config.Duration) - useAuth (bool) - clientId (string) - clientSecretLocation (string) - clientSecretEnvVar (string) - scopes (\\\[\\\]string) - useAudienceFromAdmin (bool) - audience (string) - authorizationServerUrl (string) - tokenUrl (string) - authorizationHeader (string) - pkceConfig (pkce.Config) - deviceFlowConfig (deviceflow.Config) - command (\\\[\\\]string) - proxyCommand (\\\[\\\]string) - defaultServiceConfig (string) - httpProxyURL (config.URL) - Section: auth - httpAuthorizationHeader (string) - grpcAuthorizationHeader (string) - disableForHttp (bool) - disableForGrpc (bool) - authorizedUris (\\\[\\\]config.URL) - httpProxyURL (config.URL) - userAuth (config.UserAuthConfig) - appAuth (config.OAuth2Options) - Section: catalog-cache - type (string) - endpoint (string) - insecure (bool) - max-cache-age (config.Duration) - use-admin-auth (bool) - max-retries (int) - base-scalar (int) - backoff-jitter (string) - default-service-config (string) - Section: cloudevents - enable (bool) - type (string) - aws (interfaces.AWSConfig) - gcp (interfaces.GCPConfig) - kafka (interfaces.KafkaConfig) - eventsPublisher (interfaces.EventsPublisherConfig) - reconnectAttempts (int) - reconnectDelaySeconds (int) - cloudEventVersion (uint8) - Section: cluster\_resources - templatePath (string) - templateData (map\\\[string\\\]interfaces.DataSource) - refreshInterval (config.Duration) - customData (map\\\[string\\\]map\\\[string\\\]interfaces.DataSource) - standaloneDeployment (bool) - Section: clusterpools - clusterPoolAssignments (map\\\[string\\\]interfaces.ClusterPoolAssignment) - Section: clusters - clusterConfigs (\\\[\\\]interfaces.ClusterConfig) - labelClusterMap (map\\\[string\\\]\\\[\\\]interfaces.ClusterEntity) - defaultExecutionLabel (string) - Section: database - host (string) - port (int) - dbname (string) - username (string) - password (string) - passwordPath (string) - options (string) - debug (bool) - enableForeignKeyConstraintWhenMigrating (bool) - maxIdleConnections (int) - maxOpenConnections (int) - connMaxLifeTime (config.Duration) - postgres (database.PostgresConfig) - sqlite (database.SQLiteConfig) - Section: domains - id (string) - name (string) - Section: event - type (string) - file-path (string) - rate (int64) - capacity (int) - max-retries (int) - base-scalar (int) - backoff-jitter (string) - Section: externalevents - enable (bool) - type (string) - aws (interfaces.AWSConfig) - gcp (interfaces.GCPConfig) - eventsPublisher (interfaces.EventsPublisherConfig) - reconnectAttempts (int) - reconnectDelaySeconds (int) - Section: flyteadmin - roleNameKey (string) - metricsScope (string) - metricsKeys (\\\[\\\]string) - profilerPort (int) - metadataStoragePrefix (\\\[\\\]string) - eventVersion (int) - asyncEventsBufferSize (int) - maxParallelism (int32) - labels (map\\\[string\\\]string) - annotations (map\\\[string\\\]string) - interruptible (bool) - overwriteCache (bool) - assumableIamRole (string) - k8sServiceAccount (string) - outputLocationPrefix (string) - useOffloadedWorkflowClosure (bool) - envs (map\\\[string\\\]string) - featureGates (interfaces.FeatureGates) - consoleUrl (string) - useOffloadedInputs (bool) - Section: logger - show-source (bool) - mute (bool) - level (int) - formatter (logger.FormatterConfig) - Section: namespace\_mapping - mapping (string) - template (string) - templateData (map\\\[string\\\]interfaces.DataSource) - Section: notifications - type (string) - region (string) - aws (interfaces.AWSConfig) - gcp (interfaces.GCPConfig) - publisher (interfaces.NotificationsPublisherConfig) - processor (interfaces.NotificationsProcessorConfig) - emailer (interfaces.NotificationsEmailerConfig) - reconnectAttempts (int) - reconnectDelaySeconds (int) - Section: otel - type (string) - file (otelutils.FileConfig) - jaeger (otelutils.JaegerConfig) - otlpgrpc (otelutils.OtlpGrpcConfig) - otlphttp (otelutils.OtlpHttpConfig) - sampler (otelutils.SamplerConfig) - Section: plugins - connector-service (connector.Config) - catalogcache (catalog.Config) - connector-service (connector.Config) - k8s (config.K8sPluginConfig) - k8s-array (k8s.Config) - logs (logs.LogConfig) - Section: propeller - kube-config (string) - master (string) - workers (int) - workflow-reeval-duration (config.Duration) - downstream-eval-duration (config.Duration) - limit-namespace (string) - prof-port (config.Port) - metadata-prefix (string) - rawoutput-prefix (string) - queue (config.CompositeQueueConfig) - metrics-prefix (string) - metrics-keys (\\\[\\\]string) - enable-admin-launcher (bool) - max-workflow-retries (int) - max-ttl-hours (int) - gc-interval (config.Duration) - leader-election (config.LeaderElectionConfig) - publish-k8s-events (bool) - max-output-size-bytes (int64) - enable-grpc-latency-metrics (bool) - kube-client-config (config.KubeClientConfig) - node-config (config.NodeConfig) - max-streak-length (int) - event-config (config.EventConfig) - include-shard-key-label (\\\[\\\]string) - exclude-shard-key-label (\\\[\\\]string) - include-project-label (\\\[\\\]string) - exclude-project-label (\\\[\\\]string) - include-domain-label (\\\[\\\]string) - exclude-domain-label (\\\[\\\]string) - cluster-id (string) - create-flyteworkflow-crd (bool) - node-execution-worker-count (int) - array-node-config (config.ArrayNodeConfig) - literal-offloading-config (config.LiteralOffloadingConfig) - admin-launcher (launchplan.AdminConfig) - resourcemanager (config.Config) - workflowstore (workflowstore.Config) - Section: qualityofservice - tierExecutionValues (map\\\[string\\\]interfaces.QualityOfServiceSpec) - defaultTiers (map\\\[string\\\]string) - Section: queues - executionQueues (interfaces.ExecutionQueues) - workflowConfigs (interfaces.WorkflowConfigs) - Section: registration - maxWorkflowNodes (int) - maxLabelEntries (int) - maxAnnotationEntries (int) - workflowSizeLimit (string) - Section: remotedata - scheme (string) - region (string) - signedUrls (interfaces.SignedURL) - maxSizeInBytes (int64) - inlineEventDataPolicy (int) - Section: scheduler - profilerPort (config.Port) - eventScheduler (interfaces.EventSchedulerConfig) - workflowExecutor (interfaces.WorkflowExecutorConfig) - reconnectAttempts (int) - reconnectDelaySeconds (int) - Section: secrets - secrets-prefix (string) - env-prefix (string) - Section: server - httpPort (int) - grpcPort (int) - grpcServerReflection (bool) - kube-config (string) - master (string) - security (config.ServerSecurityOptions) - grpc (config.GrpcConfig) - thirdPartyConfig (config.ThirdPartyConfigOptions) - dataProxy (config.DataProxyConfig) - readHeaderTimeoutSeconds (int) - kubeClientConfig (config.KubeClientConfig (kubeClientConfig)) - Section: storage - type (string) - connection (storage.ConnectionConfig) - stow (storage.StowConfig) - container (string) - enable-multicontainer (bool) - cache (storage.CachingConfig) - limits (storage.LimitsConfig) - defaultHttpClient (storage.HTTPClientConfig) - signedUrl (storage.SignedURLConfig) - Section: task\_resources - defaults (interfaces.TaskResourceSet) - limits (interfaces.TaskResourceSet) - Section: tasks - task-plugins (config.TaskPluginConfig) - max-plugin-phase-versions (int32) - backoff (config.BackOffConfig) - maxLogMessageLength (int) --- # Unknown \# Union.ai Documentation > Full documentation (single file): https://www.union.ai/docs/v2/union/llms-full.txt > Site: https://www.union.ai/docs/v2/union Each entry below is \`- \[Page title\](URL)\` followed by the H2/H3 headings found on that page. Pages link to individual \`page.md\` files. Sections marked with a "Section bundle" link have a \`section.md\` that concatenates all pages in the section into a single file — use it to load an entire section into context at once. ## User guide - \[Overview\](https://www.union.ai/docs/v2/union/user-guide/overview/page.md) - Pure Python, no DSL - Durability - Reproducibility - Recoverability - Built for scale - What this means in practice - \[Quickstart\](https://www.union.ai/docs/v2/union/user-guide/quickstart/page.md) - What you'll need - Install the SDK - Configure - Write your first workflow - Run it - See the results - Next steps - \[Core concepts\](https://www.union.ai/docs/v2/union/user-guide/core-concepts/page.md) - How Flyte works > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/core-concepts/section.md - \[Core concepts > TaskEnvironment\](https://www.union.ai/docs/v2/union/user-guide/core-concepts/task-environment/page.md) - A minimal example - What TaskEnvironment controls - Configuring resources - Configuring container images - Multiple tasks, one environment - Multiple environments - Next steps - \[Core concepts > Tasks\](https://www.union.ai/docs/v2/union/user-guide/core-concepts/tasks/page.md) - Defining a task - Type hints are required - Tasks calling tasks - The top-level task - Running tasks locally - Running tasks remotely - Next steps - \[Core concepts > Runs and actions\](https://www.union.ai/docs/v2/union/user-guide/core-concepts/runs-and-actions/page.md) - What is a run? - What is an action? - Runs vs actions in practice - Viewing runs in the UI - Understanding the execution graph - Checking run status - Next steps - \[Core concepts > Where your data lives\](https://www.union.ai/docs/v2/union/user-guide/core-concepts/where-data-lives/page.md) - The two stores - What goes in the database - What goes in the bucket - What "metadata" means - 1. "Metadata" as in the control-plane database (Flyte's usage) - 2. "Metadata bucket" (a deployment/ops term you may see) - Per-run customization: \`raw\_data\_path\` - What happens if the bucket is purged - The short version - \[Core concepts > Apps\](https://www.union.ai/docs/v2/union/user-guide/core-concepts/introducing-apps/page.md) - Tasks vs apps - AppEnvironment - A hello world app - Understanding the code - Serving the app - When to use apps vs tasks - Common patterns - Next steps - \[Core concepts > Projects and domains\](https://www.union.ai/docs/v2/union/user-guide/core-concepts/projects-and-domains/page.md) - How projects and domains are used - Managing projects via CLI - Create a project - List projects - Update a project - Archive a project - Unarchive a project - Listing projects programmatically - Managing projects via the UI - Domains - Targeting a domain - \[Core concepts > Key capabilities\](https://www.union.ai/docs/v2/union/user-guide/core-concepts/key-capabilities/page.md) - Environment and resources - Deployment - Data handling - Parallelism and composition - Security and automation - Durability and reliability - Apps and serving - Notebooks - Next steps - \[Core concepts > Settings\](https://www.union.ai/docs/v2/union/user-guide/core-concepts/settings/page.md) - Scope hierarchy - Editing settings - Settings for scope: DOMAIN(production) - Remove or comment out a line to inherit that setting from the parent scope. - Set a value to ~unset to explicitly clear it, blocking parent inheritance. - Local overrides - Default queue for task runs - Inherited settings (uncomment to override at this scope) - Kubernetes service account for task pods - Available settings (uncomment and edit to set at this scope) - Base path for raw data storage (e.g. s3://my-bucket/prefix) - CPU resource quantity (e.g. "500m", "2") - Memory resource quantity (e.g. "256Mi", "4Gi") - Apply settings from a file - Available settings - Inheritance rules - Explicitly clearing a value - Relationship to task configuration - Permissions - \[Core concepts > Basic project: RAG\](https://www.union.ai/docs/v2/union/user-guide/core-concepts/basic-project/page.md) - Concepts covered - Part 1: The embedding pipeline - Setting up the environment - Fetching data - Creating embeddings - Orchestrating the pipeline - Running the pipeline - Part 2: The serving application - App environment configuration - The Streamlit application - Deploying the app - Key takeaways - \[Run modes\](https://www.union.ai/docs/v2/union/user-guide/run-modes/page.md) - \[Local\](running-locally/page.md) - \[Devbox\](running-devbox/page.md) - \[Remote\](running-remote/page.md) - \[Run modes > Run locally\](https://www.union.ai/docs/v2/union/user-guide/run-modes/running-locally/page.md) - Getting started - Running tasks locally - Terminal UI - Exploring past runs - What works locally - Local to devbox/remote - Next steps - \[Run modes > Run on the devbox\](https://www.union.ai/docs/v2/union/user-guide/run-modes/running-devbox/page.md) - What you'll need - Install the SDK - Start the devbox - CPU - GPU - Configure - Run a workflow on the devbox - View results in the UI - Stop the devbox - Inline configuration - Programmatic - CLI - Delete the devbox - Using a CUDA-enabled GPU host - Next steps - \[Run modes > Run on a remote cluster\](https://www.union.ai/docs/v2/union/user-guide/run-modes/running-remote/page.md) - Prerequisites - Install the flyte package - Configuration file - Using the configuration - Explicit configuration - Programmatic - CLI - Programmatic - CLI - Check current configuration - Inline configuration - Programmatic - CLI - Next steps - \[Configure tasks\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/page.md) - Task configuration levels - Example - Task configuration parameters > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/task-configuration/section.md - \[Configure tasks > Container images\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/container-images/page.md) - Specifying your own image directly - Specifying your own image with the \`flyte.Image\` object - Example: Defining a custom image with \`Image.from\_debian\_base\` - Example: Defining an image based on uv script metadata - Image building - Configuring the \`builder\` - Local image building - Remote \`ImageBuilder\` - Install private PyPI packages - \[Configure tasks > Resources\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/resources/page.md) - Resources data class - Examples - Usage in TaskEnvironment - Usage in a task-specific override - \[Configure tasks > Secrets\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/secrets/page.md) - Creating a literal string secret - Creating a file secret - Scoping secrets - Listing secrets - Deleting secrets - Using a literal string secret - Using a file secret - \[Configure tasks > Caching\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/caching/page.md) - Overview - Basic caching usage - \`"auto"\` - Automatic versioning - \`"override"\` - \`"disable"\` - No caching - Advanced caching configuration - Ignoring specific inputs - Cache serialization - Salt for cache key variation - Cache policies - Function body policy (default) - Custom cache policies - Caching configuration at different levels - \`TaskEnvironment\` Level - \`@env.task\` decorator level - \`task.override\` level - Runtime cache control - Project and domain cache isolation - Local development caching - \[Configure tasks > Reusable containers\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/reusable-containers/page.md) - How It Works - Basic Usage - \`ReusePolicy\` parameters - Understanding parameter relationships - Key relationships - Simple example - \[Configure tasks > Pod templates\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/pod-templates/page.md) - How it works - Requirements - Basic usage - PodTemplate parameters - Volume mounts - GCS/S3 volume mounts - Sidecar containers - Image pull secrets - Cluster-specific configuration - Important notes - Best practices - Learn more - \[Configure tasks > Multiple environments\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/multiple-environments/page.md) - Constraints on multiple environments - Task \`depends\_on\` constraints - Dependency inclusion constraints - Example - \[Configure tasks > Retries and timeouts\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/retries-and-timeouts/page.md) - The action lifecycle - Retries - Retry count - Retries with exponential backoff - Skip retries for failures that can't be fixed - System retries - Timeouts - \`max\_runtime\` — bound a single attempt's execution - \`max\_queued\_time\` — fail fast when capacity isn't available - \`deadline\` — bound the total wall-clock - Combining the bounds - Combining retries and timeouts - \[Configure tasks > Triggers\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/triggers/page.md) - Triggers are set in the task decorator - \`flyte.Trigger\` - The \`automation\` parameter with \`flyte.FixedRate\` - Examples - The \`automation\` parameter with \`flyte.Cron\` - Examples - The \`inputs\` parameter - Basic Usage - Using \`flyte.TriggerTime\` - Required vs optional parameters - Complex input types - Predefined schedule triggers - Available Predefined Triggers - Trigger time in predefined triggers - Multiple triggers per task - Notifications - Execution phases - Template variables - Slack notifications - Email notifications - Microsoft Teams notifications - Custom webhook notifications - Deploying a task with triggers - Activating and deactivating triggers - Trigger run timing - Cron-based triggers - Fixed-rate triggers without \`start\_time\` - Fixed-rate triggers with \`start\_time\` - Deleting triggers - Schedule time zones - Setting time zone for a Cron schedule - \`flyte.TriggerTime\` is always in UTC - Daylight Savings Time behavior - \[Configure tasks > Interruptible tasks\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/interruptible-tasks-and-queues/page.md) - Setting at different levels - Behavior on preemption - Spot to on-demand fallback - \[Configure tasks > Queues\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/queues/page.md) - Routing work to a queue - Per-run and per-trigger routing - Overriding a queue at runtime - What a queue controls - When to use queues - Concurrency control for scheduled and automated runs - Backfill control - Multi-cluster routing and prioritizing certain workloads - Queues and timeouts - \[Configure tasks > Task plugins\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/task-plugins/page.md) - Default Execution: Containers - Compute Plugins - Available Compute Plugins - How Compute Plugins Work - Using Compute Plugins - Using Plugins on Union - Backend Integrations - Next Steps - \[Configure tasks > Additional task settings\](https://www.union.ai/docs/v2/union/user-guide/task-configuration/additional-task-settings/page.md) - Naming and metadata - \`name\` - \`short\_name\` - \`description\` - \`docs\` - \`report\` - Source-code link (automatic) - \`links\` - Default inputs - Environment variables - Inline I/O threshold - \[Build tasks\](https://www.union.ai/docs/v2/union/user-guide/task-programming/page.md) - What you'll learn - When to use these patterns > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/task-programming/section.md - \[Build tasks > Files and directories\](https://www.union.ai/docs/v2/union/user-guide/task-programming/files-and-directories/page.md) - Example usage - JSONL files - Setup - JsonlFile - Compression - JsonlDir - Error handling - Batch iteration - \[Build tasks > Data classes and structures\](https://www.union.ai/docs/v2/union/user-guide/task-programming/dataclasses-and-structures/page.md) - Example: Combining Dataclasses and Pydantic Models - \[Build tasks > DataFrames\](https://www.union.ai/docs/v2/union/user-guide/task-programming/dataframes/page.md) - Setting up the environment and sample data - Create a raw DataFrame - Create a flyte.io.DataFrame - Automatically convert between types - Downloading DataFrames - Run the example - Polars DataFrames - Setup - Eager DataFrames - Lazy DataFrames - Run the example - \[Build tasks > Custom types\](https://www.union.ai/docs/v2/union/user-guide/task-programming/handling-custom-types/page.md) - Types of extensions - Creating a type transformer - Step 1: Define your custom type - Step 2: Create the type transformer - Step 3: Register the transformer - Distributing type plugins - Configure pyproject.toml - Automatic loading - Controlling plugin loading - Using custom types in tasks - DataFrame extensions - Best practices - \[Build tasks > Custom context\](https://www.union.ai/docs/v2/union/user-guide/task-programming/custom-context/page.md) - Overview - When to use it and when not to - Setting custom context - Run-level context - Overriding inside a task (local override that affects nested tasks) - Adding new keys for nested tasks - Accessing custom context - \[Build tasks > Abort and cancel actions\](https://www.union.ai/docs/v2/union/user-guide/task-programming/abort-tasks/page.md) - Action lifetime - Canceling actions programmatically - External abort - Aborting via the CLI - Handling external aborts - \[Build tasks > Run a bioinformatics tool\](https://www.union.ai/docs/v2/union/user-guide/task-programming/container-tasks/page.md) - What are Container Tasks? - How Data Flows In and Out - Basic Usage - Template Syntax for Inputs - Using Container Tasks in Workflows - Advanced: Passing Files and Directories - Use Case: Agentic Sandbox Execution - Use Case: Legacy and Specialized Containers - Use Case: Multi-Language Workflows - Configuration Options - ContainerTask Parameters - Supported Input/Output Types - Best Practices - Local Execution - When to Use Container Tasks - \[Build tasks > Links\](https://www.union.ai/docs/v2/union/user-guide/task-programming/links/page.md) - Creating a link - Using execution metadata - Dynamic links with override - \[Build tasks > Reports\](https://www.union.ai/docs/v2/union/user-guide/task-programming/reports/page.md) - A simple example - A more complex example - Streaming example - \[Build tasks > Notebooks\](https://www.union.ai/docs/v2/union/user-guide/task-programming/notebooks/page.md) - Iterating on and running a workflow - Accessing runs and downloading logs - \[Build tasks > Remote tasks\](https://www.union.ai/docs/v2/union/user-guide/task-programming/remote-tasks/page.md) - Prerequisites - Basic usage - Understanding lazy loading - When tasks are fetched - Benefits of lazy loading - Error handling - Eager fetching with \`fetch()\` - Module-level vs dynamic loading - Complete example - Team A: Spark environment - Team B: ML environment - Team C: Orchestration - Invoke remote tasks in a script. - Why use remote tasks? - When to use remote tasks - How remote tasks work - Security model - Type system - Versioning options - Customizing remote tasks - Available overrides - Override examples - Chain overrides - Best practices - 1. Use meaningful task names - 2. Document task interfaces - 3. Prefer module-level loading - 4. Handle versioning thoughtfully - 5. Deploy remote tasks first - Limitations - Next steps - \[Build tasks > Error handling\](https://www.union.ai/docs/v2/union/user-guide/task-programming/error-handling/page.md) - \[Build tasks > Traces\](https://www.union.ai/docs/v2/union/user-guide/task-programming/traces/page.md) - What are traced functions for? - What Gets Traced - Errors are not recorded - Supported Function Types - Task Orchestration Pattern - Relationship to Caching and Checkpointing - How They Work Together - Execution Flow - Error Handling and Observability - Examples in Practice - LLM Pipeline with Traces - \[Build tasks > Grouping actions\](https://www.union.ai/docs/v2/union/user-guide/task-programming/grouping-actions/page.md) - What are groups? - The problem groups solve - How groups work - Common grouping patterns - Sequential operations - Parallel processing with groups - Multi-phase workflows - Conditional grouping - Key insights - \[Build tasks > Fanout\](https://www.union.ai/docs/v2/union/user-guide/task-programming/fanout/page.md) - Understanding fanout - Example - Parallel execution - Running the example - How Flyte handles concurrency and parallelism - \[Build tasks > Controlling parallel execution\](https://www.union.ai/docs/v2/union/user-guide/task-programming/controlling-parallelism/page.md) - The problem: unbounded parallelism - Using asyncio.Semaphore - Using flyte.map with concurrency - Running the example - When to use each approach - \[Build tasks > Human-in-the-loop\](https://www.union.ai/docs/v2/union/user-guide/task-programming/human-in-the-loop/page.md) - Setup - Automated task - Requesting human input - Wiring it together - Event options - Submitting input programmatically - \[Build tasks > Test business logic directly\](https://www.union.ai/docs/v2/union/user-guide/task-programming/unit-testing/page.md) - Understanding Task Invocation - Direct Function Invocation - Using \`flyte.run()\` - Testing Business Logic - Testing Async Tasks - Testing Nested Tasks - Testing Type Transformations and Serialization - Testing Type Restrictions - Testing Nested Tasks with Serialization - Testing Traced Functions - Best Practices - Quick Reference - Example Test Suite - Future Improvements - \[Build tasks > Regular async function (not a task)\](https://www.union.ai/docs/v2/union/user-guide/task-programming/other-features/page.md) - Task Forwarding - Passing Tasks and Functions as Arguments - Custom Action Names - Set at Task Definition - Override at Call Time - Invoking Async Functions from Sync Tasks - Async and Sync Task Interoperability - Calling Sync Tasks from Async Tasks - Using with \`flyte.map.aio()\` - Using AnyIO in Async Tasks - \[Run and deploy tasks\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/page.md) - Ephemeral deployment and immediate execution - Programmatic - CLI - Persistent deployment - Programmatic - CLI - Running already deployed tasks - Programmatic - CLI - Configuring runs with \`flyte.with\_runcontext()\` > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/task-deployment/section.md - \[Run and deploy tasks > How task run works\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/how-task-run-works/page.md) - Ephemeral deployment + run: The development shortcut - Programmatic - CLI - Running deployed tasks - Programmatic - CLI - Local execution - Programmatic - CLI - Running tasks through the Union UI - Accessing task execution in the Union UI - Execution flow and architecture - Fast registration architecture - Ephemeral preparation logic - Execution modes comparison - \[Run and deploy tasks > Interact with runs and actions\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/interacting-with-runs/page.md) - Understanding runs and actions - Key concepts - Working with runs - Retrieving a run - Programmatic - CLI - Watching run progress - Getting detailed run information - Working with actions - Retrieving an action - Programmatic - CLI - Nested actions - Getting detailed action information - Retrieving inputs and outputs - Programmatic - CLI - Handling failures - Understanding data storage - Accessing large data from cloud storage - S3 storage access - GCS storage access - Azure Blob Storage access - Complete example - API reference - Key classes - CLI commands - Storage configuration - \[Run and deploy tasks > Work with local data\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/work-with-local-data/page.md) - Local execution - Uploading local data to remote runs - Uploading DataFrames - Uploading files - Uploading directories - Passing outputs between runs - Performance considerations - Summary - \[Run and deploy tasks > Run command options\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/run-command-options/page.md) - \`--project\`, \`--domain\` - \`--run-project\`, \`--run-domain\` - \`--local\` - When to use local execution - \`--copy-style\` - Copy style options - \`--root-dir\` - \`--raw-data-path\` - Use cases - \`--service-account\` - Use cases - \`--name\` - Benefits of custom names - \`--follow\` - Behavior - \`--image\` - Image mapping formats - \`--no-sync-local-sys-paths\` - Task argument passing - SDK options - \[Run and deploy tasks > How task deployment works\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/how-task-deployment-works/page.md) - 1. Module loading and task environment discovery - Single file (default) - \`--all\` option - \`--recursive\` option - 2. Task analysis and serialization - 3. Task environment dependency resolution - 4. Code bundle creation and upload - \`--copy\_style loaded\_modules\` (default) - \`--copy\_style all\` - \`--copy\_style none\` - \`--root-dir\` option - 5. Image building - Local image building - Remote image building - 6. Source-code link discovery - How the link is built - Conditions - Understanding option relationships - \[Run and deploy tasks > Deploy command options\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/deploy-command-options/page.md) - \`--project\`, \`--domain\` - \`--version\` - When versions are used - \`--dry-run\` - \`--all\` and \`--recursive\` - \`--copy-style\` - \`--copy-style loaded\_modules\` (default) - \`--copy-style all\` - \`--copy-style none\` - \`--root-dir\` - Default behavior (without \`--root-dir\`) - Common use cases - How it works - Example with complex project structure - \`--image\` - Named image mappings - Default image mapping - How it works - \`--ignore-load-errors\` - \`--no-sync-local-sys-paths\` - Default behavior (path synchronization enabled) - When to disable path synchronization - Use cases for disabling - How it works - SDK deployment options - \[Run and deploy tasks > Code packaging for remote execution\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/packaging/page.md) - Quick comparison - Code bundling - How it works - Automatic code bundling - Manual code bundling - Including additional files with \`include\` - Controlling the root directory - Code bundling examples - When to use code bundling - Container-based deployment - How it works - Configuration - Image source copying methods - Complete container-based example - Using externally built images - Container-based best practices - When to use container-based deployment - Choosing the right approach - Decision tree - Hybrid approach - Troubleshooting - Import errors - Code changes not reflected - Files missing in container - Container build failures - Version conflicts - \[Run and deploy tasks > Running Tasks via Webhooks\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/invoke-webhook/page.md) - How passthrough authentication works - Setting up passthrough authentication - Initialize with \`flyte.init\_passthrough()\` - Passing authentication metadata - Complete example - Calling the webhook - Best practices - Troubleshooting - "FLYTE\_ENDPOINT environment variable not set" - "Authentication credentials required" - "Task not found" - Tasks run with wrong permissions - \[Run and deploy tasks > Deployment patterns\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/deployment-patterns/page.md) - Overview of deployment patterns - Simple file deployment - Example structure - Deployment commands - When to use - Custom Dockerfile deployment - Example structure - Alternative: Dockerfile in different directory - Key considerations - When to use - PyProject package deployment - Example structure - Business logic modules - Flyte orchestration layer - Entrypoint configuration - Dependencies and configuration - Key features - Key learning points - Usage patterns - What this example demonstrates - When to use - Package structure deployment - Example structure - Key concepts - Running with root directory - How \`--root-dir\` works - Alternative: Using a Python project - When to use - Full build deployment - Overview - Key configuration - Local dependency example - Critical configuration components - Configuration options - Version management best practices - Performance considerations - When to use - Troubleshooting - Python path deployment - Example structure - Implementation - Task environment dependencies - Key considerations - CLI vs Direct Python execution - Best practices - Common pitfalls - When to use - Dynamic environment deployment - Domain-based environment selection - Why this pattern works - How it works - Important constraints - Alternative: Environment variable approach - Usage patterns - When to use dynamic environments - Best practices - Project organization - Image management - Configuration management - Development workflow - Choosing the right pattern - \[Run and deploy tasks > Run context\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/run-context/page.md) - Configuring a run with \`flyte.with\_runcontext()\` - Execution target - Storage - Caching - Identity and resources - Logging - Code bundling - Context propagation - Reading context inside a task with \`flyte.ctx()\` - \`TaskContext\` fields - \`ActionID\` fields - Naming external resources - \[Run and deploy tasks > Entrypoint tasks\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/entrypoints/page.md) - When to use entrypoints - Mark a task as an entrypoint - Discover entrypoint tasks - CLI - Programmatic - UI - \[Run and deploy tasks > Run a Python script\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/run-python-script/page.md) - When to use it - Quickstart - What you can specify - Handling dependencies - 1. Standard library only - 2. Pip packages on top of the default image - 3. Bring your own image - Handling inputs and outputs - Passing inputs - Capturing outputs - Bundling: what gets uploaded - Python API - \[Run and deploy tasks > Run with notifications\](https://www.union.ai/docs/v2/union/user-guide/task-deployment/run-with-notifications/page.md) - Execution phases - Template variables - Slack notifications - Email notifications - \[Configure apps\](https://www.union.ai/docs/v2/union/user-guide/configure-apps/page.md) - Hello World example - Using fserve args - Using @app\_env.server - Differences from TaskEnvironment - Configuration topics > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/configure-apps/section.md - \[Configure apps > App environment settings\](https://www.union.ai/docs/v2/union/user-guide/configure-apps/app-environment-settings/page.md) - Shared environment settings - App-specific environment settings - Environment variable substitution in \`args\` - App startup - Server decorator via \`@app\_env.server\` - Container command via \`command\` vs \`args\` - Shared settings - \[Configure apps > Including additional files\](https://www.union.ai/docs/v2/union/user-guide/configure-apps/including-additional-files/page.md) - How include works - When to use include - Examples - Multi-file Streamlit app - Multi-file FastAPI app - App with configuration files - File discovery - Path resolution - Best practices - Limitations - \[Configure apps > Passing parameters into app environments\](https://www.union.ai/docs/v2/union/user-guide/configure-apps/passing-parameters/page.md) - Parameter types overview - Basic parameter types - Accessing parameters in your app - \`mount\` - \`env\_var\` - \`get\_parameter\` - Full example - Delayed values - RunOutput - AppEndpoint - Overriding parameters at serve time - Example: FastAPI app with configurable model - Example: Using RunOutput for model serving - Best practices - Limitations - \[Configure apps > /// script\](https://www.union.ai/docs/v2/union/user-guide/configure-apps/auto-scaling-apps/page.md) - Autoscaling apps - Scaling configuration - Basic scaling example - Scaling patterns - Idle TTL (Time To Live) - Autoscaling best practices - Autoscaling limitations - Autoscaling troubleshooting - \[Configure apps > Apps depending on other environments\](https://www.union.ai/docs/v2/union/user-guide/configure-apps/apps-depending-on-environments/page.md) - Basic usage - Example: App calling another app - Dependency chain - Multiple dependencies - Using AppEndpoint for dependency URLs - Deployment behavior - Task environment dependencies - Best practices - Example: A/B testing with dependencies - Limitations - \[Build apps\](https://www.union.ai/docs/v2/union/user-guide/build-apps/page.md) - App types - Usage patterns - Next steps > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/build-apps/section.md - \[Build apps > Single-script apps\](https://www.union.ai/docs/v2/union/user-guide/build-apps/single-script-apps/page.md) - Plain Python HTTP server - Streamlit app - FastAPI app - Running single-script apps - When to use single-script apps - \[Build apps > Multi-script apps\](https://www.union.ai/docs/v2/union/user-guide/build-apps/multi-script-apps/page.md) - FastAPI multi-script app - Project structure - Example: Multi-file FastAPI app - Automatic file discovery - Streamlit multi-script app - Project structure - Example: Multi-file Streamlit app - Deploying multi-file Streamlit app - Complex multi-file example - Project structure - Example code - Deploying complex app - Best practices - Troubleshooting - \[Build apps > Serving graphs\](https://www.union.ai/docs/v2/union/user-guide/build-apps/serving-graphs/page.md) - Core concepts: a minimal two-app chain - Deploying multiple apps together with \`depends\_on\` - Getting an upstream app's endpoint - Sizing each node independently - Example: CPU / GPU inference split - Disjoint images per node - GPU app: model.forward only - CPU app: pre/postprocess + call GPU - Deploy - Example: A/B testing with Statsig - Statsig client singleton - Variant apps - Root app with Statsig in its lifespan - App environments - Routing endpoint - Deploy - When to split into a serving graph - Best practices - \[Build apps > Hybrid app-task graphs\](https://www.union.ai/docs/v2/union/user-guide/build-apps/hybrid-graphs/page.md) - Call app from task - Example: FastAPI app called from a task - Call task from app (webhooks / APIs) - Example: Basic webhook app - Advanced webhook patterns - Webhook security and best practices - Example: GitHub webhook - Gradio agent UI - Best practices - \[Build apps > WebSocket apps\](https://www.union.ai/docs/v2/union/user-guide/build-apps/websocket-apps/page.md) - Example: Basic WebSocket app - WebSocket patterns - Using WebSockets with Flyte tasks - WebSocket client example - Best practices - \[Build apps > Browser apps\](https://www.union.ai/docs/v2/union/user-guide/build-apps/browser-apps/page.md) - Accessing browser-based apps - Common browser-based app types - Streamlit apps - Gradio apps - Custom HTML/JS apps - Best practices - \[Build apps > Secret-based authentication\](https://www.union.ai/docs/v2/union/user-guide/build-apps/secret-based-authentication/page.md) - Create the secret - Define the FastAPI app - Deploy the FastAPI app - Invoke the endpoint - Authentication for vLLM and SGLang apps - Create the authentication secret - Deploy vLLM app with authentication - Deploy SGLang app with authentication - Invoke authenticated LLM endpoints - Accessing Swagger documentation - Security best practices - Troubleshooting - Next steps - \[Build apps > Connector app\](https://www.union.ai/docs/v2/union/user-guide/build-apps/connector-app/page.md) - When to build a custom connector - Project structure - Step 1: Implement the connector - Step 2: Create the task plugin - Step 3: Deploy the connector - Step 4: Register and run tasks - How it all fits together - Secrets - Per-task secrets (per-user credentials) - Related - \[Native app integrations\](https://www.union.ai/docs/v2/union/user-guide/native-app-integrations/page.md) - When to use a native integration - Available integrations - Next steps > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/native-app-integrations/section.md - \[Native app integrations > Streamlit app\](https://www.union.ai/docs/v2/union/user-guide/native-app-integrations/streamlit-app/page.md) - Basic Streamlit app - Single-file Streamlit app - Multi-file Streamlit app - Example: Data visualization dashboard - Best practices - Troubleshooting - \[Native app integrations > FastAPI app\](https://www.union.ai/docs/v2/union/user-guide/native-app-integrations/fastapi-app/page.md) - Basic FastAPI app - Serving a machine learning model - Accessing Swagger documentation - Example: REST API with multiple endpoints - Multi-file FastAPI app - Local-to-remote model serving - Best practices - Advanced features - Troubleshooting - \[Native app integrations > vLLM app\](https://www.union.ai/docs/v2/union/user-guide/native-app-integrations/vllm-app/page.md) - Installation - Basic vLLM app - Using prefetched models - Model streaming - Custom vLLM arguments - Using the OpenAI-compatible API - Multi-GPU inference (Tensor Parallelism) - Model sharding with prefetch - Autoscaling - Best practices - Troubleshooting - \[Native app integrations > SGLang app\](https://www.union.ai/docs/v2/union/user-guide/native-app-integrations/sglang-app/page.md) - Installation - Basic SGLang app - Using prefetched models - Model streaming - Custom SGLang arguments - Using the OpenAI-compatible API - Multi-GPU inference (Tensor Parallelism) - Model sharding with prefetch - Autoscaling - Structured generation - Best practices - Troubleshooting - \[Native app integrations > Flyte webhook\](https://www.union.ai/docs/v2/union/user-guide/native-app-integrations/flyte-webhook/page.md) - Available endpoints - Basic usage - Filtering endpoints - Endpoint groups - Individual endpoints - Allow-listing - Task allow-list - App allow-list - Trigger allow-list - Calling the webhook - Authentication - Self-reference protection - \[Serve and deploy apps\](https://www.union.ai/docs/v2/union/user-guide/serve-and-deploy-apps/page.md) - Serve vs Deploy - \`flyte serve\` - \`flyte deploy\` - Using Python SDK - Serve - Deploy - Using the CLI - Serve - Deploy - Next steps > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/serve-and-deploy-apps/section.md - \[Serve and deploy apps > How app serving works\](https://www.union.ai/docs/v2/union/user-guide/serve-and-deploy-apps/how-app-serving-works/page.md) - Overview - Using the Python SDK - Overriding parameters - Advanced serving options - Using CLI - Return value - Best practices - Troubleshooting - \[Serve and deploy apps > How app custom domains work\](https://www.union.ai/docs/v2/union/user-guide/serve-and-deploy-apps/how-app-custom-domain-works/page.md) - \[Serve and deploy apps > How app deployment works\](https://www.union.ai/docs/v2/union/user-guide/serve-and-deploy-apps/how-app-deployment-works/page.md) - Overview - Using the Python SDK - Deployment plan - Overriding App configuration at deployment time - Activation/deactivation - Using the CLI - Example: Full deployment configuration - Best practices - Deployment status and return value - Troubleshooting - \[Serve and deploy apps > Activating and deactivating apps\](https://www.union.ai/docs/v2/union/user-guide/serve-and-deploy-apps/activating-and-deactivating-apps/page.md) - Activation - Activate after deployment - Activate an app - Check activation status - Deactivation - Lifecycle management - Typical deployment workflow - Blue-green deployment - Using CLI - Activate - Deactivate - Check status - Best practices - Automatic activation with serve - Example: Complete deployment and activation - Troubleshooting - \[Serve and deploy apps > Prefetching models\](https://www.union.ai/docs/v2/union/user-guide/serve-and-deploy-apps/prefetching-models/page.md) - Why prefetch? - Basic prefetch - Using Python SDK - Using CLI - Using prefetched models - Prefetch options - Custom artifact name - With HuggingFace token - With resources - Sharding models for multi-GPU - vLLM sharding - Using shard config via CLI - Using prefetched sharded models - CLI options - Complete example - Best practices - Troubleshooting - \[Build an agent\](https://www.union.ai/docs/v2/union/user-guide/build-agent/page.md) - How Union.ai maps to the agentic world - Ways to build an agent - Deploying an agent - Related > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/build-agent/section.md - \[Build an agent > Pure Python agents\](https://www.union.ai/docs/v2/union/user-guide/build-agent/python-agents/page.md) - ReAct pattern: Reason, Act, Observe (no framework needed) - Plan-and-Execute with parallel fan-out - More agentic patterns - How Union.ai's primitives map to the agent stack - Next steps - \[Build an agent > Flyte-native agents\](https://www.union.ai/docs/v2/union/user-guide/build-agent/flyte-agents/page.md) - How it works - Sync vs async - A minimal agent - Tools - Customizing a tool with \`tool(...)\` - MCP integration - Skills - Observability - Extending the Agent class - \`run\` is sync-by-default - Strategy 1: wrap the built-in loop - Strategy 2: implement \`run\` from scratch - Choosing between subclassing and composition - Next steps - \[Build an agent > Agent memory\](https://www.union.ai/docs/v2/union/user-guide/build-agent/agent-memory/page.md) - What a \`MemoryStore\` holds - Sync vs async - Keyed stores: the easy path - Path-addressed artifacts - Optimistic concurrency - Optional capabilities - Lower-level usage (without a key) - Next steps - \[Build an agent > Agent chat UI\](https://www.union.ai/docs/v2/union/user-guide/build-agent/agent-chat-ui/page.md) - Option 1: the built-in chat UI - Option 2: a custom FastAPI chat app - Next steps - \[Build an agent > Deploy an agent as a service\](https://www.union.ai/docs/v2/union/user-guide/build-agent/deploy-agent-as-service/page.md) - As a task - As a scheduled task (via \`Trigger\`) - Behind a webhook (\`AppEnvironment\`) - Chat and other app patterns - \[Agent framework integrations\](https://www.union.ai/docs/v2/union/user-guide/agent-framework-integrations/page.md) - How much control does the framework give you? - Supported frameworks - Next steps - \[Agent framework integrations > LangGraph agents\](https://www.union.ai/docs/v2/union/user-guide/agent-framework-integrations/langgraph/page.md) - A single LangGraph agent in a task - Plan-and-Execute: fan out LangGraph agents in parallel - Next steps - \[Agent framework integrations > PydanticAI agents\](https://www.union.ai/docs/v2/union/user-guide/agent-framework-integrations/pydantic-ai/page.md) - A PydanticAI agent in a task - Parallel agents - Next steps - \[Agent framework integrations > OpenAI agents\](https://www.union.ai/docs/v2/union/user-guide/agent-framework-integrations/openai-agents-sdk/page.md) - Tools that are also durable tasks - Next steps - \[Agent framework integrations > Bring your own framework\](https://www.union.ai/docs/v2/union/user-guide/agent-framework-integrations/bring-your-own-framework/page.md) - The core pattern - Make tools durable - Trace the framework's internals - Fan out across containers - Checklist - Next steps - \[Build an MCP\](https://www.union.ai/docs/v2/union/user-guide/build-mcp/page.md) - HTTP layout - Quickstart - \[Build an MCP > User-defined MCP server\](https://www.union.ai/docs/v2/union/user-guide/build-mcp/mcp\_server/page.md) - When to use it - How it works - Basic example - HTTP layout - Choosing a transport - Connecting a client - Claude Code - OpenCode - Configuration tips - Best practices - \[Build an MCP > Flyte MCP server\](https://www.union.ai/docs/v2/union/user-guide/build-mcp/flyte\_mcp\_server/page.md) - When to use it - How to run it - Running locally with \`uvx\` - Deploying remotely - Basic example - Scoping the server - 1. Tool groups - 2. Individual tools - 3. Allowlists - Enabling the search tools - Putting it together: a filtered server - Connecting a client - Claude Code — local (stdio) - Claude Code — remote (HTTP) - OpenCode — local - OpenCode — remote - Best practices - MCP tools reference - \[Sandboxing\](https://www.union.ai/docs/v2/union/user-guide/sandboxing/page.md) - Why sandboxing matters for AI - Types of sandboxes - What Flyte offers - Workflow sandbox (Monty) - Code sandbox (container) - When to use which - Learn more > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/sandboxing/section.md - \[Sandboxing > Workflow sandboxing in Flyte\](https://www.union.ai/docs/v2/union/user-guide/sandboxing/workflow-sandboxing-flyte/page.md) - Why workflow sandboxing? - How it works - Example: sandboxed orchestrator - Example: dynamic code execution - Reusable task from a code string - One-shot local execution - Parameterized code generation - Building agents with programmatic tool calling - Syntax restrictions - Allowed - Not allowed - Type restrictions - Security model - \[Sandboxing > Programmatic tool calling for agents\](https://www.union.ai/docs/v2/union/user-guide/sandboxing/code-mode/page.md) - Programmatic tool calling vs sequential tool calling - Why programmatic tool calling is powerful - Token efficiency - Performance - Natural programming patterns - Progressive tool discovery - Data privacy - Example: sequential vs programmatic tool calling - Sequential tool calling approach - Programmatic tool calling approach - Example: defining tools - Example: programmatic tool-calling agent - Example: chat app - Example: durable agent - References - \[Sandboxing > Code sandboxing\](https://www.union.ai/docs/v2/union/user-guide/sandboxing/code-sandboxing/page.md) - Execution modes - Auto-IO mode - Verbatim mode - Command mode - Executing a sandbox - Error handling - Supported types - How types are handled - Configuring the container image - Python packages - System packages - Additional Dockerfile commands - Pre-built images - Image configuration - Runtime configuration - Resources - Retries - Timeout - Environment variables - Secrets - Caching - Deploying a sandbox as a task - End-to-end example - API reference - \`flyte.sandbox.create()\` - Sandbox methods - \[Authenticating with Union\](https://www.union.ai/docs/v2/union/user-guide/authenticating/page.md) - Quick start - Programmatic - CLI - Authentication modes - PKCE authentication (browser-based) {#pkce} - Programmatic - CLI - Device flow authentication {#device-flow} - Programmatic - CLI - API key authentication (OAuth2 app credentials) {#api-key} - Programmatic - CLI - Comparison table - Switching between authentication modes - Token storage and keyring {#token-storage} - How it works - Systems without keyring support - Solution: Install keyrings.alt - Verifying keyring functionality - Troubleshooting - Browser doesn't open for PKCE - Device flow code expires - API key doesn't work - Best practices - \[User management\](https://www.union.ai/docs/v2/union/user-guide/user-management/page.md) - Built-in policies - Custom roles and policies - Prerequisites - Walkthrough: restrict a team to run workflows in production only - Programmatic - CLI - Updating roles and policies - Programmatic - CLI - Managing users in the UI - Available actions - CLI command reference - \[Project patterns\](https://www.union.ai/docs/v2/union/user-guide/project-patterns/page.md) - \[Bring your own image (BYOI)\](bring-your-own-image/page.md) - \[Monorepo with uv\](monorepo-with-uv/page.md) - \[CI/CD deployments\](cicd/page.md) - \[Resource management and multi-team scaling\](resource-management/page.md) - \[Project patterns > Bring your own image\](https://www.union.ai/docs/v2/union/user-guide/project-patterns/bring-your-own-image/page.md) - The multi-team problem - Pattern 1: Pure BYOI - Dockerfiles - Build and push - Python code - Run and deploy - Pattern 2: Remote Builder - The base images - Adapting with \`flyte.Image\` - Task definitions - Entry point - Run and deploy - Decision matrix - \[Project patterns > Structuring Flyte projects with uv\](https://www.union.ai/docs/v2/union/user-guide/project-patterns/monorepo-with-uv/page.md) - The two layers - How the image gets built - \`with\_code\_bundle()\` — one image for dev and prod - \`root\_dir\` - Monorepo patterns - Pattern A: Shared lockfile (recommended) - Pattern B: Independent packages - The full build path (production) - \[Project patterns > CI/CD deployments\](https://www.union.ai/docs/v2/union/user-guide/project-patterns/cicd/page.md) - What CI needs to do - Authentication: API keys - Mint the key - Store the key as a CI secret - Key scope and rotation - Project configuration - \`config.yaml\` - The GitHub Actions workflow - Key flag choices - Splitting build from deploy - Layering on top of an existing image build - \[Project patterns > Resource management and multi-team scaling\](https://www.union.ai/docs/v2/union/user-guide/project-patterns/resource-management/page.md) - Project-domain structure - One project per team or ML product - Domains are environments, not teams - Resource quotas - Set quotas per project-domain pair - Why quotas matter - Task-level resources - Declare resources on the task environment - Be explicit about ephemeral storage - RBAC and secrets - Roles vs policies - Scope secrets as narrowly as possible - Multi-team scaling patterns - Establish naming conventions early - Put shared utility tasks in a dedicated project - Use cluster assignment for multi-cluster deployments - Treat production as a managed service - Quick reference - \[Scale your runs\](https://www.union.ai/docs/v2/union/user-guide/run-scaling/page.md) - Understanding Flyte execution - Performance optimization - Key concepts for scaling > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/run-scaling/section.md - \[Scale your runs > Data flow\](https://www.union.ai/docs/v2/union/user-guide/run-scaling/data-flow/page.md) - Overview - Data types and transport - Passed by reference - Passed by value (inline I/O) - Task execution and data flow - Input download - Output upload - Task-to-task data flow - Caching and data hashing - Cache key computation - Inline data caching - Reference data hashing - Cache control - Traces and data flow - Object stores and latency considerations - Configuring data storage - Organization and project level - Per-run configuration - \[Scale your runs > Life of a run\](https://www.union.ai/docs/v2/union/user-guide/run-scaling/life-of-a-run/page.md) - Overview - Phase 1: Code analysis and preparation - Phase 2: Image building - Phase 3: Code bundling - Default: \`copy\_style="loaded\_modules"\` - Alternative: \`copy\_style="none"\` - Phase 4: Upload code bundle - Phase 5: Run creation and queuing - Phase 6: Task execution in data plane - Container startup - Invoking downstream tasks - Execution flow diagram - Action identifiers and crash recovery - Downstream task execution - Reusable containers - Reusable container execution flow - State replication and visualization - Queue Service to Run Service - UI limitations - Optimization opportunities - \[Scale your runs > Scale your workflows\](https://www.union.ai/docs/v2/union/user-guide/run-scaling/scale-your-workflows/page.md) - Understanding performance dimensions - Latency - Throughput - Task execution overhead - The overhead principle - System architecture and data flow - Optimization strategies - 1. Use reusable containers for concurrency - 2. Batch workloads to reduce overhead - 3. Use traces for lightweight operations - 4. Limit fanout for system stability - 5. Optimize data transfer - 6. Leverage caching - 7. Parallelize with \`flyte.map\` - Performance tuning workflow - Real-world example: PyIceberg batch processing - Example: Optimizing a data pipeline - Before optimization - After optimization - When to contact the Union team - \[Scale your runs > Maximize GPU utilization for batch inference\](https://www.union.ai/docs/v2/union/user-guide/run-scaling/batch-inference/page.md) - Why GPU utilization drops - Serving vs in-process batch inference - Solution: \`DynamicBatcher\` - Basic usage - Cost estimation - \`TokenBatcher\` for LLM inference - Combining with reusable containers - Example: batch LLM inference with vLLM - Monitoring utilization - \[Advanced project: LLM reporting agent\](https://www.union.ai/docs/v2/union/user-guide/advanced-project/page.md) - What you'll build - Concepts covered - Architecture - Prerequisites - Parts - Key takeaways > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/advanced-project/section.md - \[Advanced project: LLM reporting agent > Resilient generation\](https://www.union.ai/docs/v2/union/user-guide/advanced-project/resilient-generation/page.md) - Two environments - Reusable environment for LLM work - ReusePolicy parameters - Standard environment for orchestration - Traced LLM calls - Benefits of tracing - When to use @flyte.trace - Traced helper functions - Retry strategies - Configuring retries - Combining tracing with retries - Structured prompts - Pydantic models for structured output - Next steps - \[Advanced project: LLM reporting agent > Agentic refinement\](https://www.union.ai/docs/v2/union/user-guide/advanced-project/agentic-refinement/page.md) - The agentic pattern - Critique function - Revise function - The refinement loop - How it works - Early exit - Grouping iterations with flyte.group - Why use flyte.group? - Group context - Configuring the loop - Choosing thresholds - Best practices for agentic loops - Next steps - \[Advanced project: LLM reporting agent > Parallel outputs\](https://www.union.ai/docs/v2/union/user-guide/advanced-project/parallel-outputs/page.md) - The formatting functions - When to trace and when not to - Parallel execution with asyncio.gather - How asyncio.gather works - When to use asyncio.gather - Grouping parallel operations - Collecting outputs in a directory - The batch pipeline - Pipeline flow - Running the pipeline - Cost optimization tips - 1. Choose the right model - 2. Tune iteration parameters - 3. Use caching effectively - 4. Scale the batch - Next steps - \[Advanced project: LLM reporting agent > Serving app\](https://www.union.ai/docs/v2/union/user-guide/advanced-project/serving-app/page.md) - App environment configuration - Key configuration - Connecting to pipeline output with RunOutput - The Streamlit application - Displaying multiple reports - Generation instructions - Deploying the app - Workflow: Generate then view - Automatic updates with RunOutput - Complete example structure - Running the complete example - Summary - \[Migration\](https://www.union.ai/docs/v2/union/user-guide/migration/page.md) - \[From Flyte 1 to 2\](flyte-2/page.md) - \[From Airflow to Flyte\](from-airflow/page.md) - \[Migration > From Flyte 1 to 2\](https://www.union.ai/docs/v2/union/user-guide/migration/flyte-2/page.md) - Pure Python execution - Sync Python - Async Python - Simplified API - Fine-grained reproducibility and recoverability - Improved remote functionality - Native Notebook support - High performance engine - Enhanced UI > Section bundle (all pages): https://www.union.ai/docs/v2/union/user-guide/migration/flyte-2/section.md - \[Migration > From Flyte 1 to 2 > Pure Python\](https://www.union.ai/docs/v2/union/user-guide/migration/flyte-2/pure-python/page.md) - From \`@workflow\` DSL to pure Python - Flyte 1 - Flyte 2 - \[Migration > From Flyte 1 to 2 > Asynchronous model\](https://www.union.ai/docs/v2/union/user-guide/migration/flyte-2/async/page.md) - Why we need an async model - Understanding concurrency vs. parallelism - Python's async evolution - Parallelism in Flyte 1 vs Flyte 2 - Core async concepts - True parallelism for all workloads - Calling sync tasks from async tasks - Synchronous task support - The \`flyte.map\` function: Familiar patterns - Sync Map - Async Map - \[Migration > From Flyte 1 to 2 > Migration from Flyte 1 to Flyte 2\](https://www.union.ai/docs/v2/union/user-guide/migration/flyte-2/migration/page.md) - 1. Move task configuration to a \`TaskEnvironment\` object - 2. Replace workflow decorators - Flyte 1 - Flyte 2 Sync - Flyte 2 Async - \[Migration > From Flyte 1 to 2 > Considerations\](https://www.union.ai/docs/v2/union/user-guide/migration/flyte-2/considerations/page.md) - Non-deterministic behavior - Dealing with non-determinism - Type safety - No global state - Driver pod requirements - OOM risk from materialized I/O - \[Migration > From Airflow to Flyte\](https://www.union.ai/docs/v2/union/user-guide/migration/from-airflow/page.md) - \[Part 1 — Vanilla Operators\](part-1-vanilla-operators/page.md) - \[Migration > From Airflow to Flyte > Part 1 — Vanilla Operators\](https://www.union.ai/docs/v2/union/user-guide/migration/from-airflow/part-1-vanilla-operators/page.md) - 1. Where dependencies are specified - 2. The driver task (in place of DAGs) - 3. Triggers (in place of schedules) - 4. PythonOperator to \`@env.task\` - File and Dir — for data that doesn't fit in a return value - 5. TaskFlow to \`@env.task\` - TaskFlow decorator variants - 6. BashOperator to ContainerTask - When to use \`ContainerTask\` - How the arguments map - 7. KubernetesPodOperator to TaskEnvironment + PodTemplate - Where every KPO knob lands - What a fully-specified task looks like - 8. Orchestration: parallelism, conditionals, error handling - Parallelism - Dynamic mapping - Conditionals - Error handling - What's next - Caching - Reusable containers - Reports - Apps --- ## Tutorials - \[Biotech & Healthcare\](https://www.union.ai/docs/v2/union/tutorials/biotech-healthcare/page.md) - \[Genomic alignment\](genomic-alignment/page.md) - \[Brain tumor MRI classification\](tumor-detection/page.md) - \[Biotech & Healthcare > Genomic alignment\](https://www.union.ai/docs/v2/union/tutorials/biotech-healthcare/genomic-alignment/page.md) - Define the container image - Define the task environments - Define the data classes - Fetch assets - Quality filtering with fastp - Build the Bowtie 2 index - Align reads - Orchestrate the workflow - Run the workflow - \[Biotech & Healthcare > Brain tumor MRI classification\](https://www.union.ai/docs/v2/union/tutorials/biotech-healthcare/tumor-detection/page.md) - Define the container image - Define the task environments - Configure training - Load the dataset - Train the model - Resumable checkpointing - Generate the report - Orchestrate the pipeline - Run the pipeline - \[Geospatial\](https://www.union.ai/docs/v2/union/tutorials/geospatial/page.md) - \[GPU-accelerated climate modeling\](climate-modeling/page.md) - \[Satellite image classification\](satellite\_image\_classification/page.md) - \[Geospatial > GPU-accelerated climate modeling\](https://www.union.ai/docs/v2/union/tutorials/geospatial/climate-modeling/page.md) - Overview - Implementation - Dependencies and container image - Simulation parameters and data structures - Task environments - Data ingestion: multiple sources in parallel - Preprocessing with Dask - GPU-accelerated atmospheric simulation - Distributing across multiple GPUs - The main workflow - Running the pipeline - Key concepts - Ensemble forecasting - Adaptive mesh refinement - Real-time event detection - Where to go next - \[Geospatial > Satellite image classification\](https://www.union.ai/docs/v2/union/tutorials/geospatial/satellite\_image\_classification/page.md) - Background - Dataset - Model - Two-Phase Training - Pipeline - Task 1: Data Download (\`dataset\_env\`) - Task 2: GPU Training (\`training\_env\`) - Task 3: Report Generation (\`report\_env\`) - Task 4: Orchestration (\`pipeline\_env\`) - Running the Pipeline - What You Get - \[Financial Services & Fintech\](https://www.union.ai/docs/v2/union/tutorials/financial-services/page.md) - \[Financial research agent\](financial-research-agent/page.md) - \[Multi-agent trading simulation\](trading-agents/page.md) - \[Financial Services & Fintech > Multi-agent trading simulation\](https://www.union.ai/docs/v2/union/tutorials/financial-services/trading-agents/page.md) - TL;DR - What is an agent, anyway? - What's different here? - How it works: step-by-step walkthrough - Entry point - Analyst agents - Research agents - Trading agent - Risk agents - Retaining agent memory with S3 vectors - Running the simulation - Why Flyte? \_(A quick note before you go)\_ - \[Financial Services & Fintech > Financial research agent\](https://www.union.ai/docs/v2/union/tutorials/financial-services/financial-research-agent/page.md) - Setting up the environment - Data types - You.com Research and Search APIs - Synthesize briefings with Claude - Research one company - Orchestration - Run the agent - Create secrets - Run locally or remotely - \[Frontier AI\](https://www.union.ai/docs/v2/union/tutorials/frontier-ai/page.md) - \[Distributed LLM pretraining\](distributed-pretraining/page.md) - \[Frontier AI > Distributed LLM pretraining\](https://www.union.ai/docs/v2/union/tutorials/frontier-ai/distributed-pretraining/page.md) - Overview - Implementation - Setting up the environment - Declaring resource requirements - Model configurations - Building the GPT model - The Lightning training module - Checkpointing for fault tolerance - Real-time metrics with Flyte Reports - Streaming data at scale - Distributed training with FSDP - Tying it together - Running the pipeline - Going further - \[Computer Vision\](https://www.union.ai/docs/v2/union/tutorials/computer-vision/page.md) - \[Fine-tuning a VLM\](qwen-vl-finetuning/page.md) - \[Multimodal retrieval evaluation\](multimodal-retrieval-evaluation/page.md) - \[Computer Vision > Fine-tuning a VLM\](https://www.union.ai/docs/v2/union/tutorials/computer-vision/qwen-vl-finetuning/page.md) - Overview - Implementation - Setting up the environment - Preparing the dataset - The adapter - Multi-node training with DeepSpeed - Fault tolerance and recovery - Live observability - Evaluation - Putting it all together - Running the tutorial - Going further - \[Computer Vision > Multimodal retrieval evaluation\](https://www.union.ai/docs/v2/union/tutorials/computer-vision/multimodal-retrieval-evaluation/page.md) - Define the container image - Define the task environments - Configuration and data types - Loading, indexing, and search - Run one experiment - Compare experiments - Run the evaluation - \[Agents\](https://www.union.ai/docs/v2/union/tutorials/agents/page.md) - \[Autoresearch agent\](autoresearch/page.md) - \[Coding agent\](code-agent/page.md) - \[Competitive intelligence agent\](competitive-intelligence-agent/page.md) - \[Compliance monitoring agent\](compliance-monitoring-agent/page.md) - \[Deep research\](deep-research/page.md) - \[Field data enrichment agent\](field-data-enrichment-agent/page.md) - \[MLE Bot: autonomous ML engineer\](mle-bot/page.md) - \[Support resolution agent\](support-resolution-agent/page.md) - \[Agents > Autoresearch agent\](https://www.union.ai/docs/v2/union/tutorials/agents/autoresearch/page.md) - Define the container image - Define the task environment - What changed - All changed files - Model the result - What changed - All changed files - The autoresearch task - What changed - All changed files - What changed - All changed files - Run the agent - Create secrets - Prepare the research repository - Run remotely - \[Agents > Coding agent\](https://www.union.ai/docs/v2/union/tutorials/agents/code-agent/page.md) - What this example demonstrates - Setting up the agent environment - Retrieve docs - Code generation - Running the code agent - \[Agents > MLE Bot: an autonomous ML engineer\](https://www.union.ai/docs/v2/union/tutorials/agents/mle-bot/page.md) - TL;DR - The problem with LLMs and ML pipelines - How it works - What to expect - Declaring task environments - Building durable tool functions - Guiding the LLM with domain knowledge - Dataset context - General ML best practices - The agent loop: profile, design, execute, iterate - Dataset context - General ML best practices — apply these based on the dataset context above - Available tools - Monty sandbox restrictions - Critical patterns for using tool results - When fixing a previous error - Pipeline design — you decide the structure - Response format - Reasoning - Code - Running LLM-generated code in Flyte's sandbox - Dataset context - General ML best practices — apply these based on the dataset context above - Available tools - Monty sandbox restrictions - Critical patterns for using tool results - When fixing a previous error - Pipeline design — you decide the structure - Response format - Reasoning - Code - Streaming results to a live report - Dataset context - General ML best practices — apply these based on the dataset context above - Available tools - Monty sandbox restrictions - Critical patterns for using tool results - When fixing a previous error - Pipeline design — you decide the structure - Response format - Reasoning - Code - Running it - Why Flyte? - \[Agents > Deep research\](https://www.union.ai/docs/v2/union/tutorials/agents/deep-research/page.md) - Setting up the environment - Generate research queries - Search and summarize - Evaluate research completeness - Filter results - Generate the final answer - Orchestration - Run the deep research agent - Evaluate with Weights & Biases Weave - \[Agents > Competitive intelligence agent\](https://www.union.ai/docs/v2/union/tutorials/agents/competitive-intelligence-agent/page.md) - Setting up the environment - Data types - Search with the You.com Search API - Extract deltas with Claude - Watch one competitor - Orchestration - Run the agent - Create secrets - Run locally or remotely - \[Agents > Compliance monitoring agent\](https://www.union.ai/docs/v2/union/tutorials/agents/compliance-monitoring-agent/page.md) - Setting up the environment - Data types - Research with the You.com Research API - Triage findings with Claude - Monitor one watch item - Orchestration - Run the agent - Create secrets - Run locally or remotely - \[Agents > Field data enrichment agent\](https://www.union.ai/docs/v2/union/tutorials/agents/field-data-enrichment-agent/page.md) - Setting up the environment - Data types - Search with the You.com Search API - Enrich one event - Orchestration - Run the agent - Create secrets - Run locally or remotely - \[Agents > Support resolution agent\](https://www.union.ai/docs/v2/union/tutorials/agents/support-resolution-agent/page.md) - Setting up the environment - Data types - Ground answers with the You.com Research API - Ground one ticket - Draft a customer-ready reply - Resolve one ticket - Orchestration - Run the agent - Create secrets - Run locally or remotely - \[Context Engineering\](https://www.union.ai/docs/v2/union/tutorials/context-engineering/page.md) - \[Automatic prompt engineering\](auto\_prompt\_engineering/page.md) - \[Text-to-SQL\](text\_to\_sql/page.md) - \[Context Engineering > Text-to-SQL prompt optimization\](https://www.union.ai/docs/v2/union/tutorials/context-engineering/text\_to\_sql/page.md) - Ingesting data - From question to SQL - Vector indexing - Table retrieval and context building - SQL generation and response synthesis - Building the QA dataset - Schema extraction and chunking - Question and SQL generation - Validation and quality control - Optimizing prompts - Evaluation pipeline - Iterative optimization - Run it - What we observed - The bigger lesson - \[Context Engineering > Automatic prompt engineering\](https://www.union.ai/docs/v2/union/tutorials/context-engineering/auto\_prompt\_engineering/page.md) - Set up the environment - Prepare the evaluation dataset - Define models - Evaluate prompts - Optimize prompts - Build the full pipeline - Run it - Why this matters - Next steps - \[Model Training\](https://www.union.ai/docs/v2/union/tutorials/model-training/page.md) - \[Hyperparameter optimization\](hpo/page.md) - \[Model Training > Hyperparameter optimization\](https://www.union.ai/docs/v2/union/tutorials/model-training/hpo/page.md) - A better way to run HPO - Declare dependencies - Define the task environment - Define the optimizer - Define the objective function - Define the main optimization loop - Run the experiment - \[Data Processing\](https://www.union.ai/docs/v2/union/tutorials/data-processing/page.md) - \[Batching strategies\](micro-batching/page.md) - \[Data Processing > Batching strategies for efficient scaling\](https://www.union.ai/docs/v2/union/tutorials/data-processing/micro-batching/page.md) - Use Case - Goals - Solution Architecture - 1. Failure transparency with @flyte.trace - 2. Reusable Containers for Efficiency - Key Benefits: - Architecture Flow: - Architecture Diagram - Implementation - Step 0: Set up the runtime - Step 1: Initialize Flyte Configuration - Step 2: Define Container Image - Step 3: Define Task Environments - Step 4: Define External Service Interactions - Step 5: Implement the Batch Processing Task - Step 6: Implement the Orchestrator Workflow - Step 7: Execute the Workflow - Batch Size Selection - Summary --- ## Integrations - \[Anthropic\](https://www.union.ai/docs/v2/union/integrations/anthropic/page.md) - Installation - Quick start - API - \`function\_tool\` - \`Agent\` - \`run\_agent\` - Secrets - API reference - \[BigQuery\](https://www.union.ai/docs/v2/union/integrations/bigquery/page.md) - Installation - Quick start - Configuration - \`BigQueryConfig\` parameters - \`BigQueryTask\` parameters - Authentication - Query templating - Supported input types - Parameterized query example - Retrieving query results - API reference - \[Code generation\](https://www.union.ai/docs/v2/union/integrations/codegen/page.md) - Installation - Quick start - Two execution backends - LiteLLM (default) - Agent (Claude) - Providing data - Sample data - Schema and constraints - Inputs and outputs - Running generated code - One-shot execution with \`result.run()\` - Reusable task with \`result.as\_task()\` - Error diagnosis - Durable execution - Replay logs - Caching - Non-determinism in Agent mode - Observability - LiteLLM backend - Agent backend - Examples - Processing CSVs with different schemas - DataFrame analysis with constraints - Agent mode - Configuration - LiteLLM parameters - Image configuration - Skipping tests - Base packages - Best practices - API reference - \`AutoCoderAgent\` constructor - \`generate()\` parameters - \`CodeGenEvalResult\` fields - \`CodeGenEvalResult\` methods - \[Dask\](https://www.union.ai/docs/v2/union/integrations/dask/page.md) - When to use this plugin - Installation - Configuration - \`Dask\` parameters - \`Scheduler\` parameters - \`WorkerGroup\` parameters - Accessing the Dask client - Example - API reference - \[Databricks\](https://www.union.ai/docs/v2/union/integrations/databricks/page.md) - Installation - Quick start - Configuration - Spark fields (inherited) - Databricks-specific fields - \`databricks\_conf\` structure - Authentication - Accessing the Spark session - API reference - \[Gemini\](https://www.union.ai/docs/v2/union/integrations/gemini/page.md) - Installation - Quick start - API - \`function\_tool\` - \`Agent\` - \`run\_agent\` - Secrets - API reference - \[Hydra\](https://www.union.ai/docs/v2/union/integrations/hydra/page.md) - Installation - Requirements on tasks - A walkthrough config - Execution mode - Hydra launcher (\`@hydra.main\` scripts) - Python SDK - Single run - Grid sweep - Custom sweepers - Forwarding \`flyte.with\_runcontext\` options - Flyte CLI (\`flyte hydra run\`) - Single run - Grid sweep - App-level vs Hydra-namespace overrides - \`--follow\` and \`--no-wait\` - Shell completion - Override grammar - Sweeps - Grid sweeps (BasicSweeper) - Bayesian / TPE sweeps (Optuna) - Sweep output directories - Task environment overrides - Prebuilt images - Applying overrides to child tasks - Renaming the task-env key - What \`task\_env\` should not model - Structured configs (without YAML) - \[MLflow\](https://www.union.ai/docs/v2/union/integrations/mlflow/page.md) - Installation - Quick start - Autologging - Generic autologging - Framework-specific autologging - Run modes - Sharing a run across tasks - Creating independent runs - Nested runs - Workflow-level configuration - Per-task overrides - Configuration priority - Distributed training - MLflow UI links - Setup - Custom URL templates - Explicit links - Link behavior by run mode - Automatic Flyte tags - API reference - \`mlflow\_run\` and \`mlflow\_config\` - \`get\_mlflow\_run\` - \`get\_mlflow\_context\` - \`Mlflow\` - \[OmegaConf\](https://www.union.ai/docs/v2/union/integrations/omegaconf/page.md) - Installation - Quick start - When to use this plugin - Building a DictConfig - From a plain dict - From a YAML file - From a dataclass (structured config) - From a base config plus overrides - Variable interpolation - Nested and deeply structured configs - DictConfigs that contain lists - ListConfig as input and output - Lists of primitives - Building a schedule from another task - Nested lists (list of lists) - Lists of DictConfigs - Lists of dataclass instances - Structured configs - Basic structured config - Schema reconstruction in the receiving task - Required (\`MISSING\`) fields - Advanced field types - Merging overrides on top of a structured base - Embedding rich Python values inside a plain DictConfig - Reserved-looking keys - YAML reports - Wire format - End-to-end example - \[OpenAI\](https://www.union.ai/docs/v2/union/integrations/openai/page.md) - When to use this plugin - Installation - Usage - \`function\_tool\` - Basic pattern - Secrets - Example - API reference - \[OpenAI > Agent tools\](https://www.union.ai/docs/v2/union/integrations/openai/agent\_tools/page.md) - Define the tools - Define the agent - Run the agent - Conclusion - \[Pandera\](https://www.union.ai/docs/v2/union/integrations/pandera/page.md) - When to use this plugin - Installation - pandas - Polars - PySpark SQL - Defining schemas - Using schemas in tasks - Error handling with \`ValidationConfig\` - Image configuration - Pandas - Polars - PySpark SQL - Polars lazy frames - Examples - pandas - Polars - PySpark SQL - \[Papermill\](https://www.union.ai/docs/v2/union/integrations/papermill/page.md) - When to use this plugin - Installation - Quick start - Notebook setup - \`parameters\` cell - \`outputs\` cell - Inputs and outputs - Supported input types - Complex types: File, Dir, DataFrame - Outputs: single, multiple, none - Calling Flyte tasks from notebooks - Workflow patterns - Chaining notebooks - Mixing notebooks with regular tasks - Inline definition - Calling from sync vs. async tasks - Running a NotebookTask directly as the entrypoint - Reports and notebook artifacts - HTML report (default) - Notebook artifacts - Clean reports - Failure reports - Spark notebooks - Local testing - Execution options - \`NotebookTask\` reference - Helper functions - \[PyTorch\](https://www.union.ai/docs/v2/union/integrations/pytorch/page.md) - When to use this plugin - Installation - Configuration - \`Elastic\` parameters - \`RunPolicy\` parameters - NCCL tuning parameters - Writing a distributed training task - Example - API reference - \[Ray\](https://www.union.ai/docs/v2/union/integrations/ray/page.md) - When to use this plugin - Installation - Configuration - \`RayJobConfig\` parameters - \`WorkerNodeConfig\` parameters - \`HeadNodeConfig\` parameters - Connecting to an existing cluster - Examples - API reference - \[Snowflake\](https://www.union.ai/docs/v2/union/integrations/snowflake/page.md) - Installation - Quick start - Configuration - Required fields - Additional connection parameters - Authentication - Key-pair authentication - Password authentication - OAuth authentication - Query templating - Supported input types - Batched \`INSERT\` with list inputs - Parameterized \`SELECT\` - Multiple inputs - Retrieving query results - End-to-end example - \[Spark\](https://www.union.ai/docs/v2/union/integrations/spark/page.md) - When to use this plugin - Installation - Configuration - \`Spark\` parameters - Accessing the Spark session - Overriding configuration at runtime - Example - API reference - \[Weights & Biases\](https://www.union.ai/docs/v2/union/integrations/wandb/page.md) - Installation - Quick start - What's next - \[Weights & Biases > Experiments\](https://www.union.ai/docs/v2/union/integrations/wandb/experiments/page.md) - Basic usage - Accessing the run object - Parent-child task relationships - Run modes - Using \`run\_mode="new"\` for independent runs - Using \`run\_mode="shared"\` for explicit sharing - Configuration with \`wandb\_config\` - Workflow-level configuration - Overriding configuration for child tasks - Using traces with W&B runs - \[Weights & Biases > Distributed training\](https://www.union.ai/docs/v2/union/integrations/wandb/distributed\_training/page.md) - Quick start - Run modes in distributed training - Single-node behavior - Multi-node behavior - Choosing run mode and rank scope - Single-node multi-GPU - Basic example with \`auto\` mode - Using \`shared\` mode for per-rank metrics - Using \`new\` mode for per-rank runs - Multi-node training with \`Elastic\` - Global scope (default): Single run across all nodes - Worker scope: One run per node - Shared mode: All ranks log to the same run - New mode: Separate run per rank - How it works - Run ID patterns - \[Weights & Biases > Sweeps\](https://www.union.ai/docs/v2/union/integrations/wandb/sweeps/page.md) - Creating a sweep - Running parallel agents - Writing objective functions - \[Weights & Biases > Downloading logs\](https://www.union.ai/docs/v2/union/integrations/wandb/downloading\_logs/page.md) - Automatic download - Accessing run directories during execution - \[Weights & Biases > Constraints and best practices\](https://www.union.ai/docs/v2/union/integrations/wandb/constraints\_and\_best\_practices/page.md) - Decorator ordering - Traces cannot use decorators - Maximum sweep agents - Configuration priority - Run ID generation - Sync delay for local files - Shared run mode requirements - Objective functions for sweeps - Error handling - \[Weights & Biases > Manual integration\](https://www.union.ai/docs/v2/union/integrations/wandb/manual/page.md) - Using the Wandb link class - With a custom run ID - Adding links at runtime with override - Using the \`WandbSweep\` link class --- ## Reference - \[LLM-optimized documentation\](https://www.union.ai/docs/v2/union/api-reference/flyte-context/page.md) - Per-page Markdown (\`page.md\`) - Section bundles (\`section.md\`) - Page index (\`llms.txt\`) - Full documentation (\`llms-full.txt\`) - \[Migration from Flyte 1 to Flyte 2\](https://www.union.ai/docs/v2/union/api-reference/migration/page.md) - Key API changes at a glance - Topics - \[Philosophy and imports\](overview/page.md) - \[Container images\](images/page.md) - \[Configuration and CLI\](configuration-and-cli/page.md) - \[Tasks and workflows\](tasks-and-workflows/page.md) - \[Secrets, resources, and caching\](secrets-resources-caching/page.md) - \[Parallelism and async\](parallelism-and-async/page.md) - \[Triggers and dynamic workflows\](triggers-and-dynamic/page.md) - \[Examples and common gotchas\](examples-and-gotchas/page.md) > Section bundle (all pages): https://www.union.ai/docs/v2/union/api-reference/migration/section.md - \[Migration from Flyte 1 to Flyte 2 > Philosophy and imports\](https://www.union.ai/docs/v2/union/api-reference/migration/overview/page.md) - Key paradigm shifts - What Flyte 2 eliminates - What Flyte 2 introduces - Package imports - Basic import changes - Flyte 1 - Flyte 2 - Import mapping table - \[Migration from Flyte 1 to Flyte 2 > Container images\](https://www.union.ai/docs/v2/union/api-reference/migration/images/page.md) - Basic migration - Flyte 1 - Flyte 2 - Image constructor methods - Image builder methods (chainable) - Builder configuration (local vs remote) - Private registry with secrets - Flyte 1 - Flyte 2 - Parameter mapping - \[Migration from Flyte 1 to Flyte 2 > Configuration and CLI\](https://www.union.ai/docs/v2/union/api-reference/migration/configuration-and-cli/page.md) - Configuration files - Config file location - Config format - Flyte 1 - Flyte 2 - Key config differences - Specifying config via CLI - Flyte 1 - Flyte 2 - Specifying config in code - CLI commands - Command mapping - Running tasks - Flyte 1 - Flyte 2 - Key CLI flag differences - Deploying - Flyte 1 - Flyte 2 - Running deployed tasks - Complete Flyte 2 CLI options - \[Migration from Flyte 1 to Flyte 2 > Tasks and workflows\](https://www.union.ai/docs/v2/union/api-reference/migration/tasks-and-workflows/page.md) - Basic task migration - Flyte 1 - Flyte 2 - Workflow to task migration - Flyte 1 - Flyte 2 Sync - Flyte 2 Async - TaskEnvironment configuration - Parameter mapping: @task to TaskEnvironment + @env.task - \[Migration from Flyte 1 to Flyte 2 > Secrets, resources, and caching\](https://www.union.ai/docs/v2/union/api-reference/migration/secrets-resources-caching/page.md) - Secrets - Declaring and accessing secrets - Flyte 1 - Flyte 2 - Secret configuration options - Secret name convention changes - Creating secrets via CLI - Resources - Basic resource configuration - Flyte 1 - Flyte 2 - GPU configuration - Flyte 1 - Flyte 2 - Supported GPU types (Flyte 2) - Resource parameter mapping - Caching - Basic caching - Flyte 1 - Flyte 2 - Cache behavior options (Flyte 2) - \[Migration from Flyte 1 to Flyte 2 > Parallelism and async\](https://www.union.ai/docs/v2/union/api-reference/migration/parallelism-and-async/page.md) - Basic map\_task migration - Flyte 1 - Flyte 2 - map\_task with concurrency - Flyte 1 - Flyte 2 - Async parallel execution with asyncio.gather - Concurrency control with semaphore - Error handling with asyncio.gather - flyte.map vs asyncio.gather comparison - Recommended pattern selection - Sync and async task patterns - Sync tasks calling sync tasks - Async tasks calling async tasks - Sequential execution with await - Flyte 1 - Flyte 2 - \[Migration from Flyte 1 to Flyte 2 > Triggers and dynamic workflows\](https://www.union.ai/docs/v2/union/api-reference/migration/triggers-and-dynamic/page.md) - LaunchPlan to Trigger migration - Flyte 1 - Flyte 2 - Trigger options - Deploying triggers - Dynamic workflows - @dynamic to regular tasks - Flyte 1 - Flyte 2 Sync - Flyte 2 Async - Conditional execution - Flyte 1 - Flyte 2 - Subworkflows to nested tasks - Flyte 1 - Flyte 2 - \[Migration from Flyte 1 to Flyte 2 > Examples and common gotchas\](https://www.union.ai/docs/v2/union/api-reference/migration/examples-and-gotchas/page.md) - Complete migration examples - Example 1: Simple ML pipeline - Flyte 1 - Flyte 2 - Example 2: Parallel processing with map\_task - Flyte 1 - Flyte 2 Sync - Flyte 2 Async - Common gotchas - 1. current\_context() is replaced - 2. Workflow >> ordering notation is gone - 3. flyte.map returns a generator - 4. Image configuration location - 5. Resource configuration - 6. Cache version - 7. Entrypoint task naming - 8. Memory parameter name - 9. Retries have no platform cap - 10. Type annotations - Quick reference cheat sheet - \[Flyte CLI\](https://www.union.ai/docs/v2/union/api-reference/flyte-cli/page.md) - Union-specific functionality {#plugin-commands} - flyte - flyte abort - flyte build - flyte create - flyte delete - flyte deploy - flyte edit - flyte gen - flyte get - flyte prefetch - flyte run - flyte serve - flyte signal - flyte start - flyte stop - flyte update - flyte whoami - \[Flyte SDK\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/page.md) - \[Flyte SDK > Classes\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/classes/page.md) - \[Flyte SDK > Packages\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/page.md) - \[Flyte SDK > Packages > flyte\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/page.md) - Directory - Classes - Protocols - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte > AppHandle\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/apphandle/page.md) - Properties - Methods - activate() - deactivate() - ephemeral\_ctx() - ephemeral\_ctx\_sync() - is\_active() - is\_deactivated() - \[Flyte SDK > Packages > flyte > Backoff\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/backoff/page.md) - Parameters - Methods - compute\_delay() - \[Flyte SDK > Packages > flyte > BaseCheckpoint\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/basecheckpoint/page.md) - Properties - Methods - load() - load\_sync() - prev\_exists() - save() - save\_sync() - \[Flyte SDK > Packages > flyte > Cache\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/cache/page.md) - Parameters - Methods - get\_ignored\_inputs() - get\_version() - is\_enabled() - \[Flyte SDK > Packages > flyte > CachePolicy\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/cachepolicy/page.md) - Methods - get\_version() - \[Flyte SDK > Packages > flyte > Checkpoint\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/checkpoint/page.md) - Parameters - Properties - Methods - load() - load\_sync() - prev\_exists() - save() - save\_sync() - \[Flyte SDK > Packages > flyte > Cron\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/cron/page.md) - Parameters - Properties - \[Flyte SDK > Packages > flyte > Device\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/device/page.md) - Parameters - \[Flyte SDK > Packages > flyte > Environment\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/environment/page.md) - Parameters - Methods - add\_dependency() - clone\_with() - \[Flyte SDK > Packages > flyte > EventWebhook\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/eventwebhook/page.md) - Parameters - \[Flyte SDK > Packages > flyte > FixedRate\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/fixedrate/page.md) - Parameters - \[Flyte SDK > Packages > flyte > Image\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/image/page.md) - Parameters - Properties - Methods - clone() - from\_base() - from\_debian\_base() - from\_dockerfile() - from\_ref\_name() - from\_uv\_script() - validate() - with\_apt\_packages() - with\_code\_bundle() - with\_commands() - with\_dockerignore() - with\_env\_vars() - with\_local\_rs\_controller() - with\_local\_v2() - with\_local\_v2\_plugins() - with\_pip\_packages() - with\_poetry\_project() - with\_requirements() - with\_source\_file() - with\_source\_folder() - with\_uv\_project() - with\_workdir() - \[Flyte SDK > Packages > flyte > ImageBuild\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/imagebuild/page.md) - Parameters - \[Flyte SDK > Packages > flyte > Link\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/link/page.md) - Methods - get\_link() - \[Flyte SDK > Packages > flyte > PodTemplate\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/podtemplate/page.md) - Parameters - Methods - to\_k8s\_pod() - \[Flyte SDK > Packages > flyte > Resources\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/resources/page.md) - Parameters - Methods - get\_device() - get\_shared\_memory() - \[Flyte SDK > Packages > flyte > RetryStrategy\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/retrystrategy/page.md) - Parameters - \[Flyte SDK > Packages > flyte > ReusePolicy\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/reusepolicy/page.md) - Parameters - Properties - Methods - get\_scaledown\_ttl() - \[Flyte SDK > Packages > flyte > Secret\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/secret/page.md) - Parameters - Methods - stable\_hash() - \[Flyte SDK > Packages > flyte > TaskEnvironment\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/taskenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - from\_task() - task() - \[Flyte SDK > Packages > flyte > Timeout\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/timeout/page.md) - Parameters - \[Flyte SDK > Packages > flyte > Trigger\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte/trigger/page.md) - Parameters - Methods - daily() - hourly() - minutely() - monthly() - weekly() - \[Flyte SDK > Packages > flyte.ai.agents\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/page.md) - Directory - Classes - Protocols - Errors - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.ai.agents > AccessDenied\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/accessdenied/page.md) - \[Flyte SDK > Packages > flyte.ai.agents > Agent\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/agent/page.md) - Parameters - Properties - Methods - add\_tool() - approval\_callback() - call\_llm() - run() - tool\_descriptions() - \[Flyte SDK > Packages > flyte.ai.agents > AgentEvent\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/agentevent/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > AgentProtocol\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/agentprotocol/page.md) - Methods - run() - tool\_descriptions() - \[Flyte SDK > Packages > flyte.ai.agents > AgentResult\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/agentresult/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > AgentTool\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/agenttool/page.md) - Parameters - Methods - to\_openai\_format() - \[Flyte SDK > Packages > flyte.ai.agents > CodeModeAgent\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/codemodeagent/page.md) - Parameters - Methods - run() - tool\_descriptions() - uses\_flyte\_tools() - \[Flyte SDK > Packages > flyte.ai.agents > ConcurrencyError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/concurrencyerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > LLMMessage\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/llmmessage/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > MCPServerSpec\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/mcpserverspec/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > MemoryMeta\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/memorymeta/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents > MemoryStore\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/memorystore/page.md) - Parameters - Methods - append() - audit\_tail() - audit\_tail\_sync() - create() - current\_sha() - exists() - extend() - flush\_messages() - flush\_messages\_sync() - get\_meta() - get\_or\_create() - list\_paths() - read\_json() - read\_text() - remote\_path\_for\_key() - save() - write\_json() - write\_text() - \[Flyte SDK > Packages > flyte.ai.agents > MemoryStoreError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents/memorystoreerror/page.md) - \[Flyte SDK > Packages > flyte.ai.agents.agent\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.agent/page.md) - Directory - Classes - Variables - \[Flyte SDK > Packages > flyte.ai.agents.agent > Agent\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.agent/agent/page.md) - Parameters - Properties - Methods - add\_tool() - approval\_callback() - call\_llm() - run() - tool\_descriptions() - \[Flyte SDK > Packages > flyte.ai.agents.agent > AgentEvent\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.agent/agentevent/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents.codemode\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.codemode/page.md) - Directory - Classes - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.ai.agents.codemode > CodeModeAgent\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.codemode/codemodeagent/page.md) - Parameters - Methods - run() - tool\_descriptions() - uses\_flyte\_tools() - \[Flyte SDK > Packages > flyte.ai.agents.memory\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/page.md) - Directory - Classes - Errors - Variables - \[Flyte SDK > Packages > flyte.ai.agents.memory > AccessDenied\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/accessdenied/page.md) - \[Flyte SDK > Packages > flyte.ai.agents.memory > ConcurrencyError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/concurrencyerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents.memory > MemoryMeta\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/memorymeta/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.agents.memory > MemoryStore\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/memorystore/page.md) - Parameters - Methods - append() - audit\_tail() - audit\_tail\_sync() - create() - current\_sha() - exists() - extend() - flush\_messages() - flush\_messages\_sync() - get\_meta() - get\_or\_create() - list\_paths() - read\_json() - read\_text() - remote\_path\_for\_key() - save() - write\_json() - write\_text() - \[Flyte SDK > Packages > flyte.ai.agents.memory > MemoryStoreError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.memory/memorystoreerror/page.md) - \[Flyte SDK > Packages > flyte.ai.agents.protocol\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.protocol/page.md) - Directory - Classes - Protocols - \[Flyte SDK > Packages > flyte.ai.agents.protocol > AgentProtocol\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.protocol/agentprotocol/page.md) - Methods - run() - tool\_descriptions() - \[Flyte SDK > Packages > flyte.ai.agents.protocol > AgentResult\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.agents.protocol/agentresult/page.md) - Parameters - \[Flyte SDK > Packages > flyte.ai.chat\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.chat/page.md) - Directory - Classes - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.ai.chat > AgentChatAppEnvironment\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.chat/agentchatappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - build\_fastapi\_app() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.ai.chat > CustomTheme\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.chat/customtheme/page.md) - Parameters - Methods - to\_css() - \[Flyte SDK > Packages > flyte.ai.chat.app\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.chat.app/page.md) - Directory - Classes - Variables - \[Flyte SDK > Packages > flyte.ai.chat.app > AgentChatAppEnvironment\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.chat.app/agentchatappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - build\_fastapi\_app() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.ai.chat.app > CustomTheme\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.chat.app/customtheme/page.md) - Parameters - Methods - to\_css() - \[Flyte SDK > Packages > flyte.ai.mcp\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.mcp/page.md) - Directory - Classes - \[Flyte SDK > Packages > flyte.ai.mcp > FlyteMCPAppEnvironment\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.mcp/flytemcpappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.ai.mcp > MCPAppEnvironment\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.ai.mcp/mcpappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.app\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.app > AppEndpoint\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/appendpoint/page.md) - Parameters - Methods - check\_type() - get() - materialize() - \[Flyte SDK > Packages > flyte.app > AppEnvironment\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/appenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.app > ConnectorEnvironment\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/connectorenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.app > Domain\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/domain/page.md) - Parameters - \[Flyte SDK > Packages > flyte.app > Link\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/link/page.md) - Parameters - \[Flyte SDK > Packages > flyte.app > Parameter\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/parameter/page.md) - Parameters - \[Flyte SDK > Packages > flyte.app > Port\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/port/page.md) - Parameters - \[Flyte SDK > Packages > flyte.app > RunOutput\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/runoutput/page.md) - Parameters - Methods - check\_type() - get() - materialize() - \[Flyte SDK > Packages > flyte.app > Scaling\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/scaling/page.md) - Parameters - Methods - get\_replicas() - \[Flyte SDK > Packages > flyte.app > Timeouts\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app/timeouts/page.md) - Parameters - \[Flyte SDK > Packages > flyte.app.extras\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app.extras/page.md) - Directory - Classes - \[Flyte SDK > Packages > flyte.app.extras > FastAPIAppEnvironment\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app.extras/fastapiappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.app.extras > FastAPIPassthroughAuthMiddleware\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app.extras/fastapipassthroughauthmiddleware/page.md) - Parameters - Methods - dispatch() - extract\_authorization\_header() - extract\_cookie\_header() - extract\_custom\_header() - \[Flyte SDK > Packages > flyte.app.extras > FlyteWebhookAppEnvironment\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.app.extras/flytewebhookappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - container\_command() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Flyte SDK > Packages > flyte.config\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.config/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.config > Config\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.config/config/page.md) - Parameters - Methods - auto() - with\_params() - \[Flyte SDK > Packages > flyte.connectors\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.connectors/page.md) - Directory - Classes - \[Flyte SDK > Packages > flyte.connectors > AsyncConnector\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.connectors/asyncconnector/page.md) - Methods - create() - delete() - get() - get\_logs() - get\_metrics() - \[Flyte SDK > Packages > flyte.connectors > AsyncConnectorExecutorMixin\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.connectors/asyncconnectorexecutormixin/page.md) - Methods - execute() - \[Flyte SDK > Packages > flyte.connectors > ConnectorRegistry\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.connectors/connectorregistry/page.md) - Methods - get\_connector() - register() - \[Flyte SDK > Packages > flyte.connectors > ConnectorService\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.connectors/connectorservice/page.md) - Methods - run() - \[Flyte SDK > Packages > flyte.connectors > Resource\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.connectors/resource/page.md) - Parameters - \[Flyte SDK > Packages > flyte.connectors > ResourceMeta\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.connectors/resourcemeta/page.md) - Parameters - Methods - decode() - encode() - \[Flyte SDK > Packages > flyte.connectors.utils\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.connectors.utils/page.md) - Directory - Methods - Methods - \[Flyte SDK > Packages > flyte.durable\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.durable/page.md) - Directory - Methods - Methods - \[Flyte SDK > Packages > flyte.errors\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/page.md) - Directory - Errors - Methods - Methods - \[Flyte SDK > Packages > flyte.errors > ActionAbortedError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/actionabortederror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > ActionNotFoundError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/actionnotfounderror/page.md) - \[Flyte SDK > Packages > flyte.errors > BaseRuntimeError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/baseruntimeerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > CodeBundleError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/codebundleerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > CustomError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/customerror/page.md) - Parameters - Methods - from\_exception() - \[Flyte SDK > Packages > flyte.errors > DeploymentError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/deploymenterror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > EventAlreadyExistsError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/eventalreadyexistserror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > EventFailedError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/eventfailederror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > EventNotFoundError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/eventnotfounderror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > EventTimedoutError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/eventtimedouterror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > ImageBuildError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/imagebuilderror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > ImagePullBackOffError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/imagepullbackofferror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > InitializationError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/initializationerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > InlineIOMaxBytesBreached\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/inlineiomaxbytesbreached/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > InvalidImageNameError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/invalidimagenameerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > InvalidPackageError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/invalidpackageerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > LogsNotYetAvailableError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/logsnotyetavailableerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > ModuleLoadError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/moduleloaderror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > NonRecoverableError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/nonrecoverableerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > NotInTaskContextError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/notintaskcontexterror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > OnlyAsyncIOSupportedError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/onlyasynciosupportederror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > OOMError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/oomerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > ParameterMaterializationError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/parametermaterializationerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > PrimaryContainerNotFoundError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/primarycontainernotfounderror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RemoteTaskNotFoundError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/remotetasknotfounderror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RemoteTaskUsageError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/remotetaskusageerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RestrictedTypeError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/restrictedtypeerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RetriesExhaustedError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/retriesexhaustederror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RuntimeDataValidationError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/runtimedatavalidationerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RuntimeSystemError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/runtimesystemerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RuntimeUnknownError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/runtimeunknownerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > RuntimeUserError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/runtimeusererror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > SlowDownError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/slowdownerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > TaskInterruptedError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/taskinterruptederror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > TaskTimeoutError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/tasktimeouterror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > TraceDoesNotAllowNestedTasksError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/tracedoesnotallownestedtaskserror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.errors > UnionRpcError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.errors/unionrpcerror/page.md) - Parameters - \[Flyte SDK > Packages > flyte.extend\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extend/page.md) - Directory - Classes - Protocols - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.extend > AsyncFunctionTaskTemplate\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extend/asyncfunctiontasktemplate/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.extend > ImageBuildEngine\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extend/imagebuildengine/page.md) - \[Flyte SDK > Packages > flyte.extend > ImageBuilder\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extend/imagebuilder/page.md) - Methods - build\_image() - get\_checkers() - \[Flyte SDK > Packages > flyte.extend > ImageChecker\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extend/imagechecker/page.md) - Methods - image\_exists() - \[Flyte SDK > Packages > flyte.extend > TaskTemplate\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extend/tasktemplate/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.extras\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras/page.md) - Directory - Classes - Protocols - \[Flyte SDK > Packages > flyte.extras > BatchStats\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras/batchstats/page.md) - Parameters - Properties - \[Flyte SDK > Packages > flyte.extras > ContainerTask\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras/containertask/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.extras > CostEstimator\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras/costestimator/page.md) - Methods - estimate\_cost() - \[Flyte SDK > Packages > flyte.extras > DynamicBatcher\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras/dynamicbatcher/page.md) - Parameters - Properties - Methods - start() - stop() - submit() - submit\_batch() - \[Flyte SDK > Packages > flyte.extras > Prompt\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras/prompt/page.md) - Parameters - Methods - estimate\_tokens() - \[Flyte SDK > Packages > flyte.extras > Sleep\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras/sleep/page.md) - Parameters - \[Flyte SDK > Packages > flyte.extras > SleepTask\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras/sleeptask/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.extras > TokenBatcher\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras/tokenbatcher/page.md) - Parameters - Properties - Methods - start() - stop() - submit() - submit\_batch() - \[Flyte SDK > Packages > flyte.extras > TokenEstimator\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras/tokenestimator/page.md) - Methods - estimate\_tokens() - \[Flyte SDK > Packages > flyte.extras.shell\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras.shell/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.extras.shell > FlagSpec\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras.shell/flagspec/page.md) - Parameters - Methods - coerce() - \[Flyte SDK > Packages > flyte.extras.shell > Glob\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras.shell/glob/page.md) - Parameters - \[Flyte SDK > Packages > flyte.extras.shell > Stderr\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras.shell/stderr/page.md) - Parameters - \[Flyte SDK > Packages > flyte.extras.shell > Stdout\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.extras.shell/stdout/page.md) - Parameters - \[Flyte SDK > Packages > flyte.git\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.git/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.git > GitStatus\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.git/gitstatus/page.md) - Parameters - Methods - build\_url() - from\_current\_repo() - \[Flyte SDK > Packages > flyte.io\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.io/page.md) - IO data types - Directory - Classes - Variables - \[Flyte SDK > Packages > flyte.io > DataFrame\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.io/dataframe/page.md) - Parameters - Properties - Methods - all() - all\_sync() - column\_names() - columns() - deserialize\_dataframe() - from\_df() - from\_existing\_remote() - from\_local() - from\_local\_sync() - iter() - model\_post\_init() - open() - schema\_match() - serialize\_dataframe() - set\_literal() - wrap\_df() - \[Flyte SDK > Packages > flyte.io > Dir\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.io/dir/page.md) - Parameters - Properties - Methods - download() - download\_sync() - empty() - exists() - exists\_sync() - from\_existing\_remote() - from\_local() - from\_local\_sync() - get\_file() - get\_file\_sync() - list\_files() - list\_files\_sync() - model\_post\_init() - new\_remote() - pre\_init() - schema\_match() - walk() - walk\_sync() - \[Flyte SDK > Packages > flyte.io > EmptyDir\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.io/emptydir/page.md) - Parameters - Properties - Methods - download() - download\_sync() - empty() - exists() - exists\_sync() - from\_existing\_remote() - from\_local() - from\_local\_sync() - get\_file() - get\_file\_sync() - list\_files() - list\_files\_sync() - model\_post\_init() - new\_remote() - pre\_init() - schema\_match() - walk() - walk\_sync() - \[Flyte SDK > Packages > flyte.io > File\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.io/file/page.md) - Parameters - Properties - Methods - download() - download\_sync() - exists() - exists\_sync() - from\_existing\_remote() - from\_local() - from\_local\_sync() - model\_post\_init() - named\_remote() - new\_remote() - open() - open\_sync() - pre\_init() - schema\_match() - \[Flyte SDK > Packages > flyte.io > HashFunction\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.io/hashfunction/page.md) - Parameters - Methods - from\_fn() - reset() - result() - update() - \[Flyte SDK > Packages > flyte.io.extend\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.io.extend/page.md) - Directory - Classes - \[Flyte SDK > Packages > flyte.io.extend > DataFrameDecoder\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.io.extend/dataframedecoder/page.md) - Parameters - Properties - Methods - decode() - \[Flyte SDK > Packages > flyte.io.extend > DataFrameEncoder\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.io.extend/dataframeencoder/page.md) - Parameters - Properties - Methods - encode() - \[Flyte SDK > Packages > flyte.io.extend > DataFrameTransformerEngine\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.io.extend/dataframetransformerengine/page.md) - Parameters - Properties - Methods - assert\_type() - encode() - from\_binary\_idl() - get\_decoder() - get\_encoder() - get\_literal\_type() - get\_structured\_dataset\_type() - guess\_python\_type() - isinstance\_generic() - iter\_as() - open\_as() - register() - register\_for\_protocol() - register\_renderer() - schema\_match() - to\_html() - to\_literal() - to\_python\_value() - \[Flyte SDK > Packages > flyte.models\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/page.md) - Directory - Classes - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.models > ActionID\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/actionid/page.md) - Parameters - Methods - create\_random() - new\_sub\_action() - new\_sub\_action\_from() - unique\_id\_str() - \[Flyte SDK > Packages > flyte.models > ActionPhase\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/actionphase/page.md) - Parameters - \[Flyte SDK > Packages > flyte.models > CheckpointPaths\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/checkpointpaths/page.md) - Parameters - \[Flyte SDK > Packages > flyte.models > CodeBundle\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/codebundle/page.md) - Parameters - Methods - with\_downloaded\_path() - \[Flyte SDK > Packages > flyte.models > GroupData\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/groupdata/page.md) - Parameters - \[Flyte SDK > Packages > flyte.models > NativeInterface\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/nativeinterface/page.md) - Parameters - Properties - Methods - convert\_to\_kwargs() - from\_callable() - from\_types() - get\_input\_types() - has\_outputs() - num\_required\_inputs() - required\_inputs() - \[Flyte SDK > Packages > flyte.models > PathRewrite\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/pathrewrite/page.md) - Parameters - Methods - from\_str() - \[Flyte SDK > Packages > flyte.models > RawDataPath\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/rawdatapath/page.md) - Parameters - Methods - from\_local\_folder() - get\_random\_remote\_path() - \[Flyte SDK > Packages > flyte.models > SerializationContext\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/serializationcontext/page.md) - Parameters - Methods - get\_entrypoint\_path() - \[Flyte SDK > Packages > flyte.models > TaskContext\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.models/taskcontext/page.md) - Parameters - Properties - Methods - is\_in\_cluster() - replace() - \[Flyte SDK > Packages > flyte.notify\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.notify/page.md) - Directory - Classes - \[Flyte SDK > Packages > flyte.notify > Email\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.notify/email/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > NamedDelivery\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.notify/nameddelivery/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > NamedRule\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.notify/namedrule/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > Notification\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.notify/notification/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > Slack\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.notify/slack/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > Teams\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.notify/teams/page.md) - Parameters - \[Flyte SDK > Packages > flyte.notify > Webhook\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.notify/webhook/page.md) - Parameters - \[Flyte SDK > Packages > flyte.prefetch\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.prefetch/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.prefetch > HuggingFaceModelInfo\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.prefetch/huggingfacemodelinfo/page.md) - Parameters - \[Flyte SDK > Packages > flyte.prefetch > ShardConfig\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.prefetch/shardconfig/page.md) - Parameters - \[Flyte SDK > Packages > flyte.prefetch > StoredModelInfo\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.prefetch/storedmodelinfo/page.md) - Parameters - \[Flyte SDK > Packages > flyte.prefetch > VLLMShardArgs\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.prefetch/vllmshardargs/page.md) - Parameters - Methods - get\_vllm\_args() - \[Flyte SDK > Packages > flyte.remote\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.remote > Action\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/action/page.md) - Parameters - Properties - Methods - abort() - details() - done() - get() - get\_logs() - listall() - show\_logs() - sync() - to\_dict() - to\_json() - wait() - watch() - \[Flyte SDK > Packages > flyte.remote > ActionDetails\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/actiondetails/page.md) - Parameters - Properties - Methods - done() - get() - get\_details() - get\_phase\_transitions() - inputs() - logs\_available() - outputs() - to\_dict() - to\_json() - watch() - watch\_updates() - \[Flyte SDK > Packages > flyte.remote > ActionInputs\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/actioninputs/page.md) - Parameters - Methods - clear() - copy() - fromkeys() - get() - items() - keys() - pop() - popitem() - setdefault() - to\_dict() - to\_json() - update() - values() - \[Flyte SDK > Packages > flyte.remote > ActionOutputs\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/actionoutputs/page.md) - Parameters - Properties - Methods - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > App\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/app/page.md) - Parameters - Properties - Methods - activate() - create() - deactivate() - delete() - ephemeral\_ctx() - ephemeral\_ctx\_sync() - get() - is\_active() - is\_deactivated() - listall() - replace() - to\_dict() - to\_json() - update() - watch() - \[Flyte SDK > Packages > flyte.remote > Event\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/event/page.md) - Parameters - Properties - Methods - get() - listall() - signal() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > Project\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/project/page.md) - Parameters - Methods - archive() - create() - get() - listall() - to\_dict() - to\_json() - unarchive() - update() - \[Flyte SDK > Packages > flyte.remote > Run\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/run/page.md) - Parameters - Properties - Methods - abort() - details() - done() - get() - get\_debug\_url() - get\_logs() - inputs() - listall() - outputs() - show\_logs() - sync() - to\_dict() - to\_json() - wait() - watch() - \[Flyte SDK > Packages > flyte.remote > RunDetails\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/rundetails/page.md) - Parameters - Properties - Methods - done() - get() - get\_details() - inputs() - outputs() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > Secret\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/secret/page.md) - Parameters - Properties - Methods - create() - delete() - get() - listall() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > Settings\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/settings/page.md) - Parameters - Methods - available\_keys() - effective\_values() - get\_settings\_for\_edit() - local\_overrides() - parse\_yaml() - scope\_description() - to\_dict() - to\_json() - to\_yaml() - to\_yaml\_sections() - update\_settings() - \[Flyte SDK > Packages > flyte.remote > Task\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/task/page.md) - Parameters - Properties - Methods - get() - listall() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > TaskDetails\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/taskdetails/page.md) - Parameters - Properties - Methods - fetch() - get() - override() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.remote > TimeFilter\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/timefilter/page.md) - Parameters - \[Flyte SDK > Packages > flyte.remote > Trigger\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/trigger/page.md) - Parameters - Properties - Methods - create() - delete() - get() - get\_details() - listall() - to\_dict() - to\_json() - update() - \[Flyte SDK > Packages > flyte.remote > User\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.remote/user/page.md) - Parameters - Methods - get() - name() - subject() - to\_dict() - to\_json() - \[Flyte SDK > Packages > flyte.report\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.report/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.report > Report\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.report/report/page.md) - Parameters - Methods - get\_final\_report() - get\_tab() - \[Flyte SDK > Packages > flyte.sandbox\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.sandbox/page.md) - Directory - Classes - Methods - Variables - Methods - \[Flyte SDK > Packages > flyte.sandbox > CodeTaskTemplate\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.sandbox/codetasktemplate/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.sandbox > ImageConfig\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.sandbox/imageconfig/page.md) - Parameters - \[Flyte SDK > Packages > flyte.sandbox > SandboxedConfig\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.sandbox/sandboxedconfig/page.md) - Parameters - \[Flyte SDK > Packages > flyte.sandbox > SandboxedTaskTemplate\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.sandbox/sandboxedtasktemplate/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Flyte SDK > Packages > flyte.storage\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.storage/page.md) - Directory - Classes - Methods - Methods - \[Flyte SDK > Packages > flyte.storage > ABFS\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.storage/abfs/page.md) - Parameters - Methods - auto() - get\_fsspec\_kwargs() - \[Flyte SDK > Packages > flyte.storage > GCS\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.storage/gcs/page.md) - Parameters - Methods - auto() - get\_fsspec\_kwargs() - \[Flyte SDK > Packages > flyte.storage > S3\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.storage/s3/page.md) - Parameters - Methods - auto() - for\_sandbox() - get\_fsspec\_kwargs() - \[Flyte SDK > Packages > flyte.storage > Storage\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.storage/storage/page.md) - Parameters - Methods - auto() - get\_fsspec\_kwargs() - \[Flyte SDK > Packages > flyte.syncify\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.syncify/page.md) - Creating a Syncify Instance - How does it work? - Directory - Classes - \[Flyte SDK > Packages > flyte.syncify > Syncify\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.syncify/syncify/page.md) - Parameters - \[Flyte SDK > Packages > flyte.types\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.types/page.md) - Directory - Classes - Protocols - Errors - Methods - Methods - \[Flyte SDK > Packages > flyte.types > FlytePickle\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.types/flytepickle/page.md) - Methods - from\_pickle() - python\_type() - to\_pickle() - \[Flyte SDK > Packages > flyte.types > Renderable\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.types/renderable/page.md) - Methods - to\_html() - \[Flyte SDK > Packages > flyte.types > TypeEngine\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.types/typeengine/page.md) - Methods - dict\_to\_literal\_map() - get\_available\_transformers() - get\_transformer() - guess\_python\_type() - guess\_python\_types() - lazy\_import\_transformers() - literal\_map\_to\_kwargs() - named\_tuple\_to\_variable\_map() - register() - register\_additional\_type() - register\_restricted\_type() - to\_html() - to\_literal() - to\_literal\_checks() - to\_literal\_type() - to\_python\_value() - unwrap\_offloaded\_literal() - \[Flyte SDK > Packages > flyte.types > TypeTransformer\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.types/typetransformer/page.md) - Parameters - Properties - Methods - assert\_type() - from\_binary\_idl() - get\_literal\_type() - guess\_python\_type() - isinstance\_generic() - schema\_match() - to\_html() - to\_literal() - to\_python\_value() - \[Flyte SDK > Packages > flyte.types > TypeTransformerFailedError\](https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/packages/flyte.types/typetransformerfailederror/page.md) - \[Integrations\](https://www.union.ai/docs/v2/union/api-reference/integrations/page.md) - \[Integrations > Anthropic\](https://www.union.ai/docs/v2/union/api-reference/integrations/anthropic/page.md) - \[Integrations > Anthropic > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/anthropic/classes/page.md) - \[Integrations > Anthropic > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/anthropic/packages/page.md) - \[Integrations > Anthropic > Packages > flyteplugins.anthropic\](https://www.union.ai/docs/v2/union/api-reference/integrations/anthropic/packages/flyteplugins.anthropic/page.md) - Directory - Classes - Methods - Methods - \[Integrations > Anthropic > Packages > flyteplugins.anthropic > Agent\](https://www.union.ai/docs/v2/union/api-reference/integrations/anthropic/packages/flyteplugins.anthropic/agent/page.md) - Parameters - Methods - get\_anthropic\_tools() - \[Integrations > BigQuery\](https://www.union.ai/docs/v2/union/api-reference/integrations/bigquery/page.md) - \[Integrations > BigQuery > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/bigquery/classes/page.md) - \[Integrations > BigQuery > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/bigquery/packages/page.md) - \[Integrations > BigQuery > Packages > flyteplugins.bigquery\](https://www.union.ai/docs/v2/union/api-reference/integrations/bigquery/packages/flyteplugins.bigquery/page.md) - Directory - Classes - \[Integrations > BigQuery > Packages > flyteplugins.bigquery > BigQueryConfig\](https://www.union.ai/docs/v2/union/api-reference/integrations/bigquery/packages/flyteplugins.bigquery/bigqueryconfig/page.md) - Parameters - \[Integrations > BigQuery > Packages > flyteplugins.bigquery > BigQueryConnector\](https://www.union.ai/docs/v2/union/api-reference/integrations/bigquery/packages/flyteplugins.bigquery/bigqueryconnector/page.md) - Methods - create() - delete() - get() - get\_logs() - get\_metrics() - \[Integrations > BigQuery > Packages > flyteplugins.bigquery > BigQueryTask\](https://www.union.ai/docs/v2/union/api-reference/integrations/bigquery/packages/flyteplugins.bigquery/bigquerytask/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Integrations > Code generation\](https://www.union.ai/docs/v2/union/api-reference/integrations/codegen/page.md) - \[Integrations > Code generation > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/codegen/classes/page.md) - \[Integrations > Code generation > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/codegen/packages/page.md) - \[Integrations > Code generation > Packages > flyteplugins.codegen\](https://www.union.ai/docs/v2/union/api-reference/integrations/codegen/packages/flyteplugins.codegen/page.md) - Directory - Classes - \[Integrations > Code generation > Packages > flyteplugins.codegen > AutoCoderAgent\](https://www.union.ai/docs/v2/union/api-reference/integrations/codegen/packages/flyteplugins.codegen/autocoderagent/page.md) - Parameters - Methods - generate() - \[Integrations > Code generation > Packages > flyteplugins.codegen > CodeGenEvalResult\](https://www.union.ai/docs/v2/union/api-reference/integrations/codegen/packages/flyteplugins.codegen/codegenevalresult/page.md) - Parameters - Methods - as\_task() - run() - \[Integrations > Code generation > Packages > flyteplugins.codegen > CodePlan\](https://www.union.ai/docs/v2/union/api-reference/integrations/codegen/packages/flyteplugins.codegen/codeplan/page.md) - Parameters - \[Integrations > Code generation > Packages > flyteplugins.codegen > CodeSolution\](https://www.union.ai/docs/v2/union/api-reference/integrations/codegen/packages/flyteplugins.codegen/codesolution/page.md) - Parameters - Methods - normalize\_language() - \[Integrations > Code generation > Packages > flyteplugins.codegen > ErrorDiagnosis\](https://www.union.ai/docs/v2/union/api-reference/integrations/codegen/packages/flyteplugins.codegen/errordiagnosis/page.md) - Parameters - \[Integrations > Code generation > Packages > flyteplugins.codegen > ImageConfig\](https://www.union.ai/docs/v2/union/api-reference/integrations/codegen/packages/flyteplugins.codegen/imageconfig/page.md) - Parameters - \[Integrations > Dask\](https://www.union.ai/docs/v2/union/api-reference/integrations/dask/page.md) - \[Integrations > Dask > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/dask/classes/page.md) - \[Integrations > Dask > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/dask/packages/page.md) - \[Integrations > Dask > Packages > flyteplugins.dask\](https://www.union.ai/docs/v2/union/api-reference/integrations/dask/packages/flyteplugins.dask/page.md) - Directory - Classes - \[Integrations > Dask > Packages > flyteplugins.dask > Dask\](https://www.union.ai/docs/v2/union/api-reference/integrations/dask/packages/flyteplugins.dask/dask/page.md) - Parameters - \[Integrations > Dask > Packages > flyteplugins.dask > Scheduler\](https://www.union.ai/docs/v2/union/api-reference/integrations/dask/packages/flyteplugins.dask/scheduler/page.md) - Parameters - \[Integrations > Dask > Packages > flyteplugins.dask > WorkerGroup\](https://www.union.ai/docs/v2/union/api-reference/integrations/dask/packages/flyteplugins.dask/workergroup/page.md) - Parameters - \[Integrations > Databricks\](https://www.union.ai/docs/v2/union/api-reference/integrations/databricks/page.md) - \[Integrations > Databricks > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/databricks/classes/page.md) - \[Integrations > Databricks > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/databricks/packages/page.md) - \[Integrations > Databricks > Packages > flyteplugins.databricks\](https://www.union.ai/docs/v2/union/api-reference/integrations/databricks/packages/flyteplugins.databricks/page.md) - Directory - Classes - \[Integrations > Databricks > Packages > flyteplugins.databricks > Databricks\](https://www.union.ai/docs/v2/union/api-reference/integrations/databricks/packages/flyteplugins.databricks/databricks/page.md) - Parameters - \[Integrations > Databricks > Packages > flyteplugins.databricks > DatabricksConnector\](https://www.union.ai/docs/v2/union/api-reference/integrations/databricks/packages/flyteplugins.databricks/databricksconnector/page.md) - Methods - create() - delete() - get() - get\_logs() - get\_metrics() - \[Integrations > Gemini\](https://www.union.ai/docs/v2/union/api-reference/integrations/gemini/page.md) - \[Integrations > Gemini > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/gemini/classes/page.md) - \[Integrations > Gemini > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/gemini/packages/page.md) - \[Integrations > Gemini > Packages > flyteplugins.gemini\](https://www.union.ai/docs/v2/union/api-reference/integrations/gemini/packages/flyteplugins.gemini/page.md) - Directory - Classes - Methods - Methods - \[Integrations > Gemini > Packages > flyteplugins.gemini > Agent\](https://www.union.ai/docs/v2/union/api-reference/integrations/gemini/packages/flyteplugins.gemini/agent/page.md) - Parameters - Methods - get\_gemini\_tools() - \[Integrations > Human-in-the-Loop\](https://www.union.ai/docs/v2/union/api-reference/integrations/hitl/page.md) - \[Integrations > Human-in-the-Loop > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/hitl/classes/page.md) - \[Integrations > Human-in-the-Loop > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/hitl/packages/page.md) - \[Integrations > Human-in-the-Loop > Packages > flyteplugins.hitl\](https://www.union.ai/docs/v2/union/api-reference/integrations/hitl/packages/flyteplugins.hitl/page.md) - Basic usage: - Features: - Directory - Classes - Methods - Variables - Methods - \[Integrations > Human-in-the-Loop > Packages > flyteplugins.hitl > Event\](https://www.union.ai/docs/v2/union/api-reference/integrations/hitl/packages/flyteplugins.hitl/event/page.md) - Parameters - Properties - Methods - create() - wait() - \[Integrations > Hydra\](https://www.union.ai/docs/v2/union/api-reference/integrations/hydra/page.md) - \[Integrations > JSONL\](https://www.union.ai/docs/v2/union/api-reference/integrations/jsonl/page.md) - \[Integrations > JSONL > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/jsonl/classes/page.md) - \[Integrations > JSONL > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/jsonl/packages/page.md) - \[Integrations > JSONL > Packages > flyteplugins.jsonl\](https://www.union.ai/docs/v2/union/api-reference/integrations/jsonl/packages/flyteplugins.jsonl/page.md) - Directory - Classes - \[Integrations > JSONL > Packages > flyteplugins.jsonl > JsonlDir\](https://www.union.ai/docs/v2/union/api-reference/integrations/jsonl/packages/flyteplugins.jsonl/jsonldir/page.md) - Parameters - Properties - Methods - download() - download\_sync() - empty() - exists() - exists\_sync() - from\_existing\_remote() - from\_local() - from\_local\_sync() - get\_file() - get\_file\_sync() - iter\_arrow\_batches() - iter\_arrow\_batches\_sync() - iter\_batches() - iter\_batches\_sync() - iter\_records() - iter\_records\_sync() - list\_files() - list\_files\_sync() - model\_post\_init() - new\_remote() - pre\_init() - schema\_match() - walk() - walk\_sync() - writer() - writer\_sync() - \[Integrations > JSONL > Packages > flyteplugins.jsonl > JsonlFile\](https://www.union.ai/docs/v2/union/api-reference/integrations/jsonl/packages/flyteplugins.jsonl/jsonlfile/page.md) - Parameters - Properties - Methods - download() - download\_sync() - exists() - exists\_sync() - from\_existing\_remote() - from\_local() - from\_local\_sync() - iter\_arrow\_batches() - iter\_arrow\_batches\_sync() - iter\_records() - iter\_records\_sync() - model\_post\_init() - named\_remote() - new\_remote() - open() - open\_sync() - pre\_init() - schema\_match() - writer() - writer\_sync() - \[Integrations > MLflow\](https://www.union.ai/docs/v2/union/api-reference/integrations/mlflow/page.md) - \[Integrations > MLflow > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/mlflow/classes/page.md) - \[Integrations > MLflow > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/mlflow/packages/page.md) - \[Integrations > MLflow > Packages > flyteplugins.mlflow\](https://www.union.ai/docs/v2/union/api-reference/integrations/mlflow/packages/flyteplugins.mlflow/page.md) - Key features: - Basic usage: - Directory - Classes - Methods - Methods - \[Integrations > MLflow > Packages > flyteplugins.mlflow > Mlflow\](https://www.union.ai/docs/v2/union/api-reference/integrations/mlflow/packages/flyteplugins.mlflow/mlflow/page.md) - Parameters - Methods - get\_link() - \[Integrations > OmegaConf\](https://www.union.ai/docs/v2/union/api-reference/integrations/omegaconf/page.md) - \[Integrations > OpenAI\](https://www.union.ai/docs/v2/union/api-reference/integrations/openai/page.md) - \[Integrations > Papermill\](https://www.union.ai/docs/v2/union/api-reference/integrations/papermill/page.md) - \[Integrations > Papermill > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/papermill/classes/page.md) - \[Integrations > Papermill > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/papermill/packages/page.md) - \[Integrations > Papermill > Packages > flyteplugins.papermill\](https://www.union.ai/docs/v2/union/api-reference/integrations/papermill/packages/flyteplugins.papermill/page.md) - Directory - Classes - Methods - Methods - \[Integrations > Papermill > Packages > flyteplugins.papermill > NotebookTask\](https://www.union.ai/docs/v2/union/api-reference/integrations/papermill/packages/flyteplugins.papermill/notebooktask/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Integrations > Polars\](https://www.union.ai/docs/v2/union/api-reference/integrations/polars/page.md) - \[Integrations > Polars > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/polars/classes/page.md) - \[Integrations > Polars > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/polars/packages/page.md) - \[Integrations > Polars > Packages > flyteplugins.polars.df\_transformer\](https://www.union.ai/docs/v2/union/api-reference/integrations/polars/packages/flyteplugins.polars.df\_transformer/page.md) - Directory - Classes - Methods - Variables - Methods - \[Integrations > Polars > Packages > flyteplugins.polars.df\_transformer > ParquetToPolarsDecodingHandler\](https://www.union.ai/docs/v2/union/api-reference/integrations/polars/packages/flyteplugins.polars.df\_transformer/parquettopolarsdecodinghandler/page.md) - Parameters - Properties - Methods - decode() - \[Integrations > Polars > Packages > flyteplugins.polars.df\_transformer > ParquetToPolarsLazyFrameDecodingHandler\](https://www.union.ai/docs/v2/union/api-reference/integrations/polars/packages/flyteplugins.polars.df\_transformer/parquettopolarslazyframedecodinghandler/page.md) - Parameters - Properties - Methods - decode() - \[Integrations > Polars > Packages > flyteplugins.polars.df\_transformer > PolarsLazyFrameToParquetEncodingHandler\](https://www.union.ai/docs/v2/union/api-reference/integrations/polars/packages/flyteplugins.polars.df\_transformer/polarslazyframetoparquetencodinghandler/page.md) - Parameters - Properties - Methods - encode() - \[Integrations > Polars > Packages > flyteplugins.polars.df\_transformer > PolarsToParquetEncodingHandler\](https://www.union.ai/docs/v2/union/api-reference/integrations/polars/packages/flyteplugins.polars.df\_transformer/polarstoparquetencodinghandler/page.md) - Parameters - Properties - Methods - encode() - \[Integrations > PyTorch\](https://www.union.ai/docs/v2/union/api-reference/integrations/pytorch/page.md) - \[Integrations > PyTorch > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/pytorch/classes/page.md) - \[Integrations > PyTorch > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/pytorch/packages/page.md) - \[Integrations > PyTorch > Packages > flyteplugins.pytorch\](https://www.union.ai/docs/v2/union/api-reference/integrations/pytorch/packages/flyteplugins.pytorch/page.md) - Directory - Classes - \[Integrations > PyTorch > Packages > flyteplugins.pytorch > Elastic\](https://www.union.ai/docs/v2/union/api-reference/integrations/pytorch/packages/flyteplugins.pytorch/elastic/page.md) - Parameters - \[Integrations > Ray\](https://www.union.ai/docs/v2/union/api-reference/integrations/ray/page.md) - \[Integrations > Ray > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/ray/classes/page.md) - \[Integrations > Ray > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/ray/packages/page.md) - \[Integrations > Ray > Packages > flyteplugins.ray\](https://www.union.ai/docs/v2/union/api-reference/integrations/ray/packages/flyteplugins.ray/page.md) - Directory - Classes - \[Integrations > Ray > Packages > flyteplugins.ray > HeadNodeConfig\](https://www.union.ai/docs/v2/union/api-reference/integrations/ray/packages/flyteplugins.ray/headnodeconfig/page.md) - Parameters - \[Integrations > Ray > Packages > flyteplugins.ray > RayJobConfig\](https://www.union.ai/docs/v2/union/api-reference/integrations/ray/packages/flyteplugins.ray/rayjobconfig/page.md) - Parameters - \[Integrations > Ray > Packages > flyteplugins.ray > WorkerNodeConfig\](https://www.union.ai/docs/v2/union/api-reference/integrations/ray/packages/flyteplugins.ray/workernodeconfig/page.md) - Parameters - \[Integrations > SGLang\](https://www.union.ai/docs/v2/union/api-reference/integrations/sglang/page.md) - \[Integrations > SGLang > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/sglang/classes/page.md) - \[Integrations > SGLang > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/sglang/packages/page.md) - \[Integrations > SGLang > Packages > flyteplugins.sglang\](https://www.union.ai/docs/v2/union/api-reference/integrations/sglang/packages/flyteplugins.sglang/page.md) - Directory - Classes - Variables - \[Integrations > SGLang > Packages > flyteplugins.sglang > SGLangAppEnvironment\](https://www.union.ai/docs/v2/union/api-reference/integrations/sglang/packages/flyteplugins.sglang/sglangappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Integrations > Snowflake\](https://www.union.ai/docs/v2/union/api-reference/integrations/snowflake/page.md) - \[Integrations > Snowflake > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/snowflake/classes/page.md) - \[Integrations > Snowflake > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/snowflake/packages/page.md) - \[Integrations > Snowflake > Packages > flyteplugins.snowflake\](https://www.union.ai/docs/v2/union/api-reference/integrations/snowflake/packages/flyteplugins.snowflake/page.md) - Directory - Classes - \[Integrations > Snowflake > Packages > flyteplugins.snowflake > Snowflake\](https://www.union.ai/docs/v2/union/api-reference/integrations/snowflake/packages/flyteplugins.snowflake/snowflake/page.md) - Parameters - Properties - Methods - aio() - config() - container\_args() - custom\_config() - data\_loading\_config() - execute() - forward() - override() - post() - pre() - sql() - \[Integrations > Snowflake > Packages > flyteplugins.snowflake > SnowflakeConfig\](https://www.union.ai/docs/v2/union/api-reference/integrations/snowflake/packages/flyteplugins.snowflake/snowflakeconfig/page.md) - Parameters - \[Integrations > Snowflake > Packages > flyteplugins.snowflake > SnowflakeConnector\](https://www.union.ai/docs/v2/union/api-reference/integrations/snowflake/packages/flyteplugins.snowflake/snowflakeconnector/page.md) - Methods - create() - delete() - get() - get\_logs() - get\_metrics() - \[Integrations > Spark\](https://www.union.ai/docs/v2/union/api-reference/integrations/spark/page.md) - \[Integrations > Spark > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/spark/classes/page.md) - \[Integrations > Spark > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/spark/packages/page.md) - \[Integrations > Spark > Packages > flyteplugins.spark\](https://www.union.ai/docs/v2/union/api-reference/integrations/spark/packages/flyteplugins.spark/page.md) - Directory - Classes - \[Integrations > Spark > Packages > flyteplugins.spark > ParquetToSparkDecoder\](https://www.union.ai/docs/v2/union/api-reference/integrations/spark/packages/flyteplugins.spark/parquettosparkdecoder/page.md) - Parameters - Properties - Methods - decode() - \[Integrations > Spark > Packages > flyteplugins.spark > Spark\](https://www.union.ai/docs/v2/union/api-reference/integrations/spark/packages/flyteplugins.spark/spark/page.md) - Parameters - \[Integrations > Spark > Packages > flyteplugins.spark > SparkToParquetEncoder\](https://www.union.ai/docs/v2/union/api-reference/integrations/spark/packages/flyteplugins.spark/sparktoparquetencoder/page.md) - Parameters - Properties - Methods - encode() - \[Integrations > Union\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/page.md) - \[Integrations > Union > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/classes/page.md) - \[Integrations > Union > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/page.md) - \[Integrations > Union > Packages > flyteplugins.union.cli\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.cli/page.md) - Directory - Methods - Methods - \[Integrations > Union > Packages > flyteplugins.union.cli.queue\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.cli.queue/page.md) - Directory - Variables - \[Integrations > Union > Packages > flyteplugins.union.internal.validate.validate.validate\_pb2\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.internal.validate.validate.validate\_pb2/page.md) - Directory - Variables - \[Integrations > Union > Packages > flyteplugins.union.remote\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.remote/page.md) - Directory - Classes - \[Integrations > Union > Packages > flyteplugins.union.remote > ApiKey\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.remote/apikey/page.md) - Parameters - Properties - Methods - create() - delete() - get() - listall() - to\_dict() - to\_json() - update() - \[Integrations > Union > Packages > flyteplugins.union.remote > Assignment\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.remote/assignment/page.md) - Parameters - Properties - Methods - create() - get() - listall() - to\_dict() - to\_json() - unassign() - \[Integrations > Union > Packages > flyteplugins.union.remote > Cluster\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.remote/cluster/page.md) - Parameters - Properties - Methods - get() - listall() - to\_dict() - to\_json() - \[Integrations > Union > Packages > flyteplugins.union.remote > Member\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.remote/member/page.md) - Parameters - Properties - Methods - listall() - to\_dict() - to\_json() - \[Integrations > Union > Packages > flyteplugins.union.remote > Policy\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.remote/policy/page.md) - Parameters - Properties - Methods - create() - delete() - get() - listall() - to\_dict() - to\_json() - update() - \[Integrations > Union > Packages > flyteplugins.union.remote > Queue\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.remote/queue/page.md) - Parameters - Properties - Methods - activate() - create() - details() - drain() - get() - listall() - to\_dict() - to\_json() - update() - watch() - \[Integrations > Union > Packages > flyteplugins.union.remote > Role\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.remote/role/page.md) - Parameters - Properties - Methods - create() - delete() - get() - listall() - to\_dict() - to\_json() - update() - \[Integrations > Union > Packages > flyteplugins.union.remote > User\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.remote/user/page.md) - Parameters - Properties - Methods - create() - delete() - get() - listall() - to\_dict() - to\_json() - \[Integrations > Union > Packages > flyteplugins.union.utils.auth\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.utils.auth/page.md) - Directory - Classes - Methods - Methods - \[Integrations > Union > Packages > flyteplugins.union.utils.auth > AppClientCredentials\](https://www.union.ai/docs/v2/union/api-reference/integrations/union/packages/flyteplugins.union.utils.auth/appclientcredentials/page.md) - Parameters - \[Integrations > vLLM\](https://www.union.ai/docs/v2/union/api-reference/integrations/vllm/page.md) - \[Integrations > vLLM > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/vllm/classes/page.md) - \[Integrations > vLLM > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/vllm/packages/page.md) - \[Integrations > vLLM > Packages > flyteplugins.vllm\](https://www.union.ai/docs/v2/union/api-reference/integrations/vllm/packages/flyteplugins.vllm/page.md) - Directory - Classes - Variables - \[Integrations > vLLM > Packages > flyteplugins.vllm > VLLMAppEnvironment\](https://www.union.ai/docs/v2/union/api-reference/integrations/vllm/packages/flyteplugins.vllm/vllmappenvironment/page.md) - Parameters - Properties - Methods - add\_dependency() - clone\_with() - container\_args() - container\_cmd() - get\_port() - on\_shutdown() - on\_startup() - server() - \[Integrations > Weights & Biases\](https://www.union.ai/docs/v2/union/api-reference/integrations/wandb/page.md) - \[Integrations > Weights & Biases > Classes\](https://www.union.ai/docs/v2/union/api-reference/integrations/wandb/classes/page.md) - \[Integrations > Weights & Biases > Packages\](https://www.union.ai/docs/v2/union/api-reference/integrations/wandb/packages/page.md) - \[Integrations > Weights & Biases > Packages > flyteplugins.wandb\](https://www.union.ai/docs/v2/union/api-reference/integrations/wandb/packages/flyteplugins.wandb/page.md) - Key features: - Basic usage: - Directory - Classes - Methods - Methods - \[Integrations > Weights & Biases > Packages > flyteplugins.wandb > Wandb\](https://www.union.ai/docs/v2/union/api-reference/integrations/wandb/packages/flyteplugins.wandb/wandb/page.md) - Parameters - Methods - get\_link() - \[Integrations > Weights & Biases > Packages > flyteplugins.wandb > WandbSweep\](https://www.union.ai/docs/v2/union/api-reference/integrations/wandb/packages/flyteplugins.wandb/wandbsweep/page.md) - Parameters - Methods - get\_link() - \[Uctl CLI\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/page.md) - Installation - macOS - Linux - Windows - Configuration - Configuration file location hierarchy - Options - Commands - Entities - \[Uctl CLI > uctl\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl/page.md) - Synopsis - Options - \[Uctl CLI > uctl version\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-version/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl append\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-append/page.md) - Options - Options inherited from parent commands - \[Uctl CLI > uctl append > uctl append identityassignments\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-append/uctl-append-identityassignments/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl apply\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-apply/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl apply > uctl apply app\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-apply/uctl-apply-app/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl apply > uctl apply clusterconfig\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-apply/uctl-apply-clusterconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl apply > uctl apply clusterconfigid\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-apply/uctl-apply-clusterconfigid/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl apply > uctl apply clusterpoolconfig\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-apply/uctl-apply-clusterpoolconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl config\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-config/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl config > uctl config discover\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-config/uctl-config-discover/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl config > uctl config docs\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-config/uctl-config-docs/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl config > uctl config init\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-config/uctl-config-init/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl config > uctl config validate\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-config/uctl-config-validate/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl create\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-create/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl create > uctl create app\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-create/uctl-create-app/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl create > uctl create clusterpool\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-create/uctl-create-clusterpool/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl create > uctl create clusterpoolassignment\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-create/uctl-create-clusterpoolassignment/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl create > uctl create execution\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-create/uctl-create-execution/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl create > uctl create policy\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-create/uctl-create-policy/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl create > uctl create project\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-create/uctl-create-project/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl create > uctl create role\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-create/uctl-create-role/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete app\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-app/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete cluster\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-cluster/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete cluster-pool-attributes\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-cluster-pool-attributes/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete cluster-resource-attribute\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-cluster-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete clusterconfig\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-clusterconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete clusterpool\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-clusterpool/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete clusterpoolassignment\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-clusterpoolassignment/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete execution\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-execution/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete execution-cluster-label\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-execution-cluster-label/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete execution-queue-attribute\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-execution-queue-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete identityassignments\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-identityassignments/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete plugin-override\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-plugin-override/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete policy\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-policy/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete role\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-role/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete task-resource-attribute\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-task-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl delete > uctl delete workflow-execution-config\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-delete/uctl-delete-workflow-execution-config/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl demo\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-demo/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl demo > uctl demo exec\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-demo/uctl-demo-exec/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl demo > uctl demo reload\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-demo/uctl-demo-reload/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl demo > uctl demo start\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-demo/uctl-demo-start/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl demo > uctl demo status\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-demo/uctl-demo-status/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl demo > uctl demo teardown\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-demo/uctl-demo-teardown/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get app\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-app/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get cluster\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-cluster/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get cluster-pool-attributes\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-cluster-pool-attributes/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get cluster-resource-attribute\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-cluster-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get clusterconfig\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-clusterconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get clusterconfigs\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-clusterconfigs/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get clusterpool\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-clusterpool/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get clusterpoolconfig\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-clusterpoolconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get clusterswithconfig\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-clusterswithconfig/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get echo\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-echo/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get execution\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-execution/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get execution-cluster-label\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-execution-cluster-label/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get execution-queue-attribute\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-execution-queue-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get executionoperation\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-executionoperation/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get identityassignment\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-identityassignment/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get launchplan\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-launchplan/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get plugin-override\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-plugin-override/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get policy\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-policy/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get project\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-project/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get role\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-role/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get task\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-task/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get task-resource-attribute\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-task-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get workflow\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-workflow/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl get > uctl get workflow-execution-config\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-get/uctl-get-workflow-execution-config/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl register\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-register/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl register > uctl register examples\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-register/uctl-register-examples/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl register > uctl register files\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-register/uctl-register-files/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update cluster-pool-attributes\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-cluster-pool-attributes/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update cluster-resource-attribute\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-cluster-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update execution\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-execution/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update execution-cluster-label\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-execution-cluster-label/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update execution-queue-attribute\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-execution-queue-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update launchplan\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-launchplan/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update launchplan-meta\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-launchplan-meta/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update plugin-override\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-plugin-override/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update project\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-project/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update task-meta\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-task-meta/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update task-resource-attribute\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-task-resource-attribute/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update workflow-execution-config\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-workflow-execution-config/page.md) - Synopsis - Options - Options inherited from parent commands - \[Uctl CLI > uctl update > uctl update workflow-meta\](https://www.union.ai/docs/v2/union/api-reference/uctl-cli/uctl-update/uctl-update-workflow-meta/page.md) - Synopsis - Options - Options inherited from parent commands --- ## Community - \[Contributing docs and examples\](https://www.union.ai/docs/v2/union/community/contributing-docs/page.md) - The combined Flyte and Union docs site - Versions - Common build infrastructure - Variants - Both Flyte and Union docs are open source > Section bundle (all pages): https://www.union.ai/docs/v2/union/community/contributing-docs/section.md - \[Contributing docs and examples > Quick start\](https://www.union.ai/docs/v2/union/community/contributing-docs/quick-start/page.md) - Prerequisites - Clone the repository - Live preview - Distribution build - \[Contributing docs and examples > Variants\](https://www.union.ai/docs/v2/union/community/contributing-docs/variants/page.md) - Variants at the whole-page level - Conditional rendering within a page - {{}} - {{}} - Full example - Adding a new variant - Location - Creating a new variant - Testing the new variant - Building (just) the variant - \[Contributing docs and examples > Versions\](https://www.union.ai/docs/v2/union/community/contributing-docs/versions/page.md) - Versions are branches - How to create an archive version - How to create an archive version - Publishing an archive version - \[Contributing docs and examples > Authoring\](https://www.union.ai/docs/v2/union/community/contributing-docs/authoring/page.md) - Getting started - Target the right branch - Live preview - Pull Requests + Site Preview - Page Visibility - Page order - Page settings - Conditional Content - Linking to the API reference - Sigils for special cases - Warnings and Notices - Special Content Generation - Python Generated Content - Run on Union Instructions - Jupyter Notebooks - Mapped Keys (\`{{}}\`) - Mermaid Graphs - \[Contributing docs and examples > Shortcodes\](https://www.union.ai/docs/v2/union/community/contributing-docs/shortcodes/page.md) - How to specify a "shortcode" - Variants - Component Library - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` and \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\` - \`{{}}\`, \`{{}}\`, and \`{{}}\` - \`{{}}\` - \`{{}}\` - \[Contributing docs and examples > Redirects\](https://www.union.ai/docs/v2/union/community/contributing-docs/redirects/page.md) - \`docs.union.ai\` redirects - \`docs.flyte.org\` redirects - \[Contributing docs and examples > API docs\](https://www.union.ai/docs/v2/union/community/contributing-docs/api-docs/page.md) - API naming convention - Package Resource Resolution - Tips and Tricks - Auto-linking - Short vs. fully-qualified names - How auto-linking works - Magic-marker syntax for inline code - \[Contributing docs and examples > LLM-optimized documentation\](https://www.union.ai/docs/v2/union/community/contributing-docs/llm-docs/page.md) - Output files - Discovery hierarchy - How \`page.md\` files are generated - Enabling section bundles - The \`llms-full.txt\` link conversion - Regenerating - \[Contributing docs and examples > Publishing\](https://www.union.ai/docs/v2/union/community/contributing-docs/publishing/page.md) - Requirements - Managing the Tutorial Pages - Building and running locally - Developer Experience - Controlling Development Environment - Changing 'variants' - Troubleshootting - Identifying Problems: Missing Content - Identifying Problems: Page Visibility - Building Production - Testing Production Build --- ## Release notes - \[Release notes\](https://www.union.ai/docs/v2/union/release-notes/page.md) - June 2026 - :rocket: Retries with Backoff and Timeout Controls - :robot: Flyte-Native Agent Construct - :memo: Multi-Pod Log Streaming - :computer: Build and Deploy MCP Servers - :sparkles: Smarter Failure Classification - :zap: Faster CLI Startup and Unified Local Caches - :sparkles: Friendlier Build and Deploy Errors - :zap: More Resilient Uploads - :wrench: SDK Reliability Improvements - :wrench: Exclude Files from Code Bundles with \`.flyteignore\` - :wrench: Settings Applied at Run Creation - :sparkles: Console Run Exploration Improvements - :sparkles: Task Environment Details Drawer - :gear: Cluster and Cluster Pool Management from the CLI - :chart\_with\_upwards\_trend: GPU Configuration for Ray Clusters - :sparkles: Pydantic Union Types with Field Annotations - :sparkles: Queues with Concurrency and Depth Control (Beta) - :sparkles: Events API (Beta) - :gear: Self-Managed Dataplane Updates - May 2026 - :robot: AI Agent Components in \`flyte.ai\` - :wrench: Hierarchical Settings - :sparkles: Interactive Triggers - :hammer: Non-Root Images by Default - :sparkles: Actionable CLI Error Messages - :hammer: Image Building Enhancements - :sparkles: Console Enhancements - :computer: TUI and Shell Task Improvements - :sparkles: New Plugins and Examples - :wrench: User-Facing Logger - :wrench: Keyring Opt-Out - :gear: Cluster-Aware Data Access - :gear: OpenShift Support for Self-Managed Deployments - March 2026 - :wrench: Extended Idle Timeout for Panel Apps - :wrench: Plugin Variants Documentation - :rocket: Google Gemini Plugin Integration - :hammer: Forced Image Build Caching - :computer: LLM-Powered Code Generation - :wrench: Updated AI Plugin Examples - :wrench: Debug Mode Integration - :sparkles: Improved CLI Enum Support - :memo: Programmatic Log Access - :zap: Simplified PyTorch Example Setup - :chart\_with\_upwards\_trend: Distributed Training Evaluation - :zap: Improved Benchmark Flexibility - :computer: CLI Project Management - :robot: Anthropic Claude Integration - :hourglass\_flowing\_sand: Panel App Enhancements - :gear: AWS Config File Support - :sparkles: Improved Task Execution Reliability - :wrench: Enhanced Action Service Integration - :rocket: Async Training with Early Stopping - :wrench: Improved Include Path Handling - :zap: Enhanced Retry Management - :zap: Improved Module Loading - :zap: Dynamic Batching for Improved GPU Utilization - :sparkles: Run Cache Disabling - :computer: Vim Key Navigation for TUI - :sparkles: Clickable Image Build URLs - :sparkles: Enhanced Run Filters - :wrench: Simplified Dependency Management - :robot: MLE Agent Enhancements - :sparkles: Improved Task Command Initialization - :zap: New Example Applications & Bug Fixes - :gear: Phase Transitions Tracking - :wrench: Multiple Source Files Support - :package: Simplified Code Bundling - :wrench: Improved Error Messaging for Deployment - :wrench: Improved Debugging for Reusable Tasks - :sparkles: JSONL Plugin Support - February 2026 - :sparkles: JSON Schema Enhancement - :calculator: Panel Calculator Example - :sparkles: Spark Plugin Update - :lock: Secure Package Specification - :zap: Enum Name Acceptance in CLI - :wrench: Enhanced Pod Template Handling - :zap: Stress Testing Example Added - :bug: Correct Serialization Field - :wrench: Improved Async Task Handling - :wrench: Sync Alignment of File Upload Methods - :hourglass: Request Timeout Configuration - :wrench: Enhanced Bundling and Error Handling - :wrench: Dynamic Pydantic Model Creation - :busts\_in\_silhouette: Human-in-the-Loop Plugin - :rocket: Stateless Code Sandbox - :wrench: Improved CLI Logging Initialization - :wrench: Enhanced Ignore Handling - :whale: CI Image Builder - :wrench: TypedDict Compatibility Fix - :globe\_with\_meridians: Cross-Platform Code Bundling - :wrench: Improved CLI JSON Formatting - :wrench: Improved Pod Image Handling - :sparkles: Flyte Webhook Environment - :repeat: Retry Interceptor for gRPC - :sparkles: Orchestration Sandbox Feature - :wrench: Task Shortname Override Fix - :sparkles: NVIDIA H100 GPU Support - :zap: Enhanced Error Handling in PyTorch Elastic Jobs - :wrench: Reverse Path Priority Fix - :globe\_with\_meridians: S3 Virtual Hosted-Style Support - November 2025 - :fast\_forward: Grouped Runs - :globe\_with\_meridians: Apps (beta) - :label: Custom context - :lock: Secrets UI - Image builds now run in the same project-domain - Support for secret mounts in Poetry and UV projects - October 2025 - :infinity: Larger fanouts - :computer: Remote debugging for Ray head nodes - :zap: Triggers and audit history - :arrow\_up: Deployed tasks and input passing --- ## Getting support - \[Getting support\](https://www.union.ai/docs/v2/union/support/page.md) - Severity levels - Response time targets - Shared Slack channel - Support portal - Union Cloud console - Email --- ## Security - \[Architecture\](https://www.union.ai/docs/v2/union/security/architecture/page.md) - \[Architecture > Two-plane separation\](https://www.union.ai/docs/v2/union/security/architecture/two-plane-separation/page.md) - Control plane - Data plane - Risk mitigation - \[Architecture > Control plane\](https://www.union.ai/docs/v2/union/security/architecture/control-plane/page.md) - What it stores - Infrastructure - Capabilities - \[Architecture > Data plane\](https://www.union.ai/docs/v2/union/security/architecture/data-plane/page.md) - Components - Object store layout - Kubernetes security - Container security - IAM and workload identity - Apps & Serving security - Verification - Components - Kubernetes security - Container security - IAM and workload identity - \[Architecture > Network architecture\](https://www.union.ai/docs/v2/union/security/architecture/network/page.md) - Outbound-only model - Direct-to-DataPlane tunnel - Per-cluster routing - Sovereign Data Plane - Communication paths - Verification - Outbound-only model - Direct-to-DataPlane tunnel - \[Architecture > Sovereign Data Plane\](https://www.union.ai/docs/v2/union/security/architecture/sovereign-data-plane/page.md) - When to use it - How it differs from the default - What stays the same - Topology - Setup - Trade-offs - Verification - Network reachability - Identity continues to gate access - \[Architecture > Private connectivity (BYOC)\](https://www.union.ai/docs/v2/union/security/architecture/private-connectivity/page.md) - Verification - Private management connection - \[Data protection\](https://www.union.ai/docs/v2/union/security/data-protection/page.md) - \[Data protection > Data classification and residency\](https://www.union.ai/docs/v2/union/security/data-protection/classification-and-residency/page.md) - Data classification - Data residency - Verification - Data classification - Data residency - \[Data protection > Secrets management\](https://www.union.ai/docs/v2/union/security/data-protection/secrets/page.md) - Backends - Secret lifecycle - Verification - Write-only API - Secret lifecycle - \[Data protection > Logging and audit\](https://www.union.ai/docs/v2/union/security/data-protection/logging-and-audit/page.md) - Task logging - Observability metrics - Audit trail - Verification - Task logging - Audit trail - \[Identity and access\](https://www.union.ai/docs/v2/union/security/identity-and-access/page.md) - \[Identity and access > Authentication\](https://www.union.ai/docs/v2/union/security/identity-and-access/authentication/page.md) - Authentication methods - Single sign-on - Verification - SSO and credential lifecycle - \[Identity and access > Role-based access control\](https://www.union.ai/docs/v2/union/security/identity-and-access/rbac/page.md) - Built-in roles - Custom policies - Enforcement - Least privilege - Verification - RBAC enforcement - \[Identity and access > Human access controls\](https://www.union.ai/docs/v2/union/security/identity-and-access/human-access/page.md) - Self-managed - BYOC - Customer-side support access (optional) - Access scope - Verification - Human access controls - \[Compliance and governance\](https://www.union.ai/docs/v2/union/security/compliance/page.md) - \[Compliance and governance > Certifications and Trust Center\](https://www.union.ai/docs/v2/union/security/compliance/certifications/page.md) - Certifications overview - SOC 2 Type II - Trust Center - Verification - Certifications - \[Compliance and governance > HIPAA compliance\](https://www.union.ai/docs/v2/union/security/compliance/hipaa/page.md) - Verification - HIPAA compliance - \[Compliance and governance > GDPR alignment\](https://www.union.ai/docs/v2/union/security/compliance/gdpr/page.md) - Verification - GDPR alignment - \[Compliance and governance > Standards compliance\](https://www.union.ai/docs/v2/union/security/compliance/standards/page.md) - Verification - Standards compliance - \[Compliance and governance > Shared responsibility model\](https://www.union.ai/docs/v2/union/security/compliance/shared-responsibility/page.md) - Self-managed - BYOC shifts - Verification - Shared responsibility model - \[Compliance and governance > Organizational security\](https://www.union.ai/docs/v2/union/security/compliance/organizational-security/page.md) - Employee security lifecycle - Governance - Security development lifecycle - Verification - Organizational security - \[Compliance and governance > Vulnerability management\](https://www.union.ai/docs/v2/union/security/compliance/vulnerability-management/page.md) - Vulnerability assessment - Patch management - Incident response - Third-party dependency risk - Verification - Vulnerability management --- ## Platform deployment - \[BYOC deployment\](https://www.union.ai/docs/v2/union/deployment/byoc/page.md) - Getting started - Cloud resource integration - Additional configuration > Section bundle (all pages): https://www.union.ai/docs/v2/union/deployment/byoc/section.md - \[BYOC deployment > Platform architecture\](https://www.union.ai/docs/v2/union/deployment/byoc/platform-architecture/page.md) - Control plane - Data plane - Data plane nodes - Union.ai operator - Registry data - Execution data - Raw data - Literal data - Data privacy - \[BYOC deployment > Configuring your data plane\](https://www.union.ai/docs/v2/union/deployment/byoc/configuring-your-data-plane/page.md) - Cloud provider - Multi-cluster - Account ID - Region - VPC - Data retention policy - Worker node groups - Node group name - Node type - Minimum - Maximum - Interruptible instances - Taints - Disk - Resources held back - Example specification - After deployment - Adjusting your configuration - \[BYOC deployment > Multi-cluster and multi-cloud\](https://www.union.ai/docs/v2/union/deployment/byoc/multi-cluster/page.md) - Domain isolation - Project isolation - Data and metadata isolation - \[BYOC deployment > Data plane setup on AWS\](https://www.union.ai/docs/v2/union/deployment/byoc/data-plane-setup-on-aws/page.md) - Setting permissions through CloudFormation - Click the Launch Stack button - Confirm the details - Share the role ARN - Updating permissions through CloudFormation - Update your CloudFormation template - Setting permissions manually - Prepare the policy documents - Create the role manually - Share the role ARN - Updating permissions manually - Setting up and managing your own VPC (optional) - Private EKS endpoint - Create additional roles for ECS - Attach a new IAM policy to the Union role - Configure VPC Endpoints - \[BYOC deployment > Data plane setup on GCP\](https://www.union.ai/docs/v2/union/deployment/byoc/data-plane-setup-on-gcp/page.md) - Select or create a project - Ensure billing is linked - Create a workload identity pool and provider - In the GCP web console - On the command line using \`gcloud\` - Create a role for Union.ai admin - Create the Union.ai admin service account - In the GCP web console - On the command line using \`gcloud\` - Grant access for the Workflow Identity Pool to the Service Account - In the GCP web console - On the command line using \`gcloud\` - Enable services API - In the GCP web console - On the command line using \`gcloud\` - Setting up and managing your own VPC (optional) - Example VPC CIDR Block allocation - \[BYOC deployment > Data plane setup on Azure\](https://www.union.ai/docs/v2/union/deployment/byoc/data-plane-setup-on-azure/page.md) - Selecting Azure tenant and subscription - Create a Microsoft Entra Application Registration - Create a Microsoft Entra ID Application for Union.ai Access - Create Microsoft Entra ID Applications for Union.ai cost allocation - (Recommended) Create a Microsoft Entra group for cluster administration - (Optional) Setting up and managing your own VNet - Required Union.ai VNet permissions - Required VNet properties - Example VPC CIDR Block allocation - Union.ai Maintenance Windows - \[BYOC deployment > Data retention policy\](https://www.union.ai/docs/v2/union/deployment/byoc/data-retention-policy/page.md) - Data categories - How policies are specified - Deletion of current versions - Deletion of non-current versions - Defaults - Attempting to access deleted data - Separate sets of policies per cluster - Data retention and task caching - \[BYOC deployment > Enabling AWS resources\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-aws-resources/page.md) - Types of access - Infrastructure-level access - Task code access - Background - Enabling access - Creating a custom policy - Setting up global access - Setting up project-domain-scoped access - Create the IAM role - Configure the cluster to use the new IAM role > Section bundle (all pages): https://www.union.ai/docs/v2/union/deployment/byoc/enabling-aws-resources/section.md - \[BYOC deployment > Enabling AWS resources > Enabling AWS S3\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-aws-resources/enabling-aws-s3/page.md) - Add permissions to your custom policy - Accessing S3 from your task code - \[BYOC deployment > Enabling AWS resources > Enabling AWS ECR\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-aws-resources/enabling-aws-ecr/page.md) - Access to ECR in the same account is enabled by default - Enabling cross-account access to ECR - \[BYOC deployment > Enabling AWS resources > Enabling AWS Secrets Manager\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-aws-resources/enabling-aws-secrets-manager/page.md) - Ensure that AWS Secrets Manager is enabled - Create your secrets - Get the secret ARN - Create a policy providing access to your secrets - Bind the policy to the User Flyte Role - Using AWS secrets in your task code - \[BYOC deployment > Enabling GCP resources\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-gcp-resources/page.md) - Types of access - Infrastructure-level access - Task code access - Domain-scoped access - Globally-scoped access - Find the actual name of \`\` > Section bundle (all pages): https://www.union.ai/docs/v2/union/deployment/byoc/enabling-gcp-resources/section.md - \[BYOC deployment > Enabling GCP resources > Enabling Google Cloud Storage\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-gcp-resources/enabling-google-cloud-storage/page.md) - Grant \`\` access to the bucket - \[BYOC deployment > Enabling GCP resources > Enabling Google Artifact Registry\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-gcp-resources/enabling-google-artifact-registry/page.md) - Access to Artifact Registry in the same project is enabled by default - Enabling cross-project access to Artifact Registry - \[BYOC deployment > Enabling GCP resources > Enabling Google Secret Manager\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-gcp-resources/enabling-google-secret-manager/page.md) - Create your secrets - Same-project secrets - Cross-project secrets - \[BYOC deployment > Enabling GCP resources > Enabling BigQuery\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-gcp-resources/enabling-bigquery/page.md) - \[BYOC deployment > Enabling Azure resources\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-azure-resources/page.md) - Types of access - Infrastructure-level access - Task code access - Domain-scoped access - Globally-scoped access > Section bundle (all pages): https://www.union.ai/docs/v2/union/deployment/byoc/enabling-azure-resources/section.md - \[BYOC deployment > Enabling Azure resources > Enabling Azure Blob Storage\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-azure-resources/enabling-azure-blob-storage/page.md) - Providing permissions to Azure Blob Storage container - Union.ai-managed permissions - Manage permissions directly - \[BYOC deployment > Enabling Azure resources > Enabling Azure Container Registry (ACR)\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-azure-resources/enabling-azure-container-registry/page.md) - Creating a container registry - Creating a container registry outside of Union.ai - Creating a Union.ai-managed container registry - Enable access to ACR in a different subscription within the same Azure tenant - Allow Union.ai to manage permissions - Manage permissions directly - Enable access to ACR in a different Azure tenant - References - \[BYOC deployment > Enabling Azure resources > Enabling Azure Key Vault\](https://www.union.ai/docs/v2/union/deployment/byoc/enabling-azure-resources/enabling-azure-key-vault/page.md) - Providing permissions to Azure Key Vault - Accessing the secret within Union.ai - \[BYOC deployment > Single sign on setup\](https://www.union.ai/docs/v2/union/deployment/byoc/single-sign-on-setup/page.md) - Google OpenID Connect - Microsoft Entra ID (formerly Azure AD) - Other identity providers > Section bundle (all pages): https://www.union.ai/docs/v2/union/deployment/byoc/single-sign-on-setup/section.md - \[BYOC deployment > Single sign on setup > Google OpenID Connect\](https://www.union.ai/docs/v2/union/deployment/byoc/single-sign-on-setup/google-oidc/page.md) - Setting up OAuth 2.0 - Obtain OAuth 2.0 credentials - Share the client ID and client secret securely with Union.ai - \[BYOC deployment > Single sign on setup > Microsoft Entra ID (formerly Azure AD)\](https://www.union.ai/docs/v2/union/deployment/byoc/single-sign-on-setup/microsoft-entra-id/page.md) - Register an Entra ID application - Copy the values needed by the Union.ai team - Application (client) ID and directory (tenant) ID - Client secret - Share the client secret securely with Union.ai - Share the IDs with Union.ai - \[BYOC deployment > Single sign on setup > Other identity providers\](https://www.union.ai/docs/v2/union/deployment/byoc/single-sign-on-setup/other-identity-providers/page.md) - Share the client secret securely with the Union.ai team - Share the application (client) ID with Union.ai - \[Self-managed deployment\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/page.md) - Getting started - Configuration - Reference > Section bundle (all pages): https://www.union.ai/docs/v2/union/deployment/selfmanaged/section.md - \[Self-managed deployment > Architecture\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/architecture/page.md) > Section bundle (all pages): https://www.union.ai/docs/v2/union/deployment/selfmanaged/architecture/section.md - \[Self-managed deployment > Architecture > Overview\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/architecture/overview/page.md) - Control plane - Data plane - Data plane nodes - Union.ai operator - Registry data - Execution data - Raw data - Literal data - Data privacy - \[Self-managed deployment > Architecture > Kubernetes Access Controls\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/architecture/kubernetes-rbac/page.md) - Service account - Standard mode vs. low-privilege mode - Namespace-scoped Roles - ClusterRoles (standard mode only) - Metrics and Monitoring - Resource Management - Workflow Management - Service Access - \`operator/operator-proxy\` - \`FlytePropeller/PropellerWebhook\` - \[Self-managed deployment > Cluster recommendations\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/cluster-recommendations/page.md) - Kubernetes Versions - Networking Requirements - VPC and subnet sizing - Public vs. private subnets - NAT Gateway requirements - Service accounts - Node Pools - \[Self-managed deployment > Data plane setup on generic Kubernetes\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-generic/page.md) - \[Self-managed deployment > Data plane setup on generic Kubernetes > Prepare infrastructure\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-generic/prepare-infra/page.md) - Kubernetes Cluster - Object Storage - CORS Configuration - Data Retention - Container Registry - Identity & Access - Storage credentials - Registry credentials - \[Self-managed deployment > Data plane setup on generic Kubernetes > Deploy the dataplane\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-generic/deploy-dataplane/page.md) - Assumptions - Prerequisites - Deploy the Union.ai operator - \[Self-managed deployment > Data plane setup on AWS\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-aws/page.md) - \[Self-managed deployment > Data plane setup on AWS > Prepare infrastructure\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-aws/prepare-infra/page.md) - Environment variables - EKS Cluster - S3 - CORS Configuration - Data Retention - ECR - IAM - 1. Enable OIDC - 2. Create the IAM role - 3. Attach the S3 policy - 4. Attach the ECR policy - 5. Configure the service account annotation - \[Self-managed deployment > Data plane setup on AWS > Deploy the dataplane\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-aws/deploy-dataplane/page.md) - Assumptions - Prerequisites - Deploy the Union.ai operator - \[Self-managed deployment > Data plane setup on GKE (GCP)\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-gcp/page.md) - \[Self-managed deployment > Data plane setup on GKE (GCP) > Prepare infrastructure\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-gcp/prepare-infra/page.md) - Environment variables - GKE Cluster - BuildKit node pool - GCS - CORS Configuration - Data Retention - Artifact Registry - Workload Identity - 1. Create a Google Service Account - 2. Bind the GSA to the Kubernetes service account - 3. Grant GCS access - 4. Grant Artifact Registry access - 5. Grant token creator access - \[Self-managed deployment > Data plane setup on GKE (GCP) > Deploy the dataplane\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-gcp/deploy-dataplane/page.md) - Assumptions - Prerequisites - Deploy the Union.ai operator - \[Self-managed deployment > Data plane setup on Azure\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-azure/page.md) - \[Self-managed deployment > Data plane setup on Azure > Prepare infrastructure\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-azure/prepare-infra/page.md) - Prerequisites - Environment variables - 1. Subscription and Resource Group - 2. AKS Cluster - 3. Node Pools - System node pool - CPU worker node pool - GPU node pool (optional) - 4. Storage Account and Container - CORS Configuration - Data Retention - 5. Managed Identities - 6. Workload Identity and Federated Credentials - Backend identity (Union system components) - Worker identity (task execution pods) - 7. Role Assignments - 8. Azure Key Vault (optional) - \[Self-managed deployment > Data plane setup on Azure > Deploy the dataplane\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-azure/deploy-dataplane/page.md) - Assumptions - Prerequisites - Deploy the Union.ai operator - \[Self-managed deployment > Data plane setup on OCI\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-oci/page.md) - \[Self-managed deployment > Data plane setup on OCI > Prepare infrastructure\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-oci/prepare-infra/page.md) - OKE Cluster - Object Storage - CORS Configuration - Data Retention - Container Registry - Identity & Access - Option A: Instance Principals (recommended) - Option B: Static Credentials - \[Self-managed deployment > Data plane setup on OCI > Deploy the dataplane\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-oci/deploy-dataplane/page.md) - Assumptions - Prerequisites - Deploy the Union.ai operator - \[Self-managed deployment > Data plane setup on CoreWeave\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-coreweave/page.md) - \[Self-managed deployment > Data plane setup on CoreWeave > Prepare infrastructure\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-coreweave/prepare-infra/page.md) - CKS cluster - CoreWeave AI Object Storage - Create a bucket - Generate access credentials - Create an access policy - \[Self-managed deployment > Data plane setup on CoreWeave > Deploy the dataplane\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-coreweave/deploy-dataplane/page.md) - Assumptions - Prerequisites - Deploy the Union.ai operator - Test a workflow - Troubleshooting - Additional resources - \[Self-managed deployment > Data plane setup on Crusoe\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-crusoe/page.md) - \[Self-managed deployment > Data plane setup on Crusoe > Prepare infrastructure\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-crusoe/prepare-infra/page.md) - CMK cluster - Crusoe Cloud Storage - Create a bucket - Generate access credentials - Create an IAM / access policy - \[Self-managed deployment > Data plane setup on Crusoe > Deploy the dataplane\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-crusoe/deploy-dataplane/page.md) - Assumptions - Prerequisites - Deploy the Union.ai operator - Test a workflow - Troubleshooting - Additional resources - \[Self-managed deployment > Data plane setup on Nebius\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-nebius/page.md) - \[Self-managed deployment > Data plane setup on Nebius > Prepare infrastructure\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-nebius/prepare-infra/page.md) - Nebius Managed Kubernetes cluster - Nebius Object Storage - Create a bucket - Generate access credentials - \[Self-managed deployment > Data plane setup on Nebius > Deploy the dataplane\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/selfmanaged-nebius/deploy-dataplane/page.md) - Assumptions - Prerequisites - Deploy the Union.ai operator - GPU node configuration (Nebius-specific) - Working with the Nebius Container Registry - Test a workflow - Additional resources - \[Self-managed deployment > Advanced Configurations\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/page.md) > Section bundle (all pages): https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/section.md - \[Self-managed deployment > Advanced Configurations > Configuring Service and Worker Node Pools\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/node-pools/page.md) - \[Self-managed deployment > Advanced Configurations > Authentication\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/authentication/page.md) - Overview - Prerequisites - Configuring your Identity Provider - Authorization server setup - Application details - Control plane Helm configuration - Global variables - Flyteadmin OIDC configuration - Flyteadmin and scheduler admin SDK client - Scheduler auth secret - Service-to-service authentication - Executions service - Ingress auth annotations - Dataplane Helm configuration - Dataplane global variables - Cluster resource sync - Operator (union service auth) - Propeller admin client - In-pod control-plane authentication (EAGER\_API\_KEY) - Dataplane secrets - Secret delivery - Option A: External Secrets Operator (recommended) - Option B: Direct Kubernetes secrets - SDK and CLI authentication - Client credentials for CI/CD - Troubleshooting - Browser login redirects in a loop - SDK gets 401 Unauthenticated - Internal services get 401 - Operator or propeller cannot authenticate - Scheduler fails to start - \[Self-managed deployment > Advanced Configurations > Code Viewer\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/code-viewer/page.md) - Enable CORS policy on your fast registration bucket - AWS S3 Console - Google GCS - Azure Storage - Troubleshooting - \[Self-managed deployment > Advanced Configurations > Image Builder\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/image-builder/page.md) - Requirements - Build backends - Configuration - Authentication - AWS - Google Cloud Platform - Azure - Private registries - \[Self-managed deployment > Advanced Configurations > Multiple Clusters\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/multi-cluster/page.md) - Using cluster pools - project-domain-clusterPool mapping - project-domain-workflow-clusterPool mapping - Data sharing between cluster pools - \[Self-managed deployment > Advanced Configurations > Persistent logs\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/persistent-logs/page.md) - AWS (IRSA) - 1. Create an IAM policy - 2. Create an IAM role with a trust policy - 3. Configure the Helm values - Azure (Workload Identity Federation) - Azure prerequisites - 1. Create or reuse a Managed Identity - 2. Add a federated credential - 3. Assign a storage role - 4. Configure the Azure Helm values - GCP (Workload Identity) - GCP prerequisites - 1. Create or reuse a GCP service account - 2. Grant storage permissions - 3. Bind the Kubernetes service account to the GCP service account - 4. Configure the GCP Helm values - Disabling persistent logs - \[Self-managed deployment > Advanced Configurations > Monitoring\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/monitoring/page.md) - Architecture overview - Union features Prometheus - Scrape targets - Configuration - Internal service endpoint - Prometheus Simple (low-privilege mode) - Recording rules - Enabling cluster health monitoring - Prometheus Operator CRDs - Customizing the monitoring stack - Scraping Union services from your own Prometheus - Static scrape configs - ServiceMonitor (Prometheus Operator) - Further reading - \[Self-managed deployment > Advanced Configurations > Secrets\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/union-secrets/page.md) - \[Self-managed deployment > Advanced Configurations > Data retention policies\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/data-retention/page.md) - Where metadata vs. raw data lives - Impact of raw data loss - Applying retention deliberately - Designing lifecycle rules - \[Self-managed deployment > Advanced Configurations > Compute plugins\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/plugins/page.md) - Dask - Install the Dask operator - Configure the data plane Helm values - AWS - GCP - Ray - Install the KubeRay operator - Configure the data plane Helm values - AWS - GCP - \[Self-managed deployment > Advanced Configurations > Namespace mapping\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/configuration/namespace-mapping/page.md) - Template syntax - Examples - Data plane configuration - How it works - \[Self-managed deployment > Helm chart reference\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/helm-chart-reference/page.md) > Section bundle (all pages): https://www.union.ai/docs/v2/union/deployment/selfmanaged/helm-chart-reference/section.md - \[Self-managed deployment > Helm chart reference > Page\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/helm-chart-reference/dataplane/page.md) - Chart info - Dependencies - Values - \[Self-managed deployment > Helm chart reference > Page\](https://www.union.ai/docs/v2/union/deployment/selfmanaged/helm-chart-reference/knative-operator/page.md) - Chart info - Values - \[Managing Union with Terraform\](https://www.union.ai/docs/v2/union/deployment/terraform/page.md) - Overview - Why use Terraform? - Getting Started - \[Installation\](installation/page.md) - \[Resource Management\](management/page.md) - \[Security Best Practices\](security/page.md) - Requirements > Section bundle (all pages): https://www.union.ai/docs/v2/union/deployment/terraform/section.md - \[Managing Union with Terraform > Installing the Union Terraform Provider\](https://www.union.ai/docs/v2/union/deployment/terraform/installation/page.md) - Quick Start - Versioning - \[Managing Union with Terraform > Managing Union Resources with Terraform\](https://www.union.ai/docs/v2/union/deployment/terraform/management/page.md) - Provider Configuration - Basic Configuration - Configuration Parameters - Authentication - 1. Provider Configuration - 2. Environment Variable - Generating an API Key - Available Resources - Projects - Users - Roles - Policies - API Keys - OAuth Applications - Access Assignments - Available Data Sources - Projects - Users - Roles - Policies - API Keys - Applications - Data Plane Information - Control Plane Information - Data Plane Listings - Best Practices - Use Variables for Sensitive Data - Organize Resources with Modules - Use Organization Restrictions - Version Control Your Configuration - Use Remote State - Example: Complete Setup - Additional Resources - Requirements - Support and Contributions - \[Managing Union with Terraform > Security Best Practices\](https://www.union.ai/docs/v2/union/deployment/terraform/security/page.md) - Recommended Approaches - 1. Use Cloud Secret Managers - 2. Use HashiCorp Vault - 3. Use Environment Variables - 4. Use Terraform Variables with \`.tfvars\` Files - Additional Security Measures - Encrypt Terraform State - Use State Locking - Rotate API Keys Regularly - Restrict Provider Permissions - Use Separate API Keys per Environment - Security Checklist - CI/CD Pipeline Security - GitHub Actions - GitLab CI - Best Practices for CI/CD - Additional Resources --- # Unknown \# Flyte OSS Flyte is a free and open source platform that provides a full suite of powerful features for orchestrating AI workflows. Flyte empowers AI development teams to rapidly ship high-quality code to production by offering optimized performance, unparalleled resource efficiency, and a delightful workflow authoring experience. You deploy and manage Flyte yourself, on your own cloud infrastructure. > \[!NOTE\] > This documentation for open-source Flyte is maintained by Union.ai. > > You can switch to the documentation for the commercial versions with the selector above. ### \[Introduction\](https://www.union.ai/docs/v1/flyte/user-guide/introduction/page.md) Flyte is the leading open-source Kubernetes-native workflow orchestrator. ### \[Getting started\](https://www.union.ai/docs/v1/flyte/user-guide/getting-started/page.md) Build your first Flyte workflow, exploring the major features of the platform along the way. ### \[Core concepts\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/page.md) Understand the core concepts of the Flyte platform. ### \[Development cycle\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/page.md) Explore the Flyte development cycle from experimentation to production. ### \[Data input/output\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/page.md) Manage the input and output of data in your Flyte workflow. ### \[Programming\](https://www.union.ai/docs/v1/flyte/user-guide/programming/page.md) Learn about Flyte-specific programming constructs. ## Subpages - \[Introduction\](https://www.union.ai/docs/v1/flyte/user-guide/introduction/page.md) - Flyte - Trying out Flyte - Flyte in production - \[Getting started\](https://www.union.ai/docs/v1/flyte/user-guide/getting-started/page.md) - Try Flyte on your local machine - \[Core concepts\](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/page.md) - Defining tasks and workflows - Type annotation is required - Workflows \*are not\* full Python functions - Registering tasks and workflows - Registering on the command line with \`pyflyte\` or \`flytectl\` - Registering in Python with \`FlyteRemote\` - Results of registration - Changing tasks and workflows - Inspecting tasks and workflows - Inspecting workflows in the UI - Inspecting tasks in the UI - Inspecting workflows on the command line with \`flytectl\` - Inspecting tasks on the command line with \`flytectl\` - Inspecting tasks and workflows in Python with \`FlyteRemote\` - Running tasks and workflows - Running a task or workflow in the UI - Running a task or workflow locally on the command line with \`pyflyte\` or \`python\` - Running a task or workflow remotely on the command line with \`pyflyte\` - Running a task or workflow remotely in Python with \`FlyteRemote\` - \[Development cycle\](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/page.md) - \[Data input/output\](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/page.md) - Mapping Python to Flyte types - \[Programming\](https://www.union.ai/docs/v1/flyte/user-guide/programming/page.md) --- \*\*Source\*\*: https://github.com/unionai/unionai-docs/blob/main/content/user-guide/\_index.md \*\*HTML\*\*: https://www.union.ai/docs/v1/flyte/user-guide/ --- # Tutorials | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Tutorials ========= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/tutorials/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section provides tutorials that walk you through the process of building AI/ML applications on Flyte. The example applications range from training XGBoost models in tabular datasets to fine-tuning large language models for text generation tasks. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/tutorials/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/tutorials/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Architecture | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Architecture ============ This section covers the architecture of the Flyte system. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Reference | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Reference ========= This section provides the reference material for all Flyte APIs, SDKs and CLIs. To get started, add `flytekit` to your project $ uv add flytekit This will install the Flytekit SDKs and the `pyflyte` CLI. Flytekit SDK The Flytekit SDK provides the core Python API for building Flyte workflows. Pyflyte CLI The Pyflyte CLI is the command-line interface for interacting with your Flyte instance. Flytectl CLI The Flytectl CLI is an alternative CLI for performing administrative tasks and for use in CI/CD environments. Flyteidl Flyteidl is the specification for the Flyte language in protobuf. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/api-reference/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Platform deployment | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Platform deployment =================== Flyte is an open-source workflow orchestration platform that you deploy and manage on your own infrastructure. This section covers planning, installing, configuring, and operating a Flyte backend. The sections below cover the full scope of running Flyte in production: * [**Flyte deployment**](https://www.union.ai/docs/v1/flyte/deployment/flyte-deployment) — Planning and installing Flyte on Kubernetes (single-cluster or multi-cluster setups). * [**Flyte configuration**](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration) — Configuring authentication, secrets, notifications, monitoring, GPUs, pod templates, and other runtime settings. * [**Flyte connectors**](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors) — Integrating with external services such as Airflow, BigQuery, Databricks, Snowflake, and more. * [**Flyte plugins**](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins) — Native backend plugins for Kubernetes operators, Spark, Athena, SageMaker, and other compute backends. * [**Configuration reference**](https://www.union.ai/docs/v1/flyte/deployment/configuration-reference) — Full reference for FlyteAdmin, FlytePropeller, DataCatalog, and Scheduler config files. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/deployment/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Bioinformatics | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Bioinformatics ============== Bioinformatics encompasses all the ways we aim to solve biological problems by computational means. Flyte provides a number of excellent abstractions and features for solving such problems in a reliable, reproducible and ergonomic way. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/tutorials/bioinformatics/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Programming | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Programming =========== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/user-guide/programming/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section covers the general programming of Flyte. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/user-guide/programming/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/user-guide/programming/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Flytelab | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Flytelab ======== This section contains end-to-end ML projects using Flyte. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/tutorials/flytelab/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Feature engineering | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Feature engineering =================== **Feature Engineering** is an essential part of Machine Learning. It is the process of transforming raw data into features that better represent the underlying problem to the predictive models, resulting in improved model accuracy on unseen data. Explore how features can be engineered with the power of Flyte. | Feature Engineering Task | Description | | --- | --- | | [EDA and Feature Engineering With Papermill](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/exploratory_data_analysis) | How to use Jupyter notebook within Flyte | | [Data Cleaning and Feature Serving With Feast](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/feast_integration) | How to use Feast to serve data in Flyte | LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/tutorials/feature-engineering/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Integrations | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Integrations ============ An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/integrations/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. Flyte is designed to be highly extensible and can be customized in multiple ways. Want to contribute an integration example? Check out the [contribution guide](https://www.union.ai/docs/v1/flyte/community/contribute/contribute-examples) . [Connectors](https://www.union.ai/docs/v1/flyte/integrations/#connectors) --------------------------------------------------------------------------- Flyte supports [the following connectors out-of-the-box](https://www.union.ai/docs/v1/flyte/integrations/connectors) . If you don’t see the connector you need below, have a look at [Creating a new connector](https://www.union.ai/docs/v1/flyte/integrations/connectors#creating-a-new-connector) . | Agent | Description | | --- | --- | | [SageMaker connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/sagemaker-inference-connector) | Deploy models and create, as well as trigger inference endpoints on AWS SageMaker. | | [Airflow connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/airflow-connector) | Run Airflow jobs in your workflows with the Airflow connector. | | [BigQuery connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/bigquery-connector) | Run BigQuery jobs in your workflows with the BigQuery connector. | | [ChatGPT connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/chatgpt-connector) | Run ChatGPT jobs in your workflows with the ChatGPT connector. | | [Databricks connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/databricks-connector) | Run Databricks jobs in your workflows with the Databricks connector. | | [Memory Machine Cloud connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/mmcloud-connector) | Execute tasks using the MemVerge Memory Machine Cloud connector. | | [OpenAI Batch connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/openai-batch-connector) | Submit requests for asynchronous batch processing on OpenAI. | | [Perian connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/perian-connector) | Execute tasks on Perian Job Platform. | | [Sensor connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/sensor) | Run sensor jobs in your workflows with the sensor connector. | | [Slurm connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/slurm-connector) | Run Slurm jobs in your workflows with the Slurm connector. | | [Snowflake connector](https://www.union.ai/docs/v1/flyte/integrations/connectors//snowflake-connector) | Run Snowflake jobs in your workflows with the Snowflake connector. | [Flytekit plugins](https://www.union.ai/docs/v1/flyte/integrations/#flytekit-plugins) --------------------------------------------------------------------------------------- Flytekit plugins can be implemented purely in Python, unit tested locally, and allow extending Flytekit functionality. For comparison, these plugins can be thought of like [Airflow operators](https://airflow.apache.org/docs/apache-airflow/stable/howto/operator/index.html) . | Plugin | Description | | --- | --- | | [Comet](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/comet-ml-plugin) | `comet-ml`: Comet’s machine learning platform. | | [DBT](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/dbt-plugin) | Run and test your `dbt` pipelines in Flyte. | | [Dolt](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/dolt-plugin) | Version your SQL database with `dolt`. | | [DuckDB](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/duckdb-plugin) | Run analytical queries using DuckDB. | | [Great Expectations](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/greatexpectations-plugin) | Validate data with `great_expectations`. | | [Memray](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/memray-plugin) | `memray`: Memory profiling with memray. | | [MLFlow](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/mlflow-plugin) | `mlflow`: the open standard for model tracking. | | [Modin](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/modin-plugin) | Scale pandas workflows with `modin`. | | [Neptune](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/neptune-plugin) | `neptune`: Neptune is the MLOps stack component for experiment tracking. | | [NIM](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/nim-plugin) | Serve optimized model containers with NIM. | | [Ollama](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/ollama-plugin) | Serve fine-tuned LLMs with Ollama in a Flyte workflow. | | [ONNX](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/onnx-plugin) | Convert ML models to ONNX models seamlessly. | | [Pandera](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/pandera-plugin) | Validate pandas dataframes with `pandera`. | | [Papermill](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/papermill-plugin) | Execute Jupyter Notebooks with `papermill`. | | [SQL](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/sql-plugin) | Execute SQL queries as tasks. | | [Weights and Biases](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/wandb-plugin) | `wandb`: Machine learning platform to build better models faster. | | [WhyLogs](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/whylogs-plugin) | `whylogs`: the open standard for data logging. | ### [Using Flytekit plugins](https://www.union.ai/docs/v1/flyte/integrations/#using-flytekit-plugins) Data is automatically marshalled and unmarshalled in and out of the plugin. Users should mostly implement the `flytekit.core.base-task.PythonTask` API defined in Flytekit. Flytekit plugins are lazily loaded and can be released independently like libraries. The naming convention is `flytekitplugins-*`, where `*` indicates the package to be integrated into Flytekit. For example, `flytekitplugins-papermill` enables users to author Flytekit tasks using [Papermill](https://papermill.readthedocs.io/en/latest/) . You can find the plugins maintained by the core Flyte team [here](https://github.com/flyteorg/flytekit/tree/master/plugins) . [Native backend plugins](https://www.union.ai/docs/v1/flyte/integrations/#native-backend-plugins) --------------------------------------------------------------------------------------------------- Native backend plugins can be executed without any external service dependencies because the compute is orchestrated by Flyte itself, within its provisioned Kubernetes clusters. | Plugin | Description | | --- | --- | | [Kubeflow PyTorch](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/kfpytorch-plugin) | Run distributed PyTorch training jobs using `Kubeflow`. | | [Kubeflow TensorFlow](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/kftensorflow-plugin) | Run distributed TensorFlow training jobs using `Kubeflow`. | | [Kubernetes cluster Dask jobs](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/k8s-dask-plugin) | Run Dask jobs on a Kubernetes Cluster. | | [Kubernetes cluster Spark jobs](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/k8s-spark-plugin) | Run Spark jobs on a Kubernetes Cluster. | | [MPI Operator](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/kfmpi-plugin) | Run distributed deep learning training jobs using Horovod and MPI. | | [Ray](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/ray-plugin) | Run Ray jobs on a K8s Cluster. | [External service backend plugins](https://www.union.ai/docs/v1/flyte/integrations/#external-service-backend-plugins) ----------------------------------------------------------------------------------------------------------------------- As the term suggests, these plugins rely on external services to handle the workload defined in the Flyte task that uses the plugin. | Plugin | Description | | --- | --- | | [AWS Athena](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/athena-plugin) | Execute queries using AWS Athena | | [AWS Batch](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/aws-batch-plugin) | Running tasks and workflows on AWS batch service | | [Flyte Interactive](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/flyteinteractive-plugin) | Execute tasks using Flyte Interactive to debug. | | [Hive](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/hive-plugin) | Run Hive jobs in your workflows. | [Enabling backend plugins](https://www.union.ai/docs/v1/flyte/integrations/#enabling-backend-plugins) ------------------------------------------------------------------------------------------------------- To enable a backend plugin, you must add the `ID` of the plugin to the enabled plugins list. The `enabled-plugins` is available under the `tasks > task-plugins` section of FlytePropeller’s configuration. The plugin configuration structure is defined [here](https://pkg.go.dev/github.com/flyteorg/flytepropeller@v0.6.1/pkg/controller/nodes/task/config#TaskPluginConfig) . An example of the config follows: tasks: task-plugins: enabled-plugins: - container - sidecar - k8s-array default-for-task-types: container: container sidecar: sidecar container_array: k8s-array **Finding the `ID` of the backend plugin** To find the `ID` of the backend plugin, look at the source code of the plugin. For examples, in the case of Spark, the value of `ID` is used [here](https://github.com/flyteorg/flyteplugins/blob/v0.5.25/go/tasks/plugins/k8s/spark/spark.go#L424) , defined as [spark](https://github.com/flyteorg/flyteplugins/blob/v0.5.25/go/tasks/plugins/k8s/spark/spark.go#L41) . [Flyte operators](https://www.union.ai/docs/v1/flyte/integrations/#flyte-operators) ------------------------------------------------------------------------------------- Flyte can be integrated with other orchestrators to help you leverage Flyte’s constructs natively within other orchestration tools. | Operator | Description | | --- | --- | | [Airflow](https://www.union.ai/docs/v1/flyte/integrations/flyte-operators/airflow-plugin) | Trigger Flyte executions from Airflow. | LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/integrations/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/integrations/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Model training | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Model training ============== Understand how machine learning models can be trained from within Flyte, with an added advantage of orchestration benefits. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/tutorials/model-training/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Data input/output | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Data input/output ================= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. Flyte being a data-aware orchestration platform, types play a vital role within it. This section provides an introduction to the wide range of data types that Flyte supports. These types serve a dual-purpose by not only validating the data but also enabling seamless transfer of data between local and cloud storage. They enable: * Data lineage * Memoization * Auto parallelization * Simplifying access to data * Auto generated CLI and launch UI For a more comprehensive understanding of how Flyte manages data, refer to [Understand How Flyte Handles Data](https://www.union.ai/docs/v1/flyte/architecture/data-handling) . [Mapping Python to Flyte types](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/#mapping-python-to-flyte-types) --------------------------------------------------------------------------------------------------------------------------------- Flytekit automatically translates most Python types into Flyte types. Here’s a breakdown of these mappings: | Python Type | Flyte Type | Conversion | Comment | | --- | --- | --- | --- | | `int` | `Integer` | Automatic | Use Python 3 type hints. | | `float` | `Float` | Automatic | Use Python 3 type hints. | | `str` | `String` | Automatic | Use Python 3 type hints. | | `bool` | `Boolean` | Automatic | Use Python 3 type hints. | | `bytes`/`bytearray` | `Binary` | Not Supported | You have the option to employ your own custom typetransformer. | | `complex` | NA | Not Supported | You have the option to employ your own custom type transformer. | | `datetime.timedelta` | `Duration` | Automatic | Use Python 3 type hints. | | `datetime.datetime` | `Datetime` | Automatic | Use Python 3 type hints. | | `datetime.date` | `Datetime` | Automatic | Use Python 3 type hints. | | `typing.List[T]` / `list[T]` | `Collection [T]` | Automatic | Use `typing.List[T]` or `list[T]`, where `T` canrepresent one of the other supported types listed in the table. | | `typing.Iterator[T]` | `Collection [T]` | Automatic | Use `typing.Iterator[T]`, where `T` can represent one of the other supported types listed in the table. | | File / file-like / `os.PathLike` | `FlyteFile` | Automatic | If you’re using `file` or `os.PathLike` objects,Flyte will default to the binary protocol for the file. When using `FlyteFile["protocol"]`, it is assumedthat the file is in the specified protocol, such as ‘jpg’, ‘png’, ‘hdf5’, etc. | | Directory | `FlyteDirectory` | Automatic | When using `FlyteDirectory["protocol"]`, it is assumed that all thefiles belong to the specified protocol. | | `typing.Dict[str, V]` / `dict[str, V]` | `Map[str, V]` | Automatic | Use `typing.Dict[str, V]` or `dict[str, V`, where `V` can be one of the other supported types in the table, including a nested dictionary. |\ | `dict` | JSON (`struct.pb`) | Automatic | Use `dict`. It’s assumed that the untyped dictionary can beconverted to JSON. However, this may not always be possible and could result in a `RuntimeError`. |\ | `@dataclass` | `Struct` | Automatic | The class should be a pure value class annotated with the `@dataclass`decorator. |\ | `np.ndarray` | File | Automatic | Use `np.ndarray` as a type hint. |\ | `pandas.DataFrame` | Structured Dataset | Automatic | Use `pandas.DataFrame` as a type hint. Pandas columntypes aren’t preserved. |\ | `polars.DataFrame` | Structured Dataset | Automatic | Use `polars.DataFrame` as a type hint. Polars columntypes aren’t preserved. |\ | `polars.LazyFrame` | Structured Dataset | Automatic | Use `polars.LazyFrame` as a type hint. Polars columntypes aren’t preserved. |\ | `pyspark.DataFrame` | Structured Dataset | To utilize the type, install the `flytekitplugins-spark` plugin. | Use `pyspark.DataFrame` as a type hint. |\ | `pydantic.BaseModel` | `Map` | To utilize the type, install the `pydantic` module. | Use `pydantic.BaseModel`as a type hint. |\ | `torch.Tensor` / `torch.nn.Module` | File | To utilize the type, install the `torch` library. | Use `torchTensor` or `torch.nn.Module` as a type hint, and you can use their derived types. |\ | `tf.keras.Model` | File | To utilize the type, install the `tensorflow` library. | Use `tf.keras.Model` andits derived types. |\ | `sklearn.base.BaseEstimator` | File | To utilize the type, install the `scikit-learn` library. | Use `sklearnbase.BaseEstimator` and its derived types. |\ | User defined types | Any | Custom transformers | The `FlytePickle` transformer is the default option, but youcan also define custom transformers. For instructions on building custom type transformers, please refer to [this section](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/custom-types)
. |\ \ LLM-optimized\ \ [This page](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/page.md)\ [This section in one file](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/section.md)\ [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt)\ \ On this page\ \ 404\ \ Page not found\ \ Showing closest match --- # Flyte deployment | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Flyte deployment ================ An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/deployment/flyte-deployment/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section covers Flyte deployment. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/deployment/flyte-deployment/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/deployment/flyte-deployment/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Connector setup | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Connector setup =============== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section shows you how to set up connectors in your Flyte deployment. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Getting support | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Getting support =============== You can reach the support team through any of the channels below for all severity levels. **Urgent (Sev 1) issues** For blocking issues that require immediate attention, include **Sev 1 - Urgent** in your request through any channel. This pages the 24×7 on-call team. [Severity levels](https://www.union.ai/docs/v2/union/support/#severity-levels) -------------------------------------------------------------------------------- | Severity | Description | | --- | --- | | Sev 1 – Urgent | Platform is down, inoperable, inaccessible, or unavailable; operation has materially ceased. No workaround available. | | Sev 2 – High | Platform is severely limited or degraded; major functions are not performing properly; situation is significantly impacting operations. | | Sev 3 – Medium | Minor issues with the platform. Issues may have viable workarounds immediately available. | | Sev 4 – Low | Events have no impact on the platform; includes feature requests. Workaround works well for any functionality issues. | [Response time targets](https://www.union.ai/docs/v2/union/support/#response-time-targets) -------------------------------------------------------------------------------------------- Initial response time targets by severity level and support tier: | Severity | Standard | Premier | Advanced | Priority | | --- | --- | --- | --- | --- | | Sev 1 – Urgent | 24 hrs | 3 hrs | 2 hrs | 1 hr | | Sev 2 – High | 36 hrs | 6 hrs | 4 hrs | 2 hrs | | Sev 3 – Medium | 48 hrs | 24 hrs | 12 hrs | 4 hrs | | Sev 4 – Low | 48 hrs | 48 hrs | 24 hrs | 12 hrs | [Shared Slack channel](https://www.union.ai/docs/v2/union/support/#shared-slack-channel) ------------------------------------------------------------------------------------------ If you have a shared Slack channel with Union, opening a support ticket is easy. Simply tag `@union` or react to any message with a 🎫 emoji. Doing so will open a form where you can add additional details and set the severity. [Support portal](https://www.union.ai/docs/v2/union/support/#support-portal) ------------------------------------------------------------------------------ The support portal at [support.union.ai](https://support.union.ai/) lets you file tickets and view the status of open and closed requests. Log in with your company email address and click the **Create Ticket** button in the top right to submit a new request. [Union Cloud console](https://www.union.ai/docs/v2/union/support/#union-cloud-console) ---------------------------------------------------------------------------------------- Click the **Get help** button within the Union web console to open a support request. You can attach a screenshot or record a video directly from the console to include with your report. [Email](https://www.union.ai/docs/v2/union/support/#email) ------------------------------------------------------------ Send an email to [\[email protected\]](https://www.union.ai/cdn-cgi/l/email-protection#98ebede8e8f7eaecd8edf6f1f7f6b6f9f1) with the severity level in the subject line and all relevant details in the body: a summary and description of the issue, intended vs. actual behavior, error messages, logs, and a link to the relevant execution if applicable. LLM-optimized [This page](https://www.union.ai/docs/v2/union/support/page.md) [Full docs index](https://www.union.ai/docs/v2/union/llms.txt) On this page 404 Page not found Showing closest match --- # Community | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Community ========= Flyte is an open source project that is built and maintained by a community of contributors. Union AI is the primary maintainer of Flyte and developer of Union.ai, a closed source commercial product that is built on top of Flyte. Since the success of Flyte is essential to the success of Union.ai, the company is dedicated to building and expanding the Flyte open source project and community. For information on how to get involved and how to keep in touch, see [Joining the community](https://www.union.ai/docs/v1/flyte/community/joining-the-community) . [Contributing to the codebase](https://www.union.ai/docs/v1/flyte/community/#contributing-to-the-codebase) ------------------------------------------------------------------------------------------------------------ The full Flyte codebase is open source and available on GitHub. If you are interested, see [Contributing code](https://www.union.ai/docs/v1/flyte/community/contributing-code) . [Contributing to documentation](https://www.union.ai/docs/v1/flyte/community/#contributing-to-documentation) -------------------------------------------------------------------------------------------------------------- Union AI maintains and hosts both Flyte and Union documentation at [www.union.ai/docs](https://www.union.ai/docs/v1) . The two sets of documentation are deeply integrated, as the Union product is built on top of Flyte. To better maintain both sets of docs, they are hosted in the same repository (but rendered so that you can choose to view one or the other). Both the Flyte and Union documentation are open source. Flyte community members and Union customers are both welcome to contribute to the documentation. If you are interested, see [Contributing documentation and examples](https://www.union.ai/docs/v1/flyte/community/contributing-docs) . LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/community/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Platform configuration | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Platform configuration ====================== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section covers configuring Flyte for deeper integrations with existing infrastructure. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/deployment/flyte-configuration/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Configuration reference | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Configuration reference ======================= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/deployment/configuration-reference/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section all the supported configuration flags for all the Flyte components. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/deployment/configuration-reference/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/deployment/configuration-reference/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Plugins setup | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Plugins ======= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section includes the steps to configure integrations for the Flyte platform. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/deployment/flyte-plugins/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Native backend plugins | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Native backend plugins ====================== This section covers native backend plugins. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/integrations/native-backend-plugins/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Executions | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Executions ========== Executions are instances of workflows, nodes or tasks created in the system as a result of a user-requested execution or a scheduled execution. Execution IDs are unique within a given project domain, ensuring that no two executions within the same domain can have the same ID. [Typical Flow Using Flytectl](https://www.union.ai/docs/v1/flyte/architecture/executions/#typical-flow-using-flytectl) ------------------------------------------------------------------------------------------------------------------------ * When an execution of a workflow is triggered using UI/Flytecli/other stateless systems, the system first calls the getLaunchPlan endpoint and retrieves a launch plan matching the given version. The launch plan definition includes definitions of all input variables declared for the workflow. * The user-side component then ensures that all the required inputs are supplied and requests the FlyteAdmin service for an execution. * The FlyteAdmin service validates the inputs, ensuring that they are all specified and, if required, within the declared bounds. * FlyteAdmin then fetches the previously validated and compiled workflow closure and translates it to an executable format with all the inputs. * This executable workflow is launched on Kubernetes with an execution record in the database. ![Flyte executions overview](https://www.union.ai/docs/v1/flyte/_static/images/architecture/executions/flyte-executions-overview.svg) LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/executions/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Registration | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Registration ============ During registration, Flyte validates the workflow structure and saves the workflow. The registration process also updates the workflow graph. ![Flyte registration overview](https://www.union.ai/docs/v1/flyte/_static/images/architecture/registration/flyte-registration-overview.svg) [Typical Flow](https://www.union.ai/docs/v1/flyte/architecture/registration/#typical-flow) -------------------------------------------------------------------------------------------- The following steps elaborate on the specifics of the registration process: * Define the tasks using the Flytekit Task Definition language. * Define a workflow using the Flytekit Workflow Definition language. * Use `flytectl register` to compile the tasks into their serialized representation as described in Flyte Specification language. During this, the task representation is bound to a container that constitutes the code for the task. This associated entity is registered with FlyteAdmin using the `registerTask` API. * Use `flytectl register` to compile the workflow into their serialized representation as described in Flyte Specification language. The referenced tasks are replaced by their FlyteAdmin registered identifiers, obtained in the previous step. The associated entity is registered with FlyteAdmin using the `registerWorkflow` API. * Launch an execution using the FlyteAdmin launch execution API, which requires the necessary inputs provided. This is automatically done if the user uses `flytectl` to launch the execution. * Use the FlyteAdmin `read` APIs to get details of the execution, monitor it to completion, or retrieve a historical execution. * OR use the FlyteConsole to visualize the execution in real time as it progresses or visualize any historical execution. The console makes it easy to view debugging information for the execution. * Set specific rules such as notification on failure or success or publish all events in the execution to a pub-sub system. * Query the datastore to get a summary of all the executions and the compute resources consumed. Workflows and tasks are purely specifications and can be provided using tools like YAML, JSON, protobuf binary or any other programming language, and hence registration is possible using other tools. Contributions welcome! [Registration in the Backend](https://www.union.ai/docs/v1/flyte/architecture/registration/#registration-in-the-backend) -------------------------------------------------------------------------------------------------------------------------- When FlyteAdmin receives a workflow registration request, it uses the workflow compiler to compile and validate the workflow. It also fetches all the referenced tasks and creates a complete workflow closure, which is stored in the metastore. If the workflow compilation fails, the compiler returns an error to the client. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/registration/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Flytekit plugins | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Flytekit plugins ================ This section covers Flytekit plugins. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/integrations/flytekit-plugins/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Deprecated integrations | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Deprecated integrations ======================= This section covers deprecated integrations. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/integrations/deprecated-integrations/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Flyte operators | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Flyte operators =============== This section covers Flyte operators. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/integrations/flyte-operators/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # External service backend plugins | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) External service backend plugins ================================ This section covers external service backend plugins. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/integrations/external-service-backend-plugins/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Connectors | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Connectors ========== Connectors are long-running, stateless services that receive execution requests via gRPC and initiate jobs with appropriate external or internal services. Each connector service is a Kubernetes deployment that receives gRPC requests when users trigger a particular type of task. (For example, the BigQuery connector is tiggered by the invocation of a BigQuery tasks.) The connector service then initiates a job with the appropriate service. Connectors can be run locally as long as the appropriate connection secrets are locally available, since they are spawned in-process. Connectors are designed to be scalable and can handle large workloads efficiently, and decrease load on the core system, since they run outside it. You can also test connectors locally without having to change the backend configuration, streamlining workflow development. Connectors enable two key use cases: * **Asynchronously** launching jobs on hosted platforms (e.g. Databricks or Snowflake). * Calling external **synchronous** services, such as access control, data retrieval, or model inferencing. This section covers all currently available connectors: * [Airflow connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/airflow-connector) * [BigQuery connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/bigquery-connector) * [OpenAI ChatGPT connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/chatgpt-connector) * [OpenAI Batch connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/openai-batch-connector) * [Databricks connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/databricks-connector) * [Memory Machine Cloud connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/mmcloud-connector) * [Perian connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/perian-connector) * [Sagemaker connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/sagemaker-inference-connector) * [Sensor connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/sensor) * [Slurm connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/slurm-connector) * [Snowflake connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/snowflake-connector) [Creating a new connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/#creating-a-new-connector) ------------------------------------------------------------------------------------------------------------------ If none of the existing connectors meet your needs, you can implement your own connector. There are two types of connectors: **async** and **sync**. * **Async connectors** enable long-running jobs that execute on an external platform over time. They communicate with external services that have asynchronous APIs that support `create`, `get`, and `delete` operations. The vast majority of connectors are async connectors. * **Sync connectors** enable request/response services that return immediate outputs (e.g. calling an internal API to fetch data or communicating with the OpenAI API). While connectors can be written in any programming language since they use a protobuf interface, we currently only support Python connectors. We may support other languages in the future. ### [Async connector interface specification](https://www.union.ai/docs/v1/flyte/integrations/connectors/#async-connector-interface-specification) To create a new async connector, extend the `AsyncConnectorBase` and implement `create`, `get`, and `delete` methods. These methods must be idempotent. * `create`: This method is used to initiate a new job. Users have the flexibility to use gRPC, REST, or an SDK to create a job. * `get`: This method retrieves the job resource (job ID or output literal) associated with the task, such as a BigQuery job ID or Databricks task ID. * `delete`: Invoking this method will send a request to delete the corresponding job. For an example implementation, see the [BigQuery connector code](https://github.com/flyteorg/flytekit/blob/master/plugins/flytekit-bigquery/flytekitplugins/bigquery/connector.py) . ### [Sync connector interface specification](https://www.union.ai/docs/v1/flyte/integrations/connectors/#sync-connector-interface-specification) To create a new sync connector, extend the `SyncConnectorBase` class and implement a `do` method. This method must be idempotent. * `do`: This method is used to execute the synchronous task, and the worker in Flyte will be blocked until the method returns. For an example implementation, see the [ChatGPT connector code](https://github.com/flyteorg/flytekit/blob/master/plugins/flytekit-openai/flytekitplugins/openai/chatgpt/connector.py) . ### [Testing your connector locally](https://www.union.ai/docs/v1/flyte/integrations/connectors/#testing-your-connector-locally) To test your connector locally, create a class for the connector task that inherits from [`AsyncConnectorExecutorMixin`](https://github.com/flyteorg/flytekit/blob/1bc8302bb7a6cf4c7048a7f93627ee25fc6b88c4/flytekit/extend/backend/base_connector.py#L354) . This mixin can handle both asynchronous tasks and synchronous tasks and allows Flytekit to mimic the system’s behavior in calling the connector. For testing examples, see the [BigQuery connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/bigquery-connector#local-testing) and [Databricks connector](https://www.union.ai/docs/v1/flyte/integrations/connectors/databricks-connector#local-testing) documentation. [Enabling a connector in your Flyte deployment](https://www.union.ai/docs/v1/flyte/integrations/connectors/#enabling-a-connector-in-your-flyte-deployment) ------------------------------------------------------------------------------------------------------------------------------------------------------------ For information on setting up a connector in your Flyte deployment, see [Deployment > Connector setup](https://www.union.ai/docs/v1/flyte/deployment/flyte-connectors) LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/integrations/connectors/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Workflow timeline | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Workflow timeline ================= The illustration below shows the timeline view of a workflow execution. ![Flyte workflow timeline](https://www.union.ai/docs/v1/flyte/_static/images/architecture/workflow-timeline/flyte-workflow-timeline.svg) This illustration refers to a simple workflow, with 2 nodes `N1` an `N2`. This can be represented as follows: graph LR; Start --> N1; N1 --> N2; N2 --> End; [Acceptance Latency](https://www.union.ai/docs/v1/flyte/architecture/workflow-timeline/#acceptance-latency) ------------------------------------------------------------------------------------------------------------- Every workflow starts in the acceptance phase. Acceptance refers to the time between FlyteAdmin receiving an execution request and FlytePropeller evaluating the first round of workflow. Usually, within this phase, the Kubernetes queuing latency is the largest contributor to latency where the overall acceptance latency of less than five seconds is desirable. [Transition Latency](https://www.union.ai/docs/v1/flyte/architecture/workflow-timeline/#transition-latency) ------------------------------------------------------------------------------------------------------------- Transition latency refers to the time between successive node executions, that is, between `N1` and `N2`. For the first node `N1`, this latency also encapsulates executing the start node. Similarly, the last node also encapsulates executing end node. `Start Node` and `End Node` are capstones inserted to mark the beginning and end of the DAG. The latency involves time consumed to: 1. Gather outputs for a node after the node completes execution. 2. Send an observation event to FlyteAdmin. Failing to do so will be regarded as an error and will be tried until it succeeds or system max retries are exhausted (the number of max system retries is configured to be 30 by default and can be altered per deployment). 3. Persist data to Kubernetes. 4. Receive the persisted object back from Kubernetes (as this process is eventually consistent using informer caches). 5. Gather inputs for a node before the node starts. 6. Send a queued event for the next node to FlyteAdmin (this is what is persisted and drives the UI/CLI and historical information). [Queuing Latency](https://www.union.ai/docs/v1/flyte/architecture/workflow-timeline/#queuing-latency) ------------------------------------------------------------------------------------------------------- Queuing latency is the time taken by Kubernetes to start the pod, other services to start the job, HTTP throttle to be met, or any rate-limiting that needs to be overcome. This is usually tied to the available resources and quota, and is out of control for Flyte. [Completion Latency](https://www.union.ai/docs/v1/flyte/architecture/workflow-timeline/#completion-latency) ------------------------------------------------------------------------------------------------------------- Completion latency is the time taken to mark the workflow as complete and accumulate outputs of a workflow after the last node completes its execution. [Overview of Various Latencies in FlytePropeller](https://www.union.ai/docs/v1/flyte/architecture/workflow-timeline/#overview-of-various-latencies-in-flytepropeller) ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- Description of main events for workflow execution: | Events | Description | | --- | --- | | Acceptance | Measures the time between when we receive service call to create an Execution (Unknown) and when it has moved to Queued. | | Transition Latency | Measures the latency between two consecutive node executions, the time spent in Flyte engine. | | Queuing Latency | Measures the latency between the time a node’s been queued to the time the handler reported the executable moved to running state. | | Task Execution | Actual time spent executing user code | | Repeat steps #2 to #4 for every task | | | Transition Latency | See #2 | | Completion Latency | Measures the time between when the WF moved to succeeding/failing state and when it finally moved to a terminal state. | LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/workflow-timeline/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Control Plane | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Control Plane ============= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/architecture/control-plane/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section covers key architectural and implementation details for the Flyte control plane. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/control-plane/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/architecture/control-plane/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Workflow state transitions | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Workflow state transitions ========================== The following state diagram illustrates a high-level view of the state transitions that a workflow with a single task and node would go through as the user observes success. flowchart TD id1(( )) id1 --> Ready Ready --> Running subgraph Running id2(( )) id2 --> NodeQueued NodeQueued --> NodeRunning subgraph NodeRunning id3(( )) id3 --> TaskQueued TaskQueued --> TaskRunning TaskRunning --> TaskSuccess end TaskSuccess --> NodeSuccess end NodeSuccess --> Success The following sections explain the various observable (and some hidden) states for workflow, node, and task state transitions. [Workflow States](https://www.union.ai/docs/v1/flyte/architecture/workflow-state-transitions/#workflow-states) ---------------------------------------------------------------------------------------------------------------- flowchart TD Queued -->|On system errors more than threshold| Aborted Queued --> Ready Ready--> |Write inputs to workflow| Running Running--> |On system error| Running Running--> |On all Nodes Success| Succeeding Succeeding--> |On successful event send to Admin| Succeeded Succeeding--> |On system error| Succeeding Ready--> |On precondition failure| Failing Running--> |On any Node Failure| Failing Ready--> |On user initiated abort| Aborting Running--> |On user initiated abort| Aborting Succeeding--> |On user initiated abort| Aborting Failing--> |If Failure node exists| HandleFailureNode Failing--> |On user initiated abort| Aborting HandleFailureNode--> |On completing failure node| Failed HandleFailureNode--> |On user initiated abort| Aborting Failing--> |On successful send of Failure node| Failed Aborting--> |On successful event send to Admin| Aborted A workflow always starts in the `Ready` state and ends either in `Failed`, `Succeeded`, or `Aborted` state. Any system error within a state causes a retry on that state. These retries are capped by `system retries ` which eventually lead to an `Aborted` state if the failure persists. Every transition between states is recorded in FlyteAdmin using :std:ref:`workflowexecutionevent`. The phases in the above state diagram are captured in the admin database as specified here `workflowexecution.phase` and are sent as a part of the Execution event. The state machine specification for the illustration can be found [here](https://mermaid-js.github.io/mermaid-live-editor/#/edit/eyJjb2RlIjoic3RhdGVEaWFncmFtLXYyXG4gICAgWypdIC0tPiBBYm9ydGVkIDogT24gc3lzdGVtIGVycm9ycyBtb3JlIHRoYW4gdGhyZXNob2xkXG4gICAgWypdIC0tPiBSZWFkeVxuICAgIFJlYWR5IC0tPiBSdW5uaW5nIDogV3JpdGUgaW5wdXRzIHRvIHdvcmtmbG93XG4gICAgUnVubmluZyAtLT4gUnVubmluZyA6IE9uIHN5c3RlbSBlcnJvclxuICAgIFJ1bm5pbmcgLS0-IFN1Y2NlZWRpbmcgOiBPbiBhbGwgTm9kZXMgU3VjY2Vzc1xuICAgIFN1Y2NlZWRpbmcgLS0-IFN1Y2NlZWRlZCA6IE9uIHN1Y2Nlc3NmdWwgZXZlbnQgc2VuZCB0byBBZG1pblxuICAgIFN1Y2NlZWRpbmcgLS0-IFN1Y2NlZWRpbmcgOiBPbiBzeXN0ZW0gZXJyb3JcbiAgICBSZWFkeSAtLT4gRmFpbGluZyA6IE9uIHByZWNvbmRpdGlvbiBmYWlsdXJlXG4gICAgUnVubmluZyAtLT4gRmFpbGluZyA6IE9uIGFueSBOb2RlIEZhaWx1cmVcbiAgICBSZWFkeSAtLT4gQWJvcnRlZCA6IE9uIHVzZXIgaW5pdGlhdGVkIGFib3J0XG4gICAgUnVubmluZyAtLT4gQWJvcnRlZCA6IE9uIHVzZXIgaW5pdGlhdGVkIGFib3J0XG4gICAgU3VjY2VlZGluZyAtLT4gQWJvcnRlZCA6IE9uIHVzZXIgaW5pdGlhdGVkIGFib3J0XG5cbiAgICBGYWlsaW5nIC0tPiBIYW5kbGVGYWlsdXJlTm9kZSA6IElmIEZhaWx1cmUgbm9kZSBleGlzdHNcbiAgICBGYWlsaW5nIC0tPiBBYm9ydGVkIDogT24gdXNlciBpbml0aWF0ZWQgYWJvcnRcbiAgICBIYW5kbGVGYWlsdXJlTm9kZSAtLT4gRmFpbGVkIDogT24gY29tcGxldGluZyBmYWlsdXJlIG5vZGVcbiAgICBIYW5kbGVGYWlsdXJlTm9kZSAtLT4gQWJvcnRlZCA6IE9uIHVzZXIgaW5pdGlhdGVkIGFib3J0XG4gICAgRmFpbGluZyAtLT4gRmFpbGVkIDogT24gc3VjY2Vzc2Z1bCBzZW5kIG9mIEZhaWx1cmUgbm9kZVxuICAgICIsIm1lcm1haWQiOnt9LCJ1cGRhdGVFZGl0b3IiOmZhbHNlfQ) . [Node States](https://www.union.ai/docs/v1/flyte/architecture/workflow-state-transitions/#node-states) -------------------------------------------------------------------------------------------------------- flowchart TD id1(( )) id1-->NotYetStarted id1-->|Will stop the node execution |Aborted NotYetStarted-->|If all upstream nodes are ready, i.e, inputs are ready | Queued NotYetStarted--> |If the branch was not taken |Skipped Queued-->|Start task execution- attempt 0 | Running Running-->|If task timeout has elapsed and retry\_attempts >= max\_retries|TimingOut Running-->|Internal state|Succeeding Running-->|For dynamic nodes generating workflows| DynamicRunning DynamicRunning-->TimingOut DynamicRunning-->RetryableFailure TimingOut-->|If total node timeout has elapsed|TimedOut DynamicRunning-->Succeeding Succeeding-->|User observes the task as succeeded| Succeeded Running-->|on retryable failure| RetryableFailure RetryableFailure-->|if retry\_attempts < max\_retries|Running RetryableFailure-->|retry\_attempts >= max\_retries|Failing Failing-->Failed Succeeded-->id2(( )) Failed-->id2(( )) This state diagram illustrates the node transition through various states. This is the core finite state machine for a node. From the user’s perspective, a workflow simply consists of a sequence of tasks. But to Flyte, a workflow internally creates a meta entity known as **node**. Once a Workflow enters the `Running` state, it triggers the phantom `start node` of the workflow. The `start node` is considered to be the entry node of any workflow. The `start node` begins by executing all its child-nodes using a modified depth first search algorithm recursively. Nodes can be of different types as listed below, but all the nodes traverse through the same transitions: 1. Start Node - Only exists during the execution and is not modeled in the core spec. 2. :std:ref:`Task Node` 3. :std:ref:`Branch Node` 4. :std:ref:`Workflow Node` 5. Dynamic Node - Just a task node that does not return output but constitutes a dynamic workflow. When the task runs, it remains in the `RUNNING` state. Once the task completes and Flyte starts executing the dynamic workflow, the overarching node that contains both the original task and the dynamic workflow enters `DYNAMIC_RUNNING` state. 6. End Node - Only exists during the execution and is not modeled in the core spec Every transition between states is recorded in FlyteAdmin using `nodeexecutionevent`. Every `NodeExecutionEvent` can have any :std:ref:`nodeexecution.phase`. The state machine specification for the illustration can be found [here](https://mermaid-js.github.io/mermaid-live-editor/#/edit/eyJjb2RlIjoic3RhdGVEaWFncmFtLXYyXG4gICAgWypdIC0tPiBOb3RZZXRTdGFydGVkXG4gICAgWypdIC0tPiBBYm9ydGVkIDogV2lsbCBzdG9wIHRoZSBub2RlIGV4ZWN1dGlvblxuICAgIE5vdFlldFN0YXJ0ZWQgLS0-IFF1ZXVlZCA6IElmIGFsbCB1cHN0cmVhbSBub2RlcyBhcmUgcmVhZHkgaS5lLCBpbnB1dHMgYXJlIHJlYWR5XG4gICAgTm90WWV0U3RhcnRlZCAtLT4gU2tpcHBlZCA6IElmIHRoZSBicmFuY2ggd2FzIG5vdCB0YWtlblxuICAgIFF1ZXVlZCAtLT4gUnVubmluZyA6IFN0YXJ0IHRhc2sgZXhlY3V0aW9uIC0gYXR0ZW1wdCAwXG4gICAgUnVubmluZyAtLT4gVGltaW5nT3V0IDogSWYgdGFzayB0aW1lb3V0IGhhcyBlbGFwc2VkIGFuZCByZXRyeV9hdHRlbXB0cyA-PSBtYXhfcmV0cmllc1xuICAgIFRpbWluZ091dCAtLT4gVGltZWRPdXQgOiBJdCB0b3RhbCBub2RlIHRpbWVvdXQgaGFzIGVsYXBzZWRcbiAgICBSdW5uaW5nIC0tPiBSZXRyeWFibGVGYWlsdXJlIDogb24gcmV0cnlhYmxlIGZhaWx1cmVcbiAgICBSdW5uaW5nIC0tPiBEeW5hbWljUnVubmluZyA6IEZvciBkeW5hbWljIG5vZGVzIGdlbmVyYXRpbmcgd29ya2Zsb3dzXG4gICAgUmV0cnlhYmxlRmFpbHVyZSAtLT4gUnVubmluZyA6IGlmIHJldHJ5X2F0dGVtcHRzIDwgbWF4X3JldHJpZXNcbiAgICBSZXRyeWFibGVGYWlsdXJlIC0tPiBGYWlsaW5nIDogcmV0cnlfYXR0ZW1wdHMgPj0gbWF4X3JldHJpZXNcbiAgICBGYWlsaW5nIC0tPiBGYWlsZWRcbiAgICBSdW5uaW5nIC0tPiBTdWNjZWVkaW5nIDogSW50ZXJuYWwgc3RhdGVcbiAgICBEeW5hbWljUnVubmluZyAtLT4gU3VjY2VlZGluZ1xuICAgIER5bmFtaWNSdW5uaW5nIC0tPiBSZXRyeWFibGVGYWlsdXJlXG4gICAgRHluYW1pY1J1bm5pbmcgLS0-IFRpbWluZ091dFxuICAgIFN1Y2NlZWRpbmcgLS0-IFN1Y2NlZWRlZCA6IFVzZXIgb2JzZXJ2ZXMgdGhlIHRhc2sgYXMgc3VjY2VlZGVkXG4gICAgU3VjY2VlZGVkIC0tPiBbKl1cbiAgICBGYWlsZWQgLS0-IFsqXVxuIiwibWVybWFpZCI6e30sInVwZGF0ZUVkaXRvciI6ZmFsc2V9) . [Task States](https://www.union.ai/docs/v1/flyte/architecture/workflow-state-transitions/#task-states) -------------------------------------------------------------------------------------------------------- flowchart TD id1(( )) id1-->|Aborted by NodeHandler- timeouts, external abort, etc,.| NotReady id1-->Aborted NotReady-->|Optional-Blocked on resource quota or resource pool | WaitingForResources WaitingForResources--> |Optional- Has been submitted, but hasn't started |Queued Queued-->|Optional- Prestart initialization | Initializing Initializing-->|Actual execution of user code has started|Running Running-->|Successful execution|Success Running-->|Failed with a retryable error|RetryableFailure Running-->|Unrecoverable failure, will stop all execution|PermanentFailure Success-->id2(( )) RetryableFailure-->id2(( )) PermanentFailure-->id2(( )) The state diagram above illustrates the various states through which a task transitions. This is the core finite state machine for a task. Every transition between states is recorded in FlyteAdmin using :`taskexecutionevent`. Every `TaskExecutionEvent` can have any `taskexecution.phase`. The state machine specification for the illustration can be found [here](https://mermaid-js.github.io/mermaid-live-editor/#/edit/eyJjb2RlIjoic3RhdGVEaWFncmFtLXYyXG4gICAgWypdIC0tPiBOb3RSZWFkeVxuICAgIFsqXSAtLT4gQWJvcnRlZCA6IEFib3J0ZWQgYnkgTm9kZUhhbmRsZXIgLSB0aW1lb3V0cywgZXh0cmVuYWwgYWJvcnQsIGV0Y1xuICAgIE5vdFJlYWR5IC0tPiBXYWl0aW5nRm9yUmVzb3VyY2VzIDogQmxvY2tlZCBvbiByZXNvdXJjZSBxdW90YSBvciByZXNvdXJjZSBwb29sIChvcHRpb25hbClcbiAgICBXYWl0aW5nRm9yUmVzb3VyY2VzIC0tPiBRdWV1ZWQgOiBIYXMgYmVlbiBzdWJtaXR0ZWQsIGJ1dCBoYXMgbm90IHN0YXJ0ZWQgKG9wdGlvbmFsKVxuICAgIFF1ZXVlZCAtLT4gSW5pdGlhbGl6aW5nIDogUHJlc3RhcnQgaW5pdGlhbGl6YXRpb24gKG9wdGlvbmFsKVxuICAgIEluaXRpYWxpemluZyAtLT4gUnVubmluZyA6IEFjdHVhbCBleGVjdXRpb24gb2YgdXNlciBjb2RlIGhhcyBzdGFydGVkXG4gICAgUnVubmluZyAtLT4gU3VjY2VzcyA6IFN1Y2Nlc3NmdWwgZXhlY3V0aW9uXG4gICAgUnVubmluZyAtLT4gUmV0cnlhYmxlRmFpbHVyZSA6IEZhaWxlZCB3aXRoIGEgcmV0cnlhYmxlIGVycm9yXG4gICAgUnVubmluZyAtLT4gUGVybWFuZW50RmFpbHVyZSA6IFVucmVjb3ZlcmFibGUgZmFpbHVyZSwgd2lsbCBzdG9wIGFsbCBleGVjdXRpb25cbiAgICBTdWNjZXNzIC0tPiBbKl1cbiAgICBSZXRyeWFibGVGYWlsdXJlIC0tPiBbKl1cbiAgICBQZXJtYW5lbnRGYWlsdXJlIC0tPiBbKl1cbiIsIm1lcm1haWQiOnt9LCJ1cGRhdGVFZGl0b3IiOmZhbHNlfQ) . LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/workflow-state-transitions/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Development cycle | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Development cycle ================= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This section covers developing production-ready workflows for Flyte. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Versions | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Versions ======== One of the most important features and reasons for certain design decisions in Flyte is the need for machine learning and data practitioners to experiment. When users experiment, they do so in isolation and try multiple iterations. Unlike traditional software, the users must conduct multiple experiments concurrently with different environments, algorithms, etc. This may happen when multiple data scientists simultaneously iterate on the same workflow/pipeline. The cost of creating an independent infrastructure for each version is enormous and undesirable. It is beneficial to share the same centralized infrastructure, where the burden of maintaining the infrastructure is with a central infrastructure team, while the users can use it independently. This improves the cost of operation since the same infrastructure can be reused by multiple teams. Versioned workflows help users quickly reproduce prior results or identify the source of previous successful experiments. [Why do you need versioning?](https://www.union.ai/docs/v1/flyte/architecture/versions/#why-do-you-need-versioning) --------------------------------------------------------------------------------------------------------------------- Versioning is required to: * Work on the same project concurrently and identify the version/experiment that was successful. * Capture the environment for a version and independently launch it. * Visualize prior runs and tie them to experiment results. * Rollback to production deployments in case of failures with ease. * Execute multiple experiments in production, which may use different training or data processing algorithms. * Understand how a specific system evolved and answer questions related to the effectiveness of a specific strategy. Please note that Flyte currently does not support the use of colons in specifying versions. For instance, using “0:1” as a version number is not permitted. [Operational benefits of completely versioned workflows/pipelines](https://www.union.ai/docs/v1/flyte/architecture/versions/#operational-benefits-of-completely-versioned-workflowspipelines) ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- The entire workflow in Flyte is versioned and all tasks and entities are immutable which makes it possible to completely change the structure of a workflow between versions, without worrying about the consequences for the pipelines in production. This hermetic property makes it effortless to manage and deploy new workflow versions and is important for workflows that are long-running. If a workflow execution is in progress and another new workflow version has been activated, Flyte guarantees that the execution of the old version continues unhindered. Consider a scenario where you need to run all the previous executions if there’s a bug to be fixed. Simply fixing the bug in the task may not solve the problem. Moreover, fixing bugs involves code changes, which may affect the workflow structure. Flyte addresses this using two properties: 1. Since the entire workflow is versioned, changing the structure has no impact on the existing execution, and the workflow state won’t be corrupted. 2. Flyte provides caching/memoization of outputs. As long as the tasks and their behavior have not changed, it is possible to move them around and still recover their previous outputs, without having to rerun the tasks. This strategy will work even if the workflow changes are in a task. Let us take a sample workflow: graph TD; A-->B; B-->C; C-->D; In the above graph, let us assume that task `C` fails. It is then possible to simply fix `C` and `relaunch` the previous execution (maintaining the inputs etc). This will not re-run tasks `A`, and `B` as long as they are marked as `cache=True`. Now, let us consider that the only solution to fix the bug is to change the graph structure and introduce a new step `B1` that short circuits the execution to `D`: graph TD; A-->B; B-->B1; B1-->D; B1-->C; C-->D; The same `cache=True` will handle this complicated situation as well. [Why is versioning hard?](https://www.union.ai/docs/v1/flyte/architecture/versions/#why-is-versioning-hard) ------------------------------------------------------------------------------------------------------------- Git has become the de facto standard in version control for code, making it easy to work on branches, merge them, and revert unwanted changes. But achieving this for a live (running) algorithm usually requires the entire infrastructure to be associated and potentially re-created for every execution. [How is versioning tied to reproducibility?](https://www.union.ai/docs/v1/flyte/architecture/versions/#how-is-versioning-tied-to-reproducibility) --------------------------------------------------------------------------------------------------------------------------------------------------- Workflows can be reproduced without explicit versioning within the system. To reproduce a past experiment, users need to identify the source code and resurrect any dependencies that the code may have used (for example, TensorFlow 1.x instead of TensorFlow 2.x, or specific Python libraries). It is also required to instantiate the infrastructure that the previous version may have used. If not recorded, you’ll have to ensure that the previously used dataset (say) can be reconstructed. This is exactly how Flyte was conceived! In Flyte, every task is versioned, and it precisely captures the dependency set. For external tasks, memoization is recommended so that the constructed dataset can be cached on the Flyte side. This way, one can guarantee reproducible behavior from the external systems. Moreover, every piece of code is registered with the version of the code that was used to create the instance. Therefore, users can easily construct the data lineage for all the parts of the workflow. [What is the cost of versioning and reproducibility?](https://www.union.ai/docs/v1/flyte/architecture/versions/#what-is-the-cost-of-versioning-and-reproducibility) --------------------------------------------------------------------------------------------------------------------------------------------------------------------- One of the costs of versioning and allowing on-demand reproducibility is the need to re-instantiate the infrastructure from scratch. This may sometimes result in additional overhead. However, the advent of Docker containers and Kubernetes has made it possible to build a platform to achieve these goals. [What is the best way to version your tasks and workflows?](https://www.union.ai/docs/v1/flyte/architecture/versions/#what-is-the-best-way-to-version-your-tasks-and-workflows) --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- The best way to version tasks and workflows is to independently version every task with the Git SHA or hash of the entire code artifact. The workflows are also versioned using the Git SHA of the containing repository. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/versions/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Data catalog | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Data catalog ============ DataCatalog is a service to index parameterized, strongly-typed data artifacts across revisions. It allows clients to query artifacts based on meta information and tags. [How Flyte memoizes task executions on data catalog](https://www.union.ai/docs/v1/flyte/architecture/data-catalog/#how-flyte-memoizes-task-executions-on-data-catalog) ------------------------------------------------------------------------------------------------------------------------------------------------------------------------ Flyte memoizes task executions by creating artifacts in DataCatalog and associating meta information regarding the execution with the artifact. Let’s walk through what happens when a task execution is cached on DataCatalog. Every task instance is represented as a DataSet: Dataset { project: Flyte project the task was registered in domain: Flyte domain for the task execution name: flyte_task- version: -- } Every task execution is represented as an Artifact in the Dataset above: Artifact { id: uuid Metadata: [executionName, executionVersion] ArtifactData: [List of ArtifactData] } ArtifactData { Name: value: } To retrieve the Artifact, tag the Artifact with a hash of the input values for the memoized task execution: ArtifactTag { Name: flyte_cached- } When caching an execution, FlytePropeller will: 1. Create a dataset for the task. 2. 1. Create an artifact that represents the execution, along with the artifact data that represents the execution output. Tag the artifact with a unique hash of the input values. To ensure that the task execution is memoized, Flyte Propeller will: 1. Compute the tag by computing the hash of the input. 2. Check if a tagged artifact exists with that hash: * If it exists, we have a cache hit and the Propeller can skip the task execution. * If an artifact is not associated with the tag, Propeller needs to run the task. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/data-catalog/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Data handling | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Data handling ============= In Flyte, data is categorized into metadata and raw data to optimize data handling and improve performance and security. * **Metadata**: Small values, like integers and strings, are treated as “stack parameters” (passed by value). This metadata is globally accessible to Flyte components (FlytePropeller, FlyteAdmin, and other running pods/jobs). Each entry is limited to 10MB and is passed directly between tasks. On top of that, metadata allows in-memory computations for branches, partial outputs, and composition of multiple outputs as input for other tasks. * **Raw data**: Larger data, such as files and dataframes, are treated as “heap parameters” (passed by reference). Flyte stores raw data in an object store (e.g., S3), uploading it on first use and passing only a reference thereafter. Tasks can then access this data via Flyte’s automated download or streaming, enabling efficient access to large datasets without needing to transfer full copies. _Source code reference for auto-offloading value sizes limitation_: Flyte’s data separation avoids bottlenecks and security risks: * **Metadata** remains within Flyte’s control plane, making it accessible through the UI or CLI. * **Raw data** is accessible only by tasks, stored securely in an external blob store, preventing Flyte’s control plane from directly handling large data files. Moreover, a unique property of this separation is that all meta values are read by FlytePropeller engine and available on the UI or CLI from the control plane. [Example](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#example) ----------------------------------------------------------------------------------- Consider a basic Flyte task: @fl.task def my_task(m: int, n: str, o: FlyteFile) -> pd.DataFrame: ... In this task, `m`, `n`, and `o` are inputs: `m (int)` and `n (str)` are simple types, while `o` is a large, arbitrarily sized file. Flyte treats each differently: * **Metadata**: Small values like `m` and `n` are inlined within Flyte’s metadata and passed directly between tasks. * **Raw data**: Objects like `o` and the output `pd.DataFrame` are offloaded to an object store (e.g., S3), with only references retained in metadata. Flytekit TypeTransformers make it possible to use complex objects as if they are available locally, just like persistent filehandles. However, the Flyte backend only deals with the references. [Raw data path](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#raw-data-path) ----------------------------------------------------------------------------------------------- Every task can read/write its own data files. If `FlyteFile` or any natively supported type like `pandas.DataFrame` is used, Flyte will automatically offload and download data from the configured object-store paths. These paths are completely customizable per launch plan or execution. The default **raw output path** (prefix in an object store like S3/GCS) can be configured during registration as shown in `flytectl register` files. The argument `--outputLocationPrefix` allows us to set the destination directory for all the raw data produced. Flyte will create randomized folders in this path to store the data. To override the \*_raw output path_ (prefix in an object store like S3/GCS), you can specify an alternate location when invoking a Flyte execution in the launch form in the UI. In the local demo cluster, the default raw output path is configured to be the root of the local bucket. Hence Flyte will write all the raw data (reference types like blob, file, df/schema/parquet, etc.) under a path defined by the execution. [`LiteralType` and `Literal`](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#literaltype-and-literal) ----------------------------------------------------------------------------------------------------------------------- ### [Serialization time](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#serialization-time) When a task is declared with inputs and outputs, Flyte extracts the interface of the task and converts it to an internal representation called a TypedInterface. For each variable, a corresponding `LiteralType` is created. For example, consider the following Python function interface: @fl.task def my_task(a: int, b: str) -> FlyteFile: """ Description of my function :param a: My input integer :param b: My input string :return: My output file """ It is transformed the following: interface { inputs { variables { key: "a" value { type { simple: INTEGER } description: "My input Integer" } } variables { key: "b" value { type { simple: STRING } description: "My input string" } } } outputs { variables { key: "o0" value { type { blob { } } description: "My output File" } } } } ### [Runtime](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#runtime) At runtime, data passes through Flyte using `Literal` where the values are set. For files, the corresponding `Literal` is called `LiteralBlob (Blob)` which is a binary large object. Many different objects can be mapped to the underlying `Blob` or `Struct` types. For example, an image is a `Blob`, a `pandas.DataFrame` is a `Blob` of type `parquet`, etc. [Data movement](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#data-movement) ----------------------------------------------------------------------------------------------- Flyte is primarily a dataflow engine. It enables movement of data and provides an abstraction to enable movement of data between different languages. One implementation of Flyte is the current workflow engine. The workflow engine is responsible for moving data from a previous task to the next task. As explained previously, Flyte only deals with metadata and not the actual raw data. The illustration below explains how data flows from engine to the task and how that is transferred between tasks. The medium to transfer the data can change, and will change in the future. We could use fast metadata stores to speed up data movement or exploit locality. ### [Between FlytePropeller and tasks](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#between-flytepropeller-and-tasks) ![Flyte data movement](https://www.union.ai/docs/v1/flyte/_static/images/architecture/data-handling/flyte-data-movement.png) ### [Between tasks](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#between-tasks) ![Flyte data transfer](https://www.union.ai/docs/v1/flyte/_static/images/architecture/data-handling/flyte-data-transfer.png) ### [Practical example](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#practical-example) Let’s consider a simple example where we have some tasks that needs to operate huge dataframes. The first task reads a file from the object store, shuffles the data, saves to local disk, and passes the path to the next task: @fl.task() def task_read_and_shuffle_file(input_file: FlyteFile) -> FlyteFile: """ Reads the input file as a DataFrame, shuffles the rows, and writes the shuffled DataFrame to a new file. """ input_file.download() df = pd.read_csv(input_file.path) # Shuffle the DataFrame rows shuffled_df = df.sample(frac=1).reset_index(drop=True) output_file_path = "data_shuffle.csv" shuffled_df.to_csv(output_file_path, index=False) return FlyteFile(output_file_path) The second task reads the file from the previous task, removes a column, saves to local disk, and returns the path: @fl.task() def task_remove_column(input_file: FlyteFile, column_name: str) -> FlyteFile: """ Reads the input file as a DataFrame, removes a specified column, and outputs it as a new file. """ input_file.download() df = pd.read_csv(input_file.path) # remove column if column_name in df.columns: df = df.drop(columns=[column_name]) output_file_path = "data_finished.csv" df.to_csv(output_file_path, index=False) return FlyteFile(output_file_path) And here is the workflow: @fl.workflow def wf() -> FlyteFile: existed_file = FlyteFile("s3://custom-bucket/data.csv") shuffled_file = task_read_and_shuffle_file(input_file=existed_file) result_file = task_remove_column(input_file=shuffled_file, column_name="County") return result_file This example shows how to access an existing file in a MinIO bucket from the Flyte local demo cluster and pass it between tasks with FlyteFile. When a workflow outputs a local file as a FlyteFile, Flyte automatically uploads it to MinIO and provides an S3 URL for downstream tasks, no manual uploads needed. ### [Bringing in your own datastores for raw data](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#bringing-in-your-own-datastores-for-raw-data) Flytekit has a pluggable data persistence layer. This is driven by protocol. For example, it is theoretically possible to use S3 `s3://` for metadata and GCS `gcs://` for raw data. It is also possible to create your own protocol `my_fs://`, to change how data is stored and accessed. But for metadata, the data should be accessible to Flyte control plane. Data persistence is also pluggable. By default, it supports all major blob stores and uses an interface defined in Flytestdlib. ### [Deleting raw data in your own datastores](https://www.union.ai/docs/v1/flyte/architecture/data-handling/#deleting-raw-data-in-your-own-datastores) Flyte does not offer a direct function to delete raw data stored in external datastores like S3 or GCS. However, you can manage deletion by configuring a lifecycle policy within your datastore service. If caching is enabled in your Flyte task, ensure that the `max-cache-age` is set to be shorter than the lifecycle policy in your datastore to prevent potential data inconsistency issues. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/data-handling/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Joining the community | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Joining the community ===================== Keeping the lines of communication open is important in growing and maintain the Flyte community. Please join us on: [![Flyte Slack](https://www.union.ai/docs/v1/flyte/_static/images/community/joining-the-community/slack-chat-pink.svg)](https://slack.flyte.org/) [![GitHub Discussion](https://www.union.ai/docs/v1/flyte/_static/images/community/joining-the-community/github-discussion-badge.svg)](https://github.com/flyteorg/flyte/discussions) [![Twitter](https://www.union.ai/docs/v1/flyte/_static/images/community/joining-the-community/twitter-social-blue.svg)](https://twitter.com/flyteorg) [![LinkedIn](https://www.union.ai/docs/v1/flyte/_static/images/community/joining-the-community/linkedin-social-lightblue.svg)](https://www.linkedin.com/groups/13962256) [Community sync](https://www.union.ai/docs/v1/flyte/community/joining-the-community/#community-sync) ------------------------------------------------------------------------------------------------------ 1. **When**: First Tuesday of every month, 9:00 AM Pacific Time. 2. **Where**: Live streamed on [YouTube](https://www.youtube.com/@flyteorg/streams) and [LinkedIn](https://www.linkedin.com/company/union-ai/events/) . 3. **Watch the recordings**: [here](https://www.youtube.com/live/d81Jd4rfmzw?feature=shared) . 4. **Import the public calendar**: [here](https://lists.lfaidata.foundation/g/flyte-announce/ics/12031983/2145304139/feed.ics) to not miss any event. 5. **Want to present?** Fill out [this form](https://tally.so/r/wgN8LM) . We’re eager to learn from you! You’re welcome to join and learn from other community members sharing their experiences with Flyte or any other technology from the AI ecosystem. [Contributor’s sync](https://www.union.ai/docs/v1/flyte/community/joining-the-community/#contributors-sync) ------------------------------------------------------------------------------------------------------------- 1. **When**: Every 2 weeks on Thursdays. Alternating schedule between 11:00 AM PT and 7:00 AM PT. 2. **Where**: Live on [Zoom](https://zoom-lfx.platform.linuxfoundation.org/meeting/92309721545?password=c93d76a7-801a-47c6-9916-08e38e5a5c1f) . 3. **Purpose**: Address questions from new contributors, discuss active initiatives, and RFCs. 4. **Import the public calendar**: [here](https://lists.lfaidata.foundation/g/flyte-announce/ics/12031983/2145304139/feed.ics) to not miss any event. [Newsletter](https://www.union.ai/docs/v1/flyte/community/joining-the-community/#newsletter) ---------------------------------------------------------------------------------------------- [Join the Flyte mailing list](https://lists.lfaidata.foundation/g/flyte-announce/join) to receive the monthly newsletter. [Slack guidelines](https://www.union.ai/docs/v1/flyte/community/joining-the-community/#slack-guidelines) ---------------------------------------------------------------------------------------------------------- Flyte strives to build and maintain an open, inclusive, productive, and self-governing open source community. In consequence, we expect all community members to respect the following guidelines: ### [Abide by the](https://www.union.ai/docs/v1/flyte/community/joining-the-community/#abide-by-the-lf) [LF’s Code of Conduct](https://lfprojects.org/policies/code-of-conduct/) As a Linux Foundation project, we must enforce the rules that govern professional and positive open source communities. ### [Avoid using DMs and @mentions](https://www.union.ai/docs/v1/flyte/community/joining-the-community/#avoid-using-dms-and-mentions) Whenever possible, post your questions and responses in public channels so other community members can benefit from the conversation and outcomes. Exceptions to this are when you need to share private or sensitive information. In such a case, the outcome should still be shared publicly. Limit the use of `@mentions` of other community members to be considerate of notification noise. ### [Make use of threads](https://www.union.ai/docs/v1/flyte/community/joining-the-community/#make-use-of-threads) Threads help us keep conversations contained and organized, reducing the time it takes to give you the support you need. **Thread best practices:** * Don’t break your question into multiple messages. Put everything in one. * For long questions, write a few sentences in the first message, and put the rest in a thread. * If there’s a code snippet (more than 5 lines of code), put it inside the thread. * Avoid using the “Also send to channel” feature unless it’s really necessary. * If your question contains multiple questions, make sure to break them into multiple messages, so each could be answered in a separate thread. ### [Do not post the same question across multiple channels](https://www.union.ai/docs/v1/flyte/community/joining-the-community/#do-not-post-the-same-question-across-multiple-channels) If you consider that a question needs to be shared on other channels, ask it once and then indicate explicitly that you’re cross-posting. If you’re having a tough time getting the support you need (or aren’t sure where to go!), please DM `@David Espejo` or `@Samhita Alla` for support. ### [Do not solicit members of our Slack](https://www.union.ai/docs/v1/flyte/community/joining-the-community/#do-not-solicit-members-of-our-slack) The Flyte community exists to collaborate with, learn from, and support one another. It is not a space to pitch your products or services directly to our members via public channels, private channels, or direct messages. We are excited to have a growing presence from vendors to help answer questions from community members as they may arise, but we have a strict 3-strike policy against solicitation: * **First occurrence**: We’ll give you a friendly but public reminder that the behavior is inappropriate according to our guidelines. * **Second occurrence**: We’ll send you a DM warning that any additional violations will result in removal from the community. * **Third occurrence**: We’ll delete or ban your account. We reserve the right to ban users without notice if they are clearly spamming our community members. If you want to promote a product or service, go to the `#shameless-promotion` channel and make sure to follow these rules: * Don’t post more than two promotional posts per week. * Non-relevant topics aren’t allowed. Messages that don’t follow these rules will be deleted. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/community/joining-the-community/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Component Architecture | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Component Architecture ====================== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. This document aims to demystify how Flyte’s major components `Flyteidl`, `Flytekit`, `Flytectl`, `FlyteConsole`, `FlyteAdmin`, `FlytePropeller`, and `FlytePlugins` fit together at a high level. [FlyteIDL](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/#flyteidl) ---------------------------------------------------------------------------------------------- In Flyte, entities like “Workflows”, “Tasks”, “Launch Plans”, and “Schedules” are recognized by multiple system components. For components to communicate effectively, they need a shared understanding about the structure of these entities. Flyteidl (Interface Definition Language) is where shared Flyte entities are defined. It also defines the RPC service definition for the [core Flyte API](https://www.union.ai/docs/v1/flyte/api-reference/flyteidl#flyteidlserviceadminproto) . Flyteidl uses the [protobuf](https://developers.google.com/protocol-buffers/) schema to describe entities. Clients are generated for Python, Golang, and JavaScript and imported by Flyte components. [Planes](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/#planes) ------------------------------------------------------------------------------------------ Flyte components are separated into 3 logical planes. The planes are summarized and explained in detail below. The goal is that these planes can be replaced by alternate implementations. | **Plane** | **Description** | | --- | --- | | **User Plane** | The User Plane consists of all user tools that assist in interacting with the core Flyte API. | | | These tools include the FlyteConsole, Flytekit, and Flytectl. | | **Control Plane** | The Control Plane implements the core Flyte API. | | | It serves all client requests coming from the User Plane. | | | It stores information such as current and past running workflows, and provides that information upon request. | | | It also accepts requests to execute workflows, but offloads the work to the Data Plane. | | **Data Plane** | The sole responsibility of the Data Plane is to fulfill workflows. | | | It accepts workflow requests from the Control Plane and guides the workflow to completion, | | | launching tasks on a cluster of machines as necessary based on the workflow graph. | | | It sends status events back to the control plane so the information can be stored and surfaced to end-users. | ![Flyte Logical Architecture](https://www.union.ai/docs/v1/flyte/_static/images/architecture/component-architecture/flyte-logical-architecture.png) ### [User Plane](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/#user-plane) In Flyte, workflows are represented as a Directed Acyclic Graph (DAG) of tasks. While this representation is logical for services, managing workflow DAGs in this format is a tedious exercise for humans. The Flyte User Plane provides tools to create, manage, and visualize workflows in a format that is easily digestible to the users. These tools include: * **Flytekit**: Flytekit is an SDK that helps users design new workflows using the Python programming language. It can parse the Python code, compile it into a valid Workflow DAG, and submit it to Flyte for execution. * **FlyteConsole**: FlyteConsole provides the Web interface for Flyte. Users and administrators can use the console to view workflows, launch plans, schedules, tasks, and individual task executions. The console provides tools to visualize workflows, and surfaces relevant logs for debugging failed tasks. * **Flytectl**: Flytectl provides interactive access to Flyte to launch and access workflows via terminal. ### [Control Plane](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/#control-plane) The Control Plane supports the core REST/gRPC API defined in Flyteidl. User Plane tools like FlyteConsole and Flytekit contact the control plane on behalf of users to store and retrieve information. Currently, the entire control plane is handled by a single service called **FlyteAdmin**. FlyteAdmin is stateless. It processes requests to create entities like tasks, workflows, and schedules by persisting data in a relational database. While FlyteAdmin serves the Workflow Execution API, it does not itself execute workflows. To launch workflow executions, FlyteAdmin sends the workflow DAG to the Data Plane. For added scalability and fault-tolerance, FlyteAdmin can be configured to load-balance workflows across multiple isolated data-plane clusters. ### [Data Plane](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/#data-plane) The Data Plane is the engine that accepts DAGs, and fulfills workflow executions by launching tasks in the order defined by the graph. Requests to the Data Plane generally come via the control plane, and not from end-users. In order to support compute-intensive workflows at massive scale, the Data Plane needs to launch containers on a cluster of machines. The current implementation leverages [Kubernetes](https://kubernetes.io/) for cluster management. Unlike the user-facing Control Plane, the Data Plane does not expose a traditional REST/gRPC API. To launch an execution in the Data Plane, you create a “`flyteworkflow`” resource in Kubernetes. A “`flyteworkflow`” is a Kubernetes [Custom Resource](https://kubernetes.io/docs/concepts/extend-kubernetes/api-extension/custom-resources/) (CRD) created by our team. This custom resource represents the Flyte workflow DAG. The core state machine that processes `flyteworkflow`s is the worker known as **FlytePropeller**. FlytePropeller leverages the Kubernetes [operator pattern](https://kubernetes.io/docs/concepts/extend-kubernetes/operator/) . It polls the Kubernetes API, looking for newly created `flyteworkflow` resources. FlytePropeller understands the workflow DAG, and launches the appropriate Kubernetes pods as needed to complete tasks. It periodically checks for completed tasks, launching downstream tasks until the workflow is complete. #### [Plugins](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/#plugins) Each task in a `flyteworkflow` DAG has a specified **type**. The logic for fulfilling a task is determined by its task type. In the basic case, FlytePropeller launches a single Kubernetes pod to fulfill a task. Complex task types require workloads to be distributed across hundreds of pods. The type-specific task logic is separated into isolated code modules known as **plugins**. Each task type has an associated plugin that is responsible for handling tasks of its type. For each task in a workflow, FlytePropeller activates the appropriate plugin based on the task type in order to fulfill the task. The Flyte team has pre-built plugins for Hive, Spark, AWS Batch, and more. To support new use-cases, developers can create their own plugins and bundle them in their FlytePropeller deployment. [Component Code Architecture](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/#component-code-architecture) ------------------------------------------------------------------------------------------------------------------------------------ * [FlytePropeller](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/flytepropeller_architecture) * [Flyte Native Scheduler](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/native_scheduler_architecture) [Component Code References](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/#component-code-references) -------------------------------------------------------------------------------------------------------------------------------- * [FlyteAdmin](https://pkg.go.dev/mod/github.com/flyteorg/flyte/flyteadmin) * [FlytePropeller](https://pkg.go.dev/mod/github.com/flyteorg/flyte/flytepropeller) * [DataCatalog](https://pkg.go.dev/mod/github.com/flyteorg/flyte/datacatalog) * [FlytePlugins](https://pkg.go.dev/mod/github.com/flyteorg/flyte/flyteplugins) * [Flyte Native Scheduler](https://pkg.go.dev/github.com/flyteorg/flyte/flyteadmin/scheduler) LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/architecture/component-architecture/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Contributing docs and examples | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Contributing docs and examples ============================== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/community/contributing-docs/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. We welcome contributions to the docs and examples for both Flyte and Union. This section will explain how the docs site works, how to author and build it locally, and how to publish your changes. [The combined Flyte and Union docs site](https://www.union.ai/docs/v1/flyte/community/contributing-docs/#the-combined-flyte-and-union-docs-site) -------------------------------------------------------------------------------------------------------------------------------------------------- As the primary maintainer and contributor of the open-source Flyte project, Union AI is responsible for hosting the Flyte documentation. Additionally, Union AI is also the company behind the commercial Union.ai product, which is based on Flyte. Since Flyte and Union.ai share a lot of common functionality, much of the documentation content is common between the two. However, there are some significant differences between not only Flyte and Union.ai but also among the different Union.ai product offering (Serverless, BYOC, and Self-managed). To effectively and efficiently maintain the documentation for all of these variants, we employ a single-source-of-truth approach where: * All content is stored in a single GitHub repository, [`unionai/unionai-docs`](https://github.com/unionai/unionai-docs) * All content is published on a single website, [`www.union.ai/docs`](https://www.union.ai/docs/v2) . * The website has a variant selector at the top of the page that lets you choose which variant you want to view: * Flyte OSS * Union Serverless * Union BYOC * Union Self-managed * There is also version selector. Currently two versions are available: * v1 (the original docs for Flyte/Union 1.x) * v2 (the new docs for Flyte/Union 2.0, which is the one you are currently viewing) [Versions](https://www.union.ai/docs/v1/flyte/community/contributing-docs/#versions) -------------------------------------------------------------------------------------- The two versions of the docs are stored in separate branches of the GitHub repository: * [`v1` branch](https://github.com/unionai/unionai-docs/tree/v1) for the v1 docs. * [`main` branch](https://github.com/unionai/unionai-docs) for the v2 docs. See [Versions](https://www.union.ai/docs/v1/flyte/community/contributing-docs/versions) for more details. [Common build infrastructure](https://www.union.ai/docs/v1/flyte/community/contributing-docs/#common-build-infrastructure) ---------------------------------------------------------------------------------------------------------------------------- The build infrastructure for the docs site (Hugo configuration, layouts, themes, build scripts, and Python tools) is maintained in a separate repository, [`unionai/unionai-docs-infra`](https://github.com/unionai/unionai-docs-infra) , which is imported as a [git submodule](https://git-scm.com/book/en/v2/Git-Tools-Submodules) at `unionai-docs-infra/` in the `unionai-docs` repository. This means both the `main` (v2) and `v1` content branches share the same build infrastructure. Changes to the build system are made once in `unionai-docs-infra` and are picked up by both branches, keeping them in sync without duplicating build logic. [Variants](https://www.union.ai/docs/v1/flyte/community/contributing-docs/#variants) -------------------------------------------------------------------------------------- Within each branch the multiple variants are supported by using conditional rendering: * Each page of content has a `variants` front matter field that specifies which variants the page is applicable to. * Within each page, rendering logic can be used to include or exclude content based on the selected variant. The result is that: * Content that is common to all variants is authored and stored once. There is no need to keep multiple copies of the same content in-sync. * Content specific to a variant is conditionally rendered based on the selected variant. See [Variants](https://www.union.ai/docs/v1/flyte/community/contributing-docs/variants) for more details. [Both Flyte and Union docs are open source](https://www.union.ai/docs/v1/flyte/community/contributing-docs/#both-flyte-and-union-docs-are-open-source) -------------------------------------------------------------------------------------------------------------------------------------------------------- Since the docs are now combined in one repository, and the Flyte docs are open source, the Union docs are also open source. All the docs are available for anyone to contribute to: Flyte contributors, Union customers, and Union employees. If you are a Flyte contributor, you will be contributing docs related to Flyte features and functionality, but in many cases these features and functionality will also be available in Union. Because the docs site is a single source for all the documentation, when you make changes related to Flyte that are also valid for Union you do so in the same place. This is by design and is a key feature of the docs site. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/community/contributing-docs/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/community/contributing-docs/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Flytekit SDK | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) 1.16.23 Flytekit SDK ============ These are the Flytekit SDK API docs. Flytekit is the core Python SDK for the Union and Flyte platforms. [Developing on Flyte](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/#developing-on-flyte) ----------------------------------------------------------------------------------------------------------- For developing on the Flyte platform you need to add the `flytekit` package to your project: $ uv add flytekit This will install the Flytekit SDK and the `pyflyte` command-line tool. When working with the FLytekit SDK you will be using the `pyflyte` CLI and the Flytekit SDK docs (not the Union SDK docs). [Developing on Union](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/#developing-on-union) ----------------------------------------------------------------------------------------------------------- For developing on the Union platform you need to add the `union` package to your project: $ uv add union This will install the Union SDK, which is a superset of the Flytekit SDK. It will also install the `union` command-line tool. When working with the Union SDK you will be using the `union` CLI and both the Flytekit SDK and the Union SDK docs. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/api-reference/flytekit-sdk/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Workflow lifecycle | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Workflow lifecycle ================== Let’s understand how Flyte’s plugin machinery works and how information flows from one component to another in Flyte. Under the hood, Flyte relies on a primitive called “Plugins”. Every task that you run on Flyte is powered by a plugin. Some of these plugins are native and guaranteed by Flyte system. These native plugins, for example, run your Flyte tasks inside a k8s pod. There are three native plugins, namely, `Container`, `K8sPod`, and `Sql`. Moreover, there are plugins that are actual extensions; they create additional infrastructure and communicate with SaaS on your behalf. Examples include :ref:`Spark `, :ref:`AWS Athena `, etc. A plugin requires code to live in multiple locations. 1. Some parts of plugins logic resides in Flytekit’s SDK. This let users define tasks. You can find this logic [here](https://github.com/flyteorg/flytekit/tree/master/plugins) . Think of this as a client for an RPC service or a web service. 2. Another big chunk of plugins logic lives in [Flyteplugins](https://github.com/flyteorg/flyteplugins) . This is a library that gets loaded into [FlytePropeller](https://github.com/flyteorg/flytepropeller) . FlytePropeller (a Kubernetes operator) loads Flyteplugins upon starting. FlytePropeller is aware of the plugins and their dependency on task execution. However, FlytePropeller is unaware of how these plugins are executed. To better Illustrate how things work, lets take for example the Spark plugin and understand what is the sequence of steps that take place for it to work. The Spark plugin lets a user define a task that has access to a Spark Session. In the background Flyte will provide all the needed infrastructure such that by the time the declared task needs to run, all needed Spark infrastructure is ready and running. 1. User codes in python a task that uses Spark (See code below) @fl.task( task_config=Spark( spark_conf={ "spark.driver.memory": "1000M", "spark.executor.instances": "2", "spark.driver.cores": "1", } ) ) def hello_spark(i: int) -> float: ... As mentioned earlier some part of plugin logic lives on the SDK. In this case think of `Spark` data class here as a placeholder for all the Spark settings that we need our plugin to know. We need to pass this data across multiple places. This is the config that Flyte operator (FlytePropeller) will need in order to build the needed spark cluster. The `Spark` class also tells Flytekit’s SDK that this task will run as a `PysparkFunctionTask` because `task_config` points to a `Spark` object instance, this is clearly illustrated [in spark plugin registration step run in the background](https://github.com/flyteorg/flytekit/blob/master/plugins/flytekit-spark/flytekitplugins/spark/task.py#L129) 2. Once the user has finished writing needed Workflows. A packaging step is needed before user can run the workflows. This packaging step transforms workflows and tasks we described in python into a Protobuf representation. This protobuf representation is used by Flyte across its multiple codebases. For further details on the protobuf representation check [`FlyteIDL` repository](https://github.com/flyteorg/flyteidl) . Package step is carried out by the sdk tooling you are using. This serialization step will transform our `hello_spark` task into a Protobuf representation. It will also transform other tasks, workflows and launch plans to a protobuf representation. Our `hello_spark` protobuf representation will look as below. A Task is serialized as a [`TaskTemplate`](https://github.com/flyteorg/flyteidl/blob/master/protos/flyteidl/core/tasks.proto#L102) as defined in `FlyteIDL`. Id: Task, "example.example.hello_spark" Type: "Spark" Metadata: runtime: type: FLYTE_SDK version: 1.0.3 flavor: python interface: inputs: i : type : simple:Integer description: "i" outputs: o0: type: FLOAT description: o0 custom: executorpath: "/opt/venv/bin/python3" mainApplicationFile: /opt/venv/bin/entrypoint.py sparkConf: spark.driver.cores: 1 spark.executor.instances: 2 spark.driver.memory: 1000M Container: image: "hello_world:1" args: [\ "pyflyte-execute"\ "--inputs"\ "{{.input}}"\ "--output-prefix"\ "{{.outputPrefix}}"\ "--raw-output-data-prefix"\ "{{.rawOutputDataPrefix}}"\ "--checkpoint-path"\ "{{.checkpointOutputPrefix}}"\ "--prev-checkpoint"\ "{{.prevCheckpointPrefix}}"\ "--resolver"\ "flytekit.core.python_auto_container.default_task_resolver"\ "--"\ "task-module"\ "example.example"\ "task-name"\ "hello_spark"\ ] This representation is generated within Flytekit. Essentially the SDK is generating the instructions that Flyte’s kubernetes operator needs to know in order to run this task at a later stage. The `Type` field is really important as we will see later this will be used by Flytepropeller (Kubernetes Operator) to know “how” to execute this task. `Interface` contains information about what are the inputs and outputs of our task. Flyte uses this interface to check if tasks are composable. `Custom` is a collection of arbitrary Key/Values, think of it as a JSON dict that any plugin can define as it wishes. In this case the Spark plugin expects all its particular settings in this field i.e: Spark workers, driver memory, etc. \[`Container](https://github.com/flyteorg/flyteidl/blob/master/protos/flyteidl/core/tasks.proto#L152) is part of Flyte’s IDL primitives. Essentially any Flyte task is ran as either three primitives a` Container`a`K8sPod`or`Sql`. Every task contains a` Target`which has to be either of these. In this particular case, our Spark cluster is a`Container`target. A`Container`specifies all the needed parameters you would in a K8s`ContainerSpec\`, i.e, which docker image to run, what is the command that will be ran, args etc. It is important for the reader to note that Flyte expects to run in a container that has an entrypoint called `pyflyte-execute`. This entrypoint is provided when you `pip install flytekit`. This entrypoint and flytekit is what provides a lot of the plumbing logic inside Flyte. For example, it is this entrypoint what automagically deserializes parquet dataframes an injects them to our task’s functions if need be. It should be clear to the reader that a lot of parameters are surrounded by `{}`. These are template variables that are to be rendered at execution time. What is important from this representation is that it contains all the information that Flyte’s operator needs to know to execute this task: It is a `Spark` task, it has a function signature (inputs and outputs), it tells what docker image to run, and finally, it tells what Spark settings are needed for the cluster. For more information on why this task contains these fields check `TaskTemplate` in the [`FlyteIDL` repository](https://github.com/flyteorg/flyteidl/blob/master/protos/flyteidl/core/tasks.proto#L102) . We strongly advise you to take a look at the data structures in this file ,as they provide good insight in the interfaces used all across Flyte’s codebases. 1. Once user has packaged workflows and tasks then a registration step is needed. During registration Flyte adds these protocolbuffer files to its database, essentially making these tasks and workflows runnable for the user. Registration is done via [`Flytectl`](https://github.com/flyteorg/flytectl) . 2. At some point a Flyte user will trigger a Workflow run. The workflow run will start running the defined DAG. Eventually our Spark task will need to run. This is where the second step of a plugin kicks in. Flytepropeller (Kubernetes Operator) will realize that this is a task of type `Spark` and it will handle it differently. * FlytePropeller knows a task is of type Spark, because our `TaskTemplate` defined it so `Type: Spark`. * Flyte has a `PluginRegistry` which has a dictionary from `Task Type` to `Plugin Handlers`. * At run time Flytepropeller will run our task, Flytepropeller will figure out it is a Spark task, and then call the method `BuildResource` in Spark’s plugin implementation. `BuildResource` is a method that each plugin has to implement. * [`Plugin`](https://github.com/flyteorg/flyteplugins/blob/master/go/tasks/pluginmachinery/k8s/plugin.go#L80) is a Golang interface providing an important method `BuildResource`. * Spark has its own Plugin defined [here in the Flyteplugins repo](https://github.com/flyteorg/flyteplugins/blob/master/go/tasks/plugins/k8s/spark/spark.go) . Inside Spark’s [`BuildResource`](https://github.com/flyteorg/flyteplugins/blob/master/go/tasks/plugins/k8s/spark/spark.go#L65) method is where magic happens. At task runtime: * Flytepropeller will call `BuildResource` method. * This method will ask for the `Custom` field, tasks flagged as `type=Spark` will have a dictionary containing all sort of Spark settings. * Using these settings Flytepropeller will use Spark’s K8s Operator to spawn a spark cluster on the go and run a Spark app (Our python task). * The spark app will run a pod with `pyflyte-execute` as entrypoint. All the inputs and outputs rendered to what they need to be i.e: paths to the actual data inputs instead of `{{input}}` * For more information on Spark’s K8s operator see : [`SparkApplicationSpec`](https://github.com/GoogleCloudPlatform/spark-on-k8s-operator/blob/master/docs/api-docs.md#sparkapplicationspec) . 3. A pod with entrypoint to `pyflyte-execute` execute starts running (Spark App). * `pyflyte-execute` provides all the plumbing magic that is needed. In this particular case, it will create a SparkSession and injects it somewhere so that it is ready for when the user defined python’s code starts running. Be aware that this is part of the SDK code (Flytekit). * `pyflyte-execute` points to [`execute_task_cmd`](https://github.com/flyteorg/flytekit/blob/master/flytekit/bin/entrypoint.py#L445) . This entrypoint does a lot of things: * Resolves the function that the user wants to run, i.e., where is the needed package where this function lives? This is what `"flytekit.core.python_auto_container.default_task_resolver"` does. * Downloads needed inputs and does a transformation if need be. i.e: is this a Dataframe? If so, we need to transform it into a Pandas DF from parquet. * Calls [`dispatch_execute`](https://github.com/flyteorg/flytekit/blob/771aa8a72fbc3ded437b6ff8498404767fc438db/flytekit/core/base_task.py#L449) . This triggers the execution of our spark task. * [`PysparkFunctionTask`](https://github.com/flyteorg/flytekit/blob/master/plugins/flytekit-spark/flytekitplugins/spark/task.py#L78) defines what gets run just before the user’s task code gets executed. It essentially creates a spark session and then runs the user function (The actual code we want to run!). [Recap](https://www.union.ai/docs/v1/flyte/architecture/workflow-lifecycle/#recap) ------------------------------------------------------------------------------------ * Flyte requires coordination between multiple pieces of code. In this case the SDK and FlytePropeller (K8s operator). * [Flyte IDL (Interface Language Definition)](https://github.com/flyteorg/flyteidl) provides some primitives for services to talk with each other. Flyte uses Procolbuffer representations of these primitives. * Three important primitives are : `Container`, `K8sPod`, `Sql`. At the end of the day all tasks boil down to one of those three. * [`github.com/flyteorg/FlytePlugins`](https://github.com/flyteorg/FlytePlugins) repository contains all code for plugins: Spark, AWS Athena, BigQuery, etc. * Flyte entrypoints are the ones carrying out the heavy lifting: making sure that inputs are downloaded and/or transformed as needed. * When running workflows on Flyte, if we want to use Flyte underlying plumbing then we should include Flyte entrypoints: either Jflyte or Flytekit. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/workflow-lifecycle/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Extending Flyte | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Extending Flyte =============== An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. The core of Flyte is a container execution engine, where you can write one or more tasks and compose them together to form a data dependency DAG, called a `workflow`. If your work involves writing simple tasks that can either perform operations on their own or call out to external services, then there is _no need to extend Flyte_. [Define a Custom Type](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/#define-a-custom-type) --------------------------------------------------------------------------------------------------------------- Flyte, just like a programming language, has a core type-system, which can be extended by adding user-defined data types. For example, Flyte supports adding support for a dataframe type from a new library, a custom user data structure, or a grouping of images in a specific encoding. Flytekit natively supports structured data like [`dataclasses.dataclass`](https://docs.python.org/3/library/dataclasses.html) using JSON as the representation format. See [Dataclass](https://www.union.ai/docs/v1/flyte/user-guide/data-input-output/dataclass) . Flytekit allows users to extend Flyte’s type system and implement types in Python that are not representable as JSON documents. The user has to implement a `flytekit.extend.TypeTransformer` class to enable the translation of type from user type to Flyte-understood type. As an example, instead of using [`pandas.DataFrame`](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html) directly, you may want to use [Pandera](https://pandera.readthedocs.io/en/stable) to perform validation of an input or output dataframe. To extend the type system, refer to [Custom types](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/custom-types) . [Add a New Task Plugin](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/#add-a-new-task-plugin) ----------------------------------------------------------------------------------------------------------------- Often you want to interact with services like: * Databases (e.g., Postgres, MySQL, etc.) * DataWarehouses (e.g., Snowflake, BigQuery, Redshift etc.) * Computation (e.g., AWS EMR, Databricks etc.) You might want this interaction to be available as a template for the open-source community or in your organization. This can be done by creating a task plugin, which makes it possible to reuse the task’s underlying functionality within Flyte workflows. If you want users to write code simply using the `@fl.task` decorator, but want to provide the capability of running the function as a spark job or a sagemaker training job, then you can extend Flyte’s task system. @fl.task(task_config=MyContainerExecutionTask( plugin_specific_config_a=..., plugin_specific_config_b=..., ... )) def foo(...) -> ...: ... Alternatively, you can provide an interface like this: query_task = SnowflakeTask( query="Select * from x where x.time < {{.inputs.time}}", inputs=kwtypes(time=datetime), output_schema_type=pandas.DataFrame, ) @workflow def my_wf(t: datetime) -> ...: df = query_task(time=t) return process(df=df) There are two options when writing a new task plugin: You can write a task plugin as an extension in Flytekit or you can go deeper and write a plugin in the Flyte backend. [Flytekit-only task plugin](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/#flytekit-only-task-plugin) ------------------------------------------------------------------------------------------------------------------------- Flytekit is designed to be extremely extensible. You can add new task-types that are useful only for your use-case. Flyte does come with the capability of extending the backend, but that is only required if you want the capability to be extended to all users of Flyte, or there is a cost/visibility benefit of doing so. Writing your own Flytekit plugin is simple and is typically where you want to start when enabling custom task functionality. | Pros | Cons | | --- | --- | | Simple to write — implement in Python. Flyte will treat it like a container execution and blindly pass control to the plugin. | Limited ways of providing additional visibility in progress, external links, etc. | | Simple to publish: `flytekitplugins` can be published as independent libraries and they follow a simple API. | Has to be implemented in every language as these are SDK-side plugins only. | | Simple to perform testing: test locally in flytekit. | In case of side-effects, it could lead to resource leaks. For example, if the plugin runs a BigQuery job, it is possible that the plugin may crash after running the job and Flyte cannot guarantee that the BigQuery job will be successfully terminated. | | | Potentially expensive: in cases where the plugin runs a remote job, running a new pod for every task execution causes severe strain on Kubernetes and the task itself uses almost no CPUs. Also because of its stateful nature, using spot-instances is not trivial. | | | A bug fix to the runtime needs a new library version of the plugin. | | | Not trivial to implement resource controls, like throttling, resource pooling, etc. | [User Container vs. Pre-built Container Task Plugin](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/#user-container-vs-pre-built-container-task-plugin) -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- A Flytekit-only task plugin can be a [User container](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/user-container-task-plugins) or [Pre-built container](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/prebuilt-container-task-plugins) task plugin. | | User Container | Pre-built Container | | --- | --- | --- | | Serialization | At serialization time, a Docker container image is required. The assumption is that this Docker image has the task code. | The Docker container image is hardcoded at serialization time into the task definition by the author of that task plugin. | | Serialization | The serialized task contains instructions to the container on how to reconstitute the task. | Serialized task should contain all the information needed to run that task instance (but not necessarily to reconstitute it). | | Run-time | When Flyte runs the task, the container is launched, and the user-given instructions recreate a Python object representing the task. | When Flyte runs the task, the container is launched. The container should have an executor built into it that knows how to execute the task. | | Run-time | The task object that gets serialized at compile-time is recreated using the user’s code at run time. | The task object that gets serialized at compile-time does not exist at run time. | | Run-time | At platform-run-time, the user-decorated function is executed. | At platform-run-time, there is no user function, and the executor is responsible for producing outputs, given the inputs to the task. | [Backend Plugin](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/#backend-plugin) --------------------------------------------------------------------------------------------------- [Writing a Backend plugin](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/backend-plugins) makes it possible for users to write extensions for FlytePropeller - Flyte’s scheduling engine. This enables complete control of the visualization and availability of the plugin. | Pros | Cons | | --- | --- | | Service oriented way of deploying new plugins - strong contracts. Maintainers can deploy new versions of the backend plugin, fix bugs, without needing the users to upgrade libraries, etc. | Need to be implemented in Golang. | | Drastically cheaper and more efficient to execute. FlytePropeller is written in Golang and uses an event loop model. Each process of FlytePropeller can execute thousands of tasks concurrently. | Needs a FlytePropeller build (_currently_). | | Flyte guarantees resource cleanup. | Need to implement contract in a spec language like protobuf, OpenAPI, etc. | | Flyteconsole plugins (capability coming soon!) can be added to customize visualization and progress tracking of the execution. | Development cycle can be much slower than flytekit-only plugins. | | Resource controls and backpressure management is available. | | | Implement once, use in any SDK or language! | | [Flyte Connector Service](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/#flyte-connector-service) --------------------------------------------------------------------------------------------------------------------- The Flyte Connector service allows you to write backend plugins in Python. ### [Summary](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/#summary) flowchart LR U{Use Case} F(\[Python Flytekit Plugin\]) B(\[Golang
Backend Plugin\]) subgraph WFTP\[Writing Flytekit Task Plugins\] UCP(\[User Container Plugin\]) PCP(\[Pre-built Container Plugin\]) end subgraph WBE\[Writing Backend Extensions\] K8S(\[K8s Plugin\]) WP(\[WebAPI Plugin\]) CP(\[Complex Plugin\]) end subgraph WCFT\[Writing Custom Flyte Types\] T(\[Flytekit
Type Transformer\]) end U -- Light-weight
Extensions --> F U -- Performant
Multi-language
Extensions --> B U -- Specialized
Domain-specific Types --> T F -- Require
user-defined
container --> UCP F -- Provide
prebuilt
container --> PCP B --> K8S B --> WP B --> CP style WCFT fill:#eee,stroke:#aaa style WFTP fill:#eee,stroke:#aaa style WBE fill:#eee,stroke:#aaa style U fill:#fff2b2,stroke:#333 style B fill:#EAD1DC,stroke:#333 style K8S fill:#EAD1DC,stroke:#333 style WP fill:#EAD1DC,stroke:#333 style CP fill:#EAD1DC,stroke:#333 LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/architecture/extending-flyte/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Roadmap | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Roadmap ======= [How the Community Works](https://www.union.ai/docs/v1/flyte/community/roadmap/#how-the-community-works) ---------------------------------------------------------------------------------------------------------- Flyte is actively used in production at multiple companies. We pride ourselves on being extremely customer-focused, and care about providing a high-quality customer experience. We therefore always prioritize stability, reliability, observability, and maintainability over raw feature development. Features are usually developed in response to specific use cases and user scenarios. That being said, we are proactively thinking about the evolution of the system and how we want to keep adapting to changing requirements. Thus most of our changes reflect future development scenarios, and in cases where we feel rapid prototyping would enable us to discover potential pitfalls or uncover hidden use cases, we would proactively develop features behind feature flags. It is extremely important to let the community know about your use cases, so that we adapt parts of Flyte to meet those requirements. We welcome collaboration and contributions, but please follow our [Contribution Guidelines](https://docs.flyte.org/en/latest/community/contribute.html) . The quarterly planning meeting is also hosted publicly, please see more below. [Milestones and Release Processes](https://www.union.ai/docs/v1/flyte/community/roadmap/#milestones-and-release-processes) ---------------------------------------------------------------------------------------------------------------------------- Flyte consists of many components and services. Each service is independently iterated and coordinated by maintaining backwards compatible contracts using Protobuf messages defined in [FlyteIDL](https://docs.flyte.org/en/latest/reference_flyteidl.html) . ### [Release Cadence](https://www.union.ai/docs/v1/flyte/community/roadmap/#release-cadence) We aim to release Flyte quarterly, with the understanding that rather than being tied strictly to the calendar, we aim to have substantial features, improvements, and bug fixes at each quarter. If features slated for a given release are delayed, then the release will be delayed as well. The increased time will also give the Flyte development team more time to beta test each feature and release. ### [Versioning Scheme](https://www.union.ai/docs/v1/flyte/community/roadmap/#versioning-scheme) _Please keep in mind the CI work to implement this scheme is still in progress._ At each quarterly release, major components of Flyte and the Flyte repository itself will be released with an incremented minor version number and the version number will be aligned across those components. The major version number will remain `1` for the foreseeable future. That is, if the current version of Flyte is `1.2.x`, the next release will be `1.3.0` for Flyte and the major components. After each version is released, merges to master will be assigned beta releases of the next release version. That is, if `flytepropeller` version `v1.2.0` was just released, the next merge to master will be tagged `v1.3.0b0`. Not strictly forcing a time-constraint on the Flyte release cycle means that if a substantial number of changes is merged, perhaps due to a security issue or just a rapid pace of feature development, we can always bring up the timeline of the release. #### [Components with Versions Aligned](https://www.union.ai/docs/v1/flyte/community/roadmap/#components-with-versions-aligned) * Propeller * Admin * Console * datacatalog * flytectl * flytesnacks * Flytekit * flytekit-java The last two we are going to tie together for now, but realize that we may want to unpin in the future. #### [Components Versioned Independently](https://www.union.ai/docs/v1/flyte/community/roadmap/#components-versioned-independently) * flyteidl * flytestdlib * flyteplugins * flytecopilot #### [Helm Charts](https://www.union.ai/docs/v1/flyte/community/roadmap/#helm-charts) Helm charts deserve a special mention here. Unlike the other components which will have patch versions that differ, the Flyte release version and the Helm chart version will always be identical down to the patch. That is, a Flyte release is a Helm release and vice-versa. ### [Release Branches and Patching](https://www.union.ai/docs/v1/flyte/community/roadmap/#release-branches-and-patching) After each minor release, a release branch will be created. There will be no alignment of patch versions across the components. That is, by the end of the `1.3.x` release cycle, `flyteadmin` may be on `1.3.8` and `flytepropeller` may be on `1.3.2`. When developing bug fixes, by default we will continue to develop off of master, which will not be the stable branch. After such bug fixes are merged, it will be the responsibility of the developer to ensure that the patches are also applied to prior releases. At the current time, we propose only supporting one release back (two for security patches). That is, if `flytepropeller` has a bug fix that results in `v1.3.0b0` that patch will be applied to the `v1.2.x` release, but not the `v1.1.x` release. #### [Beta Patch Releases](https://www.union.ai/docs/v1/flyte/community/roadmap/#beta-patch-releases) We also propose that beta patch versions be merged into the release branch when patching prior releases. For example, assuming no patches have yet to be made to the `v1.2.0` release, when porting a bug fix that resulted in `v1.3.0b0` onto the `release-v1.2` branch, the developer can first release `v1.2.1b0` for testing into `release-v1.2` before releasing the `v1.2.1` release. Such beta releases should be made at the discretion of the developer. Whether or not a patch version of any of the Flyte components also creates a Flyte patch release shall also be left to the discretion of the developer. ### [Documentation Versioning](https://www.union.ai/docs/v1/flyte/community/roadmap/#documentation-versioning) We also currently have an issue with our documentation versioning. While our readthedocs page does have versioning enabled and we publish the [docs version](https://github.com/flyteorg/flyte/blob/80c098f10334b1c916d1e4274ab9f204152d9d80/rsts/conf.py#L33) , all the [intersphinx mappings](https://github.com/flyteorg/flyte/blob/80c098f10334b1c916d1e4274ab9f204152d9d80/rsts/conf.py#L219) just point to `latest`. Keep in mind that this mapping not only exists in this `flyte` repo, but also in all the other repos that that mapping points to. That is, to maintain an accurate mapping of different versions of documentation, we’ll need to update the mapping in all the repos. To remediate this, we propose the following: * Documentation should be pinned only to Major.Minor on all the repos that have their versions “aligned”. * This means that as we release patch versions of Admin, Propeller, etc., if we’re on v1.1 for instance, as Admin code/auto-generated documentation changes, the v1.1 listing of readthedocs will automatically pick it up. * Repos that are not aligned will just default to the “latest” documentation version. [Planning Process](https://www.union.ai/docs/v1/flyte/community/roadmap/#planning-process) -------------------------------------------------------------------------------------------- ### [Quarterly Planning](https://www.union.ai/docs/v1/flyte/community/roadmap/#quarterly-planning) Members of the community should feel free to join these! Core members of the Flyte team will come prepared with general initiatives in mind. We will use these meetings to prioritize these ideas, assess community interest and impact, and decide what goes into the GitHub milestone for the next release. Members of the community looking to contribute should also join. Please look for this meeting invite on the calendar - it may not be set up as a recurring meeting simply because it will likely change by a few days each quarter. ### [Change Management](https://www.union.ai/docs/v1/flyte/community/roadmap/#change-management) To ensure that changes are trackable and the history is explainable, we use a slightly cumbersome but helpful process, with the following immediate goals: * Every PR is associated with an issue (automatic searchable documentation) * Large PRs are associated with Proposals * Every major change is associated with documentation * Owner files exist for all repositories ### [Issue Lifecycle](https://www.union.ai/docs/v1/flyte/community/roadmap/#issue-lifecycle) * Incoming issues are tagged automatically as untriaged. * Periodically, members of the Flyte community will meet to triage incoming issues. We aim to do this on a weekly basis. * During this meeting we’ll attempt to assign each issue to a milestone. Some issues however will need to be investigated before we can fully assess. * Once an issue is assigned to a milestone, this means we are committed to delivering it that release. This means the burden for adding something to the milestone is relatively high. Issues that slip should only slip for good reason. [Browse Features and Issues](https://www.union.ai/docs/v1/flyte/community/roadmap/#browse-features-and-issues) ---------------------------------------------------------------------------------------------------------------- ### [Issues by Theme](https://www.union.ai/docs/v1/flyte/community/roadmap/#issues-by-theme) | Theme | Description | Open Issues | Comment | | --- | --- | --- | --- | | Bugs | Currently known and open bugs. | [Bugs](https://github.com/flyteorg/flyte/labels/bug) | We are always working on bugs. Open a new one [here](https://github.com/flyteorg/flyte/issues/new/choose)
. | | Security | Issues related to security enhancements. | [Security issues](https://github.com/flyteorg/flyte/labels/security) | | | Docs | All issues open with our documentation | [Docs issues](https://github.com/flyteorg/flyte/labels/documentation) | Starting Feb 2021, we will be completely overhauling our docs. Feedback appreciated! | | Features | All new features in development | [Features issues](https://github.com/flyteorg/flyte/labels/enhancement) | | | Plugins | New capabilities and plugins that are built into Flyte. | [Plugins issues](https://github.com/flyteorg/flyte/labels/plugins) | This is one of the best places to get started contributing to Flyte. Issues with both | | | These could be hosted services, K8s native execution, etc. | | `plugins` and `flytekit` labels refer to purely client-side plugins and are the fastest to contribute to. | | Scale | These issues deal with performance, reliability, and | [Scale issues](https://github.com/flyteorg/flyte/labels/scale) | We are always working on these issues and we would love to hear feedback about what you | | | scalability of Flyte | | would want to change or what we should prioritize. | | Contribute | If you are looking to contribute and want a great first issue, | [Contribute issues](https://github.com/flyteorg/flyte/labels/good%20first%20issue) | These are the best issues to get started with. | | | check out these issues | | | ### [Issues by Components](https://www.union.ai/docs/v1/flyte/community/roadmap/#issues-by-components) | Theme | Description | Open Issues | | --- | --- | --- | | Flyte Console | Issues concerning our web UI. | [Flyte Console issues](https://github.com/flyteorg/flyte/labels/ui) | | Flytectl | Issues concerning our standalone CLI. | [Flytectl issues](https://github.com/flyteorg/flyte/labels/flytectl) | For an overview of what we’re currently working on, check out our [live roadmap](https://github.com/orgs/flyteorg/projects/3) . LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/community/roadmap/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # FlyteKit Plugins | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) FlyteKit Plugins ================ This is a list of all plugins that are available in FlyteKit. | Plugin | Description | | --- | --- | | [Apache Airflow](https://www.union.ai/docs/v1/flyte/api-reference/plugins/airflow/) | This package holds the Airflow plugins for flytekit | | [Async FSSpec](https://www.union.ai/docs/v1/flyte/api-reference/plugins/async-fsspec/) | This package holds the data persistence plugins for flytekit | | [AWS Athena](https://www.union.ai/docs/v1/flyte/api-reference/plugins/aws-athena/) | This package holds the Athena plugins for flytekit | | [AWS Batch](https://www.union.ai/docs/v1/flyte/api-reference/plugins/aws-batch/) | This package holds the AWS Batch plugins for flytekit | | [AWS SageMaker](https://www.union.ai/docs/v1/flyte/api-reference/plugins/aws-sagemaker/) | Flytekit AWS SageMaker Plugin | | [Bigquery](https://www.union.ai/docs/v1/flyte/api-reference/plugins/bigquery/) | This package holds the Bigquery plugins for flytekit | | [Comet ML](https://www.union.ai/docs/v1/flyte/api-reference/plugins/comet-ml/) | This package enables seamless use of Comet within Flyte | | [Dask](https://www.union.ai/docs/v1/flyte/api-reference/plugins/dask/) | Dask plugin for flytekit | | [FSSpec](https://www.union.ai/docs/v1/flyte/api-reference/plugins/data-fsspec/) | This is a deprecated plugin as of flytekit 1.5 | | [DBT](https://www.union.ai/docs/v1/flyte/api-reference/plugins/dbt/) | DBT Plugin for Flytekit | | [Deck](https://www.union.ai/docs/v1/flyte/api-reference/plugins/deck-standard/) | This Plugin provides more renderers to improve task visibility | | [Dolt](https://www.union.ai/docs/v1/flyte/api-reference/plugins/dolt/) | Dolt plugin for flytekit | | [DuckDB](https://www.union.ai/docs/v1/flyte/api-reference/plugins/duckdb/) | DuckDB Plugin for Flytekit | | [Envd](https://www.union.ai/docs/v1/flyte/api-reference/plugins/envd/) | This package enables users to easily build a Docker image for tasks or workflows. | | [Flyte Interactive](https://www.union.ai/docs/v1/flyte/api-reference/plugins/flyteinteractive/) | This package holds the flyteinteractive plugins for flytekit | | [Great Expectations](https://www.union.ai/docs/v1/flyte/api-reference/plugins/greatexpectations/) | Great Expectations Plugin for Flytekit | | [Hive](https://www.union.ai/docs/v1/flyte/api-reference/plugins/hive/) | This package holds Hive plugins for flytekit | | [Hugging Face](https://www.union.ai/docs/v1/flyte/api-reference/plugins/huggingface/) | Hugging Face plugin for flytekit | | [Google IAP](https://www.union.ai/docs/v1/flyte/api-reference/plugins/identity-aware-proxy/) | External command plugin to generate ID tokens for GCP Identity Aware Proxy | | [Inference](https://www.union.ai/docs/v1/flyte/api-reference/plugins/inference/) | This package enables seamless use of model inference sidecar services within Flyte | | [Kubernetes Pod](https://www.union.ai/docs/v1/flyte/api-reference/plugins/k8s-pod/) | Flytekit plugin to support K8s Pod tasks | | [Kubeflow MPI](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-mpi/) | K8s based MPI plugin for flytekit | | [Kubeflow PyTorch](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-pytorch/) | K8s based Pytorch plugin for Flytekit | | [Kubeflow TensorFlow](https://www.union.ai/docs/v1/flyte/api-reference/plugins/kf-tensorflow/) | K8s based Tensorflow plugin for flytekit | | [Memray Profiling](https://www.union.ai/docs/v1/flyte/api-reference/plugins/memray/) | This package enables memory profiling for tasks with memray | | [MLflow](https://www.union.ai/docs/v1/flyte/api-reference/plugins/mlflow/) | This package enables seamless use of MLFlow within Flyte | | [Memory Machine Cloud](https://www.union.ai/docs/v1/flyte/api-reference/plugins/mmcloud/) | MemVerge Flyte plugin | | [Modin](https://www.union.ai/docs/v1/flyte/api-reference/plugins/modin/) | Modin plugin for flytekit | | [Neptune](https://www.union.ai/docs/v1/flyte/api-reference/plugins/neptune/) | This package enables seamless use of Neptune within Flyte | | [OmegaConf](https://www.union.ai/docs/v1/flyte/api-reference/plugins/omegaconf/) | OmegaConf plugin for Flytekit | | [ONNX PyTorch](https://www.union.ai/docs/v1/flyte/api-reference/plugins/onnx-pytorch/) | ONNX PyTorch Plugin for Flytekit | | [ONNX ScikitLearn](https://www.union.ai/docs/v1/flyte/api-reference/plugins/onnx-scikitlearn/) | ONNX ScikitLearn Plugin for Flytekit | | [ONNX TensorFlow](https://www.union.ai/docs/v1/flyte/api-reference/plugins/onnx-tensorflow/) | ONNX TensorFlow Plugin for Flytekit | | [OpenAI](https://www.union.ai/docs/v1/flyte/api-reference/plugins/openai/) | This package holds the openai plugins for flytekit | | [Optuna (wrapper)](https://www.union.ai/docs/v1/flyte/api-reference/plugins/optuna/) | Optuna plugin for flytekit | | [Pandera](https://www.union.ai/docs/v1/flyte/api-reference/plugins/pandera/) | Pandera plugin for flytekit | | [Papermill](https://www.union.ai/docs/v1/flyte/api-reference/plugins/papermill/) | This is the flytekit papermill plugin | | [Perian Job Platform](https://www.union.ai/docs/v1/flyte/api-reference/plugins/perian/) | Flyte agent for Perian Job Platform (perian.io) | | [Polars](https://www.union.ai/docs/v1/flyte/api-reference/plugins/polars/) | Polars plugin for flytekit | | [Ray](https://www.union.ai/docs/v1/flyte/api-reference/plugins/ray/) | This package holds the Ray plugins for flytekit | | [Slurm](https://www.union.ai/docs/v1/flyte/api-reference/plugins/slurm/) | This package holds the Slurm plugins for flytekit | | [Snowflake](https://www.union.ai/docs/v1/flyte/api-reference/plugins/snowflake/) | This package holds the Snowflake plugins for flytekit | | [Spark](https://www.union.ai/docs/v1/flyte/api-reference/plugins/spark/) | Spark 3 plugin for flytekit | | [SQLAlchemy](https://www.union.ai/docs/v1/flyte/api-reference/plugins/sqlalchemy/) | SQLAlchemy plugin for flytekit | | [Vaex](https://www.union.ai/docs/v1/flyte/api-reference/plugins/vaex/) | Vaex plugin for flytekit | | [Weights & Biases](https://www.union.ai/docs/v1/flyte/api-reference/plugins/wandb/) | This package enables seamless use of Weights & Biases within Flyte | | [whylogs](https://www.union.ai/docs/v1/flyte/api-reference/plugins/whylogs/) | Enable the use of whylogs profiles to be used in flyte tasks to get aggregate statistics about data. | LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/api-reference/plugins/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Troubeshooting | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Troubleshooting guide ===================== The content in this section will help Flyte users isolate the most probable causes for some of the common issues that could arise while getting started with the project. Before getting started, collect the following information from the underlying infrastructure: * Capture the `Status` column from the output of: kubectl describe pod -n Where `` will typically correspond to the node execution string that you can find in the UI. * Pay close attention to the `Events` section in the output. * Also, collect the logs from the Pod: kubectl logs pods -n Where `` will typically correspond to the Flyte `-`, e.g., `flytesnacks-development`. Depending on the contents of the logs or the `Events`, you can try different things: [Debugging Common Execution Errors](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#debugging-common-execution-errors) -------------------------------------------------------------------------------------------------------------------------------------- ### [Error: Cannot connect to the Docker daemon at unix:///var/run/docker.sock. Is the docker daemon running?](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#error-cannot-connect-to-the-docker-daemon-at-unixvarrundockersock-is-the-docker-daemon-running) This error will show if you are not running Docker with the native Docker engine in a Linux machine. Most probably you are running Docker via Docker Desktop. * If you are using Docker Desktop in macOS, run: sudo ln -s ~/Library/Containers/com.docker.docker/Data/docker.raw.sock /var/run/docker.sock * If you are using Docker Desktop in Linux, run: sudo ln -s ~$USER/.docker/desktop/docker.sock /var/run/docker.sock * If you are using another tool to run Docker, you need to make sure that `/var/run/docker.sock` is linked to the correct socket file. For example, if you are using Rancher Desktop on Linux, run: sudo ln -s ~$USER/.rd/docker.sock /var/run/docker.sock ### [message: ‘0/1 nodes are available: 1 Insufficient cpu. preemption: 0/1 nodes are available: 1 No preemption victims found for incoming pod.’](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#message-01-nodes-are-available-1-insufficient-cpu-preemption-01-nodes-are-available-1-no-preemption-victims-found-for-incoming-pod) This issue is more common on macOS devices. Make sure that your Docker daemon has allocated a minimum of 4 CPU cores and 3GB of RAM. ### [terminated with exit code (137). Reason \[OOMKilled\]](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#terminated-with-exit-code-137-reason-oomkilled) * For single binary environment deployed with Helm chart, make sure you are using [the most recent charts](https://github.com/flyteorg/flyte/tree/master/charts) . * For EKS deployments, you can adjust resource limits and requests in the `inline` section of the `eks-production.yaml` file. Example: inline: task_resources: defaults: cpu: 100m memory: 100Mi storage: 100Mi limits: memory: 1Gi * Also, the default container resource limits can be overridden from the task itself: from flytekit import Resources, task @task(limits=Resources(mem="256Mi")) def your_task(...): ... ### [Error: ImagePullBackOff](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#error-imagepullbackoff) * If your environment requires the use of a network proxy, use the `--env` option when starting the sandbox and pass the proxy configuration: flytectl demo start --env HTTP_PROXY= * If you’re building a custom Docker image, make sure to use a tag other than `latest`. Otherwise, the Kubernetes default pull policy will be changed from `IfNotPresent` to `Always`, forcing an image pull with every Pod deployment. [Issues Running Workloads](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#issues-running-workloads) -------------------------------------------------------------------------------------------------------------------- ### [OPENSSL\_internal:WRONG\_VERSION\_NUMBER](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#openssl_internalwrong_version_number) * For `flyte-binary`: make sure that the endpoint name you have set in your `config.yaml` file is included in the DNS names of the SSL certificate installed (be it self-signed or issued by a Certificate Authority). * For `sandbox`: verify the `FLYTECTL_CONFIG` environment variable has the correct value by running: export FLYTECTL_CONFIG=~/.flyte/config-sandbox.yaml ### [ModuleNotFoundError](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#modulenotfounderror) * If you’re using a custom container image and using Docker, make sure your `Dockerfile` is located at the same level of the `flyte` directory and that there is an empty `__init__.py` file in your project’s folder: myflyteapp ├── Dockerfile ├── docker_build_and_tag.sh ├── flyte │ ├── __init__.py │ └── workflows │ ├── __init__.py │ └── example.py └── requirements.txt ### [An error occurred (AccessDenied) when calling the PutObject operation in an EKS deployment](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#an-error-occurred-accessdenied-when-calling-the-putobject-operation-in-an-eks-deployment) * Make sure that the Kubernetes service account Flyte is using has the annotation that refers to the IAM Role connected to it: kubectl describe sa -n Example output: Name: Namespace: flyte Labels: app.kubernetes.io/managed-by=eksctl Annotations: eks.amazonaws.com/role-arn: arn:aws:iam:::role/flyte-system-role ... * Otherwise, obtain your IAM role’s ARN and manually annotate the service account: kubectl annotate serviceaccount -n eks.amazonaws.com/role-arn=arn:aws:iam::xxxx:role/ * Refer to this community-maintained [guide](https://github.com/davidmirror-ops/flyte-the-hard-way/blob/main/docs/03-roles-service-accounts.md) for further information about Flyte deployment on EKS. ### [FlyteScopedUserException: ‘JavaPackage’ object is not callable when running a Spark task](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#flytescopeduserexception-javapackage-object-is-not-callable-when-running-a-spark-task) Please add `spark` to the list of `enabled-plugins` in the config YAML file. For example: tasks: task-plugins: enabled-plugins: - container - sidecar - K8S-ARRAY - spark default-for-task-types: - container: container - container_array: K8S-ARRAY ### [authentication handshake failed: x509: “Kubernetes Ingress Controller Fake Certificate” certificate is not trusted when deploying flyte-core to your own Kubernetes cluster](https://www.union.ai/docs/v1/flyte/community/troubleshooting/#authentication-handshake-failed-x509-kubernetes-ingress-controller-fake-certificate-certificate-is-not-trusted-when-deploying-flyte-core-to-your-own-kubernetes-cluster) This issue is caused by TLS being disabled in your Kubernetes cluster. You can resolve the problem by following these steps: * Enable `tls` in the `values.yaml` ingress configuration of flyte-core in order to expose gRPC service at port 443: ingress: host: example.com separateGrpcIngress: true separateGrpcIngressAnnotations: ingress.kubernetes.io/backend-protocol: "grpc" annotations: ingress.kubernetes.io/app-root: "/console" ingress.kubernetes.io/default-backend-redirect: "/console" kubernetes.io/ingress.class: haproxy tls: enabled: true * Disable `insecure` in your `flytectl` client `config.yaml`: admin: endpoint: dns:///example.com authType: Pkce insecure: false insecureSkipVerify: true LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/community/troubleshooting/page.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Core concepts | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Core concepts ============= An LLM-optimized bundle of this entire section is available at [`section.md`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/section.md) . This single file contains all pages in this section, optimized for AI coding agent context. Flyte is a platform for building and orchestrating the execution of interconnected software processes across machines in a computer cluster. In Flyte terminology, the software processes are called _tasks_ and the overall organization of connections between tasks is called a _workflow_. The tasks in a workflow are connected to each other by their inputs and outputs. The output of one task becomes the input of another. More precisely, a workflow in Flyte is a _directed acyclic graph (DAG)_ of _nodes_ where each node is a unit of execution and the edges between nodes represent the flow of data between them. The most common type of node is a task node (which encapsulates a task), though there are also workflow nodes (which encapsulate subworkflows) and branch nodes. In most contexts we just say that a workflow is a DAG of tasks. You define tasks and workflows in Python using the Flytekit SDK. The Flytekit SDK provides a set of decorators and classes that allow you to define tasks and workflows in a way that is easy to understand and work with. Once defined, tasks and workflows are deployed to your Flyte instance (we say they are _registered_ to the instance), where they are compiled into a form that can be executed on your Flyte cluster. In addition to tasks and workflows, another important concept in Flyte is the [_launch plan_](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans) . A launch plan is like a template that can be used to define the inputs to a workflow. Triggering a launch plan will launch its associated workflow with the specified parameters. [Defining tasks and workflows](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#defining-tasks-and-workflows) --------------------------------------------------------------------------------------------------------------------------- Using the Flytekit SDK, tasks and workflows are defined as Python functions using the `@fl.task` and `@fl.workflow` decorators, respectively: import flytekit as fl @fl.task def task_1(a: int, b: int, c: int) -> int: return a + b + c @fl.task def task_2(m: int, n: int) -> int: return m * n @fl.task def task_3(x: int, y: int) -> int: return x - y @fl.workflow def my_workflow(a: int, b: int, c: int, m: int, n: int) -> int: x = task_1(a=a, b=b, c=c) y = task_2(m=m, n=n) return task_3(x=x, y=y) Here we see three tasks defined using the `@fl.task` decorator and a workflow defined using the `@fl.workflow` decorator. The workflow calls `task_1` and `task_2` and passes the results to `task_3` before finally outputting the result of `task_3`. When the workflow is registered, Flyte compiles the workflow into a directed acyclic graph (DAG) based on the input/output dependencies between the tasks. The DAG is then used to execute the tasks in the correct order, taking advantage of any parallelism that is possible. For example, the workflow above results in the following DAG: ![Workflow DAG](https://www.union.ai/docs/v1/flyte/_static/images/user-guide/core-concepts/workflow-dag.png) ### [Type annotation is required](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#type-annotation-is-required) One important difference between Flyte and generic Python is that in Flyte all inputs and outputs _must be type annotated_. This is because tasks are strongly typed, meaning that the types of the inputs and outputs are validated at deployment time. See [Tasks are strongly typed](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks#tasks-are-strongly-typed) for more details. ### [Workflows _are not_ full Python functions](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#workflows-are-not-full-python-functions) The definition of a workflow must be a valid Python function, so it can be run locally as a normal Python function during development, but only _a subset of Python syntax is allowed_, because it must also be compiled into a DAG that is deployed and executed on Flyte. _Technically then, the language of a workflow function is a domain-specific language (DSL) that is a subset of Python._ See [Workflows](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows) for more details. [Registering tasks and workflows](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#registering-tasks-and-workflows) --------------------------------------------------------------------------------------------------------------------------------- ### [Registering on the command line with `pyflyte` or `flytectl`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#registering-on-the-command-line-with-hahahugoshortcode42s23hbhb-or-hahahugoshortcode42s24hbhb) In most cases, workflows and tasks (and possibly other things, such as launch plans) are defined in your project code and registered as a bundle using `pyflyte` or `flytectl` For example: $ pyflyte register ./workflows --project my_project --domain development Tasks can also be registered individually, but it is more common to register alongside the workflow that uses them. See [Running your code](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/running-your-code) . ### [Registering in Python with `FlyteRemote`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#registering-in-python-with-hahahugoshortcode42s28hbhb) As with all Flyte command line actions, you can also perform registration of workflows and tasks programmatically with [`FlyteRemote`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) , specifically, [`FlyteRemote.register_script`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) , [`FlyteRemote.register_workflow`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) , and [`FlyteRemote.register_task`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) . [Results of registration](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#results-of-registration) ----------------------------------------------------------------------------------------------------------------- When the code above is registered to Flyte, it results in the creation of five objects: * The tasks `workflows.my_example.task_1`, `workflows.my_example.task_2`, and `workflows.my_example.task_3` (see [Task fundamentals](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks) for more details). * The workflow `workflows.my_example.my_workflow`. * The default launch plan `workflows.my_example.my_workflow` (see [Launch plans](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/launch-plans) for more details). Notice that the task and workflow names are derived from the path, file name and function name of the Python code that defines them: `..`. The default launch plan for a workflow always has the same name as its workflow. [Changing tasks and workflows](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#changing-tasks-and-workflows) --------------------------------------------------------------------------------------------------------------------------- Tasks and workflows are changed by altering their definition in code and re-registering. When a task or workflow with the same project, domain, and name as a preexisting one is re-registered, a new version of that entity is created. [Inspecting tasks and workflows](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#inspecting-tasks-and-workflows) ------------------------------------------------------------------------------------------------------------------------------- ### [Inspecting workflows in the UI](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#inspecting-workflows-in-the-ui) Select **Workflows** in the sidebar to display a list of all the registered workflows in the project and domain. You can search the workflows by name. Click on a workflow in the list to see the **workflow view**. The sections in this view are as follows: * **Recent Workflow Versions**: A list of recent versions of this workflow. Select a version to see the **Workflow version view**. This view shows the DAG and a list of all version of the task. You can switch between versions with the radio buttons. * **All Executions in the Workflow**: A list of all executions of this workflow. Click on an execution to go to the [Execution view](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/workflows/viewing-workflow-executions) . * **Launch Workflow button**: In the top right of the workflow view, you can click the **Launch Workflow** button to run the workflow with the default inputs. ### [Inspecting tasks in the UI](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#inspecting-tasks-in-the-ui) Select **Tasks** in the sidebar to display a list of all the registered tasks in the project and domain. You can search the launch plans by name. To filter for only those that are archived, check the **Show Only Archived Tasks** box. Click on a task in the list to see the task view The sections in the task view are as follows: * **Inputs & Outputs**: The name and type of each input and output for the latest version of this task. * **Recent Task Versions**: A list of recent versions of this task. Select a version to see the **Task version view**: This view shows the task details and a list of all version of the task. You can switch between versions with the radio buttons. See [Tasks](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/tasks) for more information. * **All Executions in the Task**: A list of all executions of this task. Click on an execution to go to the execution view. * **Launch Task button**: In the top right of the task view, you can click the **Launch Task** button to run the task with the default inputs. ### [Inspecting workflows on the command line with `flytectl`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#inspecting-workflows-on-the-command-line-with-hahahugoshortcode42s35hbhb) To view all tasks within a project and domain: $ flytectl get workflows \ --project \ --domain To view a specific workflow: $ flytectl get workflow \ --project \ --domain \ See [Flytectl CLI](https://www.union.ai/docs/v1/flyte/api-reference/uctl-cli) for more details. ### [Inspecting tasks on the command line with `flytectl`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#inspecting-tasks-on-the-command-line-with-hahahugoshortcode42s39hbhb) To view all tasks within a project and domain: $ flytectl get tasks \ --project \ --domain To view a specific task: $ flytectl get task \ --project \ --domain \ See [Flytectl CLI](https://www.union.ai/docs/v1/flyte/api-reference/uctl-cli) for more details. ### [Inspecting tasks and workflows in Python with `FlyteRemote`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#inspecting-tasks-and-workflows-in-python-with-hahahugoshortcode42s43hbhb) Use the method [`FlyteRemote.fetch_workflow`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) or [`FlyteRemote.client.get_workflow`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) to get a workflow. See [`FlyteRemote`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) for more options and details. Use the method [`FlyteRemote.fetch_task`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) or [`FlyteRemote.client.get_task`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) to get a task. See [`FlyteRemote`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) for more options and details. [Running tasks and workflows](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#running-tasks-and-workflows) ------------------------------------------------------------------------------------------------------------------------- ### [Running a task or workflow in the UI](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#running-a-task-or-workflow-in-the-ui) To run a workflow in the UI, click the **Launch Workflow** button in the workflow view. You can also run individual tasks in the UI by clicking the **Launch Task** button in the task view. ### [Running a task or workflow locally on the command line with `pyflyte` or `python`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#running-a-task-or-workflow-locally-on-the-command-line-with-hahahugoshortcode42s50hbhb-or-python) You can execute a Flyte workflow or task locally simply by calling it just like any regular Python function. For example, you can add the following to the above code: if __name__ == "__main__": my_workflow(a=1, b=2, c=3, m=4, n=5) If the file is saved as `my_example.py`, you can run it locally using the following command: $ python my_example.py Alternatively, you can run the task locally with the `pyflyte` command line tool: To run it locally, you can use the following `pyflyte run` command: $ pyflyte run my_example.py my_workflow --a 1 --b 2 --c 3 --m 4 --n 5 This has the advantage of allowing you to specify the input values as command line arguments. For more details on running workflows and tasks, see [Development cycle](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle) . ### [Running a task or workflow remotely on the command line with `pyflyte`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#running-a-task-or-workflow-remotely-on-the-command-line-with-hahahugoshortcode42s55hbhb) To run a workflow remotely on your Flyte installation, use the following command (this assumes that you have your [FLYTECTL\_CONFIG set up correctly](https://www.union.ai/docs/v1/flyte/user-guide/development-cycle/setting-up-a-project) ): $ pyflyte run --remote my_example.py my_workflow --a 1 --b 2 --c 3 --m 4 --n 5 ### [Running a task or workflow remotely in Python with `FlyteRemote`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/#running-a-task-or-workflow-remotely-in-python-with-hahahugoshortcode42s59hbhb) To run a workflow or task remotely in Python, use the method [`FlyteRemote.execute`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) . See [`FlyteRemote`](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/) for more options and details. LLM-optimized [This page](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/page.md) [This section in one file](https://www.union.ai/docs/v1/flyte/user-guide/core-concepts/section.md) [Full docs index](https://www.union.ai/docs/v1/flyte/llms.txt) On this page 404 Page not found Showing closest match --- # Pyflyte CLI | Union.ai Docs [Preview Flyte 2 for production, hosted on Union.ai\ \ Register now ↗](https://www.union.ai/try-flyte-2) Pyflyte CLI =========== The `pyflyte` CLI is the main tool developers use to interact with Flyte on the command line. [Installation](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#installation) -------------------------------------------------------------------------------------------- The recommended way to install the union CLI outside a workflow project is to use [`uv`](https://docs.astral.sh/uv/) : $ uv tool install flytekit This will install the `pyflyte` CLI globally on your system [as a `uv` tool](https://docs.astral.sh/uv/concepts/tools/) . [Configure the `pyflyte` CLI](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#configure-the-hahahugoshortcode48s5hbhb-cli) ------------------------------------------------------------------------------------------------------------------------------------------ These command will create the file `~/.flyte/config.yaml` with the configuration information to connect to the Flyte instance. See [Getting started > Local setup](https://www.union.ai/docs/v1/flyte/api-reference/user-guide/getting-started/local-setup) for more details. [Overriding the configuration file location](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#overriding-the-configuration-file-location) -------------------------------------------------------------------------------------------------------------------------------------------------------- By default, the `pyflyte` CLI will look for a configuration file at `~/.flyte/config.yaml`. You can override this behavior to specify a different configuration file by setting the `FLYTECTL_CONFIG` environment variable: export FLYTECTL_CONFIG=~/.my-config-location/my-config.yaml Alternatively, you can always specify the configuration file on the command line when invoking `pyflyte` by using the `--config` flag: $ pyflyte --config ~/.my-config-location/my-config.yaml run my_script.py my_workflow [`pyflyte` CLI configuration search path](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#hahahugoshortcode48s15hbhb-cli-configuration-search-path) ------------------------------------------------------------------------------------------------------------------------------------------------------------------- The `pyflyte` CLI will check for configuration files as follows: First, if a `--config` option is used, it will use the specified config file. Second, the config file pointed to by the `FLYTECTL_CONFIG` environment variable. Third, the following hard-coded locations (in this order): Third, the hard-coded location `~/.flyte/config.yaml`. If none of these are present, the CLI will raise an error. [`pyflyte` CLI commands](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#hahahugoshortcode48s21hbhb-cli-commands) --------------------------------------------------------------------------------------------------------------------------------- Entrypoint for all the user commands. `pyflyte [OPTIONS] COMMAND [ARGS]...` **Options** **`-v, --verbose`** Show verbose messages and exception traces **`-k, --pkgs `** Dot-delineated python packages to operate on. Multiple may be specified (can use commas, or specify the switch multiple times. Please note that this option will override the option specified in the configuration file, or environment variable **`-c, --config `** Path to config file for use within container ### [backfill](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#backfill) The backfill command generates and registers a new workflow based on the input launchplan to run an automated backfill. The workflow can be managed using the Flyte UI and can be canceled, relaunched, and recovered. > launchplan refers to the name of the Launchplanlaunchplan\_version is optional and should be a valid version for a Launchplan version. pyflyte backfill [OPTIONS] LAUNCHPLAN [LAUNCHPLAN_VERSION] **Options** **`-p, --project `** Project for workflow/launchplan. Can also be set through envvar `FLYTE_DEFAULT_PROJECT` **Default:** `'flytesnacks'` **`-d, --domain `** Domain for workflow/launchplan, can also be set through envvar `FLYTE_DEFAULT_DOMAIN` **Default:** `'development'` **`-v, --version `** Version for the registered workflow. If not specified it is auto-derived using the start and end date **`-n, --execution-name `** Create a named execution for the backfill. This can prevent launching multiple executions. **`--dry-run`** Just generate the workflow - do not register or execute **Default:** `False` **`--parallel, --serial`** All backfill steps can be run in parallel (limited by max-parallelism), if using `--parallel.` Else all steps will be run sequentially \[`--serial`\]. **Default:** `False` **`--execute, --do-not-execute`** Generate the workflow and register, do not execute **Default:** `True` **`--from-date `** Date from which the backfill should begin. Start date is inclusive. **`--to-date `** Date to which the backfill should run\_until. End date is inclusive **`--backfill-window `** Timedelta for number of days, minutes hours after the from-date or before the to-date to compute the backfills between. This is needed with from-date / to-date. Optional if both from-date and to-date are provided **`--fail-fast, --no-fail-fast`** If set to true, the backfill will fail immediately (WorkflowFailurePolicy.FAIL\_IMMEDIATELY) if any of the backfill steps fail. If set to false, the backfill will continue to run even if some of the backfill steps fail (WorkflowFailurePolicy.FAIL\_AFTER\_EXECUTABLE\_NODES\_COMPLETE). **Default:** `True` **`--overwrite-cache`** Whether to overwrite the cache if it already exists. **Default:** `False` **Arguments** **LAUNCHPLAN** Required argument **LAUNCHPLAN\_VERSION** Optional argument ### [build](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#build) This command can build an image for a workflow or a task from the command line, for fully self-contained scripts. pyflyte build [OPTIONS] COMMAND [ARGS]... **Options** **`-p, --project `** Project to register and run this workflow in. Can also be set through envvar `FLYTE_DEFAULT_PROJECT` **Default:** `'flytesnacks'` **`-d, --domain `** Domain to register and run this workflow in, can also be set through envvar `FLYTE_DEFAULT_DOMAIN` **Default:** `'development'` **`--destination-dir `** Directory inside the image where the tar file containing the code will be copied to **Default:** `'.'` **`--copy-all`** \[Deprecated, see –copy\] Copy all files in the source root directory to the destination directory. You can specify –copy all instead **Default:** `False` **`--copy `** Specifies how to detect which files to copy into image. ‘all’ will behave as the deprecated copy-all flag, ‘auto’ copies only loaded Python modules **Default:** `'auto'`\*\*Options:\*\*all | auto **`-i, --image `** Multiple values allowed.Image used to register and run. **Default:** `'cr.flyte.org/flyteorg/flytekit:py3.9-latest'` **`--service-account `** Service account used when executing this workflow **`--wait, --wait-execution`** Whether to wait for the execution to finish **Default:** `False` **`-i, --poll-interval `** Poll interval in seconds to check the status of the execution **`--dump-snippet`** Whether to dump a code snippet instructing how to load the workflow execution using flyteremote **Default:** `False` **`--overwrite-cache`** Whether to overwrite the cache if it already exists **Default:** `False` **`--envvars, --env `** Multiple values allowed.Environment variables to set in the container, of the format ENV\_NAME=ENV\_VALUE **`--tags, --tag `** Multiple values allowed.Tags to set for the execution **`--name `** Name to assign to this execution **`--labels, --label `** Multiple values allowed.Labels to be attached to the execution of the format label\_key=label\_value. **`--annotations, --annotation `** Multiple values allowed.Annotations to be attached to the execution of the format key=value. **`--raw-output-data-prefix, --raw-data-prefix `** File Path prefix to store raw output data. Examples are file://, s3://, gs:// etc as supported by fsspec. If not specified, raw data will be stored in default configured location in remote of locally to temp file system.Note, this is not metadata, but only the raw data location used to store Flytefile, Flytedirectory, Structuredataset, dataframes **`--max-parallelism `** Number of nodes of a workflow that can be executed in parallel. If not specified, project/domain defaults are used. If 0 then it is unlimited. **`--disable-notifications`** Should notifications be disabled for this execution. **Default:** `False` **`-r, --remote`** Whether to register and run the workflow on a Flyte deployment **Default:** `False` **`--limit `** Use this to limit number of entities to fetch **Default:** `50` **`--cluster-pool `** Assign newly created execution to a given cluster pool **`--execution-cluster-label, --ecl `** Assign newly created execution to a given execution cluster label **`--fast`** Use fast serialization. The image won’t contain the source code. The value is false by default. **Default:** `False` #### [conf.py](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#confpy) Build an image for \[workflow|task\] from conf.py pyflyte build conf.py [OPTIONS] COMMAND [ARGS]... ### [execution](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#execution) The execution command allows you to interact with Flyte’s execution system, such as recovering/relaunching a failed execution. pyflyte execution [OPTIONS] COMMAND [ARGS]... **Options** **`-p, --project `** Project for workflow/launchplan. Can also be set through envvar `FLYTE_DEFAULT_PROJECT` **Default:** `'flytesnacks'` **`-d, --domain `** Domain for workflow/launchplan, can also be set through envvar `FLYTE_DEFAULT_DOMAIN` **Default:** `'development'` **`--execution-id `** **Required** The execution id #### [recover](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#recover) Recover a failed execution pyflyte execution recover [OPTIONS] #### [relaunch](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#relaunch) Relaunch a failed execution pyflyte execution relaunch [OPTIONS] ### [fetch](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#fetch) Retrieve Inputs/Outputs for a Flyte Execution or any of the inner node executions from the remote server. The URI can be retrieved from the Flyte Console, or by invoking the get\_data API. pyflyte fetch [OPTIONS] FLYTE-DATA-URI (format flyte://...) DOWNLOAD-TO Local path (optional) **Options** **`-r, --recursive`** Fetch recursively, all variables in the URI. This is not needed for directories as they are automatically recursively downloaded. **Arguments** **FLYTE-DATA-URI (format flyte://…)** Required argument **DOWNLOAD-TO Local path (optional)** Optional argument ### [get](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#get) Get a single or multiple remote objects. pyflyte get [OPTIONS] COMMAND [ARGS]... #### [launchplan](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#launchplan) Interact with launchplans. pyflyte get launchplan [OPTIONS] LAUNCHPLAN-NAME LAUNCHPLAN-VERSION **Options** **`--active-only, --scheduled`** Only return active launchplans. **`-p, --project `** Project for workflow/launchplan. Can also be set through envvar `FLYTE_DEFAULT_PROJECT` **Default:** `'flytesnacks'` **`-d, --domain `** Domain for workflow/launchplan, can also be set through envvar `FLYTE_DEFAULT_DOMAIN` **Default:** `'development'` **`-l, --limit `** Limit the number of launchplans returned. **Arguments** **LAUNCHPLAN-NAME** Optional argument **LAUNCHPLAN-VERSION** Optional argument ### [info](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#info) Print out information about the current Flyte Python CLI environment - like the version of Flytekit, the version of Flyte Backend Version, backend endpoint currently configured, etc. pyflyte info [OPTIONS] ### [init](https://www.union.ai/docs/v1/flyte/api-reference/pyflyte-cli/#init) Create flyte-ready projects. pyflyte init [OPTIONS] PROJECT_NAME **Options** **`--template