Enable venv to inherit site packages from base python environment

[namannandan]

Description

The current implementation to support installation of model custom dependencies and running model workers in a python virtual environment(#2910) works when torchserve is installed to and started from system site-packages but fails when torchserve itself is installed to and started from a virtual environment.

MODEL_LOG - Traceback (most recent call last):
MODEL_LOG -   File "/home/venv/lib/python3.9/site-packages/ts/model_service_worker.py", line 15, in <module>
MODEL_LOG -     from ts.arg_parser import ArgParser
MODEL_LOG - ModuleNotFoundError: No module named 'ts'

This is because the newly created virtual environment does not inherit site-packages from the current virtual environment where torchserve is installed, which is expected behavior.

Just using the --system-site-packages flag when creating the virtual environment is not a complete solution because torchserve may have not been installed to the system site directory as shown above.

serve/frontend/server/src/main/java/org/pytorch/serve/wlm/ModelManager.java

Line 224 in bef3b63

commandParts.add("--system-site-packages");

A more generic solution would be is to identify site-packages directories setup in the python environment in which torchserve is started and inherit those into the new virtual environment.

For example:
In the base python environment(torchserve running in a virtual-environment):

$ python -m site
sys.path = [
    '/serve',
    '/usr/lib/python39.zip',
    '/usr/lib/python3.9',
    '/usr/lib/python3.9/lib-dynload',
    '/home/venv/lib/python3.9/site-packages',
]
USER_BASE: '/root/.local' (doesn't exist)
USER_SITE: '/root/.local/lib/python3.9/site-packages' (doesn't exist)
ENABLE_USER_SITE: False

In the new virtual-environment created to install model custom dependencies and run model workers:

$ /home/model-server/tmp/models/1b1c30114f0b451087e2123af09eea2e/venv/bin/python -m site
sys.path = [
    '/serve',
    '/usr/lib/python39.zip',
    '/usr/lib/python3.9',
    '/usr/lib/python3.9/lib-dynload',
    '/home/model-server/tmp/models/1b1c30114f0b451087e2123af09eea2e/venv/lib/python3.9/site-packages',
    '/home/venv/lib/python3.9/site-packages',
]
USER_BASE: '/root/.local' (doesn't exist)
USER_SITE: '/root/.local/lib/python3.9/site-packages' (doesn't exist)
ENABLE_USER_SITE: False

As can be seen above, /home/venv/lib/python3.9/site-packages is added to sys.path in the newly created environment, which is inherited from the base python environment. Also, the packages in the local venv site-packages /home/model-server/tmp/models/1b1c30114f0b451087e2123af09eea2e/venv/lib/python3.9/site-packages take precedence over the inherited site-packages directory /home/venv/lib/python3.9/site-packages.

Fixes #2942

Type of change

• Bug fix (non-breaking change which fixes an issue)

Feature/Issue validation/testing

• CI
• Docker regression: https://github.com/pytorch/serve/actions/runs/7956485287
• Manual test with installing and starting torchserve from inside a virtual environment

$ python -m venv venv_test
$ source venv_test/bin/activate
$ python ts_scripts/install_dependencies.py --environment=dev
$ python ts_scripts/install_from_src.py

$ torchserve --start --model-store model_store --models mnist=mnist.mar --ts-config custom_config.properties
.....
.....
2024-02-14T14:20:59,414 [INFO ] main org.pytorch.serve.wlm.ModelManager - Created virtual environment for model mnist: /var/folders/l4/hfy8rtpx0755sys3m8m8c48w0000gs/T/models/772f2d40c2084ff9a9ec8d1a2e27b53f/venv
2024-02-14T14:20:59,460 [INFO ] main org.pytorch.serve.wlm.ModelManager - Inherited site-packages directories to venv /var/folders/l4/hfy8rtpx0755sys3m8m8c48w0000gs/T/models/772f2d40c2084ff9a9ec8d1a2e27b53f/venv:
/Volumes/workplace/pytorch/serve/venv_test/lib/python3.8/site-packages

2024-02-14T14:21:02,004 [INFO ] main org.pytorch.serve.wlm.ModelManager - Installed custom pip packages for model mnist:
Requirement already satisfied: matplotlib in /Volumes/workplace/pytorch/serve/venv_test/lib/python3.8/site-packages (from -r /var/folders/l4/hfy8rtpx0755sys3m8m8c48w0000gs/T/models/772f2d40c2084ff9a9ec8d1a2e27b53f/requirements.txt (line 1)) (3.7.4)
.....
.....

$ python -m site
sys.path = [
    '/Volumes/workplace/pytorch/serve',
    '/Users/namannan/.pyenv/versions/3.8.13/lib/python38.zip',
    '/Users/namannan/.pyenv/versions/3.8.13/lib/python3.8',
    '/Users/namannan/.pyenv/versions/3.8.13/lib/python3.8/lib-dynload',
    '/Volumes/workplace/pytorch/serve/venv_test/lib/python3.8/site-packages',
]
USER_BASE: '/Users/namannan/.local' (exists)
USER_SITE: '/Users/namannan/.local/lib/python3.8/site-packages' (doesn't exist)
ENABLE_USER_SITE: False

$ /var/folders/l4/hfy8rtpx0755sys3m8m8c48w0000gs/T/models/772f2d40c2084ff9a9ec8d1a2e27b53f/venv/bin/python -m site
sys.path = [
    '/Volumes/workplace/pytorch/serve',
    '/Users/namannan/.pyenv/versions/3.8.13/lib/python38.zip',
    '/Users/namannan/.pyenv/versions/3.8.13/lib/python3.8',
    '/Users/namannan/.pyenv/versions/3.8.13/lib/python3.8/lib-dynload',
    '/var/folders/l4/hfy8rtpx0755sys3m8m8c48w0000gs/T/models/772f2d40c2084ff9a9ec8d1a2e27b53f/venv/lib/python3.8/site-packages',
    '/Volumes/workplace/pytorch/serve/venv_test/lib/python3.8/site-packages',
]
USER_BASE: '/Users/namannan/.local' (exists)
USER_SITE: '/Users/namannan/.local/lib/python3.8/site-packages' (doesn't exist)
ENABLE_USER_SITE: False

• Manual test with installing and starting torchserve from inside a conda environment

$ conda create -n venv_test python=3.9
$ conda activate venv_test
$ python ts_scripts/install_dependencies.py --environment=dev
$ python ts_scripts/install_from_src.py

$ torchserve --start --model-store model_store --models mnist=mnist.mar --ts-config custom_config.properties
.....
.....
2024-02-14T14:35:43,520 [INFO ] main org.pytorch.serve.wlm.ModelManager - Created virtual environment for model mnist: /var/folders/l4/hfy8rtpx0755sys3m8m8c48w0000gs/T/models/44513c42e2da4b569902ba454a8fd50b/venv
2024-02-14T14:35:43,566 [INFO ] main org.pytorch.serve.wlm.ModelManager - Inherited site-packages directories to venv /var/folders/l4/hfy8rtpx0755sys3m8m8c48w0000gs/T/models/44513c42e2da4b569902ba454a8fd50b/venv:
/Users/namannan/anaconda3/envs/venv_test/lib/python3.9/site-packages

2024-02-14T14:35:48,893 [INFO ] main org.pytorch.serve.wlm.ModelManager - Installed custom pip packages for model mnist:
Collecting matplotlib
.....
.....

$ python -m site
sys.path = [
    '/Volumes/workplace/pytorch/serve',
    '/Users/namannan/anaconda3/envs/venv_test/lib/python39.zip',
    '/Users/namannan/anaconda3/envs/venv_test/lib/python3.9',
    '/Users/namannan/anaconda3/envs/venv_test/lib/python3.9/lib-dynload',
    '/Users/namannan/anaconda3/envs/venv_test/lib/python3.9/site-packages',
]
USER_BASE: '/Users/namannan/.local' (exists)
USER_SITE: '/Users/namannan/.local/lib/python3.9/site-packages' (doesn't exist)
ENABLE_USER_SITE: True

$ /var/folders/l4/hfy8rtpx0755sys3m8m8c48w0000gs/T/models/9836a469270048148459dae28c44c6c7/venv/bin/python -m site
sys.path = [
    '/Volumes/workplace/pytorch/serve',
    '/Users/namannan/anaconda3/envs/venv_test/lib/python39.zip',
    '/Users/namannan/anaconda3/envs/venv_test/lib/python3.9',
    '/Users/namannan/anaconda3/envs/venv_test/lib/python3.9/lib-dynload',
    '/var/folders/l4/hfy8rtpx0755sys3m8m8c48w0000gs/T/models/9836a469270048148459dae28c44c6c7/venv/lib/python3.9/site-packages',
    '/Users/namannan/anaconda3/envs/venv_test/lib/python3.9/site-packages',
]
USER_BASE: '/Users/namannan/.local' (exists)
USER_SITE: '/Users/namannan/.local/lib/python3.9/site-packages' (doesn't exist)
ENABLE_USER_SITE: False

[agunapal]

LGTM