Skip to content

Commit

Permalink
breaking: Replace rescan command's --local flag with local containe…
Browse files Browse the repository at this point in the history
…r registry record
  • Loading branch information
jopemachine committed Aug 8, 2024
1 parent 31e78a2 commit 6464742
Show file tree
Hide file tree
Showing 8 changed files with 63 additions and 70 deletions.
11 changes: 11 additions & 0 deletions fixtures/manager/example-container-registries-local.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
{
"container_registries": [
{
"id": "fe878f09-06cc-4b91-9242-4c71015cce05",
"registry_name": "local",
"url": "http://localhost",
"type": "local",
"project": ""
}
]
}
8 changes: 7 additions & 1 deletion src/ai/backend/agent/docker/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -770,8 +770,14 @@ async def start_container(
container_log_size = self.local_config["agent"]["container-logs"]["max-length"]
container_log_file_count = 5
container_log_file_size = BinarySize(container_log_size // container_log_file_count)

if self.image_ref.is_local:
image = self.image_ref.canonical
else:
image = self.image_ref.short

container_config: MutableMapping[str, Any] = {
"Image": self.image_ref.canonical,
"Image": image,
"Tty": True,
"OpenStdin": True,
"Privileged": False,
Expand Down
1 change: 0 additions & 1 deletion src/ai/backend/manager/api/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -1227,7 +1227,6 @@ async def _commit_and_upload(reporter: ProgressReporter) -> None:
await rescan_images(
root_ctx.db,
new_image_ref.canonical,
local=new_image_ref.is_local,
)
await reporter.update(increment=1, message="Completed")
except BackendError:
Expand Down
2 changes: 1 addition & 1 deletion src/ai/backend/manager/cli/etcd.py
Original file line number Diff line number Diff line change
Expand Up @@ -261,7 +261,7 @@ def rescan_images(cli_ctx: CLIContext, registry: str) -> None:
Pass the name (usually hostname or "lablup") of the Docker registry configured as REGISTRY.
"""
log.warning("etcd rescan-images command is deprecated, use image rescan instead")
asyncio.run(rescan_images_impl(cli_ctx, registry, False))
asyncio.run(rescan_images_impl(cli_ctx, registry))


@cli.command()
Expand Down
10 changes: 2 additions & 8 deletions src/ai/backend/manager/cli/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,20 +82,14 @@ def set_resource_limit(

@cli.command()
@click.argument("registry_or_image", required=False, default="")
@click.option(
"--local",
is_flag=True,
default=False,
help="Scan the local Docker daemon instead of a registry",
)
@click.pass_obj
def rescan(cli_ctx, registry_or_image: str, local: bool) -> None:
def rescan(cli_ctx, registry_or_image: str) -> None:
"""
Update the kernel image metadata from all configured docker registries.
Pass the name (usually hostname or "lablup") of the Docker registry configured as REGISTRY.
"""
asyncio.run(rescan_images_impl(cli_ctx, registry_or_image, local))
asyncio.run(rescan_images_impl(cli_ctx, registry_or_image))


@cli.command()
Expand Down
6 changes: 3 additions & 3 deletions src/ai/backend/manager/cli/image_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,14 +150,14 @@ async def set_image_resource_limit(
log.exception("An error occurred.")


async def rescan_images(cli_ctx: CLIContext, registry_or_image: str, local: bool) -> None:
if not registry_or_image and not local:
async def rescan_images(cli_ctx: CLIContext, registry_or_image: str) -> None:
if not registry_or_image:
raise click.BadArgumentUsage("Please specify a valid registry or full image name.")
async with (
connect_database(cli_ctx.local_config) as db,
):
try:
await rescan_images_func(db, registry_or_image, local=local)
await rescan_images_func(db, registry_or_image)
except Exception:
log.exception("An error occurred.")

Expand Down
21 changes: 7 additions & 14 deletions src/ai/backend/manager/container_registry/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -445,20 +445,13 @@ async def _read_manifest(
except ValueError as e:
skip_reason = str(e)
continue
if self.registry_name == "local":
if image.partition("/")[1] == "":
image = "library/" + image
update_key = ImageRef(
f"{image}:{tag}",
["index.docker.io"],
architecture,
)
else:
update_key = ImageRef(
f"{self.registry_name}/{image}:{tag}",
[self.registry_name],
architecture,
)

update_key = ImageRef(
f"{self.registry_name}/{image}:{tag}",
[self.registry_name],
architecture,
)

updates = {
"config_digest": manifest["digest"],
"size_bytes": manifest["size"],
Expand Down
74 changes: 32 additions & 42 deletions src/ai/backend/manager/models/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@
from ai.backend.common.exception import UnknownImageReference
from ai.backend.common.logging import BraceStyleAdapter
from ai.backend.common.types import BinarySize, ImageAlias, ResourceSlot
from ai.backend.manager.models.container_registry import ContainerRegistryRow, ContainerRegistryType
from ai.backend.manager.models.container_registry import ContainerRegistryRow

from ..api.exceptions import ImageNotFound, ObjectNotFound
from ..container_registry import get_container_registry_cls
Expand Down Expand Up @@ -113,51 +113,41 @@ async def rescan_images(
db: ExtendedAsyncSAEngine,
registry_or_image: str | None = None,
*,
local: bool | None = False,
reporter: ProgressReporter | None = None,
) -> None:
if local:
registries = {
"local": ContainerRegistryRow(
registry_name="local",
url="http://localhost",
type=ContainerRegistryType.LOCAL,
)
}
else:
async with db.begin_readonly_session() as session:
result = await session.execute(sa.select(ContainerRegistryRow))
latest_registry_config = cast(
dict[str, ContainerRegistryRow],
{row.registry_name: row for row in result.scalars().all()},
)
async with db.begin_readonly_session() as session:
result = await session.execute(sa.select(ContainerRegistryRow))
latest_registry_config = cast(
dict[str, ContainerRegistryRow],
{row.registry_name: row for row in result.scalars().all()},
)

# TODO: delete images from registries removed from the previous config?
if registry_or_image is None:
# scan all configured registries
registries = latest_registry_config
# TODO: delete images from registries removed from the previous config?
if registry_or_image is None:
# scan all configured registries
registries = latest_registry_config
else:
# find if it's a full image ref of one of configured registries
for registry_name, registry_info in latest_registry_config.items():
if registry_or_image.startswith(registry_name + "/"):
repo_with_tag = registry_or_image.removeprefix(registry_name + "/")
log.debug(
"running a per-image metadata scan: {}, {}",
registry_name,
repo_with_tag,
)
scanner_cls = get_container_registry_cls(registry_info)
scanner = scanner_cls(db, registry_name, registry_info)
await scanner.scan_single_ref(repo_with_tag)
return
else:
# find if it's a full image ref of one of configured registries
for registry_name, registry_info in latest_registry_config.items():
if registry_or_image.startswith(registry_name + "/"):
repo_with_tag = registry_or_image.removeprefix(registry_name + "/")
log.debug(
"running a per-image metadata scan: {}, {}",
registry_name,
repo_with_tag,
)
scanner_cls = get_container_registry_cls(registry_info)
scanner = scanner_cls(db, registry_name, registry_info)
await scanner.scan_single_ref(repo_with_tag)
return
else:
# treat it as a normal registry name
registry = registry_or_image
try:
registries = {registry: latest_registry_config[registry]}
log.debug("running a per-registry metadata scan")
except KeyError:
raise RuntimeError("It is an unknown registry.", registry)
# treat it as a normal registry name
registry = registry_or_image
try:
registries = {registry: latest_registry_config[registry]}
log.debug("running a per-registry metadata scan")
except KeyError:
raise RuntimeError("It is an unknown registry.", registry)
async with aiotools.TaskGroup() as tg:
for registry_name, registry_info in registries.items():
log.info('Scanning kernel images from the registry "{0}"', registry_name)
Expand Down

0 comments on commit 6464742

Please sign in to comment.