"""

Main entrypoint for fully async training.

"""

import sys

from skyrl.train.config import SkyRLTrainConfig

from skyrl.train.entrypoints.main_base import BasePPOExp, validate_cfg

from skyrl.train.fully_async_trainer import FullyAsyncRayPPOTrainer

import asyncio

from skyrl.train.utils import initialize_ray

import ray

class FullyAsyncPPOExp(BasePPOExp):

def get_trainer(

self,

cfg,

tracker,

tokenizer,

train_dataset,

eval_dataset,

inference_engine_client,

generator,

colocate_pg,

):

return FullyAsyncRayPPOTrainer(

cfg=cfg,

tracker=tracker,

tokenizer=tokenizer,

train_dataset=train_dataset,

eval_dataset=eval_dataset,

inference_engine_client=inference_engine_client,

generator=generator,

colocate_pg=colocate_pg,

)

def run(self):

trainer = self._setup_trainer()

# Start the async training loop

asyncio.run(trainer.train())

@ray.remote(num_cpus=1)

def skyrl_entrypoint(cfg: SkyRLTrainConfig):

# make sure that the training loop is not run on the head node.

exp = FullyAsyncPPOExp(cfg)

exp.run()

def main() -> None:

cfg = SkyRLTrainConfig.from_cli_overrides(sys.argv[1:])

# validate the arguments

validate_cfg(cfg)

initialize_ray(cfg)

ray.get(skyrl_entrypoint.remote(cfg))

if __name__ == "__main__":

main()