dagster-k8s 0.27.1__py3-none-any.whl → 0.27.3__py3-none-any.whl

dagster_k8s/client.py

@@ -549,6 +549,18 @@ class DagsterKubernetesClient:
 
     ### Pod operations ###
 
+    def get_pod_by_name(self, pod_name: str, namespace: str):
+        """Get a pod by name.
+
+        Args:
+            pod_name (str): Name of the pod to get.
+            namespace (str): Namespace in which the pod is located.
+        """
+        check.str_param(pod_name, "pod_name")
+        check.str_param(namespace, "namespace")
+
+        return self.core_api.read_namespaced_pod(pod_name, namespace=namespace)
+
     def get_pods_in_job(self, job_name, namespace):
         """Get the pods launched by the job ``job_name``.
 

dagster_k8s/executor.py

@@ -1,3 +1,4 @@
+import os
 from collections.abc import Iterator
 from typing import Optional, cast
 
@@ -24,6 +25,7 @@ from dagster._core.executor.step_delegating import (
     StepHandler,
     StepHandlerContext,
 )
+from dagster._utils.cached_method import cached_method
 from dagster._utils.merger import merge_dicts
 
 from dagster_k8s.client import DagsterKubernetesClient
@@ -31,6 +33,7 @@ from dagster_k8s.container_context import K8sContainerContext
 from dagster_k8s.job import (
     USER_DEFINED_K8S_JOB_CONFIG_SCHEMA,
     DagsterK8sJobConfig,
+    OwnerReference,
     UserDefinedDagsterK8sConfig,
     construct_dagster_k8s_job,
     get_k8s_job_name,
@@ -82,6 +85,14 @@ _K8S_EXECUTOR_CONFIG_SCHEMA = merge_dicts(
             default_value={},
             description="Per op k8s configuration overrides.",
         ),
+        "enable_owner_references": Field(
+            bool,
+            is_required=False,
+            default_value=False,
+            description="Whether to insert Kubernetes owner references on step jobs to their parent run pod."
+            " This ensures that step jobs and step pods are garbage collected when the run pod is deleted."
+            " For more information, see https://kubernetes.io/docs/concepts/overview/working-with-objects/owners-dependents/",
+        ),
     },
 )
 
@@ -171,6 +182,9 @@ def k8s_job_executor(init_context: InitExecutorContext) -> Executor:
             load_incluster_config=load_incluster_config,
             kubeconfig_file=kubeconfig_file,
             per_step_k8s_config=exc_cfg.get("per_step_k8s_config", {}),
+            enable_owner_references=check.opt_bool_param(
+                exc_cfg.get("enable_owner_references"), "enable_owner_references", False
+            ),
         ),
         retries=RetryMode.from_config(exc_cfg["retries"]),  # type: ignore
         max_concurrent=check.opt_int_elem(exc_cfg, "max_concurrent"),
@@ -191,7 +205,9 @@ class K8sStepHandler(StepHandler):
         load_incluster_config: bool,
         kubeconfig_file: Optional[str],
         k8s_client_batch_api=None,
+        k8s_client_core_api=None,
         per_step_k8s_config=None,
+        enable_owner_references=False,
     ):
         super().__init__()
 
@@ -199,7 +215,7 @@ class K8sStepHandler(StepHandler):
         self._executor_container_context = check.inst_param(
             container_context, "container_context", K8sContainerContext
         )
-
+        self._kubeconfig_file = None
         if load_incluster_config:
             check.invariant(
                 kubeconfig_file is None,
@@ -209,13 +225,16 @@ class K8sStepHandler(StepHandler):
         else:
             check.opt_str_param(kubeconfig_file, "kubeconfig_file")
             kubernetes.config.load_kube_config(kubeconfig_file)
+            self._kubeconfig_file = kubeconfig_file
 
         self._api_client = DagsterKubernetesClient.production_client(
-            batch_api_override=k8s_client_batch_api
+            batch_api_override=k8s_client_batch_api,
+            core_api_override=k8s_client_core_api,
         )
         self._per_step_k8s_config = check.opt_dict_param(
             per_step_k8s_config, "per_step_k8s_config", key_type=str, value_type=dict
         )
+        self._enable_owner_references = enable_owner_references
 
     def _get_step_key(self, step_handler_context: StepHandlerContext) -> str:
         step_keys_to_execute = cast(
@@ -264,6 +283,25 @@ class K8sStepHandler(StepHandler):
 
         return f"dagster-step-{name_key}"
 
+    @cached_method
+    def _detect_current_name_and_uid(
+        self,
+    ) -> Optional[tuple[str, str]]:
+        """Get the current pod's pod name and uid, if available."""
+        from dagster_k8s.utils import detect_current_namespace
+
+        hostname = os.getenv("HOSTNAME")
+        if not hostname:
+            return None
+
+        namespace = detect_current_namespace(self._kubeconfig_file)
+        if not namespace:
+            return None
+
+        pod = self._api_client.get_pod_by_name(pod_name=hostname, namespace=namespace)
+
+        return pod.metadata.name, pod.metadata.uid
+
     def launch_step(self, step_handler_context: StepHandlerContext) -> Iterator[DagsterEvent]:
         step_key = self._get_step_key(step_handler_context)
 
