Remove use of distutils from code (not setup.py yet)

.github/workflows/test.yml

@@ -25,7 +25,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-20.04, macos-latest]
-        ver: ['3.5', '3.6', '3.7', '3.8', '3.9', '3.10', '3.11',]
+        ver: ['3.5', '3.6', '3.7', '3.8', '3.9', '3.10', '3.11', '3.12']
 
     steps:
     - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2.7.0

metaflow/__init__.py

@@ -188,11 +188,6 @@ class and related decorators.
         pass
 del globals()["_n"]
 
-import pkg_resources
+from .version import metaflow_version as _mf_version
 
-try:
-    __version__ = pkg_resources.get_distribution("metaflow").version
-except:
-    # this happens on remote environments since the job package
-    # does not have a version
-    __version__ = None
+__version__ = _mf_version

metaflow/extension_support/__init__.py

@@ -391,16 +391,15 @@ def _get_extension_packages():
             # unlikely so we are going to ignore this case.
             dist_root = dist.locate_file(EXT_PKG).as_posix()
             all_paths.discard(dist_root)
-
-            if dist.metadata["Name"] in mf_ext_packages:
+            dist_name = dist.metadata["Name"]
+            if dist_name in mf_ext_packages:
                 _ext_debug(
                     "Ignoring duplicate package '%s' (duplicate paths in sys.path? (%s))"
-                    % (dist.metadata["Name"], str(sys.path))
+                    % (dist_name, str(sys.path))
                 )
                 continue
             _ext_debug(
-                "Found extension package '%s' at '%s'..."
-                % (dist.metadata["Name"], dist_root)
+                "Found extension package '%s' at '%s'..." % (dist_name, dist_root)
             )
 
             files_to_include = []
@@ -419,7 +418,7 @@ def _get_extension_packages():
                     if parts[1] == "__init__.py":
                         raise RuntimeError(
                             "Package '%s' providing '%s' is not an implicit namespace "
-                            "package as required" % (dist.metadata["Name"], EXT_PKG)
+                            "package as required" % (dist_name, EXT_PKG)
                         )
 
                     # Record the file as a candidate for inclusion when packaging if
@@ -436,7 +435,7 @@ def _get_extension_packages():
                             "Package '%s' should conform to '%s.X.%s' and not '%s.%s' where "
                             "X is your organization's name for example"
                             % (
-                                dist.metadata["Name"],
+                                dist_name,
                                 EXT_PKG,
                                 parts[1],
                                 EXT_PKG,
@@ -453,7 +452,7 @@ def _get_extension_packages():
                                 "Package '%s' defines more than one meta configuration: "
                                 "'%s' and '%s' (at least)"
                                 % (
-                                    dist.metadata["Name"],
+                                    dist_name,
                                     meta_module,
                                     potential_meta_module,
                                 )
@@ -470,7 +469,7 @@ def _get_extension_packages():
 
                     # To give useful errors in case multiple top-level packages in
                     # one package
-                    dist_full_name = "%s[%s]" % (dist.metadata["Name"], parts[1])
+                    dist_full_name = "%s[%s]" % (dist_name, parts[1])
                     for idx, ext_list in enumerate(list_ext_points):
                         if (
                             len(parts) > len(ext_list) + 2
@@ -489,7 +488,7 @@ def _get_extension_packages():
                                 config_to_pkg[config_module].append(dist_full_name)
                             cur_pkg = (
                                 extension_points_to_pkg[_extension_points[idx]]
-                                .setdefault(dist.metadata["Name"], {})
+                                .setdefault(dist_name, {})
                                 .get(parts[1])
                             )
                             if cur_pkg is not None:
@@ -513,9 +512,9 @@ def _get_extension_packages():
                                         % (parts[1], config_module)
                                     )
                                     extension_points_to_pkg[_extension_points[idx]][
-                                        dist.metadata["Name"]
+                                        dist_name
                                     ][parts[1]] = MFExtPackage(
-                                        package_name=dist.metadata["Name"],
+                                        package_name=dist_name,
                                         tl_package=parts[1],
                                         config_module=config_module,
                                     )
@@ -525,14 +524,14 @@ def _get_extension_packages():
                                     % (parts[1], _extension_points[idx], config_module)
                                 )
                                 extension_points_to_pkg[_extension_points[idx]][
-                                    dist.metadata["Name"]
+                                    dist_name
                                 ][parts[1]] = MFExtPackage(
-                                    package_name=dist.metadata["Name"],
+                                    package_name=dist_name,
                                     tl_package=parts[1],
                                     config_module=config_module,
                                 )
                             break
