test_launcher.py

import pytest
import multiprocessing as mp
from pathlib import Path
import json
import copy

from blendtorch import btt

BLENDDIR = Path(__file__).parent / "blender"
LAUNCH_ARGS = dict(
    scene="",
    script=str(BLENDDIR / "launcher.blend.py"),
    num_instances=2,
    named_sockets=["DATA", "GYM"],
    background=True,
    instance_args=[["--x", "3"], ["--x", "4"]],
    seed=10,
)


def _validate_result(items):
    assert len(items) == 2
    first, second = 0, 1
    if items[0]["btid"] == 1:
        first, second = second, first

    assert items[first]["btargs"]["btid"] == 0
    assert items[second]["btargs"]["btid"] == 1
    assert items[first]["btargs"]["btseed"] == 10
    assert items[second]["btargs"]["btseed"] == 11
    assert items[first]["btargs"]["btsockets"]["DATA"].startswith("tcp://")
    assert items[first]["btargs"]["btsockets"]["GYM"].startswith("tcp://")
    assert items[second]["btargs"]["btsockets"]["DATA"].startswith("tcp://")
    assert items[second]["btargs"]["btsockets"]["GYM"].startswith("tcp://")
    assert items[first]["remainder"] == ["--x", "3"]
    assert items[second]["remainder"] == ["--x", "4"]


#@pytest.mark.background
def test_launcher():
    with btt.BlenderLauncher(**LAUNCH_ARGS) as bl:
        addr = bl.launch_info.addresses["DATA"]
        ds = btt.RemoteIterableDataset(addr, max_items=2)
        items = [item for item in ds]
        _validate_result(items)


def _launch(q, tmp_path):
    with btt.BlenderLauncher(**LAUNCH_ARGS) as bl:
        path = Path(tmp_path / "addresses.json")
        btt.LaunchInfo.save_json(path, bl.launch_info)
        q.put(path)
        bl.wait()


#@pytest.mark.background
def test_launcher_connected_remote(tmp_path):
    # Simulates BlenderLauncher called from a separate process and
    # shows how one can connect to already launched instances through
    # serialization of addresses.
    q = mp.Queue()
    p = mp.Process(target=_launch, args=(q, tmp_path))
    p.start()
    path = q.get()
    launch_info = btt.LaunchInfo.load_json(path)
    ds = btt.RemoteIterableDataset(launch_info.addresses["DATA"], max_items=2)
    items = [item for item in ds]
    _validate_result(items)
    p.join()


def _launch_app(tmp_path, args):
    from blendtorch.btt.apps import launch

    with open(tmp_path / "launchargs.json", "w") as fp:
        fp.write(json.dumps(args, indent=4))
    launch.main(
        [
            "--out-launch-info",
            str(tmp_path / "launchinfo.json"),
            str(tmp_path / "launchargs.json"),
        ]
    )


#@pytest.mark.background
def test_launcher_app(tmp_path):

    p = mp.Process(target=_launch_app, args=(tmp_path, LAUNCH_ARGS))
    p.start()

    import time

    path = tmp_path / "launchinfo.json"
    while not Path.exists(path):
        time.sleep(1)

    launch_info = btt.LaunchInfo.load_json(path)
    ds = btt.RemoteIterableDataset(launch_info.addresses["DATA"], max_items=2)
    items = [item for item in ds]
    _validate_result(items)

    p.join()


def test_launcher_app_primaryip(tmp_path):

    # Same with primary ip resolver
    args = copy.deepcopy(LAUNCH_ARGS)
    args["bind_addr"] = "primaryip"
    p = mp.Process(target=_launch_app, args=(tmp_path, args))
    p.start()

    import time

    path = tmp_path / "launchinfo.json"
    while not Path.exists(path):
        time.sleep(1)

    launch_info = btt.LaunchInfo.load_json(path)
    print(launch_info.addresses)
    ds = btt.RemoteIterableDataset(launch_info.addresses["DATA"], max_items=2)
    items = [item for item in ds]
    _validate_result(items)

    p.join()