client wip

Author: caroljung-cg

pyproject.toml

@@ -72,6 +72,7 @@ dependencies = [
   "langchain-anthropic>=0.3.7",
   "lox>=0.12.0",
   "httpx>=0.28.1",
+  "docker>=6.1.3",
 ]
 
 license = { text = "Apache-2.0" }

src/codegen/cli/commands/start/docker_session.py

@@ -0,0 +1,71 @@
+import docker
+
+CODEGEN_RUNNER_IMAGE = "codegen-runner"
+
+
+class DockerSession:
+    _client: docker.DockerClient
+    host: str | None
+    port: int | None
+    name: str
+
+    def __init__(self, client: docker.DockerClient, name: str, port: int | None = None, host: str | None = None):
+        self._client = client
+        self.host = host
+        self.port = port
+        self.name = name
+
+    def is_running(self) -> bool:
+        try:
+            container = self._client.containers.get(self.name)
+            return container.status == "running"
+        except docker.errors.NotFound:
+            return False
+
+    def start(self) -> bool:
+        try:
+            container = self._client.containers.get(self.name)
+            container.start()
+            return True
+        except (docker.errors.NotFound, docker.errors.APIError):
+            return False
+
+    def __str__(self) -> str:
+        return f"DockerSession(name={self.name}, host={self.host or 'unknown'}, port={self.port or 'unknown'})"
+
+
+class DockerSessions:
+    sessions: list[DockerSession]
+
+    def __init__(self, sessions: list[DockerSession]):
+        self.sessions = sessions
+
+    @classmethod
+    def load(cls) -> "DockerSessions":
+        try:
+            client = docker.from_env()
+            containers = client.containers.list(all=True, filters={"ancestor": CODEGEN_RUNNER_IMAGE})
+            sessions = []
+            for container in containers:
+                if container.attrs["Config"]["Image"] == CODEGEN_RUNNER_IMAGE:
+                    if container.status == "running":
+                        host_config = next(iter(container.ports.values()))[0]
+                        session = DockerSession(client=client, host=host_config["HostIp"], port=host_config["HostPort"], name=container.name)
+                    else:
+                        session = DockerSession(client=client, name=container.name)
+                    sessions.append(session)
+
+            return cls(sessions=sessions)
+        except docker.errors.NotFound:
+            return cls(sessions=[])
+
+    def get(self, name: str) -> DockerSession | None:
+        return next((session for session in self.sessions if session.name == name), None)
+
+    def __str__(self) -> str:
+        return f"DockerSessions(sessions={',\n'.join(str(session) for session in self.sessions)})"
+
+
+if __name__ == "__main__":
+    docker_sessions = DockerSessions.load()
+    print(docker_sessions)

src/codegen/cli/commands/start/main.py

@@ -7,18 +7,26 @@
 from rich.box import ROUNDED
 from rich.panel import Panel
 
+from codegen.cli.commands.start.docker_session import CODEGEN_RUNNER_IMAGE, DockerSession, DockerSessions
 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
 
+_default_host = "0.0.0.0"
+
 
 @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")
-@click.option("--detached", "-d", is_flag=True, default=False, help="Starts up the server as detached background process")
-def start_command(port: int | None, platform: str, detached: bool):
+def start_command(port: int | None, platform: str):
     """Starts a local codegen server"""
+    repo_path = Path.cwd().resolve()
+    repo_config = RepoConfig.from_repo_path(str(repo_path))
+    docker_sessions = DockerSessions.load()
+    if (existing_session := docker_sessions.get(repo_config.name)) is not None:
+        return _handle_existing_session(repo_config, existing_session)
+
     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
@@ -29,8 +37,9 @@ def start_command(port: int | None, platform: str, detached: bool):
         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, 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"))
+        _run_docker_container(repo_config, port)
+        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:
         rich.print(f"[bold red]Error:[/bold red] Failed to {e.cmd[0]} Docker container")
         raise click.Abort()
@@ -39,7 +48,26 @@ def start_command(port: int | None, platform: str, detached: bool):
         raise click.Abort()
 
 
-def _build_docker_image(codegen_root: Path, platform: str):
+def _handle_existing_session(repo_config: RepoConfig, docker_session: DockerSession) -> None:
+    if docker_session.is_running():
+        rich.print(
+            Panel(
+                f"[green]Codegen server for {repo_config.name} is already running at: [bold]http://{docker_session.host}:{docker_session.port}[/bold][/green]",
+                box=ROUNDED,
+                title="Codegen Server",
+            )
+        )
+        return
+
+    if docker_session.start():
+        rich.print(Panel(f"[yellow]Docker session for {repo_config.name} is not running. Restarting...[/yellow]", box=ROUNDED, title="Docker Session"))
+        return
+
+    rich.print(Panel(f"[red]Failed to restart container for {repo_config.name}[/red]", box=ROUNDED, title="Docker Session"))
+    click.Abort()
+
+
+def _build_docker_image(codegen_root: Path, platform: str) -> None:
     build_cmd = [
         "docker",
         "buildx",
@@ -57,21 +85,19 @@ def _build_docker_image(codegen_root: Path, platform: str):
     subprocess.run(build_cmd, check=True)
 
 
-def _run_docker_container(port: int, detached: bool):
-    repo_path = Path.cwd().resolve()
-    repo_config = RepoConfig.from_repo_path(repo_path)
+def _run_docker_container(repo_config: RepoConfig, port: int) -> 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_path).owner,
+        "REPOSITORY_OWNER": LocalGitRepo(repo_config.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"{port}:{port}", *mount_args, *envvars_args, "codegen-runner", entry_point]
+    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]
 
     rich.print(f"run_cmd: {str.join(' ', run_cmd)}")
     subprocess.run(run_cmd, check=True)

