Add command line arguments for k8s_glue_example.py (#196)

Co-authored-by: Meshcheryakov Ilya <i.meshcheryakov@mts.ai>
Co-authored-by: Jake Henning <59198928+jkhenning@users.noreply.github.com>

clearml_agent/glue/k8s.py

@@ -148,6 +148,10 @@ class K8sIntegration(Worker):
         :param str extra_bash_init_script: Additional bash script to run before starting the Task inside the container
         :param str namespace: K8S namespace to be used when creating the new pods (default: clearml)
         :param int max_pods_limit: Maximum number of pods that K8S glue can run at the same time
+        :param str pod_name_prefix: Define pod name prefix for k8s (default: clearml-id-)
+        :param str limit_pod_label: Define limit pod label for k8s (default: ai.allegro.agent.serial=pod-{pod_number})
+        :param bool force_system_packages: true when running tasks in containers (i.e. docker mode or k8s glue).
+            (default: true)
         """
         super(K8sIntegration, self).__init__()
         self.kind = os.environ.get("CLEARML_K8S_GLUE_KIND", "pod").strip().lower()

examples/k8s_glue_example.py

@@ -13,9 +13,19 @@ def parse_args():
     group = parser.add_mutually_exclusive_group()
 
     parser.add_argument(
-        "--queue",
-        type=str,
-        help="Queues to pull tasks from. If multiple queues, use comma separated list, e.g. 'queue1,queue2'",
+        "--k8s-pending-queue-name", type=str,
+        help="Queue name to use when task is pending in the k8s scheduler (default: %(default)s)", default="k8s_scheduler"
+    )
+    parser.add_argument(
+        "--container-bash-script", type=str,
+        help="Path to the file with container bash script to be executed in k8s", default=None
+    )
+    parser.add_argument(
+        "--debug", action="store_true", default=False,
+        help="Switch logging on (default: %(default)s)"
+    )
+    parser.add_argument(
+        "--queue", type=str, help="Queues to pull tasks from. If multiple queues, use comma separated list, e.g. 'queue1,queue2'",
     )
     group.add_argument(
         "--ports-mode",
@@ -82,6 +92,18 @@ def parse_args():
         help="Limit the maximum number of pods that this service can run at the same time."
         "Should not be used with ports-mode",
     )
+    parser.add_argument(
+        "--pod-name-prefix", type=str,
+        help="Define pod name prefix for k8s (default: %(default)s)", default="clearml-id-"
+    )
+    parser.add_argument(
+        "--limit-pod-label", type=str,
+        help="Define limit pod label for k8s (default: %(default)s)", default="ai.allegro.agent.serial=pod-{pod_number}"
+    )
+    parser.add_argument(
+        "--no-system-packages", action="store_true", default=False,
+        help="False when running tasks in containers (default: %(default)s)"
+    )
     parser.add_argument(
         "--use-owner-token",
         action="store_true",
@@ -113,7 +135,15 @@ def main():
 
         user_props_cb = k8s_user_props_cb
 
+    if args.container_bash_script:
+        with open(args.container_bash_script, "r") as file:
+            container_bash_script = file.read().splitlines()
+    else:
+        container_bash_script = None
+
     k8s = K8sIntegration(
+        k8s_pending_queue_name=args.k8s_pending_queue_name,
+        container_bash_script=container_bash_script,
         ports_mode=args.ports_mode,
         num_of_services=args.num_of_services,
         base_pod_num=args.base_pod_num,
@@ -126,6 +156,10 @@ def main():
         else None,
         namespace=args.namespace,
         max_pods_limit=args.max_pods or None,
+        pod_name_prefix=args.pod_name_prefix,
+        limit_pod_label=args.limit_pod_label,
+        force_system_packages=not args.no_system_packages,
+        debug=args.debug,
     )
     queue = [q.strip() for q in args.queue.split(",") if q.strip()] if args.queue else None