File size: 6,236 Bytes
44459bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 |
"""API custom file upload call wrappers."""
import hashlib
import os
import shutil
import tempfile
from copy import copy
from datetime import datetime
from pathlib import Path
from zipfile import ZIP_DEFLATED, ZipFile
import requests
import typer
from folding_studio_data_models import CustomFileType
from rich import (
print, # pylint:disable=redefined-builtin
)
from tqdm import tqdm
from tqdm.utils import CallbackIOWrapper
from folding_studio.config import API_URL, REQUEST_TIMEOUT
def _upload_file_to_signed_url(
signed_url: str,
src: str,
headers: dict[str, str],
) -> requests.Response:
"""Upload a local file to a GCS bucket using a signed URL.
Use a PUT request.
Args:
signed_url (str): the signed URL corresponding to the GCS path.
src (src | Path): the local file path.
headers (dict[str, str]): HTTP request headers.
Raises:
requests.exceptions.HTTPError: if something went wrong during the uploading.
Returns:
A response to the PUT request.
"""
put_headers = copy(headers)
put_headers["Content-type"] = "application/octet-stream"
file_size = os.path.getsize(src)
with open(src, "rb") as fd:
with tqdm(
desc=f"Uploading {src}",
total=file_size,
unit="B",
unit_scale=True,
unit_divisor=1024,
) as t:
reader_wrapper = CallbackIOWrapper(t.update, fd, "read")
response = requests.put(
url=signed_url,
data=reader_wrapper,
headers=put_headers,
)
response.raise_for_status()
return response
def _get_blob_name_from_file_content(src: str | Path) -> str:
"""Get a unique file name based on its content.
This file name is used as blob name when uploading the file to a bucket.
Args:
src (str | Path): Path to local file.
Returns:
The unique blob name.
"""
src = Path(src)
file_hash = hashlib.md5()
with src.open("rb") as fd:
fd.seek(0)
while chunk := fd.read(8192):
file_hash.update(chunk)
hexcode = file_hash.hexdigest()[:8]
# Take file name from src first as maybe a path with directories
# Then only extract the stem. There maybe more than 1 extension
# example: data/templates/custom_msa.a3m.pqt
file_stem = src.name.split(".")[0]
suffix = "".join(src.suffixes)
return f"{file_stem}_{hexcode}{suffix}"
def _copy_and_zip_files(
file_list: list[Path],
temp_dir: tempfile.TemporaryDirectory,
zip_name: str = "files.zip",
):
"""
Copies a list of files to a temporary directory and zips them into one archive
with the highest compression level.
Args:
file_list (list): List of file paths to be copied and zipped.
temp_dir: (TemporaryDirectory): Path to the temporary directory.
zip_name (str): Name of the resulting zip file.
Returns:
str: Path to the created zip file.
"""
to_zip = []
for file_path in file_list:
if file_path.is_file():
blob_name = _get_blob_name_from_file_content(src=file_path)
dest_file = os.path.join(temp_dir, blob_name)
shutil.copy(file_path, dest_file)
to_zip.append(dest_file)
else:
print(f"Warning: {file_path} does not exist or is not a file.")
zip_path = os.path.join(temp_dir, zip_name)
with ZipFile(zip_path, "w", compression=ZIP_DEFLATED, compresslevel=7) as zipf:
for file_name in to_zip:
zipf.write(file_name, arcname=Path(file_name).name)
return zip_path
def _get_blob_zip_name(file_type: str):
timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
return f"{file_type}_files_{timestamp}.zip"
def upload_custom_files(
headers: dict[str, str],
paths: list[Path],
file_type: CustomFileType,
) -> dict[str, str]:
"""Upload custom files to the GCS bucket.
Args:
headers (dict[str, str]): HTTP request headers.
paths (list[Path]): List of custom template files path.
file_type (CustomFileType): Type of file to upload.
Raises:
typer.Exit: If an error occurs during the API call.
Returns:
dict[str, str]: Mapping of local filenames to GCS paths.
"""
url = API_URL + "getUploadSignedURL"
paths = set(paths)
print(f"Uploading {len(paths)}: {tuple(str(p) for p in paths)}.")
blobs = [_get_blob_name_from_file_content(src=file) for file in paths]
# Zip files and upload archive
blob_zip = _get_blob_zip_name(file_type.value)
with tempfile.TemporaryDirectory() as temp_dir:
zip_path = _copy_and_zip_files(
file_list=paths,
temp_dir=temp_dir,
zip_name=blob_zip,
)
url_response = requests.get(
url,
params={
"blob_name": blob_zip,
"file_type": file_type.value,
},
headers=headers,
timeout=REQUEST_TIMEOUT,
)
if not url_response.ok:
print(f"Error while generating signed URL: {url_response.content.decode()}")
raise typer.Exit(code=1)
json_response = url_response.json()
signed_url = json_response["signed_url"]
upload_response = _upload_file_to_signed_url(
signed_url=signed_url, src=zip_path, headers=headers
)
if not upload_response.ok:
print(f"Error while uploading {zip_path}.")
raise typer.Exit(code=1)
# Unzip in dest bucket
unzip_response = requests.post(
API_URL + "unzipFileInBucket",
params={
"zip_file_path": json_response["destination_file"],
},
headers=headers,
timeout=REQUEST_TIMEOUT,
)
if not unzip_response.ok:
print(f"Error while unzip custom files: {unzip_response.content.decode()}")
raise typer.Exit(code=1)
local_to_gcs = {
str(file): f"{json_response['destination_bucket']}/{blob_name}"
for file, blob_name in zip(paths, blobs)
}
print("Custom files successfully uploaded.")
return local_to_gcs
|