# NVIDIA NIM - Generative AI Plugin

With this plugin, you can leverage and deploy [NVIDIA NIM LLMs](https://docs.nvidia.com/nim/large-language-models/latest/introduction.html).

> **⚠️ DEPRECATED**: The custom LLM connection component has been removed from this plugin. Starting with Dataiku v14.2, the LLM Mesh includes a dedicated NVIDIA NIM connection. Please use the native connection for new implementations.


This plugin provides three main components:

1. **NIM Deployment Macro** - Deploy and manage NIMs on Kubernetes clusters attached to Dataiku
2. **NemoGuard NIMs Guardrails** - Enables Nemo Guardrails (Topic Control, Content Safety, and Jailbreak) in the Dataiku LLM Mesh.
3. **Run:ai Scheduling Macro** - Validate or configure Run:ai project scheduling for Dataiku Kubernetes namespaces.
---

## 1. Macro for Deploying NIMs in Dataiku

The NIM Deployment Macro provides capabilities to deploy and manage NVIDIA NIMs on Kubernetes clusters attached to Dataiku.

### Capabilities

- Deploy, list, and remove the NVIDIA GPU Operator and NIM Operator
- Deploy, list, and remove NVIDIA NIM Cache resources
- Deploy, list, and remove NVIDIA NIM Services

**Note**: It is not mandatory to use the macro to deploy the GPU and NIM Operators. In fact, in some instances it is preferable (or even necessary) to deploy the GPU and NIM Operators externally (for example, using the OpenShift OperatorHub) instead of using the provided deployment action.

### Prerequisites

- Dataiku >= v14.2.0
- An [NVIDIA AI Enterprise](https://www.nvidia.com/en-us/data-center/products/ai-enterprise/) corporate subscription, or an [NVIDIA Developer Program](https://developer.nvidia.com) membership
- *NIM Container Registry* and *NIM Model Repository* credentials (which can both be an NVIDIA Authentication API Key if using [NGC](https://catalog.ngc.nvidia.com/))
- An attached Kubernetes cluster with:
    - An auto-provisioning Storage Class that supports "ReadWriteMany" access mode (see [NIM Operator docs](https://docs.nvidia.com/nim-operator/latest/service.html#example-create-a-pvc-instead-of-using-a-nim-cache))
    - (Optional) The Prometheus operator [installed on the cluster](https://docs.nvidia.com/nim-operator/latest/service.html#prerequisites-hpa), only required if leveraging horizontal pod autoscaling
    - (Optional) The [Nginx ingress controller](https://github.com/kubernetes/ingress-nginx) installed on the cluster; without this, the deployment macro defaults to exposing the NIM Services using a NodePort Kubernetes service

### Configuration

#### **Self-hosted Credentials** Preset
This preset stores *NIM Container Registry* and *NIM Model Repository* credentials. Use this preset when self-hosting NIMs on an attached Kubernetes cluster using the NIM Deployment Macro.

- Go to Plugins → Installed → NVIDIA NIM Plugin → Settings → Self-hosted credentials → Add a Preset.
- For Docker container registry, enter host, username and API key
- For NIM model registry:
    - If using NGC, simply enter the NGC API key
    - If using an alternative model registry (such as JFrog), select the 'override model registry' checkbox, and enter the model registry host, protocol and API key

![self-hosted credentials screenshot](assets/selfhosted-credentials.png?raw=true)

#### **NIM Environment Variables** Preset
This preset provides a mechanism to override the values of NIM environment variables. It should only be used when self-hosting NIMs on an attached Kubernetes cluster using the NIM Deployment Macro.

- Go to Plugins → Installed → NVIDIA NIM Plugin → Settings → NIM Environment Variables → Add a Preset.
- Override the value of any [NIM environment variable](https://docs.nvidia.com/nim/large-language-models/latest/configuration.html#environment-variables)

![self-hosted nim environment variables screenshot](assets/env-vars.png?raw=true)

### Using the NIM Deployment Macro

The NIM deployment macro is located in the Administration → Clusters → *Cluster name* → Actions tab of the Kubernetes cluster.

![Location of the NIM Deployment Macro in Cluster Actions tab](assets/nim-macro.png?raw=true)

#### GPU and NIM Operators

The first three macro actions provide the option to list, deploy and remove the NVIDIA NIM and GPU Operators. If these Operators are not already available on the cluster, they must be deployed prior to deploying your first NIM Service.

![NIM Deployment Macro with NIM Operator Add action selected](assets/nim-operator-add.png?raw=true)

![Output of the NIM Operator Inspect macro action](assets/nim-operator-inspect.png?raw=true)

#### NIM Cache

The NIM Cache feature allows you to pre-download and cache model files before deploying NIM Services. This significantly reduces startup time, especially for large models like BioNeMo models which can take 30+ minutes to download.


**Using NIM Cache:**

1. **Create a NIM Cache** - Use the "NIM Cache: Add" macro action to pre-download a model:
   - Select the model from the curated list or provide a custom image tag
   - Configure storage settings (storage class, volume size, access mode)
   - The cache name is auto-generated from the model tag (e.g., `meta-llama-3-1-8b-instruct-1-8-4-<timestamp>`)

2. **Monitor Cache Status** - Use "NIM Cache: Inspect" to check download progress and cache status

3. **Use Cache in NIM Services** - When creating a NIM Service with "NIM Services: Add":
   - Enable "Use NIM Cache" option
   - Select an existing cache from the dropdown
   - (Optional for LLMs) Specify a model profile (e.g., 'tp2-bf16') or leave empty for auto-selection

4. **Remove Cache** - Use "NIM Cache: Remove" to delete caches that are no longer needed

#### NIM Services

The next three macro actions provide the option to list, deploy and remove NVIDIA NIM Services. Under the hood, Dataiku leverages the NVIDIA NIM Operator, so all the options presented in the UI are those described in the NIM Operator [documentation](https://docs.nvidia.com/nim-operator/latest/service.html#about-nim-services).

- To deploy a NIM Service, use the "NIM Service: Add" macro:

![NIM Deployment Macro with NIM Services Add action selected](assets/nim-service-add.png?raw=true)

When deploying NIM Services, you can configure:

**Model Storage Options:**
- **Use NIM Cache** - Select a pre-existing NIM Cache for faster deployment (recommended for large models)
  - Choose an existing cache from the dropdown
  - Optionally specify a model profile for LLMs (e.g., 'tp2-bf16') or leave empty for auto-selection
- **Create New PVC** - The service will create its own storage and download the model on first startup

**Pod Settings:**
- **Startup Probe Failure Threshold** - Number of failures before considering the pod unhealthy. This to avoid pod restart for large NIMs
  - Default: 120 (allows ~20 minutes for startup with 10s period)
- **Startup Probe Period** - How often to check if the pod is ready (in seconds)
  - Default: 10s

- Then, use the "NIM Service: Inspect" macro to check the deployment status, and retrieve the model endpoint once the Service has been successfully rolled out. The model endpoint can then be used to create a NIM LLM Mesh connection.

![Output of the NIM Services Inspect macro showing a deployed model and its endpoint](assets/nim-service-inspect.png?raw=true)

##### Run:ai Feature

NIM services can be deployed with fraction GPUs to better optimize GPU utilization.

**Prerequisites**: Ensure that Run:ai is set up on your Kubernetes cluster and that the NIM Service Kubernetes namespace is mapped to a Run:ai project namespace with scheduler enforcement enabled. The plugin does not update Dataiku pod configuration directly.

During the NIM service creation select the following:

- Set **Kubernetes namespace** to a namespace governed by the target Run:ai project
- To use fraction GPUs, toggle the "Use fraction GPUs" button
    - Select Percentage of GPUs to use
    - Select number of GPU devices to use (Default 1 but multiple GPU can be defined)
    - Optionally select a dynamic GPU fraction; read (https://run-ai-docs.nvidia.com/self-hosted/platform-management/runai-scheduler/resource-optimization/dynamic-fractions.html) for more information

![Setting up fraction GPU with Runai](assets/runai-nim-services-fraction.png?raw=true)


---

## 2. NemoGuard NIMs Guardrails

The NemoGuard NIMs Guardrails component provides content safety, topic control, and jailbreak detection capabilities powered by NVIDIA NemoGuard microservices. These guardrails can be applied to LLM conversations to detect and prevent unsafe, off-topic, or malicious content.

### Capabilities

The plugin provides three types of guardrails:

1. **Topic Control** - Ensures conversations stay on approved topics
2. **Content Safety** - Detects inappropriate, harmful, or policy-violating content in user queries and model responses
3. **Jailbreak Detection** - Identifies attempts to bypass AI safety measures through prompt injection

### Prerequisites

- Dataiku >= v14.2.0
- Access to NVIDIA NemoGuard NIM endpoints (hosted on NVIDIA Cloud or self-hosted)
- An [NVIDIA AI Enterprise](https://www.nvidia.com/en-us/data-center/products/ai-enterprise/) corporate subscription, or an [NVIDIA Developer Program](https://developer.nvidia.com) membership

### Configuration

The configuration process depends on the type of guardrail you wish to implement.

#### For Topic Control and Content Safety

These guardrails use standard NIM models that must first be added to an NVIDIA NIM connection in the LLM Mesh.

1.  **Add the Guardrail Model to a Connection**: Navigate to **Administration > Connections** and either create a new **NVIDIA NIM** connection or edit an existing one. Add the Topic Control or Content Safety model to this connection.
2.  **Add the Guardrail to an LLM**: Go to the LLM you want to protect (in the same or a different connection) and navigate to the **Guardrails** tab.
3.  **Configure the Guardrail**:
    -   Click **Add Guardrail** and select **Nemo Guardrails NIMs**.
    -   **NIM Guardrail**: Choose **Topic Control** or **Content Safety**.
    -   **NVIDIA NIM Connection**: Select the connection where you configured the guardrail model.
    -   **Model**: Select the specific guardrail model from the dropdown.
    -   **Action to take**: Define the desired behavior (Reject, Audit, or Decline).

![Content Safety or Topic Control guardrail configured as Chat Completion models in NVIDIA NIM connection](assets/nim-model-config.png?raw=true)

#### For Jailbreak Detection

The Jailbreak Detection guardrail uses a model that is not exposed via a standard OpenAI-compatible API. Therefore, its endpoint must be configured directly within the guardrail settings.

1.  **Create an Authentication Connection**: Ensure you have an **NVIDIA NIM** connection configured in Dataiku. This connection will be used to authenticate requests to the jailbreak model endpoint, but the model itself is not added to this connection.
2.  **Add the Guardrail to an LLM**: Go to the LLM you want to protect and navigate to the **Guardrails** tab.
3.  **Configure the Guardrail**:
    -   Click **Add Guardrail** and select **Nemo Guardrails NIMs**.
    -   **NIM Guardrail**: Choose **Jailbreak Protection**.
    -   **NVIDIA NIM Connection**: Select the connection to use for authentication.
    -   **Model Endpoint**: Enter the full endpoint URL for the Jailbreak Detection NIM.
    -   **Action to take**: Define the desired behavior.

![Configuration of a Jailbreak Detection guardrail in a OpenAI LLM mesh connection](assets/jailbreak-guardrail-config.png?raw=true)

### Guardrail Types

#### Topic Control

Analyzes the full conversation context to ensure discussions remain within approved boundaries. The guardrail sends the complete message history to the NemoGuard Topic Control NIM, which determines if the conversation is on-topic. [Learn more](https://docs.nvidia.com/nim/llama-3-1-nemoguard-8b-topiccontrol/latest/prompt-template.html#system-instruction).

**Use Cases:**
- Restrict customer service chatbots to specific product domains
- Ensure educational assistants stay focused on learning topics
- Prevent brand chatbots from discussing competitor products

#### Content Safety

Checks both user queries and model responses for unsafe content including violence, hate speech, sexual content, and harassment. The guardrail analyzes the last user message and (if present) the assistant's response. [Learn more](https://docs.nvidia.com/nim/llama-3-1-nemoguard-8b-contentsafety/latest/index.html).

**Use Cases:**
- Filter inappropriate user inputs
- Prevent models from generating harmful content
- Maintain content policy compliance

#### Jailbreak Detection

Identifies prompt injection attacks and attempts to manipulate the AI system. Only operates on user queries (the last message in the conversation). [Learn more](https://docs.nvidia.com/nim/nemoguard-jailbreakdetect/latest/getting-started.html).

**Use Cases:**
- Protect against adversarial prompts
- Detect attempts to extract system prompts
- Block malicious instruction injection

### Usage

Once configured, guardrails automatically evaluate:
- **Queries**: Before they are sent to the LLM
- **Responses**: After the LLM generates a response (optional, configured via "Operate on responses")

Guardrails process the full conversation context, allowing NemoGuard NIMs to make informed decisions based on the entire dialogue history.

## 3. Run:ai Scheduling Macro

The Run:ai Scheduling macro gives Dataiku administrators an optional interface to validate and manage Run:ai projects for namespaces used by NIM workloads. It does not modify Dataiku pod templates or set Run:ai as a cluster-wide default scheduler directly.

### Capabilities

- Validate an existing Run:ai project for the selected Dataiku Kubernetes cluster and confirm scheduler enforcement is enabled
- Create a Run:ai project with an optional requested Kubernetes namespace
- Update GPU quota settings for an existing Run:ai project

### Prerequisites

- A reachable Run:ai control plane
- Run:ai API credentials stored in the plugin's **Run:ai credentials** preset
- A registered Run:ai cluster that matches the selected Dataiku Kubernetes cluster
- A Kubernetes namespace that is or will be governed by the target Run:ai project

### Configuration

#### Run:ai credentials preset

Go to Plugins → Installed → NVIDIA NIM Plugin → Settings → Run:ai credentials → Add a Preset.

Configure:

- **Run:ai API URL** - Base URL of the Run:ai control plane
- **Run:ai client ID** and **Run:ai client secret** - Credentials used to obtain short-lived API tokens
- **Dataiku to Run:ai cluster name mapping** - Optional mapping when the Dataiku cluster ID and the Run:ai cluster name differ

![runai_creds](assets/runai-credentials.png)

### Using the Run:ai Scheduling Macro

The macro is available from Administration → Clusters → *Cluster name* → Actions. Service Account with minimum permission of Editor and Department viewer on the tenant is a requirement.

#### Validate Run:ai project scheduling

Use this action to check whether a named Run:ai project exists on the Run:ai cluster mapped to the selected Dataiku cluster. The result shows the project namespace and whether `enforceRunaiScheduler` is enabled.

#### Create Run:ai project

Use this action to create a Run:ai project for a Dataiku workload namespace.

- Provide the **Run:ai project name**
- Optionally provide a **Kubernetes namespace**
- Set the **GPU quota** and optional **GPU limit**
- Optionally target a specific **Node pool ID**
- Optionally set a parent department / org unit

If you provide a namespace, the macro also ensures that namespace exists on the selected Dataiku Kubernetes cluster before creating the Run:ai project.

#### Update Run:ai project

Use this action to update the GPU quota configuration of an existing Run:ai project while keeping scheduler enforcement enabled.

### Using Run:ai with NIM Services

When submitting a NIM Service with fractional GPUs:

- Set the NIM Service **Kubernetes namespace** to a namespace mapped to the target Run:ai project
- Enable **Use fraction GPUs**
- Configure the GPU fraction, device count, and optional dynamic GPU fraction

If the namespace is not governed by the correct Run:ai project, the workload will not be scheduled through Run:ai as expected.
