polyaxon 2.1.0rc9__py3-none-any.whl → 2.6.0__py3-none-any.whl

polyaxon/_client/project.py

@@ -4,7 +4,7 @@ from datetime import datetime
 from requests import HTTPError
 from typing import Dict, List, Optional, Tuple, Union
 
-from clipped.utils.json import orjson_dumps
+from clipped.utils.json import orjson_dumps, orjson_loads
 from clipped.utils.paths import check_or_create_path, delete_path
 from clipped.utils.query_params import get_query_params
 from clipped.utils.tz import now
@@ -12,6 +12,7 @@ from clipped.utils.validation import validate_tags
 
 from polyaxon._client.client import PolyaxonClient
 from polyaxon._client.decorators import client_handler, get_global_or_inline_config
+from polyaxon._client.mixin import ClientMixin
 from polyaxon._constants.globals import DEFAULT
 from polyaxon._contexts import paths as ctx_paths
 from polyaxon._env_vars.getters.user import get_local_owner
@@ -21,13 +22,17 @@ from polyaxon._sdk.schemas.v1_list_project_versions_response import (
 )
 from polyaxon._sdk.schemas.v1_project import V1Project
 from polyaxon._sdk.schemas.v1_project_version import V1ProjectVersion
-from polyaxon._utils.fqn_utils import get_entity_full_name, get_entity_info
+from polyaxon._utils.fqn_utils import (
+    get_entity_full_name,
+    get_entity_info,
+    split_owner_team_space,
+)
 from polyaxon.exceptions import ApiException, PolyaxonClientException
 from polyaxon.logger import logger
 from traceml.artifacts import V1RunArtifact
 
 
-class ProjectClient:
+class ProjectClient(ClientMixin):
     """ProjectClient is a client to communicate with Polyaxon projects endpoints.
 
     If no values are passed to this class,
@@ -94,26 +99,13 @@ class ProjectClient:
         if not owner:
             raise PolyaxonClientException("Please provide a valid owner.")
 
+        owner, team = split_owner_team_space(owner)
         self._client = client
         self._owner = owner or DEFAULT
+        self._team = team
         self._project = project
         self._project_data = V1Project.construct()
 
-    @property
-    def client(self):
-        if self._client:
-            return self._client
-        self._client = PolyaxonClient()
-        return self._client
-
-    @property
-    def owner(self):
-        return self._owner
-
-    @property
-    def project(self):
-        return self._project
-
     @property
     def project_data(self):
         return self._project_data
@@ -139,11 +131,19 @@ class ProjectClient:
         Returns:
             V1Project, project instance from the response.
         """
-        self._project_data = self.client.projects_v1.create_project(
-            self.owner,
-            data,
-            async_req=False,
-        )
+        if self.team:
+            self._project_data = self.client.projects_v1.create_team_project(
+                self.owner,
+                self.team,
+                data,
+                async_req=False,
+            )
+        else:
+            self._project_data = self.client.projects_v1.create_project(
+                self.owner,
+                data,
+                async_req=False,
+            )
         self._project_data.owner = self.owner
         self._project = self._project_data.name
         return self._project_data
@@ -727,9 +727,12 @@ class ProjectClient:
         description: Optional[str] = None,
         tags: Optional[Union[str, List[str]]] = None,
         content: Optional[Union[str, Dict]] = None,
+        readme: Optional[str] = None,
         run: Optional[str] = None,
         connection: Optional[str] = None,
         artifacts: Optional[List[str]] = None,
+        stage: Optional[V1Stages] = None,
+        stage_conditions: Optional[List[V1StageCondition]] = None,
         force: bool = False,
     ) -> V1ProjectVersion:
         """Creates or Updates a project version based on the data passed.
@@ -745,9 +748,12 @@ class ProjectClient:
             description: str, optional, the version description.
             tags: str or List[str], optional.
             content: str or dict, optional, content/metadata (JSON object) of the version.
+            readme: str, optional, the version readme.
             run: str, optional, a uuid reference to the run.
             connection: str, optional, a uuid reference to a connection.
             artifacts: List[str], optional, list of artifacts to highlight(requires passing a run)
+            stage: V1Stages, optional, the version stage.
+            stage_conditions: List[V1StageCondition], optional, the version stage conditions.
             force: bool, optional, to force push, i.e. update if exists.
 
         Returns:
@@ -783,12 +789,18 @@ class ProjectClient:
                 version_config.tags = tags
             if content:
                 version_config.content = content  # type: ignore
+            if readme is not None:
+                version_config.readme = readme
             if run:
                 version_config.run = run
             if artifacts is not None:
                 version_config.artifacts = artifacts
             if connection is not None:
                 version_config.connection = connection
+            if stage is not None:
+                version_config.stage = stage
+            if stage_conditions is not None:
+                version_config.stage_conditions = stage_conditions
             return self.patch_version(
                 kind=kind,
                 version=version,
@@ -800,9 +812,12 @@ class ProjectClient:
                 description=description,
                 tags=tags,
                 run=run,
+                readme=readme,
                 artifacts=artifacts,
                 connection=connection,
                 content=content,
+                stage=stage,
+                stage_conditions=stage_conditions,
             )
             return self.create_version(kind=kind, data=version_config)
 
@@ -1575,3 +1590,183 @@ class ProjectClient:
             path=path,
             download_artifacts=download_artifacts,
         )
+
+    @classmethod
+    @client_handler(check_no_op=True)
+    def load_offline_version(
+        cls,
+        kind: V1ProjectVersionKind,
+        version: str,
+        path: str,
+        project_client: Optional["ProjectClient"] = None,
+        reset_project: bool = False,
+        raise_if_not_found: bool = False,
+    ) -> Optional["ProjectClient"]:
+        """Loads a project version from a local path.
+
+        Args:
+            kind: V1ProjectVersionKind, kind of the project version.
+            version: str, required, the version name/tag.
+            path: str, local path where to load the version's metadata and artifacts from.
+            project_client: ProjectClient, optional,
+                 a project client to update with the loaded version.
+            reset_project: bool, optional, to reset the project client with the loaded version.
+            raise_if_not_found: bool, optional, to raise an exception if the version is not found.
+
+        Returns:
+            ProjectClient, a project client with the loaded version.
+        """
+        path = ctx_paths.get_offline_path(
+            entity_value=version, entity_kind=kind, path=path
+        )
+        version_path = "{}/{}".format(path, ctx_paths.CONTEXT_LOCAL_VERSION)
+        if not os.path.exists(version_path):
+            if raise_if_not_found:
+                raise PolyaxonClientException(
+                    "Version not found in the provided path: {}".format(path)
+                )
+            else:
+                logger.info(f"Offline data was not found: {version_path}")
+                return None
+
+        with open(version_path, "r") as config_file:
+            config_str = config_file.read()
+            version_config = V1ProjectVersion(**orjson_loads(config_str))
+            owner = version_config.owner
+            project = version_config.project
+            if project_client:
+                if reset_project or not owner:
+                    owner = project_client.owner
+                if reset_project or not project:
+                    project = project_client.project
+                project_client._owner = owner
+                project_client._project = project
+            else:
+                project_client = cls(
+                    owner=owner,
+                    project=project,
+                )
+            logger.info("Loaded version `{}`".format(version_path))
+
+        return project_client
+
+    @client_handler(check_no_op=True, check_offline=True)
+    def push_version(
+        self,
+        kind: V1ProjectVersionKind,
+        version: str,
+        path: str,
+        force: bool = False,
+        clean: bool = False,
+    ):
+        """Pushes a local version from a local path to Polyaxon's API.
+
+        This is a generic function based on the kind passed and pushes a:
+            * component version
+            * model version
+            * artifact version
+
+        Args:
+            kind: V1ProjectVersionKind, kind of the project version.
+            version: str, required, the version name/tag.
+            path: str, optional, defaults to the offline root path,
+                 path where to load the metadata and artifacts from.
+            force: bool, optional, to force push, i.e. update if exists.
+            clean: bool, optional, to clean the version after pushing.
+        """
+        path = ctx_paths.get_offline_path(
+            entity_value=version, entity_kind=kind, path=path
+        )
+        version_path = "{}/{}".format(path, ctx_paths.CONTEXT_LOCAL_VERSION)
+        with open(version_path, "r") as config_file:
+            config_str = config_file.read()
+            version_config = V1ProjectVersion(**orjson_loads(config_str))
+
+        self.register_version(
+            kind=kind,
+            version=version,
+            description=version_config.description,
+            tags=version_config.tags,
+            content=version_config.content,
+            readme=version_config.readme,
+            run=version_config.run,
+            stage=version_config.stage,
+            stage_conditions=version_config.stage_conditions,
+            connection=version_config.connection,
+            artifacts=version_config.artifacts,
+            force=force,
+        )
+        if clean:
+            delete_path(path)
+
+    @client_handler(check_no_op=True, check_offline=True)
+    def push_component_version(
+        self,
+        version: str,
+        path: str,
+        force: bool = False,
+        clean: bool = False,
+    ):
+        """Pushes a local component version to a remove server.
+
+        Args:
+            version: str, required, the version name/tag.
+            path: str, local path where to load the metadata and artifacts from.
+            force: bool, optional, to force push, i.e. update if exists.
+            clean: bool, optional, to clean the version after pushing.
+        """
+        return self.push_version(
+            kind=V1ProjectVersionKind.COMPONENT,
+            version=version,
+            path=path,
+            force=force,
+            clean=clean,
+        )
+
+    @client_handler(check_no_op=True, check_offline=True)
+    def push_model_version(
+        self,
+        version: str,
+        path: str,
+        force: bool = True,
+        clean: bool = False,
+    ):
+        """Pushes a local model version to a remove server.
+
+        Args:
+            version: str, required, the version name/tag.
+            path: str, local path where to load the metadata and artifacts from.
+            force: bool, optional, to force push, i.e. update if exists.
+            clean: bool, optional, to clean the version after pushing.
+        """
+        return self.pull_version(
+            kind=V1ProjectVersionKind.MODEL,
+            version=version,
+            path=path,
+            force=force,
+            clean=clean,
+        )
+
+    @client_handler(check_no_op=True, check_offline=True)
+    def push_artifact_version(
+        self,
+        version: str,
+        path: str,
+        force: bool = True,
+        clean: bool = False,
+    ):
+        """Pushes a local artifact version to a remote server.
+
+        Args:
+            version: str, required, the version name/tag.
+            path: str, local path where to load the metadata and artifacts from.
+            force: bool, optional, to force push, i.e. update if exists.
+            clean: bool, optional, to clean the version after pushing.
+        """
+        return self.pull_version(
+            kind=V1ProjectVersionKind.ARTIFACT,
+            version=version,
+            path=path,
+            force=force,
+            clean=clean,
+        )

polyaxon/_client/run.py

@@ -21,6 +21,7 @@ from clipped.utils.paths import (
     get_base_filename,
     get_dirs_under_path,
     get_files_in_path_context,
+    set_permissions,
 )
 from clipped.utils.query_params import (
     get_logs_params,
@@ -35,6 +36,7 @@ from polyaxon import settings
 from polyaxon._cli.errors import handle_cli_error
 from polyaxon._client.client import PolyaxonClient
 from polyaxon._client.decorators import client_handler, get_global_or_inline_config
+from polyaxon._client.mixin import ClientMixin
 from polyaxon._client.store import PolyaxonStore
 from polyaxon._constants.metadata import META_COPY_ARTIFACTS, META_RECOMPILE
 from polyaxon._containers.names import MAIN_CONTAINER_NAMES
@@ -63,7 +65,12 @@ from polyaxon._sdk.schemas.v1_operation_body import V1OperationBody
 from polyaxon._sdk.schemas.v1_project_version import V1ProjectVersion
 from polyaxon._sdk.schemas.v1_run import V1Run
 from polyaxon._sdk.schemas.v1_run_settings import V1RunSettings
-from polyaxon._utils.fqn_utils import get_entity_full_name, to_fqn_name
+from polyaxon._sdk.schemas.v1_settings_catalog import V1SettingsCatalog
+from polyaxon._utils.fqn_utils import (
+    get_entity_full_name,
+    split_owner_team_space,
+    to_fqn_name,
+)
 from polyaxon._utils.urls_utils import get_proxy_run_url
 from polyaxon.api import K8S_V1_LOCATION, STREAMS_V1_LOCATION
 from polyaxon.exceptions import ApiException, PolyaxonClientException
@@ -80,7 +87,7 @@ if TYPE_CHECKING:
     from traceml.tracking.run import Run
 
 
-class RunClient:
+class RunClient(ClientMixin):
     """RunClient is a client to communicate with Polyaxon runs endpoints.
 
     If no values are passed to this class,
@@ -146,7 +153,7 @@ class RunClient:
             return
 
         try:
-            owner, project = get_project_or_local(
+            owner, _, project = get_project_or_local(
                 get_entity_full_name(owner=owner, entity=project)
             )
         except PolyaxonClientException:
@@ -166,8 +173,10 @@ class RunClient:
         if error_message and not self._is_offline:
             raise PolyaxonClientException(error_message)
 
+        owner, team = split_owner_team_space(owner)
         self._client = client
         self._owner = owner
+        self._team = team
         self._project = project
         self._run_uuid = (
             get_run_or_local(run_uuid)
@@ -218,12 +227,23 @@ class RunClient:
             return client.config.no_op
         return settings.CLIENT_CONFIG.no_op
 
-    @property
-    def client(self):
-        if self._client:
-            return self._client
-        self._client = PolyaxonClient()
-        return self._client
+    def _reset_agent(self, agent: Optional[str] = None):
+        if agent:
+            agent = self.client.agents_v1.get_agent(self.owner, agent)
+            self.settings.agent = V1SettingsCatalog(
+                uuid=agent.uuid,
+                name=agent.name,
+                version=agent.version,
+                url=agent.hostname,
+            )
+            self.settings.namespace = agent.namespace
+            self.settings.artifacts_store = None
+
+    def _use_agent_host(self):
+        if self.settings.agent and self.settings.agent.url:
+            self.reset_client(
+                host=self.settings.agent.url, POLYAXON_HOST=self.settings.agent.url
+            )
 
     @property
     def store(self):
@@ -256,20 +276,6 @@ class RunClient:
             return self.settings.artifacts_store.name
         return None
 
-    @property
-    def owner(self) -> str:
-        return self._owner or ""
-
-    def set_owner(self, owner: str):
-        self._owner = owner
-
-    @property
-    def project(self) -> str:
-        return self._project or ""
-
-    def set_project(self, project: str):
-        self._project = project
-
     @property
     def run_uuid(self) -> str:
         return self._run_uuid
@@ -898,6 +904,7 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
         params = get_logs_params(
             last_file=last_file, last_time=last_time, connection=self.artifacts_store
         )
@@ -921,6 +928,7 @@ class RunClient:
     def inspect(self):
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
         params = get_streams_params(connection=self.artifacts_store, status=self.status)
         return self.client.runs_v1.inspect_run(
             self.namespace, self.owner, self.project, self.run_uuid, **params
@@ -994,6 +1002,10 @@ class RunClient:
                 if not container:
                     container = pod_containers[0]
 
+        if not self.settings:
+            self.refresh_data()
+        self._use_agent_host()
+
         url = get_proxy_run_url(
             service=K8S_V1_LOCATION,
             namespace=self.namespace,
@@ -1039,6 +1051,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         params = get_streams_params(self.artifacts_store)
         return self.client.runs_v1.get_run_events(
             self.namespace,
@@ -1074,6 +1088,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         params = get_streams_params(self.artifacts_store)
         return self.client.runs_v1.get_multi_run_events(
             self.namespace,
@@ -1310,6 +1326,7 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
         params = get_streams_params(self.artifacts_store)
         return self.client.runs_v1.get_run_artifact(
             namespace=self.namespace,
@@ -1343,6 +1360,10 @@ class RunClient:
         if not self.run_uuid:
             return
 
+        if not self.settings:
+            self.refresh_data()
+        self._use_agent_host()
+
         lineage_path = lineage.path or ""
         summary = lineage.summary or {}
         is_event = summary.get("is_event")
@@ -1421,6 +1442,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         url = get_proxy_run_url(
             service=STREAMS_V1_LOCATION,
             namespace=self.namespace,
@@ -1460,6 +1483,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         url = get_proxy_run_url(
             service=STREAMS_V1_LOCATION,
             namespace=self.namespace,
@@ -1491,6 +1516,7 @@ class RunClient:
         untar: bool = False,
         overwrite: bool = True,
         show_progress: bool = True,
+        agent: Optional[str] = None,
     ):
         """Uploads a single artifact to the run's artifacts store path.
 
@@ -1501,12 +1527,17 @@ class RunClient:
                  it should be decompressed on the artifacts store.
             overwrite: bool, optional, if the file uploaded should overwrite any previous content.
             show_progress: bool, to show a progress bar.
+            agent: str, optional, uuid reference of an agent to use.
 
         Returns:
             str
         """
         if not self.settings:
             self.refresh_data()
+        if agent:
+            self._reset_agent(agent)
+
+        self._use_agent_host()
 
         params = get_streams_params(connection=self.artifacts_store)
         url = get_proxy_run_url(
@@ -1536,6 +1567,7 @@ class RunClient:
         path: str = "",
         overwrite: bool = True,
         relative_to: Optional[str] = None,
+        agent: Optional[str] = None,
     ):
         """Uploads a full directory to the run's artifacts store path.
 
@@ -1548,7 +1580,7 @@ class RunClient:
             overwrite: bool, optional, if the file uploaded should overwrite any previous content.
             relative_to: str, optional, if the path uploaded is not the current dir,
                  and you want to cancel the relative path.
-
+            agent: str, optional, uuid reference of an agent to use.
         Returns:
             str.
         """
@@ -1568,6 +1600,7 @@ class RunClient:
             path=path or "",
             overwrite=overwrite,
             relative_to=relative_to,
+            agent=agent,
         )
 
     @client_handler(check_no_op=True, check_offline=True)
@@ -1577,6 +1610,7 @@ class RunClient:
         path: str = "",
         overwrite: bool = True,
         relative_to: Optional[str] = None,
+        agent: Optional[str] = None,
     ):
         """Uploads multiple artifacts to the run's artifacts store path.
 
@@ -1586,7 +1620,7 @@ class RunClient:
             overwrite: bool, optional, if the file uploaded should overwrite any previous content.
             relative_to: str, optional, if the path uploaded is not the current dir,
                  and you want to cancel the relative path.
-
+            agent: str, optional, uuid reference of an agent to use.
         Returns:
             str.
         """
@@ -1596,6 +1630,10 @@ class RunClient:
 
         if not self.settings:
             self.refresh_data()
+        if agent:
+            self._reset_agent(agent)
+
+        self._use_agent_host()
 
         params = get_streams_params(connection=self.artifacts_store)
         url = get_proxy_run_url(
@@ -1626,6 +1664,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         params = get_streams_params(connection=self.artifacts_store)
         self.client.runs_v1.delete_run_artifact(
             namespace=self.namespace,
@@ -1645,6 +1685,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         params = get_streams_params(connection=self.artifacts_store)
         return self.client.runs_v1.delete_run_artifacts(
             namespace=self.namespace,
@@ -1667,6 +1709,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         params = get_streams_params(connection=self.artifacts_store)
         return self.client.runs_v1.get_run_artifacts_tree(
             namespace=self.namespace,
@@ -2117,9 +2161,9 @@ class RunClient:
         abspath = filepath if is_abs() else os.path.abspath(filepath)
 
         for_patterns = []
-        if getattr(self, "_artifacts_path"):
+        if getattr(self, "_artifacts_path", None):
             for_patterns.append(getattr(self, "_artifacts_path"))
-        if getattr(self, "_store_path"):
+        if getattr(self, "_store_path", None):
             for_patterns.append(getattr(self, "_store_path"))
         context_root = (
             ctx_paths.CONTEXT_OFFLINE_ROOT
@@ -2853,6 +2897,7 @@ class RunClient:
             config_file.write(
                 orjson_dumps(self.client.sanitize_for_serialization(self.run_data))
             )
+        set_permissions(run_path)
 
         if not self._artifacts_lineage:
             logger.debug("Persist offline run call did not find any lineage data. ")
@@ -2868,6 +2913,7 @@ class RunClient:
                     ]
                 )
             )
+        set_permissions(lineages_path)
 
     @classmethod
     @client_handler(check_no_op=True)
@@ -2876,6 +2922,8 @@ class RunClient:
         path: str,
         run_client: Optional[Union["RunClient", "Run"]] = None,
         reset_project: bool = False,
+        reset_uuid: bool = False,
+        name: Optional[str] = None,
         raise_if_not_found: bool = False,
     ) -> Optional[Union["RunClient", "Run"]]:
         """Loads an offline run from a local path.
@@ -2892,6 +2940,8 @@ class RunClient:
                  from the local run.
             raise_if_not_found: bool, optional, a flag to raise an error if the local path does not
                  contain a persisted run.
+            reset_uuid: bool, optional, a flag to reset the run's uuid.
+            name: str, optional, a name to set for the run.
         """
         run_path = "{}/{}".format(path, ctx_paths.CONTEXT_LOCAL_RUN)
         if not os.path.isfile(run_path):
@@ -2906,6 +2956,10 @@ class RunClient:
             run_config = V1Run(**orjson_loads(config_str))
             owner = run_config.owner
             project = run_config.project
+            if reset_uuid:
+                run_config.uuid = uuid.uuid4().hex
+            if name:
+                run_config.name = name
             if run_client:
                 if reset_project or not owner:
                     owner = run_client.owner
@@ -2964,6 +3018,7 @@ class RunClient:
         path: str,
         upload_artifacts: bool = True,
         clean: bool = False,
+        agent: Optional[str] = None,
     ):
         """Syncs an offline run to Polyaxon's API and artifacts store.
 
@@ -2972,6 +3027,7 @@ class RunClient:
             path: str, root path where the run's metadata & artifacts are stored.
             upload_artifacts: bool, optional, flag to trigger artifacts upload.
             clean: bool, optional, flag to clean local path after pushing the run.
+            agent: str, optional, uuid reference of an agent to use.
         """
         # We ensure that the is_offline is False
         is_offline = self._is_offline
@@ -3004,6 +3060,7 @@ class RunClient:
                 path="/",
                 overwrite=True,
                 relative_to=path,
+                agent=agent,
             )
             logger.info(f"Offline artifacts for run {self.run_data.uuid} uploaded")