Repository: dask/distributed

Issue:
Cannot properly upload files with multiple modules or modules with an underscore in their name
When using the client.upload_file method to upload an egg file, the upload fails if the name of the module and resulting egg file contain an underscore character. In addition, only the package whose name matches the name of the egg file is imported and reloaded.

For example, if you make an egg file from a module "test_module",

`client.upload_file("test_module-1.0-py3.6.egg") `

attempts to import and then reload a module "test-module" (and then fails as a result), even though the original module was "test_module".

In addition it does not appear import any other modules in this egg file.

This appears to be down to the implementation of the import_file method in utils.py (called from the worker). 

This uses pkg_resources.find_distributions to get the project_name from the egg file, and then assumes that this project name is the name of the module to import and reload later. However, the package name has had 'safe_name' called on it when it was created, which converts any non alpha-numeric characters into hyphens.



Discussion:
Thank you for the excellent bug report and the diagnosis.  I agree that this seems like an important bug to fix.  Do you have any interest in taking this a step further and submitting a PR with a patch?
I don't know much about the python internals, so I don't know exactly how to fix this (though there's a method called walk_packages that looks promising). If I get some time I will try and look at a patch.
@awroe have you had time to investigate a fix for this?  No worries if not, just checking in.
Sorry, not had any time so far.

Resolving patch:
diff --git a/distributed/utils.py b/distributed/utils.py
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -23,7 +23,7 @@
 import threading
 import warnings
 import weakref
-
+import pkgutil
 import six
 import tblib.pickling_support
 
@@ -1066,14 +1066,11 @@ def import_file(path):
     if ext in (".egg", ".zip", ".pyz"):
         if path not in sys.path:
             sys.path.insert(0, path)
-        if ext == ".egg":
-            import pkg_resources
-
-            pkgs = pkg_resources.find_distributions(path)
-            for pkg in pkgs:
-                names_to_import.append(pkg.project_name)
-        elif ext in (".zip", ".pyz"):
-            names_to_import.append(name)
+        if sys.version_info >= (3, 6):
+            names = (mod_info.name for mod_info in pkgutil.iter_modules([path]))
+        else:
+            names = (mod_info[1] for mod_info in pkgutil.iter_modules([path]))
+        names_to_import.extend(names)
 
     loaded = []
     if not names_to_import:


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -9,6 +9,7 @@
 import os
 import pickle
 import random
+import subprocess
 import sys
 import threading
 from threading import Semaphore
@@ -79,6 +80,7 @@
     wait_for,
     async_wait_for,
     pristine_loop,
+    save_sys_modules,
 )
 from distributed.utils_test import (  # noqa: F401
     client as c,
@@ -1484,7 +1486,7 @@ def g():
 
         return myfile.f()
 
-    try:
+    with save_sys_modules():
         for value in [123, 456]:
             with tmp_text("myfile.py", "def f():\n    return {}".format(value)) as fn:
                 yield c.upload_file(fn)
@@ -1492,10 +1494,6 @@ def g():
             x = c.submit(g, pure=False)
             result = yield x
             assert result == value
-    finally:
-        # Ensure that this test won't impact the others
-        if "myfile" in sys.modules:
-            del sys.modules["myfile"]
 
 
 @gen_cluster(client=True)
@@ -1511,28 +1509,80 @@ def g():
 
         return myfile.f()
 
-    try:
+    with save_sys_modules():
+        try:
+            for value in [123, 456]:
+                with tmp_text(
+                    "myfile.py", "def f():\n    return {}".format(value)
+                ) as fn_my_file:
+                    with zipfile.ZipFile("myfile.zip", "w") as z:
+                        z.write(fn_my_file, arcname=os.path.basename(fn_my_file))
+                    yield c.upload_file("myfile.zip")
+
+                    x = c.submit(g, pure=False)
+                    result = yield x
+                    assert result == value
+        finally:
+            if os.path.exists("myfile.zip"):
+                os.remove("myfile.zip")
+
+
+@gen_cluster(client=True)
+def test_upload_file_egg(c, s, a, b):
+    def g():
+        import package_1, package_2
+
+        return package_1.a, package_2.b
+
+    # c.upload_file tells each worker to
+    # - put this file in their local_dir
+    # - modify their sys.path to include it
+    # we don't care about the local_dir
+    # but we do care about restoring the path
+
+    with save_sys_modules():
         for value in [123, 456]:
-            with tmp_text(
-                "myfile.py", "def f():\n    return {}".format(value)
-            ) as fn_my_file:
-                with zipfile.ZipFile("myfile.zip", "w") as z:
-                    z.write(fn_my_file, arcname=os.path.basename(fn_my_file))
-                yield c.upload_file("myfile.zip")
+            with tmpfile() as dirname:
+                os.mkdir(dirname)
+
+                with open(os.path.join(dirname, "setup.py"), "w") as f:
+                    f.write("from setuptools import setup, find_packages\n")
+                    f.write(
+                        'setup(name="my_package", packages=find_packages(), version="{}")\n'.format(
+                            value
+                        )
+                    )
+
+                # test a package with an underscore in the name
+                package_1 = os.path.join(dirname, "package_1")
+                os.mkdir(package_1)
+                with open(os.path.join(package_1, "__init__.py"), "w") as f:
+                    f.write("a = {}\n".format(value))
+
+                # test multiple top-level packages
+                package_2 = os.path.join(dirname, "package_2")
+                os.mkdir(package_2)
+                with open(os.path.join(package_2, "__init__.py"), "w") as f:
+                    f.write("b = {}\n".format(value))
+
+                # compile these into an egg
+                subprocess.check_call(
+                    [sys.executable, "setup.py", "bdist_egg"], cwd=dirname
+                )
+
+                egg_root = os.path.join(dirname, "dist")
+                # first file ending with '.egg'
+                egg_name = [
+                    fname for fname in os.listdir(egg_root) if fname.endswith(".egg")
+                ][0]
+                egg_path = os.path.join(egg_root, egg_name)
+
+                yield c.upload_file(egg_path)
+                os.remove(egg_path)
 
                 x = c.submit(g, pure=False)
                 result = yield x
-                assert result == value
-    finally:
-        # Ensure that this test won't impact the others
-        if os.path.exists("myfile.zip"):
-            os.remove("myfile.zip")
-        if "myfile" in sys.modules:
-            del sys.modules["myfile"]
-        for path in sys.path:
-            if os.path.basename(path) == "myfile.zip":
-                sys.path.remove(path)
-                break
+                assert result == (value, value)
 
 
 @gen_cluster(client=True)
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -1417,6 +1417,21 @@ def gen_tls_cluster(**kwargs):
     )
 
 
+@contextmanager
+def save_sys_modules():
+    old_modules = sys.modules
+    old_path = sys.path
+    try:
+        yield
+    finally:
+        for i, elem in enumerate(sys.path):
+            if elem not in old_path:
+                del sys.path[i]
+        for elem in sys.modules.keys():
+            if elem not in old_modules:
+                del sys.modules[elem]
+
+
 @contextmanager
 def check_thread_leak():
     active_threads_start = set(threading._active)