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

dagster_cloud/util/errors.py

@@ -1,66 +1,51 @@
-from dagster._serdes import serialize_value
-from dagster._utils.error import SerializableErrorInfo
-
-ERROR_CLASS_NAME_SIZE_LIMIT = 1000
-
-
-def truncate_serialized_error(
-    error_info: SerializableErrorInfo, field_size_limit: int, max_depth: int
-):
-    if error_info.cause:
-        if max_depth == 0:
-            new_cause = (
-                error_info.cause
-                if len(serialize_value(error_info.cause)) <= field_size_limit
-                else SerializableErrorInfo(
-                    message="(Cause truncated due to size limitations)",
-                    stack=[],
-                    cls_name=None,
-                )
-            )
-        else:
-            new_cause = truncate_serialized_error(
-                error_info.cause, field_size_limit, max_depth=max_depth - 1
-            )
-        error_info = error_info._replace(cause=new_cause)
-
-    if error_info.context:
-        if max_depth == 0:
-            new_context = (
-                error_info.context
-                if len(serialize_value(error_info.context)) <= field_size_limit
-                else SerializableErrorInfo(
-                    message="(Context truncated due to size limitations)",
-                    stack=[],
-                    cls_name=None,
-                )
-            )
-        else:
-            new_context = truncate_serialized_error(
-                error_info.context, field_size_limit, max_depth=max_depth - 1
-            )
-        error_info = error_info._replace(context=new_context)
+from collections.abc import Sequence
 
-    stack_size_so_far = 0
-    truncated_stack = []
-    for stack_elem in error_info.stack:
-        stack_size_so_far += len(stack_elem)
-        if stack_size_so_far > field_size_limit:
-            truncated_stack.append("(TRUNCATED)")
-            break
-
-        truncated_stack.append(stack_elem)
-
-    error_info = error_info._replace(stack=truncated_stack)
-
-    if len(error_info.message) > field_size_limit:
-        error_info = error_info._replace(
-            message=error_info.message[:field_size_limit] + " (TRUNCATED)"
-        )
-
-    if error_info.cls_name and len(error_info.cls_name) > ERROR_CLASS_NAME_SIZE_LIMIT:
-        error_info = error_info._replace(
-            cls_name=error_info.cls_name[:ERROR_CLASS_NAME_SIZE_LIMIT] + " (TRUNCATED)"
-        )
+from dagster._utils.error import SerializableErrorInfo
 
-    return error_info
+DAGSTER_FRAMEWORK_SUBSTRINGS = [
+    "/site-packages/dagster/",
+    "/python_modules/dagster/dagster",
+]
+
+
+def remove_dagster_framework_lines_from_serializable_exc_info(error_info: SerializableErrorInfo):
+    return error_info._replace(
+        stack=remove_dagster_framework_lines_from_stack_trace(error_info.stack),
+        cause=(
+            remove_dagster_framework_lines_from_serializable_exc_info(error_info.cause)
+            if error_info.cause
+            else None
+        ),
+        context=(
+            remove_dagster_framework_lines_from_serializable_exc_info(error_info.context)
+            if error_info.context
+            else None
+        ),
+    )
+
+
+def remove_dagster_framework_lines_from_stack_trace(
+    stack: Sequence[str],
+) -> Sequence[str]:
+    # Remove lines until you find the first non-dagster framework line
+
+    for i in range(len(stack)):
+        if not _line_contains_dagster_framework_file(stack[i]):
+            return stack[i:]
+
+    # Return the full stack trace if its all Dagster framework lines,
+    # to not be left with an empty stack trace
+    return stack
+
+
+def _line_contains_dagster_framework_file(line: str):
+    # stack trace line starts with something like
+    # File "/usr/local/lib/python3.11/site-packages/dagster/_core/execution/plan/utils.py",
+    split_by_comma = line.split(",")
+    if not split_by_comma:
+        return False
+
+    file_portion = split_by_comma[0]
+    return any(
+        framework_substring in file_portion for framework_substring in DAGSTER_FRAMEWORK_SUBSTRINGS
+    )

