Skip to content

fix: tweaks to codegen docker runner #621

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Feb 24, 2025
Merged
Show file tree
Hide file tree
Changes from all 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
6 changes: 2 additions & 4 deletions Dockerfile-runner
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ ENV NVM_DIR=/root/.nvm \
PYTHONPATH="/usr/local/lib/python3.13/site-packages" \
IS_SANDBOX=True \
HATCH_BUILD_HOOKS_ENABLE=1

# Update packages lists and install git and curl
RUN apt-get update && apt-get install -y \
git \
Expand Down Expand Up @@ -41,17 +42,14 @@ RUN node --version \
&& pnpm --version \
&& python --version

# Install codegen from source instead of PyPI
# Build codegen
WORKDIR /codegen-sdk
COPY . .

# Install dependencies and build codegen with entry points
RUN --mount=type=cache,target=/root/.cache/uv \
uv venv && source .venv/bin/activate \
&& uv sync --frozen --no-dev --all-extras \
&& uv pip install --system -e . --no-deps \
&& uv pip install --system .

RUN codegen --version

# Create a non-root user for local development + debugging
Expand Down
6 changes: 3 additions & 3 deletions src/codegen/cli/commands/create/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,17 +68,17 @@
@click.command(name="create")
@requires_init
@click.argument("name", type=str)
@click.argument("path", type=click.Path(path_type=Path), default=Path.cwd())
@click.argument("path", type=click.Path(path_type=Path), default=None)
@click.option("--description", "-d", default=None, help="Description of what this codemod does.")
@click.option("--overwrite", is_flag=True, help="Overwrites function if it already exists.")
def create_command(session: CodegenSession, name: str, path: Path, description: str | None = None, overwrite: bool = False):
def create_command(session: CodegenSession, name: str, path: Path | None, description: str | None = None, overwrite: bool = False):
"""Create a new codegen function.

NAME is the name/label for the function
PATH is where to create the function (default: current directory)
"""
# Get the target path for the function
codemod_path, prompt_path = get_target_paths(name, path)
codemod_path, prompt_path = get_target_paths(name, path or Path.cwd())

# Check if file exists
if codemod_path.exists() and not overwrite:
Expand All @@ -92,8 +92,8 @@
if description:
# Use API to generate implementation
with create_spinner("Generating function (using LLM, this will take ~10s)") as status:
response = RestAPI(get_current_token()).create(name=name, query=description)

Check failure on line 95 in src/codegen/cli/commands/create/main.py

View workflow job for this annotation

GitHub Actions / mypy

error: Argument 1 to "RestAPI" has incompatible type "str | None"; expected "str" [arg-type]
code = convert_to_cli(response.code, session.config.repository.language, name)

Check failure on line 96 in src/codegen/cli/commands/create/main.py

View workflow job for this annotation

GitHub Actions / mypy

error: Argument 2 to "convert_to_cli" has incompatible type "str | None"; expected "str" [arg-type]
prompt_path.parent.mkdir(parents=True, exist_ok=True)
prompt_path.write_text(response.context)
else:
Expand Down
15 changes: 9 additions & 6 deletions src/codegen/cli/commands/start/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,24 +8,28 @@
from rich.panel import Panel

from codegen.configs.models.secrets import SecretsConfig
from codegen.git.repo_operator.local_git_repo import LocalGitRepo
from codegen.git.schemas.repo_config import RepoConfig
from codegen.shared.network.port import get_free_port


@click.command(name="start")
@click.option("--platform", "-t", type=click.Choice(["linux/amd64", "linux/arm64", "linux/amd64,linux/arm64"]), default="linux/amd64,linux/arm64", help="Target platform(s) for the Docker image")
@click.option("--port", "-p", type=int, default=8000)
@click.option("--port", "-p", type=int, default=None, help="Port to run the server on")
@click.option("--detached", "-d", is_flag=True, default=False, help="Starts up the server as detached background process")
def start_command(port: int, platform: str, detached: bool):
def start_command(port: int | None, platform: str, detached: bool):
"""Starts a local codegen server"""
codegen_version = version("codegen")
rich.print(f"[bold green]Codegen version:[/bold green] {codegen_version}")
codegen_root = Path(__file__).parent.parent.parent.parent.parent.parent
if port is None:
port = get_free_port()

try:
rich.print("[bold blue]Building Docker image...[/bold blue]")
_build_docker_image(codegen_root, platform)
rich.print("[bold blue]Starting Docker container...[/bold blue]")
_run_docker_container(port, platform, detached)
_run_docker_container(port, detached)
rich.print(Panel(f"[green]Server started successfully![/green]\nAccess the server at: [bold]http://0.0.0.0:{port}[/bold]", box=ROUNDED, title="Codegen Server"))
except subprocess.CalledProcessError as e:
rich.print(f"[bold red]Error:[/bold red] Failed to {e.cmd[0]} Docker container")
Expand Down Expand Up @@ -55,20 +59,19 @@

def _run_docker_container(port: int, detached: bool):
repo_path = Path.cwd().resolve()
repo_config = RepoConfig.from_repo_path(repo_path)

Check failure on line 62 in src/codegen/cli/commands/start/main.py

View workflow job for this annotation

GitHub Actions / mypy

