AI runtime configuration scheme

[nikos-livathinos]

Objective

• We want to introduce input parameters that control the AI runtime configuration which is used to run the Docling AI models.
• The AI runtime parameters should include the "device" and "number of threads".
• The user should be able to set these parameters either as envvars or as API/CLI parameters.

Introduce the logic to do the resolution of input parameters and decide the AI runtime configuration based on:

• The values of envvars (if set).
• The values of API/CLI parameters (if set).
• The default values of each parameter.
• The availability of AI execution devices on the system.
• The input parameters required by the models.

Proposal

A. Introduce input parameters in docling package:

• API: Extend the PipelineOptions to introduce the fields: device, num_threads.
  • The device parameter can be an Enum with values: [AUTO, CUDA, CPU, MPS]. The default value is AUTO.
  • The num_threads parameter is an integer. The default value is 4.
• Envvars: Standardize names of envvars that set the device and number of threads.
  • For example [DOCLING_DEVICE, DOCLING_NUM_THREADS] or [DOCLING_DEVICE, OMP_NUM_THREADS].
• CLI: Introduce corresponding CLI parameters device and num_threads.

B. Introduce configuration resolution logic as a utility function in docling-ibm-models package:

1. Define device precedence: [CUDA, MPS, CPU].
2. Use the user-provided API/CLI parameters, if they have been explicitly set by the user.
3. Use the envvars, if they are set.
4. Use the default values.
   5a. If device == AUTO, use the first available device according to the precedence.
   5b. If device is explicitly set by the user but it is not available in the system, replace it with the next available device.

C. Usage:

1. The user either relies on the default values or explicitly sets the input parameters.
   • The API parameters are set individually for each pipeline.
   • In case of the docling CLI, the parameters apply for all pipelines.
2. Each model receives the input runtime parameters as part of the pipeline_options.
3. The utility function from docling-ibm-models is used to resolve the runtime configuration.
   • This step can be done either by the "model predictors" inside the docling-ibm-models package or by the "models" inside the docling package.
   • The resolved runtime configuration can be further customized to feed the input parameters of each AI model.

[dolfim-ibm]

In terms of resolution priority, we should do (from higher priority to lower):

1. device set in the pipeline options (or directly in the models init)
2. CLI argument
3. ENV variable
4. Code defaults

[vagenas]

To streamline this config mgmt part, let's try to leverage Pydantic Settings as much as possible, as it already covers many of the pain points.

[vagenas]

FYI we can more precisely get inspired by my implementation of profiles in deepsearch-toolkit, which is essentially a utility for persisted, swappable configurations using Pydantic Settings and platformdirs.

[nikos-livathinos]

After taking into account the proposed ideas, here is an updated implementation plan to support AI runtime configuration:

Changes for docling-ibm-models package:

1. Refactor the constructor of the predictors (both LayoutPredictor and TFPredictor) to receive the parameters num_threads:int = 4, device: str = 'cpu'. It is the responsibility of the caller to provide the parameters.
2. Remove the device parameter from tableformer config.
3. Convert tableformer model into safe tensors.

Changes for docling package:

1. Introduce the Enum AcceleratorDevice that provides device options like ["cpu", "cuda", "mps", "auto", "directml", ...].
2. Introduce the class AcceleratorOptions as pydantic settings :
   • The fields are (num_threads: Optional[int] = None, device="auto")
   • Enable envvars: OMP_NUM_THREADS, DOCLING_DEVICE.
   • Add the field accelerator_options: AcceleratorOptions in PipelineOptions.
3. Refactor all docling model constructors to receive the accelerator_options as an additional parameter and refactor the StandardPdfPipeline to pass the accelerator_options during the model initialization.
4. Provide a utility function that implements the logic of device selection.
5. Each docling model can use the utility function and/or any custom code to properly initialize the underlying predictor or third party model.

[cau-git]

Excellent, let's proceed as you propose!

[dolfim-ibm]

As discussed, this is matching very nicely what we need for the first iteration.

Next, we expect to provide higher level of customization, where each model could get different accelerator values or inherit from the pipeline.

[cau-git]

This was implemented with docling v2.12, closing.