truefoundry 0.3.3__py3-none-any.whl → 0.4.0.dev1__py3-none-any.whl

truefoundry/ml/run_utils.py

@@ -0,0 +1,102 @@
+import argparse
+import importlib
+import json
+import os
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
+from urllib.parse import urljoin, urlsplit
+
+import numpy as np
+
+from truefoundry.ml import env_vars
+from truefoundry.ml.exceptions import MlFoundryException
+
+
+def get_module(
+    module_name: str, error_message: Optional[str] = None, required: bool = False
+):
+    try:
+        return importlib.import_module(module_name)
+    except Exception as ex:
+        msg = error_message or f"Error importing module {module_name}"
+        if required:
+            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):
+    if urlsplit(tracking_uri).netloc.startswith("localhost"):
+        return os.getenv("SERVICEFOUNDRY_SERVER_URL")
+    return urljoin(tracking_uri, "/api/svc")
+
+
+class NumpyEncoder(json.JSONEncoder):
+    """Special json encoder for numpy types"""
+
+    def default(self, obj):
+        if isinstance(obj, np.integer):
+            return int(obj)
+        elif isinstance(obj, np.floating):
+            return float(obj)
+        elif isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return json.JSONEncoder.default(self, obj)
+
+
+ParamsType = Union[Mapping[str, Any], argparse.Namespace]
+
+
+def process_params(params: ParamsType) -> Mapping[str, Any]:
+    if isinstance(params, Mapping):
+        return params
+    if isinstance(params, argparse.Namespace):
+        return vars(params)
+    # TODO: add absl support if required
+    # move to a different file then
+    raise MlFoundryException(
+        "params should be either argparse.Namespace or a Mapping (dict) type"
+    )
+
+
+def flatten_dict(
+    input_dict: Mapping[Any, Any], parent_key: str = "", sep: str = "."
+) -> Dict[str, Any]:
+    """Flatten hierarchical dict, e.g. ``{'a': {'b': 'c'}} -> {'a.b': 'c'}``.
+    All the keys will be converted to str.
+
+    Args:
+        input_dict: Dictionary containing the keys
+        parent_key: Prefix to add to the keys. Defaults to ``''``.
+        sep: Delimiter to express the hierarchy. Defaults to ``'.'``.
+
+    Returns:
+        Flattened dict.
+
+    Examples:
+        >>> flatten_dict({'a': {'b': 'c'}})
+        {'a.b': 'c'}
+        >>> flatten_dict({'a': {'b': 123}})
+        {'a.b': 123}
+        >>> flatten_dict({'a': {'b': 'c'}}, parent_key="param")
+        {'param.a.b': 'c'}
+    """
+    new_dict_items: List[Tuple[str, Any]] = []
+    for k, v in input_dict.items():
+        new_key = parent_key + sep + str(k) if parent_key else k
+        if isinstance(v, Mapping):
+            new_dict_items.extend(flatten_dict(v, new_key, sep=sep).items())
+        else:
+            new_dict_items.append((new_key, v))
+    return dict(new_dict_items)

truefoundry/ml/services/auth_service.py

