simplify

Author: JenniferWang

src/forge/actors/policy.py

@@ -17,23 +17,6 @@
 import torch
 import torch.distributed.checkpoint as dcp
 import torchstore as ts
-
-from forge.actors._torchstore_utils import (
-    extract_param_name,
-    get_dcp_whole_state_dict_key,
-    get_param_key,
-    get_param_prefix,
-    load_tensor_from_dcp,
-)
-
-from forge.controller import ForgeActor, get_proc_mesh, stop_proc_mesh
-from forge.data.sharding import VLLMSharding
-from forge.data_models.completion import Completion
-from forge.data_models.prompt import to_prompt
-from forge.interfaces import Policy as PolicyInterface
-from forge.observability.metrics import record_metric, Reduce
-from forge.observability.perf_tracker import Tracer
-from forge.types import ProcessConfig
 from monarch.actor import current_rank, endpoint, ProcMesh
 from torchstore.state_dict_utils import DELIM
 from vllm.config import VllmConfig
@@ -58,6 +41,23 @@
 from vllm.v1.structured_output import StructuredOutputManager
 from vllm.worker.worker_base import WorkerWrapperBase
 
+from forge.actors._torchstore_utils import (
+    extract_param_name,
+    get_dcp_whole_state_dict_key,
+    get_param_key,
+    get_param_prefix,
+    load_tensor_from_dcp,
+)
+
+from forge.controller import ForgeActor, get_proc_mesh, stop_proc_mesh
+from forge.data.sharding import VLLMSharding
+from forge.data_models.completion import Completion
+from forge.data_models.prompt import to_prompt
+from forge.interfaces import Policy as PolicyInterface
+from forge.observability.metrics import record_metric, Reduce
+from forge.observability.perf_tracker import Tracer
+from forge.types import ProcessConfig
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 

tests/integration_tests/test_policy_update.py

@@ -23,7 +23,6 @@
 import asyncio
 import logging
 import shutil
-import sys
 from pathlib import Path
 
 import torch
@@ -66,14 +65,9 @@ async def zero_out_model_states(self):
 
 def load_config(config_path: str) -> DictConfig:
     """Load configuration from YAML file."""
-    try:
-        cfg = OmegaConf.load(config_path)
-    except Exception as e:
-        raise RuntimeError(f"Failed to load config file {config_path}: {e}")
-
+    cfg = OmegaConf.load(config_path)
     if not isinstance(cfg, DictConfig):
         raise TypeError(f"Expected DictConfig, got {type(cfg)}")
-
     cfg = resolve_hf_hub_paths(cfg)
     return cfg
 
@@ -277,9 +271,9 @@ async def run_weight_sync_sanity_check(policy: Policy, rl_trainer: MockRLTrainer
     RESET = "\033[0m"
 
     success_message = f"""
-{GREEN}{'='*60}
+{GREEN}{'=' * 60}
 ✅ Weight sharding sanity check passed! ✅
-{'='*60}{RESET}
+{'=' * 60}{RESET}
 """
 
     print(success_message)