xmanager-slurm 0.4.5__py3-none-any.whl → 0.4.7__py3-none-any.whl

xm_slurm/execution.py

@@ -5,7 +5,12 @@ import functools
 import hashlib
 import logging
 import operator
+import os
+import pathlib
+import re
 import shlex
+import shutil
+import subprocess
 import typing as tp
 
 import asyncssh
@@ -16,36 +21,18 @@ from asyncssh.auth import KbdIntPrompts, KbdIntResponse
 from asyncssh.misc import MaybeAwait
 from rich.console import ConsoleRenderable
 from rich.rule import Rule
+from rich.text import Text
 from xmanager import xm
 
-from xm_slurm import batching, config, constants, dependencies, executors, status
+from xm_slurm import batching, config, constants, dependencies, executors, status, utils
+from xm_slurm.config import ContainerRuntime, SlurmClusterConfig, SlurmSSHConfig
 from xm_slurm.console import console
+from xm_slurm.filesystem import AsyncFileSystem, AsyncLocalFileSystem, AsyncSSHFileSystem
 from xm_slurm.job_blocks import JobArgs
 from xm_slurm.types import Descriptor
 
-SlurmClusterConfig = config.SlurmClusterConfig
-ContainerRuntime = config.ContainerRuntime
-
 logger = logging.getLogger(__name__)
 
-"""
-=== Runtime Configurations ===
-With RunC:
-    skopeo copy --dest-creds=<username>:<secret> docker://<image>@<digest> oci:<image>:<digest>
-
-    pushd $SLURM_TMPDIR
-
-    umoci raw unpack --rootless --image <image>:<digest> bundle/<digest>
-    umoci raw runtime-config --image <image>:<digest> bundle/<digest>/config.json
-
-    runc run -b bundle/<digest> <container-id>
-
-With Singularity / Apptainer:
-
-    apptainer build --fix-perms --sandbox <digest> docker://<image>@<digest>
-    apptainer run --compat <digest>
-"""
-
 _POLL_INTERVAL = 30.0
 _BATCHED_BATCH_SIZE = 16
 _BATCHED_TIMEOUT = 0.2
@@ -54,19 +41,6 @@ _BATCHED_TIMEOUT = 0.2
 class SlurmExecutionError(Exception): ...
 
 
-class NoKBAuthSSHClient(asyncssh.SSHClient):
-    """SSHClient that does not prompt for keyboard-interactive authentication."""
-
-    def kbdint_auth_requested(self) -> MaybeAwait[str | None]:
-        return ""
-
-    def kbdint_challenge_received(
-        self, name: str, instructions: str, lang: str, prompts: KbdIntPrompts
-    ) -> MaybeAwait[KbdIntResponse | None]:
-        del name, instructions, lang, prompts
-        return []
-
-
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class SlurmJob:
     job_id: str
@@ -208,7 +182,7 @@ class _BatchedSlurmHandle:
         # Reconstruct the job states in the original order
         job_states = []
         for ssh_config, slurm_job in zip(ssh_configs, slurm_jobs):
