[knobs] Fix environment propagation & scope() API #6664
Conversation
python/triton/knobs.py
Outdated
| os.unsetenv(self.key) | ||
| else: | ||
| os.putenv(self.key, value) | ||
| setenv(self.key, value) |
There was a problem hiding this comment.
But if propagate_env is False this doesn't set the value at all?
There was a problem hiding this comment.
Yeah that's right, that was intentional. In the least this is needed for testing, but I figured some users might want to opt out of the environment updates altogether (and just use the config classes as essentially a dataclass). Do you think that's ok? Or should we try to get rid of that flag (meaning I'd have to rethink the tests)?
There was a problem hiding this comment.
Oh I see there is __set__, set and setenv which all do subtly different things. Could __set__ just be:
def __set__(self, obj: object, value: Union[SetType, Env]) -> None:
if isinstance(value, Env):
obj.__dict__.pop(self.name, None)
else:
obj.__dict__[self.name] = value
setenv(self.key, repr(value) if value is not None else value)and get rid of set altogether
There was a problem hiding this comment.
Yeah, good idea I think so-- it might need to look slightly different because e.g. env_class would need a custom repr
| if env_val := toenv(value): | ||
| setenv(self.key, env_val[0]) |
There was a problem hiding this comment.
Almost exactly the same impl you gave @peterbell10 , but needed a function to convert the value (e.g. "1" instead of "True" and to ignore when e.g. value is a class for env_class)
| @@ -226,20 +245,29 @@ def test_nvidia_tool(fresh_knobs, tmp_path, monkeypatch): | |||
| triton_root = Path(fresh_knobs.__file__).parent | |||
| default_ptxas = triton_root / "backends/nvidia/bin/ptxas" | |||
|
|
|||
| assert default_ptxas.exists() | |||
There was a problem hiding this comment.
@CliveUnger I think nixing this assert should address your concern, since everything else is just validating Paths/strings (follow up from here for any other readers).
In the original PR introducing these knobs @peterbell10 raised a good point that each knob had a side effect so that there was no real way to reset them.
I thought I fixed this with the
scope()API, but it turns out I didn't-- instead I actually just broke the environment propagation. Specifically I was usingputenvandunsetenvto propagate values, but those don't updateos.environor the result ofos.getenv.I've added a test to verify that
os.environgets properly updated now, and addressed a bug inscope()which wasn't restoring the environment.I also understand that env propagation might not be desirable to everyone (since it's a side effect), so I've added a flag to disable the feature.