Use get_artifact_handle instead of smart_open directly

gretelai · Aug 9, 2023 · edddcf0 · edddcf0
1 parent ecfef78
commit edddcf0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,6 @@
 boto3~=1.20
 dask[dataframe]==2023.5.1
-gretel-client>=0.16.0
+gretel-client>=0.16.12
 jinja2~=3.1
 networkx~=3.0
 numpy~=1.20

diff --git a/src/gretel_trainer/relational/sdk_extras.py b/src/gretel_trainer/relational/sdk_extras.py
@@ -52,9 +52,8 @@ def download_file_artifact(
         artifact_name: str,
         out_path: Union[str, Path],
     ) -> None:
-        download_link = gretel_object.get_artifact_link(artifact_name)
         try:
-            with smart_open.open(download_link, "rb") as src, smart_open.open(
+            with gretel_object.get_artifact_handle("data") as src, smart_open.open(
                 out_path, "wb"
             ) as dest:
                 shutil.copyfileobj(src, dest)
@@ -80,7 +79,7 @@ def sqs_score_from_full_report(self, report: dict[str, Any]) -> Optional[int]:
                     return field_dict["value"]
 
     def get_record_handler_data(self, record_handler: RecordHandler) -> pd.DataFrame:
-        with smart_open.open(record_handler.get_artifact_link("data"), "rb") as data:
+        with record_handler.get_artifact_handle("data") as data:
             return pd.read_csv(data)
 
     def start_job_if_possible(