Subprocess output was processed line-by-line to find and extract the PyTorch URL, and the subprocess was terminated as soon as the URL was found.

This commit is contained in:
Paul Gauthier (aider) 2024-07-09 16:00:29 +01:00
parent 72f0686c66
commit 7db1c51919

View file

@ -29,19 +29,26 @@ cmd = [
"--no-deps", "--no-deps",
"--dest", "--dest",
"/dev/null", "/dev/null",
# "--no-cache-dir",
# "--ignore-installed",
"--index-url", "--index-url",
"https://download.pytorch.org/whl/cpu", "https://download.pytorch.org/whl/cpu",
] ]
result = subprocess.check_output(cmd, text=True) pytorch_url = None
print(result) try:
process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
url_match = re.search(r"Downloading (https://download\.pytorch\.org/[^\s]+\.whl)", result) for line in process.stdout:
print(line, end='') # Print each line of output
url_match = re.search(r"Downloading (https://download\.pytorch\.org/[^\s]+\.whl)", line)
if url_match: if url_match:
pytorch_url = url_match.group(1) pytorch_url = url_match.group(1)
print(f"PyTorch URL: {pytorch_url}") print(f"PyTorch URL: {pytorch_url}")
process.terminate() # Terminate the subprocess
break
process.wait() # Wait for the process to finish
except subprocess.CalledProcessError as e:
print(f"Error running pip download: {e}")
if pytorch_url:
requirements = [f"torch @ {pytorch_url}"] + requirements requirements = [f"torch @ {pytorch_url}"] + requirements
else: else:
print("PyTorch URL not found in the output") print("PyTorch URL not found in the output")