[chore]: Move run_async to _internal/utils (#2057)

jvstme · web-flow · commit acf60c9e1298 · 2024-12-05T07:44:29.000Z
Part of an upcoming refactoring, moved to a
separate PR for a smaller diff and to avoid merge
conflicts. `run_async` will be used both in
`_internal/server` and `_internal/proxy`.
diff --git a/src/dstack/_internal/core/services/ssh/tunnel.py b/src/dstack/_internal/core/services/ssh/tunnel.py
@@ -11,6 +11,7 @@
 from dstack._internal.core.models.instances import SSHConnectionParams
 from dstack._internal.core.services.ssh import get_ssh_error
 from dstack._internal.core.services.ssh.client import get_ssh_client_info
+from dstack._internal.utils.common import run_async
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.path import FilePath, FilePathOrContent, PathLike
 from dstack._internal.utils.ssh import normalize_path
@@ -200,8 +201,7 @@ def open(self) -> None:
         raise get_ssh_error(stderr)
 
     async def aopen(self) -> None:
-        loop = asyncio.get_event_loop()
-        await loop.run_in_executor(None, self._remove_log_file)
+        await run_async(self._remove_log_file)
         proc = await asyncio.create_subprocess_exec(
             *self.open_command(), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
         )
@@ -214,7 +214,7 @@ async def aopen(self) -> None:
             raise SSHError(msg) from e
         if proc.returncode == 0:
             return
-        stderr = await loop.run_in_executor(None, self._read_log_file)
+        stderr = await run_async(self._read_log_file)
         logger.debug("SSH tunnel failed: %s", stderr)
         raise get_ssh_error(stderr)
 
diff --git a/src/dstack/_internal/server/background/tasks/process_instances.py b/src/dstack/_internal/server/background/tasks/process_instances.py
@@ -89,8 +89,7 @@
 from dstack._internal.server.services.runner import client as runner_client
 from dstack._internal.server.services.runner.client import HealthStatus
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
-from dstack._internal.server.utils.common import run_async
-from dstack._internal.utils.common import get_current_datetime
+from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.network import get_ip_from_network, is_ip_among_addresses
 from dstack._internal.utils.ssh import (
@@ -262,9 +261,7 @@ async def _add_remote(instance: InstanceModel) -> None:
         authorized_keys.append(instance.project.ssh_public_key.strip())
 
         try:
-            future = asyncio.get_running_loop().run_in_executor(
-                None, _deploy_instance, remote_details, pkeys, authorized_keys
-            )
+            future = run_async(_deploy_instance, remote_details, pkeys, authorized_keys)
             deploy_timeout = 20 * 60  # 20 minutes
             result = await asyncio.wait_for(future, timeout=deploy_timeout)
             health, host_info = result
diff --git a/src/dstack/_internal/server/background/tasks/process_metrics.py b/src/dstack/_internal/server/background/tasks/process_metrics.py
@@ -13,8 +13,7 @@
 from dstack._internal.server.services.jobs import get_job_provisioning_data
 from dstack._internal.server.services.runner import client
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
-from dstack._internal.server.utils.common import run_async
-from dstack._internal.utils.common import batched, get_current_datetime
+from dstack._internal.utils.common import batched, get_current_datetime, run_async
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
diff --git a/src/dstack/_internal/server/background/tasks/process_placement_groups.py b/src/dstack/_internal/server/background/tasks/process_placement_groups.py
@@ -11,8 +11,7 @@
 from dstack._internal.server.services import backends as backends_services
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.placement import placement_group_model_to_placement_group
-from dstack._internal.server.utils.common import run_async
-from dstack._internal.utils.common import get_current_datetime
+from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
diff --git a/src/dstack/_internal/server/background/tasks/process_running_jobs.py b/src/dstack/_internal/server/background/tasks/process_running_jobs.py
@@ -47,7 +47,6 @@
     run_model_to_run,
 )
 from dstack._internal.server.services.storage import get_default_storage
-from dstack._internal.server.utils.common import run_async
 from dstack._internal.utils import common as common_utils
 from dstack._internal.utils.interpolator import VariablesInterpolator
 from dstack._internal.utils.logging import get_logger
@@ -188,7 +187,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 if job_provisioning_data.backend == BackendType.LOCAL:
                     # No need to update ~/.ssh/authorized_keys when running shim localy
                     user_ssh_key = ""
-                success = await run_async(
+                success = await common_utils.run_async(
                     _process_provisioning_with_shim,
                     server_ssh_private_key,
                     job_provisioning_data,
@@ -213,7 +212,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                     repo=repo_model,
                     code_hash=run.run_spec.repo_code_hash,
                 )
-                success = await run_async(
+                success = await common_utils.run_async(
                     _process_provisioning_no_shim,
                     server_ssh_private_key,
                     job_provisioning_data,
@@ -254,7 +253,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 repo=repo_model,
                 code_hash=run.run_spec.repo_code_hash,
             )
-            success = await run_async(
+            success = await common_utils.run_async(
                 _process_pulling_with_shim,
                 server_ssh_private_key,
                 job_provisioning_data,
@@ -268,7 +267,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
             )
         elif initial_status == JobStatus.RUNNING:
             logger.debug("%s: process running job, age=%s", fmt(job_model), job_submission.age)
-            success = await run_async(
+            success = await common_utils.run_async(
                 _process_running,
                 server_ssh_private_key,
                 job_provisioning_data,
@@ -601,7 +600,7 @@ async def _get_job_code(
     storage = get_default_storage()
     if storage is None or code_model.blob is not None:
         return code_model.blob
-    blob = await run_async(
+    blob = await common_utils.run_async(
         storage.get_code,
         project.name,
         repo.name,
diff --git a/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py b/src/dstack/_internal/server/background/tasks/process_submitted_jobs.py
@@ -66,7 +66,6 @@
 from dstack._internal.server.services.volumes import (
     volume_model_to_volume,
 )
-from dstack._internal.server.utils.common import run_async
 from dstack._internal.utils import common as common_utils
 from dstack._internal.utils.logging import get_logger
 
@@ -406,7 +405,7 @@ async def _run_job_on_new_instance(
         )
         offer_volumes = get_offer_volumes(volumes, offer)
         try:
-            job_provisioning_data = await run_async(
+            job_provisioning_data = await common_utils.run_async(
                 backend.compute().run_job,
                 run,
                 job,
@@ -587,7 +586,7 @@ async def _attach_volume(
     if volume_model.deleted:
         raise ServerClientError("Cannot attach a deleted volume")
     volume = volume_model_to_volume(volume_model)
-    attachment_data = await run_async(
+    attachment_data = await common_utils.run_async(
         backend.compute().attach_volume,
         volume=volume,
         instance_id=instance_id,
diff --git a/src/dstack/_internal/server/background/tasks/process_volumes.py b/src/dstack/_internal/server/background/tasks/process_volumes.py
@@ -9,8 +9,7 @@
 from dstack._internal.server.services import backends as backends_services
 from dstack._internal.server.services import volumes as volumes_services
 from dstack._internal.server.services.locking import get_locker
-from dstack._internal.server.utils.common import run_async
-from dstack._internal.utils.common import get_current_datetime
+from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
diff --git a/src/dstack/_internal/server/services/backends/__init__.py b/src/dstack/_internal/server/services/backends/__init__.py
@@ -28,7 +28,7 @@
 from dstack._internal.server.models import BackendModel, ProjectModel
 from dstack._internal.server.services.backends.configurators.base import Configurator
 from dstack._internal.server.settings import LOCAL_BACKEND_ENABLED
-from dstack._internal.server.utils.common import run_async
+from dstack._internal.utils.common import run_async
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
diff --git a/src/dstack/_internal/server/services/config.py b/src/dstack/_internal/server/services/config.py
@@ -34,7 +34,7 @@
     DefaultPermissions,
     set_default_permissions,
 )
-from dstack._internal.server.utils.common import run_async
+from dstack._internal.utils.common import run_async
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py
@@ -66,11 +66,8 @@
     string_to_lock_id,
 )
 from dstack._internal.server.services.logging import fmt
-from dstack._internal.server.utils.common import (
-    gather_map_async,
-    run_async,
-)
-from dstack._internal.utils.common import get_current_datetime
+from dstack._internal.server.utils.common import gather_map_async
+from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.crypto import generate_rsa_key_pair_bytes
 from dstack._internal.utils.logging import get_logger
 
diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py
@@ -39,8 +39,7 @@
 from dstack._internal.server.services.runner import client
 from dstack._internal.server.services.runner.ssh import get_runner_ports, runner_ssh_tunnel
 from dstack._internal.server.services.volumes import volume_model_to_volume
-from dstack._internal.server.utils.common import run_async
-from dstack._internal.utils.common import get_current_datetime
+from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.path import FileContent
 
diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py
@@ -24,7 +24,7 @@
 from dstack._internal.core.services.profiles import get_retry
 from dstack._internal.core.services.ssh.ports import filter_reserved_ports
 from dstack._internal.server.services.docker import ImageConfig, get_image_config
-from dstack._internal.server.utils.common import run_async
+from dstack._internal.utils.common import run_async
 
 
 def get_default_python_verison() -> str:
diff --git a/src/dstack/_internal/server/services/logs.py b/src/dstack/_internal/server/services/logs.py
@@ -20,7 +20,7 @@
 from dstack._internal.server.models import ProjectModel
 from dstack._internal.server.schemas.logs import PollLogsRequest
 from dstack._internal.server.schemas.runner import LogEvent as RunnerLogEvent
-from dstack._internal.server.utils.common import run_async
+from dstack._internal.utils.common import run_async
 from dstack._internal.utils.logging import get_logger
 
 BOTO_AVAILABLE = True
diff --git a/src/dstack/_internal/server/services/projects.py b/src/dstack/_internal/server/services/projects.py
@@ -22,8 +22,7 @@
 from dstack._internal.server.services.backends import get_configurator
 from dstack._internal.server.services.permissions import get_default_permissions
 from dstack._internal.server.settings import DEFAULT_PROJECT_NAME
-from dstack._internal.server.utils.common import run_async
-from dstack._internal.utils.common import get_current_datetime
+from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.crypto import generate_rsa_key_pair_bytes
 from dstack._internal.utils.logging import get_logger
 
diff --git a/src/dstack/_internal/server/services/repos.py b/src/dstack/_internal/server/services/repos.py
@@ -28,7 +28,7 @@
     UserModel,
 )
 from dstack._internal.server.services.storage import get_default_storage
-from dstack._internal.server.utils.common import run_async
+from dstack._internal.utils.common import run_async
 
 
 async def list_repos(
diff --git a/src/dstack/_internal/server/services/volumes.py b/src/dstack/_internal/server/services/volumes.py
@@ -29,7 +29,6 @@
     string_to_lock_id,
 )
 from dstack._internal.server.services.projects import list_project_models, list_user_project_models
-from dstack._internal.server.utils.common import run_async
 from dstack._internal.utils import common, random_names
 from dstack._internal.utils.logging import get_logger
 
@@ -350,7 +349,7 @@ async def _delete_volume(session: AsyncSession, project: ProjectModel, volume_mo
         )
         return
 
-    await run_async(
+    await common.run_async(
         backend.compute().delete_volume,
         volume=volume,
     )
diff --git a/src/dstack/_internal/server/utils/common.py b/src/dstack/_internal/server/utils/common.py
@@ -1,5 +1,4 @@
 import asyncio
-from functools import partial
 from typing import (
     Awaitable,
     Callable,
@@ -12,17 +11,6 @@
     Union,
 )
 
-from typing_extensions import ParamSpec
-
-P = ParamSpec("P")
-R = TypeVar("R")
-
-
-async def run_async(func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
-    func_with_args = partial(func, *args, **kwargs)
-    return await asyncio.get_running_loop().run_in_executor(None, func_with_args)
-
-
 ItemT = TypeVar("ItemT")
 ResultT = TypeVar("ResultT")
 
diff --git a/src/dstack/_internal/utils/common.py b/src/dstack/_internal/utils/common.py
@@ -1,12 +1,24 @@
+import asyncio
 import itertools
 import re
 import time
 from collections.abc import Callable
 from datetime import datetime, timedelta, timezone
+from functools import partial
 from pathlib import Path
 from typing import Any, Iterable, List, Optional, TypeVar
 from urllib.parse import urlparse
 
+from typing_extensions import ParamSpec
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+async def run_async(func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
+    func_with_args = partial(func, *args, **kwargs)
+    return await asyncio.get_running_loop().run_in_executor(None, func_with_args)
+
 
 def get_dstack_dir() -> Path:
     return Path.joinpath(Path.home(), ".dstack")

Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@`
`34`	`34`	`DefaultPermissions,`
`35`	`35`	`set_default_permissions,`
`36`	`36`	`)`
`37`		`-from dstack._internal.server.utils.common import run_async`
	`37`	`+from dstack._internal.utils.common import run_async`
`38`	`38`	`from dstack._internal.utils.logging import get_logger`
`39`	`39`
`40`	`40`	`logger = get_logger(__name__)`