/ mlflow / utils / download_cloud_file_chunk.py
download_cloud_file_chunk.py
 1  """
 2  This script should be executed in a fresh python interpreter process using `subprocess`.
 3  """
 4  
 5  import argparse
 6  import importlib.util
 7  import json
 8  import os
 9  import sys
10  
11  
12  def parse_args():
13      parser = argparse.ArgumentParser()
14      parser.add_argument("--range-start", required=True, type=int)
15      parser.add_argument("--range-end", required=True, type=int)
16      parser.add_argument("--headers", required=True, type=str)
17      parser.add_argument("--download-path", required=True, type=str)
18      parser.add_argument("--http-uri", required=True, type=str)
19      return parser.parse_args()
20  
21  
22  def main():
23      file_path = os.path.join(os.path.dirname(__file__), "request_utils.py")
24      module_name = "mlflow.utils.request_utils"
25  
26      spec = importlib.util.spec_from_file_location(module_name, file_path)
27      module = importlib.util.module_from_spec(spec)
28      sys.modules[module_name] = module
29      spec.loader.exec_module(module)
30      download_chunk = module.download_chunk
31  
32      args = parse_args()
33      download_chunk(
34          range_start=args.range_start,
35          range_end=args.range_end,
36          headers=json.loads(args.headers),
37          download_path=args.download_path,
38          http_uri=args.http_uri,
39      )
40  
41  
42  if __name__ == "__main__":
43      main()