JarvisLabs experiment launcher

jarvislabs_launch.py

"""
Launch new JarvisLabs instance to run a training script.

Please edit your JarvisLabs (and Weights & Biases) credentials below. Also, you first
need to upload `run_base64.sh` as a startup script in JarvisLabs UI and enter its
script id below.

The instance will have PyTorch and an RTX5000 GPU.

Usage:

    $ python jarvislabs_launch.py [-f] cleanrl/ppo_atari_envpool.py --env_id Pong-v5

You can tail logs with -f.
"""

import base64
import io
import os
import sys
import tarfile

from jlclient import jarvisclient

jarvisclient.token = "<INSERT TOKEN>"
jarvisclient.user_id = "<INSERT EMAIL>"

if sys.argv[1] == "-f":
    tail_logs = True
    del sys.argv[1]
else:
    tail_logs = False
script = sys.argv[1]
args = " ".join(sys.argv[2:])

runner = f"""
#!/usr/bin/bash

export WANDB_API_KEY="<INSERT TOKEN>"
export WANDB_ENTITY="<INSERT USERNAME>"
export OMP_NUM_THREADS="1"

# envpool 0.6.2 doesn't work with gym 0.25.0 (TypeError: metaclass conflict)
pip install envpool==0.6.2 gym==0.24.1 wandb==0.12.21

python {script} {args}
"""

buf = io.BytesIO()
tar = tarfile.open(fileobj=buf, mode="w:xz")
runner_buf = io.BytesIO(runner.encode("utf-8"))
runner_info = tarfile.TarInfo("run.sh")
runner_info.size = len(runner_buf.getvalue())
runner_info.mode = 0o644
tar.addfile(runner_info, fileobj=runner_buf)
tar.add(script)
tar.close()

payload = base64.b64encode(buf.getvalue()).decode("ascii")

instance = jarvisclient.Instance.create(
    gpu_type="RTX5000",
    num_gpus=1,
    hdd=20,
    framework_id=0,
    name=script,
    script_id=<INSERT ID>,  # run_base64.sh
    arguments=payload,
)
assert instance.status == "Running"
print(f"launched instance {instance.machine_id} {instance.name}")

if tail_logs:
    cmd = instance.ssh_str.split(" ")
    assert cmd[0] == "ssh"
    cmd = ["ssh", "-o", "StrictHostKeyChecking=no"] + cmd[1:] + ["tail", "-f", "/home/logs.txt"]
    print(" ".join(cmd))
    os.execlp("ssh", *cmd)
else:
    print(instance.ssh_str)

run_base64.sh

#!/usr/bin/bash

CODE=$1
shift 1
ARGS=$*

mkdir code
cd code
echo "$CODE" | base64 -d | tar -x --lzma
bash run.sh $ARGS

python -c "from jarviscloud import jarviscloud; jarviscloud.destroy()"