update

Author: justusschock

src/lightning/fabric/plugins/environments/kubeflow.py

@@ -35,7 +35,7 @@ class KubeflowEnvironment(ClusterEnvironment):
 
     def __init__(self) -> None:
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.environments.kubeflow import (
+        from pytorch_lightning_enterprise.plugins.environments.kubeflow import (
             KubeflowEnvironment as EnterpriseKubeflowEnvironment,
         )
 

src/lightning/fabric/plugins/environments/lsf.py

@@ -51,7 +51,7 @@ def __init__(self) -> None:
         super().__init__()
 
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.environments.lsf import (
+        from pytorch_lightning_enterprise.plugins.environments.lsf import (
             LSFEnvironment as EnterpriseLSFEnvironment,
         )
 

src/lightning/fabric/plugins/environments/slurm.py

@@ -46,7 +46,7 @@ class SLURMEnvironment(ClusterEnvironment):
     def __init__(self, auto_requeue: bool = True, requeue_signal: Optional[signal.Signals] = None) -> None:
         super().__init__()
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.environments.slurm import (
+        from pytorch_lightning_enterprise.plugins.environments.slurm import (
             SLURMEnvironment as EnterpriseSLURMEnvironment,
         )
 
@@ -88,7 +88,7 @@ def job_name() -> Optional[str]:
     @staticmethod
     def job_id() -> Optional[int]:
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.environments.slurm import (
+        from pytorch_lightning_enterprise.plugins.environments.slurm import (
             SLURMEnvironment as EnterpriseSLURMEnvironment,
         )
 
@@ -134,7 +134,7 @@ def resolve_root_node_address(nodes: str) -> str:
 
         """
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.environments.slurm import (
+        from pytorch_lightning_enterprise.plugins.environments.slurm import (
             SLURMEnvironment as EnterpriseSLURMEnvironment,
         )
 
@@ -171,7 +171,7 @@ def _validate_srun_variables() -> None:
 
         """
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.environments.slurm import (
+        from pytorch_lightning_enterprise.plugins.environments.slurm import (
             SLURMEnvironment as EnterpriseSLURMEnvironment,
         )
 

src/lightning/fabric/plugins/environments/torchelastic.py

@@ -29,7 +29,7 @@ class TorchElasticEnvironment(ClusterEnvironment):
     def __init__(self) -> None:
         super().__init__()
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.environments.torchelastic import (
+        from pytorch_lightning_enterprise.plugins.environments.torchelastic import (
             TorchElasticEnvironment as EnterpriseTorchElasticEnvironment,
         )
 

src/lightning/fabric/plugins/environments/xla.py

@@ -35,7 +35,7 @@ class XLAEnvironment(ClusterEnvironment):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__()
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.environments.xla import (
+        from pytorch_lightning_enterprise.plugins.environments.xla import (
             XLAEnvironment as EnterpriseXLAEnvironment,
         )
 

src/lightning/fabric/plugins/io/xla.py

@@ -33,7 +33,7 @@ class XLACheckpointIO(TorchCheckpointIO):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.io.xla import XLACheckpointIO as EnterpriseXLACheckpointIO
+        from pytorch_lightning_enterprise.plugins.io.xla import XLACheckpointIO as EnterpriseXLACheckpointIO
 
         self.xla_impl = EnterpriseXLACheckpointIO(*args, **kwargs)
 

src/lightning/fabric/plugins/precision/bitsandbytes.py

@@ -53,7 +53,7 @@ def __init__(
     ) -> None:
         super().__init__()
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.precision.bitsandbytes import (
+        from pytorch_lightning_enterprise.plugins.precision.bitsandbytes import (
             BitsandbytesPrecision as EnterpriseBitsandbytesPrecision,
         )
 

src/lightning/fabric/plugins/precision/deepspeed.py

@@ -44,7 +44,7 @@ class DeepSpeedPrecision(Precision):
     def __init__(self, precision: _PRECISION_INPUT) -> None:
         super().__init__()
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.precision.deepspeed import (
+        from pytorch_lightning_enterprise.plugins.precision.deepspeed import (
             DeepSpeedPrecisionFabric as EnterpriseDeepSpeedPrecision,
         )
 

src/lightning/fabric/plugins/precision/transformer_engine.py

@@ -62,7 +62,7 @@ def __init__(
     ) -> None:
         super().__init__()
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.precision.transformer_engine import (
+        from pytorch_lightning_enterprise.plugins.precision.transformer_engine import (
             TransformerEnginePrecision as EnterpriseTransformerEnginePrecision,
         )
 

src/lightning/fabric/plugins/precision/xla.py

@@ -37,7 +37,7 @@ class XLAPrecision(Precision):
     def __init__(self, precision: _PRECISION_INPUT) -> None:
         super().__init__()
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.fabric.plugins.precision.xla import XLAPrecision as EnterpriseXLAPrecision
+        from pytorch_lightning_enterprise.plugins.precision.xla import XLAPrecision as EnterpriseXLAPrecision
 
         self.xla_impl = EnterpriseXLAPrecision(precision=precision)