truefoundry 0.4.0rc1__py3-none-any.whl → 0.4.0rc3__py3-none-any.whl

truefoundry/ml/mlfoundry_api.py

@@ -15,23 +15,22 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     CreateExperimentRequestDto,
     CreateRunRequestDto,
     DatasetDto,
+    ExperimentsApi,
     ListArtifactsRequestDto,
     ListArtifactVersionsRequestDto,
     ListDatasetsRequestDto,
     ListModelVersionsRequestDto,
+    MlfoundryArtifactsApi,
     ModelDto,
+    RunsApi,
     RunTagDto,
     SearchRunsRequestDto,
 )
-from truefoundry.ml.autogen.client.api import (  # type: ignore[attr-defined]
-    ExperimentsApi,
-    MlfoundryArtifactsApi,
-    RunsApi,
-)
 from truefoundry.ml.autogen.client.exceptions import (
     ApiException,
     NotFoundException,
 )
+from truefoundry.ml.clients.servicefoundry_client import ServiceFoundryServiceClient
 from truefoundry.ml.enums import ModelFramework, ViewType
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.internal_namespace import NAMESPACE
@@ -39,10 +38,8 @@ from truefoundry.ml.log_types.artifacts.artifact import ArtifactPath, ArtifactVe
 from truefoundry.ml.log_types.artifacts.dataset import DataDirectory
 from truefoundry.ml.log_types.artifacts.general_artifact import _log_artifact_version
 from truefoundry.ml.log_types.artifacts.model import ModelVersion, _log_model_version
-from truefoundry.ml.log_types.artifacts.model_extras import CustomMetric, ModelSchema
 from truefoundry.ml.logger import logger
 from truefoundry.ml.mlfoundry_run import MlFoundryRun
-from truefoundry.ml.services.servicefoundry_service import ServicefoundryService
 from truefoundry.ml.session import (
     Session,
     _get_api_client,
@@ -126,7 +123,6 @@ class MlFoundry:
             )
             raise MlFoundryException(err_msg) from e
 
-        assert ml_repo_instance.experiment_id is not None
         return ml_repo_instance.experiment_id
 
     def list_ml_repos(self) -> List[str]:
@@ -216,7 +212,7 @@ class MlFoundry:
                 "No active session found. Perhaps you are not logged in?\n"
                 "Please log in using `tfy login [--host HOST] --relogin"
             )
-        servicefoundry_service = ServicefoundryService(
+        servicefoundry_client = ServiceFoundryServiceClient(
             tracking_uri=self.get_tracking_uri(),
             token=session.token.access_token,
         )
@@ -224,7 +220,7 @@ class MlFoundry:
         assert existing_ml_repo.storage_integration_id is not None
         try:
             existing_storage_integration = (
-                servicefoundry_service.get_integration_from_id(
+                servicefoundry_client.get_integration_from_id(
                     existing_ml_repo.storage_integration_id
                 )
             )
@@ -338,11 +334,6 @@ class MlFoundry:
             )
         )
         run = _run.run
-
-        assert run is not None
-        assert run.info.run_id is not None
-        assert run.info.fqn is not None
-
         mlf_run_id = run.info.run_id
         kwargs.setdefault("auto_end", True)
         mlf_run = MlFoundryRun(experiment_id=ml_repo_id, run_id=mlf_run_id, **kwargs)
@@ -381,9 +372,6 @@ class MlFoundry:
             return self.get_run_by_fqn(run_id)
         _run = self._runs_api.get_run_get(run_id=run_id)
         run = _run.run
-
-        assert run is not None
-
         mlfoundry_run = MlFoundryRun._from_dto(run)
         logger.info(
             f"Link to the dashboard for the run: {mlfoundry_run.dashboard_link}"
@@ -725,7 +713,6 @@ class MlFoundry:
             version=resolved_version,
         )
         model_version = _model_version.model_version
