add: codegen's modal run to postgresql/grafana

Author: clee-codegen

codegen-on-oss/Dockerfile

@@ -18,10 +18,15 @@ RUN apt-get update && apt-get install -y git \
 
 FROM python:3.12-slim
 
-ENV PATH="/venv/bin:$PATH"
+ENV PATH="/venv/bin:/app/scripts:$PATH"
 # Copy the project into the image
 COPY --from=installer /app/.venv/ /venv
 
+RUN apt-get update && apt-get install -y postgresql-client \
+    && rm -rf /var/lib/apt/lists/* \
+    && apt-get clean
+
+
 WORKDIR /app
 
 COPY . .

codegen-on-oss/codegen_modal_deploy.py

@@ -0,0 +1,74 @@
+import sys
+from pathlib import Path
+
+import modal
+from loguru import logger
+
+from codegen_on_oss.cache import cachedir
+from codegen_on_oss.metrics import MetricsProfiler
+from codegen_on_oss.outputs.sql_output import PostgresSQLOutput
+from codegen_on_oss.parser import CodegenParser
+
+app = modal.App("codegen-oss-parse")
+
+
+codegen_repo_volume = modal.Volume.from_name(
+    "codegen-oss-repo-volume",
+    create_if_missing=True,
+)
+
+
+aws_secrets = modal.Secret.from_name(
+    "codegen-oss-parse-secrets",
+)
+
+
+@app.function(
+    name="parse_repo",
+    concurrency_limit=10,
+    cpu=4,
+    memory=16384,
+    timeout=3600 * 8,
+    secrets=[aws_secrets],
+    volumes={
+        str(cachedir.absolute()): codegen_repo_volume,
+    },
+    proxy=modal.Proxy.from_name("codegen-parse-proxy"),
+    image=modal.Image.debian_slim(python_version="3.13")
+    .pip_install("uv")
+    .apt_install("git")  # required by codegen sdk
+    .env({"PATH": "/app/.venv/bin:$PATH"})
+    .workdir("/app")
+    .add_local_file("uv.lock", remote_path="/app/uv.lock", copy=True)
+    .add_local_file("pyproject.toml", remote_path="/app/pyproject.toml", copy=True)
+    .run_commands("uv sync --frozen --no-install-project --extra sql")
+    .add_local_python_source("codegen_on_oss", copy=True),
+    # .add_local_python_source("codegen_on_oss"),
+    # .add_local_dir("codegen_on_oss", remote_path="/app/codegen_on_oss"),
+)
+def parse_repo(
+    repo_url: str,
+    commit_hash: str | None,
+    language: str | None = None,
+):
+    """
+    Parse repositories on Modal.
+
+    Args:
+        repo_url: The URL of the repository to parse.
+        commit_hash: The commit hash of the repository to parse.
+    """
+    logger.add(sys.stdout, format="{time: HH:mm:ss} {level} {message}", level="DEBUG")
+
+    output = PostgresSQLOutput(modal_function_call_id=modal.current_function_call_id())
+    metrics_profiler = MetricsProfiler(output)
+    parser = CodegenParser(Path(cachedir) / "repositories", metrics_profiler)
+    # Refresh any updating repo data from other instances
+    codegen_repo_volume.reload()
+    try:
+        parser.parse(repo_url, language, commit_hash)
+    except Exception as e:
+        logger.exception(f"Error parsing repository {repo_url}: {e}")
+    finally:
+        # Commit any cache changes to the repo volume
+        codegen_repo_volume.commit()

codegen-on-oss/codegen_modal_run.py

@@ -0,0 +1,29 @@
+import modal
+
+from codegen_on_oss.sources import GithubSettings, GithubSource
+
+app = modal.App("codegen-oss-parse")
+
+
+@app.local_entrypoint()
+def main(
+    languages: str = "python,typescript",
+    heuristic: str = "stars",
+    num_repos: int = 100,
+):
+    """
+    Main entrypoint for the parse app.
+    """
+    parse_repo_on_modal_fn = modal.Function.from_name("codegen-oss-parse", "parse_repo")
+    for language in languages.split(","):
+        repo_source = GithubSource(
+            GithubSettings(
+                language=language.strip(), heuristic=heuristic, num_repos=num_repos
+            )
+        )
+        for repo_url, commit_hash in repo_source:
+            parse_repo_on_modal_fn.spawn(
+                repo_url=repo_url,
+                commit_hash=commit_hash,
+                language=language,
+            )

codegen-on-oss/codegen_on_oss/bucket_store.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 from importlib.metadata import version
 from typing import TYPE_CHECKING
 
@@ -6,35 +7,20 @@
 if TYPE_CHECKING:
     from types_boto3_s3 import S3Client
 
-    from codegen_on_oss.sources import RepoSource
-
 
 class BucketStore:
     s3_client: "S3Client"
 
     def __init__(self, bucket_name: str):
         self.bucket_name = bucket_name
         self.s3_client = client("s3")
-
-    def upload_run(
-        self,
-        repo_source: "RepoSource",
-        log_output_path: str,
-        metrics_output_path: str,
-    ):
-        codegen_version = str(version("codegen"))
-        key_prefix: str = f"{codegen_version}/{repo_source.source_type}"
-        config_key = f"{key_prefix}/config.json"
-
-        self.s3_client.put_object(
-            Bucket=self.bucket_name,
-            Key=config_key,
-            Body=repo_source.settings.model_dump_json(indent=4).encode("utf-8"),
-            ContentType="application/json",
+        self.key_prefix: str = str(version("codegen"))
+
+    def upload_file(self, local_path: str, remote_path: str) -> str:
+        key = f"{self.key_prefix}/{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}/{remote_path}"
+        self.s3_client.upload_file(
+            local_path,
+            self.bucket_name,
+            key,
         )
-
-        log_key = f"{key_prefix}/output.logs"
-        self.s3_client.upload_file(log_output_path, self.bucket_name, log_key)
-
-        metrics_key = f"{key_prefix}/metrics.csv"
-        self.s3_client.upload_file(metrics_output_path, self.bucket_name, metrics_key)
+        return key

codegen-on-oss/codegen_on_oss/cli.py

@@ -6,6 +6,7 @@
 
 from codegen_on_oss.cache import cachedir
 from codegen_on_oss.metrics import MetricsProfiler
+from codegen_on_oss.outputs.csv_output import CSVOutput
 from codegen_on_oss.parser import CodegenParser
 from codegen_on_oss.sources import RepoSource, all_sources
 
@@ -60,7 +61,8 @@ def run_one(
     """
     logger.add(error_output_path, level="ERROR")
     logger.add(sys.stdout, level="DEBUG" if debug else "INFO")
-    metrics_profiler = MetricsProfiler(output_path)
+    output = CSVOutput(MetricsProfiler.fields(), output_path)
+    metrics_profiler = MetricsProfiler(output)
 
     parser = CodegenParser(Path(cache_dir) / "repositories", metrics_profiler)
     parser.parse(url, commit_hash)
@@ -115,7 +117,8 @@ def run(
     )
 
     repo_source = RepoSource.from_source_type(source)
-    metrics_profiler = MetricsProfiler(output_path)
+    output = CSVOutput(MetricsProfiler.fields(), output_path)
+    metrics_profiler = MetricsProfiler(output)
     parser = CodegenParser(Path(cache_dir) / "repositories", metrics_profiler)
     for repo_url, commit_hash in repo_source:
         parser.parse(repo_url, commit_hash)