def test_environment_schema_with_docker_environment():
env = environments.DockerEnvironment(
base_image="a", python_dependencies=["b", "c"], registry_url="f"
)
serialized = EnvironmentSchema().dump(env)
deserialized = EnvironmentSchema().load(serialized)
assert isinstance(deserialized, environments.DockerEnvironment)
assert deserialized.registry_url == env.registry_url
assert deserialized.base_image == env.base_image
def test_deserialize_old_env_payload():
old = {
"executor": "prefect.engine.executors.LocalExecutor",
"executor_kwargs": {},
"__version__": "0.6.3",
"type": "RemoteEnvironment",
}
schema = EnvironmentSchema()
obj = schema.load(old)
assert isinstance(obj, environments.RemoteEnvironment)
assert obj.labels == set()
def test_deserialize_old_environments_still_work(cls_name):
"""Check that old removed environments can still be deserialzed in the agent"""
env = {
"type": cls_name,
"labels": ["prod"],
"executor": "prefect.engine.executors.LocalExecutor",
"__version__": "0.9.0",
"executor_kwargs": {},
}
schema = EnvironmentSchema()
obj = schema.load(env)
assert isinstance(obj, environments.Environment)
assert obj.labels == {"prod"}
assert obj.metadata == {}
def test_serialize_custom_environment():
class MyEnv(environments.Environment):
def __init__(self, x=5):
self.x = 5
super().__init__(labels=["foo", "bar"])
def custom_method(self):
pass
env = MyEnv()
schema = EnvironmentSchema()
serialized = schema.dump(env)
assert serialized["type"] == "CustomEnvironment"
assert set(serialized["labels"]) == set(["foo", "bar"])
obj = schema.load(serialized)
assert isinstance(obj, environments.Environment)
assert obj.labels == set(["foo", "bar"])
def get_flow_image(flow_run: GraphQLResult, default: str = None) -> str:
"""
Retrieve the image to use for this flow run deployment.
Args:
- flow_run (GraphQLResult): A GraphQLResult flow run object
- default (str, optional): A default image to use. If not specified,
The `prefecthq/prefect` image corresponding with the flow's prefect
version will be used.
Returns:
- str: a full image name to use for this flow run
Raises:
- ValueError: if deployment attempted on unsupported Storage type and `image` not
present in environment metadata
"""
from prefect.storage import Docker
from prefect.serialization.storage import StorageSchema
from prefect.serialization.run_config import RunConfigSchema
from prefect.serialization.environment import EnvironmentSchema
has_run_config = getattr(flow_run, "run_config", None) is not None
has_environment = getattr(flow_run.flow, "environment", None) is not None
storage = StorageSchema().load(flow_run.flow.storage)
# Not having an environment implies run-config based flow, even if
# run_config is None.
if has_run_config or not has_environment:
# Precedence:
# - Image on docker storage
# - Image on run_config
# - Provided default
# - `prefecthq/prefect` for flow's core version
if isinstance(storage, Docker):
return storage.name
if has_run_config:
run_config = RunConfigSchema().load(flow_run.run_config)
if getattr(run_config, "image", None) is not None:
return run_config.image
if default is not None:
return default
# core_version should always be present, but just in case
version = flow_run.flow.get("core_version") or "latest"
cleaned_version = version.split("+")[0]
return f"prefecthq/prefect:{cleaned_version}"
else:
environment = EnvironmentSchema().load(flow_run.flow.environment)
if hasattr(environment, "metadata") and hasattr(
environment.metadata, "image"):
return environment.metadata.get("image")
elif isinstance(storage, Docker):
return storage.name
raise ValueError(
f"Storage for flow run {flow_run.id} is not of type Docker and "
f"environment has no `image` attribute in the metadata field.")
def test_serialize_custom_environment():
class MyEnv(environments.Environment):
def __init__(self, x=5):
self.x = 5
super().__init__(labels=["b", "c", "a"], metadata={"test": "here"})
def custom_method(self):
pass
env = MyEnv()
schema = EnvironmentSchema()
serialized = schema.dump(env)
assert serialized["type"] == "CustomEnvironment"
assert serialized["labels"] == ["a", "b", "c"]
assert serialized["metadata"] == {"test": "here"}
obj = schema.load(serialized)
assert isinstance(obj, environments.Environment)
assert obj.labels == {"a", "b", "c"}
assert obj.metadata == {"test": "here"}
def get_flow_image(flow_run: GraphQLResult) -> str:
"""
Retrieve the image to use for this flow run deployment.
Args:
- flow_run (GraphQLResult): A GraphQLResult flow run object
Returns:
- str: a full image name to use for this flow run
Raises:
- ValueError: if deployment attempted on unsupported Storage type and `image` not
present in environment metadata
"""
from prefect.environments.storage import Docker
from prefect.serialization.storage import StorageSchema
from prefect.serialization.run_config import RunConfigSchema
from prefect.serialization.environment import EnvironmentSchema
has_run_config = getattr(flow_run.flow, "run_config", None) is not None
has_environment = getattr(flow_run.flow, "environment", None) is not None
storage = StorageSchema().load(flow_run.flow.storage)
# Not having an environment implies run-config based flow, even if
# run_config is None.
if has_run_config or not has_environment:
if isinstance(storage, Docker):
return storage.name
elif has_run_config:
run_config = RunConfigSchema().load(flow_run.flow.run_config)
if getattr(run_config, "image", None) is not None:
return run_config.image
# No image found on run-config, and no environment present. Use default.
# core_version should always be present, but just in case
version = flow_run.flow.get("core_version") or "latest"
cleaned_version = version.split("+")[0]
return f"prefecthq/prefect:all_extras-{cleaned_version}"
else:
environment = EnvironmentSchema().load(flow_run.flow.environment)
if hasattr(environment, "metadata") and hasattr(environment.metadata, "image"):
return environment.metadata.get("image")
elif isinstance(storage, Docker):
return storage.name
raise ValueError(
f"Storage for flow run {flow_run.id} is not of type Docker and "
f"environment has no `image` attribute in the metadata field."
)
def test_environment_schema_with_local_environment():
env = environments.LocalEnvironment(encryption_key=FERNET_KEY)
serialized = EnvironmentSchema().dump(env)
deserialized = EnvironmentSchema().load(serialized)
assert isinstance(deserialized, environments.LocalEnvironment)
assert deserialized.encryption_key == FERNET_KEY
def test_serialize_custom_environment():
class MyEnv(environments.Environment):
def __init__(self, x=5):
self.x = 5
<<<<<<< HEAD
super().__init__(labels=["b", "c", "a"], metadata={"test": "here"})
=======
super().__init__(labels=["foo", "bar"], metadata={"test": "here"})
>>>>>>> prefect clone
def custom_method(self):
pass
env = MyEnv()
schema = EnvironmentSchema()
serialized = schema.dump(env)
assert serialized["type"] == "CustomEnvironment"
<<<<<<< HEAD
assert serialized["labels"] == ["a", "b", "c"]
=======
assert set(serialized["labels"]) == set(["foo", "bar"])
>>>>>>> prefect clone
assert serialized["metadata"] == {"test": "here"}
obj = schema.load(serialized)
assert isinstance(obj, environments.Environment)
<<<<<<< HEAD
assert obj.labels == {"a", "b", "c"}
assert obj.metadata == {"test": "here"}