From fd07780c7c3238e307c88bd61a5b82aab7582373 Mon Sep 17 00:00:00 2001 From: Adithya Chakilam Date: Tue, 9 Jul 2024 16:14:58 -0500 Subject: [PATCH 1/4] kubernetes-overlord-extension: Fix tasks not being shutdown --- .../druid/k8s/overlord/KubernetesPeonLifecycle.java | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java b/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java index eaef0cba6a15..8c8b7f61efcc 100644 --- a/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java +++ b/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java @@ -254,7 +254,15 @@ protected TaskLocation getTaskLocation() if we decide we need to change this later. **/ if (taskLocation == null) { - Optional maybePod = kubernetesClient.getPeonPod(taskId.getK8sJobName()); + Optional maybePod = Optional.absent(); + try { + maybePod = kubernetesClient.getPeonPod(taskId.getK8sJobName()); + } + catch (Exception e) { + log.makeAlert("Unable to get location for task", e) + .addData("taskId", taskId); + } + if (!maybePod.isPresent()) { return TaskLocation.unknown(); } From b1523b8b4f8e520a533dd6ab5f85ddf2915e783e Mon Sep 17 00:00:00 2001 From: Adithya Chakilam Date: Mon, 15 Jul 2024 09:59:41 -0500 Subject: [PATCH 2/4] comments --- .../druid/k8s/overlord/KubernetesPeonLifecycle.java | 10 +--------- .../druid/k8s/overlord/KubernetesTaskRunner.java | 8 +++++++- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java b/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java index 8c8b7f61efcc..eaef0cba6a15 100644 --- a/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java +++ b/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java @@ -254,15 +254,7 @@ protected TaskLocation getTaskLocation() if we decide we need to change this later. **/ if (taskLocation == null) { - Optional maybePod = Optional.absent(); - try { - maybePod = kubernetesClient.getPeonPod(taskId.getK8sJobName()); - } - catch (Exception e) { - log.makeAlert("Unable to get location for task", e) - .addData("taskId", taskId); - } - + Optional maybePod = kubernetesClient.getPeonPod(taskId.getK8sJobName()); if (!maybePod.isPresent()) { return TaskLocation.unknown(); } diff --git a/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java b/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java index 080a0fdaa986..28b31d570f65 100644 --- a/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java +++ b/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java @@ -447,7 +447,13 @@ public TaskLocation getTaskLocation(String taskId) if (workItem == null) { return TaskLocation.unknown(); } else { - return workItem.getLocation(); + try { + return workItem.getLocation(); + } + catch (Exception e) { + log.warn("Unable to find location for task [%s]", taskId); + return TaskLocation.unknown(); + } } } From c887c8fe2d53f1fd0c4ce224ace8841aa9424a0e Mon Sep 17 00:00:00 2001 From: Adithya Chakilam Date: Mon, 15 Jul 2024 11:55:03 -0500 Subject: [PATCH 3/4] add test --- .../k8s/overlord/KubernetesTaskRunnerTest.java | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java b/extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java index 391db70afb28..67a5278c6a32 100644 --- a/extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java +++ b/extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java @@ -654,6 +654,24 @@ public TaskLocation getLocation() Assert.assertEquals(TaskLocation.create("host", 0, 1, false), taskLocation); } + @Test + public void test_getTaskLocation_throws() + { + KubernetesWorkItem workItem = new KubernetesWorkItem(task, null) + { + @Override + public TaskLocation getLocation() + { + throw new RuntimeException(); + } + }; + + runner.tasks.put(task.getId(), workItem); + + TaskLocation taskLocation = runner.getTaskLocation(task.getId()); + Assert.assertEquals(TaskLocation.unknown(), taskLocation); + } + @Test public void test_getTaskLocation_noTaskFound() { From 202b0361d11ad200246010d023a88b1792995893 Mon Sep 17 00:00:00 2001 From: Adithya Chakilam Date: Mon, 15 Jul 2024 14:05:59 -0500 Subject: [PATCH 4/4] comments --- .../k8s/overlord/KubernetesTaskRunner.java | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java b/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java index 28b31d570f65..c324b49e13a2 100644 --- a/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java +++ b/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java @@ -443,18 +443,18 @@ public Collection getPendingTasks() @Override public TaskLocation getTaskLocation(String taskId) { - final KubernetesWorkItem workItem = tasks.get(taskId); - if (workItem == null) { - return TaskLocation.unknown(); - } else { - try { - return workItem.getLocation(); - } - catch (Exception e) { - log.warn("Unable to find location for task [%s]", taskId); + try { + final KubernetesWorkItem workItem = tasks.get(taskId); + if (workItem == null) { return TaskLocation.unknown(); + } else { + return workItem.getLocation(); } } + catch (Exception e) { + log.warn("Unable to find location for task [%s]", taskId); + return TaskLocation.unknown(); + } } @Nullable