[packaging] Make mlflow and wandb optional extras (not hard deps); switch to mlflow-skinny by default

[piyushumate]

Problem

nemo-automodel currently lists both mlflow and wandb>=0.28.0 as hard core dependencies:

# pyproject.toml (current)
dependencies = [
    ...
    "wandb>=0.28.0",
    "mlflow",
    ...
]

This causes two practical problems:

1. mlflow conflicts with mlflow-skinny

mlflow and mlflow-skinny are mutually exclusive PyPI packages — they both provide the mlflow namespace but cannot coexist in the same environment. Many downstream tools and managed ML platforms (sagemaker-mlflow, kedro-mlflow, Vertex AI, SageMaker containers) install mlflow-skinny by default. When those users try to add nemo-automodel, the dependency resolver hard-blocks with a conflict on Requires-Dist: mlflow.

Upstream precedent: sagemaker-mlflow PR #22 solved the exact same problem by switching mlflow → mlflow-skinny and adding a [full] extra for users who need the server/UI stack.

2. Unnecessary heavy transitive dependencies

mlflow (full) pulls in pyarrow, scipy, scikit-learn, flask, sqlalchemy, alembic, matplotlib, etc. wandb>=0.28.0 also adds significant overhead. These are unconditionally installed even for users who never configure a logger backend. The loggers are already opt-in via YAML config — the packaging should match.

---

Proposed Fix

Move both logger backends to PEP 508 optional extras, using mlflow-skinny as the default (full tracking API, no server/UI deps, no conflict):

# pyproject.toml (proposed)
dependencies = [
    # wandb and mlflow removed from here
    ...
]

[project.optional-dependencies]
mlflow      = ["mlflow-skinny"]        # lightweight client — full tracking API, no conflict
mlflow-full = ["mlflow"]               # full stack (UI, SQL backend, pyarrow, etc.)
wandb       = ["wandb>=0.28.0"]
tracking    = [                        # convenience meta-extra
    "nemo-automodel[mlflow]",
    "nemo-automodel[wandb]",
]
all = [
    "nemo-automodel[cli]",
    "nemo-automodel[cuda]",
    "nemo-automodel[tracking]",        # add tracking to all
    ...
]

Install patterns after this change:

pip install nemo-automodel                  # lean install, no logger overhead
pip install nemo-automodel[mlflow]          # MLflow skinny (no conflict)
pip install nemo-automodel[mlflow-full]     # MLflow full stack
pip install nemo-automodel[wandb]           # W&B only
pip install nemo-automodel[tracking]        # both backends
pip install nemo-automodel[all]             # everything

All import mlflow / import wandb call sites in the logger modules should be guarded with try/except ImportError, raising a clear, actionable error only when the backend is actually used without the extra installed:

try:
    import wandb
    HAS_WANDB = True
except ImportError:
    HAS_WANDB = False

def init_wandb_run(wandb_cfg, full_config, default_name=""):
    if not HAS_WANDB:
        raise ImportError(
            "wandb is not installed. To enable W&B logging, run:\n"
            "  pip install nemo-automodel[wandb]"
        )
    ...

---

Impact

Scenario	Before	After
pip install nemo-automodel	Installs mlflow + wandb unconditionally	Lean install
Environment with mlflow-skinny	❌ Conflict, install fails	✅ Works (skinny satisfies [mlflow] extra)
User wants MLflow UI/SQL backend	Works but conflicts with skinny envs	[mlflow-full] extra
User wants W&B	Gets it even if unused	[wandb] extra
pip install nemo-automodel[all]	Same as today (includes both)	Same, via [tracking]

---

References

• sagemaker-mlflow PR #22 — identical fix applied upstream
• mlflow-skinny PyPI — same client API, no heavy deps
• PEP 508 – optional-dependencies
• Python Packaging – optional-dependencies guide