frogml-core 0.0.111__tar.gz → 0.0.112__tar.gz

{frogml_core-0.0.111 → frogml_core-0.0.112}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: frogml-core
-Version: 0.0.111
+Version: 0.0.112
 Summary: frogml Core contains the necessary objects and communication tools for using the Jfrog ml Platform
 License: Apache-2.0
 Keywords: mlops,ml,deployment,serving,model

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/__init__.py

@@ -1,7 +1,7 @@
 """Top-level package for frogml."""
 
 __author__ = "jfrog"
-__version__ = "0.0.111"
+__version__ = "0.0.112"
 
 from frogml_core.inner.di_configuration import wire_dependencies
 from frogml_core.model.model_version_tracking import (  # noqa: F401,E501

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/automations/automations.py

@@ -293,6 +293,7 @@ class Automation:
     create_audit: AutomationAudit = field(default_factory=AutomationAudit)
     on_error: Notification = field(default=None)
     on_success: Notification = field(default=None)
+    jfrog_token_id: str = field(default="")
 
     def to_proto(self):
         return AutomationProto(
@@ -315,10 +316,11 @@ class Automation:
             qwak_environment_id=self.environment,
             create_audit=self.create_audit.to_proto(),
             is_deleted=self.is_deleted,
+            jfrog_token_id=self.jfrog_token_id,
         )
 
-    @staticmethod
-    def from_proto(message: AutomationProto):
+    @classmethod
+    def from_proto(cls, message: AutomationProto):
         action = map_action_name_to_class(
             message.automation_spec.action.WhichOneof("action")
         )
@@ -332,7 +334,7 @@ class Automation:
             message.automation_spec.on_success.WhichOneof("notification")
         )
 
-        return Automation(
+        return cls(
             id=message.automation_id,
             name=message.automation_spec.automation_name,
             description=message.automation_spec.automation_description,
@@ -359,6 +361,7 @@ class Automation:
                 if on_success_notification
                 else None
             ),
+            jfrog_token_id=message.jfrog_token_id,
         )
 
     def __str__(self):

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/_jdbc.py

@@ -13,6 +13,7 @@ class JdbcSource(BaseBatchSource, ABC):
     url: Optional[str] = None
     db_table: Optional[str] = None
     query: Optional[str] = None
+    repository: Optional[str] = None
 
     def __post_init__(self):
         self._validate()

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/athena.py

@@ -52,6 +52,7 @@ class AthenaSource(JdbcSource):
     aws_authentication: Optional[AwsAuthentication] = None
     time_partition_columns: Optional[TimePartitionColumns] = None
     workgroup: Optional[str] = None
+    repository: Optional[str] = None
 
     def _validate(self):
         super()._validate()
@@ -67,7 +68,7 @@ class AthenaSource(JdbcSource):
                 "At least one of: `s3_output_location`, `workgroup` must be set"
             )
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         proto_date_partition_columns = None
         proto_time_fragment_partition_columns = None
 
@@ -97,6 +98,7 @@ class AthenaSource(JdbcSource):
         )
 
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -116,7 +118,7 @@ class AthenaSource(JdbcSource):
                         workgroup=self.workgroup,
                     ),
                 ),
-            )
+            ),
         )
 
     @classmethod

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/big_query.py

@@ -25,6 +25,7 @@ class BigQuerySource(BaseBatchSource):
     views_enabled: Optional[bool] = False
     materialization_dataset: Optional[str] = None
     materialization_project: Optional[str] = None
+    repository: Optional[str] = None
 
     def __post_init__(self):
         if not (self.sql or (self.table and self.dataset)):
@@ -59,8 +60,9 @@ class BigQuerySource(BaseBatchSource):
             materialization_dataset=bigquery.materialization_dataset,
         )
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -76,5 +78,5 @@ class BigQuerySource(BaseBatchSource):
                     materialization_dataset=self.materialization_dataset,
                     materialization_project=self.materialization_project,
                 ),
-            )
+            ),
         )

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/clickhouse.py

@@ -21,6 +21,7 @@ class ClickhouseSource(BaseBatchSource):
     url: Optional[str] = None
     db_table: Optional[str] = None
     query: Optional[str] = None
+    repository: Optional[str] = None
 
     def __post_init__(self):
         self._validate()
