modalflow 0.2.2__tar.gz

modalflow-0.2.2/PKG-INFO

@@ -0,0 +1,102 @@
+Metadata-Version: 2.3
+Name: modalflow
+Version: 0.2.2
+Summary: Serverless Airflow Executor on Modal
+Requires-Dist: apache-airflow>=3.0.6
+Requires-Dist: modal>=1.3.0
+Requires-Dist: click>=8.0.0
+Requires-Dist: requests>=2.31.0
+Requires-Dist: pytest>=7.0 ; extra == 'dev'
+Requires-Dist: pytest-mock ; extra == 'dev'
+Requires-Dist: pytest-timeout ; extra == 'dev'
+Requires-Dist: black ; extra == 'dev'
+Requires-Dist: ruff ; extra == 'dev'
+Requires-Python: >=3.10
+Provides-Extra: dev
+Description-Content-Type: text/markdown
+
+# Modalflow
+
+A serverless Airflow Executor for Modal.
+
+## Usage
+
+### Prerequisites
+- pip
+- A Modal account, workspace, and environment
+- Modal CLI configured to your workspace
+
+### Steps
+
+1. `pip install modalflow` on your local machine, and add it to your Airflow cluster.
+2. Run `modalflow deploy --env {environment name}` to deploy the resources into your Modal environment.
+3. Add `modalflow.executor.ModalExecutor` to your Airflow config
+
+## Networking Configuration
+
+### Local Development
+
+ModalExecutor automatically detects when you're running Airflow locally and creates a tunnel to your local Airflow instance (localhost:8080). No additional configuration needed.
+
+The executor will:
+- Detect if localhost:8080 is accessible
+- Create a Modal tunnel to expose your local Airflow API
+- Pass the tunnel URL to Modal Functions so they can "phone home"
+
+### Production / VPC Deployments
+
+If your Airflow deployment is in a VPC or behind a firewall, you need to configure the execution API URL so Modal Functions can reach it.
+
+**Option 1: Environment Variable (Recommended)**
+
+Set the execution API URL as an environment variable:
+
+```bash
+export AIRFLOW__CORE__EXECUTION_API_SERVER_URL=https://your-airflow-api.example.com/execution/
+```
+
+**Option 2: Airflow Config**
+
+Set in `airflow.cfg`:
+
+```ini
+[core]
+execution_api_server_url = https://your-airflow-api.example.com/execution/
+```
+
+**VPC Setup Considerations**
+
+If Airflow is deployed in a VPC, ensure:
+
+- **Public Endpoint**: Airflow API must be accessible via a public endpoint. Options include:
+  - Application Load Balancer (ALB) in public subnets
+  - API Gateway in front of Airflow
+  - Reverse tunnel (ngrok, Cloudflare Tunnel, AWS IoT Secure Tunneling)
+  
+- **Security**: Use authentication/authorization to protect the endpoint:
+  - API keys or bearer tokens
+  - Security groups restricting access
+  - VPN or private networking (if Modal supports VPC peering)
+
+- **Network Access**: Modal Functions run in Modal's cloud infrastructure and can reach public internet endpoints. Ensure:
+  - No firewall rules blocking Modal's IP ranges
+  - The endpoint is reachable from the internet (not just internal VPC)
+
+**Example: Using Reverse Tunnel**
+
+If you can't expose Airflow directly, use a reverse tunnel:
+
+1. Set up ngrok or similar: `ngrok http 8080`
+2. Configure the executor with the ngrok URL:
+   ```bash
+   export AIRFLOW__CORE__EXECUTION_API_SERVER_URL=https://abc123.ngrok.io/execution/
+   ```
+
+## Development
+
+We use `uv` for development. To setup:
+
+1. `cd modalflow`
+2. `uv sync`
+3. To run the CLI, use: `uv run -- modalflow [COMMAND]`
+4. To run unit tests, use: `uv run pytest`

modalflow-0.2.2/README.md

