polyaxon 2.0.0rc49__py3-none-any.whl → 2.4.0rc1__py3-none-any.whl

polyaxon/_runner/agent/sync_agent.py

@@ -13,6 +13,7 @@ from urllib3.exceptions import HTTPError
 from polyaxon import pkg, settings
 from polyaxon._env_vars.getters import get_run_info
 from polyaxon._runner.agent.base_agent import BaseAgent
+from polyaxon._utils.fqn_utils import get_run_instance
 from polyaxon.client import V1Agent, V1AgentStateResponse
 from polyaxon.exceptions import ApiException as SDKApiException
 from polyaxon.exceptions import PolyaxonAgentError, PolyaxonConverterError
@@ -25,13 +26,13 @@ class BaseSyncAgent(BaseAgent):
     def _enter(self):
         if not self.client._is_managed:
             return self
-        print("Agent is starting.")
+        logger.warning("Agent is starting.")
         try:
             agent = self.client.get_info()
             self._check_status(agent)
             self.sync()
             self.client.log_agent_running()
-            print("Agent is running.")
+            logger.warning("Agent is running.")
             return self
         except (ApiException, SDKApiException, HTTPError) as e:
             message = "Could not start the agent."
@@ -76,6 +77,49 @@ class BaseSyncAgent(BaseAgent):
             ),
         )
 
+    def reconcile(self):
+        if (
+            now() - self._last_reconciled_at
+        ).total_seconds() > self.SLEEP_AGENT_DATA_COLLECT_TIME:
+            return
+
+        # Collect data
+        self.collect_agent_data()
+
+        # Update reconcile
+        namespaces = [settings.AGENT_CONFIG.namespace]
+        namespaces += settings.AGENT_CONFIG.additional_namespaces or []
+        ops = []
+        for namespace in namespaces:
+            _ops = self.executor.list_ops(namespace=namespace)
+            if _ops:
+                ops += [
+                    (
+                        get_run_instance(
+                            owner=op["metadata"]["annotations"][
+                                "operation.polyaxon.com/owner"
+                            ],
+                            project=op["metadata"]["annotations"][
+                                "operation.polyaxon.com/project"
+                            ],
+                            run_uuid=op["metadata"]["labels"][
+                                "app.kubernetes.io/instance"
+                            ],
+                        ),
+                        op["metadata"]["annotations"]["operation.polyaxon.com/kind"],
+                        op["metadata"]["annotations"]["operation.polyaxon.com/name"],
+                        namespace,
+                    )
+                    for op in _ops
+                ]
+        if not ops:
+            return None
+
+        logger.info("Reconcile agent.")
+        return self.client.reconcile_agent(
+            reconcile={"ops": ops},
+        )
+
     def start(self):
         try:
             with sync_exit_context() as exit_event:
@@ -88,7 +132,9 @@ class BaseSyncAgent(BaseAgent):
                     while not exit_event.wait(timeout=timeout):
                         index += 1
                         self.refresh_executor()
-                        if not self._default_auth:
+                        if self._default_auth:
+                            self.reconcile()
+                        else:
                             self.cron()
                         agent_state = self.process(pool)
                         if not agent_state:
@@ -113,7 +159,7 @@ class BaseSyncAgent(BaseAgent):
                 self.sync_compatible_updates(agent_state.compatible_updates)
 
             if agent_state:
-                logger.info("Starting runs submission process.")
+                logger.info("Checking agent state.")
             else:
                 logger.info("No state was found.")
                 return V1AgentStateResponse.construct()
@@ -180,7 +226,7 @@ class BaseSyncAgent(BaseAgent):
             )
         return None
 
-    def submit_run(self, run_data: Tuple[str, str, str, str]):
+    def submit_run(self, run_data: Tuple[str, str, str, str, str]):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
         resource = self.prepare_run_resource(
             owner_name=run_owner,
@@ -194,7 +240,10 @@ class BaseSyncAgent(BaseAgent):
 
         try:
             self.executor.create(
-                run_uuid=run_uuid, run_kind=run_data[1], resource=resource
+                run_uuid=run_uuid,
+                run_kind=run_data[1],
+                resource=resource,
+                namespace=run_data[4],
             )
         except ApiException as e:
             if e.status == 409:
@@ -217,7 +266,7 @@ class BaseSyncAgent(BaseAgent):
             )
 
     def make_and_create_run(
-        self, run_data: Tuple[str, str, str, str], default_auth: bool = False
+        self, run_data: Tuple[str, str, str, str, str], default_auth: bool = False
     ):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
         resource = self.make_run_resource(
@@ -233,7 +282,10 @@ class BaseSyncAgent(BaseAgent):
 
         try:
             self.executor.create(
-                run_uuid=run_uuid, run_kind=run_data[1], resource=resource
+                run_uuid=run_uuid,
+                run_kind=run_data[1],
+                resource=resource,
+                namespace=run_data[4],
             )
         except ApiException as e:
             if e.status == 409:
@@ -247,7 +299,7 @@ class BaseSyncAgent(BaseAgent):
                 )
             )
 