-            job_states.append(job_states_by_cluster[ssh_config][slurm_job])
+            job_states.append(job_states_by_cluster[ssh_config][slurm_job])  # type: ignore
         return job_states
 
     @functools.partial(
@@ -263,52 +237,59 @@ class SlurmHandle(_BatchedSlurmHandle, tp.Generic[SlurmJobT]):
     async def logs(
         self, *, num_lines: int, block_size: int, wait: bool, follow: bool
     ) -> tp.AsyncGenerator[ConsoleRenderable, None]:
-        file = f".local/state/xm-slurm/{self.experiment_id}/slurm-{self.slurm_job.job_id}.out"
-        conn = await get_client().connection(self.ssh)
-        async with conn.start_sftp_client() as sftp:
-            if wait:
-                while not (await sftp.exists(file)):
-                    await asyncio.sleep(5)
-
-            async with sftp.open(file, "rb") as remote_file:
-                file_stat = await remote_file.stat()
-                file_size = file_stat.size
-                assert file_size is not None
-
-                data = b""
-                lines = []
-                position = file_size
-
-                while len(lines) <= num_lines and position > 0:
-                    read_size = min(block_size, position)
-                    position -= read_size
-                    await remote_file.seek(position)
-                    chunk = await remote_file.read(read_size)
-                    data = chunk + data
-                    lines = data.splitlines()
-
-                if position <= 0:
-                    yield Rule("[bold red]BEGINNING OF FILE[/bold red]")
-                for line in lines[-num_lines:]:
-                    yield line.decode("utf-8", errors="replace")
-
-                if (await self.get_state()) not in status.SlurmActiveJobStates:
-                    yield Rule("[bold red]END OF FILE[/bold red]")
-                    return
-
-                if not follow:
-                    return
-
-                await remote_file.seek(file_size)
-                while True:
-                    if new_data := (await remote_file.read(block_size)):
-                        yield new_data.decode("utf-8", errors="replace")
-                    else:
-                        await asyncio.sleep(0.25)
+        experiment_dir = await get_client().experiment_dir(self.ssh, self.experiment_id)
+        file = experiment_dir / f"slurm-{self.slurm_job.job_id}.out"
+
+        fs = await get_client().fs(self.ssh)
+
+        if wait:
+            while not (await fs.exists(file)):
+                await asyncio.sleep(5)
+
+        file_size = await fs.size(file)
+        assert file_size is not None
+
+        async with await fs.open(file, "rb") as remote_file:  # type: ignore
+            data = b""
+            lines = []
+            position = file_size
+
+            while len(lines) <= num_lines and position > 0:
+                read_size = min(block_size, position)
+                position -= read_size
+                await remote_file.seek(position)  # type: ignore
+                chunk = await remote_file.read(read_size)
+                data = chunk + data
+                lines = data.splitlines()
+
+            if position <= 0:
+                yield Rule("[bold red]BEGINNING OF FILE[/bold red]")
+            for line in lines[-num_lines:]:
+                yield Text.from_ansi(line.decode("utf-8", errors="replace"))
+
+            if (await self.get_state()) not in status.SlurmActiveJobStates:
+                yield Rule("[bold red]END OF FILE[/bold red]")
+                return
+
+            if not follow:
+                return
+
+            await remote_file.seek(file_size)  # type: ignore
+            while True:
+                if new_data := (await remote_file.read(block_size)):
+                    yield Text.from_ansi(new_data.decode("utf-8", errors="replace"))
+                else:
+                    await asyncio.sleep(0.25)
+
+
+class CompletedProcess(tp.Protocol):
+    returncode: int | None
+    stdout: bytes | str
+    stderr: bytes | str
 
 
 @functools.cache
-def get_template_env(container_runtime: ContainerRuntime) -> j2.Environment:
+def get_template_env(runtime: ContainerRuntime) -> j2.Environment:
     template_loader = j2.PackageLoader("xm_slurm", "templates/slurm")
     template_env = j2.Environment(loader=template_loader, trim_blocks=True, lstrip_blocks=False)
 
@@ -318,44 +299,222 @@ def get_template_env(container_runtime: ContainerRuntime) -> j2.Environment:
     template_env.globals["raise"] = _raise_template_exception
     template_env.globals["operator"] = operator
 
-    match container_runtime:
+    entrypoint_template = template_env.get_template("entrypoint.bash.j2")
+    template_env.globals.update(entrypoint_template.module.__dict__)
+
+    match runtime:
         case ContainerRuntime.SINGULARITY | ContainerRuntime.APPTAINER:
             runtime_template = template_env.get_template("runtimes/apptainer.bash.j2")
         case ContainerRuntime.PODMAN:
             runtime_template = template_env.get_template("runtimes/podman.bash.j2")
         case _:
-            raise NotImplementedError(f"Container runtime {container_runtime} is not implemented.")
-    # Update our global env with the runtime template's exported globals
+            raise NotImplementedError(f"Container runtime {runtime} is not implemented.")
     template_env.globals.update(runtime_template.module.__dict__)
 
     return template_env
 
 
-@functools.cache
-def get_client() -> "Client":
-    return Client()
+class NoKBAuthSSHClient(asyncssh.SSHClient):
+    """SSHClient that does not prompt for keyboard-interactive authentication."""
 
+    def kbdint_auth_requested(self) -> MaybeAwait[str | None]:
+        return ""
 
-class Client:
-    def __init__(self) -> None:
-        self._connections = dict[config.SlurmSSHConfig, asyncssh.SSHClientConnection]()
-        self._connection_lock = asyncio.Lock()
+    def kbdint_challenge_received(
+        self, name: str, instructions: str, lang: str, prompts: KbdIntPrompts
+    ) -> MaybeAwait[KbdIntResponse | None]:
+        del name, instructions, lang, prompts
+        return []
+
+
+class SlurmExecutionClient:
+    def __init__(self):
+        self._remote_connections = dict[config.SlurmSSHConfig, asyncssh.SSHClientConnection]()
+        self._remote_filesystems = dict[config.SlurmSSHConfig, AsyncSSHFileSystem]()
+        self._remote_connection_lock = asyncio.Lock()
+
+        self._local_fs = AsyncLocalFileSystem()
+
+    @backoff.on_exception(backoff.expo, asyncio.exceptions.TimeoutError, max_tries=5, max_time=60.0)
+    async def _local_run(  # type: ignore
+        self,
+        command: str,
+        *,
+        check: bool = False,
+        timeout: float | None = None,
+    ) -> subprocess.CompletedProcess[str]:
+        process = await asyncio.subprocess.create_subprocess_shell(
+            command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            # Filter out all SLURM_ environment variables as this could be running on a
+            # compute node and xm-slurm should act stateless.
+            env=dict(filter(lambda x: not x[0].startswith("SLURM_"), os.environ.items())),
+        )
+        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout)
+
+        stdout = stdout.decode("utf-8").strip() if stdout else ""
+        stderr = stderr.decode("utf-8").strip() if stderr else ""
+
+        assert process.returncode is not None
+        if check and process.returncode != 0:
+            raise RuntimeError(f"Command failed with return code {process.returncode}: {command}\n")
+
+        return subprocess.CompletedProcess[str](command, process.returncode, stdout, stderr)
 
     @backoff.on_exception(backoff.expo, asyncssh.Error, max_tries=5, max_time=60.0)
