Make NEMORUN_HOME configurable

Signed-off-by: Hemil Desai <hemild@nvidia.com>

Author: hemildesai

src/nemo_run/cli/api.py

@@ -56,7 +56,13 @@
 from nemo_run.cli import devspace as devspace_cli
 from nemo_run.cli import experiment as experiment_cli
 from nemo_run.cli.cli_parser import parse_cli_args, parse_factory
-from nemo_run.config import NEMORUN_HOME, Config, Partial, get_type_namespace, get_underlying_types
+from nemo_run.config import (
+    Config,
+    Partial,
+    get_nemorun_home,
+    get_type_namespace,
+    get_underlying_types,
+)
 from nemo_run.core.execution import LocalExecutor, SkypilotExecutor, SlurmExecutor
 from nemo_run.core.execution.base import Executor
 from nemo_run.core.frontend.console.styles import BOX_STYLE, TABLE_STYLES
@@ -758,7 +764,7 @@ def _search_workspace_file() -> str | None:
     file_names = [
         "workspace_private.py",
         "workspace.py",
-        os.path.join(NEMORUN_HOME, "workspace.py"),
+        os.path.join(get_nemorun_home(), "workspace.py"),
     ]
 
     while True:

src/nemo_run/config.py

@@ -44,12 +44,33 @@
 _BuildableT = TypeVar("_BuildableT", bound=fdl.Buildable)
 
 RECURSIVE_TYPES = (typing.Union, typing.Optional)
-NEMORUN_HOME = os.environ.get("NEMORUN_HOME", os.path.expanduser("~/.nemo_run"))
+_NEMORUN_HOME = os.environ.get("NEMORUN_HOME", os.path.expanduser("~/.nemo_run"))
 RUNDIR_NAME = "nemo_run"
 RUNDIR_SPECIAL_NAME = "/$nemo_run"
 SCRIPTS_DIR = "scripts"
 
 
+def get_nemorun_home() -> str:
+    """
+    Get the current NEMORUN_HOME directory path.
+
+    Returns:
+        The path to the NEMORUN_HOME directory.
+    """
+    return _NEMORUN_HOME
+
+
+def set_nemorun_home(path: str) -> None:
+    """
+    Set the NEMORUN_HOME directory path.
+
+    Args:
+        path: The new path for NEMORUN_HOME.
+    """
+    global _NEMORUN_HOME
+    _NEMORUN_HOME = os.path.expanduser(path)
+
+
 def get_type_namespace(typ: Type | Callable) -> str:
     """
     Get the namespace of a type or callable.

src/nemo_run/core/execution/base.py

@@ -24,7 +24,7 @@
 from torchx.specs import Role
 from typing_extensions import Self
 
-from nemo_run.config import NEMORUN_HOME, ConfigurableMixin
+from nemo_run.config import ConfigurableMixin, get_nemorun_home
 from nemo_run.core.execution.launcher import LAUNCHER_MAP, Launcher
 from nemo_run.core.packaging.base import Packager
 
@@ -201,25 +201,25 @@ def import_executor(
 
     Example:
         executor = import_executor("local", file_path="path/to/executors.py")
-        executor = import_executor("gpu")  # Uses the default location of os.path.join(NEMORUN_HOME, "executors.py")
+        executor = import_executor("gpu")  # Uses the default location of os.path.join(get_nemorun_home(), "executors.py")
 
     Args:
         name (str): The name of the executor to retrieve.
         file_path (Optional[str]): The path to the Python file containing the executor definitions.
-            Defaults to None, in which case the default location of os.path.join(NEMORUN_HOME, "executors.py") is used.
+            Defaults to None, in which case the default location of os.path.join(get_nemorun_home(), "executors.py") is used.
 
             The file_path is expected to be a string representing a file path with the following structure:
             - It should be a path to a Python file (with a .py extension).
             - The file should contain a dictionary named `EXECUTOR_MAP` that maps executor names to their corresponding instances.
-            - The file can be located anywhere in the file system, but if not provided, it defaults to `NEMORUN_HOME/executors.py`.
+            - The file can be located anywhere in the file system, but if not provided, it defaults to `get_nemorun_home()/executors.py`.
         call (bool): If True, the value from the module is called with the rest of the given kwargs.
 
     Returns:
         Executor: The executor instance corresponding to the given name.
     """
 
     if not file_path:
-        file_path = os.path.join(NEMORUN_HOME, "executors.py")
+        file_path = os.path.join(get_nemorun_home(), "executors.py")
 
     spec = importlib.util.spec_from_file_location("executors", file_path)
     assert spec

src/nemo_run/core/execution/dgxcloud.py

@@ -234,7 +234,7 @@ def assign(
                 == os.path.abspath(x["path"]),
                 self.pvcs,
             )
