def test_serialize_dask_environment_with_customized_workers():
env = environments.DaskKubernetesEnvironment(min_workers=10, max_workers=60)
schema = DaskKubernetesEnvironmentSchema()
serialized = schema.dump(env)
assert serialized
assert serialized["__version__"] == prefect.__version__
assert serialized["min_workers"] == 10
assert serialized["max_workers"] == 60
new = schema.load(serialized)
assert new.min_workers == 10
assert new.max_workers == 60
def test_serialize_dask_environment_with_private_registry():
env = environments.DaskKubernetesEnvironment(private_registry=True,
docker_secret="FOO")
schema = DaskKubernetesEnvironmentSchema()
serialized = schema.dump(env)
assert serialized
assert serialized["__version__"] == prefect.__version__
assert serialized["private_registry"] is True
assert serialized["docker_secret"] == "FOO"
new = schema.load(serialized)
assert new.private_registry is True
assert new.docker_secret == "FOO"
def test_serialize_dask_environment():
env = environments.DaskKubernetesEnvironment()
schema = DaskKubernetesEnvironmentSchema()
serialized = schema.dump(env)
assert serialized
assert serialized["__version__"] == prefect.__version__
assert serialized["docker_secret"] is None
assert serialized["min_workers"] == 1
assert serialized["max_workers"] == 2
new = schema.load(serialized)
assert new.private_registry is False
assert new.docker_secret is None
assert new.min_workers == 1
assert new.max_workers == 2
def test_serialize_dask_env_with_custom_specs():
with tempfile.TemporaryDirectory() as directory:
with open(os.path.join(directory, "scheduler.yaml"), "w+") as f:
f.write("scheduler")
with open(os.path.join(directory, "worker.yaml"), "w+") as f:
f.write("worker")
env = environments.DaskKubernetesEnvironment(
scheduler_spec_file=os.path.join(directory, "scheduler.yaml"),
worker_spec_file=os.path.join(directory, "worker.yaml"),
)
schema = DaskKubernetesEnvironmentSchema()
serialized = schema.dump(env)
deserialized = schema.load(serialized)
assert isinstance(deserialized, environments.DaskKubernetesEnvironment)
def test_serialize_dask_environment_with_labels():
env = environments.DaskKubernetesEnvironment(labels=["b", "c", "a"])
schema = DaskKubernetesEnvironmentSchema()
serialized = schema.dump(env)
assert serialized
assert serialized["__version__"] == prefect.__version__
assert serialized["docker_secret"] is None
assert serialized["min_workers"] == 1
assert serialized["max_workers"] == 2
# labels should be sorted in the serialized obj
assert serialized["labels"] == ["a", "b", "c"]
new = schema.load(serialized)
assert new.private_registry is False
assert new.docker_secret is None
assert new.min_workers == 1
assert new.max_workers == 2
assert new.labels == {"a", "b", "c"}
def test_serialize_dask_environment():
env = environments.DaskKubernetesEnvironment()
schema = DaskKubernetesEnvironmentSchema()
serialized = schema.dump(env)
assert serialized
assert serialized["__version__"] == prefect.__version__
assert serialized["docker_secret"] is None
assert serialized["min_workers"] == 1
assert serialized["max_workers"] == 2
assert serialized["labels"] == []
assert serialized["metadata"] == {}
new = schema.load(serialized)
assert new.private_registry is False
assert new.docker_secret is None
assert new.min_workers == 1
assert new.max_workers == 2
assert new.labels == set()
assert new.scheduler_spec_file is None
assert new.worker_spec_file is None
env = environments.DaskKubernetesEnvironment(
scheduler_spec_file=os.path.join(directory, "scheduler.yaml"),
worker_spec_file=os.path.join(directory, "worker.yaml"),
)
schema = DaskKubernetesEnvironmentSchema()
serialized = schema.dump(env)
deserialized = schema.load(serialized)
assert isinstance(deserialized, environments.DaskKubernetesEnvironment)
def test_serialize_dask_environment_with_labels():
<<<<<<< HEAD
env = environments.DaskKubernetesEnvironment(labels=["b", "c", "a"])
=======
env = environments.DaskKubernetesEnvironment(labels=["a", "b", "c"])
>>>>>>> prefect clone
schema = DaskKubernetesEnvironmentSchema()
serialized = schema.dump(env)
assert serialized
assert serialized["__version__"] == prefect.__version__
assert serialized["docker_secret"] is None
assert serialized["min_workers"] == 1
assert serialized["max_workers"] == 2
<<<<<<< HEAD
# labels should be sorted in the serialized obj
assert serialized["labels"] == ["a", "b", "c"]
=======
