|
| 1 | +import urllib.request |
| 2 | +import tempfile |
| 3 | +import os |
| 4 | +import uuid |
| 5 | +import shutil |
| 6 | +import glob |
| 7 | +import yaml |
| 8 | +import hashlib |
| 9 | + |
| 10 | +from zipfile import ZipFile |
| 11 | +from sys import platform |
| 12 | +from typing import Tuple, Optional, Dict, Any |
| 13 | + |
| 14 | +from mlagents_envs.logging_util import get_logger |
| 15 | + |
| 16 | +logger = get_logger(__name__) |
| 17 | + |
| 18 | +# The default logical block size is 8192 bytes (8 KB) for UFS file systems. |
| 19 | +BLOCK_SIZE = 8192 |
| 20 | + |
| 21 | + |
| 22 | +def get_local_binary_path(name: str, url: str) -> str: |
| 23 | + """ |
| 24 | + Returns the path to the executable previously downloaded with the name argument. If |
| 25 | + None is found, the executable at the url argument will be downloaded and stored |
| 26 | + under name for future uses. |
| 27 | + :param name: The name that will be given to the folder containing the extracted data |
| 28 | + :param url: The URL of the zip file |
| 29 | + """ |
| 30 | + NUMBER_ATTEMPTS = 5 |
| 31 | + path = get_local_binary_path_if_exists(name, url) |
| 32 | + if path is None: |
| 33 | + logger.debug( |
| 34 | + f"Local environment {name} not found, downloading environment from {url}" |
| 35 | + ) |
| 36 | + for attempt in range(NUMBER_ATTEMPTS): # Perform 5 attempts at downloading the file |
| 37 | + if path is not None: |
| 38 | + break |
| 39 | + try: |
| 40 | + download_and_extract_zip(url, name) |
| 41 | + except IOError: |
| 42 | + logger.debug( |
| 43 | + f"Attempt {attempt + 1} / {NUMBER_ATTEMPTS} : Failed to download" |
| 44 | + ) |
| 45 | + path = get_local_binary_path_if_exists(name, url) |
| 46 | + |
| 47 | + if path is None: |
| 48 | + raise FileNotFoundError( |
| 49 | + f"Binary not found, make sure {url} is a valid url to " |
| 50 | + "a zip folder containing a valid Unity executable" |
| 51 | + ) |
| 52 | + return path |
| 53 | + |
| 54 | + |
| 55 | +def get_local_binary_path_if_exists(name: str, url: str) -> Optional[str]: |
| 56 | + """ |
| 57 | + Recursively searches for a Unity executable in the extracted files folders. This is |
| 58 | + platform dependent : It will only return a Unity executable compatible with the |
| 59 | + computer's OS. If no executable is found, None will be returned. |
| 60 | + :param name: The name/identifier of the executable |
| 61 | + :param url: The url the executable was downloaded from (for verification) |
| 62 | + """ |
| 63 | + _, bin_dir = get_tmp_dir() |
| 64 | + extension = None |
| 65 | + |
| 66 | + if platform == "linux" or platform == "linux2": |
| 67 | + extension = "*.x86_64" |
| 68 | + if platform == "darwin": |
| 69 | + extension = "*.app" |
| 70 | + if platform == "win32": |
| 71 | + extension = "*.exe" |
| 72 | + if extension is None: |
| 73 | + raise NotImplementedError("No extensions found for this platform.") |
| 74 | + url_hash = "-" + hashlib.md5(url.encode()).hexdigest() |
| 75 | + path = os.path.join(bin_dir, name + url_hash, "**", extension) |
| 76 | + candidates = glob.glob(path, recursive=True) |
| 77 | + if len(candidates) == 0: |
| 78 | + return None |
| 79 | + else: |
| 80 | + for c in candidates: |
| 81 | + # Unity sometimes produces another .exe file that we must filter out |
| 82 | + if "UnityCrashHandler64" not in c: |
| 83 | + return c |
| 84 | + return None |
| 85 | + |
| 86 | + |
| 87 | +def get_tmp_dir() -> Tuple[str, str]: |
| 88 | + """ |
| 89 | + Returns the path to the folder containing the downloaded zip files and the extracted |
| 90 | + binaries. If these folders do not exist, they will be created. |
| 91 | + :retrun: Tuple containing path to : (zip folder, extracted files folder) |
| 92 | + """ |
| 93 | + MLAGENTS = "ml-agents-binaries" |
| 94 | + TMP_FOLDER_NAME = "tmp" |
| 95 | + BINARY_FOLDER_NAME = "binaries" |
| 96 | + mla_directory = os.path.join(tempfile.gettempdir(), MLAGENTS) |
| 97 | + if not os.path.exists(mla_directory): |
| 98 | + os.makedirs(mla_directory) |
| 99 | + os.chmod(mla_directory, 16877) |
| 100 | + zip_directory = os.path.join(tempfile.gettempdir(), MLAGENTS, TMP_FOLDER_NAME) |
| 101 | + if not os.path.exists(zip_directory): |
| 102 | + os.makedirs(zip_directory) |
| 103 | + os.chmod(zip_directory, 16877) |
| 104 | + bin_directory = os.path.join(tempfile.gettempdir(), MLAGENTS, BINARY_FOLDER_NAME) |
| 105 | + if not os.path.exists(bin_directory): |
| 106 | + os.makedirs(bin_directory) |
| 107 | + os.chmod(bin_directory, 16877) |
| 108 | + return (zip_directory, bin_directory) |
| 109 | + |
| 110 | + |
| 111 | +def download_and_extract_zip(url: str, name: str) -> None: |
| 112 | + """ |
| 113 | + Downloads a zip file under a URL, extracts its contents into a folder with the name |
| 114 | + argument and gives chmod 755 to all the files it contains. Files are downloaded and |
| 115 | + extracted into special folders in the temp folder of the machine. |
| 116 | + :param url: The URL of the zip file |
| 117 | + :param name: The name that will be given to the folder containing the extracted data |
| 118 | + """ |
| 119 | + zip_dir, bin_dir = get_tmp_dir() |
| 120 | + url_hash = "-" + hashlib.md5(url.encode()).hexdigest() |
| 121 | + binary_path = os.path.join(bin_dir, name + url_hash) |
| 122 | + if os.path.exists(binary_path): |
| 123 | + shutil.rmtree(binary_path) |
| 124 | + |
| 125 | + # Download zip |
| 126 | + try: |
| 127 | + request = urllib.request.urlopen(url, timeout=30) |
| 128 | + except urllib.error.HTTPError as e: # type: ignore |
| 129 | + e.msg += " " + url |
| 130 | + raise |
| 131 | + zip_size = int(request.headers["content-length"]) |
| 132 | + zip_file_path = os.path.join(zip_dir, str(uuid.uuid4()) + ".zip") |
| 133 | + with open(zip_file_path, "wb") as zip_file: |
| 134 | + downloaded = 0 |
| 135 | + while True: |
| 136 | + buffer = request.read(BLOCK_SIZE) |
| 137 | + if not buffer: |
| 138 | + # There is nothing more to read |
| 139 | + break |
| 140 | + downloaded += len(buffer) |
| 141 | + zip_file.write(buffer) |
| 142 | + downloaded_percent = downloaded / zip_size * 100 |
| 143 | + print_progress(f" Downloading {name}", downloaded_percent) |
| 144 | + print("") |
| 145 | + |
| 146 | + # Extraction |
| 147 | + with ZipFileWithProgress(zip_file_path, "r") as zip_ref: |
| 148 | + zip_ref.extract_zip(f" Extracting {name}", binary_path) # type: ignore |
| 149 | + print("") |
| 150 | + |
| 151 | + # Clean up zip |
| 152 | + print_progress(f" Cleaning up {name}", 0) |
| 153 | + os.remove(zip_file_path) |
| 154 | + |
| 155 | + # Give permission |
| 156 | + for f in glob.glob(binary_path + "/**/*", recursive=True): |
| 157 | + # 16877 is octal 40755, which denotes a directory with permissions 755 |
| 158 | + os.chmod(f, 16877) |
| 159 | + print_progress(f" Cleaning up {name}", 100) |
| 160 | + print("") |
| 161 | + |
| 162 | + |
| 163 | +def print_progress(prefix: str, percent: float) -> None: |
| 164 | + """ |
| 165 | + Displays a single progress bar in the terminal with value percent. |
| 166 | + :param prefix: The string that will precede the progress bar. |
| 167 | + :param percent: The percent progression of the bar (min is 0, max is 100) |
| 168 | + """ |
| 169 | + BAR_LEN = 20 |
| 170 | + percent = min(100, max(0, percent)) |
| 171 | + bar_progress = min(int(percent / 100 * BAR_LEN), BAR_LEN) |
| 172 | + bar = "|" + "\u2588" * bar_progress + " " * (BAR_LEN - bar_progress) + "|" |
| 173 | + str_percent = "%3.0f%%" % percent |
| 174 | + print(f"{prefix} : {bar} {str_percent} \r", end="", flush=True) |
| 175 | + |
| 176 | + |
| 177 | +def load_remote_manifest(url: str) -> Dict[str, Any]: |
| 178 | + """ |
| 179 | + Converts a remote yaml file into a Python dictionary |
| 180 | + """ |
| 181 | + tmp_dir, _ = get_tmp_dir() |
| 182 | + try: |
| 183 | + request = urllib.request.urlopen(url, timeout=30) |
| 184 | + except urllib.error.HTTPError as e: # type: ignore |
| 185 | + e.msg += " " + url |
| 186 | + raise |
| 187 | + manifest_path = os.path.join(tmp_dir, str(uuid.uuid4()) + ".yaml") |
| 188 | + with open(manifest_path, "wb") as manifest: |
| 189 | + while True: |
| 190 | + buffer = request.read(BLOCK_SIZE) |
| 191 | + if not buffer: |
| 192 | + # There is nothing more to read |
| 193 | + break |
| 194 | + manifest.write(buffer) |
| 195 | + try: |
| 196 | + result = load_local_manifest(manifest_path) |
| 197 | + finally: |
| 198 | + os.remove(manifest_path) |
| 199 | + return result |
| 200 | + |
| 201 | + |
| 202 | +def load_local_manifest(path: str) -> Dict[str, Any]: |
| 203 | + """ |
| 204 | + Converts a local yaml file into a Python dictionary |
| 205 | + """ |
| 206 | + with open(path) as data_file: |
| 207 | + return yaml.safe_load(data_file) |
| 208 | + |
| 209 | + |
| 210 | +class ZipFileWithProgress(ZipFile): |
| 211 | + """ |
| 212 | + This is a helper class inheriting from ZipFile that allows to display a progress |
| 213 | + bar while the files are being extracted. |
| 214 | + """ |
| 215 | + |
| 216 | + def extract_zip(self, prefix: str, path: str) -> None: |
| 217 | + members = self.namelist() |
| 218 | + path = os.fspath(path) |
| 219 | + total = len(members) |
| 220 | + n = 0 |
| 221 | + for zipinfo in members: |
| 222 | + self.extract(zipinfo, path, None) # type: ignore |
| 223 | + n += 1 |
| 224 | + print_progress(prefix, n / total * 100) |
0 commit comments