-            mf_ext_packages[dist.metadata["Name"]] = {
+            mf_ext_packages[dist_name] = {
                 "root_paths": [dist_root],
                 "meta_module": meta_module,
                 "files": files_to_include,

metaflow/metadata/util.py

@@ -1,13 +1,30 @@
 from io import BytesIO
 import os
+import shutil
 import tarfile
 
-from distutils.dir_util import copy_tree
-
 from metaflow import util
 from metaflow.plugins.datastores.local_storage import LocalStorage
 
 
+def copy_tree(src, dst, update=False):
+    if not os.path.exists(dst):
+        os.makedirs(dst)
+    for item in os.listdir(src):
+        s = os.path.join(src, item)
+        d = os.path.join(dst, item)
+        if os.path.isdir(s):
+            copy_tree(s, d, update)
+        else:
+            if (
+                update
+                and os.path.exists(d)
+                and os.path.getmtime(s) <= os.path.getmtime(d)
+            ):
+                continue
+            shutil.copy2(s, d)
+
+
 def sync_local_metadata_to_datastore(metadata_local_dir, task_ds):
     with util.TempDir() as td:
         tar_file_path = os.path.join(td, "metadata.tgz")

metaflow/metaflow_config.py

@@ -1,10 +1,7 @@
-import logging
 import os
 import sys
 import types
 
-import pkg_resources
-
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config_funcs import from_conf, get_validate_choice_fn
 
@@ -412,24 +409,6 @@
 MAX_ATTEMPTS = 6
 
 
-# the naughty, naughty driver.py imported by lib2to3 produces
-# spam messages to the root logger. This is what is required
-# to silence it:
-class Filter(logging.Filter):
-    def filter(self, record):
-        if record.pathname.endswith("driver.py") and "grammar" in record.msg:
-            return False
-        return True
-
-
-logger = logging.getLogger()
-logger.addFilter(Filter())
-
-
-def get_version(pkg):
-    return pkg_resources.get_distribution(pkg).version
-
-
 # PINNED_CONDA_LIBS are the libraries that metaflow depends on for execution
 # and are needed within a conda environment
 def get_pinned_conda_libs(python_version, datastore_type):

metaflow/multicore_utils.py

@@ -2,7 +2,6 @@
 import os
 import traceback
 from itertools import islice
-from multiprocessing import cpu_count
 from tempfile import NamedTemporaryFile
 import time
 
@@ -60,15 +59,16 @@ def _spawn(func, arg, dir):
 
 
 def parallel_imap_unordered(func, iterable, max_parallel=None, dir=None):
-
     if max_parallel is None:
+        # Lazy import to save on startup time for metaflow as a whole
+        from multiprocessing import cpu_count
+
         max_parallel = cpu_count()
 
     args_iter = iter(iterable)
     pids = [_spawn(func, arg, dir) for arg in islice(args_iter, max_parallel)]
 
     while pids:
-
         for idx, pid_info in enumerate(pids):
             pid, output_file = pid_info
             pid, exit_code = os.waitpid(pid, os.WNOHANG)

metaflow/plugins/argo/argo_workflows_cli.py

@@ -3,7 +3,6 @@
 import platform
 import re
 import sys
-from distutils.version import LooseVersion
 from hashlib import sha1
 
 from metaflow import JSONType, current, decorators, parameters
@@ -26,7 +25,7 @@
 from metaflow.plugins.environment_decorator import EnvironmentDecorator
 from metaflow.plugins.kubernetes.kubernetes_decorator import KubernetesDecorator
 from metaflow.tagging_util import validate_tags
-from metaflow.util import get_username, to_bytes, to_unicode
+from metaflow.util import get_username, to_bytes, to_unicode, version_parse
 
 from .argo_workflows import ArgoWorkflows
 
@@ -290,7 +289,7 @@ def check_metadata_service_version(obj):
     version = metadata.version()
     if version == "local":
         return
-    elif version is not None and LooseVersion(version) >= LooseVersion("2.0.2"):
+    elif version is not None and version_parse(version) >= version_parse("2.0.2"):
         # Metaflow metadata service needs to be at least at version 2.0.2
         # since prior versions did not support strings as object ids.
         return

metaflow/plugins/aws/batch/batch_cli.py

@@ -4,8 +4,6 @@
 import time
 import traceback
 
-from distutils.dir_util import copy_tree
-
 from metaflow import util
 from metaflow import R
 from metaflow.exception import CommandException, METAFLOW_EXIT_DISALLOW_RETRY

metaflow/plugins/aws/step_functions/step_functions_cli.py

@@ -3,7 +3,6 @@
 from hashlib import sha1
 import json
 import re
-from distutils.version import LooseVersion
 
 from metaflow import current, decorators, parameters, JSONType
 from metaflow.metaflow_config import (
@@ -15,7 +14,7 @@
 from metaflow.package import MetaflowPackage
 from metaflow.plugins.aws.batch.batch_decorator import BatchDecorator
 from metaflow.tagging_util import validate_tags
-from metaflow.util import get_username, to_bytes, to_unicode
+from metaflow.util import get_username, to_bytes, to_unicode, version_parse
 
 from .step_functions import StepFunctions
 from .production_token import load_token, store_token, new_token
@@ -192,7 +191,7 @@ def check_metadata_service_version(obj):
     version = metadata.version()
     if version == "local":
         return
-    elif version is not None and LooseVersion(version) >= LooseVersion("2.0.2"):
+    elif version is not None and version_parse(version) >= version_parse("2.0.2"):
         # Metaflow metadata service needs to be at least at version 2.0.2
         return
     else:

metaflow/plugins/metadata/__init__.py

@@ -1,2 +1 @@
-from .local import LocalMetadataProvider
-from .service import ServiceMetadataProvider
+

metaflow/plugins/metadata/service.py

@@ -4,8 +4,6 @@
 import requests
 import time
 
-from distutils.version import LooseVersion
-
 from metaflow.exception import (
     MetaflowException,
     MetaflowTaggingError,
@@ -20,6 +18,8 @@
 from metaflow.metadata.heartbeat import HB_URL_KEY
 from metaflow.sidecar import Message, MessageTypes, Sidecar
 
+from metaflow.util import version_parse
+
 
 # Define message enums
 class HeartbeatTypes(object):
@@ -140,7 +140,7 @@ def _start_heartbeat(
         service_version = self.version()
         payload["service_version"] = service_version
         # start sidecar
-        if service_version is None or LooseVersion(service_version) < LooseVersion(
+        if service_version is None or version_parse(service_version) < version_parse(
             "2.0.4"
         ):
             # if old version of the service is running
@@ -347,7 +347,6 @@ def _get_or_create(
         tags=None,
         sys_tags=None,
     ):
-
         if tags is None:
             tags = set()
         if sys_tags is None:

metaflow/util.py

@@ -433,3 +433,9 @@ def is_within_directory(abs_directory, target):
         raise Exception("Attempted path traversal in TAR file")
 
     tar.extractall(path, members, numeric_owner=numeric_owner)
+
+
+if sys.version_info[:2] > (3, 5):
+    from metaflow._vendor.packaging.version import parse as version_parse
+else:
+    from distutils.version import LooseVersion as version_parse

metaflow/version.py

@@ -0,0 +1 @@
+metaflow_version = "2.10.3"

setup.py

@@ -1,6 +1,7 @@
 from setuptools import setup, find_packages
 
-version = "2.10.3"
+with open("metaflow/version.py", mode="r") as f:
+    version = f.read().splitlines()[0].split("=")[1].strip(" \"'")
 
 setup(
     include_package_data=True,
@@ -25,6 +26,7 @@
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
     ],
     project_urls={
         "Source": "https://github.com/Netflix/metaflow",