Skip to content

Commit

Permalink
fix other tests
Browse files Browse the repository at this point in the history
  • Loading branch information
kevinzwang committed Oct 5, 2024
1 parent 2e7f0b2 commit df2507f
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions tests/actor_pool/test_pyactor_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from daft.execution.execution_step import StatefulUDFProject
from daft.expressions import ExpressionsProjection
from daft.runners.partitioning import PartialPartitionMetadata
from daft.runners.pyrunner import PyActorPool, PyRunner
from daft.runners.pyrunner import PyActorPool, PyRunner, PyRunnerResources
from daft.table import MicroPartition


Expand All @@ -25,7 +25,9 @@ def __call__(self, x):

def test_pyactor_pool():
projection = ExpressionsProjection([MyStatefulUDF(daft.col("x"))])
pool = PyActorPool("my-pool", 1, ResourceRequest(num_cpus=1), projection)
pool = PyActorPool(
"my-pool", 1, ResourceRequest(num_cpus=1), [PyRunnerResources(num_cpus=1, gpus={}, memory_bytes=0)], projection
)
initial_partition = MicroPartition.from_pydict({"x": [1, 1, 1]})
ppm = PartialPartitionMetadata(num_rows=None, size_bytes=None)
instr = StatefulUDFProject(projection=projection)
Expand Down Expand Up @@ -68,7 +70,7 @@ def test_pyactor_pool_not_enough_resources():
runner = get_context().runner()
assert isinstance(runner, PyRunner)

with pytest.raises(RuntimeError, match=f"Requested {float(cpu_count + 1)} CPUs but found only"):
with pytest.raises(RuntimeError, match=f"Not enough resources available to admit {cpu_count + 1} actors"):
with runner.actor_pool_context(
"my-pool", ResourceRequest(num_cpus=1), ResourceRequest(), cpu_count + 1, projection
) as _:
Expand Down

0 comments on commit df2507f

Please sign in to comment.