Home | History | Annotate | Download | only in ci_build
      1 #!/usr/bin/python
      2 # Copyright 2017 The TensorFlow Authors. All Rights Reserved.
      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 #
     17 # Automatically copy TensorFlow binaries
     18 #
     19 # Usage:
     20 #           ./tensorflow/tools/ci_build/copy_binary.py --filename
     21 # tf_nightly/tf_nightly_gpu-1.4.0.dev20170914-cp35-cp35m-manylinux1_x86_64.whl
     22 # --new_py_ver 36
     23 #
     24 """Copy binaries of TensorFlow for different python versions."""
     25 
     26 # pylint: disable=superfluous-parens
     27 
     28 import argparse
     29 import os
     30 import re
     31 import shutil
     32 import subprocess
     33 import zipfile
     34 
     35 UNZIP_CMD = "/usr/bin/unzip"
     36 ZIP_CMD = "/usr/bin/zip"
     37 SED_CMD = "/bin/sed"
     38 
     39 TF_NIGHTLY_REGEX = r"(.+)tf_nightly(|_gpu)-(\d\.\d\.\d.dev[\d]{0,8})-(.+)\.whl"
     40 BINARY_STRING_TEMPLATE = "%s-%s-%s.whl"
     41 
     42 
     43 def check_existence(filename):
     44   """Check the existence of file or dir."""
     45   if not os.path.exists(filename):
     46     raise RuntimeError("%s not found.")
     47 
     48 
     49 def copy_binary(directory, origin_tag, new_tag, version, gpu=False):
     50   """Rename and copy binaries for different python versions.
     51 
     52   Arguments:
     53     directory: string of directory
     54     origin_tag: str of the old python version tag
     55     new_tag: str of the new tag
     56     version: the version of the package
     57     gpu: bool if its a gpu build or not
     58 
     59   """
     60   print("Rename and copy binaries with %s to %s." % (origin_tag, new_tag))
     61   if gpu:
     62     package = "tf_nightly_gpu"
     63   else:
     64     package = "tf_nightly"
     65   origin_binary = BINARY_STRING_TEMPLATE % (package, version, origin_tag)
     66   new_binary = BINARY_STRING_TEMPLATE % (package, version, new_tag)
     67   zip_ref = zipfile.ZipFile(directory + origin_binary, "r")
     68   zip_ref.extractall()
     69   zip_ref.close()
     70   old_py_ver = re.search(r"(cp\d\d-cp\d\d)", origin_tag).group(1)
     71   new_py_ver = re.search(r"(cp\d\d-cp\d\d)", new_tag).group(1)
     72   subprocess.check_call(
     73       "%s -i s/%s/%s/g %s-%s.dist-info/WHEEL" % (SED_CMD, old_py_ver,
     74                                                  new_py_ver, package, version),
     75       shell=True)
     76   zout = zipfile.ZipFile(directory + new_binary, "w", zipfile.ZIP_DEFLATED)
     77   zip_these_files = [
     78       "%s-%s.dist-info" % (package, version),
     79       "%s-%s.data" % (package, version)
     80   ]
     81   for dirname in zip_these_files:
     82     for root, _, files in os.walk(dirname):
     83       for filename in files:
     84         zout.write(os.path.join(root, filename))
     85   zout.close()
     86   for dirname in zip_these_files:
     87     shutil.rmtree(dirname)
     88 
     89 
     90 def main():
     91   """This script copies binaries.
     92 
     93   Requirements:
     94     filename: The path to the whl file
     95     AND
     96     new_py_ver: Create a nightly tag with current date
     97 
     98   Raises:
     99     RuntimeError: If the whl file was not found
    100   """
    101 
    102   parser = argparse.ArgumentParser(description="Cherry picking automation.")
    103 
    104   # Arg information
    105   parser.add_argument(
    106       "--filename", help="path to whl file we are copying", required=True)
    107   parser.add_argument(
    108       "--new_py_ver", help="two digit py version eg. 27 or 33", required=True)
    109 
    110   args = parser.parse_args()
    111 
    112   # Argument checking
    113   check_existence(args.filename)
    114   regex_groups = re.search(TF_NIGHTLY_REGEX, args.filename)
    115   directory = regex_groups.group(1)
    116   gpu = regex_groups.group(2)
    117   version = regex_groups.group(3)
    118   origin_tag = regex_groups.group(4)
    119   old_py_ver = re.search(r"(cp\d\d)", origin_tag).group(1)
    120 
    121   # Create new tags
    122   new_tag = origin_tag.replace(old_py_ver, "cp" + args.new_py_ver)
    123 
    124   # Copy the binary with the info we have
    125   copy_binary(directory, origin_tag, new_tag, version, gpu)
    126 
    127 
    128 if __name__ == "__main__":
    129   main()
    130