import json import re import shlex from clearml_agent.backend_api.session import Request from clearml_agent.helper.package.requirements import ( RequirementsManager, MarkerRequirement, compare_version_rules, ) def resolve_default_container(session, task_id, container_config): container_lookup = session.config.get('agent.default_docker.match_rules', None) if not session.check_min_api_version("2.13") or not container_lookup: return container_config # check backend support before sending any more requests (because they will fail and crash the Task) try: session.verify_feature_set('advanced') except ValueError: return container_config result = session.send_request( service='tasks', action='get_all', version='2.14', json={'id': [task_id], 'only_fields': ['script.requirements', 'script.binary', 'script.repository', 'script.branch', 'project', 'container'], 'search_hidden': True}, method=Request.def_method, async_enable=False, ) try: task_info = result.json()['data']['tasks'][0] if result.ok else {} except (ValueError, TypeError): return container_config from clearml_agent.external.requirements_parser.requirement import Requirement # store tasks repository repository = task_info.get('script', {}).get('repository') or '' branch = task_info.get('script', {}).get('branch') or '' binary = task_info.get('script', {}).get('binary') or '' requested_container = task_info.get('container', {}) # get project full path project_full_name = '' if task_info.get('project', None): result = session.send_request( service='projects', action='get_all', version='2.13', json={ 'id': [task_info.get('project')], 'only_fields': ['name'], }, method=Request.def_method, async_enable=False, ) try: if result.ok: project_full_name = result.json()['data']['projects'][0]['name'] or '' except (ValueError, TypeError): pass task_packages_lookup = {} for entry in container_lookup: match = entry.get('match', None) if not match: continue if match.get('project', None): # noinspection PyBroadException try: if not re.search(match.get('project', None), project_full_name): continue except Exception: print('Failed parsing regular expression \"{}\" in rule: {}'.format( match.get('project', None), entry)) continue if match.get('script.repository', None): # noinspection PyBroadException try: if not re.search(match.get('script.repository', None), repository): continue except Exception: print('Failed parsing regular expression \"{}\" in rule: {}'.format( match.get('script.repository', None), entry)) continue if match.get('script.branch', None): # noinspection PyBroadException try: if not re.search(match.get('script.branch', None), branch): continue except Exception: print('Failed parsing regular expression \"{}\" in rule: {}'.format( match.get('script.branch', None), entry)) continue if match.get('script.binary', None): # noinspection PyBroadException try: if not re.search(match.get('script.binary', None), binary): continue except Exception: print('Failed parsing regular expression \"{}\" in rule: {}'.format( match.get('script.binary', None), entry)) continue if match.get('container', None): # noinspection PyBroadException try: if not re.search(match.get('container', None), requested_container.get('image', '')): continue except Exception: print('Failed parsing regular expression \"{}\" in rule: {}'.format( match.get('container', None), entry)) continue matched = True for req_section in ['script.requirements.pip', 'script.requirements.conda']: if not match.get(req_section, None): continue match_pip_reqs = [MarkerRequirement(Requirement.parse('{} {}'.format(k, v))) for k, v in match.get(req_section, None).items()] if not task_packages_lookup.get(req_section): req_section_parts = req_section.split('.') task_packages_lookup[req_section] = \ RequirementsManager.parse_requirements_section_to_marker_requirements( requirements=task_info.get(req_section_parts[0], {}).get( req_section_parts[1], {}).get(req_section_parts[2], None) ) matched_all_reqs = True for mr in match_pip_reqs: matched_req = False for pr in task_packages_lookup[req_section]: if mr.req.name != pr.req.name: continue if compare_version_rules(mr.specs, pr.specs): matched_req = True break if not matched_req: matched_all_reqs = False break # if ew have a match, check second section if matched_all_reqs: continue # no match stop matched = False break if matched: if not container_config.get('container'): container_config['container'] = entry.get('image', None) if not container_config.get('arguments'): container_config['arguments'] = entry.get('arguments', None) container_config['arguments'] = shlex.split(str(container_config.get('arguments') or '').strip()) print('Matching default container with rule:\n{}'.format(json.dumps(entry))) return container_config return container_config