1# (c) Stefan Countryman, 2019 

2 

3import os 

4import logging 

5import atexit 

6from subprocess import Popen, PIPE 

7from plumbum.cmd import jupyter 

8from llama.classes import COLOR 

9from llama.utils import LOGDIR, OUTPUT_DIR 

10from llama.cli import get_logging_cli, CliParser 

11from llama.serve.jupyter import ( 

12 __doc__, 

13 load_notebook, 

14 README, 

15 VIEW_LOGS, 

16 LLAMA_JUP_PORT, 

17 SERVER_LIST_FILE, 

18) 

19from llama.serve.gui.domain import LLAMA_DOMAIN 

20from llama.com.slack import alert_maintainers 

21 

22LOGGER = logging.getLogger(__name__) 

23JUPYTER_LOG = os.path.join(LOGDIR, 'jupyter.log') 

24JUPYTER_FILES = os.path.join(OUTPUT_DIR, 'jupyter') 

25DEFAULT_IP = '0.0.0.0' 

26 

27 

28def get_parser(): 

29 """Get CLI Parser.""" 

30 parser = CliParser(description=__doc__, 

31 parents=(get_logging_cli(JUPYTER_LOG, 'info'),)) 

32 parser.add_argument("--notebook-dir", default=OUTPUT_DIR, help=f""" 

33 Where to store jupyter notebook files. By default, they will be saved 

34 in the LLAMA data directory in {OUTPUT_DIR}""") 

35 parser.add_argument("--ip", default=DEFAULT_IP, help=f""" 

36 The IP address the server will listen on. (default: {DEFAULT_IP})""") 

37 parser.add_argument("-a", "--alert-maintainers", action='store_true', 

38 help=""" 

39 If provided, use ``llama.com.slack.alert_maintainers`` to 

40 message LLAMA maintainers with a list of active Jupyter notebooks and 

41 their login tokens. This allows 

42 Slack users to access the notebook at the provided URL. These tokens 

43 can be used to log in to this Jupyter notebook (and any others running 

44 on this server/container). BE CAREFUL WITH THESE TOKENS! They provide 

45 full access to the Jupyter notebook; you should probably only use this 

46 in production, and make sure not to share those tokens. You should also 

47 make sure to regenerate those tokens regularly. Note also that the 

48 script will fail if you try to alert maintainers without providing 

49 valid Slack credentials (see ``llama.com.slack``).""") 

50 parser.add_argument("-w", "--writeable-docs", action='store_true', help=""" 

51 If provided, set any documentation notebooks (like README.ipynb) to 

52 writeable. Use this for development mode and then commit the notebook 

53 with ``llama dev upload -g README.ipynb`` (from the notebook directory) 

54 and update the remote URL for ``llama.serve.jupyter.README`` with the 

55 printed URL.""") 

56 parser.add_argument("-k", "--keep-docs", action='store_true', help=""" 

57 If provided, don't bother downloading new README if it's stored 

58 locally; in other words, keep existing documentation if present so as 

59 not to accidentally delete development work. As with 

60 ``--writeable-docs``, this should probably only be used in development 

61 mode.""") 

62 parser.add_argument("--show-hidden", action='store_true', help=""" 

63 If provided, show hidden files in the file browser. This is useful for 

64 modifying metadata rider files used by the pipeline.""") 

65 return parser 

66 

67 

68def main(): 

69 """Run CLI.""" 

70 parser = get_parser() 

71 args = parser.parse_args() 

72 if not os.path.isdir(args.notebook_dir): 

73 os.makedirs(args.notebook_dir) 

74 load_notebook(README, 'README.ipynb', readonly=not args.writeable_docs, 

75 clobber=not args.keep_docs) 

76 load_notebook(VIEW_LOGS, 'VIEW_LOGS.ipynb', readonly=not args.writeable_docs, 

77 clobber=not args.keep_docs) 

78 cmd = [ 

79 'jupyter', 

80 'notebook', 

81 '--no-browser', 

82 '--ip', 

83 args.ip, 

84 '--port', 

85 LLAMA_JUP_PORT, 

86 '--port-retries=0', 

87 f'--ContentsManager.allow_hidden={args.show_hidden}', 

88 '--notebook-dir', 

89 args.notebook_dir, 

90 '--log-level', 

91 args.verbosity.upper(), 

92 '-y', 

93 '--allow-root', 

94 ] 

95 LOGGER.info("STARTING JUPYTER NOTEBOOK WITH COMMAND: %s", ' '.join(cmd)) 

96 with open(os.devnull, 'wb') as devnull: 

97 proc = Popen(cmd, stdout=devnull, stderr=PIPE) 

98 atexit.register(proc.kill) # kill the jupyter notebook on exit 

99 while True: 

100 output = proc.stderr.readline().decode().strip() 

101 if output == '' and proc.poll() is not None: 

102 break 

103 if output: 

104 LOGGER.info(COLOR.MAGENTA + output + COLOR.CLEAR) 

105 if 'Or copy and paste one of these URLs' in output: 

106 notebooks = jupyter['notebook', 'list']() 

107 notebooks = notebooks.replace( 

108 f"{args.ip}:{LLAMA_JUP_PORT}", 

109 f"{LLAMA_DOMAIN}:{LLAMA_JUP_PORT}", 

110 ) 

111 with open(SERVER_LIST_FILE, 'w') as server_list: 

112 for line in notebooks.split('\n'): 

113 if 'token=' in line: 

114 server_list.write(line+'\n') 

115 atexit.register(lambda: os.unlink(SERVER_LIST_FILE)) 

116 if args.alert_maintainers: 

117 alert_maintainers("NOTEBOOK STARTED\n" + notebooks) 

118 return proc.poll() 

119 

120 

121if __name__ == '__main__': 

122 exit(main())