Skip to content
Open
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 97 additions & 51 deletions src/together/lib/cli/api/beta/jig/jig.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import shutil
import typing
import asyncio
import tempfile
import subprocess
import concurrent.futures
from typing import TYPE_CHECKING, Any, Union, Callable, Optional, Annotated
Expand Down Expand Up @@ -549,6 +550,11 @@
"""Path for buildx --metadata-file output, used to recover the pushed digest cross-process."""
return self.config._path.parent / f".jig-{self.name}-{tag}.metadata.json"

def _custom_dockerfile(self) -> str | None:
"""Configured dockerfile path, or None when it's the default."""
path = self.config.image.dockerfile_path
return path if path != "Dockerfile" else None

Comment thread
dulaj-me marked this conversation as resolved.
Outdated
def image_with_digest(self, tag: str = "latest") -> str:
image = self.image(tag)
try:
Expand Down Expand Up @@ -667,7 +673,10 @@
if warmup:
_build_warm_image(image)

def push(self, tag: str = "latest") -> None:
def push(self, tag: str = "latest", source_image: str | None = None) -> None:
if source_image and not tag:
last = source_image.rsplit("/", 1)[-1]
tag = last.rsplit(":", 1)[1] if ":" in last else "latest"
image = self.image(tag)
host = self.registry().split("/")[0]
login_cmd = ["docker", "login", host, "--username", "user", "--password-stdin"]
Expand All @@ -676,31 +685,18 @@

console.print(f"Pushing {image}")
self._metadata_path(tag).unlink(missing_ok=True)
if source_image:
if os.getenv("JIG_DISABLE_BUILDX"):
if subprocess.run(["docker", "tag", source_image, image]).returncode != 0:
raise JigError(f"Failed to retag {source_image}")
ok = subprocess.run(["docker", "push", image]).returncode == 0
else:
ok = _push_image_as_zstd(source_image, image, self._metadata_path(tag))
# Skip buildx for warmup-baked images: a buildx rebuild would drop the warmup layer.
if _image_is_warmed(image) or os.getenv("JIG_DISABLE_BUILDX"):
elif _image_is_warmed(image) or os.getenv("JIG_DISABLE_BUILDX"):
ok = subprocess.run(["docker", "push", image]).returncode == 0
else:
builder = _ensure_zstd_builder()
if not builder:
raise JigError("`docker buildx` is required to build images.")
cmd = [
"docker",
"buildx",
"build",
"--builder",
builder,
"--platform",
"linux/amd64",
"--push",
"--output",
f"type=image,name={image},{BUILDX_OUTPUT_OPTS}",
"--metadata-file",
str(self._metadata_path(tag)),
]
if self.config.image.dockerfile_path != "Dockerfile":
cmd.extend(["-f", self.config.image.dockerfile_path])
cmd.append(".")
ok = subprocess.run(cmd).returncode == 0
ok = _buildx_push(image, self._metadata_path(tag), dockerfile=self._custom_dockerfile())
Comment thread
dulaj-me marked this conversation as resolved.
Outdated
if not ok:
raise JigError("Push failed")
console.print("\N{CHECK MARK} Pushed")
Expand All @@ -717,9 +713,6 @@
self.build(tag, False, docker_args)
self.push(tag)
return
builder = _ensure_zstd_builder()
if not builder:
raise JigError("`docker buildx` is required to build images.")

host = self.registry().split("/")[0]
login_cmd = ["docker", "login", host, "--username", "user", "--password-stdin"]
Expand All @@ -733,27 +726,10 @@

console.print(f"Building and pushing {image}")
self._metadata_path(tag).unlink(missing_ok=True)
cmd = [
"docker",
"buildx",
"build",
"--builder",
builder,
"--platform",
"linux/amd64",
"--push",
"--output",
f"type=image,name={image},{BUILDX_OUTPUT_OPTS}",
"--metadata-file",
str(self._metadata_path(tag)),
]
if self.config.image.dockerfile_path != "Dockerfile":
cmd.extend(["-f", self.config.image.dockerfile_path])
extra_args = docker_args or os.getenv("DOCKER_BUILD_EXTRA_ARGS", "")
if extra_args:
cmd.extend(shlex.split(extra_args))
cmd.append(".")
if subprocess.run(cmd).returncode != 0:
extra_args = shlex.split(docker_args or os.getenv("DOCKER_BUILD_EXTRA_ARGS", ""))

Check warning on line 729 in src/together/lib/cli/api/beta/jig/jig.py

View check run for this annotation

Broly - Code Security Scanner / Broly Security Scan

[MEDIUM] Command injection via `DOCKER_BUILD_EXTRA_ARGS` environment variable passed to docker build with only `shlex.split` (no validation).

