Add support for setting mode on files applied by the agent

Fix default docker match rules resolver (used incorrect field "container" instead of "image")
Remove "container" (image) match rule option from default docker image resolver

.gitignore

@@ -11,3 +11,6 @@ build/
 dist/
 *.egg-info
 
+# VSCode
+.vscode
+

clearml_agent/backend_api/config/default/agent.conf

@@ -69,8 +69,9 @@
         pip_version: ["<20.2 ; python_version < '3.10'", "<22.3 ; python_version >= '3.10'"],
         # specify poetry version to use (examples "<2", "==1.1.1", "", empty string will install the latest version)
         # poetry_version: "<2",
+        # poetry_install_extra_args: ["-v"]
 
-        # virtual environment inheres packages from system
+        # virtual environment inherits packages from system
         system_site_packages: false,
 
         # install with --upgrade
@@ -79,6 +80,16 @@
         # additional artifact repositories to use when installing python packages
         # extra_index_url: ["https://allegroai.jfrog.io/clearml/api/pypi/public/simple"]
 
+        # control the pytorch wheel resolving algorithm, options are: "pip", "direct"
+        # "pip" (default): would automatically detect the cuda version, and supply pip with the correct
+        # extra-index-url, based on pytorch.org tables
+        # "direct": would resolve a direct link to the pytorch wheel by parsing the pytorch.org pip repository
+        # and matching the automatically detected cuda version with the required pytorch wheel.
+        # if the exact cuda version is not found for the required pytorch wheel, it will try
+        # a lower cuda version until a match is found
+        #
+        # pytorch_resolve: "pip"
+
         # additional conda channels to use when installing with conda package manager
         conda_channels: ["pytorch", "conda-forge", "defaults", ]
 
@@ -87,24 +98,32 @@
         # force_repo_requirements_txt: false
 
         # set the priority packages to be installed before the rest of the required packages
+        # Note: this only controls the installation order of existing requirement packages (and does not add additional packages)
         # priority_packages: ["cython", "numpy", "setuptools", ]
 
         # set the optional priority packages to be installed before the rest of the required packages,
         # In case a package installation fails, the package will be ignored,
         # and the virtual environment process will continue
+        # Note: this only controls the installation order of existing requirement packages (and does not add additional packages)
         priority_optional_packages: ["pygobject", ]
 
         # set the post packages to be installed after all the rest of the required packages
+        # Note: this only controls the installation order of existing requirement packages (and does not add additional packages)
         # post_packages: ["horovod", ]
 
         # set the optional post packages to be installed after all the rest of the required packages,
         # In case a package installation fails, the package will be ignored,
         # and the virtual environment process will continue
+        # Note: this only controls the installation order of existing requirement packages (and does not add additional packages)
         # post_optional_packages: []
 
         # set to True to support torch nightly build installation,
         # notice: torch nightly builds are ephemeral and are deleted from time to time
         torch_nightly: false,
+
+        # if set to true, the agent will look for the "poetry.lock" file 
+        # in the passed current working directory instead of the repository's root directory.
+        poetry_files_from_repo_working_dir: false
     },
 
     # target folder for virtual environments builds, created when executing experiment
@@ -254,10 +273,15 @@
 
     # 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)
+    # Custom variables may be specified using the docker_container_name_format_fields option.
     # Note: resulting name must start with an alphanumeric character and
     #       continue with alphanumeric characters, underscores (_), dots (.) and/or dashes (-)
     # 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_environment: true
     # Top-level environment section is in the form of:
@@ -335,6 +359,42 @@
 
     # Disable task docker override. If true, the agent will use the default docker image and ignore any docker image
     # and arguments specified in the task's container section (setup shell script from the task container section will
-    # be used in any case, of specified).
+    # be used in any case, if specified).
     disable_task_docker_override: false
+
+    # Choose the default docker based on the Task properties,
+    # Examples: 'script.requirements', 'script.binary', 'script.repository', 'script.branch', 'project'
+    # Notice: Matching is done via regular expression, for example "^searchme$" will match exactly "searchme$" string
+    #
+    #     "default_docker": {
+    #         "image": "nvidia/cuda:10.2-cudnn7-runtime-ubuntu18.04",
+    #         # optional arguments to pass to docker image
+    #         # arguments: ["--ipc=host", ]
+    #         "match_rules": [
+    #             {
+    #                 "image": "sample_container:tag",
+    #                 "arguments": "-e VALUE=1 --ipc=host",
+    #                 "match": {
+    #                     "script": {
+    #                         "requirements": {
+    #                             "pip": {
+    #                                 "tensorflow": "~=1.6"
+    #                             }
+    #                         },
+    #                         "repository": "",
+    #                         "branch": "master"
+    #                     },
+    #                     "project": "example"
+    #                 }
+    #             },
+    #             {
+    #                 "image": "another_container:tag",
+    #                 "arguments": "",
+    #                 "match": {
+    #                     "project": "^examples", # anything that starts with "examples", e.g. "examples", "examples/sub_project"
+    #                 }
+    #             }
+    #         ]
+    #     },
+    #
 }

clearml_agent/backend_api/session/defs.py

@@ -20,6 +20,7 @@ ENV_PROPAGATE_EXITCODE = EnvEntry("CLEARML_AGENT_PROPAGATE_EXITCODE", type=bool,
 ENV_INITIAL_CONNECT_RETRY_OVERRIDE = EnvEntry(
     'CLEARML_AGENT_INITIAL_CONNECT_RETRY_OVERRIDE', default=True, converter=safe_text_to_bool
 )
+ENV_FORCE_MAX_API_VERSION = EnvEntry("CLEARML_AGENT_FORCE_MAX_API_VERSION", type=str)
 
 """
 Experimental option to set the request method for all API requests and auth login.

clearml_agent/backend_api/session/session.py

@@ -2,22 +2,25 @@
 import json as json_lib
 import os
 import sys
+import time
 import types
+from random import SystemRandom
 from socket import gethostname
 from typing import Optional
 
 import jwt
 import requests
 import six
+from requests import RequestException
 from requests.auth import HTTPBasicAuth
 from six.moves.urllib.parse import urlparse, urlunparse
 
 from clearml_agent.external.pyhocon import ConfigTree, ConfigFactory
-
 from .callresult import CallResult
 from .defs import (
     ENV_VERBOSE, ENV_HOST, ENV_ACCESS_KEY, ENV_SECRET_KEY, ENV_WEB_HOST, ENV_FILES_HOST, ENV_AUTH_TOKEN,
-    ENV_NO_DEFAULT_SERVER, ENV_DISABLE_VAULT_SUPPORT, ENV_INITIAL_CONNECT_RETRY_OVERRIDE, ENV_API_DEFAULT_REQ_METHOD, )
+    ENV_NO_DEFAULT_SERVER, ENV_DISABLE_VAULT_SUPPORT, ENV_INITIAL_CONNECT_RETRY_OVERRIDE, ENV_API_DEFAULT_REQ_METHOD,
+    ENV_FORCE_MAX_API_VERSION)
 from .request import Request, BatchRequest
 from .token_manager import TokenManager
 from ..config import load
@@ -25,6 +28,8 @@ from ..utils import get_http_session_with_retry, urllib_log_warning_setup
 from ...backend_config.environment import backward_compatibility_support
 from ...version import __version__
 
+sys_random = SystemRandom()
+
 
 class LoginError(Exception):
     pass
@@ -49,6 +54,7 @@ class Session(TokenManager):
     _session_initial_retry_connect_override = 4
     _write_session_data_size = 15000
     _write_session_timeout = (30.0, 30.)
+    _request_exception_retry_timeout = (2.0, 3.0)
 
     api_version = '2.1'
     feature_set = 'basic'
@@ -57,6 +63,7 @@ class Session(TokenManager):
     default_files = "https://demofiles.demo.clear.ml"
     default_key = "EGRTCO8JMSIGI6S39GTP43NFWXDQOW"
     default_secret = "x!XTov_G-#vspE*Y(h$Anm&DIc5Ou-F)jsl$PdOyj5wG1&E!Z8"
+    force_max_api_version = ENV_FORCE_MAX_API_VERSION.get()
 
     # TODO: add requests.codes.gateway_timeout once we support async commits
     _retry_codes = [
@@ -111,19 +118,9 @@ class Session(TokenManager):
         self._verbose = verbose if verbose is not None else ENV_VERBOSE.get()
         self._logger = logger
         self.__auth_token = None
+        self._propagate_exceptions_on_send = True
 
-        if ENV_API_DEFAULT_REQ_METHOD.get(default=None):
-            # Make sure we update the config object, so we pass it into the new containers when we map them
-            self.config.put("api.http.default_method", ENV_API_DEFAULT_REQ_METHOD.get())
-            # notice the default setting of Request.def_method are already set by the OS environment
-        elif self.config.get("api.http.default_method", None):
-            def_method = str(self.config.get("api.http.default_method", None)).strip()
-            if def_method.upper() not in ("GET", "POST", "PUT"):
-                raise ValueError(
-                    "api.http.default_method variable must be 'get' or 'post' (any case is allowed)."
-                )
-            Request.def_method = def_method
-            Request._method = Request.def_method
+        self.update_default_api_method()
 
         if ENV_AUTH_TOKEN.get(
             value_cb=lambda key, value: print("Using environment access token {}=********".format(key))
@@ -178,6 +175,10 @@ class Session(TokenManager):
         )
         # try to connect with the server
         self.refresh_token()
+
+        # for resilience, from now on we won't allow propagating exceptions when sending requests
+        self._propagate_exceptions_on_send = False
+
         # create the default session with many retries
         http_retries_config, self.__http_session = self._setup_session(http_retries_config)
 
@@ -198,6 +199,10 @@ class Session(TokenManager):
         # notice: this is across the board warning omission
         urllib_log_warning_setup(total_retries=http_retries_config.get('total', 0), display_warning_after=3)
 
+        if self.force_max_api_version and self.check_min_api_version(self.force_max_api_version):
+            print("Using forced API version {}".format(self.force_max_api_version))
+            Session.max_api_version = Session.api_version = str(self.force_max_api_version)
+
     def _setup_session(self, http_retries_config, initial_session=False, default_initial_connect_override=None):
         # type: (dict, bool, Optional[bool]) -> (dict, requests.Session)
         http_retries_config = http_retries_config or self.config.get(
@@ -223,7 +228,22 @@ class Session(TokenManager):
 
         return http_retries_config, get_http_session_with_retry(config=self.config or None, **http_retries_config)
 
+    def update_default_api_method(self):
+        if ENV_API_DEFAULT_REQ_METHOD.get(default=None):
+            # Make sure we update the config object, so we pass it into the new containers when we map them
+            self.config.put("api.http.default_method", ENV_API_DEFAULT_REQ_METHOD.get())
+            # notice the default setting of Request.def_method are already set by the OS environment
+        elif self.config.get("api.http.default_method", None):
+            def_method = str(self.config.get("api.http.default_method", None)).strip()
+            if def_method.upper() not in ("GET", "POST", "PUT"):
+                raise ValueError(
+                    "api.http.default_method variable must be 'get', 'post' or 'put' (any case is allowed)."
+                )
+            Request.def_method = def_method
+            Request._method = Request.def_method
+
     def load_vaults(self):
+        # () -> Optional[bool]
         if not self.check_min_api_version("2.15") or self.feature_set == "basic":
             return
 
@@ -244,12 +264,14 @@ class Session(TokenManager):
 
         # noinspection PyBroadException
         try:
-            res = self.send_request("users", "get_vaults", json={"enabled": True, "types": ["config"]})
+            # Use params and not data/json otherwise payload might be dropped if we're using GET with a strict firewall
+            res = self.send_request("users", "get_vaults", params="enabled=true&types=config&types=config")
             if res.ok:
                 vaults = res.json().get("data", {}).get("vaults", [])
                 data = list(filter(None, map(parse, vaults)))
                 if data:
                     self.config.set_overrides(*data)
+                    return True
             elif res.status_code != 404:
                 raise Exception(res.json().get("meta", {}).get("result_msg", res.text))
         except Exception as ex:
@@ -272,6 +294,7 @@ class Session(TokenManager):
         data=None,
         json=None,
         refresh_token_if_unauthorized=True,
+        params=None,
     ):
         """ Internal implementation for making a raw API request.
             - Constructs the api endpoint name
@@ -295,6 +318,7 @@ class Session(TokenManager):
             if version
             else "{host}/{service}.{action}"
         ).format(**locals())