@@ -29,8 +30,9 @@ class ClickhouseSource(BaseBatchSource):
         if not (bool(self.db_table) ^ bool(self.query)):
             raise FrogmlException("Only one of query and db_table must be set")
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -42,7 +44,7 @@ class ClickhouseSource(BaseBatchSource):
                     table=self.db_table,
                     sql=self.query,
                 ),
-            )
+            ),
         )
 
     @classmethod

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/csv.py

@@ -25,6 +25,7 @@ class CsvSource(BaseBatchSource):
     quote_character: str = '"'
     escape_character: str = '"'
     filesystem_configuration: Optional[FileSystemConfiguration] = None
+    repository: Optional[str] = None
 
     @classmethod
     def _from_proto(cls, proto):
@@ -42,12 +43,13 @@ class CsvSource(BaseBatchSource):
             filesystem_configuration=fs_conf,
         )
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         fs_conf = None
         if self.filesystem_configuration:
             fs_conf = self.filesystem_configuration._to_proto()
 
         return ProtoDataSourceSpec(
+            featureset_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -58,5 +60,5 @@ class CsvSource(BaseBatchSource):
                     escape_character=self.escape_character,
                     filesystem_configuration=fs_conf,
                 ),
-            )
+            ),
         )

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/elastic_search.py

@@ -23,6 +23,7 @@ class ElasticSearchSource(BaseBatchSource):
     password_secret_name: Optional[str] = None
     exclude_fields: Optional[str] = None
     parse_dates: Optional[bool] = True
+    repository: Optional[str] = None
 
     @classmethod
     def _from_proto(cls, proto):
@@ -41,8 +42,9 @@ class ElasticSearchSource(BaseBatchSource):
             parse_dates=elastic.parse_dates,
         )
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -57,5 +59,5 @@ class ElasticSearchSource(BaseBatchSource):
                     exclude_fields=self.exclude_fields,
                     parse_dates=self.parse_dates,
                 ),
-            )
+            ),
         )

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/mongodb.py

@@ -22,6 +22,7 @@ class MongoDbSource(BaseBatchSource):
     collection: str
     connection_params: str
     protocol: str = "mongodb"
+    repository: Optional[str] = None
 
     @classmethod
     def _from_proto(cls, proto):
@@ -39,8 +40,9 @@ class MongoDbSource(BaseBatchSource):
             protocol=mongo.protocol,
         )
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoBatchSource:
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -54,5 +56,5 @@ class MongoDbSource(BaseBatchSource):
                     connection_params=self.connection_params,
                     protocol=self.protocol,
                 ),
-            )
+            ),
         )

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/mysql.py

@@ -32,8 +32,9 @@ class MysqlSource(JdbcSource):
             query=mysql.query,
         )
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoJdbcSource:
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -46,5 +47,5 @@ class MysqlSource(JdbcSource):
                     query=self.query,
                     mysqlSource=ProtoMysqlSource(),
                 ),
-            )
+            ),
         )

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/parquet.py

@@ -26,6 +26,7 @@ class ParquetSource(BaseBatchSource):
     filesystem_configuration: FileSystemConfiguration = field(
         default_factory=lambda: AnonymousS3Configuration()
     )
+    repository: Optional[str] = None
 
     @classmethod
     def _from_proto(cls, proto):
@@ -40,12 +41,13 @@ class ParquetSource(BaseBatchSource):
             filesystem_configuration=fs_conf,
         )
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         fs_conf = None
         if self.filesystem_configuration:
             fs_conf = self.filesystem_configuration._to_proto()
 
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -53,5 +55,5 @@ class ParquetSource(BaseBatchSource):
                 parquetSource=ProtoParquetSource(
                     path=self.path, filesystem_configuration=fs_conf
                 ),
-            )
+            ),
         )

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/postgres.py

@@ -18,8 +18,9 @@ from frogml_core.feature_store.data_sources.batch._jdbc import JdbcSource
 
 @dataclass
 class PostgresSource(JdbcSource):
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -32,7 +33,7 @@ class PostgresSource(JdbcSource):
                     query=self.query,
                     postgresqlSource=ProtoPostgresqlSource(),
                 ),
-            )
+            ),
         )
 
     @classmethod

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/redshift.py

