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 tempfile
     33 import zipfile
     34 
     35 TF_NIGHTLY_REGEX = (r"(.+)(tf_nightly.*)-(\d\.[\d]{1,2}"
     36                     r"\.\d.dev[\d]{0,8})-(.+)\.whl")
     37 BINARY_STRING_TEMPLATE = "%s-%s-%s.whl"
     38 
     39 
     40 def check_existence(filename):
     41   """Check the existence of file or dir."""
     42   if not os.path.exists(filename):
     43     raise RuntimeError("%s not found." % filename)
     44 
     45 
     46 def copy_binary(directory, origin_tag, new_tag, version, package):
     47   """Rename and copy binaries for different python versions.
     48 
     49   Arguments:
     50     directory: string of directory
     51     origin_tag: str of the old python version tag
     52     new_tag: str of the new tag
     53     version: the version of the package
     54     package: str, name of the package
     55 
     56   """
     57   print("Rename and copy binaries with %s to %s." % (origin_tag, new_tag))
     58   origin_binary = BINARY_STRING_TEMPLATE % (package, version, origin_tag)
     59   new_binary = BINARY_STRING_TEMPLATE % (package, version, new_tag)
     60   zip_ref = zipfile.ZipFile(os.path.join(directory, origin_binary), "r")
     61 
     62   try:
     63     tmpdir = tempfile.mkdtemp()
     64     os.chdir(tmpdir)
     65 
     66     zip_ref.extractall()
     67     zip_ref.close()
     68     old_py_ver = re.search(r"(cp\d\d-cp\d\d)", origin_tag).group(1)
     69     new_py_ver = re.search(r"(cp\d\d-cp\d\d)", new_tag).group(1)
     70 
     71     wheel_file = os.path.join(
     72         tmpdir, "%s-%s.dist-info" % (package, version), "WHEEL")
     73     with open(wheel_file, "r") as f:
     74       content = f.read()
     75     with open(wheel_file, "w") as f:
     76       f.write(content.replace(old_py_ver, new_py_ver))
     77 
     78     zout = zipfile.ZipFile(directory + new_binary, "w", zipfile.ZIP_DEFLATED)
     79     zip_these_files = [
     80         "%s-%s.dist-info" % (package, version),
     81         "%s-%s.data" % (package, version),
     82     ]
     83     for dirname in zip_these_files:
     84       for root, _, files in os.walk(dirname):
     85         for filename in files:
     86           zout.write(os.path.join(root, filename))
     87     zout.close()
     88   finally:
     89     shutil.rmtree(tmpdir)
     90 
     91 
     92 def main():
     93   """This script copies binaries.
     94 
     95   Requirements:
     96     filename: The path to the whl file
     97     AND
     98     new_py_ver: Create a nightly tag with current date
     99 
    100   Raises:
    101     RuntimeError: If the whl file was not found
    102   """
    103 
    104   parser = argparse.ArgumentParser(description="Cherry picking automation.")
    105 
    106   # Arg information
    107   parser.add_argument(
    108       "--filename", help="path to whl file we are copying", required=True)
    109   parser.add_argument(
    110       "--new_py_ver", help="two digit py version eg. 27 or 33", required=True)
    111 
    112   args = parser.parse_args()
    113 
    114   # Argument checking
    115   args.filename = os.path.abspath(args.filename)
    116   check_existence(args.filename)
    117   regex_groups = re.search(TF_NIGHTLY_REGEX, args.filename)
    118   directory = regex_groups.group(1)
    119   package = regex_groups.group(2)
    120   version = regex_groups.group(3)
    121   origin_tag = regex_groups.group(4)
    122   old_py_ver = re.search(r"(cp\d\d)", origin_tag).group(1)
    123 
    124   # Create new tags
    125   new_tag = origin_tag.replace(old_py_ver, "cp" + args.new_py_ver)
    126 
    127   # Copy the binary with the info we have
    128   copy_binary(directory, origin_tag, new_tag, version, package)
    129 
    130 
    131 if __name__ == "__main__":
    132   main()
    133