chore: import upstream snapshot with attribution

nemo/utils/cloud.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import os
+import urllib.request
+from pathlib import Path
+from time import sleep
+
+from lightning.pytorch.plugins.environments import LightningEnvironment
+from lightning.pytorch.strategies import DDPStrategy, StrategyRegistry
+
+from nemo.utils import logging
+
+
+def maybe_download_from_cloud(url, filename, subfolder=None, cache_dir=None, refresh_cache=False) -> str:
+    """
+    Helper function to download pre-trained weights from the cloud
+    Args:
+        url: (str) URL of storage
+        filename: (str) what to download. The request will be issued to url/filename
+        subfolder: (str) subfolder within cache_dir. The file will be stored in cache_dir/subfolder. Subfolder can
+            be empty
+        cache_dir: (str) a cache directory where to download. If not present, this function will attempt to create it.
+            If None (default), then it will be $HOME/.cache/torch/NeMo
+        refresh_cache: (bool) if True and cached file is present, it will delete it and re-fetch
+
+    Returns:
+        If successful - absolute local path to the downloaded file
+        else - empty string
+    """
+    # try:
+    if cache_dir is None:
+        cache_location = Path.joinpath(Path.home(), ".cache/torch/NeMo")
+    else:
+        cache_location = cache_dir
+    if subfolder is not None:
+        destination = Path.joinpath(cache_location, subfolder)
+    else:
+        destination = cache_location
+
+    if not os.path.exists(destination):
+        os.makedirs(destination, exist_ok=True)
+
+    destination_file = Path.joinpath(destination, filename)
+
+    if os.path.exists(destination_file):
+        logging.info(f"Found existing object {destination_file}.")
+        if refresh_cache:
+            logging.info("Asked to refresh the cache.")
+            logging.info(f"Deleting file: {destination_file}")
+            os.remove(destination_file)
+        else:
+            logging.info(f"Re-using file from: {destination_file}")
+            return str(destination_file)
+    # download file
+    wget_uri = url + filename
+    logging.info(f"Downloading from: {wget_uri} to {str(destination_file)}")
+    # NGC links do not work everytime so we try and wait
+    i = 0
+    max_attempts = 3
+    while i < max_attempts:
+        i += 1
+        try:
+            urllib.request.urlretrieve(wget_uri, str(destination_file))
+            if os.path.exists(destination_file):
+                return destination_file
+            else:
+                return ""
+        except:
+            logging.info(f"Download from cloud failed. Attempt {i} of {max_attempts}")
+            sleep(0.05)
+            continue
+    raise ValueError("Not able to download url right now, please try again.")
+
+
+class SageMakerDDPStrategy(DDPStrategy):
+    """DDP strategy configured for AWS SageMaker distributed training."""
+
+    @property
+    def cluster_environment(self):
+        """Return a LightningEnvironment configured from SageMaker environment variables."""
+        env = LightningEnvironment()
+        env.world_size = lambda: int(os.environ["WORLD_SIZE"])
+        env.global_rank = lambda: int(os.environ["RANK"])
+        return env
+
+    @cluster_environment.setter
+    def cluster_environment(self, env):
+        """No-op setter to prevent Lightning from overriding the SageMaker environment."""
+        pass
+
+
+def initialize_sagemaker() -> None:
+    """
+    Helper function to initiate sagemaker with NeMo.
+    This function installs libraries that NeMo requires for the ASR toolkit + initializes sagemaker ddp.
+    """
+
+    logging.info("Registering SageMaker DDP strategy 'smddp'.")
+    StrategyRegistry.register(
+        name='smddp',
+        strategy=SageMakerDDPStrategy,
+        process_group_backend="smddp",
+        find_unused_parameters=False,
+    )
+
+    def _install_system_libraries() -> None:
+        import subprocess
+
+        logging.info("Installing system libraries: libsndfile1, ffmpeg")
+        try:
+            logging.info("Running apt-get update")
+            subprocess.run(
+                ["apt-get", "update"],
+                check=True,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            logging.info("Running apt-get install for libsndfile1 and ffmpeg")
+            subprocess.run(
+                ["apt-get", "install", "-y", "libsndfile1", "ffmpeg"],
+                check=True,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+        except subprocess.CalledProcessError as e:
+            logging.error(
+                "Failed to install system libraries via apt-get (command=%s, returncode=%s): %s",
+                getattr(e, "cmd", None),
+                getattr(e, "returncode", None),
+                e,
+            )
+            logging.info("System libraries installed successfully.")
+        except Exception as e:
+            logging.error(f"Failed to install system libraries: {e}")
+
+    def _patch_torch_metrics() -> None:
+        """
+        Patches torchmetrics to not rely on internal state.
+        This is because sagemaker DDP overrides the `__init__` function of the modules to do automatic-partitioning.
+        """
+        from torchmetrics import Metric
+
+        def __new_hash__(self):
+            hash_vals = [self.__class__.__name__, id(self)]
+            return hash(tuple(hash_vals))
+
+        Metric.__hash__ = __new_hash__
+
+    logging.info("Patching torchmetrics hash function for SageMaker compatibility.")
+    _patch_torch_metrics()
+
+    if os.environ.get("RANK") and os.environ.get("WORLD_SIZE"):
+        import smdistributed.dataparallel.torch.distributed as dist
+
+        # has to be imported, as it overrides torch modules and such when DDP is enabled.
+        import smdistributed.dataparallel.torch.torch_smddp  # noqa: F401
+
+        logging.info("Initializing SageMaker distributed process group.")
+        dist.init_process_group()
+
+        if dist.get_local_rank():
+            _install_system_libraries()
+        logging.info("Waiting at barrier for all processes.")
+        return dist.barrier()  # wait for main process
+    _install_system_libraries()
+    return