gunz_ml.schemas package

Configuration factory for initializing the correct trial configuration wrapper.

gunz_ml.schemas.init_trial_cfg(trial: optuna.trial.Trial | None, exp_params: Dict[str, Any]) → BaseConfig[source]

Initializes the appropriate trial configuration wrapper.

This function acts as a factory, inspecting the trial object to determine whether to configure for a single run or an Optuna trial.

Parameters:

trial (optuna.trial.Trial, optional) –
- If an optuna.trial.Trial object is provided, the function configures for an optimization trial.
- If None is provided, the function configures for a single, non-optimized run.
exp_params (dict) – A dictionary of parameters. - For an optimization, this should be the hyperparameter search space. - For a single run, this should be the fixed parameter values.

Returns:

An initialized instance of either RunConfig or OptunaConfig.

Return type:

BaseConfig

Submodules

gunz_ml.schemas.analysis module

Configuration schema for the Analysis mode.

class gunz_ml.schemas.analysis.AnalysisConfig(enable_internals: bool = True, enable_dynamics: bool = False, enable_errors: bool = True, batch_size: int = 4, output_dir: str = 'analysis_outputs', save_plots: bool = True, save_json: bool = True, ckpt_path: str = omegaconf.MISSING)[source]

Bases: object

Configuration for Deep Learning Analysis.

Parameters:

enable_internals (bool) – Whether to analyze model internals (params, stats).
enable_dynamics (bool) – Whether to analyze training dynamics (gradients).
enable_errors (bool) – Whether to analyze prediction errors (confusion matrix).
batch_size (int) – Batch size for dynamic analysis.
output_dir (str) – Directory to save analysis results.
save_plots (bool) – Whether to save visualization plots.
save_json (bool) – Whether to save raw results as JSON.
ckpt_path (str) – Path to the model checkpoint to load.

batch_size: int = 4

enable_dynamics: bool = False

enable_errors: bool = True

enable_internals: bool = True

output_dir: str = 'analysis_outputs'

save_json: bool = True

save_plots: bool = True

gunz_ml.schemas.args module

class gunz_ml.schemas.args.ArgsConfig(dry_run: bool = False, overwrite: bool = False, mode: str = 'optimize', num_trials: int = omegaconf.MISSING, study_name: str = '${study_prefix}_${hydra:runtime.choices.exp}', max_epochs: int = omegaconf.MISSING, precision: int = 32, accelerator: str = 'gpu')[source]

Bases: object

Global arguments configuration.

Defines high-level runtime parameters for experiment execution.

Parameters:

dry_run (bool, optional) – If True, simulates execution without making changes. Defaults to False.
overwrite (bool, optional) – If True, overwrites existing data. Defaults to False.
mode (str, optional) – The operation mode (e.g., ‘optimize’). Defaults to ‘optimize’.
num_trials (int) – The number of trials to run.
study_name (str) – The name of the study, interpolated from prefix and hydra choices.
max_epochs (int) – Maximum number of training epochs.
precision (int, optional) – Floating point precision (e.g., 16, 32). Defaults to 32.
accelerator (str, optional) – Hardware accelerator to use (e.g., ‘gpu’, ‘cpu’). Defaults to ‘gpu’.

accelerator: str = 'gpu'

dry_run: bool = False

mode: str = 'optimize'

overwrite: bool = False

precision: int = 32

study_name: str = '${study_prefix}_${hydra:runtime.choices.exp}'

gunz_ml.schemas.db module

class gunz_ml.schemas.db.DatabaseConfig(driver: str, host: str, port: int | None = None, db_name: str | None = None, user: str | None = None, password: str | None = None, timeout: int = 30)[source]

Bases: object

Configuration schema for database connections.

db_name: str | None = None

driver: str

host: str

password: str | None = None

port: int | None = None

timeout: int = 30

user: str | None = None

gunz_ml.schemas.experiment module

class gunz_ml.schemas.experiment.BaseExpConfig(study_prefix: str = omegaconf.MISSING, optuna_type: ~gunz_ml.consts.OptunaType = OptunaType.INT, args: ~gunz_ml.schemas.args.ArgsConfig = <factory>, node: dict = omegaconf.MISSING, db: dict | None = None, ds: dict = omegaconf.MISSING, pl: dict | None = None, optuna: dict | None = None, default: dict | None = None)[source]

Bases: object

Base configuration for experiments.

This dataclass defines the structure of an experiment configuration, integrating various sub-components like arguments, node details, and library-specific settings.

Parameters:

