feat: support multi-container cloud run jobs (#259)

* feat: support multi-container cloud run jobs

* docs: changelog

Author: hanseltime

libraries/dagster-contrib-gcp/CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## 0.0.5
+
+### Updated
+
+- (pull/259) Added the `container_name` field option to support multi-container job launching
+
 ## 0.0.4
 
 ### Updated

libraries/dagster-contrib-gcp/README.md

@@ -67,6 +67,13 @@ my-code-location-1:
   name: my-cloud-run-job-1
   project_id: gcp_123
   region: us-central1
+
+# Multi-container support - which container to override on run
+my-code-location-1: 
+  name: my-cloud-run-job-1
+  project_id: gcp_123
+  region: us-central1
+  container_name: my-dagster-container-name
 ```
 
 Additional steps may be required for configuring IAM permissions, etc. In particular:

libraries/dagster-contrib-gcp/dagster_contrib_gcp/cloud_run/run_launcher.py

@@ -98,6 +98,13 @@ def launch_run(self, context: LaunchRunContext) -> None:
             context.dagster_run.run_id, {"cloud_run_job_execution_id": execution_id}
         )
 
+    def get_container_name_for_code_location_or_default(
+        self, job_config: dict[str, Any]
+    ) -> str | None:
+        """Returns the specific container_name for multi-container cloud run jobs to override"""
+        job_config = check.dict_param(job_config, "job_config")
+        return job_config.get("container_name", None)
+
     def get_project_for_code_location_or_default(
         self, job_config: dict[str, Any]
     ) -> str:
@@ -132,6 +139,19 @@ def fully_qualified_job_name(self, code_location_name: str) -> str:
             f"projects/{project_id_for_job}/locations/{region_for_job}/jobs/{job_name}"
         )
 
+    def specific_container_name(self, code_location_name: str) -> str | None:
+        try:
+            job = self.job_name_by_code_location[code_location_name]
+        except KeyError:
+            raise Exception(
+                f"No run launcher defined for code location: {code_location_name}"
+            )
+        # if not specific, then we default to legacy 1 container behavior (no name)
+        if isinstance(job, str):
+            return None
+
+        return self.get_container_name_for_code_location_or_default(job)
+
     def resolve_secret(self, secret_name: str) -> Any:
         client = SecretManagerServiceClient()
         latest = AccessSecretVersionRequest(name=secret_name)
@@ -185,13 +205,17 @@ def env_override_for_code_location(
     def create_execution(self, code_location_name: str, args: Sequence[str]):
         job_name = self.fully_qualified_job_name(code_location_name)
         job_env = self.env_override_for_code_location(code_location_name)
-        return self.execute_job(job_name, args=args, env=job_env)
+        container_name = self.specific_container_name(code_location_name)
+        return self.execute_job(
+            job_name, args=args, env=job_env, container_name=container_name
+        )
 
     def execute_job(
         self,
         fully_qualified_job_name: str,
         args: Sequence[str] | None = None,
         env: Optional["dict[str, str]"] = None,
+        container_name: str | None = None,
     ) -> Operation:
         request = RunJobRequest(name=fully_qualified_job_name)
 
@@ -202,6 +226,8 @@ def execute_job(
             overrides["env"] = [
                 k8s_min.EnvVar(name=name, value=value) for name, value in env.items()
             ]
+        if container_name:
+            overrides["name"] = container_name
 
         container_overrides = [RunJobRequest.Overrides.ContainerOverride(**overrides)]
 
@@ -280,6 +306,8 @@ def config_type(cls) -> "UserConfigSchema":
                     " pair where the key is the code location name and the value is the job name. "
                     "Optionally, each code location key may specifiy the `job_name` and `project_id` "
                     "override value in order to the code location to a different GCP project ID."
+                    "If using multi-container jobs, you may also specify `container_name` which"
+                    "will be used to override the specific container"
                 ),
             ),
             "run_job_retry": Field(

libraries/dagster-contrib-gcp/dagster_contrib_gcp_tests/cloud_run_tests/conftest.py

@@ -73,6 +73,29 @@ def instance_with_job_configs(
         yield dagster_instance
 
 
+@pytest.fixture
+def instance_with_multicontainer_job_configs(
+    instance_cm: Callable[..., ContextManager[DagsterInstance]],
+) -> Iterator[DagsterInstance]:
+    with instance_cm(
+        {
+            "project": "test_project",
+            "region": "test_region",
+            "job_name_by_code_location": {
+                IN_PROCESS_NAME: {
+                    "name": "test_job_with_config",
+                    "project_id": "test_gcp-123",
+                    "region": "other_test_region",
+                    "container_name": "specific_container",
+                }
+            },
+            "run_job_retry": {"wait": 1, "timeout": 60},
+            "run_timeout": 7200,
+        }
+    ) as dagster_instance:
+        yield dagster_instance
+
+
 @pytest.fixture
 def workspace(instance: DagsterInstance) -> Iterator[WorkspaceRequestContext]:
     with in_process_test_workspace(
@@ -101,6 +124,21 @@ def workspace_with_job_configs(
         yield workspace
 
 
+@pytest.fixture
+def workspace_with_multicontainer_job_configs(
+    instance_with_multicontainer_job_configs: DagsterInstance,
+) -> Iterator[WorkspaceRequestContext]:
+    with in_process_test_workspace(
+        instance_with_multicontainer_job_configs,
+        loadable_target_origin=LoadableTargetOrigin(
+            python_file=repo.__file__,
+            attribute=repo.repository.name,
+        ),
+        container_image="dagster:latest",
+    ) as workspace:
+        yield workspace
+
+
 @pytest.fixture
 def job() -> JobDefinition:
     return repo.job
@@ -135,6 +173,19 @@ def run_with_job_configs(
     )
 
 
+@pytest.fixture
+def run_with_multicontainer_job_configs(
+    instance_with_multicontainer_job_configs: DagsterInstance,
+    job: JobDefinition,
+    external_job: RemoteJob,
+) -> DagsterRun:
+    return instance_with_multicontainer_job_configs.create_run_for_job(
+        job,
+        remote_job_origin=external_job.get_remote_origin(),
+        job_code_origin=external_job.get_python_origin(),
+    )
+
+
 @pytest.fixture
 def executions():
     return {}

libraries/dagster-contrib-gcp/dagster_contrib_gcp_tests/cloud_run_tests/test_run_launcher.py

@@ -68,6 +68,38 @@ def test_launch_run_with_job_config(
     assert args[0].overrides.timeout == dt.timedelta(seconds=7200)
 
 
+def test_launch_run_with_multicontainer_job_config(
+    instance_with_multicontainer_job_configs: DagsterInstance,
+    run_with_multicontainer_job_configs: DagsterRun,
+    workspace_with_multicontainer_job_configs: BaseWorkspaceRequestContext,
+    mock_jobs_client,
+    mock_executions_client,
+):
+    instance_with_multicontainer_job_configs.launch_run(
+        run_with_multicontainer_job_configs.run_id,
+        workspace_with_multicontainer_job_configs,
+    )
+    run = check.not_none(
+        instance_with_multicontainer_job_configs.get_run_by_id(
+            run_with_multicontainer_job_configs.run_id
+        )
+    )
+
+    # Assert the correct tag is set
+    assert run.tags["cloud_run_job_execution_id"] == "test_execution_id"
+
+    # Check that mock was called with correct args
+    mock_jobs_client.run_job.assert_called_once()
+    args, _ = mock_jobs_client.run_job.call_args
+    assert isinstance(args[0], RunJobRequest)
+    assert (
+        args[0].name
+        == "projects/test_gcp-123/locations/other_test_region/jobs/test_job_with_config"
+    )
+    assert args[0].overrides.timeout == dt.timedelta(seconds=7200)
+    assert args[0].overrides.container_overrides[0].name == "specific_container"
+
+
 @patch.object(run_launcher.CloudRunRunLauncher, "resolve_secret")
 def test_env_override_for_code_location(patched_resolve_secret):
     mock = MagicMock()