@@ -0,0 +1,85 @@
+# Modalflow
+
+A serverless Airflow Executor for Modal.
+
+## Usage
+
+### Prerequisites
+- pip
+- A Modal account, workspace, and environment
+- Modal CLI configured to your workspace
+
+### Steps
+
+1. `pip install modalflow` on your local machine, and add it to your Airflow cluster.
+2. Run `modalflow deploy --env {environment name}` to deploy the resources into your Modal environment.
+3. Add `modalflow.executor.ModalExecutor` to your Airflow config
+
+## Networking Configuration
+
+### Local Development
+
+ModalExecutor automatically detects when you're running Airflow locally and creates a tunnel to your local Airflow instance (localhost:8080). No additional configuration needed.
+
+The executor will:
+- Detect if localhost:8080 is accessible
+- Create a Modal tunnel to expose your local Airflow API
+- Pass the tunnel URL to Modal Functions so they can "phone home"
+
+### Production / VPC Deployments
+
+If your Airflow deployment is in a VPC or behind a firewall, you need to configure the execution API URL so Modal Functions can reach it.
+
+**Option 1: Environment Variable (Recommended)**
+
+Set the execution API URL as an environment variable:
+
+```bash
+export AIRFLOW__CORE__EXECUTION_API_SERVER_URL=https://your-airflow-api.example.com/execution/
+```
+
+**Option 2: Airflow Config**
+
+Set in `airflow.cfg`:
+
+```ini
+[core]
+execution_api_server_url = https://your-airflow-api.example.com/execution/
+```
+
+**VPC Setup Considerations**
+
+If Airflow is deployed in a VPC, ensure:
+
+- **Public Endpoint**: Airflow API must be accessible via a public endpoint. Options include:
+  - Application Load Balancer (ALB) in public subnets
+  - API Gateway in front of Airflow
+  - Reverse tunnel (ngrok, Cloudflare Tunnel, AWS IoT Secure Tunneling)
+  
+- **Security**: Use authentication/authorization to protect the endpoint:
+  - API keys or bearer tokens
+  - Security groups restricting access
+  - VPN or private networking (if Modal supports VPC peering)
+
+- **Network Access**: Modal Functions run in Modal's cloud infrastructure and can reach public internet endpoints. Ensure:
+  - No firewall rules blocking Modal's IP ranges
+  - The endpoint is reachable from the internet (not just internal VPC)
+
+**Example: Using Reverse Tunnel**
+
+If you can't expose Airflow directly, use a reverse tunnel:
+
+1. Set up ngrok or similar: `ngrok http 8080`
+2. Configure the executor with the ngrok URL:
+   ```bash
+   export AIRFLOW__CORE__EXECUTION_API_SERVER_URL=https://abc123.ngrok.io/execution/
+   ```
+
+## Development
+
+We use `uv` for development. To setup:
+
+1. `cd modalflow`
+2. `uv sync`
+3. To run the CLI, use: `uv run -- modalflow [COMMAND]`
+4. To run unit tests, use: `uv run pytest`

modalflow-0.2.2/pyproject.toml

@@ -0,0 +1,32 @@
+[build-system]
+requires = ["uv_build>=0.9.18,<0.10.0"]
+build-backend = "uv_build"
+
+[project]
+name = "modalflow"
+version = "0.2.2"
+description = "Serverless Airflow Executor on Modal"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "apache-airflow>=3.0.6",
+    "modal>=1.3.0",
+    "click>=8.0.0",
+    "requests>=2.31.0",
+]
+
+[project.scripts]
+modalflow = "modalflow.cli:cli"
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0",
+    "pytest-mock",
+    "pytest-timeout",
+    "black",
+    "ruff",
+]
+
+[tool.pytest.ini_options]
+pythonpath = "src"
+testpaths = ["tests/unit"]

modalflow-0.2.2/src/modalflow/cli.py

