(fix) Move mlflow/wandb out of core deps and into extras

.github/workflows/full-test-suite.yml

@@ -70,8 +70,10 @@ jobs:
           # The telemetry extra installs nemo-lens so the lens-gated tests under
           # tests/unit_tests/telemetry actually run; without it they skip and the
           # telemetry modules count as uncovered, failing the coverage gate below.
-          # Mirrors scripts/ci/core_unit_tests.sh, which syncs the same extra.
-          uv sync --extra dev --extra sandbox --extra openshell --extra telemetry
+          # mlflow/wandb are optional extras; the core suite covers both their
+          # installed- and missing-dependency paths, so they're synced here too.
+          # Mirrors scripts/ci/core_unit_tests.sh, which syncs the same extras.
+          uv sync --extra dev --extra sandbox --extra openshell --extra telemetry --extra mlflow --extra wandb
 
       - name: Test
         run: |

.github/workflows/unit-tests.yml

@@ -90,7 +90,7 @@ jobs:
         run: |
           ./scripts/ci/core_unit_tests.sh
           source ./scripts/ci/setup_dev.sh
-          uv sync --extra dev --extra sandbox --extra openshell
+          uv sync --extra dev --extra sandbox --extra openshell --extra mlflow --extra wandb
           source .venv/bin/activate
           pytest --cov=. --cov-append --cov-report=term-missing \
             --cov-fail-under="$(python3 scripts/ci/cov_fail_under.py)" --durations=10 -m sandbox

fern/versions/latest/pages/reference/configuration.mdx

@@ -322,6 +322,16 @@ Gym exports the resolved config, aggregate metrics, and rollouts to any tracking
 configure. Weights & Biases and MLflow are both supported, and both can be on at once. A backend
 is used only when every key it needs is set; otherwise it is silently skipped.
 
+Both SDKs are optional extras, not part of the core install.
+
+```bash
+uv sync --extra wandb   # Weights & Biases
+uv sync --extra mlflow  # MLflow
+```
+
+If a backend is configured but its extra isn't installed, Gym logs an error naming the missing
+extra and continues the run without that backend.
+
 ### Weights & Biases
 
 | Key | Required | Description |

nemo_gym/exporters/__init__.py

@@ -35,10 +35,10 @@
 
 logger = logging.getLogger(__name__)
 
