google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
# pylint: skip-file
17
from setuptools import setup
18
from distutils.command.build import build
19
from setuptools.command.install import install
20

21
from setuptools.command.develop import develop
22
from torch.utils.cpp_extension import CppExtension, BuildExtension, CUDAExtension
23

24
import os
25
import subprocess
26
BASEPATH = os.path.dirname(os.path.abspath(__file__))
27

28
extlib_path = 'lego/ext_ops'
29
compile_args = ['-Wno-deprecated-declarations']
30
link_args = []
31
ext_modules = []
32

33
class custom_develop(develop):
34
    def run(self):
35
        original_cwd = os.getcwd()
36

37
        # build custom ops
38
        folders = [
39
           os.path.join(BASEPATH, 'lego/cpp_sampler'),
40
        ]
41
        for folder in folders:
42
            os.chdir(folder)
43
            subprocess.check_call(['make'])
44

45
        os.chdir(original_cwd)
46

47
        super().run()
48

49

50
setup(name='lego',
51
      py_modules=['lego'],
52
      ext_modules=ext_modules,
53
      install_requires=[
54
          'torch>=1.7.0',
55
          'scikit-learn',
56
          'tensorboardX',
57
          'gym'
58
      ],
59
      cmdclass={
60
          'build_ext': BuildExtension,
61
          'develop': custom_develop,
62
        }
63
)
64