Notebooks

Notebook code can use Budget AI like any other application. When the notebook environment itself points AI calls at the Budget AI gateway, every notebook on a covered cluster, instance or image gets the same models, budgets, limits and guardrails, and people do not have to change their code.

There are two ways to set this up:

Notebook setup kit

The kit sets up the notebook environment so that the openai Python library sends its calls to Budget AI. A notebook that runs OpenAI() with no arguments uses the gateway and a Budget AI key instead of calling the provider directly.

How it works

The openai library reads OPENAI_BASE_URL and OPENAI_API_KEY from the environment when it is given no arguments. The kit installs a small start-up hook into the environment's Python. Every time Python starts, in notebook kernels, scripts and batch jobs alike, the hook sets:

The gateway does the enforcement: the key's model groups, expiry, rate limits and budget, and Budget AI guardrails when they are on. The ANTS Platform SDK is not needed.

Get the kit

The kit has four files:

Ask your ANTS team for the notebook setup kit.

Settings

Each environment needs two settings:

The hook reads them from the environment first, then from the settings file /etc/ants-notebook/gateway.env. The file holds KEY=VALUE lines. The setup scripts write it, because a kernel only sees the environment its parent process started with, and a file reaches every kernel.

The Databricks and SageMaker scripts run the bootstrap file as root. Set ANTS_BOOTSTRAP_SHA256 to the SHA-256 of ants_notebook_bootstrap.py (sha256sum ants_notebook_bootstrap.py) and the script refuses to run a downloaded file that differs, and stops the start. It is optional, and recommended whenever the file comes from a URL or a bucket that other people can write to.

Set up each platform

Check that it works

From a notebook terminal or a %sh cell, run:

Only the last four characters of the key are shown. Exit code 0 means healthy, 1 a configuration problem (including "not installed for this Python") and 2 that the gateway check failed.

A notebook then calls the gateway with no extra setup:

Use a model group the key is granted as model. See Getting Started with Budget AI.

Behaviour

• Personal OpenAI settings are replaced. If a person has their own OPENAI_API_KEY or OPENAI_BASE_URL set, the hook overwrites it, because leaving it would bypass the gateway. To keep values that are already set, set ANTS_GATEWAY_RESPECT_EXISTING_OPENAI=true.
• Setup fails closed. If a setting is missing or malformed, the init script or lifecycle script fails and the cluster or instance does not start, so a mistake cannot leave notebooks quietly uncovered. To start anyway with a warning, set ANTS_SETUP_FAIL_OPEN=true.
• Python always starts. The hook never stops Python from starting and never prints the key. If only one of the two settings is present, it prints a warning that AI calls are not being routed. On a machine with no settings it does nothing.
• The settings file is readable by all users on the node. Anyone on the cluster or instance can read the key, just as they could read it from the environment. Use one environment per team.
• One shared key per environment. Everyone on a cluster, instance or image deployment shares the key's budget and limits, and usage is attributed to that key. When the key expires, is blocked or runs out of budget, every notebook on the environment stops at once. Before the key expires, rotate it, update the secret and restart the environment.
• The SDK is optional. The kit does not install the ANTS Platform SDK unless you ask: ANTS_INSTALL_SDK=true for the Databricks and SageMaker scripts, or the build argument WITH_SDK=true for the Jupyter image. A failed SDK install prints a warning and does not stop the start, because the routing is already applied. If you install it, set ANTS_PLATFORM_HOST and the project key pair as in Manual setup.
• Turn it off for one process. ANTS_GATEWAY_DISABLED=true stops the hook for that process, for debugging.
• The key goes to any OpenAI-compatible endpoint. The openai library sends OPENAI_API_KEY to whatever base_url it is given, so a notebook that calls another OpenAI-compatible provider sends it the Budget AI key. Give each environment its own key with a short expiry, only the model groups the notebooks need, and allowed source networks set to the environment's outbound addresses, so a copied key is refused anywhere else.

What the gateway serves

The gateway serves chat completions (POST /v1/chat/completions), embeddings (POST /v1/embeddings) and the model list (GET /v1/models).

A call that goes directly to a provider succeeds with no budget, limit or guardrail. No kit setting changes this; use a network rule.

Make it enforcing

The kit makes the gateway the default route. It does not stop a notebook from calling a provider directly, for example with code that passes its own base_url or api_key, with the Anthropic, Bedrock or Azure OpenAI clients, with python -S (which skips the start-up hook), or with ANTS_GATEWAY_DISABLED or ANTS_GATEWAY_BASE_URL set by the user before Python starts.

To close these, add a network rule outside the notebook: allow the environment to reach the gateway, and block direct calls to api.openai.com, api.anthropic.com, the Bedrock runtime endpoints and the Azure OpenAI endpoints. Use a security group, an egress firewall, Databricks workspace network rules, or a SageMaker VPC without direct internet access. A bypass then fails instead of skipping the gateway.

What the kit does not cover

• Only OpenAI-compatible chat completions, embeddings and model listing go through the gateway. See What the gateway serves.
• It routes by default, it does not block. See Make it enforcing. Keep provider keys out of notebook environments so the gateway is the only route.
• Each Python needs its own install. A virtual environment that someone creates inside a notebook is not covered until the hook is installed into it.
• Other languages and hosted notebooks. R, Kaggle, Databricks serverless compute and SageMaker Studio are not covered. Google Colab is not covered by the kit; see Google Colab for a manual cell.

Manual setup

You can also set the variables yourself. Install the ants-platform Python SDK if you want SDK tracing and guardrail checks, and set these environment variables for every notebook kernel:

The Budget AI address is shown in Project Settings > Budget AI > Budget AI address. See Getting Started with Budget AI.

Google Colab

Colab runtimes are created by Google for each user, so there is no init script, lifecycle configuration or image for the kit to install into. Each person can add a cell instead.

Store the Budget AI key as a Colab secret (the key icon in the left bar) named ANTS_GATEWAY_KEY, and allow notebook access to it. Run this cell before any AI call:

Create the OpenAI() client after the cell runs, because the library reads both variables when the client is created. The same gateway routes apply.

• The cell is opt-in. Nothing runs it for anyone, so coverage depends on each person adding it.
• It can be bypassed: a person can skip the cell, pass their own base_url and api_key, or call a provider directly. There is no network rule you can apply to a Colab runtime.
• The person running the notebook can read the key. Use a key with a short expiry and only the model groups needed. Colab's outbound addresses are not fixed, so allowed source networks are not practical there.

What is covered

• Models, budgets and limits. Calls through the gateway use only the model groups the key grants and count against its budget and rate limits.
• Guardrails. With Budget AI guardrails on, gateway traffic is checked against your guardrail policies with no notebook code change. Code can also use the SDK's guarded clients directly. See AI Guardrails.
• Usage reporting. Gateway calls appear in Budget AI observability and as traces.

Limits.

• Code that passes its own base URL, or calls a provider's SDK with a provider key, does not go through the gateway. Keep provider keys out of notebook environments so the gateway is the only route.
• Traces from the SDK are sent in the background. Notebooks that trace with the SDK must call ants.flush() before the kernel stops. See Debugging.
• Google Colab and other hosted notebooks need the variables set by hand in each notebook. See Google Colab.

Related