dagster_cloud/version.py

@@ -1 +1 @@
-__version__ = "1.8.2"
+__version__ = "1.12.6"

dagster_cloud/workspace/config_schema/__init__.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, Optional, cast
 
 from dagster import (
     Enum as DagsterEnum,
@@ -13,15 +13,15 @@ from dagster import (
 )
 from dagster._config import EvaluationError, StringSource, validate_config
 
-from .docker import SHARED_DOCKER_CONFIG
-from .ecs import (
+from dagster_cloud.workspace.config_schema.docker import SHARED_DOCKER_CONFIG
+from dagster_cloud.workspace.config_schema.ecs import (
     ECS_CONTAINER_CONTEXT_CONFIG as ECS_CONTAINER_CONTEXT_CONFIG,
     SHARED_ECS_CONFIG as SHARED_ECS_CONFIG,
 )
-from .kubernetes import SHARED_K8S_CONFIG
+from dagster_cloud.workspace.config_schema.kubernetes import SHARED_K8S_CONFIG
 
 
-def validate_workspace_location(workspace_location) -> Optional[List[str]]:
+def validate_workspace_location(workspace_location) -> Optional[list[str]]:
     """Processes a single workspace location config. Returns a list of error
     messages if any.
     """
@@ -29,7 +29,7 @@ def validate_workspace_location(workspace_location) -> Optional[List[str]]:
     return [error.message for error in validation.errors or []]
 
 
-def validate_workspace_config(workspace_config) -> Optional[List[str]]:
+def validate_workspace_config(workspace_config) -> Optional[list[str]]:
     """Processes an entire workspace location config. Returns a list of
     error messages, if any.
     """
@@ -37,7 +37,7 @@ def validate_workspace_config(workspace_config) -> Optional[List[str]]:
     return [error.message for error in validation.errors or []]
 
 
-def process_workspace_config(workspace_config) -> Dict[str, Any]:
+def process_workspace_config(workspace_config) -> dict[str, Any]:
     """Checks a workspace config, erroring if any mismatches with config
     and migrating an input in the legacy workspace config format to the
     modern format, returning the validated input.
@@ -62,8 +62,21 @@ def process_workspace_config(workspace_config) -> Dict[str, Any]:
             python_file = config.get("python_file")
             package_name = config.get("package_name")
             module_name = config.get("module_name")
+            autoload_defs_module_name = config.get("autoload_defs_module_name")
             check.invariant(
-                len([val for val in [python_file, package_name, module_name] if val]) == 1,
+                len(
+                    [
+                        val
+                        for val in [
+                            python_file,
+                            package_name,
+                            module_name,
+                            autoload_defs_module_name,
+                        ]
+                        if val
+                    ]
+                )
+                == 1,
                 "Must supply exactly one of a file name, a package name, or a module name",
             )
 
@@ -73,7 +86,8 @@ def process_workspace_config(workspace_config) -> Dict[str, Any]:
             new_location = {
                 k: v
                 for k, v in location.items()
-                if k not in ("python_file", "package_name", "module_name")
+                if k
+                not in ("python_file", "package_name", "module_name", "autoload_defs_module_name")
             }
             new_location["code_source"] = {}
             if "python_file" in location:
@@ -82,6 +96,10 @@ def process_workspace_config(workspace_config) -> Dict[str, Any]:
                 new_location["code_source"]["package_name"] = location["package_name"]
             elif "module_name" in location:
                 new_location["code_source"]["module_name"] = location["module_name"]
+            elif "autoload_defs_module_name" in location:
+                new_location["code_source"]["autoload_defs_module_name"] = location[
+                    "autoload_defs_module_name"
+                ]
 
             new_location["location_name"] = name
             updated_locations.append(new_location)
@@ -90,10 +108,12 @@ def process_workspace_config(workspace_config) -> Dict[str, Any]:
         check.is_list(workspace_config.get("locations"))
 
         validation = validate_config(WORKSPACE_CONFIG_SCHEMA, workspace_config)
-        check.invariant(
-            validation.success,
-            ", ".join([error.message for error in cast(List[EvaluationError], validation.errors)]),
-        )
+        if not validation.success:
+            check.failed(
+                ", ".join(
+                    [error.message for error in cast("list[EvaluationError]", validation.errors)]
+                ),
+            )
         return workspace_config
 
 
@@ -139,6 +159,8 @@ K8S_CONFIG_FIELDS = {
                     ),
                     is_required=False,
                 ),
+                "deployment_metadata": Permissive(),
+                "service_metadata": Permissive(),
             }
         ),
         is_required=False,
@@ -190,6 +212,10 @@ CONFIG_SCHEMA_FIELDS = {
                     config=str,
                     description="Python module containing the target Dagster repository.",
                 ),
+                "autoload_defs_module_name": Field(
+                    config=str,
+                    description="Python module to automatically load Dagster definitions from.",
+                ),
             },
         ),
         description="Python entry point for the code location.",
@@ -250,6 +276,17 @@ CONFIG_SCHEMA_FIELDS = {
         is_required=False,
         description="Locations that specify an agent queue will only have their requests handled by agents configured to read from a matching queue. By default, requests are placed on a default queue that's handled by all agents.",
     ),
+    "defs_state_info": Field(
+        config=Shape(
+            fields={
+                "info_mapping": Map(
+                    str, Noneable(Shape(fields={"version": str, "create_timestamp": float}))
+                )
+            },
+        ),
+        is_required=False,
+        description="Defs state info for the code location.",
+    ),
 }
 
 
@@ -278,6 +315,11 @@ LEGACY_CONFIG_SCHEMA_FIELDS = {
         is_required=False,
         description="Python module containing the target Dagster repository.",
     ),
+    "autoload_defs_module_name": Field(
+        config=str,
+        is_required=False,
+        description="Python module to automatically load Dagster definitions from.",
+    ),
 }
 LEGACY_LOCATION_CONFIG_SCHEMA = Shape(fields=LEGACY_CONFIG_SCHEMA_FIELDS)
 LEGACY_NAMED_LOCATIONS_CONFIG_SCHEMA = Map(

dagster_cloud/workspace/docker/__init__.py

@@ -1,7 +1,7 @@
 import logging
 import sys
 import time
-from typing import Any, Dict, List, NamedTuple, Optional, Tuple
+from typing import Any, NamedTuple, Optional
 
 import docker
 import docker.errors
@@ -10,36 +10,44 @@ from dagster import (
     IntSource,
     _check as check,
 )
+from dagster._core.launcher.base import LaunchRunContext
 from dagster._core.utils import parse_env_var
+from dagster._grpc.types import ExecuteRunArgs
 from dagster._serdes import ConfigurableClass
 from dagster._serdes.config_class import ConfigurableClassData
 from dagster._utils import find_free_port
 from dagster._utils.merger import merge_dicts
 from dagster._vendored.dateutil.parser import parse
+from dagster_cloud_cli.core.workspace import PexMetadata
 from dagster_docker import DockerRunLauncher
 from dagster_docker.container_context import DockerContainerContext
+from dagster_shared.serdes.serdes import deserialize_value, serialize_value
 from docker.models.containers import Container
 from typing_extensions import Self
 
 from dagster_cloud.api.dagster_cloud_api import UserCodeDeploymentType
 from dagster_cloud.execution.monitoring import CloudContainerResourceLimits
-from dagster_cloud.workspace.user_code_launcher.user_code_launcher import UserCodeLauncherEntry
-
-from ..config_schema.docker import SHARED_DOCKER_CONFIG
-from ..user_code_launcher import (
+from dagster_cloud.storage.tags import PEX_METADATA_TAG
+from dagster_cloud.workspace.config_schema.docker import SHARED_DOCKER_CONFIG
+from dagster_cloud.workspace.docker.utils import unique_docker_resource_name
+from dagster_cloud.workspace.user_code_launcher import (
     DEFAULT_SERVER_PROCESS_STARTUP_TIMEOUT,
     SHARED_USER_CODE_LAUNCHER_CONFIG,
     DagsterCloudGrpcServer,
     DagsterCloudUserCodeLauncher,
     ServerEndpoint,
 )
-from ..user_code_launcher.utils import deterministic_label_for_location
-from .utils import unique_docker_resource_name
+from dagster_cloud.workspace.user_code_launcher.user_code_launcher import UserCodeLauncherEntry
+from dagster_cloud.workspace.user_code_launcher.utils import (
+    deterministic_label_for_location,
+    get_grpc_server_env,
+)
 
 GRPC_SERVER_LABEL = "dagster_grpc_server"
 MULTIPEX_SERVER_LABEL = "dagster_multipex_server"
 AGENT_LABEL = "dagster_agent_id"
 SERVER_TIMESTAMP_LABEL = "dagster_server_timestamp"
+STOP_TIMEOUT_LABEL = "dagster_stop_timeout"
 
 
 IMAGE_PULL_LOG_INTERVAL = 15
@@ -52,7 +60,7 @@ class DagsterDockerContainer(NamedTuple):
 
     container: Container
 
-    def __str__(self):
+    def __str__(self):  # pyright: ignore[reportIncompatibleMethodOverride]
         return self.container.id
 
 
@@ -76,7 +84,7 @@ class DockerUserCodeLauncher(
             container_kwargs, "container_kwargs", key_type=str
         )
 
-        super(DockerUserCodeLauncher, self).__init__(**kwargs)
+        super().__init__(**kwargs)
 
     @property
     def requires_images(self):
@@ -90,11 +98,11 @@ class DockerUserCodeLauncher(
         return self._inst_data
 
     @property
-    def env_vars(self) -> List[str]:
+    def env_vars(self) -> list[str]:
         return self._input_env_vars + self._instance.dagster_cloud_api_env_vars
 
     @property
-    def container_kwargs(self) -> Dict[str, Any]:
+    def container_kwargs(self) -> dict[str, Any]:
         return self._container_kwargs
 
     @classmethod
@@ -130,6 +138,10 @@ class DockerUserCodeLauncher(
         command,
         labels,
     ):
+        container_kwargs = {**container_context.container_kwargs}
+
+        container_kwargs.pop("stop_timeout", None)
+
         return client.containers.create(
             image,
             detach=True,
@@ -140,12 +152,12 @@ class DockerUserCodeLauncher(
             labels=labels,
             command=command,
             ports=ports,
-            **container_context.container_kwargs,
+            **container_kwargs,
         )
 
     def _get_standalone_dagster_server_handles_for_location(
         self, deployment_name: str, location_name: str
-    ) -> List[DagsterDockerContainer]:
+    ) -> list[DagsterDockerContainer]:
         client = docker.client.from_env()
         return [
             DagsterDockerContainer(container=container)
@@ -163,7 +175,7 @@ class DockerUserCodeLauncher(
 
     def _get_multipex_server_handles_for_location(
         self, deployment_name: str, location_name: str
-    ) -> List[DagsterDockerContainer]:
+    ) -> list[DagsterDockerContainer]:
         client = docker.client.from_env()
         return [
             DagsterDockerContainer(container=container)
@@ -186,13 +198,13 @@ class DockerUserCodeLauncher(
         container_name: str,
         hostname: str,
         grpc_port: int,
-        ports: Dict[int, int],
+        ports: dict[int, int],
         image: str,
         container_context: DockerContainerContext,
-        command: List[str],
-        additional_env: Dict[str, str],
-        labels: Dict[str, str],
-    ) -> Tuple[Container, ServerEndpoint]:
+        command: list[str],
+        additional_env: dict[str, str],
+        labels: dict[str, str],
+    ) -> tuple[Container, ServerEndpoint]:
         client = docker.client.from_env()
 
         self._logger.info(
@@ -304,8 +316,11 @@ class DockerUserCodeLauncher(
             command = metadata.get_grpc_server_command(
                 metrics_enabled=self._instance.user_code_launcher.code_server_metrics_enabled
             )
-            environment = metadata.get_grpc_server_env(
-                grpc_port, location_name, self._instance.ref_for_deployment(deployment_name)
+            environment = get_grpc_server_env(
+                metadata,
+                grpc_port,
+                location_name,
+                self._instance.ref_for_deployment(deployment_name),
             )
             labels = {
                 GRPC_SERVER_LABEL: "",
@@ -314,6 +329,9 @@ class DockerUserCodeLauncher(
                 SERVER_TIMESTAMP_LABEL: str(desired_entry.update_timestamp),
             }
 
+        if "stop_timeout" in container_context.container_kwargs:
+            labels[STOP_TIMEOUT_LABEL] = str(container_context.container_kwargs["stop_timeout"])
+
         container, server_endpoint = self._launch_container(
             deployment_name,
             location_name,
@@ -347,8 +365,12 @@ class DockerUserCodeLauncher(
 
     def _remove_server_handle(self, server_handle: DagsterDockerContainer) -> None:
         container = server_handle.container
+
+        stop_timeout_str = server_handle.container.labels.get(STOP_TIMEOUT_LABEL)
+        stop_timeout = int(stop_timeout_str) if stop_timeout_str else None
+
         try:
-            container.stop()
+            container.stop(timeout=stop_timeout)
         except Exception:
             self._logger.error(f"Failure stopping container {container.id}: {sys.exc_info()}")
         container.remove(force=True)
@@ -359,9 +381,9 @@ class DockerUserCodeLauncher(
 
     def get_server_create_timestamp(self, handle: DagsterDockerContainer) -> Optional[float]:
         created_time_str = handle.container.attrs["Created"]
-        return parse(created_time_str).timestamp()
+        return parse(created_time_str).timestamp()  # pyright: ignore[reportAttributeAccessIssue]
 
-    def _list_server_handles(self) -> List[DagsterDockerContainer]:
+    def _list_server_handles(self) -> list[DagsterDockerContainer]:
         client = docker.client.from_env()
         return [
             DagsterDockerContainer(container=container)
@@ -374,7 +396,7 @@ class DockerUserCodeLauncher(
         ]
 
     def run_launcher(self):
-        launcher = DockerRunLauncher(
+        launcher = CloudDockerRunLauncher(
             image=None,
             env_vars=self.env_vars,
             networks=self._networks,
@@ -383,3 +405,32 @@ class DockerUserCodeLauncher(
         launcher.register_instance(self._instance)
 
         return launcher
+
+
+class CloudDockerRunLauncher(DockerRunLauncher):
+    def launch_run(self, context: LaunchRunContext) -> None:
+        serialized_pex_metadata = context.dagster_run.tags.get(PEX_METADATA_TAG)
+
+        if serialized_pex_metadata:
+            run = context.dagster_run
+            job_origin = check.not_none(run.job_code_origin)
+
+            docker_image = self._get_docker_image(job_origin)
+
+            run_args = ExecuteRunArgs(
+                job_origin=job_origin,
+                run_id=run.run_id,
+                instance_ref=self._instance.get_ref(),
+            )
+
+            deserialize_value(serialized_pex_metadata, PexMetadata)
+            command = [
+                "dagster-cloud",
+                "pex",
+                "execute-run",
+                serialize_value(run_args),
+                serialized_pex_metadata,
+            ]
+            self._launch_container_with_command(run, docker_image, command)
+        else:
+            return super().launch_run(context)

dagster_cloud/workspace/docker/utils.py

@@ -1,6 +1,6 @@
 import re
 
-from ..user_code_launcher.utils import unique_resource_name
+from dagster_cloud.workspace.user_code_launcher.utils import unique_resource_name
 
 
 def unique_docker_resource_name(deployment_name, location_name):

dagster_cloud/workspace/ecs/__init__.py

@@ -1 +1 @@
-from .launcher import EcsUserCodeLauncher as EcsUserCodeLauncher
+from dagster_cloud.workspace.ecs.launcher import EcsUserCodeLauncher as EcsUserCodeLauncher