update install reqs if op is compatible

op_builder/builder.py

@@ -95,6 +95,18 @@ def is_compatible(self):
         '''
         return True
 
+    def python_requirements(self):
+        '''
+        Override if op wants to define special dependencies, otherwise will
+        take self.name and load requirements-<op-name>.txt if it exists.
+        '''
+        path = f'requirements/requirements-{self.name}.txt'
+        requirements = []
+        if os.path.isfile():
+            with open(path, 'r') as fd:
+                requirements = [r.strip() for r in fd.readlines()]
+        return requirements
+
     def command_exists(self, cmd):
         if '|' in cmd:
             cmds = cmd.split("|")

op_builder/cpu_adam.py

@@ -1,5 +1,6 @@
 import os
 import torch
+import warnings
 from .builder import CUDAOpBuilder
 
 

op_builder/sparse_attn.py

@@ -1,6 +1,6 @@
 import torch
 import warnings
-from .builder import OpBuilder  #, command_exists
+from .builder import OpBuilder
 
 
 class SparseAttnBuilder(OpBuilder):
@@ -33,14 +33,5 @@ def is_compatible(self):
                 f'{self.NAME} requires a torch version >= 1.5 but detected {TORCH_MAJOR}.{TORCH_MINOR}'
             )
 
-        try:
-            import triton
-            triton_installed = True
-        except ImportError:
-            triton_installed = False
-        if not triton_installed:
-            self.warning(
-                f"{self.NAME} requires the python package 'triton' to be installed")
-
         return super().is_compatible(
         ) and deps_compatible and torch_compatible and triton_installed

setup.py

@@ -26,8 +26,8 @@ def fetch_requirements(path):
 
 install_requires = fetch_requirements('requirements/requirements.txt')
 extras_require = {
-    'sparse_attn': fetch_requirements('requirements/requirements-sparse-attn.txt'),
     '1bit_adam': fetch_requirements('requirements/requirements-1bit-adam.txt'),
+    'readthedocs': fetch_requirements('requirements/requirements-readthedocs.txt'),
     'dev': fetch_requirements('requirements/requirements-dev.txt'),
 }
 
@@ -91,6 +91,8 @@ def op_enabled(op_name):
         if builder.is_compatible():
             install_ops[op_name] = op_enabled(op_name)
             ext_modules.append(builder.builder())
+            # Add any necessary python requirements for this op
+            install_requires += builder.python_requirements()
 
 compatible_ops = {op_name: op.is_compatible() for (op_name, op) in ALL_OPS.items()}