@@ -0,0 +1,109 @@
+import time
+from urllib.parse import urlparse
+
+from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+    ApiClient,
+    AuthApi,
+    Configuration,
+)
+from truefoundry.ml.exceptions import MlFoundryException
+from truefoundry.ml.logger import logger
+from truefoundry.ml.run_utils import append_path_to_rest_tracking_uri
+from truefoundry.ml.services.entities import (
+    AuthServerInfo,
+    DeviceCode,
+    HostCreds,
+    Token,
+)
+from truefoundry.ml.services.utils import http_request, http_request_safe
+
+# TODO: This will eventually go away, this is duplicate of AuthServiceClient
+
+
+class AuthService:
+    def __init__(self, url: str, tenant_name: str):
+        self._host_creds = HostCreds(host=url.rstrip("/"), token=None)
+        self._tenant_name = tenant_name
+
+    def refresh_token(self, token: Token) -> Token:
+        if not token.refresh_token:
+            # TODO: Add a way to propagate error messages without traceback to the output interface side
+            raise MlFoundryException(
+                "Unable to resume login session. Please log in again using `tfy login [--host HOST] --relogin`"
+            )
+        try:
+            response = http_request_safe(
+                method="post",
+                host_creds=self._host_creds,
+                endpoint="api/v1/oauth/token/refresh",
+                json={
+                    "tenantName": token.tenant_name,
+                    "refreshToken": token.refresh_token,
+                },
+                timeout=3,
+            )
+        except MlFoundryException as e:
+            if e.status_code and (400 <= e.status_code < 500):
+                raise MlFoundryException(
+                    "Unable to resume login session. "
+                    "Please log in again using `tfy login [--host HOST] --relogin`"
+                ) from None
+            raise
+        return Token.parse_obj(response)
+
+    def get_device_code(self) -> DeviceCode:
+        response = http_request_safe(
+            method="post",
+            host_creds=self._host_creds,
+            endpoint="api/v1/oauth/device",
+            json={"tenantName": self._tenant_name},
+            timeout=3,
+        )
+        return DeviceCode.parse_obj(response)
+
+    def get_token_from_device_code(
+        self, device_code: str, timeout: float = 60
+    ) -> Token:
+        start_time = time.monotonic()
+        while (time.monotonic() - start_time) <= timeout:
+            response = http_request(
+                method="post",
+                host_creds=self._host_creds,
+                endpoint="api/v1/oauth/device/token",
+                json={"tenantName": self._tenant_name, "deviceCode": device_code},
+                timeout=3,
+            )
+            if response.status_code == 202:
+                logger.debug("User has not authorized yet. Checking again.")
+                time.sleep(1.0)
+                continue
+            if response.status_code == 201:
+                response = response.json()
+                return Token.parse_obj(response)
+            raise MlFoundryException(
+                "Failed to get token using device code.\n"
+                f"Status Code: {response.status_code},\nResponse: {response.text}"
+            )
+        raise MlFoundryException(f"Did not get authorized within {timeout} seconds.")
+
+
+def get_auth_service(tracking_uri: str) -> AuthService:
+    tracking_uri = append_path_to_rest_tracking_uri(tracking_uri)
+    parsed_tracking_uri = urlparse(tracking_uri)
+    host = parsed_tracking_uri.netloc
+    # Anonymous api
+    api_client = ApiClient(
+        configuration=Configuration(
+            host=tracking_uri.rstrip("/"),
+            access_token=None,
+        )
+    )
+    auth_api = AuthApi(api_client=api_client)
+    auth_server_info = auth_api.get_tenant_id_get(
+        host_name=host,
+        _request_timeout=3,
+    )
+    tenant_info = AuthServerInfo.parse_obj(auth_server_info.dict())
+    return AuthService(
+        url=tenant_info.auth_server_url, tenant_name=tenant_info.tenant_name
+    )

truefoundry/ml/services/entities.py

