truefoundry 0.5.1rc7__py3-none-any.whl → 0.5.2__py3-none-any.whl

truefoundry/deploy/cli/commands/list_command.py

@@ -1,9 +1,9 @@
 import rich_click as click
 
-from truefoundry.deploy.cli.config import CliConfig
-from truefoundry.deploy.cli.const import COMMAND_CLS, GROUP_CLS
-from truefoundry.deploy.cli.display_util import print_entity_list, print_json
-from truefoundry.deploy.cli.util import handle_exception_wrapper
+from truefoundry.cli.config import CliConfig
+from truefoundry.cli.const import COMMAND_CLS, GROUP_CLS
+from truefoundry.cli.display_util import print_entity_list, print_json
+from truefoundry.cli.util import handle_exception_wrapper
 from truefoundry.deploy.io.rich_output_callback import RichOutputCallBack
 from truefoundry.deploy.lib.dao import application as application_lib
 from truefoundry.deploy.lib.dao import version as version_lib

truefoundry/deploy/cli/commands/login_command.py

@@ -1,11 +1,11 @@
 import rich_click as click
 
-from truefoundry.common.constants import TFY_HOST_ENV_KEY
-from truefoundry.deploy.cli.const import COMMAND_CLS
-from truefoundry.deploy.cli.util import (
+from truefoundry.cli.const import COMMAND_CLS
+from truefoundry.cli.util import (
     _prompt_if_no_value_and_supported,
     handle_exception_wrapper,
 )
+from truefoundry.common.constants import TFY_HOST_ENV_KEY
 from truefoundry.deploy.io.rich_output_callback import RichOutputCallBack
 from truefoundry.deploy.lib.session import login
 

truefoundry/deploy/cli/commands/logout_command.py

@@ -1,7 +1,7 @@
 import rich_click as click
 
-from truefoundry.deploy.cli.const import COMMAND_CLS
-from truefoundry.deploy.cli.util import handle_exception_wrapper
+from truefoundry.cli.const import COMMAND_CLS
+from truefoundry.cli.util import handle_exception_wrapper
 from truefoundry.deploy.io.rich_output_callback import RichOutputCallBack
 from truefoundry.deploy.lib.session import logout
 

truefoundry/deploy/cli/commands/logs_command.py

@@ -4,8 +4,8 @@ from typing import Optional
 import rich_click as click
 from dateutil.tz import tzlocal
 
-from truefoundry.deploy.cli.const import COMMAND_CLS
-from truefoundry.deploy.cli.util import handle_exception_wrapper
+from truefoundry.cli.const import COMMAND_CLS
+from truefoundry.cli.util import handle_exception_wrapper
 from truefoundry.deploy.io.rich_output_callback import RichOutputCallBack
 from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,

truefoundry/deploy/cli/commands/patch_application_command.py

@@ -3,8 +3,8 @@ import json
 import rich_click as click
 import yaml
 
-from truefoundry.deploy.cli.const import GROUP_CLS
-from truefoundry.deploy.cli.util import handle_exception_wrapper
+from truefoundry.cli.const import GROUP_CLS
+from truefoundry.cli.util import handle_exception_wrapper
 from truefoundry.deploy.lib.dao import application as application_lib
 
 

truefoundry/deploy/cli/commands/patch_command.py

@@ -3,8 +3,8 @@ from typing import Optional
 
 import rich_click as click
 
-from truefoundry.deploy.cli.const import GROUP_CLS
-from truefoundry.deploy.cli.util import handle_exception_wrapper
+from truefoundry.cli.const import GROUP_CLS
+from truefoundry.cli.util import handle_exception_wrapper
 from truefoundry.deploy.io.rich_output_callback import RichOutputCallBack
 
 

truefoundry/deploy/cli/commands/redeploy_command.py

@@ -1,7 +1,7 @@
 import rich_click as click
 
-from truefoundry.deploy.cli.const import COMMAND_CLS
-from truefoundry.deploy.cli.util import handle_exception_wrapper
+from truefoundry.cli.const import COMMAND_CLS
+from truefoundry.cli.util import handle_exception_wrapper
 from truefoundry.deploy.lib.dao import application as application_lib
 
 

truefoundry/deploy/cli/commands/terminate_comand.py

@@ -1,8 +1,8 @@
 import rich_click as click
 
-from truefoundry.deploy.cli.const import COMMAND_CLS, GROUP_CLS
-from truefoundry.deploy.cli.display_util import print_json
-from truefoundry.deploy.cli.util import handle_exception_wrapper
+from truefoundry.cli.const import COMMAND_CLS, GROUP_CLS
+from truefoundry.cli.display_util import print_json
+from truefoundry.cli.util import handle_exception_wrapper
 from truefoundry.deploy.lib.dao import application as application_lib
 
 

truefoundry/deploy/cli/commands/trigger_command.py

@@ -3,8 +3,8 @@ from typing import Optional, Sequence
 import rich_click as click
 from click import ClickException
 
-from truefoundry.deploy.cli.const import COMMAND_CLS, GROUP_CLS
-from truefoundry.deploy.cli.util import handle_exception_wrapper
+from truefoundry.cli.const import COMMAND_CLS, GROUP_CLS
+from truefoundry.cli.util import handle_exception_wrapper
 from truefoundry.deploy.lib.dao import application
 
 

truefoundry/deploy/lib/clients/servicefoundry_client.py

@@ -680,9 +680,9 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         return parse_obj_as(List[Deployment], response)
 
     @check_min_cli_version
-    def apply(self, manifest: Dict[str, Any]):
+    def apply(self, manifest: Dict[str, Any], dry_run: bool = False):
         url = f"{self._api_server_url}/{VERSION_PREFIX}/apply"
-        body = {"manifest": manifest}
+        body = {"manifest": manifest, "dryRun": dry_run}
         response = session_with_retries().put(
             url, headers=self._get_header(), json=body
         )

truefoundry/deploy/lib/const.py

@@ -2,6 +2,3 @@ from truefoundry.deploy.io.rich_output_callback import RichOutputCallBack
 
 ENTITY_JSON_DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
 RICH_OUTPUT_CALLBACK = RichOutputCallBack()
-
-TFY_DEBUG_ENV_KEY = "TFY_DEBUG"
-TFY_INTERNAL_ENV_KEY = "TFY_INTERNAL"

truefoundry/deploy/lib/dao/apply.py

@@ -15,6 +15,7 @@ def _apply_manifest(
     client: Optional[ServiceFoundryServiceClient] = None,
     filename: Optional[str] = None,
     index: Optional[int] = None,
+    dry_run: bool = False,
 ) -> ApplyResult:
     client = client or ServiceFoundryServiceClient()
 
@@ -32,29 +33,38 @@ def _apply_manifest(
             message=f"Failed to apply manifest{file_metadata}. Error: {ex}",
         )
 
+    prefix = "[Dry Run] " if dry_run else ""
+    suffix = " (No changes were applied)" if dry_run else ""
     try:
-        client.apply(manifest.dict())
+        client.apply(manifest.dict(), dry_run)
+
         return ApplyResult(
             success=True,
-            message=f"Successfully configured manifest {manifest.name} of type {manifest.type}",
+            message=(
+                f"{prefix}Successfully configured manifest {manifest.name} of type {manifest.type}.{suffix}"
+            ),
         )
     except Exception as ex:
         return ApplyResult(
             success=False,
-            message=f"Failed to apply manifest {manifest.name} of type {manifest.type}. Error: {ex}",
+            message=(
+                f"{prefix}Failed to apply manifest {manifest.name} of type {manifest.type}. Error: {ex}.{suffix}"
+            ),
         )
 
 
 def apply_manifest(
     manifest: Dict[str, Any],
     client: Optional[ServiceFoundryServiceClient] = None,
+    dry_run: bool = False,
 ) -> ApplyResult:
-    return _apply_manifest(manifest=manifest, client=client)
+    return _apply_manifest(manifest=manifest, client=client, dry_run=dry_run)
 
 
 def apply_manifest_file(
     filepath: str,
     client: Optional[ServiceFoundryServiceClient] = None,
+    dry_run: bool = False,
 ) -> Iterator[ApplyResult]:
     client = client or ServiceFoundryServiceClient()
     filename = Path(filepath).name
@@ -67,14 +77,19 @@ def apply_manifest_file(
             message=f"Failed to read file {filepath} as a valid YAML file. Error: {ex}",
         )
     else:
+        prefix = "[Dry Run] " if dry_run else ""
         for index, manifest in enumerate(manifests_it):
             if not isinstance(manifest, dict):
                 yield ApplyResult(
                     success=False,
-                    message=f"Failed to apply manifest at index {index} from file {filename}. Error: A manifest must be a dict, got type {type(manifest)}",
+                    message=f"{prefix}Failed to apply manifest at index {index} from file {filename}. Error: A manifest must be a dict, got type {type(manifest)}",
                 )
                 continue
 
             yield _apply_manifest(
-                manifest=manifest, client=client, filename=filename, index=index
+                manifest=manifest,
+                client=client,
+                filename=filename,
+                index=index,
+                dry_run=dry_run,
             )

truefoundry/deploy/lib/dao/workspace.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from truefoundry.deploy.cli.console import console
+from truefoundry.cli.console import console
 from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )

truefoundry/deploy/lib/model/entity.py

@@ -220,7 +220,7 @@ class JobRun(Base):
     status: str
 
     def list_row_data(self) -> Dict[str, Any]:
-        from truefoundry.deploy.cli.display_util import display_time_passed
+        from truefoundry.cli.display_util import display_time_passed
 
         triggered_at = (
             (datetime.datetime.now().timestamp() * 1000) - self.createdAt
@@ -238,7 +238,7 @@ class JobRun(Base):
         }
 
     def get_data(self) -> Dict[str, Any]:
-        from truefoundry.deploy.cli.display_util import display_time_passed
+        from truefoundry.cli.display_util import display_time_passed
 
         created_at = datetime.datetime.fromtimestamp(self.createdAt // 1000)
         end_time = ""

truefoundry/deploy/lib/util.py

@@ -1,20 +1,6 @@
-import os
 import re
 from typing import Union
 
-from truefoundry.deploy.lib.const import (
-    TFY_DEBUG_ENV_KEY,
-    TFY_INTERNAL_ENV_KEY,
-)
-
-
-def is_debug_env_set() -> bool:
-    return bool(os.getenv(TFY_DEBUG_ENV_KEY))
-
-
-def is_internal_env_set() -> bool:
-    return bool(os.getenv(TFY_INTERNAL_ENV_KEY))
-
 
 def get_application_fqn_from_deployment_fqn(deployment_fqn: str) -> str:
     if not re.search(r":\d+$", deployment_fqn):

truefoundry/ml/__init__.py

@@ -1,9 +1,10 @@
-from truefoundry.ml.autogen.client.models import (
-    LibraryName,
+from truefoundry.ml.autogen.client.models import (  # type: ignore[attr-defined]
+    InferMethodName,
     ModelVersionEnvironment,
     SklearnModelSchema,
     XGBoostModelSchema,
 )
+from truefoundry.ml.autogen.entities.artifacts import LibraryName
 from truefoundry.ml.enums import (
     DataSlice,
     FileFormat,
@@ -40,6 +41,7 @@ from truefoundry.ml.model_framework import (
     TransformersFramework,
     XGBoostFramework,
     sklearn_infer_schema,
+    xgboost_infer_schema,
 )
 
 __all__ = [
@@ -51,9 +53,11 @@ __all__ = [
     "DataSlice",
     "FastAIFramework",
     "FileFormat",
+    "get_client",
     "GluonFramework",
     "H2OFramework",
     "Image",
+    "InferMethodName",
     "KerasFramework",
     "LibraryName",
     "LightGBMFramework",
@@ -68,6 +72,7 @@ __all__ = [
     "PaddleFramework",
     "Plot",
     "PyTorchFramework",
+    "sklearn_infer_schema",
     "SklearnFramework",
     "SklearnModelSchema",
     "SpaCyFramework",
@@ -75,10 +80,9 @@ __all__ = [
     "TensorFlowFramework",
     "TransformersFramework",
     "ViewType",
+    "xgboost_infer_schema",
     "XGBoostFramework",
     "XGBoostModelSchema",
-    "get_client",
-    "sklearn_infer_schema",
 ]
 
 init_logger()

truefoundry/ml/cli/cli.py

@@ -1,11 +1,15 @@
 import rich_click as click
 
+from truefoundry.cli.const import GROUP_CLS
 from truefoundry.ml.cli.commands import download
 
 click.rich_click.USE_RICH_MARKUP = True
 
 
-@click.group()
+@click.group(
+    name="ml",
+    cls=GROUP_CLS,
+)
 def ml():
     """
     TrueFoundry ML CLI

truefoundry/ml/cli/commands/download.py

@@ -2,14 +2,23 @@ from typing import Optional
 
 import click
 
+from truefoundry.cli.const import COMMAND_CLS, GROUP_CLS
 from truefoundry.ml import get_client
 
 
-@click.group(help="Download files of artifact/model logged with Mlfoundry")
+@click.group(
+    name="download",
+    cls=GROUP_CLS,
+    help="Download artifact/model versions logged with TrueFoundry",
+)
 def download(): ...
 
 
-@download.command(short_help="Download files of logged model")
+@download.command(
+    name="model",
+    cls=COMMAND_CLS,
+    help="Download a model version logged with TrueFoundry",
+)
 @click.option(
     "--fqn",
     required=True,
@@ -48,7 +57,11 @@ def model(fqn: str, path: str, overwrite: bool, progress: Optional[bool] = None)
     print(f"Downloaded model files to {download_path}")
 
 
-@download.command(short_help="Download files of logged artifact")
+@download.command(
+    name="artifact",
+    cls=COMMAND_CLS,
+    short_help="Download a artifact version logged with TrueFoundry",
+)
 @click.option(
     "--fqn",
     required=True,

truefoundry/ml/cli/commands/model_init.py

@@ -3,10 +3,10 @@ from typing import Optional
 
 import rich_click as click
 
-from truefoundry.deploy.cli.console import console
-from truefoundry.deploy.cli.const import COMMAND_CLS
-from truefoundry.deploy.cli.util import handle_exception_wrapper
-from truefoundry.ml.autogen.client.models import ModelServer
+from truefoundry.cli.console import console
+from truefoundry.cli.const import COMMAND_CLS
+from truefoundry.cli.util import handle_exception_wrapper
+from truefoundry.ml.autogen.client import ModelServer  # type: ignore[attr-defined]
 from truefoundry.ml.cli.utils import (
     AppName,
     NonEmptyString,

truefoundry/ml/mlfoundry_api.py

@@ -18,7 +18,7 @@ from typing import (
 import coolname
 
 from truefoundry.common.utils import ContextualDirectoryManager, relogin_error_message
-from truefoundry.ml import ModelVersionEnvironment, constants
+from truefoundry.ml import constants
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactDto,
     ArtifactType,
@@ -35,6 +35,7 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     MlfoundryArtifactsApi,
     ModelDto,
     ModelServer,
+    ModelVersionEnvironment,
     RunsApi,
     RunTagDto,
     SearchRunsRequestDto,

truefoundry/ml/mlfoundry_run.py

@@ -18,7 +18,7 @@ from urllib.parse import urljoin, urlsplit
 
 from truefoundry import version
 from truefoundry.common.utils import relogin_error_message
-from truefoundry.ml import ModelVersionEnvironment, constants
+from truefoundry.ml import constants
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
     DeleteRunRequest,
@@ -29,6 +29,7 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     MetricDto,
     MetricsApi,
     MlfoundryArtifactsApi,
+    ModelVersionEnvironment,
     ParamDto,
     RunDataDto,
     RunDto,

truefoundry/ml/model_framework.py

@@ -35,6 +35,7 @@ from truefoundry.pydantic_v1 import BaseModel, Field
 
 if TYPE_CHECKING:
     from sklearn.base import BaseEstimator
+    from xgboost import Booster, XGBModel
 
 # Map serialization format to corresponding pip packages
 SERIALIZATION_FORMAT_TO_PACKAGES_NAME_MAP = {
@@ -467,6 +468,24 @@ def auto_update_model_framework_details(
                 )
 
 
+def _infer_schema(
+    model_input: Any,
+    model: Union["BaseEstimator", "Booster", "XGBModel"],
+    infer_method_name: str = "predict",
+) -> Dict[str, Any]:
+    if not hasattr(model, infer_method_name):
+        raise ValueError(
+            f"Model does not have the method '{infer_method_name}' to infer the schema."
+        )
+    model_infer_method = getattr(model, infer_method_name)
+    model_output = model_infer_method(model_input)
+
+    model_signature = infer_signature(
+        model_input=model_input, model_output=model_output
+    )
+    return model_signature.to_dict()
+
+
 def sklearn_infer_schema(
     model_input: Any,
     model: "BaseEstimator",
@@ -483,19 +502,37 @@ def sklearn_infer_schema(
     Returns:
         SklearnModelSchema: The inferred schema of the Sklearn model.
     """
-    if not hasattr(model, infer_method_name):
-        raise ValueError(
-            f"Model does not have the method '{infer_method_name}' to infer the schema."
-        )
-    model_infer_method = getattr(model, infer_method_name)
-    model_output = model_infer_method(model_input)
-
-    model_signature = infer_signature(
-        model_input=model_input, model_output=model_output
+    model_signature_json = _infer_schema(
+        model_input=model_input, model=model, infer_method_name=infer_method_name
     )
-    model_signature_json = model_signature.to_dict()
     return autogen_artifacts.SklearnModelSchema(
         infer_method_name=infer_method_name,
         inputs=json.loads(model_signature_json["inputs"]),
         outputs=json.loads(model_signature_json["outputs"]),
     )
+
+
+def xgboost_infer_schema(
+    model_input: Any,
+    model: Union["Booster", "XGBModel"],
+    infer_method_name: str = "predict",
+) -> autogen_artifacts.XGBoostModelSchema:
+    """
+    Infer the schema of an XGBoost model.
+
+    Args:
+        model_input (Any): The input data to be used for schema inference.
+        model (Any): The XGBoost model instance.
+        infer_method_name (str): The name of the method to be used for schema inference.
+            Eg: predict (default), predict_proba
+    Returns:
+        XGBoostModelSchema: The inferred schema of the XGBoost model.
+    """
+    model_signature_json = _infer_schema(
+        model_input=model_input, model=model, infer_method_name=infer_method_name
+    )
+    return autogen_artifacts.XGBoostModelSchema(
+        infer_method_name=infer_method_name,
+        inputs=json.loads(model_signature_json["inputs"]),
+        outputs=json.loads(model_signature_json["outputs"]),
+    )

{truefoundry-0.5.1rc7.dist-info → truefoundry-0.5.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.5.1rc7
+Version: 0.5.2
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com