-        ), f"Need to specify atleast one PVC containing {self.job_dir}.\nTo update job dir to a PVC path, you can set the NEMORUN_HOME env var."
+        ), f"Need to specify atleast one PVC containing {self.job_dir}.\nTo update job dir to a PVC path, you can use set_nemorun_home() or the NEMORUN_HOME env var."
 
     def package(self, packager: Packager, job_name: str):
         assert self.experiment_id, "Executor not assigned to an experiment."

src/nemo_run/core/execution/docker.py

@@ -31,7 +31,7 @@
     parse_app_handle,
 )
 
-from nemo_run.config import NEMORUN_HOME, RUNDIR_NAME
+from nemo_run.config import RUNDIR_NAME, get_nemorun_home
 from nemo_run.core.execution.base import Executor
 from nemo_run.core.packaging.base import Packager
 from nemo_run.core.packaging.git import GitArchivePackager
@@ -50,7 +50,7 @@
     fcntl = None
     FCNTL_AVAILABLE = False
 
-DOCKER_JOB_DIRS = os.path.join(NEMORUN_HOME, ".docker_jobs.json")
+DOCKER_JOB_DIRS = os.path.join(get_nemorun_home(), ".docker_jobs.json")
 NETWORK = "nemo_run"
 
 LABEL_EXPERIMENT_ID: str = "nemo-run/experiment-id"

src/nemo_run/core/tunnel/client.py

@@ -32,12 +32,12 @@
 from invoke.context import Context
 from invoke.runners import Result as RunResult
 
-from nemo_run.config import NEMORUN_HOME, ConfigurableMixin
+from nemo_run.config import ConfigurableMixin, get_nemorun_home
 from nemo_run.core.frontend.console.api import CONSOLE
 
 logger: logging.Logger = logging.getLogger(__name__)
 TUNNEL_DIR = ".tunnels"
-TUNNEL_FILE_SUBPATH = os.path.join(NEMORUN_HOME, TUNNEL_DIR)
+TUNNEL_FILE_SUBPATH = os.path.join(get_nemorun_home(), TUNNEL_DIR)
 
 
 def delete_tunnel_dir(file_path):
@@ -136,7 +136,7 @@ def __post_init__(self):
 
     def _set_job_dir(self, experiment_id: str):
         experiment_title, _, _ = experiment_id.rpartition("_")
-        base_job_dir = self.job_dir or os.path.join(NEMORUN_HOME, "experiments")
+        base_job_dir = self.job_dir or os.path.join(get_nemorun_home(), "experiments")
         job_dir = os.path.join(base_job_dir, experiment_title, experiment_id)
         self.job_dir = job_dir
 

src/nemo_run/run/experiment.py

@@ -41,11 +41,11 @@
 
 import nemo_run as run
 from nemo_run.config import (
-    NEMORUN_HOME,
     Config,
     ConfigurableMixin,
     Partial,
     Script,
+    get_nemorun_home,
     get_type_namespace,
 )
 from nemo_run.core.execution.base import Executor
@@ -101,7 +101,7 @@ class Experiment(ConfigurableMixin):
     The design is heavily inspired from `XManager <https://github.com/google-deepmind/xmanager/blob/main/docs/xm_launch_api_principles.md>`_.
 
     Under the hood, the Experiment metadata is stored in the local filesystem
-    inside a user specified directory controlled by NEMORUN_HOME env var.
+    inside a user specified directory controlled by get_nemorun_home() env var.
     We will explore making the metadata more persistent in the future.
 
     .. note::
@@ -206,9 +206,9 @@ def catalog(
         title: str = "",
     ) -> list[str]:
         """
-        List all experiments inside NEMORUN_HOME, optionally with the provided title.
+        List all experiments inside get_nemorun_home(), optionally with the provided title.
         """
-        parent_dir = os.path.join(NEMORUN_HOME, "experiments", title)
+        parent_dir = os.path.join(get_nemorun_home(), "experiments", title)
         return _get_sorted_dirs(parent_dir)
 
     @classmethod
