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

dagster_cloud/instance/__init__.py

@@ -20,6 +20,7 @@ from dagster._core.instance.config import config_field_for_configurable_class
 from dagster._core.instance.ref import InstanceRef, configurable_class_data
 from dagster._core.launcher import DefaultRunLauncher, RunLauncher
 from dagster._core.storage.dagster_run import DagsterRun
+from dagster._core.storage.defs_state.base import DefsStateStorage
 from dagster._serdes import ConfigurableClassData
 from dagster_cloud_cli.core.graphql_client import (
     create_agent_graphql_client,
@@ -31,14 +32,13 @@ from dagster_cloud_cli.core.headers.auth import DagsterCloudInstanceScope
 from urllib3 import Retry
 
 from dagster_cloud.agent import AgentQueuesConfig
+from dagster_cloud.auth.constants import decode_agent_token
+from dagster_cloud.opentelemetry.config import opentelemetry_config_schema
+from dagster_cloud.opentelemetry.controller import OpenTelemetryController
+from dagster_cloud.storage.client import dagster_cloud_api_config
+from dagster_cloud.util import get_env_names_from_config, is_isolated_run
 from dagster_cloud.version import __version__
 
-from ..auth.constants import get_organization_name_from_agent_token
-from ..opentelemetry.config import opentelemetry_config_schema
-from ..opentelemetry.controller import OpenTelemetryController
-from ..storage.client import dagster_cloud_api_config
-from ..util import get_env_names_from_config, is_isolated_run
-
 if TYPE_CHECKING:
     from requests import Session
 
@@ -64,6 +64,11 @@ class DagsterCloudInstance(DagsterInstance):
             "run_retries.retry_on_asset_or_op_failure is a deployment setting and can only be accessed by a DeploymentScopedHostInstance"
         )
 
