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

Code before:
            dry_run=args.dry_run,
            temp_dir=temp_dir,
            should_check_is_run=args.should_check_is_run,
            ignore_non_runs=args.ignore_non_runs,
            max_archive_size=args.max_archive_size,
        )
        if args.run_paths is not None:
            delete_bad_runs(args.run_paths, delete_bad_runs_config)
        else:
            raise ValueError("Run paths not provided for run cleaning")
    elif args.op == CleaningOperations.UNSHARD_CHECKPOINTS:
        unshard_checkpoints_config = UnshardCheckpointsConfig(
            dry_run=args.dry_run,
            temp_dir=temp_dir,
            latest_checkpoint_only=args.latest_checkpoint_only,
        )
        if args.run_path is not None:
            unshard_run_checkpoints(args.run_path, args.dest_dir, unshard_checkpoints_config)
        else:
            raise ValueError("Run path not provided for unsharding")
    else:
        raise NotImplementedError(args.op)

    if Path(temp_dir).is_dir():
        log.info("Deleting temp dir %s", temp_dir)
        shutil.rmtree(temp_dir)


def _add_delete_subparser(subparsers: _SubParsersAction):
    delete_runs_parser: ArgumentParser = subparsers.add_parser(
        "clean", help="Delete bad runs (e.g. runs with no non-trivial checkpoints)"
    )
    delete_runs_parser.set_defaults(op=CleaningOperations.DELETE_BAD_RUNS)

    delete_runs_parser.add_argument(
        "run_paths",
        nargs="+",
        help="Directory or archive file paths corresponding to runs.",
    )

    run_verification_parser = delete_runs_parser.add_mutually_exclusive_group(required=False)
    run_verification_parser.add_argument(
        "--bypass_run_verification",
        action="store_false",
        dest="should_check_is_run",
        help="(UNSAFE) Bypass the sanity check that a directory/archive being deleted is a run. This could result in a non-run directory/archive being deleted.",
    )
    run_verification_parser.add_argument(
        "--ignore_non_runs",
        action="store_true",
        help="Ignore (do not delete) directories/archives that are not runs.",

Review comment (bug):
This should probably run even if the cleaning op fails, so you might want to add a `try: ... finally: ...` block.

Code after:
            delete_bad_runs_config = DeleteBadRunsConfig(
                dry_run=args.dry_run,
                temp_dir=temp_dir,
                should_check_is_run=args.should_check_is_run,
                ignore_non_runs=args.ignore_non_runs,
                max_archive_size=args.max_archive_size,
            )
            if args.run_paths is not None:
                delete_bad_runs(args.run_paths, delete_bad_runs_config)
            else:
                raise ValueError("Run paths not provided for run cleaning")
        elif args.op == CleaningOperations.UNSHARD_CHECKPOINTS:
            unshard_checkpoints_config = UnshardCheckpointsConfig(
                dry_run=args.dry_run,
                temp_dir=temp_dir,
                latest_checkpoint_only=args.latest_checkpoint_only,
            )
            if args.run_path is not None:
                unshard_run_checkpoints(args.run_path, args.dest_dir, unshard_checkpoints_config)
            else:
                raise ValueError("Run path not provided for unsharding")
        else:
            raise NotImplementedError(args.op)
    finally:
        if Path(temp_dir).is_dir():
            log.info("Deleting temp dir %s", temp_dir)
            shutil.rmtree(temp_dir)


def _add_delete_subparser(subparsers: _SubParsersAction):
    delete_runs_parser: ArgumentParser = subparsers.add_parser(
        "clean", help="Delete bad runs (e.g. runs with no non-trivial checkpoints)"
    )
    delete_runs_parser.set_defaults(op=CleaningOperations.DELETE_BAD_RUNS)

    delete_runs_parser.add_argument(
        "run_paths",
        nargs="+",
        help="Directory or archive file paths corresponding to runs.",
    )

    run_verification_parser = delete_runs_parser.add_mutually_exclusive_group(required=False)
    run_verification_parser.add_argument(
        "--bypass_run_verification",
        action="store_false",
        dest="should_check_is_run",
        help="(UNSAFE) Bypass the sanity check that a directory/archive being deleted is a run. This could result in a non-run directory/archive being deleted.",
    )
    run_verification_parser.add_argument(
        "--ignore_non_runs",
        action="store_true",