Skip to content

Commit 03b51be

Browse files
harsh97harsh97
committed
Switching to fsclient sdk
1 parent 51d3e6d commit 03b51be

17 files changed

+114
-79
lines changed

‎ads/common/oci_mixin.py‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ class MergeStrategy(Enum):
4141
MERGE = "merge"
4242

4343

44-
class OCIModelNotExists(Exception): # pragma: no cover
44+
class OCIModelNotExists(Exception): # pragma: no cover
4545
pass
4646

4747

@@ -383,7 +383,7 @@ class OCIModelMixin(OCISerializableMixin):
383383
"""
384384

385385
# Regex pattern matching the module name of an OCI model.
386-
OCI_MODEL_PATTERN = r"oci.[^.]+\.models[\..*]?"
386+
OCI_MODEL_PATTERN = r"(oci|feature_store_client).[^.]+\.models[\..*]?"
387387
# Constants
388388
CONS_COMPARTMENT_ID = "compartment_id"
389389

@@ -937,11 +937,11 @@ def get_work_request_response(
937937
return work_request_response
938938

939939
def wait_for_progress(
940-
self,
941-
work_request_id: str,
942-
num_steps: int = DEFAULT_WORKFLOW_STEPS,
943-
max_wait_time: int = DEFAULT_WAIT_TIME,
944-
poll_interval: int = DEFAULT_POLL_INTERVAL
940+
self,
941+
work_request_id: str,
942+
num_steps: int = DEFAULT_WORKFLOW_STEPS,
943+
max_wait_time: int = DEFAULT_WAIT_TIME,
944+
poll_interval: int = DEFAULT_POLL_INTERVAL,
945945
):
946946
"""Waits for the work request progress bar to be completed.
947947
@@ -969,10 +969,10 @@ def wait_for_progress(
969969
seconds_since = time.time() - start_time
970970
exceed_max_time = max_wait_time > 0 and seconds_since >= max_wait_time
971971
if exceed_max_time:
972-
logger.error(
973-
f"Max wait time ({max_wait_time} seconds) exceeded."
974-
)
975-
while not exceed_max_time and (not work_request_logs or len(work_request_logs) < num_steps):
972+
logger.error(f"Max wait time ({max_wait_time} seconds) exceeded.")
973+
while not exceed_max_time and (
974+
not work_request_logs or len(work_request_logs) < num_steps
975+
):
976976
time.sleep(poll_interval)
977977
new_work_request_logs = []
978978

‎ads/feature_store/dataset.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from great_expectations.core import ExpectationSuite
1111

1212
from ads import deprecated
13-
from oci.feature_store.models import (
13+
from feature_store_client.feature_store.models import (
1414
DatasetFeatureGroupCollection,
1515
DatasetFeatureGroupSummary,
1616
)
@@ -720,7 +720,7 @@ def create(self, validate_sql=False, **kwargs) -> "Dataset":
720720
----------
721721
kwargs
722722
Additional kwargs arguments.
723-
Can be any attribute that `oci.feature_store.models.Dataset` accepts.
723+
Can be any attribute that `feature_store.models.Dataset` accepts.
724724
validate_sql:
725725
Boolean value indicating whether to validate sql before creating dataset
726726
@@ -821,7 +821,7 @@ def update(self, **kwargs) -> "Dataset":
821821
----------
822822
kwargs
823823
Additional kwargs arguments.
824-
Can be any attribute that `oci.feature_store.models.Dataset` accepts.
824+
Can be any attribute that `feature_store.models.Dataset` accepts.
825825
826826
Returns
827827
-------

‎ads/feature_store/dataset_job.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,7 @@ def _to_oci_fs_dataset_run(self, **kwargs):
113113
114114
kwargs
115115
Additional kwargs arguments.
116-
Can be any attribute that `oci.feature_store.models.DatasetJob` accepts.
116+
Can be any attribute that `feature_store.models.DatasetJob` accepts.
117117
118118
Returns
119119
-------
@@ -337,7 +337,7 @@ def create(self, **kwargs) -> "DatasetJob":
337337
----------
338338
kwargs
339339
Additional kwargs arguments.
340-
Can be any attribute that `oci.feature_store.models.DatasetJob` accepts.
340+
Can be any attribute that `feature_store.models.DatasetJob` accepts.
341341
342342
Returns
343343
-------
@@ -370,7 +370,7 @@ def update(self, **kwargs) -> "DatasetJob":
370370
----------
371371
kwargs
372372
Additional kwargs arguments.
373-
Can be any attribute that `oci.feature_store.models.DatasetJob` accepts.
373+
Can be any attribute that `feature_store.models.DatasetJob` accepts.
374374
375375
Returns
376376
-------

‎ads/feature_store/entity.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ def _to_oci_fs_entity(self, **kwargs):
111111
112112
kwargs
113113
Additional kwargs arguments.
114-
Can be any attribute that `oci.feature_store.models.Entity` accepts.
114+
Can be any attribute that `feature_store.models.Entity` accepts.
115115
116116
Returns
117117
-------
@@ -252,7 +252,7 @@ def create(self, **kwargs) -> "Entity":
252252
----------
253253
kwargs
254254
Additional kwargs arguments.
255-
Can be any attribute that `oci.feature_store.models.Entity` accepts.
255+
Can be any attribute that `feature_store.models.Entity` accepts.
256256
257257
Returns
258258
-------

‎ads/feature_store/feature_group.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -694,7 +694,7 @@ def create(self, **kwargs) -> "FeatureGroup":
694694
----------
695695
kwargs
696696
Additional kwargs arguments.
697-
Can be any attribute that `oci.feature_store.models.FeatureGroup` accepts.
697+
Can be any attribute that `feature_store.models.FeatureGroup` accepts.
698698
699699
Returns
700700
-------
@@ -784,7 +784,7 @@ def update(self, **kwargs) -> "FeatureGroup":
784784
----------
785785
kwargs
786786
Additional kwargs arguments.
787-
Can be any attribute that `oci.feature_store.models.FeatureGroup` accepts.
787+
Can be any attribute that `feature_store.models.FeatureGroup` accepts.
788788
789789
Returns
790790
-------

‎ads/feature_store/feature_group_job.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ def _to_oci_fs_feature_group_run(self, **kwargs):
121121
122122
kwargs
123123
Additional kwargs arguments.
124-
Can be any attribute that `oci.feature_store.models.FeatureGroupJob` accepts.
124+
Can be any attribute that `feature_store.models.FeatureGroupJob` accepts.
125125
126126
Returns
127127
-------
@@ -406,7 +406,7 @@ def create(self, **kwargs) -> "FeatureGroupJob":
406406
----------
407407
kwargs
408408
Additional kwargs arguments.
409-
Can be any attribute that `oci.feature_store.models.FeatureGroupJob` accepts.
409+
Can be any attribute that `feature_store.models.FeatureGroupJob` accepts.
410410
411411
Returns
412412
-------
@@ -441,7 +441,7 @@ def update(self, **kwargs) -> "FeatureGroupJob":
441441
----------
442442
kwargs
443443
Additional kwargs arguments.
444-
Can be any attribute that `oci.feature_store.models.FeatureGroupJob` accepts.
444+
Can be any attribute that `feature_store.models.FeatureGroupJob` accepts.
445445
446446
Returns
447447
-------

‎ads/feature_store/feature_store.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ def _to_oci_fs(self, **kwargs):
110110
111111
kwargs
112112
Additional kwargs arguments.
113-
Can be any attribute that `oci.feature_store.models.FeatureStore` accepts.
113+
Can be any attribute that `feature_store.models.FeatureStore` accepts.
114114
115115
Returns
116116
-------
@@ -271,7 +271,7 @@ def create(self, **kwargs) -> "FeatureStore":
271271
----------
272272
kwargs
273273
Additional kwargs arguments.
274-
Can be any attribute that `oci.feature_store.models.FeatureStore` accepts.
274+
Can be any attribute that `feature_store.models.FeatureStore` accepts.
275275
276276
Returns
277277
-------

‎ads/feature_store/mixin/oci_feature_store.py‎

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,21 @@
11
#!/usr/bin/env python
22
# -*- coding: utf-8; -*-
3+
from types import MethodType
4+
35
from ads.common.decorator.utils import class_or_instance_method
6+
from oci.signer import AbstractBaseSigner
47

58
# Copyright (c) 2023 Oracle and/or its affiliates.
69
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
710

811
import logging
12+
import email.utils
913
import os
14+
import oci
15+
import feature_store_client.feature_store as fs
1016

1117
logger = logging.getLogger(__name__)
1218
from ads.common.oci_mixin import OCIModelMixin
13-
import oci.feature_store
1419
import yaml
1520

1621

@@ -34,9 +39,7 @@ class OCIFeatureStoreMixin(OCIModelMixin):
3439
SERVICE_ENDPOINT = "service_endpoint"
3540

3641
@classmethod
37-
def init_client(
38-
cls, **kwargs
39-
) -> oci.feature_store.feature_store_client.FeatureStoreClient:
42+
def init_client(cls, **kwargs) -> fs.feature_store_client.FeatureStoreClient:
4043
default_kwargs: dict = cls._get_auth().get("client_kwargs", {})
4144

4245
fs_service_endpoint = (
@@ -70,13 +73,15 @@ def init_client(
7073
if fs_service_endpoint:
7174
kwargs[cls.SERVICE_ENDPOINT] = fs_service_endpoint
7275

73-
client = cls._init_client(
74-
client=oci.feature_store.feature_store_client.FeatureStoreClient, **kwargs
76+
client: fs.FeatureStoreClient = cls._init_client(
77+
client=fs.FeatureStoreClient, **kwargs
7578
)
79+
signer: oci.Signer = client.base_client.signer
80+
signer.do_request_sign = MethodType(fs_do_request_sign, signer)
7681
return client
7782

7883
@property
79-
def client(self) -> oci.feature_store.feature_store_client.FeatureStoreClient:
84+
def client(self) -> fs.feature_store_client.FeatureStoreClient:
8085
return super().client
8186

8287
@class_or_instance_method
@@ -119,3 +124,22 @@ def list_resource(
119124
**kwargs,
120125
).data
121126
return [cls.from_oci_model(item) for item in items]
127+
128+
129+
def inject_missing_headers(request):
130+
# Inject date, host, and content-type if missing
131+
date = email.utils.formatdate(usegmt=True)
132+
if request.path_url.startswith("/20230101"):
133+
request.headers.setdefault("x-date", date)
134+
request.headers.setdefault(
135+
"path", request.method.lower() + " " + request.path_url
136+
)
137+
request.headers.setdefault("date", date)
138+
139+
140+
def fs_do_request_sign(self, request, enforce_content_headers=True):
141+
inject_missing_headers(request)
142+
do_request_sign = MethodType(AbstractBaseSigner.do_request_sign, self)
143+
return do_request_sign(request, enforce_content_headers)
144+
145+
# inject_missing_headers_og(request, sign_body, enforce_content_headers)

‎ads/feature_store/service/oci_dataset.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,17 @@
66

77
import datetime
88

9-
import oci
10-
from oci.feature_store.models import CreateDatasetDetails, UpdateDatasetDetails
9+
import feature_store_client.feature_store as fs
10+
from feature_store_client.feature_store.models import (
11+
Dataset,
12+
CreateDatasetDetails,
13+
UpdateDatasetDetails,
14+
)
1115

1216
from ads.feature_store.mixin.oci_feature_store import OCIFeatureStoreMixin
1317

1418

15-
class OCIDataset(OCIFeatureStoreMixin, oci.feature_store.models.Dataset):
19+
class OCIDataset(OCIFeatureStoreMixin, Dataset):
1620
"""Represents an OCI Data Science dataset.
1721
This class contains all attributes of the `oci.data_science.models.Dataset`.
1822
The main purpose of this class is to link the `oci.data_science.models.Dataset`
@@ -59,7 +63,7 @@ class OCIDataset(OCIFeatureStoreMixin, oci.feature_store.models.Dataset):
5963
"""
6064

6165
@property
62-
def client(self) -> oci.feature_store.feature_store_client.FeatureStoreClient:
66+
def client(self) -> fs.feature_store_client.FeatureStoreClient:
6367
return super().client
6468

6569
def create(self) -> "OCIDataset":

‎ads/feature_store/service/oci_dataset_job.py‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,11 @@
88
import logging
99
import time
1010

11-
import oci.feature_store
12-
from oci.feature_store.models import CreateDatasetJobDetails, CompleteDatasetJobDetails
11+
from feature_store_client.feature_store.models import (
12+
CreateDatasetJobDetails,
13+
CompleteDatasetJobDetails,
14+
DatasetJob,
15+
)
1316

1417
from ads.feature_store.mixin.oci_feature_store import OCIFeatureStoreMixin
1518

@@ -18,7 +21,7 @@
1821
SLEEP_INTERVAL = 3
1922

2023

21-
class OCIDatasetJob(OCIFeatureStoreMixin, oci.feature_store.models.DatasetJob):
24+
class OCIDatasetJob(OCIFeatureStoreMixin, DatasetJob):
2225
"""Represents an OCI Data Science DatasetJob.
2326
This class contains all attributes of the `oci.data_science.models.DatasetJob`.
2427
The main purpose of this class is to link the `oci.data_science.models.DatasetJob`
@@ -62,8 +65,8 @@ class OCIDatasetJob(OCIFeatureStoreMixin, oci.feature_store.models.DatasetJob):
6265
"""
6366

6467
TERMINAL_STATES = [
65-
oci.feature_store.models.DatasetJob.LIFECYCLE_STATE_SUCCEEDED,
66-
oci.feature_store.models.DatasetJob.LIFECYCLE_STATE_FAILED,
68+
DatasetJob.LIFECYCLE_STATE_SUCCEEDED,
69+
DatasetJob.LIFECYCLE_STATE_FAILED,
6770
]
6871

6972
def __init__(self, **kwargs) -> None:
@@ -72,7 +75,7 @@ def __init__(self, **kwargs) -> None:
7275
Parameters
7376
----------
7477
kwargs:
75-
Same as kwargs in oci.feature_store.models.OCIDatasetJob.
78+
Same as kwargs in feature_store.models.OCIDatasetJob.
7679
Keyword arguments are passed into OCI feature group job model to initialize the properties.
7780
7881
"""

0 commit comments

Comments
 (0)