+    @property
+    def defs_state_storage(self) -> Optional[DefsStateStorage]:
+        # only DeploymentScopedHostInstance / DagsterCloudAgentInstance have a defs state storage
+        return None
+
 
 class DagsterCloudAgentInstance(DagsterCloudInstance):
     def __init__(
@@ -81,14 +86,18 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
         super().__init__(*args, **kwargs)
 
         self._unprocessed_dagster_cloud_api_config = dagster_cloud_api
-        self._dagster_cloud_api_config = self._get_processed_config(
-            "dagster_cloud_api", dagster_cloud_api, dagster_cloud_api_config()
+        self._dagster_cloud_api_config = check.not_none(
+            self._get_processed_config(
+                "dagster_cloud_api",
+                dagster_cloud_api,
+                dagster_cloud_api_config(),
+            )
         )
 
         check.invariant(
             not (
-                self._dagster_cloud_api_config.get("deployment")  # pyright: ignore[reportOptionalMemberAccess]
-                and self._dagster_cloud_api_config.get("deployments")  # pyright: ignore[reportOptionalMemberAccess]
+                self._dagster_cloud_api_config.get("deployment")
+                and self._dagster_cloud_api_config.get("deployments")
             ),
             "Cannot set both deployment and deployments in `dagster_cloud_api`",
         )
@@ -144,6 +153,17 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
 
         self._instance_uuid = str(uuid.uuid4())
 
+    @property
+    def defs_state_storage(self) -> Optional[DefsStateStorage]:
+        # temporary hack to avoid cases where the default BlobStorageStateStorage is used
+        from dagster_cloud.storage.defs_state.storage import GraphQLDefsStateStorage
+
+        return (
+            self._defs_state_storage
+            if isinstance(self._defs_state_storage, GraphQLDefsStateStorage)
+            else None
+        )
+
     def _get_processed_config(
         self, name: str, config: Optional[dict[str, Any]], config_type: dict[str, Any]
     ):
@@ -158,16 +178,16 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
         return processed_config.value
 
     def _dagster_cloud_api_config_for_deployment(self, deployment_name: Optional[str]):
-        new_api_config = dict(copy.deepcopy(self._dagster_cloud_api_config))  # pyright: ignore[reportArgumentType,reportCallIssue]
+        new_api_config = dict(copy.deepcopy(self._dagster_cloud_api_config))
         if deployment_name:
-            new_api_config["deployment"] = deployment_name  # pyright: ignore[reportArgumentType]
+            new_api_config["deployment"] = deployment_name
             if self.includes_branch_deployments:
-                del new_api_config["branch_deployments"]  # pyright: ignore
-            new_api_config.pop("deployments", None)  # pyright: ignore
-            new_api_config.pop("all_serverless_deployments", None)  # pyright: ignore
+                del new_api_config["branch_deployments"]
+            new_api_config.pop("deployments", None)
+            new_api_config.pop("all_serverless_deployments", None)
         else:
-            new_api_config.pop("deployment", None)  # pyright: ignore
-            new_api_config.pop("deployments", None)  # pyright: ignore
+            new_api_config.pop("deployment", None)
+            new_api_config.pop("deployments", None)
 
         return new_api_config
 
@@ -182,7 +202,7 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
         return create_agent_graphql_client(
             self.client_managed_retries_requests_session,
             self.dagster_cloud_graphql_url,
-            self._dagster_cloud_api_config_for_deployment(None),  # pyright: ignore[reportArgumentType]
+            self._dagster_cloud_api_config_for_deployment(None),
             scope=DagsterCloudInstanceScope.ORGANIZATION,
         )
 
@@ -190,13 +210,13 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
         return create_agent_graphql_client(
             self.client_managed_retries_requests_session,
             self.dagster_cloud_graphql_url,
-            self._dagster_cloud_api_config_for_deployment(deployment_name),  # pyright: ignore[reportArgumentType]
+            self._dagster_cloud_api_config_for_deployment(deployment_name),
             scope=DagsterCloudInstanceScope.DEPLOYMENT,
         )
 
     def headers_for_deployment(self, deployment_name: str):
         return get_agent_headers(
-            self._dagster_cloud_api_config_for_deployment(deployment_name),  # pyright: ignore[reportArgumentType]
+            self._dagster_cloud_api_config_for_deployment(deployment_name),
             DagsterCloudInstanceScope.DEPLOYMENT,
         )
 
@@ -221,11 +241,11 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
             return socket_option
 
     def _socket_options(self):
-        if self._dagster_cloud_api_config.get("socket_options") is None:  # pyright: ignore[reportOptionalMemberAccess]
+        if self._dagster_cloud_api_config.get("socket_options") is None:
             return None
 
         translated_socket_options = []
-        for socket_option in self._dagster_cloud_api_config["socket_options"]:  # pyright: ignore[reportOptionalSubscript]
+        for socket_option in self._dagster_cloud_api_config["socket_options"]:
             check.invariant(
                 len(socket_option) == 3, "Each socket option must be a list of three values"
             )
@@ -267,7 +287,7 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
                     adapter_kwargs=dict(
                         max_retries=Retry(
                             total=self.dagster_cloud_api_retries,
-                            backoff_factor=self._dagster_cloud_api_config["backoff_factor"],  # pyright: ignore[reportOptionalSubscript]
+                            backoff_factor=self._dagster_cloud_api_config["backoff_factor"],
                         ),
                         socket_options=self._socket_options(),
                     )
@@ -297,21 +317,26 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
 
     @property
     def dagster_cloud_url(self):
-        if "url" in self._dagster_cloud_api_config:  # pyright: ignore[reportOperatorIssue]
-            return self._dagster_cloud_api_config["url"]  # pyright: ignore[reportOptionalSubscript]
+        if "url" in self._dagster_cloud_api_config:
+            return self._dagster_cloud_api_config["url"]
 
-        organization = get_organization_name_from_agent_token(self.dagster_cloud_agent_token)  # pyright: ignore[reportArgumentType]
+        organization, region = decode_agent_token(self.dagster_cloud_agent_token)
         if not organization:
             raise DagsterInvariantViolationError(
                 "Could not derive Dagster Cloud URL from agent token. Create a new agent token or"
                 " set the `url` field under `dagster_cloud_api` in your `dagster.yaml`."
             )
 
-        return f"https://{organization}.agent.dagster.cloud"
+        return (
+            f"https://{organization}.agent.{region}.dagster.cloud"
+            if region
+            else f"https://{organization}.agent.dagster.cloud"
+        )
 
     @property
     def organization_name(self) -> Optional[str]:
-        return get_organization_name_from_agent_token(self.dagster_cloud_agent_token)  # pyright: ignore[reportArgumentType]
+        organization, _ = decode_agent_token(self.dagster_cloud_agent_token)
+        return organization
 
     @property
     def deployment_name(self) -> Optional[str]:
@@ -326,26 +351,31 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
 
     @property
     def deployment_names(self) -> list[str]:
-        if self._dagster_cloud_api_config.get("deployment"):  # pyright: ignore[reportOptionalMemberAccess]
-            return [self._dagster_cloud_api_config["deployment"]]  # pyright: ignore[reportOptionalSubscript]
+        if self._dagster_cloud_api_config.get("deployment"):
+            return [self._dagster_cloud_api_config["deployment"]]
 
-        return self._dagster_cloud_api_config.get("deployments", [])  # pyright: ignore[reportOptionalMemberAccess]
+        return self._dagster_cloud_api_config.get("deployments", [])
 
     @property
     def include_all_serverless_deployments(self) -> bool:
-        return self._dagster_cloud_api_config.get("all_serverless_deployments")  # pyright: ignore[reportOptionalMemberAccess,reportReturnType]
+        return self._dagster_cloud_api_config.get("all_serverless_deployments") or False
 
     @property
     def dagit_url(self):
-        organization = get_organization_name_from_agent_token(self.dagster_cloud_agent_token)  # pyright: ignore[reportArgumentType]
+        organization, region = decode_agent_token(self.dagster_cloud_agent_token)
         if not organization:
             raise Exception(
                 "Could not derive Dagster Cloud URL from agent token to generate a Dagit URL."
                 " Generate a new agent token in the Dagit UI."
             )
 
-        deployment = self._dagster_cloud_api_config.get("deployment")  # pyright: ignore[reportOptionalMemberAccess]
-        return f"https://{organization}.dagster.cloud/" + (f"{deployment}/" if deployment else "")
+        deployment = self._dagster_cloud_api_config.get("deployment")
+        base_url = (
+            f"https://{organization}.{region}.dagster.cloud/"
+            if region
+            else f"https://{organization}.dagster.cloud/"
+        )
+        return base_url + (f"{deployment}/" if deployment else "")
 
     @property
     def dagster_cloud_graphql_url(self):
@@ -395,33 +425,38 @@ class DagsterCloudAgentInstance(DagsterCloudInstance):
         return get_agent_headers(self._dagster_cloud_api_config, scope=scope)  # pyright: ignore[reportArgumentType]
 
     @property
-    def dagster_cloud_agent_token(self):
-        return self._dagster_cloud_api_config.get("agent_token")  # pyright: ignore[reportOptionalMemberAccess]
+    def dagster_cloud_agent_token(self) -> str:
+        check.invariant(
+            self._dagster_cloud_api_config.get("agent_token") is not None,
+            "No agent token found in dagster_cloud_api configuration. An agent token is required"
+            " for Dagster Cloud authentication.",
+        )
+        return self._dagster_cloud_api_config["agent_token"]
 
     @property
     def dagster_cloud_api_retries(self) -> int:
-        return self._dagster_cloud_api_config["retries"]  # pyright: ignore[reportOptionalSubscript]
+        return self._dagster_cloud_api_config["retries"]
 
     @property
     def dagster_cloud_api_timeout(self) -> int:
-        return self._dagster_cloud_api_config["timeout"]  # pyright: ignore[reportOptionalSubscript]
+        return self._dagster_cloud_api_config["timeout"]
 
     @property
     def dagster_cloud_api_proxies(self) -> Optional[dict[str, str]]:
         # Requests library modifies the proxies key so create a copy
         return (
             self._dagster_cloud_api_config.get("proxies").copy()  # pyright: ignore[reportOptionalMemberAccess]
-            if self._dagster_cloud_api_config.get("proxies")  # pyright: ignore[reportOptionalMemberAccess]
+            if self._dagster_cloud_api_config.get("proxies")
             else {}
         )
 
     @property
     def dagster_cloud_api_agent_label(self) -> Optional[str]:
-        return self._dagster_cloud_api_config.get("agent_label")  # pyright: ignore[reportOptionalMemberAccess]
+        return self._dagster_cloud_api_config.get("agent_label")
 
     @property
     def includes_branch_deployments(self) -> bool:
-        return self._dagster_cloud_api_config.get("branch_deployments", False)  # pyright: ignore[reportOptionalMemberAccess]
+        return self._dagster_cloud_api_config.get("branch_deployments", False)
 
     @property
     def instance_uuid(self) -> str:
@@ -578,6 +613,10 @@ instance_class:
             "dagster_cloud.secrets", "DagsterCloudSecretsLoader", empty_yaml
         )
 
+        defaults["defs_state_storage"] = ConfigurableClassData(  # pyright: ignore[reportIndexIssue]
+            "dagster_cloud.storage.defs_state", "GraphQLDefsStateStorage", empty_yaml
+        )
+
         return defaults
 
     def dispose(self) -> None:

dagster_cloud/metadata/source_code.py

@@ -14,7 +14,9 @@ from dagster._core.types.loadable_target_origin import LoadableTargetOrigin
 
 if TYPE_CHECKING:
     from dagster import AssetsDefinition, AssetSpec, SourceAsset
-    from dagster._core.definitions.cacheable_assets import CacheableAssetsDefinition
+    from dagster._core.definitions.assets.definition.cacheable_assets_definition import (
+        CacheableAssetsDefinition,
+    )
 
 import sys
 

dagster_cloud/opentelemetry/config/exporter.py

@@ -1,6 +1,6 @@
 from dagster import BoolSource, Enum, Field, IntSource, Map, Permissive, Shape, StringSource
 
-from ..enum import (
+from dagster_cloud.opentelemetry.enum import (
     AggregationTemporalityEnum,
     CompressionEnum,
     LoggingExporterEnum,

dagster_cloud/opentelemetry/controller.py

@@ -251,7 +251,7 @@ class OpenTelemetryController:
         )
         self._logging_handler = logs_factory.build_logging_handler(
             self._logger_provider,  # pyright: ignore[reportArgumentType]
-            logging_config.get("handler", {}),  # pyright: ignore[reportArgumentType]
+            logging_config.get("handler", {}),
         )
         self._logger.addHandler(self._logging_handler)
 

dagster_cloud/opentelemetry/observers/dagster_exception_handler.py

@@ -21,7 +21,7 @@ def extract_dagster_error_attributes(error: DagsterError) -> dict[str, str]:
         if type(attr) is list and all(isinstance(item, SerializableErrorInfo) for item in attr):
             serializable_error_infos: list[SerializableErrorInfo] = getattr(error, field)
             for serializable_error_info in serializable_error_infos:
-                if serializable_error_info.cause:
+                if serializable_error_info.cause and serializable_error_info.cause.cls_name:
                     error_attributes["cause"] = serializable_error_info.cause.cls_name
                     return error_attributes
 

dagster_cloud/opentelemetry/observers/execution_observer.py

@@ -5,8 +5,10 @@ from typing import Callable, Optional
 from dagster import DagsterError
 from dagster._time import get_current_timestamp
 
-from ..controller import OpenTelemetryController
-from .dagster_exception_handler import extract_dagster_error_attributes
+from dagster_cloud.opentelemetry.controller import OpenTelemetryController
+from dagster_cloud.opentelemetry.observers.dagster_exception_handler import (
+    extract_dagster_error_attributes,
+)
 
 ResultEvaluatorCallback = Callable[..., str]
 

dagster_cloud/pex/grpc/__init__.py

@@ -1,8 +1,8 @@
-from .client import (
+from dagster_cloud.pex.grpc.client import (
     MultiPexGrpcClient as MultiPexGrpcClient,
     wait_for_grpc_server as wait_for_grpc_server,
 )
-from .types import (
+from dagster_cloud.pex.grpc.types import (
     CreatePexServerArgs as CreatePexServerArgs,
     GetPexServersArgs as GetPexServersArgs,
     PexServerHandle as PexServerHandle,

dagster_cloud/pex/grpc/client.py

@@ -12,8 +12,8 @@ from dagster._serdes import serialize_value
 from dagster._utils.error import SerializableErrorInfo, serializable_error_info_from_exc_info
 from dagster_shared.serdes.serdes import deserialize_value
 
-from .__generated__ import MultiPexApiStub, multi_pex_api_pb2
-from .types import (
+from dagster_cloud.pex.grpc.__generated__ import MultiPexApiStub, multi_pex_api_pb2
+from dagster_cloud.pex.grpc.types import (
     CreatePexServerArgs,
     CreatePexServerResponse,
     GetCrashedPexServersArgs,

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

@@ -9,8 +9,8 @@ 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)
 

dagster_cloud/pex/grpc/server/manager.py

@@ -18,11 +18,10 @@ from dagster_shared import seven
 from dagster_shared.ipc import open_ipc_subprocess
 from pydantic import BaseModel, Extra
 
+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
 
-from ..types import PexServerHandle
-from .registry import PexS3Registry
-
 logger = logging.getLogger("dagster.multipex")
 
 
@@ -401,7 +400,9 @@ class MultiPexManager(AbstractContextManager):
             # - 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

@@ -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"]
@@ -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

@@ -14,7 +14,7 @@ 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):
@@ -189,7 +189,12 @@ class DagsterPexProxyApiServer(DagsterApiServicer):
             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)

dagster_cloud/secrets/__init__.py

@@ -1 +1 @@
-from .loader import DagsterCloudSecretsLoader as DagsterCloudSecretsLoader
+from dagster_cloud.secrets.loader import DagsterCloudSecretsLoader as DagsterCloudSecretsLoader

dagster_cloud/serverless/__init__.py

@@ -1 +1 @@
-from .io_manager import serverless_io_manager as serverless_io_manager
+from dagster_cloud.serverless.io_manager import serverless_io_manager as serverless_io_manager

dagster_cloud/serverless/io_manager.py

@@ -47,7 +47,10 @@ class PickledObjectServerlessIOManager(UPathIOManager):
             self._boto_session_expiration.tzinfo
         ) + datetime.timedelta(minutes=5):
             self._boto_session, self._boto_session_expiration = self._refresh_boto_session()
-        return self._boto_session.client("s3", region_name="us-west-2")
+        return self._boto_session.client(
+            "s3",
+            region_name=os.getenv("DAGSTER_CLOUD_SERVERLESS_REGION", "us-west-2"),
+        )
 
     def load_from_path(self, context: InputContext, path: UPath) -> Any:
         try:

dagster_cloud/storage/compute_logs/__init__.py

@@ -1 +1,3 @@
-from .compute_log_manager import CloudComputeLogManager as CloudComputeLogManager
+from dagster_cloud.storage.compute_logs.compute_log_manager import (
+    CloudComputeLogManager as CloudComputeLogManager,
+)