[ds launcher] un-hijack PYTHONPATH

src/accelerate/commands/launch.py

@@ -45,6 +45,7 @@
 )
 from accelerate.utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
 from accelerate.utils.dataclasses import SageMakerDistributedType
+from accelerate.utils.launch import env_var_path_add
 
 
 if is_rich_available():
@@ -573,7 +574,7 @@ def deepspeed_launcher(args):
         warnings.warn('--fp16 flag is deprecated. Use "--mixed_precision fp16" instead.', DeprecationWarning)
         mixed_precision = "fp16"
 
-    current_env["PYTHONPATH"] = sys.executable
+    current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", sys.executable)
     current_env["MIXED_PRECISION"] = str(mixed_precision)
     current_env["USE_DEEPSPEED"] = "true"
     current_env["DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)

src/accelerate/utils/launch.py

@@ -52,6 +52,16 @@ def _filter_args(args):
     return new_args
 
 
+def env_var_path_add(env_var_name, path_to_add):
+    """
+    Extends a path-based environment variable's value with a new path and returns the updated value. It's up to the
+    caller to set it in os.environ.
+    """
+    paths = [p for p in os.environ.get(env_var_name, "").split(":") if len(p) > 0]
+    paths.append(str(path_to_add))
+    return ":".join(paths)
+
+
 class PrepareForLaunch:
     """
     Prepare a function that will launched in a distributed setup.