##// END OF EJS Templates
Backport PR #10408: Speed up shutdown by eliminating unnecessary sleeping...
Thomas Kluyver -
Show More
@@ -1,272 +1,280 b''
1 """Magic functions for running cells in various scripts."""
1 """Magic functions for running cells in various scripts."""
2 from __future__ import print_function
2 from __future__ import print_function
3
3
4 # Copyright (c) IPython Development Team.
4 # Copyright (c) IPython Development Team.
5 # Distributed under the terms of the Modified BSD License.
5 # Distributed under the terms of the Modified BSD License.
6
6
7 import errno
7 import errno
8 import os
8 import os
9 import sys
9 import sys
10 import signal
10 import signal
11 import time
11 import time
12 from subprocess import Popen, PIPE
12 from subprocess import Popen, PIPE
13 import atexit
13 import atexit
14
14
15 from IPython.core import magic_arguments
15 from IPython.core import magic_arguments
16 from IPython.core.magic import (
16 from IPython.core.magic import (
17 Magics, magics_class, line_magic, cell_magic
17 Magics, magics_class, line_magic, cell_magic
18 )
18 )
19 from IPython.lib.backgroundjobs import BackgroundJobManager
19 from IPython.lib.backgroundjobs import BackgroundJobManager
20 from IPython.utils import py3compat
20 from IPython.utils import py3compat
21 from IPython.utils.process import arg_split
21 from IPython.utils.process import arg_split
22 from traitlets import List, Dict, default
22 from traitlets import List, Dict, default
23
23
24 #-----------------------------------------------------------------------------
24 #-----------------------------------------------------------------------------
25 # Magic implementation classes
25 # Magic implementation classes
26 #-----------------------------------------------------------------------------
26 #-----------------------------------------------------------------------------
27
27
28 def script_args(f):
28 def script_args(f):
29 """single decorator for adding script args"""
29 """single decorator for adding script args"""
30 args = [
30 args = [
31 magic_arguments.argument(
31 magic_arguments.argument(
32 '--out', type=str,
32 '--out', type=str,
33 help="""The variable in which to store stdout from the script.
33 help="""The variable in which to store stdout from the script.
34 If the script is backgrounded, this will be the stdout *pipe*,
34 If the script is backgrounded, this will be the stdout *pipe*,
35 instead of the stderr text itself.
35 instead of the stderr text itself.
36 """
36 """
37 ),
37 ),
38 magic_arguments.argument(
38 magic_arguments.argument(
39 '--err', type=str,
39 '--err', type=str,
40 help="""The variable in which to store stderr from the script.
40 help="""The variable in which to store stderr from the script.
41 If the script is backgrounded, this will be the stderr *pipe*,
41 If the script is backgrounded, this will be the stderr *pipe*,
42 instead of the stderr text itself.
42 instead of the stderr text itself.
43 """
43 """
44 ),
44 ),
45 magic_arguments.argument(
45 magic_arguments.argument(
46 '--bg', action="store_true",
46 '--bg', action="store_true",
47 help="""Whether to run the script in the background.
47 help="""Whether to run the script in the background.
48 If given, the only way to see the output of the command is
48 If given, the only way to see the output of the command is
49 with --out/err.
49 with --out/err.
50 """
50 """
51 ),
51 ),
52 magic_arguments.argument(
52 magic_arguments.argument(
53 '--proc', type=str,
53 '--proc', type=str,
54 help="""The variable in which to store Popen instance.
54 help="""The variable in which to store Popen instance.
55 This is used only when --bg option is given.
55 This is used only when --bg option is given.
56 """
56 """
57 ),
57 ),
58 ]
58 ]
59 for arg in args:
59 for arg in args:
60 f = arg(f)
60 f = arg(f)
61 return f
61 return f
62
62
63 @magics_class
63 @magics_class
64 class ScriptMagics(Magics):
64 class ScriptMagics(Magics):
65 """Magics for talking to scripts
65 """Magics for talking to scripts
66
66
67 This defines a base `%%script` cell magic for running a cell
67 This defines a base `%%script` cell magic for running a cell
68 with a program in a subprocess, and registers a few top-level
68 with a program in a subprocess, and registers a few top-level
69 magics that call %%script with common interpreters.
69 magics that call %%script with common interpreters.
70 """
70 """
71 script_magics = List(
71 script_magics = List(
72 help="""Extra script cell magics to define
72 help="""Extra script cell magics to define
73
73
74 This generates simple wrappers of `%%script foo` as `%%foo`.
74 This generates simple wrappers of `%%script foo` as `%%foo`.
75
75
76 If you want to add script magics that aren't on your path,
76 If you want to add script magics that aren't on your path,
77 specify them in script_paths
77 specify them in script_paths
78 """,
78 """,
79 ).tag(config=True)
79 ).tag(config=True)
80 @default('script_magics')
80 @default('script_magics')
81 def _script_magics_default(self):
81 def _script_magics_default(self):
82 """default to a common list of programs"""
82 """default to a common list of programs"""
83
83
84 defaults = [
84 defaults = [
85 'sh',
85 'sh',
86 'bash',
86 'bash',
87 'perl',
87 'perl',
88 'ruby',
88 'ruby',
89 'python',
89 'python',
90 'python2',
90 'python2',
91 'python3',
91 'python3',
92 'pypy',
92 'pypy',
93 ]
93 ]
94 if os.name == 'nt':
94 if os.name == 'nt':
95 defaults.extend([
95 defaults.extend([
96 'cmd',
96 'cmd',
97 ])
97 ])
98
98
99 return defaults
99 return defaults
100
100
101 script_paths = Dict(
101 script_paths = Dict(
102 help="""Dict mapping short 'ruby' names to full paths, such as '/opt/secret/bin/ruby'
102 help="""Dict mapping short 'ruby' names to full paths, such as '/opt/secret/bin/ruby'
103
103
104 Only necessary for items in script_magics where the default path will not
104 Only necessary for items in script_magics where the default path will not
105 find the right interpreter.
105 find the right interpreter.
106 """
106 """
107 ).tag(config=True)
107 ).tag(config=True)
108
108
109 def __init__(self, shell=None):
109 def __init__(self, shell=None):
110 super(ScriptMagics, self).__init__(shell=shell)
110 super(ScriptMagics, self).__init__(shell=shell)
111 self._generate_script_magics()
111 self._generate_script_magics()
112 self.job_manager = BackgroundJobManager()
112 self.job_manager = BackgroundJobManager()
113 self.bg_processes = []
113 self.bg_processes = []
114 atexit.register(self.kill_bg_processes)
114 atexit.register(self.kill_bg_processes)
115
115
116 def __del__(self):
116 def __del__(self):
117 self.kill_bg_processes()
117 self.kill_bg_processes()
118
118
119 def _generate_script_magics(self):
119 def _generate_script_magics(self):
120 cell_magics = self.magics['cell']
120 cell_magics = self.magics['cell']
121 for name in self.script_magics:
121 for name in self.script_magics:
122 cell_magics[name] = self._make_script_magic(name)
122 cell_magics[name] = self._make_script_magic(name)
123
123
124 def _make_script_magic(self, name):
124 def _make_script_magic(self, name):
125 """make a named magic, that calls %%script with a particular program"""
125 """make a named magic, that calls %%script with a particular program"""
126 # expand to explicit path if necessary:
126 # expand to explicit path if necessary:
127 script = self.script_paths.get(name, name)
127 script = self.script_paths.get(name, name)
128
128
129 @magic_arguments.magic_arguments()
129 @magic_arguments.magic_arguments()
130 @script_args
130 @script_args
131 def named_script_magic(line, cell):
131 def named_script_magic(line, cell):
132 # if line, add it as cl-flags
132 # if line, add it as cl-flags
133 if line:
133 if line:
134 line = "%s %s" % (script, line)
134 line = "%s %s" % (script, line)
135 else:
135 else:
136 line = script
136 line = script
137 return self.shebang(line, cell)
137 return self.shebang(line, cell)
138
138
139 # write a basic docstring:
139 # write a basic docstring:
140 named_script_magic.__doc__ = \
140 named_script_magic.__doc__ = \
141 """%%{name} script magic
141 """%%{name} script magic
142
142
143 Run cells with {script} in a subprocess.
143 Run cells with {script} in a subprocess.
144
144
145 This is a shortcut for `%%script {script}`
145 This is a shortcut for `%%script {script}`
146 """.format(**locals())
146 """.format(**locals())
147
147
148 return named_script_magic
148 return named_script_magic
149
149
150 @magic_arguments.magic_arguments()
150 @magic_arguments.magic_arguments()
151 @script_args
151 @script_args
152 @cell_magic("script")
152 @cell_magic("script")
153 def shebang(self, line, cell):
153 def shebang(self, line, cell):
154 """Run a cell via a shell command
154 """Run a cell via a shell command
155
155
156 The `%%script` line is like the #! line of script,
156 The `%%script` line is like the #! line of script,
157 specifying a program (bash, perl, ruby, etc.) with which to run.
157 specifying a program (bash, perl, ruby, etc.) with which to run.
158
158
159 The rest of the cell is run by that program.
159 The rest of the cell is run by that program.
160
160
161 Examples
161 Examples
162 --------
162 --------
163 ::
163 ::
164
164
165 In [1]: %%script bash
165 In [1]: %%script bash
166 ...: for i in 1 2 3; do
166 ...: for i in 1 2 3; do
167 ...: echo $i
167 ...: echo $i
168 ...: done
168 ...: done
169 1
169 1
170 2
170 2
171 3
171 3
172 """
172 """
173 argv = arg_split(line, posix = not sys.platform.startswith('win'))
173 argv = arg_split(line, posix = not sys.platform.startswith('win'))
174 args, cmd = self.shebang.parser.parse_known_args(argv)
174 args, cmd = self.shebang.parser.parse_known_args(argv)
175
175
176 try:
176 try:
177 p = Popen(cmd, stdout=PIPE, stderr=PIPE, stdin=PIPE)
177 p = Popen(cmd, stdout=PIPE, stderr=PIPE, stdin=PIPE)
178 except OSError as e:
178 except OSError as e:
179 if e.errno == errno.ENOENT:
179 if e.errno == errno.ENOENT:
180 print("Couldn't find program: %r" % cmd[0])
180 print("Couldn't find program: %r" % cmd[0])
181 return
181 return
182 else:
182 else:
183 raise
183 raise
184
184
185 if not cell.endswith('\n'):
185 if not cell.endswith('\n'):
186 cell += '\n'
186 cell += '\n'
187 cell = cell.encode('utf8', 'replace')
187 cell = cell.encode('utf8', 'replace')
188 if args.bg:
188 if args.bg:
189 self.bg_processes.append(p)
189 self.bg_processes.append(p)
190 self._gc_bg_processes()
190 self._gc_bg_processes()
191 if args.out:
191 if args.out:
192 self.shell.user_ns[args.out] = p.stdout
192 self.shell.user_ns[args.out] = p.stdout
193 if args.err:
193 if args.err:
194 self.shell.user_ns[args.err] = p.stderr
194 self.shell.user_ns[args.err] = p.stderr
195 self.job_manager.new(self._run_script, p, cell, daemon=True)
195 self.job_manager.new(self._run_script, p, cell, daemon=True)
196 if args.proc:
196 if args.proc:
197 self.shell.user_ns[args.proc] = p
197 self.shell.user_ns[args.proc] = p
198 return
198 return
199
199
200 try:
200 try:
201 out, err = p.communicate(cell)
201 out, err = p.communicate(cell)
202 except KeyboardInterrupt:
202 except KeyboardInterrupt:
203 try:
203 try:
204 p.send_signal(signal.SIGINT)
204 p.send_signal(signal.SIGINT)
205 time.sleep(0.1)
205 time.sleep(0.1)
206 if p.poll() is not None:
206 if p.poll() is not None:
207 print("Process is interrupted.")
207 print("Process is interrupted.")
208 return
208 return
209 p.terminate()
209 p.terminate()
210 time.sleep(0.1)
210 time.sleep(0.1)
211 if p.poll() is not None:
211 if p.poll() is not None:
212 print("Process is terminated.")
212 print("Process is terminated.")
213 return
213 return
214 p.kill()
214 p.kill()
215 print("Process is killed.")
215 print("Process is killed.")
216 except OSError:
216 except OSError:
217 pass
217 pass
218 except Exception as e:
218 except Exception as e:
219 print("Error while terminating subprocess (pid=%i): %s" \
219 print("Error while terminating subprocess (pid=%i): %s" \
220 % (p.pid, e))
220 % (p.pid, e))
221 return
221 return
222 out = py3compat.bytes_to_str(out)
222 out = py3compat.bytes_to_str(out)
223 err = py3compat.bytes_to_str(err)
223 err = py3compat.bytes_to_str(err)
224 if args.out:
224 if args.out:
225 self.shell.user_ns[args.out] = out
225 self.shell.user_ns[args.out] = out
226 else:
226 else:
227 sys.stdout.write(out)
227 sys.stdout.write(out)
228 sys.stdout.flush()
228 sys.stdout.flush()
229 if args.err:
229 if args.err:
230 self.shell.user_ns[args.err] = err
230 self.shell.user_ns[args.err] = err
231 else:
231 else:
232 sys.stderr.write(err)
232 sys.stderr.write(err)
233 sys.stderr.flush()
233 sys.stderr.flush()
234
234
235 def _run_script(self, p, cell):
235 def _run_script(self, p, cell):
236 """callback for running the script in the background"""
236 """callback for running the script in the background"""
237 p.stdin.write(cell)
237 p.stdin.write(cell)
238 p.stdin.close()
238 p.stdin.close()
239 p.wait()
239 p.wait()
240
240
241 @line_magic("killbgscripts")
241 @line_magic("killbgscripts")
242 def killbgscripts(self, _nouse_=''):
242 def killbgscripts(self, _nouse_=''):
243 """Kill all BG processes started by %%script and its family."""
243 """Kill all BG processes started by %%script and its family."""
244 self.kill_bg_processes()
244 self.kill_bg_processes()
245 print("All background processes were killed.")
245 print("All background processes were killed.")
246
246
247 def kill_bg_processes(self):
247 def kill_bg_processes(self):
248 """Kill all BG processes which are still running."""
248 """Kill all BG processes which are still running."""
249 if not self.bg_processes:
250 return
249 for p in self.bg_processes:
251 for p in self.bg_processes:
250 if p.poll() is None:
252 if p.poll() is None:
251 try:
253 try:
252 p.send_signal(signal.SIGINT)
254 p.send_signal(signal.SIGINT)
253 except:
255 except:
254 pass
256 pass
255 time.sleep(0.1)
257 time.sleep(0.1)
258 self._gc_bg_processes()
259 if not self.bg_processes:
260 return
256 for p in self.bg_processes:
261 for p in self.bg_processes:
257 if p.poll() is None:
262 if p.poll() is None:
258 try:
263 try:
259 p.terminate()
264 p.terminate()
260 except:
265 except:
261 pass
266 pass
262 time.sleep(0.1)
267 time.sleep(0.1)
268 self._gc_bg_processes()
269 if not self.bg_processes:
270 return
263 for p in self.bg_processes:
271 for p in self.bg_processes:
264 if p.poll() is None:
272 if p.poll() is None:
265 try:
273 try:
266 p.kill()
274 p.kill()
267 except:
275 except:
268 pass
276 pass
269 self._gc_bg_processes()
277 self._gc_bg_processes()
270
278
271 def _gc_bg_processes(self):
279 def _gc_bg_processes(self):
272 self.bg_processes = [p for p in self.bg_processes if p.poll() is None]
280 self.bg_processes = [p for p in self.bg_processes if p.poll() is None]
General Comments 0
You need to be logged in to leave comments. Login now