-
Notifications
You must be signed in to change notification settings - Fork 0
/
bot.py
308 lines (244 loc) · 9.26 KB
/
bot.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
from dotenv import load_dotenv
import os
import logging
from telegram import Update
from telegram.ext import (
filters,
MessageHandler,
ApplicationBuilder,
CommandHandler,
ContextTypes,
)
import docker
import subprocess
import io
import pathlib
import notebooks
# For all messaging, make lines ~50 columns long.
# get token from environment
load_dotenv()
TOKEN = os.getenv("TELEGRAM_TOKEN")
assert TOKEN is not None
# configure logging
logging.basicConfig(
format="%(asctime)s - %(name)s - %(levelname)s - %(name)s", level=logging.INFO
)
# Map(Alias -> CID)
running = {}
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""
create a local table of notebooks
Greet the user.
Create a new local file that saves
the map of names to notebooks.
By default, this file is called `notebooks.csv`.
"""
assert update.effective_chat is not None
# create notebooks.csv for INIT command
notebooks.CSV_FILEPATH.touch()
await context.bot.send_message(chat_id=update.effective_chat.id, text="Local table created.")
async def init(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""
add a new notebook
/init ALIAS TYPE → creates new notebook ALIAS of type TYPE
ALIAS is a short name for your new notebook.
You will use it with /run.
TYPE is one of the available types of notebooks.
You can find the whole list with /ls.
"""
assert update.effective_chat is not None
assert context.args is not None
if len(context.args) != 2:
await context.bot.send_message(
chat_id=update.effective_chat.id, text="Provide a notebook name and type."
)
return
alias, name = context.args[0], context.args[1]
if name not in notebook_types:
await context.bot.send_message(
chat_id=update.effective_chat.id, text="Invalid notebook type. Consult /ls for available notebook types."
)
return
notebook_created = notebooks.put(alias, name)
print(notebook_created)
if notebook_created:
await context.bot.send_message(
chat_id=update.effective_chat.id, text=f"{name} is now saved as {alias}."
)
else:
await context.bot.send_message(
chat_id=update.effective_chat.id, text=f"{alias} is already taken! Try again."
)
async def run(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""
run a notebook
/run ALIAS → starts an instance of a notebook
ALIAS refers to one of the existing notebooks on your machine.
These are created by calling the /init command
and stored in a local CSV file.
"""
assert update.effective_chat is not None
if context.args is None or len(context.args) == 0:
await context.bot.send_message(
chat_id=update.effective_chat.id, text="Provide a notebook name."
)
return
assert len(context.args) == 1
notebook_name = context.args[0]
#TODO: check if the cid file already exists
# this means that the notebook is already running:
# inform the user about this, send a message, don't create a new instance
# Later, find out how to handle multiple instances of the same image
# (multiple instances of a PyTorch notebook)
# print(f'Running {docker.docker_command}')
notebook_type = notebooks.get(notebook_name)
if notebook_type is None:
await context.bot.send_message(
chat_id=update.effective_chat.id, text="Notebook doesn't exist. Create it with /init."
)
return
await context.bot.send_message(
chat_id=update.effective_chat.id,
text=f'Starting notebook: {notebook_name}'
)
docker_process = subprocess.Popen(
' '.join(docker.run(HOST_PORT=60000, notebook_name=notebook_name, notebook_type=notebook_type)),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
shell=True
)
# parse jupyter token
output = ''
for line in io.TextIOWrapper(docker_process.stderr, encoding="utf-8"):
print(line)
output = line
if 'http://127.0.0.1:8888' in line:
break
output = docker.get_token(output)
# get DOCKER PID
with open(docker.CIDFILE, 'r') as cidfile:
pid = cidfile.read()
global running
running[notebook_name] = pid
# remove CIDFILE
pathlib.Path.unlink(docker.CIDFILE)
await context.bot.send_message(chat_id=update.effective_chat.id, text=f'Your token:')
await context.bot.send_message(chat_id=update.effective_chat.id, text=output)
async def kill(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""
kill a running instance of a notebook
/kill ALIAS → stops the Docker container for the notebook
"""
assert update.effective_chat is not None
if context.args is None or len(context.args) != 1:
await context.bot.send_message(
chat_id=update.effective_chat.id, text="Provide a notebook name."
)
return
notebook_name = context.args[0]
global running
if notebook_name not in running.keys():
await context.bot.send_message(
chat_id=update.effective_chat.id,
text='No such notebook currently running.'
)
return
# kill the docker container
subprocess.Popen(
docker.docker_kill_command(running[notebook_name]),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
)
# remove the key-value pair from the dictionary
del running[notebook_name]
response = f"Killed notebook: {notebook_name}"
await context.bot.send_message(chat_id=update.effective_chat.id, text=response)
async def killall(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""
kill all running instances
/killall → stops all Docker containers
"""
assert update.effective_chat is not None
global running
for notebook_name in running.keys():
# kill the docker container
subprocess.Popen(
docker.docker_kill_command(running[notebook_name]),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
)
response = f"Killed notebook: {notebook_name}"
await context.bot.send_message(chat_id=update.effective_chat.id, text=response)
running.clear()
async def ps(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""
view currently running Jupyter notebooks
/ps → return all currently running instances
"""
assert update.effective_chat is not None
response = f"Currently running: {running}"
await context.bot.send_message(chat_id=update.effective_chat.id, text=response)
async def ls(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""
view available Jupyter notebooks
/ls → return all types of Jupyter notebooks available to user
"""
assert update.effective_chat is not None
response = "Available notebooks:\n"
response += "\n".join(notebooks.read_all())
response += "\n\nNotebook types:\n"
response += "\n".join(notebook_types)
await context.bot.send_message(chat_id=update.effective_chat.id, text=response)
async def noop(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""
unrecognized command
"""
assert update.effective_chat is not None
response = "Unrecognized command. Use /man to get list of commands."
await context.bot.send_message(chat_id=update.effective_chat.id, text=response)
async def msg(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""
ignore messages
"""
assert update.effective_chat is not None
response = "Let's stick to commands! 😃"
await context.bot.send_message(chat_id=update.effective_chat.id, text=response)
async def man(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""
describe a command OR list all commands
/man → list all commands (name, summary)
/man CMD_NAME → provide full description of command
"""
assert update.effective_chat is not None
assert context.args is not None
match context.args:
case []:
response = "Commands:"
for cmd in cmds:
assert cmd.__doc__ is not None
response += f"\n/{cmd.__name__} -- {cmd.__doc__.splitlines()[1]}"
case [cmd_name] if [x for x in cmds if cmd_name.lstrip("/") == x.__name__]:
cmd = next(x for x in cmds if cmd_name.lstrip("/") == x.__name__)
response = f"\n/{cmd.__name__} -- {cmd.__doc__}"
case [cmd_name]:
response = "Command not found."
case _:
response = "Provide one command."
await context.bot.send_message(chat_id=update.effective_chat.id, text=response)
def main() -> None:
"""Run the bot."""
application = ApplicationBuilder().token(TOKEN).build()
for cmd in cmds:
application.add_handler(CommandHandler(cmd.__name__, cmd))
application.add_handler(MessageHandler(filters.COMMAND, noop))
application.add_handler(MessageHandler(None, msg))
application.run_polling()
if __name__ == "__main__":
cmds = [start, man, ls, init, run, ps, kill, killall]
notebook_types = [
"quay.io/jupyter/minimal-notebook",
"quay.io/jupyter/scipy-notebook",
"quay.io/jupyter/pytorch-notebook",
"quay.io/jupyter/tensorflow-notebook",
]
main()