code: def install(use_cuda, use_nccl):
      ext_libs = []
      ext_args = (
          [
              "-Wno-sign-compare",
              "-Wno-unused-but-set-variable",
              "-Wno-terminate",
              "-Wno-unused-function",
              "-Wno-strict-aliasing",
          ]
          if pf.system() == "Linux"
          else []
      )
      if not use_cuda:
          use_nccl = False
          extension = CppExtension
      else:
          ext_libs += ["cuda", "nvrtc"] if not IS_HIP_EXTENSION else []
          ext_args += ["-DUSE_GPU"]
          extension = CUDAExtension
  
      if use_nccl:
          if IS_HIP_EXTENSION:
              ext_libs += ["rccl"]
          else:
              ext_libs += ["nccl"]
          ext_args += ["-DUSE_NCCL"]
  
      setup(
          name="tutel",
          version="0.3",
          description="An Optimized Mixture-of-Experts Implementation.",
          url="https://github.com/microsoft/Tutel",
          author="Microsoft",
          author_email="tutel@microsoft.com",
          license="MIT",
          classifiers=[
              "Development Status :: 2 - Pre-Alpha",
              "Environment :: GPU",
              "Intended Audience :: Developers",
              "Intended Audience :: Education",
              "Intended Audience :: Science/Research",
              "License :: OSI Approved :: MIT License",
              "Programming Language :: Python :: 3",
              "Programming Language :: Python :: 3 :: Only",
              "Programming Language :: Python :: 3.5",
              "Programming Language :: Python :: 3.6",
              "Programming Language :: Python :: 3.7",
              "Programming Language :: Python :: 3.8",
              "Programming Language :: Python :: 3.9",
              "Topic :: Scientific/Engineering",
              "Topic :: Scientific/Engineering :: Artificial Intelligence",
              "Topic :: Software Development",
              "Topic :: Software Development :: Libraries",
              "Topic :: Software Development :: Libraries :: Python Modules",
          ],
          keywords=["Mixture of Experts", "MoE", "Optimization"],
          packages=find_packages(exclude=["tests"]),
          python_requires=">=3.6, <4",
          install_requires=[
              "numpy",
          ],
          zip_safe=False,
          extras_require={
              "test": [
                  "GPUtil>=1.4.0",
                  "pytest-subtests>=0.4.0",
                  "pytest>=6.2.2",
              ],
          },
          ext_modules=[
              extension(
                  "tutel_custom_kernel",
                  [
                      "./tutel/custom/custom_kernel.cpp",
                  ],
                  library_dirs=["/usr/local/cuda/lib64/stubs"],
                  libraries=ext_libs,
                  extra_compile_args={"cxx": ext_args},
              )
          ],
          cmdclass={
              "build_ext": BuildExtension,
              "test": Tester,
          },
          project_urls={
              "Source": "https://github.com/microsoft/Tutel",
              "Tracker": "https://github.com/microsoft/Tutel/issues",
          },
      )
  
file_path: tutel/setup.py
imports: import os, sys
  import subprocess
  import platform
  import shutil
  from typing import List, Tuple
  from setuptools import setup, find_packages, Command
  from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
label_desc: Use image node_id 1 for calling a global function with example usage: install(use_cuda, use_nccl) without return types