-        assert model_version.model_id is not None
         _model = self._mlfoundry_artifacts_api.get_model_get(id=model_version.model_id)
         model = _model.model
 
@@ -868,7 +855,6 @@ class MlFoundry:
         model: Optional[ModelDto] = None,
     ) -> Iterator[ModelVersion]:
         if model and not model_id:
-            assert model.id is not None
             model_id = model.id
         elif not model and model_id:
             _model = self._mlfoundry_artifacts_api.get_model_get(id=str(model_id))
@@ -943,7 +929,6 @@ class MlFoundry:
             )
         )
         artifact_version = _artifact_version.artifact_version
-        assert artifact_version.artifact_id is not None
         _artifact = self._mlfoundry_artifacts_api.get_artifact_by_id_get(
             id=artifact_version.artifact_id
         )
@@ -1072,7 +1057,6 @@ class MlFoundry:
         artifact: Optional[ArtifactDto] = None,
     ) -> Iterator[ArtifactVersion]:
         if artifact and not artifact_id:
-            assert artifact.id is not None
             artifact_id = artifact.id
         elif not artifact and artifact_id:
             _artifact = self._mlfoundry_artifacts_api.get_artifact_by_id_get(
@@ -1220,8 +1204,7 @@ class MlFoundry:
         additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
         description: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
-        model_schema: Optional[Union[ModelSchema, Dict[str, Any]]] = None,
-        custom_metrics: Optional[List[Union[CustomMetric, Dict[str, Any]]]] = None,
+        progress: Optional[bool] = None,
     ) -> ModelVersion:
         """
         Serialize and log a versioned model under the current ml_repo. Each logged model generates a new version
@@ -1277,34 +1260,7 @@ class MlFoundry:
             metadata (Optional[Dict[str, Any]], optional): arbitrary json serializable dictionary to store metadata.
                 For example, you can use this to store metrics, params, notes.
                 This field can be updated at any time after logging. Defaults to `None`
-            model_schema (Optional[Union[Dict[str, Any], ModelSchema]], optional):
-                instance of `truefoundry.ml.ModelSchema`.
-                This schema needs to be consistent with older versions of the model under the given `name` i.e.
-                a feature's value type and model's prediction type cannot be changed in the schema of new version.
-                Features can be removed or added between versions.
-                ```
-                E.g. if there exists a v1 with
-                    schema = {"features": {"name": "feat1": "int"}, "prediction": "categorical"}, then
-
-                    schema = {"features": {"name": "feat1": "string"}, "prediction": "categorical"} or
-                    schema = {"features": {"name": "feat1": "int"}, "prediction": "numerical"}
-                    are invalid because they change the types of existing features and prediction
-
-                    while
-                    schema = {"features": {"name": "feat1": "int", "feat2": "string"}, "prediction": "categorical"} or
-                    schema = {"features": {"feat2": "string"}, "prediction": "categorical"}
-                    are valid
-
-                    This field can be updated at any time after logging. Defaults to None
-                ```
-            custom_metrics: (Optional[Union[List[Dict[str, Any]], CustomMetric]], optional): list of instances of
-                `truefoundry.ml.CustomMetric`
-                The custom metrics must be added according to the prediction type of schema.
-                custom_metrics = [{
-                    "name": "mean_square_error",
-                    "type": "metric",
-                    "value_type": "float"
-                }]
+            progress (bool): value to show progress bar, defaults to None.
 
         Returns:
             truefoundry.ml.ModelVersion: an instance of `ModelVersion` that can be used to download the files,
@@ -1394,9 +1350,8 @@ class MlFoundry:
             additional_files=additional_files,
             description=description,
             metadata=metadata,
-            model_schema=model_schema,
-            custom_metrics=custom_metrics,
             step=None,
+            progress=progress,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version
@@ -1415,7 +1370,7 @@ class MlFoundry:
         Args:
             ml_repo (str): Name of the ML Repo in which you want to create data_directory
             name (str): Name of the DataDirectory to be created.
-            description (str): Description of the Datset
+            description (str): Description for the DataDirectory.
             metadata (Dict <str>: Any): Metadata about the data_directory in Dictionary form.
 
         Returns:

truefoundry/ml/mlfoundry_run.py

@@ -38,6 +38,7 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     RunTagDto,
     UpdateRunRequestDto,
 )
+from truefoundry.ml.entities import Metric
 from truefoundry.ml.enums import RunStatus
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.internal_namespace import NAMESPACE
@@ -45,7 +46,6 @@ from truefoundry.ml.log_types import Image, Plot
 from truefoundry.ml.log_types.artifacts.artifact import ArtifactPath, ArtifactVersion
 from truefoundry.ml.log_types.artifacts.general_artifact import _log_artifact_version
 from truefoundry.ml.log_types.artifacts.model import ModelVersion, _log_model_version
-from truefoundry.ml.log_types.artifacts.model_extras import CustomMetric, ModelSchema
 from truefoundry.ml.logger import logger
 from truefoundry.ml.run_utils import ParamsType, flatten_dict, process_params
 from truefoundry.ml.session import ACTIVE_RUNS, _get_api_client, get_active_session
@@ -110,8 +110,6 @@ class MlFoundryRun:
     @classmethod
     def _from_dto(cls, run_dto: RunDto) -> "MlFoundryRun":
         """classmethod to get MLfoundry run from dto instance"""
-        assert run_dto.info.experiment_id is not None
-        assert run_dto.info.run_id is not None
         run = cls(run_dto.info.experiment_id, run_dto.info.run_id)
         run._run_info = run_dto.info
         run._run_data = run_dto.data
@@ -136,7 +134,6 @@ class MlFoundryRun:
     def run_name(self) -> str:
         """Get run_name for the current `run`"""
         run_info = self._get_run_info()
-        assert run_info.name is not None
         return run_info.name
 
     @property
@@ -144,7 +141,6 @@ class MlFoundryRun:
     def fqn(self) -> str:
         """Get fqn for the current `run`"""
         run_info = self._get_run_info()
-        assert run_info.fqn is not None
         return run_info.fqn
 
     @property
@@ -162,7 +158,6 @@ class MlFoundryRun:
         _experiment = self._experiments_api.get_experiment_get(
             experiment_id=self._experiment_id
         )
-        assert _experiment.experiment.name is not None
         return _experiment.experiment.name
 
     @property
@@ -812,7 +807,7 @@ class MlFoundryRun:
     @_ensure_not_deleted
     def get_metrics(
         self, metric_names: Optional[Iterable[str]] = None
-    ) -> Dict[str, List[MetricDto]]:
+    ) -> Dict[str, List[Metric]]:
         """Get metrics logged for the current `run` grouped by metric name.
 
         Args:
@@ -860,7 +855,7 @@ class MlFoundryRun:
         unknown_metrics = metric_names - run_metric_names
         if len(unknown_metrics) > 0:
             logger.warning(f"{unknown_metrics} metrics not present in the run")
-        metrics_dict: Dict[str, List[MetricDto]] = {
+        metrics_dict: Dict[str, List[Metric]] = {
             metric_name: [] for metric_name in unknown_metrics
         }
         valid_metrics = metric_names - unknown_metrics
@@ -868,7 +863,9 @@ class MlFoundryRun:
             _metric_history = self._metrics_api.get_metric_history_get(
                 run_id=self.run_id, metric_key=metric_name
             )
-            metrics_dict[metric_name] = _metric_history.metrics
+            metrics_dict[metric_name] = [
+                Metric.from_dto(metric) for metric in _metric_history.metrics
+            ]
         return metrics_dict
 
     @_ensure_not_deleted
@@ -911,8 +908,6 @@ class MlFoundryRun:
         additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
         description: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
-        model_schema: Optional[Union[ModelSchema, Dict[str, Any]]] = None,
-        custom_metrics: Optional[List[Union[CustomMetric, Dict[str, Any]]]] = None,
         step: int = 0,
         progress: Optional[bool] = None,
     ) -> ModelVersion:
@@ -965,34 +960,6 @@ class MlFoundryRun:
             metadata (Optional[Dict[str, Any]], optional): arbitrary json serializable dictionary to store metadata.
                 For example, you can use this to store metrics, params, notes.
                 This field can be updated at any time after logging. Defaults to `None`
-            model_schema (Optional[Union[Dict[str, Any], ModelSchema]], optional):
-                instance of `truefoundry.ml.ModelSchema`.
-                This schema needs to be consistent with older versions of the model under the given `name` i.e.
-                a feature's value type and model's prediction type cannot be changed in the schema of new version.
-                Features can be removed or added between versions.
-                ```
-                E.g. if there exists a v1 with
-                    schema = {"features": {"name": "feat1": "int"}, "prediction": "categorical"}, then
-
-                    schema = {"features": {"name": "feat1": "string"}, "prediction": "categorical"} or
-                    schema = {"features": {"name": "feat1": "int"}, "prediction": "numerical"}
-                    are invalid because they change the types of existing features and prediction
-
-                    while
-                    schema = {"features": {"name": "feat1": "int", "feat2": "string"}, "prediction": "categorical"} or
-                    schema = {"features": {"feat2": "string"}, "prediction": "categorical"}
-                    are valid
-
-                    This field can be updated at any time after logging. Defaults to None
-                ```
-            custom_metrics: (Optional[Union[List[Dict[str, Any]], CustomMetric]], optional): list of instances of
-                `truefoundry.ml.CustomMetric`
-                The custom metrics must be added according to the prediction type of schema.
-                custom_metrics = [{
-                    "name": "mean_square_error",
-                    "type": "metric",
-                    "value_type": "float"
-                }]
             step (int): step/iteration at which the model is being logged, defaults to 0.
             progress (bool): value to show progress bar, defaults to None.
 
@@ -1084,8 +1051,6 @@ class MlFoundryRun:
             additional_files=additional_files,
             description=description,
             metadata=metadata,
-            model_schema=model_schema,
-            custom_metrics=custom_metrics,
             step=step,
             progress=progress,
         )

truefoundry/ml/run_utils.py

@@ -7,7 +7,6 @@ from urllib.parse import urljoin, urlsplit
 
 import numpy as np
 
-from truefoundry.ml import env_vars
 from truefoundry.ml.exceptions import MlFoundryException
 
 
@@ -22,21 +21,13 @@ def get_module(
             raise MlFoundryException(msg) from ex
 
 
-def resolve_tracking_uri(tracking_uri: Optional[str]):
-    if not tracking_uri and not os.getenv(env_vars.TRACKING_HOST_GLOBAL):
-        raise ValueError(
-            f"Either `host` should be provided by --host <value>, or `{env_vars.TRACKING_HOST_GLOBAL}` env must be set"
-        )
-    return tracking_uri or os.getenv(env_vars.TRACKING_HOST_GLOBAL)
-
-
 def append_path_to_rest_tracking_uri(tracking_uri: str):
     if urlsplit(tracking_uri).netloc.startswith("localhost"):
         return tracking_uri
     return urljoin(tracking_uri, "/api/ml")
 
 
-def append_servicefoundry_path_to_tracking_ui(tracking_uri: str):
+def append_servicefoundry_path_to_tracking_uri(tracking_uri: str):
     if urlsplit(tracking_uri).netloc.startswith("localhost"):
         return os.getenv("SERVICEFOUNDRY_SERVER_URL")
     return urljoin(tracking_uri, "/api/svc")

truefoundry/ml/session.py

@@ -1,131 +1,29 @@
-import abc
 import atexit
-import os
 import threading
 import weakref
 from typing import TYPE_CHECKING, Dict, Optional
 
+from truefoundry.common.credential_provider import (
+    CredentialProvider,
+    EnvCredentialProvider,
+    FileCredentialProvider,
+)
+from truefoundry.common.entities import Token, UserInfo
 from truefoundry.common.request_utils import urllib3_retry
-from truefoundry.ml import env_vars
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ApiClient,
     Configuration,
 )