@@ -0,0 +1,44 @@
+import click
+import os
+import subprocess
+import sys
+
+@click.group()
+def cli():
+    """Modalflow CLI."""
+    pass
+
+@cli.command()
+@click.option("--env", default="main", help="Target environment name (default: main)")
+def deploy(env):
+    """
+    Deploy the Modalflow application to the specified environment.
+    
+    This deploys the Modal App, Volume, and Dict defined in modal_app.py.
+    """
+    click.echo(f"Deploying Modalflow to environment: '{env}'...")
+    
+    # Set the environment variable to ensure modal_app.py picks up the correct name
+    env_vars = os.environ.copy()
+    env_vars["MODALFLOW_ENV"] = env
+    
+    # Run 'modal deploy' targeting the modal_app module
+    # We use sys.executable -m modal to ensure we use the same python environment
+    cmd = [sys.executable, "-m", "modal", "deploy", "modalflow.modal_app"]
+    
+    try:
+        subprocess.run(
+            cmd,
+            env=env_vars,
+            check=True
+        )
+        click.echo(f"Successfully deployed to environment '{env}'!")
+    except subprocess.CalledProcessError as e:
+        click.echo(f"Deployment failed with exit code {e.returncode}.", err=True)
+        sys.exit(e.returncode)
+    except Exception as e:
+        click.echo(f"An error occurred: {e}", err=True)
+        sys.exit(1)
+
+if __name__ == "__main__":
+    cli()

modalflow-0.2.2/src/modalflow/executor/modal_executor.py

