Add support for automatic cleanup of jobsets (#310)

AI-Hypercomputer · Jan 8, 2025 · bad83f7 · bad83f7
1 parent 51d8530
commit bad83f7
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/src/xpk/commands/workload.py b/src/xpk/commands/workload.py
@@ -77,6 +77,7 @@
   annotations:
     alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool # 1:1 job replica to node pool assignment
 spec:
+  ttlSecondsAfterFinished: {args.ttl_seconds_after_finished}
   failurePolicy:
     maxRestarts: {args.max_restarts}
   replicatedJobs:
@@ -118,6 +119,7 @@
     kueue.x-k8s.io/queue-name: multislice-queue  # Name of the LocalQueue
     xpk.google.com/workload: {args.workload}
 spec:
+  ttlSecondsAfterFinished: {args.ttl_seconds_after_finished}
   failurePolicy:
     maxRestarts: {args.max_restarts}
   replicatedJobs:
@@ -175,6 +177,7 @@
     kueue.x-k8s.io/queue-name: multislice-queue  # Name of the LocalQueue
     xpk.google.com/workload: {args.workload}
 spec:
+  ttlSecondsAfterFinished: {args.ttl_seconds_after_finished}
   failurePolicy:
     maxRestarts: {args.max_restarts}
   replicatedJobs:

diff --git a/src/xpk/parser/workload.py b/src/xpk/parser/workload.py
@@ -114,6 +114,15 @@ def set_workload_parsers(workload_parser):
       ),
   )
 
+  workload_create_parser_optional_arguments.add_argument(
+      '--ttl-seconds-after-finished',
+      type=int,
+      default=12 * 60 * 60,
+      help=(
+          'Set the number of seconds to clean up finished Jobsets (either'
+          ' Complete or Failed). This is by default set to 12 hours.'
+      ),
+  )
   workload_create_parser_optional_arguments.add_argument(
       '--num-nodes',
       type=int,