+from truefoundry.ml.clients.entities import HostCreds
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.logger import logger
-from truefoundry.ml.login import CredentialsFileContent, CredentialsFileManager
 from truefoundry.ml.run_utils import (
     append_path_to_rest_tracking_uri,
-    resolve_tracking_uri,
 )
-from truefoundry.ml.services.auth_service import get_auth_service
-from truefoundry.ml.services.entities import HostCreds, Token, UserInfo
 
 if TYPE_CHECKING:
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
 
-TOKEN_REFRESH_LOCK = threading.RLock()
-
-
-class CredentialProvider(abc.ABC):
-    @property
-    @abc.abstractmethod
-    def token(self) -> Token: ...
-
-    @staticmethod
-    @abc.abstractmethod
-    def can_provide() -> bool: ...
-
-    @property
-    @abc.abstractmethod
-    def tracking_uri(self) -> str: ...
-
-
-class EnvCredentialProvider(CredentialProvider):
-    def __init__(self):
-        logger.debug("Using env var credential provider")
-        self._tracking_uri = resolve_tracking_uri(tracking_uri=None)
-        self._auth_service = get_auth_service(tracking_uri=self._tracking_uri)
-        api_key = os.getenv(env_vars.API_KEY_GLOBAL)
-        if not api_key:
-            raise MlFoundryException(
-                f"Value of {env_vars.API_KEY_GLOBAL} env var should be non-empty string"
-            )
-        self._token: Token = Token(access_token=api_key, refresh_token=None)  # type: ignore[call-arg]
-
-    @staticmethod
-    def can_provide() -> bool:
-        return env_vars.API_KEY_GLOBAL in os.environ
-
-    @property
-    def token(self) -> Token:
-        with TOKEN_REFRESH_LOCK:
-            if self._token.is_going_to_be_expired():
-                logger.info("Refreshing access token")
-                self._token = self._auth_service.refresh_token(self._token)
-            return self._token
-
-    @property
-    def tracking_uri(self) -> str:
-        return self._tracking_uri
-
-
-class FileCredentialProvider(CredentialProvider):
-    def __init__(self):
-        logger.debug("Using file credential provider")
-        self._cred_file = CredentialsFileManager()
-
-        with self._cred_file:
-            self._last_cred_file_content = self._cred_file.read()
-            self._tracking_uri = self._last_cred_file_content.host
-            self._token = self._last_cred_file_content.to_token()
-            self._auth_service = get_auth_service(tracking_uri=self._tracking_uri)
-
-    @staticmethod
-    def can_provide() -> bool:
-        with CredentialsFileManager() as cred_file:
-            return cred_file.exists()
-
-    @property
-    def token(self) -> Token:
-        with TOKEN_REFRESH_LOCK:
-            if not self._token.is_going_to_be_expired():
-                return self._token
-
-            logger.info("Refreshing access token")
-            with self._cred_file:
-                new_cred_file_content = self._cred_file.read()
-                new_token = new_cred_file_content.to_token()
-                new_tracking_uri = new_cred_file_content.host
-
-                if new_cred_file_content == self._last_cred_file_content:
-                    self._token = self._auth_service.refresh_token(self._token)
-                    self._last_cred_file_content = CredentialsFileContent(
-                        host=self._tracking_uri,
-                        access_token=self._token.access_token,
-                        refresh_token=self._token.refresh_token,
-                    )
-                    self._cred_file.write(self._last_cred_file_content)
-                    return self._token
-
-                if (
-                    new_tracking_uri == self._tracking_uri
-                    and new_token.to_user_info() == self._token.to_user_info()
-                ):
-                    self._last_cred_file_content = new_cred_file_content
-                    self._token = new_token
-                    # recursive
-                    return self.token
-
-                raise MlFoundryException(
-                    "Credentials on disk changed while mlfoundry was running."
-                )
-
-    @property
-    def tracking_uri(self) -> str:
-        return self._tracking_uri
-
-
 SESSION_LOCK = threading.RLock()
 
 
