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

dagster_cloud/workspace/kubernetes/launcher.py

@@ -1,8 +1,9 @@
 import logging
 from collections import defaultdict
+from collections.abc import Collection, Mapping
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Any, Collection, Dict, List, Mapping, NamedTuple, Optional, Set, Tuple
+from typing import Any, NamedTuple, Optional
 
 import kubernetes
 import kubernetes.client as client
@@ -21,6 +22,7 @@ from dagster._serdes import ConfigurableClass
 from dagster._serdes.config_class import ConfigurableClassData
 from dagster._utils.merger import merge_dicts
 from dagster_cloud_cli.core.workspace import CodeLocationDeployData
+from dagster_k8s.client import PatchedApiClient
 from dagster_k8s.container_context import K8sContainerContext
 from dagster_k8s.job import UserDefinedDagsterK8sConfig
 from dagster_k8s.models import k8s_snake_case_dict
@@ -31,17 +33,7 @@ from dagster_cloud.api.dagster_cloud_api import UserCodeDeploymentType
 from dagster_cloud.constants import RESERVED_ENV_VAR_NAMES
 from dagster_cloud.execution.cloud_run_launcher.k8s import CloudK8sRunLauncher
 from dagster_cloud.execution.monitoring import CloudContainerResourceLimits
-
-from ..user_code_launcher import (
-    DEFAULT_SERVER_PROCESS_STARTUP_TIMEOUT,
-    SHARED_USER_CODE_LAUNCHER_CONFIG,
-    DagsterCloudGrpcServer,
-    DagsterCloudUserCodeLauncher,
-    ServerEndpoint,
-    UserCodeLauncherEntry,
-)
-from ..user_code_launcher.utils import deterministic_label_for_location
-from .utils import (
+from dagster_cloud.workspace.kubernetes.utils import (
     SERVICE_PORT,
     construct_code_location_deployment,
     construct_code_location_service,
@@ -49,11 +41,20 @@ from .utils import (
     unique_k8s_resource_name,
     wait_for_deployment_complete,
 )
+from dagster_cloud.workspace.user_code_launcher import (
+    DEFAULT_SERVER_PROCESS_STARTUP_TIMEOUT,
+    SHARED_USER_CODE_LAUNCHER_CONFIG,
+    DagsterCloudGrpcServer,
+    DagsterCloudUserCodeLauncher,
+    ServerEndpoint,
+    UserCodeLauncherEntry,
+)
+from dagster_cloud.workspace.user_code_launcher.utils import deterministic_label_for_location
 
 DEFAULT_DEPLOYMENT_STARTUP_TIMEOUT = 300
 DEFAULT_IMAGE_PULL_GRACE_PERIOD = 30
 
-from ..config_schema.kubernetes import SHARED_K8S_CONFIG
+from dagster_cloud.workspace.config_schema.kubernetes import SHARED_K8S_CONFIG
 
 
 class K8sHandle(NamedTuple):
@@ -94,6 +95,7 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
         security_context=None,
         only_allow_user_defined_k8s_config_fields=None,
         only_allow_user_defined_env_vars=None,
+        fail_pod_on_run_failure=None,
         **kwargs,
     ):
         self._inst_data = inst_data
@@ -176,7 +178,7 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
         self._only_allow_user_defined_k8s_config_fields = only_allow_user_defined_k8s_config_fields
         self._only_allow_user_defined_env_vars = only_allow_user_defined_env_vars
 
-        super(K8sUserCodeLauncher, self).__init__(**kwargs)
+        super().__init__(**kwargs)
 
         self._launcher = CloudK8sRunLauncher(
             dagster_home=self._dagster_home,
@@ -189,22 +191,32 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
             env_config_maps=self._env_config_maps,
             env_secrets=self._env_secrets,
             env_vars=self._env_vars,
-            job_namespace=self._namespace,
+            job_namespace=self._namespace,  # pyright: ignore[reportArgumentType]
             volume_mounts=self._volume_mounts,
             volumes=self._volumes,
             labels=self._labels,
             resources=self._resources,
             scheduler_name=self._scheduler_name,
             run_k8s_config=self._run_k8s_config,
+            fail_pod_on_run_failure=fail_pod_on_run_failure,
             kubeconfig_file=kubeconfig_file,
             load_incluster_config=not kubeconfig_file,
             security_context=self._security_context,
-            only_allow_user_defined_k8s_config_fields=self._only_allow_user_defined_k8s_config_fields,
+            # leave out the code server specific fields from only_allow_user_defined_k8s_config_fields
+            only_allow_user_defined_k8s_config_fields=(
+                {
+                    key: val
+                    for key, val in only_allow_user_defined_k8s_config_fields.items()
+                    if key not in {"service_metadata", "deployment_metadata"}
+                }
+                if only_allow_user_defined_k8s_config_fields
+                else only_allow_user_defined_k8s_config_fields
+            ),
             only_allow_user_defined_env_vars=self._only_allow_user_defined_env_vars,
         )
 
         # mutable set of observed namespaces to assist with cleanup
-        self._used_namespaces: Dict[Tuple[str, str], Set[str]] = defaultdict(set)
+        self._used_namespaces: dict[tuple[str, str], set[str]] = defaultdict(set)
 
     @property
     def requires_images(self):
@@ -272,12 +284,22 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
                     is_required=False,
                     description="Raw Kubernetes configuration for launched runs.",
                 ),
