tracking/translations_parser/publishers.py

import csv
import logging
import sys
from abc import ABC
from collections import defaultdict
from pathlib import Path
from typing import Sequence

import wandb
import yaml

from translations_parser.data import Metric, TrainingEpoch, ValidationEpoch
from translations_parser.utils import parse_task_label, parse_gcp_metric, patch_model_name

logger = logging.getLogger(__name__)

METRIC_KEYS = sorted(set(Metric.__annotations__.keys()) - {"importer", "dataset", "augmentation"})


class Publisher(ABC):
    """
    Abstract class used to publish parsed data.

    Either the `handle_*` methods can be overriden for real time
    publication (introduced later on) or the `publish` method
    with all results (including parser run date, configuration…).
    """

    def open(self, parser) -> None:
        ...

    def handle_training(self, training: TrainingEpoch) -> None:
        ...

    def handle_validation(self, validation: ValidationEpoch) -> None:
        ...

    def handle_metrics(self, metrics: Sequence[Metric]) -> None:
        ...

    def publish(self) -> None:
        ...

    def close(self) -> None:
        ...


class CSVExport(Publisher):
    def __init__(self, output_dir: Path) -> None:
        from translations_parser.parser import TrainingParser

        if not output_dir.is_dir():
            raise ValueError("Output must be a valid directory for the CSV export")
        self.output_dir = output_dir
        self.parser: TrainingParser | None = None

    def open(self, parser=None) -> None:
        self.parser = parser

    def write_data(
        self, output: Path, entries: Sequence[TrainingEpoch | ValidationEpoch], dataclass: type
    ) -> None:
        if not entries:
            logger.warning(f"No {dataclass.__name__} entry, skipping.")
        with open(output, "w") as f:
            writer = csv.DictWriter(f, fieldnames=dataclass.__annotations__)
            writer.writeheader()
            for entry in entries:
                writer.writerow(vars(entry))

    def publish(self) -> None:
        assert self.parser is not None, "Parser must be set to run CSV publication."
        training_log = self.parser.output
        training_output = self.output_dir / "training.csv"
        if training_output.exists():
            logger.warning(f"Training output file {training_output} exists, skipping.")
        else:
            self.write_data(training_output, training_log.training, TrainingEpoch)

        validation_output = self.output_dir / "validation.csv"
        if validation_output.exists():
            logger.warning(f"Validation output file {validation_output} exists, skipping.")
        else:
            self.write_data(validation_output, training_log.validation, ValidationEpoch)


