rebased and updated

Author: rithwik-db

test_single_controller_ppo.py

@@ -49,10 +49,7 @@
 from compose_rl.controllers import BaseDistributedGPUActor, SPMDActorGroup
 from compose_rl.controllers.buffer import Buffer
 from compose_rl.algorithms.online.callback_utils import preprocess_batches
-from databricks.sdk import WorkspaceClient
 
-MLFLOW_RUN_NAME=os.environ['COMPOSER_RUN_NAME'] # SHOULD BE SET BY MCLI
-MLFLOW_EXPERIMENT_NAME=f'/Users/{WorkspaceClient().current_user.me().user_name}/test_single_controller'
 
 @contextmanager
 def time_it(name: str):
@@ -193,9 +190,11 @@ def build_ppo_trainer(self):
         dummy_distributed_sampler = torch.utils.data.distributed.DistributedSampler(dummy_dataset)
         dummy_dataloader = torch.utils.data.DataLoader(dummy_dataset, sampler=dummy_distributed_sampler)
 
+        # TODO: We might be able to skip part of the setup here as some mlflow
+        # environment variables are set in the _setup_mlflow function
         mlflow_logger = MLFlowLogger(
             experiment_name=self.config.loggers.mlflow.experiment_name,
-            run_name=f'test_single_controller_ppo_async_{self.config.max_async_step}_deepseek_l8b_open_r1_48k',
+            run_name=self.config.loggers.mlflow.run_name,
             tracking_uri=self.config.loggers.mlflow.tracking_uri,
         )
 