+                "fail_pod_on_run_failure": Field(
+                    bool,
+                    is_required=False,
+                    description=(
+                        "Whether the launched Kubernetes Jobs and Pods should fail if the Dagster"
+                        " run fails"
+                    ),
+                ),
                 "server_k8s_config": Field(
                     Shape(
                         {
                             "container_config": Permissive(),
                             "pod_spec_config": Permissive(),
                             "pod_template_spec_metadata": Permissive(),
+                            "deployment_metadata": Permissive(),
+                            "service_metadata": Permissive(),
                         }
                     ),
                     is_required=False,
@@ -301,6 +323,12 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
                             "job_spec_config": Field(
                                 Map(key_type=str, inner_type=bool), is_required=False
                             ),
+                            "deployment_metadata": Field(
+                                Map(key_type=str, inner_type=bool), is_required=False
+                            ),
+                            "service_metadata": Field(
+                                Map(key_type=str, inner_type=bool), is_required=False
+                            ),
                             "namespace": Field(bool, is_required=False),
                         }
                     ),
@@ -341,7 +369,11 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
         )
 
     def _get_core_api_client(self):
-        return self._k8s_core_api_client if self._k8s_core_api_client else client.CoreV1Api()
+        return (
+            self._k8s_core_api_client
+            if self._k8s_core_api_client
+            else client.CoreV1Api(api_client=PatchedApiClient())
+        )
 
     @contextmanager
     def _get_apps_api_instance(self):
@@ -349,13 +381,13 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
             yield self._k8s_apps_api_client
             return
 
