Revert "Keep the passed in working directory in sys.path after the in… (#27718)

…itial import of user code (#27485)"

This reverts commit 199c6b3.

## Summary & Motivation

## How I Tested These Changes

## Changelog

> Insert changelog entry or delete this section.

Author: prha

CHANGES.md

@@ -29,7 +29,6 @@
 - Fixed an issue where partition keys could be dropped when converting a list of partition keys for a `MultiPartitionsDefinition` to a `PartitionSubset`.
 - Fixed an issue where the "Reload definitions" button didn't work when using `dagster dev` on Windows, starting in the 1.9.10 release.
 - Fixed an issue where dagster could not be imported alongside some other libraries using gRPC with an 'api.proto' file.
-- Fixed an issue where assets using ProcessPoolExecutor or multiprocessing with a spawn context sometimes failed to work with a ModuleNotFoundError.
 - [ui] Fixed an issue where non-`None` default config fields weren't being displayed in the Launchpad view.
 - [ui] Fixed an issue with the search bar on the Asset partitions page incorrectly filtering partitions when combined with a status filter.
 - [ui] Fixed Asset page header display of long key values.

python_modules/dagster/dagster/_core/code_pointer.py

@@ -11,7 +11,7 @@
 from dagster._core.errors import DagsterImportError, DagsterInvariantViolationError
 from dagster._serdes import whitelist_for_serdes
 from dagster._seven import get_import_error_message, import_module_from_path
-from dagster._utils import hash_collection
+from dagster._utils import alter_sys_path, hash_collection
 
 
 class CodePointer(ABC):
@@ -74,11 +74,14 @@ def load_python_file(python_file: str, working_directory: Optional[str]) -> Modu
 
     module_name = os.path.splitext(os.path.basename(python_file))[0]
 
-    if working_directory and working_directory not in sys.path:
-        sys.path.insert(0, working_directory)
-
+    # Use the passed in working directory for local imports (sys.path[0] isn't
+    # consistently set in the different entry points that Dagster uses to import code)
+    script_path = sys.path[0]
     try:
-        return import_module_from_path(module_name, python_file)
+        with alter_sys_path(
+            to_add=([working_directory] if working_directory else []), to_remove=[script_path]
+        ):
+            return import_module_from_path(module_name, python_file)
     except ImportError as ie:
         python_file = os.path.abspath(os.path.expanduser(python_file))
 
@@ -113,35 +116,43 @@ def load_python_file(python_file: str, working_directory: Optional[str]) -> Modu
 def load_python_module(
     module_name: str,
     working_directory: Optional[str],
+    remove_from_path_fn: Optional[Callable[[], Sequence[str]]] = None,
 ) -> ModuleType:
     check.str_param(module_name, "module_name")
     check.opt_str_param(working_directory, "working_directory")
-
-    if working_directory and working_directory not in sys.path:
-        sys.path.insert(0, working_directory)
-
-    try:
-        return importlib.import_module(module_name)
-    except ImportError as ie:
-        msg = get_import_error_message(ie)
-        if working_directory:
-            abs_working_directory = os.path.abspath(os.path.expanduser(working_directory))
-            raise DagsterImportError(
-                f"Encountered ImportError: `{msg}` while importing module {module_name}. "
-                "Local modules were resolved using the working "
-                f"directory `{abs_working_directory}`. If another working directory should be "
-                "used, please explicitly specify the appropriate path using the `-d` or "
-                "`--working-directory` for CLI based targets or the `working_directory` "
-                "configuration option for workspace targets. "
-            ) from ie
-        else:
-            raise DagsterImportError(
-                f"Encountered ImportError: `{msg}` while importing module {module_name}. "
-                "If relying on the working directory to resolve modules, please "
-                "explicitly specify the appropriate path using the `-d` or "
-                "`--working-directory` for CLI based targets or the `working_directory` "
-                "configuration option for workspace targets. "
-            ) from ie
+    check.opt_callable_param(remove_from_path_fn, "remove_from_path_fn")
+
+    # Use the passed in working directory for local imports (sys.path[0] isn't
+    # consistently set in the different entry points that Dagster uses to import code)
+    remove_paths: list[str] = (
+        list(remove_from_path_fn()) if remove_from_path_fn else []
+    )  # hook for tests
+    remove_paths.insert(0, sys.path[0])  # remove the script path
+    with alter_sys_path(
+        to_add=([working_directory] if working_directory else []), to_remove=remove_paths
+    ):
+        try:
+            return importlib.import_module(module_name)
+        except ImportError as ie:
+            msg = get_import_error_message(ie)
+            if working_directory:
+                abs_working_directory = os.path.abspath(os.path.expanduser(working_directory))
+                raise DagsterImportError(
+                    f"Encountered ImportError: `{msg}` while importing module {module_name}. "
+                    "Local modules were resolved using the working "
+                    f"directory `{abs_working_directory}`. If another working directory should be "
+                    "used, please explicitly specify the appropriate path using the `-d` or "
+                    "`--working-directory` for CLI based targets or the `working_directory` "
+                    "configuration option for workspace targets. "
+                ) from ie
+            else:
+                raise DagsterImportError(
+                    f"Encountered ImportError: `{msg}` while importing module {module_name}. "
+                    "If relying on the working directory to resolve modules, please "
+                    "explicitly specify the appropriate path using the `-d` or "
+                    "`--working-directory` for CLI based targets or the `working_directory` "
+                    "configuration option for workspace targets. "
+                ) from ie
 
 
 @whitelist_for_serdes

python_modules/dagster/dagster/_core/workspace/autodiscovery.py

@@ -1,7 +1,7 @@
 import inspect
 from collections.abc import Sequence
 from types import ModuleType
-from typing import NamedTuple, Optional, Union
+from typing import Callable, NamedTuple, Optional, Union
 
 from dagster import DagsterInvariantViolationError, GraphDefinition, RepositoryDefinition
 from dagster._core.code_pointer import load_python_file, load_python_module
@@ -28,19 +28,24 @@ def loadable_targets_from_python_file(
 def loadable_targets_from_python_module(
     module_name: str,
     working_directory: Optional[str],
+    remove_from_path_fn: Optional[Callable[[], Sequence[str]]] = None,
 ) -> Sequence[LoadableTarget]:
     module = load_python_module(
         module_name,
         working_directory=working_directory,
+        remove_from_path_fn=remove_from_path_fn,
     )
     return loadable_targets_from_loaded_module(module)
 
 
 def loadable_targets_from_python_package(
     package_name: str,
     working_directory: Optional[str],
+    remove_from_path_fn: Optional[Callable[[], Sequence[str]]] = None,
 ) -> Sequence[LoadableTarget]:
-    module = load_python_module(package_name, working_directory)
+    module = load_python_module(
+        package_name, working_directory, remove_from_path_fn=remove_from_path_fn
+    )
     return loadable_targets_from_loaded_module(module)
 
 

python_modules/dagster/dagster_tests/cli_tests/command_tests/asset_with_process_pool_executor.py

@@ -1,5 +1,3 @@
-import os
-import subprocess
 from typing import Optional
 
 from click.testing import CliRunner
@@ -181,23 +179,6 @@ def test_failure():
     assert result.exit_code == 1
 
 
-def test_asset_with_multiprocessing():
-    with instance_for_test():
-        subprocess.check_call(
-            [
-                "dagster",
-                "asset",
-                "materialize",
-                "-f",
-                file_relative_path(__file__, "asset_with_process_pool_executor.py"),
-                "--select",
-                "multiprocess_asset",
-                "--working-directory",
-                os.path.dirname(__file__),
-            ]
-        )
-
-
 def test_run_cli_config_json():
     with instance_for_test() as instance:
         asset_key = "asset_assert_with_config"

python_modules/dagster/dagster_tests/cli_tests/command_tests/test_materialize_command.py

@@ -223,11 +223,11 @@ def test_local_directory_module_multiple_defs() -> None:
         DagsterInvariantViolationError,
         match="Cannot have more than one Definitions object defined at module scope. Found Definitions objects: autodiscover_in_module_multiple.defs, autodiscover_in_module_multiple.defs1, autodiscover_in_module_multiple.defs2",
     ):
-        with alter_sys_path(to_add=[], to_remove=_current_test_directory_paths()):
-            loadable_targets_from_python_module(
-                "autodiscover_in_module_multiple",
-                working_directory=os.path.dirname(__file__),
-            )
+        loadable_targets_from_python_module(
+            "autodiscover_in_module_multiple",
+            working_directory=os.path.dirname(__file__),
+            remove_from_path_fn=_current_test_directory_paths,
+        )
 
 
 def _current_test_directory_paths():
@@ -251,25 +251,25 @@ def test_local_directory_module() -> None:
     if "autodiscover_in_module" in sys.modules:
         del sys.modules["autodiscover_in_module"]
 
-    with alter_sys_path(to_add=[], to_remove=_current_test_directory_paths()):
-        with pytest.raises(DagsterImportError):
-            loadable_targets_from_python_module(
-                "complete_bogus_module",
-                working_directory=os.path.dirname(__file__),
-            )
-
-    with alter_sys_path(to_add=[], to_remove=_current_test_directory_paths()):
-        with pytest.raises(DagsterImportError):
-            loadable_targets_from_python_module(
-                "autodiscover_in_module",
-                working_directory=None,
-            )
-
-    with alter_sys_path(to_add=[], to_remove=_current_test_directory_paths()):
-        loadable_targets = loadable_targets_from_python_module(
-            "autodiscover_in_module",
+    with pytest.raises(DagsterImportError):
+        loadable_targets_from_python_module(
+            "complete_bogus_module",
             working_directory=os.path.dirname(__file__),
+            remove_from_path_fn=_current_test_directory_paths,
         )
+
+    with pytest.raises(DagsterImportError):
+        loadable_targets_from_python_module(
+            "autodiscover_in_module",
+            working_directory=None,
+            remove_from_path_fn=_current_test_directory_paths,
+        )
+
+    loadable_targets = loadable_targets_from_python_module(
+        "autodiscover_in_module",
+        working_directory=os.path.dirname(__file__),
+        remove_from_path_fn=_current_test_directory_paths,
+    )
     assert len(loadable_targets) == 1
 
 
@@ -282,12 +282,8 @@ def test_local_directory_file() -> None:
 
         assert "No module named 'autodiscover_src'" in str(exc_info.value)
 
-    to_restore = [path for path in sys.path]
-    try:
-        alter_sys_path(to_add=[os.path.dirname(path)], to_remove=[])
+    with alter_sys_path(to_add=[os.path.dirname(path)], to_remove=[]):
         loadable_targets_from_python_file(path, working_directory=os.path.dirname(path))
-    finally:
-        sys.path = to_restore
 
 
 def test_lazy_definitions() -> None: