Skip to content

feat: add --daemon to codegen run #659

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 5 commits into from
Feb 26, 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
4 changes: 4 additions & 0 deletions Dockerfile-runner
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,10 @@ RUN apt-get update && apt-get install -y \
# Cleanup apt cache to reduce image size
&& rm -rf /var/lib/apt/lists/*

# Set git config
RUN git config --global user.email "[email protected]" \
&& git config --global user.name "codegen-bot"

# Install nvm and Node.js
SHELL ["/bin/bash", "-c"]
RUN curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.39.5/install.sh | bash \
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ dependencies = [
"docker>=6.1.3",
"urllib3>=2.0.0",
"datasets",
"colorlog>=6.9.0",
]

license = { text = "Apache-2.0" }
Expand Down
10 changes: 10 additions & 0 deletions src/codegen/cli/commands/run/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,22 @@
@requires_init
@click.argument("label", required=True)
@click.option("--web", is_flag=True, help="Run the function on the web service instead of locally")
@click.option("--daemon", is_flag=True, help="Run the function against a running daemon")
@click.option("--diff-preview", type=int, help="Show a preview of the first N lines of the diff")
@click.option("--arguments", type=str, help="Arguments as a json string to pass as the function's 'arguments' parameter")
def run_command(
session: CodegenSession,
label: str,
web: bool = False,
daemon: bool = False,
diff_preview: int | None = None,
arguments: str | None = None,
):
"""Run a codegen function by its label."""
if web and daemon:
msg = "Cannot enable run on both the web and daemon"
raise ValueError(msg)

# Ensure venv is initialized
venv = VenvManager(session.codegen_dir)
if not venv.is_initialized():
Expand Down Expand Up @@ -54,6 +60,10 @@ def run_command(
from codegen.cli.commands.run.run_cloud import run_cloud

run_cloud(session, codemod, diff_preview=diff_preview)
elif daemon:
from codegen.cli.commands.run.run_daemon import run_daemon

run_daemon(session, codemod, diff_preview=diff_preview)
else:
from codegen.cli.commands.run.run_local import run_local

Expand Down
87 changes: 87 additions & 0 deletions src/codegen/cli/commands/run/run_daemon.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
import rich
import rich_click as click
from rich.panel import Panel

from codegen.cli.auth.session import CodegenSession
from codegen.cli.commands.start.docker_container import DockerContainer
from codegen.cli.errors import ServerError
from codegen.cli.rich.codeblocks import format_command
from codegen.cli.rich.spinners import create_spinner
from codegen.runner.clients.docker_client import DockerClient
from codegen.runner.enums.warmup_state import WarmupState


def run_daemon(session: CodegenSession, function, diff_preview: int | None = None):
"""Run a function on the cloud service.

Args:
session: The current codegen session
function: The function to run
diff_preview: Number of lines of diff to preview (None for all)
"""
with create_spinner(f"Running {function.name}...") as status:
try:
client = _get_docker_client(session)
run_output = client.run_function(function, commit=not diff_preview)
rich.print(f"✅ Ran {function.name} successfully")

if run_output.logs:
rich.print("")
panel = Panel(run_output.logs, title="[bold]Logs[/bold]", border_style="blue", padding=(1, 2), expand=False)
rich.print(panel)

if run_output.error:
rich.print("")
panel = Panel(run_output.error, title="[bold]Error[/bold]", border_style="red", padding=(1, 2), expand=False)
rich.print(panel)

if run_output.observation:
# Only show diff preview if requested
if diff_preview:
rich.print("") # Add some spacing

# Split and limit diff to requested number of lines
diff_lines = run_output.observation.splitlines()
truncated = len(diff_lines) > diff_preview
limited_diff = "\n".join(diff_lines[:diff_preview])

if truncated:
limited_diff += "\n\n...\n\n[yellow]diff truncated to {diff_preview} lines, view the full change set on your local file system after using run with `--apply-local`[/yellow]"

panel = Panel(limited_diff, title="[bold]Diff Preview[/bold]", border_style="blue", padding=(1, 2), expand=False)
rich.print(panel)
else:
rich.print("")
rich.print("[yellow] No changes were produced by this codemod[/yellow]")

if diff_preview:
rich.print("[green]✓ Changes have been applied to your local filesystem[/green]")
rich.print("[yellow]→ Don't forget to commit your changes:[/yellow]")
rich.print(format_command("git add ."))
rich.print(format_command("git commit -m 'Applied codemod changes'"))

except ServerError as e:
status.stop()
raise click.ClickException(str(e))


def _get_docker_client(session: CodegenSession) -> DockerClient:
repo_name = session.config.repository.name
if (container := DockerContainer.get(repo_name)) is None:
msg = f"Codegen runner does not exist for {repo_name}. Please run 'codegen start' from {session.config.repository.path}."
raise click.ClickException(msg)

if not container.is_running():
msg = f"Codegen runner for {repo_name} is not running. Please run 'codegen start' from {session.config.repository.path}."
raise click.ClickException(msg)

client = DockerClient(container)
if not client.is_running():
msg = "Codebase server is not running. Please stop the container and restart."
raise click.ClickException(msg)

if client.server_info().warmup_state != WarmupState.COMPLETED:
msg = "Runner has not finished parsing the codebase. Please wait a moment and try again."
raise click.ClickException(msg)

return client
71 changes: 57 additions & 14 deletions src/codegen/cli/commands/start/docker_container.py
Original file line number Diff line number Diff line change
@@ -1,31 +1,74 @@
from functools import cached_property

import docker
from docker import DockerClient
from docker.errors import APIError, NotFound
from docker.models.containers import Container


class DockerContainer:
_client: docker.DockerClient
host: str | None
port: int | None
name: str
_client: DockerClient
_container: Container | None

def __init__(self, client: docker.DockerClient, name: str, port: int | None = None, host: str | None = None):
def __init__(self, client: DockerClient, container: Container) -> None:
self._client = client
self.host = host
self.port = port
self.name = name
self._container = container

@classmethod
def get(cls, name: str) -> "DockerContainer | None":
try:
client = docker.from_env()
container = client.containers.get(name)
return cls(client=client, container=container)
except NotFound:
return None

@cached_property
def name(self) -> str:
return self._container.name

@cached_property
def host(self) -> str | None:
if not self.is_running():
return None

host_config = next(iter(self._container.ports.values()))[0]
return host_config["HostIp"]

@cached_property
def port(self) -> int | None:
if not self.is_running():
return None

host_config = next(iter(self._container.ports.values()))[0]
return host_config["HostPort"]

def is_running(self) -> bool:
try:
container = self._client.containers.get(self.name)
return container.status == "running"
except docker.errors.NotFound:
return self._container.status == "running"
except NotFound:
return False

def start(self) -> bool:
try:
container = self._client.containers.get(self.name)
container.start()
self._container.start()
return True
except (NotFound, APIError):
return False

def stop(self) -> bool:
try:
self._container.stop()
return True
except (NotFound, APIError):
return False

def remove(self) -> bool:
try:
self.stop()
self._container.remove()
return True
except (docker.errors.NotFound, docker.errors.APIError):
except (NotFound, APIError):
return False

def __str__(self) -> str:
Expand Down
20 changes: 8 additions & 12 deletions src/codegen/cli/commands/start/docker_fleet.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import docker
from docker.errors import NotFound

from codegen.cli.commands.start.docker_container import DockerContainer

Expand All @@ -15,19 +16,14 @@ def __init__(self, containers: list[DockerContainer]):
def load(cls) -> "DockerFleet":
try:
client = docker.from_env()
containers = client.containers.list(all=True, filters={"ancestor": CODEGEN_RUNNER_IMAGE})
codegen_containers = []
for container in containers:
filters = {"ancestor": CODEGEN_RUNNER_IMAGE}
containers = []
for container in client.containers.list(all=True, filters=filters):
if container.attrs["Config"]["Image"] == CODEGEN_RUNNER_IMAGE:
if container.status == "running":
host_config = next(iter(container.ports.values()))[0]
codegen_container = DockerContainer(client=client, host=host_config["HostIp"], port=host_config["HostPort"], name=container.name)
else:
codegen_container = DockerContainer(client=client, name=container.name)
codegen_containers.append(codegen_container)

return cls(containers=codegen_containers)
except docker.errors.NotFound:
containers.append(DockerContainer(client=client, container=container))

return cls(containers=containers)
except NotFound:
return cls(containers=[])

@property
Expand Down
38 changes: 26 additions & 12 deletions src/codegen/cli/commands/start/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from rich.panel import Panel

from codegen.cli.commands.start.docker_container import DockerContainer
from codegen.cli.commands.start.docker_fleet import CODEGEN_RUNNER_IMAGE, DockerFleet
from codegen.cli.commands.start.docker_fleet import CODEGEN_RUNNER_IMAGE
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
Expand All @@ -20,13 +20,19 @@
@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=None, help="Port to run the server on")
def start_command(port: int | None, platform: str):
@click.option("--detached", "-d", is_flag=True, help="Run the server in detached mode")
@click.option("--skip-build", is_flag=True, help="Skip building the Docker image")
@click.option("--force", "-f", is_flag=True, help="Force start the server even if it is already running")
def start_command(port: int | None, platform: str, detached: bool = False, skip_build: bool = False, force: bool = False) -> None:
"""Starts a local codegen server"""
repo_path = Path.cwd().resolve()
repo_config = RepoConfig.from_repo_path(str(repo_path))
fleet = DockerFleet.load()
if (container := fleet.get(repo_config.name)) is not None:
return _handle_existing_container(repo_config, container)
if (container := DockerContainer.get(repo_config.name)) is not None:
if force:
rich.print(f"[yellow]Removing existing runner {repo_config.name} to force restart[/yellow]")
container.remove()
else:
return _handle_existing_container(repo_config, container, force)

codegen_version = version("codegen")
rich.print(f"[bold green]Codegen version:[/bold green] {codegen_version}")
Expand All @@ -35,10 +41,11 @@ def start_command(port: int | None, platform: str):
port = get_free_port()

try:
rich.print("[bold blue]Building Docker image...[/bold blue]")
_build_docker_image(codegen_root, platform)
if not skip_build:
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(repo_config, port)
_run_docker_container(repo_config, port, detached)
rich.print(Panel(f"[green]Server started successfully![/green]\nAccess the server at: [bold]http://{_default_host}:{port}[/bold]", box=ROUNDED, title="Codegen Server"))
# TODO: memory snapshot here
except subprocess.CalledProcessError as e:
Expand All @@ -49,7 +56,7 @@ def start_command(port: int | None, platform: str):
raise click.Abort()


def _handle_existing_container(repo_config: RepoConfig, container: DockerContainer) -> None:
def _handle_existing_container(repo_config: RepoConfig, container: DockerContainer, force: bool) -> None:
if container.is_running():
rich.print(
Panel(
Expand Down Expand Up @@ -86,19 +93,26 @@ def _build_docker_image(codegen_root: Path, platform: str) -> None:
subprocess.run(build_cmd, check=True)


def _run_docker_container(repo_config: RepoConfig, port: int) -> None:
def _run_docker_container(repo_config: RepoConfig, port: int, detached: bool) -> None:
container_repo_path = f"/app/git/{repo_config.name}"
name_args = ["--name", f"{repo_config.name}"]
envvars = {
"REPOSITORY_LANGUAGE": repo_config.language.value,
"REPOSITORY_OWNER": LocalGitRepo(repo_config.repo_path).owner,
"REPOSITORY_PATH": container_repo_path,
"GITHUB_TOKEN": SecretsConfig().github_token,
"PYTHONUNBUFFERED": "1", # Ensure Python output is unbuffered
}
envvars_args = [arg for k, v in envvars.items() for arg in ("--env", f"{k}={v}")]
mount_args = ["-v", f"{repo_config.repo_path}:{container_repo_path}"]
entry_point = f"uv run --frozen uvicorn codegen.runner.sandbox.server:app --host {_default_host} --port {port}"
run_cmd = ["docker", "run", "-d", "-p", f"{port}:{port}", *name_args, *mount_args, *envvars_args, CODEGEN_RUNNER_IMAGE, entry_point]
entry_point = f"uv run --frozen uvicorn codegen.runner.servers.local_daemon:app --host {_default_host} --port {port}"
port_args = ["-p", f"{port}:{port}"]
detached_args = ["-d"] if detached else []
run_cmd = ["docker", "run", *detached_args, *port_args, *name_args, *mount_args, *envvars_args, CODEGEN_RUNNER_IMAGE, entry_point]

rich.print(f"run_cmd: {str.join(' ', run_cmd)}")
subprocess.run(run_cmd, check=True)

if detached:
rich.print("[yellow]Container started in detached mode. To view logs, run:[/yellow]")
rich.print(f"[bold]docker logs -f {repo_config.name}[/bold]")
2 changes: 2 additions & 0 deletions src/codegen/git/repo_operator/repo_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -471,6 +471,8 @@ def commit_changes(self, message: str, verify: bool = False) -> bool:
staged_changes = self.git_cli.git.diff("--staged")
if staged_changes:
commit_args = ["-m", message]
if self.bot_commit:
commit_args.append(f"--author='{CODEGEN_BOT_NAME} <{CODEGEN_BOT_EMAIL}>'")
if not verify:
commit_args.append("--no-verify")
self.git_cli.git.commit(*commit_args)
Expand Down
11 changes: 11 additions & 0 deletions src/codegen/git/schemas/repo_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

from pydantic import BaseModel

from codegen.configs.models.repository import RepositoryConfig
from codegen.git.schemas.enums import RepoVisibility
from codegen.shared.enums.programming_language import ProgrammingLanguage

Expand All @@ -24,6 +25,16 @@ class RepoConfig(BaseModel):
base_path: str | None = None # root directory of the codebase within the repo
subdirectories: list[str] | None = None

@classmethod
def from_envs(cls) -> "RepoConfig":
default_repo_config = RepositoryConfig()
return RepoConfig(
name=default_repo_config.name,
full_name=default_repo_config.full_name,
base_dir=os.path.dirname(default_repo_config.path),
language=ProgrammingLanguage(default_repo_config.language.upper()),
)

@classmethod
def from_repo_path(cls, repo_path: str) -> "RepoConfig":
name = os.path.basename(repo_path)
Expand Down
Loading
Loading