-# Backend name -> (config model, "module:class").
-EXPORTER_REGISTRY: dict[str, tuple[type[ExporterConfig], str]] = {
-    "wandb": (WANDBConfig, "nemo_gym.exporters.wandb:WandbExporter"),
-    "mlflow": (MLFlowConfig, "nemo_gym.exporters.mlflow:MLflowExporter"),
+# Backend name -> (config model, "module:class", extra to install its SDK).
+EXPORTER_REGISTRY: dict[str, tuple[type[ExporterConfig], str, str]] = {
+    "wandb": (WANDBConfig, "nemo_gym.exporters.wandb:WandbExporter", "wandb"),
+    "mlflow": (MLFlowConfig, "nemo_gym.exporters.mlflow:MLflowExporter", "mlflow"),
 }
 
 _EXPORTERS: list[BaseExporter] = []
@@ -62,12 +62,18 @@ def setup_exporters(global_config_dict: DictConfig) -> list[BaseExporter]:
     """
     teardown_exporters()
 
-    for name, (config_model, class_path) in EXPORTER_REGISTRY.items():
+    for name, (config_model, class_path, extra) in EXPORTER_REGISTRY.items():
         if not config_model.model_validate(global_config_dict).is_available:
             continue
         try:
             exporter = _load_exporter_class(class_path)(global_config_dict)
             exporter.setup()
+        except ImportError as e:
+            logger.error(
+                f"Exporter {name} is configured but its SDK is not installed; continuing without it. "
+                f"Install with: pip install nemo-gym[{extra}] ({e})"
+            )
+            continue
         except Exception as e:
             logger.warning(f"Exporter {name} failed to start; continuing without it: {e}", exc_info=True)
             continue

nemo_gym/gitlab_utils.py

@@ -14,12 +14,9 @@
 # limitations under the License.
 from os import environ
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import requests
-from mlflow import MlflowClient
-from mlflow.artifacts import get_artifact_repository
-from mlflow.environment_variables import MLFLOW_TRACKING_TOKEN
-from mlflow.exceptions import RestException
 
 from nemo_gym.config_types import (
     ConfigError,
@@ -30,7 +27,21 @@
 from nemo_gym.server_utils import get_global_config_dict
 
 
-def create_mlflow_client() -> MlflowClient:  # pragma: no cover
+if TYPE_CHECKING:
+    from mlflow import MlflowClient
+
+
+def create_mlflow_client() -> "MlflowClient":  # pragma: no cover
+    # mlflow is an optional extra (`nemo-gym[mlflow]`): the GitLab dataset/model-registry
+    # integration is built on it, but a plain Gym install must not require it just to import
+    # this module. Import lazily here so the failure is scoped to this call and actionable.
+    try:
+        from mlflow import MlflowClient
+    except ImportError as e:
+        raise ConfigError(
+            "GitLab dataset storage needs the mlflow extra. Install with:\n  pip install nemo-gym[mlflow]"
+        ) from e
+
     global_config = get_global_config_dict()
     config = MLFlowConfig.model_validate(global_config)
 
@@ -52,6 +63,7 @@ def upload_jsonl_dataset(
     config: UploadJsonlDatasetGitlabConfig,
 ) -> None:  # pragma: no cover
     client = create_mlflow_client()
+    from mlflow.exceptions import RestException
 
     try:
         client.create_registered_model(config.dataset_name)
@@ -83,6 +95,8 @@ def download_jsonl_dataset(
 ) -> None:  # pragma: no cover
     # TODO: There is probably a much better way to do this, but it is not clear at the moment.
     client = create_mlflow_client()
+    from mlflow.artifacts import get_artifact_repository
+    from mlflow.environment_variables import MLFLOW_TRACKING_TOKEN
 
     model_version = client.get_model_version(config.dataset_name, config.version)
     run_id = model_version.run_id
@@ -100,6 +114,7 @@ def download_jsonl_dataset(
 
 def is_model_in_gitlab(model_name: str) -> bool:  # pragma: no cover
     client = create_mlflow_client()
+    from mlflow.exceptions import RestException
 
     # model_name in gitlab is case sensitive
     try:

nemo_gym/reward_profile.py

@@ -21,14 +21,17 @@
 from collections import Counter, defaultdict
 from numbers import Real
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 import orjson
 from pandas import DataFrame, Series, notna
 from pandas.core.groupby.generic import DataFrameGroupBy
 from pydantic import Field
 from scipy import stats
-from wandb import Histogram
+
+
+if TYPE_CHECKING:
+    from wandb import Histogram
 
 from nemo_gym.config_types import AggregateMetrics, BaseNeMoGymCLIConfig
 from nemo_gym.global_config import (
@@ -175,12 +178,28 @@ def rollout_info_from_result(self, result: Dict[str, Any]) -> Dict[str, Any]:
 
         return rollout_info
 
-    def histogram(self, data: Series) -> Optional[Histogram]:
+    def histogram(self, data: Series) -> Optional["Histogram"]:
         # W&B doesn't accept empty histograms
         data = data.dropna()
         if data.empty:
             return
 
+        # wandb is an optional extra (`nemo-gym[wandb]`). This stat is always dropped by
+        # prepare_for_serialization before it reaches any JSON output or exporter (it exists
+        # only for a wandb-native run that reads group_level_metrics/agent_metrics directly),
+        # so skipping it when wandb isn't installed changes nothing observable.
+        try:
+            from wandb import Histogram
+        except ImportError:
+            # warnings.warn's default filter shows this once per (module, lineno), so it
+            # doesn't spam once per column per describe_dataframe call.
+            warnings.warn(
+                "wandb is not installed, so the histogram/* stats are being skipped. "
+                "Install with: pip install nemo-gym[wandb]",
+                stacklevel=2,
+            )
+            return None
+
         return Histogram(data)
 
     def describe_dataframe(self, df: DataFrame) -> DataFrame:

pyproject.toml

@@ -181,13 +181,6 @@ dependencies = [
     # License: MIT https://github.com/Textualize/rich/blob/27701029082052c6541031b6f139097230253900/LICENSE
     "rich",
 
-    # mlflow-skinny: used for interacting with the Gitlab model registry. We use the "skinny" version that does not have unnecessary dependencies.
-    # mlflow dependency is necessary for importing. However, all of mlflow's dependencies that don't overlap with mlflow-skinny are excluded below.
-    # Updated Mon Sep 28, 2026 with mlflow>=3.16.1 (security bump)
-    # License: Apache 2.0 https://github.com/mlflow/mlflow/blob/c6ce84e817897e1a0461632f8a06d4e5e4ac3726/LICENSE.txt
-    "mlflow-skinny>=3.16.1",
-    "mlflow>=3.16.1",
-
     # scipy: t-distribution confidence intervals for repeat-level reward metrics in
     # nemo_gym/reward_profile.py. Lower bound is the first release with cp313 wheels;
     # deliberately NOT >=1.18, which requires numpy>=2.0 and breaks servers pinned to
@@ -244,12 +237,6 @@ dependencies = [
     # License: MIT https://raw.githubusercontent.com/jpadilla/pyjwt/master/LICENSE
     "PyJWT>=2.15.1",
 
-    # sqlparse: SQL parsing library (transitive dependency of mlflow-skinny, pinned for security)
-    # mlflow-skinny only requires sqlparse<1,>=0.4.0, so it won't pull the fixed version on its own.
-    # Updated: Wed Sep 02, 2026 with sqlparse>=0.6.0 (security bump)
-    # License: BSD 3-Clause https://github.com/andialbrecht/sqlparse/blob/master/LICENSE
-    "sqlparse>=0.6.0",
-
     # fonttools: Library to manipulate font files (transitive dependency, pinned for security)
     # Updated: Mon Feb 10, 2026 with fonttools>=4.60.2
     # License: MIT https://github.com/fonttools/fonttools/blob/main/LICENSE
@@ -260,12 +247,45 @@ dependencies = [
     # License: Apache 2.0 https://github.com/andrew-d/python-multipart/blob/master/LICENSE.txt
     "python-multipart>=0.0.27",
 
+    # pyasn1: ASN.1 library (transitive dependency via google-auth, pinned for security)
+    # Updated: Tue Jul 29, 2026 with pyasn1>=0.6.4
+    # License: BSD 2-Clause https://pyasn1.readthedocs.io/en/latest/license.html
+    "pyasn1>=0.6.4",
+
+    # pydot License: MIT https://github.com/pydot/pydot/blob/a2bee1e770698867332516fb2878fac262b4ffc5/LICENSES/MIT.txt
+    "pydot",
+]
+
+[project.optional-dependencies]
+all = [
+    "nemo-gym[dev,mlflow,sandbox,wandb]",
+]
+
+# MLflow is used both by the rollout exporter (nemo_gym/exporters/mlflow.py) and by the
+# GitLab dataset/model-registry integration (nemo_gym/gitlab_utils.py). Without this extra,
+# both call paths report an actionable error instead of failing on unrelated imports.
+mlflow = [
+    # mlflow-skinny: used for interacting with the Gitlab model registry. We use the "skinny" version that does not have unnecessary dependencies.
+    # mlflow dependency is necessary for importing. However, all of mlflow's dependencies that don't overlap with mlflow-skinny are excluded below.
+    # Updated Mon Sep 28, 2026 with mlflow>=3.16.1 (security bump)
+    # License: Apache 2.0 https://github.com/mlflow/mlflow/blob/c6ce84e817897e1a0461632f8a06d4e5e4ac3726/LICENSE.txt
+    "mlflow-skinny>=3.16.1",
+    "mlflow>=3.16.1",
+
+    # sqlparse: SQL parsing library (transitive dependency of mlflow-skinny, pinned for security)
+    # mlflow-skinny only requires sqlparse<1,>=0.4.0, so it won't pull the fixed version on its own.
+    # Updated: Wed Sep 02, 2026 with sqlparse>=0.6.0 (security bump)
+    # License: BSD 3-Clause https://github.com/andialbrecht/sqlparse/blob/master/LICENSE
+    "sqlparse>=0.6.0",
+
     # pyarrow: columnar data library (transitive dependency via mlflow, pinned for security)
     # mlflow caps pyarrow<25; versions >=23.0.1 have cp313 wheels and address use-after-free.
     # Updated: Tue Jul 29, 2026 with pyarrow>=23.0.1
     # License: Apache 2.0 https://github.com/apache/arrow/blob/main/LICENSE
     "pyarrow>=23.0.1",
+]
 
+wandb = [
     # wandb: E2E rollout collection data and metrics upload
     # Updated: Mon Sep 28, 2026 with wandb>=0.30.0 (CVE mitigation: x/crypto, x/net, grpc-go, go-git bundled in wandb-core)
     # License: MIT https://github.com/wandb/wandb/blob/f8acf479342b6aa8217dd0833bb32190b11c14bc/LICENSE
@@ -275,19 +295,6 @@ dependencies = [
     # Updated: Fri Sep 04, 2026 with GitPython>=3.1.61
     # License: BSD 3-Clause https://github.com/gitpython-developers/GitPython/blob/main/LICENSE
     "GitPython>=3.1.61",
-
-    # pyasn1: ASN.1 library (transitive dependency via google-auth, pinned for security)
-    # Updated: Tue Jul 29, 2026 with pyasn1>=0.6.4
-    # License: BSD 2-Clause https://pyasn1.readthedocs.io/en/latest/license.html
-    "pyasn1>=0.6.4",
-
-    # pydot License: MIT https://github.com/pydot/pydot/blob/a2bee1e770698867332516fb2878fac262b4ffc5/LICENSES/MIT.txt
-    "pydot",
-]
-
-[project.optional-dependencies]
-all = [
-    "nemo-gym[dev,sandbox]",
 ]
 
 vllm = [

scripts/ci/core_unit_tests.sh

@@ -19,8 +19,11 @@ source "${ci_dir}/setup_dev.sh"
 # The telemetry extra pulls in nemo-lens so the lens-present unit tests under
 # tests/unit_tests/telemetry (pytest.mark tests/unit_tests/telemetry/conftest.py::requires_lens)
 # actually run in CI instead of silently skipping, which would leave that code permanently
-# uncovered by the coverage gate below.
-uv sync --extra dev --extra telemetry
+# uncovered by the coverage gate below. mlflow and wandb are optional extras (SDKs now opt-in
+# to keep a core install lean); this suite covers both their installed- and missing-dependency
+# paths (tests/unit_tests/test_exporters.py, tests/unit_tests/test_gitlab_utils.py), so both
+# extras must be present here.
+uv sync --extra dev --extra telemetry --extra mlflow --extra wandb
 pytest_addopts='-m "not sandbox" --cov-report= --cov-fail-under=0 --color=yes'
 if [[ -n "${GYM_CI_JUNIT_DIR:-}" ]]; then
     mkdir -p "${GYM_CI_JUNIT_DIR}"

tests/unit_tests/test_exporters.py

@@ -117,7 +117,9 @@ def mlflow_config() -> DictConfig:
 def _register_recording(monkeypatch: MonkeyPatch) -> MagicMock:
     """Point the registry at RecordingExporter and hand back the (spied) lazy loader."""
     loader = MagicMock(return_value=RecordingExporter)
-    monkeypatch.setattr(exporters_module, "EXPORTER_REGISTRY", {"recording": (RecordingConfig, "recording")})
+    monkeypatch.setattr(
+        exporters_module, "EXPORTER_REGISTRY", {"recording": (RecordingConfig, "recording", "recording")}
+    )
     monkeypatch.setattr(exporters_module, "_load_exporter_class", loader)
     return loader
 
@@ -176,6 +178,24 @@ def test_setup_failure_is_skipped_rather_than_raised(
         assert setup_exporters(wandb_config) == []
         assert get_exporters() == []
 
+    def test_missing_sdk_is_skipped_with_an_actionable_install_hint(
+        self, monkeypatch: MonkeyPatch, wandb_config: DictConfig, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        monkeypatch.setattr(
+            exporters_module, "EXPORTER_REGISTRY", {"recording": (RecordingConfig, "recording", "recording")}
+        )
+        monkeypatch.setattr(
+            exporters_module,
+            "_load_exporter_class",
+            MagicMock(side_effect=ImportError("No module named 'recording_sdk'")),
+        )
+
+        with caplog.at_level("WARNING"):
+            assert setup_exporters(wandb_config) == []
+
+        assert get_exporters() == []
+        assert "pip install nemo-gym[recording]" in caplog.text
+
     def test_setup_replaces_previously_opened_exporters(
         self, monkeypatch: MonkeyPatch, wandb_config: DictConfig
     ) -> None:

tests/unit_tests/test_gitlab_utils.py

@@ -12,7 +12,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from nemo_gym.config_types import MLFlowConfig
+import sys
+
+import pytest
+from pytest import MonkeyPatch
+
+from nemo_gym.config_types import ConfigError, MLFlowConfig
+from nemo_gym.gitlab_utils import create_mlflow_client
 
 
 # TODO: Eventually we want to add more tests to ensure that the Gitlab flow does not break
@@ -24,3 +30,13 @@ def test_registry_credentials_alone_do_not_enable_the_exporter(self) -> None:
         config = MLFlowConfig(mlflow_tracking_uri="https://gitlab", mlflow_tracking_token="t")
 
         assert not config.is_available
+
+    def test_create_mlflow_client_gives_an_actionable_error_when_mlflow_is_missing(
+        self, monkeypatch: MonkeyPatch
+    ) -> None:
+        # mlflow is an optional extra (`nemo-gym[mlflow]`); force the lazy import in
+        # create_mlflow_client to fail the way it would on a plain Gym install.
+        monkeypatch.setitem(sys.modules, "mlflow", None)
+
+        with pytest.raises(ConfigError, match=r"pip install nemo-gym\[mlflow\]"):
+            create_mlflow_client()