-    async def _setup_remote_connection(self, conn: asyncssh.SSHClientConnection) -> None:
-        # Make sure the xm-slurm state directory exists
-        async with conn.start_sftp_client() as sftp_client:
-            await sftp_client.makedirs(".local/state/xm-slurm", exist_ok=True)
-
-    async def connection(self, ssh_config: config.SlurmSSHConfig) -> asyncssh.SSHClientConnection:
-        if ssh_config not in self._connections:
-            async with self._connection_lock:
+    async def _remote_run(  # type: ignore
+        self,
+        ssh_config: config.SlurmSSHConfig,
+        command: str,
+        *,
+        check: bool = False,
+        timeout: float | None = None,
+    ) -> asyncssh.SSHCompletedProcess:
+        client = await self._connection(ssh_config)
+        return await client.run(command, check=check, timeout=timeout)
+
+    @functools.cache
+    def _is_ssh_config_local(self, ssh_config: SlurmSSHConfig) -> bool:
+        """A best effort check to see if the SSH config is local so we can bypass ssh."""
+
+        # We can't verify the connection so bail out
+        if ssh_config.host_public_key is None:
+            return False
+        if "SSH_CONNECTION" not in os.environ:
+            return False
+
+        def _is_host_local(host: str) -> bool:
+            nonlocal ssh_config
+            assert ssh_config.host_public_key is not None
+
+            if shutil.which("ssh-keyscan") is None:
+                return False
+
+            keyscan_result = utils.run_command(
+                ["ssh-keyscan", "-t", ssh_config.host_public_key.algorithm, host],
+                return_stdout=True,
+            )
+
+            if keyscan_result.returncode != 0:
+                return False
+
+            try:
+                key = mit.one(
+                    filter(
+                        lambda x: not x.startswith("#"), keyscan_result.stdout.strip().split("\n")
+                    )
+                )
+                _, algorithm, key = key.split(" ")
+
+                if (
+                    algorithm == ssh_config.host_public_key.algorithm
+                    and key == ssh_config.host_public_key.key
+                ):
+                    return True
+
+            except Exception:
+                pass
+
+            return False
+
+        # 1): we're directly connected to the host
+        ssh_connection_str = os.environ["SSH_CONNECTION"]
+        _, _, server_ip, _ = ssh_connection_str.split()
+
+        logger.debug("Checking if SSH_CONNECTION server %s is local", server_ip)
+        if _is_host_local(server_ip):
+            return True
+
+        # 2): we're in a Slurm job and the submission host is the host
+        if "SLURM_JOB_ID" in os.environ and "SLURM_SUBMIT_HOST" in os.environ:
+            submit_host = os.environ["SLURM_SUBMIT_HOST"]
+            logger.debug("Checking if SLURM_SUBMIT_HOST %s is local", submit_host)
+            if _is_host_local(submit_host):
+                return True
+        elif "SLURM_JOB_ID" in os.environ and shutil.which("scontrol") is not None:
+            # Stupid edge case where if you run srun SLURM_SUBMIT_HOST isn't forwarded
+            # so we'll parse it from scontrol...
+            scontrol_result = utils.run_command(
+                ["scontrol", "show", "job", os.environ["SLURM_JOB_ID"]],
+                return_stdout=True,
+            )
+            if scontrol_result.returncode != 0:
+                return False
+
+            match = re.search(
+                r"AllocNode:Sid=(?P<host>[^ ]+):\d+", scontrol_result.stdout.strip(), re.MULTILINE
+            )
+            if match is not None:
+                host = match.group("host")
+                logger.debug("Checking if AllocNode %s is local", host)
+                if _is_host_local(host):
+                    return True
+
+        return False
+
+    @functools.cache
+    async def _state_dir(self, ssh_config: SlurmSSHConfig) -> pathlib.Path:
+        state_dirs = [
+            ("XM_SLURM_STATE_DIR", ""),
+            ("XDG_STATE_HOME", "xm-slurm"),
+            ("HOME", ".local/state/xm-slurm"),
+        ]
+
+        for env_var, subpath in state_dirs:
+            cmd = await self.run(ssh_config, f"printenv {env_var}", check=False)
+            assert isinstance(cmd.stdout, str)
+            if cmd.returncode == 0:
+                return pathlib.Path(cmd.stdout.strip()) / subpath
+
+        raise SlurmExecutionError(
+            "Failed to find a valid state directory for XManager. "
+            "We weren't able to resolve any of the following paths: "
+            f"{', '.join(env_var + ('/' + subpath if subpath else '') for env_var, subpath in state_dirs)}."
+        )
+
+    async def experiment_dir(self, ssh_config: SlurmSSHConfig, experiment_id: int) -> pathlib.Path:
+        return (await self._state_dir(ssh_config)) / f"{experiment_id:08d}"
+
+    async def run(
+        self,
+        ssh_config: SlurmSSHConfig,
+        command: xm.SequentialArgs | str | tp.Sequence[str],
+        *,
+        check: bool = False,
+        timeout: float | None = None,
+    ) -> CompletedProcess:
+        if isinstance(command, xm.SequentialArgs):
+            command = command.to_list()
+        if not isinstance(command, str) and isinstance(command, collections.abc.Sequence):
+            command = shlex.join(command)
+        assert isinstance(command, str)
+
+        if self._is_ssh_config_local(ssh_config):
+            logger.debug("Running command locally: %s", command)
+            return await self._local_run(command, check=check, timeout=timeout)  # type: ignore
+        else:
+            logger.debug("Running command on %s: %s", ssh_config.host, command)
+            return await self._remote_run(ssh_config, command, check=check, timeout=timeout)  # type: ignore
+
+    async def fs(self, ssh_config: SlurmSSHConfig) -> AsyncFileSystem:
+        if self._is_ssh_config_local(ssh_config):
+            return self._local_fs
+
+        if ssh_config not in self._remote_filesystems:
+            self._remote_filesystems[ssh_config] = AsyncSSHFileSystem(
+                await (await self._connection(ssh_config)).start_sftp_client()
+            )
+        return self._remote_filesystems[ssh_config]
+
+    async def _connection(self, ssh_config: config.SlurmSSHConfig) -> asyncssh.SSHClientConnection:
+        if ssh_config not in self._remote_connections:
+            async with self._remote_connection_lock:
                 try:
                     conn, _ = await asyncssh.create_connection(
                         NoKBAuthSSHClient, options=ssh_config.connection_options
                     )
