From 53911f2e956571fd67a62277f2b1d4d7ab505119 Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Mon, 27 Jul 2026 11:06:38 -0700 Subject: [PATCH 1/2] Switched ventis build from building sequentially to in parallel with docker buildx bake --- .dockerignore | 1 + .gitignore | 1 + README.md | 1 + tests/test_cli.py | 213 ++++++++++++++++++++++++++++----------- ventis/cli.py | 79 ++++++++++++--- ventis/stub_generator.py | 4 +- 6 files changed, 227 insertions(+), 72 deletions(-) diff --git a/.dockerignore b/.dockerignore index a5c47a7..ab0acb8 100644 --- a/.dockerignore +++ b/.dockerignore @@ -19,3 +19,4 @@ AWSCLIV2.pkg docker_container grpc_stubs stubs +tests diff --git a/.gitignore b/.gitignore index db416f0..6fc4ab9 100644 --- a/.gitignore +++ b/.gitignore @@ -40,3 +40,4 @@ docker_container/ !.env.example AWSCLIV2.pkg .python-version +uv.lock \ No newline at end of file diff --git a/README.md b/README.md index 51a2be6..4abb20a 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ Note: Installation of ventis only needs to be done on the machine where you are - **Python 3.10+** - **Docker** — Used to manage agents. +- **Docker Buildx** (optional) — If available, `ventis build` builds all agent/workflow images in a single parallel `docker buildx bake` pass; otherwise it falls back to building them sequentially. --- diff --git a/tests/test_cli.py b/tests/test_cli.py index 8bec94b..eb898ae 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,3 +1,4 @@ +import json import os import sys import tempfile @@ -98,79 +99,175 @@ def test_preflight_does_not_require_ssh_fields(self, require_docker, ensure_grpc class CliBuildTests(unittest.TestCase): - def test_build_does_not_build_global_controller_image(self): + def _run_build( + self, project_dir, agent_yaml_paths, buildx_available, platform="linux/amd64" + ): + """Run cmd_build against project_dir with docker/subprocess calls mocked. + + Returns the list of subprocess.run commands that were invoked. + """ + config_path = project_dir / "config" / "global_controller.yaml" + args = SimpleNamespace(config=str(config_path)) + docker_calls = [] + + def fake_run(cmd, check): + docker_calls.append(cmd) + return SimpleNamespace(returncode=0) + + with ( + patch( + "ventis.cli._get_package_dir", + return_value=str(project_dir / "package"), + ), + patch( + "ventis.cli.glob.glob", + side_effect=[agent_yaml_paths, ["proto/a.proto"]], + ), + patch("ventis.stub_generator.generate_stub"), + patch("ventis.stub_generator.generate_docker"), + patch("ventis.stub_generator.generate_workflow_docker"), + patch("ventis.cli.subprocess.run", side_effect=fake_run), + patch("ventis.cli._docker_available", return_value=buildx_available), + patch("ventis.cli._docker_platform", return_value=platform), + ): + cwd = os.getcwd() + os.chdir(project_dir) + try: + cli.cmd_build(args) + finally: + os.chdir(cwd) + + return docker_calls + + def _write_agent_and_workflow_config(self, project_dir): + """Scaffold a project with one agent + one workflow entry; returns the agent YAML path.""" + (project_dir / "config").mkdir() + (project_dir / "agents").mkdir() + (project_dir / "workflows").mkdir() + (project_dir / "docker").mkdir() + (project_dir / "docker" / "global-controller.Dockerfile").write_text( + "FROM scratch\n" + ) + (project_dir / "agents" / "example_agent.py").write_text("print('ok')\n") + (project_dir / "workflows" / "example_workflow.py").write_text("print('ok')\n") + agent_yaml = project_dir / "agents" / "example_agent.yaml" + agent_yaml.write_text("agent:\n name: ExampleAgent\n") + config_path = project_dir / "config" / "global_controller.yaml" + config_path.write_text( + yaml.safe_dump( + { + "agents": [ + { + "name": "ExampleAgent", + "entrypoint": "agents/example_agent.py", + "provider": "local", + }, + { + "name": "Workflow", + "type": "workflow", + "workflow_file": "workflows/example_workflow.py", + "provider": "local", + }, + ] + } + ) + ) + return agent_yaml + + def test_build_falls_back_to_sequential_docker_build_without_buildx(self): with tempfile.TemporaryDirectory() as tmpdir: project_dir = Path(tmpdir) - (project_dir / "config").mkdir() - (project_dir / "agents").mkdir() - (project_dir / "workflows").mkdir() - (project_dir / "docker").mkdir() - (project_dir / "docker" / "global-controller.Dockerfile").write_text( - "FROM scratch\n" + agent_yaml = self._write_agent_and_workflow_config(project_dir) + + docker_calls = self._run_build( + project_dir, [str(agent_yaml)], buildx_available=False ) - (project_dir / "agents" / "example_agent.py").write_text("print('ok')\n") - (project_dir / "workflows" / "example_workflow.py").write_text( - "print('ok')\n" + + flattened = [" ".join(call) for call in docker_calls] + self.assertFalse( + any("global-controller.Dockerfile" in call for call in flattened) + ) + self.assertEqual( + sum(call[:2] == ["docker", "build"] for call in docker_calls), 2 + ) + self.assertFalse( + any(call[:3] == ["docker", "buildx", "bake"] for call in docker_calls) + ) + + def test_build_uses_buildx_bake_when_available(self): + with tempfile.TemporaryDirectory() as tmpdir: + project_dir = Path(tmpdir) + agent_yaml = self._write_agent_and_workflow_config(project_dir) + + docker_calls = self._run_build( + project_dir, [str(agent_yaml)], buildx_available=True ) - agent_yaml = project_dir / "agents" / "example_agent.yaml" - agent_yaml.write_text("agent:\n name: ExampleAgent\n") + + bake_file = project_dir / "docker_container" / "docker-bake.json" + self.assertTrue(bake_file.is_file()) + with open(bake_file) as f: + bake_config = json.load(f) + + self.assertEqual( + sum(call[:2] == ["docker", "build"] for call in docker_calls), 0 + ) + bake_calls = [ + call for call in docker_calls if call[:3] == ["docker", "buildx", "bake"] + ] + self.assertEqual(len(bake_calls), 1) + self.assertIn("--file", bake_calls[0]) + self.assertEqual( + os.path.realpath(bake_calls[0][bake_calls[0].index("--file") + 1]), + os.path.realpath(bake_file), + ) + + targets = bake_config["target"] + self.assertEqual(len(targets), 2) + self.assertEqual( + os.path.realpath(targets["exampleagent"]["context"]), + os.path.realpath(project_dir / "docker_container" / "ExampleAgent"), + ) + self.assertTrue(os.path.isabs(targets["exampleagent"]["context"])) + self.assertEqual(targets["exampleagent"]["tags"], ["ventis-exampleagent"]) + self.assertEqual(targets["exampleagent"]["platforms"], ["linux/amd64"]) + self.assertEqual(targets["exampleagent"]["output"], ["type=docker"]) + self.assertEqual( + os.path.realpath(targets["workflow"]["context"]), + os.path.realpath(project_dir / "docker_container" / "Workflow"), + ) + self.assertEqual(targets["workflow"]["tags"], ["ventis-workflow"]) + + def test_build_with_no_agents_builds_nothing(self): + with tempfile.TemporaryDirectory() as tmpdir: + project_dir = Path(tmpdir) + (project_dir / "config").mkdir() + (project_dir / "agents").mkdir() + config_path = project_dir / "config" / "global_controller.yaml" + config_path.write_text(yaml.safe_dump({"agents": []})) + + docker_calls = self._run_build(project_dir, [], buildx_available=True) + + self.assertFalse(any(call[0] == "docker" for call in docker_calls)) + + def test_build_skips_agent_without_entrypoint(self): + with tempfile.TemporaryDirectory() as tmpdir: + project_dir = Path(tmpdir) + (project_dir / "config").mkdir() + (project_dir / "agents").mkdir() config_path = project_dir / "config" / "global_controller.yaml" config_path.write_text( yaml.safe_dump( { "agents": [ - { - "name": "ExampleAgent", - "entrypoint": "agents/example_agent.py", - "provider": "local", - }, - { - "name": "Workflow", - "type": "workflow", - "workflow_file": "workflows/example_workflow.py", - "provider": "local", - }, + {"name": "NoEntrypointAgent", "provider": "local"}, ] } ) ) - args = SimpleNamespace(config=str(config_path)) - docker_calls = [] - - def fake_run(cmd, check): - docker_calls.append(cmd) - return SimpleNamespace(returncode=0) - - with ( - patch( - "ventis.cli._get_package_dir", - return_value=str(project_dir / "package"), - ), - patch( - "ventis.cli.glob.glob", - side_effect=[[str(agent_yaml)], ["proto/a.proto"]], - ), - patch("ventis.stub_generator.generate_stub"), - patch("ventis.stub_generator.generate_docker"), - patch("ventis.stub_generator.generate_workflow_docker"), - patch("ventis.cli.subprocess.run", side_effect=fake_run), - patch.dict(os.environ, {}, clear=False), - ): - cwd = os.getcwd() - os.chdir(project_dir) - try: - cli.cmd_build(args) - finally: - os.chdir(cwd) + docker_calls = self._run_build(project_dir, [], buildx_available=True) - flattened = [" ".join(call) for call in docker_calls] - self.assertFalse( - any("global-controller.Dockerfile" in call for call in flattened) - ) - self.assertEqual( - sum(call[:2] == ["docker", "build"] for call in docker_calls), 2 - ) + self.assertFalse(any(call[0] == "docker" for call in docker_calls)) if __name__ == "__main__": diff --git a/ventis/cli.py b/ventis/cli.py index 31d8d5d..3aceb18 100644 --- a/ventis/cli.py +++ b/ventis/cli.py @@ -9,6 +9,7 @@ import argparse import glob +import json import logging import os import shutil @@ -60,13 +61,13 @@ def _docker_build_cmd(*args): return ["docker", "build", "--platform", _docker_platform(), *args] -def _docker_available(): +def _docker_available(probe_cmd=("docker", "info")): if not shutil.which("docker"): return False try: result = subprocess.run( - ["docker", "info"], + list(probe_cmd), capture_output=True, text=True, check=False, @@ -77,6 +78,32 @@ def _docker_available(): return result.returncode == 0 +def _write_bake_file(bake_targets, bake_file_path, platform): + """Write a docker-buildx-bake JSON file describing all build targets. + + Context paths are written absolute: `docker buildx bake` resolves relative + `context` values against the invocation cwd (not the bake file's own + directory), so an absolute path sidesteps that ambiguity entirely. + """ + bake_config = { + "target": { + target["name"]: { + "context": os.path.abspath(target["context"]), + "dockerfile": "Dockerfile", + "tags": [target["image_name"]], + "platforms": [platform], + "output": ["type=docker"], + # type=docker could be changed to tarring it up, which would be + # faster but skipped because that change would alter ventis deploy + } + for target in bake_targets + } + } + with open(bake_file_path, "w") as f: + json.dump(bake_config, f, indent=2) + return bake_file_path + + def _require_docker_for_ec2(command_name): if _docker_available(): return @@ -218,8 +245,9 @@ def cmd_build(args): ) # -------------------------------------------------------------- # - # Step 4: Generate Docker contexts and build images # + # Step 4: Generate Docker contexts # # -------------------------------------------------------------- # + bake_targets = [] for agent_cfg in agents: agent_name = agent_cfg["name"] agent_type = agent_cfg.get("type", "agent") @@ -248,12 +276,6 @@ def cmd_build(args): api_port=agent_cfg.get("api_port", 8080), ) - image_name = f"ventis-{agent_name.lower()}" - logger.info("Building Docker image: %s", image_name) - subprocess.run( - _docker_build_cmd("-t", image_name, docker_context), check=True - ) - else: # Agent container entrypoint = agent_cfg.get("entrypoint") @@ -296,10 +318,43 @@ def cmd_build(args): stub_files=stub_paths, ) - image_name = f"ventis-{agent_name.lower()}" - logger.info("Building Docker image: %s", image_name) + bake_targets.append( + { + "name": agent_name.lower(), + "context": docker_context, + "image_name": f"ventis-{agent_name.lower()}", + } + ) + + # -------------------------------------------------------------- # + # Step 5: Build all Docker images # + # -------------------------------------------------------------- # + if not bake_targets: + logger.info("No Docker images to build.") + elif _docker_available() and _docker_available(("docker", "buildx", "version")): + docker_container_dir = os.path.join(project_dir, "docker_container") + os.makedirs(docker_container_dir, exist_ok=True) + bake_file_path = os.path.join(docker_container_dir, "docker-bake.json") + _write_bake_file(bake_targets, bake_file_path, _docker_platform()) + + target_names = [target["name"] for target in bake_targets] + logger.info( + "Building %d Docker image(s) via `docker buildx bake`.", + len(bake_targets), + ) + subprocess.run( + ["docker", "buildx", "bake", "--file", bake_file_path, *target_names], + check=True, + ) + else: + logger.info( + "docker buildx unavailable; falling back to sequential `docker build`." + ) + for target in bake_targets: + logger.info("Building Docker image: %s", target["image_name"]) subprocess.run( - _docker_build_cmd("-t", image_name, docker_context), check=True + _docker_build_cmd("-t", target["image_name"], target["context"]), + check=True, ) logger.info("Build complete.") diff --git a/ventis/stub_generator.py b/ventis/stub_generator.py index eb7f1b7..5432654 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -358,7 +358,7 @@ def generate_docker( EXPOSE 50051 -CMD python local_controller.py --port 50051 +CMD ["python", "local_controller.py", "--port", "50051"] """ with open(os.path.join(output_dir, "Dockerfile"), "w") as f: f.write(dockerfile) @@ -473,7 +473,7 @@ def start_lc(): EXPOSE 50051 EXPOSE {api_port} -CMD python workflow_launcher.py +CMD ["python", "workflow_launcher.py"] """ with open(os.path.join(output_dir, "Dockerfile"), "w") as f: f.write(dockerfile) From 1085d7b5d24259daa4cec265c67fc0d88263282e Mon Sep 17 00:00:00 2001 From: Saaketh Sodanapalli Date: Wed, 29 Jul 2026 15:10:32 -0700 Subject: [PATCH 2/2] added uv for pip install + caching --- ventis/stub_generator.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/ventis/stub_generator.py b/ventis/stub_generator.py index 5432654..0fd42f0 100644 --- a/ventis/stub_generator.py +++ b/ventis/stub_generator.py @@ -344,12 +344,14 @@ def generate_docker( # ---- Dockerfile ------------------------------------------------------ agent_basename = os.path.basename(agent_file) - dockerfile = f"""FROM python:3.11-slim + dockerfile = f"""# syntax=docker/dockerfile:1 +FROM python:3.11-slim +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ WORKDIR /app COPY requirements.txt . -RUN pip install --no-cache-dir -r requirements.txt +RUN --mount=type=cache,target=/root/.cache/uv uv pip install --system -r requirements.txt COPY . . @@ -461,12 +463,14 @@ def start_lc(): f.write(launcher) # ---- Dockerfile ------------------------------------------------------ - dockerfile = f"""FROM python:3.11-slim + dockerfile = f"""# syntax=docker/dockerfile:1 +FROM python:3.11-slim +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ WORKDIR /app COPY requirements.txt . -RUN pip install --no-cache-dir -r requirements.txt +RUN --mount=type=cache,target=/root/.cache/uv uv pip install --system -r requirements.txt COPY . .