allow specifying lmndirs in .lmn.json5

lmn/cli/_config_loader.py

@@ -9,6 +9,7 @@
 from lmn.helpers import find_project_root, parse_config
 from posixpath import expandvars
 
+from lmn.config import LMNDirectories
 from lmn.machine import RemoteConfig
 
 DOCKER_ROOT_DIR = '/lmn'
@@ -52,12 +53,16 @@ class Machine:
     """Maintains machine configuration.
     - RemoteConfig (user, hostname, uri)
     """
-    def __init__(self, remote_conf: RemoteConfig, lmndir: Union[str, Path],
+    def __init__(self,
+                 remote_conf: RemoteConfig,
+                 lmndirs: Union[str, Path],
+                 container_lmndirs: Union[str, Path],
                  parsed_conf: dict,
                  startup: Union[str, List[str]] = "",
                  env: Optional[dict] = None) -> None:
         self.remote_conf = remote_conf
-        self.lmndir = Path(lmndir)
+        self.lmndirs = lmndirs
+        self.container_lmndirs = container_lmndirs
         self.env = env if env is not None else {}
         startup = startup or parsed_conf.startup
         if isinstance(startup, list):
@@ -74,14 +79,14 @@ def uri(self, path) -> str:
         """Returns user@hostname:path"""
         return f'{self.remote_conf.base_uri}:{path}'
 
-    def get_lmndirs(self, project_name: str) -> Namespace:
-        rootdir = self.lmndir / project_name
-        return Namespace(
-            codedir=str(rootdir / 'code'),
-            mountdir=str(rootdir / 'mount'),
-            outdir=str(rootdir / 'output'),
-            scriptdir=str(rootdir / 'script'),
-        )
+    # def get_lmndirs(self, project_name: str) -> Namespace:
+    #     rootdir = self.lmndir / project_name
+    #     return Namespace(
+    #         codedir=str(rootdir / 'code'),
+    #         mountdir=str(rootdir / 'mount'),
+    #         outdir=str(rootdir / 'output'),
+    #         scriptdir=str(rootdir / 'script'),
+    #     )
 
 
 def get_docker_lmndirs(lmndir: Union[Path, str], project_name: str) -> Namespace:
@@ -157,13 +162,40 @@ def load_config(machine_name: str):
 
     remote_conf = RemoteConfig(mconf.user, mconf.host)
 
+    # Setup lmndirs:
+    from lmn.cli._config_loader import DOCKER_ROOT_DIR
     if mconf.root_dir is None:
-        lmndir = f'{REMOTE_ROOT_DIR}/{remote_conf.user}/lmn'
+        lmn_rootdir = f'{REMOTE_ROOT_DIR}/{remote_conf.user}/lmn/{project.name}'
     else:
-        lmndir = f"{mconf.root_dir}/{remote_conf.user}"
+        lmn_rootdir = f'{mconf.root_dir}/{remote_conf.user}/{project.name}'
+    lmndirs = LMNDirectories(
+        codedir=str(Path(lmn_rootdir) / 'code'),
+        mountdir=str(Path(lmn_rootdir) / 'mount'),
+        outdir=str(Path(lmn_rootdir) / 'output'),
+        scriptdir=str(Path(lmn_rootdir) / 'script'),
+    )
+    container_lmndirs = LMNDirectories(
+        codedir=str(Path(DOCKER_ROOT_DIR) / 'code'),
+        mountdir=str(Path(DOCKER_ROOT_DIR) / 'mount'),
+        outdir=str(Path(DOCKER_ROOT_DIR) / 'output'),
+        scriptdir=str(Path(DOCKER_ROOT_DIR) / 'script'),
+    )
+
+    # Update lmndirs if the configurations exist in the mconf
+    remove_none_fn = lambda d: {k: v for k, v in d.items() if v is not None}
+    if mconf.lmndirs:
+        lmndirs = lmndirs.model_copy(
+            update=remove_none_fn(mconf.lmndirs.model_dump())
+        )
+    if mconf.container_lmndirs:
+        container_lmndirs = container_lmndirs.model_copy(
+            update=remove_none_fn(mconf.container_lmndirs.model_dump())
+        )
+
     machine = Machine(remote_conf,
                       parsed_conf=mconf,
-                      lmndir=lmndir,
+                      lmndirs=lmndirs,
+                      container_lmndirs=container_lmndirs,
                       env=mconf.environment)
 
     return project, machine, preset_conf

lmn/cli/brun.py

@@ -73,7 +73,7 @@ def handler(project: Project, machine: Machine, parsed: Namespace, preset: dict)
                                 cmd=cmd)
 
     env = {**project.env, **machine.env}
-    lmndirs = machine.get_lmndirs(project.name)
+    lmndirs = machine.lmndirs
 
     if parsed.with_startup:
         startup = ' ; '.join([e for e in [project.startup, machine.startup] if e.strip()])

lmn/cli/run.py

@@ -188,7 +188,7 @@ def handler(project: Project, machine: Machine, parsed: Namespace, preset: dict)
         from lmn.runner import SSHRunner
 
         startup = ' ; '.join([e for e in [project.startup, machine.startup] if e.strip()])
-        lmndirs = machine.get_lmndirs(project.name)
+        lmndirs = machine.lmndirs
         env = {**project.env, **machine.env}
 
         ssh_client = CLISSHClient(machine.remote_conf)