-                    await self._setup_remote_connection(conn)
-                    self._connections[ssh_config] = conn
+                    self._remote_connections[ssh_config] = conn
                 except asyncssh.misc.PermissionDenied as ex:
                     raise SlurmExecutionError(
                         f"Permission denied connecting to {ssh_config.host}"
@@ -375,28 +534,9 @@ class Client:
                         f"SSH connection error when connecting to {ssh_config.host}"
                     ) from ex
 
-        return self._connections[ssh_config]
-
-    @backoff.on_exception(backoff.expo, asyncssh.Error, max_tries=5, max_time=60.0)
-    async def run(
-        self,
-        ssh_config: config.SlurmSSHConfig,
-        command: xm.SequentialArgs | str | tp.Sequence[str],
-        *,
-        check: bool = False,
-        timeout: float | None = None,
-    ) -> asyncssh.SSHCompletedProcess:
-        client = await self.connection(ssh_config)
-        if isinstance(command, xm.SequentialArgs):
-            command = command.to_list()
-        if not isinstance(command, str) and isinstance(command, collections.abc.Sequence):
-            command = shlex.join(command)
-        assert isinstance(command, str)
-        logger.debug("Running command on %s: %s", ssh_config.host, command)
-
-        return await client.run(command, check=check, timeout=timeout)
+        return self._remote_connections[ssh_config]
 
