Virchow2 model

[Jurgee]

This PR introduces support for the Virchow2 foundation model (paige-ai/Virchow2) within the Ray Serve infrastructure.

New Model Deployment: Added virchow2.py implementing the Virchow2 class as a Ray Serve deployment

Summary by CodeRabbit

• New Features

  • Virchow2 foundation-model deployment with FastAPI ingress for image embeddings.
  • Automated HuggingFace downloader job and local model-download utility for offline caching.
  • PersistentVolumeClaim and shared cache mounts for HuggingFace assets.
  • New provider helper to fetch models/files from HuggingFace.

• Refactor

  • Removed thread-level inference configuration from model settings.

• Chores

  • Updated service deployment, worker scaling, resource allocations, GPU runtime packaging and environment (HF token, cache) for HuggingFace access.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds a Virchow2 Ray Serve deployment with FastAPI ingress, HuggingFace caching/support (PVC, provider helper, downloader job), Dockerfile and Ray service updates for HF/GPU integration, and removes intra_op_num_threads usage from model configs.

Changes

Cohort / File(s)	Summary
Docker Configuration docker/Dockerfile.cpu, docker/Dockerfile.gpu	Dockerfile.cpu: reformatted pip install into a multiline RUN. Dockerfile.gpu: split pip installs into multiple RUN steps, added GPU-specific packages and external indices (onnxruntime-gpu, tensorrt-cu12, torch/cu121, torchvision/cu121, timm, huggingface-hub).
Virchow2 Downloader Utilities misc/virchow2_downloader/download_virchow2.py, misc/virchow2_downloader/virchow2_downloader_job.yaml	Added script to cache HuggingFace model snapshots (supports HF_TOKEN) and a Kubernetes Job manifest to run it with PVC mount, resource limits, secure non-root context, and required env vars.
Model Config Cleanup models/binary_classifier.py, models/semantic_segmentation.py	Removed intra_op_num_threads from Config TypedDicts and stopped setting sess_options.intra_op_num_threads; small formatting tweaks.
Virchow2 Model Implementation models/virchow2.py	New Ray Serve deployment Virchow2 with embedded FastAPI ingress, reconfigure to load timm model from HF/cache, preprocessing transforms, batched async predict, LZ4-compressed request handling, and exported app = Virchow2.bind().
Model Provider Enhancement providers/model_provider.py	Added `huggingface(repo_id, filename: str
Kubernetes / Ray Serve Manifests pvc/huggingface-pvc.yaml, ray-service.yaml	Added 15Gi NFS-backed PVC huggingface-cache-pvc. Updated ray-service.yaml: new virchow2 application, HF_HOME envs, HF_TOKEN secret in GPU workers, huggingface-cache volume mounts in cpu/gpu workers, head/worker image and resource updates, MIG GPU selector, actor sizing, and virchow2 user_config. Removed intra_op_num_threads from BinaryClassifier user_config.

Sequence Diagram(s)

sequenceDiagram
    participant Client
    participant RayServe as Ray Serve (Virchow2)
    participant Provider as providers.huggingface
    participant PVC as HuggingFace PVC
    participant HFHub as HuggingFace Hub
    participant Downloader as Downloader Job

    Client->>RayServe: POST LZ4-compressed image(s)
    RayServe->>Provider: request local path for repo_id (HF_HOME)
    Provider-->>RayServe: return local cache path (exists?) 
    alt cache miss
        RayServe->>Downloader: schedule/run downloader job
        Downloader->>HFHub: snapshot_download / hf_hub_download (uses HF_TOKEN)
        HFHub-->>PVC: write model files to cache
        Downloader-->>RayServe: notify completion
    end
    RayServe->>RayServe: preprocess, batch, run timm model (GPU)
    RayServe-->>Client: return embeddings (LZ4-compressed)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Poem

🐇 I hopped to the cache and found a bright seed,
A model unrolled with embeddings to feed,
From HuggingFace burrow to PVC nest,
Ray serves each tile — the rabbits impressed!
🥕

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 22.22% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title "Virchow2 model" is vague and generic. While it mentions the model name, it does not describe the actual change or what was accomplished—such as 'adds support for', 'integrates', 'implements deployment for', or similar action verbs that clarify the scope of work.	Consider a more descriptive title like 'Add Virchow2 foundation model support to Ray Serve' or 'Implement Virchow2 model deployment with Ray Serve integration'.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feature/virchow2-model

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the model serving infrastructure by integrating the Virchow2 foundation model and optimizing existing models for GPU performance. It establishes a robust framework for deploying advanced deep learning models, leveraging TensorRT for efficient inference and Hugging Face for model management. The changes also include necessary infrastructure updates for dependency management and persistent caching, ensuring a more efficient and scalable system.

Highlights

• Virchow2 Model Integration: Introduced the Virchow2 foundation model (paige-ai/Virchow2) as a new Ray Serve deployment, enabling its use within the existing infrastructure.
• GPU Acceleration with TensorRT: Enabled GPU acceleration for existing BinaryClassifier and SemanticSegmentation models by integrating TensorRT, significantly improving inference performance.
• Infrastructure Updates for Model Caching: Added support for persistent caching of Hugging Face models and TensorRT engines through new Persistent Volume Claims (PVCs) and a dedicated model downloader script.
• Dockerfile and Dependency Management: Updated Dockerfiles to include necessary Python packages like PyTorch, timm, and huggingface_hub, and created a new GPU-specific Dockerfile with source builds for image processing libraries and TensorRT.
• Ray Serve Configuration Enhancements: Adjusted resource allocations, autoscaling parameters, and added volume mounts in ray-service.yaml to optimize performance and resource utilization across all deployments, including the new Virchow2 model.

Changelog

• builders/heatmap_builder.py
  • Added a call to mask_builder.flush() to ensure data is written before saving.
• docker/Dockerfile.cpu
  • Updated pip install commands to include torch, torchvision, timm, and huggingface-hub.
• docker/Dockerfile.gpu
  • Added a new Dockerfile for GPU environments, including building libvips and openslide from source, and installing onnxruntime-gpu, tensorrt, torch (CUDA-enabled), torchvision (CUDA-enabled), timm, and huggingface-hub.
• misc/tile_heatmap_builder.py
  • Introduced a flush method to explicitly flush image and count data.
• misc/virchow2_downloader/download_virchow2.py
  • Added a Python script to download the paige-ai/Virchow2 model from Hugging Face Hub and verify its loading.
• misc/virchow2_downloader/virchow2_downloader_job.yaml
  • Added a Kubernetes Job definition to run the download_virchow2.py script, configuring security contexts, resources, environment variables, and volume mounts for caching.
• models/binary_classifier.py
  • Imported asyncio and os.
  • Removed mean and std from Config as normalization is now handled by the model.
  • Modified reconfigure to be synchronous and include TensorRT execution provider options, cache path creation, graph optimizations, and a model warmup step.
  • Updated predict to handle uint8 input directly and flatten outputs.
  • Adjusted root method to use asyncio.to_thread for decompression and transpose for image reshaping.
• models/semantic_segmentation.py
  • Imported os.
  • Modified reconfigure to be synchronous and include TensorRT execution provider options, cache path creation, graph optimizations, and a model warmup step.
  • Updated providers list to prioritize TensorrtExecutionProvider.
• models/virchow2.py
  • Added a new Ray Serve deployment for the Virchow2 model, including __init__, reconfigure (for model loading, transformations, and warmup), predict (for inference and embedding extraction), and root (for handling requests).
• providers/model_provider.py
  • Added a new provider function for fetching models from Hugging Face Hub, supporting both single files and snapshots, with local caching.
• pvc/huggingface-pvc.yaml
  • Added a Persistent Volume Claim for the Hugging Face cache with 15Gi storage.
• pvc/tensorrt-cache-pvc.yaml
  • Added a Persistent Volume Claim for the TensorRT cache with 20Gi storage.
• pyproject.toml
  • Added 'Jiří Štípek' as an author.
• ray-service.yaml
  • Updated BinaryClassifier deployment: changed num_cpus to 4, added num_gpus to 1, reduced memory to 4Gi, decreased max_batch_size to 16, reduced batch_wait_timeout_s to 0.01, removed mean and std config, and updated artifact_uri.
  • Updated SemanticSegmentation deployment: reduced max_queued_requests to 32, max_replicas to 2, target_ongoing_requests to 8, num_cpus to 4, added num_gpus to 1, increased max_batch_size to 8, and reduced batch_wait_timeout_s to 0.1.
  • Updated HeatmapBuilder deployment: increased max_replicas to 4, num_cpus to 8, num_threads to 8, and max_concurrent_tasks to 24.
  • Added a new virchow2 service with a Virchow2 deployment, configuring runtime environment, resource limits (8 CPUs, 1 GPU, 8Gi memory), autoscaling, and user configuration for tile size, batching, and model provider.
  • Updated ray-head image to cerit.io/rationai/model-service:2.53.0 and increased memory limits/requests.
  • Added HTTPS_PROXY environment variable to ray-head and ray-worker containers.
  • Added securityContext and lifecycle hooks to ray-worker containers.
  • Added trt-cache-volume and huggingface-cache volume mounts to ray-worker containers.
  • Added gpu-workers group with specific GPU resource requests (nvidia.com/mig-2g.20gb: 1).
  • Added trt-cache-volume and huggingface-cache PVCs to the volumes section.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for the Virchow2 foundation model within the Ray Serve infrastructure, which is a significant addition. The changes include a new Ray Serve deployment for Virchow2, updates to Dockerfiles for GPU support with necessary dependencies, and Kubernetes configurations for model downloading and caching. The refactoring of existing models to leverage GPU and TensorRT is a great performance enhancement. However, the pull request introduces critical security vulnerabilities by including hardcoded secrets in configuration files. These must be addressed by using a secure secret management solution like Kubernetes Secrets. Additionally, there are minor areas for improvement regarding Docker image consistency and file permissions.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 6 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 6 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.