Combine sticky pgo key with job id (pytorch#154863)

Pull Request resolved: pytorch#154863
Approved by: https://github.com/Mingming-Ding

Author: bobrenjc93

torch/_dynamo/pgo.py

@@ -535,9 +535,6 @@ def get_cache_key() -> Optional[str]:
             )
         return f"{r}:{rank}:{tag}"
 
-    if r := torch.compiler.config.sticky_pgo_key:
-        return f"sticky:{r}:{rank}:{tag}"
-
     if (name_version := torch._utils_internal.get_mast_job_name_version()) is not None:
         mast_job_name, mast_job_version = name_version
         return f"mast:{mast_job_name}:{mast_job_version}:{rank}:{tag}"

torch/compiler/config.py

@@ -29,7 +29,10 @@
 # FB-internal note: you do NOT have to specify this explicitly specify this if
 # you run on MAST, we will automatically default this to
 # mast:MAST_JOB_NAME:MAST_JOB_VERSION.
-job_id: Optional[str] = Config(env_name_default="TORCH_COMPILE_JOB_ID", default=None)
+job_id: Optional[str] = Config(
+    env_name_default=["TORCH_COMPILE_JOB_ID", "TORCH_COMPILE_STICKY_PGO_KEY"],
+    default=None,
+)
 """
 Semantically, this should be an identifier that uniquely identifies, e.g., a
 training job.  You might have multiple attempts of the same job, e.g., if it was
@@ -74,15 +77,6 @@
 and force_parameter_static_shapes.
 """
 
-sticky_pgo_key: str = Config(
-    env_name_default="TORCH_COMPILE_STICKY_PGO_KEY", default=""
-)
-"""
-If you want to share PGO profiles across different jobs (and not just attempts), you can set
-this to a string that identifies the shared profile. This is useful if you want to share PGO profiles
-for models that are not identical, but are similar enough to share PGO profiles.
-"""
-
 unbacked_sources: str = Config(
     env_name_default="TORCH_COMPILE_UNBACKED_SOURCES", default=""
 )