def start_shell(args: Namespace) -> None:
data = {}
if args.passphrase:
data["passphrase"] = getpass.getpass("Enter new passphrase: ")
config = parse_config(args.config_file, None, args.config, args.volume)
resp = launch_command(
args.master,
"api/v1/shells",
config,
args.template,
context_path=args.context,
data=data,
)["shell"]
if args.detach:
print(resp["id"])
return
ready = False
with api.ws(args.master, "shells/{}/events".format(resp["id"])) as ws:
for msg in ws:
if msg["service_ready_event"]:
ready = True
break
render_event_stream(msg)
if ready:
shell = api.get(args.master, "api/v1/shells/{}".format(resp["id"])).json()["shell"]
check_eq(shell["state"], "STATE_RUNNING", "Shell must be in a running state")
_open_shell(args.master, shell, args.ssh_opts)
def start_notebook(args: Namespace) -> None:
config = parse_config(args.config_file, None, args.config, args.volume)
resp = launch_command(
args.master,
"api/v1/notebooks",
config,
args.template,
context_path=args.context,
preview=args.preview,
)
if args.preview:
print(render.format_object_as_yaml(resp["config"]))
return
obj = resp["notebook"]
if args.detach:
print(obj["id"])
return
with api.ws(args.master, "notebooks/{}/events".format(obj["id"])) as ws:
for msg in ws:
if msg["service_ready_event"] and not args.no_browser:
url = api.open(args.master, obj["serviceAddress"])
print(colored("Jupyter Notebook is running at: {}".format(url), "green"))
render_event_stream(msg)
def start_shell(args: Namespace) -> None:
data = {}
if args.passphrase:
data["passphrase"] = getpass.getpass("Enter new passphrase: ")
config = parse_config(args.config_file, None, args.config, args.volume)
resp = launch_command(
args.master,
"shells",
config,
args.template,
context_path=args.context,
data=data,
)
if args.detach:
print(resp["id"])
return
command = None
with api.ws(args.master, "shells/{}/events".format(resp["id"])) as ws:
for msg in ws:
if msg["service_ready_event"]:
command = render.unmarshal(Command, msg["snapshot"])
break
render_event_stream(msg)
if command:
_open_shell(args.master, command, args.ssh_opts)
def tail_logs(args: Namespace) -> None:
api_full_path = "{}/{}/events?follow={}&tail={}".format(
RemoteTaskOldAPIs[args._command],
RemoteTaskGetIDsFunc[args._command](args), # type: ignore
args.follow,
args.tail,
)
with api.ws(args.master, api_full_path) as ws:
for msg in ws:
render_event_stream(msg)
def tail_command_logs(args: Namespace) -> None:
token = api.Authentication.instance().get_session_token()
params = {"follow": args.follow, "tail": args.tail, "_auth": token}
url = "commands/{}/events?{}".format(args.command_id,
urllib.parse.urlencode(params))
with api.ws(args.master, url) as ws:
for msg in ws:
render_event_stream(msg)
def start_tensorboard(args: Namespace) -> None:
if not (args.trial_ids or args.experiment_ids):
print("Either experiment_ids or trial_ids must be specified.")
sys.exit(1)
config = parse_config(args.config_file, None, args.config, [])
req_body = {
"config": config,
"trial_ids": args.trial_ids,
"experiment_ids": args.experiment_ids,
}
if args.context is not None:
req_body["files"], _ = context.read_context(args.context,
constants.MAX_CONTEXT_SIZE)
resp = api.post(args.master, "api/v1/tensorboards",
json=req_body).json()["tensorboard"]
if args.detach:
print(resp["id"])
return
url = "tensorboard/{}/events".format(resp["id"])
with api.ws(args.master, url) as ws:
for msg in ws:
if msg["log_event"] is not None:
# TensorBoard will print a url by default. The URL is incorrect since
# TensorBoard is not aware of the master proxy address it is assigned.
if "http" in msg["log_event"]:
continue
if msg["service_ready_event"]:
if args.no_browser:
url = api.make_url(args.master, resp["serviceAddress"])
else:
url = api.browser_open(args.master, resp["serviceAddress"])
print(
colored("TensorBoard is running at: {}".format(url),
"green"))
render_event_stream(msg)
break
render_event_stream(msg)
def run_command(args: Namespace) -> None:
config = parse_config(args.config_file, args.entrypoint, args.config,
args.volume)
resp = launch_command(
args.master,
"api/v1/commands",
config,
args.template,
context_path=args.context,
)["command"]
if args.detach:
print(resp["id"])
return
url = "commands/{}/events".format(resp["id"])
with api.ws(args.master, url) as ws:
for msg in ws:
render_event_stream(msg)
def start_notebook(args: Namespace) -> None:
config = parse_config(args.config_file, None, args.config, args.volume)
files = None
if args.context is not None:
context = Context.from_local(args.context)
files = [
bindings.v1File(
content=e.content.decode("utf-8"),
gid=e.gid,
mode=e.mode,
mtime=e.mtime,
path=e.path,
type=e.type,
uid=e.uid,
) for e in context.entries
]
body = bindings.v1LaunchNotebookRequest(config, files=files, preview=False)
resp = bindings.post_LaunchNotebook(setup_session(args), body=body)
if args.preview:
print(render.format_object_as_yaml(resp.config))
return
nb = resp.notebook
if args.detach:
print(nb.id)
return
with api.ws(args.master, "notebooks/{}/events".format(nb.id)) as ws:
for msg in ws:
if msg["service_ready_event"] and nb.serviceAddress and not args.no_browser:
url = api.browser_open(args.master, nb.serviceAddress)
print(
colored("Jupyter Notebook is running at: {}".format(url),
"green"))
render_event_stream(msg)
def tail_shell_logs(args: Namespace) -> None:
url = "shells/{}/events?follow={}&tail={}".format(args.shell_id,
args.follow, args.tail)
with api.ws(args.master, url) as ws:
for msg in ws:
render_event_stream(msg)
def tail_tensorboard_logs(args: Namespace) -> None:
url = "tensorboard/{}/events?follow={}&tail={}".format(
args.tensorboard_id, args.follow, args.tail)
with api.ws(args.master, url) as ws:
for msg in ws:
render_event_stream(msg)
