{ "cells": [ { "cell_type": "markdown", "id": "356e542f", "metadata": { "tags": [ "cloud/nvidia/brev", "library/cudf", "library/xgboost", "library/prefect", "library/mlflow", "library/triton", "data-format/csv", "workflow/mlops" ] }, "source": [ "# Orchestrating a Fraud Detection Model Lifecycle with cuDF, Prefect, MLflow, and Triton\n", "\n", "_June, 2026_\n", "\n", "A production MLOps pipeline for the [NVIDIA Financial Fraud Detection AI Blueprint](https://github.com/NVIDIA-AI-Blueprints/financial-fraud-detection), built with [Prefect](https://www.prefect.io/), [MLflow](https://mlflow.org/), and [Triton Inference Server](https://github.com/triton-inference-server/server).\n", "\n", "The blueprint provides a [notebook](https://github.com/NVIDIA-AI-Blueprints/financial-fraud-detection/blob/main/notebooks/financial-fraud-usage-v2.ipynb) that manually walks through data preprocessing, GNN + XGBoost model training, Triton serving, and inference evaluation. This example takes that workflow and wraps each stage in automated orchestration with experiment tracking and continuous deployment for a production-ready model lifecycle implementation." ] }, { "cell_type": "markdown", "id": "518f17ff", "metadata": {}, "source": [ "## What the Blueprint Does\n", "\n", "The blueprint notebook can be divided into 4 steps:\n", "\n", "1. **Load and preprocess** 24M credit card transactions from the [IBM TabFormer](https://github.com/IBM/TabFormer) dataset using cuDF, building a bipartite User ↔ Merchant graph where edges represent transactions\n", "2. **Train** a 2-hop SAGEConv GNN that produces node embeddings, then an XGBoost classifier on those embeddings for edge classification (To determine fraudulent transactions)\n", "3. **Serve** the trained model on Triton Inference Server with a custom Python backend\n", "4. **Evaluate** by sending held-out test data to Triton and computing F1, precision, recall, and accuracy\n", "\n", "In the blueprint notebook, each step is a notebook cell that you run manually and inspect the output. However, in a production scenario, there is a need for orchestration, experiment tracking, model versioning, and automated deployment.\n", "\n", "The aim of this workflow example is to make this flow production-ready by adding the following components:\n", "\n", "- **Orchestration**: [Prefect](https://docs.prefect.io) flows chain the stages together, track runs, and handle failures\n", "- **Experiment tracking**: [MLflow](https://mlflow.org/docs/latest) logs hyperparameters, model artifacts, and evaluation metrics for every run\n", "- **Continuous deployment**: champion/challenger evaluation via [Triton's](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/index.html) native model versioning, with automatic promotion or rollback" ] }, { "cell_type": "markdown", "id": "4fc21661", "metadata": {}, "source": [ "## Architecture\n", "\n", "Similar to the structure of the notebook, the production pipeline has four stages, each implemented as a Prefect [flow](https://docs.prefect.io/3.0/develop/write-flows) composed of reusable [tasks](https://docs.prefect.io/3.0/develop/write-tasks):\n", "\n", "```text\n", "full_pipeline_flow\n", "├── preprocess_flow → cuDF graph formation + MLflow metadata logging\n", "├── train_flow → NGC training container + MLflow experiment logging\n", "├── evaluate_flow → Triton champion/challenger scoring + MLflow metrics\n", "└── deploy_flow → Promote or rollback + MLflow model registry\n", "```\n", "\n", "The full pipeline runs all four as subflows in sequence, passing results between them. Each stage can also run independently for ad-hoc experiments.\n", "\n", "### MLOps Components\n", "\n", "Three services support the pipeline, all running as Docker containers via `docker-compose.yml`:\n", "\n", "| Service | Role | Port |\n", "|---------|------|------|\n", "| **Prefect** | Flow orchestration, run tracking, scheduling, UI | 4200 |\n", "| **MLflow** | Experiment tracking, artifact storage, model registry | 5050 |\n", "| **Triton** | GPU inference, model versioning, model control API | 8000/8001/8002 |\n", "\n", "### Champion/Challenger via Triton Native Versioning\n", "\n", "The evaluation stage uses Triton's built-in model versioning for zero-downtime model comparison. Triton's model repository supports multiple version directories (`1/`, `2/`, `3/`...), and with `version_policy: { all: {} }` in `config.pbtxt`, all versions are served simultaneously." ] }, { "cell_type": "markdown", "id": "2930520b", "metadata": {}, "source": [ "## Component Structure\n", "\n", "```text\n", "fraud-detection-mlops/\n", "├── pipeline/\n", "│ ├── config.py Paths, service URLs, hyperparameters, thresholds\n", "│ ├── deployments.py Register flows with work pools and schedules\n", "│ ├── flows/\n", "│ │ ├── preprocess.py Stage 1: cuDF graph formation + MLflow\n", "│ │ ├── train.py Stage 2: NGC container + MLflow logging\n", "│ │ ├── evaluate.py Stage 3: Triton versioning + champion/challenger scoring\n", "│ │ ├── deploy.py Stage 4: promote or rollback\n", "│ │ └── full_pipeline.py Orchestrator: chains all stages as subflows\n", "│ └── tasks/\n", "│ ├── data.py Test data loading for evaluation\n", "│ ├── training.py Config generation + Docker container lifecycle\n", "│ ├── mlflow_utils.py Experiment tracking + model registry\n", "│ └── triton.py Version staging, model reload, inference scoring\n", "├── scripts/\n", "│ ├── preprocess_tabformer.py cuDF preprocessing (adapted from blueprint)\n", "│ └── download_data.sh TabFormer dataset download helper\n", "├── triton/\n", "│ └── Dockerfile Custom Triton image with PyTorch, XGBoost GPU, PyG\n", "├── docker-compose.yml MLflow + Prefect + Triton services\n", "├── environment.yml Conda environment (RAPIDS + pipeline deps)\n", "└── .env.example Configuration template\n", "```\n", "\n", "The `tasks/` layer contains reusable units of work (data loading, container execution, MLflow logging, Triton management). The `flows/` layer composes tasks into pipeline stages. This separation is important for scaling: tasks can be distributed across workers and [work pools](https://docs.prefect.io/3.0/deploy/infrastructure-concepts/work-pools), while flows define the orchestration logic." ] }, { "cell_type": "markdown", "id": "6fd03b8b", "metadata": {}, "source": [ "## Prerequisites\n", "\n", "- **NVIDIA GPU** with compute capability 7.5 or newer (Turing architecture or later: T4, A10, L4, A100, H100, etc.). Required by the training container.\n", "- **Docker** with the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) for GPU access.\n", "- **Conda**: [miniforge](https://github.com/conda-forge/miniforge) is recommended.\n", "- **NGC API key** from the [NVIDIA NGC API key page](https://org.ngc.nvidia.com/account/api-key), required to pull the training container.\n" ] }, { "cell_type": "markdown", "id": "7a52b98c5161", "metadata": {}, "source": [ "## Provision a GPU instance with NVIDIA Brev\n", "\n", "Use [NVIDIA Brev](https://brev.nvidia.com/) to provision a GPU instance for this example. Brev provisions GPU instances across AWS, GCP, and other cloud service providers, with NVIDIA drivers, Docker, and the NVIDIA Container Toolkit pre-installed.\n", "\n", "````{docref} /cloud/nvidia/brev\n", "Follow the instructions in [NVIDIA Brev](../../cloud/nvidia/brev) to launch an instance with:\n", "\n", "- **GPU**: a single GPU meeting the Prerequisites above.\n", "- **Host memory**: at least **64 GB of RAM** so the cuDF preprocessing comfortably fits the 24 M-row TabFormer dataset in memory.\n", "- **Disk**: at least 100 GB. The training container is ~40 GB on disk, the Triton image is ~24 GB, and the TabFormer dataset adds another ~5 GB.\n", "- **Runtime**: VM Mode w/ Jupyter. The docker-compose stack started in the Setup section below manages MLflow, Prefect, and Triton.\n", "\n", "Once the instance is running, SSH in with `brev shell ` or open it in VS Code via `brev open `.\n", "````\n", "\n", "### Install miniforge\n", "\n", "The Brev VM image ships with Docker and CUDA but not conda. Install miniforge:\n", "\n", "```console\n", "$ curl -L -O \"https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh\"\n", "$ bash Miniforge3-$(uname)-$(uname -m).sh -b -p $HOME/miniforge\n", "$ $HOME/miniforge/bin/conda init bash\n", "$ source ~/.bashrc\n", "```\n", "\n", "From here, continue with the Setup section below.\n" ] }, { "cell_type": "markdown", "id": "8048a7c1", "metadata": {}, "source": [ "## Setup\n", "\n", "Fetch the scripts and helper files onto the Brev instance and switch to the example directory:\n", "\n", "```console\n", "$ curl -sSL \"https://github.com/rapidsai/deployment/archive/refs/heads/main.tar.gz\" \\\n", " | bsdtar -xzf - --strip-components=3 \\\n", " \"deployment-main/source/examples/fraud-detection-mlops-pipeline\"\n", "$ cd fraud-detection-mlops-pipeline\n", "```\n", "\n", "### Configure\n", "\n", "Copy the environment template, set your NGC API key, and export the variables into the current shell:\n", "\n", "```console\n", "$ cp .env.example .env\n", "$ nano .env\n", "$ set -a\n", "$ source .env\n", "$ set +a\n", "```\n", "\n", "Docker Compose reads `.env` automatically, but the host Python process only sees these values after they are exported into the shell. Repeat the `set -a; source .env; set +a` step in each new shell before running the Python flows.\n", "\n", "### Install Python dependencies\n", "\n", "All dependencies (RAPIDS, Prefect, MLflow, Triton client) are specified in `environment.yml`:\n", "\n", "```console\n", "$ conda env create -f environment.yml\n", "$ conda activate fraud-mlops\n", "```\n", "\n", "### Download the TabFormer dataset\n", "\n", "The dataset is 266MB of synthetic credit card transactions (24M rows, 15 columns) from [IBM TabFormer](https://github.com/IBM/TabFormer). Download `transactions.tgz` on your local machine from the [IBM Box TabFormer data folder](https://ibm.ent.box.com/v/tabformer-data/folder/130747715605), then copy it to the Brev instance. If your browser downloads a `cred_card.zip` archive instead, unzip it locally and copy the `transactions.tgz` file inside.\n", "\n", "On the Brev instance, create the raw data directory:\n", "\n", "```console\n", "$ mkdir -p ./data/TabFormer/raw\n", "```\n", "\n", "From your local machine, use `brev copy`:\n", "\n", "```console\n", "$ brev copy ./transactions.tgz :~/fraud-detection-mlops-pipeline/data/TabFormer/raw/transactions.tgz\n", "```\n", "\n", "Alternatively, open Jupyter on the Brev instance and drag `transactions.tgz` into `fraud-detection-mlops-pipeline/data/TabFormer/raw/`.\n", "\n", "Back on the Brev instance, extract the dataset:\n", "\n", "```console\n", "$ ./scripts/download_data.sh ./data/TabFormer\n", "```\n", "\n", "The helper script exits with instructions if `transactions.tgz` is not in `./data/TabFormer/raw/` yet.\n", "\n", "### Pull the training container\n", "\n", "```console\n", "$ echo \"$NGC_API_KEY\" | docker login nvcr.io --username '$oauthtoken' --password-stdin\n", "$ docker pull nvcr.io/nvidia/cugraph/financial-fraud-training:2.0.0\n", "```\n", "\n", "### Build the Triton serving image\n", "\n", "The custom Triton image (`triton/Dockerfile`) installs PyTorch 2.7, XGBoost 3.0 (GPU), PyTorch Geometric 2.6, and Captum on top of the Triton 25.04 base image:\n", "\n", "```console\n", "$ docker build -t triton-fraud:latest triton/\n", "```\n", "\n", "This build takes 15-20 minutes and produces a ~24GB image.\n", "\n", "### Start the MLOps Components\n", "\n", "Create the host-side bind mount directories before starting Docker Compose. This keeps Docker from creating missing directories as root, which would prevent the host Python process from staging Triton model versions later.\n", "\n", "```console\n", "$ mkdir -p ./data/models ./data/trained_models ./mlflow-data ./prefect-data\n", "```\n", "\n", "For the end-to-end pipeline, start all three MLOps components Prefect, MLflow, and Triton:\n", "\n", "```console\n", "$ docker compose --profile gpu up -d\n", "```\n", "\n", "If you only want to run preprocessing or training, Triton is not required. In that case, start only Prefect and MLflow:\n", "\n", "```console\n", "$ docker compose up -d\n", "```\n", "\n", "If you change the Compose configuration or rebuild the Triton image, recreate the Triton container so the new environment is applied:\n", "\n", "```console\n", "$ docker compose --profile gpu up -d --force-recreate triton\n", "```\n", "\n", "If a previous run created root-owned `__pycache__` directories inside `./data/models`, stop Triton, remove those bytecode directories, fix ownership, and recreate Triton before rerunning the pipeline:\n", "\n", "```console\n", "$ docker compose --profile gpu stop triton\n", "$ sudo find ./data/models -type d -name __pycache__ -prune -exec rm -rf {} +\n", "$ sudo chown -R \"$USER:$USER\" ./data/models\n", "$ docker compose --profile gpu up -d --force-recreate triton\n", "```\n", "\n", "Verify the services are running from the Brev shell:\n", "\n", "```console\n", "$ curl -fsS http://127.0.0.1:4200/api/health\n", "$ curl -fsS http://127.0.0.1:5050/health\n", "$ curl -fsS http://127.0.0.1:8000/v2/health/ready\n", "```\n", "\n", "To open the UIs in your local browser, run port forwarding commands from your local machine and keep them running while you use the pages:\n", "\n", "```console\n", "$ brev port-forward -p 4200:4200\n", "$ brev port-forward -p 5050:5050\n", "```\n", "\n", "- Prefect UI: [http://localhost:4200](http://localhost:4200)\n", "- MLflow UI: [http://localhost:5050](http://localhost:5050)" ] }, { "cell_type": "markdown", "id": "d02b3bc0", "metadata": {}, "source": [ "## Running the Pipeline\n", "\n", "### Full pipeline\n", "\n", "Run all four stages end-to-end:\n", "\n", "```console\n", "$ python -m pipeline.flows.full_pipeline\n", "```\n", "\n", "This chains preprocessing → training → evaluation → deployment as subflows, passing results between stages automatically. If graph data already exists from a previous preprocessing run, skip preprocessing with:\n", "\n", "```console\n", "$ python -c 'from pipeline.flows.full_pipeline import full_pipeline_flow; full_pipeline_flow(skip_preprocess=True)'\n", "```\n", "\n", "### Individual stages\n", "\n", "Each stage can be run independently for debugging or ad-hoc experiments:\n", "\n", "```console\n", "# Preprocess raw CSV into graph data\n", "$ python -m pipeline.flows.preprocess\n", "\n", "# Train with default hyperparameters\n", "$ python -m pipeline.flows.train\n", "\n", "# Evaluate a specific training run\n", "$ python -m pipeline.flows.evaluate \n", "\n", "# Deploy based on evaluation result\n", "$ python -m pipeline.flows.deploy '{\"should_promote\": true, ...}'\n", "```\n", "\n", "### Experimenting with hyperparameters\n", "\n", "To try different model configurations, pass custom hyperparameters to the train flow:\n", "\n", "```python\n", "from pipeline.flows.train import train_flow\n", "\n", "run_id = train_flow(\n", " gnn_params={\"hidden_channels\": 64, \"n_hops\": 3, \"num_epochs\": 12, ...},\n", " xgb_params={\"max_depth\": 8, \"num_boost_round\": 1024, ...},\n", ")\n", "```\n", "\n", "Each run is tracked in MLflow with full hyperparameters and metrics, making it straightforward to compare experiments. The evaluation stage will then score this new model against the current champion." ] }, { "cell_type": "markdown", "id": "3c5fda64", "metadata": {}, "source": [ "## Test Inference\n", "\n", "After the deployment stage completes, the current champion model is served by Triton. First confirm Triton exposes the model metadata:\n", "\n", "```console\n", "$ curl -fsS http://127.0.0.1:8000/v2/models/prediction_and_shapley | python -m json.tool\n", "```\n", "\n", "Then send held-out graph data to the served model and print the inference metrics:\n", "\n", "```console\n", "$ python - <<'PY'\n", "import json\n", "import numpy as np\n", "\n", "from pipeline.config import (\n", " DECISION_THRESHOLD,\n", " TEST_DATA_DIR,\n", " TRITON_HTTP_URL,\n", " TRITON_MODEL_NAME,\n", " TRITON_MODEL_REPO,\n", ")\n", "from pipeline.tasks.data import load_test_data\n", "from pipeline.tasks.triton import get_current_version, score_model_version\n", "\n", "test_data = load_test_data.fn(TEST_DATA_DIR)\n", "labels = test_data[\"edge_label_user_to_merchant\"]\n", "if hasattr(labels, \"to_numpy\"):\n", " labels = labels.to_numpy(dtype=np.int32)\n", "if hasattr(labels, \"values\"):\n", " labels = labels.values\n", "labels = np.asarray(labels, dtype=np.int32).ravel()\n", "\n", "inference_data = {\n", " key: value for key, value in test_data.items()\n", " if not key.startswith(\"edge_label_\")\n", "}\n", "version = get_current_version.fn(TRITON_MODEL_REPO, TRITON_MODEL_NAME)\n", "metrics = score_model_version.fn(\n", " TRITON_HTTP_URL,\n", " TRITON_MODEL_NAME,\n", " version,\n", " inference_data,\n", " labels,\n", " DECISION_THRESHOLD,\n", ")\n", "\n", "print(f\"Served model: {TRITON_MODEL_NAME} version {version}\")\n", "print(json.dumps(metrics, indent=2))\n", "PY\n", "```\n", "\n", "This sends real model inputs to Triton, receives `PREDICTION`, and computes F1, precision, recall, and accuracy of the served model.\n" ] }, { "cell_type": "markdown", "id": "881b10b3", "metadata": {}, "source": [ "## Monitoring and Logging\n", "\n", "### Prefect UI\n", "\n", "The [Prefect UI](http://localhost:4200) at port 4200 tracks every flow run. Here is a completed full pipeline run showing the three subflows (train, evaluate, deploy) with their execution timeline:\n", "\n", "![Prefect flow run showing subflow hierarchy](../../images/prefect-flow-runs.png)\n", "\n", "The UI also shows task-level logs from each step, which is useful for debugging container failures or Triton connection issues.\n", "\n", "### MLflow UI\n", "\n", "The [MLflow UI](http://localhost:5050) at port 5050 tracks experiments across runs. The experiment table shows all training runs with their hyperparameters and evaluation metrics side by side:\n", "\n", "![MLflow experiment table with training runs](../../images/mlflow-experiment-table.png)\n", "\n", "Clicking into a run shows the full details: all GNN and XGBoost hyperparameters, challenger and champion metrics, deltas, and the promotion decision:\n", "\n", "![MLflow run details with parameters and metrics](../../images/mlflow-run-details.png)\n", "\n", "The model registry tracks which version is in production. The `champion` alias always points to the current production model:\n", "\n", "![MLflow model registry with champion alias](../../images/mlflow-model-registry.png)\n", "\n", "In a typical production cycle, start with a full pipeline run to establish a baseline champion. From there, the cycle is: tweak hyperparameters, run training and evaluation, and check the MLflow experiment table to compare the challenger against the current champion. Promoted models are immediately live in Triton without any need for a manual deployment." ] }, { "cell_type": "markdown", "id": "6e78c5a1", "metadata": {}, "source": [ "## Pipeline Stage Architecture\n", "\n", "The sections below explain what each pipeline stage does after you have run the end-to-end workflow and verified inference.\n", "\n", "### Preprocessing\n", "\n", "**Flow:** `pipeline/flows/preprocess.py` | **Script:** `scripts/preprocess_tabformer.py`\n", "\n", "The preprocessing stage takes the raw TabFormer CSV and produces graph-structured data for the GNN training container.\n", "\n", "The script (`scripts/preprocess_tabformer.py`, adapted from the blueprint's `preprocess_TabFormer_lp.py`) uses **cuDF** for the following GPU-accelerated DataFrame operations on the 24M-row dataset:\n", "\n", "1. **Clean column names**: strip spaces, standardize field names\n", "2. **Encode categorical features**: one-hot encoding for low-cardinality fields (< 8 categories), binary encoding for high-cardinality fields\n", "3. **Build the bipartite graph**: Users and Merchants become nodes, each transaction becomes an edge with encoded attributes and a fraud/not-fraud label\n", "4. **Undersample**: balance the dataset to a configurable fraud ratio (default 10%)\n", "5. **Split by year**: pre-2018 for training, 2018 for validation, post-2018 for testing\n", "6. **Write CSVs**: node features, edge indices, edge attributes, edge labels, and feature masks to `gnn/` and `gnn/test_gnn/`\n", "\n", "The Prefect flow wraps this script and logs preprocessing metadata (row counts, graph dimensions) to MLflow:\n", "\n", "```python\n", "# pipeline/flows/preprocess.py\n", "@flow(name=\"preprocess\", log_prints=True)\n", "def preprocess_flow(\n", " raw_csv_path: str = RAW_CSV_PATH,\n", " output_base_path: str = DATA_ROOT,\n", " fraud_ratio: float = DEFAULT_FRAUD_RATIO,\n", " under_sample: bool = DEFAULT_UNDER_SAMPLE,\n", ") -> dict:\n", " # Step 1: Run preprocessing (import here because it requires cuDF/GPU)\n", " scripts_dir = os.path.join(os.path.dirname(__file__), \"..\", \"..\", \"scripts\")\n", " sys.path.insert(0, os.path.abspath(scripts_dir))\n", " from preprocess_tabformer import preprocess_data\n", "\n", " metadata, user_mask_map, mx_mask_map, tx_mask_map = preprocess_data(\n", " raw_csv_path=raw_csv_path,\n", " output_base_path=output_base_path,\n", " fraud_ratio=fraud_ratio,\n", " under_sample=under_sample,\n", " )\n", "\n", " # Step 2: Log to MLflow\n", " experiment_id = get_or_create_experiment(MLFLOW_PREPROCESS_EXPERIMENT_NAME)\n", " with mlflow.start_run(experiment_id=experiment_id, run_name=\"preprocess\"):\n", " mlflow.log_params(\n", " {\n", " \"preprocess.fraud_ratio\": fraud_ratio,\n", " \"preprocess.under_sample\": under_sample,\n", " }\n", " )\n", " mlflow.log_metrics(\n", " {\n", " \"preprocess.row_count\": float(metadata[\"row_count\"]),\n", " \"preprocess.num_users\": float(metadata[\"num_users\"]),\n", " \"preprocess.num_merchants\": float(metadata[\"num_merchants\"]),\n", " \"preprocess.num_transactions\": float(metadata[\"num_transactions\"]),\n", " }\n", " )\n", " mlflow.set_tag(\"stage\", \"preprocess\")\n", "```" ] }, { "cell_type": "markdown", "id": "68b48cfd", "metadata": {}, "source": [ "### Training\n", "\n", "**Flow:** `pipeline/flows/train.py` | **Tasks:** `pipeline/tasks/training.py`, `pipeline/tasks/mlflow_utils.py`\n", "\n", "Training uses NVIDIA's `financial-fraud-training:2.0.0` NGC container for training. The pipeline generates the config in the format the container expects, runs the container, and collects the artifacts.\n", "\n", "#### Config generation\n", "\n", "The `generate_training_config` task builds a JSON config from the flow's hyperparameters:\n", "\n", "```python\n", "# pipeline/config.py\n", "DEFAULT_GNN_PARAMS = {\n", " \"hidden_channels\": 32,\n", " \"n_hops\": 2,\n", " \"layer\": \"SAGEConv\",\n", " \"dropout_prob\": 0.1,\n", " \"batch_size\": 4096,\n", " \"fan_out\": 10,\n", " \"num_epochs\": 8,\n", "}\n", "\n", "DEFAULT_XGB_PARAMS = {\n", " \"max_depth\": 6,\n", " \"learning_rate\": 0.2,\n", " \"num_parallel_tree\": 3,\n", " \"num_boost_round\": 512,\n", " \"gamma\": 0.0,\n", "}\n", "```\n", "\n", "#### Container execution\n", "\n", "The `run_training_container` task runs the NGC container with volume mounts for graph data, output directory, and config:\n", "\n", "```console\n", "$ docker run --rm --gpus device=0 \\\n", " --cap-add SYS_NICE --shm-size=8g --privileged \\\n", " -v :/data \\\n", " -v :/trained_models \\\n", " -v :/app/config.json \\\n", " --entrypoint bash \\\n", " nvcr.io/nvidia/cugraph/financial-fraud-training:2.0.0 \\\n", " -c \"torchrun --standalone --nproc_per_node=1 /app/main.py --config /app/config.json\"\n", "```\n", "\n", "The actual script also includes various NCCL environment variables (`NCCL_IB_DISABLE`, `NCCL_NET=Socket`, etc.) which are a workaround for instances without InfiniBand hardware. The container is bundled with UCX that is designed for faster HPC networking, which is not usually available on standard single GPU instances.\n", "\n", "```{note}\n", "If running on an InfiniBand-enabled system (DGX, HPC cluster), remove the NCCL environment variables from `pipeline/tasks/training.py`. These overrides disable InfiniBand, NVLink, and shared memory transports. On hardware that supports them, NCCL auto-detects and uses the optimal transport without any configuration.\n", "```\n", "\n", "#### MLflow logging\n", "\n", "After training completes, the flow logs everything to MLflow:\n", "\n", "- **Hyperparameters**: `gnn.hidden_channels`, `gnn.n_hops`, `xgb.max_depth`, `xgb.num_boost_round`, etc.\n", "- **Artifacts**: the training config JSON and the full model directory\n", "- **Tags**: `trigger: automated` (or custom tags for ad-hoc experiments)\n", "\n", "The flow returns the MLflow run ID, which the evaluation stage uses to associate metrics with the same run." ] }, { "cell_type": "markdown", "id": "a75c60f7", "metadata": {}, "source": [ "### Evaluation\n", "\n", "**Flow:** `pipeline/flows/evaluate.py` | **Tasks:** `pipeline/tasks/triton.py`, `pipeline/tasks/data.py`\n", "\n", "The evaluation stage compares the newly trained model (challenger) against the currently deployed model (champion) using Triton's native versioning.\n", "\n", "#### Staging the challenger\n", "\n", "The `stage_challenger_version` task copies the new model's artifacts as the next version directory in Triton's model repository. If the champion is version 1, the challenger becomes version 2:\n", "\n", "```text\n", "prediction_and_shapley/\n", "├── 1/ ← champion (currently serving)\n", "│ ├── model.py\n", "│ ├── state_dict_gnn_model.pth\n", "│ └── embedding_based_xgboost.json\n", "├── 2/ ← challenger (just staged)\n", "│ ├── model.py\n", "│ ├── state_dict_gnn_model.pth\n", "│ └── embedding_based_xgboost.json\n", "└── config.pbtxt (version_policy { all {} })\n", "```\n", "\n", "The task also copies the generated root `config.pbtxt` from the training output repository into the serving repository and ensures it has `version_policy { all {} }` so Triton serves both versions simultaneously. Without this, Triton defaults to `latest: { num_versions: 1 }` and would drop the champion when the challenger is loaded.\n", "\n", "#### Scoring\n", "\n", "After reloading the model in Triton, the flow loads the held-out test data (25,803 samples, ~8% fraud) and sends it to both versions using Triton's version-specific inference:\n", "\n", "```python\n", "# pipeline/tasks/triton.py\n", "response = client.infer(\n", " model_name,\n", " inputs=inputs,\n", " model_version=str(version),\n", " outputs=outputs,\n", ")\n", "```\n", "\n", "Each version returns fraud probability scores. The task applies a decision threshold (default 0.5) and computes F1, precision, recall, and accuracy against the ground truth labels.\n", "\n", "#### Promotion decision\n", "\n", "The flow compares the challenger's F1 score against the champion's. Promotion requires the challenger to exceed the champion by at least `min_improvement` (default 0.0):\n", "\n", "```python\n", "challenger_val = challenger_metrics.get(promotion_metric, 0)\n", "champion_val = champion_metrics.get(promotion_metric, 0)\n", "should_promote = challenger_val > champion_val + min_improvement\n", "```\n", "\n", "Both sets of metrics, the deltas, and the promotion decision are logged to the same MLflow run that was created during training." ] }, { "cell_type": "markdown", "id": "c75168ae", "metadata": {}, "source": [ "### Deployment\n", "\n", "**Flow:** `pipeline/flows/deploy.py`\n", "\n", "The deployment stage acts on the evaluation result.\n", "\n", "#### If the challenger wins (promote)\n", "\n", "1. Remove the old champion's version directory from the model repository\n", "2. Reload the model in Triton (only the challenger version remains)\n", "3. Run a health check to verify the new champion is responsive\n", "4. Register the model version in MLflow's model registry and assign the `champion` alias\n", "\n", "If the health check fails, the flow raises an error. This surfaces in Prefect as a failed run that you can investigate.\n", "\n", "#### If the champion wins (reject)\n", "\n", "1. Remove the challenger's version directory\n", "2. Reload Triton to restore the champion-only state\n", "\n", "In both cases, Triton ends up serving exactly one version: the winner. The MLflow model registry provides an authoritative record of which model version is in production via the `champion` alias:\n", "\n", "```python\n", "# pipeline/tasks/mlflow_utils.py\n", "_create_registered_model_if_missing(client, model_name)\n", "model_uri = f\"runs:/{run_id}/model\"\n", "mv = client.create_model_version(name=model_name, source=model_uri, run_id=run_id)\n", "client.set_registered_model_alias(model_name, \"champion\", mv.version)\n", "```" ] }, { "cell_type": "markdown", "id": "10b99a5d", "metadata": {}, "source": [ "## Scaling to Multiple Machines\n", "\n", "The default Brev walkthrough runs Prefect, MLflow, Triton, and all Python flows on one GPU instance. That is the simplest path for this example and is the recommended setup when you are validating the workflow end to end.\n", "\n", "For production-style orchestration, you can split the MLOps tracking components (which do not require access to a GPU) from the GPU execution machine. In that setup, `pipeline/deployments.py` registers each flow as a Prefect [deployment](https://docs.prefect.io/3.0/deploy/infrastructure-concepts/deployments) with its own work pool and schedule.\n", "\n", "### How it works\n", "\n", "The evaluation and deployment stages need both filesystem access to Triton's model repository and HTTP access to the Triton API. As a result, the GPU instance needs to run Triton and the Prefect worker. The central host only runs Prefect and MLflow so that flow orchestration, experiment tracking, and the UIs are available from the host.\n", "\n", "The split-machine setup uses two machines:\n", "\n", "- **Orchestration host**: runs Prefect and MLflow servers. This machine does not need a GPU.\n", "- **GPU instance**: runs the Prefect worker, executes all pipeline flows, and serves models with Triton.\n", "\n", "On Brev, the easiest way to view the UIs is still to keep Prefect and MLflow on the GPU instance and forward the UI ports to your local machine:\n", "\n", "```console\n", "$ brev port-forward -p 4200:4200\n", "$ brev port-forward -p 5050:5050\n", "```\n", "\n", "Use the split-machine setup only when the GPU instance can reach the orchestration host on ports 4200 and 5050. That usually requires a shared private network, VPN, public endpoint, or reverse tunnel.\n", "\n", "### Setting up deployments\n", "\n", "**1. Start Prefect and MLflow on the orchestration host:**\n", "\n", "```console\n", "$ docker compose up -d\n", "```\n", "\n", "This starts Prefect on port 4200 and MLflow on port 5050. Both UIs are available from that host.\n", "\n", "**2. Point the GPU instance at the orchestration host** by setting the API URLs:\n", "\n", "```console\n", "# On the GPU instance\n", "$ export PREFECT_API_URL=http://:4200/api\n", "$ export MLFLOW_TRACKING_URI=http://:5050\n", "```\n", "\n", "**3. Start Triton on the GPU instance:**\n", "\n", "```console\n", "$ docker compose --profile gpu up -d triton\n", "```\n", "\n", "**4. Create the work pool** (run once, from the orchestration host):\n", "\n", "```console\n", "$ prefect work-pool create gpu --type process\n", "```\n", "\n", "**5. Register and serve all deployments** from the GPU instance, after exporting `PREFECT_API_URL` and `MLFLOW_TRACKING_URI`:\n", "\n", "```console\n", "$ python -m pipeline.deployments\n", "```\n", "\n", "This registers 5 deployments. Simulating a real world scenario, preprocessing runs at midnight to rebuild graph data from the latest raw CSV, and the full pipeline runs at 2am (after preprocessing completes) to retrain, evaluate, and potentially deploy a new champion. The individual stages are on-demand for ad-hoc experiments:\n", "\n", "- `preprocess` (gpu pool, nightly at midnight)\n", "- `train` (gpu pool, on-demand)\n", "- `evaluate` (gpu pool, on-demand)\n", "- `deploy` (gpu pool, on-demand)\n", "- `full-pipeline` (gpu pool, nightly at 2am)\n", "\n", "**6. Start a worker on the GPU instance:**\n", "\n", "```console\n", "$ prefect worker start --pool gpu\n", "```\n", "\n", "You can trigger runs from the Prefect UI on the orchestration host or via the CLI:\n", "\n", "```console\n", "$ prefect deployment run full-pipeline/full-pipeline\n", "```\n", "\n", "The components of the pipeline do not change between single-machine and split-machine modes. The difference is that flow runs are tracked and triggered via the central Prefect server, while all execution still happens on the GPU instance.\n", "\n", "### Data ingestion in production\n", "\n", "This example uses a static CSV (`card_transaction.v1.csv`), but in a production scenario, new data arrives continuously. The preprocessing flow accepts a `raw_csv_path` parameter, which serves as the integration point with your data infrastructure. Common patterns on how newer data is fetched include:\n", "\n", "- **Scheduled ETL jobs.** An upstream pipeline (Spark, Airflow, dbt) lands a fresh CSV or Parquet file at a known path on a schedule. The preprocessing cron picks it up at midnight.\n", "- **Cloud storage pulls.** Transaction data accumulates in S3 or GCS. A lightweight script or Prefect task pulls the latest file to the GPU instance before preprocessing runs.\n", "- **Data warehouse queries.** A scheduled query extracts new records since the last run and writes them to the expected path.\n", "\n", "To integrate any of these, add your data fetching logic as a Prefect task at the start of the preprocessing flow. It runs as part of the same nightly cron schedule, so fresh data is pulled and processed in a single run without any additional orchestration.\n" ] }, { "cell_type": "markdown", "id": "cffc18f0", "metadata": {}, "source": [ "## Conclusion\n", "\n", "This example demonstrates how to take a ML experiment and build production infrastructure around it:\n", "\n", "- **Prefect** orchestrates the pipeline stages, tracks runs, and enables scheduling\n", "- **MLflow** tracks every experiment (hyperparameters, artifacts, metrics) and manages the model registry\n", "- **Triton** serves models with native versioning, enabling champion/challenger evaluation without separate infrastructure\n", "\n", "The same patterns (flow orchestration, experiment tracking, versioned deployment) broadly apply to any ML production scenario where you need to automate training, compare model versions, and deploy the winner.\n", "\n", "### Resources\n", "\n", "- [NVIDIA Financial Fraud Detection AI Blueprint](https://github.com/NVIDIA-AI-Blueprints/financial-fraud-detection)\n", "- [IBM TabFormer dataset](https://github.com/IBM/TabFormer)\n", "- [Prefect 3.x documentation](https://docs.prefect.io)\n", "- [MLflow documentation](https://mlflow.org/docs/latest)\n", "- [Triton Inference Server](https://github.com/triton-inference-server/server)" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.13.0" } }, "nbformat": 4, "nbformat_minor": 5 }