diff --git a/cli/polyaxon/_env_vars/getters/run.py b/cli/polyaxon/_env_vars/getters/run.py index 2de9a8cb..ca689385 100644 --- a/cli/polyaxon/_env_vars/getters/run.py +++ b/cli/polyaxon/_env_vars/getters/run.py @@ -13,6 +13,7 @@ from polyaxon._env_vars.keys import ( ENV_KEYS_COLLECT_ARTIFACTS, ENV_KEYS_COLLECT_RESOURCES, + ENV_KEYS_COLLECT_RESOURCES_DETAILED, ENV_KEYS_RUN_INSTANCE, ) from polyaxon._utils.fqn_utils import get_owner_team_space, split_owner_team_space @@ -135,6 +136,19 @@ def get_collect_resources(arg: Optional[bool] = None, default: Optional[bool] = ) +def get_collect_resources_detailed( + arg: Optional[bool] = None, default: Optional[bool] = None +): + """If set, Polyaxon will collect the detailed process and cgroup resources""" + return ( + arg + if arg is not None + else to_bool( + os.getenv(ENV_KEYS_COLLECT_RESOURCES_DETAILED, default), handle_none=True + ) + ) + + def get_log_level(): """If set on the polyaxonfile it will return the log level.""" from polyaxon import settings diff --git a/cli/polyaxon/_env_vars/keys.py b/cli/polyaxon/_env_vars/keys.py index 66cf4ff0..25f5842e 100644 --- a/cli/polyaxon/_env_vars/keys.py +++ b/cli/polyaxon/_env_vars/keys.py @@ -90,6 +90,7 @@ # Connections ENV_KEYS_COLLECT_ARTIFACTS = "POLYAXON_COLLECT_ARTIFACTS" ENV_KEYS_COLLECT_RESOURCES = "POLYAXON_COLLECT_RESOURCES" +ENV_KEYS_COLLECT_RESOURCES_DETAILED = "POLYAXON_COLLECT_RESOURCES_DETAILED" ENV_KEYS_ARTIFACTS_STORE_NAME = "POLYAXON_ARTIFACTS_STORE_NAME" ENV_KEYS_GIT_CREDENTIALS = "POLYAXON_GIT_CREDENTIALS" ENV_KEYS_GIT_CREDENTIALS_STORE = "POLYAXON_GIT_CREDENTIALS_STORE" diff --git a/cli/polyaxon/_flow/plugins/__init__.py b/cli/polyaxon/_flow/plugins/__init__.py index b7847e19..29c21e04 100644 --- a/cli/polyaxon/_flow/plugins/__init__.py +++ b/cli/polyaxon/_flow/plugins/__init__.py @@ -27,6 +27,7 @@ class V1Plugins(BaseSchemaModel): collect_artifacts: bool, optional, default: True collect_logs: bool, optional, default: True collect_resources: bool, optional, default: True + collect_resources_detailed: bool, optional, default: False sync_statuses: bool, optional, default: True auto_resume: bool, optional, default: True external_host: bool, optional, default: False @@ -47,6 +48,7 @@ class V1Plugins(BaseSchemaModel): >>> collectArtifacts: >>> collectLogs: >>> collectResources: + >>> collectResourcesDetailed: >>> syncStatuses: >>> autoResume: >>> externalHost: @@ -69,6 +71,7 @@ class V1Plugins(BaseSchemaModel): >>> collect_artifacts=False, >>> collect_logs=False, >>> collect_resources=False, + >>> collect_resources_detailed=False, >>> sync_statuses=True, >>> auto_resume=False, >>> external_host=False, @@ -275,6 +278,28 @@ class V1Plugins(BaseSchemaModel): >>> collectResources: false ``` + ### collectResourcesDetailed + +
This plugin is disabled by default.
+ + On Linux, the python client also logs a small set of container-aware resources + next to the Mem/CPU/GPU metrics: the cgroup memory percent and OOM kills, + the cgroup cpu throttling, the `/dev/shm` usage, and the number of child + processes stuck in disk sleep. This covers the usual reasons a run is + OOM-killed, throttled, or stalled inside a pod, which the host level + Mem/CPU metrics do not show. + + Enable this plugin to also log the detailed breakdown: cgroup memory current/max/shmem, + cpu quota and throttled periods, `/dev/shm` total/used, the process rss split, + threads and file descriptors, and the children's rss. + + To enable this plugin: + + ```yaml + >>> plugins: + >>> collectResourcesDetailed: true + ``` + ### autoResume
This plugin is enabled by default.
@@ -344,6 +369,9 @@ class V1Plugins(BaseSchemaModel): collect_resources: Optional[BoolOrRef] = Field( alias="collectResources", default=None ) + collect_resources_detailed: Optional[BoolOrRef] = Field( + alias="collectResourcesDetailed", default=None + ) sync_statuses: Optional[BoolOrRef] = Field(alias="syncStatuses", default=None) auto_resume: Optional[BoolOrRef] = Field(alias="autoResume", default=None) log_level: Optional[StrictStr] = Field(alias="logLevel", default=None) @@ -366,6 +394,7 @@ def get_or_create( config.set_collect_artifacts() config.set_collect_logs() config.set_collect_resources() + config.set_collect_resources_detailed() config.set_sync_statuses() config.set_auto_resume() config.set_external_host() @@ -422,6 +451,12 @@ def set_collect_resources(self, default: bool = True): elif self.collect_resources is None: self.collect_resources = default + def set_collect_resources_detailed(self, default: bool = False): + if self.no_api(): + self.collect_resources_detailed = False + elif self.collect_resources_detailed is None: + self.collect_resources_detailed = default + def set_sync_statuses(self, default: bool = True): if self.no_api(): self.sync_statuses = False diff --git a/cli/polyaxon/_runner/converter/converter.py b/cli/polyaxon/_runner/converter/converter.py index 664a99b6..81c4f4a6 100644 --- a/cli/polyaxon/_runner/converter/converter.py +++ b/cli/polyaxon/_runner/converter/converter.py @@ -21,6 +21,7 @@ ENV_KEYS_ARTIFACTS_STORE_NAME, ENV_KEYS_COLLECT_ARTIFACTS, ENV_KEYS_COLLECT_RESOURCES, + ENV_KEYS_COLLECT_RESOURCES_DETAILED, ENV_KEYS_LOG_LEVEL, ENV_KEYS_NO_API, ENV_KEYS_RUN_INSTANCE, @@ -435,6 +436,11 @@ def _get_main_env_vars( if plugins and plugins.collect_resources: env.append(self._get_env_var(name=ENV_KEYS_COLLECT_RESOURCES, value=True)) + if plugins and plugins.collect_resources_detailed: + env.append( + self._get_env_var(name=ENV_KEYS_COLLECT_RESOURCES_DETAILED, value=True) + ) + if artifacts_store_name: env.append( self._get_env_var( diff --git a/cli/tests/test_docker/test_converters/test_base/test_main/test_main_env_vars.py b/cli/tests/test_docker/test_converters/test_base/test_main/test_main_env_vars.py index 9648787a..9cc01521 100644 --- a/cli/tests/test_docker/test_converters/test_base/test_main/test_main_env_vars.py +++ b/cli/tests/test_docker/test_converters/test_base/test_main/test_main_env_vars.py @@ -11,6 +11,7 @@ ENV_KEYS_ARTIFACTS_STORE_NAME, ENV_KEYS_COLLECT_ARTIFACTS, ENV_KEYS_COLLECT_RESOURCES, + ENV_KEYS_COLLECT_RESOURCES_DETAILED, ) from polyaxon._flow.plugins import V1Plugins from polyaxon._services.values import PolyaxonServices @@ -127,6 +128,66 @@ def test_get_env_vars_with_kv_env_vars(self): + self.converter._get_additional_env_vars() ) + def test_get_env_vars_with_collect_resources_detailed(self): + base_env = self.converter._get_service_env_vars( + service_header=PolyaxonServices.RUNNER, + external_host=False, + log_level=None, + ) + # Disabled by default + plugins = V1Plugins.get_or_create( + V1Plugins( + collect_logs=False, collect_artifacts=False, collect_resources=True + ) + ) + assert plugins.collect_resources_detailed is False + assert ( + self.converter._get_main_env_vars( + plugins=plugins, + kv_env_vars=None, + artifacts_store_name=None, + connections=None, + secrets=None, + config_maps=None, + ) + == base_env + + [ + self.converter._get_env_var( + name=ENV_KEYS_COLLECT_RESOURCES, value=True + ), + ] + + self.converter._get_additional_env_vars() + ) + + plugins = V1Plugins.get_or_create( + V1Plugins( + collect_logs=False, + collect_artifacts=False, + collect_resources=True, + collect_resources_detailed=True, + ) + ) + assert ( + self.converter._get_main_env_vars( + plugins=plugins, + kv_env_vars=None, + artifacts_store_name=None, + connections=None, + secrets=None, + config_maps=None, + ) + == base_env + + [ + self.converter._get_env_var( + name=ENV_KEYS_COLLECT_RESOURCES, value=True + ), + self.converter._get_env_var( + name=ENV_KEYS_COLLECT_RESOURCES_DETAILED, value=True + ), + ] + + self.converter._get_additional_env_vars() + ) + def test_get_env_vars_with_artifacts_store(self): base_env = self.converter._get_service_env_vars( service_header=PolyaxonServices.RUNNER, diff --git a/cli/tests/test_k8s/test_converters/test_base/test_main/test_main_env_vars.py b/cli/tests/test_k8s/test_converters/test_base/test_main/test_main_env_vars.py index 1ac5e171..467c642d 100644 --- a/cli/tests/test_k8s/test_converters/test_base/test_main/test_main_env_vars.py +++ b/cli/tests/test_k8s/test_converters/test_base/test_main/test_main_env_vars.py @@ -11,6 +11,7 @@ ENV_KEYS_ARTIFACTS_STORE_NAME, ENV_KEYS_COLLECT_ARTIFACTS, ENV_KEYS_COLLECT_RESOURCES, + ENV_KEYS_COLLECT_RESOURCES_DETAILED, ) from polyaxon._flow.plugins import V1Plugins from polyaxon._services.values import PolyaxonServices @@ -120,6 +121,52 @@ def test_get_env_vars_with_kv_env_vars(self): [["key1", "val1"], ["key2", "val2"]] ) + def test_get_env_vars_with_collect_resources_detailed(self): + base_env = self.converter._get_service_env_vars( + service_header=PolyaxonServices.RUNNER, + external_host=False, + log_level=None, + ) + # Disabled by default + plugins = V1Plugins.get_or_create( + V1Plugins( + collect_logs=False, collect_artifacts=False, collect_resources=True + ) + ) + assert plugins.collect_resources_detailed is False + assert self.converter._get_main_env_vars( + plugins=plugins, + kv_env_vars=None, + artifacts_store_name=None, + connections=None, + secrets=None, + config_maps=None, + ) == base_env + [ + self.converter._get_env_var(name=ENV_KEYS_COLLECT_RESOURCES, value=True), + ] + + plugins = V1Plugins.get_or_create( + V1Plugins( + collect_logs=False, + collect_artifacts=False, + collect_resources=True, + collect_resources_detailed=True, + ) + ) + assert self.converter._get_main_env_vars( + plugins=plugins, + kv_env_vars=None, + artifacts_store_name=None, + connections=None, + secrets=None, + config_maps=None, + ) == base_env + [ + self.converter._get_env_var(name=ENV_KEYS_COLLECT_RESOURCES, value=True), + self.converter._get_env_var( + name=ENV_KEYS_COLLECT_RESOURCES_DETAILED, value=True + ), + ] + def test_get_env_vars_with_artifacts_store(self): base_env = self.converter._get_service_env_vars( service_header=PolyaxonServices.RUNNER,