Repository: allenai/OLMo
File: scripts/storage_cleaner.py (Python)

Code before:
    storage = _get_storage_adapter_for_path(run_path)
    run_dir_or_archive = _format_dir_or_archive_path(storage, run_path)
    _unshard_checkpoints(storage, run_dir_or_archive, checkpoints_dest_dir, config)


def _get_wandb_path(run_dir: str) -> str:
    run_dir_storage = _get_storage_adapter_for_path(run_dir)

    config_path = os.path.join(run_dir, CONFIG_YAML)
    if not run_dir_storage.is_file(config_path):
        raise FileNotFoundError("No config file found in run dir, cannot get wandb path")

    local_config_path = cached_path(config_path)
    config = TrainConfig.load(local_config_path, validate_paths=False)

    if config.wandb is None or config.wandb.entity is None or config.wandb.project is None:
        raise ValueError(f"Run at {run_dir} has missing wandb config, cannot get wandb run path")

    run_filters = {
        "display_name": config.wandb.name,
    }
    if config.wandb.group is not None:
        run_filters["group"] = config.wandb.group

    log.info("Wandb entity/project: %s/%s", config.wandb.entity, config.wandb.project)
    log.info("Wandb filters: %s", run_filters)

    api = wandb.Api()
    wandb_matching_runs = api.runs(path=f"{config.wandb.entity}/{config.wandb.project}", filters=run_filters)

    if len(wandb_matching_runs) == 0:
        raise RuntimeError(f"Failed to find any wandb runs for {run_dir}. Run might no longer exist")

    if len(wandb_matching_runs) > 1:
        raise RuntimeError(
            f"Found {len(wandb_matching_runs)} runs matching run dir {run_dir}, cannot determine correct run"
        )

    wandb_run = wandb_matching_runs[0]
    return "/".join(wandb_run.path)


def _append_wandb_path(
    base_dir: str, run_dir_or_archive: str, append_archive_extension: bool = False, run_dir: Optional[str] = None
) -> str:
    run_dir_or_archive_storage = _get_storage_adapter_for_path(run_dir_or_archive)
    if run_dir is None:
        run_dir = _unarchive_if_archive(run_dir_or_archive, run_dir_or_archive_storage)

    wandb_path = _get_wandb_path(run_dir)


Review comment (bug):
Can this be debug?
```suggestion
    log.debug("Wandb entity/project: %s/%s", config.wandb.entity, config.wandb.project)
    log.debug("Wandb filters: %s", run_filters)
```

Code after:
            )
            continue

        if config.dry_run:
            log.info("Would unshard sharded checkpoint %s to %s", sharded_checkpoint_directory, dest_directory)
        else:
            log.info("Unsharding sharded checkpoint %s to %s", sharded_checkpoint_directory, dest_directory)
            _unshard_checkpoint(sharded_checkpoint_directory, dest_directory, run_dir, config)


def unshard_run_checkpoints(run_path: str, checkpoints_dest_dir: str, config: UnshardCheckpointsConfig):
    storage = _get_storage_adapter_for_path(run_path)
    run_dir_or_archive = _format_dir_or_archive_path(storage, run_path)
    _unshard_checkpoints(storage, run_dir_or_archive, checkpoints_dest_dir, config)


def _get_wandb_runs_from_wandb_dir(storage: StorageAdapter, wandb_dir: str, run_config: TrainConfig) -> List:
    # For some reason, we often have a redundant nested wandb directory. Step into it here.
    nested_wandb_dir = os.path.join(wandb_dir, "wandb/")
    if storage.is_dir(nested_wandb_dir):
        wandb_dir = nested_wandb_dir

    # Wandb run directory names are stored in format <run>-<timestamp>-<id>
    # https://docs.wandb.ai/guides/track/save-restore#examples-of-wandbsave
    dir_names = storage.list_dirs(wandb_dir)
    wandb_run_dir_names = [dir_name for dir_name in dir_names if dir_name.startswith("run")]
    if len(wandb_run_dir_names) == 0:
        log.warning("No wandb run directories found in wandb dir %s", wandb_dir)
        return []

    wandb_ids = [dir_name.split("-")[2] for dir_name in wandb_run_dir_names if dir_name.count("-") >= 2]

    log.debug("Wandb ids: %s", wandb_ids)

    assert run_config.wandb is not None
    api: wandb.Api = wandb.Api()
    return [api.run(path=f"{run_config.wandb.entity}/{run_config.wandb.project}/{id}") for id in wandb_ids]


def _get_wandb_path_from_run(wandb_run) -> str:
    return "/".join(wandb_run.path)


def _get_wandb_runs_from_train_config(config: TrainConfig) -> List:
    assert config.wandb is not None

    run_filters = {
        "display_name": config.wandb.name,
    }
    if config.wandb.group is not None:
        run_filters["group"] = config.wandb.group