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

Code before:
            delete_bad_runs(args.run_paths, delete_bad_runs_config)
        else:
            raise ValueError("Run paths not provided for run cleaning")
    else:
        raise NotImplementedError(args.op)


def _add_cached_path_r2_client(r2_account_id: str):
    endpoint_url = f"https://{r2_account_id}.r2.cloudflarestorage.com"

    class R2SchemeClient(S3Client):
        """
        A class that the `cached_path` module can use to retrieve resources from
        R2. Refer to
        [cached_path docs](https://github.com/allenai/cached_path/blob/main/docs/source/overview.md#supported-url-schemes).
        """

        scheme = "r2"

        def __init__(self, resource: str) -> None:
            super().__init__(resource)
            parsed_path = urlparse(resource)
            bucket_name = parsed_path.netloc
            r2_path = parsed_path.path.lstrip("/")

            session = boto3.session.Session()
            if session.get_credentials() is None:
                # Use unsigned requests.
                s3_resource = session.resource(
                    "s3",
                    endpoint_url=endpoint_url,
                    config=botocore.client.Config(signature_version=botocore.UNSIGNED),
                )
            else:
                s3_resource = session.resource("s3", endpoint_url=endpoint_url)
            self.s3_object = s3_resource.Object(bucket_name, r2_path)  # type: ignore

    add_scheme_client(R2SchemeClient)


def _add_cached_path_scheme_clients():
    r2_account_id = os.environ.get("R2_ACCOUNT_ID")
    if r2_account_id is not None:
        _add_cached_path_r2_client(r2_account_id)


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)

Review comment (security):
How does authentication work for this? It reads environment variables? Then I assume there is no way to have both R2 and S3 work at the same time, because they read auth information from the same place?

Code after:
        """
        A class that the `cached_path` module can use to retrieve resources from
        S3 (and R2, which is S3-based).  Refer to
        [cached_path docs](https://github.com/allenai/cached_path/blob/main/docs/source/overview.md#supported-url-schemes).
        """

        # This is used by cached_path to get the schemes are handled by this client
        scheme = ("s3", "r2")

        def __init__(self, resource: str) -> None:
            super().__init__(resource)
            parsed_path = urlparse(resource)
            bucket_name = parsed_path.netloc
            key = parsed_path.path.lstrip("/")

            profile_name = util._get_s3_profile_name(parsed_path.scheme)
            endpoint_url = util._get_s3_endpoint_url(parsed_path.scheme)

            session = boto3.session.Session(profile_name=profile_name)
            s3_resource = session.resource("s3", endpoint_url=endpoint_url)
            self.s3_object = s3_resource.Object(bucket_name, key)  # type: ignore

    add_scheme_client(S3SchemeClient)


def _setup_cached_path(args: argparse.Namespace):
    if args.temp_dir is not None:
        set_cache_dir(args.temp_dir)

    _add_cached_path_s3_client()


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.",
    )