Add support for task field injection into container docker name

This commit is contained in:
allegroai 2023-07-03 11:07:12 +03:00
parent 95dadca45c
commit c77224af68
2 changed files with 47 additions and 1 deletions

View File

@ -269,10 +269,15 @@
# Name docker containers created by the daemon using the following string format (supported from Docker 0.6.5) # Name docker containers created by the daemon using the following string format (supported from Docker 0.6.5)
# Allowed variables are task_id, worker_id and rand_string (random lower-case letters string, up to 32 characters) # Allowed variables are task_id, worker_id and rand_string (random lower-case letters string, up to 32 characters)
# Custom variables may be specified using the docker_container_name_format_fields option.
# Note: resulting name must start with an alphanumeric character and # Note: resulting name must start with an alphanumeric character and
# continue with alphanumeric characters, underscores (_), dots (.) and/or dashes (-) # continue with alphanumeric characters, underscores (_), dots (.) and/or dashes (-)
# docker_container_name_format: "clearml-id-{task_id}-{rand_string:.8}" # docker_container_name_format: "clearml-id-{task_id}-{rand_string:.8}"
# Specify custom variables for the docker_container_name_format option using a mapping of variable name
# to a (nested) task field (using "." as a task field separator, digits specify array index)
# docker_container_name_format_fields: { foo: "bar.moo" }
# Apply top-level environment section from configuration into os.environ # Apply top-level environment section from configuration into os.environ
apply_environment: true apply_environment: true
# Top-level environment section is in the form of: # Top-level environment section is in the form of:

View File

@ -319,6 +319,37 @@ def get_next_task(session, queue, get_task_info=False):
return data return data
def get_task_fields(session, task_id, fields: list, log=None) -> dict:
"""
Returns dict with Task docker container setup {container: '', arguments: '', setup_shell_script: ''}
"""
result = session.send_request(
service='tasks',
action='get_all',
json={'id': [task_id], 'only_fields': list(fields), 'search_hidden': True},
method=Request.def_method,
async_enable=False,
)
# noinspection PyBroadException
try:
results = {}
result = result.json()['data']['tasks'][0]
for field in fields:
cur = result
for part in field.split("."):
if part.isdigit():
cur = cur[part]
else:
cur = cur.get(part, {})
results[field] = cur
return results
except Exception as ex:
if log:
log.error("Failed obtaining values for task fields {}: {}", fields, ex)
pass
return {}
def get_task_container(session, task_id): def get_task_container(session, task_id):
""" """
Returns dict with Task docker container setup {container: '', arguments: '', setup_shell_script: ''} Returns dict with Task docker container setup {container: '', arguments: '', setup_shell_script: ''}
@ -890,11 +921,21 @@ class Worker(ServiceCommandSection):
name_format = self._session.config.get('agent.docker_container_name_format', None) name_format = self._session.config.get('agent.docker_container_name_format', None)
if name_format: if name_format:
custom_fields = {}
name_format_fields = self._session.config.get('agent.docker_container_name_format_fields', None)
if name_format_fields:
field_values = get_task_fields(task_session, task_id, name_format_fields.values(), log=self.log)
custom_fields = {
k: field_values.get(v)
for k, v in name_format_fields.items()
}
try: try:
name = name_format.format( name = name_format.format(
task_id=re.sub(r'[^a-zA-Z0-9._-]', '-', task_id), task_id=re.sub(r'[^a-zA-Z0-9._-]', '-', task_id),
worker_id=re.sub(r'[^a-zA-Z0-9._-]', '-', worker_id), worker_id=re.sub(r'[^a-zA-Z0-9._-]', '-', worker_id),
rand_string="".join(sys_random.choice(string.ascii_lowercase) for _ in range(32)) rand_string="".join(sys_random.choice(string.ascii_lowercase) for _ in range(32)),
**custom_fields,
) )
except Exception as ex: except Exception as ex:
print("Warning: failed generating docker container name: {}".format(ex)) print("Warning: failed generating docker container name: {}".format(ex))