@@ -27,6 +27,8 @@ class RedshiftSource(JdbcSource):
     secret_access_key: Optional[str] = None
     query_group: Optional[str] = "_qwak_featurestore"
 
+    repository: Optional[str] = None
+
     def _validate(self):
         authentication_methods = sum(
             [
@@ -46,8 +48,9 @@ class RedshiftSource(JdbcSource):
                 "or user id secret name and password secret name"
             )
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -65,7 +68,7 @@ class RedshiftSource(JdbcSource):
                         secret_access_key=self.secret_access_key,
                     ),
                 ),
-            )
+            ),
         )
 
     @classmethod

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/snowflake.py

@@ -24,6 +24,7 @@ class SnowflakeSource(BaseBatchSource):
     warehouse: str
     table: Optional[str] = None
     query: Optional[str] = None
+    repository: Optional[str] = None
 
     def __post_init__(self):
         self._validate()
@@ -40,8 +41,9 @@ class SnowflakeSource(BaseBatchSource):
         if not self.password_secret_name:
             raise FrogmlException("password_secret_name must be set!")
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -56,7 +58,7 @@ class SnowflakeSource(BaseBatchSource):
                     table=self.table,
                     query=self.query,
                 ),
-            )
+            ),
         )
 
     @classmethod

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/batch/vertica.py

@@ -22,6 +22,7 @@ class VerticaSource(BaseBatchSource):
     table: str
     username_secret_name: str
     password_secret_name: str
+    repository: Optional[str] = None
 
     @classmethod
     def _from_proto(cls, proto):
@@ -39,8 +40,9 @@ class VerticaSource(BaseBatchSource):
             table=vertica.table,
         )
 
-    def _to_proto(self, artifact_url: Optional[str] = None):
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             batch_source=ProtoBatchSource(
                 name=self.name,
                 description=self.description,
@@ -54,5 +56,5 @@ class VerticaSource(BaseBatchSource):
                     port=self.port,
                     table=self.table,
                 ),
-            )
+            ),
         )

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/data_sources/streaming/kafka/kafka.py

@@ -52,6 +52,8 @@ class KafkaSource(BaseStreamingSource):
         default_factory=lambda: SslAuthentication()
     )
 
+    repository: Optional[str] = None
+
     def __post_init__(self):
         self._validate()
 
@@ -84,6 +86,7 @@ class KafkaSource(BaseStreamingSource):
 
     def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
         return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
             stream_source=ProtoStreamingSource(
                 name=self.name,
                 description=self.description,
@@ -99,7 +102,7 @@ class KafkaSource(BaseStreamingSource):
                         artifact_path=artifact_url
                     ),
                 ),
-            )
+            ),
         )
 
     @classmethod

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/feature_sets/base_feature_set.py

@@ -33,6 +33,7 @@ class BaseFeatureSet(ABC):
     data_sources: List[str]
     entity: str = str()
     key: str = str()
+    repository: Optional[str] = None
     __instance_module_path__: Optional[str] = None
 
     def _validate(self):

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/feature_sets/batch.py

@@ -72,6 +72,7 @@ def feature_set(
     name: str = None,
     entity: Optional[str] = None,
     key: Optional[str] = None,
+    repository: Optional[str] = None,
 ):
     """
     Define a batch scheduled feature set. Default scheduling policy is every 4 hours.
@@ -116,6 +117,7 @@ def feature_set(
             name=fs_name,
             entity=entity if entity else None,
             key=key if key else None,
+            repository=repository if repository else None,
             data_sources=data_sources,
             timestamp_column_name=timestamp_column_name,
             transformation=user_transformation,
@@ -323,6 +325,7 @@ class BatchFeatureSet(BaseFeatureSet):
             data_sources=[
                 ds.data_source.name for ds in batch_v1_def.feature_set_batch_sources
             ],
+            repository=proto.featureset_repository_name,
             timestamp_column_name=batch_v1_def.timestamp_column_name,
             scheduling_policy=batch_v1_def.scheduling_policy,
             transformation=BaseTransformation._from_proto(batch_v1_def.transformation),
@@ -404,6 +407,7 @@ class BatchFeatureSet(BaseFeatureSet):
                 git_commit=git_commit,
                 features=features,
                 entity=self._get_entity_definition(feature_registry),
+                featureset_repository_name=self.repository,
                 feature_set_type=ProtoFeatureSetType(
                     batch_feature_set_v1=ProtoBatchFeatureSetV1(
                         online_sink=True,

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/feature_sets/streaming.py

@@ -86,6 +86,7 @@ def feature_set(
     entity: Optional[str] = None,
     key: Optional[str] = None,
     auxiliary_sinks: List[BaseSink] = [],
+    repository: Optional[str] = None,
 ):
     """
     Creates a streaming feature set for the specified entity using the given streaming data sources.
