dagster-cloud 1.8.2__py3-none-any.whl → 1.12.6__py3-none-any.whl

dagster_cloud/pex/grpc/compile.py

@@ -69,7 +69,7 @@ def protoc(generated_dir: str):
             generated_grpc_path,
             tempfile_path,
         )
-        with open(tempfile_path, "r", encoding="utf8") as generated:
+        with open(tempfile_path, encoding="utf8") as generated:
             with open(generated_grpc_path, "w", encoding="utf8") as rewritten:
                 for line in GENERATED_HEADER:
                     rewritten.write(line)
@@ -85,7 +85,7 @@ def protoc(generated_dir: str):
             generated_pb2_path,
             tempfile_path,
         )
-        with open(tempfile_path, "r", encoding="utf8") as generated:
+        with open(tempfile_path, encoding="utf8") as generated:
             with open(generated_pb2_path, "w", encoding="utf8") as rewritten:
                 for line in GENERATED_HEADER:
                     rewritten.write(line)

dagster_cloud/pex/grpc/server/__init__.py

@@ -1,2 +1,2 @@
-from .manager import MultiPexManager as MultiPexManager
-from .server import run_multipex_server as run_multipex_server
+from dagster_cloud.pex.grpc.server.manager import MultiPexManager as MultiPexManager
+from dagster_cloud.pex.grpc.server.server import run_multipex_server as run_multipex_server

dagster_cloud/pex/grpc/server/cli/__init__.py

@@ -1,15 +1,16 @@
 import logging
 import os
+import signal
 import subprocess
 from typing import Optional
 
 from dagster._serdes import deserialize_value
-from dagster._utils.interrupts import capture_interrupts
+from dagster._utils.interrupts import setup_interrupt_handlers
 from dagster_cloud_cli.core.workspace import PexMetadata
 from typer import Option, Typer
 
-from ..registry import PexS3Registry
-from ..server import run_multipex_server
+from dagster_cloud.pex.grpc.server.registry import PexS3Registry
+from dagster_cloud.pex.grpc.server.server import run_multipex_server
 
 app = Typer(hidden=True)
 
@@ -48,20 +49,31 @@ def execute_run(
         default="/tmp/pex-files", envvar="LOCAL_PEX_FILES_DIR"
     ),
 ):
-    with capture_interrupts():
-        pex_metadata = deserialize_value(pex_metadata_json, PexMetadata)
-        executable = PexS3Registry(local_pex_files_dir).get_pex_executable(pex_metadata)
+    setup_interrupt_handlers()
+    pex_metadata = deserialize_value(pex_metadata_json, PexMetadata)
+    executable = PexS3Registry(local_pex_files_dir).get_pex_executable(pex_metadata)
 
-        subprocess.run(
-            [
-                executable.source_path,
-                "-m",
-                "dagster",
-                "api",
-                "execute_run",
-                input_json,
-            ],
-            check=True,
-            env={**os.environ.copy(), **executable.environ},
-            cwd=executable.working_directory,
-        )
+    run_process = subprocess.Popen(
+        [
+            executable.source_path,
+            "-m",
+            "dagster",
+            "api",
+            "execute_run",
+            input_json,
+        ],
+        env={**os.environ.copy(), **executable.environ},
+        cwd=executable.working_directory,
+    )
+
+    logger = logging.getLogger("dagster.pex_run")
+
+    try:
+        return_code = run_process.wait()
+        if return_code != 0:
+            raise Exception(f"PEX subprocess returned with exit code {return_code}")
+    except KeyboardInterrupt:
+        logger.info("Forwarding interrupt to PEX subprocess")
+        run_process.send_signal(signal.SIGINT)
+        run_process.wait()
+        raise

dagster_cloud/pex/grpc/server/manager.py

@@ -5,21 +5,24 @@ import sys
 import threading
 import time
 from contextlib import AbstractContextManager
-from typing import Dict, List, Optional, Set, Tuple, Union, cast
+from typing import Optional, Union, cast
 
-import dagster._seven as seven
 from dagster import _check as check
 from dagster._core.errors import DagsterUserCodeUnreachableError
 from dagster._core.instance.ref import InstanceRef
 from dagster._grpc.client import DagsterGrpcClient, client_heartbeat_thread
-from dagster._serdes.ipc import open_ipc_subprocess
 from dagster._utils import find_free_port, safe_tempfile_path_unmanaged
 from dagster._utils.error import SerializableErrorInfo, serializable_error_info_from_exc_info
 from dagster_cloud_cli.core.workspace import CodeLocationDeployData, PexMetadata
+from dagster_shared import seven
+from dagster_shared.ipc import open_ipc_subprocess
 from pydantic import BaseModel, Extra
 
-from ..types import PexServerHandle
-from .registry import PexS3Registry
+from dagster_cloud.pex.grpc.server.registry import PexS3Registry
+from dagster_cloud.pex.grpc.types import PexServerHandle
+from dagster_cloud.workspace.user_code_launcher.utils import get_grpc_server_env
+
+logger = logging.getLogger("dagster.multipex")
 
 
 class PexProcessEntry(BaseModel, frozen=True, extra=Extra.forbid, arbitrary_types_allowed=True):
@@ -43,11 +46,11 @@ class MultiPexManager(AbstractContextManager):
         enable_metrics: bool = False,
     ):
         # Keyed by hash of PexServerHandle
-        self._pex_servers: Dict[str, Union[PexProcessEntry, PexErrorEntry]] = {}
-        self._pending_startup_pex_servers: Set[str] = set()
-        self._pending_shutdown_pex_servers: Set[str] = set()
+        self._pex_servers: dict[str, Union[PexProcessEntry, PexErrorEntry]] = {}
+        self._pending_startup_pex_servers: set[str] = set()
+        self._pending_shutdown_pex_servers: set[str] = set()
         self._pex_servers_lock = threading.RLock()
-        self._pex_metadata_for_handle: Dict[
+        self._pex_metadata_for_handle: dict[
             str, Optional[PexMetadata]
         ] = {}  # maps handle id to the pex tag
         self._heartbeat_ttl = 60
@@ -62,13 +65,13 @@ class MultiPexManager(AbstractContextManager):
                 daemon=True,
             )
             self._watchdog_thread.start()
