tracker

orchard.tracking.tracker

MLflow Tracker Implementation.

Provides MLflowTracker for experiment tracking and NoOpTracker as a silent fallback when MLflow is unavailable or tracking is disabled.

TrackerProtocol

Bases: Protocol

Protocol defining the experiment tracker interface.

Both MLflowTracker and NoOpTracker implement this protocol, enabling type-safe dependency injection without inheritance.

NoOpTracker

Silent no-op tracker used when MLflow is disabled or unavailable.

All methods mirror the MLflowTracker interface but perform no operations, ensuring zero overhead when tracking is not active.

start_run(cfg, run_name, tracking_uri)

No-op: skip MLflow run initialization.

Source code in orchard/tracking/tracker.py

def start_run(self, cfg: Any, run_name: str, tracking_uri: str) -> None:  # noqa: ARG002
    """No-op: skip MLflow run initialization."""

log_epoch(epoch, train_loss, val_metrics, lr)

No-op: skip per-epoch metric logging.

Source code in orchard/tracking/tracker.py

def log_epoch(  # noqa: ARG002
    self, epoch: int, train_loss: float, val_metrics: Mapping[str, float], lr: float
) -> None:
    """No-op: skip per-epoch metric logging."""

log_test_metrics(test_acc, macro_f1)

No-op: skip test metric logging.

Source code in orchard/tracking/tracker.py

def log_test_metrics(self, test_acc: float, macro_f1: float) -> None:  # noqa: ARG002
    """No-op: skip test metric logging."""

log_artifact(path)

No-op: skip artifact logging.

Source code in orchard/tracking/tracker.py

def log_artifact(self, path: Path) -> None:  # noqa: ARG002
    """No-op: skip artifact logging."""

log_artifacts_dir(directory)

No-op: skip directory artifact logging.

Source code in orchard/tracking/tracker.py

def log_artifacts_dir(self, directory: Path) -> None:  # noqa: ARG002
    """No-op: skip directory artifact logging."""

start_optuna_trial(trial_number, params)

No-op: skip nested Optuna trial run creation.

Source code in orchard/tracking/tracker.py

def start_optuna_trial(self, trial_number: int, params: dict[str, Any]) -> None:  # noqa: ARG002
    """No-op: skip nested Optuna trial run creation."""

end_optuna_trial(best_metric)

No-op: skip nested Optuna trial run closure.

Source code in orchard/tracking/tracker.py

def end_optuna_trial(self, best_metric: float) -> None:  # noqa: ARG002
    """No-op: skip nested Optuna trial run closure."""

end_run()

No-op: skip MLflow run closure.

Source code in orchard/tracking/tracker.py

def end_run(self) -> None:
    """No-op: skip MLflow run closure."""

MLflowTracker(experiment_name='orchard-ml')

MLflow-based experiment tracker for Orchard ML runs.

Manages a single MLflow run lifecycle: parameter logging, per-epoch metrics, final test metrics, and artifact collection. Supports nested runs for Optuna trial tracking.

Attributes:

Name	Type	Description
experiment_name	str	MLflow experiment name.

Source code in orchard/tracking/tracker.py

def __init__(self, experiment_name: str = "orchard-ml") -> None:
    self.experiment_name = experiment_name
    self._parent_run_id: str | None = None

start_run(cfg, run_name, tracking_uri)

Start an MLflow run and log all config parameters.

Parameters:

Name	Type	Description	Default
cfg	Any	Config object with dump_serialized() method.	required
run_name	str	Human-readable run name (typically paths.run_id).	required
tracking_uri	str	SQLite URI for local MLflow storage.	required

Source code in orchard/tracking/tracker.py

def start_run(self, cfg: Any, run_name: str, tracking_uri: str) -> None:
    """
    Start an MLflow run and log all config parameters.

    Args:
        cfg: Config object with dump_serialized() method.
        run_name: Human-readable run name (typically paths.run_id).
        tracking_uri: SQLite URI for local MLflow storage.
    """
    # Temporarily suppress all INFO logs during MLflow/Alembic DB init
    prev_level = logging.root.manager.disable
    logging.disable(logging.WARNING)
    try:
        mlflow.set_tracking_uri(tracking_uri)
        mlflow.set_experiment(self.experiment_name)
        mlflow.start_run(run_name=run_name)
    finally:
        logging.disable(prev_level)
    active_run = mlflow.active_run()
    self._parent_run_id = active_run.info.run_id if active_run else None

    # Log flattened config as params (MLflow has 500-char value limit)
    flat_params = _flatten_dict(cfg.dump_serialized())
    safe_params = {k: str(v)[:500] for k, v in flat_params.items() if v is not None}
    mlflow.log_params(safe_params)

    logger.info("  %s MLflow run started", LogStyle.ARROW)

log_epoch(epoch, train_loss, val_metrics, lr)

Log per-epoch training metrics.

Parameters:

Name	Type	Description	Default
epoch	int	Current epoch number (1-based).	required
train_loss	float	Training loss for this epoch.	required
val_metrics	Mapping[str, float]	Validation metrics dict with 'loss', 'accuracy', 'auc', 'f1'.	required
lr	float	Current learning rate.	required

