The Factory | VIRDX

Reconciles vxData against a set of rules, and runs the inference that closes the gap.

The VirDx factory provides implementations of efficient, vxData-native inference for our various data-generating processes.

Overview

The factory serves two purposes:

It provides a unified API for inference of different workloads.
It provides rules and a loop to automatically trigger the above inference of workloads based on the state of vxData.

This factory app is structured as follows:

deployments/              # scripts to trigger image builds
src/factory/
  services/               # implementations of inference workloads
    histo_preprocessing/  # histo.preprocessing inference 
    viseg/                # viseg-anatomy inference 
  inference.py            # factory purpose 1: inference
  sync.py                 # factory purpose 2: triggering inference based on data
  registry.py             # definition of inference workloads and data trigger rules
  api.py                  # exposes the FastAPI entrypoints
  core.py                 # abstract base classes
  argo.py                 # util for talking to the argo workflows API
  vxdata.py               # util for talking to the vxData API

Adding an Inference Workload

An inference workload, as defined in services/, needs to implement factory.core.Service: a protocol that enforces the definition + implementation of:

An Input pydantic model: the atomic unit of input to its logic.
An Output pydantic model: the unit of output from its logic.
A Params pydantic model: allowing for configuration of the inference logic, such as threshold values in processing steps, or model IDs to use.
A Config pydantic model: extending Params, but adding parameters that may impact the diff/data-selection process (see below).
A dispatch function: effectively taking in a list of Input and the service’s config as input. All existing inference workloads dispatch their job to Argo Workflows. The implementation here should know best how to deal with large numbers of input cases. Different parallelization strategies may be suited for different workloads.
A diff function: for the sync of data, the inference workload needs to define what data it expects as input and what data is recognized as output of its logic. This operates on dataframes provided by vxData - it should make use of the vxData schemas here. The goal of this function is to determine what cases to run its own inference on in order to consider the state of vxData to be “fully processed”. This most likely is dependent on configurable parameters!

Add your service module to registry.py to be served as an inference service.

Take a look at existing implementations for examples. Both existing implementations submit an Argo Workflow using the Dockerfile.inference image, which has a heavy histo+viseg environment installed. In this image, the run functions are called, which perform the entire inference workflow: loading data from vxData, running inference, re-integrating results. Nothing else to do.

Note that this structure is very likely to change to some extent in the future. There are a few things that still bother me about the abstraction that I want to get rid of. There’s at least one layer of abstraction too many.

Adding a Data Sync Rule

In registry.py, extend the list of RULES with a configured inference service instantiation. See the file for examples.

Available Inference Endpoints

The factory exposes a REST API for anyone to hit with requests to process certain resources from vxData with certain processing pipelines.

We currently implement two inference modules:

Viseg Anatomy Segmentation
Histo Preprocessing

We’re planning on adding histo domain representation inference and diffsim as soon as they have inference interfaces.

Viseg Anatomy Segmentation

Given any number of Volume resources, this job runs inference of a specific Viseg model to produce a multi-class prostate anatomy segmentation based on provided MRI.

For available configuration options, see its Params object.

To request inference in Python, you can use an arbitrary HTTP client library to hit the REST API:

identifiers = ["volume/EE12345/t2_tse_tra"]  # an example volume, can be arbitrarily many
config = {
    "model_id": "DEFAULT",  # TODO: allow for loading of arbitrary ClearML model
    "compute_uncertainty_score": True,
}  # every field is optional, so `{"cases": [...]}` is a valid body
request_body = {
    "config": config,
    "cases": [{"volume": resource} for resource in identifiers],
}
response = httpx.post(
    "https://inference.fra.virdx.dev/inference/viseg", json=request_body
)
response.raise_for_status()
print(response.json())  # [{"name": "factory-viseg-abc12", "url": "https://argo.../abc12"}]

Histo Preprocessing

Given any number of raw HistoScan resources, this job runs histo.preprocessing to detect tissue and write back one component scan per gross section found on the slide.

For available configuration options, see its Params object.

identifiers = ["histoscan/EE12345/01/wsi-1"]  # a raw slide, can be arbitrarily many
request_body = {
    "config": {
        "tissue_detection_mag": 1.0,
        "tissue_detection": {
            "early_dilate_ksize": 100,
        },
    },
    "cases": [{"scan": resource} for resource in identifiers],
}
response = httpx.post(
    "https://inference.fra.virdx.dev/inference/histo-preprocessing", json=request_body
)
response.raise_for_status()

Sync Rules

The factory defines a set of data generation rules to automate large parts of our routine data generation across our entire data ecosystem in vxData.

You’ll find a live overview over deployed rules and their materialization status at https://dashboard.fra.virdx.dev/factory. Rules are synced automatically: a scheduled job POSTs /sync, which submits whatever every rule finds missing and is not already running. In case manual triggering of the sync is necessary, use the frontend or the REST-API.

Every sync is a fresh comparison rather than a queue, so nothing is lost by a missed run: the next one recomputes the gap from vxData and from the jobs currently in flight. Cases already running are subtracted, so triggering a sync twice does not submit the work twice.

Triggering a sync

# one rule, everything it finds missing
httpx.post("https://inference.fra.virdx.dev/rules/viseg-anat-t2-tse-tra/sync", json={})

# a hand-picked subset, by case UID as listed in `GET /rules/{producer}`
httpx.post(
    "https://inference.fra.virdx.dev/rules/viseg-anat-t2-tse-tra/sync",
    json={"cases": ["viseg-anat-t2-tse-tra/7444446b50c9"]},
)

# every rule; `dry_run` submits nothing, and the per-rule counts land in the API log
httpx.post("https://inference.fra.virdx.dev/sync", json={"dry_run": True})

Each call returns the workflows it started, as {"name": ..., "url": ...}, so a caller can link straight to the job in Argo.

API Reference

Consider the auto-generated https://inference.fra.virdx.dev/docs for full API docs.