Skip to content

Commit 069d64d

Browse files
StephenRiStephenRi
committed
support image mirror
1 parent a286f06 commit 069d64d

2 files changed

Lines changed: 220 additions & 0 deletions

File tree

‎rock/cli/command/image.py‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
import argparse
2+
import logging
3+
4+
from rock.cli.command.command import Command
5+
from rock.sdk.builder.image_mirror import ImageMirror
6+
7+
logger = logging.getLogger(__name__)
8+
9+
10+
class ImageCommand(Command):
11+
"""Image command"""
12+
13+
name = "image"
14+
15+
def __init__(self):
16+
super().__init__()
17+
18+
async def arun(self, args: argparse.Namespace):
19+
if "mirror" == args.image_command:
20+
await self.mirror(args)
21+
else:
22+
raise Exception(f"Unknown image command: {args.image_command}")
23+
24+
async def mirror(self, args: argparse.Namespace):
25+
"""Image mirror"""
26+
file_path = args.file
27+
mode = args.mode
28+
concurrency = args.concurrency
29+
source_registry = args.source_registry
30+
source_username = args.source_username
31+
source_password = args.source_password
32+
target_registry = args.target_registry
33+
target_username = args.target_username
34+
target_password = args.target_password
35+
auth_token = args.auth_token
36+
cluster = args.cluster
37+
38+
image_mirror = ImageMirror()
39+
if "local" == mode:
40+
await image_mirror.build_batch(
41+
dataset=file_path,
42+
source_registry=source_registry,
43+
source_username=source_username,
44+
source_password=source_password,
45+
target_registry=target_registry,
46+
target_username=target_username,
47+
target_password=target_password,
48+
)
49+
else:
50+
await image_mirror.build_remote(
51+
authorization=auth_token,
52+
cluster=cluster,
53+
dataset=file_path,
54+
concurrency=concurrency,
55+
source_registry=source_registry,
56+
source_username=source_username,
57+
source_password=source_password,
58+
target_registry=target_registry,
59+
target_username=target_username,
60+
target_password=target_password,
61+
)
62+
63+
@staticmethod
64+
async def add_parser_to(subparsers: argparse._SubParsersAction):
65+
"""Add command line parser"""
66+
image_parser = subparsers.add_parser(
67+
"image",
68+
description="Image operations",
69+
)
70+
image_subparsers = image_parser.add_subparsers(dest="image_command")
71+
72+
# rock image mirror
73+
mirror_parser = image_subparsers.add_parser("mirror", help="Mirror image")
74+
mirror_parser.add_argument(
75+
"-f",
76+
"--file",
77+
required=True,
78+
help="jsonl, source file path. each line is a swebench-like instance",
79+
)
80+
mirror_parser.add_argument(
81+
"--mode",
82+
required=False,
83+
default="local",
84+
help="mirror mode, local or remote(default local)",
85+
)
86+
mirror_parser.add_argument(
87+
"--concurrency",
88+
type=int,
89+
default=3,
90+
choices=range(1, 51),
91+
help="number of concurrent mirrors (default: 3)",
92+
)
93+
mirror_parser.add_argument("--source-registry", required=False, help="source registry url")
94+
mirror_parser.add_argument("--source-username", required=False, help="source hub username")
95+
mirror_parser.add_argument("--source-password", required=False, help="source hub password")
96+
mirror_parser.add_argument("--target-registry", required=True, help="target registry url")
97+
mirror_parser.add_argument("--target-username", required=True, help="target hub username")
98+
mirror_parser.add_argument("--target-password", required=True, help="target hub password")