Source code in orchard/tracking/tracker.py

def log_epoch(
    self, epoch: int, train_loss: float, val_metrics: Mapping[str, float], lr: float
) -> None:
    """
    Log per-epoch training metrics.

    Args:
        epoch: Current epoch number (1-based).
        train_loss: Training loss for this epoch.
        val_metrics: Validation metrics dict with 'loss', 'accuracy', 'auc', 'f1'.
        lr: Current learning rate.
    """
    mlflow.log_metrics(
        {
            "train_loss": train_loss,
            "val_loss": val_metrics[METRIC_LOSS],
            "val_accuracy": val_metrics[METRIC_ACCURACY],
            "val_auc": val_metrics.get(METRIC_AUC, 0.0),
            "val_f1": val_metrics.get(METRIC_F1, 0.0),
            "learning_rate": lr,
        },
        step=epoch,
    )

log_test_metrics(test_acc, macro_f1)

Log final test set metrics.

Parameters:

Name	Type	Description	Default
test_acc	float	Test set accuracy.	required
macro_f1	float	Test set macro-averaged F1 score.	required

Source code in orchard/tracking/tracker.py

def log_test_metrics(self, test_acc: float, macro_f1: float) -> None:
    """
    Log final test set metrics.

    Args:
        test_acc: Test set accuracy.
        macro_f1: Test set macro-averaged F1 score.
    """
    mlflow.log_metrics({"test_accuracy": test_acc, "test_macro_f1": macro_f1})

log_artifact(path)

Log a single file as an MLflow artifact.

Parameters:

Name	Type	Description	Default
path	Path	Path to file to log.	required

Source code in orchard/tracking/tracker.py

def log_artifact(self, path: Path) -> None:
    """
    Log a single file as an MLflow artifact.

    Args:
        path: Path to file to log.
    """
    if path.exists():
        mlflow.log_artifact(str(path))

log_artifacts_dir(directory)

Log all files in a directory as MLflow artifacts.

Parameters:

Name	Type	Description	Default
directory	Path	Directory whose contents to log.	required

Source code in orchard/tracking/tracker.py

def log_artifacts_dir(self, directory: Path) -> None:
    """
    Log all files in a directory as MLflow artifacts.

    Args:
        directory: Directory whose contents to log.
    """
    if directory.exists() and directory.is_dir():
        mlflow.log_artifacts(str(directory))

start_optuna_trial(trial_number, params)

Start a nested MLflow run for an Optuna trial.

Parameters:

Name	Type	Description	Default
trial_number	int	Optuna trial number.	required
params	dict[str, Any]	Sampled hyperparameters for this trial.	required

Source code in orchard/tracking/tracker.py

def start_optuna_trial(self, trial_number: int, params: dict[str, Any]) -> None:
    """
    Start a nested MLflow run for an Optuna trial.

    Args:
        trial_number: Optuna trial number.
        params: Sampled hyperparameters for this trial.
    """
    mlflow.start_run(
        run_name=f"trial_{trial_number:03d}",
        nested=True,
    )
    safe_params = {k: str(v)[:500] for k, v in params.items()}
    mlflow.log_params(safe_params)

end_optuna_trial(best_metric)

End the current nested Optuna trial run.

Parameters:

Name	Type	Description	Default
best_metric	float	Best validation metric achieved in this trial.	required

Source code in orchard/tracking/tracker.py

def end_optuna_trial(self, best_metric: float) -> None:
    """
    End the current nested Optuna trial run.

    Args:
        best_metric: Best validation metric achieved in this trial.
    """
    mlflow.log_metric("best_trial_metric", best_metric)
    mlflow.end_run()

end_run()

End the active MLflow run.

Source code in orchard/tracking/tracker.py

def end_run(self) -> None:
    """End the active MLflow run."""
    if mlflow.active_run():
        mlflow.end_run()
        logger.info("  %s MLflow run ended", LogStyle.ARROW)

create_tracker(cfg)

Factory: returns MLflowTracker if tracking is configured, else NoOpTracker.

Parameters:

Name	Type	Description	Default
cfg	Any	Config object. If cfg.tracking is set and enabled, returns MLflowTracker.	required

Returns:

Type	Description
TrackerProtocol	Active tracker instance.

Source code in orchard/tracking/tracker.py

def create_tracker(cfg: Any) -> TrackerProtocol:
    """
    Factory: returns MLflowTracker if tracking is configured, else NoOpTracker.

    Args:
        cfg: Config object. If cfg.tracking is set and enabled, returns MLflowTracker.

    Returns:
        Active tracker instance.
    """
    tracking_cfg = getattr(cfg, "tracking", None)
    if tracking_cfg is None or not tracking_cfg.enabled:
        return NoOpTracker()

    if not _MLFLOW_AVAILABLE:
        logger.warning(
            "Tracking enabled in config but mlflow is not installed. "
            "Install with: pip install mlflow"
        )
        return NoOpTracker()

    return MLflowTracker(experiment_name=tracking_cfg.experiment_name)