@@ -0,0 +1,108 @@
+import enum
+import time
+from typing import Optional
+from urllib.parse import urlparse
+
+import jwt
+from typing_extensions import TypedDict
+
+from truefoundry.pydantic_v1 import BaseModel, Field, NonEmptyStr, validator
+
+
+class HostCreds(BaseModel):
+    host: str
+    token: Optional[str] = None
+
+
+class AuthServerInfo(BaseModel):
+    tenant_name: NonEmptyStr
+    auth_server_url: str
+
+    class Config:
+        allow_mutation = False
+
+
+class ArtifactCredential(BaseModel):
+    run_id: str
+    path: str
+    signed_uri: str
+
+
+class UserType(enum.Enum):
+    user = "user"
+    serviceaccount = "serviceaccount"
+
+
+class UserInfo(BaseModel):
+    user_id: NonEmptyStr
+    user_type: UserType = UserType.user
+    email: Optional[str] = None
+    tenant_name: NonEmptyStr = Field(alias="tenantName")
+
+    class Config:
+        allow_population_by_field_name = True
+        allow_mutation = False
+
+
+class _DecodedToken(TypedDict):
+    tenantName: str
+    exp: int
+    username: str
+    email: str
+    userType: UserType
+
+
+class Token(BaseModel):
+    access_token: NonEmptyStr = Field(alias="accessToken", repr=False)
+    refresh_token: Optional[NonEmptyStr] = Field(alias="refreshToken", repr=False)
+    decoded_value: Optional[_DecodedToken] = Field(exclude=True, repr=False)
+
+    class Config:
+        allow_population_by_field_name = True
+        allow_mutation = False
+
+    @validator("decoded_value", always=True, pre=True)
+    def _decode_jwt(cls, v, values, **kwargs):
+        access_token = values["access_token"]
+        return jwt.decode(
+            access_token,
+            options={
+                "verify_signature": False,
+                "verify_aud": False,
+                "verify_exp": False,
+            },
+        )
+
+    @property
+    def tenant_name(self) -> str:
+        assert self.decoded_value is not None
+        return self.decoded_value["tenantName"]
+
+    def is_going_to_be_expired(self, buffer_in_seconds: int = 120) -> bool:
+        assert self.decoded_value is not None
+        exp = int(self.decoded_value["exp"])
+        return (exp - time.time()) < buffer_in_seconds
+
+    def to_user_info(self) -> UserInfo:
+        assert self.decoded_value is not None
+        return UserInfo(
+            user_id=self.decoded_value["username"],
+            email=(
+                self.decoded_value["email"] if "email" in self.decoded_value else None
+            ),
+            user_type=UserType(self.decoded_value.get("userType", UserType.user.value)),
+            tenant_name=self.tenant_name,  # type: ignore[call-arg]
+        )
+
+
+class DeviceCode(BaseModel):
+    user_code: str = Field(alias="userCode")
+    device_code: str = Field(alias="deviceCode")
+
+    class Config:
+        allow_population_by_field_name = True
+        allow_mutation = False
+
+    def get_user_clickable_url(self, tracking_uri: str) -> str:
+        parsed_tracking_uri = urlparse(tracking_uri)
+        return f"{parsed_tracking_uri.scheme}://{parsed_tracking_uri.netloc}/authorize/device?userCode={self.user_code}"

truefoundry/ml/services/servicefoundry_service.py

@@ -0,0 +1,35 @@
+from typing import Optional
+
+from truefoundry.common.request_utils import requests_retry_session
+from truefoundry.ml.exceptions import MlFoundryException
+from truefoundry.ml.run_utils import append_servicefoundry_path_to_tracking_ui
+from truefoundry.ml.services.entities import HostCreds
+from truefoundry.ml.services.utils import http_request_safe
+
+
+class ServicefoundryService:
+    def __init__(self, tracking_uri: str, token: Optional[str] = None):
+        self.host_creds = HostCreds(
+            host=append_servicefoundry_path_to_tracking_ui(tracking_uri), token=token
+        )
+
+    def get_integration_from_id(self, integration_id: str):
+        integration_id = integration_id or ""
+        data = http_request_safe(
+            method="get",
+            host_creds=self.host_creds,
+            endpoint="v1/provider-accounts/provider-integrations",
+            session=requests_retry_session(retries=1),
+            params={"id": integration_id, "type": "blob-storage"},
+            timeout=3,
+        )
+        if (
+            data.get("providerIntegrations")
+            and len(data["providerIntegrations"]) > 0
+            and data["providerIntegrations"][0]
+        ):
+            return data["providerIntegrations"][0]
+        else:
+            raise MlFoundryException(
+                f"Invalid storage integration id: {integration_id}"
+            )

truefoundry/ml/services/utils.py

@@ -0,0 +1,122 @@
+import json
+from contextlib import contextmanager
+from typing import Any, Optional
+from urllib.parse import urljoin
+
+import requests
+
+from truefoundry.common.request_utils import requests_retry_session
+from truefoundry.ml.exceptions import MlFoundryException
+from truefoundry.ml.services.entities import HostCreds
+
+
+# TODO: This will be moved later to truefoundry.common.request_utils
+def _make_url(host: str, endpoint: str) -> str:
+    if endpoint.startswith("/"):
+        raise ValueError("`endpoint` must not start with a leading slash (/)")
+    return urljoin(host, endpoint)
+
+
+def _http_request(
+    *, method: str, url: str, token: Optional[str] = None, session=requests, **kwargs
+) -> requests.Response:
+    headers = kwargs.pop("headers", {}) or {}
+    if token is not None:
+        headers["Authorization"] = f"Bearer {token}"
+    return session.request(method=method, url=url, **kwargs)
+
+
+def http_request(
+    *, method: str, host_creds: HostCreds, endpoint: str, session=requests, **kwargs
+) -> requests.Response:
+    url = _make_url(host=host_creds.host, endpoint=endpoint)
+    return _http_request(
+        method=method, url=url, token=host_creds.token, session=session, **kwargs
+    )
+
+
+def http_request_safe(
+    *, method: str, host_creds: HostCreds, endpoint: str, session=requests, **kwargs
+) -> Any:
+    url = _make_url(host=host_creds.host, endpoint=endpoint)
+    try:
+        response = _http_request(
+            method=method, url=url, token=host_creds.token, session=session, **kwargs
+        )
+        response.raise_for_status()
+        try:
+            return response.json()
+        except json.JSONDecodeError as je:
+            raise MlFoundryException(
+                f"Failed to parse response as json. Response: {response.text}"
+            ) from je
+    except requests.exceptions.ConnectionError as ce:
+        raise MlFoundryException("Failed to connect to TrueFoundry") from ce
+    except requests.exceptions.Timeout as te:
+        raise MlFoundryException(f"Request to {url} timed out") from te
+    except requests.exceptions.HTTPError as he:
+        raise MlFoundryException(
+            f"Request to {url} with status code {he.response.status_code}. Response: {he.response.text}",
+            status_code=he.response.status_code,
+        ) from he
+    except Exception as e:
+        raise MlFoundryException(
+            f"Request to {url} failed with an unknown error"
+        ) from e
+
+
+@contextmanager
+def cloud_storage_http_request(
+    *,
+    method,
+    url,
+    session=None,
+    timeout=None,
+    **kwargs,
+):
+    """
+    Performs an HTTP PUT/GET request using Python's `requests` module with automatic retry.
+    """
+    session = session or requests_retry_session(retries=5, backoff_factor=2)
+    kwargs["headers"] = kwargs.get("headers", {}) or {}
+    if "blob.core.windows.net" in url:
+        kwargs["headers"].update({"x-ms-blob-type": "BlockBlob"})
+    if method.lower() not in ("put", "get"):
+        raise ValueError(f"Illegal http method: {method}")
+    try:
+        yield _http_request(
+            method=method, url=url, session=session, timeout=timeout, **kwargs
+        )
+    except Exception as e:
+        raise MlFoundryException(
+            f"API request failed with exception {str(e)}"
+        ) from None
+
+
+def augmented_raise_for_status(response):
+    try:
+        response.raise_for_status()
+    except requests.exceptions.HTTPError as he:
+        raise MlFoundryException(
+            f"Request with status code {he.response.status_code}. Response: {he.response.text}",
+            status_code=he.response.status_code,
+        ) from he
+
+
+def download_file_using_http_uri(http_uri, download_path, chunk_size=100_000_000):
+    """
+    Downloads a file specified using the `http_uri` to a local `download_path`. This function
+    uses a `chunk_size` to ensure an OOM error is not raised a large file is downloaded.
+
+    Note : This function is meant to download files using presigned urls from various cloud
+            providers.
+    """
+    with cloud_storage_http_request(
+        method="get", url=http_uri, stream=True
+    ) as response:
+        augmented_raise_for_status(response)
+        with open(download_path, "wb") as output_file:
+            for chunk in response.iter_content(chunk_size=chunk_size):
+                if not chunk:
+                    break
+                output_file.write(chunk)