diff --git a/metaflow/util.py b/metaflow/util.py index 4cc96d8b0cf..9f849a4af89 100644 --- a/metaflow/util.py +++ b/metaflow/util.py @@ -573,6 +573,8 @@ def to_pod(value): # Prevent circular imports from metaflow.parameters import DeployTimeField + if value is None: + return value if isinstance(value, (str, int, float)): return value if isinstance(value, dict): diff --git a/test/unit/test_graph_structure.py b/test/unit/test_graph_structure.py index 325985ee161..d7e3c0a27ad 100644 --- a/test/unit/test_graph_structure.py +++ b/test/unit/test_graph_structure.py @@ -693,3 +693,31 @@ def test_step_mutator_non_primitive_attribute_sanitized_in_output_steps(): decorators = steps_info["only"]["decorators"] payload_deco = next(d for d in decorators if d["name"].endswith("_payload_mutator")) assert isinstance(payload_deco["attributes"]["payload"], str) + + +# --------------------------------------------------------------------------- +# Tests: unset decorator attributes stay null in _graph_info +# --------------------------------------------------------------------------- + + +class _PartialResourcesFlow(FlowSpec): + @resources(cpu=2, memory=8192) + @step(start=True, end=True) + def only(self): + pass + + +def test_unset_decorator_attributes_stay_null_in_output_steps(): + """@resources defaults gpu/disk/shared_memory to None. output_steps() runs + the attribute dict through to_pod(), which must leave None alone: the + string "None" is truthy, so a consumer of _graph_info (the DAG card, the + UI, `current.graph`) reading attributes["gpu"] would conclude a GPU was + requested when the user never asked for one.""" + steps_info, _ = _PartialResourcesFlow._graph.output_steps() + decorators = steps_info["only"]["decorators"] + attrs = next(d for d in decorators if d["name"] == "resources")["attributes"] + assert attrs["cpu"] == 2 + assert attrs["memory"] == 8192 + assert attrs["gpu"] is None + assert attrs["disk"] is None + assert attrs["shared_memory"] is None diff --git a/test/unit/test_to_pod.py b/test/unit/test_to_pod.py index 3968dc1fbbe..98ce965b668 100644 --- a/test/unit/test_to_pod.py +++ b/test/unit/test_to_pod.py @@ -62,3 +62,20 @@ def test_to_pod_lambda_uses_qualname(): fn = lambda x: x # noqa: E731 result = to_pod(fn) assert "" in result + + +def test_to_pod_none_is_preserved(): + """None is already POD (JSON null) and must not be stringified.""" + assert to_pod(None) is None + + +def test_to_pod_none_nested_in_dict_and_list(): + """Unset decorator attributes are None; they must stay null in _graph_info.""" + assert to_pod({"gpu": None, "cpu": 1}) == {"gpu": None, "cpu": 1} + assert to_pod([1, None, "x"]) == [1, None, "x"] + assert to_pod({"outer": {"inner": [None]}}) == {"outer": {"inner": [None]}} + + +def test_to_pod_none_is_falsy_after_conversion(): + """The string "None" is truthy; callers testing the attribute must not be misled.""" + assert not to_pod({"gpu": None})["gpu"]