Command injection via `DOCKER_BUILD_EXTRA_ARGS` environment variable passed to docker build with only `shlex.split` (no validation).: The `DOCKER_BUILD_EXTRA_ARGS` environment variable (or `docker_args` parameter) is split with `shlex.split` and passed directly to the `docker build`/`buildx build` command list. While `shlex.split` prevents shell-level injection (since `subprocess.run` is called with a list, not `shell=True`), it allows injection of arbitrary docker build flags. An attacker who controls the environment can inject flags like `--build-arg` to override build arguments, `--secret` to expose additional secrets, or `--label` to modify image metadata. This requires control of the build environment's environment variables.
Comment thread
github-advanced-security[bot] marked this conversation as resolved.
Fixed
if not _buildx_push(
image, self._metadata_path(tag), dockerfile=self._custom_dockerfile(), extra_args=extra_args
Comment thread
dulaj-me marked this conversation as resolved.
Outdated
):
raise JigError("Build+push failed")
console.print("\N{CHECK MARK} Built and pushed")

Expand Down Expand Up @@ -1187,9 +1163,9 @@
jig.build(tag, warmup, docker_args)


def push(jig: Jig, tag: str) -> None:
def push(jig: Jig, tag: str, source_image: str | None = None) -> None:
"""Push image to registry"""
jig.push(tag)
jig.push(tag, source_image)


def deploy(
Expand Down Expand Up @@ -1456,13 +1432,18 @@


def push_cli(
tag: Annotated[str, Parameter(help="Image tag")] = "latest",
tag: Annotated[
Optional[str], Parameter(help="Image tag (defaults to 'latest', or source image's tag when --image is used)")
] = None,
image: Annotated[
Optional[str], Parameter(name="--image", help="Existing local image to push (layers re-encoded as zstd)")
] = None,
*,
config: CLIConfigParameter,
toml_config: TomlConfigParameter = None,
) -> None:
"""Push image to registry."""
_run_jig_cmd(config, toml_config, lambda jig: push(jig, tag))
_run_jig_cmd(config, toml_config, lambda jig: push(jig, tag or ("" if image else "latest"), image))


def deploy_cli(
Expand Down Expand Up @@ -1673,6 +1654,71 @@
return r.returncode == 0 and r.stdout.strip() == "true"


def _buildx_push(
image: str,
metadata_file: Path,
dockerfile: str | None = None,
Comment thread
dulaj-me marked this conversation as resolved.
Outdated
context: str = ".",
extra_args: list[str] | None = None,
) -> bool:
"""Run `docker buildx build --push` through the zstd builder."""
builder = _ensure_zstd_builder()
if not builder:
raise JigError("`docker buildx` is required to build images.")
cmd = [
"docker",
"buildx",
"build",
"--builder",
builder,
"--platform",
"linux/amd64",
"--push",
"--output",
f"type=image,name={image},{BUILDX_OUTPUT_OPTS}",
"--metadata-file",
str(metadata_file),
]
if dockerfile:
cmd.extend(["-f", dockerfile])
Comment thread
dulaj-me marked this conversation as resolved.
Outdated
cmd.extend(extra_args or [])
cmd.append(context)
return subprocess.run(cmd).returncode == 0


def _push_image_as_zstd(source_image: str, image: str, metadata_file: Path) -> bool:
"""Push a local image with its layers re-encoded as zstd.

`docker push` through a legacy (non-containerd) image store recompresses
layers as gzip, losing the zstd cold-start benefit. Instead, export the
image as an OCI layout (`docker save`) and rebuild it through the zstd
builder with force-compression, which re-encodes every layer on push.
"""
with tempfile.TemporaryDirectory(prefix="jig-push-") as tmp:
layout = Path(tmp) / "oci"
layout.mkdir()
console.print(f"Exporting {source_image}")
save = subprocess.Popen(["docker", "save", source_image], stdout=subprocess.PIPE)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think docker save is the really slow one

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just piped to untar to make it bit faster. Not sure if there any other way to get the image to buildkit context

tar = subprocess.run(["tar", "-x", "-C", str(layout)], stdin=save.stdout)
if save.stdout:
save.stdout.close()
if save.wait() != 0 or tar.returncode != 0:
raise JigError(f"Failed to export {source_image}")
index = layout / "index.json"
if not index.exists():
raise JigError("`docker save` did not produce an OCI layout; Docker 25+ is required for --image")
digest = json.loads(index.read_text())["manifests"][0]["digest"]
ctx = Path(tmp) / "ctx"
ctx.mkdir()
(ctx / "Dockerfile").write_text("FROM source-image\n")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's sufficient to do f"FROM {source_image}" with force-compression=true

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

or even

ARG SOURCE_IMAGE
FROM ${SOURCE_IMAGE}
docker buildx build \
  --file Dockerfile.recompress \
  --build-arg SOURCE_IMAGE=registry.example.com/team/app:gzip \
  --pull \
  --output type=registry,name=...:zstd,compression=zstd,compression-level=10,force-compression=true,oci-mediatypes=true \
  .

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you mean for remote registry hosted images? I don't think you can do this with local images.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no you can

return _buildx_push(
image,
metadata_file,
context=str(ctx),
extra_args=["--build-context", f"source-image=oci-layout://{layout}:source-image@{digest}"],
)


def _ensure_zstd_builder(name: str = "jig-zstd") -> str | None:
"""Return the name of a docker-container buildx builder, creating one if needed.

Expand Down
Loading