xmanager-slurm 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

xm_slurm/experimental/parameter_controller.py

@@ -0,0 +1,200 @@
+import asyncio
+import base64
+import dataclasses
+import enum
+import functools
+import logging
+import zlib
+from typing import Awaitable, Callable, Concatenate, Coroutine, Mapping, ParamSpec, TypeVar
+
+import backoff
+import cloudpickle
+from xmanager import xm
+
+import xm_slurm
+from xm_slurm import job_blocks, status
+from xm_slurm.experiment import SlurmAuxiliaryUnit, SlurmExperiment
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+async def _monitor_parameter_controller(
+    aux_unit: SlurmAuxiliaryUnit,
+    local_parameter_controller_coro: Coroutine[None, None, T],
+    *,
+    poll_interval: float = 30.0,
+) -> None:
+    local_controller_finished = asyncio.Event()
+    local_parameter_controller = asyncio.create_task(local_parameter_controller_coro)
+
+    @local_parameter_controller.add_done_callback
+    def _(future: asyncio.Task[T]) -> None:
+        try:
+            _ = future.result()
+        except asyncio.CancelledError:
+            logging.info("Local parameter controller was cancelled, resuming on remote controller.")
+            pass
+        except Exception:
+            logging.error("Local parameter controller failed, stopping remote controller.")
+            aux_unit.stop(
+                mark_as_failed=True, mark_as_completed=False, message="Local controller failed."
+            )
+            raise
+        else:
+            logging.info(
+                "Local parameter controller finished before remote controller started, "
+                "stopping remote controller."
+            )
+            local_controller_finished.set()
+            aux_unit.stop(mark_as_completed=True, message="Local parameter controller finished.")
+
+    @backoff.on_predicate(
+        backoff.constant,
+        lambda aux_unit_status: aux_unit_status is status.SlurmWorkUnitStatusEnum.PENDING,
+        jitter=None,
+        interval=poll_interval,
+    )
+    async def wait_for_remote_controller() -> status.SlurmWorkUnitStatusEnum:
+        logging.info("Waiting for remote parameter controller to start.")
+        if local_controller_finished.is_set():
+            return status.SlurmWorkUnitStatusEnum.COMPLETED
+        return (await aux_unit.get_status()).status
+
+    logging.info("Monitoring remote parameter controller.")
+    # TODO(jfarebro): make get_status() more resiliant to errors when initially scheduling.
+    # We run into issues if we call get_status() too quickly when Slurm hasn't ingested the job.
+    await asyncio.sleep(15)
+    match await wait_for_remote_controller():
+        case status.SlurmWorkUnitStatusEnum.RUNNING:
+            logging.info("Remote parameter controller started.")
+            local_parameter_controller.cancel("Remote parameter controller started.")
+        case status.SlurmWorkUnitStatusEnum.COMPLETED:
+            if local_parameter_controller.done():
+                logging.info("Local parameter controller finished, stopping remote controller.")
+                aux_unit.stop(
+                    mark_as_completed=True, message="Local parameter controller finished."
+                )
+            else:
+                logging.info("Remote parameter controller finished, stopping local controller.")
+                local_parameter_controller.cancel()
+        case status.SlurmWorkUnitStatusEnum.FAILED:
+            logging.error("Remote parameter controller failed, stopping local controller.")
+            local_parameter_controller.cancel()
+        case status.SlurmWorkUnitStatusEnum.CANCELLED:
+            logging.info("Remote parameter controller was cancelled, stopping local controller.")
+            local_parameter_controller.cancel()
+        case status.SlurmWorkUnitStatusEnum.PENDING:
+            raise RuntimeError("Remote parameter controller is still pending, invalid state.")
+
+
+class ParameterControllerMode(enum.Enum):
+    AUTO = enum.auto()
+    REMOTE_ONLY = enum.auto()
+    # TODO(jfarebro): is it possible to get LOCAL_ONLY?
+    # We'd need to have a dummy job type as we need to return a JobType?
+
+
+def parameter_controller(
+    *,
+    executable: xm.Executable,
+    executor: xm.Executor,
+    controller_mode: ParameterControllerMode = ParameterControllerMode.AUTO,
+    controller_name: str = "parameter_controller",
+    controller_args: xm.UserArgs | None = None,
+    controller_env_vars: Mapping[str, str] | None = None,
+) -> Callable[
+    [
+        Callable[Concatenate[SlurmExperiment, P], T]
+        | Callable[Concatenate[SlurmExperiment, P], Awaitable[T]],
+    ],
+    Callable[P, xm.AuxiliaryUnitJob],
+]:
+    """Converts a function to a controller which can be added to an experiment.
+
+    Calling the wrapped function would return an xm.JobGenerator which would run
+    it as auxiliary unit on the specified executor.
+
+    Args:
+        executable: An executable that has a Python entrypoint with all the necesarry dependencies.
+        executor: The executor to launch the controller job on.
+        controller_name: Name of the parameter controller job.
+        controller_args: Mapping of flag names and values to be used by the XM
+          client running inside the parameter controller job.
+        controller_env_vars: Mapping of env variable names and values to be passed
+          to the parameter controller job.
+
+    Returns:
+        A decorator to be applied to the function.
+    """
+
+    def decorator(
+        f: Callable[Concatenate[SlurmExperiment, P], T]
+        | Callable[Concatenate[SlurmExperiment, P], Awaitable[T]],
+    ) -> Callable[P, xm.AuxiliaryUnitJob]:
+        @functools.wraps(f)
+        def make_controller(*args: P.args, **kwargs: P.kwargs) -> xm.AuxiliaryUnitJob:
+            # Modify the function to read the experiment from the API so that it can be pickled.
+
+            async def job_generator(aux_unit: SlurmAuxiliaryUnit) -> None:
+                experiment_id = aux_unit.experiment.experiment_id
+
+                async def local_controller(*args: P.args, **kwargs: P.kwargs) -> T | Awaitable[T]:
+                    if asyncio.iscoroutinefunction(f):
+                        return await f(aux_unit.experiment, *args, **kwargs)
+                    else:
+                        return f(aux_unit.experiment, *args, **kwargs)
+
+                async def remote_controller(*args: P.args, **kwargs: P.kwargs) -> T | Awaitable[T]:
+                    async with xm_slurm.get_experiment(experiment_id=experiment_id) as exp:
+                        if asyncio.iscoroutinefunction(f):
+                            return await f(exp, *args, **kwargs)
+                        else:
+                            return f(exp, *args, **kwargs)
+
+                remote_controller_serialized = base64.urlsafe_b64encode(
+                    zlib.compress(
+                        cloudpickle.dumps(
+                            functools.partial(remote_controller, *args, **kwargs),
+                        )
+                    )
+                )
+
+                parameter_controller_executable = dataclasses.replace(
+                    executable,
+                    args=xm.merge_args(
+                        job_blocks.get_args_for_python_entrypoint(
+                            xm.ModuleName("xm_slurm.scripts._cloudpickle")
+                        ),
+                        xm.SequentialArgs.from_collection({
+                            "cloudpickled_fn": remote_controller_serialized.decode("ascii"),
+                        }),
+                        xm.SequentialArgs.from_collection(controller_args),
+                    ),
+                    env_vars=controller_env_vars or {},
+                )
+
+                await aux_unit.add(
+                    xm.Job(
+                        executor=executor,
+                        executable=parameter_controller_executable,
+                        name=controller_name,
+                    )
+                )
+
+                # Launch local parameter controller and monitor for when it starts running
+                # so we can kill the local controller.
+                if controller_mode is ParameterControllerMode.AUTO:
+                    aux_unit._create_task(
+                        _monitor_parameter_controller(aux_unit, local_controller(*args, **kwargs))
+                    )
+
+            return xm.AuxiliaryUnitJob(
+                job_generator,
+                importance=xm.Importance.HIGH,
+                termination_delay_secs=0,  # TODO: add support for termination delay.?
+            )
+
+        return make_controller
+
+    return decorator

