diff --git a/environments/code_execution_server/coding_server.py b/environments/code_execution_server/coding_server.py index 5e4774fb..c2088ff4 100644 --- a/environments/code_execution_server/coding_server.py +++ b/environments/code_execution_server/coding_server.py @@ -1,19 +1,14 @@ -import random -from typing import Dict, List, Optional, Tuple, TypedDict, Union - -from datasets import load_dataset -from latex2sympy2_extended import NormalizationConfig -from math_verify import LatexExtractionConfig, parse, verify -from tqdm.asyncio import tqdm_asyncio -import regex as re - import asyncio +import os +import random +from typing import List, Optional, Tuple + +import docker import httpx -import docker, os - -from atroposlib.envs.base import BaseEnv, BaseEnvConfig, OpenaiConfig, ScoredDataGroup -from atroposlib.type_definitions import Item, number +import regex as re +from datasets import load_dataset +from atroposlib.envs.base import BaseEnv, ScoredDataGroup from atroposlib.type_definitions import GameHistory, Item from atroposlib.utils.tokenize_for_trainer import tokenize_for_trainer @@ -25,17 +20,15 @@ system_prompt = ( "tags, and then provide your solution or response to the problem.\n\n" ) + async def submit_code(client, code, test_input, language="python"): url = "http://localhost:5002/execute" - payload = { - "code": code, - "input": test_input, - "language": language - } + payload = {"code": code, "input": test_input, "language": language} response = await client.post(url, json=payload) response_json = response.json() return response_json["output"] + async def get_results(code, answer): async with httpx.AsyncClient() as client: tasks = [] @@ -45,18 +38,22 @@ async def get_results(code, answer): results = await asyncio.gather(*tasks) return [result for result in results] + def init_docker(): client = docker.from_env() + def build_docker_image(): try: # Build the Docker image print("Building Docker image...") - current_dir = os.path.dirname(os.path.abspath(__file__)) # Get the current directory of the script + current_dir = os.path.dirname( + os.path.abspath(__file__) + ) # Get the current directory of the script image, logs = client.images.build(path=current_dir, tag="code-executor") # Print the build logs for log in logs: - print(log.get('stream', '').strip()) + print(log.get("stream", "").strip()) print("Docker image built successfully.") return image @@ -67,19 +64,20 @@ def init_docker(): try: # Run the Docker container print("Running Docker container...") - container = client.containers.run("code-executor", - ports={'5002/tcp': 5002}, - detach=True) # Runs in detached mode (in the background) + container = client.containers.run( + "code-executor", ports={"5002/tcp": 5002}, detach=True + ) # Runs in detached mode (in the background) print(f"Docker container is running with ID: {container.id}") return container except docker.errors.ContainerError as e: print(f"Error during Docker container run: {e}") - + build_docker_image() container = run_docker_container() return container + class CodingEnv(BaseEnv): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -111,7 +109,7 @@ class CodingEnv(BaseEnv): item[1], ) ) - + to_postprocess = await self.score(to_score) return to_postprocess, to_backlog @@ -144,18 +142,23 @@ class CodingEnv(BaseEnv): prompt = tuple( [frozenset({"role": "user", "content": next_item["description"]}.items())] ) - answer = (tuple(next_item["private_tests"]["input"]), tuple(next_item["private_tests"]["output"]), tuple(next_item["generated_tests"]["input"]), tuple(next_item["generated_tests"]["output"])) + answer = ( + tuple(next_item["private_tests"]["input"]), + tuple(next_item["private_tests"]["output"]), + tuple(next_item["generated_tests"]["input"]), + tuple(next_item["generated_tests"]["output"]), + ) return (prompt, answer) def extract_python_code_blocks(self, text): - # Regex specifically looks for ```python\n...code...\n``` - pattern = r'^```(?:\w+)?\s*\n(.*?)(?=^```)```' + # Regex specifically looks for ```python\n...code...\n``` + pattern = r"^```(?:\w+)?\s*\n(.*?)(?=^```)```" result = re.findall(pattern, text, re.DOTALL | re.MULTILINE) python_blocks = [r for r in result] return python_blocks async def score(self, rollout_group_data) -> Optional[ScoredDataGroup]: - #print("Rollout group data", rollout_group_data) + # print("Rollout group data", rollout_group_data) scores = ScoredDataGroup() scores["tokens"] = list() scores["masks"] = list() @@ -192,5 +195,6 @@ class CodingEnv(BaseEnv): # return None # If all the same, we return None return scores + if __name__ == "__main__": CodingEnv.cli() diff --git a/environments/code_execution_server/server.py b/environments/code_execution_server/server.py index e63f6112..b3da9075 100644 --- a/environments/code_execution_server/server.py +++ b/environments/code_execution_server/server.py @@ -14,15 +14,17 @@ curl -X POST http://localhost:5002/execute \ """ + import os import subprocess import uuid -from flask import Flask, request, jsonify -import time + +from flask import Flask, jsonify, request app = Flask(__name__) -@app.route('/execute', methods=['POST']) + +@app.route("/execute", methods=["POST"]) def execute_code(): try: # Receive C++ code from API request @@ -46,16 +48,13 @@ def execute_code(): input=test_cases, capture_output=True, text=True, - timeout=5 # Prevent infinite loops + timeout=5, # Prevent infinite loops ) # Cleanup temporary files os.remove(py_filename) - return jsonify({ - "output": exec_result.stdout, - "error": exec_result.stderr - }) + return jsonify({"output": exec_result.stdout, "error": exec_result.stderr}) """ C++ stuff file_id = str(uuid.uuid4()) @@ -90,6 +89,6 @@ def execute_code(): except Exception as e: return jsonify({"error": str(e)}) -if __name__ == '__main__': - app.run(host="0.0.0.0", port=5002) +if __name__ == "__main__": + app.run(host="0.0.0.0", port=5002)