study_prefix (str) – The prefix for the study name.
optuna_type (OptunaType) – The type of Optuna optimization (e.g., INT, FLOAT). Defaults to OptunaType.INT.
args (ArgsConfig) – Runtime arguments for the experiment. Defaults to a factory-created ArgsConfig.
node (dict) – Configuration details for the compute node.
db (dict, optional) – Database configuration settings.
ds (dict) – Dataset configuration settings.
pl (dict, optional) – PyTorch Lightning specific configuration.
optuna (dict, optional) – Optuna study configuration details.
default (dict, optional) – Default values or overrides.

args: ArgsConfig

db: dict | None = None

default: dict | None = None

optuna: dict | None = None

optuna_type: OptunaType = 'int'

pl: dict | None = None

gunz_ml.schemas.mlflow_cfg module

class gunz_ml.schemas.mlflow_cfg.MLFlowConfig(exp: mlflow.entities.Run, params_cfg: dict)[source]

Bases: BaseConfig

Configuration wrapper for MLFlow runs.

This class handles parameter suggestion based on an MLFlow run’s recorded parameters.

Parameters:

exp (mlflow.entities.Run) – The MLFlow run object containing experiment data.
params_cfg (dict) – Configuration specifying parameter types and options.

suggest_param(key: str)[source]

Suggest a value for a parameter based on the MLFlow run data.

Parameters:

key (str) – The key identifying the parameter to suggest.

Returns:

The suggested parameter value. The type depends on the configuration (e.g., int, float, str).

Return type:

Any

Raises:

KeyError – If the key is not found in the parameter configuration.
RuntimeError – If an unknown parameter type is encountered.

gunz_ml.schemas.optuna_cfg module

Provides configuration handling specifically for Optuna trials.

This module defines the OptunaConfig class, which wraps an Optuna Trial object to suggest hyperparameters based on a provided configuration dictionary.

class gunz_ml.schemas.optuna_cfg.ObjectiveDirection(value)[source]

Bases: StrEnum

Enumeration for optimization directions in Optuna.

MAXIMIZE

Direction for maximizing the objective.

Type:: str

MINIMIZE

Direction for minimizing the objective.

Type:: str

MAXIMIZE = 'maximize'

MINIMIZE = 'minimize'

class gunz_ml.schemas.optuna_cfg.OptunaConfig(trial: optuna.Trial, params_cfg: dict)[source]

Bases: BaseConfig

A wrapper for an Optuna Trial object to streamline hyperparameter suggestion.

This class implements the BaseConfig interface to provide a consistent way of accessing hyperparameters during an Optuna optimization study.

trial

The active Optuna trial object.

Type:: optuna.Trial

params_cfg

The configuration dictionary defining the search space.

Type:: dict

hparams

A dictionary that logs the actual suggested hyperparameters for the trial.

Type:: dict

suggest_param(name: str, def_val: Any = None, ena_def_val: bool = False) → Any

Suggests a value for a hyperparameter based on the configuration.

This method reads the configuration for the given parameter name. If the parameter is not found and ena_def_val is True, it returns the provided def_val. Otherwise, it raises a KeyError.

Parameters:

name (str) – The name of the parameter in the configuration.
def_val (any, optional) – The default value to return if the parameter is not found. Defaults to None.
ena_def_val (bool, optional) – If True, enables the default value mechanism. Defaults to False.

Returns:

The suggested, fixed, or default value for the parameter.

Return type:

any

gunz_ml.schemas.optuna_cfg.optimize_preproc_pipeline(opt_cfg: OptunaConfig, pipeline_configs: Dict) → Dict

Optimizes hyperparameters within a preprocessing pipeline configuration.

This function searches for a special _OPTIMIZE_ tag in the pipeline configuration and replaces it with a value suggested by Optuna.

Parameters:

opt_cfg (OptunaConfig) – The OptunaConfig instance for suggesting parameters.
pipeline_configs (dict) – The preprocessing pipeline configuration dictionary.

Returns:

A new dictionary with the _OPTIMIZE_ tags replaced by suggested values.

Return type:

dict

Dispatches to the appropriate Optuna suggest_* method.

Parameters:

trial (optuna.Trial) – The Optuna trial object to suggest a value from.
type_tag (OptunaType) – An enum value that determines which suggest_* method to call.
name (str) – The name of the hyperparameter to suggest.
low (int or float, optional) – The lower bound for numerical suggestions. Defaults to None.
high (int or float, optional) – The upper bound for numerical suggestions. Defaults to None.
step (int or float, optional) – The step size for numerical suggestions. Defaults to None.
vals (sequence, optional) – A sequence of choices for categorical suggestions. Defaults to None.

