Skip to content

nexuml.training.lightning

nexuml.training.lightning

Lightning-based training backend for NexuML.

NexuLightningModule

Bases: LightningModule

Lightning module wrapping a CompiledPipeline.

LightningFeatureExtractor

Bases: Module

Partial pipeline wrapper for dataset preprocessing/export.

TrainResult dataclass

Result of a training run.

RuntimeArtifacts dataclass

Compiled runtime objects shared by training and dataset export.

NexuSession

Thin orchestration layer for building and running a NexuML session.

from_scenario classmethod

from_scenario(
    scenario: ScenarioSpec, **kwargs
) -> "NexuSession"

Create a session from a scenario definition.

Returns:

Type Description
'NexuSession'

New NexuSession instance.

from_trainer_checkpoint classmethod

from_trainer_checkpoint(
    trainer_checkpoint: str | Path,
    scenario: ScenarioSpec | None = None,
    **kwargs,
) -> "NexuSession"

Create a session from a Lightning Trainer checkpoint.

Returns:

Type Description
'NexuSession'

New NexuSession instance restored from the checkpoint.

setup

setup() -> 'NexuSession'

Build runtime objects and trainer.

Returns:

Type Description
'NexuSession'

The same NexuSession instance (for chaining).

build_runtime

build_runtime() -> RuntimeArtifacts

Create runtime artifacts lazily.

Returns:

Type Description
RuntimeArtifacts

The RuntimeArtifacts (creating them on first call).

build_trainer

build_trainer() -> L.Trainer

Create the Lightning Trainer lazily.

Returns:

Type Description
Trainer

The lightning.Trainer instance (creating it on first call).

fit

fit() -> 'NexuSession'

Run Trainer.fit() for the current session.

Returns:

Type Description
'NexuSession'

The same NexuSession instance (for chaining).

validate

validate() -> list[dict[str, float]]

Run Trainer.validate() with the session datamodule.

Returns:

Type Description
list[dict[str, float]]

List of metric dictionaries from the validation stage.

predict

predict(
    dataloaders: Any = None,
    datamodule: NexuDataModule | None = None,
    return_predictions: bool = False,
) -> Any

Run Trainer.predict() with session defaults unless explicitly overridden.

Returns:

Type Description
Any

Prediction outputs from Trainer.predict(), or None when

Any

return_predictions is False.

test

test(
    dataloaders: Any = None,
    datamodule: NexuDataModule | None = None,
) -> list[dict[str, float]]

Run Trainer.test() with session defaults unless explicitly overridden.

Returns:

Type Description
list[dict[str, float]]

List of metric dictionaries from the test stage.

run

run() -> TrainResult

Execute the standard fit → validate → post-train fit → test flow.

Returns:

Type Description
TrainResult

TrainResult containing the pipeline, module, trainer, and

TrainResult

metric results.

create_dataset_from_spec

create_dataset_from_spec(
    scenario: ScenarioSpec,
) -> NexuDataset

Create a dataset from a ScenarioSpec's data configuration.

Returns:

Type Description
NexuDataset

A NexuDataset built from the scenario's data config.

create_data_module_from_spec

create_data_module_from_spec(
    scenario: ScenarioSpec,
) -> NexuDataModule

Create a LightningDataModule from a ScenarioSpec.

Returns:

Type Description
NexuDataModule

A NexuDataModule configured for the scenario.

materialize_preprocessed_dataset

materialize_preprocessed_dataset(
    scenario: ScenarioSpec,
) -> Path

Materialize the configured preprocessing view and return its export path.

Returns:

Type Description
Path

Path to the directory containing the materialized dataset.

resolve_preprocessing_path

resolve_preprocessing_path(scenario: ScenarioSpec) -> Path

Resolve the configured preprocessing output directory.

Returns:

Type Description
Path

Resolved Path for the preprocessing output directory.

create_runtime_artifacts

create_runtime_artifacts(
    scenario: ScenarioSpec,
    apply_selective_checkpoint: bool = True,
) -> RuntimeArtifacts

Compile the pipeline and create the matching Lightning/DataModule runtime.

Returns:

Type Description
RuntimeArtifacts

RuntimeArtifacts containing the compiled pipeline, Lightning module,

RuntimeArtifacts

data module, and optional load report.

train

train(
    scenario: ScenarioSpec | None,
    accelerator: str = "auto",
    devices: int | str = "auto",
    log_dir: str | Path = ".experiments",
    enable_progress_bar: bool = True,
    trainer_checkpoint: str | Path | None = None,
    run_name: str | None = None,
) -> TrainResult

Compatibility wrapper around NexuSession.run().

Returns:

Type Description
TrainResult

TrainResult from the session run.

load_scenario_from_trainer_checkpoint

load_scenario_from_trainer_checkpoint(
    checkpoint_path: str | Path,
    fallback: ScenarioSpec | None = None,
) -> ScenarioSpec

Load serialized scenario metadata from a Lightning Trainer checkpoint.

Returns:

Type Description
ScenarioSpec

ScenarioSpec deserialized from the checkpoint, or fallback if

ScenarioSpec

the checkpoint contains no scenario metadata and fallback is provided.

Raises:

Type Description
ValueError

If the checkpoint lacks scenario metadata and no fallback is given.

create_runtime_artifacts_from_trainer_checkpoint

create_runtime_artifacts_from_trainer_checkpoint(
    checkpoint_path: str | Path,
    scenario: ScenarioSpec | None = None,
) -> RuntimeArtifacts

Rebuild runtime objects from a Lightning Trainer checkpoint.

Returns:

Type Description
RuntimeArtifacts

RuntimeArtifacts reconstructed from the checkpoint.