Add lazy import for run_training (#116)

The `main_ds` module imports DeepSpeed and PyTorch. These imports take a
very long time and slow down startup of the Instruct CLI `ilab` tool by
over a second.

Use PEP 562 hooks to implement lazy import and loading for
`run_training` function. The function is imported when the attribute is
accessed.

See: instructlab/instructlab#1467

Signed-off-by: Christian Heimes <[email protected]>

src/instructlab/training/__init__.py

@@ -1,3 +1,14 @@
+__all__ = (
+    "DataProcessArgs",
+    "DeepSpeedOffloadStrategy",
+    "DeepSpeedOptions",
+    "LoraOptions",
+    "QuantizeDataType",
+    "TorchrunArgs",
+    "TrainingArgs",
+    "run_training",  # pylint: disable=undefined-all-variable
+)
+
 # Local
 from .config import (
     DataProcessArgs,
@@ -8,4 +19,20 @@
     TorchrunArgs,
     TrainingArgs,
 )
-from .main_ds import run_training
+
+
+def __dir__():
+    return globals().keys() | {"run_training"}
+
+
+def __getattr__(name):
+    # lazy import run_training
+    if name == "run_training":
+        # pylint: disable=global-statement,import-outside-toplevel
+        global run_training
+        # Local
+        from .main_ds import run_training
+
+        return run_training
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")