Unverified Commit 25bb8647 authored by Kirthi Shankar Sivamani's avatar Kirthi Shankar Sivamani Committed by GitHub
Browse files

Get packaging module from setuptools (#190)



Use only built-ins for setup
Signed-off-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>
parent 7530b768
......@@ -10,7 +10,7 @@ import io
import re
import copy
import tempfile
from packaging.version import Version
from pkg_resources import packaging
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext
from shutil import copyfile
......@@ -172,7 +172,7 @@ class PyTorchBuilder(FrameworkBuilderBase):
@staticmethod
def install_requires():
return ["flash-attn>=1.0.2", "packaging"]
return ["flash-attn>=1.0.2"]
class TensorFlowBuilder(FrameworkBuilderBase):
......@@ -244,13 +244,13 @@ def get_cmake_bin():
try:
out = subprocess.check_output([cmake_bin, "--version"])
except OSError:
cmake_installed_version = Version("0.0")
cmake_installed_version = packaging.version.Version("0.0")
else:
cmake_installed_version = Version(
cmake_installed_version = packaging.version.Version(
re.search(r"version\s*([\d.]+)", out.decode()).group(1)
)
if cmake_installed_version < Version("3.18.0"):
if cmake_installed_version < packaging.version.Version("3.18.0"):
print(
"Could not find a recent CMake to build Transformer Engine. "
"Attempting to install CMake 3.18 to a temporary location via pip.",
......
......@@ -9,7 +9,7 @@ import warnings
from importlib.metadata import version
from contextlib import nullcontext
from typing import Any, Callable, Optional, Tuple, Union
from packaging.version import Version
from pkg_resources import packaging
import torch
......@@ -45,8 +45,8 @@ from transformer_engine.pytorch.distributed import (
)
from transformer_engine.pytorch.export import is_in_onnx_export_mode
_flash_attn_version = Version(version("flash-attn"))
_flash_attn_version_required = Version("1.0.2")
_flash_attn_version = packaging.version.Version(version("flash-attn"))
_flash_attn_version_required = packaging.version.Version("1.0.2")
warnings.filterwarnings("module", category=DeprecationWarning, module="transformer")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment