Skip to content

Commit d31bb52

Browse files
committed
refactor: update CheckpointEngineManager initialization to use dataclass configuration
Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>
1 parent 5af5eed commit d31bb52

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

tests/checkpoint_engine/test_naive_correctness.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,10 @@
2727
RayWorkerGroup,
2828
)
2929
from verl.single_controller.ray.base import create_colocated_worker_cls
30+
from verl.utils.config import omega_conf_to_dataclass
3031
from verl.utils.device import get_device_name
3132
from verl.utils.tokenizer import hf_tokenizer
33+
from verl.workers.config import CheckpointEngineConfig
3234
from verl.workers.engine_workers import ActorRolloutRefWorker
3335

3436

@@ -97,8 +99,11 @@ def test_server_adapter_colocated_weight_update(init_config):
9799
)
98100

99101
# 2. create CheckpointEngineManager
102+
checkpoint_engine_config: CheckpointEngineConfig = omega_conf_to_dataclass(
103+
init_config.actor_rollout_ref.rollout.checkpoint_engine
104+
)
100105
checkpoint_manager = CheckpointEngineManager(
101-
backend=init_config.actor_rollout_ref.rollout.checkpoint_engine.backend,
106+
config=checkpoint_engine_config,
102107
trainer=actor_rollout_wg,
103108
replicas=agent_loop_manager.rollout_replicas,
104109
)

0 commit comments

Comments
 (0)