-            logging.info(
+            logger.info(
                 "Created a watchdog thread %s for MultiPexManager with watchdog_run_interval=%s",
                 self._watchdog_thread.name,
                 watchdog_run_interval,
             )
         else:
-            logging.info(
+            logger.info(
                 "No watchdog thread started for MultiPexManager (watchdog_run_interval=%s)",
                 watchdog_run_interval,
             )
@@ -85,7 +88,7 @@ class MultiPexManager(AbstractContextManager):
                 returncode = server.grpc_server_process.poll()
                 if returncode is not None:
                     dead_server_returncodes.append((server, returncode))
-                    logging.error(
+                    logger.error(
                         "watchdog: pex subprocesss %s unexpectedly exited with returncode %s -"
                         " changing state to error",
                         server_id,
@@ -93,7 +96,7 @@ class MultiPexManager(AbstractContextManager):
                     )
             self._mark_servers_unexpected_termination(dead_server_returncodes)
             if dead_server_returncodes:
-                logging.warning(
+                logger.warning(
                     "watchdog: inspected %s active servers %s, of which %s were found unexpectedly"
                     " terminated",
                     len(active_servers),
@@ -102,7 +105,7 @@ class MultiPexManager(AbstractContextManager):
                 )
 
     def _mark_servers_unexpected_termination(
-        self, dead_server_returncodes: List[Tuple[PexProcessEntry, int]]
+        self, dead_server_returncodes: list[tuple[PexProcessEntry, int]]
     ) -> None:
         with self._pex_servers_lock:
             for server, returncode in dead_server_returncodes:
@@ -137,9 +140,9 @@ class MultiPexManager(AbstractContextManager):
             if isinstance(pex_server_or_error, PexErrorEntry):
                 return pex_server_or_error.error
 
-            return cast(PexProcessEntry, self._pex_servers[handle_id]).grpc_client
+            return cast("PexProcessEntry", self._pex_servers[handle_id]).grpc_client
 
-    def get_active_pex_servers(self) -> List[PexProcessEntry]:
+    def get_active_pex_servers(self) -> list[PexProcessEntry]:
         with self._pex_servers_lock:
             return [
                 server
@@ -147,7 +150,7 @@ class MultiPexManager(AbstractContextManager):
                 if self.is_server_active(server_id) and isinstance(server, PexProcessEntry)
             ]
 
-    def get_error_pex_servers(self) -> List[PexErrorEntry]:
+    def get_error_pex_servers(self) -> list[PexErrorEntry]:
         with self._pex_servers_lock:
             return [
                 server
@@ -157,7 +160,7 @@ class MultiPexManager(AbstractContextManager):
 
     def get_active_pex_server_handles(
         self, deployment_name, location_name: str
-    ) -> List[PexServerHandle]:
+    ) -> list[PexServerHandle]:
         return [
             server.pex_server_handle
             for server in self.get_active_pex_servers()
@@ -167,7 +170,7 @@ class MultiPexManager(AbstractContextManager):
 
     def get_error_pex_server_handles(
         self, deployment_name, location_name: str
-    ) -> List[PexServerHandle]:
+    ) -> list[PexServerHandle]:
         return [
             server.pex_server_handle
             for server in self.get_error_pex_servers()
@@ -175,7 +178,7 @@ class MultiPexManager(AbstractContextManager):
             and server.pex_server_handle.location_name == location_name
         ]
 
-    def get_all_pex_grpc_clients_map(self) -> Dict[str, DagsterGrpcClient]:
+    def get_all_pex_grpc_clients_map(self) -> dict[str, DagsterGrpcClient]:
         with self._pex_servers_lock:
             return {
                 server.pex_server_handle.get_id(): server.grpc_client
@@ -209,14 +212,14 @@ class MultiPexManager(AbstractContextManager):
                 pex_executable = self._registry.get_pex_executable(
                     check.not_none(code_location_deploy_data.pex_metadata)
                 )
-                logging.info(
+                logger.info(
                     "Installed pex executable %s at %s",
                     code_location_deploy_data.pex_metadata,
                     pex_executable.source_path,
                 )
 
                 metadata = code_location_deploy_data
-                logging.info("Launching subprocess %s", pex_executable.source_path)
+                logger.info("Launching subprocess %s", pex_executable.source_path)
                 subprocess_args = [
                     pex_executable.source_path,
                     "-m",
@@ -243,7 +246,8 @@ class MultiPexManager(AbstractContextManager):
                     port = None
                     socket = safe_tempfile_path_unmanaged()
 
-                additional_env = metadata.get_grpc_server_env(
+                additional_env = get_grpc_server_env(
+                    code_location_deploy_data=metadata,
                     port=port,
                     location_name=server_handle.location_name,
                     instance_ref=instance_ref,
@@ -274,7 +278,7 @@ class MultiPexManager(AbstractContextManager):
                     daemon=True,
                 )
                 heartbeat_thread.start()
-                logging.info(
+                logger.info(
                     "Created a heartbeat thread %s for %s",
                     heartbeat_thread.name,
                     server_handle.get_id(),
@@ -287,7 +291,7 @@ class MultiPexManager(AbstractContextManager):
                         error=serializable_error_info_from_exc_info(sys.exc_info()),
                     )
                     self._pending_startup_pex_servers.remove(server_handle.get_id())
-                logging.exception(
+                logger.exception(
                     "Creating new pex server for %s:%s failed",
                     server_handle.deployment_name,
                     server_handle.location_name,
@@ -307,7 +311,7 @@ class MultiPexManager(AbstractContextManager):
         with self._pex_servers_lock:
             handle_id = server_handle.get_id()
             if handle_id in self._pending_startup_pex_servers:
-                logging.info(
+                logger.info(
                     "Ignoring request to create pex server for %s - an identical server is"
                     " already pending start",
                     handle_id,
@@ -316,12 +320,12 @@ class MultiPexManager(AbstractContextManager):
             self._pending_startup_pex_servers.add(handle_id)
             if handle_id in self._pex_servers:
                 # clear any previous error state since we're attempting to start server again
-                logging.info(
+                logger.info(
                     "Clearing previous state for %s: %s", handle_id, self._pex_servers[handle_id]
                 )
                 del self._pex_servers[handle_id]
 
-        logging.info(
+        logger.info(
             "Creating new pex server for %s:%s",
             server_handle.deployment_name,
             server_handle.location_name,
@@ -332,7 +336,7 @@ class MultiPexManager(AbstractContextManager):
         handle_id = server_handle.get_id()
         with self._pex_servers_lock:
             if handle_id in self._pex_servers or handle_id in self._pending_startup_pex_servers:
-                logging.info("Server %s marked for shutdown", handle_id)
+                logger.info("Server %s marked for shutdown", handle_id)
                 self._pending_shutdown_pex_servers.add(handle_id)
 
     def cleanup_pending_shutdown_pex_servers(self) -> None:
@@ -343,14 +347,23 @@ class MultiPexManager(AbstractContextManager):
                 if handle_id not in self._pex_servers:
                     continue
                 server = self._pex_servers[handle_id]
-                if (
-                    isinstance(server, PexProcessEntry)
-                    and server.grpc_server_process.poll() is not None
-                ):
-                    to_remove.add(handle_id)
+                if isinstance(server, PexProcessEntry):
+                    if server.grpc_server_process.poll() is not None:
+                        # Server process shut down
+                        to_remove.add(handle_id)
+                    else:
+                        try:
+                            server.grpc_client.ping("")
+                        except DagsterUserCodeUnreachableError:
+                            logger.warning(
+                                "server process is still running but the server is unreachable - killing the process",
+                                exc_info=True,
+                            )
+                            server.grpc_server_process.kill()
+                            to_remove.add(handle_id)
 
             for handle_id in to_remove:
-                logging.info("Server %s completely shutdown, cleaning up", handle_id)
+                logger.info("Server %s completely shutdown, cleaning up", handle_id)
                 self._pending_shutdown_pex_servers.remove(handle_id)
                 del self._pex_servers[handle_id]
 
@@ -359,32 +372,37 @@ class MultiPexManager(AbstractContextManager):
                 pex_server = self._pex_servers.get(handle_id)
                 if not pex_server:
                     # still in _pending_startup_pex_servers
-                    logging.info("Server %s not up yet, will request shutdown later", handle_id)
+                    logger.info("Server %s not up yet, will request shutdown later", handle_id)
                     continue
 
                 if isinstance(pex_server, PexErrorEntry):
-                    logging.debug("Server %s was in an error state, no shutdown needed", handle_id)
+                    logger.debug("Server %s was in an error state, no shutdown needed", handle_id)
                     continue
 
                 if pex_server.heartbeat_shutdown_event.is_set():
                     # already requested shutdown
-                    logging.info("Already requested shutdown for server %s", handle_id)
                     continue
 
-                logging.info("Requesting shutdown for server %s", handle_id)
+                logger.info("Requesting shutdown for server %s", handle_id)
                 pex_server.heartbeat_shutdown_event.set()
                 pex_server.heartbeat_thread.join()
                 try:
                     pex_server.grpc_client.shutdown_server()
                 except DagsterUserCodeUnreachableError:
-                    # Server already shutdown
-                    pass
+                    logger.warning(
+                        "Server shutdown for %s over grpc failed, killing the process",
+                        handle_id,
+                        exc_info=True,
+                    )
+                    pex_server.grpc_server_process.kill()
 
             # Delete any registry files not in use anymore
             # - ensure that resources for servers starting up or shutting down are not removed
             # - important to do this while holding the lock to avoid race conditions
             running_server_ids = {
-                proc.pex_server_handle.get_id() for proc in self._pex_servers.values()
+                proc.pex_server_handle.get_id()
+                for proc in self._pex_servers.values()
+                if isinstance(proc, PexProcessEntry)
             }
             in_use_handle_ids = self._pending_startup_pex_servers.union(
                 self._pending_shutdown_pex_servers

dagster_cloud/pex/grpc/server/registry.py

@@ -9,7 +9,7 @@ import threading
 from dataclasses import dataclass
 from os.path import expanduser
 from pathlib import Path
-from typing import Dict, List, NamedTuple, Optional, Set
+from typing import NamedTuple, Optional
 from uuid import uuid4
 
 from dagster import _check as check
@@ -19,6 +19,8 @@ DEFAULT_PEX_FILES_DIR = "/tmp/pex-files"
 # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/customizations/s3.html
 MULTIPART_DOWNLOAD_THREADS = 20  # Double the boto3 default of 10
 
+logger = logging.getLogger("dagster.multipex")
+
 
 def _download_from_s3(filename: str, local_filepath: str):
     # Lazy import boto3 to avoid a hard dependency during module load
@@ -28,7 +30,11 @@ def _download_from_s3(filename: str, local_filepath: str):
 
     config = Config(retries={"max_attempts": 3, "mode": "standard"})
 
-    s3 = boto3.client("s3", region_name="us-west-2", config=config)
+    s3 = boto3.client(
+        "s3",
+        region_name=os.getenv("DAGSTER_CLOUD_SERVERLESS_REGION", "us-west-2"),
+        config=config,
+    )
 
     # TODO: move the bucket and prefix to pex_metdata
     s3_bucket_name = os.environ["DAGSTER_CLOUD_SERVERLESS_STORAGE_S3_BUCKET"]
@@ -52,22 +58,22 @@ class PexExecutable(
         "_PexExecutable",
         [
             ("source_path", str),
-            ("all_paths", List[str]),
-            ("environ", Dict[str, str]),
+            ("all_paths", list[str]),
+            ("environ", dict[str, str]),
             ("working_directory", Optional[str]),
-            ("venv_dirs", List[str]),
+            ("venv_dirs", list[str]),
         ],
     )
 ):
     def __new__(
         cls,
         source_path: str,
-        all_paths: List[str],
-        environ: Dict[str, str],
+        all_paths: list[str],
+        environ: dict[str, str],
         working_directory: Optional[str],
-        venv_dirs: List[str],
+        venv_dirs: list[str],
     ):
-        return super(PexExecutable, cls).__new__(
+        return super().__new__(
             cls,
             check.str_param(source_path, "source_path"),
             check.list_param(all_paths, "all_paths", str),
@@ -96,12 +102,12 @@ class PexS3Registry:
             local_pex_files_dir if local_pex_files_dir else DEFAULT_PEX_FILES_DIR
         )
         os.makedirs(self._local_pex_files_dir, exist_ok=True)
-        self.working_dirs: Dict[
+        self.working_dirs: dict[
             str, str
         ] = {}  # once unpacked, working dirs dont change so we cache them
 
         # keep track of local files and directories used by each pex tag
-        self.local_paths_for_pex_tag: Dict[str, Set[str]] = {}
+        self.local_paths_for_pex_tag: dict[str, set[str]] = {}
 
         # lock to do safe install and cleanup
         self._install_lock = threading.RLock()
@@ -149,7 +155,7 @@ class PexS3Registry:
                 deps_pex_filepaths.append(local_filepath)
 
         if not source_pex_filepath:
-            raise ValueError("Invalid pex_tag has no source pex: %r" % pex_metadata.pex_tag)
+            raise ValueError(f"Invalid pex_tag has no source pex: {pex_metadata.pex_tag!r}")
 
         # we unpack each pex file into its own venv
         source_venv = self.venv_for(source_pex_filepath)
@@ -188,11 +194,11 @@ class PexS3Registry:
             venv_dirs=[str(source_venv.path)] + [str(deps_venv.path) for deps_venv in deps_venvs],
         )
 
-    def cleanup_unused_files(self, in_use_pex_metadatas: List[PexMetadata]) -> None:
+    def cleanup_unused_files(self, in_use_pex_metadatas: list[PexMetadata]) -> None:
         with self._install_lock:
             return self._cleanup_unused_files(in_use_pex_metadatas)
 
-    def _cleanup_unused_files(self, in_use_pex_metadatas: List[PexMetadata]) -> None:
+    def _cleanup_unused_files(self, in_use_pex_metadatas: list[PexMetadata]) -> None:
         """Cleans up all local files and directories that are not associated with any PexMetadata provided."""
         in_use_pex_tags = [pex_metadata.pex_tag for pex_metadata in in_use_pex_metadatas]
 
@@ -207,7 +213,7 @@ class PexS3Registry:
         unused_local_paths = all_local_paths - in_use_local_paths
         unused_paths_present = [path for path in unused_local_paths if os.path.exists(path)]
         if unused_paths_present:
-            logging.info(
+            logger.info(
                 "Cleaning up %s unused local paths: %r",
                 len(unused_paths_present),
                 unused_paths_present,
@@ -219,13 +225,13 @@ class PexS3Registry:
                     else:
                         os.remove(path)
                 except OSError:
-                    logging.exception("Ignoring failure to clean up local unused path %s", path)
+                    logger.exception("Ignoring failure to clean up local unused path %s", path)
 
     def venv_for(self, pex_filepath) -> PexVenv:
         _, pex_filename = os.path.split(pex_filepath)
         venv_dir = self.venv_dir_for(pex_filepath)
         if os.path.exists(venv_dir):
-            logging.info("Reusing existing venv %r for %r", venv_dir, pex_filepath)
+            logger.info("Reusing existing venv %r for %r", venv_dir, pex_filepath)
         else:
             self.install_venv(venv_dir, pex_filepath)
             if not os.path.exists(venv_dir):
@@ -259,6 +265,7 @@ class PexS3Registry:
                     # since we combine multiple venvs, we need non hermetic scripts
                     "--non-hermetic-scripts",
                     venv_dir,
+                    "--pip",
                 ],
                 stderr=subprocess.STDOUT,
             )
@@ -267,7 +274,7 @@ class PexS3Registry:
             raise PexInstallationError(
                 f"Could not install venv. Pex output: {e.output}", pex_filepath
             ) from e
-        logging.info(
+        logger.info(
             "Unpacked pex file %r into venv at %r",
             pex_filepath,
             venv_dir,
@@ -283,7 +290,7 @@ class PexS3Registry:
         if not proc.returncode:
             return Path(proc.stdout.decode("utf-8").strip()).absolute()
         else:
-            logging.error(
+            logger.error(
                 "Cannot determine site-packages for venv at %r: %s\n%s",
                 venv_path,
                 proc.stdout.decode("utf-8"),
@@ -315,12 +322,12 @@ class PexS3Registry:
 
         except subprocess.CalledProcessError:
             # working_directory package is optional, just log a message
-            logging.info("Cannot import working_directory package - not setting current directory.")
+            logger.info("Cannot import working_directory package - not setting current directory.")
             return None
         except OSError:
             # some issue with pex not being runnable, log an error but don't fail yet
             # might fail later if we try to run this again
-            logging.exception(
+            logger.exception(
                 "Ignoring failure to run pex file to determine working_directory %r", pex_path
             )
             return None

dagster_cloud/pex/grpc/server/server.py

@@ -9,12 +9,12 @@ from typing import Optional, cast
 import dagster._check as check
 import grpc
 from dagster._core.errors import DagsterUserCodeUnreachableError
-from dagster._grpc.__generated__ import api_pb2
-from dagster._grpc.__generated__.api_pb2_grpc import (
+from dagster._grpc.__generated__ import dagster_api_pb2
+from dagster._grpc.__generated__.dagster_api_pb2_grpc import (
     DagsterApiServicer,
     add_DagsterApiServicer_to_server,
 )
-from dagster._grpc.client import DEFAULT_GRPC_TIMEOUT
+from dagster._grpc.client import DEFAULT_GRPC_TIMEOUT, DEFAULT_REPOSITORY_GRPC_TIMEOUT
 from dagster._grpc.server import server_termination_target
 from dagster._grpc.types import GetCurrentRunsResult, SensorExecutionArgs
 from dagster._grpc.utils import max_rx_bytes, max_send_bytes
@@ -22,12 +22,13 @@ from dagster._serdes import deserialize_value, serialize_value
 from dagster._utils.error import SerializableErrorInfo, serializable_error_info_from_exc_info
 from grpc_health.v1 import health, health_pb2, health_pb2_grpc
 
-from ..__generated__ import multi_pex_api_pb2
-from ..__generated__.multi_pex_api_pb2_grpc import (
+from dagster_cloud.pex.grpc.__generated__ import multi_pex_api_pb2
+from dagster_cloud.pex.grpc.__generated__.multi_pex_api_pb2_grpc import (
     MultiPexApiServicer,
     add_MultiPexApiServicer_to_server,
 )
-from ..types import (
+from dagster_cloud.pex.grpc.server.manager import MultiPexManager
+from dagster_cloud.pex.grpc.types import (
     CreatePexServerArgs,
     CreatePexServerResponse,
     GetCrashedPexServersArgs,
@@ -38,7 +39,6 @@ from ..types import (
     ShutdownPexServerArgs,
     ShutdownPexServerResponse,
 )
-from .manager import MultiPexManager
 
 
 class MultiPexApiServer(MultiPexApiServicer):
@@ -186,10 +186,15 @@ class DagsterPexProxyApiServer(DagsterApiServicer):
             self._get_handle_from_metadata(context)
         )
         if isinstance(client_or_error, SerializableErrorInfo):
-            return api_pb2.ListRepositoriesReply(
+            return dagster_api_pb2.ListRepositoriesReply(
                 serialized_list_repositories_response_or_error=serialize_value(client_or_error)
             )
-        return client_or_error._get_response("ListRepositories", request)  # noqa: SLF001
+
+        try:
+            return client_or_error._get_response("ListRepositories", request)  # noqa: SLF001
+        except grpc.RpcError as e:
+            # Surface the grpc error to the caller
+            context.abort(e.code(), e.details())
 
     def Ping(self, request, context):
         return self._query("Ping", request, context)
@@ -201,7 +206,9 @@ class DagsterPexProxyApiServer(DagsterApiServicer):
         return self._query("GetCurrentImage", request, context)
 
     def StreamingExternalRepository(self, request, context):
-        return self._streaming_query("StreamingExternalRepository", request, context)
+        return self._streaming_query(
+            "StreamingExternalRepository", request, context, timeout=DEFAULT_REPOSITORY_GRPC_TIMEOUT
+        )
 
     def Heartbeat(self, request, context):
         return self._query("Heartbeat", request, context)
@@ -228,7 +235,9 @@ class DagsterPexProxyApiServer(DagsterApiServicer):
         return self._query("ExternalPipelineSubsetSnapshot", request, context)
 
     def ExternalRepository(self, request, context):
-        return self._query("ExternalRepository", request, context)
+        return self._query(
+            "ExternalRepository", request, context, timeout=DEFAULT_REPOSITORY_GRPC_TIMEOUT
+        )
 
     def ExternalJob(self, request, context):
         return self._query("ExternalJob", request, context)
@@ -242,7 +251,7 @@ class DagsterPexProxyApiServer(DagsterApiServicer):
         except Exception as e:
             if (
                 isinstance(e, grpc.RpcError)
-                and cast(grpc.RpcError, e).code() == grpc.StatusCode.UNIMPLEMENTED
+                and cast("grpc.RpcError", e).code() == grpc.StatusCode.UNIMPLEMENTED
             ):
                 context.abort(
                     grpc.StatusCode.UNIMPLEMENTED,
@@ -278,7 +287,7 @@ class DagsterPexProxyApiServer(DagsterApiServicer):
         except Exception as e:
             if (
                 isinstance(e, grpc.RpcError)
-                and cast(grpc.RpcError, e).code() == grpc.StatusCode.UNIMPLEMENTED
+                and cast("grpc.RpcError", e).code() == grpc.StatusCode.UNIMPLEMENTED
             ):
                 context.abort(
                     grpc.StatusCode.UNIMPLEMENTED,
@@ -337,7 +346,7 @@ class DagsterPexProxyApiServer(DagsterApiServicer):
                     f"Active server hit error:\n{e}",
                 )
 
-        return api_pb2.GetCurrentRunsReply(
+        return dagster_api_pb2.GetCurrentRunsReply(
             serialized_current_runs=serialize_value(
                 GetCurrentRunsResult(current_runs=all_run_ids, serializable_error_info=None)
             )