[serve][llm] Feature: Add W&B Model Loading Callback for LLMEngine

[pjh4993]

Description

This PR introduces a WandBModelLoadingCallback to enable users to load models directly from Weights & Biases (W&B) Artifacts within Ray LLM/Serve deployments.

This feature allows users to specify a model source using the wandb:// scheme, resolving it to a local disk path before the engine configuration is finalized.

Motivation
Currently, Ray LLM primarily supports model sources from Hugging Face or direct local/S3 paths. Integrating W&B Artifacts is a common requirement for teams using W&B for model versioning and tracking. This change leverages the new callback feature to provide this integration point without modifying core LLM server logic.

Implementation Details

1. WandBArtifactHandler: A helper class is introduced to manage W&B API interaction. It handles downloading artifacts (run.use_artifact(...).download()) and supports optional custom configuration via wandb_base_url and wandb_api_key for private/custom W&B instances.

2. Symlinking: The handler creates a symlink from the W&B client's local cache location to the user-specified local_path. This is efficient as it avoids redundant copying if the artifact is already cached.

3. Configuration Model: WandBDownloaderConfig uses Pydantic validation to enforce that the necessary paths parameter is provided in the callback_kwargs. paths is expected to be a list of (wandb_uri, local_path) tuples.

4. Callback Logic (on_before_node_init): The callback initializes with its configuration validated by WandBDownloaderConfig. It iterates over the provided paths list. For each entry, it uses the WandBArtifactHandler to download the artifact and ensure it is available at the specified local_path.

Related issues

This feature was discussed in the Ray LLM Slack channel, where the callback approach was identified as the preferred path forward for integrating custom model sources like W&B Artifacts.

Additional information

How to Use
To use a W&B Artifact as the model source, a user must configure the LLMConfig to point the model source to the local path where the callback will download the artifact, and then pass the WandBDownloader and its configuration via callback_config.

Example:

from ray.llm._internal.common.callbacks.wandb_downloader import WandBDownloader
from ray.llm.config import LLMConfig, ModelLoadingConfig

# The target local path where the artifact will be placed.
LOCAL_MODEL_PATH = "/tmp/wandb-model-cache/my-llm"

config = LLMConfig(
    # 1. Point the LLM's model_source to the local path.
    model_loading_config=ModelLoadingConfig(
        model_source=LOCAL_MODEL_PATH
    ),
    # 2. Configure the callback to download the artifact to that local path.
    callback_config={
        "callback_class": WandBDownloader,
        "callback_kwargs": {
            "paths": [
                # (WandB URI, Local Target Path)
                ("wandb://my-org/my-project/llm-model:v1", LOCAL_MODEL_PATH),
            ],
            # Optional: for private/custom WandB instances
            "wandb_base_url": os.environ.get("WANDB_BASE_URL"),
            "wandb_api_key": os.environ.get("WANDB_API_KEY"),
        }
    }
)
# Deploy 'config' with Ray Serve

API Changes:

• No breaking changes

Testing:

• No additional test included