def cleanup():
from lightning_app.utilities.app_helpers import _LightningAppRef
yield
_LightningAppRef._app_instance = None
shutil.rmtree("./storage", ignore_errors=True)
shutil.rmtree(storage_root_dir(), ignore_errors=True)
shutil.rmtree("./.shared", ignore_errors=True)
if os.path.isfile(_APP_CONFIG_FILENAME):
os.remove(_APP_CONFIG_FILENAME)
_set_context(None)
def test_snapshotting(runtime_cls, tmpdir):
try:
app = CheckpointLightningApp(FlowA())
app.checkpointing = True
runtime_cls(app, start_server=False).dispatch()
except SuccessException:
pass
checkpoint_dir = os.path.join(storage_root_dir(), "checkpoints")
checkpoints = os.listdir(checkpoint_dir)
assert len(checkpoints) == 1
with open(os.path.join(checkpoint_dir, checkpoints[0]), "rb") as f:
state = pickle.load(f)
assert state["works"]["work_a"]["vars"]["counter"] == 1
assert state["works"]["work_b"]["vars"]["counter"] == 1
def test_lightning_flow_counter(runtime_cls, tmpdir):
app = LightningApp(FlowCounter())
app.checkpointing = True
runtime_cls(app, start_server=False).dispatch()
assert app.root.counter == 3
checkpoint_dir = os.path.join(storage_root_dir(), "checkpoints")
checkpoints = os.listdir(checkpoint_dir)
assert len(checkpoints) == 4
for checkpoint in checkpoints:
checkpoint_path = os.path.join(checkpoint_dir, checkpoint)
with open(checkpoint_path, "rb") as f:
app = LightningApp(FlowCounter())
app.set_state(pickle.load(f))
runtime_cls(app, start_server=False).dispatch()
assert app.root.counter == 3
def test_start_stop_server_through_frontend(process_mock):
frontend = StaticWebFrontend(serve_dir=".")
frontend.flow = MockFlow()
frontend.start_server("localhost", 5000)
log_file_root = storage_root_dir()
process_mock.assert_called_once_with(
target=lightning_app.frontend.web.start_server,
kwargs={
"host": "localhost",
"port": 5000,
"serve_dir": ".",
"path": "/root.my.flow",
"log_file": os.path.join(log_file_root, "frontend", "logs.log"),
},
)
process_mock().start.assert_called_once()
frontend.stop_server()
process_mock().kill.assert_called_once()
def test_lightning_flow_iterate(tmpdir, runtime_cls, run_once):
app = LightningApp(CFlow(run_once))
runtime_cls(app, start_server=False).dispatch()
assert app.root.looping == 0
assert app.root.tracker == 4
call_hash = list(v for v in app.root._calls if "experimental_iterate" in v)[0]
iterate_call = app.root._calls[call_hash]
assert iterate_call["counter"] == 4
assert not iterate_call["has_finished"]
checkpoint_dir = os.path.join(storage_root_dir(), "checkpoints")
app = LightningApp(CFlow(run_once))
app.load_state_dict_from_checkpoint_dir(checkpoint_dir)
app.root.restarting = True
assert app.root.looping == 0
assert app.root.tracker == 4
runtime_cls(app, start_server=False).dispatch()
assert app.root.looping == 2
assert app.root.tracker == 10 if run_once else 20
iterate_call = app.root._calls[call_hash]
assert iterate_call["has_finished"]
def get_frontend_logfile(filename: str = "logs.log") -> Path:
log_dir = Path(storage_root_dir(), "frontend")
log_dir.mkdir(parents=True, exist_ok=True)
log_file = log_dir / filename
return log_file
def checkpoint_dir(self) -> str:
return os.path.join(storage_root_dir(), "checkpoints")