remove policy config

Author: DNXie

apps/vllm/config.yaml

@@ -1,4 +1,4 @@
-policy_config:
+policy:
   worker_params:
     model: "meta-llama/Llama-3.1-8B-Instruct"
     tensor_parallel_size: 2

apps/vllm/main.py

@@ -11,38 +11,31 @@
 
 import asyncio
 import sys
-from typing import Any
 
-import yaml
-
-from forge.actors.policy import Policy, PolicyConfig
+from forge.actors.policy import Policy
+from forge.cli.config import parse
 from forge.controller.service import ServiceConfig, shutdown_service, spawn_service
+from omegaconf import DictConfig
+from vllm.outputs import RequestOutput
 
 
-def load_yaml_config(path: str) -> dict:
-    with open(path, "r") as f:
-        return yaml.safe_load(f)
-
+async def run(cfg: DictConfig):
 
-def get_configs(cfg: dict) -> tuple[PolicyConfig, ServiceConfig, str]:
-    # Instantiate PolicyConfig and ServiceConfig from nested dicts
-    policy_config = PolicyConfig.from_dict(cfg["policy_config"])
-    service_config = ServiceConfig(**cfg["service_config"])
     if "prompt" in cfg and cfg["prompt"] is not None:
         prompt = cfg["prompt"]
     else:
-        gd = policy_config.sampling_params.guided_decoding
+        gd = cfg.policy.get("sampling_params", {}).get("guided_decoding", False)
         prompt = "What is 3+5?" if gd else "Tell me a joke"
-    return policy_config, service_config, prompt
 
-
-async def run_vllm(service_config: ServiceConfig, config: PolicyConfig, prompt: str):
     print("Spawning service...")
-    policy = await spawn_service(service_config, Policy, config=config)
+
+    policy = await spawn_service(
+        ServiceConfig(**cfg.service_config), Policy, **cfg.policy
+    )
 
     async with policy.session():
         print("Requesting generation...")
-        response_output = await policy.generate.choose(prompt=prompt)
+        response_output: RequestOutput = await policy.generate.choose(prompt=prompt)
 
         print("\nGeneration Results:")
         print("=" * 80)
@@ -57,19 +50,10 @@ async def run_vllm(service_config: ServiceConfig, config: PolicyConfig, prompt:
     await shutdown_service(policy)
 
 
-def main():
-    import argparse
-
-    parser = argparse.ArgumentParser(description="vLLM Policy Inference Application")
-    parser.add_argument(
-        "--config", type=str, required=True, help="Path to YAML config file"
-    )
-    args = parser.parse_args()
-
-    cfg = load_yaml_config(args.config)
-    policy_config, service_config, prompt = get_configs(cfg)
-    asyncio.run(run_vllm(service_config, policy_config, prompt))
+@parse
+def recipe_main(cfg: DictConfig) -> None:
+    asyncio.run(run(cfg))
 
 
 if __name__ == "__main__":
-    sys.exit(main())
+    sys.exit(recipe_main())

src/forge/actors/policy.py

@@ -20,6 +20,7 @@
 from forge.interfaces import Policy as PolicyInterface
 from forge.types import ProcessConfig
 from monarch.actor import current_rank, endpoint, ProcMesh
+from omegaconf import DictConfig
 from torchstore import MultiProcessStore
 from torchstore._state_dict_utils import DELIM
 
@@ -90,28 +91,16 @@ def from_dict(cls, d: dict):
         d = dict(d)  # copy
         if "vllm_args" in d and isinstance(d["vllm_args"], dict):
             d["vllm_args"] = EngineArgs(**d["vllm_args"])
-        return cls(**d)
-
-
-@dataclass
-class PolicyConfig:
-    worker_params: WorkerConfig = field(default_factory=WorkerConfig)
-    sampling_params: SamplingOverrides = field(default_factory=SamplingOverrides)
-    available_devices: str = None
-
-    @classmethod
-    def from_dict(cls, d: dict):
-        d = dict(d)
-        if "worker_params" in d and isinstance(d["worker_params"], dict):
-            d["worker_params"] = WorkerConfig.from_dict(d["worker_params"])
-        if "sampling_params" in d and isinstance(d["sampling_params"], dict):
-            d["sampling_params"] = SamplingOverrides(**d["sampling_params"])
+        else:
+            d["vllm_args"] = None
         return cls(**d)
 
 
 @dataclass
 class Policy(PolicyInterface):
-    config: PolicyConfig
+    worker_params: WorkerConfig = field(default_factory=WorkerConfig)
+    sampling_overrides: SamplingOverrides = field(default_factory=SamplingOverrides)
+    available_devices: str | None = None
     # Gets set up by setup
     sampling_params: SamplingParams | None = None
     lora_request: LoRARequest | None = None
@@ -124,15 +113,19 @@ def __post_init__(self):
         self._policy_proc: ProcMesh | None = None
         self._worker_procs: ProcMesh | None = None
         self.weights_version: int = 0
-        if isinstance(self.config, dict):
-            self.config = PolicyConfig.from_dict(self.config)
+        if isinstance(self.worker_params, dict):
+            self.worker_params = WorkerConfig.from_dict(self.worker_params)
+        if isinstance(self.sampling_overrides, dict):
+            self.sampling_overrides = SamplingOverrides(**self.sampling_overrides)
 
     @classmethod
     async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
         cls: type["Policy"],
         *,
         process_config: ProcessConfig,
-        config: PolicyConfig,
+        worker_params: WorkerConfig | dict = WorkerConfig(),
+        sampling_overrides: SamplingOverrides | dict = SamplingOverrides(),
+        available_devices: str | None = None,
         store: MultiProcessStore | None = None,
         **kwargs,
     ) -> "Policy":