@@ -301,6 +339,13 @@ class K8sStepHandler(StepHandler):
         deployment_name_env_var = get_deployment_id_label(container_context.run_k8s_config)
         if deployment_name_env_var:
             labels["dagster/deployment-name"] = deployment_name_env_var
+
+        owner_references = []
+        if self._enable_owner_references:
+            my_pod = self._detect_current_name_and_uid()
+            if my_pod:
+                owner_references = [OwnerReference(kind="Pod", name=my_pod[0], uid=my_pod[1])]
+
         job = construct_dagster_k8s_job(
             job_config=job_config,
             args=args,
@@ -317,6 +362,7 @@ class K8sStepHandler(StepHandler):
                 },
                 {"name": "DAGSTER_RUN_STEP_KEY", "value": step_key},
             ],
+            owner_references=owner_references,
         )
 
         yield DagsterEvent.step_worker_starting(

dagster_k8s/job.py

@@ -85,6 +85,20 @@ DEFAULT_JOB_SPEC_CONFIG = {
 }
 
 
+class OwnerReference(NamedTuple):
+    kind: str
+    name: str
+    uid: str
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "api_version": "batch/v1" if self.kind == "Job" else "v1",
+            "kind": self.kind,
+            "name": self.name,
+            "uid": self.uid,
+        }
+
+
 class UserDefinedDagsterK8sConfig(
     NamedTuple(
         "_UserDefinedDagsterK8sConfig",
@@ -754,6 +768,7 @@ def construct_dagster_k8s_job(
     component: Optional[str] = None,
     labels: Optional[Mapping[str, str]] = None,
     env_vars: Optional[Sequence[Mapping[str, Any]]] = None,
+    owner_references: Optional[Sequence[OwnerReference]] = None,
 ) -> kubernetes.client.V1Job:
     """Constructs a Kubernetes Job object.
 
@@ -929,6 +944,16 @@ def construct_dagster_k8s_job(
     user_defined_job_metadata = copy.deepcopy(dict(user_defined_k8s_config.job_metadata))
     user_defined_job_labels = user_defined_job_metadata.pop("labels", {})
 
+    owner_reference_dicts = (
+        [owner_reference.to_dict() for owner_reference in owner_references]
+        if owner_references
+        else []
+    )
+    if "owner_references" in user_defined_job_metadata:
+        user_defined_job_metadata["owner_references"] = (
+            owner_reference_dicts + user_defined_job_metadata["owner_references"]
+        )
+
     job = k8s_model_from_dict(
         kubernetes.client.V1Job,
         merge_dicts(
@@ -944,6 +969,7 @@ def construct_dagster_k8s_job(
                             dagster_labels, user_defined_job_labels, job_config.labels
                         ),
                     },
+                    {"owner_references": owner_reference_dicts} if owner_reference_dicts else {},
                 ),
                 "spec": job_spec_config,
             },

dagster_k8s/pipes.py

@@ -8,7 +8,6 @@ import time
 from collections.abc import Callable, Generator, Iterator, Mapping, Sequence
 from contextlib import contextmanager
 from datetime import datetime
-from pathlib import Path
 from typing import Any, Optional, Union
 
 import kubernetes
@@ -50,7 +49,7 @@ from dagster_k8s.client import (
     WaitForPodState,
 )
 from dagster_k8s.models import k8s_model_from_dict, k8s_snake_case_dict
-from dagster_k8s.utils import get_common_labels
+from dagster_k8s.utils import detect_current_namespace, get_common_labels
 
 INIT_WAIT_TIMEOUT_FOR_READY = 1800.0  # 30mins
 INIT_WAIT_TIMEOUT_FOR_TERMINATE = 10.0  # 10s
@@ -64,7 +63,6 @@ def get_pod_name(run_id: str, op_name: str):
 
 
 DEFAULT_CONTAINER_NAME = "dagster-pipes-execution"
-_NAMESPACE_SECRET_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/namespace")
 _DEV_NULL_MESSAGE_WRITER = encode_env_var({"path": "/dev/null"})
 
 DEFAULT_CONSUME_POD_LOGS_RETRIES = 5
@@ -502,7 +500,7 @@ class PipesK8sClient(PipesClient, TreatAsResourceParam):
             context_injector=self.context_injector,
             message_reader=self.message_reader,
         ) as pipes_session:
-            namespace = namespace or _detect_current_namespace(self.kubeconfig_file) or "default"
+            namespace = namespace or detect_current_namespace(self.kubeconfig_file) or "default"
             pod_name = get_pod_name(context.run_id, context.op.name)
             pod_body = build_pod_body(
                 pod_name=pod_name,
@@ -593,29 +591,6 @@ class PipesK8sClient(PipesClient, TreatAsResourceParam):
         yield
 
 
-def _detect_current_namespace(
-    kubeconfig_file: Optional[str], namespace_secret_path: Path = _NAMESPACE_SECRET_PATH
-) -> Optional[str]:
-    """Get the current in-cluster namespace when operating within the cluster.
-
-    First attempt to read it from the `serviceaccount` secret or get it from the kubeconfig_file if it is possible.
-    It will attempt to take from the active context if it exists and returns None if it does not exist.
-    """
-    if namespace_secret_path.exists():
-        with namespace_secret_path.open() as f:
-            # We only need to read the first line, this guards us against bad input.
-            return f.read().strip()
-
-    if not kubeconfig_file:
-        return None
-
-    try:
-        _, active_context = kubernetes.config.list_kube_config_contexts(kubeconfig_file)
-        return active_context["context"]["namespace"]
-    except KeyError:
-        return None
-
-
 def build_pod_body(
     pod_name: str,
     image: Optional[str],

dagster_k8s/utils.py

@@ -1,6 +1,8 @@
 import re
-from typing import TYPE_CHECKING
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
 
+import kubernetes
 from dagster import __version__ as dagster_version
 
 if TYPE_CHECKING:
@@ -32,3 +34,29 @@ def get_deployment_id_label(user_defined_k8s_config: "UserDefinedDagsterK8sConfi
         else None
     )
     return deployment_name_env_var["value"] if deployment_name_env_var else None
+
+
+_NAMESPACE_SECRET_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/namespace")
+
+
+def detect_current_namespace(
+    kubeconfig_file: Optional[str], namespace_secret_path: Path = _NAMESPACE_SECRET_PATH
+) -> Optional[str]:
+    """Get the current in-cluster namespace when operating within the cluster.
+
+    First attempt to read it from the `serviceaccount` secret or get it from the kubeconfig_file if it is possible.
+    It will attempt to take from the active context if it exists and returns None if it does not exist.
+    """
+    if namespace_secret_path.exists():
+        with namespace_secret_path.open() as f:
+            # We only need to read the first line, this guards us against bad input.
+            return f.read().strip()
+
+    if not kubeconfig_file:
+        return None
+
+    try:
+        _, active_context = kubernetes.config.list_kube_config_contexts(kubeconfig_file)
+        return active_context["context"]["namespace"]
+    except KeyError:
+        return None

dagster_k8s/version.py

@@ -1 +1 @@
-__version__ = "0.27.1"
+__version__ = "0.27.3"

{dagster_k8s-0.27.1.dist-info → dagster_k8s-0.27.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dagster-k8s
-Version: 0.27.1
+Version: 0.27.3
 Summary: A Dagster integration for k8s
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-k8s
 Author: Dagster Labs
@@ -10,11 +10,12 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.9,<3.13
+Requires-Python: >=3.9,<3.14
 License-File: LICENSE
-Requires-Dist: dagster==1.11.1
+Requires-Dist: dagster==1.11.3
 Requires-Dist: kubernetes<33
 Requires-Dist: google-auth!=2.23.1
 Dynamic: author

{dagster_k8s-0.27.1.dist-info → dagster_k8s-0.27.3.dist-info}/RECORD

@@ -1,21 +1,21 @@
 dagster_k8s/__init__.py,sha256=1Vi9HXEW9fatVuNzWTr6Oq-JEHmyDrVkonZuELsQ1C0,751
-dagster_k8s/client.py,sha256=-Ns2saJri4Xxu-OpQ_T93e2K6aVnZvU19tj0iaMmTg8,40901
+dagster_k8s/client.py,sha256=iPv_1O7TO9S4-viB7tMAAzbd_Z8VXkJOJzrFUjlGqO8,41316
 dagster_k8s/component.py,sha256=-66T-IHktXQ8FC0INNPFWsrxpoBsEEtRAiDuhWxCTKk,1999
 dagster_k8s/container_context.py,sha256=399bdqmV3Ce6lNbfuxnhdyJSCy9keRRPJfVjDGn2U4I,22642
-dagster_k8s/executor.py,sha256=hRwzohUFfusEc84Iz9Q7N7NN5lHoCyODdARh-kpU9ec,15126
-dagster_k8s/job.py,sha256=v7RXYLZHzy2sTC7HYdz_INqbLWyZxUz9QIMwORArjJg,43065
+dagster_k8s/executor.py,sha256=Tv3huQNHS9KJoOAzC0yD1EGPG7zGh_71jSZrunnzF18,17015
+dagster_k8s/job.py,sha256=ja_AHiWBP0j4W41bRd4krZM7YUDagxKReFK37qhp8No,43899
 dagster_k8s/kubernetes_version.py,sha256=T7MDA5pOOVE_woRUkKWxtiRj4RCLAVfjSn6ZXUr6sAA,38
 dagster_k8s/launcher.py,sha256=iNvrxhF7DP0sNxydg3hU9OzVGILYUcgL9LxeXGtaXqI,16936
 dagster_k8s/models.py,sha256=REKzAY103pZw0pQFqvHaM6lL02hgbi0rxmm71qds4W8,6338
-dagster_k8s/pipes.py,sha256=kb1So6erUgjTDkicALpU1T5mG9xa9n8aeI4lPUVeU2o,34031
+dagster_k8s/pipes.py,sha256=WUEx4Lfvuz_Fy6tYzdcAz6wRzkJwuqolGuB6MYuCFEA,33041
 dagster_k8s/py.typed,sha256=la67KBlbjXN-_-DfGNcdOcjYumVpKG_Tkw-8n5dnGB4,8
 dagster_k8s/test.py,sha256=cNtcbzxytiZtd01wY5ip7KPi01y0BUQuQhohoIfAFUM,684
-dagster_k8s/utils.py,sha256=NS9HRyy67-jF6n4loEPd6xuU33t7Y27IovK92WCdgu0,1296
-dagster_k8s/version.py,sha256=lnka9HWRxSmlQAffgSpaSitns-Djhy2OArtj9IVwxrY,23
+dagster_k8s/utils.py,sha256=QQY170Ut4iBjlO8gAhdYeiwRmtxxAOBYlt4CO2uViwY,2340
+dagster_k8s/version.py,sha256=lxhjPOOCzhlHB02EzaqTtDdBFZSOLV3WLWw2HC0DYvo,23
 dagster_k8s/ops/__init__.py,sha256=ur-9GrE_DRfnsFCpYan03qOY9cWbjagC8KHZFZuiCmc,113
 dagster_k8s/ops/k8s_job_op.py,sha256=xiBxpuLnIMDd0BqhEGtfl7Z8y0bFHuFmcfep5PuyYqk,22375
-dagster_k8s-0.27.1.dist-info/licenses/LICENSE,sha256=4lsMW-RCvfVD4_F57wrmpe3vX1xwUk_OAKKmV_XT7Z0,11348
-dagster_k8s-0.27.1.dist-info/METADATA,sha256=wPe-Aya4CS0-MaB77tD8XSExT_eIEeTuZoFLxL-n9EY,910
-dagster_k8s-0.27.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-dagster_k8s-0.27.1.dist-info/top_level.txt,sha256=wFPjskoWPlk2hOLugYCaoZhSiZdUcbCA1QZe9I4dals,12
-dagster_k8s-0.27.1.dist-info/RECORD,,
+dagster_k8s-0.27.3.dist-info/licenses/LICENSE,sha256=4lsMW-RCvfVD4_F57wrmpe3vX1xwUk_OAKKmV_XT7Z0,11348
+dagster_k8s-0.27.3.dist-info/METADATA,sha256=fLxNkBGsWDycRrBuuZQyuOB3NwSkW_BTesANYn25Gv8,961
+dagster_k8s-0.27.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dagster_k8s-0.27.3.dist-info/top_level.txt,sha256=wFPjskoWPlk2hOLugYCaoZhSiZdUcbCA1QZe9I4dals,12
+dagster_k8s-0.27.3.dist-info/RECORD,,