diff --git a/traceml/tests/test_events/test_schemas.py b/traceml/tests/test_events/test_schemas.py
index 2dd178e88..d3b766cb3 100644
--- a/traceml/tests/test_events/test_schemas.py
+++ b/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,
diff --git a/traceml/tests/test_events_processing/test_event_resources.py b/traceml/tests/test_events_processing/test_event_resources.py
index 64693a91f..f04f45ed0 100644
--- a/traceml/tests/test_events_processing/test_event_resources.py
+++ b/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
 
 
diff --git a/traceml/tests/test_events_processing/test_event_values.py b/traceml/tests/test_events_processing/test_event_values.py
index 7a896dbf1..e0a2b6c62 100644
--- a/traceml/tests/test_events_processing/test_event_values.py
+++ b/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,
diff --git a/traceml/tests/test_logging/test_logging_parser.py b/traceml/tests/test_logging/test_logging_parser.py
index f816dee51..04fcdef82 100644
--- a/traceml/tests/test_logging/test_logging_parser.py
+++ b/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,
diff --git a/traceml/tests/test_logging/test_logs.py b/traceml/tests/test_logging/test_logs.py
index ff30d2d6d..0bc96a989 100644
--- a/traceml/tests/test_logging/test_logs.py
+++ b/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
 
 
diff --git a/traceml/tests/test_serialization/test_event_recorder.py b/traceml/tests/test_serialization/test_event_recorder.py
index afc24376e..0b77204e1 100644
--- a/traceml/tests/test_serialization/test_event_recorder.py
+++ b/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,
diff --git a/traceml/tests/test_tracking/test_run_tracking.py b/traceml/tests/test_tracking/test_run_tracking.py
index 1d9fe2c65..1bb7f8c24 100644
--- a/traceml/tests/test_tracking/test_run_tracking.py
+++ b/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,19 +154,19 @@ 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"
@@ -174,7 +174,7 @@ def test_run_init(self, expanduser):
         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"
diff --git a/traceml/tests/test_tracking/test_summaries.py b/traceml/tests/test_tracking/test_summaries.py
index 27af202d2..06fb3f184 100644
--- a/traceml/tests/test_tracking/test_summaries.py
+++ b/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
diff --git a/traceml/traceml/artifacts/schemas.py b/traceml/traceml/artifacts/schemas.py
index 926ed0cae..148699350 100644
--- a/traceml/traceml/artifacts/schemas.py
+++ b/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
 
 
diff --git a/traceml/traceml/events/schemas.py b/traceml/traceml/events/schemas.py
index d06f1e0bc..5360dc5c0 100644
--- a/traceml/traceml/events/schemas.py
+++ b/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
 
 
diff --git a/traceml/traceml/integrations/fastai.py b/traceml/traceml/integrations/fastai.py
index 0db486cd5..680092216 100644
--- a/traceml/traceml/integrations/fastai.py
+++ b/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
 
diff --git a/traceml/traceml/integrations/hugging_face.py b/traceml/traceml/integrations/hugging_face.py
index 150eee5e9..a03f2ca2f 100644
--- a/traceml/traceml/integrations/hugging_face.py
+++ b/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
diff --git a/traceml/traceml/integrations/keras.py b/traceml/traceml/integrations/keras.py
index 4cc077053..8283e2869 100644
--- a/traceml/traceml/integrations/keras.py
+++ b/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
diff --git a/traceml/traceml/integrations/pytorch_lightning.py b/traceml/traceml/integrations/pytorch_lightning.py
index acb63010a..04675e8ef 100644
--- a/traceml/traceml/integrations/pytorch_lightning.py
+++ b/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
 
diff --git a/traceml/traceml/logging/handler.py b/traceml/traceml/logging/handler.py
index e01088d55..f919cc197 100644
--- a/traceml/traceml/logging/handler.py
+++ b/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
 
 
diff --git a/traceml/traceml/logging/schemas.py b/traceml/traceml/logging/schemas.py
index 70ec83b12..eb509505d 100644
--- a/traceml/traceml/logging/schemas.py
+++ b/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,
diff --git a/traceml/traceml/logging/streamer.py b/traceml/traceml/logging/streamer.py
index f96c0f44c..7a6216ea9 100644
--- a/traceml/traceml/logging/streamer.py
+++ b/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
 
 
diff --git a/traceml/traceml/processors/events_processors/events_audio_processors.py b/traceml/traceml/processors/events_processors/events_audio_processors.py
index fda1f746c..70666bc41 100644
--- a/traceml/traceml/processors/events_processors/events_audio_processors.py
+++ b/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
diff --git a/traceml/traceml/processors/events_processors/events_charts_processors.py b/traceml/traceml/processors/events_processors/events_charts_processors.py
index 683e2d98b..3cdf36953 100644
--- a/traceml/traceml/processors/events_processors/events_charts_processors.py
+++ b/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 (
diff --git a/traceml/traceml/processors/events_processors/events_image_processors.py b/traceml/traceml/processors/events_processors/events_image_processors.py
index 1b9513120..0327476b4 100644
--- a/traceml/traceml/processors/events_processors/events_image_processors.py
+++ b/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 (
diff --git a/traceml/traceml/processors/events_processors/events_metrics_processors.py b/traceml/traceml/processors/events_processors/events_metrics_processors.py
index b53ca0f34..45393c0a4 100644
--- a/traceml/traceml/processors/events_processors/events_metrics_processors.py
+++ b/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,
diff --git a/traceml/traceml/processors/events_processors/events_video_processors.py b/traceml/traceml/processors/events_processors/events_video_processors.py
index 2c98c6fd5..b9bcd45ae 100644
--- a/traceml/traceml/processors/events_processors/events_video_processors.py
+++ b/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
diff --git a/traceml/traceml/tracking/__init__.py b/traceml/traceml/tracking/__init__.py
index 4309a1f58..711c4c512 100644
--- a/traceml/traceml/tracking/__init__.py
+++ b/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,
diff --git a/traceml/traceml/tracking/run.py b/traceml/traceml/tracking/run.py
index 7255d8de9..f29714a33 100644
--- a/traceml/traceml/tracking/run.py
+++ b/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
diff --git a/traceml/traceml/visualization/run_plot.py b/traceml/traceml/visualization/run_plot.py
index 7710f5a84..6ae975d25 100644
--- a/traceml/traceml/visualization/run_plot.py
+++ b/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