Repository: pytorch/pytorch

Issue:
Create a CI workflow for XLA testing using the XLA test image
### 🚀 The feature, motivation and pitch

We have onboarded XLA testing as part of one of the existing Linux workflows. This has created some issues due to diverging test environments, where the same set of tests pass in the downstream, but fails in the upstream due to dependency conflicts and environment setup, etc. This happens because XLA testing was piggybacking on the existing ciflow (config) and image that were mainly built for PyTorch tests. To avoid this and do so unobtrusively, it would be good to create a dedicated CI workflow (customize settings without affecting other tests) for XLA testing and use the XLA test image (customize dependencies and settings without affecting other tests) that's used in the downstream tests. 

Ideally, any failures in the upstream should have been caught in the downstream while testing; we want to avoid XLA upstream test failures due to test environment setup/divergence, if possible.

### Alternatives

This proposes to use the same test image we use in the downstream (XLA) circle CI tests. We maintain and publishes the image to serve XLA test needs; the test requirement change will not affect or require the upstream testing to provision/modify a new image.

Alternatively, we can provision a new upstream testing image and modify/re-provision when the XLA testing requirements change. But testing a newly provisioned image that is similar, but not identical to the one used in the downstream might run into subtle issues.

### Additional context

https://github.com/pytorch/xla/issues/3159
https://github.com/pytorch/xla/issues/3351

cc @bdhirsh


Discussion:
cc @malfet @seemethere for visibility.
@yeounoh thank you for your suggestions. If you have PR in mind that would accomplish that, please do not hesitate to add me or @seemethere as the reviewers.
Hi @malfet , here is the [PR](https://github.com/pytorch/pytorch/pull/72496) that addresses/tests the idea. I can't seem to link the PR directly to the issue, nor add the reviewers myself. It would be really great if I could get the write permission. Thank you!

Resolving patch:
diff --git a/.github/scripts/generate_ci_workflows.py b/.github/scripts/generate_ci_workflows.py
--- a/.github/scripts/generate_ci_workflows.py
+++ b/.github/scripts/generate_ci_workflows.py
@@ -629,9 +629,8 @@ def generate_workflow_file(self, workflow_template: jinja2.Template) -> None:
         num_test_shards=2,
         distributed_test=False,
         enable_noarch_test=1,
-        enable_xla_test=1,
         ciflow_config=CIFlowConfig(
-            labels={LABEL_CIFLOW_DEFAULT, LABEL_CIFLOW_LINUX, LABEL_CIFLOW_CPU, LABEL_CIFLOW_XLA, LABEL_CIFLOW_NOARCH},
+            labels={LABEL_CIFLOW_DEFAULT, LABEL_CIFLOW_LINUX, LABEL_CIFLOW_CPU, LABEL_CIFLOW_NOARCH},
         ),
     ),
     CIWorkflow(
@@ -661,6 +660,22 @@ def generate_workflow_file(self, workflow_template: jinja2.Template) -> None:
     ),
 ]
 
+XLA_WORKFLOWS = [
+    CIWorkflow(
+        arch="linux",
+        build_environment="pytorch-xla-linux-bionic-py3.7-clang8",
+        docker_image_base=f"{DOCKER_REGISTRY}/pytorch/xla_base",
+        test_runner_type=LINUX_CPU_TEST_RUNNER,
+        num_test_shards=2,
+        distributed_test=False,
+        enable_xla_test=1,
+        ciflow_config=CIFlowConfig(
+            labels={LABEL_CIFLOW_LINUX, LABEL_CIFLOW_CPU, LABEL_CIFLOW_XLA},
+        ),
+    ),
+
+]
+
 ANDROID_SHORT_WORKFLOWS = [
     CIWorkflow(
         arch="linux",
@@ -836,7 +851,6 @@ def generate_workflow_file(self, workflow_template: jinja2.Template) -> None:
 ]
 
 DOCKER_IMAGES = {
-    f"{DOCKER_REGISTRY}/pytorch/pytorch-linux-bionic-cuda10.2-cudnn7-py3.7-clang9",  # for pytorch/xla
     f"{DOCKER_REGISTRY}/pytorch/pytorch-linux-bionic-rocm4.3.1-py3.7",               # for rocm
     f"{DOCKER_REGISTRY}/pytorch/pytorch-linux-bionic-rocm4.5-py3.7",                 # for rocm
 }
@@ -960,6 +974,7 @@ def main() -> None:
     )
     template_and_workflows = [
         (jinja_env.get_template("linux_ci_workflow.yml.j2"), LINUX_WORKFLOWS),
+        (jinja_env.get_template("linux_ci_workflow.yml.j2"), XLA_WORKFLOWS),
         (jinja_env.get_template("windows_ci_workflow.yml.j2"), WINDOWS_WORKFLOWS),
         (jinja_env.get_template("bazel_ci_workflow.yml.j2"), BAZEL_WORKFLOWS),
         (jinja_env.get_template("ios_ci_workflow.yml.j2"), IOS_WORKFLOWS),


Test patch:
diff --git a/.github/scripts/generate_pytorch_test_matrix.py b/.github/scripts/generate_pytorch_test_matrix.py
--- a/.github/scripts/generate_pytorch_test_matrix.py
+++ b/.github/scripts/generate_pytorch_test_matrix.py
@@ -61,6 +61,7 @@ def run_as_if_on_trunk() -> bool:
     return current_workflow_triggered_by_label
 
 def main() -> None:
+    INCLUDE_DEFAULT_TEST = True
     TEST_RUNNER_TYPE = os.getenv('TEST_RUNNER_TYPE')
     assert TEST_RUNNER_TYPE is not None
     RUN_SMOKE_TESTS_ONLY_ON_PR = os.getenv('RUN_SMOKE_TESTS_ONLY_ON_PR')
@@ -99,6 +100,7 @@ def main() -> None:
         configs['backwards_compat'] = {'num_shards': 1, 'runner': TEST_RUNNER_TYPE}
     if os.getenv('ENABLE_XLA_TEST'):
         configs['xla'] = {'num_shards': 1, 'runner': TEST_RUNNER_TYPE}
+        INCLUDE_DEFAULT_TEST = False
     if os.getenv('ENABLE_NOARCH_TEST'):
         configs['noarch'] = {'num_shards': 1, 'runner': TEST_RUNNER_TYPE}
     if RUN_SMOKE_TESTS:
@@ -112,6 +114,7 @@ def main() -> None:
                 'runner': TEST_RUNNER_TYPE,
             }
             for shard in range(1, NUM_TEST_SHARDS + 1)
+            if INCLUDE_DEFAULT_TEST
         ] + [
             {
                 'config': name,