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

truefoundry/common/auth_service_client.py

@@ -51,7 +51,7 @@ class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
         self._api_server_url = url
 
     def refresh_token(self, token: Token, host: Optional[str] = None) -> Token:
-        host_arg_str = f"--host {host}" if host else "--host HOST"
+        host_arg_str = host if host else "HOST"
         if not token.refresh_token:
             # TODO: Add a way to propagate error messages without traceback to the output interface side
             raise Exception(
@@ -128,7 +128,7 @@ class AuthServerServiceClient(AuthServiceClient):
         self._auth_server_url = url
 
     def refresh_token(self, token: Token, host: Optional[str] = None) -> Token:
-        host_arg_str = f"--host {host}" if host else "--host HOST"
+        host_arg_str = host if host else "HOST"
         if not token.refresh_token:
             # TODO: Add a way to propagate error messages without traceback to the output interface side
             raise Exception(

truefoundry/common/constants.py

@@ -1,3 +1,4 @@
+import enum
 import os
 from pathlib import Path
 from typing import Optional
@@ -15,6 +16,11 @@ TFY_INTERNAL_SIGNED_URL_SERVER_HOST_ENV_KEY = "TFY_INTERNAL_SIGNED_URL_SERVER_HO
 TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN_ENV_KEY = "TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN"
 
 
+class PythonPackageManager(str, enum.Enum):
+    PIP = "pip"
+    UV = "uv"
+
+
 class TrueFoundrySdkEnv(BaseSettings):
     # Note: Every field in this class should have a default value
     # Never expect the user to set these values
@@ -46,6 +52,9 @@ class TrueFoundrySdkEnv(BaseSettings):
     # For local development, this enables futher configuration via _TFYServersConfig
     TFY_CLI_LOCAL_DEV_MODE: bool = False
 
+    TFY_PYTHON_BUILD_PACKAGE_MANAGER: PythonPackageManager = PythonPackageManager.PIP
+    TFY_PYTHON_BUILD_UV_IMAGE_URI: str = "ghcr.io/astral-sh/uv:latest"
+
 
 ENV_VARS = TrueFoundrySdkEnv()
 API_SERVER_RELATIVE_PATH = "api/svc"

truefoundry/common/utils.py

@@ -1,7 +1,14 @@
+import json
+import logging
+import os
+import sys
 import time
+from collections import namedtuple
 from functools import lru_cache, wraps
+from shutil import rmtree
+from subprocess import check_output
 from time import monotonic_ns
-from typing import Callable, Generator, Optional, TypeVar
+from typing import Callable, Generator, List, Optional, TypeVar
 from urllib.parse import urljoin, urlparse
 
 from truefoundry.common.constants import (
@@ -12,7 +19,10 @@ from truefoundry.common.constants import (
 )
 from truefoundry.pydantic_v1 import BaseSettings
 
+logger = logging.getLogger(__name__)
+
 T = TypeVar("T")
+InstalledPipPackage = namedtuple("InstalledPipPackage", ["name", "version"])
 
 
 class _TFYServersConfig(BaseSettings):
@@ -102,3 +112,73 @@ def resolve_tfy_host(tfy_host: Optional[str] = None) -> str:
     tfy_host = tfy_host.strip("/")
     validate_tfy_host(tfy_host)
     return tfy_host
+
+
+class ContextualDirectoryManager:
+    def __init__(self, dir_path: str, cleanup_on_error: bool = True):
+        self.dir_path = dir_path
+        self.cleanup_on_error = cleanup_on_error
+
+    def __enter__(self):
+        if os.path.exists(self.dir_path):
+            raise FileExistsError(
+                f"The directory {self.dir_path!r} already exists. "
+                "Please provide a path with a different name that does not already exist."
+            )
+
+        os.makedirs(self.dir_path, exist_ok=False)
+        return self.dir_path
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.cleanup_on_error and exc_type is not None:
+            # Only delete the directory if an exception occurred
+            if os.path.exists(self.dir_path):
+                rmtree(self.dir_path)
+
+
+def get_python_version_major_minor() -> str:
+    """
+    Returns the major.minor version of the Python interpreter
+    """
+    version_info = sys.version_info
+    return f"{version_info.major}.{version_info.minor}"
+
+
+def list_pip_packages_installed(
+    filter_package_names: Optional[List[str]] = None,
+) -> List[InstalledPipPackage]:
+    """
+    List the installed package_names, along with their versions.
+    Args:
+        filter_package_names (List[str]): A list of specific libraries to filter for.
+
+    Returns:
+        List[InstalledPipPackage]: A list of InstalledPipPackage namedtuples for each match.
+    """
+    relevant_package_names: List[InstalledPipPackage] = []
+
+    # Get the installed packages in JSON format from pip
+    try:
+        output = check_output(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "--disable-pip-version-check",
+                "list",
+                "--pre",
+                "--format=json",
+            ]
+        )
+        installed_package_names = json.loads(output.decode("utf-8"))
+    except Exception:
+        logger.exception("Failed to list installed packages using pip.")
+        return relevant_package_names
+
+    package_names_to_check = set(filter_package_names or [])
+    for package in installed_package_names:
+        if package["name"] in package_names_to_check:
+            relevant_package_names.append(
+                InstalledPipPackage(package["name"], package["version"])
+            )
+    return relevant_package_names

truefoundry/deploy/__init__.py

@@ -1,8 +1,11 @@
 from truefoundry.deploy.auto_gen.models import (
     AppProtocol,
+    Autoshutdown,
     CanaryStep,
     CapacityType,
+    Claim,
     ConcurrencyPolicy,
+    JobAlert,
     Kustomize,
     ParamType,
     Protocol,
@@ -66,6 +69,7 @@ from truefoundry.deploy.v2.lib.patched_models import (
     DynamicVolumeConfig,
     Endpoint,
     GcpTPU,
+    GitHelmRepo,
     GitSource,
     GPUType,
     HealthProbe,
@@ -73,6 +77,7 @@ from truefoundry.deploy.v2.lib.patched_models import (
     HttpProbe,
     HuggingfaceArtifactSource,
     Image,
+    JwtAuthCreds,
     KafkaInputConfig,
     KafkaMetricConfig,
     KafkaOutputConfig,

truefoundry/deploy/builder/builders/tfy_notebook_buildpack/__init__.py

@@ -8,7 +8,7 @@ from truefoundry.deploy.builder.builders.tfy_notebook_buildpack.dockerfile_templ
     NotebookImageBuild,
     generate_dockerfile_content,
 )
-from truefoundry.deploy.builder.utils import has_pip_conf_secret
+from truefoundry.deploy.builder.utils import has_python_package_manager_conf_secret
 
 __all__ = ["generate_dockerfile_content", "build"]
 
@@ -38,7 +38,9 @@ def build(
     build_configuration: NotebookImageBuild,
     extra_opts: Optional[List[str]] = None,
 ):
-    mount_pip_conf_secret = has_pip_conf_secret(extra_opts) if extra_opts else False
+    mount_pip_conf_secret = (
+        has_python_package_manager_conf_secret(extra_opts) if extra_opts else False
+    )
     with TemporaryDirectory() as local_dir:
         docker_build_configuration = _convert_to_dockerfile_build_config(
             build_configuration,

truefoundry/deploy/builder/builders/tfy_python_buildpack/__init__.py

@@ -7,7 +7,7 @@ from truefoundry.deploy.builder.builders import dockerfile
 from truefoundry.deploy.builder.builders.tfy_python_buildpack.dockerfile_template import (
     generate_dockerfile_content,
 )
-from truefoundry.deploy.builder.utils import has_pip_conf_secret
+from truefoundry.deploy.builder.utils import has_python_package_manager_conf_secret
 
 __all__ = ["generate_dockerfile_content", "build"]
 
@@ -15,11 +15,11 @@ __all__ = ["generate_dockerfile_content", "build"]
 def _convert_to_dockerfile_build_config(
     build_configuration: PythonBuild,
     dockerfile_path: str,
-    mount_pip_conf_secret: bool = False,
+    mount_python_package_manager_conf_secret: bool = False,
 ) -> DockerFileBuild:
     dockerfile_content = generate_dockerfile_content(
         build_configuration=build_configuration,
-        mount_pip_conf_secret=mount_pip_conf_secret,
+        mount_python_package_manager_conf_secret=mount_python_package_manager_conf_secret,
     )
     with open(dockerfile_path, "w", encoding="utf8") as fp:
         fp.write(dockerfile_content)
@@ -36,12 +36,14 @@ def build(
     build_configuration: PythonBuild,
     extra_opts: Optional[List[str]] = None,
 ):
-    mount_pip_conf_secret = has_pip_conf_secret(extra_opts) if extra_opts else False
+    mount_python_package_manager_conf_secret = (
+        has_python_package_manager_conf_secret(extra_opts) if extra_opts else False
+    )
     with TemporaryDirectory() as local_dir:
         docker_build_configuration = _convert_to_dockerfile_build_config(
             build_configuration,
             dockerfile_path=os.path.join(local_dir, "Dockerfile"),
-            mount_pip_conf_secret=mount_pip_conf_secret,
+            mount_python_package_manager_conf_secret=mount_python_package_manager_conf_secret,
         )
         dockerfile.build(
             tag=tag,

truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py

@@ -1,13 +1,16 @@
 import os
+import shlex
 from typing import Dict, List, Optional
 
 from mako.template import Template
 
-from truefoundry.common.constants import ENV_VARS
+from truefoundry.common.constants import ENV_VARS, PythonPackageManager
 from truefoundry.deploy.auto_gen.models import PythonBuild
 from truefoundry.deploy.builder.constants import (
     PIP_CONF_BUILDKIT_SECRET_MOUNT,
     PIP_CONF_SECRET_MOUNT_AS_ENV,
+    UV_CONF_BUILDKIT_SECRET_MOUNT,
+    UV_CONF_SECRET_MOUNT_AS_ENV,
 )
 from truefoundry.deploy.v2.lib.patched_models import CUDAVersion
 
@@ -22,8 +25,8 @@ RUN ${apt_install_command}
 % if requirements_path is not None:
 COPY ${requirements_path} ${requirements_destination_path}
 % endif
-% if pip_install_command is not None:
-RUN ${pip_config_secret_mount} ${pip_install_command}
+% if python_packages_install_command is not None:
+RUN ${package_manager_config_secret_mount} ${python_packages_install_command}
 % endif
 COPY . /app
 WORKDIR /app
@@ -114,44 +117,91 @@ def generate_pip_install_command(
     mount_pip_conf_secret: bool = False,
 ) -> Optional[str]:
     upgrade_pip_command = "python -m pip install -U pip setuptools wheel"
-    final_pip_install_command = None
-    pip_install_base_command = "python -m pip install --use-pep517 --no-cache-dir"
+    envs = []
+    if mount_pip_conf_secret:
+        envs.append(PIP_CONF_SECRET_MOUNT_AS_ENV)
+
+    command = ["python", "-m", "pip", "install", "--use-pep517", "--no-cache-dir"]
+    args = []
     if requirements_path:
-        final_pip_install_command = f"{pip_install_base_command} -r {requirements_path}"
+        args.append("-r")
+        args.append(requirements_path)
 
     if pip_packages:
-        final_pip_install_command = (
-            final_pip_install_command or pip_install_base_command
-        )
-        final_pip_install_command += " " + " ".join(
-            f"'{package}'" for package in pip_packages
-        )
+        args.extend(pip_packages)
 
-    if not final_pip_install_command:
+    if not args:
         return None
 
-    if mount_pip_conf_secret:
-        final_pip_install_command = (
-            f"{PIP_CONF_SECRET_MOUNT_AS_ENV} {final_pip_install_command}"
-        )
+    final_pip_install_command = shlex.join(envs + command + args)
+    final_docker_run_command = " && ".join(
+        [upgrade_pip_command, final_pip_install_command]
+    )
+    return final_docker_run_command
 
-    return " && ".join([upgrade_pip_command, final_pip_install_command])
+
+def generate_uv_pip_install_command(
+    requirements_path: Optional[str],
+    pip_packages: Optional[List[str]],
+    mount_uv_conf_secret: bool = False,
+) -> Optional[str]:
+    upgrade_pip_command = "python -m pip install -U pip setuptools wheel"
+    uv_mount = f"--mount=from={ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_URI},source=/uv,target=/usr/local/bin/uv"
+    envs = [
+        "UV_SYSTEM_PYTHON=true",
+        "UV_LINK_MODE=copy",
+        "UV_PYTHON_DOWNLOADS=never",
+        "UV_INDEX_STRATEGY=unsafe-best-match",
+    ]
+    if mount_uv_conf_secret:
+        envs.append(UV_CONF_SECRET_MOUNT_AS_ENV)
+
+    command = ["uv", "pip", "install", "--no-cache-dir"]
+
+    args = []
+
+    if requirements_path:
+        args.append("-r")
+        args.append(requirements_path)
+
+    if pip_packages:
+        args.extend(pip_packages)
+
+    if not args:
+        return None
+
+    uv_pip_install_command = shlex.join(envs + command + args)
+    shell_commands = " && ".join([upgrade_pip_command, uv_pip_install_command])
+    final_docker_run_command = " ".join([uv_mount, shell_commands])
+
+    return final_docker_run_command
 
 
 def generate_dockerfile_content(
     build_configuration: PythonBuild,
-    mount_pip_conf_secret: bool = False,
+    package_manager: str = ENV_VARS.TFY_PYTHON_BUILD_PACKAGE_MANAGER,
+    mount_python_package_manager_conf_secret: bool = False,
 ) -> str:
     # TODO (chiragjn): Handle recursive references to other requirements files e.g. `-r requirements-gpu.txt`
     requirements_path = resolve_requirements_txt_path(build_configuration)
     requirements_destination_path = (
         "/tmp/requirements.txt" if requirements_path else None
     )
-    pip_install_command = generate_pip_install_command(
-        requirements_path=requirements_destination_path,
-        pip_packages=build_configuration.pip_packages,
-        mount_pip_conf_secret=mount_pip_conf_secret,
-    )
+    if package_manager == PythonPackageManager.PIP.value:
+        python_packages_install_command = generate_pip_install_command(
+            requirements_path=requirements_destination_path,
+            pip_packages=build_configuration.pip_packages,
+            mount_pip_conf_secret=mount_python_package_manager_conf_secret,
+        )
+    elif package_manager == PythonPackageManager.UV.value:
+        python_packages_install_command = generate_uv_pip_install_command(
+            requirements_path=requirements_destination_path,
+            pip_packages=build_configuration.pip_packages,
+            mount_uv_conf_secret=mount_python_package_manager_conf_secret,
+        )
+    else:
+        raise ValueError(f"Unsupported package manager: {package_manager}")
+
     apt_install_command = generate_apt_install_command(
         apt_packages=build_configuration.apt_packages
     )
@@ -161,13 +211,22 @@ def generate_dockerfile_content(
         "apt_install_command": apt_install_command,
         "requirements_path": requirements_path,
         "requirements_destination_path": requirements_destination_path,
-        "pip_install_command": pip_install_command,
+        "python_packages_install_command": python_packages_install_command,
     }
 
-    if mount_pip_conf_secret:
-        template_args["pip_config_secret_mount"] = PIP_CONF_BUILDKIT_SECRET_MOUNT
+    if mount_python_package_manager_conf_secret:
+        if package_manager == PythonPackageManager.PIP.value:
+            template_args["package_manager_config_secret_mount"] = (
+                PIP_CONF_BUILDKIT_SECRET_MOUNT
+            )
+        elif package_manager == PythonPackageManager.UV.value:
+            template_args["package_manager_config_secret_mount"] = (
+                UV_CONF_BUILDKIT_SECRET_MOUNT
+            )
+        else:
+            raise ValueError(f"Unsupported package manager: {package_manager}")
     else:
-        template_args["pip_config_secret_mount"] = ""
+        template_args["package_manager_config_secret_mount"] = ""
 
     if build_configuration.cuda_version:
         template = CUDA_DOCKERFILE_TEMPLATE

truefoundry/deploy/builder/constants.py

@@ -5,3 +5,11 @@ PIP_CONF_BUILDKIT_SECRET_MOUNT = (
 PIP_CONF_SECRET_MOUNT_AS_ENV = (
     f"PIP_CONFIG_FILE=/run/secrets/{BUILDKIT_SECRET_MOUNT_PIP_CONF_ID}"
 )
+
+BUILDKIT_SECRET_MOUNT_UV_CONF_ID = "uv.toml"
+UV_CONF_BUILDKIT_SECRET_MOUNT = (
+    f"--mount=type=secret,id={BUILDKIT_SECRET_MOUNT_UV_CONF_ID}"
+)
+UV_CONF_SECRET_MOUNT_AS_ENV = (
+    f"UV_CONFIG_FILE=/run/secrets/{BUILDKIT_SECRET_MOUNT_UV_CONF_ID}"
+)

truefoundry/deploy/builder/utils.py

@@ -1,6 +1,9 @@
 from typing import List, Optional
 
-from truefoundry.deploy.builder.constants import BUILDKIT_SECRET_MOUNT_PIP_CONF_ID
+from truefoundry.deploy.builder.constants import (
+    BUILDKIT_SECRET_MOUNT_PIP_CONF_ID,
+    BUILDKIT_SECRET_MOUNT_UV_CONF_ID,
+)
 
 
 def _get_id_from_buildkit_secret_value(value: str) -> Optional[str]:
@@ -19,14 +22,16 @@ def _get_id_from_buildkit_secret_value(value: str) -> Optional[str]:
     return None
 
 
-def has_pip_conf_secret(docker_build_extra_args: List[str]) -> bool:
+def has_python_package_manager_conf_secret(docker_build_extra_args: List[str]) -> bool:
     args = [arg.strip() for arg in docker_build_extra_args]
     for i, arg in enumerate(docker_build_extra_args):
         if (
             arg == "--secret"
             and i + 1 < len(args)
-            and _get_id_from_buildkit_secret_value(args[i + 1])
-            == BUILDKIT_SECRET_MOUNT_PIP_CONF_ID
+            and (
+                _get_id_from_buildkit_secret_value(args[i + 1])
+                in (BUILDKIT_SECRET_MOUNT_PIP_CONF_ID, BUILDKIT_SECRET_MOUNT_UV_CONF_ID)
+            )
         ):
             return True
     return False

truefoundry/deploy/cli/cli.py

@@ -9,6 +9,7 @@ from truefoundry.deploy.cli.commands import (
     get_build_command,
     get_delete_command,
     get_deploy_command,
+    get_deploy_init_command,
     get_login_command,
     get_logout_command,
     get_patch_application_command,
@@ -36,6 +37,7 @@ def create_truefoundry_cli() -> click.MultiCommand:
     cli.add_command(get_logout_command())
     cli.add_command(get_apply_command())
     cli.add_command(get_deploy_command())
+    cli.add_command(get_deploy_init_command())
     cli.add_command(get_patch_application_command())
     cli.add_command(get_delete_command())
     cli.add_command(get_trigger_command())

truefoundry/deploy/cli/commands/__init__.py

@@ -4,6 +4,7 @@ from truefoundry.deploy.cli.commands.build_logs_command import get_build_logs_co
 from truefoundry.deploy.cli.commands.create_command import get_create_command
 from truefoundry.deploy.cli.commands.delete_command import get_delete_command
 from truefoundry.deploy.cli.commands.deploy_command import get_deploy_command
+from truefoundry.deploy.cli.commands.deploy_init_command import get_deploy_init_command
 from truefoundry.deploy.cli.commands.get_command import get_get_command
 from truefoundry.deploy.cli.commands.list_command import get_list_command
 from truefoundry.deploy.cli.commands.login_command import get_login_command

truefoundry/deploy/cli/commands/deploy_init_command.py

@@ -0,0 +1,22 @@
+import rich_click as click
+
+from truefoundry.deploy.cli.const import GROUP_CLS
+from truefoundry.ml.cli.commands.model_init import get_model_init_command
+
+
+@click.group(
+    name="deploy-init",
+    cls=GROUP_CLS,
+    help="Initialize the TrueFoundry deployment configuration.",
+)
+def deploy_init_command(): ...
+
+
+def get_deploy_init_command():
+    """
+    Generates the deploy-init command.
+        model: Initialize the TrueFoundry deployment configuration for a model. eg: tfy deploy-init model [--args]
+    """
+
+    deploy_init_command.add_command(get_model_init_command())
+    return deploy_init_command

truefoundry/deploy/lib/dao/application.py

@@ -179,10 +179,11 @@ def trigger_job(
         command=command_str if command_str else None,
         params=params if params else None,
     )
+    jobRunName = result.jobRunName
     previous_runs_url = f"{client.base_url.strip('/')}/deployments/{application_info.id}?tab=previousRuns"
     logger.info(
         f"{message}.\n"
-        f"You can check the status of your job run at {previous_runs_url}"
+        f"You can check the status of your job run at {previous_runs_url} with jobRunName: {jobRunName}"
     )
     return result
 

truefoundry/deploy/v2/lib/patched_models.py

@@ -186,6 +186,10 @@ class BasicAuthCreds(models.BasicAuthCreds, PatchedModelBase):
     type: Literal["basic_auth"] = "basic_auth"
 
 
+class JwtAuthCreds(models.JwtAuthCreds, PatchedModelBase):
+    type: Literal["jwt_auth"] = "jwt_auth"
+
+
 class HealthProbe(models.HealthProbe, PatchedModelBase):
     pass
 
@@ -338,6 +342,10 @@ class HelmRepo(models.HelmRepo, PatchedModelBase):
     type: Literal["helm-repo"] = "helm-repo"
 
 
+class GitHelmRepo(models.GitHelmRepo, PatchedModelBase):
+    type: Literal["git-helm-repo"] = "git-helm-repo"
+
+
 class OCIRepo(models.OCIRepo, PatchedModelBase):
     type: Literal["oci-repo"] = "oci-repo"
 

truefoundry/ml/__init__.py

@@ -1,4 +1,9 @@
-from truefoundry.ml.autogen.client.models import LibraryName, ModelVersionEnvironment
+from truefoundry.ml.autogen.client.models import (
+    LibraryName,
+    ModelVersionEnvironment,
+    SklearnModelSchema,
+    XGBoostModelSchema,
+)
 from truefoundry.ml.enums import (
     DataSlice,
     FileFormat,
@@ -14,9 +19,7 @@ from truefoundry.ml.log_types.artifacts.artifact import (
     BlobStorageDirectory,
 )
 from truefoundry.ml.log_types.artifacts.dataset import DataDirectory, DataDirectoryPath
-from truefoundry.ml.log_types.artifacts.model import (
-    ModelVersion,
-)
+from truefoundry.ml.log_types.artifacts.model import ModelVersion
 from truefoundry.ml.logger import init_logger
 from truefoundry.ml.mlfoundry_api import get_client
 from truefoundry.ml.mlfoundry_run import MlFoundryRun
@@ -36,6 +39,8 @@ from truefoundry.ml.model_framework import (
     TensorFlowFramework,
     TransformersFramework,
     XGBoostFramework,
+    sklearn_infer_schema,
+    xgboost_infer_schema,
 )
 
 __all__ = [
@@ -45,33 +50,37 @@ __all__ = [
     "DataDirectory",
     "DataDirectoryPath",
     "DataSlice",
+    "FastAIFramework",
     "FileFormat",
+    "GluonFramework",
+    "H2OFramework",
     "Image",
-    "MlFoundryRun",
+    "KerasFramework",
+    "LibraryName",
+    "LightGBMFramework",
     "MlFoundryException",
-    "ModelVersionEnvironment",
+    "MlFoundryRun",
     "ModelFramework",
+    "ModelFrameworkType",
     "ModelType",
     "ModelVersion",
-    "Plot",
-    "ViewType",
-    "get_client",
-    "FastAIFramework",
-    "GluonFramework",
-    "H2OFramework",
-    "KerasFramework",
-    "LightGBMFramework",
+    "ModelVersionEnvironment",
     "ONNXFramework",
     "PaddleFramework",
+    "Plot",
     "PyTorchFramework",
     "SklearnFramework",
+    "SklearnModelSchema",
     "SpaCyFramework",
     "StatsModelsFramework",
     "TensorFlowFramework",
     "TransformersFramework",
+    "ViewType",
     "XGBoostFramework",
-    "LibraryName",
-    "ModelFrameworkType",
+    "XGBoostModelSchema",
+    "get_client",
+    "sklearn_infer_schema",
+    "xgboost_infer_schema",
 ]
 
 init_logger()