Replace all trainer smart_open+get_artifact_link with get_artifact_handle

GitOrigin-RevId: d192abacedea7430175f4e9920530cc80c32feb6

Author: mikeknep

src/gretel_trainer/benchmark/sdk_extras.py

@@ -3,8 +3,6 @@
 
 from typing import Any
 
-import smart_open
-
 from gretel_client.projects.jobs import (
     ACTIVE_STATES,
     END_STATES,
@@ -41,9 +39,8 @@ def run_evaluate(
     job_status = await_job(run_identifier, evaluate_model, "evaluation", wait)
     if job_status in END_STATES and job_status != Status.COMPLETED:
         raise BenchmarkException("Evaluate failed")
-    return json.loads(
-        smart_open.open(evaluate_model.get_artifact_link("report_json")).read()
-    )
+    with evaluate_model.get_artifact_handle("report_json") as report:
+        return json.loads(report.read())
 
 
 def _make_evaluate_config(run_identifier: str) -> dict:

src/gretel_trainer/relational/strategies/common.py

@@ -40,9 +40,8 @@ def read_report_json_data(model: Model, report_path: Path) -> Optional[dict]:
 
 def _get_report_json(model: Model) -> Optional[dict]:
     try:
-        return json.loads(
-            smart_open.open(model.get_artifact_link("report_json")).read()
-        )
+        with model.get_artifact_handle("report_json") as report:
+            return json.loads(report.read())
     except:
         logger.warning("Failed to fetch model evaluation report JSON.")
         return None

src/gretel_trainer/relational/tasks/classify.py

@@ -151,8 +151,8 @@ def _write_results(self, job: Job, table: str) -> None:
 
         destpath = self.out_dir / filename
 
-        with smart_open.open(
-            job.get_artifact_link(artifact_name), "rb"
-        ) as src, smart_open.open(str(destpath), "wb") as dest:
+        with job.get_artifact_handle(artifact_name) as src, smart_open.open(
+            str(destpath), "wb"
+        ) as dest:
             shutil.copyfileobj(src, dest)
         self.result_filepaths.append(destpath)

src/gretel_trainer/runner.py

@@ -38,7 +38,6 @@
 from typing import List, Optional, Tuple, Union
 
 import pandas as pd
-import smart_open
 
 from gretel_client.projects import Project
 from gretel_client.projects.jobs import ACTIVE_STATES
@@ -213,9 +212,7 @@ def _update_job_status(self):
                 report = current_model.peek_report()
 
                 if report is None:
-                    with smart_open.open(
-                        current_model.get_artifact_link("report_json")
-                    ) as fin:
+                    with current_model.get_artifact_handle("report_json") as fin:
                         report = json.loads(fin.read())
 
                 sqs = report["synthetic_data_quality_score"]["score"]

tests/benchmark/conftest.py

@@ -48,12 +48,17 @@ def project():
 
 
 @pytest.fixture()
-def evaluate_report_path():
+def evaluate_report_handle():
     report = {"synthetic_data_quality_score": {"score": 95}}
     with tempfile.NamedTemporaryFile() as f:
         with open(f.name, "w") as j:
             json.dump(report, j)
-        yield f.name
+
+        ctxmgr = Mock()
+        ctxmgr.__enter__ = Mock(return_value=f)
+        ctxmgr.__exit__ = Mock(return_value=False)
+
+        yield ctxmgr
 
 
 @pytest.fixture()

tests/benchmark/test_benchmark.py

@@ -81,11 +81,11 @@ class SharedDictLstm(GretelModel):
     }
 
 
-def test_run_with_gretel_dataset(working_dir, project, evaluate_report_path, iris):
+def test_run_with_gretel_dataset(working_dir, project, evaluate_report_handle, iris):
     evaluate_model = Mock(
         status=Status.COMPLETED,
     )
-    evaluate_model.get_artifact_link.return_value = evaluate_report_path
+    evaluate_model.get_artifact_handle.return_value = evaluate_report_handle
     project.create_model_obj.side_effect = [evaluate_model]
 
     session = compare(
@@ -107,11 +107,11 @@ def test_run_with_gretel_dataset(working_dir, project, evaluate_report_path, iri
     assert result["SQS"] == 95
 
 
-def test_run_with_custom_csv_dataset(working_dir, project, evaluate_report_path, df):
+def test_run_with_custom_csv_dataset(working_dir, project, evaluate_report_handle, df):
     evaluate_model = Mock(
         status=Status.COMPLETED,
     )
-    evaluate_model.get_artifact_link.return_value = evaluate_report_path
+    evaluate_model.get_artifact_handle.return_value = evaluate_report_handle
     project.create_model_obj.side_effect = [evaluate_model]
 
     with tempfile.NamedTemporaryFile() as f:
@@ -137,11 +137,11 @@ def test_run_with_custom_csv_dataset(working_dir, project, evaluate_report_path,
     assert result["SQS"] == 95
 
 
-def test_run_with_custom_psv_dataset(working_dir, project, evaluate_report_path, df):
+def test_run_with_custom_psv_dataset(working_dir, project, evaluate_report_handle, df):
     evaluate_model = Mock(
         status=Status.COMPLETED,
     )
-    evaluate_model.get_artifact_link.return_value = evaluate_report_path
+    evaluate_model.get_artifact_handle.return_value = evaluate_report_handle
     project.create_model_obj.side_effect = [evaluate_model]
 
     with tempfile.NamedTemporaryFile() as f:
@@ -168,12 +168,12 @@ def test_run_with_custom_psv_dataset(working_dir, project, evaluate_report_path,
 
 
 def test_run_with_custom_dataframe_dataset(
-    working_dir, project, evaluate_report_path, df
+    working_dir, project, evaluate_report_handle, df
 ):
     evaluate_model = Mock(
         status=Status.COMPLETED,
     )
-    evaluate_model.get_artifact_link.return_value = evaluate_report_path
+    evaluate_model.get_artifact_handle.return_value = evaluate_report_handle
     project.create_model_obj.side_effect = [evaluate_model]
 
     dataset = create_dataset(df, datatype="tabular", name="pets")
@@ -205,7 +205,7 @@ def test_run_with_custom_dataframe_dataset(
 
 @pytest.mark.parametrize("benchmark_model", [GretelLSTM, TailoredActgan])
 def test_run_happy_path_gretel_sdk(
-    benchmark_model, working_dir, iris, project, evaluate_report_path
+    benchmark_model, working_dir, iris, project, evaluate_report_handle
 ):
     record_handler = Mock(
         status=Status.COMPLETED,
@@ -221,7 +221,7 @@ def test_run_happy_path_gretel_sdk(
     evaluate_model = Mock(
         status=Status.COMPLETED,
     )
-    evaluate_model.get_artifact_link.return_value = evaluate_report_path
+    evaluate_model.get_artifact_handle.return_value = evaluate_report_handle
 
     project.create_model_obj.side_effect = [model, evaluate_model]