+
         while True:
             if data and len(data) > self._write_session_data_size:
                 timeout = self._write_session_timeout
@@ -302,16 +326,29 @@ class Session(TokenManager):
                 timeout = self._session_initial_timeout
             else:
                 timeout = self._session_timeout
-            res = self.__http_session.request(
-                method, url, headers=headers, auth=auth, data=data, json=json, timeout=timeout)
+
+            try:
+                res = self.__http_session.request(
+                    method, url, headers=headers, auth=auth, data=data, json=json, timeout=timeout, params=params)
+            except RequestException as ex:
+                if self._propagate_exceptions_on_send:
+                    raise
+                sleep_time = sys_random.uniform(*self._request_exception_retry_timeout)
+                self._logger.error(
+                    "{} exception sending {} {}: {} (retrying in {:.1f}sec)".format(
+                        type(ex).__name__, method.upper(), url, str(ex), sleep_time
+                    )
+                )
+                time.sleep(sleep_time)
+                continue
 
             if (
                 refresh_token_if_unauthorized
                 and res.status_code == requests.codes.unauthorized
                 and not token_refreshed_on_error
             ):
-                # it seems we're unauthorized, so we'll try to refresh our token once in case permissions changed since
-                # the last time we got the token, and try again
+                # it seems we're unauthorized, so we'll try to refresh our token once in case permissions changed
+                # since the last time we got the token, and try again
                 self.refresh_token()
                 token_refreshed_on_error = True
                 # try again
@@ -348,6 +385,7 @@ class Session(TokenManager):
         data=None,
         json=None,
         async_enable=False,
