Fixed Ray compute kwargs

CuulCat · May 2, 2024 · e7227bc · e7227bc
1 parent ec67009
commit e7227bc
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/superduperdb/backends/ray/compute.py b/superduperdb/backends/ray/compute.py
@@ -55,7 +55,10 @@ def _dependable_remote_job(function, *args, **kwargs):
                     ray.wait(dependencies)
             return function(*args, **kwargs)
 
-        remote_function = ray.remote(_dependable_remote_job, **compute_kwargs)
+        if compute_kwargs:
+            remote_function = ray.remote(**compute_kwargs)(_dependable_remote_job)
+        else:
+            remote_function = ray.remote(_dependable_remote_job)
         future = remote_function.remote(function, *args, **kwargs)
         task_id = str(future.task_id().hex())
         self._futures_collection[task_id] = future

diff --git a/superduperdb/components/model.py b/superduperdb/components/model.py
@@ -190,6 +190,7 @@ class Trainer(Component):
     prefetch_size: int = 1000
     prefetch_factor: int = 100
     in_memory: bool = True
+    compute_kwargs: t.Dict = dc.field(default_factory=dict)
 
     @abstractmethod
     def fit(
@@ -232,11 +233,15 @@ def fit_in_db_job(
         db: Datalayer,
         dependencies: t.Sequence[Job] = (),
     ):
+        if self.trainer:
+            compute_kwargs = self.trainer.compute_kwargs or {}
+        else:
+            compute_kwargs = {}
         job = ComponentJob(
             component_identifier=self.identifier,
             method_name='fit_in_db',
             type_id='model',
-            kwargs={},
+            compute_kwargs=compute_kwargs,
         )
         job(db, dependencies)
         return job