print progress when installing torch

add PIP_INSTALLER_LOCATION env var to install pip if it's not installed
remove accidental call to accelerate when venv is disabled
add another env var to skip venv - SKIP_VENV
This commit is contained in:
AUTOMATIC 2023-01-24 21:13:05 +03:00
parent 5228ec8bda
commit 93fad28a97
2 changed files with 22 additions and 6 deletions

View File

@ -48,10 +48,19 @@ def extract_opt(args, name):
return args, is_present, opt
def run(command, desc=None, errdesc=None, custom_env=None):
def run(command, desc=None, errdesc=None, custom_env=None, live=False):
if desc is not None:
print(desc)
if live:
result = subprocess.run(command, shell=True, env=os.environ if custom_env is None else custom_env)
if result.returncode != 0:
raise RuntimeError(f"""{errdesc or 'Error running command'}.
Command: {command}
Error code: {result.returncode}""")
return ""
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, env=os.environ if custom_env is None else custom_env)
if result.returncode != 0:
@ -179,6 +188,8 @@ def run_extensions_installers(settings_file):
def prepare_environment():
global skip_install
pip_installer_location = os.environ.get('PIP_INSTALLER_LOCATION', None)
torch_command = os.environ.get('TORCH_COMMAND', "pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117")
requirements_file = os.environ.get('REQS_FILE', "requirements_versions.txt")
commandline_args = os.environ.get('COMMANDLINE_ARGS', "")
@ -219,9 +230,12 @@ def prepare_environment():
print(f"Python {sys.version}")
print(f"Commit hash: {commit}")
if pip_installer_location is not None and not is_installed("pip"):
run(f'"{python}" "{pip_installer_location}"', "Installing pip", "Couldn't install pip")
if reinstall_torch or not is_installed("torch") or not is_installed("torchvision"):
run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch")
run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch", live=True)
if not skip_torch_cuda_test:
run_python("import torch; assert torch.cuda.is_available(), 'Torch is not able to use GPU; add --skip-torch-cuda-test to COMMANDLINE_ARGS variable to disable this check'")

View File

@ -3,6 +3,7 @@
if not defined PYTHON (set PYTHON=python)
if not defined VENV_DIR (set "VENV_DIR=%~dp0%venv")
set ERROR_REPORTING=FALSE
mkdir tmp 2>NUL
@ -14,6 +15,7 @@ goto :show_stdout_stderr
:start_venv
if ["%VENV_DIR%"] == ["-"] goto :skip_venv
if ["%SKIP_VENV%"] == ["1"] goto :skip_venv
dir "%VENV_DIR%\Scripts\Python.exe" >tmp/stdout.txt 2>tmp/stderr.txt
if %ERRORLEVEL% == 0 goto :activate_venv
@ -28,13 +30,13 @@ goto :show_stdout_stderr
:activate_venv
set PYTHON="%VENV_DIR%\Scripts\Python.exe"
echo venv %PYTHON%
:skip_venv
if [%ACCELERATE%] == ["True"] goto :accelerate
goto :launch
:skip_venv
:accelerate
echo "Checking for accelerate"
echo Checking for accelerate
set ACCELERATE="%VENV_DIR%\Scripts\accelerate.exe"
if EXIST %ACCELERATE% goto :accelerate_launch