+        params=None,
     ):
         """
         Send a raw API request.
@@ -360,6 +398,7 @@ class Session(TokenManager):
                      content type will be application/json)
         :param data: Dictionary, bytes, or file-like object to send in the request body
         :param async_enable: whether request is asynchronous
+        :param params: additional query parameters
         :return: requests Response instance
         """
         headers = self.add_auth_headers(
@@ -376,6 +415,7 @@ class Session(TokenManager):
             headers=headers,
             data=data,
             json=json,
+            params=params,
         )
 
     def send_request_batch(
@@ -628,15 +668,14 @@ class Session(TokenManager):
 
         res = None
         try:
-            data = {"expiration_sec": exp} if exp else {}
             res = self._send_request(
                 method=Request.def_method,
                 service="auth",
                 action="login",
                 auth=auth,
-                json=data,
                 headers=headers,
                 refresh_token_if_unauthorized=False,
+                params={"expiration_sec": exp} if exp else {},
             )
             try:
                 resp = res.json()
@@ -675,3 +714,13 @@ class Session(TokenManager):
         return "{self.__class__.__name__}[{self.host}, {self.access_key}/{secret_key}]".format(
             self=self, secret_key=self.secret_key[:5] + "*" * (len(self.secret_key) - 5)
         )
+
+    @property
+    def propagate_exceptions_on_send(self):
+        # type: () -> bool
+        return self._propagate_exceptions_on_send
+
+    @propagate_exceptions_on_send.setter
+    def propagate_exceptions_on_send(self, value):
+        # type: (bool) -> None
+        self._propagate_exceptions_on_send = value

clearml_agent/backend_api/utils.py

@@ -86,7 +86,10 @@ def get_http_session_with_retry(
     session = requests.Session()
 
     if backoff_max is not None:
-        Retry.BACKOFF_MAX = backoff_max
+        if "BACKOFF_MAX" in vars(Retry):
+            Retry.BACKOFF_MAX = backoff_max
+        else:
+            Retry.DEFAULT_BACKOFF_MAX = backoff_max
 
     retry = Retry(
         total=total, connect=connect, read=read, redirect=redirect, status=status,

clearml_agent/backend_config/config.py

@@ -297,6 +297,9 @@ class Config(object):
     def put(self, key, value):
         self._config.put(key, value)
 
+    def pop(self, key, default=None):
+        return self._config.pop(key, default=default)
+
     def to_dict(self):
         return self._config.as_plain_ordered_dict()
 

clearml_agent/backend_config/utils.py

@@ -52,6 +52,7 @@ def apply_files(config):
         target_fmt = data.get("target_format", "string")
         overwrite = bool(data.get("overwrite", True))
         contents = data.get("contents")
+        mode = data.get("mode")
 
         target = Path(expanduser(expandvars(path)))
 
@@ -110,3 +111,14 @@ def apply_files(config):
         except Exception as ex:
             print("Skipped [{}]: failed saving file {} ({})".format(key, target, ex))
             continue
+
+        try:
+            if mode:
+                if isinstance(mode, int):
+                    mode = int(str(mode), 8)
+                else:
+                    mode = int(mode, 8)
+                target.chmod(mode)
+        except Exception as ex:
+            print("Skipped [{}]: failed setting mode {} for {} ({})".format(key, mode, target, ex))
+            continue

clearml_agent/commands/base.py

@@ -118,13 +118,15 @@ class ServiceCommandSection(BaseCommandSection):
         """ The name of the REST service used by this command """
         pass
 
-    def get(self, endpoint, *args, session=None, **kwargs):
+    def get(self, endpoint, *args, service=None, session=None, **kwargs):
         session = session or self._session
-        return session.get(service=self.service, action=endpoint, *args, **kwargs)
+        service = service or self.service
+        return session.get(service=service, action=endpoint, *args, **kwargs)
 
-    def post(self, endpoint, *args, session=None, **kwargs):
+    def post(self, endpoint, *args, service=None, session=None, **kwargs):
         session = session or self._session
-        return session.post(service=self.service, action=endpoint, *args, **kwargs)
+        service = service or self.service
+        return session.post(service=service, action=endpoint, *args, **kwargs)
 
     def get_with_act_as(self, endpoint, *args, **kwargs):
         return self._session.get_with_act_as(service=self.service, action=endpoint, *args, **kwargs)

clearml_agent/commands/config.py

@@ -44,7 +44,7 @@ def main():
 
     if conf_file.exists() and conf_file.is_file() and conf_file.stat().st_size > 0:
         print('Configuration file already exists: {}'.format(str(conf_file)))
-        print('Leaving setup, feel free to edit the configuration file.')
+        print('Leaving setup. If you\'ve previously initialized the ClearML SDK on this machine, manually add an \'agent\' section to this file.')
         return
 
     print(description, end='')

clearml_agent/commands/resolver.py

@@ -109,15 +109,15 @@ def resolve_default_container(session, task_id, container_config):
                     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
+        # if match.get('image', None):
+        #     # noinspection PyBroadException
+        #     try:
+        #         if not re.search(match.get('image', None), requested_container.get('image', '')):
+        #             continue
+        #     except Exception:
+        #         print('Failed parsing regular expression \"{}\" in rule: {}'.format(
+        #             match.get('image', None), entry))
+        #         continue
 
         matched = True
         for req_section in ['script.requirements.pip', 'script.requirements.conda']:
@@ -156,8 +156,8 @@ def resolve_default_container(session, task_id, container_config):
             break
 
         if matched:
-            if not container_config.get('container'):
-                container_config['container'] = entry.get('image', None)
+            if not container_config.get('image'):
+                container_config['image'] = 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())

clearml_agent/commands/worker.py

@@ -1,6 +1,7 @@
 from __future__ import print_function, division, unicode_literals
 
 import errno
+import functools
 import json
 import logging
 import os
@@ -39,6 +40,7 @@ from clearml_agent.backend_api.session import CallResult, Request
 from clearml_agent.backend_api.session.defs import (
     ENV_ENABLE_ENV_CONFIG_SECTION, ENV_ENABLE_FILES_CONFIG_SECTION,
     ENV_VENV_CONFIGURED, ENV_PROPAGATE_EXITCODE, )
+from clearml_agent.backend_config import Config
 from clearml_agent.backend_config.defs import UptimeConf
 from clearml_agent.backend_config.utils import apply_environment, apply_files
 from clearml_agent.backend_config.converters import text_to_int
@@ -57,10 +59,7 @@ from clearml_agent.definitions import (
     ENV_WORKER_ID,
     ENV_WORKER_TAGS,
     ENV_DOCKER_SKIP_GPUS_FLAG,
-    ENV_AGENT_SECRET_KEY,
     ENV_AGENT_AUTH_TOKEN,
-    ENV_AWS_SECRET_KEY,
-    ENV_AZURE_ACCOUNT_KEY,
     ENV_AGENT_DISABLE_SSH_MOUNT,
     ENV_SSH_AUTH_SOCK,
     ENV_AGENT_SKIP_PIP_VENV_INSTALL,
@@ -71,6 +70,11 @@ from clearml_agent.definitions import (
     ENV_DEBUG_INFO,
     ENV_CHILD_AGENTS_COUNT_CMD,
     ENV_DOCKER_ARGS_FILTERS,
+    ENV_FORCE_SYSTEM_SITE_PACKAGES,
+    ENV_SERVICES_DOCKER_RESTART,
+    ENV_CONFIG_BC_IN_STANDALONE,
+    ENV_FORCE_DOCKER_AGENT_REPO,
+    ENV_EXTRA_DOCKER_LABELS,
 )
 from clearml_agent.definitions import WORKING_REPOSITORY_DIR, PIP_EXTRA_INDICES
 from clearml_agent.errors import (
@@ -316,6 +320,37 @@ def get_next_task(session, queue, get_task_info=False):
     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):
     """
     Returns dict with Task docker container setup {container: '', arguments: '', setup_shell_script: ''}
@@ -337,16 +372,19 @@ def get_task_container(session, task_id):
             container = {}
     else:
         response = get_task(session, task_id, only_fields=["execution.docker_cmd"])
-        task_docker_cmd_parts = shlex.split(str(response.execution.docker_cmd or '').strip())
-        try:
-            container = dict(
-                container=task_docker_cmd_parts[0],
-                arguments=task_docker_cmd_parts[1:] if len(task_docker_cmd_parts[0]) > 1 else ''
-            )
-        except (ValueError, TypeError):
-            container = {}
+        container = {}
+        if response.execution:
+            task_docker_cmd_parts = shlex.split(str(response.execution.docker_cmd or '').strip())
+            if task_docker_cmd_parts:
+                try:
+                    container = dict(
+                        image=task_docker_cmd_parts[0],
+                        arguments=task_docker_cmd_parts[1:] if len(task_docker_cmd_parts[0]) > 1 else ''
+                    )
+                except (ValueError, TypeError):
+                    pass
 
-    if (not container or not container.get('container')) and session.check_min_api_version("2.13"):
+    if (not container or not container.get('image')) and session.check_min_api_version("2.13"):
         container = resolve_default_container(session=session, task_id=task_id, container_config=container)
 
     return container
@@ -688,7 +726,7 @@ class Worker(ServiceCommandSection):
             self._docker_args_filters = []
 
         self._task_ping_interval_sec = max(
-            0, text_to_int(self._session.config.get("agent.task_ping_interval_sec", 120.0))
+            0, text_to_int(self._session.config.get("agent.task_ping_interval_sec", 60.0))
         )
 
     @classmethod
@@ -736,8 +774,68 @@ class Worker(ServiceCommandSection):
         except Exception:
             pass
 
+    def _get_docker_restart_value(self, task_session, task_id: str):
+        try:
+            self._session.verify_feature_set('advanced')
+        except ValueError:
+            return
+
+        restart = (ENV_SERVICES_DOCKER_RESTART.get() or "").strip()
+        if not restart:
+            return
+
+        # Parse value and selector
+        restart_value, _, selector = restart.partition(";")
+
+        if restart_value not in ("unless-stopped", "no", "always") and not restart_value.startswith("on-failure"):
+            self.log.error(
+                "Invalid value \"{}\" provided for {}, ignoring".format(restart, ENV_SERVICES_DOCKER_RESTART.vars[0])
+            )
+            return
+
+        if not selector:
+            return restart_value
+
+        path, _, expected_value = selector.partition("=")
+
+        result = task_session.send_request(
+            service='tasks',
+            action='get_all',
+            json={'id': [task_id], 'only_fields': [path], 'search_hidden': True},
+            method=Request.def_method,
+        )
+        if not result.ok:
+            result_msg = self._get_path(result.json(), 'meta', 'result_msg')
+            self.log.error(
+                "Failed obtaining selector value for restart option \"{}\", ignoring: {}".format(selector, result_msg)
+            )
+            return
+
+        not_found = object()
+        try:
+            value = self._get_path(result.json(), 'data', 'tasks', 0, *path.split("."), default=not_found)
+        except (ValueError, TypeError):
+            return
+
+        if value is not_found:
+            return
+
+        if not expected_value:
+            return restart_value
+
+        # noinspection PyBroadException
+        try:
+            if (
+                (isinstance(value, bool) and value == strtobool(expected_value))  # check first - bool is also an int
+                or (isinstance(value, (int, float)) and value == float(expected_value))
+                or (str(value) == str(expected_value))
+            ):
+                return restart_value
+        except Exception as ex:
+            pass
+
     def run_one_task(self, queue, task_id, worker_args, docker=None, task_session=None):
-        # type: (Text, Text, WorkerParams, Optional[Text], Optional[Session]) -> int
+        # type: (Text, Text, WorkerParams, Optional[Text], Optional[Session]) -> Optional[int]
         """
         Run one task pulled from queue.
         :param queue: ID of queue that task was pulled from
@@ -812,6 +910,7 @@ class Worker(ServiceCommandSection):
                 docker_image=docker_image,
                 docker_arguments=docker_arguments,
                 docker_bash_setup_script=docker_setup_script,
+                restart=self._get_docker_restart_value(task_session, task_id),
             )
             if self._impersonate_as_task_owner:
                 docker_params["auth_token"] = task_session.token
@@ -826,11 +925,21 @@ class Worker(ServiceCommandSection):
 
             name_format = self._session.config.get('agent.docker_container_name_format', None)
             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:
                     name = name_format.format(
                         task_id=re.sub(r'[^a-zA-Z0-9._-]', '-', task_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:
                     print("Warning: failed generating docker container name: {}".format(ex))
@@ -1786,7 +1895,8 @@ class Worker(ServiceCommandSection):
                 if stderr:
                     stderr.flush()
 
-                if self._task_ping_interval_sec and time() - last_task_ping > self._task_ping_interval_sec:
+                if not stopping and self._task_ping_interval_sec and \
+                        time() - last_task_ping > self._task_ping_interval_sec:
                     # noinspection PyBroadException
                     try:
                         res = (session or self._session).send(tasks_api.PingRequest(task=task_id))
@@ -1795,7 +1905,7 @@ class Worker(ServiceCommandSection):
                     except Exception as ex:
                         self.log.error("Failed sending ping: %s", str(ex))
                     finally:
-                        self._task_ping_interval_sec = time()
+                        last_task_ping = time()
 
                 # get diff from previous poll
                 printed_lines, stdout_pos_count = _print_file(stdout_path, stdout_pos_count)
@@ -1968,6 +2078,11 @@ class Worker(ServiceCommandSection):
             except Exception as ex:
                 print("Error: failed applying files from configuration: {}".format(ex))
 
+        try:
+            self._session.update_default_api_method()
+        except Exception as ex:
+            print("Error: failed updating default API method: {}".format(ex))
+
     @resolve_names
     def build(
         self,
@@ -1982,6 +2097,10 @@ class Worker(ServiceCommandSection):
     ):
         if not task_id:
             raise CommandFailedError("Worker build must have valid task id")
+        
+        if target and not os.path.isabs(target):
+            # Non absolute target path will lead to errors with relative python executable
+            target = os.path.abspath(target)
 
         self._session.print_configuration()
 
@@ -2101,12 +2220,14 @@ class Worker(ServiceCommandSection):
         print('Building Task {} inside docker image: {} {} setup_script={}\n'.format(
             task_id, docker_image, docker_arguments or '', docker_setup_script or ''))
         full_docker_cmd = self.docker_image_func(
-            docker_image=docker_image, docker_arguments=docker_arguments, docker_bash_setup_script=docker_setup_script)
+            docker_image=docker_image, docker_arguments=docker_arguments, docker_bash_setup_script=docker_setup_script
+        )
 
         end_of_build_marker = "build.done=true"
         docker_cmd_suffix = ' build --id {task_id} --install-globally; ' \
-                            'echo "" >> {conf_file} ; ' \
-                            'echo {end_of_build_marker} >> {conf_file} ; ' \
+                            'ORG=$(stat -c "%u:%g" {conf_file}) ; chown $(whoami):$(whoami) {conf_file} ; ' \
+                            'echo "" >> {conf_file} ; echo {end_of_build_marker} >> {conf_file} ; ' \
+                            'chown $ORG {conf_file} ; ' \
                             'bash'.format(
                                 task_id=task_id,
                                 end_of_build_marker=end_of_build_marker,
@@ -2125,10 +2246,16 @@ class Worker(ServiceCommandSection):
 
         # now we need to wait until the line shows on our configuration file.
         while True:
-            while temp_config.stat().st_mtime == base_time_stamp:
-                sleep(5.0)
-            with open(temp_config.as_posix()) as f:
-                lines = [l.strip() for l in f.readlines()]
+            # noinspection PyBroadException
+            try:
+                while temp_config.stat().st_mtime == base_time_stamp:
+                    sleep(5.0)
+                with open(temp_config.as_posix()) as f:
+                    lines = [l.strip() for l in f.readlines()]
+            except Exception as ex:
+                # print("Failed reading status file [{}], retrying in 2 seconds".format(ex))
+                sleep(2.0)
+
             if 'build.done=true' in lines:
                 break
             base_time_stamp = temp_config.stat().st_mtime
@@ -2157,6 +2284,8 @@ class Worker(ServiceCommandSection):
         print(commit_docker(container_name=target, docker_id=docker_id, apply_change=change))
         shutdown_docker_process(docker_id=docker_id)
 
+        safe_remove_file(temp_config.as_posix())
+
         return
 
     def _get_task_python_version(self, task):
@@ -2836,8 +2965,8 @@ class Worker(ServiceCommandSection):
         # Todo: add support for poetry caching
         if not self.poetry.enabled:
             # add to cache
-            print('Adding venv into cache: {}'.format(add_venv_folder_cache))
             if add_venv_folder_cache:
+                print('Adding venv into cache: {}'.format(add_venv_folder_cache))
                 self.package_api.add_cached_venv(
                     requirements=[freeze, previous_reqs],
                     docker_cmd=execution_info.docker_cmd if execution_info else None,
@@ -2858,19 +2987,27 @@ class Worker(ServiceCommandSection):
             self.log_traceback(e)
         return freeze
 
-    def _install_poetry_requirements(self, repo_info):
-        # type: (Optional[RepoInfo]) -> Optional[PoetryAPI]
+    def _install_poetry_requirements(self, repo_info, working_dir=None):
+        # type: (Optional[RepoInfo], Optional[str]) -> Optional[PoetryAPI]
         if not repo_info:
             return None
+
+        files_from_working_dir = self._session.config.get(
+            "agent.package_manager.poetry_files_from_repo_working_dir", False)
+        lockfile_path = Path(repo_info.root) / ((working_dir or "") if files_from_working_dir else "")
+
         try:
             if not self.poetry.enabled:
                 return None
-            self.poetry.initialize(cwd=repo_info.root)
-            api = self.poetry.get_api(repo_info.root)
+
+            self.poetry.initialize(cwd=lockfile_path)
+            api = self.poetry.get_api(lockfile_path)
             if api.enabled:
                 print('Poetry Enabled: Ignoring requested python packages, using repository poetry lock file!')
                 api.install()
                 return api
+            
+            print(f"Could not find pyproject.toml or poetry.lock file in {lockfile_path} \n")
         except Exception as ex:
             self.log.error("failed installing poetry requirements: {}".format(ex))
         return None
@@ -2901,7 +3038,8 @@ class Worker(ServiceCommandSection):
          """
         if package_api:
             package_api.cwd = cwd
-        api = self._install_poetry_requirements(repo_info)
+
+        api = self._install_poetry_requirements(repo_info, execution.working_dir)
         if api:
             # update back the package manager, this hack should be fixed
             if package_api == self.package_api:
@@ -3425,6 +3563,11 @@ class Worker(ServiceCommandSection):
         requirements_manager.translator.enabled = False
         print(requirements_manager.replace(contents))
 
+    def remove_non_backwards_compatible_entries(self, config: Config):
+        if not self._standalone_mode or not ENV_CONFIG_BC_IN_STANDALONE.get() or self._session.feature_set == "basic":
+            return
+        config.pop("agent.package_manager.pip_version")  # removed due to a breaking change in v1.5.1
+
     def get_docker_config_cmd(self, docker_args, clean_api_credentials=False):
         docker_image = str(ENV_DOCKER_IMAGE.get() or
                            self._session.config.get("agent.default_docker.image", "nvidia/cuda")) \
@@ -3447,6 +3590,7 @@ class Worker(ServiceCommandSection):
             DockerArgsSanitizer.sanitize_docker_command(self._session, self._docker_arguments) or ''))
 
         temp_config = deepcopy(self._session.config)
+        self.remove_non_backwards_compatible_entries(temp_config)
         mounted_cache_dir = temp_config.get(
             "agent.docker_internal_mounts.sdk_cache", self._docker_fixed_user_cache)
         mounted_pip_dl_dir = temp_config.get(
@@ -3458,7 +3602,6 @@ class Worker(ServiceCommandSection):
         temp_config.put("sdk.storage.cache.default_base_dir", mounted_cache_dir)
         temp_config.put("agent.pip_download_cache.path", mounted_pip_dl_dir)
         temp_config.put("agent.vcs_cache.path", mounted_vcs_cache)
-        temp_config.put("agent.package_manager.system_site_packages", True)
         temp_config.put("agent.package_manager.conda_env_as_base_docker", False)
         temp_config.put("agent.default_python", "")
         temp_config.put("agent.python_binary", "")
@@ -3470,6 +3613,11 @@ class Worker(ServiceCommandSection):
         temp_config.put("agent.git_pass", (ENV_AGENT_GIT_PASS.get() or
                                            self._session.config.get("agent.git_pass", None)))
 
+        force_system_site_packages = ENV_FORCE_SYSTEM_SITE_PACKAGES.get()
+        force_system_site_packages = force_system_site_packages if force_system_site_packages is not None else True
+        if force_system_site_packages:
+            temp_config.put("agent.package_manager.system_site_packages", True)
+
         if temp_config.get("agent.venvs_cache.path", None):
             temp_config.put("agent.venvs_cache.path", '/root/.clearml/venvs-cache')
 
@@ -3697,11 +3845,20 @@ class Worker(ServiceCommandSection):
             name=None,
             mount_ssh=None, mount_ssh_ro=None, mount_apt_cache=None, mount_pip_cache=None, mount_poetry_cache=None,
             env_task_id=None,
+            restart=None,
     ):
         self.debug("Constructing docker command", context="docker")
         docker = 'docker'
 
         base_cmd = [docker, 'run', '-t']
+        use_rm = True
+        if restart:
+            if restart in ("unless-stopped", "no", "always") or restart.startswith("on-failure"):
+                base_cmd += ["--restart", restart]
+                use_rm = False
+            else:
+                self.log.error("Invalid restart value \"{}\" , ignoring".format(restart))
+
         update_scheme = ""
         dockers_nvidia_visible_devices = 'all'
         gpu_devices = Session.get_nvidia_visible_env()
@@ -3737,6 +3894,10 @@ class Worker(ServiceCommandSection):
         base_cmd += ['-l', self._worker_label.format(worker_id)]
         base_cmd += ['-l', self._parent_worker_label.format(parent_worker_id)]
 
+        extra_labels = ENV_EXTRA_DOCKER_LABELS.get()
+        for label in (extra_labels or []):
+            base_cmd += ['-l', label]
+
         self.debug("Command: {}".format(base_cmd), context="docker")
 
         # check if running inside a kubernetes
@@ -3819,6 +3980,7 @@ class Worker(ServiceCommandSection):
         # if we are running a RC version, install the same version in the docker
         # because the default latest, will be a release version (not RC)
         specify_version = ''
+        # noinspection PyBroadException
         try:
             from clearml_agent.version import __version__
             _version_parts = __version__.split('.')
@@ -3827,13 +3989,15 @@ class Worker(ServiceCommandSection):
         except:
             pass
 
+        force_agent_repo = ENV_FORCE_DOCKER_AGENT_REPO.get()
+
         if os.environ.get('FORCE_LOCAL_CLEARML_AGENT_WHEEL'):
             local_wheel = os.path.expanduser(os.environ.get('FORCE_LOCAL_CLEARML_AGENT_WHEEL'))
             docker_wheel = '/tmp/{}'.format(basename(local_wheel))
             base_cmd += ['-v', local_wheel + ':' + docker_wheel]
             clearml_agent_wheel = '\"{}\"'.format(docker_wheel)
-        elif os.environ.get('FORCE_CLEARML_AGENT_REPO'):
-            clearml_agent_wheel = os.environ.get('FORCE_CLEARML_AGENT_REPO')
+        elif force_agent_repo:
+            clearml_agent_wheel = force_agent_repo
         else:
             # clearml-agent{specify_version}
             clearml_agent_wheel = 'clearml-agent{specify_version}'.format(specify_version=specify_version)
@@ -3914,7 +4078,8 @@ class Worker(ServiceCommandSection):
             (['-v', host_cache+':'+mounted_cache] if host_cache else []) +
             (['-v', host_vcs_cache+':'+mounted_vcs_cache] if host_vcs_cache else []) +
             (['-v', host_venvs_cache + ':' + mounted_venvs_cache] if host_venvs_cache else []) +
-            ['--rm', docker_image, 'bash', '-c',
+            (['--rm'] if use_rm else []) +
+            [docker_image, 'bash', '-c',
                 update_scheme +
                 extra_shell_script +
                 "cp {} {} ; ".format(DOCKER_ROOT_CONF_FILE, DOCKER_DEFAULT_CONF_FILE) +
@@ -4120,6 +4285,15 @@ class Worker(ServiceCommandSection):
                     " found {})".format(role)
                 )
 
+    @staticmethod
+    def _get_path(d, *path, default=None):
+        try:
+            return functools.reduce(
+                lambda a, b: a[b], path, d
+            )
+        except (IndexError, KeyError):
+            return default
+
 
 if __name__ == "__main__":
     pass

clearml_agent/definitions.py

@@ -5,9 +5,9 @@ from enum import IntEnum
 from os import getenv, environ
 from typing import Text, Optional, Union, Tuple, Any
 
+import six
 from pathlib2 import Path
 
-import six
 from clearml_agent.helper.base import normalize_path
 
 PROGRAM_NAME = "clearml-agent"
@@ -69,42 +69,65 @@ ENV_AWS_SECRET_KEY = EnvironmentConfig("AWS_SECRET_ACCESS_KEY")
 ENV_AZURE_ACCOUNT_KEY = EnvironmentConfig("AZURE_STORAGE_KEY")
 
 ENVIRONMENT_CONFIG = {
-    "api.api_server": EnvironmentConfig("CLEARML_API_HOST", "TRAINS_API_HOST", ),
-    "api.files_server": EnvironmentConfig("CLEARML_FILES_HOST", "TRAINS_FILES_HOST", ),
-    "api.web_server": EnvironmentConfig("CLEARML_WEB_HOST", "TRAINS_WEB_HOST", ),
+    "api.api_server": EnvironmentConfig(
+        "CLEARML_API_HOST",
+        "TRAINS_API_HOST",
+    ),
+    "api.files_server": EnvironmentConfig(
+        "CLEARML_FILES_HOST",
+        "TRAINS_FILES_HOST",
+    ),
+    "api.web_server": EnvironmentConfig(
+        "CLEARML_WEB_HOST",
+        "TRAINS_WEB_HOST",
+    ),
     "api.credentials.access_key": EnvironmentConfig(
-        "CLEARML_API_ACCESS_KEY", "TRAINS_API_ACCESS_KEY",
+        "CLEARML_API_ACCESS_KEY",
+        "TRAINS_API_ACCESS_KEY",
     ),
     "api.credentials.secret_key": ENV_AGENT_SECRET_KEY,
-    "agent.worker_name": EnvironmentConfig("CLEARML_WORKER_NAME", "TRAINS_WORKER_NAME", ),
-    "agent.worker_id": EnvironmentConfig("CLEARML_WORKER_ID", "TRAINS_WORKER_ID", ),
-    "agent.cuda_version": EnvironmentConfig(
-        "CLEARML_CUDA_VERSION", "TRAINS_CUDA_VERSION", "CUDA_VERSION"
+    "agent.worker_name": EnvironmentConfig(
+        "CLEARML_WORKER_NAME",
+        "TRAINS_WORKER_NAME",
     ),
-    "agent.cudnn_version": EnvironmentConfig(
-        "CLEARML_CUDNN_VERSION", "TRAINS_CUDNN_VERSION", "CUDNN_VERSION"
-    ),
-    "agent.cpu_only": EnvironmentConfig(
-        names=("CLEARML_CPU_ONLY", "TRAINS_CPU_ONLY", "CPU_ONLY"), type=bool
+    "agent.worker_id": EnvironmentConfig(
+        "CLEARML_WORKER_ID",
+        "TRAINS_WORKER_ID",
     ),
+    "agent.cuda_version": EnvironmentConfig("CLEARML_CUDA_VERSION", "TRAINS_CUDA_VERSION", "CUDA_VERSION"),
+    "agent.cudnn_version": EnvironmentConfig("CLEARML_CUDNN_VERSION", "TRAINS_CUDNN_VERSION", "CUDNN_VERSION"),
+    "agent.cpu_only": EnvironmentConfig(names=("CLEARML_CPU_ONLY", "TRAINS_CPU_ONLY", "CPU_ONLY"), type=bool),
     "agent.crash_on_exception": EnvironmentConfig("CLEAMRL_AGENT_CRASH_ON_EXCEPTION", type=bool),
     "sdk.aws.s3.key": EnvironmentConfig("AWS_ACCESS_KEY_ID"),
     "sdk.aws.s3.secret": ENV_AWS_SECRET_KEY,
     "sdk.aws.s3.region": EnvironmentConfig("AWS_DEFAULT_REGION"),
-    "sdk.azure.storage.containers.0": {'account_name': EnvironmentConfig("AZURE_STORAGE_ACCOUNT"),
-                                       'account_key': ENV_AZURE_ACCOUNT_KEY},
+    "sdk.azure.storage.containers.0": {
+        "account_name": EnvironmentConfig("AZURE_STORAGE_ACCOUNT"),
+        "account_key": ENV_AZURE_ACCOUNT_KEY,
+    },
     "sdk.google.storage.credentials_json": EnvironmentConfig("GOOGLE_APPLICATION_CREDENTIALS"),
 }
 
 ENVIRONMENT_SDK_PARAMS = {
-    "task_id": ("CLEARML_TASK_ID", "TRAINS_TASK_ID", ),
-    "config_file": ("CLEARML_CONFIG_FILE", "TRAINS_CONFIG_FILE", ),
-    "log_level": ("CLEARML_LOG_LEVEL", "TRAINS_LOG_LEVEL", ),
-    "log_to_backend": ("CLEARML_LOG_TASK_TO_BACKEND", "TRAINS_LOG_TASK_TO_BACKEND", ),
+    "task_id": (
+        "CLEARML_TASK_ID",
+        "TRAINS_TASK_ID",
+    ),
+    "config_file": (
+        "CLEARML_CONFIG_FILE",
+        "TRAINS_CONFIG_FILE",
+    ),
+    "log_level": (
+        "CLEARML_LOG_LEVEL",
+        "TRAINS_LOG_LEVEL",
+    ),
+    "log_to_backend": (
+        "CLEARML_LOG_TASK_TO_BACKEND",
+        "TRAINS_LOG_TASK_TO_BACKEND",
+    ),
 }
 
-ENVIRONMENT_BACKWARD_COMPATIBLE = EnvironmentConfig(
-    names=("CLEARML_AGENT_ALG_ENV", "TRAINS_AGENT_ALG_ENV"), type=bool)
+ENVIRONMENT_BACKWARD_COMPATIBLE = EnvironmentConfig(names=("CLEARML_AGENT_ALG_ENV", "TRAINS_AGENT_ALG_ENV"), type=bool)
 
 VIRTUAL_ENVIRONMENT_PATH = {
     "python2": normalize_path(CONFIG_DIR, "py2venv"),
@@ -123,38 +146,66 @@ TOKEN_EXPIRATION_SECONDS = int(timedelta(days=2).total_seconds())
 
 METADATA_EXTENSION = ".json"
 
-DEFAULT_VENV_UPDATE_URL = (
-    "https://raw.githubusercontent.com/Yelp/venv-update/v3.2.4/venv_update.py"
-)
+DEFAULT_VENV_UPDATE_URL = "https://raw.githubusercontent.com/Yelp/venv-update/v3.2.4/venv_update.py"
 WORKING_REPOSITORY_DIR = "task_repository"
 WORKING_STANDALONE_DIR = "code"
 DEFAULT_VCS_CACHE = normalize_path(CONFIG_DIR, "vcs-cache")
-PIP_EXTRA_INDICES = [
-]
+PIP_EXTRA_INDICES = []
 DEFAULT_PIP_DOWNLOAD_CACHE = normalize_path(CONFIG_DIR, "pip-download-cache")
-ENV_DOCKER_IMAGE = EnvironmentConfig('CLEARML_DOCKER_IMAGE', 'TRAINS_DOCKER_IMAGE')
-ENV_WORKER_ID = EnvironmentConfig('CLEARML_WORKER_ID', 'TRAINS_WORKER_ID')
-ENV_WORKER_TAGS = EnvironmentConfig('CLEARML_WORKER_TAGS')
-ENV_AGENT_SKIP_PIP_VENV_INSTALL = EnvironmentConfig('CLEARML_AGENT_SKIP_PIP_VENV_INSTALL')
-ENV_AGENT_SKIP_PYTHON_ENV_INSTALL = EnvironmentConfig('CLEARML_AGENT_SKIP_PYTHON_ENV_INSTALL', type=bool)
-ENV_DOCKER_SKIP_GPUS_FLAG = EnvironmentConfig('CLEARML_DOCKER_SKIP_GPUS_FLAG', 'TRAINS_DOCKER_SKIP_GPUS_FLAG')
-ENV_AGENT_GIT_USER = EnvironmentConfig('CLEARML_AGENT_GIT_USER', 'TRAINS_AGENT_GIT_USER')
-ENV_AGENT_GIT_PASS = EnvironmentConfig('CLEARML_AGENT_GIT_PASS', 'TRAINS_AGENT_GIT_PASS')
-ENV_AGENT_GIT_HOST = EnvironmentConfig('CLEARML_AGENT_GIT_HOST', 'TRAINS_AGENT_GIT_HOST')
-ENV_AGENT_DISABLE_SSH_MOUNT = EnvironmentConfig('CLEARML_AGENT_DISABLE_SSH_MOUNT', type=bool)
-ENV_SSH_AUTH_SOCK = EnvironmentConfig('SSH_AUTH_SOCK')
-ENV_TASK_EXECUTE_AS_USER = EnvironmentConfig('CLEARML_AGENT_EXEC_USER', 'TRAINS_AGENT_EXEC_USER')
-ENV_TASK_EXTRA_PYTHON_PATH = EnvironmentConfig('CLEARML_AGENT_EXTRA_PYTHON_PATH', 'TRAINS_AGENT_EXTRA_PYTHON_PATH')
-ENV_DOCKER_HOST_MOUNT = EnvironmentConfig('CLEARML_AGENT_K8S_HOST_MOUNT', 'CLEARML_AGENT_DOCKER_HOST_MOUNT',
-                                          'TRAINS_AGENT_K8S_HOST_MOUNT', 'TRAINS_AGENT_DOCKER_HOST_MOUNT')
-ENV_VENV_CACHE_PATH = EnvironmentConfig('CLEARML_AGENT_VENV_CACHE_PATH')
-ENV_EXTRA_DOCKER_ARGS = EnvironmentConfig('CLEARML_AGENT_EXTRA_DOCKER_ARGS', type=list)
-ENV_DEBUG_INFO = EnvironmentConfig('CLEARML_AGENT_DEBUG_INFO')
-ENV_CHILD_AGENTS_COUNT_CMD = EnvironmentConfig('CLEARML_AGENT_CHILD_AGENTS_COUNT_CMD')
-ENV_DOCKER_ARGS_FILTERS = EnvironmentConfig('CLEARML_AGENT_DOCKER_ARGS_FILTERS')
-ENV_DOCKER_ARGS_HIDE_ENV = EnvironmentConfig('CLEARML_AGENT_DOCKER_ARGS_HIDE_ENV')
+ENV_DOCKER_IMAGE = EnvironmentConfig("CLEARML_DOCKER_IMAGE", "TRAINS_DOCKER_IMAGE")
+ENV_WORKER_ID = EnvironmentConfig("CLEARML_WORKER_ID", "TRAINS_WORKER_ID")
+ENV_WORKER_TAGS = EnvironmentConfig("CLEARML_WORKER_TAGS")
+ENV_AGENT_SKIP_PIP_VENV_INSTALL = EnvironmentConfig("CLEARML_AGENT_SKIP_PIP_VENV_INSTALL")
+ENV_AGENT_SKIP_PYTHON_ENV_INSTALL = EnvironmentConfig("CLEARML_AGENT_SKIP_PYTHON_ENV_INSTALL", type=bool)
+ENV_DOCKER_SKIP_GPUS_FLAG = EnvironmentConfig("CLEARML_DOCKER_SKIP_GPUS_FLAG", "TRAINS_DOCKER_SKIP_GPUS_FLAG")
+ENV_AGENT_GIT_USER = EnvironmentConfig("CLEARML_AGENT_GIT_USER", "TRAINS_AGENT_GIT_USER")
+ENV_AGENT_GIT_PASS = EnvironmentConfig("CLEARML_AGENT_GIT_PASS", "TRAINS_AGENT_GIT_PASS")
+ENV_AGENT_GIT_HOST = EnvironmentConfig("CLEARML_AGENT_GIT_HOST", "TRAINS_AGENT_GIT_HOST")
+ENV_AGENT_DISABLE_SSH_MOUNT = EnvironmentConfig("CLEARML_AGENT_DISABLE_SSH_MOUNT", type=bool)
+ENV_SSH_AUTH_SOCK = EnvironmentConfig("SSH_AUTH_SOCK")
+ENV_TASK_EXECUTE_AS_USER = EnvironmentConfig("CLEARML_AGENT_EXEC_USER", "TRAINS_AGENT_EXEC_USER")
+ENV_TASK_EXTRA_PYTHON_PATH = EnvironmentConfig("CLEARML_AGENT_EXTRA_PYTHON_PATH", "TRAINS_AGENT_EXTRA_PYTHON_PATH")
+ENV_DOCKER_HOST_MOUNT = EnvironmentConfig(
+    "CLEARML_AGENT_K8S_HOST_MOUNT",
+    "CLEARML_AGENT_DOCKER_HOST_MOUNT",
+    "TRAINS_AGENT_K8S_HOST_MOUNT",
+    "TRAINS_AGENT_DOCKER_HOST_MOUNT",
+)
+ENV_VENV_CACHE_PATH = EnvironmentConfig("CLEARML_AGENT_VENV_CACHE_PATH")
+ENV_EXTRA_DOCKER_ARGS = EnvironmentConfig("CLEARML_AGENT_EXTRA_DOCKER_ARGS", type=list)
+ENV_EXTRA_DOCKER_LABELS = EnvironmentConfig("CLEARML_AGENT_EXTRA_DOCKER_LABELS", type=list)
+ENV_DEBUG_INFO = EnvironmentConfig("CLEARML_AGENT_DEBUG_INFO")
+ENV_CHILD_AGENTS_COUNT_CMD = EnvironmentConfig("CLEARML_AGENT_CHILD_AGENTS_COUNT_CMD")
+ENV_DOCKER_ARGS_FILTERS = EnvironmentConfig("CLEARML_AGENT_DOCKER_ARGS_FILTERS")
+ENV_DOCKER_ARGS_HIDE_ENV = EnvironmentConfig("CLEARML_AGENT_DOCKER_ARGS_HIDE_ENV")
+ENV_CONFIG_BC_IN_STANDALONE = EnvironmentConfig("CLEARML_AGENT_STANDALONE_CONFIG_BC", type=bool)
+""" Maintain backwards compatible configuration when launching in standalone mode """
 
-ENV_CUSTOM_BUILD_SCRIPT = EnvironmentConfig('CLEARML_AGENT_CUSTOM_BUILD_SCRIPT')
+ENV_FORCE_DOCKER_AGENT_REPO = EnvironmentConfig("FORCE_CLEARML_AGENT_REPO", "CLEARML_AGENT_DOCKER_AGENT_REPO")
+
+ENV_SERVICES_DOCKER_RESTART = EnvironmentConfig("CLEARML_AGENT_SERVICES_DOCKER_RESTART")
+"""
+    Specify a restart value for a services agent task containers.
+    Note that when a restart value is provided, task containers will not be run with the '--rm' flag and will
+     not be cleaned up automatically when completed (this will need to be done externally using the
+     'docker container prune' command to free up resources).
+    Value format for this env var is "<restart-value>;<task-selector>", where:
+    - <restart-value> can be any valid restart value for docker-run (see https://docs.docker.com/engine/reference/commandline/run/#restart)
+    - <task-selector> is optional, allowing to restrict this behaviour to specific tasks. The format is:
+        "<path-to-task-field>=<value>" where:
+        * <path-to-task-field> is a dot-separated path to a task field (e.g. "container.image")
+        * <value> is optional. If not provided, the restart policy till be applied for the task container if the
+            path provided exists. If provided, the restart policy will be applied if the value matches the value
+            obtained from the task (value parsing and comparison is based on the type of value obtained from the task) 
+    For example:
+        CLEARML_AGENT_SERVICES_DOCKER_RESTART=unless-stopped
+        CLEARML_AGENT_SERVICES_DOCKER_RESTART=unless-stopped;container.image=some-image
+"""
+
+ENV_FORCE_SYSTEM_SITE_PACKAGES = EnvironmentConfig("CLEARML_AGENT_FORCE_SYSTEM_SITE_PACKAGES", type=bool)
+""" Force system_site_packages: true when running tasks in containers (i.e. docker mode or k8s glue) """
+
+ENV_CUSTOM_BUILD_SCRIPT = EnvironmentConfig("CLEARML_AGENT_CUSTOM_BUILD_SCRIPT")
 """
     Specifies a custom environment setup script to be executed instead of installing a virtual environment.
     If provided, this script is executed following Git cloning. Script command may include environment variable and

clearml_agent/glue/k8s.py

@@ -9,7 +9,7 @@ import os
 import re
 import subprocess
 import tempfile
-from collections import defaultdict
+from collections import defaultdict, namedtuple
 from copy import deepcopy
 from pathlib import Path
 from pprint import pformat
@@ -22,7 +22,12 @@ import yaml
 from clearml_agent.backend_api.session import Request
 from clearml_agent.commands.events import Events
 from clearml_agent.commands.worker import Worker, get_task_container, set_task_container, get_next_task
-from clearml_agent.definitions import ENV_DOCKER_IMAGE, ENV_AGENT_GIT_USER, ENV_AGENT_GIT_PASS
+from clearml_agent.definitions import (
+    ENV_DOCKER_IMAGE,
+    ENV_AGENT_GIT_USER,
+    ENV_AGENT_GIT_PASS,
+    ENV_FORCE_SYSTEM_SITE_PACKAGES,
+)
 from clearml_agent.errors import APIError
 from clearml_agent.glue.definitions import ENV_START_AGENT_SCRIPT_PATH
 from clearml_agent.helper.base import safe_remove_file
@@ -32,20 +37,20 @@ from clearml_agent.helper.resource_monitor import ResourceMonitor
 from clearml_agent.interface.base import ObjectID
 
 
-
 class K8sIntegration(Worker):
     K8S_PENDING_QUEUE = "k8s_scheduler"
 
     K8S_DEFAULT_NAMESPACE = "clearml"
     AGENT_LABEL = "CLEARML=agent"
+    QUEUE_LABEL = "clearml-agent-queue"
 
     KUBECTL_APPLY_CMD = "kubectl apply --namespace={namespace} -f"
 
-    KUBECTL_DELETE_CMD = "kubectl delete pods " \
-                         "-l={agent_label} " \
-                         "--field-selector=status.phase!=Pending,status.phase!=Running " \
-                         "--namespace={namespace} " \
-                         "--output name"
+    KUBECTL_CLEANUP_DELETE_CMD = "kubectl delete pods " \
+                                 "-l={agent_label} " \
+                                 "--field-selector=status.phase!=Pending,status.phase!=Running " \
+                                 "--namespace={namespace} " \
+                                 "--output name"
 
     BASH_INSTALL_SSH_CMD = [
         "apt-get update",
@@ -136,8 +141,11 @@ class K8sIntegration(Worker):
         self.k8s_pending_queue_name = k8s_pending_queue_name or self.K8S_PENDING_QUEUE
         self.k8s_pending_queue_id = None
         self.container_bash_script = container_bash_script or self.CONTAINER_BASH_SCRIPT
-        # Always do system packages, because by we will be running inside a docker
-        self._session.config.put("agent.package_manager.system_site_packages", True)
+        force_system_packages = ENV_FORCE_SYSTEM_SITE_PACKAGES.get()
+        self._force_system_site_packages = force_system_packages if force_system_packages is not None else True
+        if self._force_system_site_packages:
+            # Use system packages, because by we will be running inside a docker
+            self._session.config.put("agent.package_manager.system_site_packages", True)
         # Add debug logging
         if debug:
             self.log.logger.disabled = False
@@ -212,15 +220,6 @@ class K8sIntegration(Worker):
         with open(os.path.expandvars(os.path.expanduser(str(path))), 'rt') as f:
             return yaml.load(f, Loader=getattr(yaml, 'FullLoader', None))
 
-    @staticmethod
-    def _get_path(d, *path, default=None):
-        try:
-            return functools.reduce(
-                lambda a, b: a[b], path, d
-            )
-        except (IndexError, KeyError):
-            return default
-
     def _get_kubectl_options(self, command, extra_labels=None, filters=None, output="json", labels=None):
         # type: (str, Iterable[str], Iterable[str], str, Iterable[str]) -> Dict
         if not labels:
@@ -255,7 +254,7 @@ class K8sIntegration(Worker):
                 sleep(self._polling_interval)
                 continue
             pods = output_config.get('items', [])
-            task_ids = set()
+            task_id_to_details = dict()
             for pod in pods:
                 pod_name = pod.get('metadata', {}).get('name', None)
                 if not pod_name:
@@ -269,7 +268,7 @@ class K8sIntegration(Worker):
                 if not namespace:
                     continue
 
-                task_ids.add(task_id)
+                task_id_to_details[task_id] = (pod_name, namespace)
 
                 msg = None
 
@@ -329,8 +328,41 @@ class K8sIntegration(Worker):
                             )
                         )
 
+            if task_id_to_details:
+                try:
+                    result = self._session.get(
+                        service='tasks',
+                        action='get_all',
+                        json={"id": list(task_id_to_details), "status": ["stopped"], "only_fields": ["id"]},
+                        method=Request.def_method,
+                        async_enable=False,
+                    )
+                    aborted_task_ids = list(filter(None, (task.get("id") for task in result["tasks"])))
+
+                    for task_id in aborted_task_ids:
+                        pod_name, namespace = task_id_to_details.get(task_id)
+                        if not pod_name:
+                            self.log.error("Failed locating aborted task {} in pending pods list".format(task_id))
+                            continue
+                        self.log.info(
+                            "K8S Glue pods monitor: task {} was aborted by its pod {} is still pending, "
+                            "deleting pod".format(task_id, pod_name)
+                        )
+
+                        kubectl_cmd = "kubectl delete pod {pod_name} --output name {namespace}".format(
+                            namespace=f"--namespace={namespace}" if namespace else "", pod_name=pod_name,
+                        ).strip()
+                        self.log.debug("Deleting aborted task pending pod: {}".format(kubectl_cmd))
+                        output = stringify_bash_output(get_bash_output(kubectl_cmd))
+                        if not output:
+                            self.log.warning("K8S Glue pods monitor: failed deleting pod {}".format(pod_name))
+                except Exception as ex:
+                    self.log.warning(
+                        'K8S Glue pods monitor: failed checking aborted tasks for hanging pods: {}'.format(ex)
+                    )
+
             # clean up any last message for a task that wasn't seen as a pod
-            last_tasks_msgs = {k: v for k, v in last_tasks_msgs.items() if k in task_ids}
+            last_tasks_msgs = {k: v for k, v in last_tasks_msgs.items() if k in task_id_to_details}
 
             sleep(self._polling_interval)
 
@@ -377,34 +409,50 @@ class K8sIntegration(Worker):
 
         return self._agent_label
 
-    def _get_used_pods(self):
-        # type: () -> Tuple[int, Set[str]]
-        # noinspection PyBroadException
+    RunningPod = namedtuple("RunningPod", "name queue namespace")
+
+    def _get_running_pods(self):
         try:
             kubectl_cmd = self.get_kubectl_command(
                 "get pods",
-                output="jsonpath=\"{range .items[*]}{.metadata.name}{' '}{.metadata.namespace}{'\\n'}{end}\""
+                output="jsonpath=\"{{range .items[*]}}{{.metadata.name}}{{' '}}{{.metadata.namespace}}{{' '}}"
+                       "{{.metadata.labels.{}}}{{'\\n'}}{{end}}\"".format(self.QUEUE_LABEL)
             )
             self.log.debug("Getting used pods: {}".format(kubectl_cmd))
             output = stringify_bash_output(get_bash_output(kubectl_cmd, raise_error=True))
 
             if not output:
                 # No such pod exist so we can use the pod_number we found
-                return 0, set([])
+                return []
 
             try:
-                items = output.splitlines()
-                current_pod_count = len(items)
-                namespaces = {item.rpartition(" ")[-1] for item in items}
-                self.log.debug(" - found {} pods in namespaces {}".format(current_pod_count, ", ".join(namespaces)))
-            except (KeyError, ValueError, TypeError, AttributeError) as ex:
-                print("Failed parsing used pods command response for cleanup: {}".format(ex))
-                return -1, set([])
+                return [
+                    self.RunningPod(
+                        name=parts[0],
+                        namespace=parts[1],
+                        queue=parts[2]
+                    )
+                    for parts in (line.split(" ") for line in output.splitlines())
+                ]
+            except Exception as ex:
+                raise Exception("Failed parsing used pods command response for cleanup: {}".format(ex))
+        except Exception as ex:
+            raise Exception('Failed obtaining used pods information: {}'.format(ex))
 
+    def _get_used_pods(self):
+        # type: () -> Tuple[int, Set[str]]
+        # noinspection PyBroadException
+        try:
+            items = self._get_running_pods()
+            if not items:
+                return 0, set([])
+            current_pod_count = len(items)
+            namespaces = {item.namespace for item in items}
+            self.log.debug(" - found {} pods in namespaces {}".format(current_pod_count, ", ".join(namespaces)))
             return current_pod_count, namespaces
         except Exception as ex:
-            print('Failed obtaining used pods information: {}'.format(ex))
-            return -2, set([])
+            self.log.debug("Failed getting used pods: {}", ex)
+            return -1, set([])
 
     def _is_same_tenant(self, task_session):
         if not task_session or task_session is self._session:
@@ -455,14 +503,14 @@ class K8sIntegration(Worker):
         git_user = ENV_AGENT_GIT_USER.get() or self._session.config.get("agent.git_user", None)
         git_pass = ENV_AGENT_GIT_PASS.get() or self._session.config.get("agent.git_pass", None)
         extra_config_values = [
-            'agent.package_manager.system_site_packages: true',
+            'agent.package_manager.system_site_packages: true' if self._force_system_site_packages else '',
             'agent.git_user: "{}"'.format(git_user) if git_user else '',
             'agent.git_pass: "{}"'.format(git_pass) if git_pass else '',
         ]
 
         # noinspection PyProtectedMember
         config_content = (
-            self.conf_file_content or Path(session._config_file).read_text() or ""
+            self.conf_file_content or (session._config_file and Path(session._config_file).read_text()) or ""
         ) + '\n{}\n'.format('\n'.join(x for x in extra_config_values if x))
 
         hocon_config_encoded = config_content.encode("ascii")
@@ -604,6 +652,7 @@ class K8sIntegration(Worker):
                     "k8s-pod-number": pod_number,
                     "k8s-pod-label": labels[0],
                     "k8s-internal-pod-count": pod_count,
+                    "k8s-agent": self._get_agent_label(),
                 }
             )
 
@@ -625,8 +674,8 @@ class K8sIntegration(Worker):
     def _get_pod_labels(self, queue, queue_name):
         return [
             self._get_agent_label(),
-            "clearml-agent-queue={}".format(self._safe_k8s_label_value(queue)),
-            "clearml-agent-queue-name={}".format(self._safe_k8s_label_value(queue_name))
+            "{}={}".format(self.QUEUE_LABEL, self._safe_k8s_label_value(queue)),
+            "{}-name={}".format(self.QUEUE_LABEL, self._safe_k8s_label_value(queue_name))
         ]
 
     def _get_docker_args(self, docker_args, flags, target=None, convert=None):
@@ -771,7 +820,9 @@ class K8sIntegration(Worker):
             if time() - self._last_pod_cleanup_per_ns[namespace] < self._min_cleanup_interval_per_ns_sec:
                 # Do not try to cleanup the same namespace too quickly
                 continue
-            kubectl_cmd = self.KUBECTL_DELETE_CMD.format(namespace=namespace, agent_label=self._get_agent_label())
+            kubectl_cmd = self.KUBECTL_CLEANUP_DELETE_CMD.format(
+                namespace=namespace, agent_label=self._get_agent_label()
+            )
             self.log.debug("Deleting old/failed pods{} for ns {}: {}".format(
                 extra_msg or "", namespace, kubectl_cmd
             ))
@@ -788,6 +839,51 @@ class K8sIntegration(Worker):
                 self.log.error("Failed deleting old/failed pods for ns %s: %s", namespace, str(ex))
             finally:
                 self._last_pod_cleanup_per_ns[namespace] = time()
+
+        # Locate tasks belonging to deleted pods that are still marked as pending or running
+        tasks_to_abort = []
+        try:
+            task_ids = list(filter(None, (
+                pod_name[len(self.pod_name_prefix):].strip()
+                for pod_names in deleted_pods.values()
+                for pod_name in pod_names
+            )))
+            if task_ids:
+                result = self._session.get(
+                    service='tasks',
+                    action='get_all',
+                    json={"id": task_ids, "status": ["in_progress", "queued"], "only_fields": ["id", "status"]},
+                    method=Request.def_method,
+                )
+                tasks_to_abort = result["tasks"]
+        except Exception as ex:
+            self.log.warning('Failed getting running tasks for deleted pods: {}'.format(ex))
+
+        for task in tasks_to_abort:
+            task_id = task.get("id")
+            status = task.get("status")
+            if not task_id or not status:
+                self.log.warning('Failed getting task information: id={}, status={}'.format(task_id, status))
+                continue
+            try:
+                if status == "queued":
+                    self._session.get(
+                        service='tasks',
+                        action='dequeue',
+                        json={"task": task_id, "force": True, "status_reason": "Pod deleted (not pending or running)",
+                              "status_message": "Pod deleted by agent {}".format(self.worker_id or "unknown")},
+                        method=Request.def_method,
+                    )
+                self._session.get(
+                    service='tasks',
+                    action='failed',
+                    json={"task": task_id, "force": True, "status_reason": "Pod deleted (not pending or running)",
+                          "status_message": "Pod deleted by agent {}".format(self.worker_id or "unknown")},
+                    method=Request.def_method,
+                )
+            except Exception as ex:
+                self.log.warning('Failed setting task {} to status "failed": {}'.format(task_id, ex))
+
         return deleted_pods
 
     def run_tasks_loop(self, queues: List[Text], worker_params, **kwargs):
@@ -941,5 +1037,6 @@ class K8sIntegration(Worker):
         value = re.sub(r'^[^A-Za-z0-9]+', '', value)  # strip leading non-alphanumeric chars
         value = re.sub(r'[^A-Za-z0-9]+$', '', value)  # strip trailing non-alphanumeric chars
         value = re.sub(r'\W+', '-', value)  # allow only word chars (this removed "." which is supported, but nvm)
+        value = re.sub(r'_+', '-', value)  # "_" is not allowed as well
         value = re.sub(r'-+', '-', value)  # don't leave messy "--" after replacing previous chars
         return value[:63]

clearml_agent/helper/package/base.py

@@ -177,7 +177,7 @@ class PackageManager(object):
                 cls._pip_version.append("==" + version)
 
     @classmethod
-    def get_pip_versions(cls, pip="pip", wrap='"'):
+    def get_pip_versions(cls, pip="pip", wrap=''):
         return [
             (wrap + pip + version + wrap)
             for version in cls._pip_version or [pip]
@@ -192,8 +192,13 @@ class PackageManager(object):
         if not self._get_cache_manager():
             return None
 
-        keys = self._generate_reqs_hash_keys(requirements, docker_cmd, python_version, cuda_version)
-        return self._get_cache_manager().copy_cached_entry(keys, destination_folder)
+        try:
+            keys = self._generate_reqs_hash_keys(requirements, docker_cmd, python_version, cuda_version)
+            return self._get_cache_manager().copy_cached_entry(keys, destination_folder)
+        except Exception as ex:
+            print("WARNING: Failed accessing venvs cache at {}: {}".format(destination_folder, ex))
+            print("WARNING: Skipping venv cache - folder not accessible!")
+            return None
 
     def add_cached_venv(
             self,
@@ -210,9 +215,15 @@ class PackageManager(object):
         """
         if not self._get_cache_manager():
             return
-        keys = self._generate_reqs_hash_keys(requirements, docker_cmd, python_version, cuda_version)
-        return self._get_cache_manager().add_entry(
-            keys=keys, source_folder=source_folder, exclude_sub_folders=exclude_sub_folders)
+
+        try:
+            keys = self._generate_reqs_hash_keys(requirements, docker_cmd, python_version, cuda_version)
+            return self._get_cache_manager().add_entry(
+                keys=keys, source_folder=source_folder, exclude_sub_folders=exclude_sub_folders)
+        except Exception as ex:
+            print("WARNING: Failed accessing venvs cache at {}: {}".format(source_folder, ex))
+            print("WARNING: Skipping venv cache - folder not accessible!")
+            return None
 
     def get_cache_folder(self):
         # type: () -> Optional[Path]
@@ -280,12 +291,19 @@ class PackageManager(object):
 
     def _get_cache_manager(self):
         if not self._cache_manager:
-            cache_folder = ENV_VENV_CACHE_PATH.get() or self.session.config.get(self._config_cache_folder, None)
-            if not cache_folder:
+            cache_folder = None
+            try:
+                cache_folder = ENV_VENV_CACHE_PATH.get() or self.session.config.get(self._config_cache_folder, None)
+                if not cache_folder:
+                    return None
+
+                max_entries = int(self.session.config.get(self._config_cache_max_entries, 10))
+                free_space_threshold = float(self.session.config.get(self._config_cache_free_space_threshold, 0))
+                self._cache_manager = FolderCache(
+                    cache_folder, max_cache_entries=max_entries, min_free_space_gb=free_space_threshold)
+            except Exception as ex:
+                print("WARNING: Failed accessing venvs cache at {}: {}".format(cache_folder, ex))
+                print("WARNING: Skipping venv cache - folder not accessible!")
                 return None
 
-            max_entries = int(self.session.config.get(self._config_cache_max_entries, 10))
-            free_space_threshold = float(self.session.config.get(self._config_cache_free_space_threshold, 0))
-            self._cache_manager = FolderCache(
-                cache_folder, max_cache_entries=max_entries, min_free_space_gb=free_space_threshold)
         return self._cache_manager

clearml_agent/helper/package/external_req.py

@@ -50,6 +50,14 @@ class ExternalRequirements(SimpleSubstitution):
                 print("No need to reinstall \'{}\' from VCS, "
                       "the exact same version is already installed".format(req.name))
                 continue
+
+            if not req.pip_new_version:
+                # noinspection PyBroadException
+                try:
+                    freeze_base = PackageManager.out_of_scope_freeze() or dict(pip=[])
+                except Exception:
+                    freeze_base = dict(pip=[])
+
             req_line = self._add_vcs_credentials(req, session)
 
             # if we have older pip version we have to make sure we replace back the package name with the
@@ -58,14 +66,14 @@ class ExternalRequirements(SimpleSubstitution):
                 PackageManager.out_of_scope_install_package(req_line, "--no-deps")
                 # noinspection PyBroadException
                 try:
-                    freeze_post = PackageManager.out_of_scope_freeze() or ''
+                    freeze_post = PackageManager.out_of_scope_freeze() or dict(pip=[])
                     package_name = list(set(freeze_post['pip']) - set(freeze_base['pip']))
                     if package_name and package_name[0] not in self.post_install_req_lookup:
                         self.post_install_req_lookup[package_name[0]] = req.req.line
                 except Exception:
                     pass
 
-            # no need to force reinstall, pip will always rebuilt if the package comes from git
+            # no need to force reinstall, pip will always rebuild if the package comes from git
             # and make sure the required packages are installed (if they are not it will install them)
             if not PackageManager.out_of_scope_install_package(req_line):
                 raise ValueError("Failed installing GIT/HTTPs package \'{}\'".format(req_line))
@@ -84,20 +92,14 @@ class ExternalRequirements(SimpleSubstitution):
                 vcs_url = req_line[4:]
                 # reverse replace
                 vcs_url = vcs_url[::-1].replace(fragment[::-1], '', 1)[::-1]
-                # remove ssh:// or git:// prefix for git detection and credentials
-                scheme = ''
-                if vcs_url and (vcs_url.startswith('ssh://') or vcs_url.startswith('git://')):
-                    scheme = 'ssh://'  # notice git:// is actually ssh://
-                    vcs_url = vcs_url[6:]
+                # notice git:// is actually ssh://
+                if vcs_url and vcs_url.startswith('git://'):
+                    vcs_url = vcs_url.replace('git://', 'ssh://', 1)
 
                 from ..repo import Git
                 vcs = Git(session=session, url=vcs_url, location=None, revision=None)
                 vcs._set_ssh_url()
-                new_req_line = 'git+{}{}{}'.format(
-                    '' if scheme and '://' in vcs.url else scheme,
-                    vcs_url if session.config.get('agent.force_git_ssh_protocol', None) else vcs.url_with_auth,
-                    fragment
-                )
+                new_req_line = 'git+{}{}'.format(vcs.url_with_auth, fragment)
                 if new_req_line != req_line:
                     furl_line = furl(new_req_line)
                     print('Replacing original pip vcs \'{}\' with \'{}\''.format(

clearml_agent/helper/package/poetry_api.py

@@ -69,6 +69,11 @@ class PoetryConfig:
                 path = path.replace(':'+sys.base_prefix, ':'+sys.real_prefix, 1)
                 kwargs['env']['PATH'] = path
 
+        if self.session and self.session.config:
+            extra_args = self.session.config.get("agent.package_manager.poetry_install_extra_args", None)
+            if extra_args:
+                args = args + tuple(extra_args)
+
         if check_if_command_exists("poetry"):
             argv = Argv("poetry", *args)
         else:

clearml_agent/helper/package/pytorch.py

@@ -310,6 +310,12 @@ class PytorchRequirement(SimpleSubstitution):
             # yes this is for linux python 2.7 support, this is the only python 2.7 we support...
             if py_ver and py_ver[0] == '2' and len(parts) > 3 and not parts[3].endswith('u'):
                 continue
+
+            # check if this an actual match
+            if not req.compare_version(v) or \
+                    (last_v and SimpleVersion.compare_versions(last_v, '>', v, ignore_sub_versions=False)):
+                continue
+
             # update the closest matched version (from above)
             if not closest_v:
                 closest_v = v
@@ -318,10 +324,6 @@ class PytorchRequirement(SimpleSubstitution):
                     SimpleVersion.compare_versions(
                         version_a=v, op='>=', version_b=req.specs[0][1], num_parts=3):
                 closest_v = v
-            # check if this an actual match
-            if not req.compare_version(v) or \
-                    (last_v and SimpleVersion.compare_versions(last_v, '>', v, ignore_sub_versions=False)):
-                continue
 
             url = '/'.join(torch_url.split('/')[:-1] + l.split('/'))
             last_v = v
@@ -475,6 +477,23 @@ class PytorchRequirement(SimpleSubstitution):
         return self.match_version(req, base).replace(" ", "\n")
 
     def replace(self, req):
+        # we first try to resolve things ourselves because pytorch pip is not always picking the correct
+        # versions from their pip repository
+
+        resolve_algorithm = str(self.config.get("agent.package_manager.pytorch_resolve", "pip")).lower()
+        if resolve_algorithm == "direct":
+            # noinspection PyBroadException
+            try:
+                new_req = self._replace(req)
+                if new_req:
+                    self._original_req.append((req, new_req))
+                return new_req
+            except Exception:
+                pass
+        elif resolve_algorithm not in ("direct", "pip"):
+            print("Warning: `agent.package_manager.pytorch_resolve={}` "
+                  "unrecognized, default to `pip`".format(resolve_algorithm))
+
         # check if package is already installed with system packages
         self.validate_python_version()
 
@@ -493,7 +512,15 @@ class PytorchRequirement(SimpleSubstitution):
 
                 if req.specs and len(req.specs) == 1 and req.specs[0][0] == "==":
                     # remove any +cu extension and let pip resolve that
-                    line = "{} {}".format(req.name, req.format_specs(max_num_parts=3))
+                    # and add .* if we have 3 parts version to deal with nvidia container 'a' version
+                    # i.e. "1.13.0" -> "1.13.0.*" so it should match preinstalled "1.13.0a0+936e930"
+                    spec_3_parts = req.format_specs(num_parts=3)
+                    spec_max3_parts = req.format_specs(max_num_parts=3)
+                    if spec_3_parts == spec_max3_parts and not spec_max3_parts.endswith("*"):
+                        line = "{} {}.*".format(req.name, spec_max3_parts)
+                    else:
+                        line = "{} {}".format(req.name, spec_max3_parts)
+
                     if req.marker:
                         line += " ; {}".format(req.marker)
                 else:
@@ -558,6 +585,19 @@ class PytorchRequirement(SimpleSubstitution):
         :param list_of_requirements: {'pip': ['a==1.0', ]}
         :return: {'pip': ['a==1.0', ]}
         """
+        def build_specific_version_req(a_line, a_name, a_new_req):
+            try:
+                r = Requirement.parse(a_line)
+                wheel_parts = r.uri.split("/")[-1].split('-')
+                version = str(wheel_parts[1].split('%')[0].split('+')[0])
+                new_r = Requirement.parse("{} == {} # {}".format(a_name, version, str(a_new_req)))
+                if new_r.line:
+                    # great it worked!
+                    return new_r.line
+            except:  # noqa
+                pass
+            return None
+
         if not self._original_req:
             return list_of_requirements
         try:
@@ -581,9 +621,18 @@ class PytorchRequirement(SimpleSubstitution):
                                     if req.local_file:
                                         lines[i] = '{}'.format(str(new_req))
                                     else:
-                                        lines[i] = '{} # {}'.format(str(req), str(new_req))
+                                        # try to rebuild requirements with specific version:
+                                        new_line = build_specific_version_req(line, req.req.name, new_req)
+                                        if new_line:
+                                            lines[i] = new_line
+                                        else:
+                                            lines[i] = '{} # {}'.format(str(req), str(new_req))
                             else:
-                                lines[i] = '{} # {}'.format(line, str(new_req))
+                                new_line = build_specific_version_req(line, req.req.name, new_req)
+                                if new_line:
+                                    lines[i] = new_line
+                                else:
+                                    lines[i] = '{} # {}'.format(line, str(new_req))
                             break
         except:
             pass

clearml_agent/helper/package/requirements.py

@@ -240,6 +240,23 @@ class SimpleVersion:
         if not version_b:
             return True
 
+        # remove trailing "*" in both
+        if "*" in version_a:
+            ignore_sub_versions = True
+            while version_a.endswith(".*"):
+                version_a = version_a[:-2]
+            if version_a == "*":
+                version_a = ""
+            num_parts = min(len(version_a.split('.')), len(version_b.split('.')), )
+
+        if "*" in version_b:
+            ignore_sub_versions = True
+            while version_b.endswith(".*"):
+                version_b = version_b[:-2]
+            if version_b == "*":
+                version_b = ""
+            num_parts = min(len(version_a.split('.')), len(version_b.split('.')), )
+
         if not num_parts:
             num_parts = max(len(version_a.split('.')), len(version_b.split('.')), )
 

clearml_agent/helper/package/translator.py

@@ -1,3 +1,4 @@
+from tempfile import mkdtemp
 from typing import Text
 
 from furl import furl
@@ -20,7 +21,16 @@ class RequirementsTranslator(object):
         config = session.config
         self.cache_dir = cache_dir or Path(config["agent.pip_download_cache.path"]).expanduser().as_posix()
         self.enabled = config["agent.pip_download_cache.enabled"]
-        Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
+        # noinspection PyBroadException
+        try:
+            Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
+        except Exception:
+            temp_cache_folder = mkdtemp(prefix='pip_download_cache.')
+            print("Failed creating pip download cache folder at `{}` reverting to `{}`".format(
+                self.cache_dir, temp_cache_folder))
+            self.cache_dir = temp_cache_folder
+            Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
+
         self.config = Config()
         self.pip = SystemPip(interpreter=interpreter, session=self._session)
         self._translate_back = {}

clearml_agent/helper/process.py

@@ -117,10 +117,11 @@ def terminate_all_child_processes(pid=None, timeout=10., include_parent=True):
 
 
 def get_docker_id(docker_cmd_contains):
+    # noinspection PyBroadException
     try:
         containers_running = get_bash_output(cmd='docker ps --no-trunc --format \"{{.ID}}: {{.Command}}\"')
         for docker_line in containers_running.split('\n'):
-            parts = docker_line.split(':')
+            parts = docker_line.split(':', 1)
             if docker_cmd_contains in parts[-1]:
                 # we found our docker, return it
                 return parts[0]

clearml_agent/helper/repo.py

@@ -320,6 +320,7 @@ class VCS(object):
                         self.url, new_url))
                     self.url = new_url
                 return
+
             # rewrite ssh URLs only if either ssh port or ssh user are forced in config
             if parsed_url.scheme == "ssh" and (
                 self.session.config.get('agent.force_git_ssh_port', None) or
@@ -334,6 +335,9 @@ class VCS(object):
                     print("Using SSH credentials - ssh url '{}' with ssh url '{}'".format(
                         self.url, new_url))
                     self.url = new_url
+                return
+            elif parsed_url.scheme == "ssh":
+                return
 
         if not self.session.config.agent.translate_ssh:
             return
@@ -343,7 +347,7 @@ class VCS(object):
                 (ENV_AGENT_GIT_PASS.get() or self.session.config.get('agent.git_pass', None)):
             # only apply to a specific domain (if requested)
             config_domain = \
-                ENV_AGENT_GIT_HOST.get() or self.session.config.get("git_host", None)
+                ENV_AGENT_GIT_HOST.get() or self.session.config.get("agent.git_host", None)
             if config_domain and config_domain != furl(self.url).host:
                 return
 

clearml_agent/helper/resource_monitor.py

@@ -139,42 +139,45 @@ class ResourceMonitor(object):
     def _daemon(self):
         seconds_since_started = 0
         reported = 0
-        while True:
-            last_report = time()
-            current_report_frequency = (
-                self._report_frequency if reported != 0 else self._first_report_sec
-            )
-            while (time() - last_report) < current_report_frequency:
-                # wait for self._sample_frequency seconds, if event set quit
-                if self._exit_event.wait(1 / self._sample_frequency):
-                    return
-                # noinspection PyBroadException
-                try:
-                    self._update_readouts()
-                except Exception as ex:
-                    log.warning("failed getting machine stats: %s", report_error(ex))
-                    self._failure()
+        try:
+            while True:
+                last_report = time()
+                current_report_frequency = (
+                    self._report_frequency if reported != 0 else self._first_report_sec
+                )
+                while (time() - last_report) < current_report_frequency:
+                    # wait for self._sample_frequency seconds, if event set quit
+                    if self._exit_event.wait(1 / self._sample_frequency):
+                        return
+                    # noinspection PyBroadException
+                    try:
+                        self._update_readouts()
+                    except Exception as ex:
+                        log.warning("failed getting machine stats: %s", report_error(ex))
+                        self._failure()
 
-            seconds_since_started += int(round(time() - last_report))
-            # check if we do not report any metric (so it means the last iteration will not be changed)
+                seconds_since_started += int(round(time() - last_report))
+                # check if we do not report any metric (so it means the last iteration will not be changed)
 
-            # if we do not have last_iteration, we just use seconds as iteration
+                # if we do not have last_iteration, we just use seconds as iteration
 
-            # start reporting only when we figured out, if this is seconds based, or iterations based
-            average_readouts = self._get_average_readouts()
-            stats = {
-                # 3 points after the dot
-                key: round(value, 3) if isinstance(value, float) else [round(v, 3) for v in value]
-                for key, value in average_readouts.items()
-            }
+                # start reporting only when we figured out, if this is seconds based, or iterations based
+                average_readouts = self._get_average_readouts()
+                stats = {
+                    # 3 points after the dot
+                    key: round(value, 3) if isinstance(value, float) else [round(v, 3) for v in value]
+                    for key, value in average_readouts.items()
+                }
 
-            # send actual report
-            if self.send_report(stats):
-                # clear readouts if this is update was sent
-                self._clear_readouts()
+                # send actual report
+                if self.send_report(stats):
+                    # clear readouts if this is update was sent
+                    self._clear_readouts()
 
-            # count reported iterations
-            reported += 1
+                # count reported iterations
+                reported += 1
+        except Exception as ex:
+            log.exception("Error reporting monitoring info: %s", str(ex))
 
     def _update_readouts(self):
         readouts = self._machine_stats()

clearml_agent/version.py

@@ -1 +1 @@
-__version__ = '1.5.1rc0'
+__version__ = '1.5.3rc3'

docker/services/entrypoint.sh

@@ -1,16 +1,36 @@
-#!/bin/sh
+#!/bin/bash +x
 
-CLEARML_FILES_HOST=${CLEARML_FILES_HOST:-$TRAINS_FILES_HOST}
+if [ -n "$SHUTDOWN_IF_NO_ACCESS_KEY" ] && [ -z "$CLEARML_API_ACCESS_KEY" ] && [ -z "$TRAINS_API_ACCESS_KEY" ]; then
+  echo "CLEARML_API_ACCESS_KEY was not provided, service will not be started"
+  exit 0
+fi
+
+export CLEARML_FILES_HOST=${CLEARML_FILES_HOST:-$TRAINS_FILES_HOST}
 
 if [ -z "$CLEARML_FILES_HOST" ]; then
     CLEARML_HOST_IP=${CLEARML_HOST_IP:-${TRAINS_HOST_IP:-$(curl -s https://ifconfig.me/ip)}}
 fi
 
-CLEARML_FILES_HOST=${CLEARML_FILES_HOST:-${TRAINS_FILES_HOST:-"http://$CLEARML_HOST_IP:8081"}}
-CLEARML_WEB_HOST=${CLEARML_WEB_HOST:-${TRAINS_WEB_HOST:-"http://$CLEARML_HOST_IP:8080"}}
-CLEARML_API_HOST=${CLEARML_API_HOST:-${TRAINS_API_HOST:-"http://$CLEARML_HOST_IP:8008"}}
+export CLEARML_FILES_HOST=${CLEARML_FILES_HOST:-${TRAINS_FILES_HOST:-"http://$CLEARML_HOST_IP:8081"}}
+export CLEARML_WEB_HOST=${CLEARML_WEB_HOST:-${TRAINS_WEB_HOST:-"http://$CLEARML_HOST_IP:8080"}}
+export CLEARML_API_HOST=${CLEARML_API_HOST:-${TRAINS_API_HOST:-"http://$CLEARML_HOST_IP:8008"}}
 
 echo $CLEARML_FILES_HOST $CLEARML_WEB_HOST $CLEARML_API_HOST 1>&2
 
-python3 -m pip install -q -U "clearml-agent${CLEARML_AGENT_UPDATE_VERSION:-$TRAINS_AGENT_UPDATE_VERSION}"
-clearml-agent daemon --services-mode --queue services --create-queue --docker "${CLEARML_AGENT_DEFAULT_BASE_DOCKER:-$TRAINS_AGENT_DEFAULT_BASE_DOCKER}" --cpu-only ${CLEARML_AGENT_EXTRA_ARGS:-$TRAINS_AGENT_EXTRA_ARGS}
+if [[ "$CLEARML_AGENT_UPDATE_VERSION" =~ ^[0-9]{1,3}\.[0-9]{1,3}(\.[0-9]{1,3}([a-zA-Z]{1,3}[0-9]{1,3})?)?$ ]]
+then
+    CLEARML_AGENT_UPDATE_VERSION="==$CLEARML_AGENT_UPDATE_VERSION"
+fi
+
+DAEMON_OPTIONS=${CLEARML_AGENT_DAEMON_OPTIONS:---services-mode --create-queue}
+QUEUES=${CLEARML_AGENT_QUEUES:-services}
+
+if [ -z "$CLEARML_AGENT_NO_UPDATE" ]; then
+  if [ -n "$CLEARML_AGENT_UPDATE_REPO" ]; then
+    python3 -m pip install -q -U $CLEARML_AGENT_UPDATE_REPO
+  else
+    python3 -m pip install -q -U "clearml-agent${CLEARML_AGENT_UPDATE_VERSION:-$TRAINS_AGENT_UPDATE_VERSION}"
+  fi
+fi
+
+clearml-agent daemon $DAEMON_OPTIONS --queue $QUEUES --docker "${CLEARML_AGENT_DEFAULT_BASE_DOCKER:-$TRAINS_AGENT_DEFAULT_BASE_DOCKER}" --cpu-only ${CLEARML_AGENT_EXTRA_ARGS:-$TRAINS_AGENT_EXTRA_ARGS}

docs/clearml.conf

@@ -82,6 +82,7 @@ agent {
         # pip_version: ["<20.2 ; python_version < '3.10'",  "<22.3 ; python_version >= '3.10'"]
         # specify poetry version to use (examples "<2", "==1.1.1", "", empty string will install the latest version)
         # poetry_version: "<2",
+        # poetry_install_extra_args: ["-v"]
 
         # virtual environment inheres packages from system
         system_site_packages: false,
@@ -92,25 +93,39 @@ agent {
         # extra_index_url: ["https://allegroai.jfrog.io/clearml/api/pypi/public/simple"]
         extra_index_url: []
 
+        # control the pytorch wheel resolving algorithm, options are: "pip", "direct"
+        # "pip" (default): would automatically detect the cuda version, and supply pip with the correct
+        # extra-index-url, based on pytorch.org tables
+        # "direct": would resolve a direct link to the pytorch wheel by parsing the pytorch.org pip repository
+        # and matching the automatically detected cuda version with the required pytorch wheel.
+        # if the exact cuda version is not found for the required pytorch wheel, it will try
+        # a lower cuda version until a match is found
+        #
+        # pytorch_resolve: "pip"
+
         # additional conda channels to use when installing with conda package manager
         conda_channels: ["pytorch", "conda-forge", "defaults", ]
         # conda_full_env_update: false
         # conda_env_as_base_docker: false
 
         # set the priority packages to be installed before the rest of the required packages
+        # Note: this only controls the installation order of existing requirement packages (and does not add additional packages)
         # priority_packages: ["cython", "numpy", "setuptools", ]
 
         # set the optional priority packages to be installed before the rest of the required packages,
         # In case a package installation fails, the package will be ignored,
         # and the virtual environment process will continue
+        # Note: this only controls the installation order of existing requirement packages (and does not add additional packages)
         # priority_optional_packages: ["pygobject", ]
 
         # set the post packages to be installed after all the rest of the required packages
+        # Note: this only controls the installation order of existing requirement packages (and does not add additional packages)
         # post_packages: ["horovod", ]
 
         # set the optional post packages to be installed after all the rest of the required packages,
         # In case a package installation fails, the package will be ignored,
         # and the virtual environment process will continue
+        # Note: this only controls the installation order of existing requirement packages (and does not add additional packages)
         # post_optional_packages: []
 
         # set to True to support torch nightly build installation,
@@ -227,8 +242,6 @@ agent {
         #                 # no repository matching required
         #                 repository: ""
         #             }
-        #             # no container image matching required (allow to replace one requested container with another)
-        #             container: ""
         #             # no repository matching required
         #             project: ""
         #         }
@@ -468,6 +481,7 @@ sdk {
 #  target_format: format used to encode contents before writing into the target file. Supported values are json,
 #                 yaml, yml and bytes (in which case the file will be written in binary mode). Default is text mode.
 #  overwrite: overwrite the target file in case it exists. Default is true.
+#  mode: set the file mode after writing. use an integer value or a string (e.g. 600 / 777 etc.)
 #
 # Example:
 #   files {

requirements.txt

@@ -9,7 +9,7 @@ python-dateutil>=2.4.2,<2.9.0
 pyjwt>=2.4.0,<2.7.0
 PyYAML>=3.12,<6.1
 requests>=2.20.0,<2.29.0
-six>=1.13.0,<1.16.0
+six>=1.13.0,<1.17.0
 typing>=3.6.4,<3.8.0 ; python_version < '3.5'
 urllib3>=1.21.1,<1.27.0
 virtualenv>=16,<21