prefect-aca-sessions 1.0.0__tar.gz

prefect_aca_sessions-1.0.0/PKG-INFO

@@ -0,0 +1,148 @@
+Metadata-Version: 2.3
+Name: prefect-aca-sessions
+Version: 1.0.0
+Summary: Prefect worker that runs flow runs in Azure Container Apps dynamic sessions
+Author: Rafael Guimaraes Siqueira
+Author-email: Rafael Guimaraes Siqueira <rafasiqueira@gmail.com>
+Requires-Dist: azure-identity>=1.26.0
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: prefect>=3.8.8
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+
+# prefect-aca-sessions
+
+A Prefect worker that runs flow runs in
+[Azure Container Apps dynamic sessions](https://learn.microsoft.com/en-us/azure/container-apps/sessions-code-interpreter)
+(code interpreter session pools). Worker type: `azure-container-apps-sessions`.
+
+> Status: unit-tested only. It has not yet been verified against a real session pool.
+
+## How it works
+
+The worker is a long-running process that polls a Prefect work pool. For each flow run it:
+
+1. Authenticates to the pool's management endpoint with a Microsoft Entra token
+   (`DefaultAzureCredential`, audience `https://dynamicsessions.io`).
+2. Starts `prefect flow-run execute` as a detached process inside a session (after
+   `pip install`-ing `pip_packages`). The session identifier defaults to the flow run ID.
+3. Polls the session, forwards the process output to the flow run logs, and reports the exit code.
+4. Deletes the session (unless `delete_session_on_completion` is `false`).
+
+A single `/executions` call is limited to 220 seconds, which is why the flow runs as a
+background process that is polled rather than inside one call.
+
+## Prerequisites
+
+- A self-hosted Prefect server reachable from **both** the worker and the sessions.
+- An Azure Container Apps **code interpreter session pool**
+  (`az containerapp sessionpool create --container-type PythonLTS ...`). Note its
+  *pool management endpoint*:
+  `az containerapp sessionpool show -n <pool> -g <rg> --query properties.poolManagementEndpoint -o tsv`
+- An identity for the worker with the **Azure ContainerApps Session Executor** role on the pool.
+- Sessions with network egress to your Prefect API (the session pool's network status must
+  be `EgressEnabled`) and to PyPI, or a pool with the packages you need already available.
+- Python 3.12+ and [uv](https://docs.astral.sh/uv/).
+
+## 1. Install the worker
+
+The package must be installed in the environment where `prefect` runs, because Prefect
+discovers worker types through the `prefect.collections` entry point.
+
+```bash
+uv pip install git+https://github.com/<you>/prefect-aca-sessions   # or: uv pip install -e .
+prefect worker start --help                                       # sanity check
+```
+
+## 2. Point the CLI at your server
+
+```bash
+export PREFECT_API_URL=https://prefect.example.com/api
+# export PREFECT_API_KEY=...   # only if your server has basic auth / a gateway key
+```
+
+## 3. Create the work pool
+
+```bash
+prefect work-pool create aca-sessions --type azure-container-apps-sessions
+```
+
+Then open the pool in the UI (Work Pools → `aca-sessions` → Edit) and set the default
+**Pool management endpoint**, or leave it empty and set it per deployment.
+
+## 4. Authenticate the worker to Azure
+
+`DefaultAzureCredential` is used, so any of these work:
+
+- **Managed identity** (worker on an Azure VM, ACA, AKS): assign the Session Executor role to it.
+- **Service principal**: `AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_CLIENT_SECRET`.
+- **Local development**: `az login`.
+
+## 5. Start the worker
+
+```bash
+prefect worker start --pool aca-sessions --type azure-container-apps-sessions
+```
+
+Run it under a process manager (systemd, a container, etc.) for production.
+
+## 6. Deploy a flow to the pool
+
+```python
+from prefect import flow
+
+
+@flow(log_prints=True)
+def hello():
+    print("running in an ACA session")
+
+
+if __name__ == "__main__":
+    hello.from_source(
+        source="https://github.com/<you>/<repo>",
+        entrypoint="flows.py:hello",
+    ).deploy(
+        name="hello-aca",
+        work_pool_name="aca-sessions",
+        job_variables={
+            "pool_management_endpoint": "https://<region>.dynamicsessions.io/subscriptions/<sub>/resourceGroups/<rg>/sessionPools/<pool>",
+            "pip_packages": ["prefect", "pandas"],
+        },
+    )
+```
+
+The session has to fetch the flow code itself (git source, or a pull step), because nothing
+is shared with the worker's filesystem. Add private-repo credentials via the deployment's
+pull steps and `env`.
+
+## Job variables
+
+| Variable | Default | Description |
+| --- | --- | --- |
+| `pool_management_endpoint` | required | Session pool management endpoint URL |
+| `api_version` | `2025-10-02-preview` | Sessions data-plane API version |
+| `session_identifier` | flow run ID | Session to run in. A fixed value reuses one session across runs, so runs can see each other's files |
+| `pip_packages` | `["prefect"]` | Installed in the session before the run. Pin the Prefect version to match your server |
+| `poll_interval_seconds` | `10` | Seconds between status polls |
+| `delete_session_on_completion` | `true` | Delete the session after the run |
+| `env` | `{}` | Extra environment variables for the flow run process |
+| `command` | `prefect flow-run execute` | Command run in the session |
+
+## Security notes
+
+- The flow run's environment (including `PREFECT_API_URL` and any `PREFECT_API_KEY`) is sent
+  to the session inside the execution request. Use a dedicated, least-privilege API key.
+- Session identifiers are sensitive; don't share a fixed identifier across tenants.
+
+## Known limitations
+
+- Cancelling a flow run does not stop the process in the session (`kill_infrastructure` is not
+  implemented); delete the session to stop it.
+- Session lifetime and idle limits are set on the pool and apply to long flow runs.
+
+## Development
+
+```bash
+uv sync
+uv run pytest
+```

prefect_aca_sessions-1.0.0/README.md

@@ -0,0 +1,136 @@
+# prefect-aca-sessions
+
+A Prefect worker that runs flow runs in
+[Azure Container Apps dynamic sessions](https://learn.microsoft.com/en-us/azure/container-apps/sessions-code-interpreter)
+(code interpreter session pools). Worker type: `azure-container-apps-sessions`.
+
+> Status: unit-tested only. It has not yet been verified against a real session pool.
+
+## How it works
+
+The worker is a long-running process that polls a Prefect work pool. For each flow run it:
+
+1. Authenticates to the pool's management endpoint with a Microsoft Entra token
+   (`DefaultAzureCredential`, audience `https://dynamicsessions.io`).
+2. Starts `prefect flow-run execute` as a detached process inside a session (after
+   `pip install`-ing `pip_packages`). The session identifier defaults to the flow run ID.
+3. Polls the session, forwards the process output to the flow run logs, and reports the exit code.
+4. Deletes the session (unless `delete_session_on_completion` is `false`).
+
+A single `/executions` call is limited to 220 seconds, which is why the flow runs as a
+background process that is polled rather than inside one call.
+
+## Prerequisites
+
+- A self-hosted Prefect server reachable from **both** the worker and the sessions.
+- An Azure Container Apps **code interpreter session pool**
+  (`az containerapp sessionpool create --container-type PythonLTS ...`). Note its
+  *pool management endpoint*:
+  `az containerapp sessionpool show -n <pool> -g <rg> --query properties.poolManagementEndpoint -o tsv`
+- An identity for the worker with the **Azure ContainerApps Session Executor** role on the pool.
+- Sessions with network egress to your Prefect API (the session pool's network status must
+  be `EgressEnabled`) and to PyPI, or a pool with the packages you need already available.
+- Python 3.12+ and [uv](https://docs.astral.sh/uv/).
+
+## 1. Install the worker
+
+The package must be installed in the environment where `prefect` runs, because Prefect
+discovers worker types through the `prefect.collections` entry point.
+
+```bash
+uv pip install git+https://github.com/<you>/prefect-aca-sessions   # or: uv pip install -e .
+prefect worker start --help                                       # sanity check
+```
+
+## 2. Point the CLI at your server
+
+```bash
+export PREFECT_API_URL=https://prefect.example.com/api
+# export PREFECT_API_KEY=...   # only if your server has basic auth / a gateway key
+```
+
+## 3. Create the work pool
+
+```bash
+prefect work-pool create aca-sessions --type azure-container-apps-sessions
+```
+
+Then open the pool in the UI (Work Pools → `aca-sessions` → Edit) and set the default
+**Pool management endpoint**, or leave it empty and set it per deployment.
+
+## 4. Authenticate the worker to Azure
+
+`DefaultAzureCredential` is used, so any of these work:
+
+- **Managed identity** (worker on an Azure VM, ACA, AKS): assign the Session Executor role to it.
+- **Service principal**: `AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_CLIENT_SECRET`.
+- **Local development**: `az login`.
+
+## 5. Start the worker
+
+```bash
+prefect worker start --pool aca-sessions --type azure-container-apps-sessions
+```
+
+Run it under a process manager (systemd, a container, etc.) for production.
+
+## 6. Deploy a flow to the pool
+
+```python
+from prefect import flow
+
+
+@flow(log_prints=True)
+def hello():
+    print("running in an ACA session")
+
+
+if __name__ == "__main__":
+    hello.from_source(
+        source="https://github.com/<you>/<repo>",
+        entrypoint="flows.py:hello",
+    ).deploy(
+        name="hello-aca",
+        work_pool_name="aca-sessions",
+        job_variables={
+            "pool_management_endpoint": "https://<region>.dynamicsessions.io/subscriptions/<sub>/resourceGroups/<rg>/sessionPools/<pool>",
+            "pip_packages": ["prefect", "pandas"],
+        },
+    )
+```
+
+The session has to fetch the flow code itself (git source, or a pull step), because nothing
+is shared with the worker's filesystem. Add private-repo credentials via the deployment's
+pull steps and `env`.
+
+## Job variables
+
+| Variable | Default | Description |
+| --- | --- | --- |
+| `pool_management_endpoint` | required | Session pool management endpoint URL |
+| `api_version` | `2025-10-02-preview` | Sessions data-plane API version |
+| `session_identifier` | flow run ID | Session to run in. A fixed value reuses one session across runs, so runs can see each other's files |
+| `pip_packages` | `["prefect"]` | Installed in the session before the run. Pin the Prefect version to match your server |
+| `poll_interval_seconds` | `10` | Seconds between status polls |
+| `delete_session_on_completion` | `true` | Delete the session after the run |
+| `env` | `{}` | Extra environment variables for the flow run process |
+| `command` | `prefect flow-run execute` | Command run in the session |
+
+## Security notes
+
+- The flow run's environment (including `PREFECT_API_URL` and any `PREFECT_API_KEY`) is sent
+  to the session inside the execution request. Use a dedicated, least-privilege API key.
+- Session identifiers are sensitive; don't share a fixed identifier across tenants.
+
+## Known limitations
+
+- Cancelling a flow run does not stop the process in the session (`kill_infrastructure` is not
+  implemented); delete the session to stop it.
+- Session lifetime and idle limits are set on the pool and apply to long flow runs.
+
+## Development
+
+```bash
+uv sync
+uv run pytest
+```

prefect_aca_sessions-1.0.0/pyproject.toml

@@ -0,0 +1,31 @@
+[project]
+name = "prefect-aca-sessions"
+version = "1.0.0"
+description = "Prefect worker that runs flow runs in Azure Container Apps dynamic sessions"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "azure-identity>=1.26.0",
+    "httpx>=0.28.1",
+    "prefect>=3.8.8",
+]
+
+[[project.authors]]
+name = "Rafael Guimaraes Siqueira"
+email = "rafasiqueira@gmail.com"
+
+[project.entry-points."prefect.collections"]
+prefect_aca_sessions = "prefect_aca_sessions"
+
+[build-system]
+requires = ["uv_build>=0.11.30,<0.12.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=9.1.1",
+    "pytest-asyncio>=1.4.0",
+]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"

prefect_aca_sessions-1.0.0/pyproject.toml.orig

@@ -0,0 +1,30 @@
+[project]
+name = "prefect-aca-sessions"
+version = "1.0.0"
+description = "Prefect worker that runs flow runs in Azure Container Apps dynamic sessions"
+readme = "README.md"
+authors = [
+    { name = "Rafael Guimaraes Siqueira", email = "rafasiqueira@gmail.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "azure-identity>=1.26.0",
+    "httpx>=0.28.1",
+    "prefect>=3.8.8",
+]
+
+[build-system]
+requires = ["uv_build>=0.11.30,<0.12.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=9.1.1",
+    "pytest-asyncio>=1.4.0",
+]
+
+[project.entry-points."prefect.collections"]
+prefect_aca_sessions = "prefect_aca_sessions"
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"

prefect_aca_sessions-1.0.0/src/prefect_aca_sessions/__init__.py

@@ -0,0 +1,7 @@
+from prefect_aca_sessions.worker import (
+    ACASessionsJobConfiguration,
+    ACASessionsWorker,
+    ACASessionsWorkerResult,
+)
+
+__all__ = ["ACASessionsJobConfiguration", "ACASessionsWorker", "ACASessionsWorkerResult"]

prefect_aca_sessions-1.0.0/src/prefect_aca_sessions/client.py

@@ -0,0 +1,98 @@
+"""Minimal async client for the Azure Container Apps code interpreter management endpoint."""
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+from azure.core.credentials_async import AsyncTokenCredential
+from azure.identity.aio import DefaultAzureCredential
+
+SESSIONS_TOKEN_SCOPE = "https://dynamicsessions.io/.default"
+REQUEST_TIMEOUT_SECONDS = 300.0  # a single execution may run up to 220s
+
+
+class SessionExecutionError(RuntimeError):
+    """The sessions endpoint rejected a request or the executed code failed."""
+
+
+@dataclass(frozen=True)
+class ExecutionResult:
+    status: str
+    stdout: str
+    stderr: str
+
+
+class SessionsClient:
+    """Runs inline Python code in a session of a code interpreter session pool."""
+
+    def __init__(
+        self,
+        pool_management_endpoint: str,
+        api_version: str,
+        credential: AsyncTokenCredential | None = None,
+        http_client: httpx.AsyncClient | None = None,
+    ) -> None:
+        self._endpoint = pool_management_endpoint.rstrip("/")
+        self._api_version = api_version
+        self._credential = credential or DefaultAzureCredential()
+        self._http = http_client or httpx.AsyncClient(timeout=REQUEST_TIMEOUT_SECONDS)
+
+    async def __aenter__(self) -> "SessionsClient":
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        await self._http.aclose()
+        await self._credential.close()
+
+    async def _headers(self) -> dict[str, str]:
+        token = await self._credential.get_token(SESSIONS_TOKEN_SCOPE)
+        return {"Authorization": f"Bearer {token.token}"}
+
+    def _params(self, identifier: str) -> dict[str, str]:
+        return {"api-version": self._api_version, "identifier": identifier}
+
+    async def execute(self, identifier: str, code: str) -> ExecutionResult:
+        response = await self._http.post(
+            f"{self._endpoint}/executions",
+            params=self._params(identifier),
+            headers=await self._headers(),
+            json={
+                "properties": {
+                    "codeInputType": "inline",
+                    "executionType": "synchronous",
+                    "code": code,
+                }
+            },
+        )
+        if response.is_error:
+            raise SessionExecutionError(
+                f"Session pool returned HTTP {response.status_code}: {response.text}"
+            )
+        return _parse_execution(response.json())
+
+    async def delete_session(self, identifier: str) -> None:
+        response = await self._http.delete(
+            f"{self._endpoint}/session",
+            params=self._params(identifier),
+            headers=await self._headers(),
+        )
+        if response.is_error and response.status_code != 404:
+            raise SessionExecutionError(
+                f"Deleting session failed with HTTP {response.status_code}: {response.text}"
+            )
+
+
+def _parse_execution(body: dict[str, Any]) -> ExecutionResult:
+    properties = body.get("properties", {})
+    result = ExecutionResult(
+        status=properties.get("status", "Unknown"),
+        stdout=properties.get("stdout", ""),
+        stderr=properties.get("stderr", ""),
+    )
+    if result.status != "Succeeded":
+        raise SessionExecutionError(
+            f"Code execution ended with status {result.status!r}: "
+            f"{json.dumps(properties.get('result'))} {result.stderr}"
+        )
+    return result

prefect_aca_sessions-1.0.0/src/prefect_aca_sessions/snippets.py

@@ -0,0 +1,55 @@
+"""Python snippets executed inside the session to launch and observe the flow run process."""
+
+import json
+
+RUN_DIR = "/mnt/data/prefect-run"
+
+
+def build_start_code(command: str, env: dict[str, str], pip_packages: list[str]) -> str:
+    """Code that starts `command` as a detached process and returns immediately."""
+    spec = json.dumps({"command": command, "env": env, "pip": pip_packages})
+    return f"""
+import json, os, subprocess
+spec = json.loads({spec!r})
+run_dir = {RUN_DIR!r}
+os.makedirs(run_dir, exist_ok=True)
+for name in ("exit_code", "output.log"):
+    path = os.path.join(run_dir, name)
+    if os.path.exists(path):
+        os.remove(path)
+steps = []
+if spec["pip"]:
+    steps.append("pip install --quiet " + " ".join(spec["pip"]))
+steps.append(spec["command"])
+script = " && ".join("(" + s + ")" for s in steps) + "; echo $? > " + run_dir + "/exit_code"
+log = open(run_dir + "/output.log", "ab")
+proc = subprocess.Popen(
+    ["sh", "-c", script],
+    env={{**os.environ, **spec["env"]}},
+    stdout=log,
+    stderr=subprocess.STDOUT,
+    start_new_session=True,
+)
+print(proc.pid)
+"""
+
+
+def build_poll_code(log_offset: int) -> str:
+    """Code printing JSON: new log text, next log offset and the exit code (or null)."""
+    return f"""
+import json, os
+run_dir = {RUN_DIR!r}
+offset = {log_offset}
+exit_path = run_dir + "/exit_code"
+# read the exit code first so a run finishing mid-poll never loses its last log lines
+exit_code = int(open(exit_path).read().strip()) if os.path.exists(exit_path) else None
+with open(run_dir + "/output.log", "rb") as f:
+    f.seek(offset)
+    chunk = f.read()
+print(json.dumps({{
+    "log": chunk.decode("utf-8", "replace"),
+    "offset": offset + len(chunk),
+    "exit_code": exit_code,
+}}))
+"""
+

prefect_aca_sessions-1.0.0/src/prefect_aca_sessions/worker.py

@@ -0,0 +1,144 @@
+import asyncio
+import json
+
+import anyio
+import anyio.abc
+from prefect.client.schemas.objects import FlowRun
+from prefect.workers.base import (
+    BaseJobConfiguration,
+    BaseVariables,
+    BaseWorker,
+    BaseWorkerResult,
+)
+from pydantic import Field
+
+from prefect_aca_sessions.client import SessionsClient
+from prefect_aca_sessions.snippets import build_poll_code, build_start_code
+
+DEFAULT_API_VERSION = "2025-10-02-preview"
+DEFAULT_POLL_INTERVAL_SECONDS = 10
+FAILURE_EXIT_CODE = -1
+
+
+class ACASessionsJobConfiguration(BaseJobConfiguration):
+    """Job configuration for flow runs executed in an ACA code interpreter session."""
+
+    pool_management_endpoint: str = Field(
+        description=(
+            "Management endpoint of the code interpreter session pool, e.g. "
+            "https://<region>.dynamicsessions.io/subscriptions/<id>/resourceGroups/<rg>"
+            "/sessionPools/<pool>"
+        ),
+        json_schema_extra=dict(template="{{ pool_management_endpoint }}"),
+    )
+    api_version: str = Field(
+        default=DEFAULT_API_VERSION,
+        description="Sessions data-plane API version.",
+        json_schema_extra=dict(template="{{ api_version }}"),
+    )
+    session_identifier: str | None = Field(
+        default=None,
+        description="Session identifier. Defaults to the flow run ID (one session per run).",
+        json_schema_extra=dict(template="{{ session_identifier }}"),
+    )
+    pip_packages: list[str] = Field(
+        default_factory=lambda: ["prefect"],
+        description="Packages pip-installed in the session before the flow run starts.",
+        json_schema_extra=dict(template="{{ pip_packages }}"),
+    )
+    poll_interval_seconds: int = Field(
+        default=DEFAULT_POLL_INTERVAL_SECONDS,
+        ge=1,
+        description="Seconds between checks of the flow run process.",
+        json_schema_extra=dict(template="{{ poll_interval_seconds }}"),
+    )
+    delete_session_on_completion: bool = Field(
+        default=True,
+        description="Delete the session once the flow run process exits.",
+        json_schema_extra=dict(template="{{ delete_session_on_completion }}"),
+    )
+
+
+class ACASessionsVariables(BaseVariables):
+    pool_management_endpoint: str = Field(
+        description="Management endpoint of the code interpreter session pool.",
+    )
+    api_version: str = Field(default=DEFAULT_API_VERSION, description="Sessions API version.")
+    session_identifier: str | None = Field(
+        default=None, description="Session identifier. Defaults to the flow run ID."
+    )
+    pip_packages: list[str] = Field(
+        default_factory=lambda: ["prefect"],
+        description="Packages pip-installed in the session before the run.",
+    )
+    poll_interval_seconds: int = Field(
+        default=DEFAULT_POLL_INTERVAL_SECONDS, ge=1, description="Seconds between status checks."
+    )
+    delete_session_on_completion: bool = Field(
+        default=True, description="Delete the session after the run."
+    )
+
+
+class ACASessionsWorkerResult(BaseWorkerResult):
+    """Result of a flow run executed in an ACA session."""
+
+
+class ACASessionsWorker(BaseWorker):
+    type = "azure-container-apps-sessions"
+    job_configuration = ACASessionsJobConfiguration
+    job_configuration_variables = ACASessionsVariables
+    _description = "Runs flow runs in Azure Container Apps dynamic sessions (code interpreter)."
+    _documentation_url = (
+        "https://learn.microsoft.com/en-us/azure/container-apps/sessions-code-interpreter"
+    )
+
+    async def run(
+        self,
+        flow_run: FlowRun,
+        configuration: ACASessionsJobConfiguration,
+        task_status: anyio.abc.TaskStatus | None = None,
+    ) -> ACASessionsWorkerResult:
+        identifier = configuration.session_identifier or str(flow_run.id)
+        logger = self.get_flow_run_logger(flow_run)
+
+        async with SessionsClient(
+            configuration.pool_management_endpoint, configuration.api_version
+        ) as client:
+            try:
+                exit_code = await self._execute(client, identifier, configuration, logger, task_status)
+            except Exception:
+                logger.exception("Flow run failed in session %s", identifier)
+                exit_code = FAILURE_EXIT_CODE
+            finally:
+                if configuration.delete_session_on_completion:
+                    await self._delete_quietly(client, identifier, logger)
+
+        return ACASessionsWorkerResult(status_code=exit_code, identifier=identifier)
+
+    async def _execute(self, client, identifier, configuration, logger, task_status) -> int:
+        await client.execute(
+            identifier,
+            build_start_code(
+                configuration.command or "", configuration.env, configuration.pip_packages
+            ),
+        )
+        logger.info("Started flow run process in session %s", identifier)
+        if task_status is not None:
+            task_status.started(identifier)
+
+        offset = 0
+        while True:
+            await asyncio.sleep(configuration.poll_interval_seconds)
+            poll = json.loads((await client.execute(identifier, build_poll_code(offset))).stdout)
+            offset = poll["offset"]
+            if poll["log"]:
+                logger.info(poll["log"].rstrip())
+            if poll["exit_code"] is not None:
+                return poll["exit_code"]
+
+    @staticmethod
+    async def _delete_quietly(client, identifier, logger) -> None:
+        try:
+            await client.delete_session(identifier)
+        except Exception:
+            logger.warning("Could not delete session %s", identifier, exc_info=True)