‎rock/sdk/builder/image_mirror.py‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
import logging
2+
import subprocess
3+
4+
from rock.sdk.builder.base import EnvBuilder
5+
from rock.utils.docker import ImageUtil
6+
7+
logger = logging.getLogger(__name__)
8+
9+
10+
class ImageMirror(EnvBuilder):
11+
async def build(
12+
self,
13+
instance_record: dict[str, str] | None = None,
14+
target_registry: str = "",
15+
target_username: str = "",
16+
target_password: str = "",
17+
source_registry: str | None = None,
18+
source_username: str | None = None,
19+
source_password: str | None = None,
20+
):
21+
if not instance_record:
22+
raise Exception("instance_record is required")
23+
image = instance_record["docker_image"]
24+
registry_url, other_part = await ImageUtil.parse_registry_and_others(image)
25+
parsed_namespace, parsed_name, parsed_tag = await ImageUtil.split_image_name(other_part)
26+
aliyun_other_part = f"{parsed_namespace}/{parsed_name}:{parsed_tag}"
27+
target_image_full_name = f"{target_registry}/{aliyun_other_part}"
28+
29+
# check if target exist, if exist, skip
30+
# login target hub
31+
result = subprocess.run(
32+
[
33+
"docker",
34+
"login",
35+
"-u",
36+
target_username,
37+
"-p",
38+
target_password,
39+
target_registry,
40+
],
41+
capture_output=True,
42+
text=True,
43+
)
44+
if result.returncode != 0:
45+
raise Exception(
46+
f"docker login target hub failed {target_registry}, username is {target_username}, password is {target_password} error:{result.stderr}"
47+
)
48+
# check if target exist
49+
result = subprocess.run(
50+
["docker", "manifest", "inspect", target_image_full_name],
51+
capture_output=True,
52+
text=True,
53+
)
54+
if result.returncode == 0:
55+
logger.info(f"{target_image_full_name} already exists, skip")
56+
return
57+
else:
58+
logger.info(f"{target_image_full_name} not exists, start to mirror")
59+
60+
# login if source_username is not None
61+
if source_username is not None:
62+
if source_password is None:
63+
raise Exception("source_password is required")
64+
if source_registry is None or registry_url != source_registry:
65+
raise Exception(f"image {image} is not in source registry {source_registry}")
66+
result = subprocess.run(
67+
[
68+
"docker",
69+
"login",
70+
"-u",
71+
source_username,
72+
"-p",
73+
source_password,
74+
source_registry,
75+
],
76+
capture_output=True,
77+
text=True,
78+
)
79+
if result.returncode != 0:
80+
raise Exception(f"docker login source hub failed, {result.stderr}")
81+
logger.info("docker login success")
82+
83+
# pull
84+
result = subprocess.run(["docker", "pull", image], capture_output=True, text=True)
85+
if result.returncode != 0:
86+
raise Exception(f"docker pull failed, {result.stderr}")
87+
logger.info(f"docker pull {image} success")
88+
89+
# tag and push
90+
result = subprocess.run(
91+
["docker", "tag", image, target_image_full_name],
92+
capture_output=True,
93+
text=True,
94+
)
95+
if result.returncode != 0:
96+
raise Exception(f"docker tag failed, {result.stderr}")
97+
result = subprocess.run(["docker", "push", target_image_full_name], capture_output=True, text=True)
98+
if result.returncode != 0:
99+
raise Exception(f"docker push failed, {result.stderr}")
100+
logger.info(f"docker push {target_image_full_name} success")
101+
102+
async def verify(self, instance_record: dict[str, str] | None = None):
103+
pass
104+
105+
async def get_build_remote_one_split_command(self, split_filename: str, **kwargs) -> str:
106+
target_registry = kwargs.get("target_registry")
107+
target_username = kwargs.get("target_username")
108+
target_password = kwargs.get("target_password")
109+
source_registry = kwargs.get("source_registry")
110+
source_username = kwargs.get("source_username")
111+
source_password = kwargs.get("source_password")
112+
command = f"rock image mirror -f {split_filename} --target-registry={target_registry} --target-username={target_username} --target-password={target_password}"
113+
if source_registry is not None:
114+
command = f"{command} --source-registry={source_registry}"
115+
if source_username is not None:
116+
command = f"{command} --source-username={source_username}"
117+
if source_password is not None:
118+
command = f"{command} --source-password={source_password}"
119+
return command
120+
121+
async def get_env_build_image(self):
122+
return "rock-n-roll-registry.cn-hangzhou.cr.aliyuncs.com/rock/rock-env-builder:0.2.1"

0 commit comments

Comments
 (0)