@@ -210,7 +210,7 @@ def handler(project: Project, machine: Machine, parsed: Namespace, preset: dict)
             raise ValueError('Configuration must have an entry for "docker" to use docker mode.')
 
         startup = ' ; '.join([e for e in [project.startup, machine.startup] if e.strip()])
-        lmndirs = machine.get_lmndirs(project.name)
+        lmndirs = machine.lmndirs
         env = {**project.env, **machine.env}
 
         base_url = "ssh://" + machine.base_uri
@@ -246,9 +246,9 @@ def handler(project: Project, machine: Machine, parsed: Namespace, preset: dict)
         if parsed.dry_run:
             raise ValueError('dry run is not yet supported for Docker mode')
 
-        from docker.types import Mount
-        from lmn.cli._config_loader import DOCKER_ROOT_DIR, get_docker_lmndirs
-        docker_lmndirs = get_docker_lmndirs(DOCKER_ROOT_DIR, project.name)
+        # from lmn.cli._config_loader import DOCKER_ROOT_DIR, get_docker_lmndirs
+        # docker_lmndirs = get_docker_lmndirs(DOCKER_ROOT_DIR, project.name)
+        docker_lmndirs = machine.container_lmndirs
 
         docker_pconf = machine.parsed_conf.docker
         docker_pconf.name = name
@@ -259,10 +259,10 @@ def handler(project: Project, machine: Machine, parsed: Namespace, preset: dict)
 
         if not runtime_options.no_sync:
             docker_pconf.mount_from_host.update({
-                f'{lmndirs.codedir}': docker_lmndirs.codedir,
-                f'{lmndirs.outdir}': docker_lmndirs.outdir,
-                f'{lmndirs.mountdir}': docker_lmndirs.mountdir,
-                f'{lmndirs.scriptdir}': docker_lmndirs.scriptdir,
+                f'{machine.lmndirs.codedir}': machine.container_lmndirs.codedir,
+                f'{machine.lmndirs.outdir}': machine.container_lmndirs.outdir,
+                f'{machine.lmndirs.mountdir}': machine.container_lmndirs.mountdir,
+                f'{machine.lmndirs.scriptdir}': machine.container_lmndirs.scriptdir,
             })
 
         if project.mount_from_host:
@@ -334,7 +334,7 @@ def handler_scheduler(
     import randomname
 
     startup = ' ; '.join([e for e in [project.startup, machine.startup] if e.strip()])
-    lmndirs = machine.get_lmndirs(project.name)
+    lmndirs = machine.lmndirs
 
     ssh_client = CLISSHClient(machine.remote_conf)
 
@@ -409,8 +409,9 @@ def handler_scheduler(
             raise ValueError('Entry "singularity" not found in the config file.')
 
         # Overwrite lmn envvars.  Hmm I don't like this...
-        from lmn.cli._config_loader import DOCKER_ROOT_DIR, get_docker_lmndirs
-        sing_lmndirs = get_docker_lmndirs(DOCKER_ROOT_DIR, project.name)
+        # from lmn.cli._config_loader import DOCKER_ROOT_DIR, get_docker_lmndirs
+        # sing_lmndirs = get_docker_lmndirs(DOCKER_ROOT_DIR, project.name)
+        sing_lmndirs = machine.container_lmndirs
 
         # Update `pwd` entry
         sing_conf = machine.parsed_conf.singularity

lmn/cli/sync.py

@@ -29,7 +29,7 @@ def _sync_code(project: Project, machine: Machine, dry_run: bool = False):
     # rsync_options = f"--rsync-path='mkdir -p {project.remote_dir} && mkdir -p {project.remote_outdir} && mkdir -p {project.remote_mountdir} && rsync'"
 
     # A trick to create directories right before performing rsync
-    lmndirs = machine.get_lmndirs(project.name)
+    lmndirs = machine.lmndirs
     rsync_options = [f"--rsync-path='mkdir -p {lmndirs.codedir} && mkdir -p {lmndirs.outdir} && mkdir -p {lmndirs.mountdir} && mkdir -p {lmndirs.scriptdir} && rsync'"]
 
     try:
@@ -51,7 +51,7 @@ def _sync_output(project: Project, machine: Machine, dry_run: bool = False):
     # Rsync remote outdir with the local outdir.
     if project.outdir:
         try:
-            lmndirs = machine.get_lmndirs(project.name)
+            lmndirs = machine.lmndirs
             # Check if there's any output file (the first line is always 'total [num-files]')
             ssh_client = CLISSHClient(machine.remote_conf)
 

lmn/config.py

@@ -25,10 +25,27 @@ class MachineConfig(BaseModel):
     startup: Union[str, List[str]] = ''
     mode: str = 'ssh'
 
+    # LMN Directories
+    lmndirs: Optional[OptionalLMNDirectories] = None
+    container_lmndirs: Optional[OptionalLMNDirectories] = None
+
     # Container config
     docker: Optional[DockerContainerConfig] = None
     singularity: Optional[SingularityConfig] = None
 
     # Scheduler config
     slurm: Optional[SlurmConfig] = None
-    pbs: Optional[PBSConfig] = None
+    pbs: Optional[PBSConfig] = None
+
+
+class LMNDirectories(BaseModel):
+    codedir: str
+    mountdir: str
+    outdir: str
+    scriptdir: str
+
+class OptionalLMNDirectories(BaseModel):
+    codedir: Optional[str] = None
+    mountdir: Optional[str] = None
+    outdir: Optional[str] = None
+    scriptdir: Optional[str] = None