src/codegen/runner/clients/client.py

@@ -0,0 +1,47 @@
+"""Client used to abstract the weird stdin/stdout communication we have with the sandbox"""
+
+import logging
+
+import requests
+from fastapi import params
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SERVER_PORT = 4002
+
+EPHEMERAL_SERVER_PATH = "codegen.runner.sandbox.ephemeral_server:app"
+
+
+class Client:
+    """Client for interacting with the sandbox server."""
+
+    host: str
+    port: int
+    base_url: str
+
+    def __init__(self, host: str, port: int) -> None:
+        self.host = host
+        self.port = port
+        self.base_url = f"http://{host}:{port}"
+
+    def healthcheck(self, raise_on_error: bool = True) -> bool:
+        try:
+            self.get("/")
+            return True
+        except requests.exceptions.ConnectionError:
+            if raise_on_error:
+                raise
+            return False
+
+    def get(self, endpoint: str, data: dict | None = None) -> requests.Response:
+        url = f"{self.base_url}{endpoint}"
+        response = requests.get(url, json=data)
+        response.raise_for_status()
+        return response
+
+    def post(self, endpoint: str, data: dict | None = None, authorization: str | params.Header | None = None) -> requests.Response:
+        url = f"{self.base_url}{endpoint}"
+        headers = {"Authorization": str(authorization)} if authorization else None
+        response = requests.post(url, json=data, headers=headers)
+        response.raise_for_status()
+        return response

src/codegen/runner/clients/codebase_client.py

@@ -1,28 +1,70 @@
 """Client used to abstract the weird stdin/stdout communication we have with the sandbox"""
 
 import logging
+import os
+import subprocess
+import time
 
 from codegen.configs.models.secrets import SecretsConfig
 from codegen.git.schemas.repo_config import RepoConfig
-from codegen.runner.clients.server_client import LocalServerClient
+from codegen.runner.clients.client import Client
 from codegen.runner.models.apis import SANDBOX_SERVER_PORT
 
-logger = logging.getLogger(__name__)
-
+DEFAULT_SERVER_PORT = 4002
+EPHEMERAL_SERVER_PATH = "codegen.runner.sandbox.ephemeral_server:app"
 RUNNER_SERVER_PATH = "codegen.runner.sandbox.server:app"
 
 
-class CodebaseClient(LocalServerClient):
+logger = logging.getLogger(__name__)
+
+
+class CodebaseClient(Client):
     """Client for interacting with the locally hosted sandbox server."""
 
     repo_config: RepoConfig
 
-    def __init__(self, repo_config: RepoConfig, host: str = "127.0.0.1", port: int = SANDBOX_SERVER_PORT):
+    def __init__(self, repo_config: RepoConfig, host: str = "127.0.0.1", port: int = SANDBOX_SERVER_PORT, server_path: str = RUNNER_SERVER_PATH):
+        super().__init__(host=host, port=port)
         self.repo_config = repo_config
-        super().__init__(server_path=RUNNER_SERVER_PATH, host=host, port=port)
+        self._process = None
+        self._start_server(server_path)
+
+    def __del__(self):
+        """Cleanup the subprocess when the client is destroyed"""
+        if self._process is not None:
+            self._process.terminate()
+            self._process.wait()
+
+    def _start_server(self, server_path: str) -> None:
+        """Start the FastAPI server in a subprocess"""
+        envs = self._get_envs()
+        logger.info(f"Starting local server on {self.base_url} with envvars: {envs}")
+
+        self._process = subprocess.Popen(
+            [
+                "uvicorn",
+                server_path,
+                "--host",
+                self.host,
+                "--port",
+                str(self.port),
+            ],
+            env=envs,
+        )
+        self._wait_for_server()
+
+    def _wait_for_server(self, timeout: int = 30, interval: float = 0.3) -> None:
+        """Wait for the server to start by polling the health endpoint"""
+        start_time = time.time()
+        while (time.time() - start_time) < timeout:
+            if self.healthcheck(raise_on_error=False):
+                return
+            time.sleep(interval)
+        msg = "Server failed to start within timeout period"
+        raise TimeoutError(msg)
 
     def _get_envs(self) -> dict:
-        envs = super()._get_envs()
+        envs = os.environ.copy()
         codebase_envs = {
             "REPOSITORY_LANGUAGE": self.repo_config.language.value,
             "REPOSITORY_OWNER": self.repo_config.organization_name,

src/codegen/runner/clients/docker_client.py

@@ -1,10 +1,11 @@
 """Client for interacting with the locally hosted sandbox server hosted on a docker container."""
 
-from codegen.runner.clients.server_client import LocalServerClient
+from codegen.cli.commands.start.docker_session import DockerSession
+from codegen.runner.clients.client import Client
 
 
-class DockerClient(LocalServerClient):
+class DockerClient(Client):
     """Client for interacting with the locally hosted sandbox server hosted on a docker container."""
 
-    def __init__(self, repo_config: RepoConfig):
-        super().__init__(repo_config, host, port)
+    def __init__(self, docker_session: DockerSession):
+        super().__init__(docker_session.host, docker_session.port)