error: Argument 1 to "from_repo_path" of "RepoConfig" has incompatible type "Path"; expected "str" [arg-type]
container_repo_path = f"/app/git/{repo_config.name}"
envvars = {
"REPOSITORY_LANGUAGE": repo_config.language.value,
"REPOSITORY_OWNER": repo_config.organization_name,
"REPOSITORY_OWNER": LocalGitRepo(repo_path).owner,
"REPOSITORY_PATH": container_repo_path,
"GITHUB_TOKEN": SecretsConfig().github_token,
}
envvars_args = [arg for k, v in envvars.items() for arg in ("--env", f"{k}={v}")]

mount_args = ["-v", f"{repo_path}:{container_repo_path}"]
run_mode = "-d" if detached else "-it"
entry_point = f"uv run --frozen uvicorn codegen.runner.sandbox.server:app --host 0.0.0.0 --port {port}"
run_cmd = ["docker", "run", run_mode, "-p", f"8000:{port}", *mount_args, *envvars_args, "codegen-runner", entry_point]
run_cmd = ["docker", "run", run_mode, "-p", f"{port}:{port}", *mount_args, *envvars_args, "codegen-runner", entry_point]

rich.print(f"run_cmd: {str.join(' ', run_cmd)}")
subprocess.run(run_cmd, check=True)
1 change: 1 addition & 0 deletions src/codegen/git/repo_operator/local_git_repo.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ def origin_remote(self) -> Remote | None:
"""Returns the url of the first remote found on the repo, or None if no remotes are set"""
if self.has_remote():
return self.git_cli.remote("origin")
return None

@cached_property
def base_url(self) -> str | None:
Expand Down
1 change: 1 addition & 0 deletions src/codegen/git/schemas/repo_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,4 +38,5 @@ def repo_path(self) -> Path:
def organization_name(self) -> str | None:
if self.full_name is not None:
return self.full_name.split("/")[0]

return None
2 changes: 1 addition & 1 deletion src/codegen/runner/sandbox/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@

def __init__(self, repo_config: RepoConfig) -> None:
self.repo = repo_config
self.op = RepoOperator(repo_config=self.repo, setup_option=SetupOption.PULL_OR_CLONE)
self.op = RepoOperator(repo_config=self.repo, setup_option=SetupOption.PULL_OR_CLONE, bot_commit=True)
self.commit = self.op.git_cli.head.commit

async def warmup(self) -> None:
Expand Down Expand Up @@ -82,7 +82,7 @@

branch_config.custom_base_branch = branch_config.custom_base_branch or self.codebase.default_branch
self.executor.remote_repo.set_up_base_branch(branch_config.custom_base_branch)
self.executor.remote_repo.set_up_head_branch(branch_config.custom_head_branch, branch_config.force_push_head_branch)

Check failure on line 85 in src/codegen/runner/sandbox/runner.py

View workflow job for this annotation

GitHub Actions / mypy

error: Argument 1 to "set_up_head_branch" of "SandboxRepo" has incompatible type "str | None"; expected "str" [arg-type]

response = CreateBranchResponse()
if "codebase.flag_instance" in request.codemod.user_code:
Expand All @@ -92,9 +92,9 @@
response.group_segments = [group.segment for group in flag_groups]
if len(flag_groups) == 0:
logger.info("No flag groups found. Running without flagging.")
flag_groups = [None]

Check failure on line 95 in src/codegen/runner/sandbox/runner.py

View workflow job for this annotation

GitHub Actions / mypy

error: List item 0 has incompatible type "None"; expected "Group" [list-item]
else:
flag_groups = [None]

Check failure on line 97 in src/codegen/runner/sandbox/runner.py

View workflow job for this annotation

GitHub Actions / mypy

error: List item 0 has incompatible type "None"; expected "Group" [list-item]

# TODO: do this as part of find_flag_groups?
max_prs = request.grouping_config.max_prs
Expand Down
17 changes: 17 additions & 0 deletions src/codegen/shared/network/port.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
import socket
from contextlib import closing


def get_free_port() -> int:
"""Find and return a free port on localhost"""
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
s.bind(("", 0))
s.listen(1)
port = s.getsockname()[1]
return int(port)


def is_port_free(port: int, host: str = "localhost") -> bool:
"""Check if a port is free on localhost"""
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
return s.connect_ex((host, port)) != 0
17 changes: 3 additions & 14 deletions tests/integration/codegen/runner/conftest.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,4 @@
import socket
from collections.abc import Generator
from contextlib import closing
from unittest.mock import Mock

import pytest
Expand All @@ -11,16 +9,7 @@
from codegen.git.schemas.repo_config import RepoConfig
from codegen.runner.clients.codebase_client import CodebaseClient
from codegen.shared.enums.programming_language import ProgrammingLanguage


@pytest.fixture
def get_free_port():
"""Find and return a free port on localhost"""
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
s.bind(("", 0))
s.listen(1)
port = s.getsockname()[1]
return port
from codegen.shared.network.port import get_free_port


@pytest.fixture()
Expand All @@ -44,7 +33,7 @@


@pytest.fixture
def codebase_client(repo_config: RepoConfig, get_free_port) -> Generator[CodebaseClient, None, None]:
sb_client = CodebaseClient(repo_config=repo_config, port=get_free_port)
def codebase_client(repo_config: RepoConfig) -> Generator[CodebaseClient, None, None]:
sb_client = CodebaseClient(repo_config=repo_config, port=get_free_port())
sb_client.runner = Mock()

Check failure on line 38 in tests/integration/codegen/runner/conftest.py

View workflow job for this annotation

GitHub Actions / mypy

error: "CodebaseClient" has no attribute "runner" [attr-defined]
yield sb_client
Loading