Skip to content

Commit 7fc9b6b

Browse files
namannandannamannandan
authored
Enable venv to inherit site packages from base python environment (#2946)
* Enable inheriting site-packages from base python env to venv * Use a .pth file instead of sitecustomize.py * Add support to inherit user site packages if enabled * Run custom dependencies test at the end * Change info logs to debug logs and update tests * Revert test filename to original name * Update venv creation log level to info
1 parent e4735f1 commit 7fc9b6b

File tree

3 files changed

+113
-45
lines changed

3 files changed

+113
-45
lines changed

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

+52-5
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,6 @@ private void setupModelVenv(Model model)
221221
commandParts.add("-m");
222222
commandParts.add("venv");
223223
commandParts.add("--clear");
224-
commandParts.add("--system-site-packages");
225224
commandParts.add(venvPath.toString());
226225

227226
ProcessBuilder processBuilder = new ProcessBuilder(commandParts);
@@ -273,6 +272,57 @@ private void setupModelVenv(Model model)
273272
throw new ModelException(
274273
"Virtual environment creation failed for model " + model.getModelName());
275274
}
275+
276+
// Inherit site-packages directories from the current environment torchserve is running in
277+
// to the newly created virtual environment
278+
commandParts.clear();
279+
commandParts.add(configManager.getPythonExecutable());
280+
commandParts.add(
281+
Paths.get(
282+
configManager.getModelServerHome(),
283+
"ts",
284+
"utils",
285+
"inherit_site_packages.py")
286+
.toAbsolutePath()
287+
.toString());
288+
commandParts.add(venvPath.toString());
289+
290+
processBuilder = new ProcessBuilder(commandParts);
291+
processBuilder.directory(venvPath.getParentFile());
292+
environment = processBuilder.environment();
293+
envp =
294+
EnvironmentUtils.getEnvString(
295+
configManager.getModelServerHome(),
296+
model.getModelDir().getAbsolutePath(),
297+
null);
298+
for (String envVar : envp) {
299+
String[] parts = envVar.split("=", 2);
300+
if (parts.length == 2) {
301+
environment.put(parts[0], parts[1]);
302+
}
303+
}
304+
processBuilder.redirectErrorStream(true);
305+
process = processBuilder.start();
306+
exitCode = process.waitFor();
307+
brdr = new BufferedReader(new InputStreamReader(process.getInputStream()));
308+
outputString.setLength(0);
309+
while ((line = brdr.readLine()) != null) {
310+
outputString.append(line + "\n");
311+
}
312+
313+
if (exitCode == 0) {
314+
logger.debug(
315+
"Inherited site-packages directories to venv {}:\n{}",
316+
venvPath.toString(),
317+
outputString.toString());
318+
} else {
319+
logger.error(
320+
"Failed to inherit site-packages directories to venv {}:\n{}",
321+
venvPath.toString(),
322+
outputString.toString());
323+
throw new ModelException(
324+
"Failed to inherit site-packages directories to venv " + venvPath.toString());
325+
}
276326
}
277327

278328
private void setupModelDependencies(Model model)
@@ -360,10 +410,7 @@ private void setupModelDependencies(Model model)
360410
}
361411

362412
if (exitCode == 0) {
363-
logger.info(
364-
"Installed custom pip packages for model {}:\n{}",
365-
model.getModelName(),
366-
outputString.toString());
413+
logger.info("Installed custom pip packages for model {}", model.getModelName());
367414
} else {
368415
logger.error(
369416
"Custom pip package installation failed for model {}:\n{}",

‎test/pytest/test_model_custom_dependencies.py

+16-40
Original file line numberDiff line numberDiff line change
@@ -141,9 +141,7 @@ def register_model_and_make_inference_request(expect_model_load_failure=False):
141141

142142

143143
def test_install_dependencies_to_target_directory_with_requirements():
144-
# Torchserve cleanup
145-
test_utils.stop_torchserve()
146-
test_utils.delete_all_snapshots()
144+
test_utils.torchserve_cleanup()
147145

148146
try:
149147
generate_model_archive(use_requirements=True, use_venv=False)
@@ -162,14 +160,11 @@ def test_install_dependencies_to_target_directory_with_requirements():
162160
)
163161
register_model_and_make_inference_request(expect_model_load_failure=False)
164162
finally:
165-
test_utils.stop_torchserve()
166-
test_utils.delete_all_snapshots()
163+
test_utils.torchserve_cleanup()
167164

168165

169166
def test_install_dependencies_to_target_directory_without_requirements():
170-
# Torchserve cleanup
171-
test_utils.stop_torchserve()
172-
test_utils.delete_all_snapshots()
167+
test_utils.torchserve_cleanup()
173168

174169
try:
175170
generate_model_archive(use_requirements=False, use_venv=False)
@@ -188,14 +183,11 @@ def test_install_dependencies_to_target_directory_without_requirements():
188183
)
189184
register_model_and_make_inference_request(expect_model_load_failure=True)
190185
finally:
191-
test_utils.stop_torchserve()
192-
test_utils.delete_all_snapshots()
186+
test_utils.torchserve_cleanup()
193187

194188

195189
def test_disable_install_dependencies_to_target_directory_with_requirements():
196-
# Torchserve cleanup
197-
test_utils.stop_torchserve()
198-
test_utils.delete_all_snapshots()
190+
test_utils.torchserve_cleanup()
199191

200192
try:
201193
generate_model_archive(use_requirements=True, use_venv=False)
@@ -207,14 +199,11 @@ def test_disable_install_dependencies_to_target_directory_with_requirements():
207199
)
208200
register_model_and_make_inference_request(expect_model_load_failure=True)
209201
finally:
210-
test_utils.stop_torchserve()
211-
test_utils.delete_all_snapshots()
202+
test_utils.torchserve_cleanup()
212203

213204

214205
def test_disable_install_dependencies_to_target_directory_without_requirements():
215-
# Torchserve cleanup
216-
test_utils.stop_torchserve()
217-
test_utils.delete_all_snapshots()
206+
test_utils.torchserve_cleanup()
218207

219208
try:
220209
generate_model_archive(use_requirements=False, use_venv=False)
@@ -226,14 +215,11 @@ def test_disable_install_dependencies_to_target_directory_without_requirements()
226215
)
227216
register_model_and_make_inference_request(expect_model_load_failure=True)
228217
finally:
229-
test_utils.stop_torchserve()
230-
test_utils.delete_all_snapshots()
218+
test_utils.torchserve_cleanup()
231219

232220

233221
def test_install_dependencies_to_venv_with_requirements():
234-
# Torchserve cleanup
235-
test_utils.stop_torchserve()
236-
test_utils.delete_all_snapshots()
222+
test_utils.torchserve_cleanup()
237223

238224
try:
239225
generate_model_archive(use_requirements=True, use_venv=True)
@@ -252,14 +238,11 @@ def test_install_dependencies_to_venv_with_requirements():
252238
)
253239
register_model_and_make_inference_request(expect_model_load_failure=False)
254240
finally:
255-
test_utils.stop_torchserve()
256-
test_utils.delete_all_snapshots()
241+
test_utils.torchserve_cleanup()
257242

258243

259244
def test_install_dependencies_to_venv_without_requirements():
260-
# Torchserve cleanup
261-
test_utils.stop_torchserve()
262-
test_utils.delete_all_snapshots()
245+
test_utils.torchserve_cleanup()
263246

264247
try:
265248
generate_model_archive(use_requirements=False, use_venv=True)
@@ -278,14 +261,11 @@ def test_install_dependencies_to_venv_without_requirements():
278261
)
279262
register_model_and_make_inference_request(expect_model_load_failure=True)
280263
finally:
281-
test_utils.stop_torchserve()
282-
test_utils.delete_all_snapshots()
264+
test_utils.torchserve_cleanup()
283265

284266

285267
def test_disable_install_dependencies_to_venv_with_requirements():
286-
# Torchserve cleanup
287-
test_utils.stop_torchserve()
288-
test_utils.delete_all_snapshots()
268+
test_utils.torchserve_cleanup()
289269

290270
try:
291271
generate_model_archive(use_requirements=True, use_venv=True)
@@ -297,14 +277,11 @@ def test_disable_install_dependencies_to_venv_with_requirements():
297277
)
298278
register_model_and_make_inference_request(expect_model_load_failure=True)
299279
finally:
300-
test_utils.stop_torchserve()
301-
test_utils.delete_all_snapshots()
280+
test_utils.torchserve_cleanup()
302281

303282

304283
def test_disable_install_dependencies_to_venv_without_requirements():
305-
# Torchserve cleanup
306-
test_utils.stop_torchserve()
307-
test_utils.delete_all_snapshots()
284+
test_utils.torchserve_cleanup()
308285

309286
try:
310287
generate_model_archive(use_requirements=False, use_venv=True)
@@ -316,5 +293,4 @@ def test_disable_install_dependencies_to_venv_without_requirements():
316293
)
317294
register_model_and_make_inference_request(expect_model_load_failure=True)
318295
finally:
319-
test_utils.stop_torchserve()
320-
test_utils.delete_all_snapshots()
296+
test_utils.torchserve_cleanup()

‎ts/utils/inherit_site_packages.py

+45
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# This script enables the virtual environment that is passed in as an argument to inherit
2+
# site-packages directories of the environment from which this script is run
3+
4+
import glob
5+
import os
6+
import site
7+
import sys
8+
9+
10+
def inherit_site_packages(venv_path):
11+
# Identify target venv site-packages directory
12+
target_venv_glob_matches = glob.glob(
13+
os.path.join(
14+
venv_path,
15+
"lib",
16+
f"python{sys.version_info[0]}.{sys.version_info[1]}",
17+
"site-packages",
18+
)
19+
)
20+
assert (
21+
len(target_venv_glob_matches) == 1
22+
), f"{__file__} expected to find one supported python version in venv {venv_path} but found: {target_venv_glob_matches}"
23+
24+
# Create a .pth file with site-packages directories to inherit, in the target venv site-packages directory
25+
# Ref: https://docs.python.org/3/library/site.html#module-site
26+
with open(
27+
os.path.join(target_venv_glob_matches[0], "inherited-site-packages.pth"), "w"
28+
) as f:
29+
if site.ENABLE_USER_SITE:
30+
user_site_packages_dir = site.getusersitepackages()
31+
if os.path.exists(user_site_packages_dir):
32+
f.write(f"{user_site_packages_dir}\n")
33+
print(user_site_packages_dir)
34+
35+
for global_site_packages_dir in site.getsitepackages():
36+
if os.path.exists(global_site_packages_dir):
37+
f.write(f"{global_site_packages_dir}\n")
38+
print(global_site_packages_dir)
39+
40+
41+
if __name__ == "__main__":
42+
assert (
43+
len(sys.argv) == 2
44+
), f"{__file__} expects one argument: path to venv that should inherit site-packages of the current environment but got {sys.argv}"
45+
inherit_site_packages(sys.argv[1])

0 commit comments

Comments
 (0)