@@ -146,16 +139,25 @@ async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
         policy_proc_config.with_gpus = False
 
         policy_proc = await get_proc_mesh(process_config=policy_proc_config)
+
+        if isinstance(worker_params, (dict, DictConfig)):
+            worker_params = WorkerConfig.from_dict(worker_params)
+
+        if isinstance(worker_params, (dict, DictConfig)):
+            sampling_overrides = SamplingOverrides(**sampling_overrides)
+
         workers = await worker_procs.spawn(
-            "vllm_worker", PolicyWorker, **asdict(config.worker_params)
+            "vllm_worker", PolicyWorker, **asdict(worker_params)
         )
 
         # TODO - expand support so name can stick within kwargs
         actor_name = kwargs.pop("name", cls.__name__)
         policy = await policy_proc.spawn(
             actor_name,
             cls,
-            config=config,
+            worker_params=worker_params,
+            sampling_overrides=sampling_overrides,
+            available_devices=available_devices,
             policy_worker=workers,
             store=store,
         )
@@ -192,7 +194,7 @@ async def setup(self):
         self.vllm_args = await self.policy_worker.get_vllm_args.choose()
 
         # Setup sampling params
-        sampling_overrides = self.config.sampling_params
+        sampling_overrides = self.sampling_overrides
         overrides = {
             "n": sampling_overrides.num_samples,
             "guided_decoding": (

tests/unit_tests/test_policy_config.py

@@ -0,0 +1,131 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import inspect
+import tempfile
+import unittest
+from dataclasses import asdict
+
+import yaml
+
+from forge.actors.policy import Policy, SamplingOverrides, WorkerConfig
+from vllm.engine.arg_utils import EngineArgs
+
+
+class TestPolicyConfig(unittest.TestCase):
+    """Test suite for Policy configuration handling after PolicyConfig removal."""
+
+    def test_policy_default_initialization(self):
+        """Test that Policy can be initialized with default values."""
+        policy = Policy()
+
+        # Check that default factories work
+        self.assertIsInstance(policy.worker_params, WorkerConfig)
+        self.assertIsInstance(policy.sampling_overrides, SamplingOverrides)
+        self.assertIsNone(policy.available_devices)
+
+        # Check default values
+        self.assertEqual(policy.worker_params.model, "meta-llama/Llama-3.1-8B-Instruct")
+        self.assertEqual(policy.worker_params.tensor_parallel_size, 1)
+        self.assertEqual(policy.worker_params.pipeline_parallel_size, 1)
+        self.assertFalse(policy.worker_params.enforce_eager)
+
+        self.assertEqual(policy.sampling_overrides.num_samples, 1)
+        self.assertFalse(policy.sampling_overrides.guided_decoding)
+        self.assertEqual(policy.sampling_overrides.max_tokens, 512)
+
+    def test_policy_with_dict_configs(self):
+        """Test Policy initialization with dictionary configs."""
+        worker_dict = {
+            "model": "test-model-6789",
+            "tensor_parallel_size": 7777,
+            "pipeline_parallel_size": 8888,
+            "enforce_eager": True,
+            "vllm_args": {"max_model_len": 9999, "gpu_memory_utilization": 0.1234},
+        }
+
+        sampling_dict = {
+            "num_samples": 1357,
+            "guided_decoding": True,
+            "max_tokens": 2468,
+        }
+
+        policy = Policy(
+            worker_params=worker_dict,
+            sampling_overrides=sampling_dict,
+            available_devices="test-gpu-device-abcd",
+        )
+
+        # Check that dictionaries were converted to proper objects
+        self.assertIsInstance(policy.worker_params, WorkerConfig)
+        self.assertIsInstance(policy.sampling_overrides, SamplingOverrides)
+
+        self.assertEqual(policy.worker_params.model, "test-model-6789")
+        self.assertEqual(policy.worker_params.tensor_parallel_size, 7777)
+        self.assertEqual(policy.worker_params.pipeline_parallel_size, 8888)
+        self.assertTrue(policy.worker_params.enforce_eager)
+
+        self.assertEqual(policy.sampling_overrides.num_samples, 1357)
+        self.assertTrue(policy.sampling_overrides.guided_decoding)
+        self.assertEqual(policy.sampling_overrides.max_tokens, 2468)
+
+    def test_policy_yaml_config_loading(self):
+        """Test loading Policy configuration from YAML file."""
+        yaml_content = """
+        worker_params:
+          model: "yaml-test-model-9876"
+          tensor_parallel_size: 1234
+          pipeline_parallel_size: 5678
+          enforce_eager: true
+          vllm_args:
+            max_model_len: 9876
+            gpu_memory_utilization: 0.1357
+
+        sampling_overrides:
+          num_samples: 2468
+          guided_decoding: true
+          max_tokens: 1357
+
+        available_devices: "yaml-test-device-xyz"
+        """
+
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+            f.write(yaml_content)
+            f.flush()
+
+            # Load YAML and create Policy
+            with open(f.name, "r") as yaml_file:
+                config = yaml.safe_load(yaml_file)
+
+            policy = Policy(**config)
+
+            self.assertEqual(policy.worker_params.model, "yaml-test-model-9876")
+            self.assertEqual(policy.worker_params.tensor_parallel_size, 1234)
+            self.assertEqual(policy.worker_params.pipeline_parallel_size, 5678)
+            self.assertTrue(policy.worker_params.enforce_eager)
+
+            self.assertEqual(policy.sampling_overrides.num_samples, 2468)
+            self.assertTrue(policy.sampling_overrides.guided_decoding)
+            self.assertEqual(policy.sampling_overrides.max_tokens, 1357)
+
+            self.assertEqual(policy.available_devices, "yaml-test-device-xyz")
+
+    def test_invalid_worker_config_from_dict(self):
+        """Test that WorkerConfig.from_dict handles invalid vllm_args gracefully."""
+        config_dict = {
+            "model": "meta-llama/Llama-3.1-8B-Instruct",
+            "vllm_args": "invalid_string_instead_of_dict",  # This will be passed through as-is
+        }
+
+        worker_config = WorkerConfig.from_dict(config_dict)
+
+        # The invalid vllm_args gets removed and default EngineArgs is used
+        self.assertEqual(worker_config.model, "meta-llama/Llama-3.1-8B-Instruct")
+        self.assertEqual(worker_config.vllm_args, None)
+
+
+if __name__ == "__main__":
+    unittest.main()