-
Notifications
You must be signed in to change notification settings - Fork 57
/
install_and_run.py
88 lines (72 loc) · 3.05 KB
/
install_and_run.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Install and run the TensorBoard plugin for performance analysis.
Usage: python3 install_and_run.py --envdir ENVDIR --logdir LOGDIR
[--port PORT] [--version 2.2]
"""
import argparse
import os
import subprocess
NO_CHECK = '||True'
def run(*args):
"""Runs a shell command."""
subprocess.run(' '.join(args), shell=True, check=True)
class VirtualEnv(object):
"""Creates and runs programs in a virtual environment."""
def __init__(self, envdir):
self.envdir = envdir
run('virtualenv', '-p', 'python3', self.envdir)
def run(self, program, *args):
run(os.path.join(self.envdir, 'bin', program), *args)
def cleanup(self):
"""Clean up all existing TF profiler related installation."""
self.run('pip3', 'uninstall', '-q', '-y', 'tensorboard_plugin_profile',
'tensorboard', 'tensorflow-estimator', 'tensorflow', 'tbp-nightly',
'tb-nightly', 'tf-estimator-nightly', 'tf-nightly', NO_CHECK)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--envdir', help='Virtual environment', required=True)
parser.add_argument('--logdir', help='TensorBoard logdir', required=True)
parser.add_argument(
'--port',
help='TensorBoard port',
type=str,
required=False,
default='6006')
parser.add_argument(
'--version',
help='TensorFlow profiler version, e.g. nightly, 2.2',
required=False,
default='nightly')
args = parser.parse_args()
venv = VirtualEnv(args.envdir)
venv.cleanup()
if args.version == 'nightly':
venv.run('pip3', 'install', '-q', '-U', 'tf-nightly')
venv.run('pip3', 'install', '-q', '-U', 'tb-nightly')
venv.run('pip3', 'install', '-q', '-U', 'tbp-nightly')
else:
venv.run('pip3', 'install', '-q', '-U', 'tensorflow==' + args.version)
venv.run('pip3', 'install', '-q', '-U', 'tensorboard==' + args.version)
venv.run('pip3', 'install', '-q', '-U',
'tensorboard_plugin_profile==' + args.version)
tensorboard = os.path.join(args.envdir, 'bin/tensorboard')
# There is a bug that in Mac OS the shebang of tensorboard script is not
# correctly updated to the python3 of the virtual env. Directly invoke with
# python inside the virtual env to walk around.
venv.run('python3', tensorboard, '--logdir=' + args.logdir,
'--port=' + args.port, '--bind_all')
if __name__ == '__main__':
main()