temp-save,julia working

Author: wukaixingxp

apps/openenv/FIX_DOCUMENTATION.md

@@ -0,0 +1,74 @@
+# Fix for ModuleNotFoundError: No module named 'julia_utils'
+
+## Problem
+
+The application was crashing with the following error when using Monarch actors:
+
+```
+ModuleNotFoundError: No module named 'julia_utils'
+```
+
+This error occurred when remote Monarch actors tried to unpickle function references that were loaded from the `julia_utils` module.
+
+## Root Cause
+
+The issue happened because:
+
+1. The main process loads functions from `julia_utils` using `load_function_from_string()`
+2. These functions are passed as parameters to actor classes (`GenericDatasetActor`, `GenericRewardActor`)
+3. When actors are spawned as remote actors, the function objects are pickled and sent to remote processes
+4. During unpickling, Python needs to import the `julia_utils` module
+5. **The openenv directory wasn't in `sys.path` yet** because:
+   - The unpickling happens during actor initialization (when deserializing constructor parameters)
+   - The `setup()` endpoint runs AFTER actor initialization
+   - Therefore, `sys.path` wasn't modified before unpickling occurred
+
+## Solution
+
+Added module-level code to `/home/kaiwu/work/kaiwu/forge/apps/openenv/main.py` that adds the openenv directory to `sys.path` BEFORE any actor definitions:
+
+```python
+# CRITICAL: Add openenv directory to sys.path at module level
+# This ensures that when remote actors unpickle function references (e.g., julia_utils functions),
+# the module can be imported successfully. This must happen BEFORE any actor definitions.
+_openenv_dir = Path(__file__).parent
+if str(_openenv_dir) not in sys.path:
+    sys.path.insert(0, str(_openenv_dir))
+```
+
+This code runs when the module is first imported, ensuring that:
+- Remote actors that import `main.py` will have the openenv directory in their `sys.path`
+- Functions from `julia_utils` can be successfully unpickled in remote processes
+- The fix happens early enough to prevent the ModuleNotFoundError
+
+## Testing
+
+Created comprehensive tests to verify the fix:
+
+1. **test_module_import.py** - Tests basic import and pickling functionality
+2. **test_monarch_actor_simulation.py** - Simulates the exact Monarch actor scenario where a remote process receives pickled functions
+
+Both test suites pass successfully, confirming that:
+- `julia_utils` can be imported after importing `main.py`
+- Functions from `julia_utils` can be pickled and unpickled across process boundaries
+- Remote actors can successfully deserialize function references
+
+## Files Modified
+
+- `/home/kaiwu/work/kaiwu/forge/apps/openenv/main.py` - Added module-level sys.path setup
+
+## Files Added
+
+- `/home/kaiwu/work/kaiwu/forge/apps/openenv/test_module_import.py` - Basic import/pickle tests
+- `/home/kaiwu/work/kaiwu/forge/apps/openenv/test_monarch_actor_simulation.py` - Comprehensive simulation tests
+
+## Verification
+
+Run tests to verify the fix:
+```bash
+cd /home/kaiwu/work/kaiwu/forge/apps/openenv
+python test_module_import.py
+python test_monarch_actor_simulation.py
+```
+
+Both should show "✓ All tests passed!"

apps/openenv/julia_utils.py

@@ -12,7 +12,6 @@
 import re
 from typing import Dict, Any
 
-from envs.julia_env import JuliaAction
 from forge.observability.metrics import record_metric, Reduce
 
 
@@ -75,7 +74,7 @@ def build_julia_prompt(sample: Dict[str, Any], tokenizer) -> str:
     return formatted_request
 
 
-def build_julia_action(response: str, sample: Dict[str, Any]) -> JuliaAction:
+def build_julia_action(response: str, sample: Dict[str, Any]):
     """
     Build JuliaAction from model response and dataset sample.
 
@@ -86,6 +85,12 @@ def build_julia_action(response: str, sample: Dict[str, Any]) -> JuliaAction:
     Returns:
         JuliaAction instance with core code and test code
     """
+    # Import AutoAction dynamically to avoid pickle issues
+    from envs import AutoAction
+
+    # Get JuliaAction class dynamically
+    JuliaAction = AutoAction.from_env("julia")
+
     # Extract code from markdown if present
     code = extract_julia_code(response)
 

apps/openenv/llama3_8b_coding.yaml

@@ -16,9 +16,9 @@ rollout_threads: 1
 # Task-specific configuration
 task:
   env_name: "coding"  # Used to load CodingEnv and CodingAction via AutoEnv
-  build_action: !function python_utils.build_python_action
-  evaluate_response: !function python_utils.evaluate_python_response
-  transform_sample: !function python_utils.transform_python_sample
+  build_action: !function apps.openenv.python_utils.build_python_action
+  evaluate_response: !function apps.openenv.python_utils.evaluate_python_response
+  transform_sample: !function apps.openenv.python_utils.transform_python_sample
 
 # Observability configuration
 metric_logging:
@@ -31,13 +31,19 @@ metric_logging:
     log_per_rank: True
 
 # Dataset configuration
+#dataset:
+  # path: "TIGER-Lab/AceCode-87K"
+  # revision: "main"
+  # data_split: "train"
+  # streaming: true
+  # model: ${model}
 dataset:
-  path: "openai/humaneval"  # HumanEval dataset from HuggingFace
+  path: "/home/kaiwu/work/kaiwu/AceCoder/train/train_rl/OpenRLHF/scripts/data/acecode_89k/acecode_hard02.json"
+  #path: "/home/kaiwu/work/kaiwu/AceCoder/train/train_rl/OpenRLHF/data/acecode_87K/acecode_87K.json"
   revision: "main"
-  data_split: "test"
+  data_split: "train"
   streaming: false
   model: ${model}
-
 # OpenEnv configuration for GenericOpenEnvActor
 openenv_config:
   docker_image: "coding-env:latest"
@@ -80,10 +86,10 @@ trainer:
   lr_scheduler:
     warmup_steps: 0
   training:
-    local_batch_size: ${batch_size}
+    local_batch_size: ${multiply:${batch_size},${group_size}}
     seq_len: ${sum:${max_req_tokens},${max_res_tokens}}
     max_norm: 1.0
-    steps: 2000
+    steps: 100
     dtype: bfloat16
     gc_freq: 1
   compile:

apps/openenv/llama3_8b_julia.yaml

@@ -16,9 +16,9 @@ rollout_threads: 1
 # Task-specific configuration
 task:
   env_name: "julia"  # Used to load JuliaEnv and JuliaAction via AutoEnv
-  build_action: !function julia_utils.build_julia_action
-  evaluate_response: !function julia_utils.evaluate_julia_response
-  transform_sample: !function julia_utils.transform_julia_sample
+  build_action: !function apps.openenv.julia_utils.build_julia_action
+  evaluate_response: !function apps.openenv.julia_utils.evaluate_julia_response
+  transform_sample: !function apps.openenv.julia_utils.transform_julia_sample
 
 # Observability configuration
 metric_logging:
@@ -81,10 +81,10 @@ trainer:
   lr_scheduler:
     warmup_steps: 0
   training:
-    local_batch_size: ${batch_size}
+    local_batch_size: ${multiply:${batch_size},${group_size}}
     seq_len: ${sum:${max_req_tokens},${max_res_tokens}}
     max_norm: 1.0
-    steps: 3000
+    steps: 100
     dtype: bfloat16
     gc_freq: 1
   compile: