refactor python module wheel building and installation to be less prone to race conditions

2023-03-24 17:43:52 -04:00
parent bd59c3ea69
commit ab48b834f7
4 changed files with 55 additions and 52 deletions
--- a/cmake/CMakeLists.txt
+++ b/cmake/CMakeLists.txt
@ -842,7 +842,7 @@ if(BUILD_SHARED_LIBS)
  if(Python_EXECUTABLE)
    add_custom_target(
      install-python ${Python_EXECUTABLE} ${LAMMPS_PYTHON_DIR}/install.py -p ${LAMMPS_PYTHON_DIR}/lammps
-              -l ${LIBLAMMPS_SHARED_BINARY} -w ${MY_BUILD_DIR}
+      -l ${LIBLAMMPS_SHARED_BINARY} -w ${MY_BUILD_DIR} -v ${LAMMPS_SOURCE_DIR}/version.h
      COMMENT "Installing LAMMPS Python module")
  else()
    add_custom_target(
--- a/python/install.py
+++ b/python/install.py
@ -25,6 +25,8 @@ parser.add_argument("-n", "--noinstall", action="store_true", default=False,
                    help="only build a binary wheel. Don't attempt to install it")
 parser.add_argument("-w", "--wheeldir", required=False,
                    help="path to a directory where the created wheel will be stored")
+parser.add_argument("-v", "--versionfile", required=True,
+                    help="path to the LAMMPS version.h source file")

 args = parser.parse_args()

@ -32,7 +34,7 @@ args = parser.parse_args()

 if args.package:
  if not os.path.exists(args.package):
-    print("ERROR: LAMMPS package %s does not exist" % args.package)
+    print("ERROR: LAMMPS package folder %s does not exist" % args.package)
    parser.print_help()
    sys.exit(1)
  else:
@ -54,19 +56,32 @@ if args.wheeldir:
  else:
    args.wheeldir = os.path.abspath(args.wheeldir)

-# we need to switch to the folder of the python package
+if args.versionfile:
+  if not os.path.exists(args.versionfile):
+    print("ERROR: LAMMPS version file at %s does not exist" % args.versionfile)
+    parser.print_help()
+    sys.exit(1)
+  else:
+    args.versionfile = os.path.abspath(args.versionfile)
+
 olddir = os.path.abspath('.')
-os.chdir(os.path.dirname(args.package))
+pythondir = os.path.abspath(os.path.join(args.package,'..'))
+os.putenv('LAMMPS_VERSION_FILE',os.path.abspath(args.versionfile))

-# remove any wheel files left over from previous calls
-if os.path.isdir(os.path.join(olddir,"build")):
+# purge old build folder
+if os.path.isdir(os.path.join(olddir, 'build-python')):
  print("Cleaning old build directory")
-  shutil.rmtree(os.path.join(olddir,"build"))
+  shutil.rmtree(os.path.join(olddir, 'build-python'))

-#print("Purging existing wheels...")
-#for wheel in glob.glob('lammps-*.whl'):
-#  print("deleting " + wheel)
-#  os.remove(wheel)
+# remove any old wheel files left over from previous calls
+print("Purging existing wheels...")
+for wheel in glob.glob('lammps-*.whl'):
+  print("deleting " + wheel)
+  os.remove(wheel)
+
+# copy python tree to build folder
+builddir = shutil.copytree(pythondir, os.path.join(olddir, 'build-python'))
+os.chdir(builddir)

 # copy shared object to the current folder so that
 # it will show up in the installation at the expected location
@ -74,7 +89,7 @@ os.putenv('LAMMPS_SHARED_LIB',os.path.basename(args.lib))
 shutil.copy(args.lib,'lammps')

 # create a virtual environment for building the wheel
-shutil.rmtree('buildwheel',True)
+shutil.rmtree('buildwheel', True)
 try:
  txt = subprocess.check_output([sys.executable, '-m', 'venv', 'buildwheel'], stderr=subprocess.STDOUT, shell=False)
  print(txt.decode('UTF-8'))
@ -86,18 +101,19 @@ except subprocess.CalledProcessError as err:
 # there is no simple way to return from that in python.
 os.system(sys.executable + ' makewheel.py')

+# copy wheel to final location
+for wheel in glob.glob('lammps-*.whl'):
+  if args.wheeldir:
+    shutil.copy(wheel, args.wheeldir)
+
+print('wheel = ', wheel)
+
 # remove temporary folders and files
-shutil.rmtree('buildwheel',True)
-shutil.rmtree('build',True)
-shutil.rmtree('lammps.egg-info',True)
-os.remove(os.path.join('lammps',os.path.basename(args.lib)))
+os.chdir(olddir)
+shutil.rmtree('build-python',True)

 # stop here if we were asked not to install the wheel we created
 if args.noinstall:
-  if args.wheeldir:
-    for wheel in glob.glob('lammps-*.whl'):
-      shutil.copy(wheel, args.wheeldir)
-      os.remove(wheel)
  exit(0)

 # install the wheel with pip. first try to install in the default environment.
@ -122,28 +138,17 @@ else:
  print("Installing wheel into system site-packages folder")
  py_exe = sys.executable

-for wheel in glob.glob('lammps-*.whl'):
-  try:
-    txt = subprocess.check_output([py_exe, '-m', 'pip', 'install', '--force-reinstall', wheel], stderr=subprocess.STDOUT, shell=False)
-    print(txt.decode('UTF-8'))
-    if args.wheeldir:
-      shutil.copy(wheel, args.wheeldir)
-    else:
-      shutil.copy(wheel, olddir)
-    os.remove(wheel)
-    continue
-  except subprocess.CalledProcessError as err:
-    errmsg = err.output.decode('UTF-8')
-    if errmsg.find("distutils installed"):
-      sys.exit(errmsg + "You need to uninstall the LAMMPS python module manually first.\n")
-  try:
-    print('Installing wheel into system site-packages folder failed. Trying user folder now')
-    txt = subprocess.check_output([sys.executable, '-m', 'pip', 'install', '--user', '--force-reinstall', wheel], stderr=subprocess.STDOUT, shell=False)
-    print(txt.decode('UTF-8'))
-    if args.wheeldir:
-      shutil.copy(wheel, args.wheeldir)
-    else:
-      shutil.copy(wheel, olddir)
-    os.remove(wheel)
-  except:
-    sys.exit('Failed to install wheel ' + wheel)
+try:
+  txt = subprocess.check_output([py_exe, '-m', 'pip', 'install', '--force-reinstall', wheel], stderr=subprocess.STDOUT, shell=False)
+  print(txt.decode('UTF-8'))
+  sys.exit(0)
+except subprocess.CalledProcessError as err:
+  errmsg = err.output.decode('UTF-8')
+  if errmsg.find("distutils installed"):
+    sys.exit(errmsg + "You need to uninstall the LAMMPS python module manually first.\n")
+try:
+  print('Installing wheel into system site-packages folder failed. Trying user folder now')
+  txt = subprocess.check_output([sys.executable, '-m', 'pip', 'install', '--user', '--force-reinstall', wheel], stderr=subprocess.STDOUT, shell=False)
+  print(txt.decode('UTF-8'))
+except:
+  sys.exit('Failed to install wheel ' + wheel)
--- a/python/setup.py
+++ b/python/setup.py
@ -4,17 +4,15 @@ from setuptools import setup
 from setuptools.dist import Distribution
 from sys import version_info
 import os,time
-LAMMPS_PYTHON_DIR = os.path.dirname(os.path.realpath(__file__))
-LAMMPS_DIR = os.path.dirname(LAMMPS_PYTHON_DIR)
-LAMMPS_SOURCE_DIR = os.path.join(LAMMPS_DIR, 'src')

-if not os.path.exists(LAMMPS_SOURCE_DIR):
+versionfile = os.environ.get("LAMMPS_VERSION_FILE")
+if not versionfile:
    # allows installing and building wheel from current directory
    LAMMPS_DIR = os.path.realpath(os.path.join(os.environ['PWD'], '..'))
-    LAMMPS_SOURCE_DIR = os.path.join(LAMMPS_DIR, 'src')
+    versionfile = os.path.join(LAMMPS_DIR, 'src', 'version.h')

 def get_lammps_version():
-    version_h_file = os.path.join(LAMMPS_SOURCE_DIR, 'version.h')
+    version_h_file = os.path.join(versionfile)
    with open(version_h_file, 'r') as f:
        line = f.readline()
        start_pos = line.find('"')+1
--- a/src/Makefile
+++ b/src/Makefile
@ -468,7 +468,7 @@ mpi-stubs:
 sinclude ../lib/python/Makefile.lammps
 install-python:
 	@rm -rf ../python/build
-	@$(PYTHON) ../python/install.py -p ../python/lammps -l ../src/liblammps.so -w $(PWD)
+	@$(PYTHON) ../python/install.py -p ../python/lammps -l ../src/liblammps.so -w $(PWD) -v $(PWD)/version.h

 # Create a tarball of src dir and packages