@@ -0,0 +1,349 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union
+from urllib.parse import urljoin
+
+import modal
+from airflow.configuration import conf
+from airflow.executors.base_executor import BaseExecutor
+from airflow.executors import workloads as executor_workloads
+from airflow.models.taskinstance import TaskInstanceKey
+
+if TYPE_CHECKING:
+    from airflow.executors.workloads import All as ExecutorWorkload
+
+# Type alias for command - can be a list containing a workload or list of strings
+CommandType = Union[List[executor_workloads.ExecuteTask], List[str]]
+
+# Configuration - should match modal_app.py
+ENV = os.environ.get("MODALFLOW_ENV", "main")
+CONCURRENCY_LIMIT = 100
+
+
+class ModalExecutor(BaseExecutor):
+    """
+    An Airflow Executor that runs tasks as Modal Functions.
+    """
+
+    is_local: bool = True
+
+    def __init__(self):
+        # Use the same concurrency limit as the Modal function
+        super().__init__(parallelism=CONCURRENCY_LIMIT)
+        self.active_tasks: Dict[str, TaskInstanceKey] = {}
+        # These will be initialized in start()
+        self._modal_function = None
+        self._state_dict = None
+        self._tunnel_context = None  # Store the context manager
+        self._tunnel = None  # Store the tunnel object from __enter__()
+        self._execution_api_url = None
+
+    @property
+    def slots_available(self) -> int:
+        """
+        Return the number of slots available to run tasks.
+        This is checked by the scheduler to determine if more tasks can be queued.
+        """
+        return self.parallelism - len(self.running) - len(self.queued_tasks)
+
+    def start(self):
+        """
+        Initialize the executor by looking up the deployed Modal function and state dict.
+        Also sets up networking (tunnel for local or production URL).
+        """
+        self.log.info("Starting ModalExecutor")
+
+        app_name = f"modalflow-{ENV}"
+        dict_name = f"airflow-state-{ENV}"
+
+        # Look up the deployed Modal function
+        try:
+            self._modal_function = modal.Function.from_name(
+                app_name, "execute_modal_task"
+            )
+            self.log.info(f"Connected to Modal function: {app_name}/execute_modal_task")
+        except Exception as e:
+            self.log.error(
+                f"Failed to look up Modal function {app_name}/execute_modal_task: {e}"
+            )
+            raise
+
+        # Look up the state dictionary
+        try:
+            self._state_dict = modal.Dict.from_name(dict_name)
+            self.log.info(f"Connected to Modal Dict: {dict_name}")
+        except Exception as e:
+            self.log.error(f"Failed to connect to Modal Dict {dict_name}: {e}")
+            raise
+
+        # Set up execution API URL.
+        # Priority: env var > Airflow config > modal.forward() tunnel > error
+        self._execution_api_url = self._resolve_execution_api_url()
+        self.log.info(f"Execution API URL: {self._execution_api_url}")
+
+    def execute_async(
+        self,
+        key: TaskInstanceKey,
+        command: CommandType,
+        queue: Optional[str] = None,
+        executor_config: Optional[Any] = None,
+    ) -> None:
+        """
+        Trigger a task execution on Modal.
+
+        Handles two formats:
+        - New-style: command is a list containing an ExecuteTask workload object.
+        - Old-style: command is a list of strings (CLI command).
+        """
+        task_key_str = self._get_key_str(key)
+
+        if len(command) == 1 and isinstance(command[0], executor_workloads.ExecuteTask):
+            # New-style: serialize the workload to JSON
+            workload = command[0]
+            serialized_workload = workload.model_dump_json()
+            payload = {
+                "task_key": task_key_str,
+                "workload_json": serialized_workload,
+                "env": self._get_task_env(key, executor_config),
+            }
+        elif all(isinstance(c, str) for c in command):
+            # Old-style: pass the CLI command for direct execution
+            payload = {
+                "task_key": task_key_str,
+                "command": command,
+                "env": self._get_task_env(key, executor_config),
+            }
+        else:
+            raise RuntimeError(
+                f"ModalExecutor doesn't know how to handle command of type: {type(command)}"
+            )
+
+        self.log.info(f"Spawning Modal task for {task_key_str}")
+
+        if self._execution_api_url is None:
+            raise RuntimeError(
+                "Execution API URL not configured. Ensure start() was called successfully."
+            )
+
+        try:
+            self._modal_function.spawn(payload)
+            self.active_tasks[task_key_str] = key
+        except Exception as e:
+            self.log.error(f"Failed to spawn Modal task: {e}")
+            self.fail(key)
+
+    def _process_workloads(self, workloads: Sequence) -> None:
+        """
+        Process workloads from the base executor.
+
+        Handles both new-style ExecuteTask workloads and old-style
+        (command, priority, queue, executor_config) tuples queued via
+        queue_command.
+        """
+        for workload in workloads:
+            if isinstance(workload, executor_workloads.ExecuteTask):
+                ti = workload.ti
+                key = TaskInstanceKey(
+                    dag_id=ti.dag_id,
+                    task_id=ti.task_id,
+                    run_id=ti.run_id,
+                    try_number=ti.try_number,
+                    map_index=ti.map_index,
+                )
+                queue = ti.queue
+                executor_config = ti.executor_config or {}
+                command = [workload]
+            elif isinstance(workload, tuple):
+                # Old-style: (command, priority, queue, executor_config)
+                command, _priority, queue, executor_config = workload
+                # Parse TaskInstanceKey from the CLI command list
+                # Format: ['airflow', 'tasks', 'run', dag_id, task_id, run_id, ...]
+                key = TaskInstanceKey(
+                    dag_id=command[3],
+                    task_id=command[4],
+                    run_id=command[5],
+                    try_number=1,
+                    map_index=-1,
+                )
+            else:
+                self.log.error(
+                    f"Skipping unrecognized workload type: {type(workload)}"
+                )
+                continue
+
+            if key in self.queued_tasks:
+                del self.queued_tasks[key]
+
+            self.execute_async(
+                key=key,
+                command=command,
+                queue=queue,
+                executor_config=executor_config,
+            )
+            self.running.add(key)
+
+    def sync(self) -> None:
+        """
+        Check the status of running tasks.
+        """
+        if not self.active_tasks:
+            return
+
+        # Poll the state dictionary
+        # TODO: batch this or use a more efficient lookup
+
+        completed_keys = []
+
+        for task_key_str, key in self.active_tasks.items():
+            # Check if this key exists in the remote dict
+            # We use .get() to avoid errors if key is missing
+            try:
+                task_state = self._state_dict.get(task_key_str)
+            except Exception as e:
+                self.log.warning(f"Error reading state for {task_key_str}: {e}")
+                continue
+
+            if not task_state:
+                # Task not yet registered by worker, or lost
+                # TODO: implement a timeout logic here
+                continue
+
+            status = task_state.get("status")
+
+            if status == "SUCCESS":
+                self.success(key)
+                completed_keys.append(task_key_str)
+                self.log.info(f"Task {task_key_str} succeeded")
+
+            elif status == "FAILED":
+                self.fail(key)
+                completed_keys.append(task_key_str)
+                error_msg = task_state.get("error", "Unknown error")
+                self.log.error(f"Task {task_key_str} failed: {error_msg}")
+                if task_state.get("stderr"):
+                    self.log.error(f"Task {task_key_str} stderr: {task_state['stderr']}")
+                if task_state.get("stdout"):
+                    self.log.info(f"Task {task_key_str} stdout: {task_state['stdout']}")
+
+        # Cleanup local state
+        for k in completed_keys:
+            del self.active_tasks[k]
+            # Cleanup remote state
+            try:
+                self._state_dict.pop(k)
+            except Exception as e:
+                self.log.warning(f"Failed to cleanup remote state for {k}: {e}") 
+
+    def end(self) -> None:
+        """
+        Terminate the executor and cleanup resources.
+        """
+        self.log.info("Shutting down ModalExecutor")
+        self.heartbeat_interval = 0
+
+        # Cleanup tunnel if it exists
+        if self._tunnel_context is not None:
+            try:
+                # Exit the context manager
+                self._tunnel_context.__exit__(None, None, None)
+                self.log.info("Closed tunnel")
+            except Exception as e:
+                self.log.warning(f"Error closing tunnel: {e}")
+            finally:
+                self._tunnel_context = None
+                self._tunnel = None
+
+    def terminate(self) -> None:
+        """
+        Force terminate.
+        """
+        self.end()
+
+    def _get_key_str(self, key: TaskInstanceKey) -> str:
+        """
+        Serialize TaskInstanceKey to a string.
+        Format: dag_id:task_id:run_id:try_number
+        """
+        # Note: TaskInstanceKey is a named tuple, but the fields vary slightly by Airflow version
+        # We construct a stable string key
+        return f"{key.dag_id}:{key.task_id}:{key.run_id}:{key.try_number}"
+
+    def _resolve_execution_api_url(self) -> str:
+        """
+        Resolve the execution API URL.
+
+        Priority:
+        1. Environment variable AIRFLOW__CORE__EXECUTION_API_SERVER_URL
+        2. Airflow config: core.execution_api_server_url
+        3. modal.forward() tunnel (only works inside Modal Functions)
+
+        Returns:
+            Execution API URL string
+
+        Raises:
+            RuntimeError: If URL cannot be determined
+        """
+        # 1. Check environment variable (takes precedence)
+        env_url = os.environ.get("AIRFLOW__CORE__EXECUTION_API_SERVER_URL")
+        if env_url:
+            self._validate_api_url(env_url)
+            return env_url
+
+        # 2. Check Airflow config
+        try:
+            config_url = conf.get("core", "execution_api_server_url", fallback=None)
+            if config_url:
+                self._validate_api_url(config_url)
+                return config_url
+        except Exception as e:
+            self.log.warning(f"Error reading execution_api_server_url from config: {e}")
+
+        # 3. Try modal.forward() tunnel (works inside Modal Functions only)
+        self.log.info(
+            "No execution API URL configured, attempting modal.forward() tunnel"
+        )
+        try:
+            self._tunnel_context = modal.forward(8080)
+            self._tunnel = self._tunnel_context.__enter__()
+            tunnel_url = self._tunnel.url
+            return urljoin(tunnel_url, "/execution/")
+        except Exception as e:
+            raise RuntimeError(
+                f"Execution API URL not configured and tunnel creation failed: {e}. "
+                "Set AIRFLOW__CORE__EXECUTION_API_SERVER_URL or run inside a Modal container."
+            ) from e
+
+    def _validate_api_url(self, url: str) -> None:
+        """
+        Validate that the execution API URL is properly formatted.
+
+        Args:
+            url: URL string to validate
+
+        Raises:
+            ValueError: If URL is invalid
+        """
+        if not url:
+            raise ValueError("Execution API URL cannot be empty")
+
+        if not (url.startswith("http://") or url.startswith("https://")):
+            raise ValueError(
+                f"Execution API URL must start with http:// or https://: {url}"
+            )
+
+    def _get_task_env(self, key: TaskInstanceKey, executor_config: Any) -> Dict[str, str]:
+        """
+        Gather environment variables to pass to the worker.
+
+        Includes the execution API URL so Modal Functions can phone home.
+        """
+        if self._execution_api_url is None:
+            raise RuntimeError(
+                "Execution API URL not set. Ensure start() was called successfully."
+            )
+
+        env = {
+            "AIRFLOW__CORE__EXECUTION_API_SERVER_URL": self._execution_api_url,
+        }
+        return env

