qwak-core 0.1.37__py3-none-any.whl → 0.5.19__py3-none-any.whl

qwak/model/fs_info_mapping_retriever.py

@@ -0,0 +1,58 @@
+from typing import Dict, List
+
+from qwak.clients.administration.eco_system.eco_system_utils import EcosystemUtils
+from qwak.clients.feature_store import FeatureRegistryClient
+from qwak.feature_store._common.feature_set_utils import (
+    FeatureSetInfo,
+    get_env_to_featuresets_mapping,
+    get_feature_set_info,
+)
+from qwak.model.schema_entities import BaseFeature, FeatureStoreInput
+from qwak.model.utils.feature_utils import extract_env_name, extract_featureset_name
+
+
+def retrieve_fs_mapping(
+    features: List[BaseFeature],
+) -> Dict[str, Dict[str, FeatureSetInfo]]:
+    """
+    Retrieve the feature set info mapping.
+    If one of the features is a FeatureStoreInput and is in a different environment from the current one, populate it using a single API call.
+    Otherwise, retrieve each feature set info individually.
+    The reasoning behind this logic is that if all required features are in the same environment, it's more efficient to retrieve only the necessary feature set, rather than all feature sets.
+    * Assuming the feature name is fully qualified with the environment name
+    :param features:
+    :return: dict of environment name to dict of feature set name to FeatureSetInfo
+    """
+    current_env_name = EcosystemUtils().get_current_environment_name()
+    features_manager_client = FeatureRegistryClient()
+
+    if any(
+        [
+            feature
+            for feature in features
+            if isinstance(feature, FeatureStoreInput)
+            and extract_env_name(feature.name) != current_env_name
+        ]
+    ):
+        return get_env_to_featuresets_mapping(features_manager_client)
+    else:
+        return _retrieve_fs_info_one_by_one(
+            features_manager_client, current_env_name, features
+        )
+
+
+def _retrieve_fs_info_one_by_one(
+    features_manager_client: FeatureRegistryClient,
+    current_env_name: str,
+    features: List[BaseFeature],
+) -> Dict[str, Dict[str, FeatureSetInfo]]:
+    fs_info_cache: Dict[str, FeatureSetInfo] = dict()
+
+    for feature in features:
+        if isinstance(feature, FeatureStoreInput):
+            feature_set_name = extract_featureset_name(feature.name).lower()
+            if feature_set_name not in fs_info_cache:
+                fs_info_cache[feature_set_name] = get_feature_set_info(
+                    features_manager_client, feature_set_name
+                )
+    return {current_env_name: fs_info_cache}

qwak/model/schema.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import List
 
