ClearML/clearml-agent
1
0
Fork 0
mirror of https://github.com/clearml/clearml-agent synced 2025-01-31 00:56:53 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add CLEARML_TASK_ID and auth token to pod env vars in original entrypoint flow

Browse Source
This commit is contained in:
allegroai 2024-07-24 17:47:48 +03:00
parent edc333ba5f
commit 8ba4d75e80
1 changed files with 23 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
clearml_agent/glue/k8s.py
View File

@ -597,6 +597,7 @@ class K8sIntegration(Worker):
task_id=task_id, task_id=task_id,
queue=queue, queue=queue,
namespace=namespace, namespace=namespace,
task_token=task_session.token.encode("ascii") if task_session else None,
) )
print('kubectl output:\n{}\n{}'.format(error, output)) print('kubectl output:\n{}\n{}'.format(error, output))
@ -675,7 +676,7 @@ class K8sIntegration(Worker):
def _create_template_container( def _create_template_container(
self, pod_name: str, task_id: str, docker_image: str, docker_args: List[str], self, pod_name: str, task_id: str, docker_image: str, docker_args: List[str],
docker_bash: str, clearml_conf_create_script: List[str], task_worker_id: str docker_bash: str, clearml_conf_create_script: List[str], task_worker_id: str, task_token: str = None
) -> dict: ) -> dict:
container = self._get_docker_args( container = self._get_docker_args(
docker_args, docker_args,
@ -684,18 +685,27 @@ class K8sIntegration(Worker):
convert=lambda env: {'name': env.partition("=")[0], 'value': env.partition("=")[2]}, convert=lambda env: {'name': env.partition("=")[0], 'value': env.partition("=")[2]},
) )
def add_or_update_env_var(name, value):
env_vars = container.get('env', [])
for entry in env_vars:
if entry.get('name') == name:
entry['value'] = value
break
else:
container['env'] = env_vars + [{'name': name, 'value': value}]
# Set worker ID # Set worker ID
env_vars = container.get('env', []) add_or_update_env_var('CLEARML_WORKER_ID', task_worker_id)
found_worker_id = False
for entry in env_vars:
if entry.get('name') == 'CLEARML_WORKER_ID':
entry['name'] = task_worker_id
found_worker_id = True
if not found_worker_id:
container['env'] = env_vars + [{'name': 'CLEARML_WORKER_ID', 'value': task_worker_id}]
if ENV_POD_USE_IMAGE_ENTRYPOINT.get(): if ENV_POD_USE_IMAGE_ENTRYPOINT.get():
# Don't add a cmd and args, just the image # Don't add a cmd and args, just the image
# Add the task ID and token since we need it (it's usually in the init script passed to us
add_or_update_env_var('CLEARML_TASK_ID', task_id)
if task_token:
# TODO: find a way to base64 encode the token
add_or_update_env_var('CLEARML_AUTH_TOKEN', task_token)
return self._merge_containers( return self._merge_containers(
container, dict(name=pod_name, image=docker_image) container, dict(name=pod_name, image=docker_image)
) )
@ -749,7 +759,8 @@ class K8sIntegration(Worker):
task_id, task_id,
namespace, namespace,
template, template,
pod_number=None pod_number=None,
task_token=None,
): ):
if "apiVersion" not in template: if "apiVersion" not in template:
template["apiVersion"] = "batch/v1" if self.using_jobs else "v1" template["apiVersion"] = "batch/v1" if self.using_jobs else "v1"
@ -797,7 +808,8 @@ class K8sIntegration(Worker):
docker_args=docker_args, docker_args=docker_args,
docker_bash=docker_bash, docker_bash=docker_bash,
clearml_conf_create_script=clearml_conf_create_script, clearml_conf_create_script=clearml_conf_create_script,
task_worker_id=task_worker_id task_worker_id=task_worker_id,
task_token=task_token,
) )
if containers: if containers: