Sync with new rslearn

This is for PR allenai/rslearn#359

Author: uakfdotb

rslp/lightning_cli.py

@@ -14,7 +14,7 @@
 from lightning.pytorch.callbacks import Callback
 from lightning.pytorch.cli import SaveConfigCallback
 from lightning.pytorch.utilities import rank_zero_only
-from rslearn.main import RslearnLightningCLI
+from rslearn.lightning_cli import RslearnLightningCLI
 from rslearn.train.lightning_module import RslearnLightningModule
 from rslearn.utils.fsspec import open_atomic
 from upath import UPath
@@ -421,7 +421,7 @@ def before_instantiate_classes(self) -> None:
 
 
 def custom_model_handler() -> None:
-    """Overrides model_handler in rslearn.main to use CustomLightningCLI.
+    """Overrides model_handler in rslearn.lightning_cli to use CustomLightningCLI.
 
     It also sets the save_config_callback.
     """

rslp/rslearn_main.py

@@ -16,12 +16,12 @@ def main() -> None:
     from dotenv import load_dotenv
 
     load_dotenv()
-    import rslearn.main
+    import rslearn.lightning_cli
 
     from rslp.lightning_cli import CustomLightningCLI, custom_model_handler
 
-    rslearn.main.RslearnLightningCLI = CustomLightningCLI
-    rslearn.main.model_handler = custom_model_handler
+    rslearn.lightning_cli.RslearnLightningCLI = CustomLightningCLI
+    rslearn.lightning_cli.model_handler = custom_model_handler
     rslearn.main.main()