-from _qwak_proto.qwak.builds.builds_pb2 import ModelSchema as ProtoModelSchema
+from _qwak_proto.qwak.build.v1.build_pb2 import ModelSchema as ProtoModelSchema
 from qwak.model._entity_extraction import enrich_schema
 from qwak.model.schema_entities import (
     BaseFeature,

qwak/model/schema_entities.py

@@ -2,7 +2,7 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Optional, Type
 
-from _qwak_proto.qwak.builds.builds_pb2 import (
+from _qwak_proto.qwak.build.v1.build_pb2 import (
     BatchFeatureV1 as ProtoBatchFeatureV1,
     Entity as ProtoEntity,
     ExplicitFeature as ProtoExplicitFeature,

qwak/model/tools/__init__.py

@@ -21,4 +21,4 @@ def wire_runtime_with_run_local():
 if not ContainerLock.locked:
     wire_runtime_with_run_local()
 
-__all__ = ["run_local"]
+__all__ = ["run_local", "wire_runtime_with_run_local"]

qwak/model/tools/adapters/input_adapters/dataframe_input.py

@@ -1,12 +1,12 @@
 import json
 from typing import List, Optional, Union
 
-import pandas
 from qwak.exceptions import QwakHTTPException
 
 from .string_input import StringInput
 
 try:
+    import pandas
     from bentoml.adapters import DataframeInput as BentoMLDataframeInput
 except ImportError:
     pass

qwak/model/tools/adapters/output.py

@@ -46,7 +46,7 @@ def get_output_adapter(
 
             first_result = (
                 return_result[0]
-                if type(return_result) == list and len(return_result) > 0
+                if isinstance(return_result, list) and len(return_result) > 0
                 else return_result
             )
             if issubclass(type(first_result), Message):

qwak/model/tools/run_model_locally.py

@@ -1,8 +1,10 @@
 import inspect
+import os
 from typing import Union
 
 from qwak.feature_store.online.client import OnlineClient
 from qwak.model.base import QwakModel
+from contextlib import contextmanager
 
 from .adapters.input import get_input_adapter
 from .adapters.output import get_output_adapter
@@ -21,9 +23,10 @@ def run_local(model: QwakModel, payload: Union[str, bytes]):
     :return: a string or bytes with the model output
     """
     if not (hasattr(model, "initialized") and model.initialized):
-        model.build()
-        model.initialize_model()
-        setattr(model, "initialized", True)
+        with _set_local_mode():
+            model.build()
+            model.initialize_model()
+            setattr(model, "initialized", True)
     input_adapter = get_input_adapter(model)
     output_adapter = get_output_adapter(model)
     input_data = input_adapter.extract_user_func_args(payload)
@@ -54,3 +57,12 @@ def _extract_online_features(model, df):
     return OnlineClient().get_feature_values(
         schema=model.schema(), df=df, model_name="Test Model"
     )
+
+
+@contextmanager
+def _set_local_mode():
+    os.environ["QWAK_IS_RUN_LOCAL"] = "true"
+    try:
+        yield
+    finally:
+        del os.environ["QWAK_IS_RUN_LOCAL"]

qwak/model/utils/feature_utils.py

@@ -0,0 +1,69 @@
+import dataclasses
+from typing import List, cast
+
+from qwak.clients.administration.eco_system.eco_system_utils import EcosystemUtils
+from qwak.exceptions import QwakException
+from qwak.model.schema_entities import BaseFeature, FeatureStoreInput
+
+
+def validate_and_sanitize_feature_name(feature_name: str, current_env_name: str) -> str:
+    qualifiers = feature_name.split(".")
+    if len(qualifiers) <= 1:
+        raise QwakException(
+            f"Feature name '{feature_name}' must adhere to the <featureset_name>.<feature_name> or <env_name>.<featureset_name>.<feature_name> convention"
+        )
+    elif len(qualifiers) == 2:  # env name is not included in the feature name
+        feature_name = f"{current_env_name}.{feature_name}"
+    return feature_name.lower()
+
+
+def validate_and_sanitize_features_name(
+    features: List[BaseFeature],
+) -> List[BaseFeature]:
+    """
+    Verify that the features (from type FeatureStoreInput) names are valid and adhere to convention <featureset_name>.<feature_name> or <env_name>.<featureset_name>.<feature_name>
+    Add the current env qualifier if env is missing and transform name to lowercase
+    :param features:
+    :return: sanitized features with lower case (to support case-insensitive) fully qualified feature names
+    :exception QwakException: if the feature name does not adhere to the convention
+    """
+    ecosystem_utils = EcosystemUtils()
+    current_env_name = ecosystem_utils.get_current_environment_name()
+    return [
+        (
+            cast(
+                FeatureStoreInput,
+                dataclasses.replace(
+                    feature,
+                    name=validate_and_sanitize_feature_name(
+                        feature.name, current_env_name
+                    ),
+                ),
+            )
+            if isinstance(feature, FeatureStoreInput)
+            else feature
+        )
+        for feature in features
+    ]
+
+
+def extract_env_name(feature_name: str) -> str:
+    """
+    extracts env_name from featureset name
+    assuming feature name is fully qualified (passed sanitization)
+    :param feature_name:
+    :return: env_name
+    """
+    qualifiers = feature_name.split(".")
+    return ".".join(qualifiers[:-2])  # Everything except the last two parts
+
+
+def extract_featureset_name(feature_name: str) -> str:
+    qualifiers = feature_name.split(".")
+    return qualifiers[-2]
+
+
+def discard_env_from_name(
+    feature_name: str,
+) -> str:
+    return ".".join(feature_name.split(".")[-2:])

qwak/model_loggers/artifact_logger.py

@@ -1,6 +1,13 @@
-import urllib.request
+from datetime import timedelta
 from typing import Optional
+import requests
+from requests import Response
 
+from _qwak_proto.qwak.builds.build_url_pb2 import BuildVersioningTagsType
+from _qwak_proto.qwak.builds.builds_orchestrator_service_pb2 import (
+    GetBuildVersioningUploadURLResponse,
+    GetBuildVersioningDownloadURLResponse,
+)
 from qwak.clients.build_orchestrator.client import BuildOrchestratorClient
 from qwak.clients.file_versioning.client import FileVersioningManagementClient
 from qwak.exceptions import QwakException
@@ -11,6 +18,8 @@ from qwak.inner.model_loggers_utils import (
     validate_tag,
 )
 
+MAX_CHUNK_SIZE = 8 * 1024
+
 
 def log_file(
     from_path: str,
@@ -37,8 +46,13 @@ def log_file(
         # Checking if called inside a model - then build id saved as environment variable or stays
         build_id = fetch_build_id()
 
-    upload_url_response = BuildOrchestratorClient().get_build_versioning_upload_url(
-        build_id=build_id, model_id=model_id, tag=tag
+    upload_url_response: (
+        GetBuildVersioningUploadURLResponse
+    ) = BuildOrchestratorClient().get_build_versioning_upload_url(
+        build_id=build_id,
+        model_id=model_id,
+        tag=tag,
+        tag_type=BuildVersioningTagsType.FILE_TAG_TYPE,
     )
 
     FileVersioningManagementClient().register_file_tag(
@@ -46,7 +60,11 @@ def log_file(
     )
 
     with open(from_path, mode="rb") as f:
-        upload_data(upload_url_response.upload_url, f.read())
+        upload_data(
+            upload_url_response.upload_url,
+            f.read(),
+            upload_url_response.headers,
+        )
 
 
 def load_file(
@@ -67,21 +85,40 @@ def load_file(
     Returns:
         the path to the newly created data file
     """
+    print(f"Loading file to {to_path}")
     if not validate_tag(tag):
         raise QwakException(
             "Tag should contain only letters, numbers, underscore or hyphen"
         )
 
     model_id = validate_model(model_id)
-    download_url_response = BuildOrchestratorClient().get_build_versioning_download_url(
-        build_id=build_id, model_id=model_id, tag=tag
+    download_url_response: (
+        GetBuildVersioningDownloadURLResponse
+    ) = BuildOrchestratorClient().get_build_versioning_download_url(
+        build_id=build_id,
+        model_id=model_id,
+        tag=tag,
+        tag_type=BuildVersioningTagsType.FILE_TAG_TYPE,
     )
 
     try:
-        filename, headers = urllib.request.urlretrieve(  # nosec B310
-            url=download_url_response.download_url, filename=to_path
+        response: Response = requests.get(
+            download_url_response.download_url,
+            headers=download_url_response.headers,
+            stream=True,
+            timeout=(
+                timedelta(seconds=10).total_seconds(),  # timeout to connect
+                timedelta(minutes=20).total_seconds(),  # timeout to read
+            ),
         )
+        print(f"Downloading file finished with status {response.status_code}")
+        response.raise_for_status()
+
+        with open(to_path, "wb") as f:
+            for chunk in response.iter_content(chunk_size=MAX_CHUNK_SIZE):
+                if chunk:
+                    f.write(chunk)
 
-        return filename
+        return to_path
     except Exception as error:
         raise QwakException(f"Unable to load save artifact locally: {str(error)}")

qwak/model_loggers/data_logger.py

@@ -1,4 +1,5 @@
 import gzip
+import os
 from io import StringIO
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
@@ -10,6 +11,10 @@ from _qwak_proto.qwak.builds.build_pb2 import (
     DataFormat,
     DataTableDefinition,
 )
+from _qwak_proto.qwak.builds.build_url_pb2 import BuildVersioningTagsType
+from _qwak_proto.qwak.builds.builds_orchestrator_service_pb2 import (
+    GetBuildVersioningUploadURLResponse,
+)
 from qwak.clients.build_orchestrator.client import BuildOrchestratorClient
 from qwak.clients.data_versioning.client import DataVersioningManagementClient
 from qwak.exceptions import QwakException
@@ -44,6 +49,9 @@ def log_data(
         build_id: optional build id - if not given found from environment.
         encoding: Encoding for the dataframe - utf-8 is default.
     """
+    if os.getenv("QWAK_IS_RUN_LOCAL"):
+        return
+
     if not validate_tag(tag):
         raise QwakException(
             "Tag should contain only letters, numbers, underscore or hyphen"
@@ -58,8 +66,13 @@ def log_data(
         build_id = fetch_build_id()
 
     client = BuildOrchestratorClient()
-    upload_url_response = client.get_build_versioning_upload_url(
-        build_id=build_id, model_id=model_id, tag=str(data_tag)
+    upload_url_response: GetBuildVersioningUploadURLResponse = (
+        client.get_build_versioning_upload_url(
+            build_id=build_id,
+            model_id=model_id,
+            tag=str(data_tag),
+            tag_type=BuildVersioningTagsType.DATA_TAG_TYPE,
+        )
     )
 
     string_buffer = StringIO()
@@ -68,6 +81,7 @@ def log_data(
         upload_url=upload_url_response.upload_url,
         data=gzip.compress(bytes(string_buffer.getvalue(), "utf-8")),
         content_type="text/plain",
+        headers=upload_url_response.headers,
     )
 
     dataframe_definition = DataTableDefinition(

qwak/qwak_client/batch_jobs/execution.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+
+from _qwak_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
+    BatchJobDetails as BatchJobDetailsProto,
+)
+
+
+class ExecutionStatus(Enum):
+    UNDEFINED_BATCH_JOB_STATUS = 0
+    BATCH_JOB_COMMITTED_STATUS = 1
+    BATCH_JOB_PENDING_STATUS = 2
+    BATCH_JOB_RUNNING_STATUS = 3
+    BATCH_JOB_FINISHED_STATUS = 4
+    BATCH_JOB_FAILED_STATUS = 5
+    BATCH_JOB_CANCELLED_STATUS = 6
+    BATCH_JOB_TIMEOUT_STATUS = 7
+
+
+@dataclass
+class Execution:
+    environment_id: str = field(default=None)
+    build_id: str = field(default=None)
+    execution_id: str = field(default=None)
+    execution_status: ExecutionStatus = field(
+        default=ExecutionStatus.UNDEFINED_BATCH_JOB_STATUS
+    )
+    start_time: datetime = field(default=datetime.now())
+    end_time: datetime = field(default=datetime.now())
+    failure_message: str = field(default=None)
+
+    @staticmethod
+    def from_proto(batch_job_details_proto: BatchJobDetailsProto):
+        return Execution(
+            environment_id=batch_job_details_proto.environment_id,
+            build_id=batch_job_details_proto.build_id,
+            execution_id=batch_job_details_proto.job_id,
+            execution_status=ExecutionStatus(batch_job_details_proto.job_status),
+            start_time=datetime.fromtimestamp(
+                batch_job_details_proto.start_time.seconds
+                + batch_job_details_proto.start_time.nanos / 1e9
+            ),
+            end_time=datetime.fromtimestamp(
+                batch_job_details_proto.end_time.seconds
+                + batch_job_details_proto.end_time.nanos / 1e9
+            ),
+            failure_message=batch_job_details_proto.failure_message,
+        )

qwak/qwak_client/batch_jobs/task.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+
+from _qwak_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
+    TaskExecutionDetails as TaskExecutionDetailsProto,
+)
+
+
+class TaskStatus(Enum):
+    UNDEFINED_BATCH_TASK_STATUS = 0
+    BATCH_TASK_COMMITTED_STATUS = 1
+    BATCH_TASK_PENDING_START_STATUS = 2
+    BATCH_TASK_RUNNING_STATUS = 3
+    BATCH_TASK_FINISHED_STATUS = 4
+    BATCH_TASK_FAILED_STATUS = 5
+    BATCH_TASK_CANCELLED_STATUS = 6
+    BATCH_TASK_TIMEOUT_STATUS = 7
+
+
+@dataclass
+class Task:
+    task_id: str = field(default=None)
+    task_status: TaskStatus = field(default=TaskStatus.UNDEFINED_BATCH_TASK_STATUS)
+    start_time: datetime = field(default=datetime.now())
+    end_time: datetime = field(default=datetime.now())
+    filename: str = field(default=None)
+
+    @staticmethod
+    def from_proto(task_execution_details_proto: TaskExecutionDetailsProto):
+        return Task(
+            task_id=task_execution_details_proto.task_id,
+            task_status=TaskStatus(task_execution_details_proto.status),
+            start_time=datetime.fromtimestamp(
+                task_execution_details_proto.start_time.seconds
+                + task_execution_details_proto.start_time.nanos / 1e9
+            ),
+            end_time=datetime.fromtimestamp(
+                task_execution_details_proto.end_time.seconds
+                + task_execution_details_proto.end_time.nanos / 1e9
+            ),
+            filename=task_execution_details_proto.filename,
+        )

qwak/qwak_client/build_api_helpers/build_api_steps.py

@@ -0,0 +1,53 @@
+from qwak.inner.build_config.build_config_v1 import BuildConfigV1
+from qwak.inner.build_logic.interface.build_phase import BuildPhase
+from qwak.inner.build_logic.phases.phase_010_fetch_model.fetch_model_step import (
+    FetchModelStep,
+)
+from qwak.inner.build_logic.phases.phase_010_fetch_model.post_fetch_validation_step import (
+    PostFetchValidationStep,
+)
+from qwak.inner.build_logic.phases.phase_010_fetch_model.pre_fetch_validation_step import (
+    PreFetchValidationStep,
+)
+from qwak.inner.build_logic.phases.phase_010_fetch_model.set_version_step import (
+    SetVersionStep,
+)
+from qwak.inner.build_logic.phases.phase_020_remote_register_qwak_build.cleanup_step import (
+    CleanupStep,
+)
+from qwak.inner.build_logic.phases.phase_020_remote_register_qwak_build.start_remote_build_step import (
+    StartRemoteBuildStep,
+)
+from qwak.inner.build_logic.phases.phase_020_remote_register_qwak_build.upload_step import (
+    UploadStep,
+)
+from qwak.inner.build_logic.phases.phases_pipeline import PhasesPipeline
+from qwak.inner.build_logic.trigger_build_context import TriggerBuildContext
+
+FETCHING_MODEL_CODE_PHASE = "FETCHING_MODEL_CODE"
+REGISTERING_QWAK_BUILD_PHASE = "REGISTERING_QWAK_BUILD"
+
+
+def get_trigger_build_api_steps(
+    config: BuildConfigV1, context: TriggerBuildContext
+) -> PhasesPipeline:
+    steps_root = PhasesPipeline(config=config, context=context)
+    steps_root.add_phase(
+        steps=[
+            SetVersionStep(),
+            PreFetchValidationStep(),
+            FetchModelStep(),
+            PostFetchValidationStep(),
+        ],
+        build_phase=BuildPhase(phase_id=FETCHING_MODEL_CODE_PHASE),
+    )
+    steps_root.add_phase(
+        steps=[
+            UploadStep(),
+            StartRemoteBuildStep(),
+            CleanupStep(),
+        ],
+        build_phase=BuildPhase(phase_id=REGISTERING_QWAK_BUILD_PHASE),
+    )
+
+    return steps_root

qwak/qwak_client/build_api_helpers/messages.py

@@ -0,0 +1,6 @@
+SUCCESS_MSG_API = """
+Build ID {build_id} was triggered remotely
+
+To follow build logs using Qwak platform:
+https://app.qwak.ai/projects/{project_uuid}/{model_id}/build/{build_id}
+"""

qwak/qwak_client/build_api_helpers/trigger_build_api.py

@@ -0,0 +1,65 @@
+import contextlib
+import os
+import shutil
+import uuid
+from pathlib import Path
+
+from qwak.inner.build_config.build_config_v1 import BuildConfigV1
+from qwak.inner.build_logic.constants.host_resource import HOST_QWAK_HIDDEN_FOLDER
+from qwak.inner.build_logic.execute_build_pipeline import execute_build_pipeline
+from qwak.inner.build_logic.run_handlers.programmatic_phase_run_handler import (
+    ProgrammaticPhaseRunHandler,
+)
+from qwak.inner.build_logic.trigger_build_context import TriggerBuildContext
+from qwak.qwak_client.build_api_helpers.build_api_steps import (
+    get_trigger_build_api_steps,
+)
+from qwak.qwak_client.build_api_helpers.messages import SUCCESS_MSG_API
+from qwak.tools.logger import setup_qwak_logger, get_qwak_logger
+
+BUILDS_LOGS = HOST_QWAK_HIDDEN_FOLDER / "logs" / "build"
+BUILD_LOG_NAME = "build.log"
+MAX_LOGS_NUMBER = 15
+DEBUG_LEVEL = 2
+
+
+@contextlib.contextmanager
+def get_build_logger(config: BuildConfigV1, build_id: str):
+    log_path = BUILDS_LOGS / config.build_properties.model_id / build_id
+    log_path.mkdir(parents=True, exist_ok=True)
+    try:
+        (log_path / "build_config.yml").write_text(config.to_yaml())
+        setup_qwak_logger(logs_folder=log_path, log_file_name=BUILD_LOG_NAME)
+        yield get_qwak_logger()
+    finally:
+        # Cleanup - Save only x last zips
+        logs_zip_sorted_by_data = sorted(
+            BUILDS_LOGS.rglob("**/*"), key=os.path.getmtime
+        )[:-MAX_LOGS_NUMBER]
+        path: Path
+        for path in logs_zip_sorted_by_data:
+            if path.is_file():
+                os.remove(path)
+            elif path.is_dir():
+                shutil.rmtree(path, ignore_errors=True)
+
+
+def trigger_build_api(config: BuildConfigV1):
+    build_id = str(uuid.uuid4())
+    with get_build_logger(config=config, build_id=build_id) as logger:
+        context = TriggerBuildContext()
+        context.build_id = build_id
+        pipeline = get_trigger_build_api_steps(config, context)
+        build_runner = ProgrammaticPhaseRunHandler(
+            logger, config.verbose, json_logs=False
+        )
+        execute_build_pipeline(pipeline, build_runner)
+        logger.info(
+            SUCCESS_MSG_API.format(
+                build_id=build_id,
+                model_id=context.model_id,
+                project_uuid=context.project_uuid,
+            )
+        )
+
+    return build_id