Skip to content

Commit

Permalink
revert: Remove automatic auth_token issurance logic
Browse files Browse the repository at this point in the history
  • Loading branch information
jopemachine committed Jul 26, 2024
1 parent 33d543f commit 3dd31fc
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 55 deletions.
13 changes: 2 additions & 11 deletions src/ai/backend/manager/api/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,6 @@

from ai.backend.common.bgtask import ProgressReporter
from ai.backend.common.docker import ImageRef
from ai.backend.manager.container_registry.aws_ecr import AWSElasticContainerRegistry_v2
from ai.backend.manager.models.group import GroupRow
from ai.backend.manager.models.image import rescan_images

Expand Down Expand Up @@ -1188,20 +1187,12 @@ async def _commit_and_upload(reporter: ProgressReporter) -> None:
raise BackendError(extra_msg="Operation cancelled")

if not new_image_ref.is_local:
if "public.ecr" in registry_hostname or "dkr.ecr" in str(registry_hostname):
credential = AWSElasticContainerRegistry_v2.get_credential(registry_conf)
else:
credential = registry_conf

username = credential.get("username")
password = credential.get("password")

# push image to registry from local agent
image_registry = ImageRegistry(
name=registry_hostname,
url=str(registry_conf[""]),
username=username,
password=password,
username=registry_conf.get("username"),
password=registry_conf.get("password"),
)
resp = await root_ctx.registry.push_image(
session.main_kernel.agent,
Expand Down
45 changes: 1 addition & 44 deletions src/ai/backend/manager/container_registry/aws_ecr.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,10 @@
import base64
import logging
from typing import Any, AsyncIterator, Mapping
from typing import AsyncIterator

import aiohttp
import boto3

from ai.backend.common.logging import BraceStyleAdapter
from ai.backend.manager.models.utils import ExtendedAsyncSAEngine

from .base import (
BaseContainerRegistry,
Expand All @@ -16,47 +14,6 @@


class AWSElasticContainerRegistry_v2(BaseContainerRegistry):
@staticmethod
def get_credential(registry_info: Mapping[str, Any]) -> dict[str, Any]:
access_key, secret_access_key, region, type_ = (
registry_info["access_key"],
registry_info["secret_access_key"],
registry_info["region"],
registry_info["type"],
)

ecr_client = boto3.client(
type_,
region_name=region,
aws_access_key_id=access_key,
aws_secret_access_key=secret_access_key,
)

auth_token = ecr_client.get_authorization_token()["authorizationData"]["authorizationToken"]
decoded_auth_token = base64.b64decode(auth_token).decode("utf-8")
username, password = decoded_auth_token.split(":")

return {"username": username, "password": password}

def __init__(
self,
db: ExtendedAsyncSAEngine,
registry_name: str,
registry_info: Mapping[str, Any],
*,
max_concurrency_per_registry: int = 4,
ssl_verify: bool = True,
) -> None:
super().__init__(
db,
registry_name,
registry_info,
max_concurrency_per_registry=max_concurrency_per_registry,
ssl_verify=ssl_verify,
)

self.credentials = AWSElasticContainerRegistry_v2.get_credential(registry_info)

async def fetch_repositories(
self,
sess: aiohttp.ClientSession,
Expand Down

0 comments on commit 3dd31fc

Please sign in to comment.