Entity classes for stage config

An entity class is a frozen dataclass that names and types every setting one stage needs, so a stage is handed a single object whose fields are fixed, instead of reaching into a dictionary of strings.

Last updated: 19 Oct, 2026 · MLOps

Reading config with ConfigBox made the YAML files pleasant to read, but what comes out is still loose: any key can be read, a missing one fails only when it is touched, and every path is text. An entity class is the declared shape of one stage's settings, and in the video's workflow it is the file written straight after the YAML files, before any stage code exists.

Syntax: @dataclass(frozen=True) above a class of annotated fields gives a class whose constructor takes those fields by name and whose instances cannot be changed afterwards.

The clip's entity carries the deep-learning project's archive paths; this one carries the student-performance stage's CSV paths and split settings, and both are frozen dataclasses of annotated fields.

Naming one stage's settings

The module holds one class per stage and no logic at all. Two imports carry it: the decorator and Path.

from dataclasses import dataclass
from pathlib import Path

The ingestion stage needs a folder, a source URL, three file names and the two values that govern the split. Each one is a field with a type.

@dataclass(frozen=True)
class DataIngestionConfig:
    # where this stage writes, and where the raw file comes from
    root_dir: Path
    source_url: str
    local_data_file: Path

    # the two files the split step writes, and how it splits
    train_file: Path
    test_file: Path
    test_size: float
    random_state: int

Reading that class tells you what the stage is allowed to care about. A reviewer who sees a field added here knows a stage's inputs changed, which is a far better signal than a new quoted key appearing somewhere in a method.

Why the paths are Path, not str

A path kept as a string invites string surgery: concatenation, backslashes on one operating system and forward slashes on another, and a missing separator that shows up as a file in the wrong folder. Path carries the parts, joins them correctly on any operating system, and answers questions a string cannot.

from pathlib import Path

from sp.entity.config_entity import DataIngestionConfig

# one object, built by hand here, with the values config.yaml and params.yaml hold
cfg = DataIngestionConfig(
    root_dir=Path("artifacts/data_ingestion"),
    source_url="https://raw.githubusercontent.com/krishnaik06/mlproject/main/notebook/data/stud.csv",
    local_data_file=Path("artifacts/data_ingestion/stud.csv"),
    train_file=Path("artifacts/data_ingestion/train.csv"),
    test_file=Path("artifacts/data_ingestion/test.csv"),
    test_size=0.2, random_state=42)

print(cfg.local_data_file.name, "in", cfg.local_data_file.parent)
print(cfg.root_dir / "report.json")
print(cfg.test_size, type(cfg.test_size).__name__)

stud.csv in artifacts/data_ingestion
artifacts/data_ingestion/report.json
0.2 float

What the typed object answered

• The file name and its folder came apart on request, because local_data_file is a Path: .name and .parent are answers a string would need slicing to give.
• A new path was built with a slash operator, cfg.root_dir / "report.json", which joins correctly on Windows, macOS and Linux with no separator of your own.
• test_size stayed a float, so the value handed to scikit-learn is a number, not the text "0.2" that would raise deep inside the split.

Why the class is frozen

frozen=True makes the instance read-only. A stage receives its settings and cannot quietly rewrite them halfway through a run, which keeps the run reproducible: what params.yaml said is what happened. Trying to assign raises.

# a stage must not rewrite its own settings in the middle of a run
cfg.test_size = 0.5

Traceback (most recent call last):
  File "main.py", line 2, in <module>
    cfg.test_size = 0.5
dataclasses.FrozenInstanceError: cannot assign to field 'test_size'

FrozenInstanceError names the field. Without frozen=True the assignment would succeed, the run would use 0.5, and the config file would still say 0.2: two runs with the same recorded settings and different results, which is the hardest kind of difference to chase.

One entity per stage

The other three stages get their own classes in the same file. Validation needs the status file and the expected columns; transformation needs the preprocessor's file and the target column.

@dataclass(frozen=True)
class DataValidationConfig:
    root_dir: Path
    status_file: Path
    expected_columns: dict        # the columns and dtypes from schema.yaml


@dataclass(frozen=True)
class DataTransformationConfig:
    root_dir: Path
    preprocessor_file: Path
    target_column: str            # named in params.yaml, never in the code

The trainer's class is the longest, because it carries both the files it writes and the hyperparameters of the model.

@dataclass(frozen=True)
class ModelTrainerConfig:
    root_dir: Path
    model_file: Path
    metrics_file: Path

    # the four values params.yaml owns for this stage
    target_column: str
    n_estimators: int
    max_depth: int
    random_state: int

Four stages, four classes, no inheritance and no shared base. A stage cannot read another stage's settings because it was never given them, and that is the whole design.

A dict of settings against an entity class

Where an entity class earns its keep

• Adding a setting to one stage, where the field, its type and its default home in YAML are all one small diff a reviewer can read.
• Reproducing a run, where frozen settings mean the recorded config is the config that ran.
• Reading unfamiliar code, where the entity class is the shortest honest answer to what a stage depends on.

Related

• Previous: Reading config with ConfigBox
• Next: The configuration manager
• Reference: dataclasses in the Python standard library