adjust based on comments

Author: garyzhang99

trinity/common/workflows/__init__.py

@@ -46,9 +46,7 @@
         "async_simple_mm_workflow": "trinity.common.workflows.simple_mm_workflow.AsyncSimpleMMWorkflow",
         # on-policy distillation workflows
         "on_policy_distill_workflow": "trinity.common.workflows.on_policy_distill_workflow.OnPolicyDistillWorkflow",
-        "async_on_policy_distill_workflow": "trinity.common.workflows.on_policy_distill_workflow.AsyncOnPolicyDistillWorkflow",
         "on_policy_distill_math_workflow": "trinity.common.workflows.on_policy_distill_workflow.OnPolicyDistillMathWorkflow",
-        "async_on_policy_distill_math_workflow": "trinity.common.workflows.on_policy_distill_workflow.AsyncOnPolicyDistillMathWorkflow",
     },
 )
 

trinity/common/workflows/on_policy_distill_workflow.py

@@ -144,12 +144,6 @@ async def run_async(self) -> List[Experience]:
         return responses
 
 
-class AsyncOnPolicyDistillWorkflow(OnPolicyDistillWorkflow):
-    """Alias for OnPolicyDistillWorkflow (already async)."""
-
-    pass
-
-
 class OnPolicyDistillMathWorkflow(OnPolicyDistillWorkflow):
     """On-policy distillation workflow with Qwen2.5-Math style format.
 
@@ -187,9 +181,3 @@ def compute_reward(self, response: Experience) -> float:
             response.metrics["accuracy"] = accuracy
             return float(accuracy)
         return 0.0
-
-
-class AsyncOnPolicyDistillMathWorkflow(OnPolicyDistillMathWorkflow):
-    """Alias for OnPolicyDistillMathWorkflow (already async)."""
-
-    pass

trinity/common/workflows/workflow.py

@@ -4,16 +4,16 @@
 from __future__ import annotations
 
 from dataclasses import asdict, dataclass, field
-from typing import List, Optional, Type, Union
-
-import openai
+from typing import TYPE_CHECKING, Any, List, Optional, Type, Union
 
 from trinity.common.config import FormatConfig, GenerationConfig
 from trinity.common.experience import Experience
-from trinity.common.models.model import ModelWrapper
 from trinity.common.rewards.reward_fn import RewardFn
 from trinity.utils.log import get_logger
 
+if TYPE_CHECKING:
+    import openai
+    from trinity.common.models.model import ModelWrapper
 
 @dataclass
 class Task(dict):