From 0450c6a423ae806f09e4499afc13c285c5f021b5 Mon Sep 17 00:00:00 2001 From: gongsu832 Date: Wed, 28 Sep 2022 23:29:54 -0400 Subject: [PATCH] Fix onnx-mlir.py docker wrapper (#1734) * Requiring input files to have known suffix (.onnx, .json, or .mlir) greatly simplifies the onnx-mlir.py docker wrapper. Signed-off-by: Gong Su * Return exit code Signed-off-by: Gong Su * Add copyright notice Signed-off-by: Gong Su --- docker/onnx-mlir.py | 192 ++++++++++++++++++++------------------------ 1 file changed, 85 insertions(+), 107 deletions(-) diff --git a/docker/onnx-mlir.py b/docker/onnx-mlir.py index a039e97237..cdb61e2d21 100755 --- a/docker/onnx-mlir.py +++ b/docker/onnx-mlir.py @@ -1,23 +1,22 @@ #!/usr/bin/env python3 -import os -import re -import shutil -import stat -import subprocess -import sys +# SPDX-License-Identifier: Apache-2.0 -DOCKER_SOCKET = '/var/run/docker.sock' -ONNX_MLIR_IMAGE = 'onnxmlirczar/onnx-mlir' -WORK_DIR = '/workdir' -OUTPUT_DIR = '/output' -EMIT_IR_OPTS = [ '--EmitONNXBasic', - '--EmitONNXIR', - '--EmitMLIR', - '--EmitLLVMIR' ] -EMIT_BIN_OPTS = [ '--EmitLib', - '--EmitObj', - '--EmitJNI' ] +########################## onnx-mlir.py ######################################## +# +# Copyright 2022 The IBM Research Authors. +# +################################################################################ +# +# This file scans for certain patterns (listed below) and generate an md table, +# which list the operations supported, and optionally the unsupported operations. +# Among the options, we can also list the TODOs in the table. +# Invoke with the `-h` argument to list options. +# +# Limitation: currently handle at most one OP/LIMIT/TODO line per operation. +# Script currently invoked by the `onnx_mlir_supported_ops` make target. +# +################################################################################ # When running onnx-mlir inside a docker container, the directory # containing the input ONNX model file must be mounted into the @@ -25,16 +24,48 @@ # # This convenient script will do that automatically to make it # as if you are running onnx-mlir directly on the host. + +import os +import re +import shutil +import stat +import subprocess +import sys + +DOCKER_SOCKET = '/var/run/docker.sock' +ONNX_MLIR_IMAGE = 'onnxmlirczar/onnx-mlir' +KNOWN_INPUT_TYPE = ( '.onnx', '.json', '.mlir' ) + +mount_dirs = [] +mount_args = [] +onnx_mlir_args = [] + +# mount host path into container +def mount_path(path): + global mount_dirs, mount_args, onnx_mlir_args + + p = os.path.abspath(path) + d = os.path.dirname(p) + f = os.path.basename(p) + + # Haven't seen this directory before + if not d in mount_dirs: + mount_dirs += [ d ] + mount_args += [ '-v', d + ':' + d ] + onnx_mlir_args += [ p ] + else: + onnx_mlir_args += [ path ] + def main(): # Make sure docker client is installed if not shutil.which('docker'): print('docker client not found') - return + sys.exit(1) # Make sure docker daemon is running if not stat.S_ISSOCK(os.stat(DOCKER_SOCKET).st_mode): print('docker daemon not running') - return + sys.exit(1) # Pull the latest onnxmlirczar/onnx-mlir image, if image # is already up-to-date, pull will do nothing. @@ -48,112 +79,59 @@ def main(): print(line if re.match('^([0-9a-f]{12})|Error', line) else '', end='', flush=True) proc.wait() - if (proc.returncode != 0): + if proc.returncode: print("docker pull failed") - return - - # Go through the command line options and locate the - # input ONNX model file. - argi = 0 - ionnx = None - argo = 0 - obase = None - argv = sys.argv - argc = len(sys.argv) - - for i, arg in enumerate(argv): - # File specified on the first argument, defaults to --EmitLib - if i == 1 and not argv[i].startswith('-'): - argi = i - ionnx = argv[i] - # If a file is not specified on the first argument, it must be - # specified after a valid --EmitXXX option. - elif (arg in EMIT_IR_OPTS+EMIT_BIN_OPTS and i < argc-1 and - not argv[i+1].startswith('-')): - # File specified more than once, treat as not specified - if ionnx: - sys.exit("Too many --EmitXXX options") - if (arg in EMIT_BIN_OPTS and sys.platform != 'linux'): - print(('Warning: host {} is not linux, ' + - 'output not directly usable').format(sys.platform)) - argi = i + 1 - ionnx = argv[argi] - elif (arg == "-o" and i < argc-1 and - not argv[i+1].startswith('-')): - if obase: - sys.exit("Too many -o options") - argo = i + 1 - obase = argv[argo] + sys.exit(proc.returncode) # Prepare the arguments for docker run args = [ 'docker', 'run', '--rm', '-ti' ] - # Construct the mount option if an input ONNX model file is found - if ionnx: - p = os.path.abspath(ionnx) - d = os.path.dirname(p) - f = os.path.basename(p) - - # Add the mount option, directory containing the input - # ONNX model file will be mounted under /workdir inside - # the container. If /workdir doesn't exist, it will be - # created. - args.append('-v') - args.append(d + ':' + WORK_DIR + d) - - # Change directory into /workdir - #args.append('-w') - #args.append(WORK_DIR) - - # Rewrite the original input ONNX model file, which will - # reside under /workdir inside the container. - argv[argi] = WORK_DIR + p - - # Construct the mount option if -o is specified - if obase: - # Check invalid -o values such as ".", "..", "/.", "./", etc. - if re.match('(.*/)*\.*$', obase): - sys.exit("Invalid value for -o option") - - p = os.path.abspath(obase) - d = os.path.dirname(p) - f = os.path.basename(p) - - # Add the mount option, directory containing the output - # files will be mounted under /output inside the container. - # If /output/... doesn't exist, it will be - # created. - args.append('-v') - args.append(d + ':' + OUTPUT_DIR + d) - - # Rewrite the original output basename, which will - # reside under /output inside the container. - argv[argo] = OUTPUT_DIR + p + # Go through the command line options and locate the known + # file types. For each file located, construct a docker mount + # option that mounts the host directory into the container. + # + # Also do the same for the output path specified by the -o + # option. + argv = sys.argv + argc = len(sys.argv) + + global mount_dirs, mount_args, onnx_mlir_args + + verbose = False + for i in range(1, argc): + if argv[i].endswith(KNOWN_INPUT_TYPE): + mount_path(argv[i]) + elif argv[i-1] == '-o' and not argv[i].startswith('-'): + mount_path(argv[i]) + elif argv[i] == '-v': + verbose = True + onnx_mlir_args += [ argv[i] ] + else: + onnx_mlir_args += [ argv[i] ] # Add effective uid and gid - args.append('-u') - args.append(str(os.geteuid()) + ':' + str(os.getegid())) + args += [ '-u', str(os.geteuid()) + ':' + str(os.getegid()) ] + + # Add mount options + args += mount_args # Add image name - args.append(ONNX_MLIR_IMAGE) + args += [ ONNX_MLIR_IMAGE ] # Pass in all the original arguments - argv.remove(argv[0]) - args.extend(argv) - # print(args) #debug only + args += onnx_mlir_args + + if verbose: + print(args) # Run onnx-mlir in the container proc = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) for line in proc.stdout: - # Remove first occurrence of /workdir or /output before printing - print(re.sub(WORK_DIR + '|' + OUTPUT_DIR, '', - line.decode('utf-8'), 1), - end='', flush=True) + print(line.decode('utf-8'), end='', flush=True) proc.wait() - if (proc.returncode != 0): - print(os.strerror(proc.returncode)) + sys.exit(proc.returncode) if __name__ == "__main__": main()