-    def apply_run(self, run_data: Tuple[str, str, str, str]):
+    def apply_run(self, run_data: Tuple[str, str, str, str, str]):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
         resource = self.prepare_run_resource(
             owner_name=run_owner,
@@ -261,7 +313,10 @@ class BaseSyncAgent(BaseAgent):
 
         try:
             self.executor.apply(
-                run_uuid=run_uuid, run_kind=run_data[1], resource=resource
+                run_uuid=run_uuid,
+                run_kind=run_data[1],
+                resource=resource,
+                namespace=run_data[4],
             )
             self.client.log_run_running(
                 run_owner=run_owner, run_project=run_project, run_uuid=run_uuid
@@ -270,12 +325,16 @@ class BaseSyncAgent(BaseAgent):
             self.client.log_run_failed(
                 run_owner=run_owner, run_project=run_project, run_uuid=run_uuid, exc=e
             )
-            self.clean_run(run_uuid=run_uuid, run_kind=run_data[1])
+            self.clean_run(
+                run_uuid=run_uuid, run_kind=run_data[1], namespace=run_data[4]
+            )
 
-    def check_run(self, run_data: Tuple[str, str]):
+    def check_run(self, run_data: Tuple[str, str, str]):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
         try:
-            self.executor.get(run_uuid=run_uuid, run_kind=run_data[1])
+            self.executor.get(
+                run_uuid=run_uuid, run_kind=run_data[1], namespace=run_data[2]
+            )
         except ApiException as e:
             if e.status == 404:
                 logger.info(
@@ -285,10 +344,12 @@ class BaseSyncAgent(BaseAgent):
                     run_owner=run_owner, run_project=run_project, run_uuid=run_uuid
                 )
 
-    def stop_run(self, run_data: Tuple[str, str]):
+    def stop_run(self, run_data: Tuple[str, str, str]):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
         try:
-            self.executor.stop(run_uuid=run_uuid, run_kind=run_data[1])
+            self.executor.stop(
+                run_uuid=run_uuid, run_kind=run_data[1], namespace=run_data[2]
+            )
         except ApiException as e:
             if e.status == 404:
                 logger.info("Run does not exist anymore, it could have been stopped.")
@@ -304,16 +365,20 @@ class BaseSyncAgent(BaseAgent):
                 message="Agent failed stopping run.\n",
             )
 
-    def delete_run(self, run_data: Tuple[str, str, str, str]):
+    def delete_run(self, run_data: Tuple[str, str, str, str, str]):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
-        self.clean_run(run_uuid=run_uuid, run_kind=run_data[1])
+        self.clean_run(run_uuid=run_uuid, run_kind=run_data[1], namespace=run_data[4])
         if run_data[3]:
             self.make_and_create_run(run_data)
 
-    def clean_run(self, run_uuid: str, run_kind: str):
+    def clean_run(self, run_uuid: str, run_kind: str, namespace: str = None):
         try:
-            self.executor.clean(run_uuid=run_uuid, run_kind=run_kind)
-            self.executor.stop(run_uuid=run_uuid, run_kind=run_kind)
+            self.executor.clean(
+                run_uuid=run_uuid, run_kind=run_kind, namespace=namespace
+            )
+            self.executor.stop(
+                run_uuid=run_uuid, run_kind=run_kind, namespace=namespace
+            )
         except ApiException as e:
             if e.status == 404:
                 logger.info("Run does not exist.")

polyaxon/_runner/converter/converter.py

@@ -711,7 +711,9 @@ class BaseConverter:
                     V1ConnectionKind.is_ssh(connection_spec.kind)
                     and init_connection.git
                 ):