modalflow-0.2.2/src/modalflow/modal_app.py

@@ -0,0 +1,194 @@
+import os
+import subprocess
+from pathlib import Path
+
+import modal
+
+# Allow overriding the environment name via env var
+ENV = os.environ.get("MODALFLOW_ENV", "main")
+
+# Maximum number of concurrent Modal function calls
+# This should match the executor's parallelism setting
+CONCURRENCY_LIMIT = 100
+
+# Optional: path to DAGs directory to include in the image.
+# Set MODALFLOW_DAGS_DIR to include DAGs so the task worker can load them.
+DAGS_DIR = os.environ.get("MODALFLOW_DAGS_DIR", None)
+
+# Define the base image
+# We use the official Airflow image to ensure compatibility.
+# Upgrade from 3.0.6 to 3.1.5 (3.0.6 lacks queue_workload support).
+airflow_image = (
+    modal.Image.from_registry("apache/airflow:3.0.6-python3.10")
+    .run_commands(
+        'su -s /bin/bash airflow -c "pip install apache-airflow==3.1.5"'
+    )
+    .pip_install(
+        "modal",
+        "rich",
+        "click",
+        "pyyaml",
+        "psycopg2-binary",
+    )
+)
+
+# Include DAG files in the image if a DAGs directory is specified.
+# The task worker (execute_workload) needs DAG files to load task definitions.
+if DAGS_DIR:
+    airflow_image = airflow_image.add_local_dir(
+        DAGS_DIR, remote_path="/opt/airflow/dags", copy=True
+    )
+
+# Create the Modal App
+app = modal.App(f"modalflow-{ENV}", image=airflow_image)
+
+# Define the volume for logs
+# We use a dedicated volume for logs so they persist and can be read back
+log_volume = modal.Volume.from_name(f"airflow-logs-{ENV}", create_if_missing=True)
+
+# Define the dict for coordination (hot cache)
+# Maps task_key -> {status, return_code, last_updated}
+state_dict = modal.Dict.from_name(f"airflow-state-{ENV}", create_if_missing=True)
+
+
+@app.function(
+    volumes={"/opt/airflow/logs": log_volume},
+    timeout=3600,  # Default 1 hour timeout
+    max_containers=CONCURRENCY_LIMIT,
+)
+def execute_modal_task(payload: dict):
+    """
+    Executes an Airflow task either via the SDK workload API or a CLI command.
+
+    Payload structure (new-style):
+    {
+        "task_key": "dag_id:task_id:run_id:try_number",
+        "workload_json": "<serialized ExecuteTask workload JSON>",
+        "env": {"AIRFLOW__CORE__...", ...}
+    }
+
+    Payload structure (old-style):
+    {
+        "task_key": "dag_id:task_id:run_id:try_number",
+        "command": ["airflow", "tasks", "run", dag_id, task_id, run_id, ...],
+        "env": {"AIRFLOW__CORE__...", ...}
+    }
+    """
+    import json
+    import os
+
+    task_key = payload.get("task_key")
+    workload_json = payload.get("workload_json")
+    cli_command = payload.get("command")
+    env_vars = payload.get("env", {})
+
+    print(f"Starting execution for {task_key}")
+
+    if workload_json:
+        # New-style: use the Airflow SDK execute_workload module
+        command = [
+            "python",
+            "-m",
+            "airflow.sdk.execution_time.execute_workload",
+            "--json-string",
+            workload_json,
+        ]
+        print(
+            "Using SDK workload path: python -m airflow.sdk.execution_time.execute_workload --json-string <workload>"
+        )
+    elif cli_command:
+        # Old-style: run the CLI command directly
+        command = cli_command
+        print(f"Using CLI path: {' '.join(command)}")
+    else:
+        raise ValueError(
+            f"Payload must contain 'workload_json' or 'command', got keys: {list(payload.keys())}"
+        )
+
+    # Set environment variables
+    # Merge with existing env, but executor-provided vars take precedence
+    run_env = os.environ.copy()
+    run_env.update(env_vars)
+
+    # Try to extract task info for logging
+    log_file_path = None
+    try:
+        if workload_json:
+            workload_data = json.loads(workload_json)
+            ti = workload_data.get("ti", {})
+            dag_id = ti.get("dag_id", "unknown")
+            task_id = ti.get("task_id", "unknown")
+            run_id = ti.get("run_id", "unknown")
+            try_number = ti.get("try_number", 1)
+        elif cli_command and len(cli_command) >= 6:
+            # Parse from CLI args: airflow tasks run dag_id task_id run_id ...
+            dag_id = cli_command[3]
+            task_id = cli_command[4]
+            run_id = cli_command[5]
+            try_number = 1
+        else:
+            dag_id = task_id = run_id = "unknown"
+            try_number = 1
+
+        # Construct path: /opt/airflow/logs/dag_id/task_id/run_id/try_number.log
+        log_dir = os.path.join(
+            "/opt/airflow/logs",
+            f"dag_id={dag_id}",
+            f"run_id={run_id}",
+            f"task_id={task_id}",
+        )
+        os.makedirs(log_dir, exist_ok=True)
+        log_file_path = os.path.join(log_dir, f"attempt={try_number}.log")
+        print(f"Writing logs to {log_file_path}")
+    except Exception as e:
+        print(f"Warning: Failed to setup log directory structure: {e}")
+
+    # Update state to RUNNING right before execution
+    state_dict[task_key] = {
+        "status": "RUNNING",
+        "return_code": None,
+        "ts": 0,  # Timestamp placeholder
+    }
+
+    try:
+        # Run the command
+        # We use subprocess.run to block until completion
+        result = subprocess.run(
+            command,
+            env=run_env,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+
+        # Log output to Modal's centralized logging
+        print(f"Return code: {result.returncode}")
+        print(f"STDOUT (first 500): {result.stdout[:500]}")
+        print(f"STDERR (first 500): {result.stderr[:500]}")
+
+        # Write output to the log file on the volume
+        if log_file_path:
+            try:
+                with open(log_file_path, "w") as f:
+                    f.write(f"*** STDOUT ***\n{result.stdout}\n")
+                    f.write(f"*** STDERR ***\n{result.stderr}\n")
+            except Exception as e:
+                print(f"Failed to write log file: {e}")
+
+        status = "SUCCESS" if result.returncode == 0 else "FAILED"
+
+        state_dict[task_key] = {
+            "status": status,
+            "return_code": result.returncode,
+            "stdout": result.stdout[-2000:],  # Store last 2KB for quick debug
+            "stderr": result.stderr[-2000:],
+        }
+
+    except Exception as e:
+        print(f"Execution failed: {e}")
+        state_dict[task_key] = {
+            "status": "FAILED",
+            "return_code": -1,
+            "error": str(e),
+        }
+        raise