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

Code before:
        latest_checkpoint_directory = max(sharded_checkpoint_directories, default=None, key=_get_checkpoint_number)
        sharded_checkpoint_directories = (
            [latest_checkpoint_directory] if latest_checkpoint_directory is not None else []
        )

    log.info("Found %d sharded checkpoint directories for %s", len(sharded_checkpoint_directories), run_dir_or_archive)

    return sharded_checkpoint_directories


def _unshard_checkpoint(sharded_checkpoint_dir: str, dest_dir: str, run_dir_or_archive: str, unsharding_config: UnshardCheckpointsConfig):
    local_storage = LocalFileSystemAdapter()

    # Download checkpoint to a temp dir
    sharding_input_dir = local_storage.create_temp_dir()
    src_storage = _get_storage_adapter_for_path(sharded_checkpoint_dir)
    src_storage.download_folder(sharded_checkpoint_dir, sharding_input_dir)

    # Set unsharder output to a temp dir
    sharding_output_dir: str
    sharding_output_dir = local_storage.create_temp_dir()

    result = subprocess.run(
        ["python", str(unsharding_config.unshard_script_path), sharding_input_dir, sharding_output_dir],
        check=False,
    )
    if result.returncode != 0:
        log.error(
            "Unsharding from %s to %s failed with error code %d",
            sharding_input_dir,
            sharding_output_dir,
            result.returncode,
        )

        local_storage.delete_path(sharding_output_dir)
        return

    log.info(
        "Successfully unsharded from %s to %s, starting upload to %s",
        sharding_input_dir,
        sharding_output_dir,
        dest_dir,
    )

    dest_storage = _get_storage_adapter_for_path(dest_dir)
    dest_storage.upload(sharding_output_dir, dest_dir)


def _unshard_checkpoints(
    run_storage: StorageAdapter,
    run_dir_or_archive: str,

Review comment (suggestion):
The `unshard.py` script is simple enough that I would be okay with copying that logic into here instead of shelling out.

Code after:
        )

        local_storage.delete_path(sharding_input_dir)
        local_storage.delete_path(sharding_output_dir)
        return

    # model
    model_output = str(Path(sharding_output_dir) / "model.pt")
    log.info("Saving model state to %s", model_output)
    torch.save(model_state_dict, model_output)
    del model_state_dict

    # optimizer
    optim_output = str(Path(sharding_output_dir) / "optim.pt")
    log.info("Saving optimizer state to %s", optim_output)
    torch.save(optim_state_dict, optim_output)
    del optim_state_dict

    # trainer
    train_output = str(Path(sharding_output_dir) / "train.pt")
    log.info("Saving everything else to %s", train_output)
    torch.save(trainer_state_dict, train_output)
    del trainer_state_dict

    log.info("Copying config.yaml to %s", sharding_output_dir)
    shutil.copy(Path(sharding_input_dir) / "config.yaml", sharding_output_dir)

    log.info(
        "Successfully unsharded from %s to %s, starting upload to %s",
        sharding_input_dir,
        sharding_output_dir,
        dest_dir,
    )

    dest_storage = _get_storage_adapter_for_path(dest_dir)
    dest_storage.upload(sharding_output_dir, dest_dir)

    local_storage.delete_path(sharding_input_dir)
    local_storage.delete_path(sharding_output_dir)


def _unshard_checkpoints(
    run_storage: StorageAdapter,
    run_dir_or_archive: str,
    checkpoints_dest_dir: str,
    config: UnshardCheckpointsConfig,
):
    log.info("Starting unsharding checkpoints of run directory or archive %s", run_dir_or_archive)

    run_dir = _unarchive_if_archive(run_dir_or_archive, run_storage)
    run_dir_storage = _get_storage_adapter_for_path(run_dir)