[draft] proposal: generic trainer.backend registry (alternative to PR #1)

integrations/arctic-rl/arctic_rl/config.py

@@ -1,15 +1,81 @@
-"""Arctic RL client configuration builder.
+"""Arctic RL configuration types.
 
-Translates ``SkyRLTrainConfig`` into an ``ArcticRLClientConfig`` for the
-``arctic_training`` package.
+Defines:
+- ``ArcticRLTrainerConfig``: backend-specific knobs (colocate, zero_stage, ...)
+- ``ArcticTrainerConfig``: extends core ``TrainerConfig`` with ``arctic_rl`` field
+- ``ArcticSkyRLConfig``: top-level config used by the integration's entrypoint
+- ``build_rl_config(cfg)``: translates ``SkyRLTrainConfig`` → ``ArcticRLClientConfig``
 
-All shared knobs (GPU counts, colocation, vLLM settings) are derived from
-existing SkyRL config fields.  Only ARL-specific params live in
-``cfg.trainer.arctic_rl`` (see ``ArcticRLTrainerConfig``).
+These live in the integration to keep core SkyRL integration-agnostic — core only
+knows about a generic ``trainer.backend: str`` field that lazily dispatches here.
+All shared knobs (GPU counts, vLLM settings, colocation) are derived from existing
+SkyRL config fields by ``build_rl_config``.
 """
 
+from dataclasses import dataclass
+from typing import Optional
+
 from arctic_training.arctic_rl.config import ArcticRLClientConfig
 from skyrl.train.config import SkyRLTrainConfig
+from skyrl.train.config.config import BaseConfig, TrainerConfig, make_config
+
+
+# ---------------------------------------------------------------------------
+# Arctic RL backend configuration
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ArcticRLTrainerConfig(BaseConfig):
+    """Arctic RL (DeepSpeed) backend settings.
+
+    Only contains params unique to the Arctic RL server with no equivalent in
+    the standard SkyRL config.  Shared knobs are derived in ``build_rl_config``.
+    """
+
+    colocate: bool = False
+    """Share GPUs between training and inference on the ARL server.
+
+    Distinct from ``trainer.placement.colocate_all`` which controls Ray placement
+    groups.  ARL colocation is server-side GPU sharing — ``colocate_all`` must
+    stay ``false`` when using the ARL backend.
+    """
+    use_zorro: bool = False
+    """Enable ZoRRO (prompt deduplication) on the training server."""
+    zero_stage: int = 0
+    """DeepSpeed ZeRO stage (0, 2, or 3)."""
+    log_prob_gpus: int = 0
+    """Number of GPUs for log-prob computation (0 = skip separate log-prob)."""
+    offload_optimizer: bool = False
+    """Offload optimizer state to CPU when ``zero_stage >= 2``."""
+    host: str = "localhost"
+    """Server host for HTTP comm protocol; ignored for Ray."""
+    port: int = 7000
+    """Server port for HTTP comm protocol; ignored for Ray."""
+    startup_timeout: float = 300.0
+    """Seconds to wait for server jobs to come up."""
+    server_logs: bool = False
+    """Forward server logs to stdout for debugging."""
+
+
+@dataclass
+class ArcticTrainerConfig(TrainerConfig):
+    """``TrainerConfig`` extended with the Arctic RL field.  Used only when
+    ``trainer.backend == "arctic_rl"`` is set in the recipe."""
+
+    arctic_rl: Optional[ArcticRLTrainerConfig] = None
+    """Arctic RL backend settings.  ``None`` falls back to defaults."""
+
+
+# Top-level config for arctic_rl recipes.  Used by the integration's entrypoint
+# either directly (``uv run -m integrations.arctic_rl.entrypoint``) or via core
+# dispatch (``trainer.backend=arctic_rl`` from ``main_base``).
+ArcticSkyRLConfig = make_config(trainer_cls=ArcticTrainerConfig)
+
+
+# ---------------------------------------------------------------------------
+# Translation: SkyRLTrainConfig → ArcticRLClientConfig
+# ---------------------------------------------------------------------------
 
 
 def build_rl_config(cfg: SkyRLTrainConfig) -> ArcticRLClientConfig:

integrations/arctic-rl/arctic_rl/entrypoint.py

@@ -117,7 +117,12 @@ def skyrl_entrypoint(
 
 
 def main() -> None:
-    cfg = SkyRLTrainConfig.from_cli_overrides(sys.argv[1:])
+    """Arctic RL entrypoint. Reachable two ways: direct (``uv run -m
+    arctic_rl.entrypoint``) or via core dispatch (``python -m
+    skyrl.train.entrypoints.main_base trainer.backend=arctic_rl``).
+    Both paths parse with ``ArcticSkyRLConfig`` here."""
+    from arctic_rl.config import ArcticSkyRLConfig
+    cfg = ArcticSkyRLConfig.from_cli_overrides(sys.argv[1:])
     validate_cfg(cfg)
 
     rl_config = build_rl_config(cfg)
@@ -131,6 +136,9 @@ def main() -> None:
 
     from skyrl.train.utils.utils import prepare_runtime_environment
     env_vars = prepare_runtime_environment(cfg)
+    # Forward ARCTIC_* env vars to Ray workers — moved here from core utils per
+    # reviewer feedback (core stays integration-agnostic).
+    env_vars.update({k: v for k, v in os.environ.items() if k.startswith("ARCTIC_")})
     ray.init(num_gpus=0, runtime_env={"env_vars": env_vars})
     ray.get(skyrl_entrypoint.remote(cfg, reconnect_config=reconnect_cfg))
 

integrations/arctic-rl/examples/run_gsm8k_grpo_4gpu.sh

@@ -1,6 +1,10 @@
 #!/usr/bin/env bash
 # GSM8K GRPO training via Arctic RL server.
 #
+# Equivalent: (a) `python -m skyrl.train.entrypoints.main_base trainer.backend=arctic_rl ...`
+#             (b) `uv run --extra arctic-rl -m arctic_rl.entrypoint ...`
+# This script uses (a).
+#
 # Non-colocated (4 GPUs: 2 training + 2 sampling):
 #   bash integrations/arctic-rl/examples/run_gsm8k_grpo_4gpu.sh \
 #     trainer.placement.policy_num_gpus_per_node=2 \
@@ -26,6 +30,7 @@ export PYTHONUNBUFFERED=1
 python -m skyrl.train.entrypoints.main_base \
   data.train_data="['${DATA_DIR}/train.parquet']" \
   data.val_data="['${DATA_DIR}/validation.parquet']" \
+  trainer.backend=arctic_rl \
   trainer.arctic_rl={} \
   trainer.algorithm.advantage_estimator=grpo \
   trainer.policy.model.path="${MODEL}" \

pyproject.toml

@@ -2,14 +2,8 @@
 requires = ["setuptools"]
 build-backend = "setuptools.build_meta"
 
-# ``skyrl/`` is the upstream package. ``arctic-rl/`` is a top-level folder
-# (sibling of ``skyrl/``, like the legacy ``skyrl-tx/``) that hosts the
-# Arctic RL integration: the importable ``arctic_rl`` Python package and
-# its ``examples/`` subdir. Hyphenated outer folder + underscore module
-# matches the ``skyrl-tx/`` -> ``skyrl.tx`` precedent.
 [tool.setuptools.packages.find]
-where = [".", "integrations/arctic-rl"]
-include = ["skyrl*", "arctic_rl*"]
+include = ["skyrl*"]
 
 [project]
 name = "skyrl"

skyrl/train/config/config.py

@@ -192,55 +192,6 @@ class PlacementConfig(BaseConfig):
     ref_num_gpus_per_node: int = 1
 
 
-# ---------------------------------------------------------------------------
-# Arctic RL backend
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class ArcticRLTrainerConfig(BaseConfig):
-    """Arctic RL (DeepSpeed) backend settings.
-
-    Only contains params that are unique to the Arctic RL server and have
-    no equivalent in the standard SkyRL config.  Shared knobs are derived
-    automatically by ``build_rl_config``:
-
-    ============== ================================================
-    ARL concept    Derived from
-    ============== ================================================
-    training_gpus  ``trainer.placement.policy_num_gpus_per_node``
-                   ``* trainer.placement.policy_num_nodes``
-    sampling_gpus  ``generator.inference_engine.num_engines``
-    vllm mem/TP    ``generator.inference_engine.*``
-    ============== ================================================
-    """
-
-    colocate: bool = False
-    """Share GPUs between training and inference on the ARL server.
-
-    This is distinct from ``trainer.placement.colocate_all`` which controls
-    Ray placement groups.  ARL colocation is server-side GPU sharing managed
-    by the Arctic RL server — ``colocate_all`` must stay ``false`` when
-    using the ARL backend.
-    """
-    use_zorro: bool = False
-    """Enable ZoRRO (prompt deduplication) on the training server."""
-    zero_stage: int = 0
-    """DeepSpeed ZeRO stage (0, 2, or 3)."""
-    log_prob_gpus: int = 0
-    """Number of GPUs for log-prob computation (0 = skip separate log-prob)."""
-    host: str = "localhost"
-    """Arctic RL server hostname."""
-    port: int = 7000
-    """Arctic RL server port."""
-    startup_timeout: float = 600
-    """Seconds to wait for server readiness."""
-    server_logs: bool = False
-    """Enable verbose server logging."""
-    offload_optimizer: bool = False
-    """Offload optimizer to CPU. Applicable to ZeRO stage 2 and 3."""
-
-
 # ---------------------------------------------------------------------------
 # Policy / Critic / Ref
 # ---------------------------------------------------------------------------
@@ -683,8 +634,10 @@ class TrainerConfig(BaseConfig):
     rope_scaling: Optional[Dict[str, Any]] = None
     rope_theta: Optional[float] = None
 
-    arctic_rl: Optional[ArcticRLTrainerConfig] = None
-    """Config for the Arctic RL backend.  ``None`` means the backend is not used."""
+    backend: str = "fsdp"
+    """Training backend. ``"fsdp"`` is the standard SkyRL path; any other value
+    names an installed integration package (``<name>.entrypoint:main``) that
+    ``main_base`` lazily imports and dispatches to."""
 
     def __post_init__(self):
         # ref model defaults to the policy model

skyrl/train/entrypoints/main_base.py

@@ -473,21 +473,24 @@ def skyrl_entrypoint(cfg: SkyRLTrainConfig):
 
 
 def main() -> None:
-    # Parse CLI args and build typed config
+    # Peek at trainer.backend BEFORE strict config parse: integrations may add
+    # their own fields (e.g. ``trainer.arctic_rl``) that core SkyRLTrainConfig
+    # doesn't know about. If a non-default backend is selected, dispatch to the
+    # integration's entrypoint and let it parse with its own extended config.
+    # Generic — no integration-specific code lives here.
+    backend = "fsdp"
+    for arg in sys.argv[1:]:
+        if arg.startswith("trainer.backend="):
+            backend = arg.split("=", 1)[1]
+            break
+    if backend != "fsdp":
+        from importlib import import_module
+        backend_main = import_module(f"{backend}.entrypoint").main
+        return backend_main()
+
+    # Parse CLI args and build typed config (FSDP path)
     cfg = SkyRLTrainConfig.from_cli_overrides(sys.argv[1:])
 
-    # Route to Arctic RL entrypoint if arctic_rl backend is configured.
-    # This allows users to switch backends via config alone without changing
-    # the entrypoint command.  The integration code lives in a top-level
-    # ``arctic-rl/`` folder (sibling of ``skyrl/``, like the legacy
-    # ``skyrl-tx/``); the importable Python package is ``arctic_rl``.  It
-    # is distinct from the upstream ``arctic_training.arctic_rl`` sub-
-    # namespace — both coexist at import time without collision.
-    if cfg.trainer.arctic_rl is not None:
-        from arctic_rl.entrypoint import main as arctic_rl_main
-        arctic_rl_main()
-        return
-
     # validate the arguments
     validate_cfg(cfg)
 

skyrl/train/utils/utils.py

@@ -709,17 +709,9 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]:
         logger.info(f"Exporting `SKYRL_RAY_PG_TIMEOUT_IN_S` to ray runtime env: {pg_timeout}")
         env_vars["SKYRL_RAY_PG_TIMEOUT_IN_S"] = pg_timeout
 
-    _propagate_arctic_env_vars(env_vars)
     return env_vars
 
 
-def _propagate_arctic_env_vars(env_vars: dict) -> None:
-    """Propagate ARCTIC_* env vars to the Ray runtime so actors see them."""
-    for key, val in os.environ.items():
-        if key.startswith("ARCTIC_"):
-            env_vars[key] = val
-
-
 def configure_ray_worker_logging() -> None:
     """
     Configure logging for Ray workers.