Refactor enums and env vars

traceml/tests/test_events/test_schemas.py

@@ -5,21 +5,19 @@
 from clipped.utils.tz import now
 
 from polyaxon.utils.test_utils import BaseTestCase
-from traceml.events import (
+from traceml.events.schemas import (
+    LoggedEventListSpec,
     V1Event,
     V1EventArtifact,
     V1EventAudio,
     V1EventChart,
+    V1EventConfusionMatrix,
+    V1EventCurve,
+    V1EventDataframe,
     V1EventHistogram,
     V1EventImage,
     V1EventModel,
     V1Events,
-)
-from traceml.events.schemas import (
-    LoggedEventListSpec,
-    V1EventConfusionMatrix,
-    V1EventCurve,
-    V1EventDataframe,
     V1EventVideo,
 )
 

traceml/tests/test_tracking/test_run_tracking.py

@@ -19,10 +19,10 @@
 from polyaxon.contexts import paths as ctx_paths
 from polyaxon.env_vars import getters
 from polyaxon.env_vars.keys import (
-    EV_KEYS_COLLECT_ARTIFACTS,
-    EV_KEYS_COLLECT_RESOURCES,
-    EV_KEYS_LOG_LEVEL,
-    EV_KEYS_RUN_INSTANCE,
+    ENV_KEYS_COLLECT_ARTIFACTS,
+    ENV_KEYS_COLLECT_RESOURCES,
+    ENV_KEYS_LOG_LEVEL,
+    ENV_KEYS_RUN_INSTANCE,
 )
 from polyaxon.exceptions import PolyaxonClientException
 from polyaxon.lifecycle import V1ProjectFeature
@@ -42,7 +42,7 @@ def setUp(self):
 
     def test_get_collect_artifacts_return_false_out_cluster(self):
         settings.CLIENT_CONFIG.is_managed = False
-        os.environ[EV_KEYS_COLLECT_ARTIFACTS] = "false"
+        os.environ[ENV_KEYS_COLLECT_ARTIFACTS] = "false"
         assert getters.get_collect_artifacts() is False
 
     def test_empty_collect_artifacts_path(self):
@@ -52,12 +52,12 @@ def test_empty_collect_artifacts_path(self):
     def test_valid_artifacts_path(self):
         settings.CLIENT_CONFIG.is_managed = True
         self.check_valid_value(
-            EV_KEYS_COLLECT_ARTIFACTS, getters.get_collect_artifacts, "true", True
+            ENV_KEYS_COLLECT_ARTIFACTS, getters.get_collect_artifacts, "true", True
         )
 
     def test_get_collect_resources_return_false_out_cluster(self):
         settings.CLIENT_CONFIG.is_managed = False
-        os.environ[EV_KEYS_COLLECT_RESOURCES] = "false"
+        os.environ[ENV_KEYS_COLLECT_RESOURCES] = "false"
         assert getters.get_collect_resources() is False
 
     def test_empty_collect_resources_path(self):
@@ -67,16 +67,16 @@ def test_empty_collect_resources_path(self):
     def test_valid_resources_path(self):
         settings.CLIENT_CONFIG.is_managed = True
         self.check_valid_value(
-            EV_KEYS_COLLECT_RESOURCES, getters.get_collect_resources, "true", True
+            ENV_KEYS_COLLECT_RESOURCES, getters.get_collect_resources, "true", True
         )
 
     def test_get_log_level_out_cluster(self):
         settings.CLIENT_CONFIG.is_managed = False
-        self.check_empty_value(EV_KEYS_LOG_LEVEL, getters.get_log_level)
+        self.check_empty_value(ENV_KEYS_LOG_LEVEL, getters.get_log_level)
 
     def test_empty_log_level(self):
         settings.CLIENT_CONFIG.is_managed = True
-        self.check_empty_value(EV_KEYS_LOG_LEVEL, getters.get_log_level)
+        self.check_empty_value(ENV_KEYS_LOG_LEVEL, getters.get_log_level)
 
     def test_run_info_checks_is_managed(self):
         settings.CLIENT_CONFIG.is_managed = False
@@ -85,19 +85,19 @@ def test_run_info_checks_is_managed(self):
 
     def test_empty_run_info(self):
         self.check_raise_for_invalid_value(
-            EV_KEYS_RUN_INSTANCE, getters.get_run_info, None, PolyaxonClientException
+            ENV_KEYS_RUN_INSTANCE, getters.get_run_info, None, PolyaxonClientException
         )
 
     def test_non_valid_run_info(self):
         self.check_raise_for_invalid_value(
-            EV_KEYS_RUN_INSTANCE,
+            ENV_KEYS_RUN_INSTANCE,
             getters.get_run_info,
             "something random",
             PolyaxonClientException,
         )
 
         self.check_raise_for_invalid_value(
-            EV_KEYS_RUN_INSTANCE,
+            ENV_KEYS_RUN_INSTANCE,
             getters.get_run_info,
             "foo.bar",
             PolyaxonClientException,
@@ -107,7 +107,7 @@ def test_dict_run_info(self):
         uid = uuid.uuid4().hex
         run_info = "user.project_bar.runs.{}".format(uid)
         self.check_valid_value(
-            EV_KEYS_RUN_INSTANCE,
+            ENV_KEYS_RUN_INSTANCE,
             getters.get_run_info,
             run_info,
             ("user", "project_bar", uid),
@@ -167,15 +167,15 @@ def test_run_init(self, expanduser):
 
         # FQN non CE
         settings.CLI_CONFIG.installation = {"dist": dist.EE}
-        os.environ[EV_KEYS_RUN_INSTANCE] = "user.project_bar.runs.{}".format(uid)
+        os.environ[ENV_KEYS_RUN_INSTANCE] = "user.project_bar.runs.{}".format(uid)
         run = Run()
         assert run.owner == "user"
         assert run.project == "project_bar"
         assert run.run_uuid == uid
 
         # FQN CE
         settings.CLI_CONFIG.installation = {"dist": dist.CE}
-        os.environ[EV_KEYS_RUN_INSTANCE] = "user.project_bar.runs.{}".format(uid)
+        os.environ[ENV_KEYS_RUN_INSTANCE] = "user.project_bar.runs.{}".format(uid)
         run = Run()
         assert run.owner == "user"
         assert run.project == "project_bar"
@@ -240,8 +240,8 @@ def test_event_logger_from_non_managed_run(self):
         assert mock_call.call_count == 1
 
         # Set collect flag
-        os.environ[EV_KEYS_COLLECT_ARTIFACTS] = "true"
-        os.environ[EV_KEYS_COLLECT_RESOURCES] = "true"
+        os.environ[ENV_KEYS_COLLECT_ARTIFACTS] = "true"
+        os.environ[ENV_KEYS_COLLECT_RESOURCES] = "true"
         settings.CLIENT_CONFIG.is_managed = True
         with patch("traceml.tracking.run.EventFileWriter") as event_call:
             with patch("traceml.tracking.run.ResourceFileWriter") as resource_call:
@@ -269,9 +269,9 @@ def test_event_logger_from_a_managed_run(self):
         # Set managed flag
         settings.CLIENT_CONFIG.is_managed = True
         settings.CLIENT_CONFIG.is_offline = False
-        os.environ[EV_KEYS_RUN_INSTANCE] = "user.project_bar.runs.{}".format(uid)
-        os.environ[EV_KEYS_COLLECT_ARTIFACTS] = "false"
-        os.environ[EV_KEYS_COLLECT_RESOURCES] = "false"
+        os.environ[ENV_KEYS_RUN_INSTANCE] = "user.project_bar.runs.{}".format(uid)
+        os.environ[ENV_KEYS_COLLECT_ARTIFACTS] = "false"
+        os.environ[ENV_KEYS_COLLECT_RESOURCES] = "false"
 
         with patch("traceml.tracking.run.Run.refresh_data") as refresh_call:
             run = Run()
@@ -287,8 +287,8 @@ def test_event_logger_from_a_managed_run(self):
         assert run._event_logger is None
 
         # Set collect flag
-        os.environ[EV_KEYS_COLLECT_ARTIFACTS] = "true"
-        os.environ[EV_KEYS_COLLECT_RESOURCES] = "true"
+        os.environ[ENV_KEYS_COLLECT_ARTIFACTS] = "true"
+        os.environ[ENV_KEYS_COLLECT_RESOURCES] = "true"
 
         # Add run id
         uid2 = uuid.uuid4().hex
@@ -307,7 +307,7 @@ def test_event_logger_from_a_managed_run(self):
         assert exit_call.call_count == 1
 
         # Set run info
-        os.environ[EV_KEYS_RUN_INSTANCE] = "user.project_bar.runs.{}".format(uid2)
+        os.environ[ENV_KEYS_RUN_INSTANCE] = "user.project_bar.runs.{}".format(uid2)
         # Add run id
         with patch("traceml.tracking.run.Run.set_run_event_logger") as event_call:
             with patch(
@@ -324,8 +324,8 @@ def test_event_logger_from_an_offline_run(self):
         # Set managed flag
         settings.CLIENT_CONFIG.is_managed = False
         settings.CLIENT_CONFIG.is_offline = True
-        os.environ[EV_KEYS_COLLECT_ARTIFACTS] = "false"
-        os.environ[EV_KEYS_COLLECT_RESOURCES] = "false"
+        os.environ[ENV_KEYS_COLLECT_ARTIFACTS] = "false"
+        os.environ[ENV_KEYS_COLLECT_RESOURCES] = "false"
 
         with patch("traceml.tracking.run.Run._set_exit_handler") as exit_mock:
             run = Run(project="test.test", run_uuid=uid)
@@ -341,8 +341,8 @@ def test_event_logger_from_an_offline_run(self):
         assert run._event_logger is None
 
         # Set collect flag
-        os.environ[EV_KEYS_COLLECT_ARTIFACTS] = "true"
-        os.environ[EV_KEYS_COLLECT_RESOURCES] = "true"
+        os.environ[ENV_KEYS_COLLECT_ARTIFACTS] = "true"
+        os.environ[ENV_KEYS_COLLECT_RESOURCES] = "true"
 
         # Add run id
         with patch("traceml.tracking.run.Run.set_run_event_logger") as event_call:
@@ -363,8 +363,8 @@ def setUp(self):
         self.run_path = tempfile.mkdtemp()
         self.run_outputs_path = tempfile.mkdtemp()
         settings.CLIENT_CONFIG.is_managed = False
-        os.environ[EV_KEYS_COLLECT_ARTIFACTS] = "false"
-        os.environ[EV_KEYS_COLLECT_RESOURCES] = "false"
+        os.environ[ENV_KEYS_COLLECT_ARTIFACTS] = "false"
+        os.environ[ENV_KEYS_COLLECT_RESOURCES] = "false"
         with patch("traceml.tracking.run.Run._set_exit_handler") as exit_mock:
             self.run = Run(
                 project="owner.project",

traceml/tests/test_tracking/test_summaries.py

@@ -5,7 +5,10 @@
 from mock import patch
 
 from polyaxon import settings
-from polyaxon.env_vars.keys import EV_KEYS_COLLECT_ARTIFACTS, EV_KEYS_COLLECT_RESOURCES
+from polyaxon.env_vars.keys import (
+    ENV_KEYS_COLLECT_ARTIFACTS,
+    ENV_KEYS_COLLECT_RESOURCES,
+)
 from polyaxon.utils.test_utils import BaseTestCase
 from traceml.artifacts import V1RunArtifact
 from traceml.events import V1Events
@@ -18,8 +21,8 @@ def setUp(self):
         super().setUp()
         settings.CLIENT_CONFIG.is_managed = False
         settings.CLIENT_CONFIG.is_offline = True
-        os.environ[EV_KEYS_COLLECT_ARTIFACTS] = "false"
-        os.environ[EV_KEYS_COLLECT_RESOURCES] = "false"
+        os.environ[ENV_KEYS_COLLECT_ARTIFACTS] = "false"
+        os.environ[ENV_KEYS_COLLECT_RESOURCES] = "false"
         uid = uuid.uuid4().hex
         with patch("traceml.tracking.run.Run._set_exit_handler") as exit_mock:
             self.run = Run(project="test.test", run_uuid=uid)

traceml/traceml/artifacts/__init__.py

@@ -1,2 +1,2 @@
-from traceml.artifacts.kinds import V1ArtifactKind
+from traceml.artifacts.enums import V1ArtifactKind
 from traceml.artifacts.schemas import V1RunArtifact, V1RunArtifacts

traceml/traceml/artifacts/schemas.py

@@ -6,7 +6,7 @@
 from clipped.types.uuids import UUIDStr
 
 from polyaxon.schemas.base import BaseSchemaModel
-from traceml.artifacts.kinds import V1ArtifactKind
+from traceml.artifacts.enums import V1ArtifactKind
 
 
 class V1RunArtifact(BaseSchemaModel):

traceml/traceml/events/schemas.py

@@ -14,7 +14,7 @@
 from clipped.utils.tz import now
 
 from polyaxon.schemas.base import BaseSchemaModel
-from traceml.artifacts.kinds import V1ArtifactKind
+from traceml.artifacts.enums import V1ArtifactKind
 
 
 class SearchView(str, PEnum):

traceml/traceml/integrations/pytorch_lightning.py

@@ -7,7 +7,7 @@
 import packaging
 
 from polyaxon.client import RunClient
-from polyaxon.env_vars.keys import EV_KEYS_RUN_INSTANCE
+from polyaxon.env_vars.keys import ENV_KEYS_RUN_INSTANCE
 from traceml import tracking
 from traceml.exceptions import TracemlException
 
@@ -145,7 +145,7 @@ def finalize(self, status: str):
 
     def _set_run_instance_from_env_vars(self, force: bool = False):
         """Tries to extract run info from canonical env vars"""
-        run_instance = os.getenv(EV_KEYS_RUN_INSTANCE)
+        run_instance = os.getenv(ENV_KEYS_RUN_INSTANCE)
         if not run_instance:
             return
 

traceml/traceml/logging/handler.py

@@ -6,16 +6,16 @@
 from clipped.utils.env import get_user
 
 from polyaxon import settings
-from polyaxon.env_vars.keys import EV_KEYS_K8S_NODE_NAME, EV_KEYS_K8S_POD_ID
-from traceml.logging import V1Log
+from polyaxon.env_vars.keys import ENV_KEYS_K8S_NODE_NAME, ENV_KEYS_K8S_POD_ID
+from traceml.logging.schemas import V1Log
 
 
 class LogStreamHandler(logging.Handler):
     def __init__(self, add_logs, **kwargs):
         self._add_logs = add_logs
         self._container = socket.gethostname()
-        self._node = os.environ.get(EV_KEYS_K8S_NODE_NAME, "local")
-        self._pod = os.environ.get(EV_KEYS_K8S_POD_ID, get_user())
+        self._node = os.environ.get(ENV_KEYS_K8S_NODE_NAME, "local")
+        self._pod = os.environ.get(ENV_KEYS_K8S_POD_ID, get_user())
         log_level = settings.CLIENT_CONFIG.log_level
         if log_level and isinstance(log_level, str):
             log_level = log_level.upper()

traceml/traceml/tracking/run.py

@@ -29,7 +29,7 @@
     get_collect_resources,
     get_log_level,
 )
-from polyaxon.env_vars.keys import EV_KEYS_HAS_PROCESS_SIDECAR
+from polyaxon.env_vars.keys import ENV_KEYS_HAS_PROCESS_SIDECAR
 from polyaxon.lifecycle import LifeCycle, V1ProjectFeature, V1Statuses
 from polyaxon.sidecar.processor import SidecarThread
 from polyaxon.utils.fqn_utils import to_fqn_name
@@ -150,7 +150,7 @@ def __init__(
         is_new = is_new or (
             self._run_uuid is None and not settings.CLIENT_CONFIG.is_managed
         )
-        has_process_sidecar = os.environ.get(EV_KEYS_HAS_PROCESS_SIDECAR, False)
+        has_process_sidecar = os.environ.get(ENV_KEYS_HAS_PROCESS_SIDECAR, False)
 
         if auto_create and (is_new or self._is_offline):
             super().create(name=name, description=description, tags=tags)