-
Notifications
You must be signed in to change notification settings - Fork 1
/
run_docker.py
2162 lines (1832 loc) · 76.2 KB
/
run_docker.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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Copyright 2018 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.
# ============================================================================
"""Multipurpose TensorFlow Docker Helper.
- Assembles Dockerfiles
- Builds images (and optionally runs image tests)
- Pushes images to Docker Hub (provided with credentials)
Read README.md (in this directory) for instructions!
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from os.path import join as _j, abspath as _a, exists as _e, dirname as _d, basename as _b
import subprocess
import re
import os
import time
import argparse
import pwd
import collections
import pprint
import copy
import errno
import itertools
import multiprocessing
import os
import re
import shutil
import sys
import textwrap
import cerberus
import docker
from docker.models.images import Image
import yaml
# https://docker-py.readthedocs.io/en/stable/api.html#docker.api.build.BuildApiMixin.build
from io import BytesIO
from docker import APIClient
from pathlib import Path
# run_docker.py is at $REPO_ROOT/dockerfiles
REPO_ROOT = _d(_a(__file__))
sys.path.append(REPO_ROOT)
from rlscope import py_config
from rlscope.profiler.rlscope_logging import logger
from rlscope.profiler import nvidia_gpu_query
DOCKERFILES = _a(_j(REPO_ROOT, 'dockerfiles'))
NVIDIA_VISIBLE_DEVICES = [0]
assert len(NVIDIA_VISIBLE_DEVICES) > 0
PROJECT_NAME = 'rlscope'
RLSCOPE_BASH_SERVICE_NAME = 'bash'
HOME = str(Path.home())
DEFAULT_RLSCOPE_DRILL_PORT = 8129
# The tag used for a locally built "bash" RL-Scope dev environment
LOCAL_RLSCOPE_IMAGE_TAG = 'tensorflow:devel-rlscope-gpu-cuda'
DEFAULT_REMOTE_RLSCOPE_IMAGE_TAG = 'UofT-EcoSystem/rlscope:1.0.0'
RELEASE_TO_LOCAL_IMG_TAG = dict()
RELEASE_TO_LOCAL_IMG_TAG['rlscope'] = 'tensorflow:devel-rlscope-gpu-cuda'
RELEASE_TO_LOCAL_IMG_TAG['rlscope-ubuntu-20-04-cuda-11-0'] = 'tensorflow:ubuntu-20-04-devel-rlscope-cuda-11-0'
# How long should we wait for /bin/bash (rlscope_bash)
# to appear after running "docker stack deploy"?
DOCKER_DEPLOY_TIMEOUT_SEC = 10
# Schema to verify the contents of tag-spec.yml with Cerberus.
# Must be converted to a dict from yaml to work.
# Note: can add python references with e.g.
# !!python/name:builtins.str
# !!python/name:__main__.funcname
SCHEMA_TEXT = """
header:
type: string
slice_sets:
type: dict
keyschema:
type: string
valueschema:
type: list
schema:
type: dict
schema:
add_to_name:
type: string
dockerfile_exclusive_name:
type: string
dockerfile_subdirectory:
type: string
partials:
type: list
schema:
type: string
ispartial: true
test_runtime:
type: string
required: false
tests:
type: list
default: []
schema:
type: string
args:
type: list
default: []
schema:
type: string
isfullarg: true
interpolate_arg: true
run_args:
type: list
default: []
schema:
type: string
isfullarg: true
interpolate_arg: true
releases:
type: dict
keyschema:
type: string
valueschema:
type: dict
schema:
is_dockerfiles:
type: boolean
required: false
default: false
upload_images:
type: boolean
required: false
default: true
tag_specs:
type: list
required: true
schema:
type: string
"""
# https://stackoverflow.com/questions/17215400/python-format-string-unused-named-arguments
class FormatDict(dict):
def __missing__(self, key):
return '{' + key + '}'
class TfDockerTagValidator(cerberus.Validator):
"""Custom Cerberus validator for TF tag spec.
Note: Each _validate_foo function's docstring must end with a segment
describing its own validation schema, e.g. "The rule's arguments are...". If
you add a new validator, you can copy/paste that section.
"""
def __init__(self, *args, **kwargs):
# See http://docs.python-cerberus.org/en/stable/customize.html
if 'partials' in kwargs:
self.partials = kwargs['partials']
super(cerberus.Validator, self).__init__(*args, **kwargs)
def _validate_ispartial(self, ispartial, field, value):
"""Validate that a partial references an existing partial spec.
Args:
ispartial: Value of the rule, a bool
field: The field being validated
value: The field's value
The rule's arguments are validated against this schema:
{'type': 'boolean'}
"""
if ispartial and value not in self.partials:
self._error(field,
'{} is not present in the partials directory.'.format(value))
def _validate_isfullarg(self, isfullarg, field, value):
"""Validate that a string is either a FULL=arg or NOT.
Args:
isfullarg: Value of the rule, a bool
field: The field being validated
value: The field's value
The rule's arguments are validated against this schema:
{'type': 'boolean'}
"""
if isfullarg and '=' not in value:
self._error(field, '{} should be of the form ARG=VALUE.'.format(value))
if not isfullarg and '=' in value:
self._error(field, '{} should be of the form ARG (no =).'.format(value))
def _validate_interpolate_arg(self, interpolate_arg, field, value):
"""Test the oddity of a value.
The rule's arguments are validated against this schema:
{'type': 'boolean'}
"""
if interpolate_arg:
validate_interpolate_arg(value, lambda message: self._error(field, message))
def default_report_error(message):
raise RuntimeError(message)
def validate_interpolate_arg(value, report_error=default_report_error):
bracket_regex = re.compile(r'\{([^}]+)\}')
bracket_varnames = bracket_regex.findall(value)
if len(bracket_varnames) == 0:
return
for varname in bracket_varnames:
if varname not in ARG_VALUES:
report_error((
"Couldn't find {var} in --arg or in spec.yml for arg=\"{value}\"; available ARG_VALUES for use in "
"--arg or spec.yml args are:\n"
"{args}").format(
value=value,
var=varname,
args=textwrap.indent(pprint.pformat(ARG_VALUES), prefix=" ")))
def aggregate_all_slice_combinations(spec, slice_set_names):
"""Figure out all of the possible slice groupings for a tag spec."""
slice_sets = copy.deepcopy(spec['slice_sets'])
for name in slice_set_names:
for slice_set in slice_sets[name]:
slice_set['set_name'] = name
slices_grouped_but_not_keyed = [slice_sets[name] for name in slice_set_names]
all_slice_combos = list(itertools.product(*slices_grouped_but_not_keyed))
return all_slice_combos
def build_name_from_slices(format_string, slices, args, is_dockerfile=False):
"""Build the tag name (cpu-devel...) from a list of slices."""
name_formatter = copy.deepcopy(args)
name_formatter.update({s['set_name']: s['add_to_name'] for s in slices})
name_formatter.update({
s['set_name']: s['dockerfile_exclusive_name']
for s in slices
if is_dockerfile and 'dockerfile_exclusive_name' in s
})
name = format_string.format(**name_formatter)
return name
def parse_build_arg(arg):
key, sep, value = arg.partition('=')
return key, value
def update_args_dict(args_dict, updater, keep_original=False):
"""Update a dict of arg values with more values from a list or dict."""
def _get_env(env_var, default):
if env_var in os.environ:
return os.environ[env_var]
return default
if isinstance(updater, list):
for arg in updater:
key, sep, value = arg.partition('=')
if sep == '=':
# args_dict[key] = interpolate_arg(value)
if key not in args_dict or not keep_original:
args_dict[key] = _get_env(key, str(value))
if isinstance(updater, dict):
for key, value in updater.items():
# args_dict[key] = interpolate_arg(value)
if key not in args_dict or not keep_original:
args_dict[key] = _get_env(key, str(value))
return args_dict
def get_slice_sets_and_required_args(slice_sets, tag_spec):
"""Extract used-slice-sets and required CLI arguments from a spec string.
For example, {FOO}{bar}{bat} finds FOO, bar, and bat. Assuming bar and bat
are both named slice sets, FOO must be specified on the command line.
Args:
slice_sets: Dict of named slice sets
tag_spec: The tag spec string, e.g. {_FOO}{blep}
Returns:
(used_slice_sets, required_args), a tuple of lists
"""
required_args = []
used_slice_sets = []
extract_bracketed_words = re.compile(r'\{([^}]+)\}')
possible_args_or_slice_set_names = extract_bracketed_words.findall(tag_spec)
for name in possible_args_or_slice_set_names:
if name in slice_sets:
used_slice_sets.append(name)
else:
required_args.append(name)
return (used_slice_sets, required_args)
# IDEAL:
# BAZEL_BUILD_DIR can EITHER be provided via environment variable, or via --run_arg.
# If both are present, use env.BAZEL_BUILD_DIR.
def gather_tag_args(slices, cli_input_args, required_args=None, spec_field='args', cmd_opt=None):
"""Build a dictionary of all the CLI and slice-specified args for a tag."""
if cmd_opt is None:
cmd_opt = spec_field.rstrip('s')
args = dict()
for s in slices:
if spec_field in s:
args = update_args_dict(args, s[spec_field])
# Only keep environment variables that have been "declared" in the spec.yml file.
# e.g.
# CHECKOUT_TF_SRC=
for env_var, env_value in os.environ.items():
if env_var in args:
print("> Using environment variable {env}={val} for --{cmd_opt}".format(
env=env_var,
val=env_value,
cmd_opt=cmd_opt))
args[env_var] = env_value
args = update_args_dict(args, cli_input_args)
if required_args is not None:
for arg in required_args:
if arg not in args:
logger.error(('> Error: {arg} is not a valid slice_set, and also isn\'t an arg '
'provided on the command line. If it is an arg, please specify '
'it with --{cmd_opt}. If not, check the slice_sets list.'.format(
arg=arg,
cmd_opt=cmd_opt,
)))
exit(1)
return args
def gather_slice_list_items(slices, key):
"""For a list of slices, get the flattened list of all of a certain key."""
return list(itertools.chain(*[s[key] for s in slices if key in s]))
def find_first_slice_value(slices, key):
"""For a list of slices, get the first value for a certain key."""
for s in slices:
if key in s and s[key] is not None:
return s[key]
return None
def assemble_tags(spec, cli_args, cli_run_args, enabled_release, all_partials):
"""Gather all the tags based on our spec.
Args:
spec: Nested dict containing full Tag spec
cli_args: List of ARG=foo arguments to pass along to Docker build
enabled_releases: List of releases to parse. Empty list = all
all_partials: Dict of every partial, for reference
Returns:
Dict of tags and how to build them
"""
tag_data = collections.defaultdict(list)
valid_names = set(spec['releases'].keys())
if enabled_release is not None and enabled_release not in valid_names:
raise DeployError("No such --release \"{release}\" found in spec.yml; valid choices for --release are: {choices}".format(
release=enabled_release,
choices=valid_names,
))
for name, release in spec['releases'].items():
for tag_spec in release['tag_specs']:
if enabled_release is not None and name != enabled_release:
logger.info('> Skipping release {name}; (!= {enabled_release})'.format(
name=name,
enabled_release=enabled_release))
continue
used_slice_sets, required_cli_args = get_slice_sets_and_required_args(
spec['slice_sets'], tag_spec)
slice_combos = aggregate_all_slice_combinations(spec, used_slice_sets)
for slices in slice_combos:
tag_args = gather_tag_args(slices, cli_args, required_cli_args)
tag_args.update(get_implicit_build_args())
run_args = gather_tag_args(slices, cli_run_args, spec_field='run_args')
tag_name = build_name_from_slices(tag_spec, slices, tag_args,
release['is_dockerfiles'])
used_partials = gather_slice_list_items(slices, 'partials')
used_tests = gather_slice_list_items(slices, 'tests')
test_runtime = find_first_slice_value(slices, 'test_runtime')
dockerfile_subdirectory = find_first_slice_value(
slices, 'dockerfile_subdirectory')
dockerfile_contents = merge_partials(spec['header'], used_partials,
all_partials)
tag_data[tag_name].append({
'release': name,
'tag_spec': tag_spec,
'is_dockerfiles': release['is_dockerfiles'],
'upload_images': release['upload_images'],
'cli_args': tag_args,
'run_args': run_args,
'dockerfile_subdirectory': dockerfile_subdirectory or '',
'partials': used_partials,
'tests': used_tests,
'test_runtime': test_runtime,
'dockerfile_contents': dockerfile_contents,
})
return tag_data
def merge_partials(header, used_partials, all_partials):
"""Merge all partial contents with their header."""
used_partials = list(used_partials)
ret = '\n'.join([header] + [all_partials[u] for u in used_partials])
return ret
# def upload_in_background(hub_repository, dock, image, tag):
# """Upload a docker image (to be used by multiprocessing)."""
# image.tag(hub_repository, tag=tag)
# logging.info(dock.images.push(hub_repository, tag=tag))
def mkdir_p(path):
"""Create a directory and its parents, even if it already exists."""
try:
os.makedirs(path)
except OSError as e:
if e.errno != errno.EEXIST:
raise
def gather_existing_partials(partial_path):
"""Find and read all available partials.
Args:
partial_path (string): read partials from this directory.
Returns:
Dict[string, string] of partial short names (like "ubuntu/python" or
"bazel") to the full contents of that partial.
"""
partials = dict()
for path, _, files in os.walk(partial_path):
for name in files:
fullpath = os.path.join(path, name)
if not re.search(r'\.partial\.Dockerfile$', _b(fullpath)):
logger.info(('> Probably not a problem: skipping {}, which is not a '
'partial.').format(fullpath))
continue
# partial_dir/foo/bar.partial.Dockerfile -> foo/bar
simple_name = fullpath[len(partial_path) + 1:-len('.partial.dockerfile')]
with open(fullpath, 'r') as f:
partial_contents = f.read()
check_null_byte(fullpath, partial_contents)
partials[simple_name] = partial_contents
return partials
def check_null_byte(name, string):
null_idx = string.find('\x00')
if null_idx != -1:
logger.info("Found null byte in {name}:".format(name=name))
logger.info((
"> Before null byte:\n"
"{str}"
).format(str=textwrap.indent(
string[:null_idx],
prefix=' ',
)))
logger.info((
"> After null byte:\n"
"{str}"
).format(str=textwrap.indent(
string[null_idx:],
prefix=' ')))
def get_build_logfile(repo_tag):
return "{repo_tag}.build.log.txt".format(
repo_tag=repo_tag)
"""
If you use any of these in your --arg or in arg specifications in spec.yml,
they will get replaced with corresponding ARG_VALUES dict values.
"""
ARG_VALUES = {
'RLSCOPE_DIR': py_config.RLSCOPE_DIR,
}
def interpolate_arg(arg):
# bracket_regex = re.compile(r'\{([^}]+)\}')
# bracket_varnames = bracket_regex.findall(arg)
# if len(bracket_varnames) == 0:
# return arg
#
# for varname in bracket_varnames:
# if varname not in ARG_VALUES:
# raise RuntimeError((
# "Couldn't find ; available ARG_VALUES for use in "
# "--arg or spec.yml args are:\n"
# "{args}").format(
# args=textwrap.indent(pprint.pformat(ARG_VALUES), prefix=" ")))
return arg.format(**ARG_VALUES)
def get_docker_run_env(tag_def, env_list):
# Same as docker.from_env(); inherit current environment variables.
# env = dict(os.environ)
# if 'run_args' in tag_def:
# run_args = tag_def['run_args']
# else:
# run_args = dict()
# def _get_env(env_var, desc):
# if env_var not in os.environ and env_var in run_args:
# return run_args[env_var]
# elif env_var in os.environ:
# return os.environ[env_var]
# logger.info("> You must provide {env}=[ {desc} ]".format(
# env=env_var,
# desc=desc))
# sys.exit(1)
env = dict()
env.update(get_implicit_run_args())
if 'run_args' in tag_def:
env = update_args_dict(env, tag_def['run_args'], keep_original=True)
for var, value in env.items():
if value == '':
logger.error(("> ERROR: you must provide a value for --run_arg {var}=<VALUE> "
"(or define an environment variable); see {spec} for documentation.").format(
var=var,
spec="spec.yml"))
sys.exit(1)
for env_str in env_list:
assert '=' in env_str
var, value = re.split(r'=', env_str)
# Duplicate var?
assert var not in env
env[var] = value
# env['RLSCOPE_USER'] = get_username()
# env['RLSCOPE_UID'] = get_user_id()
# env['RLSCOPE_GID'] = get_group_id()
env['RLSCOPE_INSTALL_PREFIX'] = py_config.DOCKER_INSTALL_PREFIX
env['RLSCOPE_BUILD_PREFIX'] = py_config.DOCKER_BUILD_PREFIX
env['RLSCOPE_IS_DOCKER'] = 'yes'
return env
def get_docker_runtime(tag_def):
runtime = None
if tag_def['test_runtime'] == 'nvidia':
runtime = 'nvidia'
else:
# Use runtime=None for rocm.
runtime = None
return runtime
def get_docker_run_argv(argv):
"""
absl will preserve "extra" arguments that were unparsed by the FLAGS specification.
We pass all of these arguments directly to the "docker run" command.
This makes it easy to:
- Define additional environment variables (outside of required ones)
- Define additional volume mounts (outside of required ones)
:param argv:
argv = ["run_docker.py" "extra_argument[0]", ...]
:return:
"""
return argv[1:]
def main():
# Run all docker related commands from $ROOT/dockerfiles.
os.chdir(DOCKERFILES)
parser = argparse.ArgumentParser(description=__doc__)
if shutil.which('docker-compose') is None:
logger.error("Didn't find docker-compose on PATH; you must install it.")
sys.exit(1)
os.makedirs(_j(py_config.ROOT, 'dockerfiles/dockerfiles'), exist_ok=True)
# Copy requirements.txt into dockerfiles so it can be accessed by Dockerfile during its build.
shutil.copy(
_j(py_config.ROOT, "requirements.txt"),
_j(DOCKERFILES, "requirements.txt"),
)
# parser.add_argument('--hub_username',
# help='Dockerhub username, only used with --upload_to_hub')
# parser.add_argument(
# '--hub_password',
# help=('Dockerhub password, only used with --upload_to_hub. Use from an env param'
# 'so your password isn\'t in your history.'))
parser.add_argument('--hub_timeout', default=3600,
type=int,
help='Abort Hub upload if it takes longer than this.')
parser.add_argument('--pull', action='store_true',
help=textwrap.dedent("""
Pull pre-built RL-Scope dev environment image from DockerHub,
instead of building locally.
See --pull-image for specifying the image to pull.
"""))
parser.add_argument('--stop', action='store_true',
help=textwrap.dedent("""
Stop docker-compose containers (opposite of --deploy).
"""))
parser.add_argument('--reload', action='store_true',
help=textwrap.dedent("""
Restart running containers, even if the
underlying docker file hasn't changed.
"""))
parser.add_argument('--mps', action='store_true',
help=textwrap.dedent("""
Use CUDA multi-process service (MPS) daemon to allow multiple GPU processes to
share GPU execution simultaneously.
This option will add mps related setup to the generated docker-compose file (stack.yml).
"""))
parser.add_argument('--pull-image', default=DEFAULT_REMOTE_RLSCOPE_IMAGE_TAG,
help=textwrap.dedent("""
RL-Scope dev environment image to pull from DockerHub
using "docker pull <pull_img>"
"""))
parser.add_argument('--deploy-rlscope-drill-port',
default=DEFAULT_RLSCOPE_DRILL_PORT,
type=int,
help=('What port to run rlscope-drill web server on '
'(when running "docker stack deploy -c stack.yml rlscope")'))
parser.add_argument(
'--repository', default='tensorflow',
help='Tag local images as {repository}:tag (in addition to the '
'hub_repository, if uploading to hub)')
parser.add_argument(
'--volume',
action='append',
default=[],
help=textwrap.dedent("""\
Translates into docker --volume option.
We mount the path at the same path as it is in the host.
i.e.
# run_docker.py option:
--volume /one/two
#
# becomes
#
# `docker run` option:
--volume /one/two:/one/two
""").rstrip())
parser.add_argument(
'--publish',
action='append',
default=[],
help=textwrap.dedent("""\
Translates into docker --publish option; e.g.
--publish <HOST_PORT>:<CONTAINER_PORT>
""").rstrip())
parser.add_argument(
'--env',
'-e',
action='append',
default=[],
help=textwrap.dedent("""\
Translates into docker --env option.
""").rstrip())
parser.add_argument(
'--hub_repository',
help='Push tags to this Docker Hub repository, e.g. tensorflow/tensorflow')
parser.add_argument(
'--debug',
action='store_true',
help=textwrap.dedent("""
In the generated dockerfiles, print start/end markers for the partial files its composed of; for e.g.:
START: dockerfiles/partials/ubuntu/install_cuda_10_1.partial.Dockerfile
RUN ...
RUN ...
...
END: dockerfiles/partials/ubuntu/install_cuda_10_1.partial.Dockerfile
"""))
parser.add_argument(
'--pdb',
action='store_true',
help=textwrap.dedent("""
Debug: breakpoint on error.
"""))
# parser.add_argument(
# '--upload_to_hub', '-u',
# help=('Push built images to Docker Hub (you must also provide --hub_username, '
# '--hub_password, and --hub_repository)'),
# )
parser.add_argument(
'--construct_dockerfiles', '-d',
action='store_true',
help='Do not build images')
parser.add_argument(
'--keep_temp_dockerfiles', '-k',
action='store_true',
help='Retain .temp.Dockerfiles created while building images.')
parser.add_argument(
'--build_images', '-b',
action='store_true',
help='Do not build images')
parser.add_argument(
'--deploy',
action='store_true',
help=
textwrap.dedent("""\
Deploy the RL-Scope development environment using
"docker stack deploy -c stack.yml rlscope".
""".format(USER=get_username())))
parser.add_argument(
'--run',
action='store_true',
help='Run built images; use --deploy if you want to deploy the whole RL-Scope development environment')
# parser.add_argument(
# '--run_tests_path',
# help=('Execute test scripts on generated Dockerfiles before pushing them. '
# 'Flag value must be a full path to the "tests" directory, which is usually'
# ' $(realpath ./tests). A failed tests counts the same as a failed build.'))
parser.add_argument(
'--stop_on_failure',
action='store_true',
help=('Stop processing tags if any one build fails. If False or not specified, '
'failures are reported but do not affect the other images.'))
parser.add_argument(
'--dry_run', '-n',
action='store_true',
help='Do not build or deploy anything at all.',
)
parser.add_argument(
'--exclude_tags_matching', '-x',
help=('Regular expression that skips processing on any tag it matches. Must '
'match entire string, e.g. ".*gpu.*" ignores all GPU tags.'),
)
parser.add_argument(
'--only_tags_matching', '-i',
help=('Regular expression that skips processing on any tag it does not match. '
'Must match entire string, e.g. ".*gpu.*" includes only GPU tags.'),
)
parser.add_argument(
'--dockerfile_dir', '-o',
default='./dockerfiles',
help='Path to an output directory for Dockerfiles.'
' Will be created if it doesn\'t exist.'
' Existing files in this directory will be deleted when new Dockerfiles'
' are made.',
)
parser.add_argument(
'--partial_dir', '-p',
default='./partials',
help='Path to a directory containing foo.partial.Dockerfile partial files.'
' can have subdirectories, e.g. "bar/baz.partial.Dockerfile".',
)
parser.add_argument(
'--release', '-r',
default='rlscope',
help='Set of releases to build and tag. Defaults to every release type.',
)
parser.add_argument(
'--arg', '-a', default=[], action='append',
help=('Extra build arguments. These are used for expanding tag names if needed '
'(e.g. --arg _TAG_PREFIX=foo) and for using as build arguments (unused '
'args will print a warning).'),
)
parser.add_argument(
'--run_arg', default=[], action='append',
help=('Extra container run arguments (NOT build).'))
parser.add_argument(
'--spec_file', '-s',
default='./spec.yml',
help='Path to the YAML specification file',
)
parser.add_argument(
'--output_stack_yml',
default='./stack.yml',
help='Path to the generated YAML "Docker Compose" file for '
'use with "docker stack deploy -c stack.yml rlscope"',
)
argv = list(sys.argv)
args, extra_argv = parser.parse_known_args()
# NOTE: If this fails, you need to enable nvidia-persistend daemon.
nvidia_gpu_query.check_nvidia_smi(exit_if_fail=True)
if not args.stop and not args.pull and not args.reload and not args.build_images and not args.run and not args.deploy:
# Default options:
args.construct_dockerfiles = True
args.build_images = True
args.deploy = True
if args.deploy and args.run:
parser.error(
"Provide either --deploy or --run. "
"Use --deploy to deploy the RL-Scope development environment (probably what you want)")
try:
assembler = Assembler(parser, argv, args, extra_argv)
assembler.run()
except DeployError as e:
if args.debug:
raise e
print("ERROR: {e}".format(e=e), file=sys.stderr)
sys.exit(1)
class Assembler:
def __init__(self, parser, argv, args, extra_argv):
self.parser = parser
self.argv = argv
self.args = args
self.extra_argv = extra_argv
def read_spec_yml(self):
args = self.args
# Read the full spec file, used for everything
with open(args.spec_file, 'r') as spec_file:
tag_spec = yaml.load(spec_file, Loader=yaml.FullLoader)
for arg in args.arg:
key, value = parse_build_arg(arg)
validate_interpolate_arg(value)
for run_arg in args.run_arg:
key, value = parse_build_arg(run_arg)
validate_interpolate_arg(value)
# Get existing partial contents
partials = gather_existing_partials(args.partial_dir)
# Abort if spec.yaml is invalid
schema = yaml.load(SCHEMA_TEXT, Loader=yaml.FullLoader)
v = TfDockerTagValidator(schema, partials=partials)
if not v.validate(tag_spec):
logger.error('> Error: {} is an invalid spec! The errors are:'.format(
args.spec_file))
logger.error(yaml.dump(v.errors, indent=2))
exit(1)
tag_spec = v.normalized(tag_spec)
# Assemble tags and images used to build them
run_arg_required = []
for run_arg in args.run_arg:
var, value = parse_build_arg(run_arg)
if is_required_run_arg(var):
run_arg_required.append(run_arg)
all_tags = assemble_tags(tag_spec, args.arg, run_arg_required, args.release, partials)
return all_tags
def docker_build(self, dockerfile, repo_tag, tag_def, debug):
args = self.args
build_kwargs = dict(
timeout=args.hub_timeout,
path=DOCKERFILES,
# path='dockerfiles',
dockerfile=_a(dockerfile),
buildargs=tag_def['cli_args'],
tag=repo_tag,
)
if args.debug:
print("> dock.images.build")
print(textwrap.indent(pprint.pformat(build_kwargs), prefix=" "))
build_cmd = get_docker_cmdline('build', **build_kwargs)
logger.info(get_cmd_string(build_cmd, show_cwd=True))
build_output_generator = self.dock_cli.build(decode=True, **build_kwargs)
response = tee_docker(
build_output_generator,
file=get_build_logfile(repo_tag),
debug=debug)
# I've seen "docker build" fail WITHOUT any error indication
# (e.g. return code, raising dockers.errors.APIError),
# so just grep for "error" in the response['message'].
check_docker_response(response, dockerfile, repo_tag, cmd=build_cmd)
image = self.dock.images.get(repo_tag)
return image
def docker_run(self, image, tag_def, extra_argv):
# Run the container.
args = self.args
docker_run_env = get_docker_run_env(tag_def, args.env)
rlscope_volumes = get_rlscope_volumes(args, docker_run_env, args.volume)
runtime = get_docker_runtime(tag_def)
run_kwargs = dict(
image=image,
# command='/tests/' + test,
working_dir='/',
log_config={'type': 'journald'},
detach=True,
stderr=True,
stdout=True,
environment=docker_run_env,
volumes=rlscope_volumes,
# volumes={
# args.run_tests_path: {
# 'bind': '/tests',
# 'mode': 'ro'
# }
# },
remove=True,
cap_add=['SYS_ADMIN', 'SYS_PTRACE'],
security_opt=['seccomp=unconfined'],
runtime=runtime,
name="rlscope",
)
if tag_def['test_runtime'] == 'rocm':
def device_opt(path):
return '{path}:{path}:rwm'.format(path=path)
# --device=/dev/kfd --device=/dev/dri
run_kwargs['devices'] = [
device_opt('/dev/kfd'),
device_opt('/dev/dri'),
]
# --group-add video
run_kwargs['group_add'] = [
'video',
]
docker_run_argv = extra_argv
cmd = get_docker_cmdline('run', docker_run_argv, **run_kwargs)
logger.info(get_cmd_string(cmd))
subprocess.run(cmd)
# Q: Save output?
# def _is_rlscope_deployed(self):
# self.docker_ps_rlscope_bash() is not
def docker_ps_rlscope_bash(self):
# logging.info("(1) run docker ps")
p = subprocess.run("docker ps | grep rlscope_bash",
shell=True,
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
# check=True,
)
# I've seen this fail with returncode=1, but NO error message...