@@ -140,6 +141,7 @@ def feature_set(
             name=fs_name,
             entity=entity if entity else None,
             key=key if key else None,
+            repository=repository,
             data_sources=data_sources,
             timestamp_column_name=timestamp_column_name,
             transformation=user_transformation,
@@ -340,6 +342,7 @@ class StreamingFeatureSet(BaseFeatureSet):
 
         return cls(
             name=proto.name,
+            repository=proto.featureset_repository_name,
             entity=Entity._from_proto(proto.entity).name,
             data_sources=[ds.name for ds in streaming_def.data_sources],
             timestamp_column_name=streaming_def.timestamp_column_name,
@@ -420,6 +423,7 @@ class StreamingFeatureSet(BaseFeatureSet):
                 features=features,
                 entity=self._get_entity_definition(feature_registry),
                 feature_set_type=proto_featureset_type,
+                featureset_repository_name=self.repository,
             ),
             artifact_url,
         )
@@ -429,10 +433,12 @@ class StreamingFeatureSet(BaseFeatureSet):
 
         super()._validate()
 
+        # verify offline_scheduling_policy was set
+        if not self.offline_scheduling_policy:
+            raise FrogmlException("'offline_scheduling_policy' field must be set")
+
         # verify the cron expression is valid
-        if not self.offline_scheduling_policy or not croniter.croniter.is_valid(
-            self.offline_scheduling_policy
-        ):
+        if not croniter.croniter.is_valid(self.offline_scheduling_policy):
             raise FrogmlException(
                 f"offline scheduling policy "
                 f"'{self.offline_scheduling_policy}'"

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/feature_store/online/client.py

@@ -1,5 +1,6 @@
 import logging
 import math
+from io import StringIO
 from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple
 
 from frogml_proto.qwak.ecosystem.v0.ecosystem_pb2 import (
@@ -221,7 +222,7 @@ class OnlineClient:
                     request, metadata=self._metadata
                 )
                 features_df: pd.DataFrame = pd.read_json(
-                    response_df_json.pandas_df_as_json, orient="split"
+                    StringIO(response_df_json.pandas_df_as_json), orient="split"
                 )
                 results.append(
                     pd.concat(

{frogml_core-0.0.111 → frogml_core-0.0.112}/frogml_core/model_loggers/artifact_logger.py

@@ -1,5 +1,8 @@
-import urllib.request
+import logging
+from datetime import timedelta
 from typing import Optional
+import requests
+from requests import Response
 
 from frogml_proto.qwak.builds.build_url_pb2 import BuildVersioningTagsType
 from frogml_core.clients.build_orchestrator.client import BuildOrchestratorClient
@@ -13,8 +16,14 @@ from frogml_core.inner.model_loggers_utils import (
 )
 from frogml_proto.qwak.builds.builds_orchestrator_service_pb2 import (
     GetBuildVersioningUploadURLResponse,
+    GetBuildVersioningDownloadURLResponse,
 )
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+MAX_CHUNK_SIZE = 8 * 1_024
+
 
 def log_file(
     from_path: str,
@@ -80,21 +89,40 @@ def load_file(
     Returns:
         the path to the newly created data file
     """
+    logger.info(f"Loading file to {to_path}")
     if not validate_tag(tag):
         raise FrogmlException(
             "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
+            ),
         )
+        logger.info(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 FrogmlException(f"Unable to load save artifact locally: {str(error)}")

frogml_core-0.0.112/frogml_storage/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.112"

{frogml_core-0.0.111 → frogml_core-0.0.112}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "frogml-core"
-version = "0.0.111"
+version = "0.0.112"
 description = "frogml Core contains the necessary objects and communication tools for using the Jfrog ml Platform"
 authors = ["Jfrog"]
 readme = "README.md"

frogml_core-0.0.111/frogml_storage/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.0.111"