Refactor and simplify imports

traceml/tests/test_events/test_schemas.py

@@ -4,7 +4,7 @@
 from clipped.utils.dates import parse_datetime
 from clipped.utils.tz import now
 
-from polyaxon.utils.test_utils import BaseTestCase
+from polyaxon._utils.test_utils import BaseTestCase
 from traceml.events.schemas import (
     LoggedEventListSpec,
     V1Event,

traceml/tests/test_events_processing/test_event_resources.py

@@ -1,6 +1,6 @@
 import pytest
 
-from polyaxon.utils.test_utils import BaseTestCase
+from polyaxon._utils.test_utils import BaseTestCase
 from traceml.processors.events_processors import metrics_dict_to_list
 
 

traceml/tests/test_events_processing/test_event_values.py

@@ -8,7 +8,7 @@
 from PIL import Image
 from plotly import figure_factory
 
-from polyaxon.utils.test_utils import BaseTestCase, tensor_np
+from polyaxon._utils.test_utils import BaseTestCase, tensor_np
 from traceml.processors.events_processors import (
     audio,
     bokeh_chart,

traceml/tests/test_logging/test_logging_parser.py

@@ -2,7 +2,7 @@
 
 from clipped.utils.dates import parse_datetime
 
-from polyaxon.utils.test_utils import BaseTestCase
+from polyaxon._utils.test_utils import BaseTestCase
 from traceml.logging.parser import (
     DATETIME_REGEX,
     ISO_DATETIME_REGEX,

traceml/tests/test_logging/test_logs.py

@@ -3,7 +3,7 @@
 from clipped.utils.dates import parse_datetime
 from clipped.utils.tz import now
 
-from polyaxon.utils.test_utils import BaseTestCase
+from polyaxon._utils.test_utils import BaseTestCase
 from traceml.logging.schemas import V1Log, V1Logs
 
 

traceml/tests/test_serialization/test_event_recorder.py

@@ -2,7 +2,7 @@
 import pytest
 import tempfile
 
-from polyaxon.utils.test_utils import BaseTestCase
+from polyaxon._utils.test_utils import BaseTestCase
 from traceml.events.schemas import (
     LoggedEventListSpec,
     LoggedEventSpec,

traceml/tests/test_tracking/test_run_tracking.py

@@ -14,19 +14,19 @@
 from clipped.utils.paths import create_path
 from plotly import figure_factory
 
-from polyaxon import dist, settings
-from polyaxon.constants.globals import DEFAULT
-from polyaxon.contexts import paths as ctx_paths
-from polyaxon.env_vars import getters
-from polyaxon.env_vars.keys import (
+from polyaxon import _dist, settings
+from polyaxon._constants.globals import DEFAULT
+from polyaxon._contexts import paths as ctx_paths
+from polyaxon._env_vars import getters
+from polyaxon._env_vars.keys import (
     ENV_KEYS_COLLECT_ARTIFACTS,
     ENV_KEYS_COLLECT_RESOURCES,
     ENV_KEYS_LOG_LEVEL,
     ENV_KEYS_RUN_INSTANCE,
 )
+from polyaxon._utils.test_utils import TestEnvVarsCase, tensor_np
 from polyaxon.exceptions import PolyaxonClientException
-from polyaxon.lifecycle import V1ProjectFeature
-from polyaxon.utils.test_utils import TestEnvVarsCase, tensor_np
+from polyaxon.schemas import V1ProjectFeature
 from traceml.artifacts import V1ArtifactKind
 from traceml.events import V1Events, get_asset_path, get_event_path
 from traceml.serialization.writer import EventFileWriter, ResourceFileWriter
@@ -123,13 +123,13 @@ def test_run_init(self, expanduser):
             Run()
 
         # Uses default as owner in non CE
-        settings.CLI_CONFIG.installation = {"dist": dist.EE}
+        settings.CLI_CONFIG.installation = {"dist": _dist.EE}
         with self.assertRaises(PolyaxonClientException):
             Run(project="test")
 
         # Uses default as owner in CE
         settings.CLIENT_CONFIG.is_offline = True
-        settings.CLI_CONFIG.installation = {"dist": dist.CE}
+        settings.CLI_CONFIG.installation = {"dist": _dist.CE}
         with patch("traceml.tracking.run.Run._set_exit_handler") as exit_mock:
             run = Run(project="test", track_code=False, track_env=False)
         assert exit_mock.call_count == 1
@@ -154,27 +154,27 @@ def test_run_init(self, expanduser):
         with self.assertRaises(PolyaxonClientException):
             Run()
 
-        settings.CLI_CONFIG.installation = {"dist": dist.EE}
+        settings.CLI_CONFIG.installation = {"dist": _dist.EE}
         # Uses default as owner in non CE
         with self.assertRaises(PolyaxonClientException):
             Run(project="test")
 
         # Uses default as owner in CE
         settings.CLIENT_CONFIG.is_offline = True
-        settings.CLI_CONFIG.installation = {"dist": dist.CE}
+        settings.CLI_CONFIG.installation = {"dist": _dist.CE}
         run = Run(project="test")
         assert run.owner == DEFAULT
 
         # FQN non CE
-        settings.CLI_CONFIG.installation = {"dist": dist.EE}
+        settings.CLI_CONFIG.installation = {"dist": _dist.EE}
         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}
+        settings.CLI_CONFIG.installation = {"dist": _dist.CE}
         os.environ[ENV_KEYS_RUN_INSTANCE] = "user.project_bar.runs.{}".format(uid)
         run = Run()
         assert run.owner == "user"

traceml/tests/test_tracking/test_summaries.py

@@ -5,11 +5,11 @@
 from mock import patch
 
 from polyaxon import settings
-from polyaxon.env_vars.keys import (
+from polyaxon._env_vars.keys import (
     ENV_KEYS_COLLECT_ARTIFACTS,
     ENV_KEYS_COLLECT_RESOURCES,
 )
-from polyaxon.utils.test_utils import BaseTestCase
+from polyaxon._utils.test_utils import BaseTestCase
 from traceml.artifacts import V1RunArtifact
 from traceml.events import V1Events
 from traceml.tracking.run import Run

traceml/traceml/artifacts/schemas.py

@@ -5,7 +5,7 @@
 from clipped.compact.pydantic import StrictStr
 from clipped.types.uuids import UUIDStr
 
-from polyaxon.schemas.base import BaseSchemaModel
+from polyaxon._schemas.base import BaseSchemaModel
 from traceml.artifacts.enums import V1ArtifactKind
 
 

traceml/traceml/events/schemas.py

@@ -13,7 +13,7 @@
 from clipped.utils.np import sanitize_np_types
 from clipped.utils.tz import now
 
-from polyaxon.schemas.base import BaseSchemaModel
+from polyaxon._schemas.base import BaseSchemaModel
 from traceml.artifacts.enums import V1ArtifactKind
 
 

traceml/traceml/integrations/fastai.py

@@ -1,4 +1,4 @@
-from polyaxon.client.decorators import client_handler
+from polyaxon._client.decorators import client_handler
 from traceml import tracking
 from traceml.exceptions import TracemlException
 

traceml/traceml/integrations/hugging_face.py

@@ -1,4 +1,4 @@
-from polyaxon.client.decorators import client_handler
+from polyaxon._client.decorators import client_handler
 from traceml import tracking
 from traceml.exceptions import TracemlException
 from traceml.logger import logger

traceml/traceml/integrations/keras.py

@@ -4,7 +4,7 @@
 
 from clipped.utils.np import sanitize_np_types
 
-from polyaxon.client.decorators import client_handler
+from polyaxon._client.decorators import client_handler
 from traceml import tracking
 from traceml.exceptions import TracemlException
 from traceml.logger import logger

traceml/traceml/integrations/pytorch_lightning.py

@@ -6,8 +6,8 @@
 
 import packaging
 
+from polyaxon._env_vars.keys import ENV_KEYS_RUN_INSTANCE
 from polyaxon.client import RunClient
-from polyaxon.env_vars.keys import ENV_KEYS_RUN_INSTANCE
 from traceml import tracking
 from traceml.exceptions import TracemlException
 

traceml/traceml/logging/handler.py

@@ -6,7 +6,7 @@
 from clipped.utils.env import get_user
 
 from polyaxon import settings
-from polyaxon.env_vars.keys import ENV_KEYS_K8S_NODE_NAME, ENV_KEYS_K8S_POD_ID
+from polyaxon._env_vars.keys import ENV_KEYS_K8S_NODE_NAME, ENV_KEYS_K8S_POD_ID
 from traceml.logging.schemas import V1Log
 
 

traceml/traceml/logging/schemas.py

@@ -8,7 +8,7 @@
 from clipped.utils.json import orjson_dumps, orjson_loads
 from clipped.utils.tz import now
 
-from polyaxon.schemas.base import BaseSchemaModel
+from polyaxon._schemas.base import BaseSchemaModel
 from traceml.logging.parser import (
     DATETIME_REGEX,
     ISO_DATETIME_REGEX,

traceml/traceml/logging/streamer.py

@@ -8,7 +8,7 @@
 from clipped.utils.tz import local_datetime
 
 from polyaxon import settings
-from polyaxon.containers.names import MAIN_CONTAINER_NAMES
+from polyaxon._containers.names import MAIN_CONTAINER_NAMES
 from traceml.logging.schemas import V1Log, V1Logs
 
 

traceml/traceml/processors/events_processors/events_audio_processors.py

@@ -3,7 +3,7 @@
 from clipped.utils.np import to_np
 from clipped.utils.paths import check_or_create_path, copy_file_path
 
-from polyaxon.constants.globals import UNKNOWN
+from polyaxon._constants.globals import UNKNOWN
 from traceml.events import V1EventAudio
 from traceml.logger import logger
 from traceml.processors.errors import NUMPY_ERROR_MESSAGE

traceml/traceml/processors/events_processors/events_charts_processors.py

@@ -1,6 +1,6 @@
 from clipped.utils.paths import module_type
 
-from polyaxon.constants.globals import UNKNOWN
+from polyaxon._constants.globals import UNKNOWN
 from traceml.events import V1EventChart, V1EventChartKind
 from traceml.logger import logger
 from traceml.processors.errors import (

traceml/traceml/processors/events_processors/events_image_processors.py

@@ -5,7 +5,7 @@
 from clipped.utils.np import calculate_scale_factor, to_np
 from clipped.utils.paths import check_or_create_path, copy_file_path
 
-from polyaxon.constants.globals import UNKNOWN
+from polyaxon._constants.globals import UNKNOWN
 from traceml.events import V1EventImage
 from traceml.logger import logger
 from traceml.processors.errors import (

traceml/traceml/processors/events_processors/events_metrics_processors.py

@@ -2,7 +2,7 @@
 
 from clipped.utils.np import to_np
 
-from polyaxon.constants.globals import UNKNOWN
+from polyaxon._constants.globals import UNKNOWN
 from traceml.artifacts import V1ArtifactKind
 from traceml.events import (
     LoggedEventSpec,

traceml/traceml/processors/events_processors/events_video_processors.py

@@ -3,7 +3,7 @@
 from clipped.utils.np import calculate_scale_factor, to_np
 from clipped.utils.paths import check_or_create_path, copy_file_path
 
-from polyaxon.constants.globals import UNKNOWN
+from polyaxon._constants.globals import UNKNOWN
 from traceml.events import V1EventVideo
 from traceml.logger import logger
 from traceml.processors.errors import MOVIEPY_ERROR_MESSAGE, NUMPY_ERROR_MESSAGE

traceml/traceml/tracking/__init__.py

@@ -1,8 +1,7 @@
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Sequence, Union
 
-from polyaxon.client import RunClient
-from polyaxon.schemas.responses.v1_run import V1Run
+from polyaxon.client import *
 from traceml.artifacts import V1ArtifactKind, V1RunArtifact
 from traceml.tracking.run import Run
 
@@ -138,6 +137,22 @@ def update(data: Union[Dict, V1Run], async_req: bool = False):
 update.__doc__ = Run.update.__doc__
 
 
+def get_connections_catalog():
+    global TRACKING_RUN
+    return TRACKING_RUN.get_connections_catalog()
+
+
+get_connections_catalog.__doc__ = Run.get_connections_catalog.__doc__
+
+
+def get_artifacts_store_connection():
+    global TRACKING_RUN
+    return TRACKING_RUN.get_artifacts_store_connection()
+
+
+get_artifacts_store_connection.__doc__ = Run.get_artifacts_store_connection.__doc__
+
+
 def get_artifacts_path(
     rel_path: Optional[str] = None,
     ensure_path: bool = False,

traceml/traceml/tracking/run.py

@@ -18,21 +18,21 @@
 )
 
 from polyaxon import settings
-from polyaxon.client import PolyaxonClient, RunClient
-from polyaxon.client.decorators import client_handler
-from polyaxon.connections import CONNECTION_CONFIG
-from polyaxon.constants.globals import UNKNOWN
-from polyaxon.contexts import paths as ctx_paths
-from polyaxon.env_vars.getters import (
+from polyaxon._client.decorators import client_handler
+from polyaxon._connections import CONNECTION_CONFIG, V1Connection
+from polyaxon._constants.globals import UNKNOWN
+from polyaxon._contexts import paths as ctx_paths
+from polyaxon._env_vars.getters import (
     get_artifacts_store_name,
     get_collect_artifacts,
     get_collect_resources,
     get_log_level,
 )
-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
+from polyaxon._env_vars.keys import ENV_KEYS_HAS_PROCESS_SIDECAR
+from polyaxon._sidecar.processor import SidecarThread
+from polyaxon._utils.fqn_utils import to_fqn_name
+from polyaxon.client import PolyaxonClient, RunClient
+from polyaxon.schemas import LifeCycle, V1ProjectFeature, V1Statuses
 from traceml.artifacts import V1ArtifactKind
 from traceml.events import LoggedEventSpec, V1Event, get_asset_path
 from traceml.logger import logger
@@ -260,10 +260,20 @@ def create(self, **kwargs):
             "`create` method manually, please create a new instance of `Run` with `is_new=True`"
         )
 
+    def get_connections_catalog(self) -> Optional[List[V1Connection]]:
+        """Returns the current connections catalog requested by this run."""
+        catalog = CONNECTION_CONFIG.catalog
+        if catalog:
+            return catalog.connections
+
+    def get_artifacts_store_connection(self) -> Optional[V1Connection]:
+        """Returns the current artifacts store connection used by this run."""
+        return CONNECTION_CONFIG.get_connection_for(get_artifacts_store_name())
+
     def _get_store_path(self):
         if self._store_path:
             return self._store_path
-        connection = CONNECTION_CONFIG.get_connection_for(get_artifacts_store_name())
+        connection = self.get_artifacts_store_connection()
         if not connection:
             logger.warning("Artifacts store connection not detected.")
             return None

traceml/traceml/visualization/run_plot.py

@@ -2,8 +2,8 @@
 
 from typing import Dict, List, Optional, Set, Union
 
+from polyaxon._client.decorators import client_handler
 from polyaxon.client import RunClient
-from polyaxon.client.decorators import client_handler
 from traceml.artifacts import V1ArtifactKind
 from traceml.events import V1Events