class WandB(Publisher):
    def __init__(
        self,
        *,
        project: str,
        group: str,
        name: str,
        suffix: str = "",
        # Optional path to a directory containing training artifacts
        artifacts: Path | None = None,
        artifacts_name: str = "logs",
        **extra_kwargs,
    ):
        from translations_parser.parser import TrainingParser

        # Set logging of wandb module to WARNING, so we output training logs instead
        self.wandb_logger = logging.getLogger("wandb")
        self.wandb_logger.setLevel(logging.ERROR)

        self.project = project
        self.group = group
        self.suffix = suffix
        # Build a unique run identifier based on the passed suffix
        # This ID is also used as display name on W&B, as the interface expects unique display names among runs
        self.run = f"{name}{suffix}"

        self.artifacts = artifacts
        self.artifacts_name = artifacts_name
        self.extra_kwargs = extra_kwargs
        self.parser: TrainingParser | None = None
        self.wandb: wandb.sdk.wandb_run.Run | wandb.sdk.lib.disabled.RunDisabled | None = None

    def close(self) -> None:
        if self.wandb is None:
            return

        # Publish artifacts
        if self.artifacts:
            artifact = wandb.Artifact(name=self.artifacts_name, type=self.artifacts_name)
            artifact.add_dir(local_path=str(self.artifacts.resolve()))
            self.wandb.log_artifact(artifact)

        if self.parser is not None:
            # Store Marian logs as the main log artifact, instead of W&B client runtime.
            # This will be overwritten in case an unhandled exception occurs.
            for line in self.parser.parsed_logs:
                sys.stdout.write(f"{line}\n")

        self.wandb.finish()

    def open(self, parser=None) -> None:
        self.parser = parser
        config = getattr(parser, "config", {}).copy()
        config.update(self.extra_kwargs.pop("config", {}))
        # Publish datasets stats directly in the dashboard
        datasets = config.pop("datasets", None)

        try:
            self.wandb = wandb.init(
                project=self.project,
                group=self.group,
                name=self.run,
                id=self.run,
                config=config,
                # Since we use unique run names based on group ID (e.g. finetune-student_MjcJG),
                # we can use "allow" mode for resuming a stopped Taskcluster run in case of preemption.
                # It will continue logging to the same run if it exists.
                # Offline publication should handle run deletion separately (use --override-runs).
                resume="allow",
                **self.extra_kwargs,
            )
            if self.wandb.resumed:
                logger.info(f"W&B run is being resumed from existing run '{self.run}'.")
        except Exception as e:
            logger.error(f"WandB client could not be initialized: {e}. No data will be published.")

        if datasets is not None:
            # Log dataset sizes as a custom bar chart
            self.wandb.log(
                {
                    "Datasets": wandb.plot.bar(
                        wandb.Table(
                            columns=["Name", "Count"],
                            data=[[key, value] for key, value in datasets.items()],
                        ),
                        "Name",
                        "Count",
                        title="Datasets",
                    )
                }
            )

    def generic_log(self, data: TrainingEpoch | ValidationEpoch) -> None:
        if self.wandb is None:
            return
        epoch = vars(data)
        step = epoch.pop("up")
        for key, val in epoch.items():
            if val is None:
                # Do not publish null values (e.g. perplexity in Marian 1.10)
                continue
            self.wandb.log(step=step, data={key: val})

    def handle_training(self, training: TrainingEpoch) -> None:
        self.generic_log(training)

    def handle_validation(self, validation: ValidationEpoch) -> None:
        self.generic_log(validation)

    def handle_metrics(self, metrics: Sequence[Metric]) -> None:
        if self.wandb is None:
            return
        for metric in metrics:
            title = metric.importer
            if metric.augmentation:
                title = f"{title}_{metric.augmentation}"
            if metric.dataset:
                title = f"{title}_{metric.dataset}"
            # Publish a bar chart (a table with values will also be available from W&B)
            self.wandb.log(
                {
                    title: wandb.plot.bar(
                        wandb.Table(
                            columns=["Metric", "Value"],
                            data=[
                                [key, getattr(metric, key)]
                                for key in METRIC_KEYS
                                if getattr(metric, key) is not None
                            ],
                        ),
                        "Metric",
                        "Value",
                        title=title,
                    )
                }
            )

    @classmethod
    def publish_group_logs(
        cls,
        *,
        logs_parent_folder: list[str],
        project: str,
        group: str,
        suffix: str,
        existing_runs: list[str] | None = None,
        snakemake: bool = False,
    ) -> None:
        """
        Publish files within `logs_dir` to W&B artifacts for a specific group.
        A fake W&B run named `group_logs` is created to publish those artifacts
        among with all evaluation files (quantized + experiments).
        If existing run is set, runs found not specified in this list will also
        be published to W&B.
        """
        from translations_parser.parser import TrainingParser

        try:
            if (
                len(
                    wandb.Api().runs(
                        path=project, filters={"display_name": "group_logs", "group": group}
                    )
                )
                > 0
            ):
                logger.warning("Skipping group_logs fake run publication as it already exists")
                return
        except ValueError as e:
            # Project may not exist yet as group_logs is published before the first training task
            if "could not find project" not in str(e).lower():
                logger.warning(f"Detection of a previous group_logs run failed: {e}")

        logs_dir = Path("/".join([*logs_parent_folder[:-1], "logs", project, group]))
        models_dir = Path("/".join([*logs_parent_folder[:-1], "models", project, group]))
        # Old experiments use `speed` directory for quantized metrics
        quantized_metrics = sorted(
            Path(
                "/".join(
                    [*logs_parent_folder[:-1], "models", project, group, "evaluation", "speed"]
                )
            ).glob("*.metrics")
        )
        logs_metrics = sorted((logs_dir / "eval").glob("eval*.log"))
        direct_metrics = sorted((logs_dir / "metrics").glob("*.metrics"))

        taskcluster_metrics = []
        # Do not retrieve metrics from models directory for legacy Snakemake experiments
        if snakemake is False:
            taskcluster_metrics = sorted((models_dir).glob("**/*.metrics"))

        if quantized_metrics:
            logger.info(f"Found {len(quantized_metrics)} quantized metrics from speed folder")
        if logs_metrics:
            logger.info(f"Found {len(logs_metrics)} metrics from task logs")
        if direct_metrics:
            logger.info(f"Found {len(direct_metrics)} Snakemake metrics from .metrics artifacts")
        if taskcluster_metrics:
            logger.info(
                f"Found {len(taskcluster_metrics)} Taskcluster metrics from .metrics artifacts"
            )

        # Store metrics by run name
        metrics = defaultdict(list)
        # Add metrics from the speed folder
        for file in quantized_metrics:
            importer, dataset = file.stem.split("_", 1)
            metrics["quantized"].append(Metric.from_file(file, importer=importer, dataset=dataset))
        # Add metrics from tasks logs
        for file in logs_metrics:
            try:
                model_name, importer, dataset, aug = parse_task_label(file.stem)
                with file.open("r") as f:
                    lines = f.readlines()
                metrics[model_name].append(
                    Metric.from_tc_context(
                        importer=importer, dataset=dataset, lines=lines, augmentation=aug
                    )
                )
            except ValueError as e:
                logger.error(f"Could not parse metrics from {file.resolve()}: {e}")

        # Add metrics from old SnakeMake .metrics files
        for file in direct_metrics:
            model_name, importer, dataset, aug = parse_task_label(file.stem)
            try:
                metrics[model_name].append(
                    Metric.from_file(file, importer=importer, dataset=dataset, augmentation=aug)
                )
            except ValueError as e:
                logger.error(f"Could not parse metrics from {file.resolve()}: {e}")

        # Add metrics from new Taskcluster .metrics files
        for file in taskcluster_metrics:
            model_name = patch_model_name(file.parent.name)
            try:
                metric_attrs = parse_gcp_metric(file.stem)
                metrics[model_name].append(
                    Metric.from_file(
                        file,
                        importer=metric_attrs.importer,
                        dataset=metric_attrs.dataset,
                        augmentation=metric_attrs.augmentation,
                    )
                )
            except ValueError as e:
                logger.error(f"Could not parse metrics from {file.resolve()}: {e}")

        # Publish missing runs (runs without training data)
        missing_run_metrics = {}
        if existing_runs is not None:
            missing_run_metrics = {
                name: metrics for name, metrics in metrics.items() if name not in existing_runs
            }

        for model_name, model_metrics in missing_run_metrics.items():
            logger.info(f"Creating missing run {model_name} with associated metrics")
            publisher = cls(
                project=project,
                group=group,
                name=model_name,
                suffix=suffix,
            )
            publisher.open(TrainingParser(logs_iter=iter([]), publishers=[]))
            publisher.handle_metrics(model_metrics)
            publisher.close()

        # Publication of the `group_logs` fake run
        config = {}
        config_path = Path(
            "/".join([*logs_parent_folder[:-1], "experiments", project, group, "config.yml"])
        )
        if not config_path.is_file():
            logger.warning(f"No configuration file at {config_path}, skipping.")
        else:
            # Publish the YAML configuration as configuration on the group run
            with config_path.open("r") as f:
                data = f.read()
            try:
                config.update(yaml.safe_load(data))
            except Exception as e:
                logger.error(f"Config could not be read at {config_path}: {e}")

        publisher = cls(
            project=project,
            group=group,
            name="group_logs",
            suffix=suffix,
        )
        publisher.wandb = wandb.init(
            project=project,
            group=group,
            name=publisher.run,
            id=publisher.run,
            config=config,
        )

        if metrics:
            # Publish all evaluation metrics to a table
            table = wandb.Table(
                columns=["Group", "Model", "Importer", "Dataset", "Augmenation", *METRIC_KEYS],
                data=[
                    [group, run_name, metric.importer, metric.dataset, metric.augmentation]
                    + [getattr(metric, attr) for attr in METRIC_KEYS]
                    for run_name, run_metrics in metrics.items()
                    for metric in run_metrics
                ],
            )
            publisher.wandb.log({"metrics": table})

        if logs_dir.is_dir():
            # Publish logs directory content as artifacts
            artifact = wandb.Artifact(name=group, type="logs")
            artifact.add_dir(local_path=str(logs_dir.resolve()))
            publisher.wandb.log_artifact(artifact)
        publisher.wandb.finish()