Set DSTACK_RUN_ID and DSTACK_JOB_ID (#2622)

* Pass run and job ids to runner

* Set DSTACK_RUN_ID and DSTACK_JOB_ID

* Comment on container_.NetworkSettings.Ports

* Document new env vars

* Minor fixes

Author: r4victor

docs/docs/reference/environment-variables.md

@@ -19,6 +19,8 @@ tasks, and services:
 
      If `name` is not set in the configuration, it is assigned a random name (e.g. `wet-mangust-1`).
 
+- `DSTACK_RUN_ID`{ #DSTACK_RUN_ID } – The UUID of the run.
+- `DSTACK_JOB_ID`{ #DSTACK_JOB_ID } – The UUID of the job submission.
 - `DSTACK_REPO_ID`{ #DSTACK_REPO_ID } – The ID of the repo.
 - `DSTACK_GPUS_NUM`{ #DSTACK_GPUS_NUM } – The total number of GPUs in the run.
 

runner/internal/executor/executor.go

@@ -34,8 +34,9 @@ type RunExecutor struct {
 	sshPort    int
 	uid        uint32
 
-	run             schemas.RunSpec
+	run             schemas.Run
 	jobSpec         schemas.JobSpec
+	jobSubmission   schemas.JobSubmission
 	clusterInfo     schemas.ClusterInfo
 	secrets         map[string]string
 	repoCredentials *schemas.RepoCredentials
@@ -195,7 +196,8 @@ func (ex *RunExecutor) Run(ctx context.Context) (err error) {
 }
 
 func (ex *RunExecutor) SetJob(body schemas.SubmitBody) {
-	ex.run = body.RunSpec
+	ex.run = body.Run
+	ex.jobSubmission = body.JobSubmission
 	ex.jobSpec = body.JobSpec
 	ex.clusterInfo = body.ClusterInfo
 	ex.secrets = body.Secrets
@@ -256,8 +258,10 @@ func (ex *RunExecutor) execJob(ctx context.Context, jobLogFile io.Writer) error
 	gpus_num := nodes_num * gpus_per_node_num
 
 	jobEnvs := map[string]string{
-		"DSTACK_RUN_NAME":       ex.run.RunName,
-		"DSTACK_REPO_ID":        ex.run.RepoId,
+		"DSTACK_RUN_ID":         ex.run.Id,
+		"DSTACK_JOB_ID":         ex.jobSubmission.Id,
+		"DSTACK_RUN_NAME":       ex.run.RunSpec.RunName,
+		"DSTACK_REPO_ID":        ex.run.RunSpec.RepoId,
 		"DSTACK_NODES_IPS":      strings.Join(ex.clusterInfo.JobIPs, "\n"),
 		"DSTACK_MASTER_NODE_IP": ex.clusterInfo.MasterJobIP,
 		"DSTACK_NODE_RANK":      strconv.Itoa(node_rank),

runner/internal/executor/executor_test.go

@@ -132,7 +132,7 @@ func TestExecutor_RemoteRepo(t *testing.T) {
 
 	var b bytes.Buffer
 	ex := makeTestExecutor(t)
-	ex.run.RepoData = schemas.RepoData{
+	ex.run.RunSpec.RepoData = schemas.RepoData{
 		RepoType:        "remote",
 		RepoBranch:      "main",
 		RepoHash:        "2b83592e506ed6fe8e49f4eaa97c3866bc9402b1",
@@ -148,7 +148,7 @@ func TestExecutor_RemoteRepo(t *testing.T) {
 
 	err = ex.execJob(context.TODO(), io.Writer(&b))
 	assert.NoError(t, err)
-	expected := fmt.Sprintf("%s\r\n%s\r\n%s\r\n", ex.run.RepoData.RepoHash, ex.run.RepoData.RepoConfigName, ex.run.RepoData.RepoConfigEmail)
+	expected := fmt.Sprintf("%s\r\n%s\r\n%s\r\n", ex.run.RunSpec.RepoData.RepoHash, ex.run.RunSpec.RepoData.RepoConfigName, ex.run.RunSpec.RepoData.RepoConfigEmail)
 	assert.Equal(t, expected, b.String())
 }
 
@@ -161,14 +161,17 @@ func makeTestExecutor(t *testing.T) *RunExecutor {
 	require.NoError(t, err)
 
 	body := schemas.SubmitBody{
-		RunSpec: schemas.RunSpec{
-			RunName:  "red-turtle-1",
-			RepoId:   "test-000000",
-			RepoData: schemas.RepoData{RepoType: "local"},
-			Configuration: schemas.Configuration{
-				Type: "task",
+		Run: schemas.Run{
+			Id: "12346",
+			RunSpec: schemas.RunSpec{
+				RunName:  "red-turtle-1",
+				RepoId:   "test-000000",
+				RepoData: schemas.RepoData{RepoType: "local"},
+				Configuration: schemas.Configuration{
+					Type: "task",
+				},
+				ConfigurationPath: ".dstack.yml",
 			},
-			ConfigurationPath: ".dstack.yml",
 		},
 		JobSpec: schemas.JobSpec{
 			Commands:    []string{"/bin/bash", "-c"},

runner/internal/executor/repo.go

@@ -40,7 +40,7 @@ func (ex *RunExecutor) setupRepo(ctx context.Context) error {
 			err = gerrors.Wrap(err_)
 		}
 	}()
-	switch ex.run.RepoData.RepoType {
+	switch ex.run.RunSpec.RepoData.RepoType {
 	case "remote":
 		log.Trace(ctx, "Fetching git repository")
 		if err := ex.prepareGit(ctx); err != nil {
@@ -52,7 +52,7 @@ func (ex *RunExecutor) setupRepo(ctx context.Context) error {
 			return gerrors.Wrap(err)
 		}
 	default:
-		return gerrors.Newf("unknown RepoType: %s", ex.run.RepoData.RepoType)
+		return gerrors.Newf("unknown RepoType: %s", ex.run.RunSpec.RepoData.RepoType)
 	}
 	return err
 }
@@ -61,8 +61,8 @@ func (ex *RunExecutor) prepareGit(ctx context.Context) error {
 	repoManager := repo.NewManager(
 		ctx,
 		ex.repoCredentials.CloneURL,
-		ex.run.RepoData.RepoBranch,
-		ex.run.RepoData.RepoHash,
+		ex.run.RunSpec.RepoData.RepoBranch,
+		ex.run.RunSpec.RepoData.RepoHash,
 		ex.jobSpec.SingleBranch,
 	).WithLocalPath(ex.workingDir)
 	if ex.repoCredentials != nil {
@@ -92,7 +92,7 @@ func (ex *RunExecutor) prepareGit(ctx context.Context) error {
 	if err := repoManager.Checkout(); err != nil {
 		return gerrors.Wrap(err)
 	}
-	if err := repoManager.SetConfig(ex.run.RepoData.RepoConfigName, ex.run.RepoData.RepoConfigEmail); err != nil {
+	if err := repoManager.SetConfig(ex.run.RunSpec.RepoData.RepoConfigName, ex.run.RunSpec.RepoData.RepoConfigEmail); err != nil {
 		return gerrors.Wrap(err)
 	}
 

runner/internal/schemas/schemas.go

@@ -20,8 +20,9 @@ type LogEvent struct {
 }
 
 type SubmitBody struct {
-	RunSpec         RunSpec           `json:"run_spec"`
+	Run             Run               `json:"run"`
 	JobSpec         JobSpec           `json:"job_spec"`
+	JobSubmission   JobSubmission     `json:"job_submission"`
 	ClusterInfo     ClusterInfo       `json:"cluster_info"`
 	Secrets         map[string]string `json:"secrets"`
 	RepoCredentials *RepoCredentials  `json:"repo_credentials"`
@@ -37,6 +38,11 @@ type PullResponse struct {
 	// todo Result
 }
 
+type Run struct {
+	Id      string  `json:"id"`
+	RunSpec RunSpec `json:"run_spec"`
+}
+
 type RunSpec struct {
 	RunName           string        `json:"run_name"`
 	RepoId            string        `json:"repo_id"`
@@ -45,6 +51,10 @@ type RunSpec struct {
 	ConfigurationPath string        `json:"configuration_path"`
 }
 
+type JobSubmission struct {
+	Id string `json:"id"`
+}
+
 type JobSpec struct {
 	ReplicaNum     int               `json:"replica_num"`
 	JobNum         int               `json:"job_num"`

runner/internal/shim/docker.go

@@ -852,6 +852,9 @@ func (d *DockerRunner) startContainer(ctx context.Context, task *Task) error {
 	if err != nil {
 		return tracerr.Wrap(err)
 	}
+	// FIXME: container_.NetworkSettings.Ports values (bindings) are not immediately available
+	// on macOS, so ports can be empty with local backend.
+	// Workaround: restart shim after submitting the run.
 	task.ports = extractPorts(ctx, container_.NetworkSettings.Ports)
 	return nil
 }

src/dstack/_internal/core/models/runs.py

@@ -262,9 +262,9 @@ class JobRuntimeData(CoreModel):
     # or not applicable (container-based backends)
     ports: Optional[dict[int, int]] = None
     # List of volumes used by the job
-    volume_names: Optional[list[str]] = None  # None for backward compalibility
+    volume_names: Optional[list[str]] = None  # None for backward compatibility
     # Virtual shared offer
-    offer: Optional[InstanceOfferWithAvailability] = None  # None for backward compalibility
+    offer: Optional[InstanceOfferWithAvailability] = None  # None for backward compatibility
 
 
 class ClusterInfo(CoreModel):

src/dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -543,7 +543,7 @@ def _process_pulling_with_shim(
     if shim_client.is_api_v2_supported():  # raises error if shim is down, causes retry
         task = shim_client.get_task(job_model.id)
 
-        # If task goes to terminated before the job is submitted to runner, then an error occured
+        # If task goes to terminated before the job is submitted to runner, then an error occurred
         if task.status == TaskStatus.TERMINATED:
             logger.warning(
                 "shim failed to execute job %s: %s (%s)",
@@ -572,7 +572,7 @@ def _process_pulling_with_shim(
     else:
         shim_status = shim_client.pull()  # raises error if shim is down, causes retry
 
-        # If shim goes to pending before the job is submitted to runner, then an error occured
+        # If shim goes to pending before the job is submitted to runner, then an error occurred
         if (
             shim_status.state == "pending"
             and shim_status.result is not None
@@ -822,8 +822,8 @@ def _submit_job_to_runner(
         return success_if_not_available
 
     runner_client.submit_job(
-        run_spec=run.run_spec,
-        job_spec=job.job_spec,
+        run=run,
+        job=job,
         cluster_info=cluster_info,
         secrets=secrets,
         repo_credentials=repo_credentials,

src/dstack/_internal/server/schemas/runner.py

@@ -7,7 +7,14 @@
 
 from dstack._internal.core.models.common import CoreModel, NetworkMode
 from dstack._internal.core.models.repos.remote import RemoteRepoCreds
-from dstack._internal.core.models.runs import ClusterInfo, JobSpec, JobStatus, RunSpec
+from dstack._internal.core.models.runs import (
+    ClusterInfo,
+    JobSpec,
+    JobStatus,
+    JobSubmission,
+    Run,
+    RunSpec,
+)
 from dstack._internal.core.models.volumes import InstanceMountPoint, VolumeMountPoint
 
 
@@ -39,15 +46,18 @@ class PullResponse(CoreModel):
 
 
 class SubmitBody(CoreModel):
-    run_spec: Annotated[
-        RunSpec,
+    run: Annotated[
+        Run,
         Field(
             include={
-                "run_name",
-                "repo_id",
-                "repo_data",
-                "configuration",
-                "configuration_path",
+                "id": True,
+                "run_spec": {
+                    "run_name",
+                    "repo_id",
+                    "repo_data",
+                    "configuration",
+                    "configuration_path",
+                },
             }
         ),
     ]
@@ -70,9 +80,31 @@ class SubmitBody(CoreModel):
             }
         ),
     ]
+    job_submission: Annotated[
+        JobSubmission,
+        Field(
+            include={
+                "id",
+            }
+        ),
+    ]
     cluster_info: Annotated[Optional[ClusterInfo], Field(include=True)]
     secrets: Annotated[Optional[Dict[str, str]], Field(include=True)]
     repo_credentials: Annotated[Optional[RemoteRepoCreds], Field(include=True)]
+    # run_spec is deprecated in favor of run.run_spec
+    # TODO: Remove once we no longer support instances deployed with 0.19.8 or earlier.
+    run_spec: Annotated[
+        RunSpec,
+        Field(
+            include={
+                "run_name",
+                "repo_id",
+                "repo_data",
+                "configuration",
+                "configuration_path",
+            },
+        ),
+    ]
 
 
 class HealthcheckResponse(CoreModel):

src/dstack/_internal/server/services/runner/client.py

@@ -12,7 +12,7 @@
 from dstack._internal.core.models.envs import Env
 from dstack._internal.core.models.repos.remote import RemoteRepoCreds
 from dstack._internal.core.models.resources import Memory
-from dstack._internal.core.models.runs import ClusterInfo, JobSpec, RunSpec
+from dstack._internal.core.models.runs import ClusterInfo, Job, Run
 from dstack._internal.core.models.volumes import InstanceMountPoint, Volume, VolumeMountPoint
 from dstack._internal.server.schemas.runner import (
     GPUDevice,
@@ -72,15 +72,16 @@ def get_metrics(self) -> Optional[MetricsResponse]:
 
     def submit_job(
         self,
-        run_spec: RunSpec,
-        job_spec: JobSpec,
+        run: Run,
+        job: Job,
         cluster_info: ClusterInfo,
         secrets: Dict[str, str],
         repo_credentials: Optional[RemoteRepoCreds],
         instance_env: Optional[Union[Env, Dict[str, str]]] = None,
     ):
         # XXX: This is a quick-and-dirty hack to deliver InstanceModel-specific environment
         # variables to the runner without runner API modification.
+        job_spec = job.job_spec
         if instance_env is not None:
             if isinstance(instance_env, Env):
                 merged_env = instance_env.as_dict()
@@ -90,11 +91,13 @@ def submit_job(
             job_spec = job_spec.copy(deep=True)
             job_spec.env = merged_env
         body = SubmitBody(
-            run_spec=run_spec,
+            run=run,
             job_spec=job_spec,
+            job_submission=job.job_submissions[-1],
             cluster_info=cluster_info,
             secrets=secrets,
             repo_credentials=repo_credentials,
+            run_spec=run.run_spec,
         )
         resp = requests.post(
             # use .json() to encode enums