-                    patch_git(connection_spec.schema_, init_connection.git)
+                    connection_spec.schema_ = patch_git(
+                        connection_spec.schema_, init_connection.git
+                    )
                     containers.append(
                         self._get_git_init_container(
                             run_path=self.run_path,
@@ -731,7 +733,9 @@ class BaseConverter:
                     )
                 elif V1ConnectionKind.is_git(connection_spec.kind):
                     if init_connection.git:  # Update the default schema
-                        connection_spec.schema_.patch(init_connection.git)
+                        connection_spec.schema_ = patch_git(
+                            connection_spec.schema_, init_connection.git
+                        )
                     containers.append(
                         self._get_git_init_container(
                             run_path=self.run_path,

polyaxon/_runner/executor.py

@@ -50,19 +50,21 @@ class BaseExecutor:
         self._manager = None
         return self.manager
 
-    def get(self, run_uuid: str, run_kind: str):
+    def get(self, run_uuid: str, run_kind: str, namespace: str = None):
         raise NotImplementedError
 
-    def create(self, run_uuid: str, run_kind: str, resource: Any):
+    def create(
+        self, run_uuid: str, run_kind: str, resource: Any, namespace: str = None
+    ):
         raise NotImplementedError
 
-    def apply(self, run_uuid: str, run_kind: str, resource: Any):
+    def apply(self, run_uuid: str, run_kind: str, resource: Any, namespace: str = None):
         raise NotImplementedError
 
-    def stop(self, run_uuid: str, run_kind: str):
+    def stop(self, run_uuid: str, run_kind: str, namespace: str = None):
         raise NotImplementedError
 
-    def clean(self, run_uuid: str, run_kind: str):
+    def clean(self, run_uuid: str, run_kind: str, namespace: str = None):
         raise NotImplementedError
 
     def _clean_temp_execution_path(self, run_uuid: str):
@@ -154,7 +156,7 @@ class BaseExecutor:
             run_name=run_name,
             run_uuid=run_uuid,
             run_path=run_uuid,
-            namespace=agent_env.namespace,
+            namespace=compiled_operation.namespace or agent_env.namespace,
             compiled_operation=compiled_operation,
             polyaxon_init=agent_env.polyaxon_init,
             polyaxon_sidecar=agent_env.polyaxon_sidecar,
@@ -191,7 +193,7 @@ class BaseExecutor:
             params=operation.params,
         )
         return cls.get_resource(
-            namespace=resolver_obj.namespace,
+            namespace=compiled_operation.namespace or resolver_obj.namespace,
             owner_name=resolver_obj.owner_name,
             project_name=resolver_obj.project_name,
             run_name=resolver_obj.run_name,
@@ -223,4 +225,8 @@ class BaseExecutor:
             run_uuid=response.uuid,
             run_kind=response.kind,
             resource=resource,
+            namespace=response.settings.namespace if response.settings else None,
         )
+
+    def list_ops(self, namespace: str = None):
+        raise NotImplementedError

polyaxon/_schemas/agent.py

@@ -18,6 +18,7 @@ from polyaxon._connections import V1Connection, V1ConnectionKind, V1HostPathConn
 from polyaxon._contexts import paths as ctx_paths
 from polyaxon._env_vars.getters import get_artifacts_store_name
 from polyaxon._env_vars.keys import (
+    ENV_KEYS_ADDITIONAL_NAMESPACES,
     ENV_KEYS_AGENT_ARTIFACTS_STORE,
     ENV_KEYS_AGENT_CLEANER,
     ENV_KEYS_AGENT_CONNECTIONS,
@@ -35,6 +36,8 @@ from polyaxon._env_vars.keys import (
     ENV_KEYS_ARTIFACTS_STORE_NAME,
     ENV_KEYS_K8S_APP_SECRET_NAME,
     ENV_KEYS_K8S_NAMESPACE,
+    ENV_KEYS_SINGLE_NAMESPACE,
+    ENV_KEYS_WATCH_CLUSTER,
 )
 from polyaxon._fs.utils import get_store_path
 from polyaxon._schemas.base import BaseSchemaModel
@@ -185,6 +188,11 @@ class AgentConfig(BaseAgentConfig):
     _IDENTIFIER = "agent"
 
     is_replica: Optional[bool] = Field(alias=ENV_KEYS_AGENT_IS_REPLICA)
+    watch_cluster: Optional[bool] = Field(alias=ENV_KEYS_WATCH_CLUSTER)
+    single_namespace: Optional[bool] = Field(alias=ENV_KEYS_SINGLE_NAMESPACE)
+    additional_namespaces: Optional[List[StrictStr]] = Field(
+        alias=ENV_KEYS_ADDITIONAL_NAMESPACES
+    )
     sidecar: Optional[V1PolyaxonSidecarContainer] = Field(alias=ENV_KEYS_AGENT_SIDECAR)
     init: Optional[V1PolyaxonInitContainer] = Field(alias=ENV_KEYS_AGENT_INIT)
     notifier: Optional[V1PolyaxonNotifier] = Field(alias=ENV_KEYS_AGENT_NOTIFIER)
@@ -217,6 +225,24 @@ class AgentConfig(BaseAgentConfig):
             and "isReplica" in values
         ):
             values[ENV_KEYS_AGENT_IS_REPLICA] = values["isReplica"]
+        if (
+            not values.get("watch_cluster")
+            and not values.get(ENV_KEYS_WATCH_CLUSTER)
+            and "watchCluster" in values
+        ):
+            values[ENV_KEYS_WATCH_CLUSTER] = values["watchCluster"]
+        if (
+            not values.get("single_namespace")
+            and not values.get(ENV_KEYS_SINGLE_NAMESPACE)
+            and "singleNamespace" in values
+        ):
+            values[ENV_KEYS_SINGLE_NAMESPACE] = values["singleNamespace"]
+        if (
+            not values.get("additional_namespace")
+            and not values.get(ENV_KEYS_ADDITIONAL_NAMESPACES)
+            and "additionalNamespaces" in values
+        ):
+            values[ENV_KEYS_ADDITIONAL_NAMESPACES] = values["additionalNamespaces"]
         if (
             not values.get("use_proxy_env_vars_use_in_ops")
             and not values.get(ENV_KEYS_AGENT_USE_PROXY_ENV_VARS_IN_OPS)
@@ -312,7 +338,7 @@ class AgentConfig(BaseAgentConfig):
                 "Received an invalid {} `{}`".format(field.alias, v)
             ) from e
 
-    @validator("default_image_pull_secrets", pre=True)
+    @validator("additional_namespaces", "default_image_pull_secrets", pre=True)
     def validate_str_list(cls, v, field):
         try:
             return ConfigParser.parse(str)(

polyaxon/_schemas/client.py

@@ -1,3 +1,5 @@
+import os
+
 from typing import Dict, Optional, Union
 
 import urllib3
@@ -29,6 +31,7 @@ from polyaxon._env_vars.keys import (
     ENV_KEYS_NO_API,
     ENV_KEYS_NO_OP,
     ENV_KEYS_RETRIES,
+    ENV_KEYS_SECRET_INTERNAL_TOKEN,
     ENV_KEYS_SSL_CA_CERT,
     ENV_KEYS_TIME_ZONE,
     ENV_KEYS_TIMEOUT,
@@ -137,6 +140,10 @@ class ClientConfig(BaseSchemaModel):
         self.client_header["header_name"] = self.header
         self.client_header["header_value"] = self.header_service
 
+    def get_internal_header(self) -> Dict:
+        header = PolyaxonServiceHeaders.get_header(PolyaxonServiceHeaders.INTERNAL)
+        return {"header_name": header, "header_value": self.header_service}
+
     def get_full_headers(self, headers=None, auth_key="Authorization") -> Dict:
         request_headers = {}
         request_headers.update(headers or {})
@@ -151,7 +158,19 @@ class ClientConfig(BaseSchemaModel):
         return request_headers
 
     @property
-    def sdk_config(self) -> Configuration:
+    def sdk_config(self):
+        return self.get_sdk_config()
+
+    @property
+    def internal_sdk_config(self):
+        return self.get_sdk_config(
+            token=os.environ.get(ENV_KEYS_SECRET_INTERNAL_TOKEN),
+            authentication_type=AuthenticationTypes.INTERNAL_TOKEN,
+        )
+
+    def get_sdk_config(
+        self, token: str = None, authentication_type: AuthenticationTypes = None
+    ) -> Configuration:
         if not self.host and not self.in_cluster:
             raise PolyaxonClientException(
                 "Api config requires at least a host if not running in-cluster."
@@ -171,8 +190,10 @@ class ClientConfig(BaseSchemaModel):
         if self.connection_pool_maxsize:
             config.connection_pool_maxsize = self.connection_pool_maxsize
         if self.token:
-            config.api_key["ApiKey"] = self.token
-            config.api_key_prefix["ApiKey"] = self.authentication_type
+            config.api_key["ApiKey"] = token or self.token
+            config.api_key_prefix["ApiKey"] = (
+                authentication_type or self.authentication_type
+            )
         return config
 
     @property
@@ -181,6 +202,12 @@ class ClientConfig(BaseSchemaModel):
         config.connection_pool_maxsize = 100
         return config
 
+    @property
+    def async_internal_sdk_config(self) -> Configuration:
+        config = self.internal_sdk_config
+        config.connection_pool_maxsize = 100
+        return config
+
     @classmethod
     def patch_from(cls, config: "ClientConfig", **kwargs) -> "ClientConfig":
         data = {**config.to_dict(), **kwargs}

polyaxon/_schemas/installation.py

@@ -1,14 +1,15 @@
 from typing import List, Optional
 
 from clipped.compact.pydantic import StrictStr
+from clipped.config.schema import BaseAllowSchemaModel
 
-from polyaxon._schemas.base import BaseSchemaModel
 
-
-class V1Installation(BaseSchemaModel):
+class V1Installation(BaseAllowSchemaModel):
     key: Optional[StrictStr]
     version: Optional[StrictStr]
     dist: Optional[StrictStr]
     host: Optional[StrictStr]
     hmac: Optional[StrictStr]
+    mode: Optional[StrictStr]
+    org: Optional[bool]
     auth: Optional[List[StrictStr]]

polyaxon/_schemas/lifecycle.py

@@ -3,6 +3,7 @@ import datetime
 from typing import Any, Dict, List, Optional, Union
 
 from clipped.compact.pydantic import StrictStr
+from clipped.config.schema import BaseAllowSchemaModel
 from clipped.utils.dates import parse_datetime
 from clipped.utils.enums import PEnum
 from clipped.utils.tz import now
@@ -275,8 +276,8 @@ class LifeCycle:
             entity.started_at = now()
             # Update wait_time
             if entity.wait_time is None:
-                entity.wait_time = int(
-                    (entity.started_at - entity.created_at).total_seconds()
+                entity.wait_time = round(
+                    (entity.started_at - entity.created_at).total_seconds(), 2
                 )
             return True
 
@@ -290,14 +291,14 @@ class LifeCycle:
                 entity.started_at = entity.created_at
             # Update duration
             if entity.duration is None:
-                entity.duration = int(
-                    (entity.finished_at - entity.started_at).total_seconds()
+                entity.duration = round(
+                    (entity.finished_at - entity.started_at).total_seconds(), 2
                 )
             return True
         return False
 
 
-class BaseCondition(BaseSchemaModel):
+class BaseCondition(BaseAllowSchemaModel):
     @classmethod
     def get_condition(
         cls,
@@ -307,6 +308,7 @@ class BaseCondition(BaseSchemaModel):
         last_transition_time=None,
         reason=None,
         message=None,
+        meta_info=None,
     ):
         current_time = now()
         last_update_time = last_update_time or current_time
@@ -318,6 +320,7 @@ class BaseCondition(BaseSchemaModel):
             last_transition_time=last_transition_time,
             reason=reason,
             message=message,
+            meta_info=meta_info,
         )
 
     def __eq__(self, other):
@@ -335,6 +338,7 @@ class V1StatusCondition(BaseCondition):
     message: Optional[StrictStr]
     last_update_time: Optional[datetime.datetime]
     last_transition_time: Optional[datetime.datetime]
+    meta_info: Optional[Dict[str, Any]]
 
 
 class V1Status(BaseSchemaModel):
@@ -351,6 +355,7 @@ class V1StageCondition(BaseCondition):
     message: Optional[StrictStr]
     last_update_time: Optional[datetime.datetime]
     last_transition_time: Optional[datetime.datetime]
+    meta_info: Optional[Dict[str, Any]]
 
 
 class V1Stage(BaseSchemaModel):

polyaxon/_schemas/log_handler.py

@@ -3,11 +3,10 @@ import base64
 from typing import Optional
 
 from clipped.compact.pydantic import StrictStr
+from clipped.config.schema import BaseAllowSchemaModel
 
-from polyaxon._schemas.base import BaseSchemaModel
 
-
-class V1LogHandler(BaseSchemaModel):
+class V1LogHandler(BaseAllowSchemaModel):
     _IDENTIFIER = "log_handler"
 
     dsn: Optional[StrictStr]

polyaxon/_schemas/types/artifacts.py

@@ -129,8 +129,8 @@ class V1ArtifactsType(BaseTypeConfig):
 
     ```python
     >>> from polyaxon.schemas import V1Component, V1Init, V1Job
-    >>> from polyaxon.schemas.types import V1ArtifactsType
-    >>> from polyaxon.k8s import k8s_schemas
+    >>> from polyaxon.types import V1ArtifactsType
+    >>> from polyaxon import k8s
     >>> component = V1Component(
     >>>     run=V1Job(
     >>>        init=[
@@ -146,7 +146,7 @@ class V1ArtifactsType(BaseTypeConfig):
     >>>                 connection="s3-dataset"
     >>>             ),
     >>>        ],
-    >>>        container=k8s_schemas.V1Container(...)
+    >>>        container=k8s.V1Container(...)
     >>>     )
     >>> )
     ```

polyaxon/_schemas/types/dockerfile.py

@@ -111,8 +111,8 @@ class V1DockerfileType(BaseTypeConfig):
 
     ```python
     >>> from polyaxon.schemas import V1Component, V1Init, V1Job
-    >>> from polyaxon.schemas.types import V1DockerfileType
-    >>> from polyaxon.k8s import k8s_schemas
+    >>> from polyaxon.types import V1DockerfileType
+    >>> from polyaxon import k8s
     >>> component = V1Component(
     >>>     run=V1Job(
     >>>        init=[
@@ -124,7 +124,7 @@ class V1DockerfileType(BaseTypeConfig):
     >>>                 )
     >>>             ),
     >>>        ],
-    >>>        container=k8s_schemas.V1Container(...)
+    >>>        container=k8s.V1Container(...)
     >>>     )
     >>> )
     ```

polyaxon/_schemas/types/file.py

@@ -116,8 +116,8 @@ class V1FileType(BaseTypeConfig):
 
     ```python
     >>> from polyaxon.schemas import V1Component, V1Init, V1Job
-    >>> from polyaxon.schemas.types import V1FileType
-    >>> from polyaxon.k8s import k8s_schemas
+    >>> from polyaxon.types import V1FileType
+    >>> from polyaxon import k8s
     >>> component = V1Component(
     >>>     run=V1Job(
     >>>        init=[
@@ -129,7 +129,7 @@ class V1FileType(BaseTypeConfig):
     >>>                 )
     >>>             ),
     >>>        ],
-    >>>        container=k8s_schemas.V1Container(...)
+    >>>        container=k8s.V1Container(...)
     >>>     )
     >>> )
     ```

polyaxon/_schemas/types/git.py

@@ -103,8 +103,8 @@ class V1GitType(BaseTypeConfig):
     ### Usage in initializers
     ```python
     >>> from polyaxon.schemas import V1Component, V1Init, V1Job
-    >>> from polyaxon.schemas.types import V1GitType
-    >>> from polyaxon.k8s import k8s_schemas
+    >>> from polyaxon.types import V1GitType
+    >>> from polyaxon import k8s
     >>> component = V1Component(
     >>>     run=V1Job(
     >>>        init=[
@@ -116,7 +116,7 @@ class V1GitType(BaseTypeConfig):
     >>>               connection="my-git-connection",
     >>>             ),
     >>>        ],
-    >>>        container=k8s_schemas.V1Container(...)
+    >>>        container=k8s.V1Container(...)
     >>>     )
     >>> )
     ```

polyaxon/_schemas/types/tensorboard.py

@@ -99,8 +99,8 @@ class V1TensorboardType(BaseTypeConfig):
 
     ```python
     >>> from polyaxon.schemas import V1Component, V1Init, V1Job
-    >>> from polyaxon.schemas.types import V1FileType
-    >>> from polyaxon.k8s import k8s_schemas
+    >>> from polyaxon.types import V1FileType
+    >>> from polyaxon import k8s
     >>> component = V1Component(
     >>>     run=V1Job(
     >>>        init=[
@@ -114,7 +114,7 @@ class V1TensorboardType(BaseTypeConfig):
     >>>                 )
     >>>             ),
     >>>        ],
-    >>>        container=k8s_schemas.V1Container(...)
+    >>>        container=k8s.V1Container(...)
     >>>     )
     >>> )
     ```