setup.py

import setuptools
from setuptools import find_packages
from sys import platform as _platform
import sys
import glob
import os

from distutils.core import setup
from distutils.extension import Extension
from distutils.util import get_platform
from glob import glob

with open("README.md", "r") as fh:
  long_description = fh.read()

import os
import multiprocessing
import multiprocessing.pool


def parallelCCompile(self,
                     sources,
                     output_dir=None,
                     macros=None,
                     include_dirs=None,
                     debug=0,
                     extra_preargs=None,
                     extra_postargs=None,
                     depends=None):
    # those lines are copied from distutils.ccompiler.CCompiler directly
    macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
        output_dir, macros, include_dirs, sources, depends, extra_postargs)
    cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)
    # parallel code
    N = 2 * multiprocessing.cpu_count()  # number of parallel compilations
    try:
        # On Unix-like platforms attempt to obtain the total memory in the
        # machine and limit the number of parallel jobs to the number of Gbs
        # of RAM (to avoid killing smaller platforms like the Pi)
        mem = os.sysconf('SC_PHYS_PAGES') * os.sysconf('SC_PAGE_SIZE')  # bytes
    except (AttributeError, ValueError):
        # Couldn't query RAM; don't limit parallelism (it's probably a well
        # equipped Windows / Mac OS X box)
        pass
    else:
        mem = max(1, int(round(mem / 1024 ** 3)))  # convert to Gb
        N = min(mem, N)

    def _single_compile(obj):
        try:
            src, ext = build[obj]
        except KeyError:
            return
        newcc_args = cc_args
        if _platform == "darwin":
            if src.endswith('.cpp') or src.endswith('.cc'):
                newcc_args = cc_args + ["-mmacosx-version-min=10.7", "-std=c++17", "-stdlib=libc++"]
        self._compile(obj, src, ext, newcc_args, extra_postargs, pp_opts)

    # convert to list, imap is evaluated on-demand
    pool = multiprocessing.pool.ThreadPool(N)
    list(pool.imap(_single_compile, objects))
    return objects

import distutils.ccompiler

distutils.ccompiler.CCompiler.compile = parallelCCompile

platform = get_platform()
print(platform)

CXX_FLAGS = '-D__SUPPRESSANYOUTPUT__'


# libraries += [current_python]

libraries = []
include_dirs = [
  '.',
  'third_party/unitree_legged_sdk/pybind11/include', 
  'third_party/eigen3/include',
  'third_party/osqp/include',
  'third_party',
  'third_party/osqp/lin_sys/direct/qdldl',
  'third_party/osqp/lin_sys/direct/qdldl/qdldl_sources/include',
  'third_party/osqp/lin_sys/direct/qdldl/amd/include',
  'third_party/qpoases/include',
  ]

sources = [
    "mpc_implementation/mpc_osqp.cc",
	"third_party/osqp/src/auxil.c",
	"third_party/osqp/src/cs.c",
	"third_party/osqp/src/ctrlc.c",
	"third_party/osqp/src/error.c",
	"third_party/osqp/src/kkt.c",
	"third_party/osqp/src/lin_alg.c",
	"third_party/osqp/src/lin_sys.c",
	"third_party/osqp/src/osqp.c",
	"third_party/osqp/src/polish.c",
	"third_party/osqp/src/proj.c",
	"third_party/osqp/src/scaling.c",
	"third_party/osqp/src/util.c",
	"third_party/osqp/lin_sys/direct/qdldl/qdldl_interface.c",
	"third_party/osqp/lin_sys/direct/qdldl/qdldl_sources/src/qdldl.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_1.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_2.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_aat.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_control.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_defaults.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_info.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_order.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_post_tree.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_postorder.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_preprocess.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/amd_valid.c",
	"third_party/osqp/lin_sys/direct/qdldl/amd/src/SuiteSparse_config.c",
	
	
	"third_party/qpoases/src/BLASReplacement.cpp",
	"third_party/qpoases/src/Bounds.cpp",
	"third_party/qpoases/src/Constraints.cpp",
	"third_party/qpoases/src/Flipper.cpp",
	"third_party/qpoases/src/Indexlist.cpp",
	"third_party/qpoases/src/LAPACKReplacement.cpp",
	"third_party/qpoases/src/Matrices.cpp",
	"third_party/qpoases/src/MessageHandling.cpp",
	"third_party/qpoases/src/Options.cpp",
	"third_party/qpoases/src/OQPinterface.cpp",
	"third_party/qpoases/src/QProblem.cpp",
	"third_party/qpoases/src/QProblemB.cpp",
	"third_party/qpoases/src/SolutionAnalysis.cpp",
	"third_party/qpoases/src/SparseSolver.cpp",
	"third_party/qpoases/src/SQProblem.cpp",
	"third_party/qpoases/src/SQProblemSchur.cpp",
	"third_party/qpoases/src/SubjectTo.cpp",
	"third_party/qpoases/src/Utils.cpp",
	
]


include_dirs += ['third_party/osqp/include/linux']
CXX_FLAGS += '-fpermissive '
libraries = ['dl', 'pthread']
CXX_FLAGS += '-D_LINUX '
CXX_FLAGS += '-DGLEW_STATIC '
CXX_FLAGS += '-DGLEW_INIT_OPENGL11_FUNCTIONS=1 '
CXX_FLAGS += '-DGLEW_DYNAMIC_LOAD_ALL_GLX_FUNCTIONS=1 '
CXX_FLAGS += '-DDYNAMIC_LOAD_X11_FUNCTIONS '
CXX_FLAGS += '-DHAS_SOCKLEN_T '
CXX_FLAGS += '-fno-inline-functions-called-once '
CXX_FLAGS += '-fvisibility=hidden '
CXX_FLAGS += '-fvisibility-inlines-hidden '
CXX_FLAGS += '-std=c++1z '
CXX_FLAGS += '-Wno-sign-compare '
CXX_FLAGS += '-Wno-reorder '
CXX_FLAGS += '-Wno-unused-local-typedefs '
CXX_FLAGS += '-Wno-unused-variable '
CXX_FLAGS += '-Wno-unused-but-set-variable '

setup_py_dir = os.path.dirname(os.path.realpath(__file__))

extensions = []

mpc_osqp_ext = Extension(
    "mpc_osqp",
    sources=sources,
    libraries=libraries,
    extra_compile_args=CXX_FLAGS.split(),
    include_dirs=include_dirs + ["."])

extensions.append(mpc_osqp_ext)



INSTALL_REQUIRES = [
    # RL
    "gym==0.21.0",
    "torch==2.0.1",
    "pybullet==3.2.5",
    "tensorflow==2.10.0",
    "ray[rllib]==2.0.1",
    "stable-baselines3"
    ]

setup(
    name="usc_quadruped_learning",
    version="0.0.2",
    author="yiyuc",
    author_email="yiyuc@usc.edu",
    description="A framework for running reinforcement learning and imitation learning (mainly on quadrupeds).",
    long_description=long_description,
    long_description_content_type="text/markdown",
    python_requires=">=3.6",
    install_requires=INSTALL_REQUIRES,
    packages=setuptools.find_packages(),
    ext_modules=extensions,
    classifiers=[
        "Programming Language :: Python :: 3",
        "Operating System :: OS Independent",
    ],
)