xm_slurm/job_blocks.py

@@ -1,6 +1,13 @@
+from typing import Mapping, TypedDict
+
 from xmanager import xm
 
 
+class JobArgs(TypedDict, total=False):
+    args: xm.UserArgs
+    env_vars: Mapping[str, str]
+
+
 def get_args_for_python_entrypoint(
     entrypoint: xm.ModuleName | xm.CommandList,
 ) -> xm.SequentialArgs:

xm_slurm/packageables.py

@@ -1,7 +1,7 @@
 import importlib.resources as resources
 import pathlib
 import sys
-from typing import Mapping, Sequence
+from typing import Literal, Mapping, Sequence
 
 import immutabledict
 from xmanager import xm
@@ -29,8 +29,8 @@ def docker_image(
     return xm.Packageable(
         executor_spec=SlurmSpec(),
         executable_spec=DockerImage(image=image),
-        args=args,
-        env_vars=env_vars,
+        args=xm.SequentialArgs.from_collection(args),
+        env_vars=dict(env_vars),
     )
 
 
@@ -40,7 +40,7 @@ def docker_container(
     dockerfile: pathlib.Path | None = None,
     context: pathlib.Path | None = None,
     target: str | None = None,
-    ssh: list[str] | None = None,
+    ssh: Sequence[str] | Literal[True] | None = None,
     build_args: Mapping[str, str] = immutabledict.immutabledict(),
     cache_from: str | Sequence[str] | None = None,
     labels: Mapping[str, str] = immutabledict.immutabledict(),
@@ -55,7 +55,7 @@ def docker_container(
         dockerfile: The path to the dockerfile.
         context: The path to the docker context.
         target: The docker build target.
-        ssh: A list of SSH sockets/keys for the docker build step.
+        ssh: A list of SSH sockets/keys for the docker build step or `True` to use the default SSH agent.
         build_args: Build arguments to docker.
         cache_from: Where to pull the BuildKit cache from. See `--cache-from` in `docker build`.
         labels: The container labels.
@@ -71,8 +71,12 @@ def docker_container(
     if dockerfile is None:
         dockerfile = context / "Dockerfile"
     dockerfile = dockerfile.resolve()
+
     if ssh is None:
         ssh = []
+    elif ssh is True:
+        ssh = ["default"]
+
     if cache_from is None and isinstance(executor_spec, SlurmSpec):
         cache_from = executor_spec.tag
     if cache_from is None:
@@ -92,8 +96,8 @@ def docker_container(
             workdir=workdir,
             labels=labels,
         ),
-        args=args,
-        env_vars=env_vars,
+        args=xm.SequentialArgs.from_collection(args),
+        env_vars=dict(env_vars),
     )
 
 
@@ -104,14 +108,18 @@ def python_container(
     context: pathlib.Path | None = None,
     requirements: pathlib.Path | None = None,
     base_image: str = "docker.io/python:{major}.{minor}-slim",
+    extra_system_packages: Sequence[str] = (),
+    extra_python_packages: Sequence[str] = (),
     cache_from: str | Sequence[str] | None = None,
     labels: Mapping[str, str] = immutabledict.immutabledict(),
+    ssh: Sequence[str] | Literal[True] | None = None,
     args: xm.UserArgs | None = None,
     env_vars: Mapping[str, str] = immutabledict.immutabledict(),
 ) -> xm.Packageable:
     """Creates a Python container from a base image using pip from a `requirements.txt` file.
 
     NOTE: The base image will use the Python version of the current interpreter.
+    NOTE: uv is used to install packages from `requirements`.
 
     Args:
         executor_spec: The executor specification, where will the container be stored at.
@@ -119,8 +127,11 @@ def python_container(
         context: The path to the docker context.
         requirements: The path to the pip requirements file.
         base_image: The base image to use. NOTE: The base image must contain the Python runtime.
+        extra_system_packages: Additional system packages to install. NOTE: These are installed via `apt-get`.
+        extra_python_packages: Additional Python packages to install. NOTE: These are installed via `uv pip`.
         cache_from: Where to pull the BuildKit cache from. See `--cache-from` in `docker build`.
         labels: The container labels.
+        ssh: A list of SSH sockets/keys for the docker build step or `True` to use the default SSH agent.
         args: The user arguments to pass to the executable.
         env_vars: The environment variables to pass to the executable.
 
@@ -149,11 +160,11 @@ def python_container(
             executor_spec=executor_spec,
             dockerfile=dockerfile,
             context=context,
+            ssh=ssh,
             build_args={
                 "PIP_REQUIREMENTS": requirements.relative_to(context).as_posix(),
-                "PYTHON_MAJOR": str(sys.version_info.major),
-                "PYTHON_MINOR": str(sys.version_info.minor),
-                "PYTHON_MICRO": str(sys.version_info.micro),
+                "EXTRA_SYSTEM_PACKAGES": " ".join(extra_system_packages),
+                "EXTRA_PYTHON_PACKAGES": " ".join(extra_python_packages),
                 "BASE_IMAGE": base_image.format_map({
                     "major": sys.version_info.major,
                     "minor": sys.version_info.minor,
@@ -178,6 +189,7 @@ def mamba_container(
     base_image: str = "gcr.io/distroless/base-debian10",
     cache_from: str | Sequence[str] | None = None,
     labels: Mapping[str, str] = immutabledict.immutabledict(),
+    ssh: Sequence[str] | Literal[True] | None = None,
     args: xm.UserArgs | None = None,
     env_vars: Mapping[str, str] = immutabledict.immutabledict(),
 ) -> xm.Packageable:
@@ -193,6 +205,7 @@ def mamba_container(
         base_image: The base image to use.
         cache_from: Where to pull the BuildKit cache from. See `--cache-from` in `docker build`.
         labels: The container labels.
+        ssh: A list of SSH sockets/keys for the docker build step or `True` to use the default SSH agent.
         args: The user arguments to pass to the executable.
         env_vars: The environment variables to pass to the executable.
 
@@ -221,6 +234,7 @@ def mamba_container(
             executor_spec=executor_spec,
             dockerfile=dockerfile,
             context=context,
+            ssh=ssh,
             build_args={
                 "CONDA_ENVIRONMENT": environment.relative_to(context).as_posix(),
                 "BASE_IMAGE": base_image,
@@ -237,26 +251,32 @@ def mamba_container(
 conda_container = mamba_container
 
 
-def pdm_container(
+def uv_container(
     *,
     executor_spec: xm.ExecutorSpec,
     entrypoint: xm.ModuleName | xm.CommandList,
     context: pathlib.Path | None = None,
-    base_image: str = "docker.io/python:{major}.{minor}-slim",
+    base_image: str = "docker.io/python:{major}.{minor}-slim-bookworm",
+    extra_system_packages: Sequence[str] = (),
+    extra_python_packages: Sequence[str] = (),
     cache_from: str | Sequence[str] | None = None,
     labels: Mapping[str, str] = immutabledict.immutabledict(),
+    ssh: Sequence[str] | Literal[True] | None = None,
     args: xm.UserArgs | None = None,
     env_vars: Mapping[str, str] = immutabledict.immutabledict(),
 ) -> xm.Packageable:
-    """Creates a Python container from a base image using pdm from a `pdm.lock` file.
+    """Creates a Python container from a base image using uv from a `uv.lock` file.
 
     Args:
         executor_spec: The executor specification, where will the container be stored at.
         entrypoint: The entrypoint to run in the container.
         context: The path to the docker context.
         base_image: The base image to use. NOTE: The base image must contain the Python runtime.
+        extra_system_packages: Additional system packages to install. NOTE: These are installed via `apt-get`.
+        extra_python_packages: Additional Python packages to install. NOTE: These are installed via `uv pip`.
         cache_from: Where to pull the BuildKit cache from. See `--cache-from` in `docker build`.
         labels: The container labels.
+        ssh: A list of SSH sockets/keys for the docker build step or `True` to use the default SSH agent.
         args: The user arguments to pass to the executable.
         env_vars: The environment variables to pass to the executable.
 
@@ -268,20 +288,22 @@ def pdm_container(
     if context is None:
         context = utils.find_project_root()
     context = context.resolve()
-    if not (context / "pdm.lock").exists():
-        raise ValueError(f"PDM lockfile `{context / 'pdm.lock'}` doesn't exist.")
+    if not (context / "pyproject.toml").exists():
+        raise ValueError(f"Python project file `{context / 'pyproject.toml'}` doesn't exist.")
+    if not (context / "uv.lock").exists():
+        raise ValueError(f"UV lock file `{context / 'uv.lock'}` doesn't exist.")
 
     with resources.as_file(
-        resources.files("xm_slurm.templates").joinpath("docker/pdm.Dockerfile")
+        resources.files("xm_slurm.templates").joinpath("docker/uv.Dockerfile")
     ) as dockerfile:
         return docker_container(
             executor_spec=executor_spec,
             dockerfile=dockerfile,
             context=context,
+            ssh=ssh,
             build_args={
-                "PYTHON_MAJOR": str(sys.version_info.major),
-                "PYTHON_MINOR": str(sys.version_info.minor),
-                "PYTHON_MICRO": str(sys.version_info.micro),
+                "EXTRA_SYSTEM_PACKAGES": " ".join(extra_system_packages),
+                "EXTRA_PYTHON_PACKAGES": " ".join(extra_python_packages),
                 "BASE_IMAGE": base_image.format_map({
                     "major": sys.version_info.major,
                     "minor": sys.version_info.minor,
@@ -291,7 +313,7 @@ def pdm_container(
             cache_from=cache_from,
             labels=labels,
             # We must specify the workdir manually for apptainer support
-            workdir=pathlib.Path("/workspace/src"),
+            workdir=pathlib.Path("/workspace"),
             args=args,
             env_vars=env_vars,
         )

xm_slurm/packaging/docker/__init__.py

@@ -11,8 +11,8 @@ from xm_slurm.packaging import registry
 from xm_slurm.packaging.docker.abc import DockerClient
 
 FLAGS = flags.FLAGS
-REMOTE_BUILD = flags.DEFINE_enum(
-    "xm_builder", "local", ["local", "gcp", "azure"], "Remote build provider."
+DOCKER_CLIENT_PROVIDER = flags.DEFINE_enum(
+    "xm_docker_client", "docker", ["docker"], "Docker image build client."
 )
 
 IndexedContainer = registry.IndexedContainer
@@ -20,19 +20,13 @@ IndexedContainer = registry.IndexedContainer
 
 @functools.cache
 def docker_client() -> DockerClient:
-    match REMOTE_BUILD.value:
-        case "local":
+    match DOCKER_CLIENT_PROVIDER.value:
+        case "docker":
             from xm_slurm.packaging.docker.local import LocalDockerClient
 
             return LocalDockerClient()
-        case "gcp":
-            from xm_slurm.packaging.docker.cloud import GoogleCloudRemoteDockerClient
-
-            return GoogleCloudRemoteDockerClient()
-        case "azure":
-            raise NotImplementedError("Azure remote build is not yet supported.")
         case _:
-            raise ValueError(f"Unknown remote build provider: {REMOTE_BUILD.value}")
+            raise ValueError(f"Unknown build client: {DOCKER_CLIENT_PROVIDER.value}")
 
 
 @registry.register(Dockerfile)

xm_slurm/packaging/docker/local.py

@@ -9,7 +9,6 @@ import shlex
 import shutil
 import subprocess
 import tempfile
-import typing
 from typing import Sequence
 
 from xmanager import xm
@@ -37,7 +36,7 @@ class LocalDockerClient(DockerClient):
         BUILDKIT = enum.auto()
         BUILDAH = enum.auto()
 
-    def __init__(self):
+    def __init__(self) -> None:
         if "XM_DOCKER_CLIENT" in os.environ:
             client_call = shlex.split(os.environ["XM_DOCKER_CLIENT"])
         elif shutil.which("docker"):
@@ -80,7 +79,7 @@ class LocalDockerClient(DockerClient):
             return None
 
         def _parse_credentials_from_config(
-            config_path: pathlib.Path
+            config_path: pathlib.Path,
         ) -> RemoteRepositoryCredentials | None:
             """Parse credentials from the Docker configuration file."""
             if not config_path.exists():
@@ -138,9 +137,13 @@ class LocalDockerClient(DockerClient):
         targets: Sequence[IndexedContainer[xm.Packageable]],
     ) -> list[IndexedContainer[RemoteImage]]:
         executors_by_executables = packaging_utils.collect_executors_by_executable(targets)
-        executors_by_executables = typing.cast(
-            dict[Dockerfile, list[SlurmSpec]], executors_by_executables
-        )
+        for executable, executors in executors_by_executables.items():
+            assert isinstance(
+                executable, Dockerfile
+            ), "All executables must be Dockerfiles when building Docker images."
+            assert all(
+                isinstance(executor, SlurmSpec) and executor.tag for executor in executors
+            ), "All executors must be SlurmSpecs with tags when building Docker images."
 
         with tempfile.TemporaryDirectory() as tempdir:
             hcl_file = pathlib.Path(tempdir) / "docker-bake.hcl"
@@ -157,12 +160,10 @@ class LocalDockerClient(DockerClient):
             try:
                 command = DockerBakeCommand(
                     targets=list(
-                        set(
-                            [
-                                packaging_utils.hash_digest(target.value.executable_spec)
-                                for target in targets
-                            ]
-                        )
+                        set([
+                            packaging_utils.hash_digest(target.value.executable_spec)
+                            for target in targets
+                        ])
                     ),
                     files=[hcl_file],
                     metadata_file=metadata_file,

xm_slurm/packaging/utils.py

@@ -10,8 +10,7 @@ import re
 import select
 import shutil
 import subprocess
-import typing
-from typing import Callable, Concatenate, Hashable, Literal, ParamSpec, Sequence, TypeVar
+from typing import Callable, Concatenate, Hashable, ParamSpec, Sequence, TypeVar
 
 from xmanager import xm
 
@@ -23,6 +22,7 @@ ReturnT = TypeVar("ReturnT")
 
 
 def hash_digest(obj: Hashable) -> str:
+    # TODO(jfarebro): Need a better way to hash these objects
     # obj_hash = hash(obj)
     # unsigned_obj_hash = obj_hash.from_bytes(
     #     obj_hash.to_bytes((obj_hash.bit_length() + 7) // 8, "big", signed=True),
@@ -54,7 +54,7 @@ def parallel_map(
 
 
 # Cursor commands to filter out from the command data stream
-cursor_commands_regex = re.compile(
+_CURSOR_ESCAPE_SEQUENCES_REGEX = re.compile(
     rb"\x1b\[\?25[hl]"  # Matches cursor show/hide commands (CSI ?25h and CSI ?25l)
     rb"|\x1b\[[0-9;]*[Hf]"  # Matches cursor position commands (CSI n;mH and CSI n;mf)
     rb"|\x1b\[s"  # Matches cursor save position (CSI s)
@@ -64,54 +64,6 @@ cursor_commands_regex = re.compile(
 )
 
 
-@typing.overload
-def run_command(
-    args: Sequence[str] | xm.SequentialArgs,
-    env: dict[str, str] | None = ...,
-    tty: bool = ...,
-    cwd: str | os.PathLike[str] | None = ...,
-    check: bool = ...,
-    return_stdout: Literal[False] = False,
-    return_stderr: Literal[False] = False,
-) -> subprocess.CompletedProcess[None]: ...
-
-
-@typing.overload
-def run_command(
-    args: Sequence[str] | xm.SequentialArgs,
-    env: dict[str, str] | None = ...,
-    tty: bool = ...,
-    cwd: str | os.PathLike[str] | None = ...,
-    check: bool = ...,
-    return_stdout: Literal[True] = True,
-    return_stderr: Literal[False] = False,
-) -> subprocess.CompletedProcess[str]: ...
-
-
-@typing.overload
-def run_command(
-    args: Sequence[str] | xm.SequentialArgs,
-    env: dict[str, str] | None = ...,
-    tty: bool = ...,
-    cwd: str | os.PathLike[str] | None = ...,
-    check: bool = ...,
-    return_stdout: Literal[False] = False,
-    return_stderr: Literal[True] = True,
-) -> subprocess.CompletedProcess[str]: ...
-
-
-@typing.overload
-def run_command(
-    args: Sequence[str] | xm.SequentialArgs,
-    env: dict[str, str] | None = ...,
-    tty: bool = ...,
-    cwd: str | os.PathLike[str] | None = ...,
-    check: bool = ...,
-    return_stdout: Literal[True] = True,
-    return_stderr: Literal[True] = True,
-) -> subprocess.CompletedProcess[str]: ...
-
-
 def run_command(
     args: Sequence[str] | xm.SequentialArgs,
     env: dict[str, str] | None = None,
@@ -120,7 +72,7 @@ def run_command(
     check: bool = False,
     return_stdout: bool = False,
     return_stderr: bool = False,
-) -> subprocess.CompletedProcess[str] | subprocess.CompletedProcess[None]:
+) -> subprocess.CompletedProcess[str]:
     if isinstance(args, xm.SequentialArgs):
         args = args.to_list()
     args = list(args)
@@ -171,7 +123,7 @@ def run_command(
                     fds.remove(fd)
                     continue
 
-                data = re.sub(cursor_commands_regex, b"", data)
+                data = _CURSOR_ESCAPE_SEQUENCES_REGEX.sub(b"", data)
 
                 if fd == stdout_master:
                     if return_stdout:
@@ -186,8 +138,8 @@ def run_command(
                 else:
                     raise RuntimeError("Unexpected file descriptor")
 
-    stdout = stdout_data.decode(errors="replace") if stdout_data else None
-    stderr = stderr_data.decode(errors="replace") if stderr_data else None
+    stdout = stdout_data.decode(errors="replace") if stdout_data else ""
+    stderr = stderr_data.decode(errors="replace") if stderr_data else ""
 
     retcode = process.poll()
     assert retcode is not None