diff --git a/clearml_agent/glue/k8s.py b/clearml_agent/glue/k8s.py index 0d44944..d758b0b 100644 --- a/clearml_agent/glue/k8s.py +++ b/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() diff --git a/examples/k8s_glue_example.py b/examples/k8s_glue_example.py index e967e86..8a4c024 100644 --- a/examples/k8s_glue_example.py +++ b/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