Skip to content

Commit

Permalink
Restructure and refactor core module.
Browse files Browse the repository at this point in the history
  • Loading branch information
@csadorf
csadorf committed Feb 4, 2022
1 parent 7188494 commit 3145851
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 82 deletions.
4 changes: 1 addition & 3 deletions aiidalab_launch/__main__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,9 +105,7 @@ def _apply_migration_null(self):

if instance.container:
# There is already a container present, use previously used profile.
self.config.profiles[0] = instance.profile_from_container(
instance.container
)
self.config.profiles[0] = Profile.from_container(instance.container)

elif home_bind_mount_path.exists():
# Using ~/aiidalab as home directory mount point, since the
Expand Down
158 changes: 79 additions & 79 deletions aiidalab_launch/core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import docker
import toml
from docker.models.containers import Container
from packaging.version import parse as parse_version

from .util import _async_wrap_iter, get_docker_env
Expand Down Expand Up @@ -44,6 +45,59 @@ def _default_port() -> int: # explicit function required to enable test patchin
return DEFAULT_PORT


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


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


def _get_jupyter_token(container: Container) -> str:
return get_docker_env(container, "JUPYTER_TOKEN")


def _get_aiidalab_default_apps(container: Container) -> list:
try:
return get_docker_env(container, "AIIDALAB_DEFAULT_APPS").split()
except KeyError:
return []


def _find_docker_home_mount(container: Container, system_user: str) -> Path | None:
# Find the specified home bind mount path for the existing container.
try:
home_mount = [
mount
for mount in container.attrs["Mounts"]
if mount["Destination"] == f"/home/{system_user}"
][0]
except IndexError:
return None
if home_mount["Type"] == "bind":
docker_root = PurePosixPath("/host_mnt")
docker_path = PurePosixPath(home_mount["Source"])
try:
# Try Windows
drive = docker_path.relative_to(docker_root).parts[0]
return WindowsPath(
f"{drive}:",
docker_path.root,
docker_path.relative_to(docker_root, drive),
)
except NotImplementedError:
return PosixPath(docker_root.root, docker_path.relative_to(docker_root))
elif home_mount["Type"] == "volume":
return home_mount["Name"]
else:
raise RuntimeError("Unexpected mount type.")


@dataclass
class Profile:
name: str = MAIN_PROFILE_NAME
Expand Down Expand Up @@ -86,6 +140,24 @@ def dumps(self) -> str:
def loads(cls, name: str, s: str) -> Profile:
return cls(name=name, **toml.loads(s))

@classmethod
def from_container(cls, container: Container) -> Profile:
profile_name = re.sub(re.escape(CONTAINER_PREFIX), "", container.name)
if not profile_name:
raise RuntimeError(
f"Container {container.id} does not appear to be an AiiDAlab container."
)
system_user = _get_system_user(container)

return Profile(
name=profile_name,
port=_get_host_port(container),
default_apps=_get_aiidalab_default_apps(container),
home_mount=str(_find_docker_home_mount(container, system_user)),
image=container.image.tags[0],
system_user=system_user,
)


@dataclass
class Config:
Expand Down Expand Up @@ -143,7 +215,7 @@ class RequiresContainerInstance(RuntimeError):

@contextmanager
def _async_logs(
container: docker.models.containers.Container,
container: Container,
) -> Generator[AsyncGenerator[Any, None], None, None]:
logs = container.logs(stream=True, follow=False)
try:
Expand All @@ -165,15 +237,15 @@ class AiidaLabInstanceStatus(Enum):
client: docker.DockerClient
profile: Profile
_image: docker.models.images.Image = None
_container: docker.models.containers.Container = None
_container: Container = None

def _get_image(self) -> docker.models.images.Image | None:
try:
return self.client.images.get(self.profile.image)
except docker.errors.ImageNotFound:
return None

def _get_container(self) -> docker.models.containers.Container | None:
def _get_container(self) -> Container | None:
try:
return self.client.containers.get(self.profile.container_name())
except docker.errors.NotFound:
Expand All @@ -188,7 +260,7 @@ def image(self) -> docker.models.images.Image | None:
return self._image

@property
def container(self) -> docker.models.containers.Container | None:
def container(self) -> Container | None:
if self._container is None:
self._container = self._get_container()
return self._container
Expand Down Expand Up @@ -217,48 +289,6 @@ def _mounts(self) -> Generator[docker.types.Mount, None, None]:
if self.profile.home_mount:
yield self._home_mount()

def _find_docker_home_mount(
self, container: docker.models.containers.Container
) -> Path | None:
# Find the specified home bind mount path for the existing container.
try:
home_mount = [
mount
for mount in container.attrs["Mounts"]
if mount["Destination"] == f"/home/{self.profile.system_user}"
][0]
except IndexError:
return None
if home_mount["Type"] == "bind":
docker_root = PurePosixPath("/host_mnt")
docker_path = PurePosixPath(home_mount["Source"])
try:
# Try Windows
drive = docker_path.relative_to(docker_root).parts[0]
return WindowsPath(
f"{drive}:",
docker_path.root,
docker_path.relative_to(docker_root, drive),
)
except NotImplementedError:
return PosixPath(docker_root.root, docker_path.relative_to(docker_root))
elif home_mount["Type"] == "volume":
return home_mount["Name"]
else:
raise RuntimeError("Unexpected mount type.")

def profile_from_container(
self, container: docker.models.containers.Container
) -> Profile:
return Profile(
name=self.profile.name,
port=self.host_port(container),
default_apps=self._aiidalab_default_apps(container),
home_mount=str(self._find_docker_home_mount(container)),
image=container.image.tags[0],
system_user=self._system_user(container),
)

def configuration_changes(self) -> Generator[str, None, None]:
assert self.container is not None
assert self.image is not None
Expand Down Expand Up @@ -289,7 +319,7 @@ def _ensure_home_mount_exists(self) -> None:
)
home_mount_path.mkdir(exist_ok=True)

def create(self) -> docker.models.containers.Container:
def create(self) -> Container:
assert self._container is None
self._ensure_home_mount_exists()
self._container = self.client.containers.create(
Expand Down Expand Up @@ -441,38 +471,8 @@ async def status(self, timeout: float | None = 5.0) -> AiidaLabInstanceStatus:
return self.AiidaLabInstanceStatus.EXITED
return self.AiidaLabInstanceStatus.DOWN

@staticmethod
def _aiidalab_default_apps(
container: docker.models.containers.Container,
) -> list:
try:
return get_docker_env(container, "AIIDALAB_DEFAULT_APPS").split()
except KeyError:
return []

@staticmethod
def _system_user(container: docker.models.containers.Container) -> str:
return get_docker_env(container, "SYSTEM_USER")

def jupyter_token(self) -> str | None:
if self.container:
try:
return get_docker_env(self.container, "JUPYTER_TOKEN")
except KeyError:
pass
return None

@staticmethod
def host_port(container: docker.models.containers.Container) -> int | None:
try:
host_config = container.attrs["HostConfig"]
return int(host_config["PortBindings"]["8888/tcp"][0]["HostPort"])
except (KeyError, IndexError, ValueError):
pass
return None

def url(self) -> str:
assert self.container is not None
host_port = self.host_port(self.container)
jupyter_token = self.jupyter_token()
host_port = _get_host_port(self.container)
jupyter_token = _get_jupyter_token(self.container)
return f"http://localhost:{host_port}/?token={jupyter_token}"

0 comments on commit 3145851

Please sign in to comment.