Delete deepscreen/utils/hydra.py.bak

deepscreen/utils/hydra.py.bak

@@ -1,182 +0,0 @@
-from datetime import datetime
-from pathlib import Path
-import re
-from typing import Any, Tuple
-
-import pandas as pd
-from hydra import TaskFunction
-from hydra.core.hydra_config import HydraConfig
-from hydra.core.override_parser.overrides_parser import OverridesParser
-from hydra.core.utils import _save_config
-from hydra.experimental.callbacks import Callback
-from hydra.types import RunMode
-from hydra._internal.config_loader_impl import ConfigLoaderImpl
-from omegaconf import DictConfig, OmegaConf
-from omegaconf.errors import MissingMandatoryValue
-
-from deepscreen.utils import get_logger
-
-log = get_logger(__name__)
-
-
-class CSVExperimentSummary(Callback):
-    """On multirun end, aggregate the results from each job's metrics.csv and save them in metrics_summary.csv."""
-
-    def __init__(self, filename: str = 'experiment_summary.csv', prefix: str | Tuple[str] = 'test/'):
-        self.filename = filename
-        self.prefix = prefix if isinstance(prefix, str) else tuple(prefix)
-        self.input_experiment_summary = None
-        self.time = {}
-
-    def on_multirun_start(self, config: DictConfig, **kwargs: Any) -> None:
-        if config.hydra.get('overrides') and config.hydra.overrides.get('task'):
-            for i, override in enumerate(config.hydra.overrides.task):
-                if override.startswith("ckpt_path"):
-                    ckpt_path = override.split('=', 1)[1]
-                    if ckpt_path.endswith(('.csv', '.txt', '.tsv', '.ssv', '.psv')):
-                        config.hydra.overrides.task[i] = self.parse_ckpt_path_from_experiment_summary(ckpt_path)
-                        log.info(ckpt_path)
-                    break
-            if config.hydra.sweeper.get('params'):
-                if config.hydra.sweeper.params.get('ckpt_path'):
-                    ckpt_path = str(config.hydra.sweeper.params.ckpt_path).strip("'\"")
-                    if ckpt_path.endswith(('.csv', '.txt', '.tsv', '.ssv', '.psv')):
-                        config.hydra.sweeper.params.ckpt_path = self.parse_ckpt_path_from_experiment_summary(ckpt_path)
-                        log.info(ckpt_path)
-    def on_job_start(self, config: DictConfig, *, task_function: TaskFunction, **kwargs: Any) -> None:
-        self.time['start'] = datetime.now()
-
-    def on_job_end(self, config: DictConfig, job_return, **kwargs: Any) -> None:
-        # Skip callback if job is DDP subprocess
-        if "ddp" in job_return.hydra_cfg.hydra.job.name:
-            return
-
-        try:
-            self.time['end'] = datetime.now()
-            if config.hydra.mode == RunMode.RUN:
-                summary_file_path = Path(config.hydra.run.dir) / self.filename
-            elif config.hydra.mode == RunMode.MULTIRUN:
-                summary_file_path = Path(config.hydra.sweep.dir) / self.filename
-            else:
-                raise RuntimeError('Invalid Hydra `RunMode`.')
-
-            if summary_file_path.is_file():
-                summary_df = pd.read_csv(summary_file_path)
-            else:
-                summary_df = pd.DataFrame()
-
-            # Add job and override info
-            info_dict = {}
-            if job_return.overrides:
-                info_dict = dict(override.split('=', 1) for override in job_return.overrides)
-            info_dict['job_status'] = job_return.status.name
-            info_dict['job_id'] = job_return.hydra_cfg.hydra.job.id
-            info_dict['wall_time'] = str(self.time['end'] - self.time['start'])
-
-            # Add checkpoint info
-            if info_dict.get('ckpt_path'):
-                info_dict['ckpt_path'] = str(info_dict['ckpt_path']).strip("'\"")
-
-            ckpt_path = str(job_return.cfg.ckpt_path).strip("'\"")
-            if Path(ckpt_path).is_file():
-                if info_dict.get('ckpt_path') and ckpt_path != info_dict['ckpt_path']:
-                    info_dict['previous_ckpt_path'] = info_dict['ckpt_path']
-                info_dict['ckpt_path'] = ckpt_path
-            if info_dict.get('ckpt_path'):
-                info_dict['best_epoch'] = int(re.search(r'epoch_(\d+)', info_dict['ckpt_path']).group(1))
-
-            # Add metrics info
-            metrics_df = pd.DataFrame()
-            if config.get('logger'):
-                output_dir = Path(config.hydra.runtime.output_dir).resolve()
-                csv_metrics_path = output_dir / config.logger.csv.name / "metrics.csv"
-                if csv_metrics_path.is_file():
-                    log.info(f"Summarizing metrics with prefix `{self.prefix}` from {csv_metrics_path}")
-                    metrics_df = pd.read_csv(csv_metrics_path)
-                    # Find rows where 'test/' columns are not null and reset its epoch to the best model epoch
-                    test_columns = [col for col in metrics_df.columns if col.startswith('test/')]
-                    if test_columns:
-                        mask = metrics_df[test_columns].notna().any(axis=1)
-                        metrics_df.loc[mask, 'epoch'] = info_dict['best_epoch']
-                    # Group and filter by best epoch
-                    metrics_df = metrics_df.groupby('epoch').first()
-                    metrics_df = metrics_df[metrics_df.index == info_dict['best_epoch']]
-                else:
-                    log.info(f"No metrics.csv found in {output_dir}")
-
-            if metrics_df.empty:
-                metrics_df = pd.DataFrame(data=info_dict, index=[0])
-            else:
-                metrics_df = metrics_df.assign(**info_dict)
-                metrics_df.index = [0]
-
-            # Add extra info from the input batch experiment summary
-            if self.input_experiment_summary is not None and 'ckpt_path' in metrics_df.columns:
-                log.info(self.input_experiment_summary['ckpt_path'])
-                log.info(metrics_df['ckpt_path'])
-                orig_meta = self.input_experiment_summary[
-                    self.input_experiment_summary['ckpt_path'] == metrics_df['ckpt_path'][0]
-                    ].head(1)
-                if not orig_meta.empty:
-                    orig_meta.index = [0]
-                metrics_df = metrics_df.astype('O').combine_first(orig_meta.astype('O'))
-
-            summary_df = pd.concat([summary_df, metrics_df])
-
-            # Drop empty columns
-            summary_df.dropna(inplace=True, axis=1, how='all')
-            summary_df.to_csv(summary_file_path, index=False, mode='w')
-            log.info(f"Experiment summary saved to {summary_file_path}")
-        except Exception as e:
-            log.exception("Unable to save the experiment summary due to an error.", exc_info=e)
-
-    def parse_ckpt_path_from_experiment_summary(self, ckpt_path):
-        log.info(ckpt_path)
-        try:
-            self.input_experiment_summary = pd.read_csv(
-                ckpt_path, usecols=lambda col: not col.startswith(self.prefix)
-            )
-            self.input_experiment_summary['ckpt_path'] = self.input_experiment_summary['ckpt_path'].apply(
-                lambda x: x.strip("'\"")
-            )
-            ckpt_list = list(set(self.input_experiment_summary['ckpt_path']))
-            parsed_ckpt_path = ','.join([f"'{ckpt}'" for ckpt in ckpt_list])
-            return parsed_ckpt_path
-
-        except Exception as e:
-            log.exception(
-                f'Error in parsing checkpoint paths from experiment_summary file ({ckpt_path}).',
-                exc_info=e
-            )
-
-
-def checkpoint_rerun_config(config: DictConfig):
-    hydra_cfg = HydraConfig.get()
-
-    if hydra_cfg.get('output_subdir'):
-        ckpt_cfg_path = Path(config.ckpt_path).parents[1] / hydra_cfg.output_subdir / 'config.yaml'
-        hydra_output = Path(hydra_cfg.runtime.output_dir) / hydra_cfg.output_subdir
-
-        if ckpt_cfg_path.is_file():
-            log.info(f"Found config file for the checkpoint at {str(ckpt_cfg_path)}; "
-                     f"merging config overrides with checkpoint config...")
-            ckpt_cfg = OmegaConf.load(ckpt_cfg_path)
-
-            # Recompose checkpoint config with overrides
-
-            if hydra_cfg.overrides.get('task'):
-                parser = OverridesParser.create()
-                parsed_overrides = parser.parse_overrides(overrides=hydra_cfg.overrides.task)
-                filtered_overrides = []
-                for override in parsed_overrides:
-                    if not override.is_force_add():
-                        OmegaConf.update(ckpt_cfg, override.key_or_group, override.value())
-                    filtered_overrides.append(override)
-                log.info(filtered_overrides)
-                ConfigLoaderImpl._apply_overrides_to_config(filtered_overrides, config)
-
-            _save_config(config, "config.yaml", hydra_output)
-
-    return config
-
-