[buildout]
extends =
  ../openblas/buildout.cfg
  ../cmake/buildout.cfg
  ../python-cffi/buildout.cfg
  ../python-PyYAML/buildout.cfg
  ../python-cocoapi/buildout.cfg
  ../pillow/buildout.cfg
  ../scipy/buildout.cfg
  ../matplotlib/buildout.cfg
parts =
  torch

[pytorch-repository]
recipe = plone.recipe.command
stop-on-error = true
repository = https://github.com/pytorch/pytorch
commit = 4af66c43045a317b477918c503d105f565b4a66b
location = ${buildout:parts-directory}/${:_buildout_section_name_}
command =
  set -e
  export PATH=${git:location}/bin:$PATH
  ( git clone --quiet -n ${:repository} ${:location}
    cd ${:location}
    git config submodule.recurse true
    git reset --merge ${:commit}
    git submodule update --init
    git apply ${:_profile_base_location_}/pytorch.4af66c4.patch
  ) || { rm -fr ${:location}; exit 1; }

[pytorch-build-interpreter]
recipe = zc.recipe.egg
initialization =
  import scipy.spatial.ckdtree # load our own libstdc++ explicitly at the very beginning
  import sys
  sys.path.append('.')
  sys.path.append('${pytorch-repository:location}/torch/lib/ATen')
eggs =
  setuptools
  ${scipy:egg}
  ${numpy:egg}
  ${python-cffi:egg}
  ${python-PyYAML:egg}
interpreter = pytorch-build-interpreter
scripts = pytorch-build-interpreter
location = ${buildout:parts-directory}/${:_buildout_section_name_}

[torch]
recipe = slapos.recipe.build
egg = ${:_buildout_section_name_}
workdir = ${pytorch-repository:location}
egg-name = ${:egg}-${versions:torch}-py${python:version}-linux-x86_64.egg
location = ${buildout:eggs-directory}/${:egg-name}
python-bin = ${buildout:bin-directory}/${pytorch-build-interpreter:interpreter}
no-cuda = 0
cmake-bin = ${cmake:location}/bin
git-bin = ${git:location}/bin
binutils-location = ${binutils:location}
openblas-location = ${openblas:location}
install =
  import os, shutil
  from setuptools.archive_util import unpack_archive
  workdir = options['workdir']
  python_bin = options['python-bin']
  binutils_location = options['binutils-location']
  openblas_location = options['openblas-location']
  env = {'PYTHONPATH':workdir,
         'PATH':':'.join([options['cmake-bin'],
                          options['git-bin'],
                          binutils_location+'/bin',
                          os.environ['PATH']]),
         'CMAKE_INCLUDE_PATH':':'.join([binutils_location+'/include',
                                        openblas_location+'/include',
                                        ]),
         'CMAKE_LIBRARY_PATH':':'.join([binutils_location+'/lib',
                                        openblas_location+'/lib',
                                        ]),
         'PYTORCH_PYTHON':python_bin,
         'SLAPOS_COMPILE_ARGS':'-Wl,-rpath,'+openblas_location+'/lib',
         }
  if options.get('no-cuda') == '1':
    env['NO_CUDA'] = '1'
  os.chdir(workdir)
  try:
    call([python_bin, 'setup.py', 'build'], env=env)
    call([python_bin, 'setup.py', 'bdist_egg'], env=env)
    unpack_archive(os.path.join('dist', options['egg-name']), location)
  finally:
    for d in 'build', 'dist':
      if os.path.exists(d):
        shutil.rmtree(d)

[pytorch-egg]
recipe = zc.recipe.egg
initialization =
  import scipy.spatial.ckdtree # load our own libstdc++ explicitly at the very beginning
eggs = 
  ${scipy:egg}
  ${numpy:egg}
  ${python-PyYAML:egg}
  ${torch:egg}
  ${pillow-python:egg}
  ${pycocotools:egg}
  ${matplotlib:egg}
  six
  torchvision
interpreter = pytorch-python
scripts = pytorch-python

[versions]
torchvision = 0.1.6
torch = 0.2.0+4af66c4