@@ -247,7 +247,7 @@ def from_id(
         Reconstruct an experiment with the specified id.
         """
         title, _, _ = id.rpartition("_")
-        parent_dir = os.path.join(NEMORUN_HOME, "experiments", title)
+        parent_dir = os.path.join(get_nemorun_home(), "experiments", title)
         exp_dir = os.path.join(parent_dir, id)
 
         assert os.path.isdir(exp_dir), f"Experiment {id} not found."
@@ -263,7 +263,7 @@ def from_title(
         """
         Reconstruct an experiment with the specified title.
         """
-        parent_dir = os.path.join(NEMORUN_HOME, "experiments", title)
+        parent_dir = os.path.join(get_nemorun_home(), "experiments", title)
         exp_dir = _get_latest_dir(parent_dir)
 
         assert os.path.isdir(exp_dir), f"Experiment {id} not found."
@@ -303,7 +303,7 @@ def __init__(
         self._title = title
         self._id = id or f"{title}_{int(time.time())}"
 
-        base_dir = base_dir or NEMORUN_HOME
+        base_dir = base_dir or get_nemorun_home()
         self._exp_dir = os.path.join(base_dir, "experiments", title, self._id)
 
         self.log_level = log_level
@@ -967,7 +967,7 @@ def reset(self) -> "Experiment":
 
         old_id, old_exp_dir, old_launched = self._id, self._exp_dir, self._launched
         self._id = f"{self._title}_{int(time.time())}"
-        self._exp_dir = os.path.join(NEMORUN_HOME, "experiments", self._title, self._id)
+        self._exp_dir = os.path.join(get_nemorun_home(), "experiments", self._title, self._id)
         self._launched = False
         self._live_progress = None
 
@@ -1017,7 +1017,7 @@ def reset(self) -> "Experiment":
                 f"[bold magenta]Failed resetting Experiment {self._id} due to error: {e}"
             )
             # Double check exp dir is unchanged
-            new_path = os.path.join(NEMORUN_HOME, "experiments", self._title, self._id)
+            new_path = os.path.join(get_nemorun_home(), "experiments", self._title, self._id)
             if self._exp_dir == new_path and new_path != old_exp_dir:
                 shutil.rmtree(self._exp_dir)
 

src/nemo_run/run/torchx_backend/schedulers/dgxcloud.py

@@ -7,6 +7,8 @@
 from pathlib import Path
 from typing import Any, Optional
 
+import fiddle as fdl
+import fiddle._src.experimental.dataclasses as fdl_dc
 from torchx.schedulers.api import (
     AppDryRunInfo,
     DescribeAppResponse,
@@ -22,16 +24,14 @@
     runopts,
 )
 
-import fiddle as fdl
-import fiddle._src.experimental.dataclasses as fdl_dc
-from nemo_run.config import NEMORUN_HOME
+from nemo_run.config import get_nemorun_home
 from nemo_run.core.execution.base import Executor
 from nemo_run.core.execution.dgxcloud import DGXCloudExecutor, DGXCloudState
 from nemo_run.core.serialization.zlib_json import ZlibJSONSerializer
 from nemo_run.run.torchx_backend.schedulers.api import SchedulerMixin
 
 # Local placeholder for storing DGX job states
-DGX_JOB_DIRS = os.path.join(NEMORUN_HOME, ".dgx_jobs.json")
+DGX_JOB_DIRS = os.path.join(get_nemorun_home(), ".dgx_jobs.json")
 
 # Example mapping from some DGX statuses to the TorchX AppState
 DGX_STATES: dict[DGXCloudState, AppState] = {

src/nemo_run/run/torchx_backend/schedulers/local.py

@@ -46,7 +46,7 @@
 )
 from torchx.specs.api import AppDef, AppState, Role
 
-from nemo_run.config import NEMORUN_HOME
+from nemo_run.config import get_nemorun_home
 from nemo_run.core.execution.base import Executor
 from nemo_run.core.execution.local import LocalExecutor
 from nemo_run.run.torchx_backend.schedulers.api import SchedulerMixin
@@ -59,7 +59,7 @@
     fcntl = None
     FCNTL_AVAILABLE = False
 
-LOCAL_JOB_DIRS = os.path.join(NEMORUN_HOME, ".local_jobs.json")
+LOCAL_JOB_DIRS = os.path.join(get_nemorun_home(), ".local_jobs.json")
 
 
 class PersistentLocalScheduler(SchedulerMixin, LocalScheduler):  # type: ignore

src/nemo_run/run/torchx_backend/schedulers/skypilot.py

@@ -42,7 +42,7 @@
     runopts,
 )
 
-from nemo_run.config import NEMORUN_HOME
+from nemo_run.config import get_nemorun_home
 from nemo_run.core.execution.base import Executor
 from nemo_run.core.execution.skypilot import _SKYPILOT_AVAILABLE, SkypilotExecutor
 from nemo_run.run.torchx_backend.schedulers.api import SchedulerMixin
@@ -74,7 +74,7 @@
     ...
 
 log: logging.Logger = logging.getLogger(__name__)
-SKYPILOT_JOB_DIRS = os.path.join(NEMORUN_HOME, ".skypilot_jobs.json")
+SKYPILOT_JOB_DIRS = os.path.join(get_nemorun_home(), ".skypilot_jobs.json")
 
 
 @dataclass