@@ -159,8 +57,8 @@ class Session:
     def __init__(self, cred_provider: CredentialProvider):
         # Note: Whenever a new session is initialized all the active runs are ended
         self._closed = False
-        self._cred_provider: CredentialProvider = cred_provider
-        self._user_info: UserInfo = self._cred_provider.token.to_user_info()
+        self._cred_provider: Optional[CredentialProvider] = cred_provider
+        self._user_info: Optional[UserInfo] = self._cred_provider.token.to_user_info()
 
     def close(self):
         logger.debug("Closing existing session")
@@ -187,7 +85,7 @@ class Session:
 
     @property
     def tracking_uri(self) -> str:
-        return self._cred_provider.tracking_uri
+        return self._cred_provider.base_url
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, Session):
@@ -199,9 +97,7 @@ class Session:
         )
 
     def get_host_creds(self) -> HostCreds:
-        tracking_uri = append_path_to_rest_tracking_uri(
-            self._cred_provider.tracking_uri
-        )
+        tracking_uri = append_path_to_rest_tracking_uri(self._cred_provider.base_url)
         return HostCreds(
             host=tracking_uri, token=self._cred_provider.token.access_token
         )
@@ -235,6 +131,7 @@ def _get_api_client(
     )
     configuration.retries = urllib3_retry(retries=2)
     api_client = ApiClient(configuration=configuration)
+    api_client.user_agent = "truefoundry-cli"
     return api_client
 
 
@@ -247,8 +144,8 @@ def init_session() -> Session:
                 break
         if final_cred_provider is None:
             raise MlFoundryException(
-                "Please login using `mlfoundry login` command "
-                "or `truefoundry.ml.login()` function call"
+                "Please login using `tfy login` command "
+                "or `truefoundry.login()` function call"
             )
         new_session = Session(cred_provider=final_cred_provider)
 

truefoundry/pydantic_v1.py

@@ -5,7 +5,7 @@ try:
 except ImportError:
     # pydantic <=1.10.17
     from pydantic import *  # noqa: F403
-    from pydantic import utils  # noqa: F401
+    from pydantic import ConstrainedStr, utils  # noqa: F401
 
 
 class NonEmptyStr(ConstrainedStr):

truefoundry/workflow/__init__.py

@@ -1,5 +1,5 @@
 try:
-    from flytekit import task
+    from flytekit import task as _
 except ImportError:
     print("To use workflows, please run 'pip install truefoundry[workflow]'.")
 
@@ -17,3 +17,18 @@ from truefoundry.workflow.map_task import map_task
 from truefoundry.workflow.python_task import PythonFunctionTask
 from truefoundry.workflow.task import task
 from truefoundry.workflow.workflow import ExecutionConfig, workflow
+
+__all__ = [
+    "task",
+    "ContainerTask",
+    "PythonFunctionTask",
+    "map_task",
+    "workflow",
+    "conditional",
+    "FlyteDirectory",
+    "TaskDockerFileBuild",
+    "TaskPythonBuild",
+    "ContainerTaskConfig",
+    "PythonTaskConfig",
+    "ExecutionConfig",
+]

{truefoundry-0.4.0rc1.dist-info → truefoundry-0.4.0rc3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.4.0rc1
+Version: 0.4.0rc3
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com