@@ -397,7 +396,7 @@ def __init__(
         self.eval_interval_num = int(config.eval_interval.strip("iter"))
         self.num_batches_per_update = config.variables.num_batches_per_update
         self.experiment_name = config.loggers.mlflow.experiment_name
-        self.run_name = f'test_single_controller_ppo_async_{config.max_async_step}_deepseek_l8b_open_r1_48k'
+        self.run_name = config.loggers.mlflow.run_name
 
         self.callback = self.build_callback()
 
@@ -464,20 +463,19 @@ def __init__(
         self.tokenizer_pad_token_id = ray.get(self.streaming_dataset_actor.get_tokenizer_pad_token_id.remote())
         self.prompt_handler_config = ray.get(self.streaming_dataset_actor.get_prompt_handler_config.remote())
         self.max_gen_len = self.prompt_handler_config['max_gen_len']
-
-        # Load iter_num from the checkpoint
         self.save_folder = os.path.join(config.save_folder, 'RolloutAgent')
-
         self.iter_num = 0
 
-        # Load the latest checkpoint
-
-        self.latest_checkpoint = os.path.join(self.save_folder, 'latest_rollout_agent.symlink') # TODO: This might need to use the updated path
-
-        if config.autoresume and _artifact_exists(self.latest_checkpoint):
+        # Load the latest checkpoint if we are autoresuming.
+        # Note that since we are checking if the checkpoint exists with
+        # mlflow.client.list_artifacts, we need to use the relative path to
+        # the checkpoint (i.e. not include dbfs://.../{mlflow_experiment_id}/{mlflow_run_id}
+        # in the path).
+        self.latest_checkpoint_path = os.path.join(self.save_folder, 'latest_rollout_agent.symlink')
+        if config.autoresume and _artifact_exists(self.latest_checkpoint_path):
             print(f'Autoresuming from checkpoint for RolloutAgent.')
-            get_file(self.latest_checkpoint, self.latest_checkpoint, overwrite=True)
-            with open(self.latest_checkpoint, 'rb') as f:
+            get_file(self.latest_checkpoint_path, self.latest_checkpoint_path, overwrite=True)
+            with open(self.latest_checkpoint_path, 'rb') as f:
                 checkpoint = pickle.load(f)
             self.iter_num = checkpoint['iter_num']
             print(f'Loading streaming dataloader state dict for RolloutAgent.', checkpoint['streaming_dataloader'])
@@ -521,29 +519,31 @@ def get_next_iter_rollouts(self):
         processed_sequences = torch.cat([all_prompts, padded_responses], dim=-1)
         iter_data['sequences'] = processed_sequences
 
-        save_folder_iter = os.path.join(self.save_folder, f'iter_{self.iter_num}')
-        checkpoint_path = os.path.join(save_folder_iter, 'checkpoint.pt')
+        save_folder_for_curr_iter = os.path.join(self.save_folder, f'iter_{self.iter_num}')
+        checkpoint_path = os.path.join(save_folder_for_curr_iter, 'checkpoint.pt')
         self.iter_num += 1
 
         streaming_dataloader_state_dict = ray.get(self.streaming_dataset_actor.get_dataloader_state_dict.remote())
         print(f'Streaming dataloader state dict for RolloutAgent.', streaming_dataloader_state_dict)
 
         # make sure that the folder path can exist
-        os.makedirs(save_folder_iter, exist_ok=True)
+        os.makedirs(save_folder_for_curr_iter, exist_ok=True)
         with open(checkpoint_path, 'wb') as f:
             pickle.dump({
                 'iter_data': iter_data,
                 'iter_num': self.iter_num,
                 'streaming_dataloader': streaming_dataloader_state_dict,
             }, f)
 
-        mlflow.log_artifact(checkpoint_path, save_folder_iter, run_id=_get_mlflow_run_id())
+        # log the checkpoint to mlflow
+        mlflow.log_artifact(checkpoint_path, save_folder_for_curr_iter, run_id=_get_mlflow_run_id())
 
-        if os.path.exists(self.latest_checkpoint):
-            os.remove(self.latest_checkpoint)
-        create_symlink_file(checkpoint_path, self.latest_checkpoint)
-        
-        mlflow.log_artifact(self.latest_checkpoint, self.config.save_folder, run_id=_get_mlflow_run_id())
+        if os.path.exists(self.latest_checkpoint_path):
+            os.remove(self.latest_checkpoint_path)
+        create_symlink_file(checkpoint_path, self.latest_checkpoint_path)
+
+        # log the latest checkpoint to mlflow
+        mlflow.log_artifact(self.latest_checkpoint_path, self.save_folder, run_id=_get_mlflow_run_id())
         return iter_data
 
     async def run(self, num_iterations: int, experience_buffer: 'ExperienceBuffer', lock: asyncio.Lock, rollout_semaphore: asyncio.Semaphore):
@@ -743,42 +743,64 @@ async def train_async(self, max_duration: int | str):
 def _get_mlflow_run_id() -> Optional[str]:
     return os.environ.get('MLFLOW_RUN_ID', None)
 
-def _setup_mlflow():
-    print('setting up mlflow')
+def _get_valid_mlflow_experiment_name(config: Any) -> str:
+    """Fixes the experiment name to be an absolute path for mlflow.
+
+    MLflow requires the experiment name to be an absolute path.
+    If the experiment name is not an absolute path, we prepend the current
+    user's username to the experiment name.
+    """
+    mlflow_experiment_name = config.loggers.mlflow.experiment_name
+    if mlflow_experiment_name.startswith('/'):
+        return mlflow_experiment_name
+    else:
+        from databricks.sdk import WorkspaceClient
+        return f'/Users/{WorkspaceClient().current_user.me().user_name}/{mlflow_experiment_name}'
+
+def _setup_mlflow(config: Any):
     dist.init_process_group(backend='gloo')
-    # Create a new MLFlow run to be used for the entire run
     mlflow.set_tracking_uri('databricks')
 
-    # get mlflow experiment
-    experiment = mlflow.get_experiment_by_name(MLFLOW_EXPERIMENT_NAME)
-    if experiment is None:
-        experiment_id = mlflow.create_experiment(MLFLOW_EXPERIMENT_NAME)
-    else:
-        experiment_id = experiment.experiment_id
-    mlflow.set_experiment(experiment_id=experiment_id)
+    # mlflow experiment name needs to be an absolute path for databricks mlflow.
+    mlflow_experiment_name = _get_valid_mlflow_experiment_name(config)
+    setattr(config.loggers.mlflow, 'experiment_name', mlflow_experiment_name)
+    # COMPOSER_RUN_NAME is set for interactive mode as well.
+    mlflow_run_name = os.environ['COMPOSER_RUN_NAME']
+    setattr(config.loggers.mlflow, 'run_name', mlflow_run_name)
 
+    # get mlflow experiment if it exists, otherwise create it and set it to all ranks.
+    experiment_id = None
+    if composer_dist.get_global_rank() == 0:
+        experiment = mlflow.get_experiment_by_name(mlflow_experiment_name)
+        if experiment is None:
+            experiment_id = mlflow.create_experiment(mlflow_experiment_name)
+        else:
+            experiment_id = experiment.experiment_id
+    experiment_id_broadcast_list = [experiment_id]
+    composer_dist.broadcast_object_list(experiment_id_broadcast_list, src=0)
+    experiment_id = experiment_id_broadcast_list[0]
 
+    mlflow.set_experiment(experiment_id=experiment_id)
 
+    # get mlflow run if it exists and we are autoresuming, otherwise create it and set it to all ranks.
     run_id = None
     if composer_dist.get_global_rank() == 0:
-        # find a preexisting run if it exists
         existing_runs = mlflow.search_runs(
             experiment_ids=[experiment_id],
-            filter_string=f'tags.run_name = "{MLFLOW_RUN_NAME}"',
+            filter_string=f'tags.run_name = "{mlflow_run_name}"',
             output_format='list',
         ) if config.autoresume else []
         if len(existing_runs) > 0:
             run_id = existing_runs[0].info.run_id
             print(f'Resuming mlflow run with run id: {run_id}')
         else:
-            run_id = mlflow.start_run(run_name=MLFLOW_RUN_NAME).info.run_id
+            run_id = mlflow.start_run(run_name=mlflow_run_name).info.run_id
             print(f'Creating new mlflow run with run id: {run_id}')
-    broadcast_list = [run_id]
-
-    composer_dist.broadcast_object_list(broadcast_list, src=0)
+    run_id_broadcast_list = [run_id]
+    composer_dist.broadcast_object_list(run_id_broadcast_list, src=0)
+    run_id = run_id_broadcast_list[0]
 
     # set all the right enviornment variables
-    run_id = broadcast_list[0]
     assert run_id is not None and experiment_id is not None, "Run ID and experiment ID must be set"
     os.environ['MLFLOW_RUN_ID'] = run_id
     os.environ['MLFLOW_EXPERIMENT_ID'] = experiment_id
@@ -812,8 +834,7 @@ def _artifact_exists(artifact_path: str) -> bool:
 
     # If we got here, the path exists (root or found item).
     return True
-    
-    
+
 
 def _run_single_controller_ppo(
     config: Any,
@@ -830,7 +851,7 @@ def _run_single_controller_ppo(
     # Disable setting CUDA_VISIBLE_DEVICES by ray, we will set it manually
     os.environ['RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES'] = '1'
 
-    _setup_mlflow()
+    _setup_mlflow(config)
 
     with start_ray_server() as _address:
         # only rank 0 is the master controller

yamls/single-controller-grpo-workflow.yaml

@@ -50,8 +50,6 @@ parameters:
   loggers:
     mlflow:
       tags:
-        run:
-          run_name: null
         group: grpo
       tracking_uri: databricks
       experiment_name: test_single_controller_ppo
@@ -163,7 +161,7 @@ parameters:
     gradient_clipping:
       clipping_type: norm
       clipping_threshold: 0.001
-  autoresume: true
+  autoresume: false
   log_config: true
   fsdp_config:
     verbose: false
@@ -178,7 +176,7 @@ parameters:
     activation_checkpointing: true
     activation_checkpointing_reentrant: false
   max_seq_len: 10240
-  save_folder: /tmp/checkpoints
+  save_folder: artifacts/checkpoints
   dist_timeout: 1800
   max_duration: 10iter
   progress_bar: false