fix: do not use script for TFS when entry_point is not provided (#1252)

* fix: Do not set script in prepare_framework_container_def for TFS

* remove _build_airflow_workflow_tf

* fix: Set script_name and dir_name env vars to None for TFS

Co-authored-by: Lauren Yu <[email protected]>

Author: ajaykarpur

src/sagemaker/model.py

@@ -844,9 +844,12 @@ def _framework_env_vars(self):
                 dir_name = "/opt/ml/model/code"
             else:
                 dir_name = self.uploaded_code.s3_prefix
-        else:
+        elif self.entry_point is not None:
             script_name = self.entry_point
             dir_name = "file://" + self.source_dir
+        else:
+            script_name = None
+            dir_name = None
 
         return {
             SCRIPT_PARAM_NAME.upper(): script_name,

src/sagemaker/workflow/airflow.py

@@ -522,18 +522,19 @@ def prepare_framework_container_def(model, instance_type, s3_operations):
     model.name = model.name or utils.name_from_base(base_name)
 
     bucket = model.bucket or model.sagemaker_session._default_bucket
-    script = os.path.basename(model.entry_point)
-    key = "{}/source/sourcedir.tar.gz".format(model.name)
-
-    if model.source_dir and model.source_dir.lower().startswith("s3://"):
-        code_dir = model.source_dir
-        model.uploaded_code = fw_utils.UploadedCode(s3_prefix=code_dir, script_name=script)
-    else:
-        code_dir = "s3://{}/{}".format(bucket, key)
-        model.uploaded_code = fw_utils.UploadedCode(s3_prefix=code_dir, script_name=script)
-        s3_operations["S3Upload"] = [
-            {"Path": model.source_dir or script, "Bucket": bucket, "Key": key, "Tar": True}
-        ]
+    if model.entry_point is not None:
+        script = os.path.basename(model.entry_point)
+        key = "{}/source/sourcedir.tar.gz".format(model.name)
+
+        if model.source_dir and model.source_dir.lower().startswith("s3://"):
+            code_dir = model.source_dir
+            model.uploaded_code = fw_utils.UploadedCode(s3_prefix=code_dir, script_name=script)
+        else:
+            code_dir = "s3://{}/{}".format(bucket, key)
+            model.uploaded_code = fw_utils.UploadedCode(s3_prefix=code_dir, script_name=script)
+            s3_operations["S3Upload"] = [
+                {"Path": model.source_dir or script, "Bucket": bucket, "Key": key, "Tar": True}
+            ]
 
     deploy_env = dict(model.env)
     deploy_env.update(model._framework_env_vars())

tests/integ/test_airflow_config.py

@@ -572,7 +572,7 @@ def test_tf_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_inst
             path=os.path.join(TF_MNIST_RESOURCE_PATH, "data"), key_prefix="scriptmode/mnist"
         )
 
-        training_config = _build_airflow_workflow_tf(
+        training_config = _build_airflow_workflow(
             estimator=tf, instance_type=cpu_instance_type, inputs=inputs
         )
 
@@ -701,45 +701,3 @@ def _build_airflow_workflow(estimator, instance_type, inputs=None, mini_batch_si
     transform_op.set_upstream(train_op)
 
     return training_config
-
-
-def _build_airflow_workflow_tf(estimator, instance_type, inputs=None, mini_batch_size=None):
-    training_config = sm_airflow.training_config(
-        estimator=estimator, inputs=inputs, mini_batch_size=mini_batch_size
-    )
-
-    model = estimator.create_model(entry_point=estimator.entry_point)
-    assert model is not None
-
-    model_config = sm_airflow.model_config(instance_type, model)
-    assert model_config is not None
-
-    transform_config = sm_airflow.transform_config_from_estimator(
-        estimator=estimator,
-        task_id="transform_config",
-        task_type="training",
-        instance_count=SINGLE_INSTANCE_COUNT,
-        instance_type=estimator.train_instance_type,
-        data=inputs,
-        content_type="text/csv",
-    )
-
-    default_args = {
-        "owner": "airflow",
-        "start_date": airflow.utils.dates.days_ago(2),
-        "provide_context": True,
-    }
-
-    dag = DAG("tensorflow_example", default_args=default_args, schedule_interval="@once")
-
-    train_op = SageMakerTrainingOperator(
-        task_id="tf_training", config=training_config, wait_for_completion=True, dag=dag
-    )
-
-    transform_op = SageMakerTransformOperator(
-        task_id="transform_operator", config=transform_config, wait_for_completion=True, dag=dag
-    )
-
-    transform_op.set_upstream(train_op)
-
-    return training_config