Returns:

The value suggested by the Optuna trial.

Return type:

any

gunz_ml.schemas.optuna_cls module

class gunz_ml.schemas.optuna_cls.OptunaCatParamDC(*args: Any, **kwargs: Any)[source]

Bases: BaseModel

Configuration for a categorical parameter in Optuna.

Parameters:

type (int) – Type identifier.
vals (list) – List of possible values.

type: int

vals: List

class gunz_ml.schemas.optuna_cls.OptunaConfigDC(*args: Any, **kwargs: Any)[source]

Bases: BaseModel

Configuration for an Optuna study.

Parameters:

objectives (dict) – Definition of objectives.
optimize (dict) – Optimization settings.
parameters (dict) – Dictionary of parameter configurations.

objectives: Dict

optimize: Dict

parameters: Dict[str, None | int | float | str | OptunaIntParamDC | OptunaFloatParamDC | OptunaLogParamDC | OptunaCatParamDC]

class gunz_ml.schemas.optuna_cls.OptunaFloatParamDC(*args: Any, **kwargs: Any)[source]

Bases: BaseModel

Configuration for a floating-point parameter in Optuna.

Parameters:

type (int) – Type identifier.
low (float) – Lower bound of the range.
high (float) – Upper bound of the range.

high: float

low: float

type: int

class gunz_ml.schemas.optuna_cls.OptunaIntParamDC(*args: Any, **kwargs: Any)[source]

Bases: BaseModel

Configuration for an integer parameter in Optuna.

Parameters:

type (int) – Type identifier.
low (int) – Lower bound of the range.
high (int) – Upper bound of the range.
step (int, optional) – Step size. Defaults to 1.

high: int

low: int

step: int = 1

type: int

class gunz_ml.schemas.optuna_cls.OptunaLogParamDC(*args: Any, **kwargs: Any)[source]

Bases: BaseModel

Configuration for a logarithmic floating-point parameter in Optuna.

Parameters:

type (int) – Type identifier.
low (float) – Lower bound of the range.
high (float) – Upper bound of the range.

high: float

low: float

type: int

gunz_ml.schemas.run_cfg module

Provides configuration handling for a single, non-optimized run.

class gunz_ml.schemas.run_cfg.RunConfig(params_cfg: Dict[str, Any])[source]

Bases: BaseConfig

A wrapper for a fixed dictionary of hyperparameters for a single run.

This class implements the BaseConfig interface for a non-optimization scenario, where all parameter values are predefined.

get_dict_param(namespace: str, dest_dict: Dict[str, Any] | None = None, def_dict: Dict[str, Any] | None = None, ena_def_dict: bool = False) → Dict[str, Any] | None

Suggests a dictionary of parameters based on a common namespace.

This method intelligently combines parameters defined in the main configuration with a provided default dictionary.

Parameters:

namespace (str) – The prefix used to identify related parameters in the main config.
dest_dict (dict, optional) – If provided, the suggested parameters will be added to this dictionary in-place. If None, a new dictionary is returned. Defaults to None.
def_dict (dict, optional) – A dictionary of default values for the parameters in this namespace. Defaults to None.
ena_def_dict (bool, optional) – If True, enables the use of the def_dict. Defaults to False.

Returns:

A dictionary of parameters, or None if dest_dict was provided.

Return type:

dict, optional

get_param(name: str, def_val: Any = None, ena_def_val: bool = False) → Any

Retrieves a parameter value from the configuration dictionary.

This method performs a simple dictionary lookup. If the parameter is not found and ena_def_val is True, it returns the default value.

Parameters:

name (str) – The name (or key) of the parameter to retrieve.
def_val (t.Any, optional) – The default value to return if the parameter is not found. Defaults to None.
ena_def_val (bool, optional) – If True, enables the default value mechanism. Defaults to False.

Returns:

The fixed or default value for the parameter.

Return type:

t.Any

suggest_param(name: str, def_val: Any = None, ena_def_val: bool = False) → Any[source]

Retrieves a parameter value from the configuration dictionary.

This method performs a simple dictionary lookup. If the parameter is not found and ena_def_val is True, it returns the default value.

Parameters:

name (str) – The name (or key) of the parameter to retrieve.
def_val (t.Any, optional) – The default value to return if the parameter is not found. Defaults to None.
ena_def_val (bool, optional) – If True, enables the default value mechanism. Defaults to False.

Returns:

The fixed or default value for the parameter.

Return type:

t.Any