-        with client.ApiClient() as api_client:
+        with PatchedApiClient() as api_client:
             yield client.AppsV1Api(api_client)
 
     def _resolve_container_context(self, metadata: CodeLocationDeployData) -> K8sContainerContext:
         user_defined_container_context = K8sContainerContext.create_from_config(
             metadata.container_context
-        ).validate_user_k8s_config(
+        ).validate_user_k8s_config_for_code_server(
             self._only_allow_user_defined_k8s_config_fields, self._only_allow_user_defined_env_vars
         )
         return K8sContainerContext(
@@ -423,7 +455,7 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
         try:
             with self._get_apps_api_instance() as api_instance:
                 deployment_reponse = api_instance.create_namespaced_deployment(
-                    container_context.namespace,
+                    namespace=container_context.namespace,
                     body=construct_code_location_deployment(
                         self._instance,
                         deployment_name,
@@ -440,7 +472,7 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
             )
         except ApiException as e:
             self._logger.error(
-                "Exception when calling AppsV1Api->create_namespaced_deployment: %s\n" % e
+                f"Exception when calling AppsV1Api->create_namespaced_deployment: {e}\n"
             )
             raise e
 
@@ -465,7 +497,7 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
             )
         except ApiException as e:
             self._logger.error(
-                "Exception when calling AppsV1Api->create_namespaced_service: %s\n" % e
+                f"Exception when calling AppsV1Api->create_namespaced_service: {e}\n"
             )
             raise e
 
@@ -570,7 +602,7 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
                 for deployment in deployments:
                     handles.append(
                         K8sHandle(
-                            namespace=namespace,
+                            namespace=namespace,  # pyright: ignore[reportArgumentType]
                             name=deployment.metadata.name,
                             labels=deployment.metadata.labels,
                             creation_timestamp=deployment.metadata.creation_timestamp.timestamp()
@@ -581,14 +613,14 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
 
         return handles
 
-    def _list_server_handles(self) -> List[K8sHandle]:
-        namespaces: Set[str] = set()
+    def _list_server_handles(self) -> list[K8sHandle]:
+        namespaces: set[str] = set()
         if self._namespace:
             namespaces.add(self._namespace)
         for namespace_items in self._used_namespaces.values():
             for namespace in namespace_items:
                 namespaces.add(namespace)
-        handles: List[K8sHandle] = []
+        handles: list[K8sHandle] = []
         with self._get_apps_api_instance() as api_instance:
             for namespace in namespaces:
                 deployments = api_instance.list_namespaced_deployment(
@@ -623,7 +655,7 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
                 server_handle.name, server_handle.namespace
             )
         except ApiException as e:
-            if e.reason == "Not found":
+            if e.status == 404:
                 self._logger.exception(
                     f"Tried to delete service {server_handle.name} but it was not found"
                 )
@@ -636,7 +668,7 @@ class K8sUserCodeLauncher(DagsterCloudUserCodeLauncher[K8sHandle], ConfigurableC
                     server_handle.name, server_handle.namespace
                 )
             except ApiException as e:
-                if e.reason == "Not found":
+                if e.status == 404:
                     self._logger.exception(
                         f"Tried to delete deployment {server_handle.name} but it was not found"
                     )

dagster_cloud/workspace/kubernetes/utils.py

@@ -2,7 +2,8 @@ import asyncio
 import copy
 import re
 import time
-from typing import Mapping, Optional
+from collections.abc import Mapping
+from typing import Optional
 
 import kubernetes
 from dagster_k8s.client import DagsterKubernetesClient
@@ -10,9 +11,9 @@ from dagster_k8s.models import k8s_model_from_dict
 from kubernetes import client
 
 from dagster_cloud.instance import DagsterCloudAgentInstance
-
-from ..user_code_launcher.utils import (
+from dagster_cloud.workspace.user_code_launcher.utils import (
     deterministic_label_for_location,
+    get_grpc_server_env,
     get_human_readable_label,
     unique_resource_name,
 )
@@ -84,22 +85,30 @@ def construct_code_location_service(
     instance,
     server_timestamp: float,
 ):
-    labels = container_context.labels
+    user_defined_k8s_config = container_context.server_k8s_config
+    user_defined_service_metadata = copy.deepcopy(dict(user_defined_k8s_config.service_metadata))
+    user_defined_service_labels = user_defined_service_metadata.pop("labels", {})
 
-    return client.V1Service(
-        metadata=client.V1ObjectMeta(
-            name=service_name,
-            labels={
-                **labels,
-                **_get_dagster_k8s_labels(
-                    deployment_name, location_name, instance, server_timestamp
-                ),
+    return k8s_model_from_dict(
+        client.V1Service,
+        {
+            "api_version": "v1",
+            "kind": "Service",
+            "metadata": {
+                **user_defined_service_metadata,
+                "name": service_name,
+                "labels": {
+                    **user_defined_service_labels,
+                    **_get_dagster_k8s_labels(
+                        deployment_name, location_name, instance, server_timestamp
+                    ),
+                },
             },
-        ),
-        spec=client.V1ServiceSpec(
-            selector={"user-deployment": service_name},
-            ports=[client.V1ServicePort(name="grpc", protocol="TCP", port=SERVICE_PORT)],
-        ),
+            "spec": {
+                "selector": {"user-deployment": service_name},
+                "ports": [{"name": "grpc", "protocol": "TCP", "port": SERVICE_PORT}],
+            },
+        },
     )
 
 
@@ -113,12 +122,15 @@ def construct_code_location_deployment(
     args,
     server_timestamp: float,
 ):
-    env = metadata.get_grpc_server_env(
-        SERVICE_PORT, location_name, instance.ref_for_deployment(deployment_name)
+    env = get_grpc_server_env(
+        metadata, SERVICE_PORT, location_name, instance.ref_for_deployment(deployment_name)
     )
 
     user_defined_config = container_context.server_k8s_config
 
+    user_defined_deployment_metadata = copy.deepcopy(dict(user_defined_config.deployment_metadata))
+    user_defined_deployment_labels = user_defined_deployment_metadata.pop("labels", {})
+
     container_config = copy.deepcopy(user_defined_config.container_config)
 
     user_defined_env_vars = container_config.pop("env", [])
@@ -148,9 +160,10 @@ def construct_code_location_deployment(
 
     deployment_dict = {
         "metadata": {
+            **user_defined_deployment_metadata,
             "name": k8s_deployment_name,
             "labels": {
-                **container_context.labels,
+                **user_defined_deployment_labels,
                 **_get_dagster_k8s_labels(
                     deployment_name, location_name, instance, server_timestamp
                 ),
@@ -162,12 +175,11 @@ def construct_code_location_deployment(
                 "metadata": {
                     **pod_template_spec_metadata,
                     "labels": {
-                        "user-deployment": k8s_deployment_name,
-                        **container_context.labels,
                         **user_defined_pod_template_labels,
                         **_get_dagster_k8s_labels(
                             deployment_name, location_name, instance, server_timestamp
                         ),
+                        "user-deployment": k8s_deployment_name,
                     },
                 },
                 "spec": pod_spec_config,

dagster_cloud/workspace/user_code_launcher/__init__.py

@@ -1,5 +1,7 @@
-from .process import ProcessUserCodeLauncher as ProcessUserCodeLauncher
-from .user_code_launcher import (
+from dagster_cloud.workspace.user_code_launcher.process import (
+    ProcessUserCodeLauncher as ProcessUserCodeLauncher,
+)
+from dagster_cloud.workspace.user_code_launcher.user_code_launcher import (
     DEFAULT_SERVER_PROCESS_STARTUP_TIMEOUT as DEFAULT_SERVER_PROCESS_STARTUP_TIMEOUT,
     DEFAULT_SERVER_TTL_SECONDS as DEFAULT_SERVER_TTL_SECONDS,
     SHARED_USER_CODE_LAUNCHER_CONFIG as SHARED_USER_CODE_LAUNCHER_CONFIG,
@@ -8,7 +10,7 @@ from .user_code_launcher import (
     ServerEndpoint as ServerEndpoint,
     UserCodeLauncherEntry as UserCodeLauncherEntry,
 )
-from .utils import (
+from dagster_cloud.workspace.user_code_launcher.utils import (
     deterministic_label_for_location as deterministic_label_for_location,
     get_human_readable_label as get_human_readable_label,
     unique_resource_name as unique_resource_name,

dagster_cloud/workspace/user_code_launcher/process.py

@@ -4,9 +4,9 @@ import subprocess
 import sys
 import threading
 from collections import defaultdict
-from typing import Any, Collection, Dict, List, Mapping, NamedTuple, Optional, Set, Tuple, Union
+from collections.abc import Collection, Mapping
+from typing import Any, NamedTuple, Optional, Union
 
-import dagster._seven as seven
 from dagster import (
     BoolSource,
     Field,
@@ -16,17 +16,17 @@ from dagster import (
 from dagster._core.errors import DagsterUserCodeUnreachableError
 from dagster._grpc.client import DagsterGrpcClient, client_heartbeat_thread
 from dagster._serdes import ConfigurableClass, ConfigurableClassData
-from dagster._serdes.ipc import open_ipc_subprocess
 from dagster._utils import find_free_port, safe_tempfile_path_unmanaged
 from dagster._utils.merger import merge_dicts
+from dagster_shared import seven
+from dagster_shared.ipc import open_ipc_subprocess
 from typing_extensions import Self
 
 from dagster_cloud.api.dagster_cloud_api import UserCodeDeploymentType
 from dagster_cloud.execution.cloud_run_launcher.process import CloudProcessRunLauncher
 from dagster_cloud.execution.monitoring import CloudContainerResourceLimits
 from dagster_cloud.pex.grpc import MultiPexGrpcClient
-
-from .user_code_launcher import (
+from dagster_cloud.workspace.user_code_launcher.user_code_launcher import (
     DEFAULT_SERVER_PROCESS_STARTUP_TIMEOUT,
     SHARED_USER_CODE_LAUNCHER_CONFIG,
     DagsterCloudGrpcServer,
@@ -34,6 +34,7 @@ from .user_code_launcher import (
     ServerEndpoint,
     UserCodeLauncherEntry,
 )
+from dagster_cloud.workspace.user_code_launcher.utils import get_grpc_server_env
 
 CLEANUP_ZOMBIE_PROCESSES_INTERVAL = 5
 
@@ -56,7 +57,7 @@ class ProcessUserCodeEntry(
         heartbeat_shutdown_event: threading.Event,
         heartbeat_thread: threading.Thread,
     ):
-        return super(ProcessUserCodeEntry, cls).__new__(
+        return super().__new__(
             cls,
             check.inst_param(grpc_server_process, "grpc_server_process", subprocess.Popen),
             check.inst_param(grpc_client, "grpc_client", DagsterGrpcClient),
@@ -79,7 +80,7 @@ class MultipexUserCodeEntry(
         grpc_server_process: subprocess.Popen,
         grpc_client: MultiPexGrpcClient,
     ):
-        return super(MultipexUserCodeEntry, cls).__new__(
+        return super().__new__(
             cls,
             check.inst_param(grpc_server_process, "grpc_server_process", subprocess.Popen),
             check.inst_param(grpc_client, "grpc_client", MultiPexGrpcClient),
@@ -98,14 +99,14 @@ class ProcessUserCodeLauncher(DagsterCloudUserCodeLauncher, ConfigurableClass):
 
         # map from pid to server being spun up
         # (including old servers in the process of being shut down)
-        self._process_entries: Dict[int, Union[ProcessUserCodeEntry, MultipexUserCodeEntry]] = {}
+        self._process_entries: dict[int, Union[ProcessUserCodeEntry, MultipexUserCodeEntry]] = {}
 
         # map from locationname to the pid(s) for that location-metadata combination.
         # Generally there should be only one pid per location unless an exception was raised partway
         # through an update
-        self._active_pids: Dict[Tuple[str, str], Set[int]] = defaultdict(set)
+        self._active_pids: dict[tuple[str, str], set[int]] = defaultdict(set)
 
-        self._active_multipex_pids: Dict[Tuple[str, str], Set[int]] = defaultdict(set)
+        self._active_multipex_pids: dict[tuple[str, str], set[int]] = defaultdict(set)
 
         self._heartbeat_ttl = 60
         self._wait_for_processes = wait_for_processes
@@ -115,7 +116,7 @@ class ProcessUserCodeLauncher(DagsterCloudUserCodeLauncher, ConfigurableClass):
 
         self._run_launcher: Optional[CloudProcessRunLauncher] = None
 
-        super(ProcessUserCodeLauncher, self).__init__(**kwargs)
+        super().__init__(**kwargs)
 
     @property
     def requires_images(self) -> bool:
@@ -174,7 +175,7 @@ class ProcessUserCodeLauncher(DagsterCloudUserCodeLauncher, ConfigurableClass):
         return self._inst_data
 
     @classmethod
-    def config_type(cls) -> Dict:
+    def config_type(cls) -> dict:
         return merge_dicts(
             {
                 "server_process_startup_timeout": Field(
@@ -266,7 +267,8 @@ class ProcessUserCodeLauncher(DagsterCloudUserCodeLauncher, ConfigurableClass):
                 str(self._heartbeat_ttl),
             ]
 
-            additional_env = metadata.get_grpc_server_env(
+            additional_env = get_grpc_server_env(
+                metadata,
                 port=port,
                 location_name=location_name,
                 instance_ref=self._instance.ref_for_deployment(deployment_name),
@@ -384,7 +386,7 @@ class ProcessUserCodeLauncher(DagsterCloudUserCodeLauncher, ConfigurableClass):
 
         return self._run_launcher
 
-    def _list_server_handles(self) -> List[int]:
+    def _list_server_handles(self) -> list[int]:
         return list(self._process_entries.keys())
 
     def get_agent_id_for_server(self, handle: int) -> Optional[str]: