Set PYTHONHOME for inductor subprocesses using torch (#159382)

Summary:
This is needed for subprocesses that are trying to call back into torch
functionality, i.e. anything that's also setting `PYTHONPATH`.  There are more
`sys.executable` subprocesses in torch/ but it seems like they're fine.

Test Plan: Local inference runs.

Reviewed By: aorenste

Differential Revision: D79124705

Pull Request resolved: https://github.com/pytorch/pytorch/pull/159382
Approved by: https://github.com/aorenste

torch/_inductor/autotune_process.py

@@ -12,6 +12,7 @@ import queue
 import selectors
 import subprocess
 import sys
+import sysconfig
 import time
 import warnings
 from collections.abc import Iterable, Sequence
@@ -128,6 +129,8 @@ class TuningProcess:
             "PYTHONPATH": os.environ.get(
                 "TORCH_CUSTOM_PYTHONPATH", os.pathsep.join(sys.path)
             ),
+            # Need to set this for internal builds that bundle the runtime.
+            "PYTHONHOME": sysconfig.get_path("data"),
             # We shouldn't be using the Triton async compile subprocess pool,
             # but as a precaution set the env var that disables its creation.
             "TORCH_WARM_POOL": "0",

torch/_inductor/compile_worker/subproc_pool.py

@@ -8,6 +8,7 @@ import pickle
 import struct
 import subprocess
 import sys
+import sysconfig
 import threading
 import traceback
 import typing
@@ -158,6 +159,8 @@ class SubprocPool:
                 "PYTHONPATH": os.environ.get(
                     "TORCH_CUSTOM_PYTHONPATH", os.pathsep.join(sys.path)
                 ),
+                # Need to set this for internal builds that bundle the runtime.
+                "PYTHONHOME": sysconfig.get_path("data"),
                 # Safeguard against creating a SubprocPool in the subprocess.
                 "TORCH_WARM_POOL": "0",
                 # Some internal usages need a modified LD_LIBRARY_PATH.

torch/_inductor/cpu_vec_isa.py

@@ -6,6 +6,7 @@ import platform
 import re
 import subprocess
 import sys
+import sysconfig
 import warnings
 from typing import Any, Callable, Union
 
@@ -133,9 +134,12 @@ cdll.LoadLibrary("__lib_path__")
                     stderr=subprocess.DEVNULL,
                     env={
                         **os.environ,
+                        # We need to set the PYTHONPATH so the subprocess can find torch.
                         "PYTHONPATH": os.environ.get(
                             "TORCH_CUSTOM_PYTHONPATH", os.pathsep.join(sys.path)
                         ),
+                        # Need to set this for internal builds that bundle the runtime.
+                        "PYTHONHOME": sysconfig.get_path("data"),
                     },
                 )
             except Exception: