Skip to content

Commit

Permalink
Fix: Correctly Distinguish been configured and actually used host port.
Browse files Browse the repository at this point in the history
Relevant for when the host port is determined dynamically (e.g. by
setting it to 0 or None).
  • Loading branch information
csadorf committed Feb 8, 2022
1 parent 7d3b639 commit 212cb38
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 10 deletions.
4 changes: 3 additions & 1 deletion aiidalab_launch/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -487,13 +487,15 @@ async def _async_start(
else MSG_STARTUP
)
url = instance.url()
host_ports = instance.host_ports()
assert len(host_ports) > 0
click.secho(
msg_startup.format(
url=url,
home_mount=instance.profile.home_mount,
system_user=instance.profile.system_user,
user=getpass.getuser(),
port=instance.profile.port,
port=host_ports[0],
hostname=socket.getfqdn(),
).lstrip(),
fg="green",
Expand Down
48 changes: 41 additions & 7 deletions aiidalab_launch/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,15 +45,23 @@ def _default_port() -> int: # explicit function required to enable test patchin
return DEFAULT_PORT


def _get_host_port(container: Container) -> int | None:
def _get_configured_host_port(container: Container) -> int | None:
try:
host_config = container.attrs["HostConfig"]
return int(host_config["PortBindings"]["8888/tcp"][0]["HostPort"])
return int(host_config["PortBindings"]["8888/tcp"][0]["HostPort"]) or None
except (KeyError, IndexError, ValueError):
pass
return None


def _get_host_ports(container: Container) -> Generator[int, None, None]:
try:
ports = container.attrs["NetworkSettings"]["Ports"]
yield from (int(i["HostPort"]) for i in ports["8888/tcp"])
except KeyError:
pass


def _get_system_user(container: Container) -> str:
return get_docker_env(container, "SYSTEM_USER")

Expand Down Expand Up @@ -122,7 +130,7 @@ def from_container(cls, container: Container) -> Profile:
system_user = _get_system_user(container)
return Profile(
name=profile_name,
port=_get_host_port(container),
port=_get_configured_host_port(container),
default_apps=_get_aiidalab_default_apps(container),
home_mount=str(
docker_mount_for(container, PurePosixPath("/", "home", system_user))
Expand Down Expand Up @@ -186,6 +194,10 @@ class RequiresContainerInstance(RuntimeError):
"""Raised when trying to perform operation that requires a container instance."""


class NoHostPortAssigned(RuntimeError):
"""Raised when then trying to obtain the instance URL, but there is not host port."""


@contextmanager
def _async_logs(
container: Container,
Expand Down Expand Up @@ -416,13 +428,24 @@ async def _notebook_service_online(self) -> None:
else:
raise FailedToWaitForServices("Failed to reach notebook service.")

async def _host_port_assigned(self) -> None:
container = self.container
assert container is not None
while True:
container.reload()
if any(_get_host_ports(container)):
break
asyncio.sleep(1)

async def wait_for_services(self) -> None:
if self.container is None:
raise RuntimeError("Instance was not created.")

LOGGER.info(f"Waiting for services to come up ({self.container.id})...")
await asyncio.gather(
self._init_scripts_finished(), self._notebook_service_online()
self._init_scripts_finished(),
self._notebook_service_online(),
self._host_port_assigned(),
)

async def status(self, timeout: float | None = 5.0) -> AiidaLabInstanceStatus:
Expand All @@ -444,8 +467,19 @@ async def status(self, timeout: float | None = 5.0) -> AiidaLabInstanceStatus:
return self.AiidaLabInstanceStatus.EXITED
return self.AiidaLabInstanceStatus.DOWN

def host_ports(self) -> list[int]:
self._requires_container()
assert self.container is not None
self.container.reload()
return list(_get_host_ports(self.container))

def url(self) -> str:
self._requires_container()
assert self.container is not None
host_port = _get_host_port(self.container)
jupyter_token = _get_jupyter_token(self.container)
return f"http://localhost:{host_port}/?token={jupyter_token}"
self.container.reload()
host_ports = list(_get_host_ports(self.container))
if len(host_ports) > 0:
jupyter_token = _get_jupyter_token(self.container)
return f"http://localhost:{host_ports[0]}/?token={jupyter_token}"
else:
raise NoHostPortAssigned(self.container.id)
22 changes: 20 additions & 2 deletions tests/test_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,19 @@
.. moduleauthor:: Carl Simon Adorf <[email protected]>
"""
import asyncio
import re
from dataclasses import replace
from pathlib import Path
from time import sleep

import pytest

from aiidalab_launch.core import Config, Profile
from aiidalab_launch.core import (
Config,
NoHostPortAssigned,
Profile,
RequiresContainerInstance,
)

VALID_PROFILE_NAMES = ["abc", "Abc", "aBC", "a0", "a-a", "a-0"]

Expand Down Expand Up @@ -88,17 +94,29 @@ async def test_instance_home_bind_mount(instance):
@pytest.mark.slow
@pytest.mark.trylast
async def test_instance_start_stop(instance):
with pytest.raises(RequiresContainerInstance):
instance.url()
assert await instance.status() is instance.AiidaLabInstanceStatus.DOWN
instance.start()
sleep(0.1)
assert await instance.status() is instance.AiidaLabInstanceStatus.STARTING

# premature additional start should have no negative effect
# It is possible that the call below will succeed/fail non-deterministically.
assert re.match(r"http:\/\/localhost:\d+\/\?token=[a-f0-9]{64}", instance.url())

# second call to start should have no negative effect
instance.start()

await asyncio.wait_for(instance.wait_for_services(), timeout=300)
assert await instance.status() is instance.AiidaLabInstanceStatus.UP

assert re.match(r"http:\/\/localhost:\d+\/\?token=[a-f0-9]{64}", instance.url())

instance.stop()
assert await instance.status() is instance.AiidaLabInstanceStatus.EXITED

with pytest.raises(NoHostPortAssigned):
instance.url()

instance.remove()
assert await instance.status() is instance.AiidaLabInstanceStatus.DOWN

0 comments on commit 212cb38

Please sign in to comment.