-    async def template(
+    async def _submission_script_template(
         self,
         *,
         job: xm.Job | xm.JobGroup,
@@ -410,6 +550,12 @@ class Client:
             args = {}
 
         template_env = get_template_env(cluster.runtime)
+        template_context = dict(
+            dependency=dependency,
+            cluster=cluster,
+            experiment_id=experiment_id,
+            identity=identity,
+        )
 
         # Sanitize job groups
         if isinstance(job, xm.JobGroup) and len(job.jobs) == 1:
@@ -430,26 +576,14 @@ class Client:
                     )
 
                 return template.render(
-                    job=job_array,
-                    dependency=dependency,
-                    cluster=cluster,
-                    args=sequential_args,
-                    env_vars=env_vars,
-                    experiment_id=experiment_id,
-                    identity=identity,
+                    job=job_array, args=sequential_args, env_vars=env_vars, **template_context
                 )
             case xm.Job() if isinstance(args, collections.abc.Mapping):
                 template = template_env.get_template("job.bash.j2")
                 sequential_args = xm.SequentialArgs.from_collection(args.get("args", None))
                 env_vars = args.get("env_vars", None)
                 return template.render(
-                    job=job,
-                    dependency=dependency,
-                    cluster=cluster,
-                    args=sequential_args,
-                    env_vars=env_vars,
-                    experiment_id=experiment_id,
-                    identity=identity,
+                    job=job, args=sequential_args, env_vars=env_vars, **template_context
                 )
             case xm.JobGroup() as job_group if isinstance(args, collections.abc.Mapping):
                 template = template_env.get_template("job-group.bash.j2")
@@ -464,13 +598,7 @@ class Client:
                     for job_name in job_group.jobs.keys()
                 }
                 return template.render(
-                    job_group=job_group,
-                    dependency=dependency,
-                    cluster=cluster,
-                    args=sequential_args,
-                    env_vars=env_vars,
-                    experiment_id=experiment_id,
-                    identity=identity,
+                    job_group=job_group, args=sequential_args, env_vars=env_vars, **template_context
                 )
             case _:
                 raise ValueError(f"Unsupported job type: {type(job)}")
@@ -521,7 +649,7 @@ class Client:
         experiment_id: int,
         identity: str | None = None,
     ):
-        template = await self.template(
+        submission_script = await self._submission_script_template(
             job=job,
             dependency=dependency,
             cluster=cluster,
@@ -529,24 +657,19 @@ class Client:
             experiment_id=experiment_id,
             identity=identity,
         )
-        logger.debug("Slurm submission script:\n%s", template)
-
-        # Hash submission script
-        template_hash = hashlib.blake2s(template.encode()).hexdigest()[:8]
-
-        conn = await self.connection(cluster.ssh)
-        async with conn.start_sftp_client() as sftp:
-            # Write the submission script to the cluster
-            # TODO(jfarebro): SHOULD FIND A WAY TO GET THE HOME DIRECTORY
-            # INSTEAD OF ASSUMING SFTP PUTS US IN THE HOME DIRECTORY
-            await sftp.makedirs(f".local/state/xm-slurm/{experiment_id}", exist_ok=True)
-            async with sftp.open(
-                f".local/state/xm-slurm/{experiment_id}/submission-script-{template_hash}.sh", "w"
-            ) as fp:
-                await fp.write(template)
+        logger.debug("Slurm submission script:\n%s", submission_script)
+        submission_script_hash = hashlib.blake2s(submission_script.encode()).hexdigest()[:8]
+        submission_script_path = f"submission-script-{submission_script_hash}.sh"
+
+        fs = await self.fs(cluster.ssh)
+
+        template_dir = await self.experiment_dir(cluster.ssh, experiment_id)
+
+        await fs.makedirs(template_dir, exist_ok=True)
+        await fs.write(template_dir / submission_script_path, submission_script.encode())
 
         # Construct and run command on the cluster
-        command = f"sbatch --chdir .local/state/xm-slurm/{experiment_id} --parsable submission-script-{template_hash}.sh"
+        command = f"sbatch --chdir {template_dir.as_posix()} --parsable {submission_script_path}"
         result = await self.run(cluster.ssh, command)
         if result.returncode != 0:
             raise RuntimeError(f"Failed to schedule job on {cluster.ssh.host}: {result.stderr}")
@@ -596,8 +719,16 @@ class Client:
             raise ValueError(f"Unsupported job type: {type(job)}")
 
     def __del__(self):
-        for conn in self._connections.values():
+        for fs in self._remote_filesystems.values():
+            del fs
+        for conn in self._remote_connections.values():
             conn.close()
+            del conn
+
+
+@functools.cache
+def get_client() -> SlurmExecutionClient:
+    return SlurmExecutionClient()
 
 
 @tp.overload