Initialize base inference before training model creation - #2003
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the ability to initialize base-model inference before the first training model is created, specifically when not colocated. It also adds a deployment event log when the inference proxy is published and includes a corresponding unit test. The review feedback highlights two main improvements: removing a redundant call to set_inference_engine_client in skyrl_train_backend.py since it is already handled during policy building, and replacing the naive bool() conversion in engine.py with a safer check to prevent string values like "False" from being evaluated as truthy.
| if self._inference_engines_initialized: | ||
| self._dispatch.set_inference_engine_client(self._inference_engine_client) | ||
| self.init_weight_sync_state() |
There was a problem hiding this comment.
The call to self._dispatch.set_inference_engine_client(self._inference_engine_client) is redundant here. self._dispatch is instantiated inside self._build_policy (called on the line immediately preceding this block) where self._inference_engine_client is already passed to the WorkerDispatch constructor. Removing this redundant call simplifies the initialization flow.
| if self._inference_engines_initialized: | |
| self._dispatch.set_inference_engine_client(self._inference_engine_client) | |
| self.init_weight_sync_state() | |
| if self._inference_engines_initialized: | |
| self.init_weight_sync_state() |
| if hasattr(self.backend, "set_inference_state_publisher"): | ||
| self.backend.set_inference_state_publisher(self._write_inference_state_to_db) | ||
|
|
||
| is_colocated = bool(config.backend_config.get("trainer.placement.colocate_all", True)) |
There was a problem hiding this comment.
Using bool() directly on a configuration value retrieved from config.backend_config can lead to unexpected behavior if the value is parsed as a string (e.g., from command-line overrides or environment variables). In Python, bool("False") or bool("false") evaluates to True because any non-empty string is truthy. It is safer to explicitly check against common falsey string and boolean representations.
colocate_val = config.backend_config.get("trainer.placement.colocate_all", True)
is_colocated = colocate_val not in (False, "False", "false", 0, "0")
| if not ray.is_initialized(): | ||
| logger.info("Initializing Ray for base-model inference") | ||
| initialize_ray(self._cfg) | ||
| self._colocate_pg = self._create_colocate_pg() if self._cfg.trainer.placement.colocate_all else None |
There was a problem hiding this comment.
It seems like initialize_base_inference() is only ever called in non-colocated mode. What is the purpose of having this check here?
2 participants
Testing
Focused backend tests: 10 passed; two unrelated JAX engine tests require the JAX extra.