|
3 | 3 | import subprocess |
4 | 4 | from functools import lru_cache |
5 | 5 | from pathlib import Path |
6 | | -from typing import Any, cast, Optional |
| 6 | +from typing import Any, cast, Optional, Union |
7 | 7 |
|
8 | 8 | import requests |
9 | 9 |
|
@@ -120,6 +120,23 @@ def ok_changed_file(file: str) -> bool: |
120 | 120 | def check_changed_files(sha: str) -> bool: |
121 | 121 | # Return true if all the changed files are in the list of allowed files to |
122 | 122 | # be changed to reuse the old whl |
| 123 | + |
| 124 | + # Removing any files is not allowed since rysnc will not remove files |
| 125 | + removed_files = ( |
| 126 | + subprocess.check_output( |
| 127 | + ["git", "diff", "--name-only", sha, "HEAD", "--diff-filter=D"], |
| 128 | + text=True, |
| 129 | + stderr=subprocess.DEVNULL, |
| 130 | + ) |
| 131 | + .strip() |
| 132 | + .split() |
| 133 | + ) |
| 134 | + if removed_files: |
| 135 | + print( |
| 136 | + f"Removed files between {sha} and HEAD: {removed_files}, cannot reuse old whl" |
| 137 | + ) |
| 138 | + return False |
| 139 | + |
123 | 140 | changed_files = ( |
124 | 141 | subprocess.check_output( |
125 | 142 | ["git", "diff", "--name-only", sha, "HEAD"], |
@@ -179,38 +196,83 @@ def unzip_artifact_and_replace_files() -> None: |
179 | 196 | ) |
180 | 197 | os.remove("artifacts.zip") |
181 | 198 |
|
| 199 | + head_sha = get_head_sha() |
| 200 | + |
182 | 201 | # Rename wheel into zip |
183 | 202 | wheel_path = Path("artifacts/dist").glob("*.whl") |
184 | 203 | for path in wheel_path: |
185 | | - new_path = path.with_suffix(".zip") |
186 | | - os.rename(path, new_path) |
187 | | - print(f"Renamed {path} to {new_path}") |
188 | | - print(new_path.stem) |
| 204 | + # Should be of the form torch-2.0.0+git1234567-cp37-etc.whl |
| 205 | + # Should usually be the merge base sha but for the ones that didn't do |
| 206 | + # the replacement, it won't be. Can probably change it to just be merge |
| 207 | + # base later |
| 208 | + old_version = f"+git{path.stem.split('+')[1].split('-')[0][3:]}" |
| 209 | + new_version = f"+git{head_sha[:7]}" |
| 210 | + |
| 211 | + def rename_to_new_version(file: Union[str, Path]) -> None: |
| 212 | + # Rename file with old_version to new_version |
| 213 | + subprocess.check_output( |
| 214 | + ["mv", file, str(file).replace(old_version, new_version)] |
| 215 | + ) |
| 216 | + |
| 217 | + def change_content_to_new_version(file: Union[str, Path]) -> None: |
| 218 | + # Check if is a file |
| 219 | + if os.path.isdir(file): |
| 220 | + return |
| 221 | + # Replace the old version in the file with the new version |
| 222 | + with open(file) as f: |
| 223 | + content = f.read() |
| 224 | + content = content.replace(old_version, new_version) |
| 225 | + with open(file, "w") as f: |
| 226 | + f.write(content) |
| 227 | + |
| 228 | + zip_path = path.with_suffix(".zip") |
| 229 | + os.rename(path, zip_path) |
| 230 | + old_stem = zip_path.stem |
189 | 231 | # Unzip the wheel |
190 | 232 | subprocess.check_output( |
191 | | - ["unzip", "-o", new_path, "-d", f"artifacts/dist/{new_path.stem}"], |
| 233 | + ["unzip", "-o", zip_path, "-d", f"artifacts/dist/{old_stem}"], |
192 | 234 | ) |
| 235 | + |
| 236 | + # Remove the old wheel (which is now a zip file) |
| 237 | + os.remove(zip_path) |
| 238 | + |
193 | 239 | # Copy python files into the artifact |
194 | 240 | subprocess.check_output( |
195 | | - ["rsync", "-avz", "torch", f"artifacts/dist/{new_path.stem}"], |
| 241 | + ["rsync", "-avz", "torch", f"artifacts/dist/{old_stem}"], |
196 | 242 | ) |
197 | 243 |
|
| 244 | + change_content_to_new_version(f"artifacts/dist/{old_stem}/torch/version.py") |
| 245 | + |
| 246 | + for file in Path(f"artifacts/dist/{old_stem}").glob( |
| 247 | + "*.dist-info/**", |
| 248 | + ): |
| 249 | + change_content_to_new_version(file) |
| 250 | + |
| 251 | + rename_to_new_version(f"artifacts/dist/{old_stem}") |
| 252 | + new_stem = old_stem.replace(old_version, new_version) |
| 253 | + |
| 254 | + for file in Path(f"artifacts/dist/{new_stem}").glob( |
| 255 | + "*.dist-info", |
| 256 | + ): |
| 257 | + rename_to_new_version(file) |
| 258 | + |
198 | 259 | # Zip the wheel back |
199 | 260 | subprocess.check_output( |
200 | | - ["zip", "-r", f"{new_path.stem}.zip", "."], |
201 | | - cwd=f"artifacts/dist/{new_path.stem}", |
| 261 | + ["zip", "-r", f"{new_stem}.zip", "."], |
| 262 | + cwd=f"artifacts/dist/{new_stem}", |
202 | 263 | ) |
| 264 | + |
203 | 265 | subprocess.check_output( |
204 | 266 | [ |
205 | 267 | "mv", |
206 | | - f"artifacts/dist/{new_path.stem}/{new_path.stem}.zip", |
207 | | - f"artifacts/dist/{new_path.stem}.whl", |
| 268 | + f"artifacts/dist/{new_stem}/{new_stem}.zip", |
| 269 | + f"artifacts/dist/{new_stem}.whl", |
208 | 270 | ], |
209 | 271 | ) |
210 | 272 |
|
211 | 273 | # Remove the extracted folder |
212 | 274 | subprocess.check_output( |
213 | | - ["rm", "-rf", f"artifacts/dist/{new_path.stem}"], |
| 275 | + ["rm", "-rf", f"artifacts/dist/{new_stem}"], |
214 | 276 | ) |
215 | 277 |
|
216 | 278 | # Rezip the artifact |
|
0 commit comments