# cytosine-code/astrai architecture

> Architecture diagram of the GitHub repository cytosine-code/astrai, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-10-10). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/cytosine-code/astrai
- Repository: https://github.com/cytosine-code/astrai

## Overview

AstrAI covers model construction, data preparation, training, and serving/inference. A typical path is preprocessing source records into stored datasets, loading them for training, then serving a trained model through compatible chat APIs. Training-time configuration and parallel execution are shown as internal layers. The graph uses README capabilities for unsampled evaluation and advanced model features; arrows are limited to relationships visible in excerpts or the source index.

Main entry points: [`astrai/trainer/train_context.py`](https://github.com/cytosine-code/astrai/blob/main/astrai/trainer/train_context.py) (Training context), [`astrai/model/transformer.py`](https://github.com/cytosine-code/astrai/blob/main/astrai/model/transformer.py) (Transformer models), [`astrai/trainer/strategy.py`](https://github.com/cytosine-code/astrai/blob/main/astrai/trainer/strategy.py) (Training strategies), [`astrai/inference/engine.py`](https://github.com/cytosine-code/astrai/blob/main/astrai/inference/engine.py) (Inference engine)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data preparation"]
  node_preprocess_config["Preprocess config"]
  node_preprocess["Preprocessing pipeline<br/>[pipeline.py]"]
  node_tokenizer["Tokenizer and templates<br/>[tokenizer.py]"]
  node_dataset["Streaming dataset<br/>[dataset.py]"]
  node_storage[("Dataset storage<br/>[storage.py]")]
end

subgraph group_model["Models and kernels"]
  node_model_config["Model config<br/>[model_config.py]"]
  node_model["Transformer models<br/>[transformer.py]"]
  node_components["Model components<br/>[linear.py]"]
  node_lora["LoRA adapters<br/>[lora.py]"]
  node_kernels["Attention backends<br/>[attention.py]"]
  node_kernel_loader["Kernel loader<br/>[loader.py]"]
end

subgraph group_training["Training"]
  node_train_context["Training context<br/>[train_context.py]"]
  node_trainer["Training strategies<br/>[strategy.py]"]
  node_parallel["DDP/FSDP executor<br/>[executor.py]"]
  node_rollout["RL rollouts<br/>[rollout.py]"]
  node_checkpoint[("Checkpoint serialization<br/>[checkpoint.py]")]
end

subgraph group_inference["Inference and serving"]
  node_inference_engine["Inference engine<br/>[engine.py]"]
  node_scheduler["Continuous batching<br/>[scheduler.py]"]
  node_cache[("Paged KV cache<br/>[pool.py]")]
  node_runtime["Decode runtime<br/>[executor.py]"]
  node_api["HTTP API<br/>[app.py]"]
  node_protocol["Protocol adapters<br/>[protocol.py]"]
  node_evaluation["Evaluation tools"]
end

subgraph group_shared["Shared infrastructure"]
  node_config_base["Config serialization<br/>[base.py]"]
  node_factory["Component registry<br/>[factory.py]"]
end

node_user(("Data/model user"))
node_client(("API client"))

node_user -->|"configures"| node_preprocess
node_preprocess -->|"reads config"| node_preprocess_config
node_preprocess -->|"builds with"| node_tokenizer
node_preprocess -->|"writes shards"| node_storage
node_dataset -->|"uses config"| node_preprocess_config
node_dataset -->|"loads data"| node_storage
node_dataset -->|"uses tokenizer"| node_tokenizer
node_user -->|"configures"| node_train_context
node_train_context -->|"reads config"| node_model_config
node_train_context -.->|"creates model"| node_model
node_train_context -->|"creates executor"| node_parallel
node_train_context -->|"creates strategy"| node_trainer
node_train_context -->|"creates generator"| node_rollout
node_trainer -->|"runs distributed"| node_parallel
node_trainer -->|"uses rollouts"| node_rollout
node_trainer -->|"uses adapters"| node_lora
node_train_context -.->|"loads state"| node_checkpoint
node_model -->|"uses config"| node_model_config
node_model -->|"uses components"| node_components
node_model -.->|"uses cache"| node_cache
node_kernels -->|"loads operations"| node_kernel_loader
node_user -->|"sends requests"| node_api
node_api -->|"creates engine"| node_inference_engine
node_api -->|"dispatches requests"| node_protocol
node_protocol -.->|"handles generation"| node_inference_engine
node_inference_engine -->|"creates scheduler"| node_scheduler
node_scheduler -->|"executes decode"| node_runtime
node_scheduler -->|"manages pages"| node_cache
node_scheduler -.->|"runs model"| node_model
node_inference_engine -.->|"uses tokenizer"| node_tokenizer
node_client -->|"queries"| node_api
node_evaluation -.->|"evaluates generation"| node_inference_engine
node_config_base -->|"provides base"| node_preprocess_config
node_model_config -->|"inherits"| node_config_base
node_factory -.->|"registers components"| node_dataset
node_factory -.->|"registers strategies"| node_trainer

click node_preprocess_config "https://github.com/cytosine-code/astrai/blob/main/astrai/config/preprocess_config.py"
click node_preprocess "https://github.com/cytosine-code/astrai/blob/main/astrai/preprocessing/pipeline.py"
click node_tokenizer "https://github.com/cytosine-code/astrai/blob/main/astrai/tokenize/tokenizer.py"
click node_dataset "https://github.com/cytosine-code/astrai/blob/main/astrai/dataset/dataset.py"
click node_storage "https://github.com/cytosine-code/astrai/blob/main/astrai/dataset/storage.py"
click node_model_config "https://github.com/cytosine-code/astrai/blob/main/astrai/config/model_config.py"
click node_model "https://github.com/cytosine-code/astrai/blob/main/astrai/model/transformer.py"
click node_components "https://github.com/cytosine-code/astrai/blob/main/astrai/model/components/linear.py"
click node_lora "https://github.com/cytosine-code/astrai/blob/main/astrai/model/components/lora.py"
click node_kernels "https://github.com/cytosine-code/astrai/blob/main/astrai/extension/backend/attention.py"
click node_kernel_loader "https://github.com/cytosine-code/astrai/blob/main/astrai/extension/loader.py"
click node_train_context "https://github.com/cytosine-code/astrai/blob/main/astrai/trainer/train_context.py"
click node_trainer "https://github.com/cytosine-code/astrai/blob/main/astrai/trainer/strategy.py"
click node_parallel "https://github.com/cytosine-code/astrai/blob/main/astrai/parallel/executor.py"
click node_rollout "https://github.com/cytosine-code/astrai/blob/main/astrai/trainer/rollout.py"
click node_checkpoint "https://github.com/cytosine-code/astrai/blob/main/astrai/serialization/checkpoint.py"
click node_inference_engine "https://github.com/cytosine-code/astrai/blob/main/astrai/inference/engine.py"
click node_scheduler "https://github.com/cytosine-code/astrai/blob/main/astrai/inference/scheduler.py"
click node_cache "https://github.com/cytosine-code/astrai/blob/main/astrai/inference/cache/pool.py"
click node_runtime "https://github.com/cytosine-code/astrai/blob/main/astrai/inference/runtime/executor.py"
click node_api "https://github.com/cytosine-code/astrai/blob/main/astrai/inference/network/app.py"
click node_protocol "https://github.com/cytosine-code/astrai/blob/main/astrai/inference/network/protocol.py"
click node_evaluation "https://github.com/cytosine-code/astrai/tree/main/scripts/eval"
click node_config_base "https://github.com/cytosine-code/astrai/blob/main/astrai/config/base.py"
click node_factory "https://github.com/cytosine-code/astrai/blob/main/astrai/factory.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_preprocess_config,node_preprocess,node_tokenizer,node_dataset,node_storage,node_user,node_client toneBlue
class node_model_config,node_model,node_components,node_lora,node_kernels,node_kernel_loader toneAmber
class node_train_context,node_trainer,node_parallel,node_rollout,node_checkpoint toneMint
class node_inference_engine,node_scheduler,node_cache,node_runtime,node_api,node_protocol,node_evaluation toneRose
class node_config_base,node_factory toneIndigo
```

## Components

### Data preparation

- [Preprocess config](https://github.com/cytosine-code/astrai/blob/main/astrai/config/preprocess_config.py): `astrai/config/preprocess_config.py`
- [Preprocessing pipeline](https://github.com/cytosine-code/astrai/blob/main/astrai/preprocessing/pipeline.py): `astrai/preprocessing/pipeline.py`
- [Tokenizer and templates](https://github.com/cytosine-code/astrai/blob/main/astrai/tokenize/tokenizer.py): `astrai/tokenize/tokenizer.py`
- [Streaming dataset](https://github.com/cytosine-code/astrai/blob/main/astrai/dataset/dataset.py): `astrai/dataset/dataset.py`
- [Dataset storage](https://github.com/cytosine-code/astrai/blob/main/astrai/dataset/storage.py): `astrai/dataset/storage.py`

### Models and kernels

- [Model config](https://github.com/cytosine-code/astrai/blob/main/astrai/config/model_config.py): `astrai/config/model_config.py`
- [Transformer models](https://github.com/cytosine-code/astrai/blob/main/astrai/model/transformer.py): `astrai/model/transformer.py`
- [Model components](https://github.com/cytosine-code/astrai/blob/main/astrai/model/components/linear.py): `astrai/model/components/linear.py`
- [LoRA adapters](https://github.com/cytosine-code/astrai/blob/main/astrai/model/components/lora.py): `astrai/model/components/lora.py`
- [Attention backends](https://github.com/cytosine-code/astrai/blob/main/astrai/extension/backend/attention.py): `astrai/extension/backend/attention.py`
- [Kernel loader](https://github.com/cytosine-code/astrai/blob/main/astrai/extension/loader.py): `astrai/extension/loader.py`

### Training

- [Training context](https://github.com/cytosine-code/astrai/blob/main/astrai/trainer/train_context.py): `astrai/trainer/train_context.py`
- [Training strategies](https://github.com/cytosine-code/astrai/blob/main/astrai/trainer/strategy.py): `astrai/trainer/strategy.py`
- [DDP/FSDP executor](https://github.com/cytosine-code/astrai/blob/main/astrai/parallel/executor.py): `astrai/parallel/executor.py`
- [RL rollouts](https://github.com/cytosine-code/astrai/blob/main/astrai/trainer/rollout.py): `astrai/trainer/rollout.py`
- [Checkpoint serialization](https://github.com/cytosine-code/astrai/blob/main/astrai/serialization/checkpoint.py): `astrai/serialization/checkpoint.py`

### Inference and serving

- [Inference engine](https://github.com/cytosine-code/astrai/blob/main/astrai/inference/engine.py): `astrai/inference/engine.py`
- [Continuous batching](https://github.com/cytosine-code/astrai/blob/main/astrai/inference/scheduler.py): `astrai/inference/scheduler.py`
- [Paged KV cache](https://github.com/cytosine-code/astrai/blob/main/astrai/inference/cache/pool.py): `astrai/inference/cache/pool.py`
- [Decode runtime](https://github.com/cytosine-code/astrai/blob/main/astrai/inference/runtime/executor.py): `astrai/inference/runtime/executor.py`
- [HTTP API](https://github.com/cytosine-code/astrai/blob/main/astrai/inference/network/app.py): `astrai/inference/network/app.py`
- [Protocol adapters](https://github.com/cytosine-code/astrai/blob/main/astrai/inference/network/protocol.py): `astrai/inference/network/protocol.py`
- [Evaluation tools](https://github.com/cytosine-code/astrai/tree/main/scripts/eval): `scripts/eval`

### Shared infrastructure

- [Config serialization](https://github.com/cytosine-code/astrai/blob/main/astrai/config/base.py): `astrai/config/base.py`
- [Component registry](https://github.com/cytosine-code/astrai/blob/main/astrai/factory.py): `astrai/factory.py`

### Other components

- **Data/model user**
- **API client**

## Connections

- Data/model user → Preprocessing pipeline: configures
- Preprocessing pipeline → Preprocess config: reads config (evidence: `astrai/preprocessing/pipeline.py`)
- Preprocessing pipeline → Tokenizer and templates: builds with (evidence: `astrai/preprocessing/pipeline.py`)
- Preprocessing pipeline → Dataset storage: writes shards (evidence: `astrai/preprocessing/pipeline.py`)
- Streaming dataset → Preprocess config: uses config (evidence: `astrai/dataset/dataset.py`)
- Streaming dataset → Dataset storage: loads data (evidence: `astrai/dataset/dataset.py`)
- Streaming dataset → Tokenizer and templates: uses tokenizer (evidence: `astrai/dataset/dataset.py`)
- Data/model user → Training context: configures
- Training context → Model config: reads config (evidence: `astrai/trainer/train_context.py`)
- Training context → Transformer models: creates model
- Training context → DDP/FSDP executor: creates executor (evidence: `astrai/trainer/train_context.py`)
- Training context → Training strategies: creates strategy (evidence: `astrai/trainer/train_context.py`)
- Training context → RL rollouts: creates generator (evidence: `astrai/trainer/train_context.py`)
- Training strategies → DDP/FSDP executor: runs distributed (evidence: `astrai/trainer/strategy.py`)
- Training strategies → RL rollouts: uses rollouts (evidence: `astrai/trainer/strategy.py`)
- Training strategies → LoRA adapters: uses adapters (evidence: `astrai/trainer/train_context.py`)
- Training context → Checkpoint serialization: loads state
- Transformer models → Model config: uses config (evidence: `astrai/model/transformer.py`)
- Transformer models → Model components: uses components (evidence: `astrai/model/transformer.py`)
- Transformer models → Paged KV cache: uses cache
- Attention backends → Kernel loader: loads operations (evidence: `astrai/extension/backend/attention.py`)
- Data/model user → HTTP API: sends requests
- HTTP API → Inference engine: creates engine (evidence: `astrai/inference/network/app.py`)
- HTTP API → Protocol adapters: dispatches requests (evidence: `astrai/inference/network/app.py`)
- Protocol adapters → Inference engine: handles generation
- Inference engine → Continuous batching: creates scheduler (evidence: `astrai/inference/engine.py`)
- Continuous batching → Decode runtime: executes decode (evidence: `astrai/inference/scheduler.py`)
- Continuous batching → Paged KV cache: manages pages (evidence: `astrai/inference/scheduler.py`)
- Continuous batching → Transformer models: runs model
- Inference engine → Tokenizer and templates: uses tokenizer
- API client → HTTP API: queries
- Evaluation tools → Inference engine: evaluates generation
- Config serialization → Preprocess config: provides base (evidence: `astrai/config/preprocess_config.py`)
- Model config → Config serialization: inherits (evidence: `astrai/config/model_config.py`)
- Component registry → Streaming dataset: registers components (evidence: `astrai/dataset/dataset.py`)
- Component registry → Training strategies: registers strategies (evidence: `astrai/trainer/strategy.py`)